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