mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(proxy): accept Codex websocket before upstream retries (#2203)
## Description Codex Desktop could abandon Headroom's local `/v1/responses` WebSocket handshake before Headroom's upstream retry strategy had a chance to recover. The ChatGPT-auth path waited for an upstream opening handshake with a minimum 30-second timeout before sending the local 101, while the reported Codex Desktop handshake expired after about 34 seconds. This change accepts validated ChatGPT-auth Codex WebSockets before opening the upstream connection, then keeps the existing upstream retries and HTTP fallback behind the established local session. API-key sessions retain connect-before-accept behavior so upstream `x-codex-*` headers can still be attached to their client-facing 101. The change is scoped to the pre-101 timing failure and does not address the separate large-context streaming investigation in #1944. Closes #2184 ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - Accept ChatGPT-auth Codex WebSocket clients before the upstream connect and retry loop. - Preserve API-key connect-before-accept ordering and upstream `x-codex-*` handshake-header forwarding. - Keep upstream retry, relay, usage-state refresh, and WebSocket-to-HTTP fallback behavior after the local 101. - Add a deterministic regression that blocks the first upstream opening handshake and proves the local acceptance deadline is independent of it. ## Testing - [x] Unit tests pass (`uv run pytest tests/test_openai_codex_ws_lifecycle.py -q`) - [x] Linting passes (`uv run ruff check headroom/proxy/handlers/openai.py tests/test_openai_codex_ws_lifecycle.py`) - [ ] Type checking passes (`uv run mypy headroom`) - [x] New tests added for new functionality when applicable - [ ] Manual testing performed ### Test Output ```text uv run pytest tests/test_openai_codex_ws_lifecycle.py -q 28 passed in 2.02s uv run ruff check headroom/proxy/handlers/openai.py tests/test_openai_codex_ws_lifecycle.py All checks passed! uv run ruff format headroom/proxy/handlers/openai.py tests/test_openai_codex_ws_lifecycle.py --check 2 files already formatted ``` ## Real Behavior Proof - Environment: Python 3.12, synced development worktree, local fake Codex client and upstream WebSocket, no live provider - Exact command / steps: Run `uv run pytest tests/test_openai_codex_ws_lifecycle.py::test_chatgpt_ws_accepts_before_stalled_upstream_connect -q`; the fake upstream blocks its first opening handshake while the client enforces a bounded local-accept deadline. - Observed result: `1 passed in 0.46s`; the ChatGPT-auth client receives its local 101 before the blocked upstream connect is released, and the handler continues into its existing upstream recovery path. - Not tested: live Codex Desktop pre-turn compaction against ChatGPT subscription infrastructure ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [x] I have updated the CHANGELOG.md if applicable ## Additional Notes `CHANGELOG.md` is unchanged because the release pipeline generates it from conventional commits. No user documentation changes are required; the handler comments and ordered-flow docstring are updated with the auth-mode-specific behavior. The broader #1944 large-context disconnect surface remains out of scope. --------- Co-authored-by: JerrettDavis <mxjerrett@gmail.com>
This commit is contained in:
Rod Boev
GitHub
parent
2de07db281
commit
551f473e04
4 changed files with 132 additions and 65 deletions
|
|
@ -4929,7 +4929,8 @@ class OpenAIHandlerMixin:
|
|||
|
||||
Newer Codex versions use WebSocket instead of HTTP POST for the
|
||||
Responses API. This handler:
|
||||
1. Accepts the client WebSocket
|
||||
1. Validates origin and routing, then accepts ChatGPT-auth sessions
|
||||
immediately or API-key sessions after the upstream handshake
|
||||
2. Receives the first message (``response.create`` request)
|
||||
3. Opens an upstream WebSocket to OpenAI
|
||||
4. Compresses eligible `response.create` text through the Python
|
||||
|
|
@ -5214,14 +5215,57 @@ class OpenAIHandlerMixin:
|
|||
f"{[k for k in upstream_headers if k.lower() != 'authorization']}, "
|
||||
f"subprotocols={client_subprotocols}"
|
||||
)
|
||||
accept_subprotocol = client_subprotocols[0] if client_subprotocols else None
|
||||
accepted_client_ws = False
|
||||
|
||||
def _register_accepted_session() -> None:
|
||||
nonlocal session_handle
|
||||
if session_handle is not None or ws_sessions is None:
|
||||
return
|
||||
client_addr: str | None = None
|
||||
client_info = getattr(websocket, "client", None)
|
||||
if client_info is not None:
|
||||
host = getattr(client_info, "host", None)
|
||||
port = getattr(client_info, "port", None)
|
||||
if host is not None and port is not None:
|
||||
client_addr = f"{host}:{port}"
|
||||
elif host is not None:
|
||||
client_addr = str(host)
|
||||
session_handle = WSSessionHandle(
|
||||
session_id=session_id,
|
||||
request_id=request_id,
|
||||
client_addr=client_addr,
|
||||
upstream_url=upstream_url,
|
||||
)
|
||||
ws_sessions.register(session_handle)
|
||||
metrics = getattr(self, "metrics", None)
|
||||
if metrics is not None and hasattr(metrics, "inc_active_ws_sessions"):
|
||||
try:
|
||||
metrics.inc_active_ws_sessions()
|
||||
except Exception: # pragma: no cover - defensive
|
||||
pass
|
||||
|
||||
def _schedule_usage_poll() -> None:
|
||||
with contextlib.suppress(Exception):
|
||||
from headroom.subscription.codex_rate_limits import (
|
||||
maybe_schedule_usage_poll,
|
||||
)
|
||||
|
||||
maybe_schedule_usage_poll(ws_headers)
|
||||
|
||||
try:
|
||||
# ChatGPT-auth sessions no longer need upstream x-codex-* headers on
|
||||
# the client-facing 101, so accept them before the upstream retry
|
||||
# loop. API-key sessions still connect first so the allowlisted
|
||||
# handshake headers remain attachable there.
|
||||
if is_chatgpt_auth:
|
||||
async with stage_timer.measure("accept"):
|
||||
await websocket.accept(subprotocol=accept_subprotocol)
|
||||
accepted_client_ws = True
|
||||
_register_accepted_session()
|
||||
_schedule_usage_poll()
|
||||
|
||||
# --- Connect to upstream OpenAI WebSocket ---
|
||||
# NOTE: we connect *before* accepting the client. OpenAI delivers the
|
||||
# Codex subscription/rate-limit window only on the upstream WS
|
||||
# handshake response headers, so we must read them here and attach
|
||||
# the x-codex-* subset to the client-facing 101 (below). Once accept()
|
||||
# sends the 101 the headers can no longer be added.
|
||||
logger.info(f"[{request_id}] WS /v1/responses connecting to {upstream_url}")
|
||||
|
||||
# Use ssl=True to let the websockets library handle SSL natively.
|
||||
|
|
@ -5291,11 +5335,9 @@ class OpenAIHandlerMixin:
|
|||
)
|
||||
await asyncio.sleep(delay_with_jitter / 1000)
|
||||
|
||||
# Accept the client WS, forwarding OpenAI's x-codex-* subscription
|
||||
# window from the upstream handshake onto the client-facing 101 so
|
||||
# Codex, /stats, and the headroom-desktop gauge can read the live
|
||||
# window. In API-key mode the handshake carries no x-codex-* headers,
|
||||
# so accept_headers stays empty and this behaves exactly as before.
|
||||
# Preserve any upstream x-codex-* handshake headers for internal
|
||||
# rate-limit state, and forward them onto the client-facing 101 only
|
||||
# for API-key sessions that have not been accepted yet.
|
||||
accept_headers: list[tuple[bytes, bytes]] = []
|
||||
if ws_connected:
|
||||
_codex_handshake = _extract_codex_handshake_headers(upstream)
|
||||
|
|
@ -5312,46 +5354,15 @@ class OpenAIHandlerMixin:
|
|||
with contextlib.suppress(Exception):
|
||||
get_codex_rate_limit_state().update_from_headers(dict(_codex_handshake))
|
||||
|
||||
# Current Codex no longer ships x-codex-* on the handshake, so the
|
||||
# block above is usually a no-op. Pull the live subscription window
|
||||
# from the dedicated usage endpoint instead (throttled, scoped to
|
||||
# ChatGPT-session traffic, fire-and-forget so accept isn't blocked).
|
||||
with contextlib.suppress(Exception):
|
||||
from headroom.subscription.codex_rate_limits import (
|
||||
maybe_schedule_usage_poll,
|
||||
)
|
||||
|
||||
maybe_schedule_usage_poll(ws_headers)
|
||||
async with stage_timer.measure("accept"):
|
||||
await websocket.accept(
|
||||
subprotocol=client_subprotocols[0] if client_subprotocols else None,
|
||||
headers=accept_headers or None,
|
||||
)
|
||||
|
||||
# --- Unit 3: register the session as soon as accept succeeds ---
|
||||
client_addr: str | None = None
|
||||
client_info = getattr(websocket, "client", None)
|
||||
if client_info is not None:
|
||||
host = getattr(client_info, "host", None)
|
||||
port = getattr(client_info, "port", None)
|
||||
if host is not None and port is not None:
|
||||
client_addr = f"{host}:{port}"
|
||||
elif host is not None:
|
||||
client_addr = str(host)
|
||||
if ws_sessions is not None:
|
||||
session_handle = WSSessionHandle(
|
||||
session_id=session_id,
|
||||
request_id=request_id,
|
||||
client_addr=client_addr,
|
||||
upstream_url=upstream_url,
|
||||
)
|
||||
ws_sessions.register(session_handle)
|
||||
metrics = getattr(self, "metrics", None)
|
||||
if metrics is not None and hasattr(metrics, "inc_active_ws_sessions"):
|
||||
try:
|
||||
metrics.inc_active_ws_sessions()
|
||||
except Exception: # pragma: no cover - defensive
|
||||
pass
|
||||
if not accepted_client_ws:
|
||||
_schedule_usage_poll()
|
||||
async with stage_timer.measure("accept"):
|
||||
await websocket.accept(
|
||||
subprotocol=accept_subprotocol,
|
||||
headers=accept_headers or None,
|
||||
)
|
||||
accepted_client_ws = True
|
||||
_register_accepted_session()
|
||||
# Receive the first message from client (the response.create request).
|
||||
# Bound the wait with WS_FIRST_FRAME_TIMEOUT_SECONDS so a zombie
|
||||
# client that opens the WS but never sends a frame cannot hold a
|
||||
|
|
|
|||
|
|
@ -1500,9 +1500,6 @@ class KompressCompressor(Transform):
|
|||
if n == 0:
|
||||
return []
|
||||
|
||||
if self._degraded_reason is not None:
|
||||
return [self._passthrough(c, len(c.split())) for c in contents]
|
||||
|
||||
# Normalize target_ratio to a per-text list
|
||||
if isinstance(target_ratio, list):
|
||||
if len(target_ratio) != n:
|
||||
|
|
@ -1526,6 +1523,9 @@ class KompressCompressor(Transform):
|
|||
else:
|
||||
ccr_sources = [None] * n
|
||||
|
||||
if getattr(self, "_degraded_reason", None) is not None:
|
||||
return [self._passthrough(c, len(c.split())) for c in contents]
|
||||
|
||||
# Fast path: on backends where batch-dim parallelism does NOT help
|
||||
# (ONNX CPU, PyTorch CPU), fall back to sequential `compress()`
|
||||
# internally. This keeps the public API consistent while avoiding the
|
||||
|
|
|
|||
|
|
@ -138,6 +138,7 @@ def test_get_tokenizer_name_prefers_most_specific_prefix() -> None:
|
|||
assert get_tokenizer_name("qwen2-7b-instruct") == "Qwen/Qwen2-7B"
|
||||
assert get_tokenizer_name("qwen2.5-turbo") == "Qwen/Qwen2.5-7B"
|
||||
assert get_tokenizer_name("deepseek-v2.5") == "deepseek-ai/DeepSeek-V2"
|
||||
# Direct hits and the shorter family fallback still resolve as before.
|
||||
# Direct hits and shorter family fallbacks still resolve through their
|
||||
# longest matching tokenizer aliases.
|
||||
assert get_tokenizer_name("qwen-14b") == "Qwen/Qwen-14B"
|
||||
assert get_tokenizer_name("deepseek-chat") == "deepseek-ai/deepseek-llm-7b-base"
|
||||
assert get_tokenizer_name("deepseek-chat") == "deepseek-ai/DeepSeek-V3"
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ real code paths (not mocked) and assert on registry / task state.
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import json
|
||||
import logging
|
||||
import sys
|
||||
|
|
@ -145,6 +146,7 @@ class _FakeWebSocket:
|
|||
self.sent_bytes: list[bytes] = []
|
||||
self.accepted_subprotocol: str | None = None
|
||||
self.accepted_headers: list[tuple[bytes, bytes]] | None = None
|
||||
self.accepted_event = asyncio.Event()
|
||||
self.closed = False
|
||||
self.close_code: int | None = None
|
||||
self._call_log = call_log
|
||||
|
|
@ -155,6 +157,7 @@ class _FakeWebSocket:
|
|||
async def accept(self, subprotocol=None, headers=None) -> None:
|
||||
self.accepted_subprotocol = subprotocol
|
||||
self.accepted_headers = list(headers) if headers is not None else None
|
||||
self.accepted_event.set()
|
||||
if self._call_log is not None:
|
||||
self._call_log.append("accept")
|
||||
|
||||
|
|
@ -951,18 +954,19 @@ async def test_upstream_connect_failure_still_deregisters_cleanly():
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ws_connect_failure_falls_back_to_http():
|
||||
"""When every upstream connect attempt fails, the client is still
|
||||
accepted (with no x-codex-* headers, since there is no upstream
|
||||
window) and the request is served via the HTTP POST fallback with
|
||||
the client's first frame. Preserves the pre-reorder WS-upgrade-
|
||||
failure behaviour.
|
||||
"""When every ChatGPT-auth upstream connect attempt fails, the client
|
||||
still gets its local 101 immediately, then the request is served via
|
||||
the HTTP POST fallback with the first frame after retries exhaust.
|
||||
"""
|
||||
fake_ws_mod = _make_fake_websockets_module(
|
||||
None, connect_error=RuntimeError("HTTP 500 from upstream")
|
||||
)
|
||||
|
||||
first = _first_frame()
|
||||
client_ws = _FakeWebSocket(frames=[first])
|
||||
client_ws = _FakeWebSocket(
|
||||
frames=[first],
|
||||
headers=_codex_lite_headers(chatgpt=True),
|
||||
)
|
||||
handler = _DummyOpenAIHandler()
|
||||
|
||||
fallback_calls: list[tuple] = []
|
||||
|
|
@ -986,6 +990,57 @@ async def test_ws_connect_failure_falls_back_to_http():
|
|||
assert handler.ws_sessions.active_count() == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chatgpt_ws_accepts_before_stalled_upstream_connect():
|
||||
"""ChatGPT-auth sessions must send the local 101 before a stalled
|
||||
upstream opening handshake is released.
|
||||
"""
|
||||
upstream_events = [
|
||||
json.dumps({"type": "response.created", "response": {"id": "r_1"}}),
|
||||
json.dumps({"type": "response.completed", "response": {"id": "r_1"}}),
|
||||
]
|
||||
first_attempt_started = asyncio.Event()
|
||||
release_first_attempt = asyncio.Event()
|
||||
connect_calls: list[tuple[tuple, dict]] = []
|
||||
|
||||
async def _connect(*args, **kwargs):
|
||||
connect_calls.append((args, dict(kwargs)))
|
||||
if len(connect_calls) == 1:
|
||||
first_attempt_started.set()
|
||||
await release_first_attempt.wait()
|
||||
raise RuntimeError("first opening handshake stalled")
|
||||
return _FakeUpstream(list(upstream_events))
|
||||
|
||||
fake_ws_mod = MagicMock()
|
||||
fake_ws_mod.connect = _connect
|
||||
fake_ws_mod.Subprotocol = str
|
||||
|
||||
client_ws = _FakeWebSocket(
|
||||
frames=[_first_frame()],
|
||||
headers=_codex_lite_headers(chatgpt=True),
|
||||
)
|
||||
handler = _DummyOpenAIHandler()
|
||||
handler.config.retry_max_attempts = 3
|
||||
|
||||
with patch.dict(sys.modules, {"websockets": fake_ws_mod}):
|
||||
task = asyncio.create_task(handler.handle_openai_responses_ws(client_ws))
|
||||
try:
|
||||
await asyncio.wait_for(first_attempt_started.wait(), timeout=0.5)
|
||||
await asyncio.wait_for(client_ws.accepted_event.wait(), timeout=0.2)
|
||||
assert len(connect_calls) == 1
|
||||
assert client_ws.accepted_headers is None
|
||||
release_first_attempt.set()
|
||||
await asyncio.wait_for(task, timeout=2.0)
|
||||
finally:
|
||||
release_first_attempt.set()
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await task
|
||||
|
||||
assert handler.ws_sessions.active_count() == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ws_codex_responses_lite_header_is_not_forwarded_upstream():
|
||||
"""The WS upstream handshake must drop the Codex lite header only."""
|
||||
|
|
@ -1124,9 +1179,9 @@ async def test_ws_first_frame_strips_codex_lite_metadata_mirror():
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ws_connect_happens_before_accept():
|
||||
"""The upstream connect must complete before the client 101 is sent,
|
||||
so OpenAI's x-codex-* handshake headers are available to attach.
|
||||
async def test_api_key_ws_connect_happens_before_accept():
|
||||
"""API-key sessions keep the upstream connect before the client 101,
|
||||
so OpenAI's x-codex-* handshake headers remain attachable there.
|
||||
"""
|
||||
upstream_events = [
|
||||
json.dumps({"type": "response.created", "response": {"id": "r_1"}}),
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue