adagio-server 0.1.0__py3-none-any.whl

adagio_server/__init__.py

@@ -0,0 +1,8 @@
+"""Adagio's outbound self-hosted runtime server."""
+
+from importlib.metadata import PackageNotFoundError, version
+
+try:
+    __version__ = version("adagio-server")
+except PackageNotFoundError:  # Source checkout without an installed distribution.
+    __version__ = "unknown"

adagio_server/agent/__init__.py

@@ -0,0 +1 @@
+"""Outbound pull-agent core (register / claim / lease / run / report)."""

adagio_server/agent/artifacts.py

@@ -0,0 +1,118 @@
+"""QIIME ``.qza``/``.qzv`` artifact introspection.
+
+This is the **only** domain-aware code in the agent (design §5). It is kept
+verbatim in behavior from the legacy ``jobs.py`` so that when the agent relays a
+CLI ``output_saved`` event to action as an :class:`ArtifactReport`, it can fill in
+``artifact_type`` (the QIIME semantic type) and ``provenance`` (the recorded
+action graph). All parsing is best-effort and degrades gracefully - a file that
+is not a QIIME archive simply yields ``(None, None)``.
+
+Everything else in the agent treats artifacts as opaque URIs.
+"""
+
+from __future__ import annotations
+
+import zipfile
+from pathlib import Path
+from typing import Any
+
+from fastapi.encoders import jsonable_encoder
+
+try:
+    import yaml
+except Exception:  # noqa: BLE001 - yaml is optional; degrade to no introspection
+    yaml = None  # type: ignore[assignment]
+
+#: QIIME repr tags that a plain YAML loader cannot resolve.
+_Q2_REPRS = ("!ref", "!no-provenance", "!metadata", "!color", "!cite")
+
+
+def read_qiime_archive(file_path: Path) -> tuple[str | None, dict[str, Any] | None]:
+    """Return ``(artifact_type, provenance)`` for a QIIME archive, else ``(None, None)``.
+
+    ``artifact_type`` comes from ``*/metadata.yaml``'s ``type`` field; ``provenance``
+    from ``*/provenance/action/action.yaml``. Non-archives and parse failures return
+    ``(None, None)`` / partial results rather than raising.
+    """
+    if not zipfile.is_zipfile(file_path):
+        return None, None
+
+    try:
+        with zipfile.ZipFile(file_path, "r") as zf:
+            metadata = _read_yaml_member(zf, ["/metadata.yaml"])
+            provenance = _read_yaml_member(
+                zf,
+                ["/provenance/action/action.yaml", "/action/action.yaml"],
+            )
+    except Exception:  # noqa: BLE001
+        return None, None
+
+    archive_type: str | None = None
+    if isinstance(metadata, dict):
+        type_name = metadata.get("type")
+        if isinstance(type_name, str):
+            archive_type = type_name
+
+    if isinstance(provenance, dict):
+        # YAML timestamps become date/datetime objects. Normalize before the
+        # telemetry size validator calls json.dumps, not just at HTTP upload.
+        return archive_type, jsonable_encoder(provenance)
+    return archive_type, None
+
+
+def _read_yaml_member(
+    zf: zipfile.ZipFile, suffixes: list[str]
+) -> dict[str, Any] | None:
+    names = zf.namelist()
+    for suffix in suffixes:
+        match = next((name for name in names if name.endswith(suffix)), None)
+        if match is None:
+            continue
+        text = zf.read(match).decode("utf-8", errors="replace")
+        parsed = yaml_load_with_qiime_tags(text)
+        if isinstance(parsed, dict):
+            return parsed
+    return None
+
+
+def yaml_load_with_qiime_tags(text: str) -> Any:
+    """Load YAML tolerating QIIME repr tags; degrade progressively on failure."""
+    if yaml is None:
+        return None
+
+    class _Loader(yaml.SafeLoader):
+        pass
+
+    def _unknown_tag(loader: yaml.SafeLoader, node: yaml.Node) -> Any:
+        if isinstance(node, yaml.ScalarNode):
+            return loader.construct_scalar(node)
+        if isinstance(node, yaml.SequenceNode):
+            return loader.construct_sequence(node)
+        if isinstance(node, yaml.MappingNode):
+            return loader.construct_mapping(node)
+        return None
+
+    _Loader.add_constructor(None, _unknown_tag)
+    try:
+        return yaml.load(text, Loader=_Loader)
+    except Exception:  # noqa: BLE001
+        pass
+
+    stripped = text
+    for qrepr in _Q2_REPRS:
+        stripped = stripped.replace(qrepr, "")
+
+    try:
+        return yaml.safe_load(stripped)
+    except Exception:  # noqa: BLE001
+        pass
+
+    try:
+        return yaml.load(stripped, Loader=yaml.BaseLoader)
+    except Exception:  # noqa: BLE001
+        return {
+            "_parse_warning": (
+                "Unable to fully parse QIIME provenance YAML. "
+                "Some provenance details were omitted."
+            ),
+        }

