fix(relay): close failed chat streams eagerly

Signed-off-by: Alex Fournier <afournier@nvidia.com>

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.

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,