From 1617f839a197ed1f17ca2083fbd288ffa2af7820 Mon Sep 17 00:00:00 2001 From: Tejas Chopra Date: Fri, 21 Aug 2026 22:54:03 -0700 Subject: [PATCH] fix(proxy/responses): keep the Codex additional_tools carrier on the wire (#3194) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description 0.36.3 regressed Codex tool access. A user reproduced it cleanly: Codex CLI 0.149.0 + Codex TUI/app-server, terminal tools available at first (`pwd` executes), then **all shell/filesystem access disappears for the rest of the session**. The same setup on 0.36.2 works. The only functional change in 0.36.3 was #3186. ## Root cause #3186 lifted `additional_tools` definitions into top-level `tools` so the tools consumers (schema compaction, output shaper, token accounting) would engage, and dropped the carrier item. That changed the definitions' **lifetime**, not just their location: - `tools` is a **per-request parameter**, scoped to one response. - `additional_tools` is an **`input` item** — part of the conversation transcript. A stateful session declares its tools once. Codex over WebSocket sends the carrier on turn one and relies on the transcript afterwards. Forwarding the lifted shape leaves that transcript tool-less, so turn one works and every turn after it has no tool surface at all. Stateless HTTP hid this in review — it re-sends the carrier on every request, so the lift refires each turn and nothing is ever lost. That is why the original manual verification passed. ## Fix The lift stays; the savings fix it shipped for is real. It is now **symmetric**: - `_lift_codex_additional_tools` records where each carrier came from (`restore_plan`). - `_restore_codex_additional_tools` puts the post-compaction definitions back into that carrier before the payload is forwarded. Consumers still see a classic top-level array. The client still sees the shape it sent. Compaction's savings survive the round trip, because it is the *compacted* schemas that go back into the carrier. Restoration is conservative: | Situation | Behaviour | |---|---| | Compaction preserved the definition count | original per-carrier split rebuilt exactly | | A consumer rewrote the array (deferral, injection) | whole set rides the first carrier | | Array came back empty | definitions Codex sent are restored, never a tool-less forward | | Carrier cannot be put back at all | logged, never a silent lifted-shape forward | | Called twice | idempotent, no duplication | Wired into `_compress_openai_responses_payload_in_executor`, so all five call sites — HTTP, both WebSocket sites, and passthrough — are covered by construction. `HEADROOM_CODEX_ADDITIONAL_TOOLS_LIFT=0` still disables the lift entirely and remains the immediate unblock for anyone on 0.36.3 right now. ## Testing The gap in #3186 was that all nine of its tests were single-turn. These are not. - **Multi-turn regression test** — a turn-one payload is driven through the real compression entry point, and turn two is built from what was actually forwarded. On shipped `main` that turn-two transcript carries **zero** tool definitions; with this change it carries both. - **Exhaustive round trip** — 363 arrangements of messages, carriers, empty carriers, adjacent/leading/trailing carriers. Zero mismatches. This is what pins the insert-offset arithmetic. - Round-trip shape preservation, carrier position, multiple carriers, count-change fallback, emptied-array recovery, extra carrier keys, idempotence, the unrestorable-warning path, the kill switch, and untouched classic-encoding clients are each asserted. 22 tests in the file; 112 across the related suites (proxy, codex routing, passthrough, compaction); full suite 3740 passed / 156 skipped. `ruff check` and `ruff format` clean. Before/after against shipped `main`, same scenario: | | 0.36.3 (`main`) | this PR | |---|---|---| | forwarded top-level `tools` | present | absent | | carrier surviving in `input` | **0** | 1 | | tools visible to turn 2 | **none — tool loss** | `shell`, `update_plan` | ## Validation gap — please read This proves the **forwarded shape now matches what the client sent**, which is the invariant that matters regardless of the exact upstream mechanism. What is *not* directly observed here is the transcript-persistence mechanism itself — that is inferred from Responses API semantics, because there is no Codex 0.149.0 stateful WebSocket backend in CI. That is the same gap that let #3186 ship broken, so it should not be waved through twice. The reporter has a reliable reproduction and should confirm this build before it tags. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Tejas Chopra Co-authored-by: Claude Opus 5 --- headroom/proxy/handlers/openai.py | 149 +++++++++- .../test_openai_responses_additional_tools.py | 260 ++++++++++++++++++ 2 files changed, 406 insertions(+), 3 deletions(-) diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 51dfb527d..cb2f70766 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -791,7 +791,12 @@ def _codex_additional_tools_lift_enabled() -> bool: ) -def _lift_codex_additional_tools(payload: dict[str, Any], *, request_id: str | None = None) -> int: +def _lift_codex_additional_tools( + payload: dict[str, Any], + *, + request_id: str | None = None, + restore_plan: list[dict[str, Any]] | None = None, +) -> int: """Lift Codex ``additional_tools`` input items into top-level ``tools``. Codex CLI 0.149.0 stopped sending a top-level ``tools`` array on @@ -804,7 +809,13 @@ def _lift_codex_additional_tools(payload: dict[str, Any], *, request_id: str | N Mutates *payload* in place: concatenates the items' ``tools`` arrays into ``payload["tools"]`` and drops the carrier items from ``input``. Returns - the number of lifted tool definitions (0 = no-op). No-op when the payload + the number of lifted tool definitions (0 = no-op). + + This is an *internal* normalization only. The forwarded payload must keep + the shape the client sent, so callers pass ``restore_plan`` and hand it to + :func:`_restore_codex_additional_tools` before forwarding -- forwarding the + lifted shape costs a stateful session its whole tool surface (the 0.36.3 + regression from #3186). No-op when the payload already carries top-level tools, so classic-encoding clients are untouched and a future Codex reverting the change costs nothing. The classic top-level encoding is accepted upstream for these models -- @@ -823,6 +834,7 @@ def _lift_codex_additional_tools(payload: dict[str, Any], *, request_id: str | N return 0 lifted: list[Any] = [] kept: list[Any] = [] + plan: list[dict[str, Any]] = [] for item in items: if ( isinstance(item, dict) @@ -830,6 +842,16 @@ def _lift_codex_additional_tools(payload: dict[str, Any], *, request_id: str | N and isinstance(item.get("tools"), list) and item["tools"] ): + # `kept_index` is the carrier's position among the items that + # survive the lift, so the restore re-inserts it in the same + # relative slot even when compression rewrites the transcript. + plan.append( + { + "kept_index": len(kept), + "item": {k: v for k, v in item.items() if k != "tools"}, + "tools": list(item["tools"]), + } + ) lifted.extend(item["tools"]) else: kept.append(item) @@ -837,6 +859,8 @@ def _lift_codex_additional_tools(payload: dict[str, Any], *, request_id: str | N return 0 payload["tools"] = lifted payload["input"] = kept + if restore_plan is not None: + restore_plan.extend(plan) logger.info( "[%s] Lifted %d Codex additional_tools definitions to top-level tools", request_id or "-", @@ -845,6 +869,91 @@ def _lift_codex_additional_tools(payload: dict[str, Any], *, request_id: str | N return len(lifted) +def _restore_codex_additional_tools( + payload: dict[str, Any], + restore_plan: list[dict[str, Any]], + *, + request_id: str | None = None, +) -> int: + """Undo :func:`_lift_codex_additional_tools`, keeping compaction's work. + + The lift exists so Headroom's tools consumers engage; it must not change + what Codex sees on the wire. ``tools`` is a per-request parameter, while + ``additional_tools`` is an ``input`` item and therefore part of the + conversation transcript. A stateful session (Codex TUI/app-server over + WebSocket) declares its tools once and relies on the transcript for every + later turn, so forwarding the lifted shape leaves that transcript + tool-less: turn one works, then shell/filesystem access vanishes for the + rest of the session -- the 0.36.3 regression from #3186. Stateless HTTP + hid it in review: every request re-sends the carrier, so the lift refires + each turn and nothing is ever lost. + + Rewrites *payload* in place -- moves ``payload["tools"]`` (post-compaction) + back into ``additional_tools`` carriers at their original positions and + drops the top-level array. Returns the number of definitions restored. + """ + if not isinstance(payload, dict) or not restore_plan: + return 0 + items = payload.get("input") + if not isinstance(items, list): + return 0 + if any( + isinstance(item, dict) and item.get("type") == "additional_tools" and item.get("tools") + for item in items + ): + # Already in carrier form -- restoring again would duplicate the + # definitions. Keeps the restore idempotent. + return 0 + tools = payload.get("tools") + if not isinstance(tools, list): + tools = [] + + original_total = sum(len(entry.get("tools") or []) for entry in restore_plan) + if not tools: + # A consumer emptied the array. Restoring the definitions as Codex + # sent them is strictly safer than forwarding a tool-less payload. + slices = [list(entry.get("tools") or []) for entry in restore_plan] + elif len(tools) == original_total: + slices = [] + offset = 0 + for entry in restore_plan: + width = len(entry.get("tools") or []) + slices.append(tools[offset : offset + width]) + offset += width + else: + # The definition count changed (deferral, injection), so the original + # per-carrier split no longer maps. The whole set rides the first + # carrier rather than being distributed on a stale boundary. + slices = [list(tools)] + [[] for _ in restore_plan[1:]] + + restored_items = list(items) + restored = 0 + shift = 0 + for entry, tool_slice in zip(restore_plan, slices): + if not tool_slice: + continue + carrier = dict(entry.get("item") or {}) + carrier["type"] = "additional_tools" + carrier["tools"] = tool_slice + position = entry.get("kept_index") + if not isinstance(position, int) or position < 0: + position = len(restored_items) + restored_items.insert(min(position + shift, len(restored_items)), carrier) + shift += 1 + restored += len(tool_slice) + + if not restored: + return 0 + payload["input"] = restored_items + payload.pop("tools", None) + logger.debug( + "[%s] Restored %d Codex tool definitions to additional_tools", + request_id or "-", + restored, + ) + return restored + + def _allow_responses_memory_tools(is_chatgpt_auth: bool) -> bool: # Preserve the ChatGPT Codex route's existing store policy and memory-tool # exclusion while API Responses memory continuations stay stateless. @@ -2942,9 +3051,16 @@ class OpenAIHandlerMixin: # shaping/compression so every downstream tools consumer engages. # Runs once per pass, ahead of the executor closure, and never breaks # forwarding. + _restore_plan: list[dict[str, Any]] = [] try: - _lift_codex_additional_tools(payload, request_id=request_id) + _lift_codex_additional_tools( + payload, + request_id=request_id, + restore_plan=_restore_plan, + ) except Exception: # pragma: no cover - defensive; never break forwarding + # The plan is deliberately kept: if the lift raised after mutating + # the payload, the restore is what undoes it. logger.warning( "[%s] additional_tools lift failed; continuing unlifted", request_id, @@ -3018,6 +3134,33 @@ class OpenAIHandlerMixin: _compress, timeout=timeout, ) + + # Restore the carrier Codex sent the definitions in. The lift is an + # internal normalization for the tools consumers; the forwarded shape + # must match what the client sent, or a stateful session loses its + # tool surface after the first turn (0.36.3 regression from #3186). + if _restore_plan and result and isinstance(result[0], dict): + try: + if not _restore_codex_additional_tools( + result[0], + _restore_plan, + request_id=request_id, + ): + # The lifted shape is about to go out: a stateful client + # will lose its tools after this turn. Never silent. + logger.warning( + "[%s] additional_tools carrier could not be restored; " + "forwarding lifted shape (set " + "HEADROOM_CODEX_ADDITIONAL_TOOLS_LIFT=0 to opt out)", + request_id, + ) + except Exception: # pragma: no cover - defensive + logger.warning( + "[%s] additional_tools restore failed", + request_id, + exc_info=True, + ) + if len(result) == 8: return (*result, timing) return result diff --git a/tests/test_openai_responses_additional_tools.py b/tests/test_openai_responses_additional_tools.py index 3667c969d..1403a1565 100644 --- a/tests/test_openai_responses_additional_tools.py +++ b/tests/test_openai_responses_additional_tools.py @@ -9,12 +9,16 @@ sees a tool-less request and records zero tool-schema savings. from __future__ import annotations +import asyncio import copy +import json from typing import Any from headroom.proxy.handlers.openai import ( + OpenAIHandlerMixin, _compact_openai_responses_tools, _lift_codex_additional_tools, + _restore_codex_additional_tools, ) @@ -148,3 +152,259 @@ def test_lifted_tools_reach_schema_compaction() -> None: assert after_bytes < before_bytes # Compaction preserves the invocation shape the model needs. assert [t["name"] for t in compacted["tools"]] == ["shell", "update_plan"] + + +# --------------------------------------------------------------------------- +# Carrier restoration (0.36.3 regression from #3186) +# +# The lift is an internal normalization so the tools consumers engage. It must +# not change the forwarded wire shape: `tools` is a per-request parameter, +# while `additional_tools` is an `input` item and therefore part of the +# conversation transcript. Codex TUI/app-server over WebSocket declares its +# tools once and relies on the transcript for every later turn, so forwarding +# the lifted shape cost the session its whole tool surface after turn one -- +# 0.36.3 regressed shell/filesystem access while 0.36.2 worked. +# --------------------------------------------------------------------------- + + +def _handler(compress=None): # noqa: ANN001, ANN202 + """A bare mixin with the executor and compressor stubbed out.""" + handler = object.__new__(OpenAIHandlerMixin) + + async def _run_compression(fn, *, timeout): # noqa: ANN001, ANN202 + return fn() + + def _default_compress(payload, *, model, request_id, **kwargs): # noqa: ANN001, ANN202 + return (payload, True, 0, [], None, 0, 0, 0, {}) + + handler._run_compression_in_executor = _run_compression + handler._compress_openai_responses_payload = compress or _default_compress + return handler + + +def _forward(payload: dict[str, Any], compress=None) -> dict[str, Any]: # noqa: ANN001 + handler = _handler(compress) + result = asyncio.run( + handler._compress_openai_responses_payload_in_executor( + payload, + model="gpt-5.6-sol", + request_id="req-carrier", + ) + ) + return result[0] + + +def test_lift_restore_round_trip_is_shape_preserving() -> None: + payload = _codex_0149_payload() + before = copy.deepcopy(payload) + + plan: list[dict[str, Any]] = [] + _lift_codex_additional_tools(payload, restore_plan=plan) + _restore_codex_additional_tools(payload, plan) + + assert payload == before + + +def test_compressor_sees_tools_but_forwarded_payload_does_not() -> None: + """The whole point: consumers get top-level tools, the wire keeps the carrier.""" + seen: list[Any] = [] + + def _compress(payload, *, model, request_id, **kwargs): # noqa: ANN001, ANN202 + seen.append(copy.deepcopy(payload.get("tools"))) + return (payload, True, 0, [], None, 0, 0, 0, {}) + + payload = _codex_0149_payload() + before = copy.deepcopy(payload) + + forwarded = _forward(payload, _compress) + + # The savings fix (#3185) still holds: compaction saw real tools. + assert [t["name"] for t in seen[0]] == ["shell", "update_plan"] + # The regression fix: the forwarded shape is what Codex sent. + assert "tools" not in forwarded + assert forwarded["input"] == before["input"] + + +def test_stateful_second_turn_still_carries_tools() -> None: + """Reproduces the 0.36.3 session: turn one worked, then tools vanished. + + A stateful client appends to the transcript it already sent. If Headroom + forwards turn one without the carrier, the transcript the client builds + turn two from has no tool definitions at all -- and turn two carries no + top-level ``tools`` either, so the model is left with no tool surface. + """ + forwarded_turn_1 = _forward(_codex_0149_payload()) + + turn_2 = { + "model": "gpt-5.6-sol", + "input": [ + *forwarded_turn_1["input"], + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "now run pwd again"}], + }, + ], + } + + # Turn two declares no tools of its own; everything rides the transcript. + assert "tools" not in turn_2 + assert any(item.get("type") == "additional_tools" for item in turn_2["input"]), ( + "turn two lost every tool definition -- this is the 0.36.3 regression" + ) + + # And turn two survives its own trip through the proxy with tools intact. + forwarded_turn_2 = _forward(turn_2) + carriers = [i for i in forwarded_turn_2["input"] if i.get("type") == "additional_tools"] + assert [t["name"] for c in carriers for t in c["tools"]] == ["shell", "update_plan"] + + +def test_restore_keeps_the_compacted_schemas() -> None: + """Restoration returns compaction's output, not the pre-compaction copy.""" + payload = _codex_0149_payload() + plan: list[dict[str, Any]] = [] + _lift_codex_additional_tools(payload, restore_plan=plan) + + compacted, modified, _before, _after = _compact_openai_responses_tools(payload) + assert modified, "fixture should be compactable" + restored = _restore_codex_additional_tools(compacted, plan) + + assert restored == 2 + carrier = next(i for i in compacted["input"] if i.get("type") == "additional_tools") + assert [t["name"] for t in carrier["tools"]] == ["shell", "update_plan"] + # The verbose description is gone -- the savings survived the round trip. + assert len(json.dumps(carrier["tools"])) < len( + json.dumps(_codex_0149_payload()["input"][1]["tools"]) + ) + + +def test_restore_puts_the_carrier_back_in_position() -> None: + payload = _codex_0149_payload() + payload["input"].append( + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "tail"}]} + ) + before_types = [i["type"] for i in payload["input"]] + + plan: list[dict[str, Any]] = [] + _lift_codex_additional_tools(payload, restore_plan=plan) + _restore_codex_additional_tools(payload, plan) + + assert [i["type"] for i in payload["input"]] == before_types + + +def test_restore_handles_multiple_carriers() -> None: + payload = _codex_0149_payload() + payload["input"].append({"type": "additional_tools", "tools": [_verbose_tool("view_image")]}) + before = copy.deepcopy(payload) + + plan: list[dict[str, Any]] = [] + assert _lift_codex_additional_tools(payload, restore_plan=plan) == 3 + assert _restore_codex_additional_tools(payload, plan) == 3 + + assert payload == before + + +def test_restore_folds_into_first_carrier_when_the_count_changes() -> None: + """Deferral/injection rewrites the array; the split no longer maps.""" + payload = _codex_0149_payload() + payload["input"].append({"type": "additional_tools", "tools": [_verbose_tool("view_image")]}) + plan: list[dict[str, Any]] = [] + _lift_codex_additional_tools(payload, restore_plan=plan) + + payload["tools"] = [_verbose_tool("tool_search")] + assert _restore_codex_additional_tools(payload, plan) == 1 + + carriers = [i for i in payload["input"] if i.get("type") == "additional_tools"] + assert len(carriers) == 1 + assert [t["name"] for t in carriers[0]["tools"]] == ["tool_search"] + assert "tools" not in payload + + +def test_restore_recovers_the_originals_when_the_array_is_emptied() -> None: + """A tool-less forward is never the safer outcome.""" + payload = _codex_0149_payload() + plan: list[dict[str, Any]] = [] + _lift_codex_additional_tools(payload, restore_plan=plan) + + payload["tools"] = [] + assert _restore_codex_additional_tools(payload, plan) == 2 + + carrier = next(i for i in payload["input"] if i.get("type") == "additional_tools") + assert [t["name"] for t in carrier["tools"]] == ["shell", "update_plan"] + + +def test_restore_preserves_other_carrier_keys() -> None: + payload = _codex_0149_payload() + payload["input"][1]["id"] = "carrier_abc" + before = copy.deepcopy(payload) + + plan: list[dict[str, Any]] = [] + _lift_codex_additional_tools(payload, restore_plan=plan) + _restore_codex_additional_tools(payload, plan) + + assert payload == before + + +def test_restore_is_a_noop_without_a_plan() -> None: + payload = _codex_0149_payload() + before = copy.deepcopy(payload) + + assert _restore_codex_additional_tools(payload, []) == 0 + assert payload == before + assert _restore_codex_additional_tools("not-a-dict", [{"tools": []}]) == 0 # type: ignore[arg-type] + + +def test_kill_switch_leaves_the_payload_completely_untouched(monkeypatch) -> None: + monkeypatch.setenv("HEADROOM_CODEX_ADDITIONAL_TOOLS_LIFT", "0") + payload = _codex_0149_payload() + before = copy.deepcopy(payload) + + forwarded = _forward(payload) + + assert forwarded["input"] == before["input"] + assert "tools" not in forwarded + + +def test_classic_top_level_clients_are_untouched_by_the_restore() -> None: + """A non-Codex payload never enters the lift, so it never enters the restore.""" + payload = { + "model": "gpt-5.6-sol", + "tools": [_verbose_tool("shell")], + "input": [ + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]} + ], + } + before = copy.deepcopy(payload) + + forwarded = _forward(payload) + + assert forwarded["tools"] == before["tools"] + assert forwarded["input"] == before["input"] + + +def test_unrestorable_payload_warns_instead_of_failing_silently(caplog) -> None: + """If the carrier cannot go back, say so -- a stateful client will lose tools.""" + + def _mangle(payload, *, model, request_id, **kwargs): # noqa: ANN001, ANN202 + payload["input"] = "collapsed-to-a-string" + return (payload, True, 0, [], None, 0, 0, 0, {}) + + with caplog.at_level("WARNING", logger="headroom.proxy"): + forwarded = _forward(_codex_0149_payload(), _mangle) + + assert any("could not be restored" in message for message in caplog.messages) + # Degraded, not broken: this turn still carries its tools. + assert [t["name"] for t in forwarded["tools"]] == ["shell", "update_plan"] + + +def test_restore_is_idempotent() -> None: + """A second restore must not duplicate the definitions.""" + payload = _codex_0149_payload() + before = copy.deepcopy(payload) + + plan: list[dict[str, Any]] = [] + _lift_codex_additional_tools(payload, restore_plan=plan) + assert _restore_codex_additional_tools(payload, plan) == 2 + assert _restore_codex_additional_tools(payload, plan) == 0 + + assert payload == before