From 6345ccd203b17f0de0ff67f095b5dce6f0acd469 Mon Sep 17 00:00:00 2001 From: HexLab98 Date: Fri, 17 Jul 2026 07:36:48 +0700 Subject: [PATCH 1/2] fix(streaming): fence superseded streams out of the delta sink (single-writer) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When the stale-stream detector reconnects past a stream whose socket abort raced (the close never actually stopped the old worker), the superseded stream and the retry's stream both write deltas into the same turn. The persisted transcript is then two coherent responses interleaved token-by-token — de-interleaving the stored text by alternation yields two complete, independent answers to the same prompt, which is a dual-writer race in the harness, not a model/context failure (#65991). The interrupt path already positively cancels before force-closing (#6600), but the stale-kill path relies only on the socket abort, and nothing fenced late chunks from a superseded stream out of the shared delta sink. Enforce a single-writer invariant on the sink itself, guarded by attempt id rather than only socket state: every streaming attempt (chat_completions, anthropic_messages, and bedrock paths) claims a monotonic writer token before it begins consuming its stream. A newer claim supersedes any older one, so the consume loop bails the instant it is superseded and _fire_stream_delta / _fire_reasoning_delta / _record_streamed_assistant_text drop chunks from a stale writer. The token is stored per-thread, so a thread that never claimed (a non-streaming delta caller) is never fenced — the guard can only ever drop a superseded stream, never the single legitimate writer. Discards are counted and logged sparsely so a real provider problem stays visible instead of being silently swallowed. --- agent/agent_init.py | 15 ++++++ agent/chat_completion_helpers.py | 34 ++++++++++++ run_agent.py | 91 ++++++++++++++++++++++++++++++++ 3 files changed, 140 insertions(+) diff --git a/agent/agent_init.py b/agent/agent_init.py index e1d219c62c632..8249c258c974f 100644 --- a/agent/agent_init.py +++ b/agent/agent_init.py @@ -744,6 +744,21 @@ def init_agent( # assistant message. agent._current_streamed_assistant_text = "" + # Single-writer guard for the streaming delta sink (#65991). A stale/ + # superseded stream (e.g. one the stale-stream detector reconnected past, + # whose socket abort raced and never actually stopped the old worker) must + # NOT keep writing tokens into the turn alongside the retry's stream — + # otherwise two coherent responses interleave token-by-token into one + # transcript. Every streaming attempt claims a monotonic writer token; the + # delta sink drops chunks whose calling thread holds a stale token. The + # threading.local means threads that never claimed (non-streaming callers) + # are never fenced, so the guard can only ever drop a superseded stream, + # never the single legitimate writer. + agent._stream_writer_lock = threading.Lock() + agent._stream_writer_token = 0 + agent._stream_writer_tls = threading.local() + agent._stream_writer_dropped = 0 + # Optional current-turn user-message override used when the API-facing # user message intentionally differs from the persisted transcript # (e.g. CLI voice mode adds a temporary prefix for the live call only). diff --git a/agent/chat_completion_helpers.py b/agent/chat_completion_helpers.py index 45a88f8decf0d..d32f86e2be410 100644 --- a/agent/chat_completion_helpers.py +++ b/agent/chat_completion_helpers.py @@ -2109,6 +2109,10 @@ def _bedrock_call(): invalidate_runtime_client(region) raise + # Claim the delta sink for this bedrock stream (#65991) so a + # superseded attempt's callbacks are fenced by the sink guard. + agent._claim_stream_writer() + def _on_text(text): _fire_first() agent._fire_stream_delta(text) @@ -2307,6 +2311,11 @@ def _call_chat_completions(): _diag = agent._stream_diag_init() request_client_holder["diag"] = _diag stream = request_client.chat.completions.create(**stream_kwargs) + # Claim the delta sink for THIS attempt (#65991). If a prior attempt's + # stream is somehow still alive (a stale-stream reconnect whose socket + # abort raced), this claim supersedes it so its late chunks are fenced + # out of the turn instead of interleaving with ours. + _writer_token = agent._claim_stream_writer() # Some OpenAI-compatible adapters (for example copilot-acp, and the MoA # openai-codex aggregator) accept stream=True but still return a @@ -2379,6 +2388,18 @@ def _call_chat_completions(): reasoning_parts: list = [] usage_obj = None for chunk in stream: + # Stop the moment a newer attempt has claimed the delta sink + # (#65991): this attempt has been superseded, so it must neither + # fire deltas (incl. the tool-suppressed raw-callback path below) + # nor keep consuming a stream that would interleave into the turn. + if not agent._stream_writer_is_current(_writer_token): + logger.warning( + "Streaming attempt superseded by a newer stream; stopping " + "consumption to preserve the single-writer invariant " + "(model=%s).", + api_kwargs.get("model", "unknown"), + ) + break last_chunk_time["t"] = time.time() agent._touch_activity("receiving stream response") @@ -2701,7 +2722,20 @@ def _call_anthropic(): ) except Exception: pass + # Claim the delta sink for THIS attempt (#65991) — parity with the + # chat_completions path so a superseded anthropic stream is fenced. + _writer_token = agent._claim_stream_writer() for event in stream: + # Bail the instant a newer attempt supersedes this one so a + # stale stream can't interleave tokens into the turn. + if not agent._stream_writer_is_current(_writer_token): + logger.warning( + "Anthropic streaming attempt superseded by a newer " + "stream; stopping consumption to preserve the " + "single-writer invariant (model=%s).", + api_kwargs.get("model", "unknown"), + ) + break saw_stream_event = True # Update stale-stream timer on every event so the # outer poll loop knows data is flowing. Without diff --git a/run_agent.py b/run_agent.py index 4d0436247e56c..0d702cebae7ae 100644 --- a/run_agent.py +++ b/run_agent.py @@ -4732,6 +4732,12 @@ def _reset_stream_delivery_tracking(self) -> None: def _record_streamed_assistant_text(self, text: str) -> None: """Accumulate visible assistant text emitted through stream callbacks.""" + # Single-writer guard (#65991): a superseded stream must not pollute the + # turn's accumulated text (which also feeds the interim-visible-text + # de-dup comparison), even when a caller reaches this directly (the + # tool-suppressed content path) rather than through _fire_stream_delta. + if self._stream_writer_superseded(): + return if isinstance(text, str) and text: self._current_streamed_assistant_text = ( getattr(self, "_current_streamed_assistant_text", "") + text @@ -4769,8 +4775,88 @@ def _emit_interim_assistant_message(self, assistant_msg: Dict[str, Any]) -> None except Exception: logger.debug("interim_assistant_callback error", exc_info=True) + def _ensure_stream_writer_state(self) -> None: + """Lazily create the single-writer guard fields (#65991). + + The fields are normally set in ``agent_init``, but agents constructed + via ``AIAgent.__new__`` (test doubles, legacy/partially-initialized + instances) skip that path. Claiming/checking the writer must not crash + those agents, so initialize the fields on first use. + """ + if getattr(self, "_stream_writer_lock", None) is None: + self._stream_writer_lock = threading.Lock() + if not hasattr(self, "_stream_writer_token"): + self._stream_writer_token = 0 + if getattr(self, "_stream_writer_tls", None) is None: + self._stream_writer_tls = threading.local() + if not hasattr(self, "_stream_writer_dropped"): + self._stream_writer_dropped = 0 + + def _claim_stream_writer(self) -> int: + """Claim exclusive ownership of the streaming delta sink for the calling + stream attempt and return its monotonic writer token (#65991). + + Every streaming attempt (each provider path, each retry) calls this + right before it begins consuming its stream. Claiming bumps the shared + token, so any earlier attempt still alive on another thread is + immediately superseded: its cached token no longer matches and the sink + fences its late chunks out. The token is stored per-thread, so a thread + that never claimed (a non-streaming caller) is never treated as a + writer and can never be fenced. + """ + self._ensure_stream_writer_state() + with self._stream_writer_lock: + self._stream_writer_token += 1 + token = self._stream_writer_token + self._stream_writer_tls.token = token + return token + + def _stream_writer_is_current(self, token: int) -> bool: + """True when ``token`` (from a prior _claim_stream_writer) is still the + active writer — i.e. no newer stream attempt has claimed the sink since + (#65991). Lets a stream loop bail out the instant it is superseded.""" + return token == getattr(self, "_stream_writer_token", token) + + def _stream_writer_superseded(self) -> bool: + """True when the calling thread claimed the delta sink but a newer + stream attempt has since claimed it — i.e. this thread is a stale + writer whose chunks must be dropped (#65991). + + A thread that never claimed (``token is None``) is not a writer and is + never reported as superseded, so non-streaming delta callers are + unaffected. + """ + tls = getattr(self, "_stream_writer_tls", None) + token = getattr(tls, "token", None) if tls is not None else None + if token is None: + return False + return token != getattr(self, "_stream_writer_token", token) + + def _note_dropped_stream_writer(self, where: str) -> None: + """Record + log that a superseded stream's delta was discarded.""" + try: + self._stream_writer_dropped = int(getattr(self, "_stream_writer_dropped", 0)) + 1 + except Exception: + self._stream_writer_dropped = 1 + # Log sparsely (first drop, then powers of two) so a chatty superseded + # stream can't flood the log, but a real provider problem is still + # visible. A silent discard would hide genuine failures. + _n = self._stream_writer_dropped + if _n == 1 or (_n & (_n - 1)) == 0: + logger.warning( + "Dropped delta from a superseded stream writer at %s " + "(discarded=%d this turn) — a stale stream tried to write into " + "the turn after a retry superseded it.", + where, _n, + ) + def _fire_stream_delta(self, text: str) -> None: """Fire all registered stream delta callbacks (display + TTS).""" + # Single-writer guard (#65991): a superseded stream must not interleave + # its tokens into the turn alongside the retry that replaced it. + if self._stream_writer_superseded(): + self._note_dropped_stream_writer("_fire_stream_delta") + return # If a tool iteration set the break flag, prepend a single paragraph # break before the first real text delta. This prevents the original # problem (text concatenation across tool boundaries) without stacking @@ -4824,6 +4910,11 @@ def _fire_stream_delta(self, text: str) -> None: def _fire_reasoning_delta(self, text: str) -> None: """Fire reasoning callback if registered.""" + # Single-writer guard (#65991): fence out a superseded stream's + # reasoning deltas the same way as content deltas. + if self._stream_writer_superseded(): + self._note_dropped_stream_writer("_fire_reasoning_delta") + return cb = self.reasoning_callback if cb is not None: try: From 4170cb6712d130dcaf72350fb2e106705cc4bf0a Mon Sep 17 00:00:00 2001 From: HexLab98 Date: Fri, 17 Jul 2026 07:37:10 +0700 Subject: [PATCH 2/2] test(streaming): cover the single-writer invariant for superseded streams MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Assert that a superseded stream (older writer token, other thread) is fenced from the delta sink, the active writer is never fenced, a non-claiming thread is never treated as a writer, and the real consume loop stops the instant it is superseded — so two streams can never interleave into one turn (#65991). --- .../test_bedrock_interrupt_post_worker.py | 3 + .../test_stream_single_writer_65991.py | 142 ++++++++++++++++++ 2 files changed, 145 insertions(+) create mode 100644 tests/run_agent/test_stream_single_writer_65991.py diff --git a/tests/agent/test_bedrock_interrupt_post_worker.py b/tests/agent/test_bedrock_interrupt_post_worker.py index 5a5829864b5d5..b800ac944585d 100644 --- a/tests/agent/test_bedrock_interrupt_post_worker.py +++ b/tests/agent/test_bedrock_interrupt_post_worker.py @@ -26,6 +26,9 @@ class _FakeAgent: def _has_stream_consumers(self): return False + def _claim_stream_writer(self): + return 1 + def _fire_stream_delta(self, text): pass diff --git a/tests/run_agent/test_stream_single_writer_65991.py b/tests/run_agent/test_stream_single_writer_65991.py new file mode 100644 index 0000000000000..196fedcfd1d39 --- /dev/null +++ b/tests/run_agent/test_stream_single_writer_65991.py @@ -0,0 +1,142 @@ +"""Regression tests for the streaming single-writer invariant (#65991). + +A retry that supersedes a still-live SSE stream must fence the old stream out +of the delta sink; otherwise both streams write into the same turn and the +persisted transcript is two coherent responses interleaved token-by-token. + +These tests exercise the real ``AIAgent`` guard helpers and the streaming +consume-loop, asserting that exactly one writer ever reaches the turn. +""" +import threading +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + + +def _make_agent(): + from run_agent import AIAgent + + agent = AIAgent( + api_key="test-key", + base_url="https://openrouter.ai/api/v1", + model="test/model", + quiet_mode=True, + skip_context_files=True, + skip_memory=True, + ) + agent.api_mode = "chat_completions" + agent._interrupt_requested = False + return agent + + +def _chunk(content=None, finish_reason=None, model=None): + delta = SimpleNamespace(content=content, tool_calls=None, reasoning_content=None, reasoning=None) + choice = SimpleNamespace(index=0, delta=delta, finish_reason=finish_reason) + return SimpleNamespace(choices=[choice], model=model, usage=None) + + +class TestSingleWriterSink: + def test_superseded_writer_deltas_are_dropped(self): + """A stale writer (older token, other thread) is fenced; only the + newest writer reaches the callbacks and the accumulated turn text.""" + agent = _make_agent() + delivered = [] + agent.stream_delta_callback = lambda t: delivered.append(t) + agent._stream_callback = None + + a_claimed = threading.Event() + b_claimed = threading.Event() + + def writer_a(): + agent._claim_stream_writer() # token 1 + a_claimed.set() + b_claimed.wait(timeout=2) # let B supersede us first + # We are now stale — every sink call must be a no-op. + agent._fire_stream_delta("A-should-drop") + agent._fire_reasoning_delta("A-reason-drop") + agent._record_streamed_assistant_text("A-record-drop") + + def writer_b(): + a_claimed.wait(timeout=2) + agent._claim_stream_writer() # token 2 — supersedes A + b_claimed.set() + + tb = threading.Thread(target=writer_b) + ta = threading.Thread(target=writer_a) + tb.start() + ta.start() + ta.join(timeout=3) + tb.join(timeout=3) + + assert delivered == [], "a superseded stream must not deliver any deltas" + assert "A-record-drop" not in (agent._current_streamed_assistant_text or "") + assert agent._stream_writer_dropped >= 1 + + def test_current_writer_is_never_fenced(self): + """The active writer always delivers — the guard can only drop a + stream that a *newer* claim has superseded.""" + agent = _make_agent() + delivered = [] + agent.stream_delta_callback = lambda t: delivered.append(t) + agent._stream_callback = None + + agent._claim_stream_writer() + agent._fire_stream_delta("hello ") + agent._fire_stream_delta("world") + + assert "".join(delivered) == "hello world" + assert agent._stream_writer_dropped == 0 + + def test_non_claiming_thread_is_not_a_writer(self): + """A thread that never claimed (a non-streaming delta caller) is never + treated as a stale writer, even after other attempts have claimed.""" + agent = _make_agent() + delivered = [] + agent.stream_delta_callback = lambda t: delivered.append(t) + agent._stream_callback = None + + # Some other thread runs a couple of stream attempts and bumps the token. + def other(): + agent._claim_stream_writer() + agent._claim_stream_writer() + + t = threading.Thread(target=other) + t.start() + t.join(timeout=3) + + # This (main) thread never claimed → not superseded → delivers. + assert agent._stream_writer_superseded() is False + agent._fire_stream_delta("plain") + assert delivered == ["plain"] + + +class TestSingleWriterLoop: + @patch("run_agent.AIAgent._create_request_openai_client") + @patch("run_agent.AIAgent._close_request_openai_client") + def test_consume_loop_stops_when_superseded_mid_stream(self, _close, mock_create): + """The real streaming loop bails out the moment a newer attempt claims + the sink, so a superseded stream cannot interleave into the turn.""" + agent = _make_agent() + delivered = [] + agent.stream_delta_callback = lambda t: delivered.append(t) + agent._stream_callback = None + + def stream_gen(): + yield _chunk(content="first") + # A concurrent retry supersedes this stream between chunks. + agent._claim_stream_writer() + yield _chunk(content="-stale-tail", finish_reason="stop", model="m") + + mock_client = MagicMock() + mock_client.chat.completions.create.return_value = stream_gen() + mock_create.return_value = mock_client + + agent._interruptible_streaming_api_call({}) + + assert "".join(delivered) == "first" + assert "-stale-tail" not in "".join(delivered) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-q"]))