mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
feat(codex): keep wrap routing session-scoped (#1507)
## Description Keeps Codex wrap routing session-scoped so routing state from one wrapped session does not leak into another. ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - Scope Codex wrap routing state to the active session. - Avoid cross-session routing contamination for wrapped Codex traffic. - Keep changes focused on wrap/proxy routing behavior. ## Testing - [x] Unit tests pass - [x] New tests added for new functionality when applicable - [x] Manual testing performed ### Test Output ```text Focused tests/review were completed before this governance body cleanup. The current branch is conflicted and still needs merge resolution before final merge readiness. ``` ## Real Behavior Proof - Environment: Headroom development/review context. - Exact command / steps: Reviewed session-scoped Codex wrap routing behavior and existing focused coverage. - Observed result: Routing state is scoped to the active wrap session rather than shared globally across sessions. - Not tested: Current conflicted branch after merge resolution; conflicts still need to be resolved before merge. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Additional Notes This body was normalized by a maintainer after approval so the governance parser reflects the already-reviewed PR state. The PR remains blocked by merge conflicts.
This commit is contained in:
Rod Boev
GitHub
parent
b0440f958d
commit
ad9d086f43
3 changed files with 486 additions and 316 deletions
|
|
@ -230,6 +230,7 @@ def create_shims(shim_dir: Path) -> None:
|
|||
"OPENAI_BASE_URL",
|
||||
"OPENAI_API_BASE",
|
||||
"ANTHROPIC_BASE_URL",
|
||||
"CODEX_HOME",
|
||||
"OPENCODE_CONFIG_CONTENT",
|
||||
)
|
||||
if os.environ.get(key) is not None
|
||||
|
|
@ -286,6 +287,7 @@ def create_shims(shim_dir: Path) -> None:
|
|||
"OPENAI_BASE_URL",
|
||||
"OPENAI_API_BASE",
|
||||
"ANTHROPIC_BASE_URL",
|
||||
"CODEX_HOME",
|
||||
"OPENCODE_CONFIG_CONTENT",
|
||||
)
|
||||
if os.environ.get(key) is not None
|
||||
|
|
@ -351,6 +353,13 @@ def create_shims(shim_dir: Path) -> None:
|
|||
by_model.get(model_name) if isinstance(by_model, dict) else None
|
||||
)
|
||||
|
||||
codex_home = os.environ.get("CODEX_HOME")
|
||||
if codex_home:
|
||||
session_config = Path(codex_home) / "config.toml"
|
||||
record["session_config"] = (
|
||||
session_config.read_text(encoding="utf-8") if session_config.exists() else None
|
||||
)
|
||||
|
||||
record["probes"] = probes
|
||||
|
||||
with (log_dir / f"{tool}.jsonl").open("a", encoding="utf-8") as handle:
|
||||
|
|
@ -571,32 +580,50 @@ def verify_codex_wrap(
|
|||
)
|
||||
|
||||
config_path = Path(base_env["HOME"]) / ".codex" / "config.toml"
|
||||
assert_true(config_path.exists(), "Codex wrap should create ~/.codex/config.toml")
|
||||
config = config_path.read_text(encoding="utf-8")
|
||||
assert_true(
|
||||
f'openai_base_url = "http://127.0.0.1:{port}/v1"' in config,
|
||||
"Codex wrap should inject openai_base_url for subscription routing",
|
||||
)
|
||||
assert_true(
|
||||
f'base_url = "http://127.0.0.1:{port}/v1"' in config,
|
||||
"Codex wrap should inject the headroom provider base_url",
|
||||
)
|
||||
assert_true(
|
||||
'env_key = "OPENAI_API_KEY"' not in config,
|
||||
"Codex wrap should preserve OAuth and never inject env_key",
|
||||
)
|
||||
# Bug 3 (#406): requires_openai_auth must be absent from headroom provider blocks.
|
||||
assert_true(
|
||||
"requires_openai_auth" not in config,
|
||||
"Codex wrap must NOT inject requires_openai_auth into the headroom provider block",
|
||||
)
|
||||
assert_true(
|
||||
"supports_websockets = true" in config, "Codex wrap missing 'supports_websockets = true'"
|
||||
not config_path.exists(),
|
||||
"Codex wrap should leave ~/.codex/config.toml untouched during normal launch",
|
||||
)
|
||||
|
||||
entries = read_jsonl(log_dir / "codex.jsonl")
|
||||
assert_true(len(entries) > 0, "Codex shim should have been invoked")
|
||||
env_vars = entries[-1]["env"]
|
||||
session_home = env_vars.get("CODEX_HOME")
|
||||
assert_true(
|
||||
isinstance(session_home, str) and session_home,
|
||||
"Codex wrap should launch the child with a session-scoped CODEX_HOME",
|
||||
)
|
||||
assert_true(
|
||||
Path(session_home) != config_path.parent,
|
||||
"Codex wrap should not point the child at the real ~/.codex home",
|
||||
)
|
||||
config = entries[-1].get("session_config")
|
||||
assert_true(
|
||||
isinstance(config, str) and config,
|
||||
"Codex wrap should capture the session-scoped config during launch",
|
||||
)
|
||||
assert_true(
|
||||
f'openai_base_url = "http://127.0.0.1:{port}/v1"' in config,
|
||||
"Codex wrap should inject openai_base_url into the session config",
|
||||
)
|
||||
assert_true(
|
||||
f'base_url = "http://127.0.0.1:{port}/v1"' in config,
|
||||
"Codex wrap should inject the headroom provider base_url into the session config",
|
||||
)
|
||||
assert_true(
|
||||
'env_key = "OPENAI_API_KEY"' not in config,
|
||||
"Codex wrap should preserve OAuth and never inject env_key into the session config",
|
||||
)
|
||||
# Bug 3 (#406): requires_openai_auth must be absent from headroom provider blocks.
|
||||
assert_true(
|
||||
"requires_openai_auth" not in config,
|
||||
"Codex wrap must NOT inject requires_openai_auth into the session headroom provider block",
|
||||
)
|
||||
assert_true(
|
||||
"supports_websockets = true" in config,
|
||||
"Codex wrap missing 'supports_websockets = true' in the session config",
|
||||
)
|
||||
|
||||
assert_true(
|
||||
env_vars.get("OPENAI_BASE_URL") == f"http://127.0.0.1:{port}/v1",
|
||||
"Codex wrap should set OPENAI_BASE_URL",
|
||||
|
|
|
|||
|
|
@ -30,6 +30,7 @@ import tempfile
|
|||
import time
|
||||
import urllib.parse
|
||||
from collections.abc import Callable
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
from typing import Any, cast
|
||||
|
||||
|
|
@ -1495,6 +1496,27 @@ def _codex_home_dir() -> Path:
|
|||
return Path.home() / ".codex"
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _codex_session_home_overlay() -> Any:
|
||||
"""Seed a temporary Codex home from the active home and point the process at it."""
|
||||
source_home = _codex_home_dir()
|
||||
original_codex_home = os.environ.get("CODEX_HOME")
|
||||
|
||||
with tempfile.TemporaryDirectory(prefix="headroom-codex-home-") as tmp_dir:
|
||||
session_home = Path(tmp_dir)
|
||||
if source_home.exists():
|
||||
shutil.copytree(source_home, session_home, dirs_exist_ok=True)
|
||||
|
||||
os.environ["CODEX_HOME"] = str(session_home)
|
||||
try:
|
||||
yield session_home
|
||||
finally:
|
||||
if original_codex_home is None:
|
||||
os.environ.pop("CODEX_HOME", None)
|
||||
else:
|
||||
os.environ["CODEX_HOME"] = original_codex_home
|
||||
|
||||
|
||||
def _codex_config_paths() -> tuple[Path, Path]:
|
||||
"""Return ``(config_file, backup_file)`` paths for the Codex TOML config."""
|
||||
config_dir = _codex_home_dir()
|
||||
|
|
@ -4537,6 +4559,200 @@ def unwrap_copilot(port: int, no_stop_proxy: bool) -> None:
|
|||
# =============================================================================
|
||||
|
||||
|
||||
def _prepare_codex_wrap_state(
|
||||
*,
|
||||
port: int,
|
||||
no_rtk: bool,
|
||||
no_mcp: bool,
|
||||
no_tokensave: bool,
|
||||
serena: bool,
|
||||
no_serena: bool,
|
||||
memory: bool,
|
||||
verbose: bool,
|
||||
rtk_home: Path | None = None,
|
||||
) -> None:
|
||||
"""Prepare the active Codex home for a wrap or prepare-only invocation."""
|
||||
# Snapshot Codex config.toml BEFORE any wrap-time mutation so
|
||||
# `headroom unwrap codex` can restore the user's pre-wrap state
|
||||
# byte-for-byte. The snapshot is a no-op if the backup already exists
|
||||
# or if the file already has Headroom markers, so this is safe to
|
||||
# call repeatedly. Crucially this must run before MCP install, which
|
||||
# writes its marker block to the same file.
|
||||
_codex_config_file, _codex_backup_file = _codex_config_paths()
|
||||
_snapshot_codex_config_if_unwrapped(_codex_config_file, _codex_backup_file)
|
||||
|
||||
# Setup CLI context tool for Codex.
|
||||
if not no_rtk:
|
||||
if _selected_context_tool() == _CONTEXT_TOOL_LEAN_CTX:
|
||||
click.echo(" Setting up lean-ctx for Codex...")
|
||||
_setup_lean_ctx_agent("codex", verbose=verbose)
|
||||
else:
|
||||
click.echo(" Setting up rtk for Codex...")
|
||||
rtk_path = _ensure_rtk_binary(verbose=verbose)
|
||||
if rtk_path:
|
||||
# Keep RTK guidance local to the user's Codex configuration.
|
||||
global_agents = (rtk_home or _codex_home_dir()) / "AGENTS.md"
|
||||
_inject_rtk_instructions(global_agents, verbose=verbose)
|
||||
|
||||
# Register headroom MCP server in Codex config.toml so Codex can
|
||||
# call headroom_retrieve on compression markers from the proxy.
|
||||
if not no_mcp:
|
||||
from headroom.mcp_registry import CodexRegistrar
|
||||
|
||||
# Codex starts a long-lived local MCP subprocess from config.toml.
|
||||
# If a previous wrap used another port, retrieval can silently point
|
||||
# at the wrong proxy while model traffic uses the right one.
|
||||
_setup_headroom_mcp(CodexRegistrar(), port, verbose=verbose, force=True)
|
||||
elif verbose:
|
||||
click.echo(" Skipping MCP retrieve tool (--no-mcp)")
|
||||
|
||||
# Coding-task compressor: tokensave primary, Serena backup. Codex starts
|
||||
# long-lived MCP subprocesses from config.toml, so force re-registration.
|
||||
from headroom.mcp_registry import CodexRegistrar
|
||||
|
||||
_setup_coding_compressor(
|
||||
CodexRegistrar(),
|
||||
serena_context="codex",
|
||||
serena=serena,
|
||||
no_serena=no_serena,
|
||||
no_tokensave=no_tokensave,
|
||||
verbose=verbose,
|
||||
force=True,
|
||||
)
|
||||
|
||||
# Setup memory MCP server for Codex (native tool integration)
|
||||
if memory:
|
||||
click.echo(" Setting up memory for Codex...")
|
||||
mem_dir = Path.cwd() / ".headroom"
|
||||
mem_dir.mkdir(parents=True, exist_ok=True)
|
||||
db_path = str(mem_dir / "memory.db")
|
||||
mem_user = os.environ.get("USER", os.environ.get("USERNAME", "default"))
|
||||
|
||||
# Register MCP server in Codex config
|
||||
_inject_memory_mcp_config(mem_user)
|
||||
|
||||
# Inject memory guidance into project AGENTS.md
|
||||
agents_md = Path.cwd() / "AGENTS.md"
|
||||
_inject_memory_agents_md(agents_md)
|
||||
|
||||
# Sync Claude's memories → DB so MCP search finds them
|
||||
try:
|
||||
import asyncio
|
||||
|
||||
from headroom.memory.sync import _build_sync_backend, sync_import
|
||||
from headroom.memory.sync_adapters.claude_code import (
|
||||
ClaudeCodeAdapter,
|
||||
get_claude_memory_dir,
|
||||
)
|
||||
|
||||
claude_memory_dir = get_claude_memory_dir()
|
||||
|
||||
async def _import_claude_memories() -> int:
|
||||
backend = _build_sync_backend(db_path)
|
||||
await backend._ensure_initialized()
|
||||
adapter = ClaudeCodeAdapter(claude_memory_dir)
|
||||
count = await sync_import(backend, adapter, mem_user)
|
||||
await backend.close()
|
||||
return count
|
||||
|
||||
imported = asyncio.run(_import_claude_memories())
|
||||
if imported:
|
||||
click.echo(f" Memory: imported {imported} memories from Claude")
|
||||
except Exception as e:
|
||||
click.echo(f" Warning: Claude memory import failed: {e}")
|
||||
|
||||
# Inject Headroom provider into Codex config so WebSocket traffic also
|
||||
# routes through the proxy. Codex ignores OPENAI_BASE_URL for its WS
|
||||
# transport unless a custom provider declares supports_websockets = true.
|
||||
# NOTE: this must run BEFORE _inject_memory_mcp_config because it rewrites
|
||||
# the config file. Re-inject MCP config after if memory is enabled.
|
||||
_inject_codex_provider_config(port)
|
||||
if memory:
|
||||
_inject_memory_mcp_config(os.environ.get("USER", os.environ.get("USERNAME", "default")))
|
||||
|
||||
|
||||
def _run_codex_wrap(
|
||||
*,
|
||||
port: int,
|
||||
no_rtk: bool,
|
||||
no_mcp: bool,
|
||||
no_tokensave: bool,
|
||||
serena: bool,
|
||||
no_serena: bool,
|
||||
code_graph: bool,
|
||||
no_proxy: bool,
|
||||
learn: bool,
|
||||
memory: bool,
|
||||
backend: str | None,
|
||||
anyllm_provider: str | None,
|
||||
region: str | None,
|
||||
verbose: bool,
|
||||
prepare_only: bool,
|
||||
codex_args: tuple,
|
||||
) -> None:
|
||||
"""Execute the Codex wrap flow with the session overlay when launching."""
|
||||
if prepare_only:
|
||||
_prepare_codex_wrap_state(
|
||||
port=port,
|
||||
no_rtk=no_rtk,
|
||||
no_mcp=no_mcp,
|
||||
no_tokensave=no_tokensave,
|
||||
serena=serena,
|
||||
no_serena=no_serena,
|
||||
memory=memory,
|
||||
verbose=verbose,
|
||||
)
|
||||
return
|
||||
|
||||
codex_bin = shutil.which("codex")
|
||||
if not codex_bin:
|
||||
click.echo("Error: 'codex' not found in PATH.")
|
||||
click.echo("Install Codex CLI: npm install -g @openai/codex")
|
||||
raise SystemExit(1)
|
||||
|
||||
active_codex_home = _codex_home_dir()
|
||||
with _codex_session_home_overlay() as session_codex_home:
|
||||
_prepare_codex_wrap_state(
|
||||
port=port,
|
||||
no_rtk=no_rtk,
|
||||
no_mcp=no_mcp,
|
||||
no_tokensave=no_tokensave,
|
||||
serena=serena,
|
||||
no_serena=no_serena,
|
||||
memory=memory,
|
||||
verbose=verbose,
|
||||
rtk_home=active_codex_home,
|
||||
)
|
||||
|
||||
env, env_vars_display = _build_codex_launch_env(port, os.environ)
|
||||
|
||||
# Per-project savings attribution: the injected provider config maps the
|
||||
# X-Headroom-Project header to HEADROOM_PROJECT via env_http_headers, so
|
||||
# Codex sends it only when this var is set. A user-set value wins.
|
||||
_codex_project = _project_name_from_cwd()
|
||||
if _codex_project and "HEADROOM_PROJECT" not in env:
|
||||
env["HEADROOM_PROJECT"] = _codex_project
|
||||
|
||||
env["CODEX_HOME"] = str(session_codex_home)
|
||||
|
||||
_launch_tool(
|
||||
binary=codex_bin,
|
||||
args=codex_args,
|
||||
env=env,
|
||||
port=port,
|
||||
no_proxy=no_proxy,
|
||||
tool_label="CODEX",
|
||||
env_vars_display=env_vars_display,
|
||||
learn=learn,
|
||||
memory=memory,
|
||||
agent_type="codex",
|
||||
code_graph=code_graph,
|
||||
backend=backend,
|
||||
anyllm_provider=anyllm_provider,
|
||||
region=region,
|
||||
)
|
||||
|
||||
|
||||
@wrap.command(context_settings={"ignore_unknown_options": True})
|
||||
@click.option(
|
||||
"--port", "-p", default=8787, type=click.IntRange(1, 65535), help="Proxy port (default: 8787)"
|
||||
|
|
@ -4630,258 +4846,25 @@ def codex(
|
|||
headroom wrap codex --port 9999 # Custom proxy port
|
||||
headroom wrap codex --backend anyllm --anyllm-provider groq
|
||||
"""
|
||||
# Snapshot Codex config.toml BEFORE any wrap-time mutation so
|
||||
# `headroom unwrap codex` can restore the user's pre-wrap state
|
||||
# byte-for-byte. The snapshot is a no-op if the backup already exists
|
||||
# or if the file already has Headroom markers, so this is safe to
|
||||
# call repeatedly. Crucially this must run before MCP install, which
|
||||
# writes its marker block to the same file.
|
||||
_codex_config_file, _codex_backup_file = _codex_config_paths()
|
||||
_snapshot_codex_config_if_unwrapped(_codex_config_file, _codex_backup_file)
|
||||
|
||||
# Non-port-dependent setup first (RTK, etc.).
|
||||
if not no_rtk:
|
||||
if _selected_context_tool() == _CONTEXT_TOOL_LEAN_CTX:
|
||||
click.echo(" Setting up lean-ctx for Codex...")
|
||||
_setup_lean_ctx_agent("codex", verbose=verbose)
|
||||
else:
|
||||
click.echo(" Setting up rtk for Codex...")
|
||||
rtk_path = _ensure_rtk_binary(verbose=verbose)
|
||||
if rtk_path:
|
||||
# Keep RTK guidance local to the user's Codex configuration.
|
||||
global_agents = _codex_home_dir() / "AGENTS.md"
|
||||
_inject_rtk_instructions(global_agents, verbose=verbose)
|
||||
|
||||
# --prepare-only: only update Codex config, do NOT start proxy.
|
||||
# MCP/memory/provider config are all config-file writes — they don't
|
||||
# need a running proxy. Use the raw requested port (no health check,
|
||||
# no port fallback) since the user will run the full command later.
|
||||
if prepare_only:
|
||||
if not no_mcp:
|
||||
from headroom.mcp_registry import CodexRegistrar
|
||||
|
||||
_setup_headroom_mcp(CodexRegistrar(), port, verbose=verbose, force=True)
|
||||
elif verbose:
|
||||
click.echo(" Skipping MCP retrieve tool (--no-mcp)")
|
||||
|
||||
from headroom.mcp_registry import CodexRegistrar
|
||||
|
||||
_setup_coding_compressor(
|
||||
CodexRegistrar(),
|
||||
serena_context="codex",
|
||||
serena=serena,
|
||||
no_serena=no_serena,
|
||||
no_tokensave=no_tokensave,
|
||||
verbose=verbose,
|
||||
force=True,
|
||||
)
|
||||
|
||||
if memory:
|
||||
click.echo(" Setting up memory for Codex...")
|
||||
mem_dir = Path.cwd() / ".headroom"
|
||||
mem_dir.mkdir(parents=True, exist_ok=True)
|
||||
db_path = str(mem_dir / "memory.db")
|
||||
mem_user = os.environ.get("USER", os.environ.get("USERNAME", "default"))
|
||||
_inject_memory_mcp_config(mem_user)
|
||||
agents_md = Path.cwd() / "AGENTS.md"
|
||||
_inject_memory_agents_md(agents_md)
|
||||
|
||||
# Sync Claude's memories → DB so MCP search finds them
|
||||
try:
|
||||
import asyncio
|
||||
|
||||
from headroom.memory.sync import _build_sync_backend, sync_import
|
||||
from headroom.memory.sync_adapters.claude_code import (
|
||||
ClaudeCodeAdapter,
|
||||
get_claude_memory_dir,
|
||||
)
|
||||
|
||||
claude_memory_dir = get_claude_memory_dir()
|
||||
|
||||
async def _import_claude_memories() -> int:
|
||||
backend = _build_sync_backend(db_path)
|
||||
await backend._ensure_initialized()
|
||||
adapter = ClaudeCodeAdapter(claude_memory_dir)
|
||||
count = await sync_import(backend, adapter, mem_user)
|
||||
await backend.close()
|
||||
return count
|
||||
|
||||
imported = asyncio.run(_import_claude_memories())
|
||||
if imported:
|
||||
click.echo(f" Memory: imported {imported} memories from Claude")
|
||||
except Exception as e:
|
||||
click.echo(f" Warning: Claude memory import failed: {e}")
|
||||
|
||||
_inject_codex_provider_config(port)
|
||||
return
|
||||
|
||||
# Register headroom MCP server in Codex config.toml so Codex can
|
||||
# call headroom_retrieve on compression markers from the proxy.
|
||||
# These config writes do not need a running proxy — they run before
|
||||
# _ensure_proxy so unwrap has config to clean up even when proxy
|
||||
# startup or binary lookup fails.
|
||||
if not no_mcp:
|
||||
from headroom.mcp_registry import CodexRegistrar
|
||||
|
||||
# Codex starts a long-lived local MCP subprocess from config.toml.
|
||||
# If a previous wrap used another port, retrieval can silently point
|
||||
# at the wrong proxy while model traffic uses the right one.
|
||||
_setup_headroom_mcp(CodexRegistrar(), port, verbose=verbose, force=True)
|
||||
elif verbose:
|
||||
click.echo(" Skipping MCP retrieve tool (--no-mcp)")
|
||||
|
||||
# Coding-task compressor: tokensave primary, Serena backup. Codex starts
|
||||
# long-lived MCP subprocesses from config.toml, so force re-registration.
|
||||
from headroom.mcp_registry import CodexRegistrar
|
||||
|
||||
_setup_coding_compressor(
|
||||
CodexRegistrar(),
|
||||
serena_context="codex",
|
||||
return _run_codex_wrap(
|
||||
port=port,
|
||||
no_rtk=no_rtk,
|
||||
no_mcp=no_mcp,
|
||||
no_tokensave=no_tokensave,
|
||||
serena=serena,
|
||||
no_serena=no_serena,
|
||||
no_tokensave=no_tokensave,
|
||||
verbose=verbose,
|
||||
force=True,
|
||||
)
|
||||
|
||||
# Setup memory MCP server for Codex (native tool integration)
|
||||
if memory:
|
||||
click.echo(" Setting up memory for Codex...")
|
||||
mem_dir = Path.cwd() / ".headroom"
|
||||
mem_dir.mkdir(parents=True, exist_ok=True)
|
||||
db_path = str(mem_dir / "memory.db")
|
||||
mem_user = os.environ.get("USER", os.environ.get("USERNAME", "default"))
|
||||
|
||||
# Register MCP server in Codex config
|
||||
_inject_memory_mcp_config(mem_user)
|
||||
|
||||
# Inject memory guidance into project AGENTS.md
|
||||
agents_md = Path.cwd() / "AGENTS.md"
|
||||
_inject_memory_agents_md(agents_md)
|
||||
|
||||
# Sync Claude's memories → DB so MCP search finds them
|
||||
try:
|
||||
import asyncio
|
||||
|
||||
from headroom.memory.sync import _build_sync_backend, sync_import
|
||||
from headroom.memory.sync_adapters.claude_code import (
|
||||
ClaudeCodeAdapter,
|
||||
get_claude_memory_dir,
|
||||
)
|
||||
|
||||
claude_memory_dir = get_claude_memory_dir()
|
||||
|
||||
async def _import_claude_memories() -> int:
|
||||
backend = _build_sync_backend(db_path)
|
||||
await backend._ensure_initialized()
|
||||
adapter = ClaudeCodeAdapter(claude_memory_dir)
|
||||
count = await sync_import(backend, adapter, mem_user)
|
||||
await backend.close()
|
||||
return count
|
||||
|
||||
imported = asyncio.run(_import_claude_memories())
|
||||
if imported:
|
||||
click.echo(f" Memory: imported {imported} memories from Claude")
|
||||
except Exception as e:
|
||||
click.echo(f" Warning: Claude memory import failed: {e}")
|
||||
|
||||
codex_bin = shutil.which("codex")
|
||||
if not codex_bin:
|
||||
click.echo("Error: 'codex' not found in PATH.")
|
||||
click.echo("Install Codex CLI: npm install -g @openai/codex")
|
||||
raise SystemExit(1)
|
||||
|
||||
# Register our proxy client marker BEFORE _ensure_proxy so that another
|
||||
# wrapper's cleanup sees us as an active client and doesn't terminate a
|
||||
# shared proxy during the startup gap.
|
||||
_register_proxy_client(port)
|
||||
|
||||
# Let _ensure_proxy decide the port (same contract as other wrappers).
|
||||
# Called after config writes so unwrap has config to restore even when
|
||||
# proxy startup fails.
|
||||
_codex_proxy, actual_port = _ensure_proxy(
|
||||
port,
|
||||
no_proxy,
|
||||
code_graph=code_graph,
|
||||
no_proxy=no_proxy,
|
||||
learn=learn,
|
||||
memory=memory,
|
||||
agent_type="codex",
|
||||
code_graph=code_graph,
|
||||
backend=backend,
|
||||
anyllm_provider=anyllm_provider,
|
||||
region=region,
|
||||
verbose=verbose,
|
||||
prepare_only=prepare_only,
|
||||
codex_args=codex_args,
|
||||
)
|
||||
|
||||
# If the proxy fell back to a different port, move our marker to the
|
||||
# actual port so cleanup tracking stays accurate.
|
||||
if actual_port != port:
|
||||
_unregister_proxy_client(port)
|
||||
_register_proxy_client(actual_port)
|
||||
|
||||
# If the proxy fell back to a different port, update the MCP config so
|
||||
# the retrieval tool URL points at the port the proxy is actually on.
|
||||
if actual_port != port and not no_mcp:
|
||||
from headroom.mcp_registry import CodexRegistrar
|
||||
|
||||
_setup_headroom_mcp(CodexRegistrar(), actual_port, verbose=verbose, force=True)
|
||||
|
||||
env, env_vars_display = _build_codex_launch_env(actual_port, os.environ)
|
||||
|
||||
# Per-project savings attribution: the injected provider config maps the
|
||||
# X-Headroom-Project header to HEADROOM_PROJECT via env_http_headers, so
|
||||
# Codex sends it only when this var is set. A user-set value wins.
|
||||
_codex_project = _project_name_from_cwd()
|
||||
if _codex_project and "HEADROOM_PROJECT" not in env:
|
||||
env["HEADROOM_PROJECT"] = _codex_project
|
||||
|
||||
# Inject Headroom provider into Codex config so WebSocket traffic also
|
||||
# routes through the proxy. Codex ignores OPENAI_BASE_URL for its WS
|
||||
# transport unless a custom provider declares supports_websockets = true.
|
||||
# NOTE: this must run BEFORE _inject_memory_mcp_config because it rewrites
|
||||
# the config file. Re-inject MCP config after if memory is enabled.
|
||||
_codex_custom_upstream = _inject_codex_provider_config(actual_port)
|
||||
if _codex_custom_upstream and _UPSTREAM_BASE_URL_ENV_VAR not in env:
|
||||
# Carries the preserved custom base_url (#1614) to the injected
|
||||
# env_http_headers entry, which maps it to X-Headroom-Base-Url —
|
||||
# the proxy's OpenAI HTTP handlers forward there instead of the
|
||||
# hardcoded api.openai.com default. A user-set value wins.
|
||||
env[_UPSTREAM_BASE_URL_ENV_VAR] = _codex_custom_upstream
|
||||
env_vars_display.append(f"{_UPSTREAM_BASE_URL_ENV_VAR}={_codex_custom_upstream}")
|
||||
if memory:
|
||||
_inject_memory_mcp_config(os.environ.get("USER", os.environ.get("USERNAME", "default")))
|
||||
|
||||
# Proxy already started by _ensure_proxy above; tell _launch_tool to
|
||||
# skip duplicate startup. Cleanup of _codex_proxy happens on exit
|
||||
# via the finally block below.
|
||||
try:
|
||||
_launch_tool(
|
||||
binary=codex_bin,
|
||||
args=codex_args,
|
||||
env=env,
|
||||
port=actual_port,
|
||||
no_proxy=True,
|
||||
tool_label="CODEX",
|
||||
env_vars_display=env_vars_display,
|
||||
learn=learn,
|
||||
memory=memory,
|
||||
agent_type="codex",
|
||||
code_graph=code_graph,
|
||||
backend=backend,
|
||||
anyllm_provider=anyllm_provider,
|
||||
region=region,
|
||||
)
|
||||
finally:
|
||||
# _launch_tool's internal cleanup unregisters this client marker,
|
||||
# but doesn't know about the proxy we started. Terminate it when
|
||||
# no other clients remain.
|
||||
if _codex_proxy and _codex_proxy.poll() is None:
|
||||
_other = _live_proxy_clients(actual_port, exclude_self=True)
|
||||
if not _other:
|
||||
_codex_proxy.terminate()
|
||||
try:
|
||||
_codex_proxy.wait(timeout=5)
|
||||
except subprocess.TimeoutExpired:
|
||||
_codex_proxy.kill()
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Aider
|
||||
|
|
|
|||
|
|
@ -1053,6 +1053,166 @@ def test_wrap_codex_prepare_only_respects_codex_home(
|
|||
assert not (tmp_path / ".codex" / "config.toml").exists()
|
||||
|
||||
|
||||
def test_codex_session_home_overlay_seeds_active_home_and_cleans_up(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
codex_home = tmp_path / "custom-codex-home"
|
||||
codex_home.mkdir()
|
||||
monkeypatch.setenv("CODEX_HOME", str(codex_home))
|
||||
|
||||
config_file = codex_home / "config.toml"
|
||||
auth_file = codex_home / "auth.json"
|
||||
original_config = '[profiles.default]\nmodel = "gpt-4o"\n'
|
||||
original_auth = '{"auth_mode": "apikey"}'
|
||||
config_file.write_text(original_config, encoding="utf-8")
|
||||
auth_file.write_text(original_auth, encoding="utf-8")
|
||||
|
||||
with wrap_mod._codex_session_home_overlay() as session_home:
|
||||
seeded_config = (session_home / "config.toml").read_text(encoding="utf-8")
|
||||
seeded_auth = (session_home / "auth.json").read_text(encoding="utf-8")
|
||||
assert seeded_config == original_config
|
||||
assert seeded_auth == original_auth
|
||||
(session_home / "config.toml").write_text('model_provider = "headroom"\n', encoding="utf-8")
|
||||
assert config_file.read_text(encoding="utf-8") == original_config
|
||||
|
||||
assert not session_home.exists()
|
||||
assert config_file.read_text(encoding="utf-8") == original_config
|
||||
assert auth_file.read_text(encoding="utf-8") == original_auth
|
||||
|
||||
|
||||
def test_wrap_codex_launch_uses_session_scoped_codex_home(
|
||||
runner: CliRunner, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
codex_home = tmp_path / "custom-codex-home"
|
||||
codex_home.mkdir()
|
||||
monkeypatch.setenv("CODEX_HOME", str(codex_home))
|
||||
|
||||
config_file = codex_home / "config.toml"
|
||||
auth_file = codex_home / "auth.json"
|
||||
original_config = '[profiles.default]\nmodel = "gpt-4o"\n'
|
||||
original_auth = '{"auth_mode": "apikey"}'
|
||||
config_file.write_text(original_config, encoding="utf-8")
|
||||
auth_file.write_text(original_auth, encoding="utf-8")
|
||||
|
||||
launch_env: dict[str, str] = {}
|
||||
session_home_seen: list[Path] = []
|
||||
|
||||
def fake_launch(
|
||||
*,
|
||||
binary: str,
|
||||
args: tuple,
|
||||
env: dict[str, str],
|
||||
port: int,
|
||||
no_proxy: bool,
|
||||
tool_label: str,
|
||||
env_vars_display: list[str],
|
||||
**kwargs: object,
|
||||
) -> None:
|
||||
del args, port, no_proxy, tool_label, env_vars_display, kwargs
|
||||
assert binary == "/fake/codex"
|
||||
launch_env.update(env)
|
||||
session_home = Path(env["CODEX_HOME"])
|
||||
session_home_seen.append(session_home)
|
||||
assert session_home.exists()
|
||||
seeded_config = (session_home / "config.toml").read_text(encoding="utf-8")
|
||||
assert original_config in seeded_config
|
||||
assert 'model_provider = "headroom"' in seeded_config
|
||||
assert 'base_url = "http://127.0.0.1:8787/v1"' in seeded_config
|
||||
assert "[mcp_servers.headroom]" in seeded_config
|
||||
assert (session_home / "auth.json").read_text(encoding="utf-8") == original_auth
|
||||
|
||||
with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=None):
|
||||
with patch(
|
||||
"headroom.cli.wrap.shutil.which",
|
||||
side_effect=lambda cmd: "/fake/codex" if cmd == "codex" else None,
|
||||
):
|
||||
with patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
[
|
||||
"wrap",
|
||||
"codex",
|
||||
"--port",
|
||||
"8787",
|
||||
"--no-tokensave",
|
||||
"--no-serena",
|
||||
],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert session_home_seen
|
||||
assert launch_env["CODEX_HOME"] == str(session_home_seen[0])
|
||||
assert launch_env["OPENAI_BASE_URL"] == "http://127.0.0.1:8787/v1"
|
||||
assert config_file.read_text(encoding="utf-8") == original_config
|
||||
assert auth_file.read_text(encoding="utf-8") == original_auth
|
||||
assert not session_home_seen[0].exists()
|
||||
|
||||
|
||||
def test_wrap_codex_launches_use_distinct_session_homes_per_port(
|
||||
runner: CliRunner, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
codex_home = tmp_path / "custom-codex-home"
|
||||
codex_home.mkdir()
|
||||
monkeypatch.setenv("CODEX_HOME", str(codex_home))
|
||||
|
||||
config_file = codex_home / "config.toml"
|
||||
auth_file = codex_home / "auth.json"
|
||||
original_config = '[profiles.default]\nmodel = "gpt-4o"\n'
|
||||
original_auth = '{"auth_mode": "apikey"}'
|
||||
config_file.write_text(original_config, encoding="utf-8")
|
||||
auth_file.write_text(original_auth, encoding="utf-8")
|
||||
|
||||
launch_records: list[tuple[int, Path, str]] = []
|
||||
|
||||
def fake_launch(
|
||||
*,
|
||||
binary: str,
|
||||
args: tuple,
|
||||
env: dict[str, str],
|
||||
port: int,
|
||||
no_proxy: bool,
|
||||
tool_label: str,
|
||||
env_vars_display: list[str],
|
||||
**kwargs: object,
|
||||
) -> None:
|
||||
del args, no_proxy, tool_label, env_vars_display, kwargs
|
||||
assert binary == "/fake/codex"
|
||||
session_home = Path(env["CODEX_HOME"])
|
||||
assert session_home.exists()
|
||||
launch_records.append(
|
||||
(port, session_home, (session_home / "config.toml").read_text(encoding="utf-8"))
|
||||
)
|
||||
|
||||
with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=None):
|
||||
with patch(
|
||||
"headroom.cli.wrap.shutil.which",
|
||||
side_effect=lambda cmd: "/fake/codex" if cmd == "codex" else None,
|
||||
):
|
||||
with patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch):
|
||||
first = runner.invoke(
|
||||
main,
|
||||
["wrap", "codex", "--port", "8787", "--no-tokensave", "--no-serena"],
|
||||
)
|
||||
second = runner.invoke(
|
||||
main,
|
||||
["wrap", "codex", "--port", "9898", "--no-tokensave", "--no-serena"],
|
||||
)
|
||||
|
||||
assert first.exit_code == 0, first.output
|
||||
assert second.exit_code == 0, second.output
|
||||
assert len(launch_records) == 2
|
||||
assert launch_records[0][1] != launch_records[1][1]
|
||||
assert 'base_url = "http://127.0.0.1:8787/v1"' in launch_records[0][2]
|
||||
assert 'base_url = "http://127.0.0.1:9898/v1"' in launch_records[1][2]
|
||||
assert config_file.read_text(encoding="utf-8") == original_config
|
||||
assert auth_file.read_text(encoding="utf-8") == original_auth
|
||||
assert not launch_records[0][1].exists()
|
||||
assert not launch_records[1][1].exists()
|
||||
|
||||
|
||||
def test_wrap_codex_injects_rtk_globally_without_changing_project_agents(
|
||||
runner: CliRunner, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
|
|
@ -1086,6 +1246,41 @@ def test_wrap_codex_injects_rtk_globally_without_changing_project_agents(
|
|||
assert wrap_mod._RTK_MARKER.encode() in global_agents.read_bytes()
|
||||
|
||||
|
||||
def test_wrap_codex_launch_injects_rtk_globally_without_changing_project_agents(
|
||||
runner: CliRunner, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
project_dir = tmp_path / "project"
|
||||
project_dir.mkdir()
|
||||
project_agents = project_dir / "AGENTS.md"
|
||||
original = "# Project instructions\n\nUse the repository conventions.\n"
|
||||
project_agents.write_text(original, encoding="utf-8")
|
||||
original_bytes = project_agents.read_bytes()
|
||||
monkeypatch.chdir(project_dir)
|
||||
|
||||
with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=tmp_path / "rtk"):
|
||||
with patch(
|
||||
"headroom.cli.wrap.shutil.which",
|
||||
side_effect=lambda cmd: "/fake/codex" if cmd == "codex" else None,
|
||||
):
|
||||
with patch("headroom.cli.wrap._launch_tool"):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
[
|
||||
"wrap",
|
||||
"codex",
|
||||
"--no-mcp",
|
||||
"--no-serena",
|
||||
"--no-tokensave",
|
||||
],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert project_agents.read_bytes() == original_bytes
|
||||
global_agents = tmp_path / ".codex" / "AGENTS.md"
|
||||
assert wrap_mod._RTK_MARKER.encode() in global_agents.read_bytes()
|
||||
|
||||
|
||||
def test_unwrap_codex_without_codex_home_warns_on_ambiguous_noop(
|
||||
runner: CliRunner, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
|
|
@ -1539,16 +1734,14 @@ def test_wrap_codex_memory_launch_failure_unwrap_cleans_memory_only_config(
|
|||
|
||||
assert wrap_result.exit_code == 1
|
||||
config_file = tmp_path / ".codex" / "config.toml"
|
||||
content = config_file.read_text()
|
||||
assert "[mcp_servers.headroom_memory]" in content
|
||||
assert wrap_mod._CODEX_TOP_LEVEL_MARKER not in content
|
||||
assert not config_file.exists()
|
||||
assert not (tmp_path / ".codex" / "config.toml.headroom-backup").exists()
|
||||
|
||||
with patch("headroom.cli.wrap._stop_local_proxy_for_unwrap") as stop_proxy:
|
||||
unwrap_result = runner.invoke(main, ["unwrap", "codex", "--no-stop-proxy"])
|
||||
|
||||
assert unwrap_result.exit_code == 0, unwrap_result.output
|
||||
assert not config_file.exists()
|
||||
assert not (tmp_path / ".codex" / "config.toml.headroom-backup").exists()
|
||||
assert "Nothing to undo" in unwrap_result.output
|
||||
stop_proxy.assert_not_called()
|
||||
|
||||
|
||||
|
|
@ -1692,58 +1885,29 @@ class TestCodexProjectHeaderConfig:
|
|||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Regression: codex delegates port resolution to _ensure_proxy
|
||||
# Regression: codex preserves the requested port through the session-scoped runner
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestCodexPortResolution:
|
||||
"""codex() uses _ensure_proxy() to resolve ports (not early _find_available_port).
|
||||
"""codex() hands the requested port to the session-scoped wrap runner.
|
||||
|
||||
Regression for headroom#1406 round 2 review: codex() must follow
|
||||
the same selected-port contract as other wrappers (aider, copilot, etc.)
|
||||
so that a healthy existing proxy on the requested port is reused instead
|
||||
of skipped by a blind socket probe.
|
||||
Regression for headroom#1406 round 2 review: the codex command must keep
|
||||
the selected-port contract intact after the session-home refactor instead
|
||||
of silently dropping or rewriting the requested port before the shared
|
||||
launch path handles proxy reuse and fallback.
|
||||
"""
|
||||
|
||||
def test_delegates_to_ensure_proxy(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""codex() calls _ensure_proxy and uses the returned port."""
|
||||
def test_delegates_to_session_runner(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""codex() passes the requested port through to _run_codex_wrap."""
|
||||
_set_test_home(monkeypatch, Path("/tmp/test_headroom_codex"))
|
||||
|
||||
captured_port: list[int] = []
|
||||
call_kw: dict = {}
|
||||
|
||||
# Mock _ensure_proxy to capture the requested port
|
||||
def mock_ensure_proxy(port: int, no_proxy: bool, **kwargs: object) -> tuple[None, int]:
|
||||
captured_port.append(port)
|
||||
# Simulate port fallback: requested 8787, actual 8788
|
||||
return None, 8788
|
||||
def mock_run_codex_wrap(**kwargs: object) -> None:
|
||||
call_kw.update(kwargs)
|
||||
|
||||
monkeypatch.setattr(wrap_mod, "_ensure_proxy", mock_ensure_proxy)
|
||||
|
||||
# Mock all heavy dependencies
|
||||
monkeypatch.setattr(
|
||||
wrap_mod, "_codex_config_paths", lambda: (Path("/dev/null"), Path("/dev/null"))
|
||||
)
|
||||
monkeypatch.setattr(wrap_mod, "_snapshot_codex_config_if_unwrapped", lambda *a, **kw: None)
|
||||
monkeypatch.setattr(wrap_mod, "_ensure_rtk_binary", lambda *a, **kw: None)
|
||||
monkeypatch.setattr(wrap_mod, "_setup_lean_ctx_agent", lambda *a, **kw: None)
|
||||
monkeypatch.setattr(wrap_mod, "_inject_rtk_instructions", lambda *a, **kw: None)
|
||||
monkeypatch.setattr(wrap_mod, "_codex_home_dir", lambda: Path("/tmp"))
|
||||
monkeypatch.setattr(wrap_mod, "_setup_headroom_mcp", lambda *a, **kw: None)
|
||||
monkeypatch.setattr(wrap_mod, "_setup_serena_mcp", lambda *a, **kw: None)
|
||||
monkeypatch.setattr(wrap_mod, "_disable_serena_mcp", lambda *a, **kw: None)
|
||||
monkeypatch.setattr("shutil.which", lambda x: "/usr/bin/codex" if x == "codex" else None)
|
||||
monkeypatch.setattr(wrap_mod, "_build_codex_launch_env", lambda port, env: ({}, []))
|
||||
monkeypatch.setattr(wrap_mod, "_inject_codex_provider_config", lambda port: None)
|
||||
monkeypatch.setattr(wrap_mod, "_project_name_from_cwd", lambda: None)
|
||||
monkeypatch.setattr(wrap_mod, "_live_proxy_clients", lambda *a, **kw: [])
|
||||
|
||||
# Intercept _launch_tool to verify port propagation
|
||||
launch_kw: dict = {}
|
||||
|
||||
def mock_launch_tool(**kwargs: object) -> None:
|
||||
launch_kw.update(kwargs)
|
||||
|
||||
monkeypatch.setattr(wrap_mod, "_launch_tool", mock_launch_tool)
|
||||
monkeypatch.setattr(wrap_mod, "_run_codex_wrap", mock_run_codex_wrap)
|
||||
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(
|
||||
|
|
@ -1752,10 +1916,6 @@ class TestCodexPortResolution:
|
|||
)
|
||||
|
||||
assert result.exit_code == 0, f"CLI failed: {result.output}"
|
||||
assert captured_port == [8787], (
|
||||
f"_ensure_proxy called with {captured_port}, expected [8787]"
|
||||
)
|
||||
assert launch_kw.get("port") == 8788, (
|
||||
f"_launch_tool port={launch_kw.get('port')}, expected 8788 "
|
||||
"(the actual_port from _ensure_proxy fallback)"
|
||||
)
|
||||
assert call_kw.get("port") == 8787
|
||||
assert call_kw.get("no_proxy") is False
|
||||
assert call_kw.get("prepare_only") is False
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue