diff --git a/src/youtube_extension/services/agents/adapters/agent_orchestrator.py b/src/youtube_extension/services/agents/adapters/agent_orchestrator.py index 6ffd1b345..1103ec43a 100644 --- a/src/youtube_extension/services/agents/adapters/agent_orchestrator.py +++ b/src/youtube_extension/services/agents/adapters/agent_orchestrator.py @@ -16,6 +16,7 @@ from typing import Any, Optional from ..base_agent import AgentRequest, AgentResult, BaseAgent +from ..registry import get as get_agent_class @dataclass @@ -47,20 +48,18 @@ class OrchestrationResult: timestamp: datetime = field(default_factory=datetime.now) -from ..registry import get as get_agent_class - - class AgentOrchestrator: """ Centralized orchestration for AI agents. Handles task delegation, parallel processing, and result aggregation. """ - def __init__(self): + def __init__(self, audit_store: Any | None = None): """Initialize agent orchestrator""" self.logger = logging.getLogger("agent_orchestrator") self._agents: dict[str, BaseAgent] = {} self._agent_types: dict[str, type[BaseAgent]] = {} + self._audit_store = audit_store # Bounded: the module-level `orchestrator` singleton lives for the whole # process and every dispatch appends here, so an unbounded list would # grow without limit. maxlen evicts the oldest entries automatically. @@ -517,6 +516,38 @@ async def execute_antigravity_backend( }, ) ) + if self._audit_store is not None: + input_tokens = receipt.usage.get("input_tokens") + output_tokens = receipt.usage.get("output_tokens") + self._audit_store.append( + receipt.receipt_id, + agent_id="google_antigravity", + action="managed_backend_dispatch", + success=receipt.status == "completed" and receipt.error is None, + duration_ms=max(receipt.elapsed_seconds * 1000.0, 0.0), + details={ + "backend": "google_antigravity", + "provider": receipt.provider, + "agent": receipt.agent, + "status": receipt.status, + "receipt_id": receipt.receipt_id, + "request_sha256": receipt.request_sha256, + "interaction_id": receipt.interaction_id, + "environment_id": receipt.environment_id, + "budget_exceeded": receipt.budget_exceeded, + "mcp_servers": list(receipt.mcp_servers), + "policy": dict(receipt.policy), + "error": receipt.error, + }, + input_tokens=( + int(input_tokens) if isinstance(input_tokens, (int, float)) else None + ), + output_tokens=( + int(output_tokens) + if isinstance(output_tokens, (int, float)) + else None + ), + ) return receipt # --- Legacy local Antigravity-pattern workflow --- diff --git a/tests/unit/test_antigravity_backend.py b/tests/unit/test_antigravity_backend.py index 2286bde84..7b028f052 100644 --- a/tests/unit/test_antigravity_backend.py +++ b/tests/unit/test_antigravity_backend.py @@ -2,7 +2,8 @@ from __future__ import annotations -from typing import Any, Mapping +from collections.abc import Mapping +from typing import Any import pytest @@ -39,6 +40,36 @@ async def create_interaction(self, payload: dict[str, Any]) -> Mapping[str, Any] return self.response +class FakeAuditStore: + def __init__(self) -> None: + self.entries: list[dict[str, Any]] = [] + + def append( + self, + run_id: str, + *, + agent_id: str, + action: str, + success: bool, + duration_ms: float, + details: dict[str, Any] | None = None, + input_tokens: int | None = None, + output_tokens: int | None = None, + ) -> None: + self.entries.append( + { + "run_id": run_id, + "agent_id": agent_id, + "action": action, + "success": success, + "duration_ms": duration_ms, + "details": details or {}, + "input_tokens": input_tokens, + "output_tokens": output_tokens, + } + ) + + def config(**changes: Any) -> AntigravityBackendConfig: values = { "enabled": True, @@ -213,6 +244,32 @@ async def test_orchestrator_records_managed_backend_dispatch() -> None: assert "context" not in entry +@pytest.mark.asyncio +async def test_orchestrator_records_managed_backend_receipt_durably() -> None: + from youtube_extension.services.agents.adapters.agent_orchestrator import ( + AgentOrchestrator, + ) + + audit_store = FakeAuditStore() + orchestrator = AgentOrchestrator(audit_store=audit_store) + receipt = await orchestrator.execute_antigravity_backend( + backend=AntigravityBackend(config(), FakeTransport()), + task="review", + context={"video_pack_id": "pack-3"}, + ) + + assert len(audit_store.entries) == 1 + entry = audit_store.entries[0] + assert entry["run_id"] == receipt.receipt_id + assert entry["agent_id"] == "google_antigravity" + assert entry["action"] == "managed_backend_dispatch" + assert entry["success"] is True + assert entry["duration_ms"] >= 0 + assert entry["details"]["receipt_id"] == receipt.receipt_id + assert entry["details"]["request_sha256"] == receipt.request_sha256 + assert "context" not in entry["details"] + + @pytest.mark.asyncio async def test_comparison_artifact_is_provider_neutral() -> None: transport = FakeTransport()