diff --git a/headroom/cli/wrap.py b/headroom/cli/wrap.py index 59194da31..fd89f0c75 100644 --- a/headroom/cli/wrap.py +++ b/headroom/cli/wrap.py @@ -1757,7 +1757,97 @@ def _apply_project_header_env(env: dict[str, str]) -> None: env["ANTHROPIC_CUSTOM_HEADERS"] = header_line -def _inject_codex_provider_config(port: int) -> None: +# Codex's own built-in providers plus Headroom's injected one — never treated +# as a "custom upstream to preserve" by _detect_custom_codex_upstream_base_url. +_CODEX_BUILTIN_PROVIDER_NAMES = frozenset({"openai", "anthropic", "azure", "headroom"}) + +# Header carrying a preserved custom upstream (freemodel.dev, LiteLLM, vLLM, +# ...) so the proxy forwards to it instead of the hardcoded OpenAI default. +# Codex's env_http_headers only accepts an env-var *name* per header (not a +# literal value), so the detected URL is exported into this env var by the +# `wrap codex` launch path — see its use in `codex()` below. +_UPSTREAM_BASE_URL_HEADER_NAME = "X-Headroom-Base-Url" +_UPSTREAM_BASE_URL_ENV_VAR = "HEADROOM_CODEX_UPSTREAM_BASE_URL" + + +def _codex_custom_provider_base_urls(content: str) -> dict[str, str]: + """Return ``{provider_name: base_url}`` for user-declared custom providers. + + Excludes Codex's built-ins (``openai``/``anthropic``/``azure``) and + Headroom's own ``headroom`` table. A table with no ``base_url`` line, or + one already pointing at Headroom's own localhost proxy (a leftover from a + prior wrap this pass hasn't stripped yet), is excluded too. + """ + import re + + tables: dict[str, str] = {} + for match in re.finditer( + r"(?ms)^[ \t]*\[model_providers\.(?P[^\]\s]+)\][ \t]*\n" + r"(?P.*?)(?=^[ \t]*\[|\Z)", + content, + ): + name = match.group("name") + if name in _CODEX_BUILTIN_PROVIDER_NAMES: + continue + base_match = re.search( + r'(?m)^[ \t]*base_url[ \t]*=[ \t]*"(?P[^"\n]*)"', match.group("body") + ) + if not base_match: + continue + url = base_match.group("url").strip().rstrip("/") + if not url or url.startswith(("http://127.0.0.1", "http://localhost")): + continue + tables[name] = url + return tables + + +def _detect_custom_codex_upstream_base_url(content: str) -> str | None: + """Return a user-configured custom provider ``base_url`` to preserve, if any. + + Codex lets users declare OpenAI-compatible gateways (LiteLLM, vLLM, + freemodel.dev, ...) under ``[model_providers.]`` and select one via + the top-level ``model_provider`` key. Before this, ``headroom wrap codex`` + unconditionally pointed the proxy's upstream OpenAI route at + ``api.openai.com``, silently discarding that selection — the user's + gateway API key then gets sent to OpenAI, which rejects it (#1614). + + If the top-level ``model_provider`` names one of the detected custom + tables, that selection wins unambiguously. This also covers the + already-wrapped case: once wrap has run once, the top-level key reads + ``model_provider = "headroom" # was: `` (see + ``_redirect_existing_top_level_keys``), so the original selection is + recovered from that trailing comment on re-wrap / port changes. + + Falls back to the sole candidate when exactly one custom table exists + and there is no (or no matching) top-level selection — the common case + from the bug report, where the table is declared but selection happens + via ``--profile`` rather than a static top-level key. Returns ``None`` + when there are multiple, un-selected candidates (ambiguous — guessing + wrong is worse than the prior default behavior) or none at all. + """ + import re + + candidates = _codex_custom_provider_base_urls(content) + if not candidates: + return None + + selected = re.search( + r'(?m)^[ \t]*model_provider[ \t]*=[ \t]*"(?P[^"\n]*)"' + r"(?:[ \t]*#[ \t]*was:[ \t]*(?P[^\r\n]*))?", + content, + ) + if selected: + was = (selected.group("was") or "").strip() + chosen = was or selected.group("name") + if chosen in candidates: + return candidates[chosen] + + if len(candidates) == 1: + return next(iter(candidates.values())) + return None + + +def _inject_codex_provider_config(port: int) -> str | None: """Inject a Headroom model provider into Codex's config.toml. Two keys need to be in effect for the proxy to route all traffic: @@ -1784,10 +1874,32 @@ def _inject_codex_provider_config(port: int) -> None: Before the first injection, the pre-wrap file is snapshotted to ``config.toml.headroom-backup`` so ``headroom unwrap codex`` can restore it byte-for-byte. + + Returns the custom upstream ``base_url`` preserved from an existing + ``[model_providers.*]`` table, if one was detected (#1614); ``None`` + otherwise. Callers that go on to launch Codex should export this value + into ``HEADROOM_CODEX_UPSTREAM_BASE_URL`` (the injected + ``env_http_headers`` entry maps it to the ``X-Headroom-Base-Url`` header, + which the proxy's OpenAI HTTP handlers honor over the hardcoded + ``api.openai.com`` default) — see its use in ``codex()`` below. """ config_file, backup_file = _codex_config_paths() config_dir = config_file.parent + # Detect an existing custom OpenAI-compatible provider BEFORE building the + # injected block below, so it can be preserved as the upstream the proxy + # forwards to instead of silently rerouting to api.openai.com (#1614). + # Best-effort: any read/parse failure just means nothing is preserved, + # matching prior behavior. + custom_upstream_base_url: str | None = None + if config_file.exists(): + try: + custom_upstream_base_url = _detect_custom_codex_upstream_base_url( + _read_text(config_file) + ) + except OSError: + custom_upstream_base_url = None + # The injected content is split into two self-contained, marker-delimited # blocks: a top-level key block (at the start of the file, because bare # TOML keys must precede any [section]) and a provider-table block (at @@ -1803,6 +1915,15 @@ def _inject_codex_provider_config(port: int) -> None: requires_openai_auth = ( "requires_openai_auth = true\n" if codex_uses_chatgpt_auth(config_dir / "auth.json") else "" ) + # Per-project savings: Codex sends the X-Headroom-Project header only + # when the mapped env var (HEADROOM_PROJECT, set by `headroom wrap + # codex`) exists at Codex runtime. When a custom upstream was detected, + # add a second entry so Codex also sends X-Headroom-Base-Url — the proxy + # forwards there instead of api.openai.com (#1614). + env_http_headers_map = {_PROJECT_HEADER_NAME: "HEADROOM_PROJECT"} + if custom_upstream_base_url: + env_http_headers_map[_UPSTREAM_BASE_URL_HEADER_NAME] = _UPSTREAM_BASE_URL_ENV_VAR + env_http_headers_toml = ", ".join(f'"{k}" = "{v}"' for k, v in env_http_headers_map.items()) provider_section = ( f"{_CODEX_TOP_LEVEL_MARKER}\n" "[model_providers.headroom]\n" @@ -1810,11 +1931,9 @@ def _inject_codex_provider_config(port: int) -> None: f'base_url = "http://127.0.0.1:{port}/v1"\n' f"supports_websockets = true\n" f"{requires_openai_auth}" - # Per-project savings: Codex sends the header only when the mapped - # env var (HEADROOM_PROJECT, set by `headroom wrap codex`) exists at - # Codex runtime. Inline table keeps the key inside this section so + # Inline table keeps the key inside this section so # _strip_codex_headroom_blocks removes it with the rest of the block. - f'env_http_headers = {{ "{_PROJECT_HEADER_NAME}" = "HEADROOM_PROJECT" }}\n' + f"env_http_headers = {{ {env_http_headers_toml} }}\n" f"{_CODEX_END_MARKER}\n" ) @@ -1893,11 +2012,20 @@ def _inject_codex_provider_config(port: int) -> None: _write_text(config_file, content) click.echo(f" Codex config: injected Headroom provider (WS + HTTP) into {config_file}") + if custom_upstream_base_url: + click.echo( + f" Codex config: preserving existing custom upstream " + f"{custom_upstream_base_url} (from a pre-existing [model_providers.*] " + "base_url)" + ) # Pull existing native threads into the headroom-provider menu so Codex's # history list stays whole once it routes through Headroom. Best-effort. retag_to_headroom(_codex_home_dir()) except Exception as e: click.echo(f" Warning: could not update Codex config: {e}") + return None + + return custom_upstream_base_url def _restore_codex_provider_config() -> tuple[str, Path]: @@ -4701,7 +4829,14 @@ def codex( # 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(actual_port) + _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"))) diff --git a/tests/test_cli/test_wrap_codex.py b/tests/test_cli/test_wrap_codex.py index 7a3122a55..2d004e6d7 100644 --- a/tests/test_cli/test_wrap_codex.py +++ b/tests/test_cli/test_wrap_codex.py @@ -673,6 +673,157 @@ class TestSubscriptionRouting: assert "env_key" not in content +# --------------------------------------------------------------------------- +# Custom upstream preservation (#1614): wrap must not silently reroute a +# pre-existing custom [model_providers.*] base_url to api.openai.com. +# --------------------------------------------------------------------------- + + +class TestDetectCustomCodexUpstreamBaseUrl: + """Unit tests for the detection helper used by ``_inject_codex_provider_config``.""" + + def test_no_config_returns_none(self) -> None: + assert wrap_mod._detect_custom_codex_upstream_base_url("") is None + + def test_no_custom_provider_returns_none(self) -> None: + content = ( + 'model_provider = "openai"\n\n' + "[model_providers.openai]\n" + 'base_url = "https://api.openai.com/v1"\n' + ) + assert wrap_mod._detect_custom_codex_upstream_base_url(content) is None + + def test_sole_candidate_used_without_explicit_selection(self) -> None: + """Matches the #1614 repro: a custom table with no static top-level pin.""" + content = ( + "[model_providers.freemodel]\n" + 'base_url = "https://api.freemodel.dev"\n' + 'wire_api = "responses"\n' + ) + assert ( + wrap_mod._detect_custom_codex_upstream_base_url(content) == "https://api.freemodel.dev" + ) + + def test_explicit_top_level_selection_wins(self) -> None: + content = ( + 'model_provider = "freemodel"\n\n' + "[model_providers.freemodel]\n" + 'base_url = "https://api.freemodel.dev"\n\n' + "[model_providers.other]\n" + 'base_url = "https://api.other.example"\n' + ) + assert ( + wrap_mod._detect_custom_codex_upstream_base_url(content) == "https://api.freemodel.dev" + ) + + def test_was_comment_recovers_selection_on_rewrap(self) -> None: + """After a prior wrap, model_provider reads 'headroom # was: freemodel'.""" + content = ( + 'model_provider = "headroom" # was: freemodel\n\n' + "[model_providers.freemodel]\n" + 'base_url = "https://api.freemodel.dev"\n' + ) + assert ( + wrap_mod._detect_custom_codex_upstream_base_url(content) == "https://api.freemodel.dev" + ) + + def test_ambiguous_multiple_candidates_returns_none(self) -> None: + content = ( + "[model_providers.freemodel]\n" + 'base_url = "https://api.freemodel.dev"\n\n' + "[model_providers.other]\n" + 'base_url = "https://api.other.example"\n' + ) + assert wrap_mod._detect_custom_codex_upstream_base_url(content) is None + + def test_builtin_provider_tables_excluded(self) -> None: + content = ( + 'model_provider = "openai"\n\n' + "[model_providers.openai]\n" + 'base_url = "https://api.openai.com/v1"\n\n' + "[model_providers.anthropic]\n" + 'base_url = "https://api.anthropic.com/v1"\n' + ) + assert wrap_mod._detect_custom_codex_upstream_base_url(content) is None + + def test_own_headroom_table_excluded(self) -> None: + content = ( + 'model_provider = "headroom"\n\n' + "[model_providers.headroom]\n" + 'base_url = "http://127.0.0.1:8787/v1"\n' + ) + assert wrap_mod._detect_custom_codex_upstream_base_url(content) is None + + +class TestInjectPreservesCustomUpstreamBaseUrl: + """``_inject_codex_provider_config`` must preserve a pre-existing custom + provider's ``base_url`` instead of silently rerouting to api.openai.com.""" + + def test_inject_returns_and_carries_custom_base_url( + 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" + config_file.write_text( + "[model_providers.freemodel]\n" + 'base_url = "https://api.freemodel.dev"\n' + 'wire_api = "responses"\n', + encoding="utf-8", + ) + + result = wrap_mod._inject_codex_provider_config(8787) + + assert result == "https://api.freemodel.dev" + content = config_file.read_text(encoding="utf-8") + parsed = tomllib.loads(content) + headers = parsed["model_providers"]["headroom"]["env_http_headers"] + assert ( + headers[wrap_mod._UPSTREAM_BASE_URL_HEADER_NAME] == wrap_mod._UPSTREAM_BASE_URL_ENV_VAR + ) + # The user's own table is left untouched — only headroom's own is managed. + assert parsed["model_providers"]["freemodel"]["base_url"] == "https://api.freemodel.dev" + + def test_inject_without_custom_provider_returns_none( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + _set_test_home(monkeypatch, tmp_path) + + result = wrap_mod._inject_codex_provider_config(8787) + + assert result is None + content = (tmp_path / ".codex" / "config.toml").read_text(encoding="utf-8") + assert wrap_mod._UPSTREAM_BASE_URL_HEADER_NAME not in content + + def test_preserved_upstream_survives_rewrap_and_port_change( + 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" + config_file.write_text( + 'model_provider = "freemodel"\n\n' + "[model_providers.freemodel]\n" + 'base_url = "https://api.freemodel.dev"\n', + encoding="utf-8", + ) + + first = wrap_mod._inject_codex_provider_config(8787) + second = wrap_mod._inject_codex_provider_config(9999) # port change / re-wrap + + assert first == "https://api.freemodel.dev" + assert second == "https://api.freemodel.dev" + content = config_file.read_text(encoding="utf-8") + parsed = tomllib.loads(content) + assert parsed["model_providers"]["headroom"]["base_url"] == "http://127.0.0.1:9999/v1" + headers = parsed["model_providers"]["headroom"]["env_http_headers"] + assert ( + headers[wrap_mod._UPSTREAM_BASE_URL_HEADER_NAME] == wrap_mod._UPSTREAM_BASE_URL_ENV_VAR + ) + + class TestInjectAvoidsDuplicateTopLevelKeys: """Wrap must not produce a TOML-validity-breaking duplicate-key error.