From 4e7258ad69da563a59aa93fd9c0fe2e3f4e84f65 Mon Sep 17 00:00:00 2001 From: Rahul Johny <116638720+johnyrahul@users.noreply.github.com> Date: Wed, 7 Oct 2026 03:38:39 +0000 Subject: [PATCH 1/4] UN-4237 [FIX] Close streamed LLM responses so garbage collection cannot deadlock executor workers on httpcore's pool lock LiteLLM's sync CustomStreamWrapper has no close() and stops iterating before the HTTP body ends, so every streamed completion left its response checked out of httpcore's pool until the garbage collector finalised it. When that happened on a thread inside ConnectionPool.handle_request, the finaliser's PoolByteStream.close() waited on the same non-reentrant lock the thread held, and the worker hung forever with its lease still renewing. Add close_stream(), which closes the stream and the provider response behind it (litellm's ModelResponseIterator, the OpenAI SDK Stream, plain generators), and call it in a finally in collect_with_retry and iter_with_retry: on success, on error and on early exit. The retry backoff now runs after the stream is closed. Tests run the real LLM.complete() path against a local Anthropic SSE server with GC disabled, and force a collection inside the pool lock to reproduce the deadlock; they fail against the previous code. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01VQoJFFiRZyDKyu14pg5hDb --- .../src/unstract/sdk1/utils/retry_utils.py | 70 +++- .../tests/test_stream_connection_release.py | 327 ++++++++++++++++++ .../tests/utils/test_collect_with_retry.py | 153 +++++++- 3 files changed, 536 insertions(+), 14 deletions(-) create mode 100644 unstract/sdk1/tests/test_stream_connection_release.py diff --git a/unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py b/unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py index 794082e218..cef14f31dd 100644 --- a/unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py +++ b/unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py @@ -234,6 +234,50 @@ async def acall_with_retry[T]( raise RuntimeError("unreachable") # for type-checker: loop always returns or raises +# Attributes on a provider stream that hold the HTTP response (or the line +# generator over it). LiteLLM's ``ModelResponseIterator`` (Anthropic, Vertex, +# Gemini, the generic HTTP handler) keeps ``streaming_response`` and, for +# Vertex, ``response``; the OpenAI SDK's ``Stream`` keeps ``response``. +_STREAM_RESPONSE_ATTRS = ("streaming_response", "response") + + +def close_stream(stream: object) -> None: + """Release the HTTP response behind a streamed completion, best effort. + + A stream left open is only closed when the garbage collector finalises + it, and that can happen on a thread holding httpcore's connection-pool + lock: the finaliser's ``PoolByteStream.close()`` takes the same + non-reentrant lock and the thread deadlocks on itself (UN-4237). Closing + here, outside any pool lock, leaves nothing for the collector. + + LiteLLM's sync ``CustomStreamWrapper`` has no ``close()`` and stops + iterating before the HTTP body ends, so even a fully drained stream is + still open; the wrapper's ``completion_stream`` and the response it + holds are closed directly. Duck-typed so this module never imports + litellm. Never raises: it runs in ``finally`` blocks and must not mask + the original error. + """ + if stream is None: + return + targets: list[object] = [stream] + inner = getattr(stream, "completion_stream", None) + if inner is not None: + targets.append(inner) + for obj in list(targets): + for attr in _STREAM_RESPONSE_ATTRS: + handle = getattr(obj, attr, None) + if handle is not None: + targets.append(handle) + for target in targets: + close = getattr(target, "close", None) + if not callable(close): + continue + try: + close() + except Exception as e: + logger.debug("Ignoring error while closing stream %r: %s", target, e) + + def iter_with_retry[T]( fn: Callable[[], Iterable[T]], *, @@ -258,12 +302,6 @@ def iter_with_retry[T]( yield item return except Exception as e: - # Close generator to release in-flight HTTP/socket resources - # before retrying — otherwise streaming providers leak sockets - # until GC. - close = getattr(gen, "close", None) - if callable(close): - close() if has_yielded: raise delay = _get_retry_delay( @@ -271,7 +309,13 @@ def iter_with_retry[T]( ) if delay is None: raise - time.sleep(delay) + finally: + # On success, on error and when the caller abandons this + # generator: see ``close_stream``. + close_stream(gen) + # After ``finally`` so the connection is released, not held + # through the backoff. + time.sleep(delay) def collect_with_retry[T]( @@ -310,10 +354,6 @@ def collect_with_retry[T]( has_content = has_content or is_content(item) return items except Exception as e: - # Release the in-flight HTTP/socket resources before retrying. - close = getattr(gen, "close", None) - if callable(close): - close() if has_content: raise delay = _get_retry_delay( @@ -321,7 +361,13 @@ def collect_with_retry[T]( ) if delay is None: raise - time.sleep(delay) + finally: + # On success too: a drained stream is still open (see + # ``close_stream``). + close_stream(gen) + # After ``finally`` so the connection is released, not held + # through the backoff. + time.sleep(delay) raise RuntimeError("unreachable") # for type-checker: loop always returns or raises diff --git a/unstract/sdk1/tests/test_stream_connection_release.py b/unstract/sdk1/tests/test_stream_connection_release.py new file mode 100644 index 0000000000..555cfc07a7 --- /dev/null +++ b/unstract/sdk1/tests/test_stream_connection_release.py @@ -0,0 +1,327 @@ +"""A streamed completion releases its pooled HTTP connection (UN-4237). + +Executor workers deadlocked under load inside httpcore's connection pool. +LiteLLM's sync stream wrapper has no ``close()`` and stops iterating before +the HTTP body ends, so every streamed completion left its response checked +out of the pool. The garbage collector finalised it later, and when that +happened on a thread already inside ``ConnectionPool.handle_request``, the +finaliser's ``PoolByteStream.close()`` waited on the same non-reentrant lock +the thread held: a deadlock no timeout could break. + +These tests run the real litellm Anthropic path against a local server +speaking Anthropic's SSE format, with the garbage collector disabled so it +cannot be the one releasing anything. Then they reproduce the deadlock's +trigger directly: a garbage collection inside the pool lock on the next call. +""" + +from __future__ import annotations + +import gc +import json +import threading +from concurrent.futures import ThreadPoolExecutor +from functools import lru_cache +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from importlib import import_module +from typing import TYPE_CHECKING +from unittest.mock import patch + +import httpcore +import litellm +import pytest +from litellm.litellm_core_utils import litellm_logging, streaming_handler +from litellm.llms.custom_httpx.http_handler import HTTPHandler + +if TYPE_CHECKING: + from collections.abc import Iterator + +ANTHROPIC_ADAPTER_ID = "anthropic|90ebd4cd-2f19-4cef-a884-9eeb6ac0f203" +TEXT = "hello over a real socket" +# Bounds how long a regressed test hangs before failing. +DEADLOCK_TIMEOUT_SECONDS = 10 +_REAL_COMPLETION = litellm.completion + +_EVENTS: list[tuple[str, dict[str, object]]] = [ + ( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 10, "output_tokens": 1}, + }, + }, + ), + ( + "content_block_start", + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text", "text": ""}, + }, + ), + ( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": TEXT}, + }, + ), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 5}, + }, + ), + ("message_stop", {"type": "message_stop"}), +] +# Content arrives, then Anthropic reports a mid-stream error. The connection +# is still open when litellm raises, so only an explicit close releases it. A +# dropped connection would not test this: httpcore closes those itself. +_ERROR_AFTER = 3 +_ERROR_EVENT = { + "type": "error", + "error": {"type": "overloaded_error", "message": "Overloaded"}, +} + + +def _sse(event: str, data: dict[str, object]) -> bytes: + return f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() + + +class _AnthropicSSEHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + mid_stream_error = False + + def log_message(self, *_: object) -> None: + pass + + def _chunk(self, body: bytes) -> None: + self.wfile.write(b"%x\r\n%s\r\n" % (len(body), body)) + self.wfile.flush() + + def do_POST(self) -> None: # noqa: N802 - BaseHTTPRequestHandler API + self.rfile.read(int(self.headers.get("content-length", 0))) + self.send_response(200) + self.send_header("content-type", "text/event-stream") + self.send_header("transfer-encoding", "chunked") + self.end_headers() + for i, (event, data) in enumerate(_EVENTS): + if self.mid_stream_error and i == _ERROR_AFTER: + self._chunk(_sse("error", _ERROR_EVENT)) + self._chunk(_sse(event, data)) + self._chunk(b"") + + +class _QuietServer(ThreadingHTTPServer): + daemon_threads = True + + def handle_error(self, *_: object) -> None: + # A client closing mid-stream is the behaviour under test, not a fault. + pass + + +@pytest.fixture +def server() -> Iterator[ThreadingHTTPServer]: + handler = type("_Handler", (_AnthropicSSEHandler,), {}) + srv = _QuietServer(("127.0.0.1", 0), handler) + thread = threading.Thread(target=srv.serve_forever, daemon=True) + thread.start() + try: + yield srv + finally: + srv.shutdown() + srv.server_close() + + +@pytest.fixture +def pools() -> Iterator[set[httpcore.ConnectionPool]]: + """Every httpcore pool a request went through during the test.""" + seen: set[httpcore.ConnectionPool] = set() + real = httpcore.ConnectionPool.handle_request + + def handle_request( + pool: httpcore.ConnectionPool, request: httpcore.Request + ) -> httpcore.Response: + seen.add(pool) + return real(pool, request) + + with patch.object(httpcore.ConnectionPool, "handle_request", handle_request): + yield seen + + +@pytest.fixture +def client() -> HTTPHandler: + """A pool of the test's own. + + litellm otherwise reuses one module-level client, so a stream leaked by a + regressed test would deadlock whichever later test the collector ran in, + hanging the suite instead of failing the test. Not closed on teardown: + closing takes the pool lock, which a regressed deadlock test never frees. + """ + return HTTPHandler() + + +class _LoggingExecutor: + """litellm's background logging, made something a test can wait for. + + Its tasks hold the stream wrapper, so a leaked stream only becomes + garbage once they finish. + """ + + _MODULES = (streaming_handler, litellm_logging) + + def __init__(self) -> None: + self._originals = [module.executor for module in self._MODULES] + self._install(ThreadPoolExecutor(max_workers=2)) + + def _install(self, executor: ThreadPoolExecutor) -> None: + self._executor = executor + for module in self._MODULES: + module.executor = executor + + def drain(self) -> None: + """Wait for every task submitted so far, then accept new ones.""" + self._executor.shutdown(wait=True) + self._install(ThreadPoolExecutor(max_workers=2)) + + def restore(self) -> None: + self._executor.shutdown(wait=False) + for module, original in zip(self._MODULES, self._originals, strict=True): + module.executor = original + + +@pytest.fixture +def logging_executor() -> Iterator[_LoggingExecutor]: + executor = _LoggingExecutor() + try: + yield executor + finally: + executor.restore() + + +@pytest.fixture +def no_gc() -> Iterator[None]: + """Keep the garbage collector from releasing what the code should.""" + was_enabled = gc.isenabled() + gc.disable() + try: + yield + finally: + if was_enabled: + gc.enable() + + +@lru_cache(maxsize=1) +def _load_llm_module() -> object: + import sys + from types import ModuleType + + sys.modules.setdefault("magic", ModuleType("magic")) + return import_module("unstract.sdk1.llm") + + +def _checked_out(pools: set[httpcore.ConnectionPool]) -> int: + """Requests whose response the pool still considers open.""" + return sum(len(pool._requests) for pool in pools) + + +def _complete(srv: ThreadingHTTPServer, client: HTTPHandler) -> dict[str, object]: + """Run ``LLM.complete()`` for Anthropic against the local server.""" + llm_module = _load_llm_module() + llm = llm_module.LLM( + adapter_id=ANTHROPIC_ADAPTER_ID, + adapter_metadata={"model": "claude-sonnet-4-6", "api_key": "test-key"}, + ) + api_base = f"http://127.0.0.1:{srv.server_port}" + + def completion(**kwargs: object) -> object: + return _REAL_COMPLETION(**kwargs, api_base=api_base, client=client) + + with ( + patch.object(llm_module.litellm, "completion", completion), + patch.object(llm_module.litellm, "cost_per_token", return_value=(0.0, 0.0)), + ): + return llm.complete("hi", max_retries=0) + + +def test_completed_stream_releases_its_connection( + server: ThreadingHTTPServer, + client: HTTPHandler, + pools: set[httpcore.ConnectionPool], + no_gc: None, +) -> None: + result = _complete(server, client) + + assert result["response"].text == TEXT + assert pools, "the request never reached httpcore" + assert _checked_out(pools) == 0 + + +def test_stream_failing_after_content_releases_its_connection( + server: ThreadingHTTPServer, + client: HTTPHandler, + pools: set[httpcore.ConnectionPool], + no_gc: None, +) -> None: + server.RequestHandlerClass.mid_stream_error = True + llm_module = _load_llm_module() + + with pytest.raises(llm_module.LLMError): + _complete(server, client) + + assert pools, "the request never reached httpcore" + assert _checked_out(pools) == 0 + + +def test_gc_inside_the_pool_lock_does_not_deadlock( + server: ThreadingHTTPServer, + client: HTTPHandler, + logging_executor: _LoggingExecutor, + no_gc: None, +) -> None: + """The production deadlock, triggered on purpose. + + Under load the collector ran on an allocation inside ``handle_request``'s + locked section. Forcing a collection there, after a completed stream, is + what deadlocked every time before the fix. Takes no ``pools`` fixture: + pytest reprs a failing test's arguments, and ``ConnectionPool.__repr__`` + takes the very lock a regressed run leaves held. + """ + _complete(server, client) + # Until litellm's logging tasks finish they still reference the first + # stream, and a collection would find nothing to finalise. + logging_executor.drain() + + real_assign = httpcore.ConnectionPool._assign_requests_to_connections + + def assign(pool: httpcore.ConnectionPool) -> object: + gc.collect() # runs under ``pool._optional_thread_lock`` + return real_assign(pool) + + outcome: list[object] = [] + + def second_call() -> None: + try: + outcome.append(_complete(server, client)["response"].text) + except Exception as e: # surfaced by the assertion below + outcome.append(e) + + with patch.object(httpcore.ConnectionPool, "_assign_requests_to_connections", assign): + # Daemon: if this regresses, the stuck thread must not block exit. + worker = threading.Thread(target=second_call, daemon=True) + worker.start() + worker.join(DEADLOCK_TIMEOUT_SECONDS) + + assert not worker.is_alive(), "deadlocked on httpcore's pool lock" + assert outcome == [TEXT] diff --git a/unstract/sdk1/tests/utils/test_collect_with_retry.py b/unstract/sdk1/tests/utils/test_collect_with_retry.py index 974a6e6e00..e15d114a1a 100644 --- a/unstract/sdk1/tests/utils/test_collect_with_retry.py +++ b/unstract/sdk1/tests/utils/test_collect_with_retry.py @@ -141,7 +141,12 @@ def close(self) -> None: closed.append(True) attempts = iter([_Gen(fail=True), _Gen(fail=False)]) - with patch.object(retry_utils.time, "sleep"): + + def _sleep(_: float) -> None: + # The failed attempt is released before the backoff, not after. + assert closed == [True] + + with patch.object(retry_utils.time, "sleep", side_effect=_sleep): result = collect_with_retry( lambda: next(attempts), max_retries=1, @@ -149,4 +154,148 @@ def close(self) -> None: is_content=_is_content, ) assert result == ["c1"] - assert closed == [True] + # Both attempts: the failed one and the successful one (UN-4237). + assert closed == [True, True] + + +# ── Releasing the HTTP response (UN-4237) ──────────────────────────────────── +# +# A stream left open is closed by the garbage collector, possibly while the +# thread holds httpcore's non-reentrant pool lock, which deadlocks the worker. +# LiteLLM's sync stream wrapper has no ``close()`` and stops before the HTTP +# body ends, so every stream must be closed explicitly, through the handles +# the wrapper holds. ``test_stream_connection_release.py`` covers this against +# a real socket. + + +class _Closable: + def __init__(self, name: str, log: list[str]) -> None: + self._name = name + self._log = log + + def close(self) -> None: + self._log.append(self._name) + + +class _ProviderIterator: + """Shaped like litellm's ``ModelResponseIterator``.""" + + def __init__(self, log: list[str]) -> None: + self.streaming_response = _Closable("lines", log) + self.response = _Closable("response", log) + + +class _Wrapper: + """Shaped like litellm's sync ``CustomStreamWrapper``: no ``close()``.""" + + def __init__(self, items: list[str | Exception], log: list[str]) -> None: + self._items = items + self.completion_stream = _ProviderIterator(log) + + def __iter__(self) -> Iterator[str]: + for item in self._items: + if isinstance(item, Exception): + raise item + yield item + + +def test_close_stream_reaches_the_response_behind_a_wrapper() -> None: + log: list[str] = [] + retry_utils.close_stream(_Wrapper([], log)) + assert sorted(log) == ["lines", "response"] + + +def test_close_stream_ignores_none_and_objects_without_close() -> None: + retry_utils.close_stream(None) + retry_utils.close_stream(iter(["c1"])) + retry_utils.close_stream(object()) + + +def test_close_stream_never_raises() -> None: + class _Broken: + def close(self) -> None: + raise RuntimeError("already torn down") + + retry_utils.close_stream(_Broken()) + + +def test_closes_drained_stream_on_success() -> None: + log: list[str] = [] + result = collect_with_retry( + lambda: _Wrapper(["meta", "c1"], log), + max_retries=0, + retry_predicate=lambda _: True, + is_content=_is_content, + ) + assert result == ["meta", "c1"] + assert sorted(log) == ["lines", "response"] + + +def test_closes_stream_when_failure_after_content_is_raised() -> None: + log: list[str] = [] + with pytest.raises(TimeoutError): + collect_with_retry( + lambda: _Wrapper(["c1", TimeoutError()], log), + max_retries=2, + retry_predicate=lambda _: True, + is_content=_is_content, + ) + assert sorted(log) == ["lines", "response"] + + +def test_close_failure_does_not_mask_the_stream_error() -> None: + class _Gen: + def __iter__(self) -> Iterator[str]: + raise ValueError("bad request") + + def close(self) -> None: + raise RuntimeError("close failed") + + with pytest.raises(ValueError, match="bad request"): + collect_with_retry( + _Gen, max_retries=0, retry_predicate=lambda _: False, is_content=_is_content + ) + + +def test_iter_with_retry_closes_stream_on_success() -> None: + log: list[str] = [] + result = list( + retry_utils.iter_with_retry( + lambda: _Wrapper(["meta", "c1"], log), + max_retries=0, + retry_predicate=lambda _: True, + ) + ) + assert result == ["meta", "c1"] + assert sorted(log) == ["lines", "response"] + + +def test_iter_with_retry_closes_stream_when_caller_stops_early() -> None: + log: list[str] = [] + gen = retry_utils.iter_with_retry( + lambda: _Wrapper(["c1", "c2"], log), + max_retries=0, + retry_predicate=lambda _: True, + ) + assert next(gen) == "c1" + gen.close() + assert sorted(log) == ["lines", "response"] + + +def test_iter_with_retry_closes_failed_attempt_before_backoff() -> None: + log: list[str] = [] + attempts = iter([_Wrapper([TimeoutError()], log), _Wrapper(["c1"], log)]) + + def _sleep(_: float) -> None: + assert sorted(log) == ["lines", "response"] + + with patch.object(retry_utils.time, "sleep", side_effect=_sleep): + result = list( + retry_utils.iter_with_retry( + lambda: next(attempts), + max_retries=1, + retry_predicate=lambda _: True, + ) + ) + assert result == ["c1"] + assert len(log) == 4 From 7146e07fd0a3905d32fa8a019058eef8e2be5a86 Mon Sep 17 00:00:00 2001 From: Rahul Johny <116638720+johnyrahul@users.noreply.github.com> Date: Wed, 7 Oct 2026 07:10:22 +0000 Subject: [PATCH 2/4] UN-4237 Close the Responses-API stream too, and warn when a close fails Addresses the standardized review on #2318. - close_stream now walks the stream's handle attributes to a bounded depth with a visited set, instead of one fixed level. On litellm's Responses-API bridge (responses/ and the models routed there) the HTTP response sits under an iterator with no close(), one level below where the old walk stopped, so an error after content or an early close left the pool slot checked out. - A close that raises is logged at WARNING with the traceback: httpcore marks its stream closed before closing the socket, so that failure can hold the slot for good. - Comments and docstrings no longer claim an abandoned generator is covered: only a caller that closes it is. stream_complete's docstring says callers that stop early must close it. - Tests: socket tests for the Responses-API bridge (error after content, early close, drain), and unit tests for a generator completion_stream (Bedrock), an OpenAI-SDK-shaped Stream, the Responses chain, cycles, and the warning. Each was checked against a mutation of the fix. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01VQoJFFiRZyDKyu14pg5hDb --- unstract/sdk1/src/unstract/sdk1/llm.py | 6 + .../src/unstract/sdk1/utils/retry_utils.py | 104 +++++++--- .../tests/test_stream_connection_release.py | 191 ++++++++++++++++-- .../tests/utils/test_collect_with_retry.py | 137 +++++++++++++ 4 files changed, 392 insertions(+), 46 deletions(-) diff --git a/unstract/sdk1/src/unstract/sdk1/llm.py b/unstract/sdk1/src/unstract/sdk1/llm.py index 894a9b7d97..326a1b39c0 100644 --- a/unstract/sdk1/src/unstract/sdk1/llm.py +++ b/unstract/sdk1/src/unstract/sdk1/llm.py @@ -827,6 +827,12 @@ def stream_complete( Chunks arrive as they stream from the provider. ``cache_prefix`` behaves as in :meth:`complete` — a stable prefix cached ahead of ``prompt`` when prompt caching is active for a supported provider. + + A caller that stops iterating early, on purpose or because its own + code raised, must close the generator, e.g. with + ``contextlib.closing(llm.stream_complete(...))``. Otherwise the + provider's HTTP response is released only when the garbage collector + finalises it, which can deadlock httpcore's connection pool (UN-4237). """ try: messages = self._build_messages(prompt, cache_prefix=cache_prefix) diff --git a/unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py b/unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py index cef14f31dd..a26605ce53 100644 --- a/unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py +++ b/unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py @@ -234,11 +234,48 @@ async def acall_with_retry[T]( raise RuntimeError("unreachable") # for type-checker: loop always returns or raises -# Attributes on a provider stream that hold the HTTP response (or the line -# generator over it). LiteLLM's ``ModelResponseIterator`` (Anthropic, Vertex, -# Gemini, the generic HTTP handler) keeps ``streaming_response`` and, for -# Vertex, ``response``; the OpenAI SDK's ``Stream`` keeps ``response``. -_STREAM_RESPONSE_ATTRS = ("streaming_response", "response") +# Attributes through which a provider stream reaches its HTTP response, at +# any depth. LiteLLM's sync ``CustomStreamWrapper`` holds the provider +# iterator in ``completion_stream``; its ``ModelResponseIterator`` (Anthropic, +# Vertex/Gemini, the generic HTTP handler) holds the line generator in +# ``streaming_response`` and, for Vertex/Gemini, the response in ``response``. +# On the Responses-API bridge (``responses/`` and models litellm routes +# there), ``streaming_response`` is an iterator with no ``close()`` that holds +# ``response`` and ``stream_iterator`` one level further down. The OpenAI +# SDK's ``Stream`` keeps ``response``. +_STREAM_HANDLE_ATTRS = ( + "completion_stream", + "streaming_response", + "response", + "http_response", + "stream_iterator", +) +# Levels walked below the stream itself; the Responses-API bridge is the +# deepest known chain (wrapper -> bridge -> iterator -> response). +_STREAM_WALK_DEPTH = 3 + + +def _stream_handles(stream: object) -> list[object]: + """``stream`` and everything reachable from it through the handle attrs.""" + found: list[object] = [] + seen: set[int] = set() + level = [stream] + for depth in range(_STREAM_WALK_DEPTH + 1): + below: list[object] = [] + for obj in level: + if obj is None or id(obj) in seen: + continue + seen.add(id(obj)) + found.append(obj) + if depth == _STREAM_WALK_DEPTH: + continue + for attr in _STREAM_HANDLE_ATTRS: + try: + below.append(getattr(obj, attr, None)) + except Exception: # a misbehaving property must not stop the walk + continue + level = below + return found def close_stream(stream: object) -> None: @@ -247,35 +284,32 @@ def close_stream(stream: object) -> None: A stream left open is only closed when the garbage collector finalises it, and that can happen on a thread holding httpcore's connection-pool lock: the finaliser's ``PoolByteStream.close()`` takes the same - non-reentrant lock and the thread deadlocks on itself (UN-4237). Closing - here, outside any pool lock, leaves nothing for the collector. - - LiteLLM's sync ``CustomStreamWrapper`` has no ``close()`` and stops - iterating before the HTTP body ends, so even a fully drained stream is - still open; the wrapper's ``completion_stream`` and the response it - holds are closed directly. Duck-typed so this module never imports - litellm. Never raises: it runs in ``finally`` blocks and must not mask - the original error. + non-reentrant lock and the thread deadlocks on itself (UN-4237). Called + by the consumer, outside any pool lock, this leaves nothing for the + collector. + + LiteLLM's sync ``CustomStreamWrapper`` has no ``close()``, and on some + providers (Anthropic among them) it stops iterating before the HTTP body + ends, so even a fully drained stream can still be open. Every object + reachable through ``_STREAM_HANDLE_ATTRS`` that has a ``close()`` is + closed. Duck-typed so this module never imports litellm. Never raises: + it runs in ``finally`` blocks and must not mask the original error. """ - if stream is None: - return - targets: list[object] = [stream] - inner = getattr(stream, "completion_stream", None) - if inner is not None: - targets.append(inner) - for obj in list(targets): - for attr in _STREAM_RESPONSE_ATTRS: - handle = getattr(obj, attr, None) - if handle is not None: - targets.append(handle) - for target in targets: + for target in _stream_handles(stream): close = getattr(target, "close", None) if not callable(close): continue try: close() - except Exception as e: - logger.debug("Ignoring error while closing stream %r: %s", target, e) + except Exception: + # httpcore marks its stream closed before closing the socket, so + # a failure here can leave the pool slot checked out for good: + # the leak this function exists to prevent. + logger.warning( + "Failed to close %s; its pooled HTTP connection may stay checked out", + type(target).__name__, + exc_info=True, + ) def iter_with_retry[T]( @@ -290,6 +324,10 @@ def iter_with_retry[T]( Once items have been yielded to the caller a mid-iteration failure is raised immediately — partial output can't be un-yielded. + + A caller that stops before the end must ``close()`` the generator (for + example with ``contextlib.closing``) so the stream is released on its own + thread rather than by the garbage collector (see ``close_stream``). """ _validate_max_retries(max_retries) log = logger_instance or logger @@ -310,8 +348,10 @@ def iter_with_retry[T]( if delay is None: raise finally: - # On success, on error and when the caller abandons this - # generator: see ``close_stream``. + # On success, on error, and when the caller closes this generator + # early. A generator dropped without ``close()`` is finalised by + # the garbage collector instead, possibly inside httpcore's pool + # lock, so callers that stop early must close it (UN-4237). close_stream(gen) # After ``finally`` so the connection is released, not held # through the backoff. @@ -362,8 +402,8 @@ def collect_with_retry[T]( if delay is None: raise finally: - # On success too: a drained stream is still open (see - # ``close_stream``). + # On success too: on some providers a drained stream is still + # open (see ``close_stream``). close_stream(gen) # After ``finally`` so the connection is released, not held # through the backoff. diff --git a/unstract/sdk1/tests/test_stream_connection_release.py b/unstract/sdk1/tests/test_stream_connection_release.py index 555cfc07a7..ece77011bf 100644 --- a/unstract/sdk1/tests/test_stream_connection_release.py +++ b/unstract/sdk1/tests/test_stream_connection_release.py @@ -1,17 +1,20 @@ """A streamed completion releases its pooled HTTP connection (UN-4237). Executor workers deadlocked under load inside httpcore's connection pool. -LiteLLM's sync stream wrapper has no ``close()`` and stops iterating before -the HTTP body ends, so every streamed completion left its response checked -out of the pool. The garbage collector finalised it later, and when that +LiteLLM's sync stream wrapper has no ``close()``, so a streamed completion +that failed, was abandoned, or (on Anthropic) simply finished before the HTTP +body ended left its response checked out of the pool. The garbage collector +finalised it later, and when that happened on a thread already inside ``ConnectionPool.handle_request``, the finaliser's ``PoolByteStream.close()`` waited on the same non-reentrant lock the thread held: a deadlock no timeout could break. -These tests run the real litellm Anthropic path against a local server -speaking Anthropic's SSE format, with the garbage collector disabled so it -cannot be the one releasing anything. Then they reproduce the deadlock's -trigger directly: a garbage collection inside the pool lock on the next call. +These tests run real litellm against a local server speaking the provider's +SSE format, with the garbage collector disabled so it cannot be the one +releasing anything: the Anthropic path through ``LLM.complete()``, and the +OpenAI Responses-API bridge, whose response sits two iterators deep. Then they +reproduce the deadlock's trigger directly: a garbage collection inside the +pool lock on the next call. """ from __future__ import annotations @@ -31,9 +34,10 @@ import pytest from litellm.litellm_core_utils import litellm_logging, streaming_handler from litellm.llms.custom_httpx.http_handler import HTTPHandler +from unstract.sdk1.utils.retry_utils import collect_with_retry, iter_with_retry if TYPE_CHECKING: - from collections.abc import Iterator + from collections.abc import Callable, Iterator ANTHROPIC_ADAPTER_ID = "anthropic|90ebd4cd-2f19-4cef-a884-9eeb6ac0f203" TEXT = "hello over a real socket" @@ -95,12 +99,68 @@ } +_RESPONSE: dict[str, object] = { + "id": "resp_1", + "object": "response", + "created_at": 1, + "status": "in_progress", + "model": "gpt-4o", + "output": [], + "usage": None, +} +_MESSAGE: dict[str, object] = { + "id": "msg_1", + "type": "message", + "role": "assistant", + "status": "in_progress", + "content": [], +} +_TEXT_PART = {"item_id": "msg_1", "output_index": 0, "content_index": 0} + + +def _event(kind: str, **fields: object) -> tuple[str, dict[str, object]]: + return kind, {"type": kind, **fields} + + +# OpenAI's Responses API, which litellm bridges to chat completions for +# ``responses/`` and the models it routes there. +_RESPONSES_EVENTS: list[tuple[str, dict[str, object]]] = [ + _event("response.created", response=_RESPONSE), + _event("response.output_item.added", output_index=0, item=_MESSAGE), + _event( + "response.content_part.added", + **_TEXT_PART, + part={"type": "output_text", "text": "", "annotations": []}, + ), + _event("response.output_text.delta", **_TEXT_PART, delta=TEXT), + _event("response.output_text.done", **_TEXT_PART, text=TEXT), + _event( + "response.completed", + response={ + **_RESPONSE, + "status": "completed", + "output": [ + { + **_MESSAGE, + "status": "completed", + "content": [{"type": "output_text", "text": TEXT, "annotations": []}], + } + ], + "usage": {"input_tokens": 5, "output_tokens": 2, "total_tokens": 7}, + }, + ), +] + + def _sse(event: str, data: dict[str, object]) -> bytes: return f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() class _AnthropicSSEHandler(BaseHTTPRequestHandler): protocol_version = "HTTP/1.1" + events = _EVENTS + error_event: tuple[str, dict[str, object]] = ("error", _ERROR_EVENT) + error_after = _ERROR_AFTER mid_stream_error = False def log_message(self, *_: object) -> None: @@ -116,13 +176,19 @@ def do_POST(self) -> None: # noqa: N802 - BaseHTTPRequestHandler API self.send_header("content-type", "text/event-stream") self.send_header("transfer-encoding", "chunked") self.end_headers() - for i, (event, data) in enumerate(_EVENTS): - if self.mid_stream_error and i == _ERROR_AFTER: - self._chunk(_sse("error", _ERROR_EVENT)) + for i, (event, data) in enumerate(self.events): + if self.mid_stream_error and i == self.error_after: + self._chunk(_sse(*self.error_event)) self._chunk(_sse(event, data)) self._chunk(b"") +class _ResponsesSSEHandler(_AnthropicSSEHandler): + events = _RESPONSES_EVENTS + error_event = _event("error", code="server_error", message="Overloaded", param=None) + error_after = 4 # after the text delta + + class _QuietServer(ThreadingHTTPServer): daemon_threads = True @@ -131,9 +197,10 @@ def handle_error(self, *_: object) -> None: pass -@pytest.fixture -def server() -> Iterator[ThreadingHTTPServer]: - handler = type("_Handler", (_AnthropicSSEHandler,), {}) +def _serve(base: type[_AnthropicSSEHandler]) -> Iterator[ThreadingHTTPServer]: + # A subclass per test, so a test setting ``mid_stream_error`` changes + # only its own server. + handler = type("_Handler", (base,), {}) srv = _QuietServer(("127.0.0.1", 0), handler) thread = threading.Thread(target=srv.serve_forever, daemon=True) thread.start() @@ -144,6 +211,16 @@ def server() -> Iterator[ThreadingHTTPServer]: srv.server_close() +@pytest.fixture +def server() -> Iterator[ThreadingHTTPServer]: + yield from _serve(_AnthropicSSEHandler) + + +@pytest.fixture +def responses_server() -> Iterator[ThreadingHTTPServer]: + yield from _serve(_ResponsesSSEHandler) + + @pytest.fixture def pools() -> Iterator[set[httpcore.ConnectionPool]]: """Every httpcore pool a request went through during the test.""" @@ -325,3 +402,89 @@ def second_call() -> None: assert not worker.is_alive(), "deadlocked on httpcore's pool lock" assert outcome == [TEXT] + + +# ── OpenAI Responses-API bridge ────────────────────────────────────────────── +# +# The wrapper's ``completion_stream`` is litellm's Responses-to-chat bridge, +# whose ``streaming_response`` is an iterator with no ``close()``; the HTTP +# response sits one level further down. A drained stream reads to the end of +# the body and releases itself, so only failures and early exits show a leak. + + +def _responses_stream( + srv: ThreadingHTTPServer, client: HTTPHandler +) -> Callable[[], Iterator[object]]: + def open_stream() -> Iterator[object]: + return _REAL_COMPLETION( + model="openai/responses/gpt-4o", + api_base=f"http://127.0.0.1:{srv.server_port}/v1", + api_key="test-key", + messages=[{"role": "user", "content": "hi"}], + stream=True, + stream_options={"include_usage": True}, + client=client, + ) + + return open_stream + + +def _has_text(chunk: object) -> bool: + choices = getattr(chunk, "choices", None) + return bool(choices and choices[0].delta.content) + + +def test_responses_api_stream_failing_after_content_releases_its_connection( + responses_server: ThreadingHTTPServer, + client: HTTPHandler, + pools: set[httpcore.ConnectionPool], + no_gc: None, +) -> None: + responses_server.RequestHandlerClass.mid_stream_error = True + + # litellm rewrites the provider message; the type is what it raises here. + with pytest.raises(litellm.exceptions.MidStreamFallbackError): + collect_with_retry( + _responses_stream(responses_server, client), + max_retries=0, + retry_predicate=lambda _: False, + is_content=_has_text, + ) + + assert pools, "the request never reached httpcore" + assert _checked_out(pools) == 0 + + +def test_responses_api_stream_closed_early_releases_its_connection( + responses_server: ThreadingHTTPServer, + client: HTTPHandler, + pools: set[httpcore.ConnectionPool], + no_gc: None, +) -> None: + stream = iter_with_retry( + _responses_stream(responses_server, client), + max_retries=0, + retry_predicate=lambda _: False, + ) + next(stream) + stream.close() + + assert pools, "the request never reached httpcore" + assert _checked_out(pools) == 0 + + +def test_responses_api_stream_drained_returns_the_text( + responses_server: ThreadingHTTPServer, + client: HTTPHandler, + pools: set[httpcore.ConnectionPool], + no_gc: None, +) -> None: + chunks = collect_with_retry( + _responses_stream(responses_server, client), + max_retries=0, + retry_predicate=lambda _: False, + is_content=_has_text, + ) + + assert "".join(c.choices[0].delta.content or "" for c in chunks if c.choices) == TEXT + assert _checked_out(pools) == 0 diff --git a/unstract/sdk1/tests/utils/test_collect_with_retry.py b/unstract/sdk1/tests/utils/test_collect_with_retry.py index e15d114a1a..e0896f7c86 100644 --- a/unstract/sdk1/tests/utils/test_collect_with_retry.py +++ b/unstract/sdk1/tests/utils/test_collect_with_retry.py @@ -9,6 +9,7 @@ from __future__ import annotations +import logging from typing import TYPE_CHECKING from unittest.mock import patch @@ -219,6 +220,142 @@ def close(self) -> None: retry_utils.close_stream(_Broken()) +def test_failed_close_is_logged_as_a_warning(caplog: pytest.LogCaptureFixture) -> None: + """A close that raises can leave the pool slot checked out for good.""" + + class _Broken: + def close(self) -> None: + raise RuntimeError("already torn down") + + with caplog.at_level(logging.WARNING, logger=retry_utils.logger.name): + retry_utils.close_stream(_Broken()) + + [record] = caplog.records + assert record.levelno == logging.WARNING + assert "_Broken" in record.getMessage() + assert record.exc_info is not None + + +def _bedrock_shaped(log: list[str], fail_after_first: bool = False) -> object: + """A wrapper whose ``completion_stream`` is a plain generator, as on Bedrock. + + There is no response handle to reach: closing the generator itself is the + only thing that releases the connection. With ``fail_after_first`` the + wrapper raises while the provider generator is still suspended, as + litellm does on a mid-stream error. + """ + + def provider_stream() -> Iterator[str]: + try: + yield "c1" + yield "c2" + finally: + log.append("generator") + + class _Wrapper: + def __init__(self) -> None: + self.completion_stream = provider_stream() + + def __iter__(self) -> Iterator[str]: + for item in self.completion_stream: + yield item + if fail_after_first: + raise TimeoutError("mid-stream") + + return _Wrapper() + + +def test_close_stream_closes_a_generator_completion_stream() -> None: + log: list[str] = [] + wrapper = _bedrock_shaped(log) + next(iter(wrapper)) + + retry_utils.close_stream(wrapper) + + assert log == ["generator"] + + +def test_collect_with_retry_closes_a_generator_completion_stream_on_error() -> None: + log: list[str] = [] + # Held here so reference counting cannot finalise the generator once the + # call returns: only an explicit close may run its ``finally``. + opened: list[object] = [] + + def open_stream() -> object: + opened.append(_bedrock_shaped(log, fail_after_first=True)) + return opened[-1] + + with pytest.raises(TimeoutError): + collect_with_retry( + open_stream, + max_retries=0, + retry_predicate=lambda _: True, + is_content=_is_content, + ) + + assert log == ["generator"] + + +def test_close_stream_closes_an_openai_sdk_shaped_stream() -> None: + """The OpenAI SDK ``Stream``: its own ``close()`` plus ``.response``.""" + log: list[str] = [] + + class _Stream: + def __init__(self) -> None: + self.response = _Closable("response", log) + + def close(self) -> None: + log.append("stream") + + class _Wrapper: + completion_stream = _Stream() + + retry_utils.close_stream(_Wrapper()) + + assert sorted(log) == ["response", "stream"] + + +def test_close_stream_reaches_the_responses_api_response() -> None: + """The Responses-API bridge keeps the response two iterators down. + + Shaped like litellm 1.104.0: the wrapper's ``completion_stream`` is the + bridge, whose ``streaming_response`` is an iterator with no ``close()`` + holding ``response`` and ``stream_iterator``. + """ + log: list[str] = [] + + class _ResponsesIterator: # no close() + def __init__(self) -> None: + self.response = _Closable("response", log) + self.stream_iterator = _Closable("stream_iterator", log) + + class _Bridge: # no close() + streaming_response = _ResponsesIterator() + + class _Wrapper: + completion_stream = _Bridge() + + retry_utils.close_stream(_Wrapper()) + + assert sorted(log) == ["response", "stream_iterator"] + + +def test_close_stream_closes_each_object_once_and_terminates_on_cycles() -> None: + log: list[str] = [] + + class _Loop: + def __init__(self) -> None: + self.response = self + self.completion_stream = self + + def close(self) -> None: + log.append("loop") + + retry_utils.close_stream(_Loop()) + + assert log == ["loop"] + + def test_closes_drained_stream_on_success() -> None: log: list[str] = [] result = collect_with_retry( From 330ff265f46ce0b2de555fd564906c1b27faa5bf Mon Sep 17 00:00:00 2001 From: Rahul Johny <116638720+johnyrahul@users.noreply.github.com> Date: Wed, 7 Oct 2026 08:38:02 +0000 Subject: [PATCH 3/4] UN-4237 Keep close_stream from raising when reading close itself fails Addresses Greptile's review on #2318: the close attribute was read outside the try, so a handle whose close property raised would stop the cleanup before the remaining handles and could replace the original stream error in the caller's finally. The lookup now sits inside the guarded block, with a test that fails without it. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01VQoJFFiRZyDKyu14pg5hDb --- .../src/unstract/sdk1/utils/retry_utils.py | 6 +++--- .../tests/utils/test_collect_with_retry.py | 20 +++++++++++++++++++ 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py b/unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py index a26605ce53..1558d3a3d7 100644 --- a/unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py +++ b/unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py @@ -296,10 +296,10 @@ def close_stream(stream: object) -> None: it runs in ``finally`` blocks and must not mask the original error. """ for target in _stream_handles(stream): - close = getattr(target, "close", None) - if not callable(close): - continue try: + close = getattr(target, "close", None) + if not callable(close): + continue close() except Exception: # httpcore marks its stream closed before closing the socket, so diff --git a/unstract/sdk1/tests/utils/test_collect_with_retry.py b/unstract/sdk1/tests/utils/test_collect_with_retry.py index e0896f7c86..fbd0f44382 100644 --- a/unstract/sdk1/tests/utils/test_collect_with_retry.py +++ b/unstract/sdk1/tests/utils/test_collect_with_retry.py @@ -236,6 +236,26 @@ def close(self) -> None: assert record.exc_info is not None +def test_close_stream_survives_a_close_attribute_that_raises( + caplog: pytest.LogCaptureFixture, +) -> None: + """Reading ``close`` can raise too; the remaining handles still close.""" + log: list[str] = [] + + class _Wrapper: + completion_stream = _Closable("inner", log) + + @property + def close(self) -> object: + raise RuntimeError("torn down") + + with caplog.at_level(logging.WARNING, logger=retry_utils.logger.name): + retry_utils.close_stream(_Wrapper()) + + assert log == ["inner"] + assert "_Wrapper" in caplog.records[0].getMessage() + + def _bedrock_shaped(log: list[str], fail_after_first: bool = False) -> object: """A wrapper whose ``completion_stream`` is a plain generator, as on Bedrock. From b1301c4f9b6aa530cd0466c480231d9492f2a8a1 Mon Sep 17 00:00:00 2001 From: Rahul Johny <116638720+johnyrahul@users.noreply.github.com> Date: Fri, 9 Oct 2026 04:30:16 +0000 Subject: [PATCH 4/4] UN-4237 Address review: traceable close-failure warning, version-agnostic comments From Chandru's review on #2318: - A failed close logs only its own traceback: close_stream runs in a finally while the provider error is propagating, so Python chained that error onto the close failure and exc_info printed it a second time. - close_stream takes the caller's logger and adapter description, and both retry helpers pass theirs, so a leaked pool slot can be traced to the provider that leaked it. - Comments: the explanation lives once, in close_stream's docstring (the only ticket reference left); the rest point at it. Dropped litellm version pins and "before the fix" wording, and described the Responses-API chain the same way everywhere (wrapper -> completion_stream -> streaming_response -> response). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01VQoJFFiRZyDKyu14pg5hDb --- unstract/sdk1/src/unstract/sdk1/llm.py | 3 +- .../src/unstract/sdk1/utils/retry_utils.py | 64 ++++++++++------- .../tests/test_stream_connection_release.py | 40 ++++------- .../tests/utils/test_collect_with_retry.py | 70 +++++++++++++++---- 4 files changed, 113 insertions(+), 64 deletions(-) diff --git a/unstract/sdk1/src/unstract/sdk1/llm.py b/unstract/sdk1/src/unstract/sdk1/llm.py index 2332670f9d..2b461abd6a 100644 --- a/unstract/sdk1/src/unstract/sdk1/llm.py +++ b/unstract/sdk1/src/unstract/sdk1/llm.py @@ -832,7 +832,8 @@ def stream_complete( code raised, must close the generator, e.g. with ``contextlib.closing(llm.stream_complete(...))``. Otherwise the provider's HTTP response is released only when the garbage collector - finalises it, which can deadlock httpcore's connection pool (UN-4237). + finalises it, which can deadlock the connection pool (see + ``retry_utils.close_stream``). """ try: messages = self._build_messages(prompt, cache_prefix=cache_prefix) diff --git a/unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py b/unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py index 1558d3a3d7..ba98cdd851 100644 --- a/unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py +++ b/unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py @@ -234,15 +234,11 @@ async def acall_with_retry[T]( raise RuntimeError("unreachable") # for type-checker: loop always returns or raises -# Attributes through which a provider stream reaches its HTTP response, at -# any depth. LiteLLM's sync ``CustomStreamWrapper`` holds the provider -# iterator in ``completion_stream``; its ``ModelResponseIterator`` (Anthropic, -# Vertex/Gemini, the generic HTTP handler) holds the line generator in -# ``streaming_response`` and, for Vertex/Gemini, the response in ``response``. -# On the Responses-API bridge (``responses/`` and models litellm routes -# there), ``streaming_response`` is an iterator with no ``close()`` that holds -# ``response`` and ``stream_iterator`` one level further down. The OpenAI -# SDK's ``Stream`` keeps ``response``. +# Attributes through which a provider stream reaches its HTTP response: +# litellm's stream wrapper (completion_stream), provider iterators +# (streaming_response / response / http_response), the Responses-API bridge +# (stream_iterator) and the OpenAI SDK Stream (response). The shapes are pinned +# by tests/utils/test_collect_with_retry.py; update both together. _STREAM_HANDLE_ATTRS = ( "completion_stream", "streaming_response", @@ -250,8 +246,9 @@ async def acall_with_retry[T]( "http_response", "stream_iterator", ) -# Levels walked below the stream itself; the Responses-API bridge is the -# deepest known chain (wrapper -> bridge -> iterator -> response). +# Levels walked below the stream itself. The deepest known chain is the +# Responses-API bridge: wrapper -> completion_stream -> streaming_response +# -> response. _STREAM_WALK_DEPTH = 3 @@ -278,7 +275,12 @@ def _stream_handles(stream: object) -> list[object]: return found -def close_stream(stream: object) -> None: +def close_stream( + stream: object, + *, + log: logging.Logger | None = None, + description: str = "", +) -> None: """Release the HTTP response behind a streamed completion, best effort. A stream left open is only closed when the garbage collector finalises @@ -291,24 +293,36 @@ def close_stream(stream: object) -> None: LiteLLM's sync ``CustomStreamWrapper`` has no ``close()``, and on some providers (Anthropic among them) it stops iterating before the HTTP body ends, so even a fully drained stream can still be open. Every object - reachable through ``_STREAM_HANDLE_ATTRS`` that has a ``close()`` is - closed. Duck-typed so this module never imports litellm. Never raises: - it runs in ``finally`` blocks and must not mask the original error. + reachable through ``_STREAM_HANDLE_ATTRS`` within ``_STREAM_WALK_DEPTH`` + levels that has a ``close()`` is closed. Duck-typed so this module never + imports litellm. Never raises: it runs in ``finally`` blocks and must not + mask the original error. + + ``log`` and ``description`` (the adapter info, as the retry helpers take + it) let a failed close be traced back to the provider that leaked. """ + log = log or logger + source = f" ({description})" if description else "" for target in _stream_handles(stream): try: close = getattr(target, "close", None) if not callable(close): continue close() - except Exception: + except Exception as close_err: # httpcore marks its stream closed before closing the socket, so # a failure here can leave the pool slot checked out for good: - # the leak this function exists to prevent. - logger.warning( - "Failed to close %s; its pooled HTTP connection may stay checked out", + # the leak this function exists to prevent. This runs in a + # ``finally`` while the provider error may still be propagating; + # suppress that chained context so the traceback shows only the + # close failure (the provider error is logged by the caller). + close_err.__suppress_context__ = True + log.warning( + "Failed to close %s%s; its pooled HTTP connection may stay " + "checked out", type(target).__name__, - exc_info=True, + source, + exc_info=close_err, ) @@ -348,11 +362,9 @@ def iter_with_retry[T]( if delay is None: raise finally: - # On success, on error, and when the caller closes this generator - # early. A generator dropped without ``close()`` is finalised by - # the garbage collector instead, possibly inside httpcore's pool - # lock, so callers that stop early must close it (UN-4237). - close_stream(gen) + # Release the stream on success, on error, and when the caller + # close()s this generator early. + close_stream(gen, log=log, description=description) # After ``finally`` so the connection is released, not held # through the backoff. time.sleep(delay) @@ -404,7 +416,7 @@ def collect_with_retry[T]( finally: # On success too: on some providers a drained stream is still # open (see ``close_stream``). - close_stream(gen) + close_stream(gen, log=log, description=description) # After ``finally`` so the connection is released, not held # through the backoff. time.sleep(delay) diff --git a/unstract/sdk1/tests/test_stream_connection_release.py b/unstract/sdk1/tests/test_stream_connection_release.py index ece77011bf..ad9657b01c 100644 --- a/unstract/sdk1/tests/test_stream_connection_release.py +++ b/unstract/sdk1/tests/test_stream_connection_release.py @@ -1,20 +1,10 @@ -"""A streamed completion releases its pooled HTTP connection (UN-4237). - -Executor workers deadlocked under load inside httpcore's connection pool. -LiteLLM's sync stream wrapper has no ``close()``, so a streamed completion -that failed, was abandoned, or (on Anthropic) simply finished before the HTTP -body ended left its response checked out of the pool. The garbage collector -finalised it later, and when that -happened on a thread already inside ``ConnectionPool.handle_request``, the -finaliser's ``PoolByteStream.close()`` waited on the same non-reentrant lock -the thread held: a deadlock no timeout could break. - -These tests run real litellm against a local server speaking the provider's -SSE format, with the garbage collector disabled so it cannot be the one -releasing anything: the Anthropic path through ``LLM.complete()``, and the -OpenAI Responses-API bridge, whose response sits two iterators deep. Then they -reproduce the deadlock's trigger directly: a garbage collection inside the -pool lock on the next call. +"""Streamed completions release their pooled HTTP connection explicitly. + +If the garbage collector finalises an open response while the thread holds +httpcore's non-reentrant pool lock, the thread deadlocks. These tests run real +litellm against a local SSE server with GC disabled, covering the Anthropic +path via ``LLM.complete()`` and the OpenAI Responses-API bridge. Then they +force a collection inside the pool lock to reproduce the deadlock's trigger. """ from __future__ import annotations @@ -367,11 +357,11 @@ def test_gc_inside_the_pool_lock_does_not_deadlock( logging_executor: _LoggingExecutor, no_gc: None, ) -> None: - """The production deadlock, triggered on purpose. + """The pool-lock deadlock, triggered on purpose. - Under load the collector ran on an allocation inside ``handle_request``'s - locked section. Forcing a collection there, after a completed stream, is - what deadlocked every time before the fix. Takes no ``pools`` fixture: + Forcing a collection inside ``handle_request``'s locked section, after a + completed stream, deadlocks unless every stream was closed explicitly. + Takes no ``pools`` fixture: pytest reprs a failing test's arguments, and ``ConnectionPool.__repr__`` takes the very lock a regressed run leaves held. """ @@ -406,10 +396,10 @@ def second_call() -> None: # ── OpenAI Responses-API bridge ────────────────────────────────────────────── # -# The wrapper's ``completion_stream`` is litellm's Responses-to-chat bridge, -# whose ``streaming_response`` is an iterator with no ``close()``; the HTTP -# response sits one level further down. A drained stream reads to the end of -# the body and releases itself, so only failures and early exits show a leak. +# The HTTP response sits at wrapper -> completion_stream -> streaming_response +# -> response, and streaming_response has no close(). A drained stream reads +# to the end of the body and releases itself, so only failures and early +# exits show a leak. def _responses_stream( diff --git a/unstract/sdk1/tests/utils/test_collect_with_retry.py b/unstract/sdk1/tests/utils/test_collect_with_retry.py index fbd0f44382..1ec51fa349 100644 --- a/unstract/sdk1/tests/utils/test_collect_with_retry.py +++ b/unstract/sdk1/tests/utils/test_collect_with_retry.py @@ -155,18 +155,13 @@ def _sleep(_: float) -> None: is_content=_is_content, ) assert result == ["c1"] - # Both attempts: the failed one and the successful one (UN-4237). + # Both attempts: the failed one and the successful one. assert closed == [True, True] -# ── Releasing the HTTP response (UN-4237) ──────────────────────────────────── +# ── Releasing the HTTP response (see retry_utils.close_stream) ────────────── # -# A stream left open is closed by the garbage collector, possibly while the -# thread holds httpcore's non-reentrant pool lock, which deadlocks the worker. -# LiteLLM's sync stream wrapper has no ``close()`` and stops before the HTTP -# body ends, so every stream must be closed explicitly, through the handles -# the wrapper holds. ``test_stream_connection_release.py`` covers this against -# a real socket. +# Real-socket coverage: test_stream_connection_release.py. class _Closable: @@ -236,6 +231,57 @@ def close(self) -> None: assert record.exc_info is not None +def test_failed_close_traceback_omits_the_propagating_provider_error( + caplog: pytest.LogCaptureFixture, +) -> None: + """``close_stream`` runs in a ``finally`` while the provider error unwinds. + + The warning must show the close failure alone, not the provider error + Python chains onto it as context; the caller logs that one itself. + """ + + class _Broken: + def close(self) -> None: + raise RuntimeError("torn down") + + with caplog.at_level(logging.WARNING, logger=retry_utils.logger.name): + try: + raise ValueError("provider returned 429") + except ValueError: + retry_utils.close_stream(_Broken()) + + assert "torn down" in caplog.text + assert "provider returned 429" not in caplog.text + + +def test_failed_close_names_the_adapter_on_the_callers_logger( + caplog: pytest.LogCaptureFixture, +) -> None: + """The retry helpers pass their logger and adapter description through.""" + caller_log = logging.getLogger("tests.caller") + + class _Stream: + def __iter__(self) -> Iterator[str]: + yield "c1" + + def close(self) -> None: + raise RuntimeError("torn down") + + with caplog.at_level(logging.WARNING, logger=caller_log.name): + collect_with_retry( + _Stream, + max_retries=0, + retry_predicate=lambda _: True, + is_content=_is_content, + description="anthropic/claude-sonnet-4-6", + logger_instance=caller_log, + ) + + [record] = caplog.records + assert record.name == caller_log.name + assert "anthropic/claude-sonnet-4-6" in record.getMessage() + + def test_close_stream_survives_a_close_attribute_that_raises( caplog: pytest.LogCaptureFixture, ) -> None: @@ -336,11 +382,11 @@ class _Wrapper: def test_close_stream_reaches_the_responses_api_response() -> None: - """The Responses-API bridge keeps the response two iterators down. + """The Responses-API bridge's response, reached through a non-closable iterator. - Shaped like litellm 1.104.0: the wrapper's ``completion_stream`` is the - bridge, whose ``streaming_response`` is an iterator with no ``close()`` - holding ``response`` and ``stream_iterator``. + Shaped like litellm's Responses-API bridge: the response sits at wrapper + -> ``completion_stream`` -> ``streaming_response`` -> ``response``, and + ``streaming_response`` has no ``close()``. """ log: list[str] = []