Browse Source

feat(agent): add atomic mission observation API

SamLee 1 day ago
parent
commit
fda4d4a74b

+ 12 - 1
agent/agent/orchestration/__init__.py

@@ -67,7 +67,14 @@ from .validation_policy import (
     ValidationContext,
     ValidationPolicy,
 )
-from .wire import OperationView, ProblemDetails
+from .wire import (
+    ArtifactSnapshotView,
+    MissionSnapshotView,
+    OperationView,
+    PlannerDecisionView,
+    ProblemDetails,
+    RootCompletionView,
+)
 
 __all__ = [
     "AcceptanceCriterion",
@@ -76,6 +83,7 @@ __all__ = [
     "ArtifactConflict",
     "ArtifactRef",
     "ArtifactSnapshot",
+    "ArtifactSnapshotView",
     "ArtifactStore",
     "AttemptStatus",
     "AttemptSubmission",
@@ -105,6 +113,7 @@ __all__ = [
     "FileSystemArtifactStore",
     "FileSystemTaskStore",
     "InvalidTaskTransition",
+    "MissionSnapshotView",
     "OperationKind",
     "OperationStatus",
     "OperationView",
@@ -114,8 +123,10 @@ __all__ = [
     "OrchestrationError",
     "OrchestrationEvent",
     "PlannerDecision",
+    "PlannerDecisionView",
     "ProblemDetails",
     "RevisionConflict",
+    "RootCompletionView",
     "RuleBasedDeterministicValidator",
     "TaskAttempt",
     "TaskConflict",

+ 102 - 0
agent/agent/orchestration/api_v2.py

@@ -19,14 +19,18 @@ from .store import (
     TaskStoreNotFound,
 )
 from .wire import (
+    ArtifactSnapshotView,
     AttemptList,
     AttemptView,
     DispatchOperationRequest,
     EventPageView,
     EventView,
+    MissionSnapshotView,
     OperationRequest,
     OperationView,
+    PlannerDecisionView,
     ProblemDetails,
+    RootCompletionView,
     TaskList,
     TaskView,
     ValidationList,
@@ -82,6 +86,104 @@ def create_orchestration_router(coordinator: TaskCoordinator) -> Any:
         prefix="/api/v2", tags=["orchestration-v2"], route_class=ProblemRoute
     )
 
+    @router.get(
+        "/roots/{root_trace_id}/snapshot",
+        response_model=MissionSnapshotView,
+        responses=_ERROR_RESPONSES,
+    )
+    async def get_mission_snapshot(root_trace_id: str) -> Any:
+        path = f"/api/v2/roots/{root_trace_id}/snapshot"
+
+        async def action() -> MissionSnapshotView:
+            # This is deliberately the only load in the endpoint: every
+            # collection must describe one immutable ledger revision.
+            ledger = await coordinator.task_store.load(root_trace_id)
+            return MissionSnapshotView(
+                revision=ledger.revision,
+                root_trace_id=ledger.root_trace_id,
+                root_task_id=ledger.root_task_id,
+                root_objective=ledger.root_objective,
+                focused_task_id=ledger.focused_task_id,
+                tasks=[
+                    TaskView.from_domain(item)
+                    for item in sorted(
+                        ledger.tasks.values(),
+                        key=lambda value: (value.created_at, value.task_id),
+                    )
+                ],
+                attempts=[
+                    AttemptView.from_domain(item)
+                    for item in sorted(
+                        ledger.attempts.values(),
+                        key=lambda value: (value.created_at, value.attempt_id),
+                    )
+                ],
+                validations=[
+                    ValidationView.from_domain(item)
+                    for item in sorted(
+                        ledger.validations.values(),
+                        key=lambda value: (value.created_at, value.validation_id),
+                    )
+                ],
+                decisions=[
+                    PlannerDecisionView.from_domain(item)
+                    for item in sorted(
+                        ledger.decisions.values(),
+                        key=lambda value: (value.created_at, value.decision_id),
+                    )
+                ],
+                operations=[
+                    OperationView.from_domain(item)
+                    for item in sorted(
+                        ledger.operations.values(),
+                        key=lambda value: (value.created_at, value.operation_id),
+                    )
+                ],
+            )
+
+        return await call(action, path)
+
+    @router.get(
+        "/roots/{root_trace_id}/completion",
+        response_model=RootCompletionView,
+        responses=_ERROR_RESPONSES,
+    )
+    async def get_root_completion(root_trace_id: str) -> Any:
+        path = f"/api/v2/roots/{root_trace_id}/completion"
+
+        async def action() -> RootCompletionView:
+            value = await coordinator.root_completion(root_trace_id)
+            return RootCompletionView(
+                root_task_id=value["root_task_id"],
+                root_objective=value["root_objective"],
+                status=value["status"],
+                blocked_reason=value.get("blocked_reason"),
+                result_summary=value.get("result_summary"),
+            )
+
+        return await call(action, path)
+
+    @router.get(
+        "/roots/{root_trace_id}/snapshots/{snapshot_id}",
+        response_model=ArtifactSnapshotView,
+        responses=_ERROR_RESPONSES,
+    )
+    async def get_artifact_snapshot(root_trace_id: str, snapshot_id: str) -> Any:
+        path = f"/api/v2/roots/{root_trace_id}/snapshots/{snapshot_id}"
+
+        async def action() -> ArtifactSnapshotView:
+            try:
+                value = await coordinator.artifact_store.get(
+                    root_trace_id, snapshot_id
+                )
+            except FileNotFoundError as exc:
+                raise _ResourceNotFound(
+                    f"Artifact snapshot not found: {snapshot_id}"
+                ) from exc
+            return ArtifactSnapshotView.from_domain(value)
+
+        return await call(action, path)
+
     def problem(exc: Exception, instance: str) -> JSONResponse:
         status, title, code, detail = _problem_fields(exc)
         value = ProblemDetails(

+ 25 - 0
agent/agent/orchestration/client_v2.py

@@ -9,12 +9,15 @@ from urllib.parse import quote
 import httpx
 
 from .wire import (
+    ArtifactSnapshotView,
     AttemptList,
     AttemptView,
     EventPageView,
     EventView,
+    MissionSnapshotView,
     OperationView,
     ProblemDetails,
+    RootCompletionView,
     TaskList,
     TaskView,
     ValidationList,
@@ -144,6 +147,28 @@ class OrchestrationClient:
             await self._request("GET", f"{self._root(root_trace_id)}/tasks")
         )
 
+    async def get_mission_snapshot(
+        self, root_trace_id: str
+    ) -> MissionSnapshotView:
+        return MissionSnapshotView.model_validate(
+            await self._request("GET", f"{self._root(root_trace_id)}/snapshot")
+        )
+
+    async def get_root_completion(self, root_trace_id: str) -> RootCompletionView:
+        return RootCompletionView.model_validate(
+            await self._request("GET", f"{self._root(root_trace_id)}/completion")
+        )
+
+    async def get_artifact_snapshot(
+        self, root_trace_id: str, snapshot_id: str
+    ) -> ArtifactSnapshotView:
+        return ArtifactSnapshotView.model_validate(
+            await self._request(
+                "GET",
+                f"{self._root(root_trace_id)}/snapshots/{self._part(snapshot_id)}",
+            )
+        )
+
     async def get_task(self, root_trace_id: str, task_id: str) -> TaskView:
         return TaskView.model_validate(
             await self._request(

+ 67 - 0
agent/agent/orchestration/wire.py

@@ -8,12 +8,15 @@ from typing import Annotated, Any, Dict, List, Literal, Optional, Union
 from pydantic import BaseModel, ConfigDict, Field
 
 from .models import (
+    ArtifactSnapshot,
     BackgroundOperation,
     AttemptStatus,
+    DecisionAction,
     FailureCode,
     OperationKind,
     OperationStatus,
     OrchestrationEvent,
+    PlannerDecision,
     TaskAttempt,
     TaskRecord,
     TaskStatus,
@@ -156,6 +159,7 @@ class AttemptView(WireModel):
     worker_trace_id: str
     worker_preset: str
     execution_mode: str
+    accepted_child_decision_ids: Optional[List[str]] = None
     status: AttemptStatus
     operation_id: Optional[str] = None
     execution_epoch: int
@@ -213,6 +217,65 @@ class ValidationView(WireModel):
         return cls.model_validate(data)
 
 
+class PlannerDecisionView(WireModel):
+    decision_id: str
+    task_id: str
+    action: DecisionAction
+    reason: str
+    from_status: TaskStatus
+    to_status: TaskStatus
+    attempt_id: Optional[str] = None
+    validation_id: Optional[str] = None
+    payload: Dict[str, Any]
+    created_at: str
+
+    @classmethod
+    def from_domain(cls, value: PlannerDecision) -> "PlannerDecisionView":
+        return cls.model_validate(json_values(asdict(value)))
+
+
+class MissionSnapshotView(WireModel):
+    schema_version: Literal[1] = 1
+    revision: int
+    root_trace_id: str
+    root_task_id: Optional[str] = None
+    root_objective: str
+    focused_task_id: Optional[str] = None
+    tasks: List[TaskView]
+    attempts: List[AttemptView]
+    validations: List[ValidationView]
+    decisions: List[PlannerDecisionView]
+    operations: List[OperationView]
+
+
+class RootCompletionView(WireModel):
+    root_task_id: str
+    root_objective: str
+    status: TaskStatus
+    blocked_reason: Optional[str] = None
+    result_summary: Optional[str] = None
+
+
+class ArtifactSnapshotView(WireModel):
+    snapshot_id: str
+    attempt_id: str
+    sha256: str
+    artifact_refs: List[ArtifactRefView]
+    evidence_refs: List[ArtifactRefView]
+    created_at: str
+
+    @classmethod
+    def from_domain(cls, value: ArtifactSnapshot) -> "ArtifactSnapshotView":
+        return cls(
+            snapshot_id=value.snapshot_id,
+            attempt_id=value.attempt_id,
+            sha256=value.sha256,
+            artifact_refs=json_values([asdict(item) for item in value.artifact_refs]),
+            evidence_refs=json_values([asdict(item) for item in value.evidence_refs]),
+            created_at=value.created_at,
+        )
+
+
 class TaskList(WireModel):
     revision: int
     root_task_id: Optional[str] = None
@@ -262,16 +325,20 @@ class ProblemDetails(WireModel):
 
 
 __all__ = [
+    "ArtifactSnapshotView",
     "AttemptList",
     "AttemptView",
     "DispatchOperationRequest",
     "EventPageView",
     "EventView",
     "ExecutionStatsView",
+    "MissionSnapshotView",
     "OperationRequest",
     "OperationView",
+    "PlannerDecisionView",
     "ProblemDetails",
     "RevalidateOperationRequest",
+    "RootCompletionView",
     "TaskList",
     "TaskView",
     "ValidationList",

+ 102 - 2
agent/tests/test_orchestration_v2_api.py

@@ -10,11 +10,15 @@ from agent.orchestration.client_v2 import OrchestrationAPIError, OrchestrationCl
 from agent.orchestration.coordinator import TaskConflict
 from agent.orchestration.models import (
     AcceptanceCriterion,
+    ArtifactRef,
+    ArtifactSnapshot,
     BackgroundOperation,
+    DecisionAction,
     EventPage,
     OperationKind,
     OperationStatus,
     OrchestrationEvent,
+    PlannerDecision,
     TaskAttempt,
     TaskLedger,
     TaskRecord,
@@ -32,8 +36,10 @@ class StubStore:
     def __init__(self, ledger: TaskLedger, event: OrchestrationEvent) -> None:
         self.ledger = ledger
         self.event = event
+        self.load_calls = 0
 
     async def load(self, root_trace_id):
+        self.load_calls += 1
         if root_trace_id != self.ledger.root_trace_id:
             raise TaskStoreNotFound(f"No task ledger for root trace {root_trace_id}")
         return self.ledger
@@ -47,8 +53,9 @@ class StubStore:
 
 
 class StubCoordinator:
-    def __init__(self, store, operation):
+    def __init__(self, store, operation, artifact_store):
         self.task_store = store
+        self.artifact_store = artifact_store
         self.operation = operation
         self.calls = []
 
@@ -77,6 +84,28 @@ class StubCoordinator:
         self.calls.append(("resume", idempotency_key))
         return replace(self.operation, status=OperationStatus.PENDING)
 
+    async def root_completion(self, root_trace_id):
+        ledger = await self.task_store.load(root_trace_id)
+        root = ledger.tasks[ledger.root_task_id]
+        return {
+            "root_task_id": root.task_id,
+            "root_objective": ledger.root_objective,
+            "status": root.status.value,
+            "blocked_reason": root.blocked_reason,
+            "result_summary": None,
+            "pending_tasks": ["must not leak"],
+        }
+
+
+class StubArtifactStore:
+    def __init__(self, snapshot):
+        self.snapshot = snapshot
+
+    async def get(self, root_trace_id, snapshot_id):
+        if root_trace_id != "root" or snapshot_id != self.snapshot.snapshot_id:
+            raise FileNotFoundError(snapshot_id)
+        return self.snapshot
+
 
 @pytest.fixture
 def api_fixture():
@@ -106,6 +135,7 @@ def api_fixture():
         worker_trace_id="worker-1",
         worker_preset="worker",
         execution_mode="new",
+        accepted_child_decision_ids=(),
     )
     validation = ValidationReport(
         validation_id="validation-1",
@@ -132,6 +162,17 @@ def api_fixture():
     ledger.tasks[task.task_id] = task
     ledger.attempts[attempt.attempt_id] = attempt
     ledger.validations[validation.validation_id] = validation
+    decision = PlannerDecision(
+        decision_id="decision-1",
+        task_id=task.task_id,
+        action=DecisionAction.RETRY,
+        reason="retry",
+        from_status=TaskStatus.NEEDS_REPLAN,
+        to_status=TaskStatus.PENDING,
+        created_at=task.created_at,
+    )
+    ledger.decisions[decision.decision_id] = decision
+    ledger.operations[operation.operation_id] = operation
     event = OrchestrationEvent(
         schema_version=1,
         event_id="event-1",
@@ -142,7 +183,15 @@ def api_fixture():
         occurred_at="2026-07-18T00:00:00+00:00",
     )
     store = StubStore(ledger, event)
-    coordinator = StubCoordinator(store, operation)
+    snapshot = ArtifactSnapshot(
+        snapshot_id="snapshot-1",
+        attempt_id=attempt.attempt_id,
+        normalized_content={"summary": "must not leak"},
+        sha256="sha256",
+        artifact_refs=[ArtifactRef(uri="memory://artifact", version="1")],
+        evidence_refs=[],
+    )
+    coordinator = StubCoordinator(store, operation, StubArtifactStore(snapshot))
     app = FastAPI()
     app.include_router(create_orchestration_router(coordinator))
     transport = httpx.ASGITransport(app=app)
@@ -270,6 +319,14 @@ async def test_async_client_round_trip_and_watch(api_fixture):
         await client.get_validation("root", "validation-1")
     ).validation_id == "validation-1"
     assert (await client.list_events("root", cursor=None)).next_cursor == "end"
+    snapshot = await client.get_mission_snapshot("root")
+    assert snapshot.revision == 3
+    assert snapshot.root_objective == "test"
+    assert snapshot.attempts[0].accepted_child_decision_ids == []
+    completion = await client.get_root_completion("root")
+    assert completion.root_task_id == "task-1"
+    artifact = await client.get_artifact_snapshot("root", "snapshot-1")
+    assert artifact.sha256 == "sha256"
 
     stream = client.watch_events("root", wait_seconds=0)
     assert (await anext(stream)).event_id == "event-1"
@@ -277,6 +334,49 @@ async def test_async_client_round_trip_and_watch(api_fixture):
     await http.aclose()
 
 
+@pytest.mark.asyncio
+async def test_atomic_mission_snapshot_completion_and_artifact_schema(api_fixture):
+    coordinator, http = api_fixture
+    coordinator.task_store.load_calls = 0
+    response = await http.get("/api/v2/roots/root/snapshot")
+    assert response.status_code == 200
+    assert coordinator.task_store.load_calls == 1
+    body = response.json()
+    assert set(body) == {
+        "schema_version",
+        "revision",
+        "root_trace_id",
+        "root_task_id",
+        "root_objective",
+        "focused_task_id",
+        "tasks",
+        "attempts",
+        "validations",
+        "decisions",
+        "operations",
+    }
+    assert body["schema_version"] == 1
+    assert "request" not in body["operations"][0]
+    assert "request_fingerprint" not in body["operations"][0]
+    assert body["attempts"][0]["accepted_child_decision_ids"] == []
+
+    completion = await http.get("/api/v2/roots/root/completion")
+    assert set(completion.json()) == {
+        "root_task_id",
+        "root_objective",
+        "status",
+        "blocked_reason",
+        "result_summary",
+    }
+    artifact = await http.get("/api/v2/roots/root/snapshots/snapshot-1")
+    assert artifact.status_code == 200
+    assert "normalized_content" not in artifact.json()
+    missing = await http.get("/api/v2/roots/root/snapshots/missing")
+    assert missing.status_code == 404
+    assert missing.json()["code"] == "not_found"
+    await http.aclose()
+
+
 @pytest.mark.asyncio
 async def test_client_transport_retry_and_non_problem_error():
     calls = 0