diff --git a/unstract/sdk1/src/unstract/sdk1/llm.py b/unstract/sdk1/src/unstract/sdk1/llm.py index 241601d168..2b461abd6a 100644 --- a/unstract/sdk1/src/unstract/sdk1/llm.py +++ b/unstract/sdk1/src/unstract/sdk1/llm.py @@ -827,6 +827,13 @@ 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 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 794082e218..ba98cdd851 100644 --- a/unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py +++ b/unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py @@ -234,6 +234,98 @@ 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: +# 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", + "response", + "http_response", + "stream_iterator", +) +# 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 + + +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, + *, + 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 + 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). 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`` 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 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. 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__, + source, + exc_info=close_err, + ) + + def iter_with_retry[T]( fn: Callable[[], Iterable[T]], *, @@ -246,6 +338,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 @@ -258,12 +354,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 +361,13 @@ def iter_with_retry[T]( ) if delay is None: raise - time.sleep(delay) + finally: + # 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) def collect_with_retry[T]( @@ -310,10 +406,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 +413,13 @@ def collect_with_retry[T]( ) if delay is None: raise - time.sleep(delay) + finally: + # On success too: on some providers a drained stream is still + # open (see ``close_stream``). + close_stream(gen, log=log, description=description) + # 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..ad9657b01c --- /dev/null +++ b/unstract/sdk1/tests/test_stream_connection_release.py @@ -0,0 +1,480 @@ +"""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 + +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 +from unstract.sdk1.utils.retry_utils import collect_with_retry, iter_with_retry + +if TYPE_CHECKING: + from collections.abc import Callable, 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"}, +} + + +_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: + 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(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 + + def handle_error(self, *_: object) -> None: + # A client closing mid-stream is the behaviour under test, not a fault. + pass + + +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() + try: + yield srv + finally: + srv.shutdown() + 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.""" + 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 pool-lock deadlock, triggered on purpose. + + 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. + """ + _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] + + +# ── OpenAI Responses-API bridge ────────────────────────────────────────────── +# +# 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( + 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 974a6e6e00..1ec51fa349 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 @@ -141,7 +142,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 +155,350 @@ def close(self) -> None: is_content=_is_content, ) assert result == ["c1"] - assert closed == [True] + # Both attempts: the failed one and the successful one. + assert closed == [True, True] + + +# ── Releasing the HTTP response (see retry_utils.close_stream) ────────────── +# +# Real-socket coverage: test_stream_connection_release.py. + + +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_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 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: + """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. + + 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's response, reached through a non-closable 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] = [] + + 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( + 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