from __future__ import annotations import asyncio import os from collections.abc import Mapping from typing import Any import httpx _TERMINAL_TRACE_STATUSES = {"completed", "failed", "stopped", "cancelled"} class HostApiError(RuntimeError): def __init__(self, status_code: int, message: str) -> None: super().__init__(message) self.status_code = status_code class HostClient: def __init__( self, base_url: str | None = None, *, transport: httpx.AsyncBaseTransport | None = None, ) -> None: resolved = base_url or os.getenv("SCRIPT_BUILD_HOST_API_BASE") or "http://127.0.0.1:8080" self._client = httpx.AsyncClient( base_url=resolved.rstrip("/"), timeout=20.0, transport=transport, ) self._inputs: dict[int, dict[str, Any]] = {} self._contracts_by_spec: dict[tuple[int, str, int], dict[str, Any]] = {} self._events: dict[int, list[dict[str, Any]]] = {} self._event_cursors: dict[int, str] = {} self._messages: dict[tuple[int, str], list[dict[str, Any]]] = {} async def aclose(self) -> None: await self._client.aclose() async def list_runs(self, headers: Mapping[str, str]) -> dict[str, Any]: return await self._get("/api/pattern/script_builds?page=1&page_size=100", headers) async def journey_source( self, script_build_id: int, headers: Mapping[str, str] ) -> dict[str, Any]: prefix = f"/api/pattern/script_builds/{script_build_id}" snapshot, traces = await asyncio.gather( self._get(f"{prefix}/mission", headers), self._get(f"{prefix}/traces", headers), ) input_summary, events = await asyncio.gather( self._input_summary(script_build_id, prefix, headers), self._incremental_events(script_build_id, f"{prefix}/mission/events", headers), ) tasks = snapshot.get("tasks") or [] contract_tasks = [ task for task in tasks if any( str(ref).startswith("script-build://task-contracts/sha256/") for ref in (task.get("current_spec") or {}).get("context_refs") or [] ) ] trace_values = traces.get("traces") or [] contracts, messages = await asyncio.gather( self._contracts(script_build_id, prefix, contract_tasks, headers), self._trace_messages(script_build_id, prefix, trace_values, headers), ) return { "snapshot": snapshot, "input_summary": input_summary, "events": events, "messages": messages, "contracts": contracts, } async def artifact( self, script_build_id: int, artifact_version_id: int, headers: Mapping[str, str] ) -> dict[str, Any]: return await self._get( f"/api/pattern/script_builds/{script_build_id}/artifacts/{artifact_version_id}", headers, ) async def _input_summary( self, script_build_id: int, prefix: str, headers: Mapping[str, str] ) -> dict[str, Any]: cached = self._inputs.get(script_build_id) if cached is None: cached = await self._get(f"{prefix}/input-summary", headers) self._inputs[script_build_id] = cached return cached async def _contracts( self, script_build_id: int, prefix: str, tasks: list[dict[str, Any]], headers: Mapping[str, str], ) -> dict[str, dict[str, Any]]: async def read(task: dict[str, Any]) -> dict[str, Any]: task_id = str(task["task_id"]) version = int((task.get("current_spec") or {}).get("version") or 0) key = (script_build_id, task_id, version) cached = self._contracts_by_spec.get(key) if cached is None: cached = await self._get(f"{prefix}/tasks/{task_id}/contract", headers) self._contracts_by_spec[key] = cached return cached values = await asyncio.gather(*(read(task) for task in tasks)) return {str(value["task_id"]): value for value in values} async def _trace_messages( self, script_build_id: int, prefix: str, traces: list[dict[str, Any]], headers: Mapping[str, str], ) -> dict[str, list[dict[str, Any]]]: async def read(trace: dict[str, Any]) -> tuple[str, list[dict[str, Any]]]: trace_id = str(trace["trace_id"]) key = (script_build_id, trace_id) cached = self._messages.get(key, []) status = str(trace.get("status") or "").lower() if cached and status in _TERMINAL_TRACE_STATUSES: return trace_id, cached after = max((int(item.get("sequence") or 0) for item in cached), default=0) page = await self._get( f"{prefix}/traces/{trace_id}/messages?after_sequence={after}", headers ) merged = _merge_by_sequence(cached, list(page.get("messages") or [])) self._messages[key] = merged return trace_id, merged values = await asyncio.gather(*(read(trace) for trace in traces if trace.get("trace_id"))) return dict(values) async def _incremental_events( self, script_build_id: int, path: str, headers: Mapping[str, str] ) -> list[dict[str, Any]]: events = list(self._events.get(script_build_id, [])) cursor = self._event_cursors.get(script_build_id) for _ in range(100): query = "?limit=500" if cursor: query += f"&after={cursor}" page = await self._get(path + query, headers) additions = list(page.get("events") or []) events.extend(additions) next_cursor = page.get("next_cursor") if next_cursor: cursor = str(next_cursor) self._event_cursors[script_build_id] = cursor if len(additions) < 500: break self._events[script_build_id] = events return events async def _get(self, path: str, headers: Mapping[str, str]) -> dict[str, Any]: forwarded = { key: value for key, value in headers.items() if key.lower() in {"authorization", "cookie", "x-request-id"} } response = await self._client.get(path, headers=forwarded) if not response.is_success: try: payload = response.json() message = str(payload.get("message") or payload.get("detail") or response.text) except ValueError: message = response.text raise HostApiError(response.status_code, message or "Host request failed") value = response.json() if not isinstance(value, dict): raise HostApiError(502, "Host returned a non-object response") return value def _merge_by_sequence( current: list[dict[str, Any]], additions: list[dict[str, Any]] ) -> list[dict[str, Any]]: merged = {int(item.get("sequence") or 0): item for item in [*current, *additions]} return [merged[key] for key in sorted(merged)]