diff --git a/CHANGELOG.md b/CHANGELOG.md index fc2bbb9c..0a0032f0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -141,6 +141,15 @@ Earlier stable lines (`stable/1.6`, `stable/1.5`) are frozen. ### Changed +- **Log event rename: `drain_stream.post_finish_blip` is now + `stream.post_finish_blip`, without the `usage_captured` field.** The + single-shot drain normalizes mid-body transport deaths through the same + `transport_guarded` wrapper the interactive loop uses, so its + post-finish-blip tolerance logs under the wrapper's event name. Update + any external log filters pinned to the old name; the drained result's + possible `usage=None` on a post-finish blip is unchanged and documented + on `drain_stream`. + - **Breaking (1.8): compaction feedback moved from `info` events to the typed `compaction` SSE event.** Pre-1.8 SSE/SDK clients that ignore unknown event types no longer see compaction lines (they are diff --git a/tests/test_drain_stream.py b/tests/test_drain_stream.py index 262fc833..fafa24ef 100644 --- a/tests/test_drain_stream.py +++ b/tests/test_drain_stream.py @@ -351,6 +351,22 @@ class TestTransportGuarded: class TestErrorPropagation: + def test_post_finish_blip_keeps_completed_result(self): + # The generation completed (finish reason in hand) — a trailing + # transport blip forfeits only trailing metadata (here: the usage + # chunk), never the completed result. + import httpx + + def chunks(): + yield StreamChunk(content_delta="whole answer") + yield StreamChunk(finish_reason="stop") + raise httpx.ReadError("late blip") + + result = drain_stream(chunks()) + assert result.content == "whole answer" + assert result.finish_reason == "stop" + assert result.usage is None + 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 diff --git a/turnstone/core/providers/__init__.py b/turnstone/core/providers/__init__.py index 99950ab3..d4ff36bb 100644 --- a/turnstone/core/providers/__init__.py +++ b/turnstone/core/providers/__init__.py @@ -18,6 +18,7 @@ from turnstone.core.providers._protocol import ( UsageInfo, accumulate_tool_call_delta, drain_stream, + merge_usage, transport_guarded, ) from turnstone.core.providers._xai import XAI_DEFAULT_BASE_URL, XAIProvider @@ -40,6 +41,7 @@ __all__ = [ "drain_stream", "list_known_models", "lookup_model_capabilities", + "merge_usage", "transport_guarded", ] diff --git a/turnstone/core/providers/_protocol.py b/turnstone/core/providers/_protocol.py index d6663fab..aa8c1042 100644 --- a/turnstone/core/providers/_protocol.py +++ b/turnstone/core/providers/_protocol.py @@ -92,9 +92,11 @@ def merge_usage(acc: UsageInfo | None, new: UsageInfo) -> UsageInfo: 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). + Serves :func:`drain_stream` and ``ChatSession``'s inline chunk + consumer alike — the interactive loop re-projects the merged + accumulator into its ``_last_usage`` dict on every usage chunk (that + dict has mid-stream readers), so the two lanes cannot drift on the + merge rule. """ if acc is None: return replace(new) @@ -221,7 +223,9 @@ def drain_stream(chunks: Iterator[StreamChunk]) -> CompletionResult: never be handed to a caller that stores it as a complete result (a compaction summary, a title). A transport blip AFTER the finish reason keeps the completed result and forfeits only trailing - metadata. + metadata — on the chat lane the usage chunk trails the finish + reason, so that result may report usage=None and the call's spend + goes missing from usage accounting. - ``usage`` merges via :func:`merge_usage` — Anthropic splits prompt and completion tokens across separate events. - Tool calls accumulate by ``ToolCallDelta.index`` via @@ -241,15 +245,14 @@ 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 — 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. + transport — EXCEPT httpx transport failures, normalized by + :func:`transport_guarded` (the one conversion rule, shared with the + interactive loop): a mid-body death before the finish reason is + re-raised (chained) as :class:`IncompleteStreamError`, restoring the + wire-blip retryability the non-streaming transport had, and a + post-finish blip ends the stream cleanly so the completed result is + kept. """ - 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] = [] @@ -258,32 +261,7 @@ def drain_stream(chunks: Iterator[StreamChunk]) -> CompletionResult: finish_reason: str | None = None provider_blocks: list[dict[str, Any]] = [] - iterator = iter(chunks) - while True: - try: - sc = next(iterator) - except StopIteration: - break - except httpx.TransportError as exc: - if finish_reason is not None: - # The generation already completed (finish reason in hand); - # the blip only cost trailing metadata — a usage-only chunk - # or the citation footer. Keep the complete result rather - # than discarding it for a retry that re-pays the tokens — - # but say so: on the chat lane the usage chunk trails the - # finish reason, so this result may report usage=None and - # the call's spend goes missing from usage accounting. - import structlog # noqa: PLC0415 — deferred with httpx off the type-module path - - structlog.get_logger(__name__).warning( - "drain_stream.post_finish_blip", - error_type=type(exc).__name__, - usage_captured=usage is not None, - ) - break - raise IncompleteStreamError( - f"stream transport failed mid-response ({type(exc).__name__}: {exc})" - ) from exc + for sc in transport_guarded(chunks): if sc.content_delta: content_parts.append(sc.content_delta) if sc.reasoning_delta: diff --git a/turnstone/core/session.py b/turnstone/core/session.py index 8b3b3590..32fb8417 100644 --- a/turnstone/core/session.py +++ b/turnstone/core/session.py @@ -175,6 +175,7 @@ from turnstone.core.preview import ( from turnstone.core.providers import ( accumulate_tool_call_delta, create_provider, + merge_usage, transport_guarded, ) from turnstone.core.ratelimit import TokenBucket @@ -7786,6 +7787,7 @@ class ChatSession: reasoning_parts: list[str] = [] tool_calls_acc: dict[int, dict[str, Any]] = {} provider_blocks: list[dict[str, Any]] = [] + usage_acc: UsageInfo | None = None first_token = True in_think = False # inside a ... block path1_reasoning = False # last reasoning came via reasoning_delta field @@ -7858,6 +7860,31 @@ class ChatSession: self.ui.on_thinking_stop() first_token = False + def _record_cancelled_partial() -> None: + """Flush buffered text, finalize the stream in the UI, and stash + the partial for send()'s cancel handler — the one sequence both + cancel arms (cooperative and stream-close-converted) must run. + ``tool_calls`` and ``_provider_content`` are DELIBERATELY + OMITTED from the partial: + + * ``tool_calls`` — incomplete, no matching tool_result; + re-emitting on the next turn would orphan them. + * ``_provider_content`` — the Anthropic provider reads this + lane verbatim ahead of plain ``content`` (see + ``providers/_anthropic.py``), and a cancellation can leave + partial tool_use blocks here too. Keeping it would also + cause the next-turn replay to bypass the ``[generation + cancelled before completion]`` marker the cancel handler + appends to ``content``, hiding the partial-output signal + from the model. + """ + if pending: + _flush_text(pending, in_think) + self.ui.on_stream_end() + partial: dict[str, Any] = {"role": "assistant"} + partial["content"] = "".join(content_parts) or "" + self._cancelled_partial_msg = partial + finish_reason = None try: for chunk in stream: @@ -7873,35 +7900,20 @@ class ChatSession: finish_reason = chunk.finish_reason # Accumulate usage (Anthropic sends prompt tokens in message_start - # and completion tokens in message_delta as separate events) + # and completion tokens in message_delta as separate events) via + # THE shared max-merge rule (merge_usage, drain_stream's twin). + # self._last_usage is re-written on EVERY usage chunk: it is + # read mid-stream (_estimated_prompt_tokens, the status line), + # so the dict write cannot defer to stream end. if chunk.usage: - if self._last_usage is None: - self._last_usage = { - "prompt_tokens": chunk.usage.prompt_tokens, - "completion_tokens": chunk.usage.completion_tokens, - "total_tokens": chunk.usage.total_tokens, - "cache_creation_tokens": chunk.usage.cache_creation_tokens, - "cache_read_tokens": chunk.usage.cache_read_tokens, - } - else: - self._last_usage["prompt_tokens"] = max( - self._last_usage["prompt_tokens"], chunk.usage.prompt_tokens - ) - self._last_usage["completion_tokens"] = max( - self._last_usage["completion_tokens"], chunk.usage.completion_tokens - ) - self._last_usage["total_tokens"] = ( - self._last_usage["prompt_tokens"] - + self._last_usage["completion_tokens"] - ) - self._last_usage["cache_creation_tokens"] = max( - self._last_usage.get("cache_creation_tokens", 0), - chunk.usage.cache_creation_tokens, - ) - self._last_usage["cache_read_tokens"] = max( - self._last_usage.get("cache_read_tokens", 0), - chunk.usage.cache_read_tokens, - ) + usage_acc = merge_usage(usage_acc, chunk.usage) + self._last_usage = { + "prompt_tokens": usage_acc.prompt_tokens, + "completion_tokens": usage_acc.completion_tokens, + "total_tokens": usage_acc.total_tokens, + "cache_creation_tokens": usage_acc.cache_creation_tokens, + "cache_read_tokens": usage_acc.cache_read_tokens, + } if self.debug: parts = [] @@ -7960,26 +7972,7 @@ class ChatSession: if chunk.provider_blocks: provider_blocks = chunk.provider_blocks except GenerationCancelled: - # Flush whatever was buffered and build a partial message. - # Both ``tool_calls`` and ``_provider_content`` are - # DELIBERATELY OMITTED: - # * ``tool_calls`` — incomplete, no matching tool_result; - # re-emitting on the next turn would orphan them. - # * ``_provider_content`` — the Anthropic provider reads - # this lane verbatim ahead of plain ``content`` (see - # ``providers/_anthropic.py``), and a cancellation can - # leave partial tool_use blocks here too. Keeping it - # would also cause the next-turn replay to bypass the - # ``[generation cancelled before completion]`` marker - # the cancel handler appends to ``content``, hiding - # the partial-output signal from the model. - if pending: - _flush_text(pending, in_think) - self.ui.on_stream_end() - partial: dict[str, Any] = {"role": "assistant"} - partial_content = "".join(content_parts) - partial["content"] = partial_content or "" - self._cancelled_partial_msg = partial + _record_cancelled_partial() raise except Exception: # cancel() closed the underlying SDK stream, aborting the HTTP @@ -7987,17 +7980,7 @@ class ChatSession: # transport-level error (httpx, httpcore, etc.). Convert to # GenerationCancelled if a cancel was requested. if self._cancel_event.is_set(): - if pending: - _flush_text(pending, in_think) - self.ui.on_stream_end() - partial = {"role": "assistant"} - partial["content"] = "".join(content_parts) or "" - # Same reasoning as the cooperative-cancel branch - # above: ``_provider_content`` is omitted so the - # next-turn replay reads from the marker-bearing - # plain content and any partial tool_use blocks - # inside provider_blocks don't leak through. - self._cancelled_partial_msg = partial + _record_cancelled_partial() raise GenerationCancelled() from None raise