Skip to content

Commit 5f5f261

Browse files
chopratejaschopratejas
authored
Merge pull request #555 from ashishpatel26/fix/531-codex-wrap-fail-open-on-compression
fix(codex): auto-enable fail-open on compression timeout in headroom wrap codex (#531)
2 parents e87cece + e8ecd08 commit 5f5f261

4 files changed

Lines changed: 88 additions & 5 deletions

File tree

‎headroom/proxy/handlers/openai.py‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2957,7 +2957,11 @@ async def handle_openai_responses(
29572957
decide_compression_failure_action,
29582958
)
29592959

2960-
_http_action = decide_compression_failure_action(_e, _http_body_bytes)
2960+
_http_action = decide_compression_failure_action(
2961+
_e,
2962+
_http_body_bytes,
2963+
client=client,
2964+
)
29612965
if _http_action.refuse:
29622966
logger.error(
29632967
"[%s] /v1/responses REFUSING to forward request "
@@ -4125,7 +4129,11 @@ def _prepare_ws_performance_metrics() -> tuple[float, float, dict[str, float]]:
41254129
decide_compression_failure_action,
41264130
)
41274131

4128-
_ws_action = decide_compression_failure_action(_ce, _ws_frame_bytes)
4132+
_ws_action = decide_compression_failure_action(
4133+
_ce,
4134+
_ws_frame_bytes,
4135+
client=client,
4136+
)
41294137
if _ws_action.refuse:
41304138
logger.error(
41314139
"[%s] WS /v1/responses REFUSING to forward "

‎headroom/proxy/helpers.py‎

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -794,7 +794,8 @@ class CompressionFailureAction:
794794
reason: str
795795
"""Short machine-readable label for telemetry. One of:
796796
``timeout``, ``oversize:bytes=<n>>threshold=<m>``,
797-
``small_frame_transient``, or ``env_override:fail_open``."""
797+
``small_frame_transient``, ``client_override:codex``, or
798+
``env_override:fail_open``."""
798799

799800
frame_bytes: int
800801
"""Original frame size in bytes (for logging / metrics)."""
@@ -803,6 +804,8 @@ class CompressionFailureAction:
803804
def decide_compression_failure_action(
804805
exception: BaseException,
805806
frame_bytes: int,
807+
*,
808+
client: str | None = None,
806809
) -> CompressionFailureAction:
807810
"""Decide whether to refuse-and-close vs forward-original after the
808811
proxy's compression pipeline fails on a Realtime WebSocket frame
@@ -812,6 +815,10 @@ def decide_compression_failure_action(
812815
813816
* env :data:`WS_COMPRESSION_FAIL_OPEN_ENV` truthy → forward (legacy
814817
behaviour, opt-in for debugging or strict compatibility).
818+
* Codex client compression timeout → forward. Codex currently treats
819+
the proxy's 1009/413 refusal path as a hard connection failure, so
820+
fail-open is safer for Codex sessions even when the proxy is run
821+
standalone rather than through ``headroom wrap codex``.
815822
* exception is :class:`asyncio.TimeoutError` → refuse (the compression
816823
stage hit its own timeout, which only fires on frames Headroom
817824
thought were big enough to need compression in the first place).
@@ -834,6 +841,15 @@ def decide_compression_failure_action(
834841
frame_bytes=frame_bytes,
835842
)
836843

844+
if (client or "").strip().lower() == "codex" and isinstance(
845+
exception, asyncio.TimeoutError
846+
):
847+
return CompressionFailureAction(
848+
refuse=False,
849+
reason="client_override:codex",
850+
frame_bytes=frame_bytes,
851+
)
852+
837853
threshold = WS_COMPRESSION_OVERSIZE_BYTES_DEFAULT
838854
raw_threshold = os.environ.get(WS_COMPRESSION_OVERSIZE_BYTES_ENV, "").strip()
839855
if raw_threshold:

‎tests/test_openai_codex_routing.py‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -340,6 +340,40 @@ async def _timeout_wait_for(awaitable, timeout):
340340
assert body.get("instructions") is None
341341

342342

343+
def test_codex_responses_timeout_fails_open_in_standalone_proxy(monkeypatch):
344+
"""Codex users running only the proxy still get fail-open on timeout."""
345+
request = _build_request(
346+
{
347+
"model": "gpt-5.4",
348+
"input": [
349+
{
350+
"type": "function_call_output",
351+
"call_id": "call-1",
352+
"output": "large tool output",
353+
}
354+
],
355+
},
356+
{"Authorization": "Bearer sk-test", "x-client": "codex"},
357+
)
358+
handler = _DummyOpenAIHandler()
359+
handler.config.optimize = True
360+
361+
monkeypatch.setattr("headroom.tokenizers.get_tokenizer", lambda model: _DummyTokenizer())
362+
monkeypatch.setattr(
363+
handler,
364+
"_compress_openai_responses_payload",
365+
lambda *args, **kwargs: (_ for _ in ()).throw(TimeoutError()),
366+
)
367+
368+
response = anyio.run(handler.handle_openai_responses, request)
369+
370+
assert response.status_code == 200
371+
assert handler.captured_request is not None
372+
_, url, _, body = handler.captured_request
373+
assert url == "https://api.openai.com/v1/responses"
374+
assert body["input"][0]["output"] == "large tool output"
375+
376+
343377
class _DummyWebSocket:
344378
def __init__(self, headers: dict[str, str]):
345379
self.headers = headers

‎tests/test_proxy/test_compression_failure_action.py‎

Lines changed: 27 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -49,8 +49,8 @@ def _env(**overrides: str | None) -> Iterator[None]:
4949
os.environ[key] = prior
5050

5151

52-
def test_timeout_always_refuses_regardless_of_frame_size() -> None:
53-
"""asyncio.TimeoutError → refuse, even for a tiny frame.
52+
def test_timeout_refuses_without_client_override_regardless_of_frame_size() -> None:
53+
"""asyncio.TimeoutError → refuse, even for a tiny non-Codex frame.
5454
5555
Compression timeout fires after ``COMPRESSION_TIMEOUT_SECONDS``, which
5656
means the pipeline already started work on the frame. A small frame
@@ -64,6 +64,31 @@ def test_timeout_always_refuses_regardless_of_frame_size() -> None:
6464
assert action.frame_bytes == 128
6565

6666

67+
def test_codex_client_timeout_fails_open_without_env_override() -> None:
68+
"""Codex direct-proxy traffic should keep flowing on compression timeout."""
69+
with _env(**{WS_COMPRESSION_FAIL_OPEN_ENV: None, WS_COMPRESSION_OVERSIZE_BYTES_ENV: None}):
70+
action = decide_compression_failure_action(
71+
asyncio.TimeoutError(),
72+
frame_bytes=128,
73+
client="codex",
74+
)
75+
assert action.refuse is False
76+
assert action.reason == "client_override:codex"
77+
assert action.frame_bytes == 128
78+
79+
80+
def test_non_codex_timeout_still_refuses_without_env_override() -> None:
81+
"""The Codex override must not restore global fail-open behavior."""
82+
with _env(**{WS_COMPRESSION_FAIL_OPEN_ENV: None, WS_COMPRESSION_OVERSIZE_BYTES_ENV: None}):
83+
action = decide_compression_failure_action(
84+
asyncio.TimeoutError(),
85+
frame_bytes=128,
86+
client="claude-code",
87+
)
88+
assert action.refuse is True
89+
assert action.reason == "timeout"
90+
91+
6792
def test_small_transient_error_falls_through_to_passthrough() -> None:
6893
"""Non-timeout error on a small frame: forward original (legacy)."""
6994
with _env(**{WS_COMPRESSION_FAIL_OPEN_ENV: None, WS_COMPRESSION_OVERSIZE_BYTES_ENV: None}):

0 commit comments

Comments
 (0)