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 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:
                submission = self.scm.submit(assignment, request)
                pull = str(submission.get("pull_request") or "")
                if pull and pull not in structured["artifacts"]:
                    structured["artifacts"].append(pull)
            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())