From 259220e2d8f8ff9a9ee1278410495677d1d08e48 Mon Sep 17 00:00:00 2001 From: mateo Date: Fri, 11 Sep 2026 05:43:10 +0000 Subject: [PATCH 1/6] feat(proxy): add admin-only /debug/asyncio-tasks/stacks endpoint Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- gateway/routes/allowlist.py | 1 + litellm/proxy/common_utils/debug_utils.py | 97 ++++++++++++++++++- .../proxy/common_utils/test_debug_utils.py | 68 +++++++++++++ .../proxy/test_component_allowlists.py | 47 +++++---- .../proxy/test_sensitive_route_auth.py | 15 +-- 5 files changed, 191 insertions(+), 37 deletions(-) create mode 100644 tests/test_litellm/proxy/common_utils/test_debug_utils.py diff --git a/gateway/routes/allowlist.py b/gateway/routes/allowlist.py index 92b73867e67..65a0a259111 100644 --- a/gateway/routes/allowlist.py +++ b/gateway/routes/allowlist.py @@ -109,6 +109,7 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = ( # Health & ops "/health", "/metrics", + "/debug/asyncio-tasks", "/watsonx", ) diff --git a/litellm/proxy/common_utils/debug_utils.py b/litellm/proxy/common_utils/debug_utils.py index 554a6ae8d1a..942cd0d82a6 100644 --- a/litellm/proxy/common_utils/debug_utils.py +++ b/litellm/proxy/common_utils/debug_utils.py @@ -7,7 +7,8 @@ import sys import tracemalloc from collections import Counter from collections.abc import Mapping, Sequence -from typing import Any, Final, NamedTuple, Protocol, TypedDict +from types import FrameType +from typing import Any, Final, NamedTuple, Protocol, TypeAlias, TypedDict from fastapi import APIRouter, Depends, HTTPException, Query from typing_extensions import ReadOnly @@ -15,12 +16,83 @@ from typing_extensions import ReadOnly 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 LitellmUserRoles, UserAPIKeyAuth 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 _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)) + + +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(stack_key: _TaskStackKey, records: tuple[_TaskStackRecord, ...]) -> _TaskStackGroup: + matching_records: Final = tuple(record for record in records if record[0] == stack_key) + sample_record: Final = matching_records[0] + sample_tasks: Final = tuple(record[2] for record in matching_records) + result: Final[_TaskStackGroup] = { + "count": len(matching_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( + (stack_key, stack, task) + for task in tasks + for stack in (_task_stack(task, max_frames),) + for stack_key in (_task_stack_key(stack),) + ) + stack_keys: Final = tuple(dict.fromkeys(record[0] for record in records)) + groups: Final = tuple(_task_stack_group(stack_key, records) for stack_key in stack_keys) + 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 +159,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 user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + 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 diff --git a/tests/test_litellm/proxy/common_utils/test_debug_utils.py b/tests/test_litellm/proxy/common_utils/test_debug_utils.py new file mode 100644 index 00000000000..3f39f8f152c --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_debug_utils.py @@ -0,0 +1,68 @@ +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) + + +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_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_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 + assert all(len(group["stack"]) <= 1 for group in response.json()["groups"]) diff --git a/tests/test_litellm/proxy/test_component_allowlists.py b/tests/test_litellm/proxy/test_component_allowlists.py index 0fdb43d60da..c2beebf40e6 100644 --- a/tests/test_litellm/proxy/test_component_allowlists.py +++ b/tests/test_litellm/proxy/test_component_allowlists.py @@ -108,12 +108,8 @@ def test_gateway_plus_backend_covers_full_app(): for r in app.router.routes if not isinstance(r, Mount) and getattr(r, "path", None) is not None } - gateway_paths = _component_paths( - app.router.routes, GATEWAY_EXACT_PATHS, GATEWAY_PATH_PREFIXES - ) - backend_paths = _component_paths( - app.router.routes, BACKEND_EXACT_PATHS, BACKEND_PATH_PREFIXES - ) + gateway_paths = _component_paths(app.router.routes, GATEWAY_EXACT_PATHS, GATEWAY_PATH_PREFIXES) + backend_paths = _component_paths(app.router.routes, BACKEND_EXACT_PATHS, BACKEND_PATH_PREFIXES) uncovered = all_paths - (gateway_paths | backend_paths) @@ -126,16 +122,15 @@ def test_gateway_plus_backend_covers_full_app(): def test_backend_mount_paths_defined(): """BACKEND_MOUNT_PATHS constant must exist and be a frozenset.""" - assert isinstance(BACKEND_MOUNT_PATHS, frozenset), \ + assert isinstance(BACKEND_MOUNT_PATHS, frozenset), ( f"BACKEND_MOUNT_PATHS must be a frozenset, got {type(BACKEND_MOUNT_PATHS)}" - assert len(BACKEND_MOUNT_PATHS) > 0, \ - "BACKEND_MOUNT_PATHS must contain at least one Mount path" + ) + assert len(BACKEND_MOUNT_PATHS) > 0, "BACKEND_MOUNT_PATHS must contain at least one Mount path" def test_swagger_mount_in_backend_allowlist(): """The /swagger Mount must be in BACKEND_MOUNT_PATHS.""" - assert "/swagger" in BACKEND_MOUNT_PATHS, \ - "/swagger Mount path must be in BACKEND_MOUNT_PATHS" + assert "/swagger" in BACKEND_MOUNT_PATHS, "/swagger Mount path must be in BACKEND_MOUNT_PATHS" def test_backend_keeps_swagger_mount(): @@ -145,32 +140,31 @@ def test_backend_keeps_swagger_mount(): for r in app.router.routes if isinstance(r, Mount) and getattr(r, "path", None) in BACKEND_MOUNT_PATHS } - assert "/swagger" in backend_mounts, \ + assert "/swagger" in backend_mounts, ( "/swagger Mount is expected on the proxy app and should be in BACKEND_MOUNT_PATHS" + ) def test_backend_drops_non_allowlisted_mounts(): """Verify that Mounts NOT in BACKEND_MOUNT_PATHS would be dropped from backend.""" all_mounts = { - getattr(r, "path") - for r in app.router.routes - if isinstance(r, Mount) and getattr(r, "path", None) is not None + getattr(r, "path") for r in app.router.routes if isinstance(r, Mount) and getattr(r, "path", None) is not None } non_backend_mounts = all_mounts - BACKEND_MOUNT_PATHS - assert len(non_backend_mounts) > 0, \ + assert len(non_backend_mounts) > 0, ( "Expected at least one non-backend Mount (e.g., /ui, /_next) to verify filtering logic" + ) for mount_path in non_backend_mounts: - assert mount_path not in BACKEND_MOUNT_PATHS, \ - f"Mount {mount_path} should not be in BACKEND_MOUNT_PATHS" + assert mount_path not in BACKEND_MOUNT_PATHS, f"Mount {mount_path} should not be in BACKEND_MOUNT_PATHS" def test_gateway_mount_paths_defined(): """GATEWAY_MOUNT_PATHS constant must exist and expose /metrics.""" - assert isinstance(GATEWAY_MOUNT_PATHS, frozenset), \ + assert isinstance(GATEWAY_MOUNT_PATHS, frozenset), ( f"GATEWAY_MOUNT_PATHS must be a frozenset, got {type(GATEWAY_MOUNT_PATHS)}" - assert "/metrics" in GATEWAY_MOUNT_PATHS, \ - "/metrics Mount path must be in GATEWAY_MOUNT_PATHS" + ) + assert "/metrics" in GATEWAY_MOUNT_PATHS, "/metrics Mount path must be in GATEWAY_MOUNT_PATHS" def test_gateway_trim_keeps_metrics_mount(): @@ -185,15 +179,20 @@ def test_gateway_trim_keeps_metrics_mount(): metrics_mount = Mount("/metrics", app=make_asgi_app()) routes = [*app.router.routes, metrics_mount] trimmed = [r for r in routes if _is_gateway_route(r)] - assert metrics_mount in trimmed, \ - "/metrics Mount must survive the gateway route trim" + assert metrics_mount in trimmed, "/metrics Mount must survive the gateway route trim" def test_gateway_drops_ui_and_swagger_mounts(): """UI static and swagger Mounts must still be trimmed from the gateway.""" for path in ("/ui", "/_next", "/litellm-asset-prefix/_next", "/swagger"): - assert not _is_gateway_route(Mount(path, app=make_asgi_app())), \ + assert not _is_gateway_route(Mount(path, app=make_asgi_app())), ( 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) def test_every_app_mount_is_assigned_to_a_component(): diff --git a/tests/test_litellm/proxy/test_sensitive_route_auth.py b/tests/test_litellm/proxy/test_sensitive_route_auth.py index 19998e52779..fa91b09b46b 100644 --- a/tests/test_litellm/proxy/test_sensitive_route_auth.py +++ b/tests/test_litellm/proxy/test_sensitive_route_auth.py @@ -9,11 +9,7 @@ from litellm.proxy.spend_tracking.spend_management_endpoints import ( def _get_route_dependency_calls(router, path: str, method: str): for route in router.routes: - if ( - isinstance(route, APIRoute) - and route.path == path - and method in route.methods - ): + if isinstance(route, APIRoute) and route.path == path and method in route.methods: return [dependency.call for dependency in route.dependant.dependencies] raise AssertionError(f"Route {method} {path} not found") @@ -21,14 +17,11 @@ 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( - debug_router, path, method - ) + assert user_api_key_auth in _get_route_dependency_calls(debug_router, path, method) def test_provider_budgets_requires_auth_dependency(): - assert user_api_key_auth in _get_route_dependency_calls( - spend_router, "/provider/budgets", "GET" - ) + assert user_api_key_auth in _get_route_dependency_calls(spend_router, "/provider/budgets", "GET") From dcea5fbe67d93befbf6e3753dff48b82c1202c4d Mon Sep 17 00:00:00 2001 From: mateo Date: Fri, 11 Sep 2026 05:46:27 +0000 Subject: [PATCH 2/6] refactor(proxy): group asyncio task stacks in one pass Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_utils/debug_utils.py | 28 ++++++++----- .../proxy/test_component_allowlists.py | 42 +++++++++++-------- .../proxy/test_sensitive_route_auth.py | 14 +++++-- 3 files changed, 52 insertions(+), 32 deletions(-) diff --git a/litellm/proxy/common_utils/debug_utils.py b/litellm/proxy/common_utils/debug_utils.py index 942cd0d82a6..d9d7b70fe3a 100644 --- a/litellm/proxy/common_utils/debug_utils.py +++ b/litellm/proxy/common_utils/debug_utils.py @@ -1,6 +1,7 @@ # Start tracing memory allocations import asyncio import gc +import itertools import json import os import sys @@ -68,12 +69,11 @@ def _task_coroutine_name(task: asyncio.Task[object]) -> str: return coroutine_name if isinstance(coroutine_name, str) else repr(coroutine) -def _task_stack_group(stack_key: _TaskStackKey, records: tuple[_TaskStackRecord, ...]) -> _TaskStackGroup: - matching_records: Final = tuple(record for record in records if record[0] == stack_key) - sample_record: Final = matching_records[0] - sample_tasks: Final = tuple(record[2] for record in matching_records) +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(matching_records), + "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], @@ -83,13 +83,19 @@ def _task_stack_group(stack_key: _TaskStackKey, records: tuple[_TaskStackRecord, def _group_task_stacks(tasks: tuple[asyncio.Task[object], ...], max_frames: int) -> tuple[_TaskStackGroup, ...]: records: Final = tuple( - (stack_key, stack, task) - for task in tasks - for stack in (_task_stack(task, max_frames),) - for stack_key in (_task_stack_key(stack),) + 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]) ) - stack_keys: Final = tuple(dict.fromkeys(record[0] for record in records)) - groups: Final = tuple(_task_stack_group(stack_key, records) for stack_key in stack_keys) return tuple(sorted(groups, key=lambda group: group["count"], reverse=True)) diff --git a/tests/test_litellm/proxy/test_component_allowlists.py b/tests/test_litellm/proxy/test_component_allowlists.py index c2beebf40e6..4e65f190a14 100644 --- a/tests/test_litellm/proxy/test_component_allowlists.py +++ b/tests/test_litellm/proxy/test_component_allowlists.py @@ -108,8 +108,12 @@ def test_gateway_plus_backend_covers_full_app(): for r in app.router.routes if not isinstance(r, Mount) and getattr(r, "path", None) is not None } - gateway_paths = _component_paths(app.router.routes, GATEWAY_EXACT_PATHS, GATEWAY_PATH_PREFIXES) - backend_paths = _component_paths(app.router.routes, BACKEND_EXACT_PATHS, BACKEND_PATH_PREFIXES) + gateway_paths = _component_paths( + app.router.routes, GATEWAY_EXACT_PATHS, GATEWAY_PATH_PREFIXES + ) + backend_paths = _component_paths( + app.router.routes, BACKEND_EXACT_PATHS, BACKEND_PATH_PREFIXES + ) uncovered = all_paths - (gateway_paths | backend_paths) @@ -122,15 +126,16 @@ def test_gateway_plus_backend_covers_full_app(): def test_backend_mount_paths_defined(): """BACKEND_MOUNT_PATHS constant must exist and be a frozenset.""" - assert isinstance(BACKEND_MOUNT_PATHS, frozenset), ( + assert isinstance(BACKEND_MOUNT_PATHS, frozenset), \ f"BACKEND_MOUNT_PATHS must be a frozenset, got {type(BACKEND_MOUNT_PATHS)}" - ) - assert len(BACKEND_MOUNT_PATHS) > 0, "BACKEND_MOUNT_PATHS must contain at least one Mount path" + assert len(BACKEND_MOUNT_PATHS) > 0, \ + "BACKEND_MOUNT_PATHS must contain at least one Mount path" def test_swagger_mount_in_backend_allowlist(): """The /swagger Mount must be in BACKEND_MOUNT_PATHS.""" - assert "/swagger" in BACKEND_MOUNT_PATHS, "/swagger Mount path must be in BACKEND_MOUNT_PATHS" + assert "/swagger" in BACKEND_MOUNT_PATHS, \ + "/swagger Mount path must be in BACKEND_MOUNT_PATHS" def test_backend_keeps_swagger_mount(): @@ -140,31 +145,32 @@ def test_backend_keeps_swagger_mount(): for r in app.router.routes if isinstance(r, Mount) and getattr(r, "path", None) in BACKEND_MOUNT_PATHS } - assert "/swagger" in backend_mounts, ( + assert "/swagger" in backend_mounts, \ "/swagger Mount is expected on the proxy app and should be in BACKEND_MOUNT_PATHS" - ) def test_backend_drops_non_allowlisted_mounts(): """Verify that Mounts NOT in BACKEND_MOUNT_PATHS would be dropped from backend.""" all_mounts = { - getattr(r, "path") for r in app.router.routes if isinstance(r, Mount) and getattr(r, "path", None) is not None + getattr(r, "path") + for r in app.router.routes + if isinstance(r, Mount) and getattr(r, "path", None) is not None } non_backend_mounts = all_mounts - BACKEND_MOUNT_PATHS - assert len(non_backend_mounts) > 0, ( + assert len(non_backend_mounts) > 0, \ "Expected at least one non-backend Mount (e.g., /ui, /_next) to verify filtering logic" - ) for mount_path in non_backend_mounts: - assert mount_path not in BACKEND_MOUNT_PATHS, f"Mount {mount_path} should not be in BACKEND_MOUNT_PATHS" + assert mount_path not in BACKEND_MOUNT_PATHS, \ + f"Mount {mount_path} should not be in BACKEND_MOUNT_PATHS" def test_gateway_mount_paths_defined(): """GATEWAY_MOUNT_PATHS constant must exist and expose /metrics.""" - assert isinstance(GATEWAY_MOUNT_PATHS, frozenset), ( + assert isinstance(GATEWAY_MOUNT_PATHS, frozenset), \ f"GATEWAY_MOUNT_PATHS must be a frozenset, got {type(GATEWAY_MOUNT_PATHS)}" - ) - assert "/metrics" in GATEWAY_MOUNT_PATHS, "/metrics Mount path must be in GATEWAY_MOUNT_PATHS" + assert "/metrics" in GATEWAY_MOUNT_PATHS, \ + "/metrics Mount path must be in GATEWAY_MOUNT_PATHS" def test_gateway_trim_keeps_metrics_mount(): @@ -179,15 +185,15 @@ def test_gateway_trim_keeps_metrics_mount(): metrics_mount = Mount("/metrics", app=make_asgi_app()) routes = [*app.router.routes, metrics_mount] trimmed = [r for r in routes if _is_gateway_route(r)] - assert metrics_mount in trimmed, "/metrics Mount must survive the gateway route trim" + assert metrics_mount in trimmed, \ + "/metrics Mount must survive the gateway route trim" def test_gateway_drops_ui_and_swagger_mounts(): """UI static and swagger Mounts must still be trimmed from the gateway.""" for path in ("/ui", "/_next", "/litellm-asset-prefix/_next", "/swagger"): - assert not _is_gateway_route(Mount(path, app=make_asgi_app())), ( + assert not _is_gateway_route(Mount(path, app=make_asgi_app())), \ f"Mount {path} must not be served by the gateway" - ) def test_gateway_keeps_asyncio_task_stacks_route(): diff --git a/tests/test_litellm/proxy/test_sensitive_route_auth.py b/tests/test_litellm/proxy/test_sensitive_route_auth.py index fa91b09b46b..e420041da45 100644 --- a/tests/test_litellm/proxy/test_sensitive_route_auth.py +++ b/tests/test_litellm/proxy/test_sensitive_route_auth.py @@ -9,7 +9,11 @@ from litellm.proxy.spend_tracking.spend_management_endpoints import ( def _get_route_dependency_calls(router, path: str, method: str): for route in router.routes: - if isinstance(route, APIRoute) and route.path == path and method in route.methods: + if ( + isinstance(route, APIRoute) + and route.path == path + and method in route.methods + ): return [dependency.call for dependency in route.dependant.dependencies] raise AssertionError(f"Route {method} {path} not found") @@ -20,8 +24,12 @@ def test_sensitive_debug_routes_require_auth_dependency(): ("/debug/asyncio-tasks/stacks", "GET"), ("/otel-spans", "GET"), ): - assert user_api_key_auth in _get_route_dependency_calls(debug_router, path, method) + assert user_api_key_auth in _get_route_dependency_calls( + debug_router, path, method + ) def test_provider_budgets_requires_auth_dependency(): - assert user_api_key_auth in _get_route_dependency_calls(spend_router, "/provider/budgets", "GET") + assert user_api_key_auth in _get_route_dependency_calls( + spend_router, "/provider/budgets", "GET" + ) From 7e7354cc73a63f2d1543a3664709d5f2955778a7 Mon Sep 17 00:00:00 2001 From: mateo Date: Fri, 11 Sep 2026 05:58:02 +0000 Subject: [PATCH 3/6] fix(proxy): restore task stack CI compatibility Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_utils/debug_utils.py | 4 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 77 +++++++++++++++++++ 2 files changed, 79 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/common_utils/debug_utils.py b/litellm/proxy/common_utils/debug_utils.py index d9d7b70fe3a..f496b9718a4 100644 --- a/litellm/proxy/common_utils/debug_utils.py +++ b/litellm/proxy/common_utils/debug_utils.py @@ -9,10 +9,10 @@ import tracemalloc from collections import Counter from collections.abc import Mapping, Sequence from types import FrameType -from typing import Any, Final, NamedTuple, Protocol, TypeAlias, TypedDict +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 diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 29435c31aee..dabf81e7344 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -4042,6 +4042,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; @@ -39628,6 +39645,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 */ @@ -45869,6 +45915,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?: { From a9131b216e8f369870f153baa5cae674853de2a8 Mon Sep 17 00:00:00 2001 From: mateo Date: Fri, 11 Sep 2026 06:14:55 +0000 Subject: [PATCH 4/6] 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 | 38 +++++++++++++++++-- .../proxy/common_utils/test_debug_utils.py | 37 +++++++++++++++++- .../proxy/test_component_allowlists.py | 30 +++++++++++++-- 3 files changed, 98 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/common_utils/debug_utils.py b/litellm/proxy/common_utils/debug_utils.py index f496b9718a4..2adc61cc5ca 100644 --- a/litellm/proxy/common_utils/debug_utils.py +++ b/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: diff --git a/tests/test_litellm/proxy/common_utils/test_debug_utils.py b/tests/test_litellm/proxy/common_utils/test_debug_utils.py index 3f39f8f152c..9a4a42f6ef6 100644 --- a/tests/test_litellm/proxy/common_utils/test_debug_utils.py +++ b/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) diff --git a/tests/test_litellm/proxy/test_component_allowlists.py b/tests/test_litellm/proxy/test_component_allowlists.py index 4e65f190a14..a937fd8e2bd 100644 --- a/tests/test_litellm/proxy/test_component_allowlists.py +++ b/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(): From 35f81e6491653fc2f389f34fd29ac2f909692489 Mon Sep 17 00:00:00 2001 From: mateo Date: Fri, 11 Sep 2026 06:21:35 +0000 Subject: [PATCH 5/6] fix(proxy): avoid recursive task stack traversal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_utils/debug_utils.py | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/common_utils/debug_utils.py b/litellm/proxy/common_utils/debug_utils.py index 2adc61cc5ca..ebc4d211560 100644 --- a/litellm/proxy/common_utils/debug_utils.py +++ b/litellm/proxy/common_utils/debug_utils.py @@ -75,14 +75,21 @@ def _awaited(awaitable: object) -> object | None: return None -def _awaited_chain(awaitable: object) -> Iterator[FrameType]: +def _next_awaitable(awaitable: object, _: object) -> object | None: + return _awaited(awaitable) + + +def _awaitable_frame_tuple(awaitable: object) -> tuple[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) + 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, ...]: From 45e08ffbf5cfbe29e83209e1528a42e9cf8617b8 Mon Sep 17 00:00:00 2001 From: mateo Date: Fri, 11 Sep 2026 07:27:09 +0000 Subject: [PATCH 6/6] fix(proxy): admit view-only admins and test the gateway trim for task stacks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_utils/debug_utils.py | 4 ++-- .../proxy/common_utils/test_debug_utils.py | 10 ++++++++++ .../proxy/test_component_allowlists.py | 17 +++++++++++------ 3 files changed, 23 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/common_utils/debug_utils.py b/litellm/proxy/common_utils/debug_utils.py index ebc4d211560..f500526d52e 100644 --- a/litellm/proxy/common_utils/debug_utils.py +++ b/litellm/proxy/common_utils/debug_utils.py @@ -17,7 +17,7 @@ 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 LitellmUserRoles, 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() @@ -209,7 +209,7 @@ 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 user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + 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 diff --git a/tests/test_litellm/proxy/common_utils/test_debug_utils.py b/tests/test_litellm/proxy/common_utils/test_debug_utils.py index 9a4a42f6ef6..75b2950a9d9 100644 --- a/tests/test_litellm/proxy/common_utils/test_debug_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_debug_utils.py @@ -42,6 +42,16 @@ async def test_task_stacks_require_proxy_admin() -> None: 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) diff --git a/tests/test_litellm/proxy/test_component_allowlists.py b/tests/test_litellm/proxy/test_component_allowlists.py index a937fd8e2bd..8618f2f4438 100644 --- a/tests/test_litellm/proxy/test_component_allowlists.py +++ b/tests/test_litellm/proxy/test_component_allowlists.py @@ -28,6 +28,7 @@ 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 @@ -85,7 +86,6 @@ _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(): @@ -204,25 +204,30 @@ def test_gateway_drops_ui_and_swagger_mounts(): @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) + 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=gateway_app), + 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: - gateway_app.dependency_overrides.pop(user_api_key_auth, None) + app.dependency_overrides.pop(user_api_key_auth, None) else: - gateway_app.dependency_overrides[user_api_key_auth] = previous_override + 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():