mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
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>
This commit is contained in:
parent
35f81e6491
commit
45e08ffbf5
3 changed files with 23 additions and 8 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue