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()