diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 57f5f32bf..a598ca85b 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -1915,9 +1915,14 @@ class OpenAIHandlerMixin: # deregister / metrics / stage-timings emission as usual. return - # PR-C5: WebSocket /v1/responses no longer compresses in Python. - # The Rust handler covers item-aware compression; the WS first - # frame is now passed through unmodified. + # PR-C5 retired Python compression on this path expecting the + # standalone Rust proxy binary to take over. That binary is not + # deployed by the CLI today (`headroom proxy` runs only Python + # via uvicorn). Hot-fix follow-up to PR #406: the first frame + # is now compressed via the inline PyO3 binding right before + # upstream send (see the compression block further below). + # Subsequent client→upstream frames in the relay loop remain + # unmodified — multi-frame compression is a separate follow-up. body: dict[str, Any] = {} tokens_saved = 0 try: @@ -2063,6 +2068,76 @@ class OpenAIHandlerMixin: except Exception as e: logger.warning(f"[{request_id}] WS Memory injection failed: {e}") + # Hot-fix follow-up to PR #406 — inline Rust compression on the + # WS first frame before forwarding upstream. PR #406 enabled + # the same call for HTTP /v1/responses; PR-C5's "WS-side + # compression is a follow-up" note is closed here. Codex + # subscription users default to WebSocket transport for + # /v1/responses (proxy-confirmed via #409 reviewer testing), + # so without this call subscription traffic flows through + # Headroom uncompressed. + # + # The first frame may be either: + # • {"type": "response.create", "response": {...payload...}} + # • the payload directly (older shapes) + # We unwrap, compress the inner payload via the PyO3 dispatcher, + # and re-wrap so both shapes work. + # + # Re-parses from `first_msg_raw` rather than reusing `body` + # because `body` may be partially mutated if memory injection + # raised an exception above (in which case `first_msg_raw` is + # the canonical pre-memory bytes that will actually be sent + # upstream). The PyO3 binding never raises (passthrough on + # internal errors), but we wrap the call site in try/except + # anyway so a JSON-shape edge case can never break the WS + # session. + if self.config.optimize: + try: + from headroom._core import ( + compress_openai_responses_live_zone as _rust_compress_responses, + ) + + _ws_auth_mode = classify_auth_mode(ws_headers) + try: + _send_body = json.loads(first_msg_raw) + except json.JSONDecodeError: + _send_body = None + + if isinstance(_send_body, dict): + _wrapped = "response" in _send_body and isinstance( + _send_body["response"], dict + ) + _inner = _send_body["response"] if _wrapped else _send_body + _model = (_inner.get("model") if isinstance(_inner, dict) else None) or "" + + _inner_bytes = json.dumps(_inner).encode("utf-8") + _new_bytes, _modified = _rust_compress_responses( + _inner_bytes, + _ws_auth_mode.value, + _model, + ) + if _modified: + try: + _new_inner = json.loads(_new_bytes) + except json.JSONDecodeError: + _new_inner = None + if isinstance(_new_inner, dict): + if _wrapped: + _send_body["response"] = _new_inner + else: + _send_body = _new_inner + first_msg_raw = json.dumps(_send_body) + logger.info( + f"[{request_id}] WS /v1/responses compressed " + f"{len(_inner_bytes):,}→{len(_new_bytes):,} bytes " + f"(auth_mode={_ws_auth_mode.value})" + ) + except Exception as _ce: + logger.warning( + f"[{request_id}] WS /v1/responses compression failed; " + f"forwarding original frame: {type(_ce).__name__}: {_ce}" + ) + # --- Connect to upstream OpenAI WebSocket --- logger.info(f"[{request_id}] WS /v1/responses connecting to {upstream_url}") diff --git a/tests/test_responses_ws_pyo3_compression.py b/tests/test_responses_ws_pyo3_compression.py new file mode 100644 index 000000000..759e9e259 --- /dev/null +++ b/tests/test_responses_ws_pyo3_compression.py @@ -0,0 +1,205 @@ +"""WebSocket /v1/responses compression integration tests. + +PR-C5 retired Python compression on the WS path expecting the standalone +Rust proxy binary to take over (it isn't deployed by the CLI). PR #406 +re-enabled compression on the HTTP path via the inline PyO3 binding. +This module pins the WS-side equivalent: the first frame from the client +must be compressed via the same PyO3 dispatcher before forwarding to +the upstream WebSocket. + +The tests exercise the compression *transformation logic* in isolation — +they replicate the body-shape handling the WS handler does (envelope +detect, compress inner, re-wrap) without spinning up a full WebSocket +session. Full session-lifecycle coverage already exists in +`test_openai_codex_ws_lifecycle.py`. +""" + +from __future__ import annotations + +import json +from typing import Any + +import pytest + + +def _ensure_binding(): + """Skip if the Rust extension hasn't been built (mirrors the pattern + in `test_responses_pyo3_compression.py`).""" + try: + from headroom._core import compress_openai_responses_live_zone + + return compress_openai_responses_live_zone + except ImportError: + pytest.skip("headroom._core not built — run scripts/build_rust_extension.sh") + + +def _ws_compress_first_frame( + first_msg_raw: str, + auth_mode_value: str = "payg", +) -> tuple[str, bool]: + """Replicates the WS-handler compression block as a pure function. + + Returns ``(new_first_msg_raw, modified)``. The real handler embeds + this logic inline in `handle_openai_responses_ws`; pulling it out + here lets us pin the exact byte-shape contract without standing + up a full WebSocket fixture. If you change the handler's + compression block, mirror it here so the tests catch the drift. + """ + compress = _ensure_binding() + + try: + send_body: Any = json.loads(first_msg_raw) + except json.JSONDecodeError: + return first_msg_raw, False + + if not isinstance(send_body, dict): + return first_msg_raw, False + + wrapped = "response" in send_body and isinstance(send_body["response"], dict) + inner = send_body["response"] if wrapped else send_body + model = (inner.get("model") if isinstance(inner, dict) else None) or "" + + inner_bytes = json.dumps(inner).encode("utf-8") + new_bytes, modified = compress(inner_bytes, auth_mode_value, model) + if not modified: + return first_msg_raw, False + + try: + new_inner = json.loads(new_bytes) + except json.JSONDecodeError: + return first_msg_raw, False + + if not isinstance(new_inner, dict): + return first_msg_raw, False + + if wrapped: + send_body["response"] = new_inner + else: + send_body = new_inner + return json.dumps(send_body), True + + +class TestWrappedEnvelopeShape: + """Codex's WebSocket protocol wraps the Responses payload in a + ``response.create`` envelope. The WS handler must unwrap to compress + and re-wrap to forward.""" + + def test_passthrough_when_inner_has_no_input_array(self): + # No `input` array → dispatcher's NoMessagesArray path → passthrough. + first_msg = json.dumps( + { + "type": "response.create", + "response": {"model": "gpt-5"}, + } + ) + out, modified = _ws_compress_first_frame(first_msg) + assert modified is False + assert out == first_msg + + def test_envelope_preserved_on_passthrough(self): + first_msg = json.dumps( + { + "type": "response.create", + "response": { + "model": "gpt-5", + "input": [{"type": "message", "role": "user", "content": "hi"}], + }, + } + ) + out, modified = _ws_compress_first_frame(first_msg) + # Single small user message → no compression applies. + assert modified is False + assert json.loads(out) == json.loads(first_msg) + + +class TestUnwrappedShape: + """Older Codex versions (and some test fixtures) send the Responses + payload directly as the first frame, without a `response.create` + envelope. The handler must work for both shapes.""" + + def test_passthrough_when_no_input_array(self): + first_msg = json.dumps({"model": "gpt-5"}) + out, modified = _ws_compress_first_frame(first_msg) + assert modified is False + assert out == first_msg + + def test_passthrough_when_empty_input(self): + first_msg = json.dumps({"model": "gpt-5", "input": []}) + out, modified = _ws_compress_first_frame(first_msg) + assert modified is False + assert out == first_msg + + +class TestNonJsonFirstFrame: + """If the first frame isn't JSON, we forward it byte-for-byte rather + than crashing the WS session.""" + + def test_garbage_passthrough(self): + out, modified = _ws_compress_first_frame("not actually json") + assert modified is False + assert out == "not actually json" + + def test_json_array_passthrough(self): + # Top-level array isn't a Responses envelope. + first_msg = json.dumps([1, 2, 3]) + out, modified = _ws_compress_first_frame(first_msg) + assert modified is False + assert out == first_msg + + def test_json_string_passthrough(self): + first_msg = json.dumps("a string at the top level") + out, modified = _ws_compress_first_frame(first_msg) + assert modified is False + assert out == first_msg + + +class TestAuthModeForwarded: + """Every F1 AuthMode value reaches the dispatcher without raising. + The dispatcher itself currently treats all modes identically (per-mode + tuning is F2.2 follow-up), but the call must not fail on any value + the F1 classifier produces.""" + + @pytest.mark.parametrize( + "auth_mode_value", + ["payg", "oauth", "subscription", "unknown"], + ) + def test_all_auth_modes_accepted(self, auth_mode_value: str): + first_msg = json.dumps({"model": "gpt-5", "input": []}) + out, modified = _ws_compress_first_frame(first_msg, auth_mode_value) + assert modified is False + assert out == first_msg + + +class TestNoExceptionLeak: + """The WS handler wraps the compression block in try/except so a + JSON-shape edge case can never crash the WS session. This pins the + contract that no input shape produces an exception in the + transformation function.""" + + @pytest.mark.parametrize( + "first_msg", + [ + "", + "{", + "{", + "}", + "[", + "null", + "true", + "0", + json.dumps({}), + json.dumps({"response": "not a dict"}), + json.dumps({"response": []}), + json.dumps({"response": None}), + json.dumps({"response": {"input": "not an array"}}), + json.dumps({"input": "string instead of array"}), + json.dumps({"input": None}), + ], + ) + def test_no_exception_for_garbage_shapes(self, first_msg: str): + # Should never raise — return passthrough on anything malformed. + out, modified = _ws_compress_first_frame(first_msg) + # Regardless of result, no exception leaked. modified might be + # False here (garbage input → no compression). + assert isinstance(out, str) + assert isinstance(modified, bool)