fix(session): scope the arm-duty gate the way the rest of the file scopes generations

The consumer's arm hook and cancel-partial recorder compared generations
with a bare !=, while the ref that fires them treats generation 0 as
UNSCOPED — so for a direct seam caller the ref armed and fired the hook
and the hook refused to act. On a session whose generation had ever been
claimed, that left the previous turn's usage in place as this turn's
estimate and dropped the serving lane's health success. Both now ask the
consumer's own _superseded(), which mirrors the ref's predicate, so the
two halves of one decision cannot disagree.

The which-errors-speak-for-the-backend policy gets one spelling
(_speaks_for_backend over _NON_BACKEND_ERRORS) instead of a matching
isinstance in each walk arm, and the length arm stops calling
finalize_provider_blocks over an empty list only to discard the result.

Tests: the fourteen hand-rolled FakeChunk dataclasses in the cancel suite
are replaced by the real StreamChunk its sibling suites already use, so
the fakes cannot drift from the shape production emits.
This commit is contained in:
Patrick Buckley
2026-08-05 23:37:29 -07:00
parent 5de54147e1
commit 4dd92d150b
3 changed files with 92 additions and 183 deletions
+53 -165
View File
@@ -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()
+3 -3
View File
@@ -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).
+34 -13
View File
@@ -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
native = None
if old_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
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),