diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index ca26246aa..116d3c12c 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -331,6 +331,35 @@ RESPONSES_CONTEXT_SEARCH_TIMEOUT_SECONDS = 2.0 WS_FIRST_FRAME_TIMEOUT_SECONDS = 60.0 +def _extract_codex_handshake_headers(upstream: Any) -> list[tuple[str, str]]: + """Return the ``x-codex-*`` headers from an upstream WS handshake response. + + OpenAI delivers the Codex subscription/rate-limit window only on the + WebSocket handshake response headers (not in data frames). We forward + that subset onto the client-facing 101 so Codex, ``/stats``, and the + headroom-desktop gauge can all read the live window. Filtered strictly + to ``x-codex-*`` -- never ``set-cookie``/``authorization``/etc. + """ + resp = getattr(upstream, "response", None) + headers = getattr(resp, "headers", None) + if headers is None: + return [] + raw_items = getattr(headers, "raw_items", None) + try: + items = list(raw_items()) if callable(raw_items) else list(headers.items()) + except Exception: + return [] + out: list[tuple[str, str]] = [] + for name, value in items: + name_str = name.decode("latin-1") if isinstance(name, (bytes, bytearray)) else str(name) + if name_str.lower().startswith("x-codex-"): + value_str = ( + value.decode("latin-1") if isinstance(value, (bytes, bytearray)) else str(value) + ) + out.append((name_str, value_str)) + return out + + def _infer_openai_cache_write_tokens(input_tokens: int, cache_read_tokens: int) -> int: """Infer OpenAI automatic prompt-cache writes from uncached input tokens. @@ -3346,38 +3375,6 @@ class OpenAIHandlerMixin: if raw_protocol: client_subprotocols = [p.strip() for p in raw_protocol.split(",") if p.strip()] - # Accept client connection with the requested subprotocol - async with stage_timer.measure("accept"): - if client_subprotocols: - await websocket.accept(subprotocol=client_subprotocols[0]) - else: - await websocket.accept() - - # --- 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=None, # set below once upstream_url is computed - ) - 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 - # Forward all client headers except hop-by-hop / per-connection headers. # These are WebSocket handshake mechanics that the `websockets` library # generates fresh for the upstream connection — forwarding them would conflict. @@ -3451,10 +3448,6 @@ class OpenAIHandlerMixin: }, ) - # Unit 3: attach the resolved upstream URL to the session handle. - if session_handle is not None: - session_handle.upstream_url = upstream_url - logger.info( "[%s] WS /v1/responses accepted (route=%s, auth_mode=%s, subprotocols=%s)", request_id, @@ -3554,6 +3547,118 @@ class OpenAIHandlerMixin: ) try: + # --- 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. + # Manual ssl.create_default_context() + certifi doesn't load the + # Windows system cert store, causing HTTP 500 on wss:// connections. + use_ssl: bool | None = True if upstream_url.startswith("wss://") else None + + ws_connected = False + ws_connect_attempts = max(1, getattr(self.config, "retry_max_attempts", 3)) + ws_last_err: Exception | None = None + _upstream_connect_started = time.perf_counter() + _upstream_connect_recorded = False + _upstream_first_event_started: float | None = None + upstream: Any = None + + for ws_attempt in range(ws_connect_attempts): + try: + upstream = await websockets.connect( + upstream_url, + additional_headers=upstream_headers, + subprotocols=( + [websockets.Subprotocol(p) for p in client_subprotocols] + if client_subprotocols and hasattr(websockets, "Subprotocol") + else client_subprotocols or None + ), + ssl=use_ssl, + open_timeout=max(30, self.config.connect_timeout_seconds * 3), + close_timeout=10, + ping_interval=20, + ping_timeout=20, + ) + ws_connected = True + if not _upstream_connect_recorded: + stage_timer.record( + "upstream_connect", + (time.perf_counter() - _upstream_connect_started) * 1000.0, + ) + _upstream_connect_recorded = True + _upstream_first_event_started = time.perf_counter() + break + except Exception as ws_err: + ws_last_err = ws_err + if ws_attempt >= ws_connect_attempts - 1: + break + delay_with_jitter = jitter_delay_ms( + self.config.retry_base_delay_ms, + self.config.retry_max_delay_ms, + ws_attempt, + ) + logger.warning( + f"[{request_id}] WS upstream connect failed " + f"(attempt {ws_attempt + 1}/{ws_connect_attempts}): {ws_err}; " + f"retrying in {delay_with_jitter:.0f}ms" + ) + 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. + accept_headers: list[tuple[bytes, bytes]] = [] + if ws_connected: + _codex_handshake = _extract_codex_handshake_headers(upstream) + if _codex_handshake: + accept_headers = [ + (name.encode("latin-1"), value.encode("latin-1")) + for name, value in _codex_handshake + ] + # Parity with the HTTP path: also refresh Python /stats state. + from headroom.subscription.codex_rate_limits import ( + get_codex_rate_limit_state, + ) + + with contextlib.suppress(Exception): + get_codex_rate_limit_state().update_from_headers(dict(_codex_handshake)) + 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 # 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 @@ -4190,171 +4295,152 @@ class OpenAIHandlerMixin: }, ) - # --- Connect to upstream OpenAI WebSocket --- - logger.info(f"[{request_id}] WS /v1/responses connecting to {upstream_url}") + if ws_connected: + async with upstream: + await upstream.send(first_msg_raw) - # Use ssl=True to let the websockets library handle SSL natively. - # Manual ssl.create_default_context() + certifi doesn't load the - # Windows system cert store, causing HTTP 500 on wss:// connections. - use_ssl: bool | None = True if upstream_url.startswith("wss://") else None + # Unit 3: flag the upstream side flips on seeing + # ``response.completed`` so the outer cause + # classifier can prefer it over the raw + # "upstream iterator ended" default. + response_completed_seen = False + # Captures the first exception surfaced by the + # inner relay ``except`` blocks so the outer + # classifier can still tell ``upstream_error`` + # from ``upstream_disconnect`` / ``response_completed`` + # even though the halves swallow and log. + upstream_relay_error: BaseException | None = None + client_relay_error: BaseException | None = None - ws_connected = False - ws_connect_attempts = max(1, getattr(self.config, "retry_max_attempts", 3)) - ws_last_err: Exception | None = None - _upstream_connect_started = time.perf_counter() - _upstream_connect_recorded = False - _upstream_first_event_started: float | None = None - - for ws_attempt in range(ws_connect_attempts): - try: - async with websockets.connect( - upstream_url, - additional_headers=upstream_headers, - subprotocols=( - [websockets.Subprotocol(p) for p in client_subprotocols] - if client_subprotocols and hasattr(websockets, "Subprotocol") - else client_subprotocols or None - ), - ssl=use_ssl, - open_timeout=max(30, self.config.connect_timeout_seconds * 3), - close_timeout=10, - ping_interval=20, - ping_timeout=20, - ) as upstream: - ws_connected = True - if not _upstream_connect_recorded: - stage_timer.record( - "upstream_connect", - (time.perf_counter() - _upstream_connect_started) * 1000.0, + async def _maybe_compress_response_create_frame( + raw_msg: str, + *, + frame_index: int, + ) -> tuple[str, bool, str | None]: + """Compress a single client→upstream frame + when its `type` is `response.create`. Other + event types (response.cancel, session.update, + etc.) pass through unchanged. Errors are + warned and the original frame is returned — + fail loud in logs, fail safe on the wire. + Updates outer-scope ``tokens_saved``, + ``transforms_applied``, and + ``ws_frames_compressed`` so the session-end + log reports cumulative savings across all + frames in the WS session. + """ + nonlocal tokens_saved, transforms_applied, attempted_input_tokens_total + nonlocal ws_frames_compressed + if _ws_bypass: + _log_ws_passthrough( + "bypass_header", + frame_index=frame_index, + raw_bytes=len(raw_msg.encode("utf-8", errors="replace")), ) - _upstream_connect_recorded = True - _upstream_first_event_started = time.perf_counter() - await upstream.send(first_msg_raw) - - # Unit 3: flag the upstream side flips on seeing - # ``response.completed`` so the outer cause - # classifier can prefer it over the raw - # "upstream iterator ended" default. - response_completed_seen = False - # Captures the first exception surfaced by the - # inner relay ``except`` blocks so the outer - # classifier can still tell ``upstream_error`` - # from ``upstream_disconnect`` / ``response_completed`` - # even though the halves swallow and log. - upstream_relay_error: BaseException | None = None - client_relay_error: BaseException | None = None - - async def _maybe_compress_response_create_frame( - raw_msg: str, - *, - frame_index: int, - ) -> tuple[str, bool, str | None]: - """Compress a single client→upstream frame - when its `type` is `response.create`. Other - event types (response.cancel, session.update, - etc.) pass through unchanged. Errors are - warned and the original frame is returned — - fail loud in logs, fail safe on the wire. - Updates outer-scope ``tokens_saved``, - ``transforms_applied``, and - ``ws_frames_compressed`` so the session-end - log reports cumulative savings across all - frames in the WS session. - """ - nonlocal tokens_saved, transforms_applied, attempted_input_tokens_total - nonlocal ws_frames_compressed - if _ws_bypass: - _log_ws_passthrough( - "bypass_header", - frame_index=frame_index, - raw_bytes=len(raw_msg.encode("utf-8", errors="replace")), - ) - return raw_msg, False, "bypass_header" - if not self.config.optimize: - _log_ws_passthrough( - "optimize_disabled", - frame_index=frame_index, - raw_bytes=len(raw_msg.encode("utf-8", errors="replace")), - ) - return raw_msg, False, "optimize_disabled" - _preflight_started = time.perf_counter() + return raw_msg, False, "bypass_header" + if not self.config.optimize: + _log_ws_passthrough( + "optimize_disabled", + frame_index=frame_index, + raw_bytes=len(raw_msg.encode("utf-8", errors="replace")), + ) + return raw_msg, False, "optimize_disabled" + _preflight_started = time.perf_counter() + try: + parsed_frame = json.loads(raw_msg) + except json.JSONDecodeError: + _log_ws_passthrough( + "non_json", + frame_index=frame_index, + raw_bytes=len(raw_msg.encode("utf-8", errors="replace")), + ) + return raw_msg, False, "non_json" + if ( + not isinstance(parsed_frame, dict) + or parsed_frame.get("type") != "response.create" + ): + _log_ws_passthrough( + "not_response_create", + frame_index=frame_index, + raw_bytes=len(raw_msg.encode("utf-8", errors="replace")), + frame_type=( + parsed_frame.get("type") + if isinstance(parsed_frame, dict) + else type(parsed_frame).__name__ + ), + ) + return raw_msg, False, "not_response_create" + wrapped_frame = isinstance(parsed_frame.get("response"), dict) + inner_payload = parsed_frame["response"] if wrapped_frame else parsed_frame + if not isinstance(inner_payload, dict): + _log_ws_passthrough( + "invalid_inner_payload", + frame_index=frame_index, + raw_bytes=len(raw_msg.encode("utf-8", errors="replace")), + frame_type="response.create", + ) + return raw_msg, False, "invalid_inner_payload" + frame_compression_elapsed_ms = 0.0 + try: + model_for_frame = inner_payload.get("model") or "" + _frame_auth_mode = classify_auth_mode(ws_headers) + _preflight_ms = (time.perf_counter() - _preflight_started) * 1000.0 + _record_ws_compression_timing( + "compression_preflight_serialization", + _preflight_ms, + ) + _record_ws_compression_overhead(_preflight_ms) + _compression_started = time.perf_counter() try: - parsed_frame = json.loads(raw_msg) - except json.JSONDecodeError: - _log_ws_passthrough( - "non_json", - frame_index=frame_index, - raw_bytes=len(raw_msg.encode("utf-8", errors="replace")), + ( + new_inner, + modified, + frame_saved, + frame_transforms, + frame_reason, + bytes_before, + bytes_after, + frame_attempted_tokens, + frame_compression_timing, + ) = await self._compress_openai_responses_payload_in_executor( + inner_payload, + model=model_for_frame, + request_id=request_id, ) - return raw_msg, False, "non_json" - if ( - not isinstance(parsed_frame, dict) - or parsed_frame.get("type") != "response.create" - ): - _log_ws_passthrough( - "not_response_create", - frame_index=frame_index, - raw_bytes=len(raw_msg.encode("utf-8", errors="replace")), - frame_type=( - parsed_frame.get("type") - if isinstance(parsed_frame, dict) - else type(parsed_frame).__name__ + for ( + _timing_name, + _timing_ms, + ) in frame_compression_timing.items(): + _record_ws_compression_timing(_timing_name, _timing_ms) + finally: + frame_compression_elapsed_ms = ( + time.perf_counter() - _compression_started + ) * 1000.0 + _record_ws_compression_timing( + "compression_executor_wait_run", + frame_compression_elapsed_ms, + ) + _record_ws_compression_overhead(frame_compression_elapsed_ms) + record_frame = getattr( + getattr(self, "metrics", None), + "record_codex_ws_frame", + None, + ) + if record_frame is not None: + record_frame( + elapsed_ms=frame_compression_elapsed_ms, + bytes_before=bytes_before, + bytes_after=bytes_after, + attempted_tokens=frame_attempted_tokens, + tokens_saved=frame_saved, + modified=modified, + strategy_chain=_codex_ws_strategy_chain(frame_transforms), + final_strategies=_codex_ws_final_strategies( + frame_compression_timing ), ) - return raw_msg, False, "not_response_create" - wrapped_frame = isinstance(parsed_frame.get("response"), dict) - inner_payload = ( - parsed_frame["response"] if wrapped_frame else parsed_frame - ) - if not isinstance(inner_payload, dict): - _log_ws_passthrough( - "invalid_inner_payload", - frame_index=frame_index, - raw_bytes=len(raw_msg.encode("utf-8", errors="replace")), - frame_type="response.create", - ) - return raw_msg, False, "invalid_inner_payload" - frame_compression_elapsed_ms = 0.0 - try: - model_for_frame = inner_payload.get("model") or "" - _frame_auth_mode = classify_auth_mode(ws_headers) - _preflight_ms = (time.perf_counter() - _preflight_started) * 1000.0 - _record_ws_compression_timing( - "compression_preflight_serialization", - _preflight_ms, - ) - _record_ws_compression_overhead(_preflight_ms) - _compression_started = time.perf_counter() - try: - ( - new_inner, - modified, - frame_saved, - frame_transforms, - frame_reason, - bytes_before, - bytes_after, - frame_attempted_tokens, - frame_compression_timing, - ) = await self._compress_openai_responses_payload_in_executor( - inner_payload, - model=model_for_frame, - request_id=request_id, - ) - for ( - _timing_name, - _timing_ms, - ) in frame_compression_timing.items(): - _record_ws_compression_timing(_timing_name, _timing_ms) - finally: - frame_compression_elapsed_ms = ( - time.perf_counter() - _compression_started - ) * 1000.0 - _record_ws_compression_timing( - "compression_executor_wait_run", - frame_compression_elapsed_ms, - ) - _record_ws_compression_overhead(frame_compression_elapsed_ms) + except Exception as _frame_err: + if frame_compression_elapsed_ms > 0: record_frame = getattr( getattr(self, "metrics", None), "record_codex_ws_frame", @@ -4363,807 +4449,749 @@ class OpenAIHandlerMixin: if record_frame is not None: record_frame( elapsed_ms=frame_compression_elapsed_ms, - bytes_before=bytes_before, - bytes_after=bytes_after, - attempted_tokens=frame_attempted_tokens, - tokens_saved=frame_saved, - modified=modified, - strategy_chain=_codex_ws_strategy_chain(frame_transforms), - final_strategies=_codex_ws_final_strategies( - frame_compression_timing - ), + bytes_before=len(raw_msg.encode("utf-8", errors="replace")), + failed=True, ) - except Exception as _frame_err: - if frame_compression_elapsed_ms > 0: - record_frame = getattr( - getattr(self, "metrics", None), - "record_codex_ws_frame", - None, - ) - if record_frame is not None: - record_frame( - elapsed_ms=frame_compression_elapsed_ms, - bytes_before=len( - raw_msg.encode("utf-8", errors="replace") - ), - failed=True, - ) - logger.warning( - "[%s] WS /v1/responses frame compression " - "failed; forwarding original: %s: %s", - request_id, - type(_frame_err).__name__, - _frame_err, - ) - _log_ws_passthrough( - "compression_exception", - frame_index=frame_index, - raw_bytes=len(raw_msg.encode("utf-8", errors="replace")), - frame_type="response.create", - model=str(inner_payload.get("model") or "unknown"), - ) - return raw_msg, False, "compression_exception" - if not modified: - reason = frame_reason or "no_compression" - _log_ws_passthrough( - reason, - frame_index=frame_index, - raw_bytes=bytes_before, - frame_type="response.create", - model=str(inner_payload.get("model") or "unknown"), - ) - return raw_msg, False, reason - if not isinstance(new_inner, dict): - _log_ws_passthrough( - "compressed_payload_not_dict", - frame_index=frame_index, - raw_bytes=len(raw_msg.encode("utf-8", errors="replace")), - frame_type="response.create", - model=str(inner_payload.get("model") or "unknown"), - ) - return raw_msg, False, "compressed_payload_not_dict" - if wrapped_frame: - _rewrite_started = time.perf_counter() - parsed_frame["response"] = new_inner - rewritten = json.dumps(parsed_frame) - else: - _rewrite_started = time.perf_counter() - rewritten = json.dumps(new_inner) - _rewrite_ms = (time.perf_counter() - _rewrite_started) * 1000.0 - _record_ws_compression_timing( - "compression_payload_rewrite_json_dump", - _rewrite_ms, - ) - _record_ws_compression_overhead(_rewrite_ms) - tokens_saved += int(frame_saved) - attempted_input_tokens_total += int(frame_attempted_tokens) - for t in frame_transforms: - if t not in transforms_applied: - transforms_applied.append(t) - ws_frames_compressed += 1 - logger.info( - "[%s] WS /v1/responses frame compressed " - "%d→%d bytes (%d tokens saved, " - "auth_mode=%s, frame=%d)", + logger.warning( + "[%s] WS /v1/responses frame compression " + "failed; forwarding original: %s: %s", request_id, - bytes_before, - bytes_after, - int(frame_saved), - _frame_auth_mode.value, - ws_frames_compressed, + type(_frame_err).__name__, + _frame_err, ) - return rewritten, True, frame_reason or "compressed" + _log_ws_passthrough( + "compression_exception", + frame_index=frame_index, + raw_bytes=len(raw_msg.encode("utf-8", errors="replace")), + frame_type="response.create", + model=str(inner_payload.get("model") or "unknown"), + ) + return raw_msg, False, "compression_exception" + if not modified: + reason = frame_reason or "no_compression" + _log_ws_passthrough( + reason, + frame_index=frame_index, + raw_bytes=bytes_before, + frame_type="response.create", + model=str(inner_payload.get("model") or "unknown"), + ) + return raw_msg, False, reason + if not isinstance(new_inner, dict): + _log_ws_passthrough( + "compressed_payload_not_dict", + frame_index=frame_index, + raw_bytes=len(raw_msg.encode("utf-8", errors="replace")), + frame_type="response.create", + model=str(inner_payload.get("model") or "unknown"), + ) + return raw_msg, False, "compressed_payload_not_dict" + if wrapped_frame: + _rewrite_started = time.perf_counter() + parsed_frame["response"] = new_inner + rewritten = json.dumps(parsed_frame) + else: + _rewrite_started = time.perf_counter() + rewritten = json.dumps(new_inner) + _rewrite_ms = (time.perf_counter() - _rewrite_started) * 1000.0 + _record_ws_compression_timing( + "compression_payload_rewrite_json_dump", + _rewrite_ms, + ) + _record_ws_compression_overhead(_rewrite_ms) + tokens_saved += int(frame_saved) + attempted_input_tokens_total += int(frame_attempted_tokens) + for t in frame_transforms: + if t not in transforms_applied: + transforms_applied.append(t) + ws_frames_compressed += 1 + logger.info( + "[%s] WS /v1/responses frame compressed " + "%d→%d bytes (%d tokens saved, " + "auth_mode=%s, frame=%d)", + request_id, + bytes_before, + bytes_after, + int(frame_saved), + _frame_auth_mode.value, + ws_frames_compressed, + ) + return rewritten, True, frame_reason or "compressed" - async def _client_to_upstream() -> None: - nonlocal client_relay_error, ws_response_create_frames - nonlocal ws_client_frames_total, ws_cancel_frames - nonlocal ws_last_client_frame_type, ws_client_disconnect_seen - client_frame_index = 1 - try: - while True: - msg = await websocket.receive_text() - client_frame_index += 1 - ws_client_frames_total += 1 - if session_handle is not None: - session_handle.mark_activity() - _inbound_frame_body: Any = None - try: - _inbound_frame_body = json.loads(msg) - except json.JSONDecodeError: - _inbound_frame_body = None - ws_last_client_frame_type = ( - str(_inbound_frame_body.get("type") or "unknown") - if isinstance(_inbound_frame_body, dict) - else "non_json" - ) - if ws_last_client_frame_type == "response.cancel": - ws_cancel_frames += 1 - logger.info( - "[%s] WS client sent response.cancel " - "session_id=%s frame=%d cancels=%d", - request_id, - session_id, - client_frame_index, - ws_cancel_frames, - ) - else: - logger.debug( - "[%s] WS client frame session_id=%s frame=%d type=%s", - request_id, - session_id, - client_frame_index, - ws_last_client_frame_type, - ) - capture_codex_wire_debug( - "ws_inbound_client_frame", - request_id=request_id, - session_id=session_id, - transport="websocket", - direction="client_to_headroom", - url=_ws_url, - body=_inbound_frame_body, - raw_text=None if _inbound_frame_body is not None else msg, - metadata={"frame": client_frame_index}, - ) - if ( - isinstance(_inbound_frame_body, dict) - and _inbound_frame_body.get("type") == "response.create" - ): - ws_response_create_frames += 1 - ( - msg, - _frame_modified, - _frame_reason, - ) = await _maybe_compress_response_create_frame( - msg, - frame_index=client_frame_index, - ) - _outbound_frame_body: Any = None - try: - _outbound_frame_body = json.loads(msg) - except json.JSONDecodeError: - _outbound_frame_body = None - capture_codex_wire_debug( - "ws_upstream_client_frame", - request_id=request_id, - session_id=session_id, - transport="websocket", - direction="headroom_to_upstream", - url=upstream_url, - body=_outbound_frame_body, - raw_text=None if _outbound_frame_body is not None else msg, - metadata={ - "frame": client_frame_index, - "tokens_saved_total": tokens_saved, - "transforms_applied": transforms_applied, - }, - ) - await upstream.send(msg) - except asyncio.CancelledError: - # Explicit cancel from the outer - # orchestrator — re-raise so - # ``t.cancelled()`` and ``t.exception()`` - # behave correctly in the caller. - raise - except Exception as relay_err: - # Surface real errors to the classifier - # without re-raising (existing fork - # behavior: log and return so the - # partner task can be cancelled - # deterministically). - if "WebSocketDisconnect" not in type(relay_err).__name__: - client_relay_error = relay_err - logger.debug( - f"[{request_id}] WS client→upstream relay ended: {relay_err}" - ) - else: - ws_client_disconnect_seen = True + async def _client_to_upstream() -> None: + nonlocal client_relay_error, ws_response_create_frames + nonlocal ws_client_frames_total, ws_cancel_frames + nonlocal ws_last_client_frame_type, ws_client_disconnect_seen + client_frame_index = 1 + try: + while True: + msg = await websocket.receive_text() + client_frame_index += 1 + ws_client_frames_total += 1 + if session_handle is not None: + session_handle.mark_activity() + _inbound_frame_body: Any = None + try: + _inbound_frame_body = json.loads(msg) + except json.JSONDecodeError: + _inbound_frame_body = None + ws_last_client_frame_type = ( + str(_inbound_frame_body.get("type") or "unknown") + if isinstance(_inbound_frame_body, dict) + else "non_json" + ) + if ws_last_client_frame_type == "response.cancel": + ws_cancel_frames += 1 logger.info( - "[%s] WS client disconnected session_id=%s " - "frames=%d cancels=%d last_type=%s", + "[%s] WS client sent response.cancel " + "session_id=%s frame=%d cancels=%d", request_id, session_id, - ws_client_frames_total, + client_frame_index, ws_cancel_frames, + ) + else: + logger.debug( + "[%s] WS client frame session_id=%s frame=%d type=%s", + request_id, + session_id, + client_frame_index, ws_last_client_frame_type, ) - with contextlib.suppress(Exception): - await upstream.close() + capture_codex_wire_debug( + "ws_inbound_client_frame", + request_id=request_id, + session_id=session_id, + transport="websocket", + direction="client_to_headroom", + url=_ws_url, + body=_inbound_frame_body, + raw_text=None if _inbound_frame_body is not None else msg, + metadata={"frame": client_frame_index}, + ) + if ( + isinstance(_inbound_frame_body, dict) + and _inbound_frame_body.get("type") == "response.create" + ): + ws_response_create_frames += 1 + ( + msg, + _frame_modified, + _frame_reason, + ) = await _maybe_compress_response_create_frame( + msg, + frame_index=client_frame_index, + ) + _outbound_frame_body: Any = None + try: + _outbound_frame_body = json.loads(msg) + except json.JSONDecodeError: + _outbound_frame_body = None + capture_codex_wire_debug( + "ws_upstream_client_frame", + request_id=request_id, + session_id=session_id, + transport="websocket", + direction="headroom_to_upstream", + url=upstream_url, + body=_outbound_frame_body, + raw_text=None if _outbound_frame_body is not None else msg, + metadata={ + "frame": client_frame_index, + "tokens_saved_total": tokens_saved, + "transforms_applied": transforms_applied, + }, + ) + await upstream.send(msg) + except asyncio.CancelledError: + # Explicit cancel from the outer + # orchestrator — re-raise so + # ``t.cancelled()`` and ``t.exception()`` + # behave correctly in the caller. + raise + except Exception as relay_err: + # Surface real errors to the classifier + # without re-raising (existing fork + # behavior: log and return so the + # partner task can be cancelled + # deterministically). + if "WebSocketDisconnect" not in type(relay_err).__name__: + client_relay_error = relay_err + logger.debug( + f"[{request_id}] WS client→upstream relay ended: {relay_err}" + ) + else: + ws_client_disconnect_seen = True + logger.info( + "[%s] WS client disconnected session_id=%s " + "frames=%d cancels=%d last_type=%s", + request_id, + session_id, + ws_client_frames_total, + ws_cancel_frames, + ws_last_client_frame_type, + ) + with contextlib.suppress(Exception): + await upstream.close() - async def _upstream_to_client() -> None: - """Relay upstream→client with transparent memory tool handling. + async def _upstream_to_client() -> None: + """Relay upstream→client with transparent memory tool handling. - Uses a buffer-then-decide approach: - 1. Buffer events until first output item arrives - 2. If first output is a memory tool → suppress entire response, - execute tools silently, send continuation upstream - 3. If first output is non-memory → flush buffer, stream normally - 4. Continuation response events are relayed to Codex seamlessly + Uses a buffer-then-decide approach: + 1. Buffer events until first output item arrives + 2. If first output is a memory tool → suppress entire response, + execute tools silently, send continuation upstream + 3. If first output is non-memory → flush buffer, stream normally + 4. Continuation response events are relayed to Codex seamlessly - This prevents orphaned response.created events from confusing Codex. - """ - from headroom.proxy.memory_handler import MEMORY_TOOL_NAMES + This prevents orphaned response.created events from confusing Codex. + """ + from headroom.proxy.memory_handler import MEMORY_TOOL_NAMES - # Unit 3: surface response.completed observation - # to the outer scope so the termination-cause - # classifier can prefer ``response_completed`` - # over ``upstream_disconnect``. - nonlocal response_completed_seen - nonlocal upstream_relay_error - nonlocal ws_input_tokens_total, ws_output_tokens_total - nonlocal ws_cache_read_tokens_total, ws_cache_write_tokens_total - nonlocal ws_uncached_input_tokens_total + # Unit 3: surface response.completed observation + # to the outer scope so the termination-cause + # classifier can prefer ``response_completed`` + # over ``upstream_disconnect``. + nonlocal response_completed_seen + nonlocal upstream_relay_error + nonlocal ws_input_tokens_total, ws_output_tokens_total + nonlocal ws_cache_read_tokens_total, ws_cache_write_tokens_total + nonlocal ws_uncached_input_tokens_total + nonlocal ws_recorded_input_tokens_total + nonlocal ws_recorded_output_tokens_total + nonlocal ws_recorded_cache_read_tokens_total + nonlocal ws_recorded_cache_write_tokens_total + nonlocal ws_recorded_uncached_input_tokens_total + nonlocal ws_recorded_tokens_saved_total + nonlocal ws_recorded_overhead_ms_total, ws_recorded_ttfb_ms + nonlocal ws_upstream_frames_total, ws_last_upstream_frame_type + nonlocal ws_ttfb_ms + + memory_enabled = bool(self.memory_handler and memory_user_id) + + # Per-response state (reset after each response.completed) + event_buffer: list[str] = [] + decided = False + suppress_response = False + pending_fcs: list[dict[str, Any]] = [] + resp_id: str | None = None + + def _reset() -> None: + nonlocal decided, suppress_response, resp_id + event_buffer.clear() + decided = False + suppress_response = False + pending_fcs.clear() + resp_id = None + + response_started_ms: float | None = None + + async def _record_ws_response_metrics() -> None: + """Record one completed Responses turn on long-lived WS sessions.""" nonlocal ws_recorded_input_tokens_total nonlocal ws_recorded_output_tokens_total nonlocal ws_recorded_cache_read_tokens_total nonlocal ws_recorded_cache_write_tokens_total nonlocal ws_recorded_uncached_input_tokens_total nonlocal ws_recorded_tokens_saved_total + nonlocal ws_recorded_attempted_input_tokens_total nonlocal ws_recorded_overhead_ms_total, ws_recorded_ttfb_ms - nonlocal ws_upstream_frames_total, ws_last_upstream_frame_type - nonlocal ws_ttfb_ms - memory_enabled = bool(self.memory_handler and memory_user_id) + input_delta = ws_input_tokens_total - ws_recorded_input_tokens_total + output_delta = ws_output_tokens_total - ws_recorded_output_tokens_total + cache_read_delta = ( + ws_cache_read_tokens_total - ws_recorded_cache_read_tokens_total + ) + cache_write_delta = ( + ws_cache_write_tokens_total - ws_recorded_cache_write_tokens_total + ) + uncached_delta = ( + ws_uncached_input_tokens_total + - ws_recorded_uncached_input_tokens_total + ) + saved_delta = tokens_saved - ws_recorded_tokens_saved_total + attempted_delta = ( + attempted_input_tokens_total + - ws_recorded_attempted_input_tokens_total + ) + ( + overhead_delta_ms, + ttfb_for_record_ms, + dashboard_pipeline_timing, + ) = _prepare_ws_performance_metrics() + if ( + input_delta <= 0 + and output_delta <= 0 + and cache_read_delta <= 0 + and cache_write_delta <= 0 + and uncached_delta <= 0 + and saved_delta <= 0 + and attempted_delta <= 0 + and overhead_delta_ms <= 0 + and ttfb_for_record_ms <= 0 + ): + return - # Per-response state (reset after each response.completed) - event_buffer: list[str] = [] - decided = False - suppress_response = False - pending_fcs: list[dict[str, Any]] = [] - resp_id: str | None = None - - def _reset() -> None: - nonlocal decided, suppress_response, resp_id - event_buffer.clear() - decided = False - suppress_response = False - pending_fcs.clear() - resp_id = None - - response_started_ms: float | None = None - - async def _record_ws_response_metrics() -> None: - """Record one completed Responses turn on long-lived WS sessions.""" - nonlocal ws_recorded_input_tokens_total - nonlocal ws_recorded_output_tokens_total - nonlocal ws_recorded_cache_read_tokens_total - nonlocal ws_recorded_cache_write_tokens_total - nonlocal ws_recorded_uncached_input_tokens_total - nonlocal ws_recorded_tokens_saved_total - nonlocal ws_recorded_attempted_input_tokens_total - nonlocal ws_recorded_overhead_ms_total, ws_recorded_ttfb_ms - - input_delta = ws_input_tokens_total - ws_recorded_input_tokens_total - output_delta = ( - ws_output_tokens_total - ws_recorded_output_tokens_total - ) - cache_read_delta = ( - ws_cache_read_tokens_total - ws_recorded_cache_read_tokens_total - ) - cache_write_delta = ( - ws_cache_write_tokens_total - - ws_recorded_cache_write_tokens_total - ) - uncached_delta = ( - ws_uncached_input_tokens_total - - ws_recorded_uncached_input_tokens_total - ) - saved_delta = tokens_saved - ws_recorded_tokens_saved_total - attempted_delta = ( - attempted_input_tokens_total - - ws_recorded_attempted_input_tokens_total - ) - ( - overhead_delta_ms, - ttfb_for_record_ms, - dashboard_pipeline_timing, - ) = _prepare_ws_performance_metrics() - if ( - input_delta <= 0 - and output_delta <= 0 - and cache_read_delta <= 0 - and cache_write_delta <= 0 - and uncached_delta <= 0 - and saved_delta <= 0 - and attempted_delta <= 0 - and overhead_delta_ms <= 0 - and ttfb_for_record_ms <= 0 - ): - return - - model_for_metrics = str(body.get("model") or "unknown") - latency_ms = ( - (time.perf_counter() * 1000.0 - response_started_ms) - if response_started_ms is not None - else 0.0 - ) - # Per-turn record: delta values capture - # this turn's contribution since the - # Codex WS handler accumulates session - # totals. Pre-refactor this site - # emitted only metrics + cost_tracker - # — no RequestLog, no PERF — so Codex - # traffic was invisible to - # ``headroom perf`` and the recent- - # requests feed. Funnel restores all - # four effects uniformly per turn. Per- - # turn outcomes carry ``ws_tags`` (the - # `x-headroom-tag-*` headers extracted - # at the WS upgrade) so dashboards can - # slice WS turns by tag — same surface - # as HTTP turns. - await self._record_request_outcome( - RequestOutcome( - request_id=request_id, - provider="openai", - model=model_for_metrics, - original_tokens=max(0, input_delta) + max(0, saved_delta), - optimized_tokens=max(0, input_delta), - output_tokens=max(0, output_delta), - tokens_saved=max(0, saved_delta), - attempted_input_tokens=max(0, attempted_delta), - cache_read_tokens=max(0, cache_read_delta), - cache_write_tokens=max(0, cache_write_delta), - uncached_input_tokens=max(0, uncached_delta), - total_latency_ms=latency_ms, - overhead_ms=overhead_delta_ms, - ttfb_ms=ttfb_for_record_ms, - pipeline_timing=dashboard_pipeline_timing, - transforms_applied=tuple(transforms_applied), - num_messages=len( - body.get("messages") or body.get("input") or [] - ) - if isinstance(body, dict) - else 0, - tags=ws_tags, - client=client, + model_for_metrics = str(body.get("model") or "unknown") + latency_ms = ( + (time.perf_counter() * 1000.0 - response_started_ms) + if response_started_ms is not None + else 0.0 + ) + # Per-turn record: delta values capture + # this turn's contribution since the + # Codex WS handler accumulates session + # totals. Pre-refactor this site + # emitted only metrics + cost_tracker + # — no RequestLog, no PERF — so Codex + # traffic was invisible to + # ``headroom perf`` and the recent- + # requests feed. Funnel restores all + # four effects uniformly per turn. Per- + # turn outcomes carry ``ws_tags`` (the + # `x-headroom-tag-*` headers extracted + # at the WS upgrade) so dashboards can + # slice WS turns by tag — same surface + # as HTTP turns. + await self._record_request_outcome( + RequestOutcome( + request_id=request_id, + provider="openai", + model=model_for_metrics, + original_tokens=max(0, input_delta) + max(0, saved_delta), + optimized_tokens=max(0, input_delta), + output_tokens=max(0, output_delta), + tokens_saved=max(0, saved_delta), + attempted_input_tokens=max(0, attempted_delta), + cache_read_tokens=max(0, cache_read_delta), + cache_write_tokens=max(0, cache_write_delta), + uncached_input_tokens=max(0, uncached_delta), + total_latency_ms=latency_ms, + overhead_ms=overhead_delta_ms, + ttfb_ms=ttfb_for_record_ms, + pipeline_timing=dashboard_pipeline_timing, + transforms_applied=tuple(transforms_applied), + num_messages=len( + body.get("messages") or body.get("input") or [] ) - ) - - # Structured PERF log line so ``headroom perf`` - # counts this Codex turn. Pre-P2 this emit was - # missing, which is why Codex traffic showed up - # as ``Requests: 0`` in the perf report even - # under heavy load — the same visibility bug - # class as #327's "Cache write: 0" report. - _perf_input_tokens = max(0, input_delta) - _perf_cache_read = max(0, cache_read_delta) - _perf_cache_write = max(0, cache_write_delta) - _perf_cache_hit_pct = ( - round( - _perf_cache_read - / (_perf_cache_read + _perf_cache_write) - * 100 - ) - if (_perf_cache_read + _perf_cache_write) > 0 - else 0 - ) - _perf_tok_before = _perf_input_tokens + max(0, saved_delta) - _perf_num_msgs = ( - len(body.get("messages") or body.get("input") or []) if isinstance(body, dict) - else 0 - ) - logger.info( - f"[{request_id}] PERF " - f"model={model_for_metrics} msgs={_perf_num_msgs} " - f"tok_before={_perf_tok_before} " - f"tok_after={_perf_input_tokens} " - f"tok_saved={max(0, saved_delta)} " - f"cache_read={_perf_cache_read} " - f"cache_write={_perf_cache_write} " - f"cache_hit_pct={_perf_cache_hit_pct} " - f"opt_ms={overhead_delta_ms:.0f} " - f"transforms={_summarize_transforms(transforms_applied)}" + else 0, + tags=ws_tags, + client=client, ) + ) - ws_recorded_input_tokens_total = ws_input_tokens_total - ws_recorded_output_tokens_total = ws_output_tokens_total - ws_recorded_cache_read_tokens_total = ws_cache_read_tokens_total - ws_recorded_cache_write_tokens_total = ws_cache_write_tokens_total - ws_recorded_uncached_input_tokens_total = ( - ws_uncached_input_tokens_total + # Structured PERF log line so ``headroom perf`` + # counts this Codex turn. Pre-P2 this emit was + # missing, which is why Codex traffic showed up + # as ``Requests: 0`` in the perf report even + # under heavy load — the same visibility bug + # class as #327's "Cache write: 0" report. + _perf_input_tokens = max(0, input_delta) + _perf_cache_read = max(0, cache_read_delta) + _perf_cache_write = max(0, cache_write_delta) + _perf_cache_hit_pct = ( + round( + _perf_cache_read / (_perf_cache_read + _perf_cache_write) * 100 ) - ws_recorded_tokens_saved_total = tokens_saved - ws_recorded_attempted_input_tokens_total = ( - attempted_input_tokens_total - ) - ws_recorded_overhead_ms_total = _current_ws_overhead_ms() - ws_recorded_compression_timing_totals.update( - ws_compression_timing_totals - ) - if ttfb_for_record_ms > 0: - ws_recorded_ttfb_ms = True + if (_perf_cache_read + _perf_cache_write) > 0 + else 0 + ) + _perf_tok_before = _perf_input_tokens + max(0, saved_delta) + _perf_num_msgs = ( + len(body.get("messages") or body.get("input") or []) + if isinstance(body, dict) + else 0 + ) + logger.info( + f"[{request_id}] PERF " + f"model={model_for_metrics} msgs={_perf_num_msgs} " + f"tok_before={_perf_tok_before} " + f"tok_after={_perf_input_tokens} " + f"tok_saved={max(0, saved_delta)} " + f"cache_read={_perf_cache_read} " + f"cache_write={_perf_cache_write} " + f"cache_hit_pct={_perf_cache_hit_pct} " + f"opt_ms={overhead_delta_ms:.0f} " + f"transforms={_summarize_transforms(transforms_applied)}" + ) - # The retry-loop variable is safe to close over here: - # ``_upstream_to_client`` is defined and awaited within - # a single iteration and never escapes. - _first_event_started_at = _upstream_first_event_started # noqa: B023 + ws_recorded_input_tokens_total = ws_input_tokens_total + ws_recorded_output_tokens_total = ws_output_tokens_total + ws_recorded_cache_read_tokens_total = ws_cache_read_tokens_total + ws_recorded_cache_write_tokens_total = ws_cache_write_tokens_total + ws_recorded_uncached_input_tokens_total = ws_uncached_input_tokens_total + ws_recorded_tokens_saved_total = tokens_saved + ws_recorded_attempted_input_tokens_total = attempted_input_tokens_total + ws_recorded_overhead_ms_total = _current_ws_overhead_ms() + ws_recorded_compression_timing_totals.update( + ws_compression_timing_totals + ) + if ttfb_for_record_ms > 0: + ws_recorded_ttfb_ms = True - try: - upstream_frame_index = 0 - async for msg in upstream: - upstream_frame_index += 1 - ws_upstream_frames_total += 1 - if session_handle is not None: - session_handle.mark_activity() - if ( - _first_event_started_at is not None - and "upstream_first_event" not in stage_timer - ): - if ws_ttfb_ms is None: - ws_ttfb_ms = ( - time.perf_counter() - session_started_at - ) * 1000.0 - stage_timer.record( - "upstream_first_event", - (time.perf_counter() - _first_event_started_at) - * 1000.0, - ) - if isinstance(msg, bytes): - ws_last_upstream_frame_type = "binary" - capture_codex_wire_debug( - "ws_upstream_binary_frame", - request_id=request_id, - session_id=session_id, - transport="websocket", - direction="upstream_to_headroom", - url=upstream_url, - metadata={ - "frame": upstream_frame_index, - "byte_count": len(msg), - }, - ) - await websocket.send_bytes(msg) - continue - msg_str = msg if isinstance(msg, str) else str(msg) - _upstream_frame_body: Any = None - try: - _upstream_frame_body = json.loads(msg_str) - except json.JSONDecodeError: - _upstream_frame_body = None + # The retry-loop variable is safe to close over here: + # ``_upstream_to_client`` is defined and awaited within + # a single iteration and never escapes. + _first_event_started_at = _upstream_first_event_started # noqa: B023 + + try: + upstream_frame_index = 0 + async for msg in upstream: + upstream_frame_index += 1 + ws_upstream_frames_total += 1 + if session_handle is not None: + session_handle.mark_activity() + if ( + _first_event_started_at is not None + and "upstream_first_event" not in stage_timer + ): + if ws_ttfb_ms is None: + ws_ttfb_ms = ( + time.perf_counter() - session_started_at + ) * 1000.0 + stage_timer.record( + "upstream_first_event", + (time.perf_counter() - _first_event_started_at) * 1000.0, + ) + if isinstance(msg, bytes): + ws_last_upstream_frame_type = "binary" capture_codex_wire_debug( - "ws_upstream_text_frame", + "ws_upstream_binary_frame", request_id=request_id, session_id=session_id, transport="websocket", direction="upstream_to_headroom", url=upstream_url, - body=_upstream_frame_body, - raw_text=None - if _upstream_frame_body is not None - else msg_str, - metadata={"frame": upstream_frame_index}, + metadata={ + "frame": upstream_frame_index, + "byte_count": len(msg), + }, ) + await websocket.send_bytes(msg) + continue + msg_str = msg if isinstance(msg, str) else str(msg) + _upstream_frame_body: Any = None + try: + _upstream_frame_body = json.loads(msg_str) + except json.JSONDecodeError: + _upstream_frame_body = None + capture_codex_wire_debug( + "ws_upstream_text_frame", + request_id=request_id, + session_id=session_id, + transport="websocket", + direction="upstream_to_headroom", + url=upstream_url, + body=_upstream_frame_body, + raw_text=None if _upstream_frame_body is not None else msg_str, + metadata={"frame": upstream_frame_index}, + ) - # Parse event - try: - event = json.loads(msg_str) - except (json.JSONDecodeError, TypeError): - ws_last_upstream_frame_type = "non_json" - await websocket.send_text(msg_str) - continue - - event_type = event.get("type", "") - ws_last_upstream_frame_type = str(event_type or "unknown") - logger.debug( - "[%s] WS upstream frame session_id=%s frame=%d type=%s", - request_id, - session_id, - upstream_frame_index, - ws_last_upstream_frame_type, - ) - if event_type == "response.created": - response_started_ms = time.perf_counter() * 1000.0 - ( - usage_input_tokens, - usage_output_tokens, - usage_cache_read_tokens, - usage_cache_write_tokens, - usage_uncached_tokens, - ) = _extract_responses_usage(event) - if usage_input_tokens or usage_output_tokens: - ws_input_tokens_total += usage_input_tokens - ws_output_tokens_total += usage_output_tokens - ws_cache_read_tokens_total += usage_cache_read_tokens - ws_cache_write_tokens_total += usage_cache_write_tokens - ws_uncached_input_tokens_total += usage_uncached_tokens - - if not memory_enabled: - if event_type == "response.completed": - response_completed_seen = True - await _record_ws_response_metrics() - await websocket.send_text(msg_str) - continue - - # --- Phase 1: Buffer until first output item --- - if not decided: - event_buffer.append(msg_str) - - if event_type == "response.output_item.added": - item = event.get("item", {}) - if ( - item.get("type") == "function_call" - and item.get("name") in MEMORY_TOOL_NAMES - ): - # Memory tool first → suppress entire response - suppress_response = True - decided = True - event_buffer.clear() - logger.info( - f"[{request_id}] WS Memory: Detected " - f"{item.get('name')} — suppressing response" - ) - else: - # Non-memory first → flush buffer, pass through - decided = True - for buf in event_buffer: - await websocket.send_text(buf) - event_buffer.clear() - - elif event_type == "response.completed": - # No output items at all — flush - decided = True - for buf in event_buffer: - await websocket.send_text(buf) - event_buffer.clear() - await _record_ws_response_metrics() - _reset() - response_completed_seen = True - + # Parse event + try: + event = json.loads(msg_str) + except (json.JSONDecodeError, TypeError): + ws_last_upstream_frame_type = "non_json" + await websocket.send_text(msg_str) continue - # --- Phase 2a: Suppress mode (memory response) --- - if suppress_response: - if event_type == "response.output_item.done": - item = event.get("item", {}) - if ( - item.get("type") == "function_call" - and item.get("name") in MEMORY_TOOL_NAMES - ): - pending_fcs.append(item) + event_type = event.get("type", "") + ws_last_upstream_frame_type = str(event_type or "unknown") + logger.debug( + "[%s] WS upstream frame session_id=%s frame=%d type=%s", + request_id, + session_id, + upstream_frame_index, + ws_last_upstream_frame_type, + ) + if event_type == "response.created": + response_started_ms = time.perf_counter() * 1000.0 + ( + usage_input_tokens, + usage_output_tokens, + usage_cache_read_tokens, + usage_cache_write_tokens, + usage_uncached_tokens, + ) = _extract_responses_usage(event) + if usage_input_tokens or usage_output_tokens: + ws_input_tokens_total += usage_input_tokens + ws_output_tokens_total += usage_output_tokens + ws_cache_read_tokens_total += usage_cache_read_tokens + ws_cache_write_tokens_total += usage_cache_write_tokens + ws_uncached_input_tokens_total += usage_uncached_tokens - elif event_type == "response.completed": - response_completed_seen = True - await _record_ws_response_metrics() - resp = event.get("response", {}) - resp_id = resp.get("id") - - if pending_fcs: - logger.info( - f"[{request_id}] WS Memory: Executing " - f"{len(pending_fcs)} tool(s) transparently" - ) - - # Execute memory tool calls - tool_outputs: list[dict[str, Any]] = [] - for fc in pending_fcs: - call_id = fc.get("call_id", fc.get("id", "")) - fc_name = fc.get("name", "") - args_str = fc.get("arguments", "{}") - try: - fc_args = json.loads(args_str) - except json.JSONDecodeError: - fc_args = {} - - await self.memory_handler._ensure_initialized() - if self.memory_handler._backend: - result = await self.memory_handler._execute_memory_tool( - fc_name, - fc_args, - memory_user_id, - "openai", - ) - else: - result = json.dumps( - {"error": "backend not ready"} - ) - - tool_outputs.append( - { - "type": "function_call_output", - "call_id": call_id, - "output": result, - } - ) - logger.info( - f"[{request_id}] WS Memory: Executed " - f"{fc_name} for user {memory_user_id}" - ) - - # Send continuation upstream - cont: dict[str, Any] = { - "type": "response.create", - "response": {"input": tool_outputs}, - } - if resp_id: - cont["response"]["previous_response_id"] = ( - resp_id - ) - await upstream.send(json.dumps(cont)) - logger.info( - f"[{request_id}] WS Memory: Sent continuation " - f"with {len(tool_outputs)} result(s)" - ) - - _reset() - # All events suppressed in this mode - continue - - # --- Phase 2b: Pass-through mode --- + if not memory_enabled: + if event_type == "response.completed": + response_completed_seen = True + await _record_ws_response_metrics() await websocket.send_text(msg_str) + continue - except asyncio.CancelledError: - raise - except Exception as relay_err: - if "WebSocketDisconnect" not in type(relay_err).__name__: - # Capture for the outer classifier - # so ``upstream_error`` can be - # distinguished from a clean - # upstream disconnect. - upstream_relay_error = relay_err - logger.debug( - f"[{request_id}] WS upstream→client relay ended: {relay_err}" - ) - finally: - with contextlib.suppress(Exception): - await websocket.close() + # --- Phase 1: Buffer until first output item --- + if not decided: + event_buffer.append(msg_str) - # --- Unit 3: deterministic relay-task cancellation --- - # Spawn each half as a named task so we can: - # (a) attach them to the session registry for - # ``/debug/ws-sessions``, - # (b) cancel the survivor explicitly when the - # first one exits, and - # (c) classify the termination cause for the - # duration histogram. - client_task = asyncio.create_task( - _client_to_upstream(), - name=f"codex-ws-c2u-{session_id}", - ) - upstream_task = asyncio.create_task( - _upstream_to_client(), - name=f"codex-ws-u2c-{session_id}", - ) - relay_tasks = [client_task, upstream_task] - if ws_sessions is not None: - ws_sessions.attach_tasks(session_id, relay_tasks) - metrics_for_tasks = getattr(self, "metrics", None) - if metrics_for_tasks is not None and hasattr( - metrics_for_tasks, "inc_active_relay_tasks" - ): - try: - metrics_for_tasks.inc_active_relay_tasks(len(relay_tasks)) - except Exception: # pragma: no cover - defensive - pass - - try: - done, pending = await asyncio.wait( - {client_task, upstream_task}, - return_when=asyncio.FIRST_COMPLETED, - ) - # Cancel the survivor so we don't leak the - # partner task. Suppress the CancelledError - # we just raised ourselves — any *other* - # exception from the cancelled task is - # already logged inside its own try/except. - for t in pending: - t.cancel() - if pending: - with contextlib.suppress(asyncio.CancelledError): - await asyncio.gather(*pending, return_exceptions=True) - - # Classify termination cause from whichever - # task completed first. ``CancelledError`` - # can show up on the "done" side if the - # handler itself was cancelled from outside - # (e.g. server shutdown). - for t in done: - exc = None - # Cancelled tasks raise CancelledError from - # .exception(); surface it explicitly so the - # downstream ``isinstance(exc, CancelledError)`` - # branches actually run. For any other - # unexpected state (``InvalidStateError`` if - # the task somehow isn't done — shouldn't - # happen post-gather but defensive), we - # suppress and leave ``exc`` as ``None``. - if t.cancelled(): - exc = asyncio.CancelledError() - else: - with contextlib.suppress(asyncio.InvalidStateError): - exc = t.exception() - task_name = t.get_name() or "" - if t is client_task: - if client_relay_error is not None: - termination_cause = "client_error" - elif exc is None: - termination_cause = "client_disconnect" - elif isinstance(exc, asyncio.CancelledError): - termination_cause = "client_disconnect" - else: - # Distinguish legitimate client - # disconnect exceptions from - # real errors: WebSocketDisconnect - # is a normal client exit. - if "WebSocketDisconnect" in type(exc).__name__: - termination_cause = "client_disconnect" + if event_type == "response.output_item.added": + item = event.get("item", {}) + if ( + item.get("type") == "function_call" + and item.get("name") in MEMORY_TOOL_NAMES + ): + # Memory tool first → suppress entire response + suppress_response = True + decided = True + event_buffer.clear() + logger.info( + f"[{request_id}] WS Memory: Detected " + f"{item.get('name')} — suppressing response" + ) else: - termination_cause = "client_error" - elif t is upstream_task: - if upstream_relay_error is not None: - termination_cause = "upstream_error" - logger.debug( - f"[{request_id}] WS relay {task_name} " - f"raised: {upstream_relay_error!r}" - ) - elif exc is None: - termination_cause = ( - "response_completed" - if response_completed_seen - else "upstream_disconnect" - ) - elif isinstance(exc, asyncio.CancelledError): - termination_cause = "upstream_disconnect" - else: - termination_cause = "upstream_error" - logger.debug( - f"[{request_id}] WS relay {task_name} raised: {exc!r}" - ) - if ( - ws_cancel_frames > 0 - and not response_completed_seen - and termination_cause - in {"upstream_disconnect", "client_disconnect", "unknown"} - ): - termination_cause = "client_cancel" + # Non-memory first → flush buffer, pass through + decided = True + for buf in event_buffer: + await websocket.send_text(buf) + event_buffer.clear() + + elif event_type == "response.completed": + # No output items at all — flush + decided = True + for buf in event_buffer: + await websocket.send_text(buf) + event_buffer.clear() + await _record_ws_response_metrics() + _reset() + response_completed_seen = True + + continue + + # --- Phase 2a: Suppress mode (memory response) --- + if suppress_response: + if event_type == "response.output_item.done": + item = event.get("item", {}) + if ( + item.get("type") == "function_call" + and item.get("name") in MEMORY_TOOL_NAMES + ): + pending_fcs.append(item) + + elif event_type == "response.completed": + response_completed_seen = True + await _record_ws_response_metrics() + resp = event.get("response", {}) + resp_id = resp.get("id") + + if pending_fcs: + logger.info( + f"[{request_id}] WS Memory: Executing " + f"{len(pending_fcs)} tool(s) transparently" + ) + + # Execute memory tool calls + tool_outputs: list[dict[str, Any]] = [] + for fc in pending_fcs: + call_id = fc.get("call_id", fc.get("id", "")) + fc_name = fc.get("name", "") + args_str = fc.get("arguments", "{}") + try: + fc_args = json.loads(args_str) + except json.JSONDecodeError: + fc_args = {} + + await self.memory_handler._ensure_initialized() + if self.memory_handler._backend: + result = await self.memory_handler._execute_memory_tool( + fc_name, + fc_args, + memory_user_id, + "openai", + ) + else: + result = json.dumps( + {"error": "backend not ready"} + ) + + tool_outputs.append( + { + "type": "function_call_output", + "call_id": call_id, + "output": result, + } + ) + logger.info( + f"[{request_id}] WS Memory: Executed " + f"{fc_name} for user {memory_user_id}" + ) + + # Send continuation upstream + cont: dict[str, Any] = { + "type": "response.create", + "response": {"input": tool_outputs}, + } + if resp_id: + cont["response"]["previous_response_id"] = resp_id + await upstream.send(json.dumps(cont)) + logger.info( + f"[{request_id}] WS Memory: Sent continuation " + f"with {len(tool_outputs)} result(s)" + ) + + _reset() + # All events suppressed in this mode + continue + + # --- Phase 2b: Pass-through mode --- + await websocket.send_text(msg_str) + + except asyncio.CancelledError: + raise + except Exception as relay_err: + if "WebSocketDisconnect" not in type(relay_err).__name__: + # Capture for the outer classifier + # so ``upstream_error`` can be + # distinguished from a clean + # upstream disconnect. + upstream_relay_error = relay_err + logger.debug( + f"[{request_id}] WS upstream→client relay ended: {relay_err}" + ) finally: - # In case anything above raised before the - # cancel-and-await loop ran. - for t in relay_tasks: - if not t.done(): - t.cancel() - with contextlib.suppress(asyncio.CancelledError): - await asyncio.gather(*relay_tasks, return_exceptions=True) + with contextlib.suppress(Exception): + await websocket.close() - logger.info( - "[%s] WS /v1/responses completed " - "(tokens_saved=%d, cause=%s, client_frames=%d, upstream_frames=%d, " - "cancel_frames=%d, client_disconnect=%s, last_client_type=%s, " - "last_upstream_type=%s)", - request_id, - tokens_saved, - termination_cause, - ws_client_frames_total, - ws_upstream_frames_total, - ws_cancel_frames, - ws_client_disconnect_seen, - ws_last_client_frame_type, - ws_last_upstream_frame_type, + # --- Unit 3: deterministic relay-task cancellation --- + # Spawn each half as a named task so we can: + # (a) attach them to the session registry for + # ``/debug/ws-sessions``, + # (b) cancel the survivor explicitly when the + # first one exits, and + # (c) classify the termination cause for the + # duration histogram. + client_task = asyncio.create_task( + _client_to_upstream(), + name=f"codex-ws-c2u-{session_id}", + ) + upstream_task = asyncio.create_task( + _upstream_to_client(), + name=f"codex-ws-u2c-{session_id}", + ) + relay_tasks = [client_task, upstream_task] + if ws_sessions is not None: + ws_sessions.attach_tasks(session_id, relay_tasks) + metrics_for_tasks = getattr(self, "metrics", None) + if metrics_for_tasks is not None and hasattr( + metrics_for_tasks, "inc_active_relay_tasks" + ): + try: + metrics_for_tasks.inc_active_relay_tasks(len(relay_tasks)) + except Exception: # pragma: no cover - defensive + pass + + try: + done, pending = await asyncio.wait( + {client_task, upstream_task}, + return_when=asyncio.FIRST_COMPLETED, ) - break - except Exception as ws_err: - if ws_connected: - # WS was established but broke mid-stream — re-raise - raise + # Cancel the survivor so we don't leak the + # partner task. Suppress the CancelledError + # we just raised ourselves — any *other* + # exception from the cancelled task is + # already logged inside its own try/except. + for t in pending: + t.cancel() + if pending: + with contextlib.suppress(asyncio.CancelledError): + await asyncio.gather(*pending, return_exceptions=True) - ws_last_err = ws_err - if ws_attempt >= ws_connect_attempts - 1: - break + # Classify termination cause from whichever + # task completed first. ``CancelledError`` + # can show up on the "done" side if the + # handler itself was cancelled from outside + # (e.g. server shutdown). + for t in done: + exc = None + # Cancelled tasks raise CancelledError from + # .exception(); surface it explicitly so the + # downstream ``isinstance(exc, CancelledError)`` + # branches actually run. For any other + # unexpected state (``InvalidStateError`` if + # the task somehow isn't done — shouldn't + # happen post-gather but defensive), we + # suppress and leave ``exc`` as ``None``. + if t.cancelled(): + exc = asyncio.CancelledError() + else: + with contextlib.suppress(asyncio.InvalidStateError): + exc = t.exception() + task_name = t.get_name() or "" + if t is client_task: + if client_relay_error is not None: + termination_cause = "client_error" + elif exc is None: + termination_cause = "client_disconnect" + elif isinstance(exc, asyncio.CancelledError): + termination_cause = "client_disconnect" + else: + # Distinguish legitimate client + # disconnect exceptions from + # real errors: WebSocketDisconnect + # is a normal client exit. + if "WebSocketDisconnect" in type(exc).__name__: + termination_cause = "client_disconnect" + else: + termination_cause = "client_error" + elif t is upstream_task: + if upstream_relay_error is not None: + termination_cause = "upstream_error" + logger.debug( + f"[{request_id}] WS relay {task_name} " + f"raised: {upstream_relay_error!r}" + ) + elif exc is None: + termination_cause = ( + "response_completed" + if response_completed_seen + else "upstream_disconnect" + ) + elif isinstance(exc, asyncio.CancelledError): + termination_cause = "upstream_disconnect" + else: + termination_cause = "upstream_error" + logger.debug( + f"[{request_id}] WS relay {task_name} raised: {exc!r}" + ) + if ( + ws_cancel_frames > 0 + and not response_completed_seen + and termination_cause + in {"upstream_disconnect", "client_disconnect", "unknown"} + ): + termination_cause = "client_cancel" + finally: + # In case anything above raised before the + # cancel-and-await loop ran. + for t in relay_tasks: + if not t.done(): + t.cancel() + with contextlib.suppress(asyncio.CancelledError): + await asyncio.gather(*relay_tasks, return_exceptions=True) - delay_with_jitter = jitter_delay_ms( - self.config.retry_base_delay_ms, - self.config.retry_max_delay_ms, - ws_attempt, + logger.info( + "[%s] WS /v1/responses completed " + "(tokens_saved=%d, cause=%s, client_frames=%d, upstream_frames=%d, " + "cancel_frames=%d, client_disconnect=%s, last_client_type=%s, " + "last_upstream_type=%s)", + request_id, + tokens_saved, + termination_cause, + ws_client_frames_total, + ws_upstream_frames_total, + ws_cancel_frames, + ws_client_disconnect_seen, + ws_last_client_frame_type, + ws_last_upstream_frame_type, ) - logger.warning( - f"[{request_id}] WS upstream connect failed " - f"(attempt {ws_attempt + 1}/{ws_connect_attempts}): {ws_err}; " - f"retrying in {delay_with_jitter:.0f}ms" - ) - await asyncio.sleep(delay_with_jitter / 1000) - - if not ws_connected: + else: # WS upgrade failed (HTTP 500 from OpenAI is common). # Fall back to HTTP POST streaming and relay SSE events # back over the client WebSocket transparently. @@ -5371,6 +5399,12 @@ class OpenAIHandlerMixin: "total_session", (time.perf_counter() - session_started_at) * 1000.0, ) + # Close the upstream WS on early-return paths (e.g. first-frame + # timeout after we connected). The relay path closes it via + # `async with upstream`; this idempotent backstop covers the rest. + with contextlib.suppress(Exception): + if upstream is not None: + await upstream.close() # Unit 3: deregister the session before (or independently # of) the stage-timings log so a failure there cannot leak # the registry entry. ``deregister`` is idempotent, so a @@ -5548,6 +5582,17 @@ class OpenAIHandlerMixin: await websocket.send_text(json.dumps(error_event)) return + # Refresh Codex /stats from the fallback response + # headers. We can't forward them onto the client 101 + # (already accepted headerless on this arm), but /stats + # parity is still worth keeping on a WS->HTTP fallback. + with contextlib.suppress(Exception): + from headroom.subscription.codex_rate_limits import ( + get_codex_rate_limit_state, + ) + + get_codex_rate_limit_state().update_from_headers(dict(response.headers)) + # Relay SSE events as WS text messages buffer = "" async for chunk in response.aiter_text(): diff --git a/headroom/proxy/handlers/streaming.py b/headroom/proxy/handlers/streaming.py index 08f68d6d2..df8e7aa69 100644 --- a/headroom/proxy/handlers/streaming.py +++ b/headroom/proxy/handlers/streaming.py @@ -967,6 +967,22 @@ class StreamingMixin: return StreamingResponse(_error_gen(), media_type="text/event-stream") + # Capture Codex rate-limit window data from the upstream response + # headers, for *every* status. Codex (gpt-5.x) almost always streams, so + # without this the session/weekly windows surfaced in ``/stats`` and the + # dashboard would only refresh on the rare non-streaming reply. We do this + # *before* the error early-return below so a streaming 429/5xx — the moment + # usage is most relevant — still refreshes the windows, matching the + # non-streaming HTTP handlers which capture on all statuses. + # ``update_from_headers`` is a no-op when the response carries no + # ``x-codex-*`` headers (e.g. the Anthropic streaming path), so this is + # safe to call unconditionally. + from headroom.subscription.codex_rate_limits import ( + get_codex_rate_limit_state, + ) + + get_codex_rate_limit_state().update_from_headers(dict(upstream_response.headers)) + if upstream_response.status_code >= 400: logger.warning( "[%s] Forwarding upstream streaming error status=%s url=%s", @@ -1050,9 +1066,15 @@ class StreamingMixin: headers=response_headers, ) - # Forward upstream ratelimit headers to the client + # Forward upstream rate-limit headers to the client. We pass both the + # generic ``*ratelimit*`` headers (Anthropic) and Codex's ``x-codex-*`` + # window/credit headers — the latter do not contain the ``ratelimit`` + # substring, so without the second clause the Codex CLI's own + # session/weekly display would stop updating on the streaming path. forwarded_headers = { - k: v for k, v in upstream_response.headers.items() if "ratelimit" in k.lower() + k: v + for k, v in upstream_response.headers.items() + if "ratelimit" in k.lower() or k.lower().startswith("x-codex") } async def generate(): diff --git a/tests/e2e_ws_codex_usage_headers.py b/tests/e2e_ws_codex_usage_headers.py new file mode 100644 index 000000000..99d6a822a --- /dev/null +++ b/tests/e2e_ws_codex_usage_headers.py @@ -0,0 +1,247 @@ +"""End-to-end verification that Codex x-codex-* usage headers are forwarded +onto the client-facing WebSocket handshake (101). + +Unit tests in tests/test_openai_codex_ws_lifecycle.py prove the handler +*logic* (it builds the right accept-header list), but they stub starlette's +``WebSocket.accept`` -- so they cannot prove the one risky assumption: that +starlette + uvicorn actually WRITE ``accept(headers=...)`` onto the real 101. + +This e2e closes that gap with real wire traffic and no OpenAI quota: + + 1. Stand up a *fake upstream* WS server whose handshake response carries + several ``x-codex-*`` headers PLUS a ``set-cookie`` and an + ``authorization`` header (which must NOT be forwarded). + 2. Boot the real proxy pointed at the fake upstream via --openai-api-url. + 3. Connect a real ``websockets`` client to the proxy and read + ``client.response.headers`` -- i.e. the client-facing 101. + +Asserts: + - every ``x-codex-*`` header from the upstream handshake is present on the + client 101 (original casing preserved), + - ``set-cookie`` and ``authorization`` are NOT forwarded, + - the proxy's /stats reflects the Codex window (update_from_headers parity). + +Run via: + + .venv/bin/python tests/e2e_ws_codex_usage_headers.py +""" + +from __future__ import annotations + +import asyncio +import json +import os +import socket +import subprocess +import sys +import time +import urllib.request +from pathlib import Path + +import websockets + +REPO_ROOT = Path(__file__).resolve().parent.parent + +# The x-codex-* window the fake upstream advertises on its handshake. +UPSTREAM_CODEX_HEADERS = { + "x-codex-primary-used-percent": "42", + "x-codex-primary-window-minutes": "300", + "x-codex-secondary-used-percent": "7", + "x-codex-secondary-window-minutes": "10080", +} +# Sensitive headers that must NEVER reach the client 101. +UPSTREAM_LEAK_HEADERS = { + "set-cookie": "session=should-not-forward", + "authorization": "Bearer upstream-secret-should-not-forward", +} + + +def free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +def wait_ready(port: int, timeout_s: float = 60.0) -> None: + deadline = time.time() + timeout_s + while time.time() < deadline: + try: + with urllib.request.urlopen(f"http://127.0.0.1:{port}/livez", timeout=2) as r: + if r.status == 200: + return + except Exception: + time.sleep(0.5) + raise TimeoutError("proxy not ready") + + +# ── Fake upstream WS server ───────────────────────────────────────────── +# +# The handshake response (101) carries the x-codex-* window plus sensitive +# headers, mirroring what OpenAI's Codex WS endpoint returns. + + +class FakeUpstream: + def __init__(self) -> None: + self.server: websockets.asyncio.server.Server | None = None + self.port: int = 0 + + async def _handler(self, ws): + try: + async for msg in ws: + if isinstance(msg, str): + await ws.send( + json.dumps( + { + "type": "response.completed", + "response": {"id": "fake_resp", "output": []}, + } + ) + ) + except websockets.exceptions.ConnectionClosed: + pass + + def _process_response(self, connection, request, response): + # Inject the x-codex-* window + sensitive headers onto the 101. + for name, value in {**UPSTREAM_CODEX_HEADERS, **UPSTREAM_LEAK_HEADERS}.items(): + response.headers[name] = value + return response + + async def start(self) -> int: + self.port = free_port() + self.server = await websockets.serve( + self._handler, + "127.0.0.1", + self.port, + process_response=self._process_response, + ) + return self.port + + async def stop(self) -> None: + if self.server: + self.server.close() + await self.server.wait_closed() + + +async def main_async() -> int: + fake = FakeUpstream() + upstream_port = await fake.start() + upstream_url = f"http://127.0.0.1:{upstream_port}" + + proxy_port = free_port() + print(f"[codex-hdr-e2e] fake upstream at ws://127.0.0.1:{upstream_port}") + print(f"[codex-hdr-e2e] starting proxy on :{proxy_port}") + + log_fp = open("/tmp/e2e_ws_codex_headers_proxy.log", "w") + proc = subprocess.Popen( + [ + str(REPO_ROOT / ".venv/bin/headroom"), + "proxy", + "--port", + str(proxy_port), + "--no-telemetry", + "--openai-api-url", + upstream_url, + ], + env={ + **os.environ, + "OPENAI_API_KEY": os.environ.get("OPENAI_API_KEY", "sk-fake-for-test"), + "ANTHROPIC_API_KEY": os.environ.get("ANTHROPIC_API_KEY", "sk-ant-fake"), + }, + stdout=log_fp, + stderr=subprocess.STDOUT, + cwd=str(REPO_ROOT), + ) + + failures: list[str] = [] + try: + wait_ready(proxy_port) + print("[codex-hdr-e2e] proxy ready") + + proxy_ws_url = f"ws://127.0.0.1:{proxy_port}/v1/responses" + # NOTE: API-key auth (no ChatGPT-Account-ID) so the upstream routes to + # --openai-api-url (our fake upstream). The x-codex forwarding code is + # auth-mode-agnostic -- it forwards whatever x-codex-* headers the + # upstream handshake returns -- so this exercises the exact same path + # the real chatgpt.com subscription handshake would, without leaving + # localhost. + async with websockets.connect( + proxy_ws_url, + additional_headers={ + "Authorization": "Bearer sk-fake", + "OpenAI-Beta": "responses_websockets=2026-02-06", + }, + ) as ws: + # The client-facing 101 response headers — the thing under test. + client_101 = {k.lower(): v for k, v in ws.response.headers.raw_items()} + print("[codex-hdr-e2e] client 101 headers:") + for k, v in sorted(client_101.items()): + if k.startswith("x-codex-") or k in ("set-cookie", "authorization"): + print(f" {k}: {v}") + + # 1. Every x-codex-* header forwarded. + for name, value in UPSTREAM_CODEX_HEADERS.items(): + if client_101.get(name) != value: + failures.append( + f"x-codex header not forwarded to client 101: " + f"{name} (got {client_101.get(name)!r}, want {value!r})" + ) + + # 2. Sensitive headers NOT forwarded. + for name in UPSTREAM_LEAK_HEADERS: + if name in client_101: + failures.append(f"sensitive header leaked to client 101: {name}") + + # Drive one frame so the session is real (upstream echoes completed). + await ws.send( + json.dumps( + { + "type": "response.create", + "response": {"model": "gpt-5.4", "input": "hi"}, + } + ) + ) + await asyncio.sleep(1.0) + + # 3. /stats reflects the Codex window (update_from_headers parity). + # Secondary, shape-dependent signal: look for the actual forwarded + # value, not just the word "codex" (which can appear as a null key). + await asyncio.sleep(0.5) + try: + with urllib.request.urlopen(f"http://127.0.0.1:{proxy_port}/stats", timeout=5) as r: + stats_text = r.read().decode("utf-8", errors="replace") + want = UPSTREAM_CODEX_HEADERS["x-codex-primary-used-percent"] + if want in stats_text: + print(f"[codex-hdr-e2e] /stats reflects codex window (primary-used={want})") + else: + print( + "[codex-hdr-e2e] note: /stats did not surface the codex value " + f"({want!r}); shape may differ — 101 forwarding is the primary check" + ) + except Exception as exc: # noqa: BLE001 - best-effort secondary check + print(f"[codex-hdr-e2e] /stats check skipped: {exc}") + + finally: + print("[codex-hdr-e2e] terminating proxy") + proc.terminate() + try: + proc.wait(timeout=10) + except subprocess.TimeoutExpired: + proc.kill() + log_fp.close() + await fake.stop() + + if failures: + print("\n=== CODEX-HDR E2E FAILURES ===") + for f in failures: + print(" -", f) + return 1 + print("\n=== CODEX-HDR E2E ALL GREEN ===") + return 0 + + +def main() -> int: + return asyncio.run(main_async()) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/parity/fixtures/codex_openai_contracts/codex-openai-interaction.schema.json b/tests/parity/fixtures/codex_openai_contracts/codex-openai-interaction.schema.json new file mode 100644 index 000000000..bbb6df33e --- /dev/null +++ b/tests/parity/fixtures/codex_openai_contracts/codex-openai-interaction.schema.json @@ -0,0 +1,259 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "$id": "https://github.com/chopratejas/headroom/blob/main/tests/parity/fixtures/codex_openai_contracts/codex-openai-interaction.schema.json", + "title": "Codex <-> OpenAI interaction contracts (PR #794)", + "description": "Enshrines the OpenAI interaction expectations this changeset depends on, so drift is detectable. Header values are transported as strings on the wire; the `x-headroom-parsed-type` annotation on each records the semantic type the parser (headroom/subscription/codex_rate_limits.py) coerces them to. Sources: codex_rate_limits.parse_codex_rate_limits (header family + gating), openai._extract_codex_handshake_headers (WS-101 forward filter), streaming.py (SSE forward filter).", + "$defs": { + "OpenAICodexWindowHeaders": { + "title": "x-codex-*-{primary,secondary} window headers", + "description": "A rolling rate-limit/subscription window. A window is materialized iff its `*-used-percent` header is present and numeric; `*-window-minutes` and `*-reset-at` are optional. `primary` and `secondary` are independent and either may be absent.", + "type": "object", + "properties": { + "x-codex-primary-used-percent": { + "type": "string", + "pattern": "^\\d+(?:\\.\\d+)?$", + "x-headroom-parsed-type": "float (0-100, NaN-guarded)", + "description": "Percent of the primary window consumed. Gates creation of the primary window." + }, + "x-codex-primary-window-minutes": { + "type": "string", + "pattern": "^\\d+$", + "x-headroom-parsed-type": "int", + "description": "Primary window size in minutes." + }, + "x-codex-primary-reset-at": { + "type": "string", + "pattern": "^\\d+$", + "x-headroom-parsed-type": "int (Unix epoch seconds)", + "description": "Absolute reset time of the primary window." + }, + "x-codex-secondary-used-percent": { + "type": "string", + "pattern": "^\\d+(?:\\.\\d+)?$", + "x-headroom-parsed-type": "float (0-100, NaN-guarded)", + "description": "Percent of the secondary window consumed. Gates creation of the secondary window." + }, + "x-codex-secondary-window-minutes": { + "type": "string", + "pattern": "^\\d+$", + "x-headroom-parsed-type": "int" + }, + "x-codex-secondary-reset-at": { + "type": "string", + "pattern": "^\\d+$", + "x-headroom-parsed-type": "int (Unix epoch seconds)" + } + }, + "additionalProperties": true + }, + "OpenAICodexCreditsHeaders": { + "title": "x-codex-credits-* headers", + "description": "OpenAI credits balance. A credits snapshot is materialized iff `x-codex-credits-has-credits` is present; `unlimited` defaults to false; `balance` is optional.", + "type": "object", + "properties": { + "x-codex-credits-has-credits": { + "type": "string", + "pattern": "^(?:[Tt][Rr][Uu][Ee]|[Ff][Aa][Ll][Ss][Ee]|[01])$", + "x-headroom-parsed-type": "bool (true|false|1|0, case-insensitive)", + "description": "Gates creation of the credits snapshot." + }, + "x-codex-credits-unlimited": { + "type": "string", + "pattern": "^(?:[Tt][Rr][Uu][Ee]|[Ff][Aa][Ll][Ss][Ee]|[01])$", + "x-headroom-parsed-type": "bool (defaults false when absent/unparseable)" + }, + "x-codex-credits-balance": { + "type": "string", + "x-headroom-parsed-type": "str (empty -> null)", + "description": "Free-form server string, e.g. \"$5.00\"." + } + }, + "additionalProperties": true + }, + "OpenAICodexMetaHeaders": { + "title": "x-codex meta headers", + "type": "object", + "properties": { + "x-codex-limit-name": { + "type": "string", + "x-headroom-parsed-type": "str (empty -> null)", + "description": "Active limit/model label, e.g. \"gpt-5.2-codex-sonic\"." + }, + "x-codex-promo-message": { + "type": "string", + "x-headroom-parsed-type": "str (empty -> null)", + "description": "Server announcement. Also gates snapshot creation when present." + } + }, + "additionalProperties": true + }, + "OpenAICodexRateLimitHeaders": { + "title": "Full x-codex-* header family OpenAI may emit", + "description": "Superset of every x-codex-* header headroom reads. parse_codex_rate_limits returns a snapshot iff at least one of: a primary window, a secondary window, a credits snapshot, or a non-empty promo message is present; otherwise null (treated as a non-Codex response). All members are individually optional.", + "type": "object", + "allOf": [ + { "$ref": "#/$defs/OpenAICodexWindowHeaders" }, + { "$ref": "#/$defs/OpenAICodexCreditsHeaders" }, + { "$ref": "#/$defs/OpenAICodexMetaHeaders" } + ], + "additionalProperties": true + }, + "WSUpstreamHandshakeResponse": { + "title": "OpenAI WS handshake (101) response headers consumed by the WS fix", + "description": "On the Codex WebSocket transport the x-codex-* window is delivered ONLY on the upstream handshake response (never in data frames). handle_openai_responses_ws reads upstream.response.headers here. This is the contract the connect-before-accept reorder depends on: if OpenAI ever moves these headers off the handshake (e.g. into a frame), the WS half of the fix goes stale.", + "$ref": "#/$defs/OpenAICodexRateLimitHeaders" + }, + "StreamingUpstreamResponseHeaders": { + "title": "OpenAI streaming/HTTP response headers consumed by the SSE fix", + "description": "On the streaming SSE/HTTP transport the same x-codex-* headers ride the HTTP response. streaming.py captures them on ALL statuses (including >=400) via update_from_headers, and forwards a wider set to the client (see ClientForwardedStreamingHeaders).", + "$ref": "#/$defs/OpenAICodexRateLimitHeaders" + }, + "ClientForwardedHandshakeHeaders": { + "title": "Headers forwarded onto the CLIENT-facing WS 101 (allow/deny contract)", + "description": "_extract_codex_handshake_headers forwards ONLY headers whose (lowercased) name starts with `x-codex-`. Every other upstream handshake header - notably set-cookie and authorization - MUST NOT appear on the client 101. Enforced by propertyNames below and asserted by the unit tests + tests/e2e_ws_codex_usage_headers.py.", + "type": "object", + "propertyNames": { + "pattern": "^[Xx]-[Cc][Oo][Dd][Ee][Xx]-" + }, + "not": { + "anyOf": [ + { "required": ["set-cookie"] }, + { "required": ["Set-Cookie"] }, + { "required": ["authorization"] }, + { "required": ["Authorization"] } + ] + }, + "additionalProperties": { "type": "string" } + }, + "ClientForwardedStreamingHeaders": { + "title": "Headers forwarded to the client on the streaming SSE path", + "description": "streaming.py forwards a header iff `\"ratelimit\" in name.lower()` OR `name.lower().startswith(\"x-codex\")`. This is a SUPERSET of the WS allow-list: it additionally passes generic *ratelimit* headers (e.g. the Anthropic streaming path) which do not contain the x-codex prefix.", + "type": "object", + "propertyNames": { + "pattern": "(?:[Rr][Aa][Tt][Ee][Ll][Ii][Mm][Ii][Tt])|^[Xx]-[Cc][Oo][Dd][Ee][Xx]" + }, + "additionalProperties": { "type": "string" } + }, + "WSClientRequestFrame": { + "title": "Client -> proxy WS data frame (Responses API over WS)", + "description": "Codex sends the request as a response.create envelope. The HTTP fallback unwraps `.response` for the POST body, forces stream=true, and strips any top-level `type`. A flattened variant (no envelope, fields at top level) is also tolerated by the fallback.", + "type": "object", + "properties": { + "type": { "const": "response.create" }, + "response": { + "type": "object", + "properties": { + "model": { "type": "string", "description": "e.g. gpt-5.4" }, + "input": { + "description": "String prompt or Responses-API structured input array.", + "type": ["string", "array"] + }, + "stream": { "type": "boolean" } + }, + "required": ["model"], + "additionalProperties": true + } + }, + "required": ["type", "response"], + "additionalProperties": true + }, + "WSRelayEvent": { + "title": "proxy -> client WS data frame (relayed Responses API event)", + "description": "SSE `data:` payloads relayed verbatim as WS text frames. `[DONE]` sentinels are dropped (not relayed). Every relayed event is a JSON object carrying a `type`. response.completed additionally carries usage under `response.usage`. anyOf (not oneOf): an error event also satisfies the looser lifecycle shape, which is fine.", + "anyOf": [ + { + "title": "lifecycle event", + "type": "object", + "properties": { + "type": { + "type": "string", + "examples": [ + "response.created", + "response.output_item.added", + "response.completed" + ] + }, + "response": { "type": "object", "additionalProperties": true } + }, + "required": ["type"], + "additionalProperties": true + }, + { + "title": "error event", + "type": "object", + "properties": { + "type": { "const": "error" }, + "error": { + "type": "object", + "properties": { "message": { "type": "string" } }, + "required": ["message"], + "additionalProperties": true + } + }, + "required": ["type", "error"], + "additionalProperties": true + } + ] + }, + "HTTPFallbackRequestBody": { + "title": "proxy -> OpenAI HTTP POST body on WS->HTTP fallback", + "description": "Derived from WSClientRequestFrame: the inner `.response` object, with `stream` forced to true and any top-level `type` removed.", + "type": "object", + "properties": { + "model": { "type": "string" }, + "stream": { "const": true }, + "input": { "type": ["string", "array"] } + }, + "required": ["model", "stream"], + "not": { "required": ["type"] }, + "additionalProperties": true + }, + "CodexRateLimitStatsOutput": { + "title": "headroom /stats output for the codex tracker (CodexRateLimitSnapshot.to_dict)", + "description": "Internal (headroom-emitted) shape produced from the headers above; the WS and SSE update_from_headers parity tests assert this is refreshed. Included so drift in our own surface is also caught.", + "type": "object", + "properties": { + "limit_id": { "const": "codex" }, + "limit_name": { "type": ["string", "null"] }, + "primary": { "$ref": "#/$defs/CodexWindowDict" }, + "secondary": { "$ref": "#/$defs/CodexWindowDict" }, + "credits": { + "oneOf": [ + { "type": "null" }, + { + "type": "object", + "properties": { + "has_credits": { "type": "boolean" }, + "unlimited": { "type": "boolean" }, + "balance": { "type": ["string", "null"] } + }, + "required": ["has_credits", "unlimited", "balance"], + "additionalProperties": false + } + ] + }, + "promo_message": { "type": ["string", "null"] }, + "captured_at": { "type": "number", "description": "Unix epoch seconds (float)." } + }, + "required": ["limit_id", "limit_name", "primary", "secondary", "credits", "promo_message", "captured_at"], + "additionalProperties": false + }, + "CodexWindowDict": { + "oneOf": [ + { "type": "null" }, + { + "type": "object", + "properties": { + "used_percent": { "type": "number" }, + "window_minutes": { "type": ["integer", "null"] }, + "window_label": { "type": "string", "description": "e.g. \"5h\", \"7d\"-style label; \"unknown\" when window_minutes is null." }, + "resets_at": { "type": ["integer", "null"], "description": "Unix epoch seconds." }, + "seconds_until_reset": { "type": ["integer", "null"] } + }, + "required": ["used_percent", "window_minutes", "window_label", "resets_at", "seconds_until_reset"], + "additionalProperties": false + } + ] + } + } +} diff --git a/tests/test_codex_openai_contract_parity.py b/tests/test_codex_openai_contract_parity.py new file mode 100644 index 000000000..53c612483 --- /dev/null +++ b/tests/test_codex_openai_contract_parity.py @@ -0,0 +1,161 @@ +"""Parity test: the committed Codex<->OpenAI contract schema vs the live code. + +The schema at ``tests/parity/fixtures/codex_openai_contracts/`` enshrines the +OpenAI interaction expectations the Codex usage-header fix depends on (PR #577): +the ``x-codex-*`` header family, and the WS-101 forward allow/deny rule. + +Rather than validate golden instances against the schema (which would only check +the instances, and would pull in ``jsonschema`` as a new dep), this test binds +the schema to the *live code* in both directions, so drift in either the schema +or the parser/filter fails CI: + + - every ``x-codex-*`` header the schema declares is actually consumed by + ``parse_codex_rate_limits`` (rename/removal upstream -> this test fails -> + update schema + parser together); + - ``_extract_codex_handshake_headers`` forwards exactly the ``x-codex-*`` + subset the schema's allow/deny ``$def`` permits, and never ``set-cookie`` / + ``authorization`` (the security half of the contract). +""" + +from __future__ import annotations + +import json +from pathlib import Path +from types import SimpleNamespace + +from headroom.proxy.handlers.openai import _extract_codex_handshake_headers +from headroom.subscription.codex_rate_limits import parse_codex_rate_limits + +_SCHEMA_PATH = ( + Path(__file__).parent + / "parity" + / "fixtures" + / "codex_openai_contracts" + / "codex-openai-interaction.schema.json" +) + + +def _load_schema() -> dict: + return json.loads(_SCHEMA_PATH.read_text()) + + +def _declared_codex_header_names(schema: dict) -> set[str]: + """Every ``x-codex-*`` header name declared across the header ``$defs``.""" + names: set[str] = set() + for def_name in ( + "OpenAICodexWindowHeaders", + "OpenAICodexCreditsHeaders", + "OpenAICodexMetaHeaders", + ): + props = schema["$defs"][def_name].get("properties", {}) + names.update(k for k in props if k.lower().startswith("x-codex-")) + return names + + +class _FakeHeaders: + def __init__(self, items: list[tuple[str, str]]): + self._items = items + + def raw_items(self): + return list(self._items) + + +def _fake_upstream(items: list[tuple[str, str]]): + return SimpleNamespace(response=SimpleNamespace(headers=_FakeHeaders(items))) + + +# A valid wire value for every declared header (string-typed, as on the wire). +_VALID_VALUES = { + "x-codex-primary-used-percent": "42", + "x-codex-primary-window-minutes": "300", + "x-codex-primary-reset-at": "1900000000", + "x-codex-secondary-used-percent": "7", + "x-codex-secondary-window-minutes": "10080", + "x-codex-secondary-reset-at": "1900000000", + "x-codex-credits-has-credits": "true", + "x-codex-credits-unlimited": "false", + "x-codex-credits-balance": "$5.00", + "x-codex-limit-name": "gpt-5.2-codex-sonic", + "x-codex-promo-message": "hello", +} + + +def test_schema_is_wellformed_and_has_expected_defs(): + """Guard against accidental corruption/deletion of the committed schema.""" + schema = _load_schema() + assert schema["$schema"].endswith("2020-12/schema") + for required_def in ( + "OpenAICodexRateLimitHeaders", + "WSUpstreamHandshakeResponse", + "StreamingUpstreamResponseHeaders", + "ClientForwardedHandshakeHeaders", + "ClientForwardedStreamingHeaders", + "WSClientRequestFrame", + "WSRelayEvent", + "HTTPFallbackRequestBody", + "CodexRateLimitStatsOutput", + ): + assert required_def in schema["$defs"], f"missing $def: {required_def}" + + +def test_every_declared_test_value_covers_the_schema(): + """The test's fixture values must cover exactly the declared header set, + so a header added to the schema without a value here is caught here rather + than silently skipped by the parity assertion below.""" + declared = _declared_codex_header_names(_load_schema()) + assert set(_VALID_VALUES) == declared, ( + f"_VALID_VALUES drifted from the schema header set; " + f"missing={declared - set(_VALID_VALUES)} extra={set(_VALID_VALUES) - declared}" + ) + + +def test_declared_headers_are_consumed_by_parser(): + """Every x-codex-* header the schema declares is actually parsed into the + snapshot. Catches an upstream rename/removal or a schema/parser divergence.""" + declared = _declared_codex_header_names(_load_schema()) + headers = {name: _VALID_VALUES[name] for name in declared} + snap = parse_codex_rate_limits(headers) + assert snap is not None + + # Each schema group must materialize from its declared headers. + assert snap.primary is not None, "primary window not parsed" + assert snap.primary.used_percent == 42.0 + assert snap.primary.window_minutes == 300 + assert snap.primary.resets_at == 1900000000 + assert snap.secondary is not None, "secondary window not parsed" + assert snap.secondary.used_percent == 7.0 + assert snap.credits is not None, "credits not parsed" + assert snap.credits.has_credits is True + assert snap.credits.balance == "$5.00" + assert snap.limit_name == "gpt-5.2-codex-sonic" + assert snap.promo_message == "hello" + + +def test_non_codex_headers_do_not_form_a_snapshot(): + """A response with no recognized x-codex-* headers yields no snapshot, + matching the schema's gate (snapshot iff a window/credits/promo is present).""" + assert parse_codex_rate_limits({"content-type": "text/event-stream"}) is None + assert parse_codex_rate_limits({"x-codex-unknown-future-field": "1"}) is None + + +def test_handshake_forward_obeys_allow_deny_contract(): + """_extract_codex_handshake_headers forwards exactly the x-codex-* subset + and never set-cookie / authorization (ClientForwardedHandshakeHeaders).""" + declared = _declared_codex_header_names(_load_schema()) + upstream_items = [(name, _VALID_VALUES[name]) for name in sorted(declared)] + upstream_items += [ + ("set-cookie", "session=should-not-forward"), + ("authorization", "Bearer upstream-secret"), + ("content-type", "application/json"), + ] + + forwarded = _extract_codex_handshake_headers(_fake_upstream(upstream_items)) + forwarded_names = {name.lower() for name, _ in forwarded} + + assert forwarded_names == declared, ( + f"forwarded set != declared x-codex set; " + f"missing={declared - forwarded_names} extra={forwarded_names - declared}" + ) + assert "set-cookie" not in forwarded_names + assert "authorization" not in forwarded_names + assert "content-type" not in forwarded_names diff --git a/tests/test_openai_codex_ws_lifecycle.py b/tests/test_openai_codex_ws_lifecycle.py index 3674b939a..fc55e2486 100644 --- a/tests/test_openai_codex_ws_lifecycle.py +++ b/tests/test_openai_codex_ws_lifecycle.py @@ -121,6 +121,7 @@ class _FakeWebSocket: *, disconnect_after_n_sends: int | None = None, hold_after_initial: bool = False, + call_log: list[str] | None = None, ) -> None: self.headers = {"authorization": "Bearer test"} self._frames = list(frames or []) @@ -129,14 +130,19 @@ class _FakeWebSocket: self.sent_text: list[str] = [] self.sent_bytes: list[bytes] = [] self.accepted_subprotocol: str | None = None + self.accepted_headers: list[tuple[bytes, bytes]] | None = None self.closed = False self.close_code: int | None = None + self._call_log = call_log # "client" can trip this event to simulate mid-stream disconnect. self._disconnect_event = asyncio.Event() self.client = SimpleNamespace(host="127.0.0.1", port=12345) - async def accept(self, subprotocol=None) -> None: + 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 + if self._call_log is not None: + self._call_log.append("accept") async def receive_text(self) -> str: if self._frames: @@ -169,6 +175,26 @@ class _FakeWebSocket: self._disconnect_event.set() +class _FakeHeaders: + """Minimal stand-in for websockets' handshake ``Headers``. + + Exposes both ``raw_items()`` (preferred by the production header + extractor to survive duplicate names like ``set-cookie``) and + ``items()``. + """ + + def __init__(self, pairs) -> None: + if isinstance(pairs, dict): + pairs = list(pairs.items()) + self._pairs = [(str(k), str(v)) for k, v in pairs] + + def raw_items(self): + return list(self._pairs) + + def items(self): + return list(self._pairs) + + class _FakeUpstream: """Upstream that streams scripted events then optionally blocks. @@ -187,12 +213,16 @@ class _FakeUpstream: *, hold_after_events: bool = False, raise_mid_stream: Exception | None = None, + response_headers=None, ) -> None: self._events = list(events) self._hold_after_events = hold_after_events self._raise_mid_stream = raise_mid_stream self.sent: list[str] = [] self.closed = False + # Mirror websockets' ClientConnection.response.headers, which is the + # only place OpenAI delivers the Codex x-codex-* subscription window. + self.response = SimpleNamespace(headers=_FakeHeaders(response_headers or [])) async def __aenter__(self) -> _FakeUpstream: return self @@ -219,9 +249,29 @@ class _FakeUpstream: await asyncio.Event().wait() -def _make_fake_websockets_module(upstream: _FakeUpstream): +def _make_fake_websockets_module( + upstream: _FakeUpstream | None, + *, + call_log: list[str] | None = None, + connect_error: Exception | None = None, +): + """Build a fake ``websockets`` module. + + Production now does ``upstream = await websockets.connect(...)`` (then + ``async with upstream``), so ``connect`` must return an awaitable that + resolves to the connection. ``connect_error`` makes the await raise to + simulate an upstream handshake failure. + """ module = MagicMock() - module.connect = MagicMock(return_value=upstream) + + async def _connect(*args, **kwargs): + if call_log is not None: + call_log.append("connect") + if connect_error is not None: + raise connect_error + return upstream + + module.connect = _connect module.Subprotocol = str return module @@ -572,17 +622,7 @@ async def test_upstream_connect_failure_still_deregisters_cleanly(): registered+deregistered cleanly (or never registered). Either way, no leak. """ - - class _BoomUpstream: - async def __aenter__(self): - raise RuntimeError("upstream refused") - - async def __aexit__(self, exc_type, exc, tb): - return None - - fake_ws_mod = MagicMock() - fake_ws_mod.connect = MagicMock(return_value=_BoomUpstream()) - fake_ws_mod.Subprotocol = str + fake_ws_mod = _make_fake_websockets_module(None, connect_error=RuntimeError("upstream refused")) client_ws = _FakeWebSocket(frames=[_first_frame()]) handler = _DummyOpenAIHandler() @@ -598,6 +638,157 @@ async def test_upstream_connect_failure_still_deregisters_cleanly(): assert handler.ws_sessions.active_count() == 0 +@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. + """ + fake_ws_mod = _make_fake_websockets_module( + None, connect_error=RuntimeError("HTTP 500 from upstream") + ) + + first = _first_frame() + client_ws = _FakeWebSocket(frames=[first]) + handler = _DummyOpenAIHandler() + + fallback_calls: list[tuple] = [] + + async def _fallback(websocket, body, first_msg_raw, upstream_headers, request_id): + fallback_calls.append((body, first_msg_raw)) + + handler._ws_http_fallback = _fallback # type: ignore[assignment] + + with patch.dict(sys.modules, {"websockets": fake_ws_mod}): + await handler.handle_openai_responses_ws(client_ws) + + # Client was accepted with no upstream window to forward. + assert client_ws.accepted_headers is None + # Fallback ran with the first frame. + assert len(fallback_calls) == 1 + _body, _first_raw = fallback_calls[0] + assert _first_raw == first + assert _body == json.loads(first) + # Clean teardown. + assert handler.ws_sessions.active_count() == 0 + + +@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. + """ + upstream_events = [ + json.dumps({"type": "response.created", "response": {"id": "r_1"}}), + json.dumps({"type": "response.completed", "response": {"id": "r_1"}}), + ] + call_log: list[str] = [] + upstream = _FakeUpstream(upstream_events) + fake_ws_mod = _make_fake_websockets_module(upstream, call_log=call_log) + + client_ws = _FakeWebSocket(frames=[_first_frame()], call_log=call_log) + handler = _DummyOpenAIHandler() + + with patch.dict(sys.modules, {"websockets": fake_ws_mod}): + await handler.handle_openai_responses_ws(client_ws) + + assert "connect" in call_log and "accept" in call_log + assert call_log.index("connect") < call_log.index("accept"), ( + f"connect must precede accept, got {call_log}" + ) + + +@pytest.mark.asyncio +async def test_ws_forwards_codex_headers_to_client_accept(): + """OpenAI's x-codex-* subscription window from the upstream WS + handshake must be forwarded onto the client-facing 101 (and only + that subset — never set-cookie/authorization), and Python /stats + state must be refreshed. + """ + upstream_events = [ + json.dumps({"type": "response.created", "response": {"id": "r_1"}}), + json.dumps({"type": "response.completed", "response": {"id": "r_1"}}), + ] + # Include duplicate set-cookie to ensure raw_items() is used (a plain + # dict-style .items() on real websockets Headers raises on dupes). + handshake_headers = [ + ("x-codex-primary-used-percent", "42"), + ("X-Codex-Primary-Window-Minutes", "300"), + ("set-cookie", "a=1"), + ("set-cookie", "b=2"), + ("authorization", "Bearer leak"), + ] + upstream = _FakeUpstream(upstream_events, response_headers=handshake_headers) + fake_ws_mod = _make_fake_websockets_module(upstream) + + client_ws = _FakeWebSocket(frames=[_first_frame()]) + handler = _DummyOpenAIHandler() + + captured: dict = {} + + def _fake_state(): + class _S: + def update_from_headers(self, headers): + captured.update(headers) + + return _S() + + with ( + patch.dict(sys.modules, {"websockets": fake_ws_mod}), + patch( + "headroom.subscription.codex_rate_limits.get_codex_rate_limit_state", + _fake_state, + ), + ): + await handler.handle_openai_responses_ws(client_ws) + + assert client_ws.accepted_headers is not None + names = {name.decode("latin-1").lower() for name, _ in client_ws.accepted_headers} + assert names == {"x-codex-primary-used-percent", "x-codex-primary-window-minutes"} + assert "set-cookie" not in names + assert "authorization" not in names + # Original-case names preserved on the wire. + sent = {name.decode("latin-1") for name, _ in client_ws.accepted_headers} + assert "X-Codex-Primary-Window-Minutes" in sent + # Python /stats state refreshed with the same x-codex-* subset. + assert captured == { + "x-codex-primary-used-percent": "42", + "X-Codex-Primary-Window-Minutes": "300", + } + + +@pytest.mark.asyncio +async def test_ws_first_frame_timeout_after_connect_closes_upstream(): + """If the client never sends its first frame after we connected, the + upstream WS must be closed (no leak) and the session deregistered. + """ + upstream = _FakeUpstream([], hold_after_events=True) + fake_ws_mod = _make_fake_websockets_module(upstream) + + # No frames + hold => receive_text blocks until disconnect; we force a + # short first-frame timeout so the handler hits the timeout branch. + client_ws = _FakeWebSocket(frames=[], hold_after_initial=True) + handler = _DummyOpenAIHandler() + + with ( + patch.dict(sys.modules, {"websockets": fake_ws_mod}), + patch( + "headroom.proxy.handlers.openai.WS_FIRST_FRAME_TIMEOUT_SECONDS", + 0.05, + ), + ): + await asyncio.wait_for( + handler.handle_openai_responses_ws(client_ws), + timeout=2.0, + ) + + assert upstream.closed, "upstream not closed on first-frame timeout" + assert client_ws.closed and client_ws.close_code == 1001 + assert handler.ws_sessions.active_count() == 0 + + @pytest.mark.asyncio async def test_many_concurrent_sessions_cleanly_drained(): """50 concurrent sessions: all drain; registry and named tasks go to 0.""" diff --git a/tests/test_openai_codex_ws_timings.py b/tests/test_openai_codex_ws_timings.py index f1362439a..a3b76b2a6 100644 --- a/tests/test_openai_codex_ws_timings.py +++ b/tests/test_openai_codex_ws_timings.py @@ -59,11 +59,13 @@ class _FakeWebSocket: self.sent_text: list[str] = [] self.sent_bytes: list[bytes] = [] self.accepted_subprotocol = None + self.accepted_headers: list[tuple[bytes, bytes]] | None = None self.closed = False self.close_code: int | None = None - async def accept(self, subprotocol=None) -> None: + 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 async def receive_text(self) -> str: if not self._frames: @@ -112,7 +114,13 @@ class _FakeUpstream: def _make_fake_websockets_module(upstream: _FakeUpstream): module = MagicMock() - module.connect = MagicMock(return_value=upstream) + + # Production now does ``upstream = await websockets.connect(...)`` then + # ``async with upstream`` — so connect must return an awaitable. + async def _connect(*args, **kwargs): + return upstream + + module.connect = _connect module.Subprotocol = str # the handler wraps client subprotocols if present return module @@ -221,15 +229,12 @@ def test_codex_ws_upstream_connect_failure_still_logs_timings(stage_log_capture) """A session that never connects upstream still logs a timing line with ``upstream_first_event`` absent (null).""" - class _BoomUpstream: - async def __aenter__(self): - raise RuntimeError("upstream refused") - - async def __aexit__(self, exc_type, exc, tb): - return None - fake_ws_mod = MagicMock() - fake_ws_mod.connect = MagicMock(return_value=_BoomUpstream()) + + async def _boom_connect(*args, **kwargs): + raise RuntimeError("upstream refused") + + fake_ws_mod.connect = _boom_connect fake_ws_mod.Subprotocol = str first_frame = json.dumps( @@ -251,12 +256,14 @@ def test_codex_ws_upstream_connect_failure_still_logs_timings(stage_log_capture) payload = _parse_stage_log(stage_log_capture) stages = payload["stages"] - # upstream_first_event never fired because connect failed on entry. + # upstream_first_event never fired because connect failed. assert stages.get("upstream_first_event") is None - # upstream_connect is also None because we record it only after the - # context manager successfully enters. + # upstream_connect is also None because we record it only after a + # successful ``await websockets.connect(...)``. assert stages.get("upstream_connect") is None - # But the envelope is still complete. + # But the envelope is still complete: the client is accepted and its + # first frame is read before falling back to HTTP, even on connect + # failure. assert stages["accept"] is not None assert stages["first_client_frame"] is not None assert stages["total_session"] > 0.0 diff --git a/tests/test_proxy_streaming_ratelimit_headers.py b/tests/test_proxy_streaming_ratelimit_headers.py index c6a69857c..a4f771938 100644 --- a/tests/test_proxy_streaming_ratelimit_headers.py +++ b/tests/test_proxy_streaming_ratelimit_headers.py @@ -18,6 +18,24 @@ import headroom.proxy.handlers.streaming as streaming_module from headroom.proxy.server import HeadroomProxy +@pytest.fixture(autouse=True) +def _reset_codex_rate_limit_singleton(): + """Isolate the process-global CodexRateLimitState across tests. + + The tracker is a module singleton; save/restore ``_latest`` around every + test so a captured snapshot never leaks into (or depends on) another test. + """ + from headroom.subscription.codex_rate_limits import get_codex_rate_limit_state + + state = get_codex_rate_limit_state() + saved = state._latest + state._latest = None + try: + yield + finally: + state._latest = saved + + class TestStreamingRatelimitHeaderForwarding: """Test that upstream ratelimit headers are forwarded in streaming responses.""" @@ -397,3 +415,153 @@ class TestStreamingRatelimitHeaderForwarding: assert attempts["count"] == 2 assert chunks + + @pytest.mark.asyncio + async def test_codex_rate_limit_headers_captured_and_forwarded_in_streaming(self): + """Codex x-codex-* headers must refresh /stats state AND reach the client. + + Regression guard for the bug where Codex session/weekly usage never + updated on the streaming SSE transport: the proxy neither captured the + ``x-codex-*`` headers into ``CodexRateLimitState`` nor forwarded them to + the client (the old ``"ratelimit" in k`` filter dropped them, so the + Codex CLI's own usage display also went stale). + """ + from headroom.subscription.codex_rate_limits import get_codex_rate_limit_state + + state = get_codex_rate_limit_state() + + proxy = self._create_mock_proxy() + mock_response = self._create_mock_upstream_response( + extra_headers={ + "x-codex-primary-used-percent": "42.0", + "x-codex-primary-window-minutes": "300", + "x-codex-secondary-used-percent": "8.0", + "x-codex-secondary-window-minutes": "10080", + "x-codex-limit-name": "gpt-5.4-codex", + } + ) + + mock_request = MagicMock() + proxy.http_client.build_request = MagicMock(return_value=mock_request) + proxy.http_client.send = AsyncMock(return_value=mock_response) + + result = await proxy._stream_response( + url="https://chatgpt.com/backend-api/codex/responses", + headers={"authorization": "Bearer sk-test"}, + body={"model": "gpt-5.4", "stream": True, "input": "hi"}, + provider="openai", + model="gpt-5.4", + request_id="test-codex-sse", + original_tokens=10, + optimized_tokens=10, + tokens_saved=0, + transforms_applied=[], + tags={}, + optimization_latency=0.0, + ) + + # 1. Rate-limit state refreshed from the *streaming* response. + snap = state.latest + assert snap is not None + assert snap.primary is not None + assert snap.primary.used_percent == 42.0 + assert snap.primary.window_minutes == 300 + assert snap.secondary is not None + assert snap.secondary.used_percent == 8.0 + assert snap.secondary.window_minutes == 10080 + assert snap.limit_name == "gpt-5.4-codex" + + # 2. x-codex headers forwarded so the Codex CLI's native usage display + # keeps working through the proxy on the streaming path. + assert result.headers.get("x-codex-primary-used-percent") == "42.0" + assert result.headers.get("x-codex-limit-name") == "gpt-5.4-codex" + # 3. Generic ratelimit headers still forwarded; unrelated headers dropped. + assert result.headers.get("anthropic-ratelimit-tokens-limit") == "80000" + assert result.headers.get("x-request-id") is None + + @pytest.mark.asyncio + async def test_codex_rate_limit_captured_on_streaming_429(self): + """A streaming 429 carrying x-codex-* must still refresh /stats. + + The capture runs *before* the >=400 early-return, matching the + non-streaming HTTP handlers (which capture on all statuses). A 429 is + exactly when the session/weekly windows are most worth surfacing, so the + previous success-only placement left the most important update missing. + """ + from headroom.subscription.codex_rate_limits import get_codex_rate_limit_state + + state = get_codex_rate_limit_state() + + proxy = self._create_mock_proxy() + mock_response = self._create_mock_upstream_response() + mock_response.status_code = 429 + mock_response.headers = httpx.Headers( + { + "content-type": "application/json", + "x-codex-primary-used-percent": "99.5", + "x-codex-primary-window-minutes": "300", + } + ) + mock_response.aread = AsyncMock(return_value=b'{"error":{"message":"rate limited"}}') + mock_response.aclose = AsyncMock() + + mock_request = MagicMock() + proxy.http_client.build_request = MagicMock(return_value=mock_request) + proxy.http_client.send = AsyncMock(return_value=mock_response) + + result = await proxy._stream_response( + url="https://chatgpt.com/backend-api/codex/responses", + headers={"authorization": "Bearer sk-test"}, + body={"model": "gpt-5.4", "stream": True, "input": "hi"}, + provider="openai", + model="gpt-5.4", + request_id="test-codex-429", + original_tokens=10, + optimized_tokens=10, + tokens_saved=0, + transforms_applied=[], + tags={}, + optimization_latency=0.0, + ) + + assert result.status_code == 429 + snap = state.latest + assert snap is not None + assert snap.primary is not None + assert snap.primary.used_percent == 99.5 + + @pytest.mark.asyncio + async def test_anthropic_stream_leaves_codex_state_untouched(self): + """The now-unconditional capture must be a no-op for non-Codex streams.""" + from headroom.subscription.codex_rate_limits import get_codex_rate_limit_state + + state = get_codex_rate_limit_state() + + proxy = self._create_mock_proxy() + mock_response = self._create_mock_upstream_response() # anthropic-ratelimit-* only + + mock_request = MagicMock() + proxy.http_client.build_request = MagicMock(return_value=mock_request) + proxy.http_client.send = AsyncMock(return_value=mock_response) + + await proxy._stream_response( + url="https://api.anthropic.com/v1/messages", + headers={"x-api-key": "sk-test"}, + body={ + "model": "claude-sonnet-4-20250514", + "max_tokens": 100, + "stream": True, + "messages": [{"role": "user", "content": "hi"}], + }, + provider="anthropic", + model="claude-sonnet-4-20250514", + request_id="test-anthropic-noop", + original_tokens=10, + optimized_tokens=10, + tokens_saved=0, + transforms_applied=[], + tags={}, + optimization_latency=0.0, + ) + + assert state.latest is None diff --git a/tests/test_ws_http_fallback.py b/tests/test_ws_http_fallback.py index b633cd129..69da0d654 100644 --- a/tests/test_ws_http_fallback.py +++ b/tests/test_ws_http_fallback.py @@ -31,9 +31,15 @@ class FakeWebSocket: class FakeStreamResponse: """Mock httpx streaming response.""" - def __init__(self, status_code: int = 200, sse_events: list[str] | None = None): + def __init__( + self, + status_code: int = 200, + sse_events: list[str] | None = None, + headers: dict[str, str] | None = None, + ): self.status_code = status_code self._events = sse_events or [] + self.headers = headers or {} async def aiter_text(self): for event in self._events: @@ -279,3 +285,37 @@ class TestWsHttpFallback: asyncio.run(handler._ws_http_fallback(ws, body, json.dumps(body), headers, "req_6")) assert "api.openai.com" in captured_url["url"] + + def test_fallback_refreshes_codex_rate_limit_state(self, monkeypatch): + """A successful fallback refreshes Codex /stats from response headers. + + The fallback can't forward headers onto the (already-accepted) client + 101, but it should still keep Python /stats in sync so the gauge does + not go stale when the WS upgrade fails and we drop to HTTP. + """ + handler = _make_handler() + ws = FakeWebSocket() + captured: dict[str, dict[str, str]] = {} + + class _FakeState: + def update_from_headers(self, hdrs): + captured["headers"] = dict(hdrs) + + import headroom.subscription.codex_rate_limits as crl + + monkeypatch.setattr(crl, "get_codex_rate_limit_state", lambda: _FakeState()) + + response = FakeStreamResponse( + 200, + ['data: {"type":"response.completed"}\n\n', "data: [DONE]\n\n"], + headers={ + "x-codex-primary-used-percent": "42", + "content-type": "text/event-stream", + }, + ) + handler.http_client = FakeHttpClient(response) + + body = {"model": "gpt-5.4", "input": "hi"} + asyncio.run(handler._ws_http_fallback(ws, body, json.dumps(body), {}, "req_capture")) + + assert captured["headers"]["x-codex-primary-used-percent"] == "42"