UN-4237 [FIX] Close streamed LLM responses so executor workers no longer deadlock on httpcore's pool lock

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)

unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py

@@ -234,6 +234,98 @@
     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 @@
 
     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,20 +354,20 @@
                 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(
                 e, attempt, max_retries, retry_predicate, description, log
             )
             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,18 +406,20 @@
                 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(
                 e, attempt, max_retries, retry_predicate, description, log
             )
             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