diff --git a/agent/chat_completion_helpers.py b/agent/chat_completion_helpers.py index 882a846456ea..826d3b54885f 100644 --- a/agent/chat_completion_helpers.py +++ b/agent/chat_completion_helpers.py @@ -2429,6 +2429,67 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= # resolved, so the builder degrades to its plain default if it ever runs # first. _stream_stale_timeout = None + stream_attempt_lock = threading.Lock() + stream_attempt_state = { + "current": 0, + "cancelled": set(), + "discarded_chunks": 0, + "discarded_bytes": 0, + } + + def _start_stream_attempt() -> int: + with stream_attempt_lock: + stream_attempt_state["current"] += 1 + return int(stream_attempt_state["current"]) + + def _cancel_current_stream_attempt(reason: str) -> None: + with stream_attempt_lock: + current = int(stream_attempt_state.get("current") or 0) + if current: + stream_attempt_state["cancelled"].add(current) + if current: + logger.debug( + "Marked stream attempt %s cancelled: %s", + current, + reason, + ) + + def _stream_attempt_is_active(stream_attempt_id: int) -> bool: + with stream_attempt_lock: + return ( + stream_attempt_id == int(stream_attempt_state.get("current") or 0) + and stream_attempt_id not in stream_attempt_state["cancelled"] + ) + + def _stream_attempt_was_cancelled(stream_attempt_id: int) -> bool: + with stream_attempt_lock: + return stream_attempt_id in stream_attempt_state["cancelled"] + + def _discard_stale_stream_chunk(stream_attempt_id: int, chunk) -> None: + try: + chunk_bytes = len(repr(chunk)) + except Exception: + chunk_bytes = 0 + with stream_attempt_lock: + stream_attempt_state["discarded_chunks"] += 1 + stream_attempt_state["discarded_bytes"] += chunk_bytes + discarded_chunks = stream_attempt_state["discarded_chunks"] + discarded_bytes = stream_attempt_state["discarded_bytes"] + if discarded_chunks == 1: + logger.warning( + "Discarding chunk from superseded stream attempt %s " + "(discarded_chunks=%s discarded_bytes=%s)", + stream_attempt_id, + discarded_chunks, + discarded_bytes, + ) + logger.debug( + "Discarded stale stream chunk from attempt %s " + "(discarded_chunks=%s discarded_bytes=%s)", + stream_attempt_id, + discarded_chunks, + discarded_bytes, + ) def _fire_first_delta(): if not first_delta_fired["done"] and on_first_delta: @@ -2438,7 +2499,7 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= except Exception: pass - def _call_chat_completions(): + def _call_chat_completions(stream_attempt_id: int): """Stream a chat completions response.""" import httpx as _httpx # Per-provider / per-model request_timeout_seconds (from config.yaml) @@ -2635,6 +2696,10 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= if agent._interrupt_requested: break + if not _stream_attempt_is_active(stream_attempt_id): + _discard_stale_stream_chunk(stream_attempt_id, chunk) + continue + if not chunk.choices: if hasattr(chunk, "model") and chunk.model: model_name = chunk.model @@ -2760,6 +2825,11 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= if hasattr(chunk, "usage") and chunk.usage: usage_obj = chunk.usage + if _stream_attempt_was_cancelled(stream_attempt_id): + raise _httpx.RemoteProtocolError( + f"stream attempt {stream_attempt_id} was superseded" + ) + # Build mock response matching non-streaming shape full_content = "".join(content_parts) or None mock_tool_calls = None @@ -3048,6 +3118,7 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= try: for _stream_attempt in range(_max_stream_retries + 1): + stream_attempt_id = _start_stream_attempt() # Check for interrupt before each retry attempt. Without # this, /stop closes the HTTP connection (outer poll loop), # but the retry loop opens a FRESH connection — negating the @@ -3055,13 +3126,14 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= # retry can block for the full stream-read timeout (120s+), # causing multi-minute delays between /stop and response. if agent._interrupt_requested: + _cancel_current_stream_attempt("interrupt_before_stream_retry") raise InterruptedError("Agent interrupted before stream retry") try: if agent.api_mode == "anthropic_messages": agent._try_refresh_anthropic_client_credentials() result["response"] = _call_anthropic() else: - result["response"] = _call_chat_completions() + result["response"] = _call_chat_completions(stream_attempt_id) return # success except Exception as e: # If the main poll loop force-closed this request because @@ -3183,6 +3255,7 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= mid_tool_call=True, diag=request_client_holder.get("diag"), ) + _cancel_current_stream_attempt("stream_mid_tool_retry_cleanup") _close_request_client_once("stream_mid_tool_retry_cleanup") if agent.api_mode == "anthropic_messages": try: @@ -3247,6 +3320,7 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= diag=request_client_holder.get("diag"), ) # Close the stale request client before retry + _cancel_current_stream_attempt("stream_retry_cleanup") _close_request_client_once("stream_retry_cleanup") # Also rebuild the primary client to purge # any dead connections from the pool. @@ -3486,6 +3560,7 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= f"Reconnecting..." ) try: + _cancel_current_stream_attempt("stale_stream_kill") _close_request_client_once("stale_stream_kill") except Exception: pass @@ -3527,6 +3602,7 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= "(not a network error)." ) try: + _cancel_current_stream_attempt("stream_interrupt_abort") if agent.api_mode == "anthropic_messages": agent._anthropic_client.close() agent._rebuild_anthropic_client() diff --git a/tests/run_agent/test_stream_interrupt_retry.py b/tests/run_agent/test_stream_interrupt_retry.py index 84f6e2a7979d..09ceeef86d7a 100644 --- a/tests/run_agent/test_stream_interrupt_retry.py +++ b/tests/run_agent/test_stream_interrupt_retry.py @@ -160,3 +160,83 @@ class TestStreamInterruptBeforeRetry: result = agent._interruptible_streaming_api_call({}) assert result is not None assert attempts[0] == 3 + + @pytest.mark.filterwarnings( + "ignore::pytest.PytestUnhandledThreadExceptionWarning" + ) + @patch("run_agent.AIAgent._replace_primary_openai_client") + @patch("run_agent.AIAgent._abort_request_openai_client") + @patch("run_agent.AIAgent._create_request_openai_client") + @patch("run_agent.AIAgent._close_request_openai_client") + def test_stale_stream_attempt_cannot_emit_late_chunks_after_retry( + self, + mock_close, + mock_create, + mock_abort, + mock_replace, + monkeypatch, + ): + """A stale attempt must not keep writing deltas after it is killed. + + This reproduces the race where the outer stale detector aborts an SSE + connection, but the old iterator still yields one more chunk before + surfacing the connection error that triggers the retry. + """ + import httpx + import time + + from tests.run_agent.test_streaming import ( + _make_stream_chunk, + _make_tool_call_delta, + ) + + monkeypatch.setenv("HERMES_STREAM_STALE_TIMEOUT", "0.05") + monkeypatch.setenv("HERMES_STREAM_RETRIES", "1") + + class LateChunkAfterStaleStream: + response = SimpleNamespace(headers={}) + + def __iter__(self): + yield _make_stream_chunk(content="old start ") + yield _make_stream_chunk( + tool_calls=[ + _make_tool_call_delta( + index=0, + tc_id="call_1", + name="terminal", + ) + ] + ) + time.sleep(0.45) + yield _make_stream_chunk(content="old late ") + raise httpx.RemoteProtocolError("peer closed connection") + + retry_chunks = [ + _make_stream_chunk(content="new final"), + _make_stream_chunk(finish_reason="stop", model="test/model"), + ] + class RetryStream: + response = SimpleNamespace(headers={}) + + def __iter__(self): + return iter(retry_chunks) + + mock_client = MagicMock() + mock_client.chat.completions.create.side_effect = [ + LateChunkAfterStaleStream(), + RetryStream(), + ] + mock_create.return_value = mock_client + + agent = _make_agent() + agent._interrupt_requested = False + deltas = [] + agent.stream_delta_callback = deltas.append + + response = agent._interruptible_streaming_api_call({}) + + delivered = "".join(deltas) + assert "old late" not in delivered + assert "new final" in delivered + assert response.choices[0].message.content == "new final" + assert mock_abort.called