adagio_server/agent/client.py

@@ -0,0 +1,305 @@
+"""Typed HTTP client for the ``action`` runner endpoints (design §6).
+
+Every call carries ``Authorization: Bearer {AGENT_TOKEN}`` and is typed with the
+vendored contract models. The agent authenticates **outward** only; nothing here
+ever listens inbound.
+
+Transport: ``httpx`` (sync ``Client``, driven from ``asyncio.to_thread`` in the
+supervisor so the event loop is never blocked). Bounded retries with backoff make
+the client robust to transient network failures; the long-poll ``claim`` uses a
+longer read timeout and treats 204 as "no work".
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from typing import Any
+from uuid import UUID
+
+import httpx
+
+from adagio_server.contracts.agent import (
+    AgentClaimRequest,
+    AgentEventBatch,
+    AgentInfo,
+    AgentRegister,
+    ArtifactReport,
+    ClaimedJob,
+    JobStatusReport,
+    LeaseRenew,
+    LeaseState,
+    LogBatch,
+    LogBatchAccepted,
+    NodeStatusReport,
+    RunnerEnrollRequest,
+    RunnerEnrollResponse,
+)
+
+logger = logging.getLogger("adagio.agent.client")
+
+#: HTTP statuses worth retrying (transient / server-side).
+_RETRYABLE_STATUS = frozenset({408, 425, 429, 500, 502, 503, 504})
+
+
+class ActionClientError(RuntimeError):
+    """Raised when a call to action fails after exhausting retries."""
+
+
+class ActionAuthError(ActionClientError):
+    """Raised immediately when action rejects the agent token (401/403).
+
+    Auth rejection is terminal for the sidecar: retrying can never succeed
+    until a NEW token is minted, and only the desktop shell can trigger that.
+    The supervisor loop propagates this and the process exits with
+    :data:`AGENT_AUTH_EXIT_CODE` so the shell can tell "re-mint the token"
+    apart from a crash (design §7).
+    """
+
+
+class ActionConflictError(ActionClientError):
+    """Raised when Action reports that this agent no longer owns a job."""
+
+
+class ActionCompatibilityError(ActionClientError):
+    """Registration requires a runner update, not a credential refresh."""
+
+
+AGENT_COMPATIBILITY_EXIT_CODE = 42
+
+
+#: Process exit code meaning "action rejected the agent token". This is a
+#: contract with the desktop shell supervisor - mirrored as
+#: AGENT_AUTH_EXIT_CODE in adagio-desktop/src/agent-supervisor.ts, which
+#: reacts by discarding the token and re-entering 'awaiting agent token'
+#: instead of restart-with-backoff.
+AGENT_AUTH_EXIT_CODE = 41
+
+
+class ActionClient:
+    """Thin, typed, retrying client over the ``action`` agent API."""
+
+    def __init__(
+        self,
+        *,
+        base_url: str,
+        token: str | None,
+        verify_tls: bool = True,
+        timeout: float = 15.0,
+        long_poll_timeout: float = 40.0,
+        max_retries: int = 4,
+        backoff_base: float = 0.5,
+        claim_wait: int | None = None,
+    ) -> None:
+        self._base_url = base_url.rstrip("/")
+        self._token = token
+        self._max_retries = max_retries
+        self._backoff_base = backoff_base
+        self._long_poll_timeout = long_poll_timeout
+        self._claim_wait = claim_wait
+        headers = {"Content-Type": "application/json"}
+        if token:
+            headers["Authorization"] = f"Bearer {token}"
+        self._client = httpx.Client(
+            headers=headers,
+            verify=verify_tls,
+            timeout=timeout,
+        )
+
+    @property
+    def claim_wait(self) -> int | None:
+        """Server-side hold (seconds) requested on ``claim`` (``?wait=N``)."""
+        return self._claim_wait
+
+    @claim_wait.setter
+    def claim_wait(self, value: int | None) -> None:
+        self._claim_wait = value
+
+    def close(self) -> None:
+        """Close the underlying connection pool."""
+        self._client.close()
+
+    def __enter__(self) -> ActionClient:
+        return self
+
+    def __exit__(self, *_exc: object) -> None:
+        self.close()
+
+    # -- internal request helper --------------------------------------------------
+    def _request(
+        self,
+        method: str,
+        path: str,
+        *,
+        json: Any | None = None,
+        params: dict[str, Any] | None = None,
+        timeout: float | None = None,
+        allow_204: bool = False,
+    ) -> httpx.Response:
+        url = f"{self._base_url}{path}"
+        last_exc: Exception | None = None
+        for attempt in range(self._max_retries):
+            try:
+                response = self._client.request(
+                    method, url, json=json, params=params, timeout=timeout
+                )
+            except httpx.HTTPError as exc:  # network / timeout
+                last_exc = exc
+                logger.warning(
+                    "action call %s %s failed (attempt %d/%d): %r",
+                    method,
+                    path,
+                    attempt + 1,
+                    self._max_retries,
+                    exc,
+                )
+                self._sleep_backoff(attempt)
+                continue
+
+            if response.status_code in _RETRYABLE_STATUS:
+                last_exc = ActionClientError(
+                    f"{method} {path} -> {response.status_code}"
+                )
+                logger.warning(
+                    "action call %s %s returned %d (attempt %d/%d)",
+                    method,
+                    path,
+                    response.status_code,
+                    attempt + 1,
+                    self._max_retries,
+                )
+                self._sleep_backoff(attempt)
+                continue
+
+            if response.status_code in (401, 403):
+                raise ActionAuthError(
+                    f"{method} {path} -> {response.status_code}: agent token rejected"
+                )
+            if response.status_code == 409:
+                if path == "/agent/register":
+                    raise ActionCompatibilityError(
+                        "Runtime server must be updated to match this Adagio deployment"
+                    )
+                raise ActionConflictError(
+                    f"{method} {path} -> 409: job ownership lost"
+                )
+            if response.status_code >= 400:
+                raise ActionClientError(
+                    f"{method} {path} -> {response.status_code}: {response.text[:500]}"
+                )
+            return response
+
+        raise ActionClientError(
+            f"{method} {path} failed after {self._max_retries} attempts"
+        ) from last_exc
+
+    def _sleep_backoff(self, attempt: int) -> None:
+        time.sleep(self._backoff_base * (2**attempt))
+
+    # -- lifecycle & work pull ----------------------------------------------------
+    def enroll(self, request: RunnerEnrollRequest) -> RunnerEnrollResponse:
+        """Exchange a one-time enrollment token for the persistent credential."""
+        resp = self._request(
+            "POST",
+            "/agent/enroll",
+            json=request.model_dump(mode="json"),
+        )
+        return RunnerEnrollResponse.model_validate(resp.json())
+
+    def register(self, register: AgentRegister) -> AgentInfo:
+        """POST /agent/register -> target-scoped runner identity and timing."""
+        resp = self._request(
+            "POST", "/agent/register", json=register.model_dump(mode="json")
+        )
+        return AgentInfo.model_validate(resp.json())
+
+    def claim(self, claim: AgentClaimRequest) -> ClaimedJob | None:
+        """POST /agent/claim (long-poll). Returns a job or ``None`` on 204.
+
+        When ``claim_wait`` is set (the supervisor copies action's
+        ``claim_long_poll_seconds`` here after register), the ``?wait=N`` query
+        parameter asks action to hold the request until work arrives or the
+        window elapses; the read timeout is stretched to cover the hold.
+        """
+        params: dict[str, Any] | None = None
+        timeout = self._long_poll_timeout
+        if self._claim_wait:
+            params = {"wait": self._claim_wait}
+            timeout = max(self._long_poll_timeout, self._claim_wait + 15.0)
+        resp = self._request(
+            "POST",
+            "/agent/claim",
+            json=claim.model_dump(mode="json"),
+            params=params,
+            timeout=timeout,
+            allow_204=True,
+        )
+        if resp.status_code == 204 or not resp.content:
+            return None
+        return ClaimedJob.model_validate(resp.json())
+
+    def renew_lease(self, job_id: UUID, renew: LeaseRenew) -> LeaseState:
+        """POST /agent/jobs/{job_id}/lease -> LeaseState (with cancel_requested)."""
+        resp = self._request(
+            "POST",
+            f"/agent/jobs/{job_id}/lease",
+            json=renew.model_dump(mode="json"),
+        )
+        return LeaseState.model_validate(resp.json())
+
+    def get_job_control(self, job_id: UUID) -> LeaseState:
+        """GET active-job flags without renewing the lease."""
+        resp = self._request("GET", f"/agent/jobs/{job_id}/control")
+        return LeaseState.model_validate(resp.json())
+
+    # -- reporting ----------------------------------------------------------------
+    def report_job_status(self, job_id: UUID, report: JobStatusReport) -> None:
+        """POST /agent/jobs/{job_id}/status (terminal or transitional job status)."""
+        self._request(
+            "POST",
+            f"/agent/jobs/{job_id}/status",
+            json=report.model_dump(mode="json"),
+            allow_204=True,
+        )
+
+    def report_node_status(
+        self, job_id: UUID, report: NodeStatusReport
+    ) -> None:
+        """POST /agent/jobs/{job_id}/nodes/{node_id}/status."""
+        self._request(
+            "POST",
+            f"/agent/jobs/{job_id}/nodes/{report.node_id}/status",
+            json=report.model_dump(mode="json"),
+            allow_204=True,
+        )
+
+    def post_events(self, job_id: UUID, batch: AgentEventBatch) -> None:
+        """POST /agent/jobs/{job_id}/events (batched AgentEvents)."""
+        if not batch.events:
+            return
+        self._request(
+            "POST",
+            f"/agent/jobs/{job_id}/events",
+            json=batch.model_dump(mode="json"),
+            allow_204=True,
+        )
+
+    def post_logs(
+        self, job_id: UUID, node_id: UUID, batch: LogBatch
+    ) -> LogBatchAccepted:
+        """POST /agent/jobs/{job_id}/nodes/{node_id}/logs -> {accepted, max_seq}."""
+        resp = self._request(
+            "POST",
+            f"/agent/jobs/{job_id}/nodes/{node_id}/logs",
+            json=batch.model_dump(mode="json"),
+        )
+        return LogBatchAccepted.model_validate(resp.json())
+
+    def report_artifact(self, job_id: UUID, report: ArtifactReport) -> None:
+        """POST /agent/jobs/{job_id}/artifacts."""
+        self._request(
+            "POST",
+            f"/agent/jobs/{job_id}/artifacts",
+            json=report.model_dump(mode="json"),
+            allow_204=True,
+        )

adagio_server/agent/identity.py

@@ -0,0 +1,10 @@
+"""Stable host identity used by Desktop and self-hosted server enrollment."""
+
+import platform
+import socket
+
+
+def host_fingerprint() -> str:
+    """Return the existing deterministic installation fingerprint."""
+    return f"{socket.gethostname()}:{platform.node()}:{platform.machine()}"
+

adagio_server/agent/queues.py

@@ -0,0 +1,52 @@
+"""In-process fan-out of CLI loopback events to the supervisor.
+
+The loopback FastAPI router (``adagio_server/api/endpoints/loopback.py``) receives the
+CLI's ``--connected`` events and enqueues them onto a per-job
+:class:`asyncio.Queue`. The supervisor's relay task drains that queue, enriches
+each event to the action contract, and posts it outward.
+
+Keeping this as a module-level registry (rather than app state) lets both the
+router and the supervisor - which live in the same process/event loop - share it
+without threading FastAPI dependencies through the agent core.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+from uuid import UUID
+
+
+class LoopbackRegistry:
+    """Registry of per-job event queues, keyed by ``job_id``."""
+
+    def __init__(self) -> None:
+        self._queues: dict[UUID, asyncio.Queue[dict[str, Any]]] = {}
+
+    def register(self, job_id: UUID) -> asyncio.Queue[dict[str, Any]]:
+        """Create (or return) the queue for ``job_id``."""
+        queue = self._queues.get(job_id)
+        if queue is None:
+            queue = asyncio.Queue()
+            self._queues[job_id] = queue
+        return queue
+
+    def get(self, job_id: UUID) -> asyncio.Queue[dict[str, Any]] | None:
+        """Return the queue for ``job_id`` if one is registered."""
+        return self._queues.get(job_id)
+
+    def unregister(self, job_id: UUID) -> None:
+        """Drop the queue for ``job_id`` (after the job completes)."""
+        self._queues.pop(job_id, None)
+
+    async def publish(self, job_id: UUID, event: dict[str, Any]) -> bool:
+        """Enqueue ``event`` for ``job_id``. Returns ``False`` if no queue exists."""
+        queue = self._queues.get(job_id)
+        if queue is None:
+            return False
+        await queue.put(event)
+        return True
+
+
+#: Process-global registry shared by the loopback router and the supervisor.
+registry = LoopbackRegistry()