From 97bc9989435d0afd6a2757945b380381c62a1b46 Mon Sep 17 00:00:00 2001
From: Ryan Crabbe <rcrabbe@berkeley.edu>
Date: Wed, 18 Feb 2026 14:51:06 -0800
Subject: [PATCH] fix: add cancellation shielding to async_data_generator
 cleanup per Greptile review

---
 litellm/proxy/proxy_server.py | 19 ++++++++++++-------
 1 file changed, 12 insertions(+), 7 deletions(-)

diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 390d68813e6..9d9734148b7 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -5146,13 +5146,18 @@ async def async_data_generator(
         # Close the response stream to release the underlying HTTP connection
         # back to the connection pool. This prevents pool exhaustion when
         # clients disconnect mid-stream.
-        if hasattr(response, "aclose"):
-            try:
-                await response.aclose()
-            except Exception as e:
-                verbose_proxy_logger.debug(
-                    "async_data_generator: error closing response stream: %s", e
-                )
+        # Shield from cancellation so the close awaits can complete.
+        import anyio
+
+        with anyio.CancelScope(shield=True):
+            if hasattr(response, "aclose"):
+                try:
+                    await response.aclose()
+                except Exception as e:
+                    verbose_proxy_logger.debug(
+                        "async_data_generator: error closing response stream: %s",
+                        e,
+                    )
 
 
 def select_data_generator(