#!/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())