diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md index 84cd057..02ec88f 100644 --- a/docs/CHANGELOG.md +++ b/docs/CHANGELOG.md @@ -1,3 +1,39 @@ +## v1.47.0 — 2026-10-21 + +**Opt-in: one shared Codex host for a flow walk.** A Codex worker starts its +own `codex exec` process: the launcher, the native binary, a login-shell +snapshot and git probes. Native Codex runs its subagents as threads in one +process. Codex lab R200 measured about 114 s from `codex exec` launch to the +first model turn for each cold builder, because the lab guards every process. +R202 measured 0.05 s from a native subagent spawn to its own turn. + +- Set `"defaults": {"codex_host": true}` in the graph, or + `PUPPETMASTER_CODEX_HOST=1` for the walker. A graph value wins. Off by + default. +- The top walker starts one `codex app-server` on a Unix socket + (`codex_host.py`) and stops it at walk end. Child walkers and workers find + it through `PUPPETMASTER_CODEX_HOST_RECORD`. The next walk of the run stops + a host that a dead walker left. +- Each Codex agent and judge runs as a thread on the host + (`codex_session.run_host_turn`, stdlib WebSocket client `codex_ws.py`). A + repair resumes its thread. A cut interrupts only its own turn. An approval + request is declined and fails the turn. +- The adapter reports file changes and commands as `codex exec --json` + items, so the write-scope gate, rollout usage and resume do not change. + Each worker gets its own shell environment through the thread's + `shell_environment_policy`. +- These runs use `codex exec`: a fork, a bypass run, a native-steer run, a + custom `executable`, extra arguments other than `-c key=value`, a thread in + another Codex home, and Windows (no Unix sockets). A host that does not + start, or a connection that fails before the turn, also falls back. +- Each host turn writes `codex_host_turn.json`: host start and ready times, + the request time, the thread start and the turn start. +- Replica (R179 regions graph, Sol 6.1 high, 3 interleaved pairs): host on + 548.6 / 292.2 / 370.6 s, host off 345.4 / 600.4 / 372.8 s. All runs + passed the judge. Host runs: 69 of 69 turns on the host, 0 `codex exec` + dispatches, request to own-turn start median 1.13 s. A local `codex exec` + start is about 1.7 s, so the wall time does not change on this Mac. + ## v1.46.2 — 2026-10-19 **A nested builder's `revise` no longer blocks follow-up reuse.** Codex R196 diff --git a/docs/FLOWS.md b/docs/FLOWS.md index ea004fa..e208c6e 100644 --- a/docs/FLOWS.md +++ b/docs/FLOWS.md @@ -444,5 +444,33 @@ a resumed walker kills the command the dead walker left running (when the pid still names that command) and runs it again. A pid that now names another process is left alone. +### One shared Codex host + +A Codex worker usually starts its own `codex exec` process: the launcher, the +native binary, a login-shell snapshot and git probes. Native Codex runs its +subagents as threads in one process. With the host on, the top walker starts +one `codex app-server` on a Unix socket. Every Codex agent and judge of the +walk, in map items too, then runs as a thread on that host. + +- Turn it on with `"defaults": {"codex_host": true}` in the graph, or set + `PUPPETMASTER_CODEX_HOST=1` for the walker. A graph value wins over the + environment. `false` turns it off. +- A repair resumes its thread on the host. A cut stops only its own turn. +- Each worker gets its own shell environment (ports, worker flags) through the + thread's `shell_environment_policy`. +- The write-scope gate, rollout usage and session resume work as with + `codex exec`. The adapter reports the host's file changes and commands as + `codex exec --json` items. +- These runs use `codex exec`: a fork (`fork`), a bypass run, a native-steer + run, a custom `executable`, extra arguments other than `-c key=value`, and + a thread that another Codex home holds. +- When the host does not start, or a connection fails before the turn, the + worker runs `codex exec`. The walker stops the host at walk end, and the + next walk of the run stops a host that a dead walker left. +- POSIX only. Windows runs `codex exec`. + +Each host turn writes `codex_host_turn.json` in its capture directory: host +start and ready times, the request time, the thread start and the turn start. + Runs live under `/flows/runs//`. Pass `--cwd` (CLI) or `cwd` (MCP) so the run's state follows the workspace. diff --git a/puppetmaster/__init__.py b/puppetmaster/__init__.py index 0d71a4a..fcef00a 100644 --- a/puppetmaster/__init__.py +++ b/puppetmaster/__init__.py @@ -1,4 +1,4 @@ """Puppetmaster: distributed agent workers with stitched shared memory.""" __all__ = ["__version__"] -__version__ = "1.46.2" +__version__ = "1.47.0" diff --git a/puppetmaster/adapters/codex.py b/puppetmaster/adapters/codex.py index e694153..e56e1eb 100644 --- a/puppetmaster/adapters/codex.py +++ b/puppetmaster/adapters/codex.py @@ -2,6 +2,7 @@ import json import os +import time from pathlib import Path from typing import Any, Mapping, Optional, Union @@ -314,6 +315,14 @@ def _invoke_cli( cwd: Path, timeout_seconds: int, ) -> StreamedProcess: + host = _host_for(task, prepared) + if host is not None: + from puppetmaster.invocation import check_external_dispatch + + check_external_dispatch() + hosted = self._invoke_host(task, prepared, cwd, timeout_seconds, host) + if hosted is not None: + return hosted if prepared.extras.get("resumed") or not bool((task.payload or {}).get("native_steer")): home = self._worker_home(prepared) from puppetmaster import codex_home @@ -346,6 +355,92 @@ def _invoke_cli( except Exception: return super()._invoke_cli(task, prepared, cwd, timeout_seconds) + def _invoke_host(self, task: Task, prepared: CliInvocation, cwd: Path, + timeout_seconds: int, host: dict) -> Optional[StreamedProcess]: + """Run the turn as a thread on the flow's shared host (see codex_host). + + Returns None when the host fails before the turn starts: the caller + then runs ``codex exec``. The turn is reported as ``codex exec --json`` + events, with the file changes and commands Codex made. + """ + from puppetmaster.adapters.codex_session import run_host_turn + from puppetmaster.adapters.codex_ws import HostSocket + from puppetmaster.cancellation import JobCancelled, is_cancelled + from puppetmaster.fs_permissions import mkdir_private, write_private_text + from puppetmaster.invocation import current_attempt + from puppetmaster.swarm_reasoning import caller_pinned_effort + from ._streaming import _resolve_sidecar_state_dir, _write_dispatch_receipt, capture_dir + + config = _host_config(task.payload.get("extra_args")) + prompt = str(prepared.extras.get("prompt") or task.instruction or "") + env = prepared.env + if env is None: + env = inject_worker_cli_env(apply_worktree_ports(os.environ.copy(), cwd)) + config["shell_environment_policy"] = _host_shell_policy(env) + resume_record = prepared.extras.get("resume") if prepared.extras.get("resumed") else None + session_id = _resumed_session_id(resume_record) + home = Path(host["home"]) if host.get("home") else _codex_home().user_home(env) + baseline = _rollout_baseline(home, resume_record) + state_dir = _resolve_sidecar_state_dir() + log_dir = capture_dir(state_dir, task) if state_dir is not None else None + log_path = None + if log_dir is not None: + mkdir_private(log_dir) + log_path = log_dir / "codex_host.log" + receipt = _write_dispatch_receipt( + log_dir, task=task, sidecar_name="codex_host", + command=["codex-host", str(host.get("socket")), "resume" if session_id else "start"], + cwd=str(cwd), env=env, stdin_data=prompt, pid=int(host["pid"])) + budget = task.payload.get("max_output_bytes") + requested_at = time.time() + result = run_host_turn( + lambda: HostSocket(str(host["socket"]), timeout=10.0), + cwd, prompt, + model=str(prepared.extras.get("model") or "") or None, + sandbox=str(prepared.extras.get("sandbox") or "workspace-write"), + approval_policy=str(prepared.extras.get("approval_policy") or "never"), + config=config, + effort=caller_pinned_effort(task.payload), + ephemeral=bool(prepared.extras.get("ephemeral", True)), + resume_thread=session_id, + timeout=float(timeout_seconds), + cancellation_check=lambda: is_cancelled(task.job_id), + protocol_log_path=log_path, + max_output_bytes=budget if isinstance(budget, int) and not isinstance(budget, bool) and budget > 0 else None, + ) + if result.thread_id is None and result.status == "failed": + return None + if log_dir is not None: + # Lab observation: host start, this request and its own turn start. + write_private_text(log_dir / "codex_host_turn.json", json.dumps({ + "host_pid": host.get("pid"), "host_started_at": host.get("started_at"), + "host_ready_at": host.get("ready_at"), "requested_at": requested_at, + "thread_started_at": result.thread_started_at, + "turn_started_at": result.turn_started_at, "finished_at": time.time(), + "thread_id": result.thread_id, "resumed": session_id is not None, + "status": result.status}), lock=False) + if result.status == "cancelled": + raise JobCancelled(task.job_id) + stdout = _app_server_events(result) + usage = None + if session_id is not None or not prepared.extras.get("ephemeral", True): + usage = _rollout_attempt_usage(home, resume_record, baseline, stdout) + if usage is None and session_id is None: + usage = _app_server_attempt_usage(result) + output_limit = result.status == "output_limit" + return StreamedProcess( + returncode=0 if result.status == "completed" else 1, + stdout=stdout, + stderr=result.error or "", + timed_out=result.status == "timeout" or output_limit, + live_log_path=str(log_path) if log_path is not None else None, + spawn_error=None, + output_limit_hit=output_limit, + attempt_id=current_attempt(task), + dispatch_receipt=receipt, + attempt_usage=usage, + ) + def _invoke_app_server(self, task: Task, prepared: CliInvocation, cwd: Path, timeout_seconds: int, cancelled) -> StreamedProcess: """Run the native_steer turn and report it as a ``codex exec --json`` stream.""" @@ -734,6 +829,7 @@ def _app_server_events(result: Any) -> str: if result.thread_id: events.append({"type": "thread.started", "thread_id": result.thread_id}) events.append({"type": "turn.started"}) + events.extend({"type": "item.completed", "item": item} for item in getattr(result, "items", ())) events.extend({"type": "item.completed", "item": {"type": "agent_message", "text": text}} for text in result.messages) if result.status == "failed": @@ -743,6 +839,58 @@ def _app_server_events(result: Any) -> str: return "\n".join(json.dumps(event) for event in events) + "\n" +def _host_for(task: Task, prepared: CliInvocation) -> Optional[dict]: + """The flow's live shared host when this run can be a thread on it, else None. + + A fork, a native-steer run, a bypass run, extra arguments other than + ``-c key=value`` and a thread held in another home run ``codex exec``. + """ + from puppetmaster import codex_host + + host = codex_host.current() + if host is None or prepared.extras.get("fork") or prepared.extras.get("bypass") \ + or (task.payload or {}).get("executable"): + return None + if bool((task.payload or {}).get("native_steer")) or _host_config(task.payload.get("extra_args")) is None: + return None + resume = prepared.extras.get("resume") if prepared.extras.get("resumed") else None + session_id = _resumed_session_id(resume) + if session_id is not None: + home = _codex_home().home_for_session(session_id) + if home is None or str(home) != str(host.get("home")): + return None + return host + + +def _host_config(extra_args: object) -> Optional[dict]: + """``-c key=value`` arguments as a thread config, or None for any other argument.""" + parts = command_parts(extra_args) if extra_args else [] + config: dict = {} + index = 0 + while index < len(parts): + if parts[index] != "-c" or index + 1 >= len(parts) or "=" not in parts[index + 1]: + return None + key, raw = parts[index + 1].split("=", 1) + try: + value: Any = json.loads(raw) + except ValueError: + value = raw.strip("'") + config[key.strip()] = value + index += 2 + return config + + +def _host_shell_policy(env: Mapping[str, str]) -> dict: + """The thread shell environment: this worker's env where it differs from the host's.""" + host_env = inject_worker_cli_env(os.environ.copy()) + changed = {key: value for key, value in env.items() if host_env.get(key) != value} + policy: dict = {"inherit": "all", "set": changed} + dropped = sorted(key for key in host_env if key not in env) + if dropped: + policy["exclude"] = dropped + return policy + + def _app_server_attempt_usage(result: Any) -> Optional[dict]: """The app-server's token counts as attempt usage, or None when it sent none. diff --git a/puppetmaster/adapters/codex_session.py b/puppetmaster/adapters/codex_session.py index 0baba82..06f0108 100644 --- a/puppetmaster/adapters/codex_session.py +++ b/puppetmaster/adapters/codex_session.py @@ -78,7 +78,8 @@ def number(name: str) -> Optional[int]: def run_codex_session( command_prefix: Union[str, Sequence[str]], cwd: Union[str, Path], prompt: str, - model: Optional[str] = None, sandbox: str = "workspace-write", config: Optional[Dict[str, Any]] = None, + model: Optional[str] = None, sandbox: str = "workspace-write", approval_policy: str = "never", + config: Optional[Dict[str, Any]] = None, effort: Optional[str] = None, ephemeral: bool = True, timeout: float = 600, env: Optional[Dict[str, str]] = None, log: Optional[Callable[[str], None]] = None, event_sink: Optional[Callable[[Dict[str, Any]], None]] = None, cancellation_check: Optional[Callable[[], bool]] = None, @@ -239,3 +240,186 @@ def deny_server_request(obj: Dict[str, Any]) -> None: if protocol_log_path: Path(protocol_log_path).write_text(result.raw_protocol, encoding="utf-8") return result + + +@dataclass +class HostTurnResult(CodexSessionResult): + """One thread turn on a shared host, with the items the write-scope gate reads.""" + items: List[Dict[str, Any]] = field(default_factory=list) + rollout_path: Optional[str] = None + thread_started_at: Optional[float] = None + turn_started_at: Optional[float] = None + + +def _exec_item(item: Dict[str, Any]) -> Optional[Dict[str, Any]]: + """An app-server item in the ``codex exec --json`` item shape, or None.""" + kind = item.get("type") + if kind == "fileChange": + changes = [{"path": change.get("path"), "kind": (change.get("kind") or {}).get("type")} + for change in item.get("changes") or [] if isinstance(change, dict)] + return {"type": "file_change", "changes": changes, "status": item.get("status")} + if kind == "commandExecution": + return {"type": "command_execution", "command": item.get("command"), + "exit_code": item.get("exitCode"), "status": item.get("status")} + return None + + +def run_host_turn( + connect: Callable[[], Any], cwd: Union[str, Path], prompt: str, *, + model: Optional[str] = None, sandbox: str = "workspace-write", approval_policy: str = "never", + config: Optional[Dict[str, Any]] = None, + effort: Optional[str] = None, ephemeral: bool = False, resume_thread: Optional[str] = None, + timeout: float = 600, cancellation_check: Optional[Callable[[], bool]] = None, + protocol_log_path: Optional[Union[str, Path]] = None, max_output_bytes: Optional[int] = None, +) -> HostTurnResult: + """Run one turn as a thread on a shared ``codex app-server`` host. + + ``connect`` returns a connected :class:`codex_ws.HostSocket`. A fresh run + starts a thread; a repair resumes ``resume_thread``. A cut or the deadline + interrupts only this thread's turn: the host and its other threads go on. + An approval request is declined and fails the turn, as in + :func:`run_codex_session`. Notifications for another thread are ignored. + """ + import socket + + from puppetmaster.adapters.codex_ws import HostClosed + + result = HostTurnResult(status="failed", + protocol_log_path=str(protocol_log_path) if protocol_log_path else None) + captured: List[str] = [] + deadline = time.monotonic() + max(0.01, timeout) + try: + conn = connect() + except OSError as exc: + result.error = f"host connect failed: {exc}" + return result + next_id = [0] + waiting: Dict[int, str] = {} + output_bytes = 0 + interrupted = "" + + def request(method: str, params: Dict[str, Any], tag: str) -> None: + next_id[0] += 1 + message = {"jsonrpc": "2.0", "id": next_id[0], "method": method, "params": params} + _bounded(captured, "> " + json.dumps(message, separators=(",", ":"))) + conn.send(message) + waiting[next_id[0]] = tag + + def interrupt(reason: str) -> None: + nonlocal interrupted + if not interrupted and result.thread_id and result.turn_id: + request("turn/interrupt", {"threadId": result.thread_id, "turnId": result.turn_id}, "interrupt") + interrupted = interrupted or reason + + try: + request("initialize", {"clientInfo": {"name": "puppetmaster", "version": "1"}, "capabilities": {}}, + "initialize") + while True: + now = time.monotonic() + if now >= deadline: + if interrupted: + result.status, result.error = "timeout", "host turn timeout" + break + interrupt("timeout") + deadline = now + 10 + if not interrupted and cancellation_check and cancellation_check(): + interrupt("cancelled") + conn.settimeout(max(0.05, min(0.5, deadline - time.monotonic()))) + try: + obj = conn.recv() + except socket.timeout: + continue + line = json.dumps(obj, separators=(",", ":")) + output_bytes += len(line) + _bounded(captured, "< " + line) + if max_output_bytes and output_bytes > max_output_bytes and not interrupted: + interrupt("output_limit") + if not isinstance(obj, dict): + continue + if "id" in obj and "method" in obj: + # The run never approves: refuse, then fail like the stdio session. + conn.send({"jsonrpc": "2.0", "id": obj["id"], "result": {"decision": "decline"}}) + result.error = "unexpected approval/server request" + interrupt("approval") + continue + if "id" in obj: + tag = waiting.pop(obj["id"], "") + payload = obj.get("result") if isinstance(obj.get("result"), dict) else {} + if "error" in obj: + if tag != "interrupt": + result.error = str(obj["error"]) + break + continue + if tag == "initialize": + conn.send({"jsonrpc": "2.0", "method": "initialized", "params": {}}) + if resume_thread: + request("thread/resume", {"threadId": resume_thread, "cwd": str(Path(cwd).resolve()), + "sandbox": sandbox, "approvalPolicy": approval_policy, + "config": config or {}}, "thread") + else: + request("thread/start", {"model": model, "cwd": str(Path(cwd).resolve()), + "sandbox": sandbox, "approvalPolicy": approval_policy, + "ephemeral": bool(ephemeral), "config": config or {}}, "thread") + elif tag == "thread": + thread = payload.get("thread") if isinstance(payload.get("thread"), dict) else {} + result.thread_id = thread.get("id") or payload.get("threadId") + result.rollout_path = thread.get("path") + result.thread_started_at = time.time() + if not result.thread_id: + result.error = "thread response missing thread id" + break + params: Dict[str, Any] = {"threadId": result.thread_id, + "input": [{"type": "text", "text": prompt}]} + if model: + params["model"] = model + if effort: + params["effort"] = effort + request("turn/start", params, "turn") + elif tag == "turn": + turn = payload.get("turn") if isinstance(payload.get("turn"), dict) else {} + result.turn_id = result.turn_id or turn.get("id") or payload.get("turnId") + continue + method, params = obj.get("method"), obj.get("params") + if not isinstance(params, dict): + continue + if params.get("threadId") not in (None, result.thread_id): + continue + if method == "turn/started": + turn = params.get("turn") if isinstance(params.get("turn"), dict) else {} + result.turn_id = result.turn_id or turn.get("id") or params.get("turnId") + result.turn_started_at = result.turn_started_at or time.time() + elif method == "thread/tokenUsage/updated": + result.input_tokens, result.output_tokens, result.cached_input_tokens = _usage(params) + elif method == "item/completed": + result.messages.extend(_text_from(params)) + item = params.get("item") + mapped = _exec_item(item) if isinstance(item, dict) else None + if mapped is not None: + result.items.append(mapped) + elif method == "turn/completed": + turn = params.get("turn") if isinstance(params.get("turn"), dict) else {} + state = str(turn.get("status") or "completed") + if interrupted: + result.status = {"cancelled": "cancelled", "timeout": "timeout", + "output_limit": "output_limit"}.get(interrupted, "failed") + result.error = result.error or f"turn interrupted: {interrupted}" + elif state in ("failed", "interrupted"): + result.status, result.error = "failed", str(turn.get("error") or state) + else: + result.status = "completed" + break + elif method in ("turn/failed", "error") and not params.get("willRetry"): + result.status, result.error = "failed", str(params)[:2000] + break + except (HostClosed, OSError, ValueError) as exc: + result.status = "failed" + result.error = f"host connection lost: {exc}" + finally: + try: + conn.close() + except OSError: + pass + result.raw_protocol = "\n".join(captured) + if protocol_log_path: + Path(protocol_log_path).write_text(result.raw_protocol, encoding="utf-8") + return result diff --git a/puppetmaster/adapters/codex_ws.py b/puppetmaster/adapters/codex_ws.py new file mode 100644 index 0000000..7913698 --- /dev/null +++ b/puppetmaster/adapters/codex_ws.py @@ -0,0 +1,115 @@ +"""A WebSocket client over a Unix socket, for the shared Codex host. + +``codex app-server --listen unix://PATH`` speaks WebSocket on that socket. +The client sends masked text frames of JSON-RPC and reads text frames back. +It answers ping and stops on close. Standard library only. +""" +from __future__ import annotations + +import base64 +import json +import os +import socket +import struct +from typing import Any, Optional + + +class HostClosed(OSError): + """The host closed the connection or the socket ended.""" + + +class HostSocket: + def __init__(self, path: str, timeout: float = 30.0) -> None: + if not hasattr(socket, "AF_UNIX"): + raise OSError("this platform has no Unix sockets") + self._sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + self._sock.settimeout(timeout) + try: + self._sock.connect(path) + self._buffer = self._handshake() + self._parts: list = [] + except Exception: + self._sock.close() + raise + + def _handshake(self) -> bytes: + key = base64.b64encode(os.urandom(16)).decode("ascii") + self._sock.sendall(( + "GET / HTTP/1.1\r\nHost: localhost\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n" + f"Sec-WebSocket-Key: {key}\r\nSec-WebSocket-Version: 13\r\n\r\n").encode("ascii")) + data = b"" + while b"\r\n\r\n" not in data: + chunk = self._sock.recv(4096) + if not chunk: + raise HostClosed("host closed during the WebSocket handshake") + data += chunk + if len(data) > 65536: + raise OSError("WebSocket handshake reply too long") + head, rest = data.split(b"\r\n\r\n", 1) + status = head.split(b"\r\n", 1)[0].split() + if len(status) < 2 or status[1] != b"101": + raise OSError("host refused the WebSocket upgrade") + return rest + + def settimeout(self, timeout: Optional[float]) -> None: + self._sock.settimeout(timeout) + + def _next_frame(self) -> tuple: + """The next whole frame as (fin, opcode, payload). Bytes stay buffered until a frame is whole.""" + while True: + data = self._buffer + if len(data) >= 2: + size, offset = data[1] & 0x7F, 2 + if size == 126 and len(data) >= 4: + size, offset = struct.unpack("!H", data[2:4])[0], 4 + elif size == 127 and len(data) >= 10: + size, offset = struct.unpack("!Q", data[2:10])[0], 10 + if (data[1] & 0x7F) < 126 or offset > 2: + if data[1] & 0x80: + offset += 4 # A server never masks; skip a mask key if one is there. + if len(data) >= offset + size: + self._buffer = data[offset + size:] + return bool(data[0] & 0x80), data[0] & 0x0F, data[offset:offset + size] + # A timeout here keeps the partial frame for the next call. + chunk = self._sock.recv(65536) + if not chunk: + raise HostClosed("host socket ended") + self._buffer += chunk + + def _frame(self, opcode: int, payload: bytes) -> None: + mask = os.urandom(4) + size = len(payload) + if size < 126: + head = struct.pack("!BB", 0x80 | opcode, 0x80 | size) + elif size < 65536: + head = struct.pack("!BBH", 0x80 | opcode, 0x80 | 126, size) + else: + head = struct.pack("!BBQ", 0x80 | opcode, 0x80 | 127, size) + masked = bytes(byte ^ mask[index % 4] for index, byte in enumerate(payload)) + self._sock.sendall(head + mask + masked) + + def send(self, message: dict) -> None: + self._frame(0x1, json.dumps(message, separators=(",", ":")).encode("utf-8")) + + def recv(self) -> Any: + """The next JSON message. Raises ``socket.timeout`` when none arrives in time.""" + while True: + fin, opcode, payload = self._next_frame() + if opcode == 0x8: + raise HostClosed("host closed the WebSocket") + if opcode == 0x9: + self._frame(0xA, payload) + continue + if opcode == 0xA: + continue + self._parts.append(payload) + if fin: + message, self._parts = b"".join(self._parts), [] + return json.loads(message.decode("utf-8")) + + def close(self) -> None: + try: + self._frame(0x8, b"") + except OSError: + pass + self._sock.close() diff --git a/puppetmaster/codex_host.py b/puppetmaster/codex_host.py new file mode 100644 index 0000000..3368911 --- /dev/null +++ b/puppetmaster/codex_host.py @@ -0,0 +1,192 @@ +"""One shared ``codex app-server`` for a flow walk. + +A ``codex exec`` worker starts its own Codex process: the Node launcher, the +native binary, a login-shell snapshot and git probes. Native Codex runs its +subagents as threads in one process. With the host on, the top flow walker +starts one ``codex app-server`` on a Unix socket, and every Codex agent and +judge of the walk (map items too) runs as a thread on it. + +The walker exports the record path in ``PUPPETMASTER_CODEX_HOST_RECORD``. +Child walkers and workers inherit it. The adapter uses the host only while +the recorded process is alive; otherwise it runs ``codex exec`` as before. +POSIX only: the socket is AF_UNIX. +""" +from __future__ import annotations + +import json +import os +import signal +import socket +import subprocess +import time +from pathlib import Path +from typing import Any, Mapping, Optional + +from puppetmaster.proc_identity import pid_reused, process_identity + +OPT_IN_ENV = "PUPPETMASTER_CODEX_HOST" +RECORD_ENV = "PUPPETMASTER_CODEX_HOST_RECORD" +RECORD_NAME = "codex-host.json" +# Bounded workers never fork subagents, and must not accrue memories. +HOST_FLAGS = ("--disable", "memories", "--disable", "multi_agent") +_START_SECONDS = 20.0 + + +def supported() -> bool: + return hasattr(socket, "AF_UNIX") and os.name != "nt" + + +def requested(graph: Mapping[str, Any], env: Optional[Mapping[str, str]] = None) -> bool: + """Whether a walk of ``graph`` runs its Codex workers on a shared host.""" + env = os.environ if env is None else env + defaults = graph.get("defaults") if isinstance(graph.get("defaults"), dict) else {} + if "codex_host" in defaults: + return defaults.get("codex_host") is True and supported() + return str(env.get(OPT_IN_ENV, "")).strip().lower() in ("1", "true", "on", "yes") and supported() + + +def _alive(record: Mapping[str, Any]) -> bool: + try: + pid = int(record.get("pid") or 0) + except (TypeError, ValueError): + return False + if pid <= 0: + return False + try: + os.kill(pid, 0) + except OSError: + return False + return not pid_reused(pid, record.get("proc")) + + +def read(path: Path) -> Optional[dict]: + try: + record = json.loads(Path(path).read_text(encoding="utf-8")) + except (OSError, ValueError): + return None + return record if isinstance(record, dict) else None + + +def current(env: Optional[Mapping[str, str]] = None) -> Optional[dict]: + """The live host this process belongs to, or None.""" + env = os.environ if env is None else env + path = (env.get(RECORD_ENV) or "").strip() + if not path or not supported(): + return None + record = read(Path(path)) + return record if record is not None and _alive(record) and record.get("socket") else None + + +def start(run_dir: Path, command_base: list, env: Optional[Mapping[str, str]] = None) -> Optional[dict]: + """Start the host for a walk and write its record in ``run_dir``. + + Returns None when the host cannot start; workers then use ``codex exec``. + """ + from puppetmaster import codex_home + from puppetmaster.codegraph import inject_worker_cli_env + + run_dir = Path(run_dir) + reap(run_dir) + host_env = inject_worker_cli_env(dict(os.environ if env is None else env)) + home: Optional[Path] = None + if codex_home.enabled(host_env): + try: + home = codex_home.prepare(host_env) + except Exception: + home = None # Best effort: the host then runs in the user's own home. + if home is not None: + host_env["CODEX_HOME"] = str(home) + try: + codex_home.ensure_system_skills(home, list(command_base)) + except Exception: + pass # Best effort: stock Codex still installs the bundle itself. + link = run_dir / "codex-host.sock" + try: + link.unlink() + except OSError: + pass + command = [*command_base, "app-server", "--listen", f"unix://{link}", *HOST_FLAGS] + try: + with open(run_dir / "codex-host.log", "ab") as log: + process = subprocess.Popen(command, cwd=str(run_dir), env=host_env, stdin=subprocess.DEVNULL, + stdout=log, stderr=log, start_new_session=True) + except OSError: + return None + started = time.time() + record = {"pid": process.pid, "proc": process_identity(process.pid), "link": str(link), + "socket": None, "home": str(home) if home is not None else None, + "started_at": started, "command": command} + deadline = time.monotonic() + _START_SECONDS + while time.monotonic() < deadline and process.poll() is None: + # Codex links the given path to a short socket path; AF_UNIX paths are + # limited to about 104 bytes, so clients connect to the target. + target = os.path.realpath(link) if link.is_symlink() else str(link) + if Path(target).is_socket(): + record["socket"] = target + break + time.sleep(0.05) + if record["socket"] is None: + _signal(record, signal.SIGKILL) + return None + record["ready_at"] = time.time() + path = run_dir / RECORD_NAME + tmp = path.with_suffix(".tmp") + tmp.write_text(json.dumps(record), encoding="utf-8") + os.replace(tmp, path) + record["record"] = str(path) + return record + + +def _signal(record: Mapping[str, Any], sig: int) -> None: + if not _alive(record): + return + try: + os.killpg(int(record["pid"]), sig) + except OSError: + try: + os.kill(int(record["pid"]), sig) + except OSError: + pass + + +def _reap_child(pid: int) -> None: + """Reap the host when this process started it: a zombie still answers kill(pid, 0).""" + try: + os.waitpid(pid, os.WNOHANG) + except (ChildProcessError, OSError): + pass + + +def stop(record: Optional[Mapping[str, Any]], grace: float = 5.0) -> None: + """Stop the host process group and copy a refreshed login back.""" + if not record: + return + from puppetmaster import codex_home + + _signal(record, signal.SIGTERM) + deadline = time.monotonic() + grace + while _alive(record) and time.monotonic() < deadline: + _reap_child(int(record["pid"])) + time.sleep(0.05) + _signal(record, signal.SIGKILL) + _reap_child(int(record["pid"])) + home = record.get("home") + if home and Path(home) == codex_home.worker_home_root(): + try: + codex_home.sync_back() + except Exception: + pass # Best effort: the next worker start copies auth in again. + path = record.get("record") + if path: + try: + Path(path).unlink() + except OSError: + pass + + +def reap(run_dir: Path) -> None: + """Stop a host left by a walker of this run that died.""" + path = Path(run_dir) / RECORD_NAME + record = read(path) + if record is not None: + stop({**record, "record": str(path)}, grace=1.0) diff --git a/puppetmaster/flow.py b/puppetmaster/flow.py index 33b5234..228d49d 100644 --- a/puppetmaster/flow.py +++ b/puppetmaster/flow.py @@ -47,6 +47,7 @@ import os import re import signal +import shutil import subprocess import sys import time @@ -134,6 +135,8 @@ def validate_graph(graph: Any, *, child: bool = False, depth: int = 0, problems.append("defaults.payload must be an object") own_defaults = {key: value for key, value in own_defaults.items() if key != "payload"} problems.extend(_effort_problems(own_defaults, "defaults")) + if "codex_host" in own_defaults and not isinstance(own_defaults["codex_host"], bool): + problems.append("defaults.codex_host must be true or false") lanes = own_defaults.get("lanes") if lanes is not None: if not isinstance(lanes, dict) or not set(lanes) <= set(LANES): @@ -879,6 +882,7 @@ def walk(state_dir: Path, run_id: str, *, execute: Optional[NodeExecutor] = None lock.acquire() except TimeoutError: raise FlowError(f"flow run {run_id} is already being walked") from None + host: Optional[dict] = None try: run = load_run(state_dir, run_id) if not _prepare(state_dir, run, answer=answer, restart=restart, @@ -888,14 +892,47 @@ def walk(state_dir: Path, run_id: str, *, execute: Optional[NodeExecutor] = None save_run(state_dir, run) if run.parent is None: _mark_walking(state_dir, run) + host = _start_codex_host(state_dir, run) executor = execute or JobNodeExecutor(state_dir, backend=run.backend, worker_mode=run.worker_mode) return _walk(state_dir, run, executor) finally: + _stop_codex_host(host) _clear_marker(state_dir, run_id, "walker.pid") _unmark_walking(run_id) lock.release() +def _start_codex_host(state_dir: Path, run: FlowRun) -> Optional[dict]: + """Start the walk's shared Codex host when the graph or the operator asks for it. + + Child walkers and workers find it through the inherited record path. A + host that does not start leaves every worker on ``codex exec``. + """ + from puppetmaster import codex_host + from puppetmaster.adapters._base import command_parts + + if not codex_host.requested(run.graph) or codex_host.current() is not None: + return None + command = command_parts(os.environ.get("CODEX_COMMAND") or "codex") + resolved = shutil.which(command[0]) + if resolved is None: + return None + record = codex_host.start(run_path(state_dir, run.run_id).parent, [resolved, *command[1:]]) + if record is not None: + os.environ[codex_host.RECORD_ENV] = record["record"] + return record + + +def _stop_codex_host(record: Optional[dict]) -> None: + if record is None: + return + from puppetmaster import codex_host + + if os.environ.get(codex_host.RECORD_ENV) == record.get("record"): + os.environ.pop(codex_host.RECORD_ENV, None) + codex_host.stop(record) + + def _mark_walking(state_dir: Path, run: FlowRun) -> None: from puppetmaster.fs_permissions import write_private_text from puppetmaster.state import walking_runs_dir diff --git a/pyproject.toml b/pyproject.toml index bc4b397..c8dd9e4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "puppetmaster-ai" -version = "1.46.2" +version = "1.47.0" description = "Gunicorn-style agent swarms with Redis-like shared memory and artifact stitching. (Imported as `puppetmaster`; published as `puppetmaster-ai` because the bare PyPI name is held by an abandoned 2019 project, name-reassignment pending.)" readme = "README.md" requires-python = ">=3.9" diff --git a/tests/test_codex_host.py b/tests/test_codex_host.py new file mode 100644 index 0000000..47aefee --- /dev/null +++ b/tests/test_codex_host.py @@ -0,0 +1,308 @@ +from __future__ import annotations + +import base64 +import hashlib +import json +import os +import shutil +import socket +import struct +import tempfile +import threading +import sys +import unittest +from pathlib import Path +from unittest.mock import patch + +from puppetmaster import codex_host +from puppetmaster.adapters import codex as codex_adapter +from puppetmaster.adapters.codex_session import run_host_turn +from puppetmaster.adapters.codex_ws import HostSocket +from puppetmaster import flow +from puppetmaster.flow import validate_graph +from puppetmaster.proc_identity import process_identity +from puppetmaster.worker_attribution import codex_event_references + + +class FakeHost: + """A one-connection WebSocket JSON-RPC server on a Unix socket. + + ``script(method, params, send)`` answers each client request. + """ + + def __init__(self, script): + self.dir = tempfile.mkdtemp(prefix="pmh", dir="/tmp" if os.path.isdir("/tmp") else None) + self.path = os.path.join(self.dir, "s") + self.script = script + self.received: list = [] + self.server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + self.server.bind(self.path) + self.server.listen(1) + self.thread = threading.Thread(target=self._serve, daemon=True) + self.thread.start() + + def close(self): + self.server.close() + shutil.rmtree(self.dir, ignore_errors=True) + + def _serve(self): + try: + conn, _ = self.server.accept() + except OSError: + return + self.conn = conn + with conn: + data = b"" + while b"\r\n\r\n" not in data: + data += conn.recv(4096) + key = [line.split(b":", 1)[1].strip() for line in data.split(b"\r\n") + if line.lower().startswith(b"sec-websocket-key")][0] + accept = base64.b64encode(hashlib.sha1(key + b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11").digest()) + conn.sendall(b"HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\n" + b"Connection: Upgrade\r\nSec-WebSocket-Accept: " + accept + b"\r\n\r\n") + # A ping first: the client must answer it and go on. + conn.sendall(b"\x89\x00") + buffer = b"" + + def read(count): + nonlocal buffer + while len(buffer) < count: + chunk = conn.recv(65536) + if not chunk: + raise OSError("closed") + buffer += chunk + out, buffer = buffer[:count], buffer[count:] + return out + + def send(message): + payload = json.dumps(message).encode() + head = struct.pack("!BB", 0x81, len(payload)) if len(payload) < 126 else \ + struct.pack("!BBH", 0x81, 126, len(payload)) + conn.sendall(head + payload) + + while True: + try: + first, second = read(2) + except OSError: + return + size = second & 0x7F + if size == 126: + size = struct.unpack("!H", read(2))[0] + elif size == 127: + size = struct.unpack("!Q", read(8))[0] + mask = read(4) + payload = bytes(b ^ mask[i % 4] for i, b in enumerate(read(size))) + if first & 0x0F == 0x8: + return + if first & 0x0F != 0x1: + continue + message = json.loads(payload) + self.received.append(message) + if "method" in message: + self.script(message, send) + + +def turn_script(extra=None): + def script(message, send): + method, ident = message["method"], message.get("id") + if method == "initialize": + send({"id": ident, "result": {}}) + elif method in ("thread/start", "thread/resume"): + send({"id": ident, "result": {"thread": {"id": "t1", "path": "/rollouts/t1.jsonl"}}}) + elif method == "turn/start": + send({"id": ident, "result": {"turn": {"id": "u1"}}}) + send({"method": "turn/started", "params": {"threadId": "t1", "turn": {"id": "u1"}}}) + if extra: + extra(send) + return + send({"method": "item/completed", "params": {"threadId": "other", "item": { + "type": "agentMessage", "text": "not mine"}}}) + send({"method": "item/completed", "params": {"threadId": "t1", "item": { + "type": "fileChange", "status": "completed", + "changes": [{"path": "/w/out/a.txt", "kind": {"type": "add"}, "diff": "ok"}]}}}) + send({"method": "item/completed", "params": {"threadId": "t1", "item": { + "type": "commandExecution", "command": "/bin/zsh -lc 'cat out/b.txt'", + "exitCode": 0, "status": "completed"}}}) + send({"method": "item/completed", "params": {"threadId": "t1", "item": { + "type": "agentMessage", "text": "VERDICT: PASS", "phase": "final_answer"}}}) + send({"method": "thread/tokenUsage/updated", "params": {"threadId": "t1", "tokenUsage": { + "total": {"inputTokens": 120, "outputTokens": 7, "cachedInputTokens": 64}}}}) + send({"method": "turn/completed", "params": {"threadId": "t1", "turn": { + "id": "u1", "status": "completed"}}}) + elif method == "turn/interrupt": + send({"id": ident, "result": {}}) + send({"method": "turn/completed", "params": {"threadId": "t1", "turn": { + "id": "u1", "status": "interrupted"}}}) + return script + + +@unittest.skipUnless(codex_host.supported(), "the shared host needs Unix sockets") +class HostTurnTests(unittest.TestCase): + def host(self, script): + fake = FakeHost(script) + self.addCleanup(fake.close) + return fake + + def test_turn_collects_its_own_thread_items(self): + fake = self.host(turn_script()) + result = run_host_turn(lambda: HostSocket(fake.path, timeout=5), "/w", "build it", + model="m", effort="high", config={"x": 1}, timeout=10) + self.assertEqual(result.status, "completed") + self.assertEqual((result.thread_id, result.turn_id), ("t1", "u1")) + self.assertEqual(result.messages, ["VERDICT: PASS"]) + self.assertEqual(result.rollout_path, "/rollouts/t1.jsonl") + self.assertEqual((result.input_tokens, result.output_tokens, result.cached_input_tokens), + (120, 7, 64)) + self.assertEqual([item["type"] for item in result.items], ["file_change", "command_execution"]) + start = next(m for m in fake.received if m.get("method") == "thread/start")["params"] + self.assertEqual((start["model"], start["approvalPolicy"], start["config"]), ("m", "never", {"x": 1})) + turn = next(m for m in fake.received if m.get("method") == "turn/start")["params"] + self.assertEqual((turn["threadId"], turn["effort"]), ("t1", "high")) + + def test_events_feed_the_write_scope_attribution(self): + fake = self.host(turn_script()) + result = run_host_turn(lambda: HostSocket(fake.path, timeout=5), "/w", "build it", timeout=10) + events = codex_adapter.parse_codex_events(codex_adapter._app_server_events(result)) + self.assertEqual(codex_adapter.observed_thread_id(events), "t1") + paths = codex_event_references(events, Path("/w")).paths + self.assertIn("out/a.txt", paths) + self.assertIn("out/b.txt", paths) + + def test_resume_sends_thread_resume(self): + fake = self.host(turn_script()) + result = run_host_turn(lambda: HostSocket(fake.path, timeout=5), "/w", "repair", + resume_thread="t1", timeout=10) + self.assertEqual(result.status, "completed") + methods = [m.get("method") for m in fake.received] + self.assertIn("thread/resume", methods) + self.assertNotIn("thread/start", methods) + + def test_cut_interrupts_only_its_turn(self): + cut = threading.Event() + fake = self.host(turn_script(extra=lambda send: cut.set())) + result = run_host_turn(lambda: HostSocket(fake.path, timeout=5), "/w", "build it", + timeout=10, cancellation_check=cut.is_set) + self.assertEqual(result.status, "cancelled") + interrupt = next(m for m in fake.received if m.get("method") == "turn/interrupt")["params"] + self.assertEqual(interrupt, {"threadId": "t1", "turnId": "u1"}) + + def test_approval_request_is_declined_and_fails(self): + def ask(send): + send({"id": 99, "method": "commandExecution/requestApproval", "params": {"threadId": "t1"}}) + fake = self.host(turn_script(extra=ask)) + result = run_host_turn(lambda: HostSocket(fake.path, timeout=5), "/w", "build it", timeout=10) + self.assertEqual(result.status, "failed") + self.assertIn({"jsonrpc": "2.0", "id": 99, "result": {"decision": "decline"}}, fake.received) + + def test_a_frame_split_across_a_timeout_stays_whole(self): + gate = threading.Event() + + def split(message, send): + if message["method"] == "initialize": + frame = struct.pack("!BB", 0x81, 9) + b'{"id": 1}' + fake.conn.sendall(frame[:4]) + gate.wait(5) + fake.conn.sendall(frame[4:]) + fake = self.host(split) + sock = HostSocket(fake.path, timeout=5) + self.addCleanup(sock.close) + sock.send({"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {}}) + sock.settimeout(0.2) + with self.assertRaises(socket.timeout): + sock.recv() + gate.set() + sock.settimeout(5) + self.assertEqual(sock.recv(), {"id": 1}) + + def test_no_host_fails_before_a_thread(self): + result = run_host_turn(lambda: HostSocket("/tmp/pm-no-such-host.sock", timeout=1), "/w", "x", timeout=5) + self.assertEqual(result.status, "failed") + self.assertIsNone(result.thread_id) + + +class HostConfigTests(unittest.TestCase): + def test_only_config_overrides_can_ride_a_thread(self): + self.assertEqual(codex_adapter._host_config(None), {}) + self.assertEqual(codex_adapter._host_config(["-c", 'model_reasoning_effort="high"', "-c", "n=2"]), + {"model_reasoning_effort": "high", "n": 2}) + self.assertIsNone(codex_adapter._host_config(["--search"])) + self.assertIsNone(codex_adapter._host_config(["-c"])) + + def test_shell_policy_sends_only_this_worker_differences(self): + base = codex_adapter.inject_worker_cli_env(os.environ.copy()) + dropped = next(iter(base)) + env = {key: value for key, value in base.items() if key != dropped} + env["PM_PORT_TEST"] = "4100" + policy = codex_adapter._host_shell_policy(env) + self.assertEqual(policy["set"], {"PM_PORT_TEST": "4100"}) + self.assertEqual(policy["exclude"], [dropped]) + + +class HostLifecycleTests(unittest.TestCase): + def test_graph_or_operator_requests_the_host(self): + on = {"defaults": {"codex_host": True}} + self.assertEqual(codex_host.requested(on, env={}), codex_host.supported()) + self.assertFalse(codex_host.requested({"defaults": {"codex_host": False}}, + env={codex_host.OPT_IN_ENV: "1"})) + self.assertEqual(codex_host.requested({}, env={codex_host.OPT_IN_ENV: "1"}), codex_host.supported()) + self.assertFalse(codex_host.requested({}, env={})) + + def test_codex_host_must_be_a_boolean(self): + graph = {"id": "g", "entry": "a", "defaults": {"codex_host": "yes"}, + "nodes": [{"id": "a", "kind": "end", "status": "pass"}], "edges": []} + self.assertIn("defaults.codex_host must be true or false", validate_graph(graph)) + + @unittest.skipUnless(codex_host.supported(), "the shared host needs Unix sockets") + def test_current_needs_a_live_recorded_process(self): + folder = Path(tempfile.mkdtemp()) + self.addCleanup(shutil.rmtree, folder, True) + path = folder / codex_host.RECORD_NAME + record = {"pid": os.getpid(), "proc": process_identity(os.getpid()), "socket": "/tmp/s"} + path.write_text(json.dumps(record), encoding="utf-8") + self.assertEqual(codex_host.current({codex_host.RECORD_ENV: str(path)})["pid"], os.getpid()) + path.write_text(json.dumps({**record, "proc": "reused"}), encoding="utf-8") + self.assertIsNone(codex_host.current({codex_host.RECORD_ENV: str(path)})) + self.assertIsNone(codex_host.current({})) + + +FAKE_CODEX = """#!{python} +import os, signal, socket, sys, tempfile, time +link = sys.argv[sys.argv.index("--listen") + 1][len("unix://"):] +short = os.path.join(tempfile.mkdtemp(prefix="pmf", dir="/tmp"), "s") +server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) +server.bind(short) +os.symlink(short, link) +signal.signal(signal.SIGTERM, lambda *_: sys.exit(0)) +while True: + time.sleep(0.1) +""" + + +@unittest.skipUnless(codex_host.supported(), "the shared host needs Unix sockets") +class WalkerHostTests(unittest.TestCase): + def test_top_walk_owns_one_host_for_its_nodes(self): + folder = Path(tempfile.mkdtemp()) + self.addCleanup(shutil.rmtree, folder, True) + fake = folder / "codex" + fake.write_text(FAKE_CODEX.format(python=sys.executable), encoding="utf-8") + fake.chmod(0o755) + check = (f'"{sys.executable}" -c "import json,os;' + f'r=json.load(open(os.environ[\'{codex_host.RECORD_ENV}\']));' + f'os.kill(r[\'pid\'],0);print(\'VERDICT: PASS\')"') + graph = {"id": "host-walk", "entry": "check", "defaults": {"codex_host": True}, + "nodes": [{"id": "check", "kind": "shell", "command": check}, + {"id": "done", "kind": "end", "status": "pass"}], + "edges": [{"from": "check", "to": "done", "when": "ok"}]} + env = {"CODEX_COMMAND": str(fake), "PUPPETMASTER_CODEX_LEAN_HOME": "0"} + with patch.dict(os.environ, env): + os.environ.pop(codex_host.RECORD_ENV, None) + run = flow.new_run(folder / "state", graph, cwd=str(folder)) + run = flow.walk(folder / "state", run.run_id) + self.assertNotIn(codex_host.RECORD_ENV, os.environ) + self.assertEqual((run.status, run.reason), ("done", "pass"), run.steps) + run_dir = flow.run_path(folder / "state", run.run_id).parent + self.assertFalse((run_dir / codex_host.RECORD_NAME).exists()) + + +if __name__ == "__main__": + unittest.main()