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 os
import threading
import urllib.error
import urllib.request
from http.server import BaseHTTPRequestHandler
from pathlib import Path
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"}
)


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 payload


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.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"},
        )
        with urllib.request.urlopen(request, timeout=60) as response:
            body = response.read(MAX_WIRE_BYTES + 1)
        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")
            checkout = self.scm.checkout(response)
            self.current = response
            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 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],
    ) -> 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:
            submission = self.scm.submit(assignment, request)
        except (OSError, ProtocolError, RuntimeError, ValueError) as error:
            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
        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)

    def finish(self, request: dict[str, Any]) -> dict[str, Any]:
        payload = _validate_result(request.get("payload"))
        with self.lock:
            assignment, binding = self._current_for(request.get("binding"))
            structured = payload["structured"]
            if structured["status"] == "completed" and int(payload.get("returncode", 1)) == 0:
                self._submit(assignment, request, structured)
            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
            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")
                operation = str(request.get("operation") or "")
                routes = {
                    "poll": boundary.poll,
                    "heartbeat": lambda: boundary.heartbeat(request),
                    "finish": lambda: boundary.finish(request),
                }
                if operation not in routes:
                    raise ProtocolError("unsupported local mediator operation")
                self._reply(200, routes[operation]())
            except (ProtocolError, OSError, ValueError, urllib.error.URLError) as 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())