From 3ffa8b9057b54f705f7607ba5bb9af03e932851b Mon Sep 17 00:00:00 2001 From: Patrick Buckley Date: Mon, 13 Jul 2026 11:22:56 -0700 Subject: [PATCH] =?UTF-8?q?fix(providers):=20review=20round=202=20?= =?UTF-8?q?=E2=80=94=20complete-or-error=20drain,=20code-gated=20retries,?= =?UTF-8?q?=20truncation-safe=20blocks?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Correctness (3 confirmed + 2 plausible, all fixed): - drain_stream now raises typed, retryable IncompleteStreamError when a stream exhausts without any finish reason — every adapter emits one on a healthy stream, so its absence means the generation died mid-response behind a cleanly-closing proxy. This restores the retired transport's complete-or-error contract (a half-generated compaction summary was previously returned as finish=stop and stored, silently replacing real history) and DELETES round 1's suffix-info fold: with no finish-less success path there is nothing to classify, so a trailing status ping can never be stored as content either. - Index-degenerate parallel tool calls get distinct slots: a delta whose id differs from its slot's opens a new call (id-less fragments still follow their index's current call), so historical compat servers that emit every parallel call at index 0 no longer fuse distinct calls into concatenated garbage arguments. Result order stays index-sorted (stable) like the retired array parse. - response.failed retryability is code-gated: only transient codes (server_error, rate_limit_exceeded) raise the retryable typed error; deterministic rejections (invalid prompt, image fetch, policy) raise plain RuntimeError and stop retry loops on attempt zero instead of running the full backoff ladder against a doomed request. - Terminal Responses events rebuild provider_blocks from response.output when present: the item being generated at max_output_tokens truncation never receives output_item.done, and storing a reasoning item without its required following item made the next turn's replay a 400. - merge_usage's base case uses dataclasses.replace so a future UsageInfo field can't be silently zeroed on drained lanes. Cleanup: run_abortable_with_deadline bundles the three-point abort wiring (ref + cancel_ref + on_abandon) so it cannot be half-wired — both judges converted; scripted_chat_client hoists the 14 chat-lane fake_create closures (call scripts + .calls recording replace per-test counter cells); fake_chat_stream gains reasoning=, collapsing the reasoning-capture suite's hand-rolled chunk shape; FakeAnthropicBlock hoists the duplicated _Block test class; the class and judge PlantUML diagrams drop the retired create_completion flow. Also converts test_model_registry's agent-model fakes, which returned legacy response objects that iterated as EMPTY streams — they only passed through the old drain's silent finish=stop default, exactly the hazard the new gate exists to catch. --- docs/diagrams/03-core-engine-classes.puml | 3 +- docs/diagrams/22-judge-architecture.puml | 4 +- tests/_session_helpers.py | 47 ++- tests/test_drain_stream.py | 101 +++-- tests/test_model_registry.py | 15 +- tests/test_providers.py | 94 +++-- tests/test_session.py | 374 +++++++----------- turnstone/core/deadline.py | 31 ++ turnstone/core/judge.py | 35 +- turnstone/core/output_guard_judge.py | 17 +- turnstone/core/providers/_anthropic.py | 2 + turnstone/core/providers/_openai_common.py | 3 + turnstone/core/providers/_openai_responses.py | 48 ++- turnstone/core/providers/_protocol.py | 119 +++--- 14 files changed, 494 insertions(+), 399 deletions(-) diff --git a/docs/diagrams/03-core-engine-classes.puml b/docs/diagrams/03-core-engine-classes.puml index 74aae8a6..d7cdee64 100644 --- a/docs/diagrams/03-core-engine-classes.puml +++ b/docs/diagrams/03-core-engine-classes.puml @@ -66,8 +66,7 @@ class "NullUI" as NullUI { interface "LLMProvider" as LLMProvider <> { + provider_name: str {property} + get_capabilities(model) → ModelCapabilities - + create_streaming(client, model, messages, ..., replay_reasoning_to_model) → Iterator[StreamChunk] - + create_completion(client, model, messages, ..., replay_reasoning_to_model) → CompletionResult + + create_streaming(client, model, messages, ..., cancel_ref, replay_reasoning_to_model) → Iterator[StreamChunk] + convert_tools(tools) → list[dict] + extract_reasoning_text(provider_blocks) → str + retryable_error_names: frozenset[str] {property} diff --git a/docs/diagrams/22-judge-architecture.puml b/docs/diagrams/22-judge-architecture.puml index e5c2d3e4..5a2a9997 100644 --- a/docs/diagrams/22-judge-architecture.puml +++ b/docs/diagrams/22-judge-architecture.puml @@ -84,8 +84,8 @@ end note loop up to 3 turns (timeout budget) - Judge -> LLM : create_completion(\nmodel, judge_messages,\ntools=[read_file, list_directory]) - LLM --> Judge : CompletionResult + Judge -> LLM : model_turn(lane, judge_turns,\ntools=[read_file, list_directory])\nvia drained create_streaming + LLM --> Judge : ModelTurnResult alt tool_calls present (turn < 3) Judge -> Judge : _exec_read_only_tool() diff --git a/tests/_session_helpers.py b/tests/_session_helpers.py index 86d03e01..69a9f0e6 100644 --- a/tests/_session_helpers.py +++ b/tests/_session_helpers.py @@ -80,6 +80,7 @@ def fake_chat_stream( prompt_tokens: int = 10, completion_tokens: int = 5, reasoning_content: str | None = None, + reasoning: str | None = None, ) -> list[Any]: """Fake OpenAI Chat Completions SSE chunks for driving the REAL ``OpenAIChatCompletionsProvider`` through a fake SDK client:: @@ -102,20 +103,25 @@ def fake_chat_stream( content_val: str | None = None, tcs: list[Any] | None = None, rc: str | None = None, + rsn: str | None = None, ) -> SimpleNamespace: return SimpleNamespace( content=content_val, tool_calls=tcs, - reasoning=None, + reasoning=rsn, reasoning_content=rc, annotations=None, ) chunks: list[Any] = [] - if reasoning_content is not None: + if reasoning_content is not None or reasoning is not None: chunks.append( SimpleNamespace( - choices=[SimpleNamespace(finish_reason=None, delta=_delta(rc=reasoning_content))], + choices=[ + SimpleNamespace( + finish_reason=None, delta=_delta(rc=reasoning_content, rsn=reasoning) + ) + ], usage=None, ) ) @@ -164,6 +170,41 @@ def fake_chat_stream( return chunks +def scripted_chat_client(*scripts: Any) -> Any: + """A fake ``client.chat.completions.create`` that follows a script. + + Call N returns the stream described by ``scripts[N]``; the last script + repeats for any further calls. Each script is a dict of + :func:`fake_chat_stream` kwargs or a pre-built chunk list. The + returned callable records every call's kwargs on ``.calls`` — read + ``len(fn.calls)`` where a test previously kept its own counter cell, + and ``fn.calls[i]["messages"]`` where it captured request bodies. + """ + + def _create(**kwargs: Any) -> Any: + _create.calls.append(kwargs) # type: ignore[attr-defined] + i = min(len(_create.calls) - 1, len(scripts) - 1) # type: ignore[attr-defined] + script = scripts[i] + return fake_chat_stream(**script) if isinstance(script, dict) else script + + _create.calls = [] # type: ignore[attr-defined] + return _create + + +class FakeAnthropicBlock: + """A full-content Anthropic content-block fake for + :func:`fake_anthropic_stream` — plain attributes plus the + ``model_dump()`` the provider's block capture reads.""" + + def __init__(self, **fields: Any) -> None: + self._fields = fields + for key, value in fields.items(): + setattr(self, key, value) + + def model_dump(self, **_kw: Any) -> dict[str, Any]: + return dict(self._fields) + + def fake_anthropic_stream( blocks: list[Any], *, diff --git a/tests/test_drain_stream.py b/tests/test_drain_stream.py index 79cf43c9..63d4c91b 100644 --- a/tests/test_drain_stream.py +++ b/tests/test_drain_stream.py @@ -17,6 +17,8 @@ from turnstone.core.providers import ( UsageInfo, drain_stream, ) +from turnstone.core.providers._openai_common import RETRYABLE_ERROR_NAMES +from turnstone.core.providers._protocol import IncompleteStreamError class TestContentAndReasoning: @@ -40,20 +42,25 @@ class TestContentAndReasoning: StreamChunk(reasoning_delta="think "), StreamChunk(reasoning_delta="hard"), StreamChunk(content_delta="answer"), + StreamChunk(finish_reason="stop"), ] ) ) assert result.reasoning == "think hard" assert result.content == "answer" - def test_empty_stream_yields_defaults(self): - result = drain_stream(iter([])) - assert result.content == "" - assert result.reasoning == "" - assert result.tool_calls is None - assert result.finish_reason == "stop" - assert result.usage is None - assert result.provider_blocks == [] + def test_stream_without_finish_reason_raises_incomplete(self): + # Complete-or-error: every adapter emits a finish reason on a + # healthy stream, so its absence means the generation died + # mid-response — partial text must never be stored as a complete + # result (compaction summary, title). Typed and retryable. + assert "IncompleteStreamError" in RETRYABLE_ERROR_NAMES + with pytest.raises(IncompleteStreamError): + drain_stream(iter([StreamChunk(content_delta="half a summar")])) + + def test_empty_stream_raises_incomplete(self): + with pytest.raises(IncompleteStreamError): + drain_stream(iter([])) class TestToolCallAssembly: @@ -70,6 +77,7 @@ class TestToolCallAssembly: StreamChunk( tool_call_deltas=[ToolCallDelta(index=0, arguments_delta='"x.py"}')] ), + StreamChunk(finish_reason="tool_calls"), ] ) ) @@ -95,6 +103,7 @@ class TestToolCallAssembly: ToolCallDelta(index=1, arguments_delta='{"k": 1}'), ] ), + StreamChunk(finish_reason="tool_calls"), ] ) ) @@ -105,10 +114,44 @@ class TestToolCallAssembly: # Google compat can stream blank tool ids — the drain must hand them # through untouched so model_turn's pairwise blank-id repair sees them. result = drain_stream( - iter([StreamChunk(tool_call_deltas=[ToolCallDelta(index=0, name="f")])]) + iter( + [ + StreamChunk(tool_call_deltas=[ToolCallDelta(index=0, name="f")]), + StreamChunk(finish_reason="tool_calls"), + ] + ) ) assert result.tool_calls[0]["id"] == "" + def test_index_degenerate_parallel_calls_get_distinct_slots(self): + # Historical compat servers (older vLLM, some llama.cpp builds) + # stream every parallel call at index 0 as whole deltas. A delta + # whose id differs from its slot's opens a NEW call — without this, + # distinct calls fuse into one entry with concatenated garbage + # arguments. Id-less fragments keep following their index's + # current call. + result = drain_stream( + iter( + [ + StreamChunk( + tool_call_deltas=[ + ToolCallDelta(index=0, id="a", name="read", arguments_delta='{"p": 1}') + ] + ), + StreamChunk( + tool_call_deltas=[ + ToolCallDelta(index=0, id="b", name="write", arguments_delta='{"p": ') + ] + ), + StreamChunk(tool_call_deltas=[ToolCallDelta(index=0, arguments_delta="2}")]), + StreamChunk(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}' + class TestUsageMerge: def test_anthropic_split_emission_max_merges(self): @@ -144,6 +187,7 @@ class TestUsageMerge: iter( [ StreamChunk(content_delta="x"), + StreamChunk(finish_reason="stop"), StreamChunk( usage=UsageInfo(prompt_tokens=10, completion_tokens=5, total_tokens=15) ), @@ -233,33 +277,20 @@ 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"), - ] + def test_finishless_stream_raises_even_with_trailing_info(self): + # A stream that dies after a status ping must NOT return the ping + # as content (nor the partial body as a clean result) — the + # complete-or-error gate turns the whole stream into a retryable + # error instead of guessing which trailing info was a citation. + with pytest.raises(IncompleteStreamError): + drain_stream( + iter( + [ + StreamChunk(content_delta="body"), + StreamChunk(info_delta="[Searching: kubernetes CVEs]"), + ] + ) ) - ) - 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: diff --git a/tests/test_model_registry.py b/tests/test_model_registry.py index 04db8cd9..20ace732 100644 --- a/tests/test_model_registry.py +++ b/tests/test_model_registry.py @@ -8,6 +8,7 @@ from unittest.mock import MagicMock, patch import pytest +from tests._session_helpers import fake_chat_stream from turnstone.core.model_registry import ( ModelConfig, ModelRegistry, @@ -1301,16 +1302,11 @@ class TestSessionAgentModel: # Mock the API to capture what model was used captured_model = None - mock_response = MagicMock() - mock_response.choices = [MagicMock()] - mock_response.choices[0].message.content = "done" - mock_response.choices[0].message.tool_calls = None - mock_response.choices[0].finish_reason = "stop" def fake_create(**kwargs: Any) -> Any: nonlocal captured_model captured_model = kwargs.get("model") - return mock_response + return fake_chat_stream(content="done") # Get the agent client from the registry and patch it agent_client = reg.get_client("agent") @@ -1327,15 +1323,10 @@ class TestSessionAgentModel: def _capture_on(client: Any) -> dict[str, Any]: """Patch *client* (registry-resolved or session.client) to capture kwargs.""" captured: dict[str, Any] = {} - mock_response = MagicMock() - mock_response.choices = [MagicMock()] - mock_response.choices[0].message.content = "done" - mock_response.choices[0].message.tool_calls = None - mock_response.choices[0].finish_reason = "stop" def fake_create(**kwargs: Any) -> Any: captured.update(kwargs) - return mock_response + return fake_chat_stream(content="done") client.chat.completions.create = fake_create return captured diff --git a/tests/test_providers.py b/tests/test_providers.py index bb7134fc..854f5146 100644 --- a/tests/test_providers.py +++ b/tests/test_providers.py @@ -4415,31 +4415,9 @@ class TestOpenAIChatReasoningCapture: @staticmethod def _client(*, reasoning: Any = None, reasoning_content: Any = None) -> MagicMock: - delta = SimpleNamespace( - content="ok", - tool_calls=None, - annotations=None, - reasoning=reasoning, - reasoning_content=reasoning_content, + chunks = fake_chat_stream( + content="ok", reasoning=reasoning, reasoning_content=reasoning_content ) - chunks = [ - SimpleNamespace(choices=[SimpleNamespace(finish_reason=None, delta=delta)], usage=None), - SimpleNamespace( - choices=[ - SimpleNamespace( - finish_reason="stop", - delta=SimpleNamespace( - content=None, - tool_calls=None, - annotations=None, - reasoning=None, - reasoning_content=None, - ), - ) - ], - usage=None, - ), - ] client = MagicMock() client.chat.completions.create.return_value = chunks return client @@ -5252,21 +5230,77 @@ class TestResponsesDrainedStream: 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). + def test_truncation_rebuilds_blocks_from_terminal_response_output(self) -> None: + # The item being generated at max_output_tokens truncation never + # receives output_item.done; only the terminal response.output has + # it. Storing the .done-collected list alone would keep a + # reasoning item without its required following item — the next + # turn's replay 400s. The terminal event's own output wins. + def _item(d: dict) -> SimpleNamespace: + item = SimpleNamespace(**{k: v for k, v in d.items() if k != "model_dump"}) + item.model_dump = lambda d=d, **_kw: d # type: ignore[method-assign] + return item + + reasoning_item = _item({"type": "reasoning", "id": "rs_1", "summary": []}) + partial_msg = _item({"type": "message", "content": [], "status": "incomplete"}) + usage = SimpleNamespace( + input_tokens=10, + output_tokens=5, + total_tokens=15, + input_tokens_details=SimpleNamespace(cached_tokens=0), + ) + events = [ + # Only the reasoning item completed before truncation. + SimpleNamespace(type="response.output_item.done", item=reasoning_item), + SimpleNamespace( + type="response.incomplete", + response=SimpleNamespace( + status="incomplete", + usage=usage, + output=[reasoning_item, partial_msg], + ), + ), + ] + result = self._drain(events) + assert result.finish_reason == "length" + assert [b["type"] for b in result.provider_blocks] == ["reasoning", "message"] + + def test_transient_failed_event_raises_typed_retryable_error(self) -> None: + # A TRANSIENT in-band response.failed (server_error / rate limit) + # raises the typed error the provider advertises as retryable — + # retry loops re-run it like the wire errors it stands in for. from turnstone.core.providers._openai_responses import ResponsesStreamFailedError events = [ SimpleNamespace( type="response.failed", response=SimpleNamespace( - status="failed", error=SimpleNamespace(message="model overloaded") + status="failed", + error=SimpleNamespace(message="model overloaded", code="server_error"), ), ) ] with pytest.raises(ResponsesStreamFailedError, match="model overloaded"): self._drain(events) assert "ResponsesStreamFailedError" in self.provider.retryable_error_names + + def test_deterministic_failed_event_is_not_retryable(self) -> None: + # Deterministic in-band rejections (invalid prompt, image fetch, + # policy) re-fail identically on every attempt — they surface as + # plain RuntimeError so retry loops stop on attempt zero instead + # of running the whole backoff ladder against a doomed request. + from turnstone.core.providers._openai_responses import ResponsesStreamFailedError + + events = [ + SimpleNamespace( + type="response.failed", + response=SimpleNamespace( + status="failed", + error=SimpleNamespace(message="prompt was rejected", code="invalid_prompt"), + ), + ) + ] + with pytest.raises(RuntimeError, match="invalid_prompt") as excinfo: + self._drain(events) + assert not isinstance(excinfo.value, ResponsesStreamFailedError) + assert type(excinfo.value).__name__ not in self.provider.retryable_error_names diff --git a/tests/test_session.py b/tests/test_session.py index 10a5221d..bc0c81e5 100644 --- a/tests/test_session.py +++ b/tests/test_session.py @@ -12,10 +12,11 @@ from unittest.mock import MagicMock, patch import pytest from tests._session_helpers import ( + FakeAnthropicBlock, as_stream, fake_anthropic_stream, - fake_chat_stream, mock_completion_result, + scripted_chat_client, ) from turnstone.core.session import _IMAGE_EXTENSIONS, _IMAGE_SIZE_CAP, ChatSession from turnstone.core.trajectory import ( @@ -2059,26 +2060,20 @@ class TestAgentOutputGuard: session, "_evaluate_output", wraps=lambda cid, o, fn, **_kw: (o, None) ) as mock_eval: # Simulate _run_agent getting a tool call response then a text response - call_count = [0] - - def fake_create(**kwargs): - call_count[0] += 1 - if call_count[0] == 1: - # First call: model returns a tool call - return fake_chat_stream( - tool_calls=[ - { - "id": "call_1", - "name": "read_file", - "arguments": '{"path": "/tmp/test"}', - } - ], - finish_reason="tool_calls", - ) - # Second call: model returns text (done) - return fake_chat_stream(content="Done") - - session.client.chat.completions.create = fake_create + # Script: a tool-call turn, then text (done). + session.client.chat.completions.create = scripted_chat_client( + { + "tool_calls": [ + { + "id": "call_1", + "name": "read_file", + "arguments": '{"path": "/tmp/test"}', + } + ], + "finish_reason": "tool_calls", + }, + {"content": "Done"}, + ) # Mock tool preparation to return a simple output def fake_prepare(tc_dict, **kwargs): @@ -2117,24 +2112,19 @@ class TestAgentOutputGuard: session._provider = OpenAIChatCompletionsProvider() with patch.object(session, "_evaluate_output") as mock_eval: - call_count = [0] - - def fake_create(**kwargs): - call_count[0] += 1 - if call_count[0] == 1: - return fake_chat_stream( - tool_calls=[ - { - "id": "call_1", - "name": "read_file", - "arguments": '{"path": "/tmp/test"}', - } - ], - finish_reason="tool_calls", - ) - return fake_chat_stream(content="Done") - - session.client.chat.completions.create = fake_create + session.client.chat.completions.create = scripted_chat_client( + { + "tool_calls": [ + { + "id": "call_1", + "name": "read_file", + "arguments": '{"path": "/tmp/test"}', + } + ], + "finish_reason": "tool_calls", + }, + {"content": "Done"}, + ) def fake_prepare(tc_dict, **kwargs): return { @@ -2172,11 +2162,7 @@ class TestAgentOutputGuard: with patch.object( session, "_evaluate_output", wraps=lambda cid, o, fn, **_kw: (o, None) ) as mock_eval: - - def fake_create(**_kwargs): - return fake_chat_stream(content=synth) - - session.client.chat.completions.create = fake_create + session.client.chat.completions.create = scripted_chat_client({"content": synth}) result = session._run_agent( [Turn.user("test")], @@ -2204,11 +2190,9 @@ class TestAgentOutputGuard: with patch.object( session, "_evaluate_output", wraps=lambda cid, o, fn: (o, None) ) as mock_eval: - - def fake_create(**_kwargs): - return fake_chat_stream(content=partial, finish_reason="length") - - session.client.chat.completions.create = fake_create + session.client.chat.completions.create = scripted_chat_client( + {"content": partial, "finish_reason": "length"} + ) result = session._run_agent( [Turn.user("test")], tools=[{"type": "function", "function": {"name": "read_file"}}], @@ -2323,28 +2307,20 @@ class TestAgentOutputGuard: session.agent_max_turns = 1 # one tool turn, then forced synthesis forced = "Forced synthesis after hitting the tool-turn ceiling." - call_count = [0] with patch.object( session, "_evaluate_output", wraps=lambda cid, o, fn, **_kw: (o, None) ) as mock_eval: - - def fake_create(**_kwargs): - call_count[0] += 1 - if call_count[0] == 1: - # First call: tool call, eats the turn budget. - return fake_chat_stream( - tool_calls=[ - { - "id": "call_1", - "name": "read_file", - "arguments": '{"path": "/tmp/x"}', - } - ], - finish_reason="tool_calls", - ) - # Forced synthesis turn. - return fake_chat_stream(content=forced) + # Script: a tool call eats the turn budget, then forced synthesis. + fake_create = scripted_chat_client( + { + "tool_calls": [ + {"id": "call_1", "name": "read_file", "arguments": '{"path": "/tmp/x"}'} + ], + "finish_reason": "tool_calls", + }, + {"content": forced}, + ) session.client.chat.completions.create = fake_create @@ -2383,20 +2359,15 @@ class TestAgentChildRegistration: session._provider = OpenAIChatCompletionsProvider() session.ui.note_agent_child = MagicMock() - call_count = [0] - - def fake_create(**_kwargs): - call_count[0] += 1 - if call_count[0] == 1: - return fake_chat_stream( - tool_calls=[ - {"id": "call_1", "name": "read_file", "arguments": '{"path": "/tmp/x"}'} - ], - finish_reason="tool_calls", - ) - return fake_chat_stream(content="done") - - session.client.chat.completions.create = fake_create + session.client.chat.completions.create = scripted_chat_client( + { + "tool_calls": [ + {"id": "call_1", "name": "read_file", "arguments": '{"path": "/tmp/x"}'} + ], + "finish_reason": "tool_calls", + }, + {"content": "done"}, + ) def fake_prepare(tc_dict, **_kwargs): return { @@ -2433,28 +2404,22 @@ class TestAgentChildRegistration: session._provider = OpenAIChatCompletionsProvider() session.ui.note_agent_child = MagicMock() - call_count = [0] + def _reused_call(path: str) -> dict: + # id reused verbatim across turns — the local-server shape. + return { + "tool_calls": [{"id": "call_0", "name": "read_file", "arguments": path}], + "finish_reason": "tool_calls", + } - def fake_create(**_kwargs): - call_count[0] += 1 - if call_count[0] <= 2: - return fake_chat_stream( - tool_calls=[ - { - # reused verbatim across turns - "id": "call_0", - "name": "read_file", - "arguments": f'{{"path": "/tmp/f{call_count[0]}"}}', - } - ], - finish_reason="tool_calls", - ) - return fake_chat_stream(content="done") - - session.client.chat.completions.create = fake_create + client_fn = scripted_chat_client( + _reused_call('{"path": "/tmp/f1"}'), + _reused_call('{"path": "/tmp/f2"}'), + {"content": "done"}, + ) + session.client.chat.completions.create = client_fn def fake_prepare(tc_dict, **_kwargs): - n = call_count[0] + n = len(client_fn.calls) return { "call_id": tc_dict["id"], "func_name": "read_file", @@ -2488,22 +2453,16 @@ class TestAgentChildRegistration: @staticmethod def _reusing_provider(session, tool_turns: int = 1): """Fake create() reissuing id "call_0" for ``tool_turns`` turns, then - stopping — the local-server id-reuse shape. Returns the counter.""" - call_count = [0] - - def fake_create(**_kwargs): - call_count[0] += 1 - if call_count[0] <= tool_turns: - return fake_chat_stream( - tool_calls=[ - {"id": "call_0", "name": "read_file", "arguments": '{"path": "/tmp/x"}'} - ], - finish_reason="tool_calls", - ) - return fake_chat_stream(content="done") - - session.client.chat.completions.create = fake_create - return call_count + stopping — the local-server id-reuse shape.""" + reused = { + "tool_calls": [ + {"id": "call_0", "name": "read_file", "arguments": '{"path": "/tmp/x"}'} + ], + "finish_reason": "tool_calls", + } + session.client.chat.completions.create = scripted_chat_client( + *([reused] * tool_turns), {"content": "done"} + ) def test_parent_id_reuse_across_runs_mints_distinct_child_ids(self): # A local provider reuses "call_0" for the PARENT task_agent call too: @@ -2565,28 +2524,22 @@ class TestAgentChildRegistration: session._provider = OpenAIChatCompletionsProvider() session.ui.note_agent_child = MagicMock() - seen_messages: list[list[dict]] = [] - call_count = [0] - - def fake_create(**kwargs): - seen_messages.append(kwargs.get("messages") or []) - call_count[0] += 1 - if call_count[0] == 1: - return fake_chat_stream( - tool_calls=[ - { - "id": "call_0", - "name": "read_file", - # Malformed: unterminated JSON with a non-"length" - # finish reason — the sanitize pass's reason to exist. - "arguments": '{"path": "/tmp/x"', - } - ], - finish_reason="tool_calls", - ) - return fake_chat_stream(content="done") - - session.client.chat.completions.create = fake_create + client_fn = scripted_chat_client( + { + "tool_calls": [ + { + "id": "call_0", + "name": "read_file", + # Malformed: unterminated JSON with a non-"length" + # finish reason — the sanitize pass's reason to exist. + "arguments": '{"path": "/tmp/x"', + } + ], + "finish_reason": "tool_calls", + }, + {"content": "done"}, + ) + session.client.chat.completions.create = client_fn def fake_prepare(tc_dict, **_kwargs): return { @@ -2612,7 +2565,7 @@ class TestAgentChildRegistration: # shape the provider-native tool_use block also holds, so a native # replay and a rebuild agree); arguments are legalized to a JSON # object. - replay = seen_messages[1] + replay = client_fn.calls[1].get("messages") or [] wire_calls = [tc for m in replay if m.get("tool_calls") for tc in m["tool_calls"]] wire_results = [m for m in replay if m.get("role") == "tool"] assert wire_calls and wire_results @@ -2631,15 +2584,6 @@ class TestAgentChildRegistration: # thinking-enabled tool_use turn without its thinking block). from turnstone.core.providers._anthropic import AnthropicProvider - class _Block: - def __init__(self, **d): - self._d = d - for k, v in d.items(): - setattr(self, k, v) - - def model_dump(self, **_kw): - return dict(self._d) - session = _make_session() session._provider = AnthropicProvider() session.ui.note_agent_child = MagicMock() @@ -2653,17 +2597,17 @@ class TestAgentChildRegistration: if call_count[0] == 1: return fake_anthropic_stream( [ - _Block( + FakeAnthropicBlock( type="thinking", thinking="check the file first", signature="sig_v1" ), - _Block(type="text", text="reading"), - _Block( + FakeAnthropicBlock(type="text", text="reading"), + FakeAnthropicBlock( type="tool_use", id="toolu_01AB", name="read_file", input={"path": "x"} ), ], stop_reason="tool_use", ) - return fake_anthropic_stream([_Block(type="text", text="done")]) + return fake_anthropic_stream([FakeAnthropicBlock(type="text", text="done")]) session.client.messages.stream = fake_stream @@ -2724,15 +2668,6 @@ class TestAgentChildRegistration: # fallback (pinned in test_model_turn). from turnstone.core.providers._anthropic import AnthropicProvider - class _Block: - def __init__(self, **d): - self._d = d - for k, v in d.items(): - setattr(self, k, v) - - def model_dump(self, **_kw): - return dict(self._d) - session = _make_session() session._provider = AnthropicProvider() session.ui.note_agent_child = MagicMock() @@ -2746,13 +2681,15 @@ class TestAgentChildRegistration: if call_count[0] == 1: return fake_anthropic_stream( [ - _Block(type="thinking", thinking="hm", signature="sig_b"), + FakeAnthropicBlock(type="thinking", thinking="hm", signature="sig_b"), # Blank provider id — the back-fill case. - _Block(type="tool_use", id="", name="read_file", input={"path": "x"}), + FakeAnthropicBlock( + type="tool_use", id="", name="read_file", input={"path": "x"} + ), ], stop_reason="tool_use", ) - return fake_anthropic_stream([_Block(type="text", text="done")]) + return fake_anthropic_stream([FakeAnthropicBlock(type="text", text="done")]) session.client.messages.stream = fake_stream @@ -2833,28 +2770,19 @@ class TestAgentChildRegistration: ) session.ui.note_agent_child = MagicMock() - call_count = [0] - - def fake_create(**kwargs): - call_count[0] += 1 - if call_count[0] == 1: - return fake_chat_stream( - tool_calls=[ - { - # blank id — the back-fill case - "id": "", - "name": "read_file", - "arguments": '{"path": "x"}', - } - ], - finish_reason="tool_calls", - reasoning_content="work it out", - prompt_tokens=1, - completion_tokens=1, - ) - return fake_chat_stream(content="done", prompt_tokens=1, completion_tokens=1) - - session.client.chat.completions.create = fake_create + session.client.chat.completions.create = scripted_chat_client( + { + "tool_calls": [ + # blank id — the back-fill case + {"id": "", "name": "read_file", "arguments": '{"path": "x"}'} + ], + "finish_reason": "tool_calls", + "reasoning_content": "work it out", + "prompt_tokens": 1, + "completion_tokens": 1, + }, + {"content": "done", "prompt_tokens": 1, "completion_tokens": 1}, + ) def fake_prepare(tc_dict, **_kwargs): return { @@ -2906,23 +2834,15 @@ class TestAgentChildRegistration: ) session.ui.note_agent_child = MagicMock() - seen_messages: list[list[dict]] = [] - call_count = [0] - - def fake_create(**kwargs): - seen_messages.append(kwargs.get("messages") or []) - call_count[0] += 1 - if call_count[0] == 1: - return fake_chat_stream( - tool_calls=[ - {"id": "call_0", "name": "read_file", "arguments": '{"path": "x"}'} - ], - finish_reason="tool_calls", - reasoning_content="scan the repo first", - ) - return fake_chat_stream(content="done") - - session.client.chat.completions.create = fake_create + client_fn = scripted_chat_client( + { + "tool_calls": [{"id": "call_0", "name": "read_file", "arguments": '{"path": "x"}'}], + "finish_reason": "tool_calls", + "reasoning_content": "scan the repo first", + }, + {"content": "done"}, + ) + session.client.chat.completions.create = client_fn def fake_prepare(tc_dict, **_kwargs): return { @@ -2956,7 +2876,7 @@ class TestAgentChildRegistration: # The replay request carries the vLLM ``reasoning`` field on the # assistant turn; the internal ``_provider_content`` key is stripped # by the provider's sanitize before the wire. - replay = seen_messages[1] + replay = client_fn.calls[1].get("messages") or [] assistant_wire = next(m for m in replay if m.get("role") == "assistant") assert assistant_wire.get("reasoning") == "scan the repo first" assert "_provider_content" not in assistant_wire @@ -2976,20 +2896,15 @@ class TestRunAgentDenialMessage: session = _make_session() session._provider = OpenAIChatCompletionsProvider() - call_count = [0] - - def fake_create(**_kwargs): - call_count[0] += 1 - if call_count[0] == 1: - return fake_chat_stream( - tool_calls=[ - {"id": "call_1", "name": "notify", "arguments": '{"message": "hi"}'} - ], - finish_reason="tool_calls", - ) - return fake_chat_stream(content="done") - - session.client.chat.completions.create = fake_create + session.client.chat.completions.create = scripted_chat_client( + { + "tool_calls": [ + {"id": "call_1", "name": "notify", "arguments": '{"message": "hi"}'} + ], + "finish_reason": "tool_calls", + }, + {"content": "done"}, + ) # approve_tools is the real two-phase gate: on denial it stamps a # specific denial_msg on the item AND returns the reason as its 2nd # value. The sub-agent must honour both, not overwrite them. @@ -3275,20 +3190,15 @@ class TestSubAgentErrorRecall: session = _make_session() session._provider = OpenAIChatCompletionsProvider() - calls = [0] - - def fake_create(**_kwargs): - calls[0] += 1 - if calls[0] == 1: - return fake_chat_stream( - tool_calls=[ - {"id": "call_1", "name": "bash", "arguments": '{"command":"false"}'} - ], - finish_reason="tool_calls", - ) - return fake_chat_stream(content="done") - - session.client.chat.completions.create = fake_create + session.client.chat.completions.create = scripted_chat_client( + { + "tool_calls": [ + {"id": "call_1", "name": "bash", "arguments": '{"command":"false"}'} + ], + "finish_reason": "tool_calls", + }, + {"content": "done"}, + ) def fake_prepare(tc_dict, **_kwargs): cid = tc_dict["id"] diff --git a/turnstone/core/deadline.py b/turnstone/core/deadline.py index 3bf4cc8b..3441c81a 100644 --- a/turnstone/core/deadline.py +++ b/turnstone/core/deadline.py @@ -77,6 +77,37 @@ class StreamAbortRef(list[Any]): stream.close() +def run_abortable_with_deadline( + fn: Callable[[StreamAbortRef], _T], + *, + timeout: float, + cancel_event: threading.Event | None = None, + poll: float = 1.0, + thread_name: str = "deadline-worker", +) -> _T: + """:func:`run_with_deadline` with the stream-abort wiring built in. + + Mints a :class:`StreamAbortRef`, hands it to *fn* (thread it into the + provider call as ``cancel_ref``), and aborts it on either abandonment + path — the three-point pairing (ref + ``cancel_ref`` + ``on_abandon``) + cannot be half-wired. The canonical deadline-bounded sampling shape:: + + run_abortable_with_deadline( + lambda ref: model_turn(lane, turns, cancel_ref=ref, ...), + timeout=..., + ) + """ + abort_ref = StreamAbortRef() + return run_with_deadline( + lambda: fn(abort_ref), + timeout=timeout, + cancel_event=cancel_event, + poll=poll, + thread_name=thread_name, + on_abandon=abort_ref.abort, + ) + + def run_with_deadline( fn: Callable[[], _T], *, diff --git a/turnstone/core/judge.py b/turnstone/core/judge.py index 82cb8020..84f8483c 100644 --- a/turnstone/core/judge.py +++ b/turnstone/core/judge.py @@ -16,15 +16,13 @@ import threading import time import uuid from dataclasses import dataclass, field -from functools import partial from pathlib import Path from typing import TYPE_CHECKING, Any from turnstone.core.deadline import ( DeadlineCancelledError, DeadlineExceededError, - StreamAbortRef, - run_with_deadline, + run_abortable_with_deadline, ) from turnstone.core.log import get_logger from turnstone.core.model_turn import model_turn, resolve_capabilities, resolve_lane @@ -33,6 +31,8 @@ from turnstone.core.trajectory import Turn if TYPE_CHECKING: from collections.abc import Callable + from turnstone.core.deadline import StreamAbortRef + from turnstone.core.model_turn import ModelTurnResult from turnstone.core.providers._protocol import LLMProvider, ModelCapabilities log = get_logger(__name__) @@ -1387,34 +1387,39 @@ class IntentJudge: # Per-turn timeout: each turn gets a fresh budget so local # 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. 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). # A timeout or cancel abandons the call without pinning a # non-daemon thread that would block interpreter exit — the old # single-slot ThreadPoolExecutor left a stuck worker that # poisoned the pool, which is why the restart dance existed. + # The abort wiring (fresh ref per turn) closes the abandoned + # worker's HTTP stream so the read raises promptly. # Temperature is deliberately NOT pinned (house rule): the # lane inherits the judge model's configured temperature — # many modern models misbehave below 1.0, so the model's own # configuration beats a hard determinism pin. - result = run_with_deadline( - partial( - model_turn, + turn_tools = None if is_last_turn else tools + + # Bound default: the call runs synchronously within this + # iteration; the binding makes the per-turn capture explicit + # (and satisfies B023 in the loop). + def _sample( + ref: StreamAbortRef, _tools: list[dict[str, Any]] | None = turn_tools + ) -> ModelTurnResult: + return model_turn( lane, judge_turns, - tools=None if is_last_turn else tools, + tools=_tools, max_tokens=2048, - cancel_ref=abort_ref, - ), + cancel_ref=ref, + ) + + result = run_abortable_with_deadline( + _sample, timeout=per_call_timeout, cancel_event=cancel_event, thread_name="judge-api", - on_abandon=abort_ref.abort, ) except DeadlineCancelledError: return None diff --git a/turnstone/core/output_guard_judge.py b/turnstone/core/output_guard_judge.py index 04aac758..98a8e2fe 100644 --- a/turnstone/core/output_guard_judge.py +++ b/turnstone/core/output_guard_judge.py @@ -13,7 +13,7 @@ Design: already in hand. - JSON-in-content verdict. 4-strategy parser inlined from :meth:`IntentJudge._parse_verdict`. -- Wall-clock deadline via :func:`turnstone.core.deadline.run_with_deadline`, +- Wall-clock deadline via :func:`turnstone.core.deadline.run_abortable_with_deadline`, which runs the call on a *daemon* worker and polls the cancel event each second. A timeout or cancel abandons the call rather than waiting it out, and the daemon worker can never block process or interpreter exit — unlike @@ -46,8 +46,7 @@ from turnstone.core import fence from turnstone.core.deadline import ( DeadlineCancelledError, DeadlineExceededError, - StreamAbortRef, - run_with_deadline, + run_abortable_with_deadline, ) from turnstone.core.judge import ( _CHARS_PER_TOKEN, @@ -457,7 +456,7 @@ class OutputGuardJudge: field leave it at its default — the prompt skips empty sections. Timeout enforcement is real wall-clock: the upstream call runs on a - daemon worker via :func:`~turnstone.core.deadline.run_with_deadline` + daemon worker via :func:`~turnstone.core.deadline.run_abortable_with_deadline` and is abandoned on the timeout / cancel path, so a hung upstream LLM call neither blocks return nor pins interpreter exit. """ @@ -548,23 +547,21 @@ class OutputGuardJudge: # unbounded, a pass that consumes the whole 512-token cap parses # to a labelled llm_error verdict (heuristic tier stands) — the # remediation is an effort value on the guard's model alias. - # The abort ref closes the abandoned worker's HTTP stream on the + # The abort wiring closes the abandoned worker's HTTP stream on the # timeout/cancel paths so the daemon thread exits promptly instead # of blocking on the read until the upstream's next chunk. - abort_ref = StreamAbortRef() try: - result = run_with_deadline( - lambda: model_turn( + result = run_abortable_with_deadline( + lambda ref: model_turn( lane, judge_turns, tools=None, max_tokens=512, - cancel_ref=abort_ref, + cancel_ref=ref, ), timeout=timeout, cancel_event=cancel_event, thread_name="output-guard-judge", - on_abandon=abort_ref.abort, ) except DeadlineCancelledError: return self._error_verdict(verdict_id, call_id, start, "cancelled") diff --git a/turnstone/core/providers/_anthropic.py b/turnstone/core/providers/_anthropic.py index 8775930f..624027af 100644 --- a/turnstone/core/providers/_anthropic.py +++ b/turnstone/core/providers/_anthropic.py @@ -1103,6 +1103,8 @@ class AnthropicProvider: "InternalServerError", "APIError", "OverloadedError", + # Transport-level: drained stream ended without a stop reason. + "IncompleteStreamError", } ) diff --git a/turnstone/core/providers/_openai_common.py b/turnstone/core/providers/_openai_common.py index 5f9d013b..2f276722 100644 --- a/turnstone/core/providers/_openai_common.py +++ b/turnstone/core/providers/_openai_common.py @@ -839,5 +839,8 @@ RETRYABLE_ERROR_NAMES: frozenset[str] = frozenset( "RateLimitError", "Timeout", "APITimeoutError", + # Transport-level: the drained stream ended with no finish signal + # (generation died mid-response) — re-run like a wire error. + "IncompleteStreamError", } ) diff --git a/turnstone/core/providers/_openai_responses.py b/turnstone/core/providers/_openai_responses.py index f9dacb15..82c151ee 100644 --- a/turnstone/core/providers/_openai_responses.py +++ b/turnstone/core/providers/_openai_responses.py @@ -42,15 +42,23 @@ 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). +# response.failed error codes that indicate a transient server-side +# condition — the only ones worth retrying (the API's other codes are +# deterministic request rejections). +_TRANSIENT_FAILURE_CODES = frozenset({"server_error", "rate_limit_exceeded"}) - 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). + +class ResponsesStreamFailedError(RuntimeError): + """A TRANSIENT in-band ``response.failed`` terminal event. + + Raised only for ``_TRANSIENT_FAILURE_CODES`` (server error, rate + limit) — and listed in the provider's ``retryable_error_names`` — so + retry loops treat those like the wire-level errors they stand in + for. Deterministic in-band failures (invalid prompt, image fetch, + policy) raise plain ``RuntimeError`` and stop retry loops on attempt + zero, exactly as the retired non-streaming lane's HTTP errors did. + Callers that give up keep their degrade paths (judges fall back to + the heuristic tier). """ @@ -666,6 +674,19 @@ class OpenAIResponsesProvider: 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 sc = StreamChunk( finish_reason=last_finish, usage=usage, @@ -680,7 +701,16 @@ 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 ResponsesStreamFailedError(f"Responses API error: {error_msg}") + 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}") log.debug( "openai.responses.response", diff --git a/turnstone/core/providers/_protocol.py b/turnstone/core/providers/_protocol.py index 2e923847..438a8664 100644 --- a/turnstone/core/providers/_protocol.py +++ b/turnstone/core/providers/_protocol.py @@ -7,7 +7,7 @@ knowing provider-specific details. from __future__ import annotations -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable if TYPE_CHECKING: @@ -67,6 +67,21 @@ class CompletionResult: reasoning: str = "" +class IncompleteStreamError(RuntimeError): + """The stream ended without any terminal/finish signal. + + Every adapter emits a finish reason on a healthy stream (Chat + Completions' final choice chunk, Anthropic's ``message_delta`` stop + reason, Responses' terminal event); a stream that exhausts without + one is a generation that died mid-response behind a proxy/ASGI layer + that closed the body cleanly. Typed and listed in every provider's + ``retryable_error_names`` so callers re-run it like the wire errors + it stands in for — restoring the retired non-streaming transport's + complete-or-error contract for single-shot lanes (the interactive + loop keeps showing partial output live; this gate is drain-only). + """ + + def merge_usage(acc: UsageInfo | None, new: UsageInfo) -> UsageInfo: """Merge one stream-chunk usage report into an accumulator, per-field max. @@ -82,13 +97,7 @@ def merge_usage(acc: UsageInfo | None, new: UsageInfo) -> UsageInfo: 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, - ) + return replace(new) prompt = max(acc.prompt_tokens, new.prompt_tokens) completion = max(acc.completion_tokens, new.completion_tokens) return UsageInfo( @@ -106,25 +115,33 @@ def drain_stream(chunks: Iterator[StreamChunk]) -> CompletionResult: The ONE non-streaming transport: single-shot callers (``model_turn``) sample through the provider's streaming entry and accumulate here, so the streaming and non-streaming lanes cannot drift apart per adapter. - Accumulation mirrors the main loop's chunk consumer (``ChatSession``): + Accumulation mirrors the main loop's chunk consumer (``ChatSession``), + plus the complete-or-error gate the interactive loop doesn't need: - - ``usage`` merges per-field ``max`` across chunks — Anthropic splits - prompt tokens (``message_start``) and completion tokens - (``message_delta``) into separate events, so neither first-wins nor - last-wins sees both. + - A stream that exhausts with NO finish reason raises + :class:`IncompleteStreamError` (retryable) — every adapter emits one + on a healthy stream, so its absence means the generation died + mid-response. Partial text must never be handed to a caller that + stores it as a complete result (a compaction summary, a title). + - ``usage`` merges via :func:`merge_usage` — Anthropic splits prompt + and completion tokens across separate events. - Tool calls accumulate by ``ToolCallDelta.index``: ``id``/``name`` - are whole values (last truthy wins), ``arguments_delta`` concatenates. + are whole values (last truthy wins), ``arguments_delta`` + concatenates. A delta carrying an id DIFFERENT from its slot's + opens a new call instead — index-degenerate compat servers + (historical vLLM/llama.cpp builds emit every parallel call at + index 0) would otherwise fuse distinct calls into garbage + arguments. Deltas without ids keep routing to their index's + current call: fragments follow their call's announcement. - ``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`` 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, 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. + - ``info_delta`` before the finish reason is transient status (server- + side search pings) that the non-streaming lane never surfaced — + dropped. ``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. Raises whatever the underlying stream raises — retry/deadline/fallback policy stays with the caller, exactly as with the old non-streaming @@ -133,8 +150,12 @@ 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]] = {} + # Slots in arrival order, each remembering its wire index — the result + # sorts by (index, arrival) so well-formed streams keep the array order + # the retired non-streaming body had, and collision-opened slots stay + # in arrival order behind their shared index. + tool_slots: list[tuple[int, dict[str, Any]]] = [] + slot_for_index: dict[int, int] = {} usage: UsageInfo | None = None finish_reason: str | None = None provider_blocks: list[dict[str, Any]] = [] @@ -145,10 +166,19 @@ def drain_stream(chunks: Iterator[StreamChunk]) -> CompletionResult: if sc.reasoning_delta: reasoning_parts.append(sc.reasoning_delta) for tcd in sc.tool_call_deltas: - tc = tool_calls_acc.setdefault( - tcd.index, - {"id": "", "type": "function", "function": {"name": "", "arguments": ""}}, - ) + slot = slot_for_index.get(tcd.index) + if slot is None or ( + tcd.id and tool_slots[slot][1]["id"] and tool_slots[slot][1]["id"] != tcd.id + ): + slot = len(tool_slots) + slot_for_index[tcd.index] = slot + tool_slots.append( + ( + tcd.index, + {"id": "", "type": "function", "function": {"name": "", "arguments": ""}}, + ) + ) + tc = tool_slots[slot][1] if tcd.id: tc["id"] = tcd.id if tcd.name: @@ -161,34 +191,25 @@ def drain_stream(chunks: Iterator[StreamChunk]) -> CompletionResult: finish_reason = sc.finish_reason if sc.provider_blocks: provider_blocks = sc.provider_blocks - 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() + # 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 finish_reason is None: + raise IncompleteStreamError( + "stream ended without a finish reason — generation died mid-response" + ) content = "".join(content_parts) - for info in trailing_info_parts + suffix_info_parts: + for info in trailing_info_parts: content += "\n\n" + info - tool_calls = [tool_calls_acc[i] for i in sorted(tool_calls_acc)] + tool_calls = [tc for _, tc in sorted(tool_slots, key=lambda pair: pair[0])] return CompletionResult( content=content, tool_calls=tool_calls or None, - finish_reason=finish_reason or "stop", + finish_reason=finish_reason, usage=usage, provider_blocks=provider_blocks, reasoning="".join(reasoning_parts),