From 6bd5580425c4d64dfc9082de286ac179dbda9f27 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 14 Oct 2026 12:14:04 -0700 Subject: [PATCH] fix(streaming): add CustomStreamWrapper.close() and release the connection when a sync stream ends (#46354) * fix(streaming): close sync streams so the connection goes back to the pool CustomStreamWrapper gets a sync close() (plus context manager support) that closes the wrapped provider stream, and both the sync and async wrappers close the provider stream when the stream reaches its end. The Anthropic stream iterator keeps the httpx response it reads so close()/aclose() can release the socket after message_stop instead of leaving it checked out of the pool Co-authored-by: Blair Jordan Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(streaming): keep the basedpyright ceiling flat in CustomStreamWrapper.close() Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(streaming): close the response on the shared sync stream path too BaseModelResponseIterator gets a sync close() next to its aclose(), and the sync branch of BaseLLMHTTPHandler hands it the httpx response the way the async branch already does, so providers built on it release the socket on close() instead of waiting for garbage collection Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(streaming): unit tests for CustomStreamWrapper.close() and the shared iterator close() Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(streaming): close the provider stream before re-raising a mid-stream error Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Blair Jordan --- .../litellm_core_utils/streaming_handler.py | 41 +++- litellm/llms/anthropic/chat/handler.py | 21 ++ litellm/llms/base_llm/base_model_iterator.py | 4 + litellm/llms/custom_httpx/llm_http_handler.py | 2 + tests/integration/_support/wire.py | 10 + ...est_sync_stream_connection_release_wire.py | 214 ++++++++++++++++++ .../test_streaming_handler.py | 48 ++++ .../llms/base_llm/test_base_model_iterator.py | 9 + 8 files changed, 347 insertions(+), 2 deletions(-) create mode 100644 tests/integration/sdk/test_sync_stream_connection_release_wire.py diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index a78e6f17cc9..236286a493b 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1,6 +1,7 @@ import asyncio import collections.abc import datetime +import inspect import json import logging import threading @@ -234,6 +235,7 @@ class CustomStreamWrapper: self.completion_stream = completion_stream self.sent_first_chunk = False self.sent_last_chunk = False + self._closed = False self._stream_created_time: float = time.time() litellm_params: Final[GenericLiteLLMParams] = GenericLiteLLMParams.model_validate( @@ -328,6 +330,18 @@ class CustomStreamWrapper: def __aiter__(self) -> AsyncIterator["ModelResponseStream"]: return self + def __enter__(self) -> "CustomStreamWrapper": + return self + + def __exit__(self, *exc_info: object) -> None: + self.close() + + async def __aenter__(self) -> "CustomStreamWrapper": + return self + + async def __aexit__(self, *exc_info: object) -> None: + await self.aclose() + def _restore_consumer_correlation_context(self, *, guarded: bool = False) -> None: """Restore trace_id/session_id in the *consuming* thread/task/context. @@ -392,11 +406,26 @@ class CustomStreamWrapper: """ self._restore_consumer_correlation_context(guarded=True) + def close(self) -> None: + stream_to_close: Final = cast(object, self.completion_stream) # cast-ok: completion_stream is untyped + self.completion_stream = None + self._closed = True + close_stream: Final = getattr(stream_to_close, "close", None) + if callable(close_stream): + try: + result: Final = close_stream() + if inspect.iscoroutine(result): + result.close() + except Exception as e: # noqa: BLE001 # best-effort cleanup; must not raise into the caller + verbose_logger.debug("CustomStreamWrapper.close: error closing completion_stream: %s", e) + self._restore_consumer_correlation_context() + async def aclose(self): # Restore the consumer's outer context only after the underlying # provider stream's own close (and its diagnostic logging below, if # closing fails) completes - not before - so those log lines still # carry this closing stream's own trace_id/session_id. + self._closed = True if self.completion_stream is not None: stream_to_close: Final = self.completion_stream self.completion_stream = None @@ -1780,6 +1809,8 @@ class CustomStreamWrapper: response.hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] = _cost def __next__(self) -> "ModelResponseStream": + if self._closed: + raise StopIteration cache_hit = False if self.custom_llm_provider is not None and self.custom_llm_provider == "cached_response": cache_hit = True @@ -1950,7 +1981,7 @@ class CustomStreamWrapper: if self.sent_stream_usage is False and self.send_stream_usage is True: self.sent_stream_usage = True return response - self._restore_consumer_correlation_context() + self.close() raise # Re-raise StopIteration else: self.sent_last_chunk = True @@ -1986,6 +2017,7 @@ class CustomStreamWrapper: traceback_exception: Final = traceback.format_exc() # LOG FAILURE - handle streaming failure logging in the _next_ object, remove `handle_failure` once it's deprecated threading.Thread(target=self.logging_obj.failure_handler, args=(e, traceback_exception)).start() + self.close() self._handle_stream_fallback_error(e) def fetch_sync_stream(self): @@ -2005,6 +2037,8 @@ class CustomStreamWrapper: return self.completion_stream async def __anext__(self) -> "ModelResponseStream": + if self._closed: + raise StopAsyncIteration cache_hit = False if self.custom_llm_provider is not None and self.custom_llm_provider == "cached_response": cache_hit = True @@ -2125,12 +2159,15 @@ class CustomStreamWrapper: asyncio.create_task( self.logging_obj.dispatch_failure_handlers(e, traceback_exception, prefer_async_handlers=True) ) + await self.aclose() self._handle_stream_fallback_error(e) except (httpx.ReadError, httpx.RemoteProtocolError) as e: if self.received_finish_reason is None: + await self.aclose() self._log_stream_failure_and_raise(e) return await self._finalize_completed_stream(cache_hit=cache_hit) except Exception as e: + await self.aclose() self._log_stream_failure_and_raise(e) async def _finalize_completed_stream(self, cache_hit: bool) -> "ModelResponseStream": @@ -2224,7 +2261,7 @@ class CustomStreamWrapper: ) ) - self._restore_consumer_correlation_context() + await self.aclose() raise StopAsyncIteration # Re-raise StopIteration else: self.sent_last_chunk = True diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index ecd11279374..dff10473ae7 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -232,6 +232,7 @@ async def make_call( speed=speed, tool_name_reverse_map=tool_name_reverse_map, response_format_schema=response_format_schema, + response=response, ) # LOGGING @@ -302,6 +303,7 @@ def make_sync_call( speed=speed, tool_name_reverse_map=tool_name_reverse_map, response_format_schema=response_format_schema, + response=response, ) # LOGGING @@ -675,9 +677,12 @@ class ModelResponseIterator: speed: str | None = None, tool_name_reverse_map: dict[str, str] | None = None, response_format_schema: ResponseFormatSchema | None = None, + response: httpx.Response | None = None, ): self.streaming_response = streaming_response self.response_iterator = self.streaming_response + self.sync_stream = sync_stream + self._response = response self.content_blocks: list[ContentBlockDelta] = [] self.tool_index = -1 self.json_mode = json_mode @@ -1375,6 +1380,22 @@ class ModelResponseIterator: return self._handle_accumulated_json_chunk(data_str) # Sync iterator + def close(self) -> None: + response: Final = self._response + self._response = None + if response is not None: + response.close() + + async def aclose(self) -> None: + response: Final = self._response + self._response = None + if response is None: + return + if self.sync_stream: + response.close() + return + await response.aclose() + def __iter__(self): return self diff --git a/litellm/llms/base_llm/base_model_iterator.py b/litellm/llms/base_llm/base_model_iterator.py index 9cb5c2c55dd..aede8aedc60 100644 --- a/litellm/llms/base_llm/base_model_iterator.py +++ b/litellm/llms/base_llm/base_model_iterator.py @@ -81,6 +81,10 @@ class BaseModelResponseIterator: if self.http_response is not None: await self.http_response.aclose() + def close(self) -> None: + if self.http_response is not None: + self.http_response.close() + def chunk_parser(self, chunk: dict) -> GenericStreamingChunk | ModelResponseStream: return GenericStreamingChunk( text="", diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index a8b14b83ef3..cfafca4d174 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1008,6 +1008,8 @@ class BaseLLMHTTPHandler: json_mode=json_mode, optional_params=_OPTIONAL_PARAMS.validate_python(optional_params), ) + if isinstance(completion_stream, BaseModelResponseIterator): + completion_stream.http_response = response # LOGGING logging_obj.post_call( diff --git a/tests/integration/_support/wire.py b/tests/integration/_support/wire.py index e6a166e8457..b261973973c 100644 --- a/tests/integration/_support/wire.py +++ b/tests/integration/_support/wire.py @@ -70,6 +70,7 @@ class Wire: received: SimpleQueue[Request] disconnected: SimpleQueue[str] connected: SimpleQueue[str] + closed: SimpleQueue[str] def drain(self) -> tuple[Request, ...]: return tuple(self.received.get_nowait() for _ in range(self.received.qsize())) @@ -77,6 +78,9 @@ class Wire: def connections(self) -> int: return self.connected.qsize() + def open_connections(self) -> int: + return self.connected.qsize() - self.closed.qsize() + def _await_release(gate: threading.Event, closing: threading.Event) -> bool: while not closing.is_set(): @@ -98,6 +102,7 @@ def wire_server( errors: Final[SimpleQueue[Exception]] = SimpleQueue() disconnected: Final[SimpleQueue[str]] = SimpleQueue() connected: Final[SimpleQueue[str]] = SimpleQueue() + closed: Final[SimpleQueue[str]] = SimpleQueue() closing: Final = threading.Event() class Handler(BaseHTTPRequestHandler): @@ -108,6 +113,10 @@ def wire_server( super().setup() connected.put(f"{self.client_address[0]}:{self.client_address[1]}") + def finish(self) -> None: + super().finish() + closed.put(f"{self.client_address[0]}:{self.client_address[1]}") + def respond(self) -> None: headers: Final = {name.lower(): value for name, value in self.headers.items()} try: @@ -194,6 +203,7 @@ def wire_server( received, disconnected, connected, + closed, ) finally: closing.set() diff --git a/tests/integration/sdk/test_sync_stream_connection_release_wire.py b/tests/integration/sdk/test_sync_stream_connection_release_wire.py new file mode 100644 index 00000000000..a8b49d285c3 --- /dev/null +++ b/tests/integration/sdk/test_sync_stream_connection_release_wire.py @@ -0,0 +1,214 @@ +from __future__ import annotations + +import asyncio +import json +import threading +import time +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Final + +import openai +import pytest +from integration._support.client import eventually +from integration._support.wire import Reply, Request, Wire, wire_server + +import litellm +from litellm.types.utils import ModelResponseStream + +_API_KEY: Final = "synthetic-stream-release-key" +_PROMPT: Final = "does the socket go back to the pool once the stream is done with it?" +_FIRST_WORDS: Final = "hello " +_LAST_WORD: Final = "world" +_RELEASE_SECONDS: Final = 2 + + +def _anthropic_event(event: str, payload: Mapping[str, object]) -> bytes: + return f"event: {event}\ndata: {json.dumps(payload)}\n\n".encode() + + +def _openai_delta(text: str, finish_reason: str | None) -> bytes: + chunk: Final = { + "id": "chatcmpl-stream-release-wire", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "gpt-5.6", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": text}, "finish_reason": finish_reason}], + } + return b"data: " + json.dumps(chunk).encode() + b"\n\n" + + +_ANTHROPIC_OPENING: Final = b"".join( + ( + _anthropic_event( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_stream_release_wire", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 12, "output_tokens": 1}, + }, + }, + ), + _anthropic_event( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + _anthropic_event( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": _FIRST_WORDS}}, + ), + ) +) +_ANTHROPIC_CLOSING: Final = ( + _anthropic_event( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": _LAST_WORD}}, + ), + _anthropic_event("content_block_stop", {"type": "content_block_stop", "index": 0}), + _anthropic_event( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 3}, + }, + ), + _anthropic_event("message_stop", {"type": "message_stop"}), +) +_OPENAI_OPENING: Final = _openai_delta(_FIRST_WORDS, None) +_OPENAI_CLOSING: Final = (_openai_delta(_LAST_WORD, None), _openai_delta("", "stop"), b"data: [DONE]\n\n") + + +@dataclass(frozen=True, slots=True) +class _Shape: + name: str + model: str + opening: bytes + closing: tuple[bytes, ...] + + +_SHAPES: Final = ( + _Shape("anthropic", "anthropic/claude-sonnet-5-5", _ANTHROPIC_OPENING, _ANTHROPIC_CLOSING), + _Shape("openai", "openai/gpt-5.6", _OPENAI_OPENING, _OPENAI_CLOSING), + _Shape("databricks", "databricks/databricks-llama", _OPENAI_OPENING, _OPENAI_CLOSING), +) + + +def _peer(shape: _Shape, gate_after_first: threading.Event | None = None): + def respond(request: Request) -> Reply: + return Reply( + content_type="text/event-stream", chunks=(shape.opening, *shape.closing), gate_after_first=gate_after_first + ) + + return respond + + +def _request_kwargs(shape: _Shape, wire: Wire) -> Mapping[str, object]: + client: Final = openai.OpenAI(api_key=_API_KEY, base_url=wire.url) if shape.name == "openai" else None + return { + "model": shape.model, + "api_base": wire.url, + "api_key": _API_KEY, + "messages": ({"role": "user", "content": _PROMPT},), + "stream": True, + **({"client": client} if client is not None else {}), + } + + +def _text_of(chunk: ModelResponseStream) -> str: + return chunk.choices[0].delta.content or "" + + +def _open_sessions_settling_to(wire: Wire, at_most: int) -> int: + return eventually(wire.open_connections, lambda open_sessions: open_sessions <= at_most, seconds=_RELEASE_SECONDS) + + +async def _open_sessions_settling_to_while_the_loop_runs(wire: Wire, at_most: int) -> int: + deadline: Final = time.monotonic() + _RELEASE_SECONDS + while wire.open_connections() > at_most and time.monotonic() < deadline: + await asyncio.sleep(0.05) + return wire.open_connections() + + +@pytest.mark.parametrize("shape", _SHAPES, ids=lambda shape: shape.name) +def test_sync_stream_read_to_its_last_chunk_releases_the_connection(shape: _Shape) -> None: + with wire_server(_peer(shape), keep_alive=True) as wire: + request: Final = _request_kwargs(shape, wire) + first_text: Final = "".join(_text_of(chunk) for chunk in litellm.completion(**request)) + second_text: Final = "".join(_text_of(chunk) for chunk in litellm.completion(**request)) + assert (first_text, second_text) == (_FIRST_WORDS + _LAST_WORD, _FIRST_WORDS + _LAST_WORD) + assert _open_sessions_settling_to(wire, at_most=1) <= 1, "a finished stream left its TCP session open" + + +@pytest.mark.parametrize("shape", _SHAPES, ids=lambda shape: shape.name) +def test_sync_stream_closed_after_its_first_chunk_releases_the_connection(shape: _Shape) -> None: + gate: Final = threading.Event() + with wire_server(_peer(shape, gate_after_first=gate), keep_alive=True) as wire: + request: Final = _request_kwargs(shape, wire) + abandoned: Final = litellm.completion(**request) + first_chunk: Final = next(iter(abandoned)) + assert _text_of(first_chunk) == _FIRST_WORDS + abandoned.close() + gate.set() + assert _open_sessions_settling_to(wire, at_most=0) == 0, "close() left the TCP session open" + gate.clear() + follow_up: Final = litellm.completion(**request) + assert _text_of(next(iter(follow_up))) == _FIRST_WORDS + gate.set() + assert "".join(_text_of(chunk) for chunk in follow_up) == _LAST_WORD + + +@pytest.mark.asyncio +async def test_async_stream_read_to_its_last_chunk_releases_the_connection() -> None: + shape: Final = _SHAPES[0] + with wire_server(_peer(shape), keep_alive=True) as wire: + request: Final = _request_kwargs(shape, wire) + first_text: Final = "".join([_text_of(chunk) async for chunk in await litellm.acompletion(**request)]) + second_text: Final = "".join([_text_of(chunk) async for chunk in await litellm.acompletion(**request)]) + assert (first_text, second_text) == (_FIRST_WORDS + _LAST_WORD, _FIRST_WORDS + _LAST_WORD) + assert await _open_sessions_settling_to_while_the_loop_runs(wire, at_most=1) <= 1, ( + "a finished stream left its TCP session open" + ) + + +@pytest.mark.asyncio +async def test_async_stream_aclosed_after_its_first_chunk_releases_the_connection() -> None: + shape: Final = _SHAPES[0] + gate: Final = threading.Event() + with wire_server(_peer(shape, gate_after_first=gate), keep_alive=True) as wire: + request: Final = _request_kwargs(shape, wire) + abandoned: Final = await litellm.acompletion(**request) + assert _text_of(await abandoned.__anext__()) == _FIRST_WORDS + await abandoned.aclose() + gate.set() + assert await _open_sessions_settling_to_while_the_loop_runs(wire, at_most=0) == 0, ( + "aclose() left the TCP session open" + ) + gate.clear() + follow_up: Final = await litellm.acompletion(**request) + assert _text_of(await follow_up.__anext__()) == _FIRST_WORDS + gate.set() + assert "".join([_text_of(chunk) async for chunk in follow_up]) == _LAST_WORD + + +_ANTHROPIC_ERROR_MIDDLE: Final = _anthropic_event( + "error", {"type": "error", "error": {"type": "overloaded_error", "message": "peer overloaded mid stream"}} +) + + +def test_sync_stream_that_raises_mid_stream_releases_the_connection() -> None: + shape: Final = _SHAPES[0] + broken: Final = _Shape(shape.name, shape.model, shape.opening, (_ANTHROPIC_ERROR_MIDDLE, *shape.closing)) + with wire_server(_peer(broken), keep_alive=True) as wire: + request: Final = _request_kwargs(broken, wire) + with pytest.raises(Exception, match="overloaded"): + "".join(_text_of(chunk) for chunk in litellm.completion(**request)) + assert _open_sessions_settling_to(wire, at_most=0) == 0, ( + "a stream that raised mid-way left its TCP session open" + ) diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index b60544e6606..0172f0e9aeb 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -6415,3 +6415,51 @@ def validate_last_format(chunk): assert isinstance( choice["finish_reason"], str ), "'finish_reason' should be a string." + + +class _RecordingProviderStream: + def __init__(self) -> None: + self.close_calls = 0 + + def __iter__(self) -> "_RecordingProviderStream": + return self + + def __next__(self) -> object: + raise StopIteration + + def close(self) -> None: + self.close_calls += 1 + + +def _wrapper_over(stream: _RecordingProviderStream) -> CustomStreamWrapper: + return CustomStreamWrapper( + completion_stream=stream, + model="gpt-5.6", + logging_obj=MagicMock(), + custom_llm_provider="openai", + ) + + +def test_close_releases_the_provider_stream_once_and_ends_iteration(): + stream: Final = _RecordingProviderStream() + wrapper: Final = _wrapper_over(stream) + wrapper.close() + wrapper.close() + assert stream.close_calls == 1 + with pytest.raises(StopIteration): + next(wrapper) + + +def test_leaving_a_with_block_closes_the_provider_stream(): + stream: Final = _RecordingProviderStream() + with _wrapper_over(stream): + pass + assert stream.close_calls == 1 + + +@pytest.mark.asyncio +async def test_leaving_an_async_with_block_closes_the_provider_stream(): + stream: Final = _RecordingProviderStream() + async with _wrapper_over(stream): + pass + assert stream.close_calls == 1 diff --git a/tests/unit/llms/base_llm/test_base_model_iterator.py b/tests/unit/llms/base_llm/test_base_model_iterator.py index 36747c52416..936bf06bc18 100644 --- a/tests/unit/llms/base_llm/test_base_model_iterator.py +++ b/tests/unit/llms/base_llm/test_base_model_iterator.py @@ -7,6 +7,7 @@ import json from collections.abc import AsyncIterator from typing import Final +import httpx import pytest from pydantic import BaseModel @@ -305,3 +306,11 @@ async def test_bytes_chunk_with_text_before_data_prefix_is_parsed(sync_mode: boo assert chunks[0].id == "chatcmpl-bytes" assert chunks[0].choices[0].delta.content == "Hello" assert chunks[1]["is_finished"] is True + + +def test_close_releases_the_attached_http_response(): + response: Final = httpx.Response(200, stream=httpx.ByteStream(b"data: [DONE]\n\n")) + iterator: Final = BaseModelResponseIterator(streaming_response=iter(()), sync_stream=True) + iterator.http_response = response + iterator.close() + assert response.is_closed