fix(proxy): inline WebSocket /v1/responses compression via PyO3

PR-C5 retired Python compression on the WebSocket /v1/responses path
expecting the standalone Rust proxy binary to take over. That binary
isn't deployed by the CLI today (`headroom proxy` runs only the Python
proxy via uvicorn). PR #406 closed the equivalent gap on the HTTP path;
this commit closes the gap on WebSocket.

Subscription users matter most here. The PR #409 reviewer confirmed
empirically that ChatGPT-subscription Codex CLI defaults to WebSocket
transport for /v1/responses. After PR #409 the routing reaches Headroom;
before this commit, every byte was forwarded uncompressed.

# What this does

In handle_openai_responses_ws, after memory injection finalises the
first-frame body, we re-parse first_msg_raw, detect the wrap shape
(Codex sends either {"type": "response.create", "response": {...}} or
the payload directly), call the PyO3 binding from PR #406 on the
inner payload, and re-wrap on modified return.

The compression engine runs in Rust — the binding exposes it inline
so the Python WS handler can call it without a process chain.

# Failure mode

Wrapped in try/except — passthrough on any unexpected condition.

# Subsequent client→upstream frames

Out of scope. Multi-frame compression is a separate follow-up.

# Tests

26 new tests in tests/test_responses_ws_pyo3_compression.py pinning
the body-shape contract: wrapped + unwrapped envelopes, garbage shapes
(no exception leak), every F1 AuthMode accepted. Existing WS
lifecycle + timings tests still pass.

Refs: PR #406 (HTTP path), PR-C5 (the retirement that needs closing)

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}")
 

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)