Spaces:
Paused
Paused
| """Codex through `codex app-server` (JSON-RPC 2.0 over stdio, newline-delimited, the | |
| "jsonrpc" field omitted on the wire). One app-server process per working agent. | |
| Protocol as observed with 0.162.1 (tests/fixtures/codex_appserver.jsonl): | |
| `initialize` -> `initialized` notification -> `thread/start` | `thread/resume` -> | |
| `turn/start` (no active turn) or `turn/steer` with `expectedTurnId` (active turn). | |
| Notifications: turn/started, item/started|completed, thread/tokenUsage/updated | |
| (`total` is cumulative for the thread, across resumes; on resume the previous total is | |
| replayed once, tagged with the previous turn id), turn/completed, error. | |
| `thread/resume` in 0.162.1 takes only `threadId` (no `path`); the rollout file under | |
| $CODEX_HOME/sessions is enough to resume (verified: rollout-only CODEX_HOME resumed with | |
| context intact). | |
| """ | |
| from __future__ import annotations | |
| import asyncio | |
| import json | |
| import logging | |
| import os | |
| import time | |
| from collections.abc import AsyncIterator | |
| from pathlib import Path | |
| from typing import Any | |
| from flock.models import TokenUsage, TypeSpec | |
| from flock.worker.adapters.base import AdapterEvent, EventQueue, ResumeError | |
| log = logging.getLogger(__name__) | |
| def codex_home() -> Path: | |
| return Path(os.environ.get("CODEX_HOME") or Path.home() / ".codex") | |
| def find_rollout(home: Path, thread_id: str) -> Path | None: | |
| """$CODEX_HOME/sessions/YYYY/MM/DD/rollout-<ts>-<thread_id>.jsonl (newest if several).""" | |
| hits = sorted((home / "sessions").glob(f"*/*/*/rollout-*-{thread_id}.jsonl")) | |
| if not hits: | |
| hits = sorted((home / "sessions").rglob(f"rollout-*-{thread_id}.jsonl")) | |
| return hits[-1] if hits else None | |
| class RpcError(RuntimeError): | |
| pass | |
| def _tokens(b: dict[str, Any]) -> TokenUsage: | |
| inp = int(b.get("inputTokens") or 0) | |
| cached = int(b.get("cachedInputTokens") or 0) | |
| return TokenUsage( | |
| input=max(0, inp - cached), | |
| output=int(b.get("outputTokens") or 0), | |
| cache_read=cached, | |
| cache_write=int(b.get("cacheWriteInputTokens") or 0), | |
| reasoning=int(b.get("reasoningOutputTokens") or 0), | |
| ) | |
| def _ms() -> int: | |
| return int(time.time() * 1000) | |
| class CodexMapper: | |
| """Maps app-server notifications to AdapterEvents; tracks the active turn.""" | |
| def __init__(self, model: str | None = None) -> None: | |
| self.model = model | |
| self.active_turn: str | None = None | |
| self.started: set[str] = set() | |
| def _assistant(self, content: list[dict], stop: str = "stop") -> dict[str, Any]: | |
| return {"role": "assistant", "content": content, "api": "openai-responses", | |
| "provider": "openai", "model": self.model, "stopReason": stop, "timestamp": _ms()} | |
| def _tool(self, item: dict, name: str, args: dict, output: str, is_error: bool) -> list[AdapterEvent]: | |
| call = self._assistant([{"type": "toolCall", "id": item["id"], "name": name, "arguments": args}], | |
| "toolUse") | |
| res = {"role": "toolResult", "toolCallId": item["id"], "toolName": name, | |
| "content": [{"type": "text", "text": output}], "isError": is_error, "timestamp": _ms()} | |
| return [AdapterEvent("message", message=call, raw=item), | |
| AdapterEvent("message", message=res, raw=item)] | |
| def on_turn_id(self, turn_id: str) -> list[AdapterEvent]: | |
| self.active_turn = turn_id | |
| if turn_id in self.started: | |
| return [] | |
| self.started.add(turn_id) | |
| return [AdapterEvent("turn_started", data={"harness_turn_id": turn_id})] | |
| def on_notification(self, method: str, p: dict[str, Any]) -> list[AdapterEvent]: | |
| if method == "turn/started": | |
| return self.on_turn_id(p["turn"]["id"]) | |
| if method == "turn/completed": | |
| turn = p.get("turn") or {} | |
| if turn.get("id") == self.active_turn: | |
| self.active_turn = None | |
| err = turn.get("error") | |
| return [AdapterEvent("turn_ended", status=turn.get("status") or "completed", | |
| error=(err or {}).get("message") if isinstance(err, dict) else None, | |
| data={"harness_turn_id": turn.get("id")}, raw=p)] | |
| if method == "thread/tokenUsage/updated": | |
| total = (p.get("tokenUsage") or {}).get("total") or {} | |
| baseline = p.get("turnId") not in self.started | |
| return [AdapterEvent("usage", model=self.model, usage=_tokens(total), cumulative=True, | |
| baseline=baseline, raw=p)] | |
| if method == "error": | |
| err = p.get("error") or {} | |
| return [AdapterEvent("error", error=str(err.get("message") or err), | |
| data={"willRetry": p.get("willRetry")}, raw=p)] | |
| if method == "item/completed": | |
| return self._item(p.get("item") or {}) | |
| return [] | |
| def _item(self, item: dict[str, Any]) -> list[AdapterEvent]: | |
| t = item.get("type") | |
| if t == "userMessage": | |
| text = "\n".join(c.get("text", "") for c in item.get("content") or [] if c.get("type") == "text") | |
| return [AdapterEvent("message", message={"role": "user", "content": text, "timestamp": _ms()}, | |
| raw=item)] | |
| if t == "agentMessage": | |
| return [AdapterEvent("message", message=self._assistant([{"type": "text", "text": item.get("text") or ""}]), | |
| model=self.model, raw=item)] | |
| if t == "reasoning": | |
| text = "\n".join(str(s) for s in (item.get("summary") or [])) | |
| if not text: | |
| return [] | |
| return [AdapterEvent("message", message=self._assistant([{"type": "thinking", "thinking": text}]), | |
| raw=item)] | |
| if t == "commandExecution": | |
| code = item.get("exitCode") | |
| return self._tool(item, "shell", {"command": item.get("command")}, | |
| item.get("aggregatedOutput") or "", code not in (0, None)) | |
| if t == "fileChange": | |
| return self._tool(item, "apply_patch", {"changes": item.get("changes")}, | |
| str(item.get("status")), item.get("status") == "failed") | |
| if t in ("mcpToolCall", "dynamicToolCall"): | |
| out = item.get("result") if t == "mcpToolCall" else item.get("contentItems") | |
| failed = item.get("status") == "failed" or bool(item.get("error")) | |
| return self._tool(item, str(item.get("tool")), item.get("arguments") or {}, | |
| json.dumps(out, default=str), failed) | |
| if t == "webSearch": | |
| return self._tool(item, "web_search", {"query": item.get("query")}, | |
| json.dumps(item.get("results"), default=str), False) | |
| return [] | |
| class CodexAdapter: | |
| harness = "codex" | |
| def __init__(self, workdir: Path, *, home: Path | None = None, codex_bin: str | None = None, | |
| extra_globs: tuple[str, ...] = ()) -> None: | |
| self.workdir = Path(workdir) | |
| self.home = home or codex_home() | |
| self.codex_bin = codex_bin or os.environ.get("FLOCK_CODEX_CLI") or "codex" | |
| self.extra_globs = extra_globs # e.g. ('*.sqlite*',) if a future version needs more | |
| self.thread_id: str | None = None | |
| self.rollout: Path | None = None | |
| self._proc: asyncio.subprocess.Process | None = None | |
| self._queue = EventQueue() | |
| self._mapper = CodexMapper() | |
| self._pending: dict[int, asyncio.Future] = {} | |
| self._next_id = 0 | |
| self._tasks: list[asyncio.Task] = [] | |
| self._stderr: list[str] = [] | |
| self.effort: str | None = os.environ.get("FLOCK_CODEX_EFFORT") or None | |
| # ------------------------------------------------------------------ json-rpc | |
| async def _write(self, msg: dict[str, Any]) -> None: | |
| if self._proc is None or self._proc.stdin is None: | |
| raise RpcError("app-server is not running") | |
| self._proc.stdin.write((json.dumps(msg) + "\n").encode()) | |
| await self._proc.stdin.drain() | |
| async def request(self, method: str, params: dict[str, Any], timeout: float = 120) -> dict[str, Any]: | |
| self._next_id += 1 | |
| rid = self._next_id | |
| fut: asyncio.Future = asyncio.get_running_loop().create_future() | |
| self._pending[rid] = fut | |
| await self._write({"method": method, "id": rid, "params": params}) | |
| try: | |
| msg = await asyncio.wait_for(fut, timeout) | |
| finally: | |
| self._pending.pop(rid, None) | |
| if "error" in msg: | |
| raise RpcError(f"{method}: {msg['error'].get('message') or msg['error']}") | |
| return msg.get("result") or {} | |
| async def _read(self) -> None: | |
| assert self._proc and self._proc.stdout | |
| try: | |
| while line := await self._proc.stdout.readline(): | |
| try: | |
| msg = json.loads(line) | |
| except ValueError: | |
| continue | |
| await self._dispatch(msg) | |
| except asyncio.CancelledError: | |
| pass | |
| except Exception as e: # noqa: BLE001 | |
| self._queue.put(AdapterEvent("error", error=f"app-server reader: {e}")) | |
| finally: | |
| for f in self._pending.values(): | |
| if not f.done(): | |
| f.set_result({"error": {"message": "app-server exited"}}) | |
| if self._proc is not None: | |
| try: | |
| await asyncio.wait_for(self._proc.wait(), 5) | |
| except (TimeoutError, ProcessLookupError): | |
| pass | |
| self._queue.put(AdapterEvent("exited", data={"stderr": self._stderr[-5:]})) | |
| async def _dispatch(self, msg: dict[str, Any]) -> None: | |
| method = msg.get("method") | |
| if method is None and "id" in msg: # response | |
| fut = self._pending.get(msg["id"]) | |
| if fut and not fut.done(): | |
| fut.set_result(msg) | |
| return | |
| if method and "id" in msg: # server request; approvals are off, so decline anything | |
| await self._write({"id": msg["id"], "error": {"code": -32601, "message": "unsupported"}}) | |
| return | |
| if method: | |
| for ev in self._mapper.on_notification(method, msg.get("params") or {}): | |
| self._queue.put(ev) | |
| async def _drain_stderr(self) -> None: | |
| assert self._proc and self._proc.stderr | |
| while line := await self._proc.stderr.readline(): | |
| self._stderr = (self._stderr + [line.decode(errors="replace").rstrip()])[-20:] | |
| # ------------------------------------------------------------------ adapter | |
| async def start(self, workdir: Path, session_id: str | None, spec: TypeSpec, *, | |
| env: dict[str, str] | None = None) -> str: | |
| self.workdir = Path(workdir) | |
| if session_id and not find_rollout(self.home, session_id): | |
| raise ResumeError(f"no rollout for thread {session_id} under {self.home}/sessions") | |
| self._queue = EventQueue() | |
| self._mapper = CodexMapper(spec.model) | |
| self._proc = await asyncio.create_subprocess_exec( | |
| self.codex_bin, "app-server", | |
| stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.PIPE, | |
| stderr=asyncio.subprocess.PIPE, cwd=str(self.workdir), limit=64 * 1024 * 1024, | |
| env={**os.environ, **(env or {}), "CODEX_HOME": str(self.home)}, | |
| ) | |
| self._tasks = [asyncio.create_task(self._read()), asyncio.create_task(self._drain_stderr())] | |
| try: | |
| await self.request("initialize", {"clientInfo": {"name": "flock", "version": "0.0.1"}}) | |
| await self._write({"method": "initialized", "params": {}}) | |
| params = {"cwd": str(self.workdir), "model": spec.model, "approvalPolicy": "never", | |
| "sandbox": "danger-full-access"} | |
| if session_id: | |
| try: | |
| r = await self.request("thread/resume", {**params, "threadId": session_id, | |
| "excludeTurns": True}) | |
| except RpcError as e: | |
| raise ResumeError(str(e)) from e | |
| else: | |
| r = await self.request("thread/start", {**params, "ephemeral": False}) | |
| except BaseException: | |
| await self.stop() | |
| raise | |
| thread = r.get("thread") or {} | |
| self.thread_id = thread["id"] | |
| self._mapper.model = r.get("model") or thread.get("model") or spec.model | |
| self.rollout = Path(thread["path"]) if thread.get("path") else None | |
| return self.thread_id | |
| async def send(self, text: str) -> None: | |
| inp = [{"type": "text", "text": text, "text_elements": []}] | |
| turn = self._mapper.active_turn | |
| if turn: | |
| try: | |
| await self.request("turn/steer", {"threadId": self.thread_id, "expectedTurnId": turn, | |
| "input": inp}) | |
| return | |
| except RpcError as e: # the turn ended under us; start a new one | |
| log.info("steer failed (%s); starting a new turn", e) | |
| params: dict[str, Any] = {"threadId": self.thread_id, "input": inp} | |
| if self.effort: | |
| params["effort"] = self.effort | |
| r = await self.request("turn/start", params) | |
| for ev in self._mapper.on_turn_id(r["turn"]["id"]): | |
| self._queue.put(ev) | |
| async def stop(self) -> None: | |
| proc, self._proc = self._proc, None | |
| if proc is not None and proc.returncode is None: | |
| try: | |
| if proc.stdin: | |
| proc.stdin.close() | |
| await asyncio.wait_for(proc.wait(), 5) | |
| except (TimeoutError, ProcessLookupError, ConnectionError): | |
| try: | |
| proc.kill() | |
| await proc.wait() | |
| except ProcessLookupError: | |
| pass | |
| for t in self._tasks: | |
| try: | |
| await asyncio.wait_for(t, 5) | |
| except (TimeoutError, asyncio.CancelledError): | |
| t.cancel() | |
| self._tasks = [] | |
| async def events(self) -> AsyncIterator[AdapterEvent]: | |
| async for ev in self._queue: | |
| yield ev | |
| def session_files(self, session_id: str | None) -> list[Path]: | |
| """The thread's rollout, $CODEX_HOME/sessions/YYYY/MM/DD/rollout-*-<id>.jsonl. That is the | |
| verified resume set; Codex rebuilds its SQLite index itself. `extra_globs` (relative to | |
| CODEX_HOME, off by default) adds more files should a future version need them.""" | |
| files: list[Path] = [] | |
| sid = session_id or self.thread_id | |
| r = find_rollout(self.home, sid) if sid else None | |
| if r is None and self.rollout and self.rollout.is_file(): | |
| r = self.rollout | |
| if r: | |
| files.append(r) | |
| for g in self.extra_globs: | |
| files += sorted(p for p in self.home.glob(g) if p.is_file()) | |
| return files | |