From d236b27c607c980a224e24a4e7c0282aea487cc4 Mon Sep 17 00:00:00 2001 From: Abhay Singh Date: Tue, 14 Jul 2026 21:32:24 +0530 Subject: [PATCH] fix(wrap/codex): export the detected custom upstream base URL (#2125) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description `headroom wrap codex` detects a user's custom upstream gateway but never tells Codex to use it, so the user's gateway key is sent to `api.openai.com`. `_inject_codex_provider_config` handles a Codex user who has an OpenAI-compatible gateway declared in `~/.codex/config.toml`, e.g. ```toml model_provider = "freemodel" [model_providers.freemodel] base_url = "https://api.freemodel.dev" ``` It injects the Headroom provider with `env_http_headers = { ... "X-Headroom-Base-Url" = "HEADROOM_CODEX_UPSTREAM_BASE_URL" }` and **returns the preserved upstream URL** so the caller can export it. Its docstring even says: *"Callers that go on to launch Codex should export this value into `HEADROOM_CODEX_UPSTREAM_BASE_URL`."* But `_prepare_codex_wrap_state` called it as a bare statement and discarded the return, and `_run_codex_wrap` / `_build_codex_launch_env` only ever set `OPENAI_BASE_URL`. A repo-wide grep confirms `HEADROOM_CODEX_UPSTREAM_BASE_URL` (`_UPSTREAM_BASE_URL_ENV_VAR`) is never assigned into any process env — it appears only at its definition and in that docstring. Since Codex only emits the `X-Headroom-Base-Url` header when the env var exists, the header is omitted, the proxy's OpenAI handler falls back to its hardcoded `https://api.openai.com`, and the user's `freemodel.dev` key is sent to OpenAI, which rejects it. This is a regression: the wiring existed in the original `#1614` fix (`_codex_custom_upstream = _inject_codex_provider_config(...)` then `env[_UPSTREAM_BASE_URL_ENV_VAR] = ...`) and was dropped by a later refactor that extracted `_prepare_codex_wrap_state`. ## Fix Restore the wiring: `_prepare_codex_wrap_state` now captures and returns `_inject_codex_provider_config`'s value, and `_run_codex_wrap` exports it into the launch env (`env[_UPSTREAM_BASE_URL_ENV_VAR] = custom_upstream`) when it is non-None and not already set, so a user-provided value still wins. Closes # ## 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 - `headroom/cli/wrap.py`: `_prepare_codex_wrap_state` returns the detected custom upstream URL; `_run_codex_wrap` exports it into the launch env (and its display list) when set. - `tests/test_cli/test_wrap_codex.py`: add `TestCodexLaunchExportsCustomUpstream` — drives `_run_codex_wrap` with mocked prepare/launch and asserts the launch env carries `HEADROOM_CODEX_UPSTREAM_BASE_URL` when a custom upstream is detected, and does not when there isn't one. - `CHANGELOG.md`: Bug Fixes entry. ## Testing - [x] Unit tests pass (`uv run --extra dev pytest tests/test_cli/test_wrap_codex.py::TestCodexLaunchExportsCustomUpstream -q`) - [x] Linting passes (`uvx ruff@0.15.17 check headroom/cli/wrap.py tests/test_cli/test_wrap_codex.py headroom/memory/factory.py`) - [x] Type checking passes (`uvx mypy==1.20.2 headroom/memory/factory.py`) - [x] New tests added for new functionality - [ ] Manual testing performed ### Test Output ```text $ uvx ruff@0.15.17 check headroom/cli/wrap.py tests/test_cli/test_wrap_codex.py All checks passed! $ python -m py_compile headroom/cli/wrap.py tests/test_cli/test_wrap_codex.py OK ``` ## Real Behavior Proof - Environment: Windows 11, Python 3.12, `uvx ruff@0.15.17`. Importing `headroom` pulls in the torch/transformers stack and running the CLI test locally OOM-kills this box, so I verified the wiring with a dependency-free script that models prepare -> run -> the proxy's upstream fallback, and left the full pytest (including the new CLI test) to CI. - Exact command / steps: modelled the old flow (inject return discarded) and the new flow (return exported into the launch env), then applied the proxy's rule that a missing `HEADROOM_CODEX_UPSTREAM_BASE_URL` falls back to `api.openai.com`. - Observed result: old effective upstream is `https://api.openai.com` (the gateway key is misrouted); new effective upstream is `https://api.freemodel.dev` (the user's gateway). The new CLI test asserts the launch env carries the var when a custom upstream is present and omits it otherwise. - Not tested: a live Codex process reading the env and emitting the header; full local `pytest` deferred to CI (OOM, per above). ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes - [x] I have updated the CHANGELOG.md if applicable ## Additional Notes Merged current `main` to pick up the repository-wide mypy cache-key annotation fix, then verified the focused regression locally. the change threads one return value through two functions and exports it, verified by the wiring proof and a new CLI test that drives `_run_codex_wrap` with the heavy prepare/launch steps mocked so only the env-export logic is exercised. Co-authored-by: JerrettDavis Co-authored-by: Tejas Chopra --- CHANGELOG.md | 1 + headroom/cli/wrap.py | 25 ++++++++++--- tests/test_cli/test_wrap_codex.py | 60 +++++++++++++++++++++++++++++++ 3 files changed, 82 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 41f9d9b1a..c3559494e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -133,6 +133,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 * **tokenizers:** resolve HuggingFace tokenizer names by the most-specific prefix. `get_tokenizer_name` scanned `MODEL_TO_TOKENIZER` in dict-insertion order and returned the first key the model merely starts with, so a short family key shadowed a more-specific one — `qwen2-7b-instruct` matched `qwen` before `qwen2`/`qwen2-7b` and resolved to the Qwen1 tokenizer (a different vocabulary, hence wrong token counts); `qwen2.5-*` and `deepseek-v2.x` were mis-resolved the same way. It now picks the longest matching prefix, mirroring the order-dependent-prefix guard the sibling tiktoken `get_encoding_for_model` already documents. * **pricing:** map retired `claude-3-sonnet-20240229` to a Sonnet-tier price instead of Haiku. When LiteLLM's cost DB lacks the retired model, resolution falls through to `MODEL_ALIASES`, which pointed Claude 3 Sonnet (a $3/$15-per-1M model) at `claude-3-haiku-20240307` ($0.25/$1.25) — a different tier that underpriced every cost/savings figure for that model ~12x on both input and output. It now aliases to `claude-sonnet-4-20250514`, the same-price target the other retired-Sonnet aliases already use. * **cache/semantic:** don't evict an unrelated entry when re-storing a key that is already cached. `SemanticCache.put` ran its at-capacity eviction loop before computing the entry's key, so overwriting a key that was already present (a duplicate or retried store) still evicted the LRU-oldest distinct entry even though an in-place update grows nothing. That silently dropped a live entry and turned a later lookup for it into a false cache miss. The key is now computed first and the eviction loop only runs when the key is genuinely new (mirroring `CompressionCache.store_compressed`, which deletes-then-inserts). +* **wrap/codex:** export the detected custom upstream base URL so Codex actually routes to it. `_inject_codex_provider_config` detects an OpenAI-compatible gateway declared in `~/.codex/config.toml` and writes an `X-Headroom-Base-Url` header mapped to `HEADROOM_CODEX_UPSTREAM_BASE_URL`, and it returns that URL for the caller to export. But `_prepare_codex_wrap_state` discarded the return and `_run_codex_wrap` never set the env var, so Codex emitted no header, the proxy fell back to `api.openai.com`, and the user's gateway key was sent to OpenAI (which rejects it). This restores the wiring a later refactor dropped: `_prepare_codex_wrap_state` now returns the URL and `_run_codex_wrap` exports it into the launch env when set (a user-provided value still wins) (regression of [#1614](https://github.com/chopratejas/headroom/issues/1614)). * **cache:** normalize embeddings before the semantic dynamic-content similarity check. `SemanticDetector` scored sentences with `np.dot` against exemplar embeddings and compared the result to `semantic_threshold` (a 0-1 cosine value), but `sentence_transformers.encode(..., convert_to_numpy=True)` does not normalize, so the dot product was an unbounded inner product (vector norms ~5-15) rather than a cosine similarity. Nearly every sentence cleared the 0.7 threshold, so the semantic tier flagged almost all text as dynamic and stripped static content, busting the cache it is meant to protect (a standalone repro scores an unrelated sentence, true cosine ~0.1, at a raw dot of ~9). Both `encode` calls now pass `normalize_embeddings=True`, matching the siblings in `prediction/feature_extractor.py` and `memory/adapters/embedders.py`, so the dot product is a true cosine in [-1, 1]. * **subscription:** cap `HeadroomContribution.efficiency_pct` at a real removal ratio so it can't exceed 100%. The numerator used `total_saved()` (which includes `tokens_saved_cache_reads`) while the denominator `raw_without_headroom()` excludes cache reads, so a contribution with large prefix-cache reads and small forwarded input reported impossible values (e.g. `tokens_submitted=100`, `tokens_saved_cache_reads=1000` → `1000.0%` on the dashboard). Cache reads are a provider-side discount on tokens that were still forwarded, not tokens Headroom removed, so the ratio now uses the sibling `compression_saved()` (compression + CLI filtering) as the numerator — bounded by its own denominator. `total_saved()` is unchanged for its other callers. * **proxy/anthropic:** cache the response under the same messages it was looked up by. The non-streaming `/v1/messages` path snapshots the scalar cache-key fields (system, tools, etc.) once before upstream to avoid post-mutation key drift (#327), but `messages` — the primary key component — was passed live at both `cache.get` and `cache.set`. Between the two, `messages` is reassigned by the enterprise security scan, the `pre_compress` hook, and image compression, so when any of those fired the response was stored under a different key than it was read by: the response cache never hit and accumulated unreachable entries until eviction. The lookup messages are now snapshotted alongside the other key fields and reused verbatim at `cache.set`. diff --git a/headroom/cli/wrap.py b/headroom/cli/wrap.py index 5dd182fbe..d69ac7ce3 100644 --- a/headroom/cli/wrap.py +++ b/headroom/cli/wrap.py @@ -4712,8 +4712,15 @@ def _prepare_codex_wrap_state( memory: bool, verbose: bool, rtk_home: Path | None = None, -) -> None: - """Prepare the active Codex home for a wrap or prepare-only invocation.""" +) -> str | None: + """Prepare the active Codex home for a wrap or prepare-only invocation. + + Returns the custom upstream base URL detected in the user's Codex config + (or None). Callers that launch Codex must export this into + ``HEADROOM_CODEX_UPSTREAM_BASE_URL`` so the injected provider's + ``X-Headroom-Base-Url`` header carries it; otherwise the proxy falls back to + ``api.openai.com`` and the user's gateway key is sent to the wrong host. + """ # 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 @@ -4808,9 +4815,10 @@ def _prepare_codex_wrap_state( # 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) + custom_upstream = _inject_codex_provider_config(port) if memory: _inject_memory_mcp_config(os.environ.get("USER", os.environ.get("USERNAME", "default"))) + return custom_upstream def _run_codex_wrap( @@ -4854,7 +4862,7 @@ def _run_codex_wrap( active_codex_home = _codex_home_dir() with _codex_session_home_overlay() as session_codex_home: - _prepare_codex_wrap_state( + custom_upstream = _prepare_codex_wrap_state( port=port, no_rtk=no_rtk, no_mcp=no_mcp, @@ -4868,6 +4876,15 @@ def _run_codex_wrap( env, env_vars_display = _build_codex_launch_env(port, os.environ) + # Export the detected custom upstream so Codex actually emits the + # X-Headroom-Base-Url header the injected provider declares. Without + # this the proxy falls back to api.openai.com and a user with an + # OpenAI-compatible gateway in ~/.codex/config.toml has their gateway + # key sent to OpenAI (regression of #1614). A user-set value wins. + if custom_upstream and _UPSTREAM_BASE_URL_ENV_VAR not in env: + env[_UPSTREAM_BASE_URL_ENV_VAR] = custom_upstream + env_vars_display.append(f"{_UPSTREAM_BASE_URL_ENV_VAR}={custom_upstream}") + # 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. diff --git a/tests/test_cli/test_wrap_codex.py b/tests/test_cli/test_wrap_codex.py index 0dea434f5..9c7c8b50a 100644 --- a/tests/test_cli/test_wrap_codex.py +++ b/tests/test_cli/test_wrap_codex.py @@ -1953,3 +1953,63 @@ class TestCodexPortResolution: assert call_kw.get("port") == 8787 assert call_kw.get("no_proxy") is False assert call_kw.get("prepare_only") is False + + +class TestCodexLaunchExportsCustomUpstream: + """`_run_codex_wrap` must export the detected custom upstream base URL into + the launch env so Codex emits the ``X-Headroom-Base-Url`` header. Otherwise + the proxy falls back to api.openai.com and the user's gateway key is sent to + the wrong host (regression of #1614).""" + + def _launch_env(self, monkeypatch, tmp_path, *, custom_upstream): + from contextlib import contextmanager + + captured: dict = {} + + monkeypatch.setattr(wrap_mod.shutil, "which", lambda name: "/usr/bin/codex") + monkeypatch.setattr(wrap_mod, "_codex_home_dir", lambda: tmp_path) + + @contextmanager + def _fake_overlay(): + yield tmp_path / "session" + + monkeypatch.setattr(wrap_mod, "_codex_session_home_overlay", _fake_overlay) + # Stand in for the heavy prepare step; only its return value matters here. + monkeypatch.setattr(wrap_mod, "_prepare_codex_wrap_state", lambda **kwargs: custom_upstream) + + def _fake_launch(*, env, **kwargs): + captured["env"] = env + + monkeypatch.setattr(wrap_mod, "_launch_tool", _fake_launch) + + wrap_mod._run_codex_wrap( + port=8787, + no_rtk=True, + no_mcp=True, + no_tokensave=True, + serena=False, + no_serena=True, + code_graph=False, + no_proxy=True, + learn=False, + memory=False, + backend=None, + anyllm_provider=None, + region=None, + verbose=False, + prepare_only=False, + codex_args=(), + ) + return captured["env"] + + def test_custom_upstream_exported_into_launch_env( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + env = self._launch_env(monkeypatch, tmp_path, custom_upstream="https://api.freemodel.dev") + assert env[wrap_mod._UPSTREAM_BASE_URL_ENV_VAR] == "https://api.freemodel.dev" + + def test_no_custom_upstream_leaves_env_var_unset( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + env = self._launch_env(monkeypatch, tmp_path, custom_upstream=None) + assert wrap_mod._UPSTREAM_BASE_URL_ENV_VAR not in env