"""Entrypoint for both supervised roles. Python >=3.11, Unix control socket.""" from __future__ import annotations # Source-tree launch works with Python -I -B; no inherited PYTHONPATH required. if __package__ in {None, ""}: import sys from pathlib import Path sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) __package__ = "klbr_runtime" import argparse import asyncio import contextlib import json import sys import traceback from pathlib import Path from typing import Any from klbr.hooks import HookContext from . import skill_index from .context import Bridge, RoutedText from .kernel import Kernel from .policy import Policy from .wire import ( MAX_CODE, BindingToken, ProtocolError, RELAY_PROTOCOL_VERSION, VERSION, Wire, fields, identifier, measured_runtime_identity, text, ) async def serve(args: argparse.Namespace) -> None: # Read the host's skill payload before anything else: the file lives beside the # control socket in a directory the host owns, and the host drops that directory # once this process has connected. Leaving the read until a cell asks would race # the very handshake that makes the directory disposable. if args.skills: skill_index.install(args.skills) if args.behavior: behavior_path = str(Path(args.behavior).resolve()) if behavior_path not in sys.path: sys.path.insert(0, behavior_path) policy = Policy(Path(args.behavior), args.revision) if args.role == "hooks" else None reader, writer = await asyncio.open_unix_connection(args.socket) # The --target value is the host-issued launch binding: a binding token # this process echoes so the host can refuse a foreign or stale worker. # It is NOT measured identity; the runtime identity below is measured # on this machine by the endpoint itself. binding = BindingToken.parse(json.loads(args.target)) if args.target else None wire = Wire(reader, writer, args.session, args.generation, target=binding.as_stamp() if binding else None) bridge, kernel, dispatcher = Bridge(wire), None, None if args.role == "workbench": # The tree vocabulary is bound before the first cell runs, so ordinary work does not # begin with an import line the model has to remember. The package stays importable for # anything it does not pre-bind. from klbr import workspace as tree kernel = Kernel(bridge, prelude=tree.prelude()) elif args.role == "remote": # An auxiliary remote namespace: a persistent target-side Python heap that is # deliberately NOT the workbench. It shares no globals with a local cell, and # the dispatcher is the only place that knows what an operation name means. from klbr_runtime.remote_dispatch import RemoteDispatchError, RemoteDispatcher dispatcher = RemoteDispatcher(workspace=args.workspace) hello: dict[str, Any] = {"kind": "hello", "token": args.token, "role": args.role, "revision": args.revision, "slots": sorted(policy.handlers) if policy else [], # Echoed, not assumed: the host records what this generation # actually loaded rather than what the roots hold now. "skills": skill_index.revision() if skill_index.loaded() else None} if binding is not None: runtime = measured_runtime_identity() if dispatcher is not None: # Capabilities belong INSIDE the measured runtime identity, because that is # the endpoint's own probe of this machine. The host reads them from # identity.runtime.capabilities to decide admission, so an advertised-but- # unusable capability is refused before a call rather than at use. runtime["capabilities"] = list(dispatcher.capabilities) hello["identity"] = { "protocol_version": VERSION, "relay_protocol_version": RELAY_PROTOCOL_VERSION, "target": {"stamp": binding.as_stamp()}, "runtime": runtime, } await wire.send(hello) original_out, original_err = sys.stdout, sys.stderr sys.stdout, sys.stderr = RoutedText("stdout", original_out), RoutedText("stderr", original_err) active: asyncio.Task | None = None active_id: str | None = None async def run(message: dict) -> None: try: if message["kind"] == "execute": assert kernel is not None outcome = await kernel.execute( message["id"], message["code"], purpose=message.get("purpose", "interactive"), ) elif message["kind"] == "remote_dispatch": assert dispatcher is not None # The runner frames, bounds and identifies; it never interprets the # operation name. Dispatch owns the whole vocabulary. try: value = await dispatcher.dispatch( message["operation"], message["payload"], operation_id=message["id"], ) outcome = {"kind": "remote", "value": value} except RemoteDispatchError as error: outcome = {"kind": "remote_error", "code": error.code, "message": error.message} else: assert policy is not None ctx = HookContext(message["id"], args.session, args.revision) value = await policy.invoke(message["slot"], message["input"], ctx) outcome = {"kind": "hook", "value": value} except asyncio.CancelledError: outcome = {"kind": "failed", "error": "hook invocation interrupted"} except BaseException as error: outcome = {"kind": "failed", "error": "".join(traceback.format_exception(error))[-16_384:]} await wire.send({"kind": "completed", "id": message["id"], "outcome": outcome}) try: while True: message = await wire.receive() kind = message["kind"] if kind == "host_reply": bridge.reply(message) elif kind == "interrupt": fields(message, {"kind", "id"}) if active and not active.done() and message["id"] == active_id: active.cancel() elif kind == "shutdown": fields(message, {"kind"}) break elif kind in {"execute", "invoke", "remote_dispatch"}: if kind == "remote_dispatch": fields(message, {"kind", "id", "operation", "payload"}) if dispatcher is None: raise ProtocolError("only a remote worker can dispatch operations") text(message["operation"], 128) elif kind == "execute": fields(message, {"kind", "id", "code"}, {"purpose"}) if kernel is None: raise ProtocolError("hook worker cannot execute cells") text(message["code"], MAX_CODE) purpose = message.get("purpose", "interactive") if purpose not in {"interactive", "maintenance", "eval"}: raise ProtocolError(f"invalid execution purpose: {purpose}") else: fields(message, {"kind", "id", "slot", "input"}) if policy is None: raise ProtocolError("workbench cannot run required policy hooks") text(message["slot"], 128) identifier(message["id"]) if active and not active.done(): raise ProtocolError("concurrent foreground operation rejected") if active: active.result() # Propagate previous transport errors. active_id = message["id"] active = asyncio.create_task(run(message)) else: raise ProtocolError(f"unexpected host message: {kind}") finally: if active and not active.done(): active.cancel() # Host enforces a hard process deadline; do not depend on cooperative # cancellation when arbitrary Python catches CancelledError. sys.stdout, sys.stderr = original_out, original_err writer.close() with contextlib.suppress(ConnectionError): await writer.wait_closed() def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--socket", required=True) parser.add_argument("--role", choices=["workbench", "hooks", "remote"], required=True) parser.add_argument("--session", required=True) parser.add_argument("--generation", required=True) parser.add_argument("--token", required=True) # Host-issued launch binding (a TargetStamp JSON). Echoed in Hello as a # binding token; verified identity always comes from a measured probe. parser.add_argument("--target") parser.add_argument("--behavior") parser.add_argument("--workspace") parser.add_argument("--revision") # Content-identity payload of the skill bodies this worker may offer. Written by # the host in the same directory as the control socket; absent means no skills. parser.add_argument("--skills") args = parser.parse_args() if args.role == "hooks" and (not args.behavior or not args.revision): parser.error("hook workers need a behavior snapshot and revision") sys.dont_write_bytecode = True try: asyncio.run(serve(args)) except (ProtocolError, ConnectionError, asyncio.IncompleteReadError) as error: print(f"klbr worker disconnected: {error}", file=sys.stderr) raise SystemExit(2) from error if __name__ == "__main__": main()