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" + )