From 45e08ffbf5cfbe29e83209e1528a42e9cf8617b8 Mon Sep 17 00:00:00 2001 From: mateo Date: Fri, 11 Sep 2026 07:27:09 +0000 Subject: [PATCH] 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():