From ce9f45b590968121671e9a1f5bb4f21e05d9c9f5 Mon Sep 17 00:00:00 2001 From: ogt Date: Thu, 9 Jul 2026 17:48:26 +0800 Subject: [PATCH] fix(agents): budget runtime receipt readback queries --- .../ai_agent_autonomous_runtime_control.py | 32 +++++- ...est_ai_agent_autonomous_runtime_control.py | 104 ++++++++++++++++++ 2 files changed, 135 insertions(+), 1 deletion(-) 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 a4cec5676..9eb42ec09 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,8 @@ _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_READBACK_CACHE_TTL_SECONDS = 20.0 # CD cancel-stale-cd no-op triggers must not change runtime payloads. _EXECUTOR_OPERATION_TYPES = ( @@ -4113,16 +4115,42 @@ async def _load_ai_agent_autonomous_runtime_receipt_readback_uncached( partial_query_failures: list[dict[str, str]] = [] try: async with get_db_context(project_id) as db: + query_deadline = time.monotonic() + min( + _RUNTIME_RECEIPT_QUERY_BUDGET_SECONDS, + max(0.5, _LIVE_RUNTIME_RECEIPT_TIMEOUT_SECONDS - 0.75), + ) + async def _set_statement_timeout() -> None: - await db.execute(text("SET LOCAL statement_timeout = '5000ms'")) + await db.execute( + text( + "SET LOCAL statement_timeout = " + f"'{int(_RUNTIME_RECEIPT_STATEMENT_TIMEOUT_MS)}ms'" + ) + ) await _set_statement_timeout() + def _query_budget_exhausted(query_name: str) -> bool: + if time.monotonic() < query_deadline: + return False + partial_query_failures.append({ + "query_name": query_name, + "error_type": "RuntimeReceiptQueryBudgetExceeded", + }) + logger.warning( + "ai_agent_autonomous_runtime_trace_query_budget_exhausted", + project_id=project_id, + query_name=query_name, + ) + return True + async def _safe_rows( query_name: str, sql: str, fallback_sql: str | None = None, ) -> list[Mapping[str, Any]]: + if _query_budget_exhausted(query_name): + return [] try: return (await db.execute(text(sql), params)).mappings().all() except Exception as exc: # pragma: no cover - depends on live schema drift @@ -4149,6 +4177,8 @@ async def _load_ai_agent_autonomous_runtime_receipt_readback_uncached( error_type=type(reset_exc).__name__, ) if fallback_sql: + if _query_budget_exhausted(f"{query_name}_fallback"): + return [] try: return (await db.execute(text(fallback_sql), params)).mappings().all() except Exception as fallback_exc: # pragma: no cover - live schema drift 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 d5140c074..0e58b6991 100644 --- a/apps/api/tests/test_ai_agent_autonomous_runtime_control.py +++ b/apps/api/tests/test_ai_agent_autonomous_runtime_control.py @@ -30,6 +30,49 @@ class _FailingRuntimeDbContext: return False +class _FakeMappingResult: + def __init__(self, rows: list[dict]): + self._rows = rows + + def mappings(self): + return self + + def all(self): + return self._rows + + +class _BudgetRuntimeDb: + def __init__(self): + self.statements: list[str] = [] + + async def execute(self, statement, *_args, **_kwargs): + sql = str(statement) + self.statements.append(sql) + if "SET LOCAL statement_timeout" in sql: + return _FakeMappingResult([]) + if "FROM automation_operation_log" in sql and "GROUP BY 1, status" in sql: + return _FakeMappingResult([ + { + "operation_type": "ansible_apply_executed", + "status": "success", + "total": 1, + "recent": 1, + } + ]) + return _FakeMappingResult([]) + + +class _BudgetRuntimeDbContext: + def __init__(self, db: _BudgetRuntimeDb): + self._db = db + + async def __aenter__(self) -> _BudgetRuntimeDb: + return self._db + + 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" @@ -241,6 +284,67 @@ async def test_live_runtime_receipt_keeps_consumer_readback_when_trace_query_fai ) +@pytest.mark.asyncio +async def test_live_runtime_receipt_readback_returns_partial_when_query_budget_expires( + monkeypatch, +): + runtime_control_module._clear_runtime_receipt_readback_cache() + fake_db = _BudgetRuntimeDb() + tick = {"value": 100.0} + + def _fake_monotonic(): + tick["value"] += 0.05 + return tick["value"] + + async def _fake_consumer_readback(*, project_id: str): + assert project_id == "awoooi" + return _log_controlled_writeback_consumer_readback() + + monkeypatch.setattr(runtime_control_module.time, "monotonic", _fake_monotonic) + monkeypatch.setattr( + runtime_control_module, + "_RUNTIME_RECEIPT_QUERY_BUDGET_SECONDS", + 0.11, + ) + monkeypatch.setattr( + runtime_control_module, + "_RUNTIME_RECEIPT_STATEMENT_TIMEOUT_MS", + 25, + ) + monkeypatch.setattr( + runtime_control_module, + "get_db_context", + lambda project_id: _BudgetRuntimeDbContext(fake_db), + ) + 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"] == "partial" + assert readback["ansible_apply_executed"]["total"] == 1 + assert any( + item["error_type"] == "RuntimeReceiptQueryBudgetExceeded" + for item in readback["partial_query_failures"] + ) + assert readback["runtime_receipt_readback_recovery"]["status"] == ( + "degraded_runtime_receipt_partial_query_failures" + ) + assert any( + "SET LOCAL statement_timeout = '25ms'" in sql for sql in fake_db.statements + ) + + @pytest.mark.asyncio async def test_live_runtime_receipt_readback_uses_short_cache_to_protect_db_pool( monkeypatch,