diff --git a/tests/test_drain_stream.py b/tests/test_drain_stream.py index dae5e63c..bfb2bbe3 100644 --- a/tests/test_drain_stream.py +++ b/tests/test_drain_stream.py @@ -271,6 +271,21 @@ class TestInfoDelta: class TestErrorPropagation: + def test_httpx_transport_error_becomes_retryable_incomplete(self): + # Streaming moves the body read out of the SDK's wrapped request: + # a mid-body wire failure surfaces as a raw httpx.TransportError + # no retry predicate recognizes. The drain re-raises it (chained, + # message preserved) as the retryable IncompleteStreamError. + import httpx + + def chunks(): + yield StreamChunk(content_delta="partial") + raise httpx.RemoteProtocolError("peer closed connection") + + with pytest.raises(IncompleteStreamError, match="RemoteProtocolError") as excinfo: + drain_stream(chunks()) + assert isinstance(excinfo.value.__cause__, httpx.RemoteProtocolError) + def test_mid_stream_exception_propagates_verbatim(self): # Retry/deadline/fallback policy is the caller's — the drain adds # no exception translation, exactly like the old transport. diff --git a/tests/test_providers.py b/tests/test_providers.py index 633acb09..648b9f39 100644 --- a/tests/test_providers.py +++ b/tests/test_providers.py @@ -743,6 +743,15 @@ class TestOpenAIProvider: assert result.tool_calls[0]["function"]["arguments"] == '{"path": "foo.py"}' assert result.finish_reason == "tool_calls" + def _drain_chunks(self, chunks: list[Any]): + client = MagicMock() + client.chat.completions.create.return_value = chunks + return drain_stream( + self.provider.create_streaming( + client=client, model="m", messages=[{"role": "user", "content": "x"}] + ) + ) + def test_streaming_remaps_index_degenerate_parallel_calls(self) -> None: # Historical compat servers (older vLLM, some llama.cpp builds) # stream every parallel call at index 0 as whole deltas. The @@ -750,53 +759,71 @@ class TestOpenAIProvider: # index's current call, so BOTH consumers (drain_stream and the # chat loop's accumulator) see distinct calls; id-less argument # fragments keep following their index's current slot. - def _tc_chunk(tc_id: str, name: str, args: str) -> SimpleNamespace: - tc = SimpleNamespace( - index=0, id=tc_id, function=SimpleNamespace(name=name, arguments=args) - ) - delta = SimpleNamespace( - content=None, - tool_calls=[tc], - reasoning=None, - reasoning_content=None, - annotations=None, - ) - return SimpleNamespace( - choices=[SimpleNamespace(finish_reason=None, delta=delta)], usage=None - ) - - finish = SimpleNamespace( - choices=[ - SimpleNamespace( - finish_reason="tool_calls", - delta=SimpleNamespace( - content=None, - tool_calls=None, - reasoning=None, - reasoning_content=None, - annotations=None, - ), - ) - ], - usage=None, - ) - chunks = [ - _tc_chunk("a", "read", '{"p": 1}'), - _tc_chunk("b", "write", '{"p": 2}'), - finish, - ] - client = MagicMock() - client.chat.completions.create.return_value = chunks - - result = drain_stream( - self.provider.create_streaming( - client=client, model="m", messages=[{"role": "user", "content": "x"}] - ) + result = self._drain_chunks( + [ + _openai_stream_chunk( + tool_calls=[ + _openai_tool_call_delta( + index=0, tc_id="a", name="read", arguments='{"p": 1}' + ) + ] + ), + _openai_stream_chunk( + tool_calls=[ + _openai_tool_call_delta( + index=0, tc_id="b", name="write", arguments='{"p": 2}' + ) + ] + ), + _openai_stream_chunk(finish_reason="tool_calls"), + ] ) assert [tc["id"] for tc in result.tool_calls] == ["a", "b"] assert result.tool_calls[0]["function"]["arguments"] == '{"p": 1}' assert result.tool_calls[1]["function"]["arguments"] == '{"p": 2}' + def test_streaming_splits_idless_degenerate_parallel_calls(self) -> None: + # The same degenerate servers may omit ids entirely: a delta that + # ANNOUNCES a name for a slot that already accumulated arguments is + # a second whole call, not a fragment — without the split, two + # id-less calls fuse into one with concatenated garbage arguments. + result = self._drain_chunks( + [ + _openai_stream_chunk( + tool_calls=[_openai_tool_call_delta(index=0, name="read", arguments='{"a": 1}')] + ), + _openai_stream_chunk( + tool_calls=[ + _openai_tool_call_delta(index=0, name="write", arguments='{"b": 2}') + ] + ), + _openai_stream_chunk(finish_reason="tool_calls"), + ] + ) + assert [tc["function"]["name"] for tc in result.tool_calls] == ["read", "write"] + assert result.tool_calls[0]["function"]["arguments"] == '{"a": 1}' + assert result.tool_calls[1]["function"]["arguments"] == '{"b": 2}' + + def test_fragmented_single_call_does_not_split(self) -> None: + # The normal well-behaved shape — name announced once, arguments + # streamed in fragments — must stay ONE call. + result = self._drain_chunks( + [ + _openai_stream_chunk( + tool_calls=[_openai_tool_call_delta(index=0, tc_id="a", name="read")] + ), + _openai_stream_chunk( + tool_calls=[_openai_tool_call_delta(index=0, arguments='{"p": ')] + ), + _openai_stream_chunk( + tool_calls=[_openai_tool_call_delta(index=0, arguments='"x"}')] + ), + _openai_stream_chunk(finish_reason="tool_calls"), + ] + ) + assert len(result.tool_calls) == 1 + assert result.tool_calls[0]["function"]["arguments"] == '{"p": "x"}' + def test_drained_stream_usage(self) -> None: client = MagicMock() client.chat.completions.create.return_value = fake_chat_stream( @@ -2262,6 +2289,39 @@ class TestGoogleProviderFidelity: assert fc.provider_blocks[0]["id"] == "call_abc" assert fc.provider_blocks[0]["function"]["name"] == "write_file" + def test_tap_slots_degenerate_calls_like_the_mirror(self) -> None: + # The raw fidelity tap and the base iterator slot the SAME delta + # sequence identically: two wire-index-0 calls with distinct ids + # yield TWO raw dicts, each keeping its own thought_signature — a + # fused single dict would fail _prepare_messages' length gate and + # silently drop the signature lane from the replay. + from turnstone.core.providers._google import GoogleProvider + + prov = GoogleProvider() + + def _tc(tc_id: str, name: str, args: str, sig: str) -> MagicMock: + tcd = _openai_tool_call_delta(index=0, tc_id=tc_id, name=name, arguments=args) + tcd.__pydantic_extra__ = {"thought_signature": sig} + return tcd + + chunks = [ + _openai_stream_chunk(tool_calls=[_tc("c1", "read", '{"a": 1}', "sig_a")]), + _openai_stream_chunk(tool_calls=[_tc("c2", "write", '{"b": 2}', "sig_b")]), + _openai_stream_chunk(finish_reason="tool_calls"), + ] + client = MagicMock() + client.chat.completions.create.return_value = chunks + + result = drain_stream( + prov.create_streaming( + client=client, model="gemini-2.5-pro", messages=[{"role": "user", "content": "x"}] + ) + ) + assert len(result.tool_calls) == 2 + assert len(result.provider_blocks) == 2 + assert [b["thought_signature"] for b in result.provider_blocks] == ["sig_a", "sig_b"] + assert [b["id"] for b in result.provider_blocks] == ["c1", "c2"] + def test_base_chat_lane_emits_no_provider_blocks(self) -> None: """The base chat lane carries NO provider_blocks for tool calls — only the Google subclass's tap captures raw dicts. Pinned on the @@ -5358,6 +5418,48 @@ class TestResponsesDrainedStream: with pytest.raises(RuntimeError, match="bad tool schema"): self._drain(events) + def test_terminal_without_payload_keeps_collected_blocks(self) -> None: + # The collected output_item.done blocks came from the stream, not + # the missing terminal payload — a payload-less terminal must not + # drop them (reasoning items lost = replay degradation). + item = SimpleNamespace(type="reasoning", summary=[]) + item.model_dump = lambda **_kw: {"type": "reasoning", "summary": []} # type: ignore[method-assign] + events = [ + SimpleNamespace(type="response.output_item.done", item=item), + SimpleNamespace(type="response.completed", response=None), + ] + result = self._drain(events) + assert result.finish_reason == "stop" + assert result.provider_blocks == [{"type": "reasoning", "summary": []}] + + def test_truncation_rebuild_recovers_annotations(self) -> None: + # The message item open at truncation never got output_item.done, + # so its annotations were never collected — the terminal rebuild + # walks them so truncated web-search turns keep their Sources. + ann = MagicMock() + ann.type = "url_citation" + ann.url_citation = MagicMock(title="Cite", url="https://cite.test") + part = SimpleNamespace(type="output_text", text="truncated bod", annotations=[ann]) + msg_item = SimpleNamespace(type="message", content=[part], status="incomplete") + msg_item.model_dump = lambda **_kw: {"type": "message", "content": []} # type: ignore[method-assign] + usage = SimpleNamespace( + input_tokens=10, + output_tokens=5, + total_tokens=15, + input_tokens_details=SimpleNamespace(cached_tokens=0), + ) + events = [ + SimpleNamespace(type="response.output_text.delta", delta="truncated bod"), + SimpleNamespace( + type="response.incomplete", + response=SimpleNamespace(status="incomplete", usage=usage, output=[msg_item]), + ), + ] + result = self._drain(events) + assert result.finish_reason == "length" + assert "Sources:" in result.content + assert "[Cite](https://cite.test)" in result.content + def test_transient_error_event_is_retryable(self) -> None: from turnstone.core.providers._openai_responses import ResponsesStreamFailedError diff --git a/turnstone/core/providers/__init__.py b/turnstone/core/providers/__init__.py index 14b77fe3..1a270104 100644 --- a/turnstone/core/providers/__init__.py +++ b/turnstone/core/providers/__init__.py @@ -10,6 +10,7 @@ from turnstone.core.providers._openai_chat import OpenAIChatCompletionsProvider from turnstone.core.providers._openai_responses import OpenAIResponsesProvider from turnstone.core.providers._protocol import ( CompletionResult, + IncompleteStreamError, LLMProvider, ModelCapabilities, StreamChunk, @@ -21,6 +22,7 @@ from turnstone.core.providers._xai import XAI_DEFAULT_BASE_URL, XAIProvider __all__ = [ "CompletionResult", + "IncompleteStreamError", "LLMProvider", "ModelCapabilities", "OpenAIChatCompletionsProvider", diff --git a/turnstone/core/providers/_google.py b/turnstone/core/providers/_google.py index 3caedf62..f6fe7f5a 100644 --- a/turnstone/core/providers/_google.py +++ b/turnstone/core/providers/_google.py @@ -24,7 +24,10 @@ if TYPE_CHECKING: from collections.abc import Iterator from turnstone.core.lowering import legalize_tool_call_entry -from turnstone.core.providers._openai_chat import OpenAIChatCompletionsProvider +from turnstone.core.providers._openai_chat import ( + OpenAIChatCompletionsProvider, + ToolCallSlotter, +) from turnstone.core.providers._openai_common import sanitize_messages from turnstone.core.providers._protocol import ModelCapabilities, StreamChunk @@ -137,8 +140,17 @@ class GoogleProvider(OpenAIChatCompletionsProvider): delegates all chunk processing to the base class. The accumulated raw tool-call dicts are emitted as ``provider_blocks`` on the final chunk so the session stores them as ``_provider_content``. + + The tap keys its raw dicts through its OWN :class:`ToolCallSlotter` + over the same delta sequence the base iterator slots, so the raw + fidelity lane and the normalized ``tool_calls`` mirror assign call + identity identically — a desync (one fused raw dict vs two mirror + calls on an index-degenerate stream) would fail + ``_prepare_messages``' length gate and silently drop + ``thought_signature`` from the replay. """ raw_tool_calls: dict[int, dict[str, Any]] = {} + slotter = ToolCallSlotter() def _tap(raw_stream: Any) -> Any: """Pass-through iterator that captures tool-call extras.""" @@ -147,21 +159,27 @@ class GoogleProvider(OpenAIChatCompletionsProvider): delta = chunk.choices[0].delta if delta.tool_calls: for tc_delta in delta.tool_calls: - idx = tc_delta.index - if idx not in raw_tool_calls: - raw_tool_calls[idx] = { + fn = tc_delta.function + slot = slotter.slot_for( + tc_delta.index, + tc_delta.id or "", + has_name=bool(fn and fn.name), + has_args=bool(fn and fn.arguments), + ) + if slot not in raw_tool_calls: + raw_tool_calls[slot] = { "id": "", "type": "function", "function": {"name": "", "arguments": ""}, } - raw_tc = raw_tool_calls[idx] + raw_tc = raw_tool_calls[slot] if tc_delta.id: raw_tc["id"] = tc_delta.id - if tc_delta.function: - if tc_delta.function.name: - raw_tc["function"]["name"] = tc_delta.function.name - if tc_delta.function.arguments: - raw_tc["function"]["arguments"] += tc_delta.function.arguments + if fn: + if fn.name: + raw_tc["function"]["name"] = fn.name + if fn.arguments: + raw_tc["function"]["arguments"] += fn.arguments # Capture provider-specific extras (e.g. thought_signature) extras = getattr(tc_delta, "__pydantic_extra__", None) if extras: diff --git a/turnstone/core/providers/_openai_chat.py b/turnstone/core/providers/_openai_chat.py index 6afd3bdb..902498a1 100644 --- a/turnstone/core/providers/_openai_chat.py +++ b/turnstone/core/providers/_openai_chat.py @@ -55,6 +55,47 @@ def _reasoning_text(obj: Any) -> str: log = structlog.get_logger(__name__) +class ToolCallSlotter: + """Wire-index → logical-slot assignment for streamed tool-call deltas. + + Shared by the base iterator and ``GoogleProvider``'s raw fidelity tap + (both observe the same delta sequence, so their assignments agree and + the normalized mirror can never desync from the raw ``provider_blocks`` + lane). Index-degenerate compat servers (historical vLLM/llama.cpp + builds) emit every parallel call at index 0 — a delta opens a NEW slot + when its id contradicts the slot's id, or when it announces a name for + a slot that already accumulated arguments (the id-less whole-delta + shape: name+arguments per call, so a second announcement after + arguments is a second call). Id-less argument fragments keep + following their index's current slot. + """ + + def __init__(self) -> None: + self._slot_for_index: dict[int, int] = {} + self._slot_ids: dict[int, str] = {} + self._slot_has_args: set[int] = set() + self._next_slot = 0 + + def slot_for(self, wire_index: int, tc_id: str, *, has_name: bool, has_args: bool) -> int: + slot = self._slot_for_index.get(wire_index) + id_conflict = ( + slot is not None + and tc_id + and self._slot_ids.get(slot, "") + and self._slot_ids[slot] != tc_id + ) + reannounce = slot is not None and has_name and slot in self._slot_has_args + if slot is None or id_conflict or reannounce: + slot = self._next_slot + self._next_slot += 1 + self._slot_for_index[wire_index] = slot + if tc_id: + self._slot_ids[slot] = tc_id + if has_args: + self._slot_has_args.add(slot) + return slot + + class OpenAIChatCompletionsProvider: """Provider for local OpenAI-compatible servers (vLLM, llama.cpp, SGLang). @@ -202,16 +243,10 @@ class OpenAIChatCompletionsProvider: tool_call_count = 0 last_finish_reason: str | None = None completion_tokens: int | None = None - # Remap wire indexes onto logical slots: index-degenerate compat - # servers (historical vLLM/llama.cpp builds) emit every parallel - # tool call at index 0 — a delta whose id contradicts its index's - # current call opens a new slot, mirroring the Anthropic iterator's - # per-block index assignment. Id-less argument fragments keep - # following their index's current slot, so downstream accumulators - # (drain_stream, the chat loop) can key by index safely. - slot_for_index: dict[int, int] = {} - slot_ids: dict[int, str] = {} - next_slot = 0 + # Remap wire indexes onto logical slots (see ToolCallSlotter) so + # downstream accumulators (drain_stream, the chat loop) can key by + # index safely even on index-degenerate compat servers. + slotter = ToolCallSlotter() for chunk in stream: sc = StreamChunk() @@ -246,25 +281,22 @@ class OpenAIChatCompletionsProvider: # Tool calls if delta.tool_calls: for tc_delta in delta.tool_calls: - wire_index = tc_delta.index tc_id = tc_delta.id or "" - slot = slot_for_index.get(wire_index) - if slot is None or ( - tc_id and slot_ids.get(slot, "") and slot_ids[slot] != tc_id - ): - slot = next_slot - next_slot += 1 - slot_for_index[wire_index] = slot - if tc_id: - slot_ids[slot] = tc_id + fn = tc_delta.function + slot = slotter.slot_for( + tc_delta.index, + tc_id, + has_name=bool(fn and fn.name), + has_args=bool(fn and fn.arguments), + ) tcd = ToolCallDelta(index=slot) if tc_id: tcd.id = tc_id - if tc_delta.function: - if tc_delta.function.name: - tcd.name = tc_delta.function.name - if tc_delta.function.arguments: - tcd.arguments_delta = tc_delta.function.arguments + if fn: + if fn.name: + tcd.name = fn.name + if fn.arguments: + tcd.arguments_delta = fn.arguments sc.tool_call_deltas.append(tcd) tool_call_count += 1 diff --git a/turnstone/core/providers/_openai_responses.py b/turnstone/core/providers/_openai_responses.py index 9c023b0c..850c8bba 100644 --- a/turnstone/core/providers/_openai_responses.py +++ b/turnstone/core/providers/_openai_responses.py @@ -48,6 +48,16 @@ log = structlog.get_logger(__name__) _TRANSIENT_FAILURE_CODES = frozenset({"server_error", "rate_limit_exceeded"}) +def _raise_responses_failure(error_code: str, error_msg: str) -> None: + """One classification ladder for BOTH in-band failure shapes (`error` + events and ``response.failed``) — a one-sided edit would make the same + API failure retryable through one event type and fatal through the + other.""" + if error_code in _TRANSIENT_FAILURE_CODES: + raise ResponsesStreamFailedError(f"Responses API error ({error_code}): {error_msg}") + raise RuntimeError(f"Responses API error ({error_code or 'unknown'}): {error_msg}") + + class ResponsesStreamFailedError(RuntimeError): """A TRANSIENT in-band ``response.failed`` terminal event. @@ -671,38 +681,49 @@ class OpenAIResponsesProvider: if not response: # A terminal event without its payload (lax compat # server) is still a terminal signal — emit the finish - # reason implied by the event type so the drain's - # complete-or-error gate sees a completed stream, just - # without usage/blocks. - last_finish = "stop" if event_type == "response.completed" else "length" - yield StreamChunk(finish_reason=last_finish) - continue - if response: - status = getattr(response, "status", "completed") - last_finish = "stop" if status == "completed" else "length" - usage = extract_usage(getattr(response, "usage", None)) - if usage: - completion_tokens = usage.completion_tokens - # Prefer the terminal response's own output items over - # the incrementally collected ones: an item still being - # generated at truncation never receives its - # ``output_item.done`` event, and storing a reasoning - # item without its required following item makes the - # next turn's replay a 400. - final_items = [ - item.model_dump() - for item in (getattr(response, "output", None) or []) - if hasattr(item, "model_dump") - ] - if final_items: - provider_blocks = final_items + # reason implied by the event type, keeping the blocks + # already collected from output_item.done events (they + # came from the stream, not the missing payload); only + # usage is genuinely unavailable. sc = StreamChunk( - finish_reason=last_finish, - usage=usage, + finish_reason="stop" if event_type == "response.completed" else "length" ) if provider_blocks: sc.provider_blocks = provider_blocks yield sc + continue + status = getattr(response, "status", "completed") + last_finish = "stop" if status == "completed" else "length" + usage = extract_usage(getattr(response, "usage", None)) + if usage: + completion_tokens = usage.completion_tokens + # Prefer the terminal response's own output items over the + # incrementally collected ones: an item still being + # generated at truncation never receives its + # ``output_item.done`` event, and storing a reasoning item + # without its required following item makes the next + # turn's replay a 400. Its annotations were likewise + # never collected — walk them here (format_citations + # dedupes by URL, so re-seeing .done'd items is harmless). + out_items = getattr(response, "output", None) or [] + final_items = [ + item.model_dump() for item in out_items if hasattr(item, "model_dump") + ] + if final_items: + provider_blocks = final_items + for item in out_items: + if getattr(item, "type", "") == "message": + for content_part in getattr(item, "content", []) or []: + part_anns = getattr(content_part, "annotations", None) + if part_anns: + annotations.extend(part_anns) + sc = StreamChunk( + finish_reason=last_finish, + usage=usage, + ) + if provider_blocks: + sc.provider_blocks = provider_blocks + yield sc continue # -- in-band error event (ResponseErrorEvent) -- @@ -711,29 +732,19 @@ class OpenAIResponsesProvider: # stream exhausts finish-less and the real API message is lost # behind a misleading IncompleteStreamError. if event_type == "error": - error_msg = getattr(event, "message", "Unknown error") or "Unknown error" - error_code = getattr(event, "code", "") or "" - if error_code in _TRANSIENT_FAILURE_CODES: - raise ResponsesStreamFailedError( - f"Responses API error ({error_code}): {error_msg}" - ) - raise RuntimeError(f"Responses API error ({error_code or 'unknown'}): {error_msg}") + _raise_responses_failure( + getattr(event, "code", "") or "", + getattr(event, "message", "Unknown error") or "Unknown error", + ) # -- error -- if event_type == "response.failed": 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" - error_code = getattr(error, "code", "") if error else "" - # Only transient in-band failures are worth the caller's - # backoff ladder; a deterministic rejection (invalid prompt, - # image fetch, policy) re-fails identically on every retry - # and must surface immediately, as it did pre-#831. - if error_code in _TRANSIENT_FAILURE_CODES: - raise ResponsesStreamFailedError( - f"Responses API error ({error_code}): {error_msg}" - ) - raise RuntimeError(f"Responses API error ({error_code or 'unknown'}): {error_msg}") + _raise_responses_failure( + (getattr(error, "code", "") if error else "") or "", + (getattr(error, "message", "") if error else "") or "Unknown error", + ) log.debug( "openai.responses.response", diff --git a/turnstone/core/providers/_protocol.py b/turnstone/core/providers/_protocol.py index 76a191f8..1213796f 100644 --- a/turnstone/core/providers/_protocol.py +++ b/turnstone/core/providers/_protocol.py @@ -142,8 +142,15 @@ def drain_stream(chunks: Iterator[StreamChunk]) -> CompletionResult: Raises whatever the underlying stream raises — retry/deadline/fallback policy stays with the caller, exactly as with the old non-streaming - transport. + transport — EXCEPT httpx transport failures: streaming moves the body + read out of the SDK's ``APIConnectionError``-wrapped request into raw + iteration, so a mid-body connection drop or read timeout surfaces as a + bare ``httpx.TransportError`` no retry predicate recognizes. Those + are re-raised (chained) as :class:`IncompleteStreamError`, restoring + the wire-blip retryability the non-streaming transport had. """ + import httpx # noqa: PLC0415 — heavyweight; deferred off the type-module import path + content_parts: list[str] = [] reasoning_parts: list[str] = [] trailing_info_parts: list[str] = [] @@ -152,7 +159,16 @@ def drain_stream(chunks: Iterator[StreamChunk]) -> CompletionResult: finish_reason: str | None = None provider_blocks: list[dict[str, Any]] = [] - for sc in chunks: + iterator = iter(chunks) + while True: + try: + sc = next(iterator) + except StopIteration: + break + except httpx.TransportError as exc: + raise IncompleteStreamError( + f"stream transport failed mid-response ({type(exc).__name__}: {exc})" + ) from exc if sc.content_delta: content_parts.append(sc.content_delta) if sc.reasoning_delta: