From a1f0ef2a9975313c9cfb00b060ea01b1790f369e Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 17 May 2026 22:47:47 +0000 Subject: [PATCH] fix(claude_code): verify streaming wire in basic_messaging_streaming cells Address the Greptile concern that basic_messaging_streaming and basic_messaging_non_streaming used the same implementation, so a proxy that buffered the upstream stream would silently show green for the streaming row. The fix: - _basic_messaging.run_basic_messaging_cell accepts verify_streaming=True, which passes --include-partial-messages to the claude CLI. That flag causes the CLI to emit one stream_event record per upstream SSE event (message_start, content_block_delta, message_stop, ...). A buffering proxy collapses the stream to a single non-streaming response, so zero stream_event records are emitted. - The cell rejects any model whose stream_event count is below MIN_STREAM_DELTA_EVENTS (2) -- safely above the buffered case for any non-trivial reply. Same all-must-pass shape as the existing tool_use_streaming row. - All five basic_messaging_streaming/test_*.py per-provider cells now pass verify_streaming=True; the non-streaming variants are unchanged. - New unit tests cover the helper, the partial-messages flag wiring, the streamed/buffered branching, and the all-models-must-stream contract. Co-authored-by: Cursor Agent --- tests/claude_code/_basic_messaging.py | 80 ++++++- .../test_basic_messaging.py | 201 ++++++++++++++++++ .../test_anthropic.py | 24 +-- .../basic_messaging_streaming/test_azure.py | 1 + .../test_bedrock_converse.py | 1 + .../test_bedrock_invoke.py | 1 + .../test_vertex_ai.py | 1 + 7 files changed, 285 insertions(+), 24 deletions(-) create mode 100644 tests/claude_code/_driver_unit_tests/test_basic_messaging.py diff --git a/tests/claude_code/_basic_messaging.py b/tests/claude_code/_basic_messaging.py index 4e23d5f1ed2..03213c9ced8 100644 --- a/tests/claude_code/_basic_messaging.py +++ b/tests/claude_code/_basic_messaging.py @@ -10,6 +10,11 @@ Every basic_messaging cell follows the same skeleton: 4. Surface a joined failure message via `pytest.fail(...)` so the pytest run also goes red. +The streaming variant additionally passes `verify_streaming=True`, +which adds the `--include-partial-messages` CLI flag and asserts that +the proxy actually streamed the response (see the helper docstring for +the wire-level rationale). + The conftest infers `(feature_id, provider)` purely from the test file path, so each per-provider file just declares its model list and calls `run_basic_messaging_cell(...)`. This keeps all cell logic in one place @@ -23,7 +28,7 @@ collecting this module as a test file. from __future__ import annotations import os -from typing import Sequence +from typing import Any, Mapping, Sequence import pytest @@ -36,23 +41,63 @@ from tests.claude_code.cli_driver import ( PROXY_BASE_URL_ENV = "LITELLM_PROXY_BASE_URL" PROXY_API_KEY_ENV = "LITELLM_PROXY_API_KEY" +# Floor on the number of `stream_event` records (with delta payloads) +# we expect to see when the proxy actually streams. With +# `--include-partial-messages`, the CLI emits one `stream_event` per +# raw upstream SSE event — a fully-streamed response produces many +# (`message_start`, multiple `content_block_delta`s, `content_block_stop`, +# `message_delta`, `message_stop`); a proxy that buffers the upstream +# and returns a single non-streaming chunk produces 0 or 1. Floor of 2 +# is safely above the buffered case for any non-trivial reply, which +# is why the streaming cells use a "count from 1 to 5" style prompt. +MIN_STREAM_DELTA_EVENTS = 2 + + +def _count_stream_event_deltas(events: Sequence[Mapping[str, Any]]) -> int: + """Count `stream_event` records that carry an SSE event payload. + + With `--include-partial-messages`, Claude Code wraps every upstream + SSE event in a `{"type": "stream_event", "event": {...}}` record. + A buffering proxy collapses the upstream stream into a single + non-streaming response, so these records vanish. Counting them + (rather than just `len(events)`) is the wire-level signal that + "did the proxy preserve streaming?" — independent of the `system` + /`assistant`/`result` boilerplate records the CLI always emits. + """ + count = 0 + for event in events: + if event.get("type") != "stream_event": + continue + if isinstance(event.get("event"), Mapping): + count += 1 + return count + def run_basic_messaging_cell( *, compat_result, models: Sequence[str], prompt: str, + verify_streaming: bool = False, ) -> None: """Run the shared `basic_messaging_*` × cell body. - The streaming and non-streaming variants share this body because - the CLI driver consumes stdout via `subprocess.run(capture_output=True)` - after the process exits — we can only observe that events arrived, - not *when* they arrived. A wire-level "did the proxy buffer the - full response before flushing?" check therefore can't live here; - it belongs in a driver that streams stdout incrementally. Until - that exists, the streaming cells exercise the same shape and - check the same per-model outcomes as the non-streaming cells. + When ``verify_streaming=True``, the cell additionally asserts that + the proxy streamed the response end-to-end. The check works by + passing ``--include-partial-messages`` to the `claude` CLI, which + causes it to emit one ``stream_event`` record per raw upstream SSE + event (``message_start``, ``content_block_delta``, ``message_stop``, + etc.). A proxy that buffers the upstream stream and returns a + single non-streaming response collapses those records to zero — + so a floor of ``MIN_STREAM_DELTA_EVENTS`` ``stream_event`` records + catches the buffering regression without needing a streaming-aware + driver. + + This is the same shape of check as ``tool_use_streaming`` uses, + just keyed off the explicit partial-message flag so it works for + plain assistant replies (where the CLI would otherwise collapse a + streamed reply to a single ``assistant`` event in + ``--print --output-format stream-json`` mode). """ base_url = os.environ.get(PROXY_BASE_URL_ENV) api_key = os.environ.get(PROXY_API_KEY_ENV) @@ -71,11 +116,16 @@ def run_basic_messaging_cell( pytrace=False, ) + extra_args: Sequence[str] = ( + ("--include-partial-messages",) if verify_streaming else () + ) + outcomes = run_claude_models_parallel( models=models, prompt=prompt, base_url=base_url, api_key=api_key, + extra_args=extra_args, ) failures = [] @@ -99,6 +149,18 @@ def run_basic_messaging_cell( failures.append(error) continue + if verify_streaming: + stream_event_count = _count_stream_event_deltas(outcome.events) + if stream_event_count < MIN_STREAM_DELTA_EVENTS: + error = ( + f"[{model}] only {stream_event_count} stream_event records " + f"observed (< {MIN_STREAM_DELTA_EVENTS}); proxy likely " + f"buffered the upstream response instead of streaming it" + ) + compat_result.add({"status": "fail", "error": error}) + failures.append(error) + continue + compat_result.add({"status": "pass"}) if failures: diff --git a/tests/claude_code/_driver_unit_tests/test_basic_messaging.py b/tests/claude_code/_driver_unit_tests/test_basic_messaging.py new file mode 100644 index 00000000000..87ddaca8281 --- /dev/null +++ b/tests/claude_code/_driver_unit_tests/test_basic_messaging.py @@ -0,0 +1,201 @@ +"""Unit tests for the shared `run_basic_messaging_cell` helper. + +These tests mock `run_claude_models_parallel` so they exercise the +helper's branching (env-missing guard, per-model pass/fail/empty-text, +streaming wire check) without spawning the real CLI. The streaming +check is the regression we care about: a proxy that buffers the +upstream stream must turn the cell red, not green. +""" + +from __future__ import annotations + +import os +from typing import Any, Dict, List, Mapping, Optional, Sequence + +import pytest + +from tests.claude_code import _basic_messaging +from tests.claude_code._basic_messaging import ( + MIN_STREAM_DELTA_EVENTS, + _count_stream_event_deltas, + run_basic_messaging_cell, +) +from tests.claude_code.cli_driver import DriverResult + + +class _FakeResult: + """Stand-in for the test's `compat_result` fixture. + + Records every `set` / `add` payload so assertions can inspect what + the cell reported, in order, without needing the real + `pytest_runtest_logreport` plumbing from `conftest.py`. + """ + + def __init__(self) -> None: + self.rows: List[Dict[str, Any]] = [] + self.single: Optional[Dict[str, Any]] = None + + def set(self, payload: Mapping[str, Any]) -> None: + self.single = dict(payload) + + def add(self, payload: Mapping[str, Any]) -> None: + self.rows.append(dict(payload)) + + +def _streamed_events(n_deltas: int = 5) -> List[Dict[str, Any]]: + """Build a stream-json event list that *looks* streamed. + + Includes `n_deltas` `stream_event` records (matching what + `--include-partial-messages` produces) plus the usual + `system`/`assistant`/`result` boilerplate the CLI always emits. + """ + events: List[Dict[str, Any]] = [{"type": "system", "subtype": "init"}] + for i in range(n_deltas): + events.append( + { + "type": "stream_event", + "event": { + "type": "content_block_delta", + "delta": {"type": "text_delta", "text": str(i)}, + }, + } + ) + events.append( + { + "type": "assistant", + "message": {"content": [{"type": "text", "text": "1\n2\n3"}]}, + } + ) + events.append({"type": "result"}) + return events + + +def _buffered_events() -> List[Dict[str, Any]]: + """Event list a buffering proxy would produce: zero `stream_event`s.""" + return [ + {"type": "system", "subtype": "init"}, + { + "type": "assistant", + "message": {"content": [{"type": "text", "text": "1\n2\n3"}]}, + }, + {"type": "result"}, + ] + + +def _install_fake_runner(monkeypatch, *, outcomes_by_model): + """Patch `run_claude_models_parallel` to return canned outcomes. + + Captures the kwargs the cell passed in so tests can assert on + `extra_args` (which is how the streaming variant opts into + `--include-partial-messages`). + """ + captured: Dict[str, Any] = {} + + def fake(*, models, prompt, base_url, api_key, extra_args=None, **_kwargs): + captured["models"] = list(models) + captured["prompt"] = prompt + captured["base_url"] = base_url + captured["api_key"] = api_key + captured["extra_args"] = list(extra_args) if extra_args else [] + return {model: outcomes_by_model[model] for model in models} + + monkeypatch.setattr(_basic_messaging, "run_claude_models_parallel", fake) + return captured + + +@pytest.fixture(autouse=True) +def _proxy_env(monkeypatch): + monkeypatch.setenv("LITELLM_PROXY_BASE_URL", "http://localhost:4000") + monkeypatch.setenv("LITELLM_PROXY_API_KEY", "sk-test") + + +def test_count_stream_event_deltas_only_counts_records_with_event_payload(): + events = [ + {"type": "system"}, + {"type": "stream_event", "event": {"type": "message_start"}}, + {"type": "stream_event", "event": {"type": "content_block_delta"}}, + {"type": "stream_event"}, + {"type": "stream_event", "event": None}, + {"type": "stream_event", "event": "not-a-dict"}, + {"type": "assistant"}, + {"type": "result"}, + ] + assert _count_stream_event_deltas(events) == 2 + + +def test_verify_streaming_passes_when_proxy_streams(monkeypatch): + fake_result = _FakeResult() + model = "claude-haiku-4-5" + outcome = DriverResult(text="1\n2\n3", events=_streamed_events(n_deltas=5)) + captured = _install_fake_runner(monkeypatch, outcomes_by_model={model: outcome}) + + run_basic_messaging_cell( + compat_result=fake_result, + models=[model], + prompt="Count from 1 to 5, one number per line.", + verify_streaming=True, + ) + + assert captured["extra_args"] == ["--include-partial-messages"] + assert fake_result.rows == [{"status": "pass"}] + + +def test_verify_streaming_fails_when_proxy_buffers(monkeypatch): + fake_result = _FakeResult() + model = "claude-haiku-4-5" + outcome = DriverResult(text="1\n2\n3", events=_buffered_events()) + _install_fake_runner(monkeypatch, outcomes_by_model={model: outcome}) + + with pytest.raises(BaseException): + run_basic_messaging_cell( + compat_result=fake_result, + models=[model], + prompt="Count from 1 to 5, one number per line.", + verify_streaming=True, + ) + + assert len(fake_result.rows) == 1 + row = fake_result.rows[0] + assert row["status"] == "fail" + assert "stream_event" in row["error"] + assert f"< {MIN_STREAM_DELTA_EVENTS}" in row["error"] + + +def test_non_streaming_variant_omits_partial_messages_flag(monkeypatch): + """Default `verify_streaming=False` keeps the non-streaming wire identical.""" + fake_result = _FakeResult() + model = "claude-haiku-4-5" + outcome = DriverResult(text="pong", events=_buffered_events()) + captured = _install_fake_runner(monkeypatch, outcomes_by_model={model: outcome}) + + run_basic_messaging_cell( + compat_result=fake_result, + models=[model], + prompt="Reply with the single word 'pong' and nothing else.", + ) + + assert captured["extra_args"] == [] + assert fake_result.rows == [{"status": "pass"}] + + +def test_verify_streaming_requires_all_models_to_stream(monkeypatch): + """If any one tier buffers, the cell fails — same all-must-pass shape as + the non-streaming check.""" + fake_result = _FakeResult() + outcomes = { + "claude-haiku-4-5": DriverResult(text="ok", events=_streamed_events(5)), + "claude-sonnet-4-6": DriverResult(text="ok", events=_buffered_events()), + "claude-opus-4-7": DriverResult(text="ok", events=_streamed_events(5)), + } + _install_fake_runner(monkeypatch, outcomes_by_model=outcomes) + + with pytest.raises(BaseException): + run_basic_messaging_cell( + compat_result=fake_result, + models=list(outcomes.keys()), + prompt="Count from 1 to 5, one number per line.", + verify_streaming=True, + ) + + statuses = [row["status"] for row in fake_result.rows] + assert statuses == ["pass", "fail", "pass"] diff --git a/tests/claude_code/basic_messaging_streaming/test_anthropic.py b/tests/claude_code/basic_messaging_streaming/test_anthropic.py index 7d7a5bb32e2..21d8c3a0757 100644 --- a/tests/claude_code/basic_messaging_streaming/test_anthropic.py +++ b/tests/claude_code/basic_messaging_streaming/test_anthropic.py @@ -1,21 +1,14 @@ """basic_messaging_streaming x Anthropic. -Drive the real `claude` CLI in headless `--output-format stream-json` -mode against a running LiteLLM proxy that routes to Anthropic, and -report the outcome via `compat_result`. +Drive the real `claude` CLI in headless `--output-format stream-json +--include-partial-messages` mode against a running LiteLLM proxy that +routes to Anthropic, and report the outcome via `compat_result`. -The CLI is run with `--print --output-format stream-json`, which streams -incremental events as the upstream produces tokens. The cell goes green -only when every Claude tier returns a non-empty reply. - -Note: a true "did the proxy buffer the full response before flushing?" -check would require observing event arrival times on the wire, which -the `cli_driver` cannot do today — it consumes stdout via -`subprocess.run(capture_output=True)` after the process exits, so a -buffered-then-flushed response is indistinguishable from a truly -streamed one. That regression check belongs in a streaming-aware -driver; until then this cell verifies the same shape as the -non-streaming variant. +The cell goes green only when every Claude tier (a) returns a non-empty +reply and (b) the proxy actually streamed it — i.e. the CLI observed +multiple `stream_event` records carrying raw SSE deltas. A proxy that +buffers the upstream stream and returns a single non-streaming chunk +emits zero such records, which is the regression this row catches. The (feature, provider) for this cell is inferred from the file path by `tests/claude_code/conftest.py`: @@ -49,4 +42,5 @@ def test_basic_messaging_streaming_anthropic(compat_result): compat_result=compat_result, models=ANTHROPIC_MODELS, prompt="Count from 1 to 5, one number per line.", + verify_streaming=True, ) diff --git a/tests/claude_code/basic_messaging_streaming/test_azure.py b/tests/claude_code/basic_messaging_streaming/test_azure.py index f7d344615c8..5eff74fa1a1 100644 --- a/tests/claude_code/basic_messaging_streaming/test_azure.py +++ b/tests/claude_code/basic_messaging_streaming/test_azure.py @@ -36,4 +36,5 @@ def test_basic_messaging_streaming_azure(compat_result): compat_result=compat_result, models=AZURE_MODELS, prompt="Count from 1 to 5, one number per line.", + verify_streaming=True, ) diff --git a/tests/claude_code/basic_messaging_streaming/test_bedrock_converse.py b/tests/claude_code/basic_messaging_streaming/test_bedrock_converse.py index 3a8a81bfc5b..3f7b7a55863 100644 --- a/tests/claude_code/basic_messaging_streaming/test_bedrock_converse.py +++ b/tests/claude_code/basic_messaging_streaming/test_bedrock_converse.py @@ -32,4 +32,5 @@ def test_basic_messaging_streaming_bedrock_converse(compat_result): compat_result=compat_result, models=BEDROCK_CONVERSE_MODELS, prompt="Count from 1 to 5, one number per line.", + verify_streaming=True, ) diff --git a/tests/claude_code/basic_messaging_streaming/test_bedrock_invoke.py b/tests/claude_code/basic_messaging_streaming/test_bedrock_invoke.py index a24bb7f1ef4..5fc1f49442a 100644 --- a/tests/claude_code/basic_messaging_streaming/test_bedrock_invoke.py +++ b/tests/claude_code/basic_messaging_streaming/test_bedrock_invoke.py @@ -32,4 +32,5 @@ def test_basic_messaging_streaming_bedrock_invoke(compat_result): compat_result=compat_result, models=BEDROCK_INVOKE_MODELS, prompt="Count from 1 to 5, one number per line.", + verify_streaming=True, ) diff --git a/tests/claude_code/basic_messaging_streaming/test_vertex_ai.py b/tests/claude_code/basic_messaging_streaming/test_vertex_ai.py index a369c780ebc..1fa2092beb4 100644 --- a/tests/claude_code/basic_messaging_streaming/test_vertex_ai.py +++ b/tests/claude_code/basic_messaging_streaming/test_vertex_ai.py @@ -32,4 +32,5 @@ def test_basic_messaging_streaming_vertex_ai(compat_result): compat_result=compat_result, models=VERTEX_AI_MODELS, prompt="Count from 1 to 5, one number per line.", + verify_streaming=True, )