"""Codex-specific MCP transport helpers."""
from __future__ import annotations
import atexit
import collections
import json
import re
import shutil
import tempfile
import tomllib
from pathlib import Path
from typing import cast
from loguru import logger
from ralph.mcp.tools.names import (
CODEX_NATIVE_FEATURE_OVERRIDES,
RALPH_MCP_SERVER_NAME,
)
from ralph.mcp.transport.common import merge_existing_upstreams
from ralph.mcp.upstream.config import UpstreamMcpServer, normalize_upstream_mcp_servers
#: Sane upper bound for the in-process Codex-home registry. The deque
#: provides FIFO eviction so a long-lived process that spawns many
#: codex invocations cannot grow the in-memory registry unboundedly.
#: On-disk bound is provided by:
#: - ``release_codex_home``: the production release path invoked from
#: the per-invocation ``cleanup`` hook on
#: ``ResolvedInvocationRuntime`` (see
#: ``ralph.agents.invoke._runtime_resolvers.CodexRuntimeResolver``)
#: in ``invoke_agent``'s ``finally`` block.
#: - ``cleanup_codex_homes``: the atexit-registered net that rmtree's
#: any orphan homes a crashed interpreter left behind.
#: An earlier version of this module used ``_allocate_codex_home_dir``
#: to rmtree the FIFO-evicted oldest entry on every append past the
#: cap, but that could delete a home that was still active in a live
#: Codex subprocess (analysis feedback wt-024 round 2). The active-home
#: invariant is now: eviction only REMOVES from the registry; the
#: on-disk directory survives until ``release_codex_home`` is called
#: by the owner or ``atexit`` reaps it.
_DEFAULT_CODEX_HOME_CAP: int = 64
# bounded-accumulator-ok: deque(maxlen=_DEFAULT_CODEX_HOME_CAP) provides
# real FIFO eviction of the in-memory registry. The on-disk bound is
# provided by ``release_codex_home`` (production release path via the
# ResolvedInvocationRuntime.cleanup hook) and ``cleanup_codex_homes``
# (atexit net); see _allocate_codex_home_dir for the active-home
# invariant that prevents evicting a home whose subprocess is still
# running.
_allocated_codex_homes: collections.deque[str] = collections.deque(maxlen=_DEFAULT_CODEX_HOME_CAP) # bounded-accumulator-ok # noqa: E501 # type: ignore[var-annotated] # reason: autogenerated code has no type support, see docs/agents/type-ignore-policy.md#autogenerated-code
#: Set of every Codex home ever allocated by this process (FIFO-evicted
#: homes from ``_allocated_codex_homes`` are still tracked here so the
#: ``cleanup_codex_homes`` atexit net can find and rmtree them, even
## after the bookkeeping deque has wrapped past them).
#:
#: Without this set, ``cleanup_codex_homes`` only iterates the bounded
#: deque and misses any FIFO-evicted home whose owning session crashed
#: before calling ``release_codex_home`` (analysis-feedback wt-024
#: round 3). The set lives for the lifetime of the process so a
#: crashed interpreter can still reap all of its orphan homes on
#: shutdown.
_all_allocated_codex_homes: set[str] = set() # bounded-accumulator-ok: lifetime tracking
[docs]
def cleanup_codex_homes() -> None:
"""Remove every Codex home dir this process ever allocated.
Iterates ``_all_allocated_codex_homes`` (NOT the bounded
``_allocated_codex_homes`` deque) so FIFO-evicted homes are also
reaped on interpreter shutdown. Earlier versions of this function
only iterated ``_allocated_codex_homes``, which (after the
analysis-feedback wt-024 round 2 active-home fix changed
``_allocate_codex_home_dir`` to NOT rmtree on FIFO eviction)
meant homes evicted from the bookkeeping deque could leak past
``atexit``. The fix is to maintain a separate ``_all_allocated_codex_homes``
set that tracks every allocation regardless of deque membership;
``release_codex_home`` discards from BOTH collections.
Standalone importable function so tests can invoke cleanup
directly without depending on ``atexit`` timing. ``ignore_errors``
makes the function robust to partial interpreter shutdown and
already-removed dirs.
"""
for home in list(_all_allocated_codex_homes):
shutil.rmtree(home, ignore_errors=True)
_all_allocated_codex_homes.clear()
_allocated_codex_homes.clear()
[docs]
def release_codex_home(codex_home: str) -> bool:
"""Release a single Codex home during the normal runtime lifecycle.
Removes the path from both ``_all_allocated_codex_homes`` (so the
atexit net will not re-rrmtree it) AND ``_allocated_codex_homes``
(the bounded deque). Rmtree's the on-disk directory with
``ignore_errors=True`` so the operation is idempotent.
Returns ``True`` if the home was in the bounded deque (active
bookkeeping member); ``False`` if the home was not in the deque
(already released, never registered, or FIFO-evicted before this
call). The returned boolean preserves the documented contract
used by callers that want to know whether they were the first
releaser; an evicted-but-unreleased home returns ``False``
even though the on-disk rmtree still happens.
This is the production release path: callers that allocate a
Codex home, use it for a bounded operation (e.g. the MCP probe
synthesizes a config + runs a handshake, then has no further
use for the home), and want to release it BEFORE interpreter
shutdown MUST call this function. ``atexit``-only cleanup leaves
every allocated home on disk for the entire interpreter lifetime
and grows the registry unboundedly across a long run.
"""
_all_allocated_codex_homes.discard(codex_home)
try:
_allocated_codex_homes.remove(codex_home)
except ValueError:
shutil.rmtree(codex_home, ignore_errors=True)
return False
shutil.rmtree(codex_home, ignore_errors=True)
return True
atexit.register(cleanup_codex_homes)
[docs]
def prepare_codex_home(
endpoint: str | None,
*,
workspace_path: Path | None,
existing_home: str | None,
system_prompt_file: str | None,
unsafe_mode: bool = False,
) -> str:
"""Prepare an isolated Codex home directory and return its path."""
codex_home, _upstreams = prepare_codex_home_with_upstreams(
endpoint,
workspace_path=workspace_path,
existing_home=existing_home,
system_prompt_file=system_prompt_file,
unsafe_mode=unsafe_mode,
)
return codex_home
def _flat_dict_to_toml_servers(flat_dict: dict[str, object]) -> str:
"""Convert a flat dict with mcp_servers.X keys to TOML server sections."""
lines: list[str] = []
for key, value in sorted(flat_dict.items()):
if not isinstance(key, str) or not key.startswith("mcp_servers."):
continue
lines.append(f"[{key}]")
if isinstance(value, dict):
for k, v in sorted(value.items()):
lines.append(f"{k} = {json.dumps(v)}")
return "\n".join(lines)
[docs]
def prepare_codex_home_with_upstreams(
endpoint: str | None,
*,
workspace_path: Path | None,
existing_home: str | None,
system_prompt_file: str | None,
unsafe_mode: bool = False,
) -> tuple[str, tuple[UpstreamMcpServer, ...]]:
"""Prepare an isolated Codex home directory and return its path with upstream servers."""
codex_root = _allocate_codex_home_dir(workspace_path)
codex_root.mkdir(parents=True, exist_ok=True)
source_home = Path(existing_home).expanduser() if existing_home else Path.home() / ".codex"
if source_home.exists():
_mirror_codex_home(source_home, codex_root)
source_config = source_home / "config.toml"
base_config = source_config.read_text(encoding="utf-8") if source_config.exists() else ""
upstreams = _extract_codex_upstream_servers(base_config)
prefix_sections: list[str] = []
appended_sections: list[str] = []
if endpoint:
logger.warning(
"Codex MCP tool restriction is best-effort: apply_patch and core "
"editing primitives cannot be disabled. See "
"ralph-workflow/docs/mcp-tool-restriction.md."
)
existing_from_base: dict[str, object] = {}
if base_config.strip():
try:
parsed: object = tomllib.loads(base_config)
if isinstance(parsed, dict):
existing_from_base = {
key: value
for key, value in parsed.items()
if isinstance(key, str) and key.startswith("mcp_servers.")
}
except Exception:
pass
if not unsafe_mode:
base_config = _remove_all_toml_mcp_server_tables(base_config)
merged = merge_existing_upstreams(
"codex",
existing_from_base,
unsafe_mode=unsafe_mode,
)
merged_toml = _flat_dict_to_toml_servers(merged)
if merged_toml:
appended_sections.append(merged_toml + "\n")
ralph_section = (
f'[mcp_servers.{RALPH_MCP_SERVER_NAME}]\nurl = "{endpoint}"\nenabled = true\n'
)
if ralph_section.strip() not in merged_toml:
appended_sections.append(ralph_section)
features_in_base = "[features]" in base_config
feature_lines = [
f"{key.split('.', 1)[1]} = {value}" for key, value in CODEX_NATIVE_FEATURE_OVERRIDES
]
feature_block = "\n".join(feature_lines) + "\n"
if features_in_base:
base_config = base_config.replace("[features]\n", "[features]\n" + feature_block, 1)
if not features_in_base:
appended_sections.append("[features]\n" + feature_block)
if system_prompt_file:
prefix_sections.append(f"model_instructions_file = {json.dumps(system_prompt_file)}\n")
config_suffix = "\n".join(section.rstrip() for section in appended_sections if section.strip())
prefix_text = "\n".join(section.rstrip() for section in prefix_sections if section.strip())
config_text = "\n\n".join(
part for part in [prefix_text, base_config.rstrip(), config_suffix] if part
)
(codex_root / "config.toml").write_text(config_text, encoding="utf-8")
return str(codex_root), upstreams
def _remove_toml_table(config_text: str, table_name: str) -> str:
pattern = re.compile(
rf"(?ms)^\[{re.escape(table_name)}\]\n.*?(?=^\[|\Z)",
)
return pattern.sub("", config_text).strip()
def _remove_all_toml_mcp_server_tables(config_text: str) -> str:
pattern = re.compile(r"(?ms)^\[mcp_servers(?:\.[^\]]+)?\]\n.*?(?=^\[|\Z)")
return pattern.sub("", config_text).strip()
def _mirror_codex_home(source_home: Path, codex_root: Path) -> None:
for entry in source_home.iterdir():
if entry.name == "config.toml":
continue
destination = codex_root / entry.name
try:
destination.symlink_to(entry, target_is_directory=entry.is_dir())
except OSError:
if entry.is_dir():
shutil.copytree(entry, destination, dirs_exist_ok=True)
else:
shutil.copy2(entry, destination)
def _allocate_codex_home_dir(workspace_path: Path | None) -> Path:
if workspace_path is None:
codex_root = Path(tempfile.mkdtemp(prefix="ralph-codex-home-"))
else:
tmp_root = workspace_path / ".agent" / "tmp"
tmp_root.mkdir(parents=True, exist_ok=True)
codex_root = Path(tempfile.mkdtemp(prefix="codex-home-", dir=str(tmp_root)))
# The deque is bounded; when appending past the cap, FIFO eviction
# drops the oldest entry from the in-memory registry so the
# registry itself cannot grow unboundedly. We deliberately do
# NOT rmtree the evicted entry on-disk: the evicted home may
# still be ACTIVE (its owning Codex subprocess is still running).
# Earlier code performed the rmtree here, but analysis feedback
# wt-024 round 2 found that this could delete the
# CODEX_HOME directory out from under a running Codex agent when
# more than ``_DEFAULT_CODEX_HOME_CAP`` concurrent invocations
# occurred. The on-disk bound is now provided by:
# 1. ``release_codex_home`` invoked from the
# ``ResolvedInvocationRuntime.cleanup`` hook (the production
# release path) when the owning subprocess finishes; AND
# 2. ``cleanup_codex_homes`` (the atexit net, iterating
# ``_all_allocated_codex_homes`` so FIFO-evicted homes are
# still reaped) for orphans.
# Use ``_allocated_codex_homes.maxlen`` (not the module-level
# ``_DEFAULT_CODEX_HOME_CAP``) so a test that swaps the deque with
# a smaller-maxlen one (see
# ``tests/integration/test_codex_home_release_path.py``) still
# bounds the registry at the test's cap. Tests MUST NOT be able
# to opt out of the registry bound by shrinking the cap.
maxlen = _allocated_codex_homes.maxlen
if maxlen is not None and len(_allocated_codex_homes) >= maxlen:
# FIFO eviction from the REGISTRY ONLY. The on-disk directory
# is intentionally left in place; it is rmtree'd only when
# ``release_codex_home`` (production release path) or
# ``cleanup_codex_homes`` (atexit net) runs. The evicted path
# remains in ``_all_allocated_codex_homes`` so the atexit net
# can still find and rmtree it if no owner releases it first.
_allocated_codex_homes.popleft()
home_str = str(codex_root)
_allocated_codex_homes.append(home_str)
_all_allocated_codex_homes.add(home_str)
return codex_root
def _extract_codex_upstream_servers(config_text: str) -> tuple[UpstreamMcpServer, ...]:
if not config_text.strip():
return ()
try:
parsed: object = tomllib.loads(config_text)
except Exception:
return ()
if not isinstance(parsed, dict):
return ()
mcp_servers = parsed.get("mcp_servers")
if not isinstance(mcp_servers, dict):
return ()
return normalize_upstream_mcp_servers(cast("dict[str, object]", mcp_servers))
__all__ = [
"cleanup_codex_homes",
"prepare_codex_home",
"prepare_codex_home_with_upstreams",
"release_codex_home",
]