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