From 08580f25f953f5e42a2bb17bcd3b6950c1136728 Mon Sep 17 00:00:00 2001 From: Patrick Buckley Date: Mon, 13 Jul 2026 10:23:22 -0700 Subject: [PATCH] =?UTF-8?q?fix(providers):=20review=20round=201=20?= =?UTF-8?q?=E2=80=94=20streaming=20parity=20gaps=20the=20collapse=20expose?= =?UTF-8?q?d?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Correctness (4 confirmed + 1 plausible fixed, 2 accepted+documented): - Anthropic _iter_anthropic_stream handles citations_delta: text-block citations now ride the raw block into provider_blocks, as replay requires (the retired non-streaming lane preserved them via model_dump; the streaming lane dropped them — a pre-existing main-loop gap the collapse would have extended to single-shot lanes). - Anthropic text blocks separate with "\n" at each subsequent block start, restoring the retired lane's "\n".join rendering on drained lanes AND un-fusing streamed web-search responses in the chat loop. - response.failed raises typed ResponsesStreamFailedError, listed in the provider's retryable_error_names — retry loops treat an in-band failure like the wire errors it stands in for instead of hard-stopping on a bare RuntimeError (judges keep their heuristic fallback after retries). - drain_stream folds a finish-less stream's terminal citations footer (suffix rule: pre-finish info invalidated by any later payload), so lax compat servers that never send finish_reason keep their Sources. - usage max-merge extracted as merge_usage() in _protocol.py — the one definition drain uses now and the session's inline consumer adopts on #832. Accepted + release-noted instead of coded around: strict pre-2024 compat servers that 400 on stream_options (such a server already cannot serve the chat loop; CHANGELOG caveat extended), and repeated-index parallel tool-call merging on legacy compat servers (identical to the main loop's accumulator semantics; a shared guard belongs in the #832 unification). Cleanup: run_with_deadline grows on_abandon (best-effort, cannot mask the deadline error) and both judges drop the copy-pasted abort choreography; StreamAbortRef documents the _CancelRef adoption plan; test_model_turn's fake replays through the shared as_stream adapter; docs/architecture.md drops the retired Protocol row. Tests: refusal handler pinned (was advertised, untested); typed-failed retryability; citations capture; text-block separator (plus the mixed text+search expectation updated for the separator chunk); finish-less citation fold; on_abandon firing matrix; StreamAbortRef arrival race. --- CHANGELOG.md | 16 +++- docs/architecture.md | 3 +- tests/test_deadline.py | 71 ++++++++++++++ tests/test_drain_stream.py | 28 ++++++ tests/test_model_turn.py | 27 +----- tests/test_providers.py | 96 ++++++++++++++++++- turnstone/core/deadline.py | 23 ++++- turnstone/core/judge.py | 9 +- turnstone/core/output_guard_judge.py | 3 +- turnstone/core/providers/_anthropic.py | 24 +++++ turnstone/core/providers/_openai_responses.py | 16 +++- turnstone/core/providers/_protocol.py | 87 ++++++++++++----- 12 files changed, 333 insertions(+), 70 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f546364f..c1b1d051 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -29,11 +29,17 @@ Earlier stable lines (`stable/1.6`, `stable/1.5`) are frozen. blocking read that can hit client read-timeouts — the same reason the Anthropic adapter already streamed internally — and judge timeouts now *abort* the underlying HTTP read instead of abandoning a worker thread - on a dead call. Caveats: OpenAI-compatible servers old enough to ignore - `stream_options.include_usage` stop producing usage rows on these - lanes, and legacy models that reject streaming requests outright - (o1-era) need a model alias pointing at a current model — the unread - `supports_streaming` capability flag (and its admin tile) is gone. + on a dead call. Caveats: these lanes now carry the same + `stream_options: {include_usage: true}` the chat loop always sent — + OpenAI-compatible servers old enough to *ignore* it stop producing + usage rows on these lanes, and servers strict enough to *reject* + unknown fields (pre-2024 llama.cpp/proxy builds) will 400 — such a + server already couldn't serve turnstone's chat loop, but a + judge/utility alias pointed at one worked on 1.7 and needs to move to + a current server. Legacy models that reject streaming requests + outright (o1-era) likewise need a model alias pointing at a current + model — the unread `supports_streaming` capability flag (and its + admin tile) is gone. - **One turn interface for every model call: `core/model_turn.py` (#827).** Judges (intent + output guard), perception, title generation, compaction, diff --git a/docs/architecture.md b/docs/architecture.md index 7035d9c7..54d0816b 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -609,8 +609,7 @@ LLMProvider (protocol) | Method | Purpose | |--------|---------| -| `create_streaming()` | Streaming request, yields normalized `StreamChunk` objects | -| `create_completion()` | Non-streaming request, returns `CompletionResult` | +| `create_streaming()` | The one transport: streaming request, yields normalized `StreamChunk` objects (single-shot callers accumulate via `drain_stream()` into a `CompletionResult`) | | `get_capabilities()` | Per-model flags (`ModelCapabilities`) | | `convert_tools()` | Translate OpenAI tool schemas to provider format | | `retryable_error_names` | Exception class names that trigger retry | diff --git a/tests/test_deadline.py b/tests/test_deadline.py index 6750d072..b8f97d8c 100644 --- a/tests/test_deadline.py +++ b/tests/test_deadline.py @@ -64,3 +64,74 @@ def test_cancel_returns_promptly() -> None: assert time.monotonic() - start < 1.0 stragglers = [t for t in threading.enumerate() if t.name == "dl-cancel" and not t.daemon] assert stragglers == [], f"non-daemon worker survived: {stragglers}" + + +def test_on_abandon_fires_on_timeout_and_cancel_but_not_success() -> None: + calls: list[str] = [] + + with pytest.raises(DeadlineExceededError): + run_with_deadline( + lambda: time.sleep(2.0), + timeout=0.1, + poll=0.05, + thread_name="dl-abandon-t", + on_abandon=lambda: calls.append("timeout"), + ) + assert calls == ["timeout"] + + cancel = threading.Event() + cancel.set() + with pytest.raises(DeadlineCancelledError): + run_with_deadline( + lambda: time.sleep(2.0), + timeout=10.0, + cancel_event=cancel, + poll=0.05, + thread_name="dl-abandon-c", + on_abandon=lambda: calls.append("cancel"), + ) + assert calls == ["timeout", "cancel"] + + assert run_with_deadline(lambda: 7, timeout=1.0, on_abandon=lambda: calls.append("no")) == 7 + assert calls == ["timeout", "cancel"] + + +def test_on_abandon_errors_do_not_mask_the_deadline_error() -> None: + def _boom() -> None: + raise RuntimeError("abort hook broke") + + with pytest.raises(DeadlineExceededError): + run_with_deadline( + lambda: time.sleep(2.0), + timeout=0.1, + poll=0.05, + thread_name="dl-abandon-e", + on_abandon=_boom, + ) + + +class TestStreamAbortRef: + def test_abort_closes_captured_stream(self) -> None: + from unittest.mock import MagicMock + + from turnstone.core.deadline import StreamAbortRef + + ref = StreamAbortRef() + stream = MagicMock() + ref.append(stream) + stream.close.assert_not_called() + ref.abort() + stream.close.assert_called_once() + + def test_late_arriving_stream_closes_on_append(self) -> None: + # The arrival race: abort fires while the worker is still inside the + # SDK connect — the handle must close the moment it is captured. + from unittest.mock import MagicMock + + from turnstone.core.deadline import StreamAbortRef + + ref = StreamAbortRef() + ref.abort() + stream = MagicMock() + ref.append(stream) + stream.close.assert_called_once() diff --git a/tests/test_drain_stream.py b/tests/test_drain_stream.py index 620c88dc..79cf43c9 100644 --- a/tests/test_drain_stream.py +++ b/tests/test_drain_stream.py @@ -233,6 +233,34 @@ class TestInfoDelta: ) assert result.content == "\n\nSources:\n- x" + def test_finishless_stream_still_folds_terminal_citations(self): + # A lax compat server may never send finish_reason (a shape the + # adapters tolerate); the adapters' post-loop citations footer is + # then the stream's final suffix and must still fold into content. + result = drain_stream( + iter( + [ + StreamChunk(content_delta="body"), + StreamChunk(info_delta="Sources:\n- x"), + ] + ) + ) + assert result.content == "body\n\nSources:\n- x" + assert result.finish_reason == "stop" + + def test_finishless_interleaved_ping_still_dropped(self): + # Info followed by real payload is a status ping, not the terminal + # footer — dropped even when the stream never sends finish_reason. + result = drain_stream( + iter( + [ + StreamChunk(info_delta="[Searching: x]"), + StreamChunk(content_delta="answer"), + ] + ) + ) + assert result.content == "answer" + class TestErrorPropagation: def test_mid_stream_exception_propagates_verbatim(self): diff --git a/tests/test_model_turn.py b/tests/test_model_turn.py index 289b5b2c..deb683b2 100644 --- a/tests/test_model_turn.py +++ b/tests/test_model_turn.py @@ -14,6 +14,7 @@ from unittest.mock import MagicMock import pytest +from tests._session_helpers import as_stream from turnstone.core.model_turn import ( ModelLane, finalize_provider_blocks, @@ -26,7 +27,6 @@ from turnstone.core.providers._protocol import ( CompletionResult, ModelCapabilities, StreamChunk, - ToolCallDelta, UsageInfo, ) from turnstone.core.trajectory import Role, ToolCall, Turn @@ -34,8 +34,9 @@ from turnstone.core.trajectory import Role, ToolCall, Turn class _FakeProvider: """Records every ``create_streaming`` call; replays scripted results - as single-chunk streams (multi-chunk accumulation is pinned by the - dedicated ``drain_stream`` unit tests).""" + as single-chunk streams via the shared ``as_stream`` adapter + (multi-chunk accumulation is pinned by the dedicated ``drain_stream`` + unit tests).""" provider_name = "openai-compatible" @@ -48,25 +49,7 @@ class _FakeProvider: def create_streaming(self, **kwargs: Any) -> list[StreamChunk]: self.calls.append(kwargs) - result = self.results.pop(0) - return [ - StreamChunk( - content_delta=result.content or "", - reasoning_delta=result.reasoning or "", - tool_call_deltas=[ - ToolCallDelta( - index=i, - id=tc.get("id", ""), - name=tc.get("function", {}).get("name", ""), - arguments_delta=tc.get("function", {}).get("arguments", ""), - ) - for i, tc in enumerate(result.tool_calls or []) - ], - usage=result.usage, - finish_reason=result.finish_reason or "stop", - provider_blocks=list(result.provider_blocks or []), - ) - ] + return as_stream(self.results.pop(0)) def _fake_registry( diff --git a/tests/test_providers.py b/tests/test_providers.py index 5ab4ccde..bb7134fc 100644 --- a/tests/test_providers.py +++ b/tests/test_providers.py @@ -1117,6 +1117,63 @@ class TestAnthropicProvider: assert tc["function"]["name"] == "read_file" assert json.loads(tc["function"]["arguments"]) == {"path": "foo.py"} + @patch("turnstone.core.providers._anthropic._ensure_anthropic") + def test_drained_stream_separates_text_blocks(self, mock_ensure: MagicMock) -> None: + # The retired non-streaming lane joined text blocks with "\n"; the + # iterator now emits the separator at each subsequent text block + # start, so drained content keeps the block boundary (web-search + # responses interleave text / server-tool / text). + client = MagicMock() + client.messages.stream.return_value = fake_anthropic_stream( + [ + SimpleNamespace(type="text", text="Before the search."), + SimpleNamespace(type="text", text="After the results."), + ] + ) + + result = drain_stream( + self.provider.create_streaming( + client=client, + model="claude-sonnet-4-20250514", + messages=[{"role": "user", "content": "hi"}], + ) + ) + assert result.content == "Before the search.\nAfter the results." + + @patch("turnstone.core.providers._anthropic._ensure_anthropic") + def test_streaming_captures_text_block_citations(self, mock_ensure: MagicMock) -> None: + # citations_delta events must land on the raw block: Anthropic + # requires citations to replay unmodified alongside their + # web_search_tool_result blocks on later turns, and the retired + # non-streaming lane preserved them via model_dump. + start = _anthropic_event("content_block_start", block_type="text", index=0) + # A real dict from model_dump so the raw block accepts the + # citations append (a MagicMock auto-dict would swallow it). + start.content_block.model_dump.return_value = {"type": "text", "text": ""} + cite = _anthropic_event("content_block_delta", delta_type="citations_delta", index=0) + cite.delta.citation = {"type": "web_search_result_location", "url": "https://x.test"} + text = _anthropic_event( + "content_block_delta", delta_type="text_delta", text="cited claim", index=0 + ) + finish = _anthropic_event("message_delta", stop_reason="end_turn", usage_output_tokens=1) + + stream_ctx = MagicMock() + stream_ctx.__enter__ = MagicMock(return_value=iter([start, cite, text, finish])) + stream_ctx.__exit__ = MagicMock(return_value=False) + client = MagicMock() + client.messages.stream.return_value = stream_ctx + + result = drain_stream( + self.provider.create_streaming( + client=client, + model="claude-sonnet-4-20250514", + messages=[{"role": "user", "content": "hi"}], + ) + ) + assert result.provider_blocks, "expected the text block in provider_blocks" + citations = result.provider_blocks[0].get("citations") + assert citations == [{"type": "web_search_result_location", "url": "https://x.test"}] + @patch("turnstone.core.providers._anthropic._ensure_anthropic") def test_drained_stream_usage(self, mock_ensure: MagicMock) -> None: client = MagicMock() @@ -2767,9 +2824,13 @@ class TestAnthropicWebSearch: chunks = list(self.provider._iter_anthropic_stream(events)) text_chunks = [c for c in chunks if c.content_delta] info_chunks = [c for c in chunks if c.info_delta] - assert len(text_chunks) == 2 - assert text_chunks[0].content_delta == "Let me search." - assert text_chunks[1].content_delta == "Based on the results..." + # Three content chunks: the second text BLOCK opens with the "\n" + # separator (matching the retired non-streaming join), then its text. + assert [c.content_delta for c in text_chunks] == [ + "Let me search.", + "\n", + "Based on the results...", + ] assert len(info_chunks) == 1 assert "test query" in info_chunks[0].info_delta @@ -5180,3 +5241,32 @@ class TestResponsesDrainedStream: assert result.usage is not None assert result.usage.prompt_tokens == 10 assert result.usage.completion_tokens == 5 + + def test_refusal_renders_in_content(self) -> None: + # The response.refusal.done handler (CHANGELOG "refusals render in + # content") — a refused turn must not drain to empty content. + events = [ + SimpleNamespace(type="response.refusal.done", refusal="cannot help with that"), + *self._make_events(), + ] + result = self._drain(events) + assert result.content == "[Refused: cannot help with that]" + + def test_failed_event_raises_typed_retryable_error(self) -> None: + # An in-band response.failed terminal event raises the TYPED error, + # whose name the provider advertises as retryable — retry loops keep + # retrying a transient in-band failure instead of hard-stopping on a + # bare RuntimeError (the retired non-streaming lane degraded instead). + from turnstone.core.providers._openai_responses import ResponsesStreamFailedError + + events = [ + SimpleNamespace( + type="response.failed", + response=SimpleNamespace( + status="failed", error=SimpleNamespace(message="model overloaded") + ), + ) + ] + with pytest.raises(ResponsesStreamFailedError, match="model overloaded"): + self._drain(events) + assert "ResponsesStreamFailedError" in self.provider.retryable_error_names diff --git a/turnstone/core/deadline.py b/turnstone/core/deadline.py index 5c9c1004..3bf4cc8b 100644 --- a/turnstone/core/deadline.py +++ b/turnstone/core/deadline.py @@ -52,7 +52,9 @@ class StreamAbortRef(list[Any]): worker is still inside the SDK's connect (no handle captured yet), the handle is closed the moment it arrives. Both paths tolerate double close (SDK ``close()`` is idempotent) so no lock is needed — mirrors - ``ChatSession``'s ``_CancelRef``. + ``ChatSession``'s ``_CancelRef``, which adopts this class when the + main loop moves onto ``model_turn`` (#832); until then a hardening + fix here must be mirrored there. """ __slots__ = ("_aborted",) @@ -82,6 +84,7 @@ def run_with_deadline( cancel_event: threading.Event | None = None, poll: float = 1.0, thread_name: str = "deadline-worker", + on_abandon: Callable[[], None] | None = None, ) -> _T: """Run ``fn()`` on a daemon thread, bounded by ``timeout``/``cancel_event``. @@ -91,6 +94,14 @@ def run_with_deadline( abort the worker thread is abandoned; being a daemon it cannot block process or interpreter exit. + ``on_abandon`` runs (best-effort) right before either abandonment raise — + the one hook for releasing whatever the worker is blocked on, so callers + can't wire one abort path and forget the other. The canonical use is + ``on_abandon=abort_ref.abort`` with a :class:`StreamAbortRef` threaded + into the provider call as ``cancel_ref``: the abandoned worker's blocked + HTTP read raises promptly instead of pinning the thread until the next + upstream chunk. + ``poll`` bounds how often ``cancel_event`` is checked (and thus the worst- case latency from a cancel to this function returning). """ @@ -104,6 +115,12 @@ def run_with_deadline( threading.Thread(target=_runner, name=thread_name, daemon=True).start() + def _abandon(exc: Exception) -> None: + if on_abandon is not None: + with contextlib.suppress(Exception): + on_abandon() + raise exc + deadline = time.monotonic() + timeout while True: # Prefer a result that has already arrived over a deadline or cancel @@ -119,10 +136,10 @@ def run_with_deadline( raise payload # type: ignore[misc] # ok=False ⇒ payload is the raised exc if cancel_event is not None and cancel_event.is_set(): - raise DeadlineCancelledError + _abandon(DeadlineCancelledError()) remaining = deadline - time.monotonic() if remaining <= 0: - raise DeadlineExceededError + _abandon(DeadlineExceededError()) try: ok, payload = box.get(timeout=min(remaining, poll)) except queue.Empty: diff --git a/turnstone/core/judge.py b/turnstone/core/judge.py index ac520e95..82cb8020 100644 --- a/turnstone/core/judge.py +++ b/turnstone/core/judge.py @@ -1388,9 +1388,9 @@ class IntentJudge: # models aren't penalised for slow earlier turns. per_call_timeout = max(self._config.timeout, 5.0) # at least 5s # Fresh per turn — aborting turn N's stream must never touch a - # later turn's. On the abandon paths below, closing the stream - # makes the daemon worker's blocked HTTP read raise promptly - # instead of pinning the thread until the next upstream chunk. + # later turn's. run_with_deadline's abandon hook closes it so + # the daemon worker's blocked HTTP read raises promptly instead + # of pinning the thread until the next upstream chunk. abort_ref = StreamAbortRef() try: # Each turn runs on its own daemon worker (1s cancel polling). @@ -1414,12 +1414,11 @@ class IntentJudge: timeout=per_call_timeout, cancel_event=cancel_event, thread_name="judge-api", + on_abandon=abort_ref.abort, ) except DeadlineCancelledError: - abort_ref.abort() return None except DeadlineExceededError: - abort_ref.abort() log.info("judge.turn.timeout", turn=turn + 1, timeout=per_call_timeout) # Safety net: if we have a partial result from a previous turn, # try to parse a verdict from it before giving up. diff --git a/turnstone/core/output_guard_judge.py b/turnstone/core/output_guard_judge.py index d8166944..04aac758 100644 --- a/turnstone/core/output_guard_judge.py +++ b/turnstone/core/output_guard_judge.py @@ -564,12 +564,11 @@ class OutputGuardJudge: timeout=timeout, cancel_event=cancel_event, thread_name="output-guard-judge", + on_abandon=abort_ref.abort, ) except DeadlineCancelledError: - abort_ref.abort() return self._error_verdict(verdict_id, call_id, start, "cancelled") except DeadlineExceededError: - abort_ref.abort() return self._error_verdict(verdict_id, call_id, start, "timeout") except Exception as e: return self._error_verdict( diff --git a/turnstone/core/providers/_anthropic.py b/turnstone/core/providers/_anthropic.py index 727e95cd..8775930f 100644 --- a/turnstone/core/providers/_anthropic.py +++ b/turnstone/core/providers/_anthropic.py @@ -925,6 +925,7 @@ class AnthropicProvider: server_tool_blocks: dict[int, dict[str, str]] = {} # Capture raw content blocks for multi-turn preservation raw_blocks: dict[int, dict[str, Any]] = {} + saw_text_block = False for event in stream: sc = StreamChunk() @@ -933,6 +934,15 @@ class AnthropicProvider: if event_type == "content_block_start": block = event.content_block raw_blocks[event.index] = _block_to_dict(block) + if block.type == "text": + # Separate consecutive text blocks the way the Messages + # API's non-streaming shape reads when joined — without + # this, a web-search turn's post-results text block fuses + # onto the pre-search sentence. Emitted as a plain + # content delta so it never lands in the raw block. + if saw_text_block: + sc.content_delta = "\n" + saw_text_block = True if block.type == "tool_use": idx = next_tool_index tool_block_to_index[event.index] = idx @@ -983,6 +993,20 @@ class AnthropicProvider: raw_blocks[event.index]["signature"] = ( raw_blocks[event.index].get("signature", "") + delta.signature ) + elif delta.type == "citations_delta": + # Text-block citations (web-search models) arrive one + # citation object per delta. They must ride the raw + # block for replay — Anthropic requires citations to be + # passed back unmodified alongside their + # web_search_tool_result blocks on later turns. + if event.index in raw_blocks: + citation = getattr(delta, "citation", None) + if citation is not None: + raw_blocks[event.index].setdefault("citations", []).append( + citation.model_dump() + if hasattr(citation, "model_dump") + else citation + ) elif delta.type == "input_json_delta": if event.index in server_tool_blocks: # Accumulate server tool input (search query) diff --git a/turnstone/core/providers/_openai_responses.py b/turnstone/core/providers/_openai_responses.py index 1965eb40..f9dacb15 100644 --- a/turnstone/core/providers/_openai_responses.py +++ b/turnstone/core/providers/_openai_responses.py @@ -42,6 +42,18 @@ from turnstone.core.trajectory import materialize_attachments log = structlog.get_logger(__name__) +class ResponsesStreamFailedError(RuntimeError): + """An in-band ``response.failed`` terminal event (HTTP 200 stream). + + Typed (and listed in the provider's ``retryable_error_names``) so + retry loops treat an in-band failure — typically a transient + server-side error delivered inside an otherwise healthy stream — + like the wire-level errors it stands in for, instead of stopping on + a bare ``RuntimeError``. Callers that give up after retries keep + their normal degrade paths (judges fall back to the heuristic tier). + """ + + def convert_content_parts(parts: list[Any]) -> list[dict[str, Any]]: """Convert Chat Completions content parts to Responses API format. @@ -668,7 +680,7 @@ class OpenAIResponsesProvider: response = getattr(event, "response", None) error = getattr(response, "error", None) if response else None error_msg = getattr(error, "message", "Unknown error") if error else "Unknown error" - raise RuntimeError(f"Responses API error: {error_msg}") + raise ResponsesStreamFailedError(f"Responses API error: {error_msg}") log.debug( "openai.responses.response", @@ -697,7 +709,7 @@ class OpenAIResponsesProvider: @property def retryable_error_names(self) -> frozenset[str]: - return RETRYABLE_ERROR_NAMES + return RETRYABLE_ERROR_NAMES | {"ResponsesStreamFailedError"} # -- reasoning extraction ------------------------------------------------ diff --git a/turnstone/core/providers/_protocol.py b/turnstone/core/providers/_protocol.py index b9d564e1..2e923847 100644 --- a/turnstone/core/providers/_protocol.py +++ b/turnstone/core/providers/_protocol.py @@ -67,6 +67,39 @@ class CompletionResult: reasoning: str = "" +def merge_usage(acc: UsageInfo | None, new: UsageInfo) -> UsageInfo: + """Merge one stream-chunk usage report into an accumulator, per-field max. + + Anthropic splits a request's usage across events (``message_start`` + carries prompt tokens with completion 0; ``message_delta`` carries + completion tokens and may omit prompt tokens), so neither first-wins + nor last-wins sees both — the max-merge does. ``total_tokens`` is + recomputed from the merged parts. Returns a fresh ``UsageInfo`` and + never mutates ``new`` (the provider's object). + + ``ChatSession``'s inline chunk consumer implements the same rule over + its dict-shaped accumulator; it adopts this helper when the main loop + moves onto ``model_turn`` (#832). + """ + if acc is None: + return UsageInfo( + prompt_tokens=new.prompt_tokens, + completion_tokens=new.completion_tokens, + total_tokens=new.total_tokens, + cache_creation_tokens=new.cache_creation_tokens, + cache_read_tokens=new.cache_read_tokens, + ) + prompt = max(acc.prompt_tokens, new.prompt_tokens) + completion = max(acc.completion_tokens, new.completion_tokens) + return UsageInfo( + prompt_tokens=prompt, + completion_tokens=completion, + total_tokens=prompt + completion, + cache_creation_tokens=max(acc.cache_creation_tokens, new.cache_creation_tokens), + cache_read_tokens=max(acc.cache_read_tokens, new.cache_read_tokens), + ) + + def drain_stream(chunks: Iterator[StreamChunk]) -> CompletionResult: """Drain a ``create_streaming`` iterator into a ``CompletionResult``. @@ -84,12 +117,14 @@ def drain_stream(chunks: Iterator[StreamChunk]) -> CompletionResult: - ``provider_blocks`` replaces on each non-empty emission — every adapter attaches its full block list exactly once, on or after the terminal chunk. - - ``info_delta`` BEFORE the finish reason is transient status (server- + - ``info_delta`` interleaved with the data is transient status (server- side search pings) that the non-streaming lane never surfaced — drop - it. ``info_delta`` AFTER the finish reason is the citations footer - (``format_citations("", annotations).strip()``); folding it back as - ``content + "\\n\\n" + info`` byte-matches the non-streaming lane's - ``format_citations(content, annotations)`` append. + it. ``info_delta`` AFTER the finish reason, or forming the stream's + final suffix when a lax server never sent a finish reason, is the + citations footer (``format_citations("", annotations).strip()``); + folding it back as ``content + "\\n\\n" + info`` byte-matches the + non-streaming lane's ``format_citations(content, annotations)`` + append. Raises whatever the underlying stream raises — retry/deadline/fallback policy stays with the caller, exactly as with the old non-streaming @@ -98,6 +133,7 @@ def drain_stream(chunks: Iterator[StreamChunk]) -> CompletionResult: content_parts: list[str] = [] reasoning_parts: list[str] = [] trailing_info_parts: list[str] = [] + suffix_info_parts: list[str] = [] tool_calls_acc: dict[int, dict[str, Any]] = {} usage: UsageInfo | None = None finish_reason: str | None = None @@ -120,33 +156,32 @@ def drain_stream(chunks: Iterator[StreamChunk]) -> CompletionResult: if tcd.arguments_delta: tc["function"]["arguments"] += tcd.arguments_delta if sc.usage is not None: - if usage is None: - usage = UsageInfo( - prompt_tokens=sc.usage.prompt_tokens, - completion_tokens=sc.usage.completion_tokens, - total_tokens=sc.usage.total_tokens, - cache_creation_tokens=sc.usage.cache_creation_tokens, - cache_read_tokens=sc.usage.cache_read_tokens, - ) - else: - usage.prompt_tokens = max(usage.prompt_tokens, sc.usage.prompt_tokens) - usage.completion_tokens = max(usage.completion_tokens, sc.usage.completion_tokens) - usage.total_tokens = usage.prompt_tokens + usage.completion_tokens - usage.cache_creation_tokens = max( - usage.cache_creation_tokens, sc.usage.cache_creation_tokens - ) - usage.cache_read_tokens = max(usage.cache_read_tokens, sc.usage.cache_read_tokens) + usage = merge_usage(usage, sc.usage) if sc.finish_reason: finish_reason = sc.finish_reason if sc.provider_blocks: provider_blocks = sc.provider_blocks - # Pre-finish info is transient status — intentionally dropped; - # only the trailing (post-finish) citations footer folds back. - if sc.info_delta and finish_reason is not None: - trailing_info_parts.append(sc.info_delta) + if sc.info_delta: + if finish_reason is not None: + trailing_info_parts.append(sc.info_delta) + else: + # Candidate citations footer on a finish-less stream — + # kept only while nothing but info follows it (below). + suffix_info_parts.append(sc.info_delta) + if ( + sc.content_delta + or sc.reasoning_delta + or sc.tool_call_deltas + or sc.usage is not None + or sc.finish_reason + or sc.provider_blocks + ): + # Real payload after a pre-finish info chunk: that info was an + # interleaved status ping, not the terminal citations footer. + suffix_info_parts.clear() content = "".join(content_parts) - for info in trailing_info_parts: + for info in trailing_info_parts + suffix_info_parts: content += "\n\n" + info tool_calls = [tool_calls_acc[i] for i in sorted(tool_calls_acc)]