From 763bd5b97a9c920ae6cb2938c179cbd424fa1aae Mon Sep 17 00:00:00 2001 From: ogt Date: Thu, 9 Jul 2026 18:35:55 +0800 Subject: [PATCH] fix(agents): bound runtime receipt db entry --- .../ai_agent_autonomous_runtime_control.py | 66 +++++++++++++++- ...est_ai_agent_autonomous_runtime_control.py | 77 +++++++++++++++++++ 2 files changed, 139 insertions(+), 4 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 2b286c39d..6f1d3ebc7 100644 --- a/apps/api/src/services/ai_agent_autonomous_runtime_control.py +++ b/apps/api/src/services/ai_agent_autonomous_runtime_control.py @@ -46,6 +46,9 @@ _LIVE_READBACK_SCHEMA_VERSION = "ai_agent_autonomous_runtime_receipt_readback_v1 _DEFAULT_PROJECT_ID = "awoooi" _DEFAULT_LOOKBACK_HOURS = 24 _LIVE_RUNTIME_RECEIPT_TIMEOUT_SECONDS = 5.0 +_RUNTIME_RECEIPT_LOCK_TIMEOUT_SECONDS = 0.25 +_RUNTIME_RECEIPT_DB_CONTEXT_TIMEOUT_SECONDS = 0.7 +_RUNTIME_RECEIPT_DB_CONTEXT_EXIT_TIMEOUT_SECONDS = 0.35 _RUNTIME_RECEIPT_QUERY_BUDGET_SECONDS = 2.75 _RUNTIME_RECEIPT_SINGLE_QUERY_TIMEOUT_SECONDS = 0.7 _RUNTIME_RECEIPT_STATEMENT_TIMEOUT_MS = 650 @@ -4090,7 +4093,14 @@ async def load_ai_agent_autonomous_runtime_receipt_readback( if cached is not None: return cached - async with _get_runtime_receipt_readback_lock(): + lock = _get_runtime_receipt_readback_lock() + lock_acquired = False + try: + await asyncio.wait_for( + lock.acquire(), + timeout=_RUNTIME_RECEIPT_LOCK_TIMEOUT_SECONDS, + ) + lock_acquired = True cached = _runtime_receipt_readback_cache_get(cache_key) if cached is not None: return cached @@ -4101,6 +4111,21 @@ async def load_ai_agent_autonomous_runtime_receipt_readback( ) _runtime_receipt_readback_cache_store(cache_key, readback) return readback + except TimeoutError: + logger.warning( + "ai_agent_autonomous_runtime_receipt_readback_lock_timeout", + project_id=project_id, + timeout_seconds=_RUNTIME_RECEIPT_LOCK_TIMEOUT_SECONDS, + ) + return build_runtime_receipt_readback_from_rows( + project_id=project_id, + lookback_hours=normalized_lookback_hours, + db_read_status="unavailable", + error_type="RuntimeReceiptReadbackLockTimeout", + ) + finally: + if lock_acquired: + lock.release() async def _load_ai_agent_autonomous_runtime_receipt_readback_uncached( @@ -4118,8 +4143,17 @@ async def _load_ai_agent_autonomous_runtime_receipt_readback_uncached( } log_controlled_writeback_consumer: dict[str, Any] | None = None partial_query_failures: list[dict[str, str]] = [] + db_context = get_db_context(project_id) + db_context_entered = False + db_context_stage = "db_context_enter" try: - async with get_db_context(project_id) as db: + db = await asyncio.wait_for( + db_context.__aenter__(), + timeout=_RUNTIME_RECEIPT_DB_CONTEXT_TIMEOUT_SECONDS, + ) + db_context_entered = True + db_context_stage = "db_queries" + try: query_deadline = time.monotonic() + min( _RUNTIME_RECEIPT_QUERY_BUDGET_SECONDS, max(0.5, _LIVE_RUNTIME_RECEIPT_TIMEOUT_SECONDS - 0.75), @@ -4300,11 +4334,35 @@ async def _load_ai_agent_autonomous_runtime_receipt_readback_uncached( "grouped_alert_event_counts", _RUNTIME_GROUPED_ALERT_EVENT_COUNTS_SQL, ) + finally: + if db_context_entered: + db_context_stage = "db_context_exit" + try: + await asyncio.wait_for( + db_context.__aexit__(None, None, None), + timeout=_RUNTIME_RECEIPT_DB_CONTEXT_EXIT_TIMEOUT_SECONDS, + ) + except Exception as exit_exc: # pragma: no cover - live DB state + partial_query_failures.append({ + "query_name": "db_context_exit", + "error_type": type(exit_exc).__name__, + }) + logger.warning( + "ai_agent_autonomous_runtime_trace_db_context_exit_failed", + project_id=project_id, + error_type=type(exit_exc).__name__, + ) except Exception as exc: + error_type = ( + "RuntimeReceiptDbContextTimeout" + if isinstance(exc, TimeoutError) and db_context_stage == "db_context_enter" + else type(exc).__name__ + ) logger.warning( "ai_agent_autonomous_runtime_receipt_readback_failed", project_id=project_id, - error_type=type(exc).__name__, + error_type=error_type, + db_context_stage=db_context_stage, ) log_controlled_writeback_consumer = ( await _load_log_controlled_writeback_consumer_readback( @@ -4315,7 +4373,7 @@ async def _load_ai_agent_autonomous_runtime_receipt_readback_uncached( project_id=project_id, lookback_hours=params["lookback_hours"], db_read_status="unavailable", - error_type=type(exc).__name__, + error_type=error_type, log_controlled_writeback_consumer=log_controlled_writeback_consumer, ) diff --git a/apps/api/tests/test_ai_agent_autonomous_runtime_control.py b/apps/api/tests/test_ai_agent_autonomous_runtime_control.py index c36241376..c4f2e5635 100644 --- a/apps/api/tests/test_ai_agent_autonomous_runtime_control.py +++ b/apps/api/tests/test_ai_agent_autonomous_runtime_control.py @@ -76,6 +76,15 @@ class _BudgetRuntimeDbContext: return False +class _SlowEnterRuntimeDbContext: + async def __aenter__(self): + await asyncio.sleep(1) + return _FailingRuntimeDb() + + async def __aexit__(self, _exc_type, _exc, _tb) -> bool: + return False + + def _assert_log_controlled_writeback_executor(payload: dict): assert payload["schema_version"] == "ai_agent_log_controlled_writeback_executor_readback_v1" assert payload["priority"] == "P1-LOG-KM-RAG-MCP-PLAYBOOK" @@ -287,6 +296,74 @@ async def test_live_runtime_receipt_keeps_consumer_readback_when_trace_query_fai ) +@pytest.mark.asyncio +async def test_live_runtime_receipt_db_context_timeout_returns_degraded_payload( + monkeypatch, +): + runtime_control_module._clear_runtime_receipt_readback_cache() + + async def _fake_consumer_readback(*, project_id: str): + assert project_id == "awoooi" + return _log_controlled_writeback_consumer_readback() + + monkeypatch.setattr( + runtime_control_module, + "_RUNTIME_RECEIPT_DB_CONTEXT_TIMEOUT_SECONDS", + 0.01, + ) + monkeypatch.setattr( + runtime_control_module, + "get_db_context", + lambda project_id: _SlowEnterRuntimeDbContext(), + ) + monkeypatch.setattr( + runtime_control_module, + "_load_log_controlled_writeback_consumer_readback", + _fake_consumer_readback, + ) + + try: + loader = getattr( + runtime_control_module, + "_load_ai_agent_autonomous_runtime_receipt_readback_uncached", + ) + readback = await loader() + finally: + runtime_control_module._clear_runtime_receipt_readback_cache() + + assert readback["db_read_status"] == "unavailable" + assert readback["error"]["type"] == "RuntimeReceiptDbContextTimeout" + _assert_log_controlled_writeback_consumer( + readback["log_controlled_writeback_consumer"] + ) + + +@pytest.mark.asyncio +async def test_live_runtime_receipt_lock_timeout_returns_degraded_payload( + monkeypatch, +): + runtime_control_module._clear_runtime_receipt_readback_cache() + lock = asyncio.Lock() + await lock.acquire() + monkeypatch.setattr(runtime_control_module, "_runtime_receipt_readback_lock", lock) + monkeypatch.setattr( + runtime_control_module, + "_RUNTIME_RECEIPT_LOCK_TIMEOUT_SECONDS", + 0.01, + ) + + try: + readback = ( + await runtime_control_module.load_ai_agent_autonomous_runtime_receipt_readback() + ) + finally: + lock.release() + runtime_control_module._clear_runtime_receipt_readback_cache() + + assert readback["db_read_status"] == "unavailable" + assert readback["error"]["type"] == "RuntimeReceiptReadbackLockTimeout" + + @pytest.mark.asyncio async def test_live_runtime_receipt_readback_returns_partial_when_query_budget_expires( monkeypatch,