atlas-iac/services/hermes/scripts/execution_pool_client.py

#!/usr/bin/env python3
"""Ordinal-local signing and SCM mediation boundary for a model worker."""

from __future__ import annotations

import json
import logging
import os
import threading
import urllib.error
import urllib.request
from http.server import BaseHTTPRequestHandler
from pathlib import Path
from collections.abc import Callable
from typing import Any

import cli_lane_goal
from execution_pool_protocol import (
    MAX_WIRE_BYTES,
    PROTOCOL_VERSION,
    BoundedHTTPServer,
    ProtocolError,
    canonical_json,
    parse_wire,
    read_key,
    sign_envelope,
    verify_envelope,
)
from execution_pool_scm import Boundary as SCMBoundary


KEY_PATH = Path(
    os.environ.get("HERMES_EXECUTION_POOL_KEY_FILE", "/pool-access/execution-pool-key")
)
COORDINATOR = os.environ.get(
    "HERMES_EXECUTION_POOL_URL",
    "http://hermes-execution-pool.hermes.svc.cluster.local:9007",
).rstrip("/")
ORDINAL = int(os.environ.get("HERMES_WORKER_ORDINAL", "-1"))
PORT = int(os.environ.get("HERMES_EXECUTION_CLIENT_PORT", "9009"))
RESULT_FIELDS = frozenset(
    {"status", "summary", "changed_files", "tests_run", "artifacts", "findings", "blockers"}
)
MEDIATOR_TERMINAL_FIELDS = frozenset({
    "scm_submission", "scm_resume", "publication_retry_transient",
    "publication_retry_lease_fence", "lease_fence",
})
LOG = logging.getLogger(__name__)
COORDINATOR_REJECTION_CATEGORIES = {
    "assignment": (
        "assignment is unknown or stale", "worker ordinal does not own this assignment",
        "assignment attempt is stale",
    ),
    "lease": ("assignment lease expired", "Kanban run no longer owns this worker"),
    "state": ("assignment is no longer running", "assignment cannot accept a result"),
    "delivery": ("delivery identifier was reused", "conflicting result for completed delivery"),
}


class CoordinatorRejected(ProtocolError):
    """A coordinator conflict reduced to one non-sensitive fixed category."""

    def __init__(self, category: str):
        self.category = category
        super().__init__(f"coordinator rejected request: {category}")


def _coordinator_rejection(error: urllib.error.HTTPError) -> CoordinatorRejected:
    """Read one bounded coordinator conflict body and retain no free-form detail."""
    try:
        raw = error.read(MAX_WIRE_BYTES + 1)
        value = json.loads(raw) if len(raw) <= MAX_WIRE_BYTES else {}
        detail = value.get("error") if isinstance(value, dict) else None
    except (OSError, UnicodeDecodeError, json.JSONDecodeError):
        detail = None
    for category, messages in COORDINATOR_REJECTION_CATEGORIES.items():
        if detail in messages:
            return CoordinatorRejected(category)
    return CoordinatorRejected("policy")


def _rejection_category(operation: str, error: Exception) -> str:
    """Classify local mediator rejections without logging request or error text."""
    if isinstance(error, CoordinatorRejected):
        return f"coordinator-{error.category}"
    if isinstance(error, urllib.error.URLError):
        return "transport"
    if isinstance(error, OSError):
        return "transport"
    if not isinstance(error, ProtocolError):
        return "input"
    if operation != "resume":
        return "policy"
    message = str(error)
    if message.startswith("SCM resume source") or message.startswith("SCM resume artifact"):
        return "source"
    if message.startswith("SCM resume evidence") or message.startswith("SCM resume metadata"):
        return "evidence"
    if message.startswith("preserved SCM workspace") or message.startswith("preserved SCM baseline"):
        return "workspace"
    if "lease" in message:
        return "lease"
    if message.startswith("SCM resume publication"):
        return "scm"
    return "policy"


def _binding(value: dict[str, Any]) -> dict[str, Any]:
    return {
        name: value[name]
        for name in ("board", "task_id", "run_id", "worker_ordinal", "attempt")
    }


