From 41fdcae6881512cb3de842877a7589173fc9ce32 Mon Sep 17 00:00:00 2001 From: brooklyn! Date: Fri, 17 Jul 2026 14:31:09 -0400 Subject: [PATCH] fix(streaming): make the single-writer fence best-effort so a missing guard can't crash a turn (#66448) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A cron job ("Daily Buzz Report") died with 'AIAgent' object has no attribute '_claim_stream_writer'. The #65991 single-writer fence lives on AIAgent (run_agent.py), but the streaming paths that use it live in other modules — chat_completion_helpers (chat / anthropic / bedrock) and codex_runtime (codex responses) — and called it directly as agent._claim_stream_writer() / agent._stream_writer_is_current(). That makes those modules hard-depend on the method being present on whatever object is passed as agent. The fence is an *additive* safety net that may only ever drop a provably superseded stream, never the sole legitimate writer. But the direct calls turned any agent that doesn't expose it — a version-skewed checkout (the streaming helper module newer than run_agent), a hot-reloaded gateway mid git-pull, a duck-typed agent, or a test double — into a fatal AttributeError that aborts the whole turn (and, on cron, fails the job). Route every cross-module claim/check through agent/stream_single_writer.py. claim_stream_writer(agent) returns 0 when the fence is unavailable (or raises), and stream_writer_is_current(agent, token) treats a 0 token or an absent guard as "current" — so a guard-less agent degrades to "no fence" instead of crashing, while a real AIAgent keeps the full single-writer protection. Internal self.* uses inside run_agent are unchanged (self is always a full AIAgent there). --- agent/chat_completion_helpers.py | 11 +-- agent/codex_runtime.py | 6 +- agent/stream_single_writer.py | 70 ++++++++++++++++ .../agent/test_stream_single_writer_guard.py | 79 +++++++++++++++++++ 4 files changed, 159 insertions(+), 7 deletions(-) create mode 100644 agent/stream_single_writer.py create mode 100644 tests/agent/test_stream_single_writer_guard.py diff --git a/agent/chat_completion_helpers.py b/agent/chat_completion_helpers.py index da00e1053e48f..7badd2a087125 100644 --- a/agent/chat_completion_helpers.py +++ b/agent/chat_completion_helpers.py @@ -36,6 +36,7 @@ from agent.message_sanitization import ( _sanitize_surrogates, _repair_tool_call_arguments, ) +from agent.stream_single_writer import claim_stream_writer, stream_writer_is_current from tools.terminal_tool import is_persistent_env from utils import base_url_host_matches, base_url_hostname, env_float, env_int @@ -2146,7 +2147,7 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= # 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() + claim_stream_writer(agent) def _on_text(text): _fire_first() @@ -2350,7 +2351,7 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= # 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() + _writer_token = claim_stream_writer(agent) # Some OpenAI-compatible adapters (for example copilot-acp, and the MoA # openai-codex aggregator) accept stream=True but still return a @@ -2427,7 +2428,7 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= # (#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): + if not stream_writer_is_current(agent, _writer_token): logger.warning( "Streaming attempt superseded by a newer stream; stopping " "consumption to preserve the single-writer invariant " @@ -2759,11 +2760,11 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= 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() + _writer_token = claim_stream_writer(agent) 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): + if not stream_writer_is_current(agent, _writer_token): logger.warning( "Anthropic streaming attempt superseded by a newer " "stream; stopping consumption to preserve the " diff --git a/agent/codex_runtime.py b/agent/codex_runtime.py index 394ada176e22f..4629eff1802ea 100644 --- a/agent/codex_runtime.py +++ b/agent/codex_runtime.py @@ -23,6 +23,8 @@ import time from types import SimpleNamespace from typing import Any, Callable, Dict, List +from agent.stream_single_writer import claim_stream_writer, stream_writer_is_current + logger = logging.getLogger(__name__) @@ -1190,12 +1192,12 @@ def run_codex_stream(agent, api_kwargs: dict, client: Any = None, on_first_delta # late deltas are fenced out of the turn; conversely, a newer # attempt supersedes us and the interrupt_check below stops our # consumption immediately. - _writer_token = agent._claim_stream_writer() + _writer_token = claim_stream_writer(agent) def _interrupt_or_superseded(_tok=_writer_token) -> bool: if agent._interrupt_requested: return True - if not agent._stream_writer_is_current(_tok): + if not stream_writer_is_current(agent, _tok): logger.warning( "Codex streaming attempt superseded by a newer stream; " "stopping consumption to preserve the single-writer " diff --git a/agent/stream_single_writer.py b/agent/stream_single_writer.py new file mode 100644 index 0000000000000..c0b18245187dd --- /dev/null +++ b/agent/stream_single_writer.py @@ -0,0 +1,70 @@ +"""Best-effort accessors for the single-writer stream fence (#65991). + +The fence itself lives on ``AIAgent`` (``_claim_stream_writer`` / +``_stream_writer_is_current`` in ``run_agent.py``), but the streaming code paths +that use it live in *other* modules — ``chat_completion_helpers`` (chat / +anthropic / bedrock) and ``codex_runtime`` (codex responses). Calling the fence +directly as ``agent._claim_stream_writer()`` from those modules makes them +hard-depend on the method being present on whatever object is passed in as +``agent``. + +That coupling is a latent crash: a partially-updated checkout (the streaming +helper module newer than ``run_agent``), a hot-reloaded gateway, a duck-typed +agent, or a test double without the method turns an *additive* safety net into a +fatal ``AttributeError`` that aborts the whole turn. A cron job died exactly +this way with ``'AIAgent' object has no attribute '_claim_stream_writer'``. + +The fence is only ever allowed to drop a *provably* superseded stream — never +the sole legitimate writer. So when the guard is unavailable (or raises), the +correct degradation is "no fence": keep streaming. These helpers make the +claim/check best-effort to guarantee that. +""" + +from __future__ import annotations + +import logging +from typing import Any + +logger = logging.getLogger(__name__) + + +def claim_stream_writer(agent: Any) -> int: + """Claim the delta sink for the calling stream attempt, best-effort. + + Returns the agent's monotonic writer token when the fence is available, or + ``0`` when the agent doesn't expose it (or the claim raised). A ``0`` token + pairs with :func:`stream_writer_is_current` always returning ``True``, so a + guard-less agent is simply never fenced instead of crashing the turn. + """ + claim = getattr(agent, "_claim_stream_writer", None) + if callable(claim): + try: + return int(claim()) + except Exception: + logger.debug( + "stream single-writer: claim failed; proceeding unfenced", + exc_info=True, + ) + return 0 + + +def stream_writer_is_current(agent: Any, token: int) -> bool: + """True when ``token`` is still the active writer, best-effort. + + A falsy token (from a claim that no-oped) or an agent without the fence + means we cannot prove supersession, so the stream is treated as current and + never fenced. This preserves the single-writer invariant's one-way promise: + only a demonstrably stale writer is ever stopped. + """ + if not token: + return True + is_current = getattr(agent, "_stream_writer_is_current", None) + if callable(is_current): + try: + return bool(is_current(token)) + except Exception: + logger.debug( + "stream single-writer: is_current check failed; treating as current", + exc_info=True, + ) + return True diff --git a/tests/agent/test_stream_single_writer_guard.py b/tests/agent/test_stream_single_writer_guard.py new file mode 100644 index 0000000000000..b807dd686fb29 --- /dev/null +++ b/tests/agent/test_stream_single_writer_guard.py @@ -0,0 +1,79 @@ +"""Regression tests for the best-effort single-writer fence accessors. + +The streaming paths in ``chat_completion_helpers`` and ``codex_runtime`` reach +the #65991 single-writer fence through :mod:`agent.stream_single_writer` instead +of calling ``agent._claim_stream_writer()`` directly. That indirection exists so +an agent object that doesn't expose the fence (a version-skewed checkout, a +duck-typed agent, a test double) degrades to "no fence" rather than aborting the +whole turn with ``'AIAgent' object has no attribute '_claim_stream_writer'`` — +the exact AttributeError that killed a cron job. + +These tests assert the fence's *contract*: it may drop a provably superseded +stream, but it must never fence (or crash) the sole legitimate writer. +""" + +import run_agent +from agent.stream_single_writer import claim_stream_writer, stream_writer_is_current + + +class _NoFenceAgent: + """An agent-like object that predates / lacks the single-writer fence.""" + + +class _RaisingFenceAgent: + """An agent whose fence methods exist but blow up when called.""" + + def _claim_stream_writer(self): + raise RuntimeError("boom") + + def _stream_writer_is_current(self, token): + raise RuntimeError("boom") + + +def _real_agent(): + """A real AIAgent without running the heavy __init__ (fields self-heal).""" + return object.__new__(run_agent.AIAgent) + + +def test_claim_on_fenceless_agent_does_not_raise(): + # Regression: this is the cron crash path — the streaming helper must not + # explode when the agent lacks _claim_stream_writer. + assert claim_stream_writer(_NoFenceAgent()) == 0 + + +def test_is_current_on_fenceless_agent_is_always_current(): + agent = _NoFenceAgent() + # A no-op claim (token 0) must never report as superseded, regardless of + # what token value a caller threads through. + assert stream_writer_is_current(agent, 0) is True + assert stream_writer_is_current(agent, 7) is True + + +def test_zero_token_is_never_fenced_even_with_a_real_fence(): + # Invariant: a claim that no-oped (token 0) is not a writer and can never be + # fenced, even against an agent that does implement the fence. + assert stream_writer_is_current(_real_agent(), 0) is True + + +def test_claim_swallows_fence_exceptions(): + assert claim_stream_writer(_RaisingFenceAgent()) == 0 + + +def test_is_current_swallows_fence_exceptions_as_current(): + assert stream_writer_is_current(_RaisingFenceAgent(), 123) is True + + +def test_real_agent_fence_still_supersedes_and_preserves_sole_writer(): + agent = _real_agent() + + first = claim_stream_writer(agent) + assert first > 0 + # Sole writer so far — still current. + assert stream_writer_is_current(agent, first) is True + + # A newer attempt claims the sink: the older token is now superseded, the + # newer one is current. The fence drops only the provably stale writer. + second = claim_stream_writer(agent) + assert second > first + assert stream_writer_is_current(agent, first) is False + assert stream_writer_is_current(agent, second) is True