mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
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 <mxjerrett@gmail.com>
This commit is contained in:
yaowei
GitHub
parent
02c77640a9
commit
ac7ee4e0bf
2 changed files with 157 additions and 4 deletions
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue