From ac7ee4e0bf25a49f05df44a38d40ecb806cf3972 Mon Sep 17 00:00:00 2001 From: yaowei Date: Thu, 16 Jul 2026 04:36:27 +0800 Subject: [PATCH] fix(proxy): support Codex WS compatible gateways (#1281) Adds opt-in compatibility for OpenAI-compatible WebSocket gateways used behind Codex /v1/responses. - HEADROOM_OPENAI_WS_FLATTEN_RESPONSE_CREATE=1 flattens Codex response.create frames before upstream send. - HEADROOM_OPENAI_WS_PROPAGATE_UPSTREAM_CLOSE=1 propagates upstream close code/reason back to the client. - Default behavior is unchanged. Tested: python -m pytest tests/test_openai_codex_ws_lifecycle.py -q 18 passed Co-authored-by: JerrettDavis --- headroom/proxy/handlers/openai.py | 77 ++++++++++++++++++++++- tests/test_openai_codex_ws_lifecycle.py | 84 ++++++++++++++++++++++++- 2 files changed, 157 insertions(+), 4 deletions(-) diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 329b2d4e4..0854a6495 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -5537,6 +5537,40 @@ class OpenAIHandlerMixin: model or "unknown", ) + def _normalize_ws_response_create_for_upstream(raw_msg: str) -> str: + """Optionally flatten Codex WS response.create frames for gateways. + + Codex sends {"type":"response.create","response":{...}}. Some + OpenAI-compatible gateways accept the WebSocket upgrade but expect + model/input/tools at the top level of response.create frames. + Preserve default upstream behavior unless explicitly enabled. + """ + if os.environ.get( + "HEADROOM_OPENAI_WS_FLATTEN_RESPONSE_CREATE", "" + ).strip().lower() not in ( + "1", + "true", + "yes", + "on", + ): + return raw_msg + try: + parsed = json.loads(raw_msg) + except (json.JSONDecodeError, TypeError): + return raw_msg + if not isinstance(parsed, dict) or parsed.get("type") != "response.create": + return raw_msg + inner = parsed.get("response") + if not isinstance(inner, dict): + return raw_msg + + flattened = dict(inner) + flattened["type"] = "response.create" + for key, value in parsed.items(): + if key not in {"type", "response"} and key not in flattened: + flattened[key] = value + return json.dumps(flattened, ensure_ascii=False) + body: dict[str, Any] = {} tokens_saved = 0 # Session-scoped accumulator for tokens we *attempted* to @@ -6140,6 +6174,7 @@ class OpenAIHandlerMixin: _shape_labels, ) + first_msg_raw = _normalize_ws_response_create_for_upstream(first_msg_raw) _first_upstream_body: Any = None try: _first_upstream_body = json.loads(first_msg_raw) @@ -6516,6 +6551,7 @@ class OpenAIHandlerMixin: _shape_labels, ) + msg = _normalize_ws_response_create_for_upstream(msg) _outbound_frame_body: Any = None try: _outbound_frame_body = json.loads(msg) @@ -6965,7 +7001,11 @@ class OpenAIHandlerMixin: } if resp_id: cont["response"]["previous_response_id"] = resp_id - await upstream.send(json.dumps(cont)) + await upstream.send( + _normalize_ws_response_create_for_upstream( + json.dumps(cont) + ) + ) logger.info( f"[{request_id}] WS Memory: Sent continuation " f"with {len(tool_outputs)} result(s)" @@ -6987,9 +7027,40 @@ class OpenAIHandlerMixin: # distinguished from a clean # upstream disconnect. upstream_relay_error = relay_err - logger.debug( - f"[{request_id}] WS upstream→client relay ended: {relay_err}" + _upstream_close_code = getattr(relay_err, "code", None) or getattr( + getattr(relay_err, "rcvd", None), "code", None ) + _upstream_close_reason = ( + getattr(relay_err, "reason", None) + or getattr(getattr(relay_err, "rcvd", None), "reason", None) + or str(relay_err) + ) + logger.warning( + "[%s] WS upstream→client relay ended: %s code=%s reason=%s", + request_id, + type(relay_err).__name__, + _upstream_close_code, + _upstream_close_reason, + ) + if os.environ.get( + "HEADROOM_OPENAI_WS_PROPAGATE_UPSTREAM_CLOSE", "" + ).strip().lower() in ( + "1", + "true", + "yes", + "on", + ): + _client_close_code = ( + int(_upstream_close_code) + if isinstance(_upstream_close_code, int) + and 1000 <= int(_upstream_close_code) <= 4999 + else 1011 + ) + with contextlib.suppress(Exception): + await websocket.close( + code=_client_close_code, + reason=str(_upstream_close_reason)[:120], + ) finally: with contextlib.suppress(Exception): await websocket.close() diff --git a/tests/test_openai_codex_ws_lifecycle.py b/tests/test_openai_codex_ws_lifecycle.py index 4e1d99a21..d08c9fb38 100644 --- a/tests/test_openai_codex_ws_lifecycle.py +++ b/tests/test_openai_codex_ws_lifecycle.py @@ -127,6 +127,13 @@ class _FakeWebSocketDisconnect(Exception): _FakeWebSocketDisconnect.__name__ = "WebSocketDisconnect_Fake" +class _FakeUpstreamClose(Exception): + def __init__(self, code: int, reason: str) -> None: + super().__init__(reason) + self.code = code + self.reason = reason + + class _FakeWebSocket: """Scripted client WebSocket that can delay / disconnect mid-stream.""" @@ -150,6 +157,7 @@ class _FakeWebSocket: self.accepted_event = asyncio.Event() self.closed = False self.close_code: int | None = None + self.close_reason: str | None = None self._call_log = call_log # "client" can trip this event to simulate mid-stream disconnect. self._disconnect_event = asyncio.Event() @@ -187,7 +195,10 @@ class _FakeWebSocket: async def close(self, code: int | None = None, reason: str | None = None) -> None: self.closed = True - self.close_code = code + if code is not None or self.close_code is None: + self.close_code = code + if reason is not None or self.close_reason is None: + self.close_reason = reason def trigger_disconnect(self) -> None: self._disconnect_event.set() @@ -758,6 +769,77 @@ async def test_ws_session_metrics_include_dashboard_performance_timings(): ) +@pytest.mark.asyncio +async def test_ws_opt_in_flattens_response_create_for_openai_compatible_upstream(monkeypatch): + """Some OpenAI-compatible WS gateways expect top-level response.create payloads.""" + + monkeypatch.setenv("HEADROOM_OPENAI_WS_FLATTEN_RESPONSE_CREATE", "1") + upstream_events = [ + json.dumps({"type": "response.created", "response": {"id": "r_1"}}), + json.dumps({"type": "response.completed", "response": {"id": "r_1"}}), + ] + upstream = _FakeUpstream(upstream_events) + fake_ws_mod = _make_fake_websockets_module(upstream) + + first = json.dumps( + { + "type": "response.create", + "event_id": "evt_flatten", + "response": { + "model": "gpt-5.4", + "input": "hello", + "instructions": "be concise", + "tools": [{"type": "function", "name": "shell"}], + }, + } + ) + client_ws = _FakeWebSocket(frames=[first]) + handler = _DummyOpenAIHandler() + + with patch.dict(sys.modules, {"websockets": fake_ws_mod}): + await handler.handle_openai_responses_ws(client_ws) + + assert upstream.sent + sent = json.loads(upstream.sent[0]) + assert sent == { + "model": "gpt-5.4", + "input": "hello", + "instructions": "be concise", + "tools": [{"type": "function", "name": "shell"}], + "type": "response.create", + "event_id": "evt_flatten", + } + + +@pytest.mark.asyncio +async def test_ws_opt_in_propagates_upstream_close_code_and_reason(monkeypatch): + """Expose upstream close details to Codex instead of swallowing them in debug logs.""" + + monkeypatch.setenv("HEADROOM_OPENAI_WS_PROPAGATE_UPSTREAM_CLOSE", "1") + upstream = _FakeUpstream( + [], + raise_mid_stream=_FakeUpstreamClose(4001, "bad request shape"), + ) + fake_ws_mod = _make_fake_websockets_module(upstream) + + client_ws = _FakeWebSocket( + frames=[_first_frame()], + hold_after_initial=True, + ) + handler = _DummyOpenAIHandler() + + with patch.dict(sys.modules, {"websockets": fake_ws_mod}): + await asyncio.wait_for( + handler.handle_openai_responses_ws(client_ws), + timeout=2.0, + ) + + assert client_ws.closed + assert client_ws.close_code == 4001 + assert client_ws.close_reason == "bad request shape" + assert handler.metrics.termination_causes[-1] == "upstream_error" + + @pytest.mark.asyncio async def test_client_disconnect_cancels_upstream_relay_within_100ms(): """**Failing-test-first** scenario from the plan.