Merge 45e08ffbf5 into 83ab0113f0

gateway/routes/allowlist.py

@@ -109,6 +109,7 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = (
     # Health & ops
     "/health",
     "/metrics",
+    "/debug/asyncio-tasks",
     "/watsonx",
 )
 

litellm/proxy/common_utils/debug_utils.py

@@ -1,26 +1,143 @@
 # Start tracing memory allocations
 import asyncio
 import gc
+import itertools
 import json
 import os
 import sys
 import tracemalloc
 from collections import Counter
-from collections.abc import Mapping, Sequence
-from typing import Any, Final, NamedTuple, Protocol, TypedDict
+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
-from typing_extensions import ReadOnly
+from typing_extensions import ReadOnly, TypedDict
 
 from litellm import get_secret_str
 from litellm._logging import verbose_proxy_logger
 from litellm.constants import PYTHON_GC_THRESHOLD
-from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy._types import UserAPIKeyAuth, user_api_key_has_admin_view
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 
 router: Final = APIRouter()
 
 
+class _Frame(TypedDict):
+    file: ReadOnly[str]
+    line: ReadOnly[int]
+    function: ReadOnly[str]
+
+
+class _TaskStackGroup(TypedDict):
+    count: ReadOnly[int]
+    coroutine: ReadOnly[str]
+    task_names: ReadOnly[tuple[str, ...]]
+    stack: ReadOnly[tuple[_Frame, ...]]
+
+
+class _TaskStackDump(TypedDict):
+    worker_pid: ReadOnly[int]
+    total_active_tasks: ReadOnly[int]
+    groups: ReadOnly[tuple[_TaskStackGroup, ...]]
+
+
+_TaskStackKey: TypeAlias = tuple[tuple[str, int, str], ...]
+_TaskStackRecord: TypeAlias = tuple[_TaskStackKey, tuple[_Frame, ...], asyncio.Task[object]]
+
+
+def _frame_from_stack_frame(frame: FrameType) -> _Frame:
+    result: Final[_Frame] = {
+        "file": frame.f_code.co_filename,
+        "line": frame.f_lineno,
+        "function": frame.f_code.co_name,
+    }
+    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 _next_awaitable(awaitable: object, _: object) -> object | None:
+    return _awaited(awaitable)
+
+
+def _awaitable_frame_tuple(awaitable: object) -> tuple[FrameType, ...]:
+    frame: Final = _awaitable_frame(awaitable)
+    return (frame,) if frame is not None else ()
+
+
+def _awaited_chain(awaitable: object) -> Iterator[FrameType]:
+    awaitables: Final = itertools.takewhile(
+        lambda current: _awaitable_frame(current) is not None,
+        itertools.accumulate(itertools.repeat(None), _next_awaitable, initial=awaitable),
+    )
+    return itertools.chain.from_iterable(_awaitable_frame_tuple(current) for current in awaitables)
+
+
+def _task_stack(task: asyncio.Task[object], max_frames: int) -> tuple[_Frame, ...]:
+    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:
+    return tuple((frame["file"], frame["line"], frame["function"]) for frame in stack)
+
+
+def _task_coroutine_name(task: asyncio.Task[object]) -> str:
+    coroutine: Final = task.get_coro()
+    coroutine_name: Final = getattr(coroutine, "__qualname__", None)
+    return coroutine_name if isinstance(coroutine_name, str) else repr(coroutine)
+
+
+def _task_stack_group(records: tuple[_TaskStackRecord, ...]) -> _TaskStackGroup:
+    sample_record: Final = records[0]
+    sample_tasks: Final = tuple(record[2] for record in records)
+    result: Final[_TaskStackGroup] = {
+        "count": len(records),
+        "coroutine": _task_coroutine_name(sample_tasks[0]),
+        "task_names": tuple(task.get_name() for task in sample_tasks[:5]),
+        "stack": sample_record[1],
+    }
+    return result
+
+
+def _group_task_stacks(tasks: tuple[asyncio.Task[object], ...], max_frames: int) -> tuple[_TaskStackGroup, ...]:
+    records: Final = tuple(
+        sorted(
+            (
+                (stack_key, stack, task)
+                for task in tasks
+                for stack in (_task_stack(task, max_frames),)
+                for stack_key in (_task_stack_key(stack),)
+            ),
+            key=lambda record: record[0],
+        )
+    )
+    groups: Final = tuple(
+        _task_stack_group(tuple(group)) for _, group in itertools.groupby(records, key=lambda record: record[0])
+    )
+    return tuple(sorted(groups, key=lambda group: group["count"], reverse=True))
+
+
 # Configure garbage collection thresholds from environment variables
 def configure_gc_thresholds():
     """Configure Python garbage collection thresholds from environment variables."""
