"""Bridge state for native Codex TUI sessions.""" from __future__ import annotations import hashlib import json import os import secrets import sys import tempfile from dataclasses import dataclass from pathlib import Path import tomllib CODEX_NATIVE_BRIDGE_ID_LABEL_KEY = "omnigent.codex_native.bridge_id" CODEX_NATIVE_BRIDGE_DIR_ENV_VAR = "HARNESS_CODEX_NATIVE_BRIDGE_DIR" CODEX_NATIVE_REQUEST_SESSION_ID_ENV_VAR = "HARNESS_CODEX_NATIVE_REQUEST_SESSION_ID" _STATE_FILE = "state.json" _STARTUP_ERROR_FILE = "startup_error.json" # Per-MCP-server startup state mirrored from Codex's # ``mcpServer/startupStatus/updated`` notifications. Written by the # forwarder (and by ``wait_for_thread_started`` while it drains startup # events), read by the executor's first-turn gate and the runner's # Stop handler. _MCP_STARTUP_FILE = "mcp_startup.json" # Startup states mirrored from Codex's ``McpServerStartupState`` enum. MCP_STARTUP_STARTING = "starting" MCP_STARTUP_READY = "ready" MCP_STARTUP_FAILED = "failed" MCP_STARTUP_CANCELLED = "cancelled" MCP_STARTUP_STATES = frozenset( {MCP_STARTUP_STARTING, MCP_STARTUP_READY, MCP_STARTUP_FAILED, MCP_STARTUP_CANCELLED} ) # Must match ``_CONFIG_FILE`` in ``claude_native_bridge.py`` because # ``serve-mcp`` reads this filename for the token. _MCP_CONFIG_FILE = "bridge.json" # Config the codex-native PreToolUse/PostToolUse policy hook subprocess # reads to reach the Omnigent server. Mirrors Claude-native's # ``permission_hook.json`` (see ``claude_native_bridge``). Kept in a # separate file from ``state.json`` because it is written once at bridge # prep time (the Omnigent URL + auth do not change across thread rotations), # whereas ``state.json`` mutates on every turn/thread change. _POLICY_HOOK_FILE = "policy_hook.json" _BRIDGE_ROOT = Path.home() / ".omnigent" / "codex-native" def bridge_root() -> Path: """ Return the configured Codex-native bridge root. Tests may monkeypatch :data:`_BRIDGE_ROOT` to isolate bridge files. :returns: Absolute root for Codex-native bridge directories, e.g. ``Path("~/.omnigent/codex-native")``. """ return _BRIDGE_ROOT @dataclass(frozen=True) class CodexNativeBridgeState: """ Runtime state shared by the native Codex wrapper and harness. :param session_id: Omnigent conversation id, e.g. ``"conv_abc123"``. :param socket_path: Unix socket path for the Codex app-server, e.g. ``"/home/user/.omnigent/codex-native/x/app-server.sock"``. :param thread_id: Codex app-server thread id, e.g. ``"0196..."``. :param codex_home: Private per-session ``CODEX_HOME`` path, e.g. ``"/home/user/.omnigent/codex-native/x/codex-home"``. :param active_turn_id: Current Codex turn id, if one is running, e.g. ``"turn_abc123"``. """ session_id: str socket_path: str thread_id: str codex_home: str active_turn_id: str | None = None def bridge_dir_for_bridge_id(bridge_id: str) -> Path: """ Return the bridge directory for a native Codex bridge id. :param bridge_id: Opaque bridge id, e.g. ``"bridge_abc123"``. :returns: Absolute bridge directory under ``~/.omnigent/codex-native``. """ digest = hashlib.sha256(bridge_id.encode("utf-8")).hexdigest()[:32] return _BRIDGE_ROOT / digest def build_codex_native_spawn_env( conversation_id: str, *, bridge_id: str | None = None, ) -> dict[str, str]: """ Build spawn env for the ``codex-native`` harness process. :param conversation_id: Omnigent conversation id, e.g. ``"conv_abc123"``. :param bridge_id: Opaque bridge id from :data:`CODEX_NATIVE_BRIDGE_ID_LABEL_KEY`, e.g. ``"bridge_abc123"``. ``None`` uses *conversation_id*. :returns: Environment variables needed by the Codex-native harness executor. """ resolved_bridge_id = bridge_id or conversation_id return { CODEX_NATIVE_BRIDGE_DIR_ENV_VAR: str(bridge_dir_for_bridge_id(resolved_bridge_id)), CODEX_NATIVE_REQUEST_SESSION_ID_ENV_VAR: conversation_id, } def prepare_bridge_dir(bridge_id: str) -> Path: """ Create the bridge directory for *bridge_id*. :param bridge_id: Opaque bridge id, e.g. ``"bridge_abc123"``. :returns: Prepared absolute bridge directory. """ bridge_dir = bridge_dir_for_bridge_id(bridge_id) bridge_dir.mkdir(mode=0o700, parents=True, exist_ok=True) os.chmod(bridge_dir, 0o700) return bridge_dir def write_mcp_bridge_config(bridge_dir: Path) -> None: """ Write a minimal ``bridge.json`` so ``serve-mcp`` can boot. The config contains only an authentication token (no ``workspace`` key), so the MCP server serves **relay tools only** (from ``tool_relay.json``) — no ``sys_os_*`` tools. This is correct for codex-native: Codex owns its own filesystem tools. Idempotent: skips if a config already exists (avoids overwriting a token that the relay HTTP server was started with). :param bridge_dir: Codex bridge directory, e.g. ``Path("~/.omnigent/codex-native/")``. """ config_path = bridge_dir / _MCP_CONFIG_FILE if config_path.exists(): return bridge_dir.mkdir(mode=0o700, parents=True, exist_ok=True) payload = {"token": secrets.token_urlsafe(32)} fd, tmp_name = tempfile.mkstemp(prefix=f"{_MCP_CONFIG_FILE}.", dir=str(bridge_dir)) try: with os.fdopen(fd, "w", encoding="utf-8") as handle: json.dump(payload, handle, sort_keys=True) handle.write("\n") os.replace(tmp_name, config_path) finally: if os.path.exists(tmp_name): os.unlink(tmp_name) def codex_mcp_config_overrides( bridge_dir: Path, *, python_executable: str | None = None, ) -> list[str]: """ Return ``-c`` config overrides that register the Omnigent MCP server. The overrides configure codex to launch ``serve-mcp`` from :mod:`omnigent.claude_native_bridge` as a stdio MCP server. ``serve-mcp`` reads ``tool_relay.json`` from *bridge_dir* dynamically on every ``tools/list`` call, so relay tools appear as soon as the runner writes the file. :param bridge_dir: Codex bridge directory containing ``bridge.json`` and (eventually) ``tool_relay.json``. :param python_executable: Python executable to run, e.g. ``"/path/to/python"``. ``None`` uses :data:`sys.executable`. :returns: Codex ``-c`` config override strings, e.g. ``['mcp_servers.omnigent.command="python"', ...]``. """ python = python_executable or sys.executable args_toml = json.dumps( ["-m", "omnigent.claude_native_bridge", "serve-mcp", "--bridge-dir", str(bridge_dir)] ) return [ f'mcp_servers.omnigent.command="{python}"', f"mcp_servers.omnigent.args={args_toml}", ] def write_policy_hook_config( bridge_dir: Path, *, ap_server_url: str, ap_auth_headers: dict[str, str], ) -> None: """ Write the Omnigent coordinates the codex-native policy hook needs. The ``PreToolUse`` / ``PostToolUse`` command hook runs as a short subprocess that must POST to ``/v1/sessions/{id}/policies/evaluate`` on the Omnigent server. It cannot inherit the long-lived forwarder's in-memory client, so the Omnigent base URL and auth headers are persisted here and read by :func:`read_policy_hook_config` at hook time. :param bridge_dir: Native Codex bridge directory, e.g. ``Path("~/.omnigent/codex-native/")``. :param ap_server_url: Omnigent server base URL the hook POSTs to, e.g. ``"http://127.0.0.1:8787"``. :param ap_auth_headers: Outbound auth headers for Omnigent requests, e.g. ``{"Authorization": "Bearer "}``. Empty dict for local-server mode with no auth provider. :returns: None. """ bridge_dir.mkdir(mode=0o700, parents=True, exist_ok=True) path = bridge_dir / _POLICY_HOOK_FILE payload = {"ap_server_url": ap_server_url, "ap_auth_headers": ap_auth_headers} fd, tmp_name = tempfile.mkstemp(prefix=f"{_POLICY_HOOK_FILE}.", dir=str(bridge_dir)) try: with os.fdopen(fd, "w", encoding="utf-8") as handle: json.dump(payload, handle, sort_keys=True) handle.write("\n") os.replace(tmp_name, path) finally: if os.path.exists(tmp_name): os.unlink(tmp_name) def read_policy_hook_config(bridge_dir: Path) -> dict[str, object] | None: """ Read the Omnigent coordinates for the codex-native policy hook. :param bridge_dir: Native Codex bridge directory. :returns: Parsed config, e.g. ``{"ap_server_url": "http://127.0.0.1:8787", "ap_auth_headers": {"Authorization": "Bearer "}}``, or ``None`` when no config has been written (no Omnigent server configured for this session). """ path = bridge_dir / _POLICY_HOOK_FILE if not path.is_file(): return None try: raw = json.loads(path.read_text(encoding="utf-8")) except (OSError, json.JSONDecodeError): return None if not isinstance(raw, dict): return None return raw def socket_path_for_bridge_dir(bridge_dir: Path) -> Path: """ Return the Codex app-server socket path for *bridge_dir*. :param bridge_dir: Native Codex bridge directory. :returns: Absolute Unix socket path for the app-server. """ return bridge_dir / "app-server.sock" def codex_home_for_bridge_dir(bridge_dir: Path) -> Path: """ Return the private ``CODEX_HOME`` path for *bridge_dir*. :param bridge_dir: Native Codex bridge directory. :returns: Absolute per-session ``CODEX_HOME`` directory. """ return bridge_dir / "codex-home" def read_codex_config_model(bridge_dir: Path) -> str | None: """ Read the active model from this session's Codex ``config.toml``. The top-level ``model`` key is exactly what an in-TUI ``/model`` writes (codex's ``config/batchWrite``), so it is the source of truth for which model the user has selected. Reading it from the hook at evaluation time is race-free: unlike the forwarder's async ``external_model_change`` mirror to ``model_override``, the value is read synchronously the instant the cost gate needs it, so a ``/model`` switch takes effect on the very next tool call. Best-effort + fail-safe: a missing / unreadable / unparsable file (or a config with no top-level ``model``) returns ``None``, so the caller falls back to the server-resolved model rather than crashing. Per-session isolation: ``config.toml`` is **copied** (not symlinked) into each session's private ``CODEX_HOME`` by ``_populate_codex_home_config`` (see ``_CODEX_HOME_COPY_FILES`` in ``omnigent.inner.codex_executor``), then seeded with the session's launch model by ``_pin_codex_config_model`` in ``omnigent.codex_native_app_server``. An in-TUI ``/model`` writes only to that session's copy, so concurrent sessions do not interfere. :param bridge_dir: The session's native-Codex bridge directory. :returns: The top-level ``model`` from ``config.toml`` (e.g. ``"gpt-5.4"``), or ``None`` when undeterminable. """ config_path = codex_home_for_bridge_dir(bridge_dir) / "config.toml" try: data = tomllib.loads(config_path.read_text()) except (OSError, tomllib.TOMLDecodeError): return None model = data.get("model") return model if isinstance(model, str) and model else None def write_bridge_state(bridge_dir: Path, state: CodexNativeBridgeState) -> None: """ Persist shared native Codex state atomically. :param bridge_dir: Native Codex bridge directory. :param state: State payload to persist. :returns: None. """ bridge_dir.mkdir(mode=0o700, parents=True, exist_ok=True) path = bridge_dir / _STATE_FILE fd, tmp_name = tempfile.mkstemp(prefix=f"{_STATE_FILE}.", dir=str(bridge_dir)) try: with os.fdopen(fd, "w", encoding="utf-8") as handle: json.dump( { "session_id": state.session_id, "socket_path": state.socket_path, "thread_id": state.thread_id, "codex_home": state.codex_home, "active_turn_id": state.active_turn_id, }, handle, sort_keys=True, ) handle.write("\n") os.replace(tmp_name, path) finally: if os.path.exists(tmp_name): os.unlink(tmp_name) def clear_bridge_state(bridge_dir: Path) -> None: """ Remove stale native Codex runtime state for a bridge directory. New app-server launches reuse the same bridge directory for a conversation id, but the old ``state.json`` may point at a thread from a previous app-server process. Clear it before starting the new server so web message forwarding waits for the new launch to publish its current transport and thread instead of injecting into stale state. :param bridge_dir: Native Codex bridge directory. :returns: None. """ for name in (_STATE_FILE, _STARTUP_ERROR_FILE, _MCP_STARTUP_FILE): try: (bridge_dir / name).unlink() except FileNotFoundError: continue def write_bridge_startup_error(bridge_dir: Path, message: str) -> None: """ Record why a native Codex app-server never started its thread (issue #59). :param bridge_dir: Native Codex bridge directory. :param message: Human-readable failure cause. :returns: None. """ try: bridge_dir.mkdir(mode=0o700, parents=True, exist_ok=True) path = bridge_dir / _STARTUP_ERROR_FILE fd, tmp_name = tempfile.mkstemp(prefix=f"{_STARTUP_ERROR_FILE}.", dir=str(bridge_dir)) try: with os.fdopen(fd, "w", encoding="utf-8") as handle: json.dump({"message": message}, handle, sort_keys=True) handle.write("\n") os.replace(tmp_name, path) finally: if os.path.exists(tmp_name): os.unlink(tmp_name) except OSError: return # best-effort; the real failure is already logged def read_bridge_startup_error(bridge_dir: Path) -> str | None: """ Read a recorded native Codex startup-failure message, if any. :param bridge_dir: Native Codex bridge directory. :returns: The recorded failure cause, or ``None`` if absent/unreadable. """ path = bridge_dir / _STARTUP_ERROR_FILE if not path.is_file(): return None try: raw = json.loads(path.read_text(encoding="utf-8")) except (OSError, json.JSONDecodeError): return None if not isinstance(raw, dict): return None message = raw.get("message") return message if isinstance(message, str) and message else None def read_mcp_startup(bridge_dir: Path) -> dict[str, dict[str, str | None]]: """ Read the recorded per-MCP-server startup state. :param bridge_dir: Native Codex bridge directory. :returns: Mapping of server name to its latest startup record, e.g. ``{"safe": {"status": "starting", "error": None}}``. Empty when no state has been recorded or the file is unreadable. """ path = bridge_dir / _MCP_STARTUP_FILE if not path.is_file(): return {} try: raw = json.loads(path.read_text(encoding="utf-8")) except (OSError, json.JSONDecodeError): return {} servers = raw.get("servers") if isinstance(raw, dict) else None if not isinstance(servers, dict): return {} parsed: dict[str, dict[str, str | None]] = {} for name, record in servers.items(): if not (isinstance(name, str) and name and isinstance(record, dict)): continue status = record.get("status") if status not in MCP_STARTUP_STATES: continue error = record.get("error") parsed[name] = { "status": status, "error": error if isinstance(error, str) and error else None, } return parsed def _write_mcp_startup(bridge_dir: Path, servers: dict[str, dict[str, str | None]]) -> None: """ Persist the per-MCP-server startup map atomically (best-effort). :param bridge_dir: Native Codex bridge directory. :param servers: Full startup map, e.g. ``{"safe": {"status": "ready", "error": None}}``. :returns: None. """ try: bridge_dir.mkdir(mode=0o700, parents=True, exist_ok=True) path = bridge_dir / _MCP_STARTUP_FILE fd, tmp_name = tempfile.mkstemp(prefix=f"{_MCP_STARTUP_FILE}.", dir=str(bridge_dir)) try: with os.fdopen(fd, "w", encoding="utf-8") as handle: json.dump({"servers": servers}, handle, sort_keys=True) handle.write("\n") os.replace(tmp_name, path) finally: if os.path.exists(tmp_name): os.unlink(tmp_name) except OSError: return # best-effort; surfacing MCP state must never sink startup def update_mcp_server_startup( bridge_dir: Path, name: str, status: str, error: str | None = None, ) -> dict[str, dict[str, str | None]]: """ Record one Codex MCP-server startup update. :param bridge_dir: Native Codex bridge directory. :param name: MCP server name, e.g. ``"storage-console"``. :param status: One of :data:`MCP_STARTUP_STATES`. :param error: Failure detail when ``status == "failed"``, e.g. ``"handshaking with MCP server failed"``. ``None`` otherwise. :returns: The full startup map after the update. """ servers = read_mcp_startup(bridge_dir) servers[name] = {"status": status, "error": error} _write_mcp_startup(bridge_dir, servers) return servers def pending_mcp_servers(servers: dict[str, dict[str, str | None]]) -> list[str]: """ Return the MCP servers still reported as ``starting``. :param servers: Startup map from :func:`read_mcp_startup`. :returns: Sorted server names whose latest status is ``starting``. """ return sorted( name for name, record in servers.items() if record.get("status") == MCP_STARTUP_STARTING ) def cancel_pending_mcp_startup(bridge_dir: Path) -> list[str]: """ Mark every still-``starting`` MCP server as ``cancelled``. Used by the Stop path so the executor's first-turn gate unblocks immediately, even when Codex's own ``cancelled`` notifications are delayed or lost. :param bridge_dir: Native Codex bridge directory. :returns: Sorted names of the servers that were flipped, e.g. ``["storage-console"]``. Empty when nothing was pending. """ servers = read_mcp_startup(bridge_dir) pending = pending_mcp_servers(servers) if not pending: return [] for name in pending: servers[name] = {"status": MCP_STARTUP_CANCELLED, "error": servers[name].get("error")} _write_mcp_startup(bridge_dir, servers) return pending def settle_pending_mcp_startup(bridge_dir: Path) -> tuple[dict[str, dict[str, str | None]], bool]: """ Drop every still-``starting`` MCP server from the recorded map. Codex delivers per-server terminal states (ready/failed) only to the connection that owns the thread — never to Omnigent's observer connection — so when a settle signal arrives (the thread went idle after a turn, or the startup window elapsed) the round is known to be over but the per-server outcomes are not. Unresolved entries are removed rather than guessed; locally-known terminal states (``cancelled`` from a Stop) are preserved. :param bridge_dir: Native Codex bridge directory. :returns: ``(map_after, changed)`` — the settled map and whether any entry was dropped. """ # The read→write below is not locked across processes: a runner Stop # can flip an entry to ``cancelled`` in between, and this write drops # it. Cosmetic only — both outcomes end the round, and the Stop path # publishes its cancelled map independently. servers = read_mcp_startup(bridge_dir) pending = pending_mcp_servers(servers) if not pending: return servers, False for name in pending: servers.pop(name, None) _write_mcp_startup(bridge_dir, servers) return servers, True def mcp_startup_waiting_detail(servers: dict[str, dict[str, str | None]]) -> str | None: """ Describe the MCP servers a startup wait is still blocked on. :param servers: Startup map from :func:`read_mcp_startup`. :returns: Text naming the pending servers, e.g. ``"MCP startup still waiting on storage-console"``, or ``None`` when nothing is pending. """ pending = pending_mcp_servers(servers) if not pending: return None return f"MCP startup still waiting on {', '.join(pending)}" def read_bridge_state(bridge_dir: Path) -> CodexNativeBridgeState | None: """ Read shared native Codex bridge state. :param bridge_dir: Native Codex bridge directory. :returns: Parsed state, or ``None`` when no state exists. """ path = bridge_dir / _STATE_FILE if not path.is_file(): return None try: raw = json.loads(path.read_text(encoding="utf-8")) except (OSError, json.JSONDecodeError): return None if not isinstance(raw, dict): return None session_id = raw.get("session_id") socket_path = raw.get("socket_path") thread_id = raw.get("thread_id") codex_home = raw.get("codex_home") active_turn_id = raw.get("active_turn_id") required = (session_id, socket_path, thread_id, codex_home) if not all(isinstance(value, str) and value for value in required): return None parsed_active_turn_id = ( active_turn_id if isinstance(active_turn_id, str) and active_turn_id else None ) return CodexNativeBridgeState( session_id=session_id, socket_path=socket_path, thread_id=thread_id, codex_home=codex_home, active_turn_id=parsed_active_turn_id, ) def update_active_turn_id(bridge_dir: Path, active_turn_id: str | None) -> None: """ Update the active Codex turn id in bridge state. :param bridge_dir: Native Codex bridge directory. :param active_turn_id: Active turn id, e.g. ``"turn_abc123"``, or ``None`` when no turn is running. :returns: None. """ state = read_bridge_state(bridge_dir) if state is None: return write_bridge_state( bridge_dir, CodexNativeBridgeState( session_id=state.session_id, socket_path=state.socket_path, thread_id=state.thread_id, codex_home=state.codex_home, active_turn_id=active_turn_id, ), ) def update_thread_id(bridge_dir: Path, thread_id: str, active_turn_id: str | None = None) -> None: """ Update the Codex thread id in bridge state. Used when a native Codex action creates a fresh thread while the Omnigent session stays the same. :param bridge_dir: Native Codex bridge directory. :param thread_id: New Codex thread id, e.g. ``"thread_abc123"``. :param active_turn_id: Active turn id for the new thread, e.g. ``"turn_abc123"``, or ``None`` when no turn is running yet. :returns: None. """ state = read_bridge_state(bridge_dir) if state is None: return write_bridge_state( bridge_dir, CodexNativeBridgeState( session_id=state.session_id, socket_path=state.socket_path, thread_id=thread_id, codex_home=state.codex_home, active_turn_id=active_turn_id, ), ) def clear_active_turn_id_if_matches(bridge_dir: Path, completed_turn_id: str | None) -> bool: """ Clear the active Codex turn id if a terminal event matches it. Terminal Codex notifications can race with a newer ``turn/started`` notification under rapid web sends. A stale terminal event must not erase the newer active turn id, or later web messages stop steering the running native Codex turn. A terminal event without a turn id is ambiguous: it cannot be correlated to the active turn, so when a turn is live it is ignored rather than clearing the turn. Clearing it would post a premature ``idle`` to the session and hide the "working" spinner while Codex is still mid-turn. (Codex includes a turn id on real terminal events; the id-less shape is a legacy/malformed edge case.) :param bridge_dir: Native Codex bridge directory. :param completed_turn_id: Completed or failed turn id, e.g. ``"turn_abc123"``. ``None`` means Codex did not include an id; if a turn is live it is left intact (returns ``False``), and if no turn is live the call is a no-op (returns ``True``). :returns: ``True`` when bridge state was cleared or did not exist, ``False`` when a stale or ambiguous terminal event was ignored. """ state = read_bridge_state(bridge_dir) if state is None: return True if completed_turn_id is None: # No-id terminal mid-turn is ambiguous — ignore (clearing posts a premature idle). if state.active_turn_id is not None: return False elif state.active_turn_id != completed_turn_id: return False write_bridge_state( bridge_dir, CodexNativeBridgeState( session_id=state.session_id, socket_path=state.socket_path, thread_id=state.thread_id, codex_home=state.codex_home, active_turn_id=None, ), ) return True