feat(proxy): walk nested awaits in asyncio task stacks

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/proxy/common_utils/debug_utils.py

@@ -7,8 +7,8 @@ import os
 import sys
 import tracemalloc
 from collections import Counter
-from collections.abc import Mapping, Sequence
-from types import FrameType
+from collections.abc import Iterator, Mapping, Sequence
+from types import AsyncGeneratorType, CoroutineType, FrameType, GeneratorType
 from typing import Any, Final, NamedTuple, Protocol, TypeAlias
 
 from fastapi import APIRouter, Depends, HTTPException, Query
@@ -55,8 +55,40 @@ def _frame_from_stack_frame(frame: FrameType) -> _Frame:
     return result
 
 
+def _awaitable_frame(awaitable: object) -> FrameType | None:
+    if isinstance(awaitable, CoroutineType):
+        return awaitable.cr_frame
+    if isinstance(awaitable, GeneratorType):
+        return awaitable.gi_frame
+    if isinstance(awaitable, AsyncGeneratorType):
+        return awaitable.ag_frame
+    return None
+
+
+def _awaited(awaitable: object) -> object | None:
+    if isinstance(awaitable, CoroutineType):
+        return awaitable.cr_await
+    if isinstance(awaitable, GeneratorType):
+        return awaitable.gi_yieldfrom
+    if isinstance(awaitable, AsyncGeneratorType):
+        return awaitable.ag_await
+    return None
+
+
+def _awaited_chain(awaitable: object) -> Iterator[FrameType]:
+    frame: Final = _awaitable_frame(awaitable)
+    if frame is None:
+        return
+    yield frame
+    awaited: Final = _awaited(awaitable)
+    if awaited is not None:
+        yield from _awaited_chain(awaited)
+
+
 def _task_stack(task: asyncio.Task[object], max_frames: int) -> tuple[_Frame, ...]:
-    return tuple(_frame_from_stack_frame(frame) for frame in task.get_stack(limit=max_frames))
+    awaited_frames: Final = tuple(itertools.islice(_awaited_chain(task.get_coro()), max_frames))
+    frames: Final = awaited_frames or tuple(task.get_stack(limit=max_frames))
+    return tuple(_frame_from_stack_frame(frame) for frame in frames)
 
 
 def _task_stack_key(stack: tuple[_Frame, ...]) -> _TaskStackKey:

tests/test_litellm/proxy/common_utils/test_debug_utils.py

@@ -13,6 +13,14 @@ async def _park_for_test() -> None:
     await asyncio.sleep(30)
 
 
+async def _park_inner(event: asyncio.Event) -> None:
+    await event.wait()
+
+
+async def _park_outer(event: asyncio.Event) -> None:
+    await _park_inner(event)
+
+
 def _test_app(user_role: LitellmUserRoles) -> FastAPI:
     app = FastAPI()
     app.include_router(debug_router)
@@ -53,6 +61,29 @@ async def test_task_stacks_include_parked_task() -> None:
     assert body["total_active_tasks"] >= 1
 
 
+@pytest.mark.asyncio
+async def test_task_stacks_include_nested_await_frames() -> None:
+    app = _test_app(LitellmUserRoles.PROXY_ADMIN)
+    event = asyncio.Event()
+    parked_task = asyncio.create_task(_park_outer(event), name="nested-parked-test-task")
+    await asyncio.sleep(0)
+    try:
+        async with _client(app) as client:
+            response = await client.get("/debug/asyncio-tasks/stacks")
+    finally:
+        parked_task.cancel()
+        await asyncio.gather(parked_task, return_exceptions=True)
+
+    assert response.status_code == 200
+    nested_group = next(
+        group for group in response.json()["groups"] if group["task_names"] == ["nested-parked-test-task"]
+    )
+    frames = nested_group["stack"]
+    functions = [frame["function"] for frame in frames]
+    assert functions.index("_park_outer") < functions.index("_park_inner")
+    assert any(frame["file"].endswith("asyncio/locks.py") or frame["function"] == "wait" for frame in frames)
+
+
 @pytest.mark.asyncio
 async def test_task_stacks_respect_max_frames() -> None:
     app = _test_app(LitellmUserRoles.PROXY_ADMIN)
@@ -65,4 +96,8 @@ async def test_task_stacks_respect_max_frames() -> None:
         await asyncio.gather(parked_task, return_exceptions=True)
 
     assert response.status_code == 200
-    assert all(len(group["stack"]) <= 1 for group in response.json()["groups"])
+    groups = response.json()["groups"]
+    parked_group = next(group for group in groups if "_park_for_test" in group["coroutine"])
+    assert len(parked_group["stack"]) == 1
+    assert parked_group["stack"][0]["function"] == "_park_for_test"
+    assert all(len(group["stack"]) <= 1 for group in groups)

tests/test_litellm/proxy/test_component_allowlists.py

@@ -26,6 +26,9 @@ RDS IAM token when ``IAM_TOKEN_DB_AUTH`` is set).
 import os
 import sys
 
+import httpx
+import pytest
+
 # Importing ``litellm.proxy.proxy_server`` runs its module-level setup, which
 # reads ``DATABASE_URL`` (Prisma) and ``LITELLM_MASTER_KEY``. Tier-zero CI
 # runners don't set these. We pin throwaway values before the import so the
@@ -59,6 +62,8 @@ from gateway.routes.allowlist import (
     GATEWAY_MOUNT_PATHS,
     GATEWAY_PATH_PREFIXES,
 )
+from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
+from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 from litellm.proxy.proxy_server import app
 
 for _key, _previous in _PRE_EXISTING_ENV.items():
@@ -80,6 +85,7 @@ _DB_ENV_KEYS = (
 _PRE_DB_ENV = {_key: os.environ.pop(_key, None) for _key in _DB_ENV_KEYS}
 _PRE_COMPONENT_LIFESPAN = app.router.lifespan_context
 from gateway.main import _is_gateway_route
+from gateway.main import app as gateway_app
 
 app.router.lifespan_context = _PRE_COMPONENT_LIFESPAN
 for _key, _previous in _PRE_DB_ENV.items():
@@ -196,9 +202,27 @@ def test_gateway_drops_ui_and_swagger_mounts():
             f"Mount {path} must not be served by the gateway"
 
 
-def test_gateway_keeps_asyncio_task_stacks_route():
-    route = next(route for route in app.router.routes if getattr(route, "path", None) == "/debug/asyncio-tasks/stacks")
-    assert _is_gateway_route(route)
+@pytest.mark.asyncio
+async def test_gateway_serves_asyncio_task_debug_routes() -> None:
+    previous_override = gateway_app.dependency_overrides.get(user_api_key_auth)
+    gateway_app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
+    try:
+        async with httpx.AsyncClient(
+            transport=httpx.ASGITransport(app=gateway_app),
+            base_url="http://testserver",
+        ) as client:
+            stacks_response = await client.get("/debug/asyncio-tasks/stacks")
+            count_response = await client.get("/debug/asyncio-tasks")
+    finally:
+        if previous_override is None:
+            gateway_app.dependency_overrides.pop(user_api_key_auth, None)
+        else:
+            gateway_app.dependency_overrides[user_api_key_auth] = previous_override
+
+    assert stacks_response.status_code == 200
+    assert stacks_response.json()["worker_pid"]
+    assert "groups" in stacks_response.json()
+    assert count_response.status_code == 200
 
 
 def test_every_app_mount_is_assigned_to_a_component():