From 5a5743188b8a6bae0550d068dad131709cb68f16 Mon Sep 17 00:00:00 2001 From: Alex Fournier Date: Thu, 23 Jul 2026 14:25:14 -0700 Subject: [PATCH] fix(relay): close failed chat streams eagerly Signed-off-by: Alex Fournier --- agent/chat_completion_helpers.py | 114 +++++++++++------- .../test_stream_single_writer_65991.py | 15 +++ 2 files changed, 83 insertions(+), 46 deletions(-) diff --git a/agent/chat_completion_helpers.py b/agent/chat_completion_helpers.py index 8b2dd47bc6b64..409df6604a6bb 100644 --- a/agent/chat_completion_helpers.py +++ b/agent/chat_completion_helpers.py @@ -2639,6 +2639,22 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= "discarded_chunks": 0, "discarded_bytes": 0, } + managed_stream_holder = {"stream": None} + + def _set_managed_stream(stream: Any) -> Any: + managed_stream_holder["stream"] = stream + return stream + + def _close_managed_stream() -> None: + stream = managed_stream_holder.pop("stream", None) + if stream is None: + return + close = getattr(stream, "close", None) + if callable(close): + try: + close() + except Exception: + logger.debug("Managed provider stream cleanup failed", exc_info=True) def _start_stream_attempt() -> int: with stream_attempt_lock: @@ -2856,28 +2872,30 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= from agent import relay_llm - stream = relay_llm.stream( - api_kwargs, - _open_stream, - session_id=str(getattr(agent, "session_id", "") or ""), - name=str(getattr(agent, "provider", "") or "provider"), - model_name=str(getattr(agent, "model", "") or ""), - finalizer=_relay_final_response, - on_stream_created=_stream_created, - accept_chunk=_accept_stream_chunk, - completed_response_predicate=lambda value: hasattr(value, "choices"), - metadata={ - "api_mode": "chat_completions", - "api_request_id": getattr(agent, "_current_api_request_id", None), - "call_role": ( - "delegated" - if getattr(agent, "is_subagent", False) - else "fallback" - if int(getattr(agent, "_fallback_index", 0) or 0) > 0 - else "primary" - ), - }, - defer_logical_completion=True, + stream = _set_managed_stream( + relay_llm.stream( + api_kwargs, + _open_stream, + session_id=str(getattr(agent, "session_id", "") or ""), + name=str(getattr(agent, "provider", "") or "provider"), + model_name=str(getattr(agent, "model", "") or ""), + finalizer=_relay_final_response, + on_stream_created=_stream_created, + accept_chunk=_accept_stream_chunk, + completed_response_predicate=lambda value: hasattr(value, "choices"), + metadata={ + "api_mode": "chat_completions", + "api_request_id": getattr(agent, "_current_api_request_id", None), + "call_role": ( + "delegated" + if getattr(agent, "is_subagent", False) + else "fallback" + if int(getattr(agent, "_fallback_index", 0) or 0) > 0 + else "primary" + ), + }, + defer_logical_completion=True, + ) ) for chunk in stream: last_chunk_time["t"] = time.time() @@ -3033,7 +3051,7 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= if hasattr(chunk, "usage") and chunk.usage: usage_obj = chunk.usage - stream.close() + _close_managed_stream() if _stream_attempt_was_cancelled(stream_attempt_id): raise _httpx.RemoteProtocolError( @@ -3282,28 +3300,30 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= ) return False - stream = relay_llm.stream( - api_kwargs, - _open_anthropic_stream, - session_id=str(getattr(agent, "session_id", "") or ""), - name=str(getattr(agent, "provider", "") or "anthropic"), - model_name=str(getattr(agent, "model", "") or ""), - finalizer=accumulator.finalize, - on_stream_created=_anthropic_stream_created, - on_chunk=accumulator.observe, - accept_chunk=_accept_anthropic_event, - metadata={ - "api_mode": "anthropic_messages", - "api_request_id": getattr(agent, "_current_api_request_id", None), - "call_role": ( - "delegated" - if getattr(agent, "is_subagent", False) - else "fallback" - if int(getattr(agent, "_fallback_index", 0) or 0) > 0 - else "primary" - ), - }, - defer_logical_completion=True, + stream = _set_managed_stream( + relay_llm.stream( + api_kwargs, + _open_anthropic_stream, + session_id=str(getattr(agent, "session_id", "") or ""), + name=str(getattr(agent, "provider", "") or "anthropic"), + model_name=str(getattr(agent, "model", "") or ""), + finalizer=accumulator.finalize, + on_stream_created=_anthropic_stream_created, + on_chunk=accumulator.observe, + accept_chunk=_accept_anthropic_event, + metadata={ + "api_mode": "anthropic_messages", + "api_request_id": getattr(agent, "_current_api_request_id", None), + "call_role": ( + "delegated" + if getattr(agent, "is_subagent", False) + else "fallback" + if int(getattr(agent, "_fallback_index", 0) or 0) > 0 + else "primary" + ), + }, + defer_logical_completion=True, + ) ) try: for event in stream: @@ -3358,7 +3378,7 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= raise finally: try: - stream.close() + _close_managed_stream() finally: manager = _stream_context["manager"] if manager is not None: @@ -3421,6 +3441,7 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= result["response"] = _call_chat_completions(stream_attempt_id) return # success except Exception as e: + _close_managed_stream() # If the main poll loop force-closed this request because # of an interrupt, the resulting transport error is the # expected consequence of our own close — NOT a transient @@ -3717,6 +3738,7 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta= result["error"] = e return finally: + _close_managed_stream() _close_request_client_once("stream_request_complete") # Provider-configured stale timeout takes priority over env default. diff --git a/tests/run_agent/test_stream_single_writer_65991.py b/tests/run_agent/test_stream_single_writer_65991.py index 00d324fac930e..4fbf74ae2514b 100644 --- a/tests/run_agent/test_stream_single_writer_65991.py +++ b/tests/run_agent/test_stream_single_writer_65991.py @@ -137,6 +137,21 @@ class TestSingleWriterLoop: assert "".join(delivered) == "first" assert "-stale-tail" not in "".join(delivered) + def test_chat_parser_failure_closes_managed_stream(self): + agent = _make_agent() + managed_stream = MagicMock() + managed_stream.__iter__.return_value = iter([object()]) + managed_stream.final_response = None + + with patch( + "agent.relay_llm.stream", + return_value=managed_stream, + ): + with pytest.raises(AttributeError): + agent._interruptible_streaming_api_call({}) + + managed_stream.close.assert_called_once() + class TestCodexSingleWriter: """The codex_responses path claims the sink and stops when superseded,