diff --git a/apps/api/src/services/platform_operator_service.py b/apps/api/src/services/platform_operator_service.py index b60850102..ea1eaf481 100644 --- a/apps/api/src/services/platform_operator_service.py +++ b/apps/api/src/services/platform_operator_service.py @@ -106,6 +106,8 @@ _REMEDIATION_HISTORY_SCHEMA_VERSIONS = { "adr100_remediation_dry_run_history_v1", "adr100_remediation_approval_history_v1", } +_MAX_REMEDIATION_HISTORY_BATCH_INCIDENTS = _MAX_LIST_CONTEXT_ROWS +_MAX_REMEDIATION_HISTORY_BATCH_ROWS = 5_000 _ADR100_GATE5_PROJECTION_TRIGGER = "adr100_runtime_replay_gate5" _CALLBACK_REPLY_CACHE_TTL_SECONDS = int( os.getenv("AWOOOP_CALLBACK_REPLY_CACHE_TTL_SECONDS", "20") @@ -7617,17 +7619,21 @@ async def _fetch_run_remediation_histories_by_incident( db: Any | None = None, ) -> tuple[dict[str, list[dict[str, Any]]], dict[str, dict[str, str]]]: """Fetch ADR-100 list evidence with one bounded query for all incidents.""" - normalized_incident_ids = list(dict.fromkeys(incident_ids)) + requested_incident_ids = list(dict.fromkeys(incident_ids)) + normalized_incident_ids = requested_incident_ids[ + :_MAX_REMEDIATION_HISTORY_BATCH_INCIDENTS + ] if not normalized_incident_ids: return {}, {} safe_limit = max(1, min(limit, 200)) per_event_fetch_limit = min(max(safe_limit * 4, 50), 200) - batch_limit = ( + requested_batch_limit = ( len(normalized_incident_ids) * len(_REMEDIATION_HISTORY_EVENT_TYPES) * per_event_fetch_limit ) + batch_limit = min(requested_batch_limit, _MAX_REMEDIATION_HISTORY_BATCH_ROWS) event_rank = func.row_number().over( partition_by=( AlertOperationLog.incident_id, @@ -7654,7 +7660,7 @@ async def _fetch_run_remediation_histories_by_incident( ) .where(ranked_rows.c.event_rank <= per_event_fetch_limit) .order_by(AlertOperationLog.created_at.desc()) - .limit(batch_limit) + .limit(batch_limit + 1) ) try: @@ -7665,17 +7671,19 @@ async def _fetch_run_remediation_histories_by_incident( result = await db.execute(stmt) rows = list(result.scalars().all()) except Exception as exc: - error = str(exc) + error = f"remediation_history_batch_fetch_failed:{type(exc).__name__}" logger.warning( "run_list_remediation_history_batch_fetch_failed", incident_count=len(normalized_incident_ids), - error=error, + error_type=type(exc).__name__, ) return {}, { incident_id: {"incident_id": incident_id, "error": error} - for incident_id in normalized_incident_ids + for incident_id in requested_incident_ids } + row_limit_reached = len(rows) > batch_limit + rows = rows[:batch_limit] rows_by_incident: dict[str, list[AlertOperationLog]] = defaultdict(list) for row in rows: incident_id = str(getattr(row, "incident_id", "") or "") @@ -7686,6 +7694,23 @@ async def _fetch_run_remediation_histories_by_incident( histories_by_incident: dict[str, list[dict[str, Any]]] = {} errors_by_incident: dict[str, dict[str, str]] = {} + if row_limit_reached: + errors_by_incident.update({ + incident_id: { + "incident_id": incident_id, + "error": "remediation_history_batch_row_limit_reached", + } + for incident_id in normalized_incident_ids + }) + errors_by_incident.update({ + incident_id: { + "incident_id": incident_id, + "error": "remediation_history_batch_incident_limit_reached", + } + for incident_id in requested_incident_ids[ + _MAX_REMEDIATION_HISTORY_BATCH_INCIDENTS: + ] + }) for incident_id in normalized_incident_ids: try: items: list[dict[str, Any]] = [] @@ -7706,11 +7731,11 @@ async def _fetch_run_remediation_histories_by_incident( break histories_by_incident[incident_id] = items except Exception as exc: - error = str(exc) + error = f"remediation_history_row_decode_failed:{type(exc).__name__}" logger.warning( "run_list_remediation_history_row_decode_failed", incident_id=incident_id, - error=error, + error_type=type(exc).__name__, ) errors_by_incident[incident_id] = { "incident_id": incident_id, @@ -7754,9 +7779,16 @@ async def _build_run_remediation_summaries( db=db, ) ) + legacy_incident_ids = all_incident_ids[ + :_MAX_REMEDIATION_HISTORY_BATCH_INCIDENTS + ] legacy_mcp_by_incident = await _fetch_legacy_mcp_by_incident_ids( - all_incident_ids, - limit=min(max(len(all_incident_ids) * _REMEDIATION_HISTORY_LIMIT, 100), 5_000), + legacy_incident_ids, + limit=min( + max(len(legacy_incident_ids) * _REMEDIATION_HISTORY_LIMIT, 100), + 5_000, + ), + db=db, ) summaries: dict[UUID, dict[str, Any]] = {} @@ -7881,19 +7913,24 @@ async def _fetch_legacy_mcp_by_incident_ids( incident_ids: list[str], *, limit: int, + db: Any | None = None, ) -> dict[str, list[dict[str, Any]]]: """Fetch legacy/self-built MCP rows for list evidence summaries.""" if not incident_ids: return {} - async with get_db_context("awoooi") as db: - result = await db.execute( - select(MCPAuditLog) - .where(MCPAuditLog.incident_id.in_(incident_ids)) - .order_by(MCPAuditLog.created_at.desc()) - .limit(limit) - ) - rows = list(result.scalars().all()) + stmt = ( + select(MCPAuditLog) + .where(MCPAuditLog.incident_id.in_(incident_ids)) + .order_by(MCPAuditLog.created_at.desc()) + .limit(limit) + ) + if db is None: + async with get_db_context("awoooi") as legacy_db: + result = await legacy_db.execute(stmt) + else: + result = await db.execute(stmt) + rows = list(result.scalars().all()) by_incident: dict[str, list[dict[str, Any]]] = defaultdict(list) for row in rows: diff --git a/apps/api/tests/test_platform_operator_run_remediation_batch.py b/apps/api/tests/test_platform_operator_run_remediation_batch.py index 01ed35019..491bf08f7 100644 --- a/apps/api/tests/test_platform_operator_run_remediation_batch.py +++ b/apps/api/tests/test_platform_operator_run_remediation_batch.py @@ -29,19 +29,22 @@ class _FakeDb: self, rows: list[SimpleNamespace], *, - error: Exception | None = None, + followup_rows: list[list[SimpleNamespace]] | None = None, + errors: list[Exception | None] | None = None, ) -> None: - self._rows = rows - self._error = error + self._row_batches = [rows, *(followup_rows or [])] + self._errors = errors or [] self.execute_calls = 0 self.statements: list[object] = [] async def execute(self, statement: object) -> _FakeResult: self.execute_calls += 1 self.statements.append(statement) - if self._error is not None: - raise self._error - return _FakeResult(self._rows) + call_index = self.execute_calls - 1 + if call_index < len(self._errors) and self._errors[call_index] is not None: + raise self._errors[call_index] + row_index = min(call_index, len(self._row_batches) - 1) + return _FakeResult(self._row_batches[row_index]) class _FakeDbContext: @@ -84,15 +87,6 @@ def _history_row(incident_id: str, sequence: int) -> SimpleNamespace: ) -async def _empty_legacy_history( - _incident_ids: list[str], - *, - limit: int, -) -> dict[str, list[dict[str, object]]]: - assert limit > 0 - return {} - - @pytest.mark.asyncio async def test_run_remediation_summaries_use_one_bounded_batch_query( monkeypatch: pytest.MonkeyPatch, @@ -114,7 +108,7 @@ async def test_run_remediation_summaries_use_one_bounded_batch_query( for incident_id in incident_by_run.values() for sequence in range(25) ] - db = _FakeDb(rows) + db = _FakeDb(rows, followup_rows=[[]]) enter_count = [0] monkeypatch.setattr( @@ -127,11 +121,6 @@ async def test_run_remediation_summaries_use_one_bounded_batch_query( "_collect_run_incident_ids", lambda *, run, inbound_events, outbound_messages: [incident_by_run[run.run_id]], ) - monkeypatch.setattr( - platform_operator_service, - "_fetch_legacy_mcp_by_incident_ids", - _empty_legacy_history, - ) async def forbidden_per_incident_history( *_args: object, **_kwargs: object @@ -154,8 +143,8 @@ async def test_run_remediation_summaries_use_one_bounded_batch_query( ) assert enter_count == [0] - assert db.execute_calls == 1 - assert len(db.statements) == 1 + assert db.execute_calls == 2 + assert len(db.statements) == 2 assert summaries[run_a.run_id]["total"] == 20 assert summaries[run_b.run_id]["total"] == 20 assert summaries[run_a.run_id]["status"] == "read_only_dry_run" @@ -180,7 +169,11 @@ async def test_run_remediation_batch_failure_is_reported_for_each_incident( run_a.run_id: "INC-20260714-A001", run_b.run_id: "INC-20260714-B001", } - db = _FakeDb([], error=RuntimeError("database pool exhausted")) + db = _FakeDb( + [], + followup_rows=[[]], + errors=[RuntimeError("database pool exhausted"), None], + ) enter_count = [0] monkeypatch.setattr( @@ -193,12 +186,6 @@ async def test_run_remediation_batch_failure_is_reported_for_each_incident( "_collect_run_incident_ids", lambda *, run, inbound_events, outbound_messages: [incident_by_run[run.run_id]], ) - monkeypatch.setattr( - platform_operator_service, - "_fetch_legacy_mcp_by_incident_ids", - _empty_legacy_history, - ) - summaries = await platform_operator_service._build_run_remediation_summaries( runs=[run_a, run_b], inbound_by_run={}, @@ -207,10 +194,68 @@ async def test_run_remediation_batch_failure_is_reported_for_each_incident( ) assert enter_count == [0] - assert db.execute_calls == 1 + assert db.execute_calls == 2 for run in (run_a, run_b): incident_id = incident_by_run[run.run_id] assert summaries[run.run_id]["status"] == "no_evidence" assert summaries[run.run_id]["errors"] == [ - {"incident_id": incident_id, "error": "database pool exhausted"} + { + "incident_id": incident_id, + "error": "remediation_history_batch_fetch_failed:RuntimeError", + } ] + + +@pytest.mark.asyncio +async def test_run_remediation_batch_has_global_incident_and_row_limits() -> None: + requested_count = ( + platform_operator_service._MAX_REMEDIATION_HISTORY_BATCH_INCIDENTS + 2 + ) + incident_ids = [f"INC-20260714-{index:04d}" for index in range(requested_count)] + db = _FakeDb([]) + + histories, errors = await ( + platform_operator_service._fetch_run_remediation_histories_by_incident( + incident_ids, + db=db, + ) + ) + + assert len(histories) == ( + platform_operator_service._MAX_REMEDIATION_HISTORY_BATCH_INCIDENTS + ) + assert all(items == [] for items in histories.values()) + assert db.execute_calls == 1 + statement = db.statements[0] + assert statement._limit_clause.value == ( + platform_operator_service._MAX_REMEDIATION_HISTORY_BATCH_ROWS + 1 + ) + assert incident_ids[0] not in errors + assert errors[incident_ids[-1]]["error"] == ( + "remediation_history_batch_incident_limit_reached" + ) + + +@pytest.mark.asyncio +async def test_run_remediation_batch_reports_only_an_observed_row_truncation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + incident_id = "INC-20260714-A001" + monkeypatch.setattr( + platform_operator_service, + "_MAX_REMEDIATION_HISTORY_BATCH_ROWS", + 2, + ) + db = _FakeDb([_history_row(incident_id, sequence) for sequence in range(3)]) + + _histories, errors = await ( + platform_operator_service._fetch_run_remediation_histories_by_incident( + [incident_id], + db=db, + ) + ) + + assert db.statements[0]._limit_clause.value == 3 + assert errors[incident_id]["error"] == ( + "remediation_history_batch_row_limit_reached" + )