mirror of
https://github.com/turnstonelabs/turnstone.git
synced 2026-08-12 23:12:23 -06:00
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:
+53
-165
@@ -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()
|
||||
|
||||
@@ -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
@@ -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),
|
||||
|
||||
Reference in New Issue
Block a user