from __future__ import annotations from dataclasses import replace import httpx import pytest from agent.orchestration.api_v2 import create_orchestration_router from agent.orchestration.client_v2 import OrchestrationAPIError, OrchestrationClient from agent.orchestration.coordinator import TaskConflict from agent.orchestration.models import ( AcceptanceCriterion, ArtifactRef, ArtifactSnapshot, BackgroundOperation, DecisionAction, EventPage, OperationKind, OperationStatus, OrchestrationEvent, PlannerDecision, TaskAttempt, TaskLedger, TaskRecord, TaskSpec, TaskStatus, ValidationReport, ) from agent.orchestration.store import TaskStoreError, TaskStoreNotFound FastAPI = pytest.importorskip("fastapi").FastAPI 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 async def list_events(self, root_trace_id, cursor=None, limit=100): await self.load(root_trace_id) if cursor == "bad": raise ValueError("Invalid event cursor") events = [] if cursor == "end" else [self.event] return EventPage(events=events, next_cursor="end", has_more=False) class StubCoordinator: def __init__(self, store, operation, artifact_store): self.task_store = store self.artifact_store = artifact_store self.operation = operation self.calls = [] async def start_operation(self, root_trace_id, kind, **kwargs): self.calls.append(("start", root_trace_id, kind, kwargs)) return replace( self.operation, root_trace_id=root_trace_id, kind=OperationKind(kind), ) async def get_operation(self, root_trace_id, operation_id): if operation_id == "conflict": raise TaskConflict("operation conflict") if operation_id == "unavailable": raise TaskStoreError("storage unavailable") if operation_id != self.operation.operation_id: raise ValueError(f"Operation not found: {operation_id}") return self.operation async def stop_operation(self, root_trace_id, operation_id, idempotency_key=None): self.calls.append(("stop", idempotency_key)) return replace(self.operation, status=OperationStatus.STOPPED) async def resume_operation(self, root_trace_id, operation_id, idempotency_key=None): 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(): task = TaskRecord( task_id="task-1", goal_id=None, parent_task_id=None, display_path="1", specs=[ TaskSpec( version=1, objective="test", acceptance_criteria=[ AcceptanceCriterion( criterion_id="done", description="The task is complete", ) ], ) ], status=TaskStatus.PENDING, ) attempt = TaskAttempt( attempt_id="attempt-1", task_id=task.task_id, spec_version=1, worker_trace_id="worker-1", worker_preset="worker", execution_mode="new", accepted_child_decision_ids=(), ) validation = ValidationReport( validation_id="validation-1", task_id=task.task_id, attempt_id=attempt.attempt_id, spec_version=1, snapshot_id="snapshot-1", validator_trace_id="validator-1", ) operation = BackgroundOperation( operation_id="operation-1", root_trace_id="root", kind=OperationKind.DISPATCH, request={"kind": "dispatch", "task_ids": [task.task_id], "secret": "hidden"}, request_fingerprint="private-fingerprint", task_ids=[task.task_id], ) ledger = TaskLedger( root_trace_id="root", mission="test", root_task_id=task.task_id, revision=3, ) 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", sequence=1, root_trace_id="root", ledger_revision=3, event_type="task_created", occurred_at="2026-07-18T00:00:00+00:00", ) store = StubStore(ledger, event) 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) http = httpx.AsyncClient(transport=transport, base_url="http://test") return coordinator, http @pytest.mark.asyncio async def test_router_start_queries_and_problem_details(api_fixture): coordinator, http = api_fixture response = await http.post( "/api/v2/roots/root/operations", headers={"Idempotency-Key": "dispatch-once"}, json={"kind": "dispatch", "task_ids": ["task-1"]}, ) assert response.status_code == 202 assert response.json()["operation_id"] == "operation-1" assert "request" not in response.json() assert "request_fingerprint" not in response.json() assert coordinator.calls[0][-1]["idempotency_key"] == "dispatch-once" task_list = (await http.get("/api/v2/roots/root/tasks")).json() assert task_list["revision"] == 3 assert task_list["root_task_id"] == "task-1" assert (await http.get("/api/v2/roots/root/tasks/task-1")).json()["task_id"] == "task-1" assert (await http.get("/api/v2/roots/root/attempts")).json()["items"][0]["attempt_id"] == "attempt-1" assert (await http.get("/api/v2/roots/root/attempts/attempt-1")).json()["attempt_id"] == "attempt-1" assert (await http.get("/api/v2/roots/root/validations")).json()["items"][0]["validation_id"] == "validation-1" assert ( await http.get("/api/v2/roots/root/validations/validation-1") ).json()["validation_id"] == "validation-1" assert (await http.get("/api/v2/roots/root/events")).json()["events"][0]["sequence"] == 1 missing = await http.get("/api/v2/roots/root/tasks/missing") assert missing.status_code == 404 assert missing.json()["code"] == "not_found" invalid = await http.post( "/api/v2/roots/root/operations", json={"kind": "unknown"} ) assert invalid.status_code == 422 assert invalid.json()["code"] == "invalid_request" invalid_shape = await http.post( "/api/v2/roots/root/operations", json=["not", "an", "object"] ) assert invalid_shape.status_code == 422 assert invalid_shape.json()["code"] == "invalid_request" missing_body = await http.post("/api/v2/roots/root/operations") assert missing_body.status_code == 422 assert missing_body.json()["code"] == "invalid_request" for path in ( "/api/v2/roots/root/attempts/missing", "/api/v2/roots/root/validations/missing", "/api/v2/roots/missing/tasks", ): response = await http.get(path) assert response.status_code == 404 assert response.json()["code"] == "not_found" for query in ("limit=nope", "limit=0", "wait_seconds=31"): response = await http.get(f"/api/v2/roots/root/events?{query}") assert response.status_code == 422 assert response.json()["code"] == "invalid_request" empty = await http.get( "/api/v2/roots/root/events?cursor=end&wait_seconds=0.001" ) assert empty.status_code == 200 assert empty.json()["events"] == [] bad_cursor = await http.get("/api/v2/roots/root/events?cursor=bad") assert bad_cursor.status_code == 422 for operation_id, status, code in ( ("conflict", 409, "conflict"), ("unavailable", 503, "storage_unavailable"), ): response = await http.get( f"/api/v2/roots/root/operations/{operation_id}" ) assert response.status_code == status assert response.json()["code"] == code for key in ("", "x" * 201, "contains space"): response = await http.post( "/api/v2/roots/root/operations", headers={"Idempotency-Key": key}, json={"kind": "dispatch", "task_ids": ["task-1"]}, ) assert response.status_code == 422 assert response.json()["code"] == "invalid_request" schema = (await http.get("/openapi.json")).json() operation_schema = schema["paths"]["/api/v2/roots/{root_trace_id}/operations"]["post"] assert "discriminator" in operation_schema["requestBody"]["content"]["application/json"]["schema"] assert operation_schema["responses"]["202"]["content"]["application/json"]["schema"] assert operation_schema["responses"]["409"]["content"]["application/json"]["schema"] await http.aclose() @pytest.mark.asyncio async def test_async_client_round_trip_and_watch(api_fixture): _, http = api_fixture client = OrchestrationClient("http://test", client=http) operation = await client.start_revalidation( "root", "task-1", "attempt-1", deadline_at="2026-07-19T00:00:00+00:00", idempotency_key="revalidate-once", ) assert operation.kind == "revalidate" assert (await client.poll("root", "operation-1")).operation_id == "operation-1" assert (await client.stop("root", "operation-1")).status == "stopped" assert (await client.resume("root", "operation-1")).status == "pending" task_list = await client.list_tasks("root") assert task_list.root_task_id == "task-1" assert task_list.items[0].task_id == "task-1" assert (await client.get_task("root", "task-1")).task_id == "task-1" assert (await client.list_attempts("root")).items[0].attempt_id == "attempt-1" assert (await client.get_attempt("root", "attempt-1")).attempt_id == "attempt-1" assert ( await client.list_validations("root") ).items[0].validation_id == "validation-1" assert ( 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" await stream.aclose() 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" assert (await http.get("/api/v2/capabilities")).json() == { "api_version": "v2", "mission_snapshot_schema_version": 1, "event_schema_version": 2, "capabilities": [ "mission_snapshot", "root_completion", "artifact_snapshot", "orchestration_event_cursor", "trace_correlation_filters", ], } await http.aclose() @pytest.mark.asyncio async def test_client_transport_retry_and_non_problem_error(): calls = 0 async def flaky(request): nonlocal calls calls += 1 if calls == 1: raise httpx.ConnectError("temporary", request=request) return httpx.Response(418, text="not problem details") http = httpx.AsyncClient( transport=httpx.MockTransport(flaky), base_url="http://test" ) client = OrchestrationClient("http://test", client=http, max_retries=1) with pytest.raises(OrchestrationAPIError) as caught: await client.poll("root", "operation-1") assert calls == 2 assert caught.value.code == "http_error" assert caught.value.status_code == 418 await http.aclose() @pytest.mark.asyncio async def test_client_retries_mutation_only_with_idempotency_key(): attempts = {"unsafe": 0, "safe": 0} async def handler(request): key = "safe" if request.headers.get("Idempotency-Key") else "unsafe" attempts[key] += 1 if attempts[key] < 2: return httpx.Response( 503, json={ "title": "Unavailable", "status": 503, "detail": "retry", "code": "service_unavailable", }, ) return httpx.Response( 200, json={ "operation_id": "operation-1", "root_trace_id": "root", "kind": "dispatch", "status": "pending", "task_ids": ["task-1"], "attempt_ids": [], "validation_ids": [], "result_ref": None, "execution_epoch": 0, "deadline_at": None, "error": None, "started_at": None, "completed_at": None, "created_at": "now", "updated_at": "now", }, ) http = httpx.AsyncClient( transport=httpx.MockTransport(handler), base_url="http://test" ) client = OrchestrationClient("http://test", client=http, max_retries=2) with pytest.raises(OrchestrationAPIError): await client.start_dispatch("root", ["task-1"]) assert attempts["unsafe"] == 1 result = await client.start_dispatch( "root", ["task-1"], idempotency_key="safe-on-retry" ) assert result.operation_id == "operation-1" assert attempts["safe"] == 2 await http.aclose()