diff --git a/tests/_session_helpers.py b/tests/_session_helpers.py index 16446649..d9e224a2 100644 --- a/tests/_session_helpers.py +++ b/tests/_session_helpers.py @@ -547,7 +547,10 @@ def arm_session( provider.provider_name = name provider.get_capabilities.return_value = ModelCapabilities() provider.retryable_error_names = retryable - provider._armed_handle = ArmedHandle() + # One handle PER CREATE (the real adapters' rule): `handles` records + # them all, `_armed_handle` is the latest. + provider._armed_handle = None + provider.handles = [] remaining = list(streams) def _create(**kwargs: Any): @@ -557,7 +560,10 @@ def arm_session( raise nxt ref = kwargs.get("cancel_ref") if ref is not None: - ref.append(provider._armed_handle) + handle = ArmedHandle() + provider.handles.append(handle) + provider._armed_handle = handle + ref.append(handle) return iter(nxt) if not hasattr(nxt, "__next__") else nxt provider.create_streaming = MagicMock(side_effect=_create) diff --git a/tests/test_cancel.py b/tests/test_cancel.py index f5062cad..ab705982 100644 --- a/tests/test_cancel.py +++ b/tests/test_cancel.py @@ -1284,3 +1284,39 @@ class TestNeverArmedStopLeavesNoRow: assert any("cancelled" in i.lower() for i in ui.infos) assistant = [m for m in dicts_from_turns(session.messages) if m["role"] == "assistant"] assert assistant == [] + + +class TestOrphanArmDutiesGate: + def test_superseded_arrival_in_toctou_window_fires_no_duties(self, tmp_db): + """The ref's supersession read and the arm hook are two lockless + steps; a force-cancel can claim a new generation between them. + The duties self-gate on the generation, so even an append whose + supersession read went stale cannot null the successor's usage + slots or record health for the abandoned lane.""" + from turnstone.core.session import _StreamTurnConsumer + + session = _make_session() + consumer = _StreamTurnConsumer(session, my_generation=1) + tracker = MagicMock() + + class _StaleReadRef(_CancelRef): + # The stale supersession read the accepted bytecode-width + # window produces — the hook itself must hold the line. + def _superseded(self) -> bool: + return False + + ref = _StaleReadRef(session, 1, on_first_append=consumer.on_stream_armed) + consumer.begin_attempt(ref, tracker, MagicMock()) + + session._generation = 1 + sentinel = {"prompt_tokens": 7} + session._last_usage = sentinel + session._assistant_pending_tokens = 42 + + # The force-cancel claims a newer generation before the append. + session._generation = 2 + ref.append(MagicMock()) + + assert session._last_usage is sentinel + assert session._assistant_pending_tokens == 42 + tracker.record_success.assert_not_called() diff --git a/tests/test_midstream_retry.py b/tests/test_midstream_retry.py index 81c6cdc7..d6326d0a 100644 --- a/tests/test_midstream_retry.py +++ b/tests/test_midstream_retry.py @@ -887,9 +887,44 @@ class TestFallbackFailureRedaction: from turnstone.core.session import _StreamTurnConsumer consumer = _StreamTurnConsumer(session, 0) - result = session._try_fallback_lane("fb", consumer, lambda w: w, 0) + result = session._try_fallback_lane("fb", consumer, lambda w, lane: w, 0) assert result is None infos = ui.of("info") assert any("Fallback fb also failed: ConnectError" in i for i in infos) assert not any("SECRETKEY" in i for i in infos) + + +class TestPrepareWireLaneCaps: + """The per-attempt wire prep folds with the SERVING lane's + capabilities — a fallback whose template rejects mid-conversation + system roles gets the folded shape even when the primary keeps them + inline.""" + + def test_caps_override_controls_fold_posture(self, tmp_db): + from turnstone.core.providers import ModelCapabilities + + session = _make_session(RecordingUI()) + msgs = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "yo"}, + {"role": "system", "content": "nudge", "_source": "idle_nudge"}, + {"role": "user", "content": "next"}, + ] + native = session._prepare_wire_messages( + list(msgs), caps=ModelCapabilities(supports_mid_conversation_system=True) + ) + folded = session._prepare_wire_messages( + list(msgs), caps=ModelCapabilities(supports_mid_conversation_system=False) + ) + assert any(m["role"] == "system" for m in native[1:]) + assert not any(m["role"] == "system" for m in folded[1:]) + + def test_stream_prepare_passes_serving_lane_caps(self, tmp_db): + session = _make_session(RecordingUI()) + arm_session(session, _good_stream("ok")) + with patch.object( + session, "_prepare_wire_messages", wraps=session._prepare_wire_messages + ) as prep: + session.send("test") + assert prep.call_args.kwargs["caps"] is session._get_capabilities() diff --git a/tests/test_model_provider_obo.py b/tests/test_model_provider_obo.py index e1a2a4dc..7016842c 100644 --- a/tests/test_model_provider_obo.py +++ b/tests/test_model_provider_obo.py @@ -1577,7 +1577,7 @@ class TestModelOboToken: # build is the one place the alias enters. sess = MagicMock() sess._model_alias = "oboagent" - ChatSession._model_turn_with_fallback(sess, MagicMock(), lambda wire: wire) + ChatSession._model_turn_with_fallback(sess, MagicMock(), lambda wire, lane: wire) sess._build_main_lane.assert_called_once() assert sess._build_main_lane.call_args.kwargs["alias"] == "oboagent" @@ -1590,7 +1590,7 @@ class TestModelOboToken: sess._get_health_tracker.return_value = tracker with pytest.raises(BackendAuthUnavailableError): - ChatSession._model_turn_with_fallback(sess, MagicMock(), lambda wire: wire) + ChatSession._model_turn_with_fallback(sess, MagicMock(), lambda wire, lane: wire) sess._try_fallback_lane.assert_not_called() # An auth refusal is never reinterpreted as backend health. diff --git a/tests/test_tool_truncation.py b/tests/test_tool_truncation.py index 951c6da1..306517ea 100644 --- a/tests/test_tool_truncation.py +++ b/tests/test_tool_truncation.py @@ -200,7 +200,14 @@ class TestRemainingTokenBudget: class TestContextOverflowRecovery: """Test that context-length errors trigger compact-and-retry.""" - def test_openai_context_length_error_triggers_compact(self, session): + @pytest.mark.parametrize( + "overflow_text", + [ + pytest.param("maximum context length exceeded", id="openai"), + pytest.param("prompt is too long: 250000 tokens > 200000 maximum", id="anthropic"), + ], + ) + def test_context_overflow_triggers_compact(self, session, overflow_text): session.messages = turns_from_dicts([{"role": "user", "content": "hi"}]) session._msg_tokens = [1] @@ -210,7 +217,7 @@ class TestContextOverflowRecovery: nonlocal call_count call_count += 1 if call_count == 1: - raise Exception("maximum context length exceeded") + raise Exception(overflow_text) return make_result(content="ok") compact_mock = MagicMock() @@ -230,33 +237,6 @@ class TestContextOverflowRecovery: compact_mock.assert_called_once_with(auto=True, my_generation=session._generation) assert call_count == 2 - def test_anthropic_prompt_too_long_triggers_compact(self, session): - session.messages = turns_from_dicts([{"role": "user", "content": "hi"}]) - session._msg_tokens = [1] - - call_count = 0 - - def mock_stream_response(my_generation=0): - nonlocal call_count - call_count += 1 - if call_count == 1: - raise Exception("prompt is too long: 250000 tokens > 200000 maximum") - return make_result(content="ok") - - compact_mock = MagicMock() - with ( - patch.object(session, "_stream_response", side_effect=mock_stream_response), - patch.object(session, "_compact_messages", compact_mock), - patch.object(session, "_full_messages", return_value=[]), - patch.object(session, "_update_token_table"), - patch.object(session, "_print_status_line"), - patch.object(session, "_emit_state"), - patch("turnstone.core.session.save_message"), - ): - session.send("hello") - - compact_mock.assert_called_once_with(auto=True, my_generation=session._generation) - def test_non_context_error_propagates(self, session): session.messages = turns_from_dicts([{"role": "user", "content": "hi"}]) session._msg_tokens = [1] diff --git a/turnstone/core/model_turn.py b/turnstone/core/model_turn.py index 742d5376..2139f396 100644 --- a/turnstone/core/model_turn.py +++ b/turnstone/core/model_turn.py @@ -886,7 +886,7 @@ def model_turn( cancel_ref: list[Any] | None = None, backend_auth_token: str | None = None, deferred_names: frozenset[str] | None = None, - prepare_wire: Callable[[list[dict[str, Any]]], list[dict[str, Any]]] | None = None, + prepare_wire: Callable[[list[dict[str, Any]], ModelLane], list[dict[str, Any]]] | None = None, on_chunk: Callable[[StreamChunk], None] | None = None, ) -> ModelTurnResult: """Advance a trajectory by one model turn: lower, sample, re-ingest. @@ -962,7 +962,9 @@ def model_turn( reads the error at all, which is what keeps a Stop mid-summary off the red-error path. - *prepare_wire* is the caller's OWN deterministic lowering, composed + *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 @@ -1036,7 +1038,10 @@ def model_turn( ) if prepare_wire is not None: try: - wire = prepare_wire(wire) + # The serving lane rides along so caller lowering can be + # capability-correct per attempt — a fallback's fold posture + # is its own, not the primary's. + wire = prepare_wire(wire, lane) except Exception as prep_err: # A caller-data fault, never a backend signal — typed so the # retry and fallback ladders cannot treat it as one. diff --git a/turnstone/core/session.py b/turnstone/core/session.py index 5a93e281..aa90fde1 100644 --- a/turnstone/core/session.py +++ b/turnstone/core/session.py @@ -544,8 +544,16 @@ class _StreamTurnConsumer: stale-usage leak (a post-finish blip can lose the trailing usage chunk, and a stale completion count would be recycled as the next turn's estimate) and the reconnect status-bar blackout. + + Generation-gated: the ref's supersession read and this hook are + 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. """ s = self._session + if s._generation != self._my_generation: + return s._last_usage = None s._assistant_pending_tokens = 0 if self.tracker: @@ -5492,6 +5500,8 @@ class ChatSession: def _prepare_wire_messages( self, messages: list[dict[str, Any]], + *, + caps: ModelCapabilities | None = None, ) -> list[dict[str, Any]]: """Return a transient copy of *messages* prepared for the provider wire. @@ -5543,11 +5553,16 @@ class ChatSession: messages = self._inject_sender_labels(messages) folded = messages if self._provider is not None: + # *caps* is the SERVING lane's capabilities when the streaming + # wrapper prepares per attempt — a fallback whose template + # rejects mid-conversation system roles must get the folded + # shape even when the primary keeps them inline. Callers + # without a lane in hand (the token-table re-fold) default to + # the primary binding. + fold_caps = caps if caps is not None else self._get_capabilities() folded = fold_system_turns( messages, - supports_mid_conversation_system=( - self._get_capabilities().supports_mid_conversation_system - ), + supports_mid_conversation_system=fold_caps.supports_mid_conversation_system, nonce=self._envelope_nonce, ) dropped = drop_empty_user_turns(folded) @@ -6212,7 +6227,7 @@ class ChatSession: def _model_turn_with_fallback( self, consumer: _StreamTurnConsumer, - prepare_wire: Callable[[list[dict[str, Any]]], list[dict[str, Any]]], + prepare_wire: Callable[[list[dict[str, Any]], ModelLane], list[dict[str, Any]]], my_generation: int = 0, ) -> ModelTurnResult: """Run one plant call with lane-swap fallback: one ``model_turn`` @@ -6284,7 +6299,7 @@ class ChatSession: self, alias: str, consumer: _StreamTurnConsumer, - prepare_wire: Callable[[list[dict[str, Any]]], list[dict[str, Any]]], + prepare_wire: Callable[[list[dict[str, Any]], ModelLane], list[dict[str, Any]]], my_generation: int, ) -> ModelTurnResult | None: """Attempt a single fallback lane. Returns the result or ``None``. @@ -6367,7 +6382,7 @@ class ChatSession: lane: ModelLane, tracker: BackendHealthTracker | None, consumer: _StreamTurnConsumer, - prepare_wire: Callable[[list[dict[str, Any]]], list[dict[str, Any]]], + prepare_wire: Callable[[list[dict[str, Any]], ModelLane], list[dict[str, Any]]], my_generation: int = 0, ) -> ModelTurnResult: """One lane's creation ladder around ``model_turn``. @@ -8227,7 +8242,7 @@ class ChatSession: debug_printed = False - def _prepare(lowered: list[dict[str, Any]]) -> list[dict[str, Any]]: + def _prepare(lowered: list[dict[str, Any]], lane: ModelLane) -> list[dict[str, Any]]: """The main loop's ``prepare_wire``: system prepend + the session lowering passes, plus the debug request dump behind a once-per-invocation latch, so re-issues and fallback lanes @@ -8236,7 +8251,9 @@ class ChatSession: post-compaction the wire CHANGED, and the re-prepared dump is the one that diagnoses the recovery (a named #832 delta).""" nonlocal debug_printed - wire = self._prepare_wire_messages([*self.system_messages, *lowered]) + wire = self._prepare_wire_messages( + [*self.system_messages, *lowered], caps=lane.capabilities + ) if self.debug and not debug_printed: debug_printed = True self._debug_print_request(wire)