From dd22cfd72ad9265c25a95ef5536dc3d17e85dbbf Mon Sep 17 00:00:00 2001 From: wangxiangyu7 Date: Tue, 16 Jun 2026 00:04:29 +0800 Subject: [PATCH] fix(wrap): avoid duplicate top-level keys when injecting codex provider (#884) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description `_inject_codex_provider_config` in `headroom/cli/wrap.py` unconditionally prepended a top-level block to `~/.codex/config.toml`: ```toml # --- Headroom proxy (auto-injected by headroom wrap codex) --- model_provider = "headroom" openai_base_url = "http://127.0.0.1:8787/v1" # --- end Headroom --- ``` If the user already had a top-level `model_provider` (or `openai_base_url`), the result was two top-level keys with the same name. That violates the TOML spec, and Codex refuses to start with `duplicate key`. This change makes the injector rewrite any pre-existing top-level `model_provider` / `openai_base_url` in place to the headroom values (keeping the user's original value in a `# was: …` trailing comment) and only emit the marker-delimited top-level block for keys the user has not declared. The pre-wrap snapshot mechanism is unchanged, so `headroom unwrap codex` still restores the file byte-for-byte. Closes #883 ## 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` - New helper `_redirect_existing_top_level_keys(content, port)`: rewrites existing top-level `model_provider` / `openai_base_url` lines to the headroom values and preserves the previous value in a trailing `# was: …` comment. - New helper `_has_redirectable_top_level_key(content, key)`: cheap predicate for the two redirectable keys. - New helper `_build_top_level_block(user_content)`: emits a marker-delimited block containing only the redirectable keys the user has **not** already declared (declared ones are rewritten in place instead, avoiding the TOML duplicate-key error). - `_inject_codex_provider_config` now rewrites declared keys in place and only prepends the marker block for the remaining keys; `requires_openai_auth` handling (#406) is preserved. - `tests/test_cli/test_wrap_codex.py` - New class `TestInjectAvoidsDuplicateTopLevelKeys` (TOML-validity after wrap on a config already declaring a provider, original-value preservation in a `# was:` comment, idempotent re-wrap with a port change, marker-block fallback on an empty file, snapshot-based unwrap restoration). The TOML-validity test parses the wrapped file with `tomllib.loads`, which fails before the fix and passes after. - `CHANGELOG.md` - Added entry under `## Unreleased` → `### Bug Fixes`. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom --ignore-missing-imports`) - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text $ uv run pytest tests/test_cli/test_wrap_codex.py -q ======================== 52 passed, 1 warning in 5.28s ========================= $ uv run ruff check headroom/cli/wrap.py tests/test_cli/test_wrap_codex.py All checks passed! $ uv run ruff format --check headroom/cli/wrap.py tests/test_cli/test_wrap_codex.py 2 files already formatted $ mypy headroom --ignore-missing-imports Success: no issues found in 358 source files ``` ## Real Behavior Proof - Environment: macOS 24.6.0, Python 3.13.3, branch `fix/wrap-codex-no-duplicate-keys` rebased onto upstream `main`, Codex CLI config at `~/.codex/config.toml`. - Exact command / steps: Seed a user config matching the bug report (`model_provider = "ccswitch"` + `openai_base_url = "…"` + `[model_providers.ccswitch]`), run the same path `headroom wrap codex` takes (`_inject_codex_provider_config(8787)`), then parse the result with `tomllib.loads(...)` and run `headroom unwrap codex`. - Observed result: On patched code the wrapped `config.toml` parses cleanly — exactly one `model_provider` and one `openai_base_url` remain (the user's prior value preserved in a `# was: …` comment) and the `[model_providers.headroom]` table is present; `unwrap` restores the file byte-for-byte. On the unpatched code the same file raises `tomllib.TOMLDecodeError` (duplicate key). Full suite: 52 passed, ruff lint + format clean (see Test Output). - Not tested: End-to-end launch of the Codex CLI against a live proxy (no Codex e2e harness in this sandbox); Windows / `$CODEX_HOME` override paths (covered only by existing tests). ## 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 - [x] 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 - [x] New and existing unit tests pass locally with my changes - [x] I have updated the CHANGELOG.md if applicable ## Screenshots (if applicable) N/A — CLI/config change, no UI. ## Additional Notes `ruff check .`, `ruff format --check .`, and `mypy headroom --ignore-missing-imports` all pass on the rebased branch. The diff stays narrow: `headroom/cli/wrap.py` + its tests + a CHANGELOG entry. Co-authored-by: wangxiangyu7 --- CHANGELOG.md | 1 + headroom/cli/wrap.py | 144 +++++++++++++++++++++++++++--- tests/test_cli/test_wrap_codex.py | 125 ++++++++++++++++++++++++++ 3 files changed, 256 insertions(+), 14 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e67e3d517..dba10dff1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,6 +20,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 * **ccr:** make retrieval store TTL configurable with `HEADROOM_CCR_TTL_SECONDS`, expose the effective TTL in `/v1/retrieve/stats`, and distinguish expired retrievals from missing hashes. * **proxy:** add native Bedrock `/model/{id}/converse-stream` route and forward it through the existing streaming EventStream/SSE pipeline. +* **wrap (codex):** fix `headroom wrap codex` producing a `config.toml` with duplicate top-level `model_provider` / `openai_base_url` keys (TOML-spec error) when the user had already configured their own provider. The injector now rewrites pre-existing top-level `model_provider` and `openai_base_url` lines in place — the previous value is kept in a `# was: …` trailing comment — instead of unconditionally prepending a duplicate, so `codex` can start against the proxy. The pre-wrap snapshot mechanism continues to byte-for-byte restore the original file on `headroom unwrap codex`. ## [0.25.0](https://github.com/chopratejas/headroom/compare/v0.24.0...v0.25.0) (2026-06-12) diff --git a/headroom/cli/wrap.py b/headroom/cli/wrap.py index e29c0f099..5f85b0c4f 100644 --- a/headroom/cli/wrap.py +++ b/headroom/cli/wrap.py @@ -943,6 +943,69 @@ def _strip_codex_headroom_blocks(content: str, *, remove_mcp: bool = False) -> s return content.lstrip("\n").rstrip() + "\n" if content.strip() else "" +# Top-level bare keys we redirect to headroom values when the user already +# has them set. Match the entire line (including any trailing comment) so +# we can rewrite it cleanly. Bare keys must precede any [section] in TOML, +# so a `^` anchor combined with `^[ \t]*key` is sufficient — table lines +# start with `[`, not with the key name. +_REDIRECTABLE_KEYS: tuple[str, ...] = ("model_provider", "openai_base_url") + + +def _redirect_existing_top_level_keys(content: str, port: int) -> str: + """Rewrite user-defined top-level keys so wrap does not create duplicates. + + Codex's ``config.toml`` rejects duplicate top-level keys (TOML spec), + which would break ``codex`` startup after ``headroom wrap codex`` runs + on a config that already declares its own ``model_provider`` or + ``openai_base_url``. + + For each redirectable key, if the user's line already sets it, replace + the value with the headroom one and append ``# was: `` + so the user can still see and recover their previous setting. The + snapshot taken in ``_snapshot_codex_config_if_unwrapped`` ensures the + pre-wrap file can be restored byte-for-byte on ``headroom unwrap + codex``. + + Returns the modified content. If no redirectable keys are present, + the content is returned unchanged and the caller should fall back to + prepending the marker-delimited top-level block (current behavior). + """ + import re # local import to match the module's existing convention + + if not content.strip(): + return content + + def _make_replacer(current_key: str, current_port: int) -> Callable[[re.Match[str]], str]: + def _replace(match: re.Match[str]) -> str: + original_value = match.group("value") + if current_key == "model_provider": + new_value = "headroom" + else: # openai_base_url + new_value = f"http://127.0.0.1:{current_port}/v1" + if original_value == new_value: + return match.group(0) + # Keep the user's original value in a trailing comment so they + # can see what was changed. This is metadata, not a TOML + # duplicate. + return f'{current_key} = "{new_value}" # was: {original_value}' + + return _replace + + redirected = content + for key in _REDIRECTABLE_KEYS: + pattern = re.compile(rf'(?m)^[ \t]*{re.escape(key)}[ \t]*=[ \t]*"(?P[^"\n]*)"[^\n]*') + redirected = pattern.sub(_make_replacer(key, port), redirected, count=1) + return redirected + + +def _has_redirectable_top_level_key(content: str, key: str) -> bool: + """Return True if ``content`` declares ``key = "..."`` as a top-level key.""" + import re # local import to match the module's existing convention + + pattern = re.compile(rf'(?m)^[ \t]*{key}[ \t]*=[ \t]*"[^"\n]*"') + return pattern.search(content) is not None + + def _snapshot_codex_config_if_unwrapped(config_file: Path, backup_file: Path) -> None: """Snapshot ``config.toml`` to ``backup_file`` before the first injection. @@ -1046,7 +1109,7 @@ def _apply_project_header_env(env: dict[str, str]) -> None: def _inject_codex_provider_config(port: int) -> None: """Inject a Headroom model provider into Codex's config.toml. - Two keys are written in the top-level block: + Two keys need to be in effect for the proxy to route all traffic: * ``model_provider = "headroom"`` — selects the custom provider for API-key mode traffic. @@ -1057,6 +1120,14 @@ def _inject_codex_provider_config(port: int) -> None: ``model_provider``, so without this override it bypasses the proxy and hits ``https://chatgpt.com/backend-api/codex`` directly. + If the user has not already declared these top-level keys, they are + added in a marker-delimited block at the top of the file. If the + user *has* declared one or both, the existing lines are rewritten + in place to the headroom values (with the previous value kept in a + ``# was: …`` trailing comment) so the resulting file stays TOML-valid + — TOML rejects duplicate top-level keys, which would break + ``codex`` startup. + Safe to call multiple times — the injected block is fully replaced on each call, so re-running with a different ``port`` updates the config. Before the first injection, the pre-wrap file is snapshotted to @@ -1071,13 +1142,10 @@ def _inject_codex_provider_config(port: int) -> None: # TOML keys must precede any [section]) and a provider-table block (at # the end). Each block has its own matching begin/end marker pair so # stripping them is unambiguous and never consumes user content that - # happens to sit between the two. - top_level_block = ( - 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" - ) + # happens to sit between the two. The top-level block is built + # dynamically below — it contains only keys the user has not already + # declared (we rewrite the existing ones in place to avoid TOML + # duplicate-key errors). # Emit requires_openai_auth only for ChatGPT-OAuth users (restores the # account menu); omitting it for API-key users avoids forcing an OAuth # login (#406). @@ -1099,6 +1167,29 @@ def _inject_codex_provider_config(port: int) -> None: f"{_CODEX_END_MARKER}\n" ) + # The two redirectable keys and their headroom target values. + _REDIRECT_TARGETS = { + "model_provider": "headroom", + "openai_base_url": f"http://127.0.0.1:{port}/v1", + } + + def _build_top_level_block(user_content: str) -> str: + """Build a marker-delimited block containing only the keys the user + has not already declared at the top level. For keys the user + *has* declared, the in-place rewrite below handles them. + """ + lines = [_CODEX_TOP_LEVEL_MARKER] + for key, value in _REDIRECT_TARGETS.items(): + if _has_redirectable_top_level_key(user_content, key): + continue + lines.append(f'{key} = "{value}"') + if len(lines) == 1: + # User already declared every redirectable key — no marker + # block needed (it would be empty). + return "" + lines.append(_CODEX_END_MARKER) + return "\n".join(lines) + "\n" + try: config_dir.mkdir(parents=True, exist_ok=True) @@ -1112,16 +1203,41 @@ def _inject_codex_provider_config(port: int) -> None: # the operation is idempotent and supports port changes. content = _strip_codex_headroom_blocks(content) - # Place the top-level key block at the very beginning of the file - # (bare TOML keys must precede any [section]) and the provider - # table at the end. User content, if any, sits between them. + # Bare top-level keys must precede any [section] in TOML, and + # TOML rejects duplicate top-level keys. Rewrite any existing + # top-level ``model_provider`` / ``openai_base_url`` in place + # to the headroom values; for keys the user has not declared, + # add them in a marker-delimited block at the top of the + # file. The original values are kept in a trailing ``# was: + # `` comment, and the snapshot mechanism guarantees + # byte-for-byte restoration on unwrap. user_content = content.strip() if user_content: - content = top_level_block + "\n" + user_content + "\n\n" + provider_section + redirected = _redirect_existing_top_level_keys(user_content, port) + top_block = _build_top_level_block(user_content) + if top_block: + content = top_block + "\n" + redirected + "\n\n" + provider_section + else: + content = redirected + "\n\n" + provider_section else: - content = top_level_block + "\n" + provider_section + # Empty user content — no keys to rewrite in place; emit + # the full marker block with both redirectable keys. + content = ( + f"{_CODEX_TOP_LEVEL_MARKER}\n" + f'model_provider = "{_REDIRECT_TARGETS["model_provider"]}"\n' + f'openai_base_url = "{_REDIRECT_TARGETS["openai_base_url"]}"\n' + f"{_CODEX_END_MARKER}\n" + f"\n{provider_section}" + ) else: - content = top_level_block + "\n" + provider_section + # No config file yet — same as the empty-content path. + content = ( + f"{_CODEX_TOP_LEVEL_MARKER}\n" + f'model_provider = "{_REDIRECT_TARGETS["model_provider"]}"\n' + f'openai_base_url = "{_REDIRECT_TARGETS["openai_base_url"]}"\n' + f"{_CODEX_END_MARKER}\n" + f"\n{provider_section}" + ) config_file.write_text(content) click.echo(f" Codex config: injected Headroom provider (WS + HTTP) into {config_file}") diff --git a/tests/test_cli/test_wrap_codex.py b/tests/test_cli/test_wrap_codex.py index d97fb4446..e451c0314 100644 --- a/tests/test_cli/test_wrap_codex.py +++ b/tests/test_cli/test_wrap_codex.py @@ -429,6 +429,131 @@ class TestSubscriptionRouting: assert "env_key" not in content +class TestInjectAvoidsDuplicateTopLevelKeys: + """Wrap must not produce a TOML-validity-breaking duplicate-key error. + + Codex's ``config.toml`` is parsed strictly: two top-level + ``model_provider = …`` (or two ``openai_base_url = …``) declarations + cause ``codex`` to refuse to start with + ``Error loading config.toml: …: …:1: duplicate key``. The injector + used to unconditionally prepend a top-level block, breaking any user + who had already configured their own provider (e.g. ``ccswitch``). + """ + + def test_inject_does_not_create_duplicate_model_provider( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + import tomllib # Python 3.11+ stdlib + + _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 = "ccswitch"\n' + 'openai_base_url = "http://llm-gateway-proxy/v1"\n' + 'model = "azure-gpt-5_5"\n' + "\n" + "[model_providers.ccswitch]\n" + 'name = "OpenAI"\n' + 'base_url = "http://llm-gateway-proxy/v1"\n' + 'wire_api = "responses"\n' + ) + + wrap_mod._inject_codex_provider_config(8787) + + content = config_file.read_text() + # The wrapped file must be TOML-parseable — duplicate keys were + # the failure mode the user reported. + tomllib.loads(content) + # No duplicate top-level key for either redirectable key. + assert content.count("model_provider =") == 1 + assert content.count("openai_base_url =") == 1 + # And the rewritten values are the headroom ones. + assert 'model_provider = "headroom"' in content + assert 'openai_base_url = "http://127.0.0.1:8787/v1"' in content + + @pytest.mark.parametrize("blank", ["", " ", "\n\t\n"]) + def test_redirect_existing_top_level_keys_noop_on_blank(self, blank: str) -> None: + # No redirectable keys to rewrite in blank/whitespace content — the + # helper returns it unchanged so the caller falls back to prepending + # the marker-delimited top-level block. + assert wrap_mod._redirect_existing_top_level_keys(blank, 8787) == blank + + def test_inject_preserves_user_value_in_trailing_comment( + 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 = "ccswitch"\nopenai_base_url = "http://llm-gateway-proxy/v1"\n' + ) + + wrap_mod._inject_codex_provider_config(8787) + + content = config_file.read_text() + # Original value kept in a comment so the user can recover it. + # The comment intentionally drops the surrounding quotes — the + # value is a single TOML string and the comment is human-facing. + assert "was: ccswitch" in content + assert "was: http://llm-gateway-proxy/v1" in content + + def test_inject_rewrap_updates_existing_redirected_keys( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + """Idempotent re-wrap on a config that already has top-level keys.""" + import tomllib + + _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 = "ccswitch"\n') + + wrap_mod._inject_codex_provider_config(8787) + wrap_mod._inject_codex_provider_config(9999) # port change + + content = config_file.read_text() + tomllib.loads(content) + assert content.count("model_provider =") == 1 + assert 'model_provider = "headroom"' in content + # Port updated in the openai_base_url we injected. + 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_inject_empty_file_still_uses_marker_block( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + """No existing top-level keys → fall back to the marker-delimited block.""" + _set_test_home(monkeypatch, tmp_path) + wrap_mod._inject_codex_provider_config(8787) + + content = (tmp_path / ".codex" / "config.toml").read_text() + assert wrap_mod._CODEX_TOP_LEVEL_MARKER in content + assert 'model_provider = "headroom"' in content + assert 'openai_base_url = "http://127.0.0.1:8787/v1"' in content + assert "[model_providers.headroom]" in content + + def test_unwrap_restores_prior_model_provider_after_rewrite( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + """The snapshot mechanism must still restore the pre-wrap state byte-for-byte.""" + _set_test_home(monkeypatch, tmp_path) + config_dir = tmp_path / ".codex" + config_dir.mkdir() + config_file = config_dir / "config.toml" + original = 'model_provider = "ccswitch"\nopenai_base_url = "http://llm-gateway-proxy/v1"\n' + config_file.write_text(original) + + wrap_mod._inject_codex_provider_config(8787) + + status, _ = wrap_mod._restore_codex_provider_config() + assert status == "restored" + assert config_file.read_text() == original + + # --------------------------------------------------------------------------- # Integration tests: full `headroom wrap codex` / `headroom unwrap codex` # ---------------------------------------------------------------------------