From 961a2017dc59635739cfe424239f23f18ace3d7d Mon Sep 17 00:00:00 2001 From: Patrick Buckley Date: Tue, 4 Aug 2026 03:09:48 -0700 Subject: [PATCH] fix(streaming): delete the retry window's shared slots and gate the send epilogue MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fourth review round. The recurring defect family — cross-frame session slots racing an orphanable window — is removed structurally instead of gated again: - The wire-fold slot is deleted. The fold the stream was actually created from rides the returned message dict on the underscore lane (like _provider_content) and is popped at the single calibration site before commit, so a superseding generation can never alias it and there is nothing left to clear. Plain-dict test fakes fall through the pop to the frame-local fold. - The stream-provider slot is demoted to a creation-time handoff register: _try_stream stamps it, _stream_response copies it into a frame-local immediately after each create returns, and only that local feeds the retry gate. The fatal formatter returns to the consistent PRIMARY identity triple — pairing a fallback's provider name with the primary's base_url and alias sent operators to debug the wrong backend; stamping the full producing identity is #964. - send()'s epilogue is generation-gated: a superseded thread's escaped death no longer records a fatal error over the healthy successor turn (error banner, buffer-wiping error-state drain, wrong last_error for the coord), and a Ctrl-C on an orphan no longer mutates history. - The terminal arm discards as well as finalizes. Keeping the buffers bought nothing — the fatal path's error-state drain wipes them on every server lane — and the skipped discard let a mid-consumption overflow recovered by compact-and-retry concatenate the dead attempt's text with the recovered answer in the idle payload. Pinned with real-buffer tests for the overflow-recovery and orphan-epilogue paths. - stream.post_finish_blip regains usage_captured, tracked by transport_guarded from the chunks it forwards, restoring missing-spend attribution on both lanes. - TerminalUI.on_thinking_start is idempotent at the callee (a live spinner is stopped before being replaced), removing the caller-side stop-first dance and the leak the next unaware call site would have reintroduced. - The think-tag vocabulary in _strip_reasoning and the title lane is derived from ThinkTagSplitter, closing the drift channel that would leak raw reasoning into compaction summaries and titles. - on_stream_discarded's docstring states the true pending-batch semantics (defensive drop; the shipped sequence flushes via the preceding stream_end), and the live-suite recording fake gains the protocol method. --- tests/test_cancel.py | 3 + tests/test_drain_stream.py | 6 +- tests/test_midstream_retry.py | 85 ++++++++++++--- tests/test_server_live.py | 4 + turnstone/cli.py | 5 + turnstone/core/providers/_protocol.py | 9 ++ turnstone/core/session.py | 149 ++++++++++++++------------ turnstone/core/session_ui_base.py | 8 +- 8 files changed, 183 insertions(+), 86 deletions(-) diff --git a/tests/test_cancel.py b/tests/test_cancel.py index 0fc8ffae..f4f5e783 100644 --- a/tests/test_cancel.py +++ b/tests/test_cancel.py @@ -39,6 +39,9 @@ class NullUI: def on_turn_committed(self): pass + def on_stream_discarded(self): + pass + def on_thinking_start(self): pass diff --git a/tests/test_drain_stream.py b/tests/test_drain_stream.py index fafa24ef..3a7bb1ff 100644 --- a/tests/test_drain_stream.py +++ b/tests/test_drain_stream.py @@ -322,7 +322,11 @@ class TestTransportGuarded: out = list(transport_guarded(chunks())) assert [c.content_delta for c in out] == ["done", ""] assert out[-1].finish_reason == "stop" - assert any("stream.post_finish_blip" in r.message for r in caplog.records) + blips = [r.message for r in caplog.records if "stream.post_finish_blip" in r.message] + assert blips + # usage_captured is the missing-spend attribution signal: this + # stream never delivered a usage chunk, so the blip must say so. + assert any("usage_captured" in m and "False" in m for m in blips) def test_chunks_pass_through_untouched(self): src = [ diff --git a/tests/test_midstream_retry.py b/tests/test_midstream_retry.py index 6cb9b02d..128ca471 100644 --- a/tests/test_midstream_retry.py +++ b/tests/test_midstream_retry.py @@ -229,14 +229,13 @@ class TestMidStreamRetry: assert persisted and "Backend stream died mid-response" in persisted fatal = [r for r in caplog.records if "session.fatal.recorded" in r.message] assert fatal and any(r.levelno == logging.ERROR for r in fatal) - # Every dead attempt is finalized client-side — two retry-path - # stream_ends plus the terminal arm's (the CLI markdown fence reset - # rides on it). The retry arms DISCARD their dead segments; no - # commit ever happens, and the terminal arm deliberately keeps the - # in-progress snapshot, the last partial's only surviving copy, for - # refresh-replay. + # Every dead attempt is finalized AND discarded — the terminal arm + # included (keeping its buffers bought nothing: the fatal path's + # error-state drain wipes them anyway, and an overflow recovered by + # compact-and-retry would otherwise concatenate dead text into the + # idle payload). No commit ever happens. assert ui.kinds().count("stream_end") == 3 - assert ui.kinds().count("stream_discarded") == 2 + assert ui.kinds().count("stream_discarded") == 3 assert ui.kinds().count("turn_committed") == 0 def test_cancel_during_backoff_stops_without_recreate(self, tmp_db): @@ -329,7 +328,7 @@ class TestMidStreamRetry: # hardening). assert refresh.call_count == 3 - def test_retry_window_stops_then_restarts_spinner(self, tmp_db): + def test_retry_window_restarts_spinner(self, tmp_db): ui = RecordingUI() session = _make_session(ui) streams = [ @@ -343,12 +342,12 @@ class TestMidStreamRetry: session.send("test") # A pre-first-token death leaves the spinner RUNNING; the CLI's - # on_thinking_start replaces the spinner object without stopping - # the old one (thread leak), so the retry arm must stop-then-start. + # on_thinking_start is idempotent at the callee (it stops a live + # spinner before replacing it), so the retry arm restarts with a + # single call. notice = [i for i, (k, d) in enumerate(ui.events) if k == "info" and "stream died" in d] assert notice - after = [k for k, _ in ui.events[notice[0] + 1 : notice[0] + 3]] - assert after == ["thinking_stop", "thinking_start"] + assert ui.events[notice[0] + 1][0] == "thinking_start" def test_pretoken_death_stop_persists_marker_row(self, tmp_db): ui = RecordingUI() @@ -638,6 +637,63 @@ class TestMidStreamRetry: # retried attempt's text reaches the IDLE payload. assert "".join(ui._ws_turn_content) == "final answer" + def test_overflow_recovery_discards_dead_text_from_turn_buffer(self, tmp_db): + # A mid-consumption overflow is TERMINAL for the retry ladder but + # RECOVERED by send()'s compact-and-retry — the dead attempt's text + # must not concatenate with the recovered answer in the idle + # payload (the terminal arm discards, same as the retry arm). + class _BufferUI(NullUI): + def on_state_change(self, state): + pass + + ui = _BufferUI() + session = _make_session(ui) + calls = {"n": 0} + + def create(msgs): + calls["n"] += 1 + if calls["n"] == 1: + return _dying_stream( + "dead overflow text", + exc=RuntimeError("maximum context length exceeded"), + ) + return _good_stream("recovered") + + with ( + patch.object(session, "_create_stream_with_retry", side_effect=create), + patch.object(session, "_compact_messages"), + patch.object(session, "_full_messages", return_value=[]), + ): + session.send("test") + + assert calls["n"] == 2 + assert "".join(ui._ws_turn_content) == "recovered" + + def test_orphan_death_records_no_fatal_over_successor(self, tmp_db): + ui = RecordingUI() + session = _make_session(ui) + + def dying_superseded(): + yield StreamChunk(content_delta="Hel") + # A force-cancel claims a successor generation before the + # orphan's death propagates (its cancel event was replaced, so + # cancel conversion cannot fire). + session._generation += 1 + raise ValueError("orphan death") + + with ( + patch.object(session, "_create_stream_with_retry", side_effect=[dying_superseded()]), + patch.object(session, "_full_messages", return_value=[]), + pytest.raises(ValueError, match="orphan death"), + ): + session.send("test") + + # The orphan must not flash an error banner over the live + # successor turn or persist a wrong last_error for the coord. + assert ("state", "error") not in ui.events + assert not ui.of("error") + assert not load_last_error(session._ws_id) + def test_non_retryable_error_is_immediately_fatal(self, tmp_db): ui = RecordingUI() session = _make_session(ui) @@ -655,10 +711,9 @@ class TestMidStreamRetry: assert create.call_count == 1 # zero retries assert ("state", "error") in ui.events assert not any("stream died mid-response" in d for d in ui.of("info")) - # The terminal arm finalizes even a zero-retry death — client-side - # only; the snapshot survives for refresh-replay of the partial. + # The terminal arm finalizes and discards even a zero-retry death. assert ui.kinds().count("stream_end") == 1 - assert ui.kinds().count("stream_discarded") == 0 + assert ui.kinds().count("stream_discarded") == 1 assert ui.kinds().count("turn_committed") == 0 def test_recreate_failure_surfaces_original_stream_death(self, tmp_db, caplog): diff --git a/tests/test_server_live.py b/tests/test_server_live.py index 2622a0f1..19373806 100644 --- a/tests/test_server_live.py +++ b/tests/test_server_live.py @@ -76,6 +76,10 @@ class RecordingUI: def on_turn_committed(self): self.events.append(("turn_committed",)) + def on_stream_discarded(self): + # A live-wire blip entering the retry arm must not crash the fake. + self.events.append(("stream_discarded",)) + def on_thinking_start(self): self.events.append(("thinking_start",)) diff --git a/turnstone/cli.py b/turnstone/cli.py index be00e7e1..def9389d 100644 --- a/turnstone/cli.py +++ b/turnstone/cli.py @@ -127,6 +127,11 @@ class TerminalUI(SessionUI): pass def on_thinking_start(self) -> None: + # Idempotent: replacing a RUNNING spinner would orphan its render + # thread (painting frames over all later output for the process + # lifetime) — callers must not need a stop-first protocol. + if self.spinner: + self.spinner.stop() self.spinner = Spinner("Thinking") self.spinner.start() diff --git a/turnstone/core/providers/_protocol.py b/turnstone/core/providers/_protocol.py index aa8c1042..59456281 100644 --- a/turnstone/core/providers/_protocol.py +++ b/turnstone/core/providers/_protocol.py @@ -182,6 +182,7 @@ def transport_guarded(chunks: Iterator[StreamChunk]) -> Iterator[StreamChunk]: import httpx # noqa: PLC0415 — heavyweight; deferred off the type-module import path finish_seen = False + usage_seen = False iterator = iter(chunks) while True: try: @@ -192,9 +193,15 @@ def transport_guarded(chunks: Iterator[StreamChunk]) -> Iterator[StreamChunk]: if finish_seen: import structlog # noqa: PLC0415 — deferred with httpx off the type-module path + # usage_captured distinguishes "completed result kept but + # its spend went missing from usage accounting" (the chat + # lane's usage chunk trails the finish) from a harmless + # citation-footer loss — the one log signal that lets a + # missing-spend incident be attributed afterward. structlog.get_logger(__name__).warning( "stream.post_finish_blip", error_type=type(exc).__name__, + usage_captured=usage_seen, ) return raise IncompleteStreamError( @@ -202,6 +209,8 @@ def transport_guarded(chunks: Iterator[StreamChunk]) -> Iterator[StreamChunk]: ) from exc if sc.finish_reason: finish_seen = True + if sc.usage is not None: + usage_seen = True yield sc diff --git a/turnstone/core/session.py b/turnstone/core/session.py index 1ff0f7ce..2c404009 100644 --- a/turnstone/core/session.py +++ b/turnstone/core/session.py @@ -1897,15 +1897,14 @@ class ChatSession: # the model detached on purpose. close() reaps everything. self._background_shells = BackgroundShellRegistry(on_exit=self._on_background_shell_exit) self._cancelled_partial_msg: dict[str, Any] | None = None - # The provider that created the LIVE stream (primary or fallback) — - # the mid-stream retry gate consults ITS retryable set, and the - # fatal formatter labels errors with it. Best-effort under - # supersede races; None falls back to the primary binding. + # Creation-time HANDOFF REGISTER: _try_stream stamps the provider + # that owns the stream it is about to return (fallback walk + # included), and _stream_response copies it into a frame local + # immediately after each create returns. Nothing else reads it — + # a late read would race a superseding generation's creation — and + # it is deliberately never cleared (stale values are unreachable + # by construction). self._active_stream_provider: LLMProvider | None = None - # The wire-message fold the live stream was actually created from — - # a mid-retry rebind re-prepares it, and send()'s token-table - # calibration must count the fold the provider counted. - self._active_wire_msgs: list[dict[str, Any]] | None = None self._pending_retry: str | None = None # True when a fatal exception's text has been persisted to # workstream_config["last_error"] for the coord's inspect/wait @@ -3945,7 +3944,7 @@ class ChatSession: # tag from lanes that pre-inject the opening ```` into the # prompt (only ```` reaches ``content``). See ``_TITLE_*``. stripped = self._strip_reasoning(raw) - for _close in ("", ""): + for _close in ThinkTagSplitter.CLOSE_TAGS: _pos = stripped.rfind(_close) if _pos != -1: stripped = stripped[_pos + len(_close) :] @@ -5461,12 +5460,12 @@ class ChatSession: log.debug("session.fatal.base_url_lookup_failed", exc_info=True) provider_label = "?" try: - # Prefer the provider that actually owned the live stream — a - # fallback-carried turn's failure must not be labeled with the - # primary binding it never touched. getattr: this helper is - # bound to lightweight stubs in tests (see the module-scope - # note on the _BACKEND_* sets) that predate the field. - prov = getattr(self, "_active_stream_provider", None) or self._provider + # The PRIMARY binding, deliberately: base_url and model_label + # above come from the primary too, and a mixed identity (a + # fallback's provider name over the primary's URL and alias) + # sends the operator to debug the wrong backend. Stamping the + # full producing identity onto turns/errors is #964. + prov = self._provider provider_label = ( getattr(prov, "provider_name", None) or type(prov).__name__ if prov else "?" ) @@ -6907,11 +6906,16 @@ class ChatSession: # (perf-2); passing the already-prepared list keeps the # calibration char count aligned with what the provider # actually counted. - # The wire fold the provider ACTUALLY counted — a mid-retry - # rebind re-prepares it inside _stream_response, invisibly - # to this frame's msgs local (the calibration contract: - # char counts align with the counted fold). - self._update_token_table(assistant_msg, msgs=self._active_wire_msgs or msgs) + # The wire fold the provider ACTUALLY counted rides the + # returned message (a mid-retry rebind re-prepares it + # inside _stream_response, invisibly to this frame's msgs + # local) — popped BEFORE the message is committed so the + # carrier key never persists. Frame-owned, so a + # superseding generation cannot alias it; plain-dict fakes + # without the key fall through to the frame-local fold. + self._update_token_table( + assistant_msg, msgs=assistant_msg.pop("_wire_msgs", None) or msgs + ) self._print_status_line() # Report usage for EVERY API call self.messages.append(turn_from_dict(assistant_msg)) # Clear per-turn inflight buffers — the assistant @@ -7409,12 +7413,23 @@ class ChatSession: # Do NOT re-raise — return normally so server worker thread # completes cleanly. except KeyboardInterrupt as exc: + if self._generation != my_generation: + raise # orphaned: no history mutation or fatal over the live turn self._synthesize_cancelled_results("Interrupted by user.") self._flush_queued_messages() self._drain_pending_advisories() self._record_fatal_error(exc) raise except Exception as exc: + # Orphan gate: a superseded thread's stream death can escape + # cancel conversion (the successor replaced the cancel event + # before the blocked read died) and reach here — recording it + # would flash an error banner over the HEALTHY successor turn, + # wipe its buffers via the error-state drain, and persist a + # wrong last_error for the coord's inspect/wait. The wrapper's + # orphan arm re-raises for exactly this gate to absorb. + if self._generation != my_generation: + raise self._flush_queued_messages() self._drain_pending_advisories() self._record_fatal_error(exc) @@ -7425,15 +7440,6 @@ class ChatSession: # an idle session until the next send. Restores the "None outside a # send" invariant on every exit (success, cancel, or error). self._wire_part_cache = None - # Send-scoped stream bookkeeping: the live-stream provider label - # must not leak into a LATER lane's fatal formatting (title/ - # summary/task_agent fatals would otherwise wear a stale - # interactive binding), and the wire fold is a full-context- - # sized copy that must not outlive its calibration use above. - # Cleared HERE (after the except arms' _record_fatal_error ran, - # which is what reads the provider on the fatal path). - self._active_stream_provider = None - self._active_wire_msgs = None # Consume this generation's cancel signal on exit so a cancel that # targeted THIS send can't later abort an unrelated idle operation # (e.g. a manual /compact between sends would otherwise inherit the @@ -7709,11 +7715,16 @@ class ChatSession: @staticmethod def _strip_reasoning(text: str) -> str: - """Remove / tags and their content.""" - for open_t, close_t in [ - ("", ""), - ("", ""), - ]: + """Remove think-tag blocks and their content. + + The tag vocabulary is ThinkTagSplitter's — deriving it keeps this + strip (and the title lane's) from going blind when a variant is + added to the splitter, which would leak raw chain-of-thought into + compaction summaries. + """ + for open_t, close_t in zip( + ThinkTagSplitter.OPEN_TAGS, ThinkTagSplitter.CLOSE_TAGS, strict=True + ): while open_t in text: start = text.find(open_t) end = text.find(close_t, start) @@ -7778,14 +7789,25 @@ class ChatSession: "content": dead_partial, } - # Exported for send()'s token-table calibration — the char count - # must match the fold the provider actually counted, and a - # mid-retry rebind can re-prepare it below. - self._active_wire_msgs = msgs stream = self._create_stream_with_retry(msgs) + # FRAME-LOCAL copy of the creation-time handoff register, taken + # immediately after the create returns: the retry gate must judge a + # death by the provider that owns THIS stream (a fallback's set can + # differ), and reading the shared register later would race a + # superseding generation's own creation. + live_provider = self._active_stream_provider or self._provider while True: try: - return self._stream_attempt(transport_guarded(stream), my_generation) + result = self._stream_attempt(transport_guarded(stream), my_generation) + # The fold this turn was ACTUALLY created from rides the + # returned message (popped by send() at calibration, before + # commit — the message-dict underscore lane, like + # _provider_content): a mid-retry rebind re-prepares msgs, + # and send()'s frame-local copy cannot see that. Carried on + # the frame-owned dict rather than a session slot so a + # superseding generation can never alias it. + result["_wire_msgs"] = msgs + return result except GenerationCancelled: # A Stop during an attempt (incl. the re-create/TTFT window # after a death) — preserve the window's best partial. @@ -7811,16 +7833,11 @@ class ChatSession: # emitted here would clobber the NEW generation's # in-flight stream state. raise - # The provider whose stream actually died: after - # _try_fallback the live stream belongs to the FALLBACK - # provider, whose retryable set can differ (e.g. - # ResponsesStreamFailedError) — judging it by the primary - # binding's set would declare fallback transients terminal. - # Best-effort field (a racing creation could re-write it); - # at worst one gate decision consults a sibling set. - live_provider = self._active_stream_provider or self._provider # The terminal predicate is the SHARED _stop_retrying, - # capped at _MID_STREAM_RETRIES: its overflow arm applies + # capped at _MID_STREAM_RETRIES, judged by the FRAME-LOCAL + # live_provider (the provider that owns this stream — a + # fallback's retryable set can differ, e.g. + # ResponsesStreamFailedError). The overflow arm applies # here too — an overflow can surface mid-consumption # (error-frame lanes), and it must fall through to send()'s # compact-and-retry arm rather than burn re-issues on a @@ -7828,18 +7845,16 @@ class ChatSession: if self._stop_retrying( e, attempt, live_provider, max_retries=self._MID_STREAM_RETRIES ): - # Terminal: client-side finalize only (the CLI's - # markdown fence resets in on_stream_end; the server - # /send workers emit their own after a fatal, the - # CLI's direct send() does not). on_turn_committed is - # DELIBERATELY absent here, unlike the retry arm - # below: the dead partial is in neither history nor - # storage, so the in-progress snapshot is its only - # surviving copy — wiping it would silently lose - # visible text on the most reconnect-prone path. It - # clears at the next send, the pre-#937 fatal - # behavior. + # Terminal: finalize AND discard, exactly like the + # retry arm. Keeping the buffers bought nothing — the + # fatal path's _emit_state("error") drains and wipes + # them anyway on every SessionUIBase lane — and an + # overflow that send()'s compact-and-retry RECOVERS + # re-streams into buffers that would otherwise still + # hold the dead attempt's text, concatenating the two + # in the idle payload. self.ui.on_stream_end() + self.ui.on_stream_discarded() raise # fatal path otherwise unchanged # Delay from the PRE-increment attempt index — the same # convention as the three sibling ladders' range loops. @@ -7878,13 +7893,10 @@ class ChatSession: f"[stream died mid-response ({cause}) — retrying in " f"{delay:.0f}s ({attempt}/{self._MID_STREAM_RETRIES})]" ) - # Stop-then-start: a pre-first-token death leaves the - # spinner RUNNING (_stop_spinner_once never fired), and the - # CLI's on_thinking_start is not idempotent — it replaces - # the spinner object without stopping the old one, leaking - # its render thread (the compact-retry arm guards the same - # way). - self.ui.on_thinking_stop() + # A pre-first-token death leaves the spinner RUNNING + # (_stop_spinner_once never fired) — on_thinking_start is + # idempotent at the callee (TerminalUI stops a live spinner + # before replacing it), so no stop-first dance here. self.ui.on_thinking_start() try: self._backoff_or_cancelled(delay, my_generation) @@ -7909,9 +7921,12 @@ class ChatSession: # client-identity check alone would re-issue the # old model's fold against the new model. msgs = self._prepare_wire_messages(self._full_messages()) - self._active_wire_msgs = msgs try: stream = self._create_stream_with_retry(msgs) + # Refresh the frame-local gate identity: the + # re-create may have walked to a different + # provider (fallback, rebind). + live_provider = self._active_stream_provider or self._provider except Exception as recreate_exc: if _is_ctx_overflow(recreate_exc): # Deterministic — surface as ITSELF so send()'s diff --git a/turnstone/core/session_ui_base.py b/turnstone/core/session_ui_base.py index 745496a3..494c42bd 100644 --- a/turnstone/core/session_ui_base.py +++ b/turnstone/core/session_ui_base.py @@ -3194,9 +3194,11 @@ class SessionUIBase: payload drains — which :meth:`on_turn_committed` deliberately does NOT clear (earlier segments of a tool-looping turn must survive commits). Truncating back to the segment watermark removes exactly - the dead attempt's contribution; the pending token batch is - dropped, not flushed (its text was never displayed and must not - be). + the dead attempt's contribution. The pending-batch reset is + DEFENSIVE: in the shipped sequence the preceding ``stream_end`` + already flushed (and broadcast) any pending tail, so the reset + no-ops — it matters only for a caller that discards without + finalizing first, whose tail genuinely was never displayed. """ with self._ws_lock: self._reset_pending_locked()