github-featured/chopratejas--headroom
2
0
Fork 0
Code Issues Pull Requests Actions 16 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

fix(proxy/responses): keep the Codex additional_tools carrier on the wire (#3194)

Browse Source 
## 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 <tejas@Tejass-MacBook-Pro.local>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Tejas Chopra
2026-08-21 22:54:03 -07:00
committed by GitHub
parent b4857685ff
commit 1617f839a1
2 changed files with 406 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
headroom/proxy/handlers/openai.py
+146 -3
View File
@@ -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
tests/test_openai_responses_additional_tools.py
+260
View File
@@ -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