From 415e03c1688a5a1a919872335257df1b6b424c42 Mon Sep 17 00:00:00 2001 From: Abhay Singh Date: Mon, 13 Jul 2026 19:07:32 +0530 Subject: [PATCH] fix(mcp/codex): don't clobber an unparseable/non-table config.toml (#2062) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description `CodexRegistrar.register_server` (`headroom/mcp_registry/codex.py`) guards against clobbering a user-managed `[mcp_servers.]` entry — but **only inside the `if existing is not None` branches**. `existing` comes from `get_server`, which returns `None` in two cases that are *not* "nothing there": 1. the `config.toml` is **unparseable** (`_load_toml` catches `TOMLDecodeError` and returns `{}`), and 2. `mcp_servers` (or `mcp_servers.`) is present but **not a table** (`get_server` returns `None` via its `isinstance` guards). With `existing is None`, all three protection branches are skipped and control falls straight to `_write_block`, which blindly appends a fresh `[mcp_servers.]` table. So for a **valid** TOML file like: ```toml [mcp_servers] headroom = "not-a-table" ``` `register_server(headroom_spec)` appends `[mcp_servers.headroom]`, producing a file that defines `mcp_servers.headroom` **both** as a string and as a table — a duplicate key that `tomllib`/codex then reject, **corrupting a previously-valid user config**. The unparseable-file case similarly appends our block into a file that can't be parsed. This is the exact Codex sibling of the claude (#1660) and opencode (#1661) clobber-guard fixes; codex never received it. Closes: no issue filed — found while auditing the registrars for the #1660/#1661 class. ## Fix Add `_unmergeable_reason(name)` — returns why the existing file can't be safely merged (present but unparseable, or a non-table `mcp_servers` / `mcp_servers.`), else `None`. In `register_server`, when `existing is None`, refuse with `RegisterStatus.FAILED` (leaving the file untouched) instead of appending. Absent/empty/valid configs are unaffected. ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) ## Changes Made - `headroom/mcp_registry/codex.py`: add `_unmergeable_reason`; refuse in `register_server` when the existing config is unparseable or defines a non-table `mcp_servers`/`mcp_servers.`. - `tests/test_mcp_registry/test_codex_registrar.py`: add tests for unparseable TOML, non-table `mcp_servers.headroom`, and non-table `mcp_servers` (all refuse + file untouched). ## Testing - [x] New regression tests added (`tests/test_mcp_registry/test_codex_registrar.py`) - [x] Linting/formatting clean — run with the CI-pinned `ruff==0.15.17` - [ ] Full `pytest` deferred to CI (local-OOM reason below). ```text $ uvx ruff@0.15.17 check headroom/mcp_registry/codex.py tests/test_mcp_registry/test_codex_registrar.py All checks passed! ``` ## Real Behavior Proof - Environment: Windows 11, Python 3.10, headroom from this branch. Importing `headroom` pulls in the torch/transformers stack and a full `pytest` gets OOM-killed on this box, so I verified the `_unmergeable_reason` logic with a dependency-free script (stdlib `tomllib`) and left the full pytest to CI. - Exact command / steps: ran the two clobber cases (non-table entry, unparseable TOML) and the safe cases (absent/empty/valid/other-server) through the guard. - Observed result: the guard refuses exactly the two corrupting cases and allows every valid config: ```text REFUSE [non-table entry (valid TOML)]: non-table mcp_servers.headroom REFUSE [unparseable TOML]: not valid TOML (Invalid value (at line 1, column 8)) ALLOW [absent]: reason=None ALLOW [empty]: reason=None ALLOW [valid, no mcp_servers]: reason=None ALLOW [valid, mcp_servers table w/ other server]: reason=None CODEX CLOBBER-GUARD VERIFIED (refuses non-table/unparseable; allows valid configs) ``` - Not tested: a live `codex` launch reading the config (mocked in the registrar tests). 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 — ran lint + a standalone logic check; full pytest deferred to CI (local OOM, disclosed above) - [x] I have updated the CHANGELOG.md if applicable ## Additional Notes - Completes the registrar clobber-guard trio (claude #1660, opencode #1661, codex here); no new dependencies. - @JerrettDavis tagging you — same class you already reviewed for claude/opencode, just the codex side. Thanks! --------- Co-authored-by: JerrettDavis --- CHANGELOG.md | 1 + headroom/mcp_registry/codex.py | 46 +++++++++++++++++++ .../test_mcp_registry/test_codex_registrar.py | 44 ++++++++++++++++++ 3 files changed, 91 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1d9b4a6ec..e578229d3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -104,6 +104,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 * **tokenizers:** price dense scripts (CJK/Kana/Hangul) in the fixed-ratio estimator path. `EstimatingTokenCounter.count_text` applied the CJK correction only on the auto-detect path; its fixed-ratio early return divided by the Latin ratio with no adjustment. The registry builds every provider-calibrated counter with a fixed ratio (Anthropic 3.5, Google 4.0, Cohere 4.0, Moonshot 3.1), and the Anthropic/Gemini proxy handlers count via `get_tokenizer(model).count_messages`, so a CJK-heavy context read as ~40-55% of its true token size — it could fall under the size/backpressure gates and skip compression, and every `x-headroom-tokens-before` metric for CJK traffic was materially wrong. The fixed-ratio path now applies the same dense-script split. * **ccr:** don't crash `parse_tool_call` on a CCR tool call whose arguments aren't an object. For the OpenAI/`openai_responses` shape the arguments are `json.loads`-decoded and only `JSONDecodeError` was caught, so a model that emitted `arguments='[]'`/`'"abc"'`/`'123'` (decoding to a list/str/number) — or a non-dict Anthropic `input` — reached `input_data.get("hash")` and raised `AttributeError`; a null `arguments` raised an uncaught `TypeError` from `json.loads(None)`. Both are now handled: the decode also catches `TypeError`, and a non-dict `input_data` returns `None` (not a valid CCR call) instead of crashing CCR response processing. +* **mcp/codex:** don't corrupt an unparseable or non-table `config.toml` on register. `CodexRegistrar.register_server` only guarded against clobbering a user-managed entry when `get_server` returned one, but `get_server` returns `None` both for an unparseable TOML file and for an `mcp_servers`/`mcp_servers.` that is present but not a table. In those cases `register_server` fell through to `_write_block`, which blindly appended a `[mcp_servers.]` table — appending into an unparseable file, or creating a duplicate `[mcp_servers.headroom]` key alongside a non-table entry (e.g. `headroom = "..."`), which `tomllib`/codex then reject, destroying a previously-valid config. It now refuses (`FAILED`) and leaves the file untouched, mirroring the claude (#1660) and opencode (#1661) guards. * **proxy/vertex:** route Vertex `publisher=google` (Gemini) requests to the region matching the request path. `vertex_generate_content`, `vertex_stream_generate_content`, and `vertex_count_tokens` discarded the path's `location` and forwarded to the single fixed host from `_api_target(proxy, "vertex")` (default `us-central1`), instead of the region-aware `_vertex_target_for_location` the sibling Anthropic `rawPredict` route already uses. So a request to `.../locations/europe-west1/publishers/google/...` was sent to a `us-central1` host, which Vertex rejects on the region/host mismatch. The three google routes now derive the host from the request's `location` (operator-pinned upstreams are still honored). * **proxy/anthropic:** give each Anthropic conversation its own session id. `SessionTrackerStore.compute_session_id` derived its fallback id from `model` + system text harvested only from `role:"system"` entries inside `messages` — but Anthropic carries the system prompt as a top-level `body["system"]` field, so genuine Anthropic requests (which never carry `x-headroom-session-id`) collapsed to `md5(model:[])` and every conversation on the same model shared one `PrefixCacheTracker`. That let session-sticky state cross-contaminate: conversation A's sticky `headroom_retrieve`/memory tools and `anthropic-beta` headers were injected into conversation B, and frozen-prefix/compression-cache state mixed across conversations. The Anthropic handler now folds the top-level `system` into the session-id inputs (prepending a synthetic `role:"system"` message used only to derive the id), giving distinct conversations distinct ids. * **cache/semantic:** key entries by the full-context hash, not the trailing query text. `SemanticCache.put` stored each response under `sha256(query)[:16]` where `query` is only the last user message, and the exact-match branch of `get` returned the slot without checking the stored entry's `messages_hash`. Two requests that share a trailing message ("continue", "yes", "run the tests") but differ in earlier context therefore collided on one slot — the second overwrote the first, and the first's hash then resolved to the second's cached response (wrong data served). Entries are now keyed by `messages_hash` when present, and `get` verifies `entry.messages_hash` before returning. diff --git a/headroom/mcp_registry/codex.py b/headroom/mcp_registry/codex.py index 190ad96df..8e7c4b25c 100644 --- a/headroom/mcp_registry/codex.py +++ b/headroom/mcp_registry/codex.py @@ -110,6 +110,22 @@ class CodexRegistrar(MCPRegistrar): # Drop any prior Headroom block before re-writing. self.unregister_server(spec.name) + # `existing is None` here can also mean the file is present but + # unparseable, or defines mcp_servers[.] as a non-table. + # _write_block appends a `[mcp_servers.]` table, so appending into + # an unparseable file corrupts it further, and appending alongside a + # non-table entry creates a duplicate `[mcp_servers.]` key that + # tomllib/codex then reject — destroying a previously-valid user config. + # Refuse rather than clobber, mirroring the claude (#1660) / opencode + # (#1661) guards. + if existing is None: + reason = self._unmergeable_reason(spec.name) + if reason is not None: + return RegisterResult( + RegisterStatus.FAILED, + f"{reason}; refusing to overwrite. Fix or remove the file, then re-run.", + ) + return self._write_block(spec) def unregister_server(self, server_name: str) -> bool: @@ -155,6 +171,36 @@ class CodexRegistrar(MCPRegistrar): return {} return data if isinstance(data, dict) else {} + def _unmergeable_reason(self, name: str) -> str | None: + """Return why the existing config cannot be safely merged, or ``None``. + + ``_write_block`` appends a ``[mcp_servers.]`` table. That is only + safe when the file is absent/empty or parses as a TOML table whose + ``mcp_servers`` (and ``mcp_servers.``) are tables. A present-but- + unparseable file, or a non-table ``mcp_servers`` / ``mcp_servers.``, + would be corrupted (unparseable) or made to hold a duplicate key + (non-table entry) by a blind append. + """ + if not self._config_file.exists(): + return None + raw = self._read_text() + if not raw.strip(): + return None + try: + data = tomllib.loads(fsutil.read_text(self._config_file)) + except (tomllib.TOMLDecodeError, OSError) as exc: + return f"{self._config_file} is not valid TOML ({exc})" + if not isinstance(data, dict): + return f"{self._config_file} top-level TOML is not a table" + servers = data.get("mcp_servers") + if servers is not None and not isinstance(servers, dict): + return f"{self._config_file} has a non-table mcp_servers" + if isinstance(servers, dict): + entry = servers.get(name) + if entry is not None and not isinstance(entry, dict): + return f"{self._config_file} has a non-table mcp_servers.{name}" + return None + def _read_text(self) -> str: return fsutil.read_text(self._config_file, default="") diff --git a/tests/test_mcp_registry/test_codex_registrar.py b/tests/test_mcp_registry/test_codex_registrar.py index 937777610..efe3cb58f 100644 --- a/tests/test_mcp_registry/test_codex_registrar.py +++ b/tests/test_mcp_registry/test_codex_registrar.py @@ -144,6 +144,50 @@ def test_get_server_robust_to_unparseable_toml(tmp_path: Path) -> None: assert _make_registrar(tmp_path).get_server("headroom") is None +def test_register_refuses_unparseable_config(tmp_path: Path) -> None: + """An unparseable config.toml must not be appended to (that would corrupt it + further); refuse and leave it byte-for-byte untouched.""" + cfg = _config_path(tmp_path) + cfg.parent.mkdir() + original = "this = is = not = valid\n" + cfg.write_text(original) + + result = _make_registrar(tmp_path).register_server(_spec()) + + assert result.status == RegisterStatus.FAILED + assert "not valid TOML" in result.detail + assert cfg.read_text() == original + + +def test_register_refuses_non_table_mcp_servers_entry(tmp_path: Path) -> None: + """A valid config whose mcp_servers.headroom is a non-table must not get a + duplicate `[mcp_servers.headroom]` table appended (which tomllib rejects).""" + cfg = _config_path(tmp_path) + cfg.parent.mkdir() + original = '[mcp_servers]\nheadroom = "not-a-table"\n' + cfg.write_text(original) + + result = _make_registrar(tmp_path).register_server(_spec()) + + assert result.status == RegisterStatus.FAILED + assert "non-table" in result.detail + # Untouched — still the original single (string) definition. + assert cfg.read_text() == original + + +def test_register_refuses_non_table_mcp_servers(tmp_path: Path) -> None: + """A non-table top-level mcp_servers is also refused, not clobbered.""" + cfg = _config_path(tmp_path) + cfg.parent.mkdir() + original = 'mcp_servers = "oops"\n' + cfg.write_text(original) + + result = _make_registrar(tmp_path).register_server(_spec()) + + assert result.status == RegisterStatus.FAILED + assert cfg.read_text() == original + + # ---------------------------------------------------------------------- # register_server() — happy paths # ----------------------------------------------------------------------