From ff56251555f73d835d0424635303c3cb5a2a46ce Mon Sep 17 00:00:00 2001 From: Burke Autrey Date: Fri, 17 Jul 2026 14:27:23 -0500 Subject: [PATCH] fix(streaming): Bedrock liveness watchdog (into #58962 breaker) + finite local timeout MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Re-applied against v0.18.2. Upstream now covers the OpenAI-path stall watchdog (cross-turn give-up breaker #58962) and the tool-batch deadline, so those are dropped. Two gaps remain: - Bedrock streaming had NO liveness watchdog and was excluded from the #58962 breaker (it returns before _check_stale_giveup). Add an on_event hook to stream_converse_with_callbacks (fires per yielded event = true wire-level liveness), drive a stale timer from it, and wire Bedrock INTO the existing breaker: entry _check_stale_giveup, _bump_stale_streak on stall, raise to end the call (invalidate_runtime_client can't abort the in-flight botocore stream, so the streak escalates across turns like the OpenAI path), and reset the streak on success. - local providers still got float('inf') stale timeout (watchdog disabled) — give a finite ceiling (HERMES_LOCAL_STREAM_STALE_TIMEOUT, default 900s). - tests: on_event per-event + swallow; Bedrock stall bumps streak + aborts; pre-elevated streak aborts at entry; success resets streak. Co-Authored-By: Claude Opus 4.8 (1M context) --- agent/bedrock_adapter.py | 16 ++ agent/chat_completion_helpers.py | 119 ++++++++++++- .../test_bedrock_interrupt_post_worker.py | 8 + tests/run_agent/test_streaming.py | 168 ++++++++++++++++++ 4 files changed, 307 insertions(+), 4 deletions(-) diff --git a/agent/bedrock_adapter.py b/agent/bedrock_adapter.py index 51d0afbe3cbe..c0a039aa0323 100644 --- a/agent/bedrock_adapter.py +++ b/agent/bedrock_adapter.py @@ -789,6 +789,7 @@ def stream_converse_with_callbacks( on_tool_start=None, on_reasoning_delta=None, on_interrupt_check=None, + on_event=None, ) -> SimpleNamespace: """Process a Bedrock ConverseStream event stream with real-time callbacks. @@ -808,6 +809,12 @@ def stream_converse_with_callbacks( on supported models (Claude 4.6+). on_interrupt_check: Called on each event. Should return True if the agent has been interrupted and streaming should stop. + on_event: Called once at the top of the loop body for EVERY yielded + Bedrock event (text/tool-input/reasoning/metadata deltas alike), + before any branching. Provides a wire-level liveness signal so an + external watchdog can distinguish "still receiving events" from + "stream wedged with no data". Errors raised by the callback are + swallowed so a liveness hook can never abort the stream. Returns: An OpenAI-compatible SimpleNamespace response, identical in shape to @@ -823,6 +830,15 @@ def stream_converse_with_callbacks( usage_data: Dict[str, int] = {} for event in event_stream.get("stream", []): + # Wire-level liveness signal: fire on EVERY yielded event (text, tool + # input, reasoning, metadata) before branching so an external watchdog + # can tell a still-flowing stream from a wedged one. Best-effort — a + # liveness callback must never be able to abort the stream. + if on_event is not None: + try: + on_event() + except Exception: + pass # Check for interrupt if on_interrupt_check and on_interrupt_check(): break diff --git a/agent/chat_completion_helpers.py b/agent/chat_completion_helpers.py index 7badd2a08712..38ed1428cc14 100644 --- a/agent/chat_completion_helpers.py +++ b/agent/chat_completion_helpers.py @@ -264,6 +264,35 @@ def _check_stale_giveup(agent) -> None: ) +def _derive_stream_stale_timeout(agent, api_kwargs: dict) -> float: + """Stale-stream patience for a provider that is never a local endpoint. + + Mirrors the main streaming path's derivation — provider config → env base + → context-size scaling → reasoning-model floor — minus the local-endpoint + ``float('inf')``/900s disable branch, which cannot apply to Bedrock (its + endpoint is always the AWS cloud). Factored so the Bedrock streaming + watchdog shares the exact same patience budget as the OpenAI/Anthropic + stale-stream detector below. + """ + _cfg_stale = get_provider_stale_timeout(agent.provider, agent.model) + if _cfg_stale is not None: + _base = _cfg_stale + else: + _base = env_float("HERMES_STREAM_STALE_TIMEOUT", 180.0) + _est_tokens = estimate_request_context_tokens(api_kwargs) + if _est_tokens > 100_000: + _timeout = max(_base, 300.0) + elif _est_tokens > 50_000: + _timeout = max(_base, 240.0) + else: + _timeout = _base + from agent.reasoning_timeouts import get_reasoning_stale_timeout_floor + _reasoning_floor = get_reasoning_stale_timeout_floor(api_kwargs.get("model")) + if _reasoning_floor is not None: + _timeout = max(_timeout, _reasoning_floor) + return _timeout + + def _dispatch_nonstreaming_api_request(agent, api_kwargs: dict, *, make_client): """Run one non-streaming LLM request for the active api_mode and return it. @@ -2091,6 +2120,24 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= result = {"response": None, "error": None} first_delta_fired = {"done": False} deltas_were_sent = {"yes": False} + # Wire-level liveness for the boto3 converse_stream worker: the worker + # thread blocks inside ``for event in event_stream`` with NO read + # timeout, so a provider that opens the stream then stops yielding + # events wedges the thread forever. on_event stamps this on EVERY + # yielded Bedrock event (text/tool/metadata) — the poll loop below + # trips a watchdog when the gap exceeds the stale timeout. + _bedrock_last_event = {"t": time.time()} + # Region captured for the poll-loop client eviction below. Read + # (not popped) here so the worker's own pop inside _bedrock_call still + # resolves the same value. + _bedrock_region = api_kwargs.get("__bedrock_region__", "us-east-1") + # Same patience budget as the OpenAI/Anthropic stale detector. + _bedrock_stale_timeout = _derive_stream_stale_timeout(agent, api_kwargs) + + # Cross-turn stale-stream circuit breaker (#58962): a pre-elevated + # streak from prior wedged turns aborts before we even start — mirrors + # the entry check on the OpenAI/Anthropic path below. + _check_stale_giveup(agent) def _fire_first(): if not first_delta_fired["done"] and on_first_delta: @@ -2168,6 +2215,7 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= on_tool_start=_on_tool, on_reasoning_delta=_on_reasoning if agent.reasoning_callback or agent.stream_delta_callback else None, on_interrupt_check=lambda: agent._interrupt_requested, + on_event=lambda: _bedrock_last_event.__setitem__("t", time.time()), ) except Exception as e: result["error"] = e @@ -2178,6 +2226,56 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= t.join(timeout=0.3) if agent._interrupt_requested: raise InterruptedError("Agent interrupted during Bedrock API call") + # Liveness watchdog: no Bedrock event for longer than the stale + # timeout means the stream has wedged (open socket, keep-alives but + # no data, or a silently hung provider). Without this the worker + # blocks in ``for event in event_stream`` indefinitely. + _stale_elapsed = time.time() - _bedrock_last_event["t"] + if _stale_elapsed > _bedrock_stale_timeout: + logger.warning( + "Bedrock stream stale for %.0fs (threshold %.0fs) — no events " + "received. region=%s model=%s. Aborting call.", + _stale_elapsed, _bedrock_stale_timeout, + _bedrock_region, api_kwargs.get("modelId", "unknown"), + ) + agent._buffer_status( + f"⚠️ No events from Bedrock for {int(_stale_elapsed)}s " + f"(model: {api_kwargs.get('modelId', 'unknown')}). Aborting..." + ) + # Count the stale kill in the SAME cross-turn breaker as the + # OpenAI/Anthropic path (#58962). + _bump_stale_streak(agent) + # Best-effort: evict the region's cached bedrock-runtime client + # so the NEXT call reconnects with a fresh pool. NOTE: this does + # NOT abort the in-flight botocore EventStream the worker thread + # is blocked on — botocore exposes no external cancellation for + # it — so the daemon worker keeps reading until its socket read + # ultimately errors. We therefore end THIS call by raising + # below and let the streak+give-up breaker escalate across turns. + try: + from agent.bedrock_adapter import invalidate_runtime_client + invalidate_runtime_client(_bedrock_region) + except Exception as _inval_exc: + logger.debug( + "bedrock: stale client eviction failed: %s", _inval_exc + ) + # Reset the timer so a repeated trip (should the worker somehow + # survive) waits a fresh interval rather than re-firing instantly. + _bedrock_last_event["t"] = time.time() + # Escalate across turns: raises RuntimeError once the streak + # crosses HERMES_STREAM_STALE_GIVEUP, so a persistently wedged + # Bedrock provider aborts fast instead of re-waiting the timeout. + _check_stale_giveup(agent) + # Streak still under the give-up threshold: end THIS call with a + # TimeoutError so the outer retry loop / next turn re-evaluates + # and the streak carries forward. Break rather than keep polling + # a worker we cannot abort. + result["error"] = TimeoutError( + f"Bedrock stream produced no events for {int(_stale_elapsed)}s " + f"(threshold {int(_bedrock_stale_timeout)}s) — aborting stalled " + f"stream so the retry/fallback path can recover." + ) + break # Worker exited before the poll loop observed the interrupt flag. The # Bedrock stream callback breaks out and returns a PARTIAL response # without raising on interrupt (see bedrock_adapter.py @@ -2190,6 +2288,11 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= raise InterruptedError("Agent interrupted during Bedrock API call (post-worker)") if result["error"] is not None: raise result["error"] + # Success — clear the cross-turn breaker (#58962): Bedrock proved + # responsive. Mirrors the OpenAI/Anthropic success reset below so a + # recovered provider doesn't carry a stale streak into later turns. + if result["response"] is not None: + _reset_stale_streak(agent) return result["response"] result = {"response": None, "error": None, "partial_tool_names": []} @@ -3196,11 +3299,19 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= else: _stream_stale_timeout_base = env_float("HERMES_STREAM_STALE_TIMEOUT", 180.0) # Local providers (Ollama, oMLX, llama-cpp) can take 300+ seconds - # for prefill on large contexts. Disable the stale detector unless - # the user explicitly set HERMES_STREAM_STALE_TIMEOUT. + # for prefill on large contexts, so tolerate far longer silence than + # the cloud default — but a wedged local server must EVENTUALLY trip the + # detector rather than hang forever (an infinite timeout meant a crashed + # or deadlocked local endpoint stalled the session indefinitely). 900s + # tolerates slow prefill while still bounding a hung endpoint. Applies + # unless the user explicitly set HERMES_STREAM_STALE_TIMEOUT; override the + # local ceiling with HERMES_LOCAL_STREAM_STALE_TIMEOUT. if _stream_stale_timeout_base == 180.0 and agent.base_url and is_local_endpoint(agent.base_url): - _stream_stale_timeout = float("inf") - logger.debug("Local provider detected (%s) — stale stream timeout disabled", agent.base_url) + _stream_stale_timeout = env_float("HERMES_LOCAL_STREAM_STALE_TIMEOUT", 900.0) + logger.debug( + "Local provider detected (%s) — stale stream timeout set to %.0fs", + agent.base_url, _stream_stale_timeout, + ) else: # Scale the stale timeout for large contexts: slow models (like Opus) # can legitimately think for minutes before producing the first token diff --git a/tests/agent/test_bedrock_interrupt_post_worker.py b/tests/agent/test_bedrock_interrupt_post_worker.py index b800ac944585..0ee5e4fce33d 100644 --- a/tests/agent/test_bedrock_interrupt_post_worker.py +++ b/tests/agent/test_bedrock_interrupt_post_worker.py @@ -22,10 +22,18 @@ class _FakeAgent: _disable_streaming = False reasoning_callback = None stream_delta_callback = None + # Real AIAgent always carries these; the streaming stale-timeout derivation + # (chat_completion_helpers._derive_stream_stale_timeout) reads them. + provider = "bedrock" + model = "anthropic.claude-3-sonnet-20240229-v1:0" + _consecutive_stale_streams = 0 def _has_stream_consumers(self): return False + def _buffer_status(self, *a, **k): + pass + def _claim_stream_writer(self): return 1 diff --git a/tests/run_agent/test_streaming.py b/tests/run_agent/test_streaming.py index cdbcd51828bd..e9cb80a08aa4 100644 --- a/tests/run_agent/test_streaming.py +++ b/tests/run_agent/test_streaming.py @@ -1966,3 +1966,171 @@ class TestBedrockIamStreamingFallback: client.converse.assert_not_called() assert getattr(agent, "_disable_streaming", False) is False + + +class _BlockingEventStream: + """Mock boto3 ``converse_stream()`` response whose event iterator blocks + forever — simulates a provider that opens the stream then stops yielding + events. The worker thread sits inside ``for event in event_stream`` exactly + as a wedged Bedrock stream would, giving the liveness watchdog something to + trip on.""" + + def __init__(self, release): + self._release = release + + def get(self, key, default=None): + if key == "stream": + return self + return default + + def __iter__(self): + return self + + def __next__(self): + # Never yields — blocks until the test releases it (teardown) so the + # daemon worker can exit instead of leaking a truly-hung thread. + self._release.wait(timeout=30) + raise StopIteration + + +def test_on_event_fires_per_bedrock_event(): + """FIX 1: on_event fires once for EVERY yielded Bedrock event — text, + tool-input delta, messageStop, and metadata alike — providing wire-level + liveness (not just text deltas).""" + from agent.bedrock_adapter import stream_converse_with_callbacks + + events = [ + {"contentBlockDelta": {"delta": {"text": "a"}}}, + {"contentBlockStart": {"start": {"toolUse": {"toolUseId": "t1", "name": "x"}}}}, + {"contentBlockDelta": {"delta": {"toolUse": {"input": "{}"}}}}, + {"contentBlockStop": {}}, + {"messageStop": {"stopReason": "end_turn"}}, + {"metadata": {"usage": {"inputTokens": 1, "outputTokens": 1}}}, + ] + calls = {"n": 0} + + stream_converse_with_callbacks( + {"stream": iter(events)}, + on_event=lambda: calls.__setitem__("n", calls["n"] + 1), + ) + + assert calls["n"] == len(events) + + +def test_on_event_exception_is_swallowed(): + """FIX 1: a raising on_event callback must never abort the stream.""" + from agent.bedrock_adapter import stream_converse_with_callbacks + + events = [{"messageStop": {"stopReason": "end_turn"}}] + + def _boom(): + raise ValueError("liveness hook blew up") + + resp = stream_converse_with_callbacks({"stream": iter(events)}, on_event=_boom) + assert resp is not None + assert resp.choices[0].finish_reason == "stop" + + +class TestBedrockStreamLivenessWatchdog: + """FIX 1: Bedrock streaming participates in the #58962 cross-turn stale + breaker and no longer hangs when the stream stops yielding events.""" + + def _make_bedrock_agent(self): + from run_agent import AIAgent + + agent = AIAgent( + api_key="test-key", + base_url="https://openrouter.ai/api/v1", + model="anthropic.claude-3-sonnet-20240229-v1:0", + quiet_mode=True, + skip_context_files=True, + skip_memory=True, + ) + agent.api_mode = "bedrock_converse" + agent._interrupt_requested = False + return agent + + def test_stalled_stream_bumps_streak_and_aborts(self, monkeypatch): + """A Bedrock stream that opens then stops yielding events trips the + watchdog: it bumps the cross-turn stale streak and raises TimeoutError + instead of hanging forever.""" + pytest.importorskip("botocore", reason="botocore required for Bedrock tests") + import threading as _t + + # Tiny stale timeout so the watchdog trips quickly; give-up threshold + # kept above 1 so a single call raises TimeoutError (not the breaker). + monkeypatch.setenv("HERMES_STREAM_STALE_TIMEOUT", "0.5") + monkeypatch.setenv("HERMES_STREAM_STALE_GIVEUP", "5") + + agent = self._make_bedrock_agent() + agent._consecutive_stale_streams = 0 + release = _t.Event() + + client = MagicMock() + client.converse_stream.return_value = _BlockingEventStream(release) + + try: + with patch( + "agent.bedrock_adapter._get_bedrock_runtime_client", + return_value=client, + ): + with pytest.raises(TimeoutError): + agent._interruptible_streaming_api_call( + {"modelId": agent.model, "messages": []} + ) + finally: + release.set() + + # Watchdog counted exactly one stale kill in the cross-turn breaker. + assert agent._consecutive_stale_streams == 1 + + def test_pre_elevated_streak_aborts_before_streaming(self, monkeypatch): + """A streak already past the give-up threshold aborts at entry with + RuntimeError — Bedrock never even opens a stream (cross-turn breaker).""" + pytest.importorskip("botocore", reason="botocore required for Bedrock tests") + + monkeypatch.setenv("HERMES_STREAM_STALE_GIVEUP", "5") + + agent = self._make_bedrock_agent() + agent._consecutive_stale_streams = 5 + + client = MagicMock() + with patch( + "agent.bedrock_adapter._get_bedrock_runtime_client", + return_value=client, + ): + with pytest.raises(RuntimeError, match="unresponsive"): + agent._interruptible_streaming_api_call( + {"modelId": agent.model, "messages": []} + ) + + client.converse_stream.assert_not_called() + + def test_successful_stream_resets_streak(self, monkeypatch): + """A Bedrock stream that completes normally clears any prior stale + streak so a recovered provider doesn't carry it into later turns.""" + pytest.importorskip("botocore", reason="botocore required for Bedrock tests") + + monkeypatch.setenv("HERMES_STREAM_STALE_TIMEOUT", "60") + + agent = self._make_bedrock_agent() + agent._consecutive_stale_streams = 3 # simulate a prior wedged streak + + events = [ + {"contentBlockDelta": {"delta": {"text": "hi"}}}, + {"messageStop": {"stopReason": "end_turn"}}, + {"metadata": {"usage": {"inputTokens": 1, "outputTokens": 1}}}, + ] + client = MagicMock() + client.converse_stream.return_value = {"stream": iter(events)} + + with patch( + "agent.bedrock_adapter._get_bedrock_runtime_client", + return_value=client, + ): + response = agent._interruptible_streaming_api_call( + {"modelId": agent.model, "messages": []} + ) + + assert response.choices[0].message.content == "hi" + assert agent._consecutive_stale_streams == 0