fix(streaming): delete the retry window's shared slots and gate the send epilogue

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.
2026-08-12 23:12:23 -06:00 · 2026-08-04 03:09:48 -07:00
parent 476cce2e58
commit 961a2017dc
8 changed files with 183 additions and 86 deletions
@@ -39,6 +39,9 @@ class NullUI:
    def on_turn_committed(self):
        pass

+    def on_stream_discarded(self):
+        pass
+
    def on_thinking_start(self):
        pass

@@ -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 = [
@@ -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):
@@ -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",))

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

@@ -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


@@ -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 ``<think>`` into the
            # prompt (only ``</think>`` reaches ``content``). See ``_TITLE_*``.
            stripped = self._strip_reasoning(raw)
-            for _close in ("</think>", "</reasoning>"):
+            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 <think>/<reasoning> tags and their content."""
-        for open_t, close_t in [
-            ("<think>", "</think>"),
-            ("<reasoning>", "</reasoning>"),
-        ]:
+        """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
@@ -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()