akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix: preserve Codex OAuth proxy delivery

Browse source 
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>
This commit is contained in:
JerrettDavis 2026-05-05 21:03:42 -05:00
parent 4745219901
commit 06428d20fd
10 changed files with 679 additions and 30 deletions
Download patch file Download diff file Expand all files Collapse all files

10
e2e/init/run.py
View file

@ -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"))

144
e2e/wrap/run.py
View file

@ -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)

14
headroom/cli/init.py
View file

@ -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(

38
headroom/cli/wrap.py
View file

@ -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}")

34
headroom/providers/codex/install.py
View file

@ -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()

2
tests/test_cli/test_init_cli.py
View file

@ -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:

80
tests/test_cli/test_wrap_codex.py
View file

@ -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`
# ---------------------------------------------------------------------------

4
tests/test_install/test_providers.py
View file

@ -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(

27
tests/test_provider_codex_install.py Normal file
View file

@ -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

356
tests/test_provider_codex_runtime.py
View file

@ -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.",
)