script-build-and-framework-tmp/agent/tests/test_orchestration_v2_api.py

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 (
    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")],
        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", 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"

    assert (await http.get("/api/v2/roots/root/tasks")).json()["revision"] == 3
    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"
    assert (await client.list_tasks("root")).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()