From 5bcd3b53c94b44ca4eeabb9206140e2b9bcfbb98 Mon Sep 17 00:00:00 2001 From: Julio Quinteros Pro Date: Fri, 13 Feb 2026 05:25:33 -0300 Subject: [PATCH] fix: import LiteLLM_ObjectPermissionTable from _types instead of proxy_server Per maintainer feedback, FastAPI should always be available in proxy code. The issue was that MCP tests were importing from proxy_server unnecessarily, pulling in all proxy dependencies including policy_resolve_endpoints. Fix: - Revert policy_resolve_endpoints.py to use direct FastAPI imports - Update MCP tests to import LiteLLM_ObjectPermissionTable from litellm.proxy._types instead of litellm.proxy.proxy_server This avoids importing the entire proxy_server module with all its dependencies when tests only need specific types. Addresses: https://github.com/BerriAI/litellm/pull/21075/changes#r2802201174 --- .../policy_engine/policy_resolve_endpoints.py | 53 ++++--------------- tests/mcp_tests/test_mcp_logging.py | 3 +- tests/mcp_tests/test_mcp_server.py | 2 +- 3 files changed, 11 insertions(+), 47 deletions(-) diff --git a/litellm/proxy/policy_engine/policy_resolve_endpoints.py b/litellm/proxy/policy_engine/policy_resolve_endpoints.py index 8f61369450..318e990ff1 100644 --- a/litellm/proxy/policy_engine/policy_resolve_endpoints.py +++ b/litellm/proxy/policy_engine/policy_resolve_endpoints.py @@ -6,32 +6,8 @@ Policy resolve and attachment impact estimation endpoints. """ import json -from typing import TYPE_CHECKING - -# FastAPI imports - may not be available in all environments (e.g. during testing) -try: - from fastapi import APIRouter, Depends, HTTPException, Query -except ImportError: - # Provide stubs for type checking only - if TYPE_CHECKING: - from fastapi import APIRouter, Depends, HTTPException, Query - else: - # Create mock classes that won't be used - class APIRouter: # type: ignore[no-redef] - def post(self, *args, **kwargs): - def decorator(func): - return func - - return decorator - - def Depends(func): - return func # type: ignore[misc] - - HTTPException = Exception # type: ignore[misc,assignment] - - def Query(*args, **kwargs): - return None # type: ignore[misc] +from fastapi import APIRouter, Depends, HTTPException, Query from litellm._logging import verbose_proxy_logger from litellm.constants import MAX_POLICY_ESTIMATE_IMPACT_ROWS @@ -101,9 +77,7 @@ def _get_tags_from_metadata(metadata: object, json_metadata: object = None) -> l async def _fetch_all_teams(prisma_client: object) -> list: """Fetch teams from DB once. Reuse the result across tag and alias lookups.""" return await prisma_client.db.litellm_teamtable.find_many( # type: ignore - where={}, - order={"created_at": "desc"}, - take=MAX_POLICY_ESTIMATE_IMPACT_ROWS, + where={}, order={"created_at": "desc"}, take=MAX_POLICY_ESTIMATE_IMPACT_ROWS, ) @@ -185,8 +159,7 @@ async def _find_affected_by_team_patterns( if matched_team_ids: keys = await prisma_client.db.litellm_verificationtoken.find_many( # type: ignore where={"team_id": {"in": matched_team_ids}}, - order={"created_at": "desc"}, - take=MAX_POLICY_ESTIMATE_IMPACT_ROWS, + order={"created_at": "desc"}, take=MAX_POLICY_ESTIMATE_IMPACT_ROWS, ) for key in keys: key_alias = key.key_alias or "" @@ -208,8 +181,7 @@ async def _find_affected_keys_by_alias( keys = await prisma_client.db.litellm_verificationtoken.find_many( # type: ignore where=_build_alias_where("key_alias", key_patterns), - order={"created_at": "desc"}, - take=MAX_POLICY_ESTIMATE_IMPACT_ROWS, + order={"created_at": "desc"}, take=MAX_POLICY_ESTIMATE_IMPACT_ROWS, ) for key in keys: key_alias = key.key_alias or "" @@ -392,24 +364,19 @@ async def estimate_attachment_impact( # Tag-based impact if tag_patterns: keys = await prisma_client.db.litellm_verificationtoken.find_many( # type: ignore - where={}, - order={"created_at": "desc"}, + where={}, order={"created_at": "desc"}, take=MAX_POLICY_ESTIMATE_IMPACT_ROWS, ) affected_keys, unnamed_keys = _filter_keys_by_tags(keys, tag_patterns) affected_teams, unnamed_teams = _filter_teams_by_tags( - all_teams, - tag_patterns, + all_teams, tag_patterns, ) # Team-based impact (alias matching + keys belonging to those teams) if team_patterns: new_teams, new_keys, new_unnamed = await _find_affected_by_team_patterns( - prisma_client, - all_teams, - team_patterns, - affected_teams, - affected_keys, + prisma_client, all_teams, team_patterns, + affected_teams, affected_keys, ) affected_teams.extend(new_teams) affected_keys.extend(new_keys) @@ -419,9 +386,7 @@ async def estimate_attachment_impact( key_patterns = request.keys or [] if key_patterns: new_keys = await _find_affected_keys_by_alias( - prisma_client, - key_patterns, - affected_keys, + prisma_client, key_patterns, affected_keys, ) affected_keys.extend(new_keys) diff --git a/tests/mcp_tests/test_mcp_logging.py b/tests/mcp_tests/test_mcp_logging.py index aeebb2e913..d9ecb594b7 100644 --- a/tests/mcp_tests/test_mcp_logging.py +++ b/tests/mcp_tests/test_mcp_logging.py @@ -19,8 +19,7 @@ from litellm.proxy._experimental.mcp_server.server import ( from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( MCPServerManager, ) -from litellm.proxy.proxy_server import LiteLLM_ObjectPermissionTable -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import LiteLLM_ObjectPermissionTable, UserAPIKeyAuth from litellm.types.mcp import MCPPostCallResponseObject from litellm.types.utils import HiddenParams from mcp.types import Tool as MCPTool, CallToolResult, TextContent diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index 9718d714cf..98cf3c40c8 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -1,7 +1,6 @@ # Create server parameters for stdio connection import os import sys -from litellm.proxy.proxy_server import LiteLLM_ObjectPermissionTable import pytest from unittest.mock import AsyncMock, MagicMock, patch from contextlib import asynccontextmanager @@ -15,6 +14,7 @@ from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( MCPServer, MCPTransport, ) +from litellm.proxy._types import LiteLLM_ObjectPermissionTable from mcp.types import Tool as MCPTool, CallToolResult, ListToolsResult from mcp.types import TextContent