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():