chopratejas--headroom/tests/test_responses_ws_pyo3_compression.py

"""WebSocket-shaped `/v1/responses` Rust binding tests.

The default Python CLI runtime now compresses WS `response.create` frames
through its CompressionUnit + ContentRouter path. These tests keep the
lower-level PyO3 live-zone binding covered on WebSocket-shaped envelopes
so Rust migration work cannot silently break the exposed bridge.

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",
    bypass: bool = False,
) -> 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.
    """
    if bypass:
        return first_msg_raw, False

    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, _saved, _transforms, _reason = 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)

    def test_bypass_header_short_circuits_first_frame(self):
        first_msg = json.dumps(
            {
                "type": "response.create",
                "response": {
                    "model": "gpt-5",
                    "input": [
                        {
                            "type": "function_call_output",
                            "call_id": "call_1",
                            "output": json.dumps(
                                [
                                    {
                                        "id": i,
                                        "name": f"Item {i}",
                                        "desc": "large repeated payload " * 20,
                                    }
                                    for i in range(100)
                                ]
                            ),
                        }
                    ],
                },
            }
        )

        out, modified = _ws_compress_first_frame(first_msg, bypass=True)

        assert modified is False
        assert out == 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)