@@ -87,6 +204,27 @@ async def get_active_tasks_stats():
     }
 
 
+@router.get("/debug/asyncio-tasks/stacks", include_in_schema=False)
+async def get_active_task_stacks(
+    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
+    max_frames: int = Query(default=40, ge=1, le=200),
+) -> _TaskStackDump:
+    if not user_api_key_has_admin_view(user_api_key_dict):
+        raise HTTPException(status_code=403, detail="Only proxy admins can read asyncio task stacks")
+
+    max_tasks_to_check: Final = 5000
+    active_tasks: Final = tuple(task for task in asyncio.all_tasks() if not task.done())
+    current_task: Final = asyncio.current_task()
+    tasks: Final = tuple(task for task in active_tasks if task is not current_task)[:max_tasks_to_check]
+    groups: Final = _group_task_stacks(tasks, max_frames)
+    result: Final[_TaskStackDump] = {
+        "worker_pid": os.getpid(),
+        "total_active_tasks": len(active_tasks),
+        "groups": groups,
+    }
+    return result
+
+
 if os.environ.get("LITELLM_PROFILE", "false").lower() == "true":
     try:
         import objgraph

tests/test_litellm/proxy/common_utils/test_debug_utils.py

@@ -0,0 +1,113 @@
+import asyncio
+
+import httpx
+import pytest
+from fastapi import FastAPI
+
+from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
+from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+from litellm.proxy.common_utils.debug_utils import router as debug_router
+
+
+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)
+    app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=user_role)
+    return app
+
+
+def _client(app: FastAPI) -> httpx.AsyncClient:
+    return httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://testserver")
+
+
+@pytest.mark.asyncio
+async def test_task_stacks_require_proxy_admin() -> None:
+    app = _test_app(LitellmUserRoles.INTERNAL_USER)
+    async with _client(app) as client:
+        response = await client.get("/debug/asyncio-tasks/stacks")
+
+    assert response.status_code == 403
+    assert response.json()["detail"] == "Only proxy admins can read asyncio task stacks"
+
+
+@pytest.mark.asyncio
+async def test_task_stacks_allow_proxy_admin_view_only() -> None:
+    app = _test_app(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY)
+    async with _client(app) as client:
+        response = await client.get("/debug/asyncio-tasks/stacks")
+
+    assert response.status_code == 200
+    assert "groups" in response.json()
+
+
+@pytest.mark.asyncio
+async def test_task_stacks_include_parked_task() -> None:
+    app = _test_app(LitellmUserRoles.PROXY_ADMIN)
+    parked_task = asyncio.create_task(_park_for_test(), name="parked-test-task")
+    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
+    body = response.json()
+    parked_group = next(group for group in body["groups"] if "_park_for_test" in group["coroutine"])
+    assert any(frame["function"] == "_park_for_test" for frame in parked_group["stack"])
+    assert any(frame["file"].endswith("test_debug_utils.py") for frame in parked_group["stack"])
+    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)
+    parked_task = asyncio.create_task(_park_for_test(), name="parked-test-task")
+    try:
+        async with _client(app) as client:
+            response = await client.get("/debug/asyncio-tasks/stacks?max_frames=1")
+    finally:
+        parked_task.cancel()
+        await asyncio.gather(parked_task, return_exceptions=True)
+
+    assert response.status_code == 200
+    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,10 @@ RDS IAM token when ``IAM_TOKEN_DB_AUTH`` is set).
 import os
 import sys
 
+import httpx
+import pytest
+from fastapi import FastAPI
+
 # 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 +63,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():
@@ -196,6 +202,34 @@ def test_gateway_drops_ui_and_swagger_mounts():
             f"Mount {path} must not be served by the gateway"
 
 
