diff --git a/tests/test_cancel.py b/tests/test_cancel.py index ab705982..c4ade926 100644 --- a/tests/test_cancel.py +++ b/tests/test_cancel.py @@ -4,12 +4,13 @@ import contextlib import json import threading import time -from dataclasses import dataclass, field +from dataclasses import dataclass from unittest.mock import MagicMock, patch import pytest from tests._session_helpers import arm_session, make_session +from turnstone.core.providers import StreamChunk from turnstone.core.session import ( GenerationCancelled, _CancelRef, @@ -137,17 +138,7 @@ class TestCancelEvent: session = _make_session(ui=ui) session.cancel() # Set stale flag - @dataclass - class FakeChunk: - content_delta: str = "" - reasoning_delta: str = "" - tool_call_deltas: list = field(default_factory=list) - usage: None = None - finish_reason: str = "stop" - info_delta: str = "" - provider_blocks: list = field(default_factory=list) - - fake_stream = iter([FakeChunk(content_delta="Hello", finish_reason="stop")]) + fake_stream = iter([StreamChunk(content_delta="Hello", finish_reason="stop")]) arm_session(session, fake_stream) session.send("test") @@ -164,22 +155,12 @@ class TestCancelDuringStreaming: ui = NullUI() session = _make_session(ui=ui) - @dataclass - class FakeChunk: - content_delta: str = "" - reasoning_delta: str = "" - tool_call_deltas: list = field(default_factory=list) - usage: None = None - finish_reason: str = "" - info_delta: str = "" - provider_blocks: list = field(default_factory=list) - def cancelling_stream(): """Yield a few chunks then cancel.""" - yield FakeChunk(content_delta="Hello ") - yield FakeChunk(content_delta="world") + yield StreamChunk(content_delta="Hello ") + yield StreamChunk(content_delta="world") session.cancel() - yield FakeChunk(content_delta=" — this should not appear") + yield StreamChunk(content_delta=" — this should not appear") arm_session(session, cancelling_stream()) session.send("test") @@ -212,16 +193,6 @@ class TestCancelDuringToolExecution: ui = NullUI() session = _make_session(ui=ui) - @dataclass - class FakeChunk: - content_delta: str = "" - reasoning_delta: str = "" - tool_call_deltas: list = field(default_factory=list) - usage: None = None - finish_reason: str = "" - info_delta: str = "" - provider_blocks: list = field(default_factory=list) - @dataclass class FakeToolDelta: index: int = 0 @@ -231,11 +202,11 @@ class TestCancelDuringToolExecution: # First call: return content with a tool call def stream_with_tool(): - yield FakeChunk( + yield StreamChunk( tool_call_deltas=[FakeToolDelta(index=0, id="tc_1", name="bash")], finish_reason="", ) - yield FakeChunk( + yield StreamChunk( tool_call_deltas=[FakeToolDelta(index=0, arguments_delta='{"command":"echo hi"}')], finish_reason="tool_calls", ) @@ -274,17 +245,7 @@ class TestCancelWhenIdle: session.cancel() # Next send should work normally (cancel cleared at start) - @dataclass - class FakeChunk: - content_delta: str = "" - reasoning_delta: str = "" - tool_call_deltas: list = field(default_factory=list) - usage: None = None - finish_reason: str = "stop" - info_delta: str = "" - provider_blocks: list = field(default_factory=list) - - fake_stream = iter([FakeChunk(content_delta="ok", finish_reason="stop")]) + fake_stream = iter([StreamChunk(content_delta="ok", finish_reason="stop")]) arm_session(session, fake_stream) session.send("hello") @@ -301,23 +262,13 @@ class TestCancelThreadSafety: ui = NullUI() session = _make_session(ui=ui) - @dataclass - class FakeChunk: - content_delta: str = "" - reasoning_delta: str = "" - tool_call_deltas: list = field(default_factory=list) - usage: None = None - finish_reason: str = "" - info_delta: str = "" - provider_blocks: list = field(default_factory=list) - barrier = threading.Event() def slow_stream(): - yield FakeChunk(content_delta="Start") + yield StreamChunk(content_delta="Start") barrier.set() # Signal that streaming has started time.sleep(2) # Simulate slow streaming - yield FakeChunk(content_delta=" end", finish_reason="stop") + yield StreamChunk(content_delta=" end", finish_reason="stop") arm_session(session, slow_stream()) # Run send() in a thread @@ -374,16 +325,6 @@ class TestStreamFlushBeforeToolCalls: ui = TrackingUI() session = _make_session(ui=ui) - @dataclass - class FakeChunk: - content_delta: str = "" - reasoning_delta: str = "" - tool_call_deltas: list = field(default_factory=list) - usage: None = None - finish_reason: str = "" - info_delta: str = "" - provider_blocks: list = field(default_factory=list) - @dataclass class FakeToolDelta: index: int = 0 @@ -395,11 +336,11 @@ class TestStreamFlushBeforeToolCalls: # Content long enough to leave chars in the tag-scan carry # buffer (ThinkTagSplitter retains the last MAX_TAG_LEN = 12 # chars until a flush) - yield FakeChunk(content_delta="Hello world, this is a test message") - yield FakeChunk( + yield StreamChunk(content_delta="Hello world, this is a test message") + yield StreamChunk( tool_call_deltas=[FakeToolDelta(index=0, id="tc_1", name="bash")], ) - yield FakeChunk( + yield StreamChunk( tool_call_deltas=[FakeToolDelta(index=0, arguments_delta='{"command":"echo hi"}')], finish_reason="tool_calls", ) @@ -411,7 +352,7 @@ class TestStreamFlushBeforeToolCalls: arm_session( session, stream_content_then_tool(), - iter([FakeChunk(finish_reason="stop")]), + iter([StreamChunk(finish_reason="stop")]), ) with ( # Prevent real tool execution (e.g., bash) during this test. @@ -465,16 +406,6 @@ class TestStreamAbort: ui = NullUI() session = _make_session(ui=ui) - @dataclass - class FakeChunk: - content_delta: str = "" - reasoning_delta: str = "" - tool_call_deltas: list = field(default_factory=list) - usage: None = None - finish_reason: str = "stop" - info_delta: str = "" - provider_blocks: list = field(default_factory=list) - seen: dict = {} def observing_stream(): @@ -482,7 +413,7 @@ class TestStreamAbort: # registered (append happens inside create_streaming, before # the iterator is handed back). seen["handle_at_first_chunk"] = session._cancel_stream - yield FakeChunk(content_delta="hi", finish_reason="stop") + yield StreamChunk(content_delta="hi", finish_reason="stop") provider = arm_session(session, observing_stream()) session.send("test") @@ -497,18 +428,8 @@ class TestStreamAbort: ui = NullUI() session = _make_session(ui=ui) - @dataclass - class FakeChunk: - content_delta: str = "" - reasoning_delta: str = "" - tool_call_deltas: list = field(default_factory=list) - usage: None = None - finish_reason: str = "" - info_delta: str = "" - provider_blocks: list = field(default_factory=list) - def stream_that_errors(): - yield FakeChunk(content_delta="Hello") + yield StreamChunk(content_delta="Hello") session._cancel_event.set() raise ConnectionError("stream closed") @@ -532,18 +453,8 @@ class TestStreamAbort: ui = NullUI() session = _make_session(ui=ui) - @dataclass - class FakeChunk: - content_delta: str = "" - reasoning_delta: str = "" - tool_call_deltas: list = field(default_factory=list) - usage: None = None - finish_reason: str = "" - info_delta: str = "" - provider_blocks: list = field(default_factory=list) - def stream_that_errors(): - yield FakeChunk(content_delta="Hello") + yield StreamChunk(content_delta="Hello") raise ValueError("unexpected error") arm_session(session, stream_that_errors()) @@ -633,17 +544,7 @@ class TestCancelRef: ui = NullUI() session = _make_session(ui=ui) - @dataclass - class FakeChunk: - content_delta: str = "" - reasoning_delta: str = "" - tool_call_deltas: list = field(default_factory=list) - usage: None = None - finish_reason: str = "stop" - info_delta: str = "" - provider_blocks: list = field(default_factory=list) - - arm_session(session, iter([FakeChunk(content_delta="hi", finish_reason="stop")])) + arm_session(session, iter([StreamChunk(content_delta="hi", finish_reason="stop")])) session.send("test") # During the stream the armed handle was registered; after send() @@ -694,18 +595,8 @@ class TestForceCancelOrphanNoReissue: ui = NullUI() session = _make_session(ui=ui) - @dataclass - class FakeChunk: - content_delta: str = "" - reasoning_delta: str = "" - tool_call_deltas: list = field(default_factory=list) - usage: None = None - finish_reason: str = "" - info_delta: str = "" - provider_blocks: list = field(default_factory=list) - def dying_orphan_stream(): - yield FakeChunk(content_delta="old ") + yield StreamChunk(content_delta="old ") # Force-cancel: a successor claims the generation (bumped # counter + fresh UNSET event) while this stream is mid-body. session._claim_generation() @@ -760,19 +651,9 @@ class TestForceCancelGeneration: ui = NullUI() session = _make_session(ui=ui) - @dataclass - class FakeChunk: - content_delta: str = "" - reasoning_delta: str = "" - tool_call_deltas: list = field(default_factory=list) - usage: None = None - finish_reason: str = "stop" - info_delta: str = "" - provider_blocks: list = field(default_factory=list) - original_event = session._cancel_event - arm_session(session, iter([FakeChunk(content_delta="hi", finish_reason="stop")])) + arm_session(session, iter([StreamChunk(content_delta="hi", finish_reason="stop")])) session.send("test") # After send() completes, _cancel_event should be a NEW Event @@ -790,24 +671,14 @@ class TestForceCancelThreaded: ui = NullUI() session = _make_session(ui=ui) - @dataclass - class FakeChunk: - content_delta: str = "" - reasoning_delta: str = "" - tool_call_deltas: list = field(default_factory=list) - usage: None = None - finish_reason: str = "" - info_delta: str = "" - provider_blocks: list = field(default_factory=list) - barrier = threading.Event() old_done = threading.Event() def slow_stream(): - yield FakeChunk(content_delta="Old content") + yield StreamChunk(content_delta="Old content") barrier.set() # signal: first chunk delivered time.sleep(2) # simulate stuck stream - yield FakeChunk(content_delta=" more", finish_reason="stop") + yield StreamChunk(content_delta=" more", finish_reason="stop") # Start generation 1 (will get stuck) arm_session(session, slow_stream()) @@ -843,23 +714,13 @@ class TestForceCancelThreaded: ui = NullUI() session = _make_session(ui=ui) - @dataclass - class FakeChunk: - content_delta: str = "" - reasoning_delta: str = "" - tool_call_deltas: list = field(default_factory=list) - usage: None = None - finish_reason: str = "stop" - info_delta: str = "" - provider_blocks: list = field(default_factory=list) - barrier = threading.Event() def stuck_stream(): - yield FakeChunk(content_delta="stuck") + yield StreamChunk(content_delta="stuck", finish_reason="stop") barrier.set() time.sleep(2) - yield FakeChunk(content_delta=" end", finish_reason="stop") + yield StreamChunk(content_delta=" end", finish_reason="stop") # Start stuck generation arm_session(session, stuck_stream()) @@ -871,7 +732,9 @@ class TestForceCancelThreaded: session.cancel() # New generation should work - arm_session(session, iter([FakeChunk(content_delta="Fresh response")])) + arm_session( + session, iter([StreamChunk(content_delta="Fresh response", finish_reason="stop")]) + ) session.send("new message") # The new generation should have completed successfully @@ -1320,3 +1183,28 @@ class TestOrphanArmDutiesGate: assert session._last_usage is sentinel assert session._assistant_pending_tokens == 42 tracker.record_success.assert_not_called() + + def test_unscoped_generation_still_runs_arm_duties(self, tmp_db): + """Generation 0 means UNSCOPED, the convention ``_check_cancelled`` + and ``_CancelRef._superseded`` share: the ref fires the hook for a + direct seam caller, so the hook's own gate must not refuse it. A + bare ``!=`` compare skips the duties on any session whose + generation was ever claimed, silently recycling the previous + turn's usage into this turn's estimate and losing the lane's + health success.""" + from turnstone.core.session import _StreamTurnConsumer + + session = _make_session() + session._generation = 3 # a prior send claimed generations + consumer = _StreamTurnConsumer(session, my_generation=0) + tracker = MagicMock() + ref = _CancelRef(session, 0, on_first_append=consumer.on_stream_armed) + consumer.begin_attempt(ref, tracker, MagicMock()) + + session._last_usage = {"prompt_tokens": 99} + session._assistant_pending_tokens = 42 + ref.append(MagicMock()) + + assert session._last_usage is None + assert session._assistant_pending_tokens == 0 + tracker.record_success.assert_called_once() diff --git a/turnstone/core/model_turn.py b/turnstone/core/model_turn.py index a2d63be8..f5364a08 100644 --- a/turnstone/core/model_turn.py +++ b/turnstone/core/model_turn.py @@ -968,9 +968,9 @@ def model_turn( *prepare_wire* is the caller's OWN deterministic lowering, called with the serving lane so per-lane capability posture is available, composed after the seam passes and before the Phase-5 attach: the - main loop's - system-message prepend, sender labels, capability-sensitive - system-turn fold, empty-user drop, and orphan repair live here, each + main loop's system-message prepend, sender labels, + capability-sensitive system-turn fold, empty-user drop, and orphan + repair live here, each a ``lowering.py``-composed pass. It must be pure lowering — no learned selection, no provider calls, no side effects (the session's debug request dump, a read-only latch, is the tolerated exception). diff --git a/turnstone/core/session.py b/turnstone/core/session.py index 4e0edeff..3e353778 100644 --- a/turnstone/core/session.py +++ b/turnstone/core/session.py @@ -537,6 +537,15 @@ class _StreamTurnConsumer: duties (health success, usage-slot resets).""" return (self.ref is not None and self.ref.armed) or self._saw_chunk + def _superseded(self) -> bool: + """Whether a newer generation has claimed this session — the + consumer's copy of :meth:`_CancelRef._superseded`, scoped the same + way (generation 0 is unscoped, so a direct seam caller is never + superseded). The two must agree: the ref decides whether to fire + the arm hook, and the hook's own gate decides whether to act.""" + gen = self._my_generation + return bool(gen and self._session._generation != gen) + def on_stream_armed(self) -> None: """`_CancelRef.on_first_append` — the request-accepted instant. @@ -550,10 +559,13 @@ class _StreamTurnConsumer: two lockless steps, so a force-cancel can claim a new generation between them — an orphan's late-arriving registration must not null the SUCCESSOR's usage slots or record health for an - abandoned lane. + abandoned lane. Scoping matches the ref's own ``_superseded`` + and ``_check_cancelled``: generation 0 is UNSCOPED (a direct + seam caller), and the ref fires the hook for it, so the gate + must not refuse it. """ s = self._session - if s._generation != self._my_generation: + if self._superseded(): return s._last_usage = None s._assistant_pending_tokens = 0 @@ -706,7 +718,7 @@ class _StreamTurnConsumer: ``tool_calls`` and the native lane are DELIBERATELY omitted — incomplete calls would orphan their results, and the marker-as-message contract needs plain content.""" - if self._session._generation != self._my_generation: + if self._superseded(): return content = self.partial_content() self._flush_terminal_carries() @@ -1827,6 +1839,18 @@ _SELF_SURFACING_ERRORS: tuple[type[Exception], ...] = ( WirePreparationError, ) +# Creation failures that say nothing about the BACKEND: the caller's own +# lowering raised. Recording them would paint a cluster-wide outage over +# one session's malformed history. +_NON_BACKEND_ERRORS: tuple[type[Exception], ...] = (WirePreparationError,) + + +def _speaks_for_backend(err: BaseException) -> bool: + """Whether a creation failure is a health signal for its lane — THE + predicate both walk arms use, so primary and fallback can never + classify the same error differently.""" + return not isinstance(err, _NON_BACKEND_ERRORS) + def _mint_refusal_cause( prefix: str, @@ -6276,7 +6300,7 @@ class ChatSession: # health — but it DOES enter the walk: prepare runs per lane # (fold posture follows lane.capabilities), so another lane's # posture may serve a turn the primary's could not. - if tracker and not isinstance(primary_err, WirePreparationError): + if tracker and _speaks_for_backend(primary_err): tracker.record_failure() if not self._registry or not self._registry.fallback: raise @@ -6352,7 +6376,7 @@ class ChatSession: # This lane's wire-preparation fault is not its health signal, # and it must not abort the walk: prepare is lane-variant, so # the next alias may still serve the turn. - if fb_tracker and not isinstance(fb_err, WirePreparationError): + if fb_tracker and _speaks_for_backend(fb_err): fb_tracker.record_failure() # Class name only in the UI line: a ConnectError's text can # carry a credential-bearing base_url, and this string lands in @@ -8499,16 +8523,13 @@ class ChatSession: count=len(dropped), ) old_native = result.turn.native - blocks = finalize_provider_blocks( - list(old_native.blocks) if old_native else [], - [], - has_tool_calls=False, - ) - native = ( - ProviderNative(producer=old_native.producer, blocks=tuple(blocks)) - if blocks and old_native - else None - ) + native = None + if old_native: + blocks = finalize_provider_blocks( + list(old_native.blocks), [], has_tool_calls=False + ) + if blocks: + native = ProviderNative(producer=old_native.producer, blocks=tuple(blocks)) result = dataclasses.replace( result, turn=Turn.assistant(result.turn.text, native=native),