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 <cursoragent@cursor.com>

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_*` × <provider> 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:

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"]

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,
     )

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,
     )

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,
     )

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,
     )

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,
     )