def _validate_result(payload: Any) -> dict[str, Any]:
    if not isinstance(payload, dict):
        raise ProtocolError("terminal result payload must be an object")
    structured = payload.get("structured")
    if not isinstance(structured, dict) or set(structured) != RESULT_FIELDS:
        raise ProtocolError("terminal result fields do not match the reviewed schema")
    if structured.get("status") not in cli_lane_goal.RESULT_STATUSES:
        raise ProtocolError("terminal result status is invalid")
    if not isinstance(structured.get("summary"), str) or not structured["summary"].strip():
        raise ProtocolError("terminal result summary is required")
    for name in RESULT_FIELDS - {"status", "summary"}:
        value = structured.get(name)
        if not isinstance(value, list) or any(not isinstance(item, str) for item in value):
            raise ProtocolError(f"terminal result {name} must be a text list")
    return dict(payload)


def _worker_terminal(payload: dict[str, Any]) -> dict[str, Any]:
    """Copy ordinary worker output while removing mediator-only authority fields."""
    return {name: value for name, value in payload.items() if name not in MEDIATOR_TERMINAL_FIELDS}


class ClientBoundary:
    """Keep HMAC and SCM authority outside the model-facing container."""

    def __init__(self, key: bytes, scm: SCMBoundary | None = None):
        self.key = key
        self.scm = scm or SCMBoundary(key)
        self.current: dict[str, Any] | None = None
        self.resumed: dict[str, str] | None = None
        self.resumed_evidence: dict[str, Any] | None = None
        self.resume_transient = False
        self.lock = threading.RLock()

    def _post(self, path: str, envelope: dict[str, Any]) -> dict[str, Any]:
        request = urllib.request.Request(
            COORDINATOR + path,
            data=canonical_json(envelope),
            method="POST",
            headers={"Content-Type": "application/json", "Cache-Control": "no-store"},
        )
        try:
            with urllib.request.urlopen(request, timeout=60) as response:
                body = response.read(MAX_WIRE_BYTES + 1)
        except urllib.error.HTTPError as error:
            if error.code == 409:
                raise _coordinator_rejection(error) from error
            raise
        if len(body) > MAX_WIRE_BYTES:
            raise ProtocolError("coordinator response exceeds the wire limit")
        return verify_envelope(self.key, json.loads(body))

    def poll(self) -> dict[str, Any]:
        poll_binding = {
            "board": "",
            "task_id": "",
            "run_id": "",
            "worker_ordinal": ORDINAL,
            "attempt": 0,
        }
        response = self._post(
            "/v1/poll",
            sign_envelope(self.key, "poll", poll_binding, {"ready": True}),
        )
        with self.lock:
            if response["kind"] == "ack":
                self.current = None
                return {"assignment": None}
            if response["kind"] != "assignment" or response["worker_ordinal"] != ORDINAL:
                raise ProtocolError("coordinator returned a foreign assignment")
            publication_only = isinstance(response.get("payload", {}).get("scm_resume"), dict)
            checkout = {"workspace": "", "baseline_sha": ""} if publication_only else self.scm.checkout(response)
            self.current = response
            self.resumed = None
            self.resumed_evidence = None
            self.resume_transient = False
            assignment = {
                **_binding(response),
                "payload": response["payload"],
                "workspace": checkout["workspace"],
                "baseline_sha": checkout["baseline_sha"],
                "protocol_version": PROTOCOL_VERSION,
            }
        return {"assignment": assignment}

    def _current_for(self, supplied: Any) -> tuple[dict[str, Any], dict[str, Any]]:
        if not isinstance(supplied, dict):
            raise ProtocolError("local request binding must be an object")
        if self.current is None or supplied != _binding(self.current):
            raise ProtocolError("local request does not own the current assignment")
        return self.current, dict(supplied)

    def _publication_lease(self, binding: dict[str, Any]) -> tuple[Callable[[], None], Callable[[], None]]:
        """Keep a resume lease alive outside the SCM lock and fail before another step."""
        stop, failures = threading.Event(), []

        def renew() -> None:
            try:
                response = self._post("/v1/heartbeat", sign_envelope(self.key, "heartbeat", binding, {"note": "SCM publication in progress"}))
                if response["kind"] != "ack" or _binding(response) != binding or not response["payload"].get("accepted"):
                    raise ProtocolError("coordinator rejected SCM publication lease")
            except (OSError, ProtocolError, ValueError, urllib.error.URLError) as error:
                failures.append(error)

        renew()
        def keepalive() -> None:
            while not stop.wait(20):
                renew()

        thread = threading.Thread(target=keepalive, daemon=True)
        thread.start()

        def checkpoint() -> None:
            if failures:
                error = failures[0]
                if isinstance(error, CoordinatorRejected):
                    raise error
                raise ProtocolError("SCM publication lease was lost") from error

        return checkpoint, lambda: (stop.set(), thread.join(timeout=1))

    def heartbeat(self, request: dict[str, Any]) -> dict[str, Any]:
        payload = request.get("payload")
        if not isinstance(payload, dict):
            raise ProtocolError("heartbeat payload must be an object")
        with self.lock:
            _assignment, binding = self._current_for(request.get("binding"))
            response = self._post(
                "/v1/heartbeat",
                sign_envelope(self.key, "heartbeat", binding, payload),
            )
            if response["kind"] != "ack" or _binding(response) != binding:
                raise ProtocolError("coordinator acknowledgement binding changed")
        return {"ack": response["payload"]}

    def _submit(
        self, assignment: dict[str, Any], request: dict[str, Any],
        structured: dict[str, Any],
    ) -> dict[str, str] | None:
        """Publish the run's work, downgrading the result rather than losing it.

        A refused push or draft used to unwind the whole worker, so the run's
        actual output was discarded and the task was blocked with only the SCM
        error. The commits live on the durable workspace volume, so the result is
        still reported here -- as blocked, carrying the reason -- which keeps the
        prior work recoverable and reviewable by a human.
        """
        try:
            if isinstance(self.scm, SCMBoundary):
                checkpoint, close = self._publication_lease(_binding(assignment))
                try:
                    submission = self.scm.submit(assignment, request, checkpoint=checkpoint)
                    checkpoint()
                finally:
                    close()
            else:
                submission = self.scm.submit(assignment, request)
        except (OSError, ProtocolError, RuntimeError, ValueError) as error:
            payload = assignment.get("payload")
            resume = None
            if isinstance(payload, dict) and payload.get("continuation_kind") == "repair":
                try:
                    resume = self.scm.resume_artifact(assignment, request, structured)
                except (OSError, ProtocolError, RuntimeError, ValueError):
                    resume = None
            structured["status"] = "blocked"
            reason = (
                "Distributed SCM submission failed; the commits remain on this "
                f"ordinal's workspace: {type(error).__name__}: {error}"
            )
            if reason not in structured["blockers"]:
                structured["blockers"].append(reason)
            return {"resume": resume} if resume is not None else None
        pull = str(submission.get("pull_request") or "")
        branch = str(submission.get("branch") or "")
        for artifact in (pull, f"branch:{branch}" if branch else ""):
            if artifact and artifact not in structured["artifacts"]:
                structured["artifacts"].append(artifact)
        head = str(submission.get("head") or "")
        return {"branch": branch, "pull_request": pull, "head": head}

    def resume(self, request: dict[str, Any]) -> dict[str, Any]:
        """Publish the coordinator-carried artifact without exposing it to the worker."""
        with self.lock:
            assignment, _binding = self._current_for(request.get("binding"))
            artifact = assignment.get("payload", {}).get("scm_resume")
            if not isinstance(artifact, dict):
                raise ProtocolError("assignment has no SCM resume artifact")
            evidence = _validate_result({"structured": artifact.get("structured")})["structured"]
            if evidence["status"] != "completed":
                raise ProtocolError("SCM resume evidence is not completed")
            checkpoint, close = self._publication_lease(_binding)
            try:
                self.resume_transient = False
                try:
                    self.resumed = self.scm.resume(assignment, artifact, checkpoint=checkpoint)
                    checkpoint()
                except (OSError, RuntimeError, urllib.error.URLError) as error:
                    detail = str(error).lower()
                    if any(marker in detail for marker in ("timed out", "connection", "could not resolve", "http 502", "http 503", "http 504")):
                        self.resume_transient = True
                        return {"publication_retry_transient": True}
                    raise ProtocolError("SCM resume publication was rejected") from error
            finally:
                close()
            self.resumed_evidence = json.loads(canonical_json(evidence))
            return {"scm_submission": self.resumed}

    def finish(self, request: dict[str, Any]) -> dict[str, Any]:
        payload = _worker_terminal(_validate_result(request.get("payload")))
        # SCM evidence and retry control are private mediator outputs. A worker
        # may report ordinary terminal fields but cannot carry authority forward.
        with self.lock:
            assignment, binding = self._current_for(request.get("binding"))
            structured = payload["structured"]
            if self.resume_transient:
                payload["publication_retry_transient"] = True
            if structured["status"] == "completed" and int(payload.get("returncode", 1)) == 0:
                if self.resumed is not None:
                    if self.resumed_evidence is None:
                        raise ProtocolError("SCM resume evidence is unavailable")
                    structured = json.loads(canonical_json(self.resumed_evidence))
                    payload["structured"] = structured
                    payload["returncode"] = 0
                    payload["scm_submission"] = self.resumed
                else:
                    submission = self._submit(assignment, request, structured)
                    if submission is not None and "resume" not in submission:
                        # The mediator, not the model-facing request, records the
                        # broker-confirmed PR/branch in the signed terminal wire.
                        payload["scm_submission"] = submission
                    elif submission is not None:
                        # The exact clean commit remains on the durable worker volume.
                        # A broker refusal is execution infrastructure, not a provider
                        # capability verdict, so let the coordinator retry it after the
                        # scoped SCM condition is corrected.
                        payload["capacity_failure"] = True
                        payload["scm_resume"] = submission["resume"]
                    else:
                        payload["capacity_failure"] = True
            response = self._post(
                "/v1/result", sign_envelope(self.key, "result", binding, payload)
            )
            if response["kind"] != "ack" or _binding(response) != binding:
                raise ProtocolError("coordinator acknowledgement binding changed")
            if response["payload"].get("accepted"):
                self.current = None
        return {"ack": response["payload"], "structured": structured}


