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 9eb42ec09..2b286c39d 100644 --- a/apps/api/src/services/ai_agent_autonomous_runtime_control.py +++ b/apps/api/src/services/ai_agent_autonomous_runtime_control.py @@ -46,8 +46,10 @@ _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_QUERY_BUDGET_SECONDS = 4.0 -_RUNTIME_RECEIPT_STATEMENT_TIMEOUT_MS = 900 +_RUNTIME_RECEIPT_QUERY_BUDGET_SECONDS = 2.75 +_RUNTIME_RECEIPT_SINGLE_QUERY_TIMEOUT_SECONDS = 0.7 +_RUNTIME_RECEIPT_STATEMENT_TIMEOUT_MS = 650 +_LOG_CONTROLLED_WRITEBACK_CONSUMER_TIMEOUT_SECONDS = 0.75 _RUNTIME_RECEIPT_READBACK_CACHE_TTL_SECONDS = 20.0 # CD cancel-stale-cd no-op triggers must not change runtime payloads. _EXECUTOR_OPERATION_TYPES = ( @@ -439,8 +441,11 @@ async def _load_log_controlled_writeback_consumer_readback( """Attach LOG consumer bindings without writing KM/RAG/PlayBook/MCP targets.""" try: - return await load_latest_ai_agent_log_controlled_writeback_consumer_readback( - project_id=project_id, + return await asyncio.wait_for( + load_latest_ai_agent_log_controlled_writeback_consumer_readback( + project_id=project_id, + ), + timeout=_LOG_CONTROLLED_WRITEBACK_CONSUMER_TIMEOUT_SECONDS, ) except Exception as exc: # pragma: no cover - keeps runtime control API visible logger.warning( @@ -4121,11 +4126,14 @@ async def _load_ai_agent_autonomous_runtime_receipt_readback_uncached( ) async def _set_statement_timeout() -> None: - await db.execute( - text( - "SET LOCAL statement_timeout = " - f"'{int(_RUNTIME_RECEIPT_STATEMENT_TIMEOUT_MS)}ms'" - ) + await asyncio.wait_for( + db.execute( + text( + "SET LOCAL statement_timeout = " + f"'{int(_RUNTIME_RECEIPT_STATEMENT_TIMEOUT_MS)}ms'" + ) + ), + timeout=_RUNTIME_RECEIPT_SINGLE_QUERY_TIMEOUT_SECONDS, ) await _set_statement_timeout() @@ -4152,7 +4160,11 @@ async def _load_ai_agent_autonomous_runtime_receipt_readback_uncached( if _query_budget_exhausted(query_name): return [] try: - return (await db.execute(text(sql), params)).mappings().all() + result = await asyncio.wait_for( + db.execute(text(sql), params), + timeout=_RUNTIME_RECEIPT_SINGLE_QUERY_TIMEOUT_SECONDS, + ) + return result.mappings().all() except Exception as exc: # pragma: no cover - depends on live schema drift partial_query_failures.append({ "query_name": query_name, @@ -4180,7 +4192,11 @@ async def _load_ai_agent_autonomous_runtime_receipt_readback_uncached( if _query_budget_exhausted(f"{query_name}_fallback"): return [] try: - return (await db.execute(text(fallback_sql), params)).mappings().all() + result = await asyncio.wait_for( + db.execute(text(fallback_sql), params), + timeout=_RUNTIME_RECEIPT_SINGLE_QUERY_TIMEOUT_SECONDS, + ) + return result.mappings().all() except Exception as fallback_exc: # pragma: no cover - live schema drift partial_query_failures.append({ "query_name": f"{query_name}_fallback", 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 0e58b6991..c36241376 100644 --- a/apps/api/tests/test_ai_agent_autonomous_runtime_control.py +++ b/apps/api/tests/test_ai_agent_autonomous_runtime_control.py @@ -44,6 +44,7 @@ class _FakeMappingResult: class _BudgetRuntimeDb: def __init__(self): self.statements: list[str] = [] + self.expire_budget_after_operation_counts = None async def execute(self, statement, *_args, **_kwargs): sql = str(statement) @@ -51,6 +52,8 @@ class _BudgetRuntimeDb: if "SET LOCAL statement_timeout" in sql: return _FakeMappingResult([]) if "FROM automation_operation_log" in sql and "GROUP BY 1, status" in sql: + if self.expire_budget_after_operation_counts is not None: + self.expire_budget_after_operation_counts() return _FakeMappingResult([ { "operation_type": "ansible_apply_executed", @@ -293,9 +296,13 @@ async def test_live_runtime_receipt_readback_returns_partial_when_query_budget_e tick = {"value": 100.0} def _fake_monotonic(): - tick["value"] += 0.05 return tick["value"] + def _expire_budget(): + tick["value"] = 101.0 + + fake_db.expire_budget_after_operation_counts = _expire_budget + async def _fake_consumer_readback(*, project_id: str): assert project_id == "awoooi" return _log_controlled_writeback_consumer_readback() @@ -345,6 +352,37 @@ async def test_live_runtime_receipt_readback_returns_partial_when_query_budget_e ) +@pytest.mark.asyncio +async def test_log_controlled_writeback_consumer_timeout_returns_fallback( + monkeypatch, +): + async def _slow_consumer_readback(*, project_id: str): + assert project_id == "awoooi" + await asyncio.sleep(1) + return _log_controlled_writeback_consumer_readback() + + monkeypatch.setattr( + runtime_control_module, + "_LOG_CONTROLLED_WRITEBACK_CONSUMER_TIMEOUT_SECONDS", + 0.01, + ) + monkeypatch.setattr( + runtime_control_module, + "load_latest_ai_agent_log_controlled_writeback_consumer_readback", + _slow_consumer_readback, + ) + + readback = await runtime_control_module._load_log_controlled_writeback_consumer_readback( + project_id="awoooi", + ) + + assert readback["status"] == "blocked_waiting_controlled_writeback_consumer_receipts" + assert "log_controlled_writeback_consumer_readback_unavailable" in readback[ + "active_blockers" + ] + assert readback["readback"]["error_type"] == "TimeoutError" + + @pytest.mark.asyncio async def test_live_runtime_receipt_readback_uses_short_cache_to_protect_db_pool( monkeypatch,