|
|
@@ -0,0 +1,185 @@
|
|
|
+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)]
|