From 89a0fd4914cabadb6562feab3ce3abd59b99287f Mon Sep 17 00:00:00 2001 From: ogt Date: Thu, 9 Jul 2026 21:55:45 +0800 Subject: [PATCH] fix(agents): stabilize log consumer readback --- .../ai_agent_autonomous_runtime_control.py | 2 +- ..._controlled_writeback_consumer_readback.py | 322 ++++++++++++++---- ...trolled_writeback_consumer_readback_api.py | 44 +++ 3 files changed, 298 insertions(+), 70 deletions(-) diff --git a/apps/api/src/services/ai_agent_autonomous_runtime_control.py b/apps/api/src/services/ai_agent_autonomous_runtime_control.py index a7d2ca6e1..30bb6f1fc 100644 --- a/apps/api/src/services/ai_agent_autonomous_runtime_control.py +++ b/apps/api/src/services/ai_agent_autonomous_runtime_control.py @@ -52,7 +52,7 @@ _RUNTIME_RECEIPT_DB_CONTEXT_EXIT_TIMEOUT_SECONDS = 0.25 _RUNTIME_RECEIPT_QUERY_BUDGET_SECONDS = 2.0 _RUNTIME_RECEIPT_SINGLE_QUERY_TIMEOUT_SECONDS = 0.45 _RUNTIME_RECEIPT_STATEMENT_TIMEOUT_MS = 400 -_LOG_CONTROLLED_WRITEBACK_CONSUMER_TIMEOUT_SECONDS = 0.5 +_LOG_CONTROLLED_WRITEBACK_CONSUMER_TIMEOUT_SECONDS = 2.0 _RUNTIME_RECEIPT_READBACK_CACHE_TTL_SECONDS = 20.0 _AUXILIARY_RUNTIME_RECEIPT_QUERY_NAMES = { "alert_operation_counts", diff --git a/apps/api/src/services/ai_agent_log_controlled_writeback_consumer_readback.py b/apps/api/src/services/ai_agent_log_controlled_writeback_consumer_readback.py index 0fbda0f70..b0f4dda24 100644 --- a/apps/api/src/services/ai_agent_log_controlled_writeback_consumer_readback.py +++ b/apps/api/src/services/ai_agent_log_controlled_writeback_consumer_readback.py @@ -16,6 +16,7 @@ from typing import Any from sqlalchemy import text +from src.core.config import settings from src.core.logging import get_logger from src.db.base import get_db_context from src.services.ai_agent_log_controlled_writeback_dispatch import ( @@ -41,7 +42,12 @@ _CONSUMER_SURFACES = { "verifier": "post_apply_verifier_feedback_context", "ai_agent": "autonomous_runtime_decision_context", } -_DB_RETRY_DELAYS_SECONDS = (0.0, 0.25, 1.0) +_DB_RETRY_DELAYS_SECONDS = (0.0,) +_CONSUMER_DB_CONTEXT_TIMEOUT_SECONDS = 0.45 +_CONSUMER_DB_CONTEXT_EXIT_TIMEOUT_SECONDS = 0.25 +_CONSUMER_SINGLE_QUERY_TIMEOUT_SECONDS = 0.45 +_CONSUMER_TELEGRAM_CONTEXT_TIMEOUT_SECONDS = 0.75 +_CONSUMER_STATEMENT_TIMEOUT_MS = 400 logger = get_logger(__name__) @@ -54,6 +60,17 @@ async def load_latest_ai_agent_log_controlled_writeback_consumer_readback( try: rows, consumer_rows = await _load_consumer_rows_with_retry(project_id=project_id) except Exception as exc: # pragma: no cover - live DB pool pressure + direct_rows = await _load_consumer_rows_direct(project_id=project_id) + if direct_rows is not None: + rows, consumer_rows = direct_rows + telegram_alert_learning_context = await _load_telegram_alert_learning_context( + project_id=project_id + ) + return _build_consumer_readback( + rows=rows, + consumer_rows=consumer_rows, + telegram_alert_learning_context=telegram_alert_learning_context, + ) logger.warning( "log_controlled_writeback_consumer_readback_db_unavailable", project_id=project_id, @@ -101,74 +118,238 @@ async def _load_consumer_rows_once( *, project_id: str, ) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: - async with get_db_context(project_id) as db: - await db.execute(text("SET LOCAL statement_timeout = '5000ms'"), {}) - result = await db.execute( - text(""" - SELECT - op_id::text AS op_id, - operation_type, - actor, - status, - created_at, - input ->> 'semantic_operation_type' AS semantic_operation_type, - input ->> 'ledger_operation_type' AS ledger_operation_type, - input ->> 'dispatch_receipt_id' AS dispatch_receipt_id, - input ->> 'batch_id' AS batch_id, - input ->> 'target' AS target, - input ->> 'target_surface' AS target_surface, - input ->> 'risk_tier' AS risk_tier, - input ->> 'project_id' AS project_id, - input ->> 'raw_payload_included' AS raw_payload_included, - output ->> 'next_action' AS next_action, - output ->> 'km_write_performed' AS km_write_performed, - output ->> 'rag_index_write_performed' AS rag_index_write_performed, - output ->> 'playbook_trust_write_performed' AS playbook_trust_write_performed, - output ->> 'mcp_tool_call_performed' AS mcp_tool_call_performed, - output ->> 'telegram_send_performed' AS telegram_send_performed, - output -> 'post_apply_verifier_refs' AS post_apply_verifier_refs - FROM automation_operation_log - WHERE coalesce(input ->> 'semantic_operation_type', operation_type) - = :operation_type - ORDER BY created_at DESC, op_id DESC - LIMIT 50 - """), - {"operation_type": OPERATION_TYPE}, + db_context = get_db_context(project_id) + db_context_entered = False + try: + db = await asyncio.wait_for( + db_context.__aenter__(), + timeout=_CONSUMER_DB_CONTEXT_TIMEOUT_SECONDS, ) - consumer_result = await db.execute( - text(""" - SELECT - op_id::text AS op_id, - operation_type, - actor, - status, - created_at, - input ->> 'semantic_operation_type' AS semantic_operation_type, - input ->> 'ledger_operation_type' AS ledger_operation_type, - input ->> 'consumer_receipt_id' AS consumer_receipt_id, - input ->> 'source_dispatch_receipt_id' AS source_dispatch_receipt_id, - input ->> 'target' AS target, - input ->> 'target_surface' AS target_surface, - input ->> 'consumer_surface' AS consumer_surface, - input ->> 'runtime_target_write_performed' AS runtime_target_write_performed, - input ->> 'raw_payload_included' AS raw_payload_included, - output ->> 'consumer_context_receipt_recorded' - AS consumer_context_receipt_recorded, - output ->> 'target_context_receipt_write_performed' - AS target_context_receipt_write_performed, - output ->> 'next_action' AS next_action, - output -> 'post_apply_verifier_refs' AS post_apply_verifier_refs - FROM automation_operation_log - WHERE coalesce(input ->> 'semantic_operation_type', operation_type) - = :operation_type - ORDER BY created_at DESC, op_id DESC - LIMIT 50 - """), - {"operation_type": CONSUMER_OPERATION_TYPE}, + db_context_entered = True + await asyncio.wait_for( + db.execute( + text( + "SET LOCAL statement_timeout = " + f"'{int(_CONSUMER_STATEMENT_TIMEOUT_MS)}ms'" + ), + {}, + ), + timeout=_CONSUMER_SINGLE_QUERY_TIMEOUT_SECONDS, ) + result = await asyncio.wait_for( + db.execute( + text(""" + SELECT + op_id::text AS op_id, + operation_type, + actor, + status, + created_at, + input ->> 'semantic_operation_type' AS semantic_operation_type, + input ->> 'ledger_operation_type' AS ledger_operation_type, + input ->> 'dispatch_receipt_id' AS dispatch_receipt_id, + input ->> 'batch_id' AS batch_id, + input ->> 'target' AS target, + input ->> 'target_surface' AS target_surface, + input ->> 'risk_tier' AS risk_tier, + input ->> 'project_id' AS project_id, + input ->> 'raw_payload_included' AS raw_payload_included, + output ->> 'next_action' AS next_action, + output ->> 'km_write_performed' AS km_write_performed, + output ->> 'rag_index_write_performed' AS rag_index_write_performed, + output ->> 'playbook_trust_write_performed' AS playbook_trust_write_performed, + output ->> 'mcp_tool_call_performed' AS mcp_tool_call_performed, + output ->> 'telegram_send_performed' AS telegram_send_performed, + output -> 'post_apply_verifier_refs' AS post_apply_verifier_refs + FROM automation_operation_log + WHERE coalesce(input ->> 'semantic_operation_type', operation_type) + = :operation_type + ORDER BY created_at DESC, op_id DESC + LIMIT 50 + """), + {"operation_type": OPERATION_TYPE}, + ), + timeout=_CONSUMER_SINGLE_QUERY_TIMEOUT_SECONDS, + ) + consumer_result = await asyncio.wait_for( + db.execute( + text(""" + SELECT + op_id::text AS op_id, + operation_type, + actor, + status, + created_at, + input ->> 'semantic_operation_type' AS semantic_operation_type, + input ->> 'ledger_operation_type' AS ledger_operation_type, + input ->> 'consumer_receipt_id' AS consumer_receipt_id, + input ->> 'source_dispatch_receipt_id' AS source_dispatch_receipt_id, + input ->> 'target' AS target, + input ->> 'target_surface' AS target_surface, + input ->> 'consumer_surface' AS consumer_surface, + input ->> 'runtime_target_write_performed' AS runtime_target_write_performed, + input ->> 'raw_payload_included' AS raw_payload_included, + output ->> 'consumer_context_receipt_recorded' + AS consumer_context_receipt_recorded, + output ->> 'target_context_receipt_write_performed' + AS target_context_receipt_write_performed, + output ->> 'next_action' AS next_action, + output -> 'post_apply_verifier_refs' AS post_apply_verifier_refs + FROM automation_operation_log + WHERE coalesce(input ->> 'semantic_operation_type', operation_type) + = :operation_type + ORDER BY created_at DESC, op_id DESC + LIMIT 50 + """), + {"operation_type": CONSUMER_OPERATION_TYPE}, + ), + timeout=_CONSUMER_SINGLE_QUERY_TIMEOUT_SECONDS, + ) + finally: + if db_context_entered: + await asyncio.wait_for( + db_context.__aexit__(None, None, None), + timeout=_CONSUMER_DB_CONTEXT_EXIT_TIMEOUT_SECONDS, + ) return _result_rows(result), _result_rows(consumer_result) +async def _load_consumer_rows_direct( + *, + project_id: str, +) -> tuple[list[dict[str, Any]], list[dict[str, Any]]] | None: + """Read consumer rows through a short-lived connection when the pool is busy.""" + + db_url = settings.DATABASE_URL.replace("postgresql+asyncpg://", "postgresql://") + try: + import asyncpg + + conn = await asyncio.wait_for( + asyncpg.connect(db_url), + timeout=_CONSUMER_DB_CONTEXT_TIMEOUT_SECONDS, + ) + try: + await asyncio.wait_for( + conn.execute("SELECT set_config('app.project_id', $1, TRUE)", project_id), + timeout=_CONSUMER_SINGLE_QUERY_TIMEOUT_SECONDS, + ) + await asyncio.wait_for( + conn.execute( + "SET statement_timeout = " + f"'{int(_CONSUMER_STATEMENT_TIMEOUT_MS)}ms'" + ), + timeout=_CONSUMER_SINGLE_QUERY_TIMEOUT_SECONDS, + ) + dispatch_rows = await asyncio.wait_for( + conn.fetch(_CONSUMER_DISPATCH_DIRECT_SQL, OPERATION_TYPE), + timeout=_CONSUMER_SINGLE_QUERY_TIMEOUT_SECONDS, + ) + consumer_rows = await asyncio.wait_for( + conn.fetch(_CONSUMER_APPLY_DIRECT_SQL, CONSUMER_OPERATION_TYPE), + timeout=_CONSUMER_SINGLE_QUERY_TIMEOUT_SECONDS, + ) + return _result_rows(dispatch_rows), _result_rows(consumer_rows) + finally: + try: + await asyncio.wait_for( + conn.close(), + timeout=_CONSUMER_DB_CONTEXT_EXIT_TIMEOUT_SECONDS, + ) + except Exception as close_exc: # pragma: no cover - live connection cleanup + logger.warning( + "log_controlled_writeback_consumer_direct_close_failed", + project_id=project_id, + error_type=type(close_exc).__name__, + ) + except Exception as exc: # pragma: no cover - live DB pool pressure + logger.warning( + "log_controlled_writeback_consumer_direct_readback_failed", + project_id=project_id, + error_type=type(exc).__name__, + ) + return None + + +_CONSUMER_DISPATCH_DIRECT_SQL = """ + SELECT + op_id::text AS op_id, + operation_type, + actor, + status, + created_at, + input ->> 'semantic_operation_type' AS semantic_operation_type, + input ->> 'ledger_operation_type' AS ledger_operation_type, + input ->> 'dispatch_receipt_id' AS dispatch_receipt_id, + input ->> 'batch_id' AS batch_id, + input ->> 'target' AS target, + input ->> 'target_surface' AS target_surface, + input ->> 'risk_tier' AS risk_tier, + input ->> 'project_id' AS project_id, + input ->> 'raw_payload_included' AS raw_payload_included, + output ->> 'next_action' AS next_action, + output ->> 'km_write_performed' AS km_write_performed, + output ->> 'rag_index_write_performed' AS rag_index_write_performed, + output ->> 'playbook_trust_write_performed' AS playbook_trust_write_performed, + output ->> 'mcp_tool_call_performed' AS mcp_tool_call_performed, + output ->> 'telegram_send_performed' AS telegram_send_performed, + COALESCE( + ARRAY( + SELECT jsonb_array_elements_text( + CASE + WHEN jsonb_typeof(output -> 'post_apply_verifier_refs') = 'array' + THEN output -> 'post_apply_verifier_refs' + ELSE '[]'::jsonb + END + ) + ), + ARRAY[]::text[] + ) AS post_apply_verifier_refs + FROM automation_operation_log + WHERE coalesce(input ->> 'semantic_operation_type', operation_type) = $1 + ORDER BY created_at DESC, op_id DESC + LIMIT 50 +""" + +_CONSUMER_APPLY_DIRECT_SQL = """ + SELECT + op_id::text AS op_id, + operation_type, + actor, + status, + created_at, + input ->> 'semantic_operation_type' AS semantic_operation_type, + input ->> 'ledger_operation_type' AS ledger_operation_type, + input ->> 'consumer_receipt_id' AS consumer_receipt_id, + input ->> 'source_dispatch_receipt_id' AS source_dispatch_receipt_id, + input ->> 'target' AS target, + input ->> 'target_surface' AS target_surface, + input ->> 'consumer_surface' AS consumer_surface, + input ->> 'runtime_target_write_performed' AS runtime_target_write_performed, + input ->> 'raw_payload_included' AS raw_payload_included, + output ->> 'consumer_context_receipt_recorded' + AS consumer_context_receipt_recorded, + output ->> 'target_context_receipt_write_performed' + AS target_context_receipt_write_performed, + output ->> 'next_action' AS next_action, + COALESCE( + ARRAY( + SELECT jsonb_array_elements_text( + CASE + WHEN jsonb_typeof(output -> 'post_apply_verifier_refs') = 'array' + THEN output -> 'post_apply_verifier_refs' + ELSE '[]'::jsonb + END + ) + ), + ARRAY[]::text[] + ) AS post_apply_verifier_refs + FROM automation_operation_log + WHERE coalesce(input ->> 'semantic_operation_type', operation_type) = $1 + ORDER BY created_at DESC, op_id DESC + LIMIT 50 +""" + + def _build_consumer_readback( *, rows: list[dict[str, Any]], @@ -606,11 +787,14 @@ async def _load_telegram_alert_learning_context( ) -> dict[str, Any]: """Read Telegram alert learning registry as AI Loop consumer context.""" try: - readback = await list_ai_alert_card_delivery_readback( - project_id=project_id, - page=1, - per_page=TELEGRAM_ALERT_CONTEXT_PER_PAGE, - refresh=True, + readback = await asyncio.wait_for( + list_ai_alert_card_delivery_readback( + project_id=project_id, + page=1, + per_page=TELEGRAM_ALERT_CONTEXT_PER_PAGE, + refresh=True, + ), + timeout=_CONSUMER_TELEGRAM_CONTEXT_TIMEOUT_SECONDS, ) except Exception as exc: # pragma: no cover - live DB pool pressure logger.warning( diff --git a/apps/api/tests/test_ai_agent_log_controlled_writeback_consumer_readback_api.py b/apps/api/tests/test_ai_agent_log_controlled_writeback_consumer_readback_api.py index 77340da4a..ffa48621b 100644 --- a/apps/api/tests/test_ai_agent_log_controlled_writeback_consumer_readback_api.py +++ b/apps/api/tests/test_ai_agent_log_controlled_writeback_consumer_readback_api.py @@ -242,6 +242,16 @@ async def test_log_controlled_writeback_consumer_loader_degrades_on_db_pressure( lambda project_id: _FailingContext(), ) + async def fake_direct_unavailable(*, project_id: str): + assert project_id == "awoooi" + return None + + monkeypatch.setattr( + consumer_module, + "_load_consumer_rows_direct", + fake_direct_unavailable, + ) + payload = await load_latest_ai_agent_log_controlled_writeback_consumer_readback() assert payload["status"] == ( @@ -258,6 +268,40 @@ async def test_log_controlled_writeback_consumer_loader_degrades_on_db_pressure( assert payload["operation_boundaries"]["runtime_target_write_performed"] is False +@pytest.mark.asyncio +async def test_log_controlled_writeback_consumer_loader_uses_direct_fallback( + monkeypatch, +): + monkeypatch.setattr( + consumer_module, + "get_db_context", + lambda project_id: _FailingContext(), + ) + + async def fake_direct_fallback(*, project_id: str): + assert project_id == "awoooi" + return _ledger_rows(), _consumer_receipt_rows() + + monkeypatch.setattr( + consumer_module, + "_load_consumer_rows_direct", + fake_direct_fallback, + ) + + payload = await load_latest_ai_agent_log_controlled_writeback_consumer_readback() + + assert payload["status"] == "controlled_writeback_consumer_readback_ready" + assert payload["active_blockers"] == [] + assert payload["controlled_consume"]["controlled_consume_allowed"] is True + assert payload["rollups"]["consumer_apply_receipt_row_count"] == 6 + assert payload["rollups"]["consumer_binding_count"] == 6 + assert payload["rollups"]["ready_consumer_binding_count"] == 6 + assert payload["rollups"]["ready_target_count"] == 6 + assert payload["rollups"]["target_context_receipt_write_count"] == 6 + assert payload["rollups"]["runtime_target_write_performed"] is True + assert payload["operation_boundaries"]["metadata_ledger_read_performed"] is True + + def test_log_controlled_writeback_consumer_endpoint_returns_readback(monkeypatch): async def fake_loader(): fake_db = _FakeDb(_ledger_rows())