From 06428d20fd0514375024a6784b30b873a0f8d82d Mon Sep 17 00:00:00 2001 From: JerrettDavis Date: Tue, 5 May 2026 21:03:42 -0500 Subject: [PATCH] fix: preserve Codex OAuth proxy delivery Preserve Codex OAuth-safe provider config across init, wrap, and persistent install paths, and strengthen coverage so Codex requests are proven to reach Headroom and the mock upstream. The wrap e2e now sends a real chat-completions probe and checks Headroom /stats. Runtime tests cover temporary launch env, install env, init config, provider-scope config delivery, and the Python 3.11 ws bootstrap path. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- e2e/init/run.py | 10 + e2e/wrap/run.py | 144 ++++++++++- headroom/cli/init.py | 14 +- headroom/cli/wrap.py | 38 ++- headroom/providers/codex/install.py | 34 ++- tests/test_cli/test_init_cli.py | 2 + tests/test_cli/test_wrap_codex.py | 80 +++++- tests/test_install/test_providers.py | 4 +- tests/test_provider_codex_install.py | 27 ++ tests/test_provider_codex_runtime.py | 356 +++++++++++++++++++++++++++ 10 files changed, 679 insertions(+), 30 deletions(-) create mode 100644 tests/test_provider_codex_install.py diff --git a/e2e/init/run.py b/e2e/init/run.py index d7931704c..87a986ec3 100644 --- a/e2e/init/run.py +++ b/e2e/init/run.py @@ -137,6 +137,11 @@ def _verify_codex_local(ctx: CaseContext) -> None: if 'base_url = "http://127.0.0.1:9012/v1"' not in config: raise AssertionError("Codex config should point at the requested proxy port (9012)") + if 'env_key = "OPENAI_API_KEY"' in config: + raise AssertionError("Codex local init should preserve OAuth and never inject env_key") + for expected in ("requires_openai_auth = true", "supports_websockets = true"): + if expected not in config: + raise AssertionError(f"Codex local init missing {expected!r}") if config.count("[features]") != 1: raise AssertionError("Codex config should keep a single [features] table") if "codex_hooks = true" not in config: @@ -170,6 +175,11 @@ def _verify_codex_global(ctx: CaseContext) -> None: config = (ctx.home / ".codex" / "config.toml").read_text(encoding="utf-8") if 'base_url = "http://127.0.0.1:8787/v1"' not in config: raise AssertionError("Codex user config should point at port 8787 by default") + if 'env_key = "OPENAI_API_KEY"' in config: + raise AssertionError("Codex global init should preserve OAuth and never inject env_key") + for expected in ("requires_openai_auth = true", "supports_websockets = true"): + if expected not in config: + raise AssertionError(f"Codex global init missing {expected!r}") if "codex_hooks = true" not in config: raise AssertionError("Codex user config should enable codex_hooks") hooks = json.loads((ctx.home / ".codex" / "hooks.json").read_text(encoding="utf-8")) diff --git a/e2e/wrap/run.py b/e2e/wrap/run.py index 9b914a765..47f077745 100644 --- a/e2e/wrap/run.py +++ b/e2e/wrap/run.py @@ -247,6 +247,98 @@ def create_shims(shim_dir: Path) -> None: raise SystemExit(0) """ ) + codex_shim = textwrap.dedent( + """\ + #!/usr/bin/env python3 + from __future__ import annotations + + import json + import os + import sys + import urllib.request + from pathlib import Path + + tool = Path(sys.argv[0]).name + log_dir = Path(os.environ["HEADROOM_E2E_LOG_DIR"]) + log_dir.mkdir(parents=True, exist_ok=True) + record = { + "tool": tool, + "argv": sys.argv[1:], + "cwd": os.getcwd(), + "env": { + key: os.environ.get(key) + for key in ("OPENAI_BASE_URL", "OPENAI_API_BASE", "ANTHROPIC_BASE_URL") + if os.environ.get(key) is not None + }, + } + + probes = [] + + def request_json( + url: str, + *, + payload: dict[str, object] | None = None, + headers: dict[str, str] | None = None, + ) -> tuple[int, dict[str, object]]: + data = None if payload is None else json.dumps(payload).encode("utf-8") + request = urllib.request.Request(url, data=data, headers=headers or {}) + if payload is not None: + request.add_header("Content-Type", "application/json") + with urllib.request.urlopen(request, timeout=10) as response: + raw = response.read() + body = json.loads(raw.decode("utf-8") or "{}") if raw else {} + return response.status, body + + openai_base = os.environ.get("OPENAI_BASE_URL") or os.environ.get("OPENAI_API_BASE") + if openai_base: + auth_headers = {"Authorization": "Bearer test-key"} + models_status, _models_body = request_json( + f"{openai_base.rstrip('/')}/models", + headers=auth_headers, + ) + probes.append({"url": f"{openai_base.rstrip('/')}/models", "status": models_status}) + + model_name = "headroom-wrap-e2e" + chat_status, chat_body = request_json( + f"{openai_base.rstrip('/')}/chat/completions", + payload={ + "model": model_name, + "messages": [ + { + "role": "user", + "content": "Confirm Headroom received this wrapped Codex message.", + } + ], + }, + headers=auth_headers, + ) + probes.append( + {"url": f"{openai_base.rstrip('/')}/chat/completions", "status": chat_status} + ) + record["chat_completion"] = ( + chat_body.get("choices", [{}])[0].get("message", {}).get("content") + ) + + stats_url = openai_base.rstrip("/").removesuffix("/v1") + "/stats" + stats_status, stats_body = request_json(stats_url, headers=auth_headers) + probes.append({"url": stats_url, "status": stats_status}) + requests = stats_body.get("requests", {}) + by_model = requests.get("by_model", {}) if isinstance(requests, dict) else {} + record["headroom_request_total"] = ( + requests.get("total") if isinstance(requests, dict) else None + ) + record["headroom_model_count"] = ( + by_model.get(model_name) if isinstance(by_model, dict) else None + ) + + record["probes"] = probes + + with (log_dir / f"{tool}.jsonl").open("a", encoding="utf-8") as handle: + handle.write(json.dumps(record) + "\\n") + print(f"{tool} shim executed") + raise SystemExit(0) + """ + ) rtk_shim = textwrap.dedent( """\ #!/usr/bin/env python3 @@ -262,7 +354,7 @@ def create_shims(shim_dir: Path) -> None: """ ) write_executable(shim_dir / "claude", generic_shim) - write_executable(shim_dir / "codex", generic_shim) + write_executable(shim_dir / "codex", codex_shim) write_executable(shim_dir / "aider", generic_shim) write_executable(shim_dir / "rtk", rtk_shim) @@ -437,7 +529,9 @@ def verify_proxy_round_trip(base_env: dict[str, str], mock_server: MockOpenAISer stop_process(proc) -def verify_codex_wrap(base_env: dict[str, str], project_dir: Path, log_dir: Path) -> None: +def verify_codex_wrap( + base_env: dict[str, str], project_dir: Path, log_dir: Path, mock_server: MockOpenAIServer +) -> None: port = CODEX_PORT run( ["headroom", "wrap", "codex", "--port", str(port), "--", "--help"], @@ -454,6 +548,24 @@ def verify_codex_wrap(base_env: dict[str, str], project_dir: Path, log_dir: Path RTK_MARKER in global_agents.read_text(encoding="utf-8"), "Missing global RTK marker" ) + 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", + ) + for expected in ("requires_openai_auth = true", "supports_websockets = true"): + assert_true(expected in config, f"Codex wrap missing {expected!r}") + entries = read_jsonl(log_dir / "codex.jsonl") assert_true(len(entries) > 0, "Codex shim should have been invoked") env_vars = entries[-1]["env"] @@ -462,8 +574,30 @@ def verify_codex_wrap(base_env: dict[str, str], project_dir: Path, log_dir: Path "Codex wrap should set OPENAI_BASE_URL", ) assert_true( - entries[-1]["probes"] == [{"url": f"http://127.0.0.1:{port}/v1/models", "status": 200}], - "Codex shim should prove OPENAI_BASE_URL points at a live proxy", + entries[-1]["probes"] + == [ + {"url": f"http://127.0.0.1:{port}/v1/models", "status": 200}, + {"url": f"http://127.0.0.1:{port}/v1/chat/completions", "status": 200}, + {"url": f"http://127.0.0.1:{port}/stats", "status": 200}, + ], + "Codex shim should prove OPENAI_BASE_URL points at a live proxy and that Headroom logged the wrapped message", + ) + assert_true( + entries[-1].get("chat_completion") == "mock completion from upstream", + "Codex wrap should receive the mock upstream completion through Headroom", + ) + assert_true( + entries[-1].get("headroom_model_count", 0) >= 1, + "Codex wrap should appear in Headroom request stats", + ) + assert_true( + any( + item["path"] == "/v1/chat/completions" + and isinstance(item.get("body"), dict) + and item["body"].get("model") == "headroom-wrap-e2e" + for item in mock_server.requests + ), + "Codex wrap should forward the wrapped message upstream through Headroom", ) @@ -683,7 +817,7 @@ def main() -> None: try: verify_proxy_round_trip(base_env, mock_server) verify_claude_wrap(base_env, project_dir, log_dir) - verify_codex_wrap(base_env, project_dir, log_dir) + verify_codex_wrap(base_env, project_dir, log_dir, mock_server) verify_aider_wrap(base_env, project_dir, log_dir) verify_cursor_wrap(base_env, project_dir) local_plugin_dir = prepare_local_openclaw_plugin(base_env, tmp_dir) diff --git a/headroom/cli/init.py b/headroom/cli/init.py index 6c57581e9..d4cc59cdd 100644 --- a/headroom/cli/init.py +++ b/headroom/cli/init.py @@ -28,6 +28,7 @@ from headroom.install.runtime import ( ) from headroom.install.state import load_manifest, save_manifest from headroom.install.supervisors import start_supervisor +from headroom.providers.codex.install import build_provider_section as build_codex_provider_section from .main import main @@ -212,13 +213,12 @@ def _ensure_codex_provider(path: Path, port: int) -> None: block = ( f"{_CODEX_PROVIDER_MARKER_START}\n" 'model_provider = "headroom"\n\n' - "[model_providers.headroom]\n" - 'name = "Headroom init proxy"\n' - f'base_url = "http://127.0.0.1:{port}/v1"\n' - 'env_key = "OPENAI_API_KEY"\n' - "requires_openai_auth = true\n" - "supports_websockets = true\n" - f"{_CODEX_PROVIDER_MARKER_END}" + + build_codex_provider_section( + port=port, + name="Headroom init proxy", + include_markers=False, + ) + + f"{_CODEX_PROVIDER_MARKER_END}" ) content = path.read_text(encoding="utf-8") if path.exists() else "" content = _replace_marker_block( diff --git a/headroom/cli/wrap.py b/headroom/cli/wrap.py index fe8c697cb..aba347b0c 100644 --- a/headroom/cli/wrap.py +++ b/headroom/cli/wrap.py @@ -457,9 +457,14 @@ def _strip_codex_headroom_blocks(content: str) -> str: content = content.replace(_CODEX_TOP_LEVEL_MARKER + "\n", "") content = content.replace(_CODEX_END_MARKER + "\n", "") - # Strip any leftover top-level `model_provider = "headroom"` line, which - # older versions of `wrap codex` wrote outside the marker block. + # Strip any leftover top-level keys that older (or crashed) versions of + # `wrap codex` may have written outside the marker block. content = re.sub(r'(?m)^[ \t]*model_provider[ \t]*=[ \t]*"headroom"[ \t]*\r?\n', "", content) + content = re.sub( + r'(?m)^[ \t]*openai_base_url[ \t]*=[ \t]*"http://127\.0\.0\.1:\d+/v1"[ \t]*\r?\n', + "", + content, + ) # Strip any orphaned `[model_providers.headroom]` table with the fields we # write. We only remove it if the table is recognisably ours (base_url @@ -537,10 +542,20 @@ def _prepare_wrap_rtk(verbose: bool = False, *, label: str | None = None) -> Pat def _inject_codex_provider_config(port: int) -> None: """Inject a Headroom model provider into Codex's config.toml. - Codex ignores OPENAI_BASE_URL for WebSocket transport unless a custom - provider declares ``supports_websockets = true``. This writes a - ``[model_providers.headroom]`` section that routes both HTTP and WS - through the proxy, and sets ``model_provider = "headroom"``. + Two keys are written in the top-level block: + + * ``model_provider = "headroom"`` — selects the custom provider for + API-key mode traffic. + * ``openai_base_url = "http://127.0.0.1:{port}/v1"`` — overrides the + built-in ``openai`` provider's base URL. This is the critical key for + **subscription (ChatGPT plan) users**: Codex detects subscription auth + and routes through the built-in ``openai`` provider regardless of + ``model_provider``, so without this override it bypasses the proxy and + hits ``https://chatgpt.com/backend-api/codex`` directly. + + A ``[model_providers.headroom]`` section is also written with + ``supports_websockets = true`` so that WebSocket transport (used by + default) also flows through the proxy for API-key users. Safe to call multiple times — the injected block is fully replaced on each call, so re-running with a different ``port`` updates the config. @@ -558,14 +573,16 @@ def _inject_codex_provider_config(port: int) -> None: # stripping them is unambiguous and never consumes user content that # happens to sit between the two. top_level_block = ( - f'{_CODEX_TOP_LEVEL_MARKER}\nmodel_provider = "headroom"\n{_CODEX_END_MARKER}\n' + f"{_CODEX_TOP_LEVEL_MARKER}\n" + f'model_provider = "headroom"\n' + f'openai_base_url = "http://127.0.0.1:{port}/v1"\n' + f"{_CODEX_END_MARKER}\n" ) provider_section = ( f"{_CODEX_TOP_LEVEL_MARKER}\n" "[model_providers.headroom]\n" 'name = "OpenAI via Headroom proxy"\n' f'base_url = "http://127.0.0.1:{port}/v1"\n' - f'env_key = "OPENAI_API_KEY"\n' f"requires_openai_auth = true\n" f"supports_websockets = true\n" f"{_CODEX_END_MARKER}\n" @@ -596,7 +613,10 @@ def _inject_codex_provider_config(port: int) -> None: content = top_level_block + "\n" + provider_section config_file.write_text(content) - click.echo(f" Codex config: injected Headroom provider (WS + HTTP) into {config_file}") + click.echo( + f" Codex config: injected Headroom provider (WS + HTTP, API-key + subscription)" + f" into {config_file}" + ) except Exception as e: click.echo(f" Warning: could not update Codex config: {e}") diff --git a/headroom/providers/codex/install.py b/headroom/providers/codex/install.py index b8f718602..72775b350 100644 --- a/headroom/providers/codex/install.py +++ b/headroom/providers/codex/install.py @@ -18,6 +18,27 @@ _CODEX_PATTERN = re.compile( ) +def build_provider_section( + *, + port: int, + name: str, + marker_start: str = _CODEX_MARKER_START, + marker_end: str = _CODEX_MARKER_END, + include_markers: bool = True, +) -> str: + """Build a managed Codex provider block that preserves OpenAI OAuth.""" + body = ( + "[model_providers.headroom]\n" + f'name = "{name}"\n' + f'base_url = "{proxy_base_url(port)}"\n' + "requires_openai_auth = true\n" + "supports_websockets = true\n" + ) + if not include_markers: + return body + return f"{marker_start}\n{body}{marker_end}\n" + + def build_install_env(*, port: int, backend: str) -> dict[str, str]: """Build the persistent install environment for Codex.""" del backend @@ -34,13 +55,12 @@ def apply_provider_scope(manifest: DeploymentManifest) -> ManagedMutation | None section = ( f"{_CODEX_MARKER_START}\n" 'model_provider = "headroom"\n\n' - "[model_providers.headroom]\n" - 'name = "Headroom persistent proxy"\n' - f'base_url = "{proxy_base_url(manifest.port)}"\n' - 'env_key = "OPENAI_API_KEY"\n' - "requires_openai_auth = true\n" - "supports_websockets = true\n" - f"{_CODEX_MARKER_END}\n" + + build_provider_section( + port=manifest.port, + name="Headroom persistent proxy", + include_markers=False, + ) + + f"{_CODEX_MARKER_END}\n" ) if path.exists(): existing = path.read_text() diff --git a/tests/test_cli/test_init_cli.py b/tests/test_cli/test_init_cli.py index 151e4057f..6ec6b5e8d 100644 --- a/tests/test_cli/test_init_cli.py +++ b/tests/test_cli/test_init_cli.py @@ -203,6 +203,7 @@ def test_init_codex_merges_feature_flag_into_existing_table(monkeypatch, tmp_pat assert 'base_url = "http://127.0.0.1:9000/v1"' in content assert content.count("[features]") == 1 assert "codex_hooks = true" in content + assert 'env_key = "OPENAI_API_KEY"' not in content hooks = json.loads((tmp_path / ".codex" / "hooks.json").read_text(encoding="utf-8")) assert "--profile init-local-demo" in hooks["hooks"]["SessionStart"][0]["hooks"][0]["command"] assert "init hook ensure" in hooks["hooks"]["SessionStart"][0]["hooks"][0]["command"] @@ -436,6 +437,7 @@ def test_ensure_codex_provider_replaces_existing_marker(monkeypatch, tmp_path: P assert content.count(init_cli._CODEX_PROVIDER_MARKER_START) == 1 assert 'base_url = "http://127.0.0.1:9100/v1"' in content assert "old = true" not in content + assert 'env_key = "OPENAI_API_KEY"' not in content def test_ensure_codex_feature_flag_replaces_existing_marker(monkeypatch, tmp_path: Path) -> None: diff --git a/tests/test_cli/test_wrap_codex.py b/tests/test_cli/test_wrap_codex.py index 9944b63bd..b5e612232 100644 --- a/tests/test_cli/test_wrap_codex.py +++ b/tests/test_cli/test_wrap_codex.py @@ -195,9 +195,11 @@ class TestInjectAndRestoreRoundTrip: # provider-table block. Re-wrapping must not duplicate them. assert content.count(wrap_mod._CODEX_TOP_LEVEL_MARKER) == 2 assert content.count(wrap_mod._CODEX_END_MARKER) == 2 - # Latest port is honoured. + # Latest port is honoured in both keys. assert 'base_url = "http://127.0.0.1:9999/v1"' in content + assert 'openai_base_url = "http://127.0.0.1:9999/v1"' in content assert 'base_url = "http://127.0.0.1:8787/v1"' not in content + assert 'openai_base_url = "http://127.0.0.1:8787/v1"' not in content # User's original content is preserved. assert 'model = "gpt-4o"' in content @@ -256,6 +258,82 @@ class TestInjectAndRestoreRoundTrip: assert config_file.read_text() == malformed +# --------------------------------------------------------------------------- +# Subscription routing: openai_base_url intercepts ChatGPT plan traffic +# --------------------------------------------------------------------------- + + +class TestSubscriptionRouting: + """Codex subscription (ChatGPT plan) bypasses OPENAI_BASE_URL and the + custom model_provider; it uses the built-in ``openai`` provider whose + base_url defaults to ``https://chatgpt.com/backend-api/codex``. + Setting ``openai_base_url`` overrides that default for all auth modes.""" + + def test_inject_writes_openai_base_url( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + _set_test_home(monkeypatch, tmp_path) + + wrap_mod._inject_codex_provider_config(8787) + + content = (tmp_path / ".codex" / "config.toml").read_text() + assert 'openai_base_url = "http://127.0.0.1:8787/v1"' in content + + def test_openai_base_url_port_updates_on_rewrap( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + _set_test_home(monkeypatch, tmp_path) + + wrap_mod._inject_codex_provider_config(8787) + wrap_mod._inject_codex_provider_config(9999) + + content = (tmp_path / ".codex" / "config.toml").read_text() + assert 'openai_base_url = "http://127.0.0.1:9999/v1"' in content + assert 'openai_base_url = "http://127.0.0.1:8787/v1"' not in content + + def test_openai_base_url_removed_on_unwrap( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + _set_test_home(monkeypatch, tmp_path) + config_dir = tmp_path / ".codex" + config_dir.mkdir() + config_file = config_dir / "config.toml" + original = '[profiles.default]\nmodel = "gpt-4o"\n' + config_file.write_text(original) + + wrap_mod._inject_codex_provider_config(8787) + assert 'openai_base_url = "http://127.0.0.1:8787/v1"' in config_file.read_text() + + wrap_mod._restore_codex_provider_config() + assert config_file.read_text() == original + + def test_strip_cleans_orphaned_openai_base_url(self) -> None: + """Safety net: orphaned openai_base_url lines are cleaned up.""" + content = ( + '[profiles.default]\nmodel = "gpt-4o"\nopenai_base_url = "http://127.0.0.1:8787/v1"\n' + ) + cleaned = wrap_mod._strip_codex_headroom_blocks(content) + assert "openai_base_url" not in cleaned + assert 'model = "gpt-4o"' in cleaned + + def test_no_env_key_in_injected_provider( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + """env_key must be absent so Codex doesn't require OPENAI_API_KEY. + + Codex treats env_key as a hard requirement — if the env var is missing + it throws "Missing environment variable" at startup. Subscription + (ChatGPT Plus) users don't have OPENAI_API_KEY set, so injecting + env_key breaks them (issue #393). + """ + _set_test_home(monkeypatch, tmp_path) + + wrap_mod._inject_codex_provider_config(8787) + + content = (tmp_path / ".codex" / "config.toml").read_text() + assert "env_key" not in content + + # --------------------------------------------------------------------------- # Integration tests: full `headroom wrap codex` / `headroom unwrap codex` # --------------------------------------------------------------------------- diff --git a/tests/test_install/test_providers.py b/tests/test_install/test_providers.py index 5c9fba42c..e7e5925e5 100644 --- a/tests/test_install/test_providers.py +++ b/tests/test_install/test_providers.py @@ -70,6 +70,8 @@ def test_apply_and_revert_codex_provider_scope(monkeypatch, tmp_path: Path) -> N content = config_path.read_text() assert 'model_provider = "headroom"' in content assert 'base_url = "http://127.0.0.1:8787/v1"' in content + assert 'env_key = "OPENAI_API_KEY"' not in content + assert "requires_openai_auth = true" in content assert mutation is not None revert_codex_provider_scope(mutation, manifest) @@ -107,7 +109,6 @@ def test_apply_codex_provider_scope_replaces_existing_managed_block( "[model_providers.headroom]\n" 'name = "Headroom persistent proxy"\n' 'base_url = "http://127.0.0.1:1111/v1"\n' - 'env_key = "OPENAI_API_KEY"\n' "requires_openai_auth = true\n" "supports_websockets = true\n" "# --- end Headroom persistent provider ---\n" @@ -122,6 +123,7 @@ def test_apply_codex_provider_scope_replaces_existing_managed_block( assert content.count("# --- Headroom persistent provider ---") == 1 assert 'base_url = "http://127.0.0.1:9999/v1"' in content assert 'base_url = "http://127.0.0.1:1111/v1"' not in content + assert 'env_key = "OPENAI_API_KEY"' not in content def test_apply_codex_provider_scope_creates_new_config_when_missing( diff --git a/tests/test_provider_codex_install.py b/tests/test_provider_codex_install.py new file mode 100644 index 000000000..2b3704f58 --- /dev/null +++ b/tests/test_provider_codex_install.py @@ -0,0 +1,27 @@ +from __future__ import annotations + +from headroom.providers.codex.install import build_provider_section + + +def test_codex_provider_section_preserves_openai_oauth() -> None: + section = build_provider_section(port=8787, name="OpenAI via Headroom proxy") + + assert 'name = "OpenAI via Headroom proxy"' in section + assert 'base_url = "http://127.0.0.1:8787/v1"' in section + assert "requires_openai_auth = true" in section + assert "supports_websockets = true" in section + assert 'env_key = "OPENAI_API_KEY"' not in section + + +def test_codex_provider_section_supports_custom_markers() -> None: + section = build_provider_section( + port=9100, + name="Headroom init proxy", + marker_start="# --- start ---", + marker_end="# --- end ---", + ) + + assert section.startswith("# --- start ---\n") + assert section.endswith("# --- end ---\n") + assert 'base_url = "http://127.0.0.1:9100/v1"' in section + assert 'env_key = "OPENAI_API_KEY"' not in section diff --git a/tests/test_provider_codex_runtime.py b/tests/test_provider_codex_runtime.py index 1c3859e74..c7929afd0 100644 --- a/tests/test_provider_codex_runtime.py +++ b/tests/test_provider_codex_runtime.py @@ -1,6 +1,285 @@ from __future__ import annotations +import json +import os +import socket +import threading +import time +from collections.abc import Iterator +from dataclasses import dataclass +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path + +import httpx +import pytest +import uvicorn + +try: + import tomllib +except ModuleNotFoundError: # pragma: no cover + import tomli as tomllib # type: ignore[no-redef] + +from headroom.cli import init as init_cli +from headroom.install.models import ConfigScope, DeploymentManifest from headroom.providers.codex import build_launch_env, proxy_base_url +from headroom.providers.codex.install import apply_provider_scope, build_install_env +from headroom.proxy.server import ProxyConfig, create_app + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +class _MockOpenAIServer(ThreadingHTTPServer): + allow_reuse_address = True + daemon_threads = True + + def __init__(self, server_address: tuple[str, int]) -> None: + super().__init__(server_address, _MockOpenAIHandler) + self.requests: list[dict[str, object]] = [] + + +class _MockOpenAIHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def log_message(self, format: str, *args: object) -> None: # noqa: A003 + return + + def _write_json(self, status_code: int, payload: dict[str, object]) -> None: + body = json.dumps(payload).encode("utf-8") + self.send_response(status_code) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def _record(self, body: dict[str, object] | None = None) -> None: + server = self.server + assert isinstance(server, _MockOpenAIServer) + server.requests.append( + { + "method": self.command, + "path": self.path, + "authorization": self.headers.get("Authorization"), + "body": body, + } + ) + + def do_GET(self) -> None: # noqa: N802 + self._record() + if self.path == "/v1/models": + self._write_json( + 200, + { + "object": "list", + "data": [{"id": "gpt-4o-mini", "object": "model", "owned_by": "openai"}], + }, + ) + return + self._write_json(404, {"error": {"message": "not found"}}) + + def do_POST(self) -> None: # noqa: N802 + length = int(self.headers.get("Content-Length", "0")) + raw_body = self.rfile.read(length) if length else b"" + payload = json.loads(raw_body.decode("utf-8") or "{}") + self._record(body=payload) + if self.path == "/v1/chat/completions": + self._write_json( + 200, + { + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 0, + "model": payload.get("model", "gpt-4o-mini"), + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "message": { + "role": "assistant", + "content": "mock completion from upstream", + }, + } + ], + "usage": { + "prompt_tokens": 12, + "completion_tokens": 5, + "total_tokens": 17, + }, + }, + ) + return + self._write_json(404, {"error": {"message": "not found"}}) + + +class _ProxyThread: + def __init__(self, port: int, upstream_port: int) -> None: + app = create_app( + ProxyConfig( + host="127.0.0.1", + port=port, + optimize=False, + cache_enabled=False, + rate_limit_enabled=False, + openai_api_url=f"http://127.0.0.1:{upstream_port}", + ) + ) + config = uvicorn.Config( + app, + host="127.0.0.1", + port=port, + log_level="warning", + loop="asyncio", + lifespan="on", + ws="none", + ) + self.server = uvicorn.Server(config) + self.thread = threading.Thread(target=self.server.run, name="codex-proxy", daemon=True) + + def start(self) -> None: + self.thread.start() + deadline = time.perf_counter() + 10.0 + while time.perf_counter() < deadline: + if self.server.started: + return + time.sleep(0.05) + raise RuntimeError("proxy failed to start within 10s") + + def stop(self) -> None: + self.server.should_exit = True + self.thread.join(timeout=10.0) + + +@dataclass +class _CodexProxyStack: + proxy_port: int + upstream: _MockOpenAIServer + proxy: _ProxyThread + + @property + def base_url(self) -> str: + return proxy_base_url(self.proxy_port) + + @property + def stats_url(self) -> str: + return f"http://127.0.0.1:{self.proxy_port}/stats" + + +@pytest.fixture(scope="module") +def codex_proxy_stack(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_CodexProxyStack]: + temp_home = tmp_path_factory.mktemp("codex-proxy-home") + previous_env = {name: os.environ.get(name) for name in ("HEADROOM_REQUIRE_RUST_CORE", "HOME")} + os.environ["HEADROOM_REQUIRE_RUST_CORE"] = "false" + os.environ["HOME"] = str(temp_home) + + upstream_port = _free_port() + upstream = _MockOpenAIServer(("127.0.0.1", upstream_port)) + upstream_thread = threading.Thread(target=upstream.serve_forever, daemon=True) + upstream_thread.start() + + proxy_port = _free_port() + proxy = _ProxyThread(proxy_port, upstream_port) + proxy.start() + + try: + _wait_for_stats(f"http://127.0.0.1:{proxy_port}/stats") + yield _CodexProxyStack(proxy_port=proxy_port, upstream=upstream, proxy=proxy) + finally: + proxy.stop() + upstream.shutdown() + upstream_thread.join(timeout=5.0) + for name, value in previous_env.items(): + if value is None: + os.environ.pop(name, None) + else: + os.environ[name] = value + + +def _wait_for_stats(url: str) -> None: + deadline = time.time() + 10.0 + while time.time() < deadline: + try: + response = httpx.get(url, timeout=2.0) + if response.status_code == 200: + return + except Exception: # pragma: no cover - best effort poll + pass + time.sleep(0.1) + raise RuntimeError(f"Timed out waiting for {url}") + + +def _model_count(stats_url: str, model: str) -> int: + response = httpx.get(stats_url, timeout=5.0) + response.raise_for_status() + payload = response.json() + return int(payload["requests"]["by_model"].get(model, 0)) + + +def _send_probe(base_url: str, *, model: str, content: str) -> dict[str, object]: + response = httpx.post( + f"{base_url}/chat/completions", + headers={"Authorization": "Bearer test-key"}, + json={ + "model": model, + "messages": [{"role": "user", "content": content}], + }, + timeout=10.0, + ) + response.raise_for_status() + return response.json() + + +def _assert_delivery( + stack: _CodexProxyStack, + *, + base_url: str, + model: str, + content: str, +) -> None: + before = _model_count(stack.stats_url, model) + payload = _send_probe(base_url, model=model, content=content) + + deadline = time.time() + 10.0 + while time.time() < deadline: + if _model_count(stack.stats_url, model) >= before + 1: + break + time.sleep(0.1) + else: + raise AssertionError(f"Headroom never recorded model {model!r} in /stats") + + assert payload["choices"][0]["message"]["content"] == "mock completion from upstream" + assert any( + item["path"] == "/v1/chat/completions" + and isinstance(item.get("body"), dict) + and item["body"].get("model") == model + and item["body"].get("messages") == [{"role": "user", "content": content}] + for item in stack.upstream.requests + ) + + +def _codex_base_url_from_config(path: Path) -> str: + parsed = tomllib.loads(path.read_text(encoding="utf-8")) + assert parsed["model_provider"] == "headroom" + return str(parsed["model_providers"]["headroom"]["base_url"]) + + +def _manifest(tmp_path: Path, *, port: int) -> DeploymentManifest: + return DeploymentManifest( + profile="default", + preset="persistent-service", + runtime_kind="python", + supervisor_kind="service", + scope=ConfigScope.PROVIDER.value, + provider_mode="manual", + targets=["codex"], + port=port, + host="127.0.0.1", + backend="openai", + memory_db_path=str(tmp_path / "memory.db"), + tool_envs={"codex": {"OPENAI_BASE_URL": proxy_base_url(port)}}, + ) def test_codex_proxy_base_url_and_launch_env() -> None: @@ -10,3 +289,80 @@ def test_codex_proxy_base_url_and_launch_env() -> None: assert env["OPENAI_API_KEY"] == "sk-test" assert env["OPENAI_BASE_URL"] == "http://127.0.0.1:9999/v1" assert lines == ["OPENAI_BASE_URL=http://127.0.0.1:9999/v1"] + + +def test_codex_launch_env_routes_messages_through_headroom( + codex_proxy_stack: _CodexProxyStack, +) -> None: + env, _ = build_launch_env(codex_proxy_stack.proxy_port, {"OPENAI_API_KEY": "sk-test"}) + + _assert_delivery( + codex_proxy_stack, + base_url=str(env["OPENAI_BASE_URL"]), + model="wrap-launch-delivery-probe", + content="Verify temporary launch env traffic reaches Headroom.", + ) + + +def test_codex_install_env_routes_messages_through_headroom( + codex_proxy_stack: _CodexProxyStack, +) -> None: + env = build_install_env(port=codex_proxy_stack.proxy_port, backend="ignored") + + assert env == {"OPENAI_BASE_URL": codex_proxy_stack.base_url} + _assert_delivery( + codex_proxy_stack, + base_url=env["OPENAI_BASE_URL"], + model="persistent-install-delivery-probe", + content="Verify persistent install env traffic reaches Headroom.", + ) + + +def test_init_codex_config_routes_messages_through_headroom( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + codex_proxy_stack: _CodexProxyStack, +) -> None: + monkeypatch.chdir(tmp_path) + + init_cli._init_codex( + global_scope=False, + profile="init-local-demo", + port=codex_proxy_stack.proxy_port, + ) + + config_path = tmp_path / ".codex" / "config.toml" + content = config_path.read_text(encoding="utf-8") + assert 'env_key = "OPENAI_API_KEY"' not in content + assert "requires_openai_auth = true" in content + + _assert_delivery( + codex_proxy_stack, + base_url=_codex_base_url_from_config(config_path), + model="init-config-delivery-probe", + content="Verify init-generated Codex config sends traffic to Headroom.", + ) + + +def test_provider_scope_codex_config_routes_messages_through_headroom( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + codex_proxy_stack: _CodexProxyStack, +) -> None: + config_path = tmp_path / "config.toml" + config_path.write_text('model = "gpt-4o"\n', encoding="utf-8") + monkeypatch.setattr("headroom.providers.codex.install.codex_config_path", lambda: config_path) + + mutation = apply_provider_scope(_manifest(tmp_path, port=codex_proxy_stack.proxy_port)) + + assert mutation is not None + content = config_path.read_text(encoding="utf-8") + assert 'env_key = "OPENAI_API_KEY"' not in content + assert 'model_provider = "headroom"' in content + + _assert_delivery( + codex_proxy_stack, + base_url=_codex_base_url_from_config(config_path), + model="provider-scope-delivery-probe", + content="Verify persistent provider config sends traffic to Headroom.", + )