def handler_factory(boundary: ClientBoundary) -> type[BaseHTTPRequestHandler]:
    class Handler(BaseHTTPRequestHandler):
        server_version = f"hermes-execution-mediator/{PROTOCOL_VERSION}"

        def _reply(self, status: int, value: dict[str, Any]) -> None:
            body = canonical_json(value)
            self.send_response(status)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(body)))
            self.send_header("Cache-Control", "no-store")
            self.end_headers()
            self.wfile.write(body)

        def do_GET(self) -> None:  # noqa: N802
            value = {"ready": True, "protocol_version": PROTOCOL_VERSION}
            self._reply(200, value) if self.path == "/ready" else self._reply(
                404, {"error": "not found"}
            )

        def do_POST(self) -> None:  # noqa: N802
            operation = "unknown"
            try:
                length = int(self.headers.get("Content-Length", "0"))
                request = (
                    parse_wire(self.rfile.read(length))
                    if 0 < length <= MAX_WIRE_BYTES
                    else None
                )
                allowed = {"operation", "binding", "payload", "title", "body"}
                if not request or set(request) - allowed:
                    raise ProtocolError("invalid local mediator request")
                requested = str(request.get("operation") or "")
                routes = {
                    "poll": boundary.poll,
                    "heartbeat": lambda: boundary.heartbeat(request),
                    "resume": lambda: boundary.resume(request),
                    "finish": lambda: boundary.finish(request),
                }
                if requested not in routes:
                    raise ProtocolError("unsupported local mediator operation")
                operation = requested
                self._reply(200, routes[operation]())
            except (ProtocolError, OSError, ValueError, urllib.error.URLError) as error:
                LOG.warning(
                    "mediator_rejected operation=%s category=%s",
                    operation,
                    _rejection_category(operation, error),
                )
                self._reply(409, {"error": str(error)[:2000]})

        def log_message(self, _format: str, *_arguments: Any) -> None:
            return

    return Handler


def main() -> int:
    if ORDINAL not in range(3):
        raise SystemExit("HERMES_WORKER_ORDINAL must be 0, 1, or 2")
    key = read_key(KEY_PATH)
    BoundedHTTPServer(
        ("0.0.0.0", PORT),
        handler_factory(ClientBoundary(key)),
        max_workers=4,
    ).serve_forever()
    return 0


if __name__ == "__main__":  # pragma: no cover - exercised by the container entrypoint
    raise SystemExit(main())