+@pytest.mark.asyncio
+async def test_gateway_serves_asyncio_task_debug_routes() -> None:
+    previous_override = app.dependency_overrides.get(user_api_key_auth)
+    app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
+    trimmed = FastAPI(routes=[route for route in app.router.routes if _is_gateway_route(route)])
+    management_route = next(route for route in app.router.routes if getattr(route, "path", None) == "/key/info")
+    try:
+        async with httpx.AsyncClient(
+            transport=httpx.ASGITransport(app=trimmed),
+            base_url="http://testserver",
+        ) as client:
+            stacks_response = await client.get("/debug/asyncio-tasks/stacks")
+            count_response = await client.get("/debug/asyncio-tasks")
+            management_response = await client.get("/key/info")
+    finally:
+        if previous_override is None:
+            app.dependency_overrides.pop(user_api_key_auth, None)
+        else:
+            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
+    assert not _is_gateway_route(management_route)
+    assert management_response.status_code == 404
+
+
 def test_every_app_mount_is_assigned_to_a_component():
     """Every Mount on the proxy app must be consciously assigned to a component.
 

tests/test_litellm/proxy/test_sensitive_route_auth.py

@@ -21,6 +21,7 @@ def _get_route_dependency_calls(router, path: str, method: str):
 def test_sensitive_debug_routes_require_auth_dependency():
     for path, method in (
         ("/debug/asyncio-tasks", "GET"),
+        ("/debug/asyncio-tasks/stacks", "GET"),
         ("/otel-spans", "GET"),
     ):
         assert user_api_key_auth in _get_route_dependency_calls(

ui/litellm-dashboard/src/lib/http/schema.d.ts

@@ -4067,6 +4067,23 @@ export interface paths {
         patch?: never;
         trace?: never;
     };
+    "/debug/asyncio-tasks/stacks": {
+        parameters: {
+            query?: never;
+            header?: never;
+            path?: never;
+            cookie?: never;
+        };
+        /** Get Active Task Stacks */
+        get: operations["get_active_task_stacks_debug_asyncio_tasks_stacks_get"];
+        put?: never;
+        post?: never;
+        delete?: never;
+        options?: never;
+        head?: never;
+        patch?: never;
+        trace?: never;
+    };
     "/debug/memory/details": {
         parameters: {
             query?: never;
@@ -39578,6 +39595,35 @@ export interface components {
             /** Status */
             status?: ("pending" | "running" | "paused" | "completed" | "failed") | null;
         };
+        /** _Frame */
+        _Frame: {
+            /** File */
+            file: string;
+            /** Function */
+            function: string;
+            /** Line */
+            line: number;
+        };
+        /** _TaskStackDump */
+        _TaskStackDump: {
+            /** Groups */
+            groups: components["schemas"]["_TaskStackGroup"][];
+            /** Total Active Tasks */
+            total_active_tasks: number;
+            /** Worker Pid */
+            worker_pid: number;
+        };
+        /** _TaskStackGroup */
+        _TaskStackGroup: {
+            /** Coroutine */
+            coroutine: string;
+            /** Count */
+            count: number;
+            /** Stack */
+            stack: components["schemas"]["_Frame"][];
+            /** Task Names */
+            task_names: string[];
+        };
         /** ModelInfo */
         litellm__proxy___types__ModelInfo: {
             /** Base Model */
@@ -45851,6 +45897,37 @@ export interface operations {
             };
         };
     };
+    get_active_task_stacks_debug_asyncio_tasks_stacks_get: {
+        parameters: {
+            query?: {
+                max_frames?: number;
+            };
+            header?: never;
+            path?: never;
+            cookie?: never;
+        };
+        requestBody?: never;
+        responses: {
+            /** @description Successful Response */
+            200: {
+                headers: {
+                    [name: string]: unknown;
+                };
+                content: {
+                    "application/json": components["schemas"]["_TaskStackDump"];
+                };
+            };
+            /** @description Validation Error */
+            422: {
+                headers: {
+                    [name: string]: unknown;
+                };
+                content: {
+                    "application/json": components["schemas"]["HTTPValidationError"];
+                };
+            };
+        };
+    };
     get_memory_details_debug_memory_details_get: {
         parameters: {
             query?: {