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(wrap): avoid duplicate top-level keys when injecting codex provider (#884)

Browse source 
## 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 <wangxiangyu7@lixiang.com>
This commit is contained in:
wangxiangyu7 2026-06-16 00:04:29 +08:00 • committed by GitHub
parent 6b4a101740
commit dd22cfd72a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 256 additions and 14 deletions
Download patch file Download diff file Expand all files Collapse all files

1
CHANGELOG.md
View file

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

144
headroom/cli/wrap.py
View file

@ -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: <original-value>``
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<value>[^"\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:
# <value>`` 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}")

125
tests/test_cli/test_wrap_codex.py
View file

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