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, BackgroundOperation, EventPage, OperationKind, OperationStatus, OrchestrationEvent, 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 async def load(self, root_trace_id): 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): self.task_store = 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) @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", ) 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 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) coordinator = StubCoordinator(store, operation) 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" 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_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()