From 7776cc0c2fc75d58fe1af97e86189590bff1fa2f Mon Sep 17 00:00:00 2001 From: Patrick Buckley Date: Tue, 4 Aug 2026 04:48:12 -0700 Subject: [PATCH] fix(streaming): probe on_stream_discarded for pre-existing UIs and format the hoisted fake MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PR feedback round: - on_stream_discarded now follows on_compaction's compat pattern for a hook added after UIs exist in the wild: the protocol member carries a REAL no-op default (an explicit subclass inherits a correct implementation — a UI without server-side turn buffers has nothing to truncate), and both call sites route through a getattr probe, so a duck-typed UI predating the hook degrades to no-truncate instead of raising an AttributeError from the very arm that is handling a stream death — which would replace the wire failure with the attribute error in the retry gate. Pinned with a hook-less-UI retry test. - tests/_session_helpers.py gains the formatting pass the RecordingUI hoist bypassed (the CI lint failure). --- tests/_session_helpers.py | 1 + tests/test_midstream_retry.py | 24 ++++++++++++++++++++++++ turnstone/core/session.py | 33 ++++++++++++++++++++++++++++++--- 3 files changed, 55 insertions(+), 3 deletions(-) diff --git a/tests/_session_helpers.py b/tests/_session_helpers.py index 5c4b537c..19965323 100644 --- a/tests/_session_helpers.py +++ b/tests/_session_helpers.py @@ -348,6 +348,7 @@ def as_stream(result: Any) -> list[StreamChunk]: ) ] + class RecordingUI: """UI adapter recording the ordered event stream ``send()`` emits.""" diff --git a/tests/test_midstream_retry.py b/tests/test_midstream_retry.py index c6a3121c..ec8d45b0 100644 --- a/tests/test_midstream_retry.py +++ b/tests/test_midstream_retry.py @@ -597,6 +597,30 @@ class TestMidStreamRetry: assistant = _assistant_msgs(session) assert assistant and assistant[-1]["content"].startswith("a dead attempt") + def test_pre937_ui_without_discard_hook_survives_retry(self, tmp_db): + # A duck-typed UI predating on_stream_discarded must degrade to + # "no server-buffer truncate", not crash the retry arm with an + # AttributeError that replaces the stream death being handled. + class _Pre937UI(RecordingUI): + # property() with no getter raises AttributeError on access — + # simulating the hook's absence on an inheriting fake. + on_stream_discarded = property() + + ui = _Pre937UI() + session = _make_session(ui) + streams = [ + _dying_stream("x", exc=httpx.ReadError("wire died")), + _good_stream("ok"), + ] + with ( + patch.object(session, "_create_stream_with_retry", side_effect=streams), + patch.object(session, "_full_messages", return_value=[]), + ): + session.send("test") + + assert _assistant_msgs(session)[-1]["content"] == "ok" + assert ("state", "idle") in ui.events + 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 diff --git a/turnstone/core/session.py b/turnstone/core/session.py index 8ce29354..784d3c1f 100644 --- a/turnstone/core/session.py +++ b/turnstone/core/session.py @@ -1256,7 +1256,19 @@ class SessionUI(Protocol): def on_turn_start(self) -> None: ... def on_turn_committed(self) -> None: ... - def on_stream_discarded(self) -> None: ... + def on_stream_discarded(self) -> None: + """Drop a dead stream attempt's text from any server-side buffers. + + A REAL no-op default, not a bare ``...`` stub, for the same reason + ``on_compaction`` below carries one: an explicit pre-#937 + ``SessionUI`` subclass inherits this as its implementation, and + for a discard the no-op IS correct — a UI without server-side + turn buffers has nothing to truncate. ``SessionUIBase`` overrides + with the real truncation; duck-typed UIs that never subclass are + covered by the getattr probe in ``_ui_stream_discarded``. + """ + return + def on_thinking_start(self) -> None: ... def on_thinking_stop(self) -> None: ... def on_reasoning_token(self, text: str) -> None: ... @@ -7730,6 +7742,21 @@ class ChatSession: text = text[:start] + text[end + len(close_t) :] if end != -1 else text[:start] return text.strip() + def _ui_stream_discarded(self) -> None: + """Best-effort dead-segment discard across UI generations. + + Probed, never called directly: a pre-#937 duck-typed UI may lack + the hook, and an AttributeError raised from the retry/terminal + arms would REPLACE the stream death being handled — the retry + gate would then judge the attribute error instead of the wire + failure. Missing hook degrades to no server-buffer truncate, + which is correct for UIs without server-side buffers (the same + compat posture as ``_compaction_event``'s probe). + """ + discard = getattr(self.ui, "on_stream_discarded", None) + if discard is not None: + discard() + def _stream_response( self, msgs: list[dict[str, Any]], my_generation: int = 0 ) -> dict[str, Any]: @@ -7854,7 +7881,7 @@ class ChatSession: # hold the dead attempt's text, concatenating the two # in the idle payload. self.ui.on_stream_end() - self.ui.on_stream_discarded() + self._ui_stream_discarded() raise # fatal path otherwise unchanged # Delay from the PRE-increment attempt index — the same # convention as the three sibling ladders' range loops. @@ -7901,7 +7928,7 @@ class ChatSession: # source) and reset the inflight snapshot BEFORE any # retried token lands, or every consumer appends the # retried text onto the dead attempt's. - self.ui.on_stream_discarded() + self._ui_stream_discarded() # Spinner for the recreate+TTFT window, and the fresh # segment watermark — AFTER the truncate, so a later # discard cannot resurrect this dead segment. A