From 21ed38971d244c0a034604f6439c0584d55b4d20 Mon Sep 17 00:00:00 2001 From: Michael-RZ-Berri Date: Tue, 28 Apr 2026 17:04:40 -0700 Subject: [PATCH] lazy-load optional feature routers on first request (#26534) Co-authored-by: Michael Riad Zaky --- litellm/proxy/_lazy_features.py | 307 ++++++++++++++++++ litellm/proxy/proxy_server.py | 126 ++----- tests/proxy_unit_tests/test_proxy_routes.py | 14 + tests/test_litellm/proxy/test_proxy_server.py | 252 ++++++++++++++ .../test_vector_store_endpoints.py | 15 + 5 files changed, 609 insertions(+), 105 deletions(-) create mode 100644 litellm/proxy/_lazy_features.py diff --git a/litellm/proxy/_lazy_features.py b/litellm/proxy/_lazy_features.py new file mode 100644 index 00000000000..450c9483f31 --- /dev/null +++ b/litellm/proxy/_lazy_features.py @@ -0,0 +1,307 @@ +""" +Lazy registration for optional feature routers. Each LAZY_FEATURES entry +imports its module only on the first request matching its path prefix, +saving ~700 MB at idle for deployments that don't use these features. +First hit pays the import cost (1-3 s for heavy modules); /openapi.json +omits each feature's routes until the feature is warmed. +""" + +import asyncio +import importlib +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Callable, Tuple + +from starlette.types import Receive, Scope, Send + +from litellm._logging import verbose_proxy_logger + +if TYPE_CHECKING: + from fastapi import FastAPI + + +def _include_router(attr_name: str = "router") -> Callable[["FastAPI", object], None]: + def _register(app: "FastAPI", module: object) -> None: + app.include_router(getattr(module, attr_name)) + + return _register + + +def _mount_app( + prefix: str, attr_name: str = "app" +) -> Callable[["FastAPI", object], None]: + def _register(app: "FastAPI", module: object) -> None: + app.mount(path=prefix, app=getattr(module, attr_name)) + + return _register + + +@dataclass(frozen=True) +class LazyFeature: + name: str + module_path: str + path_prefixes: Tuple[str, ...] + register_fn: Callable[["FastAPI", object], None] = field( + default_factory=lambda: _include_router("router") + ) + # For routes whose path has a leading parameter (e.g. /{server}/authorize) + # — startswith can't match those, so the matcher also checks endswith. + path_suffixes: Tuple[str, ...] = () + + +LAZY_FEATURES: Tuple[LazyFeature, ...] = ( + LazyFeature( + name="guardrails", + module_path="litellm.proxy.guardrails.guardrail_endpoints", + path_prefixes=( + "/guardrails", + "/v2/guardrails", + "/apply_guardrail", + "/policies/usage", + ), + ), + LazyFeature( + name="policies", + module_path="litellm.proxy.management_endpoints.policy_endpoints", + # Trailing slash to avoid matching /policies/... (policy_engine). + path_prefixes=("/policy/", "/utils/test_policies_and_guardrails"), + ), + LazyFeature( + name="policy_engine", + module_path="litellm.proxy.policy_engine.policy_endpoints", + path_prefixes=("/policies",), + ), + LazyFeature( + name="policy_resolve", + module_path="litellm.proxy.policy_engine.policy_resolve_endpoints", + path_prefixes=("/policies/resolve", "/policies/attachments/estimate-impact"), + ), + LazyFeature( + name="agents", + module_path="litellm.proxy.agent_endpoints.endpoints", + path_prefixes=("/v1/agents", "/agents", "/agent/"), + ), + LazyFeature( + name="a2a", + module_path="litellm.proxy.agent_endpoints.a2a_endpoints", + path_prefixes=("/a2a", "/v1/a2a"), + ), + LazyFeature( + name="vector_stores", + module_path="litellm.proxy.vector_store_endpoints.endpoints", + path_prefixes=("/v1/vector_stores", "/vector_stores", "/v1/indexes"), + ), + LazyFeature( + name="vector_store_management", + module_path="litellm.proxy.vector_store_endpoints.management_endpoints", + # Trailing slash to avoid matching /vector_stores/... (vector_stores). + path_prefixes=("/vector_store/", "/v1/vector_store/"), + ), + LazyFeature( + name="vector_store_files", + # Routes appear under both /v1/vector_stores/{id}/files and the + # un-versioned form, so both prefixes must trigger the load. + module_path="litellm.proxy.vector_store_files_endpoints.endpoints", + path_prefixes=("/v1/vector_stores", "/vector_stores"), + ), + LazyFeature( + name="tools", + module_path="litellm.proxy.management_endpoints.tool_management_endpoints", + path_prefixes=("/v1/tool", "/tool"), + ), + LazyFeature( + name="search_tools", + module_path="litellm.proxy.search_endpoints.search_tool_management", + path_prefixes=("/search_tools",), + ), + # mcp_management owns most /v1/mcp/* admin routes; mcp_app is the mounted + # streaming sub-app at /mcp. + LazyFeature( + name="mcp_management", + module_path="litellm.proxy.management_endpoints.mcp_management_endpoints", + path_prefixes=("/v1/mcp/",), + ), + LazyFeature( + # Also serves /.well-known/oauth-* (OAuth metadata discovery). + # No /mcp/oauth prefix here: the mounted /mcp sub-app would + # shadow it, and there are no actual routes there anyway. + name="mcp_byok_oauth", + module_path="litellm.proxy._experimental.mcp_server.byok_oauth_endpoints", + path_prefixes=("/v1/mcp/oauth", "/.well-known/oauth-"), + ), + LazyFeature( + # Serves OAuth dance endpoints (/authorize, /token, /callback, + # /register) plus several /.well-known/ discovery URLs at the proxy + # root — needed for MCP-over-OAuth flows even before /mcp is hit. + name="mcp_discoverable", + module_path="litellm.proxy._experimental.mcp_server.discoverable_endpoints", + path_prefixes=( + "/.well-known/oauth-", + "/.well-known/openid-configuration", + "/.well-known/jwks.json", + "/authorize", + "/token", + "/callback", + "/register", + ), + # Catches the /{mcp_server_name}/authorize|token|register variants. + path_suffixes=("/authorize", "/token", "/register"), + ), + LazyFeature( + name="mcp_rest", + module_path="litellm.proxy._experimental.mcp_server.rest_endpoints", + path_prefixes=("/mcp-rest",), + ), + LazyFeature( + # Hardcoded /mcp matches BASE_MCP_ROUTE; importing the constant + # here would defeat lazy loading. + name="mcp_app", + module_path="litellm.proxy._experimental.mcp_server.server", + path_prefixes=("/mcp",), + register_fn=_mount_app("/mcp", attr_name="app"), + ), + LazyFeature( + name="config_overrides", + module_path="litellm.proxy.management_endpoints.config_override_endpoints", + path_prefixes=("/config_overrides",), + ), + LazyFeature( + name="realtime", + module_path="litellm.proxy.realtime_endpoints.endpoints", + path_prefixes=("/openai/v1/realtime", "/v1/realtime", "/realtime"), + ), + LazyFeature( + name="anthropic_passthrough", + module_path="litellm.proxy.anthropic_endpoints.endpoints", + path_prefixes=("/v1/messages", "/anthropic", "/api/event_logging"), + ), + LazyFeature( + name="anthropic_skills", + module_path="litellm.proxy.anthropic_endpoints.skills_endpoints", + path_prefixes=("/v1/skills", "/skills"), + ), + LazyFeature( + name="langfuse_passthrough", + module_path="litellm.proxy.vertex_ai_endpoints.langfuse_endpoints", + path_prefixes=("/langfuse",), + ), + LazyFeature( + name="evals", + module_path="litellm.proxy.openai_evals_endpoints.endpoints", + path_prefixes=("/v1/evals", "/evals"), + ), + LazyFeature( + name="claude_code_marketplace", + module_path="litellm.proxy.anthropic_endpoints.claude_code_endpoints", + path_prefixes=("/claude-code",), + register_fn=_include_router("claude_code_marketplace_router"), + ), + LazyFeature( + name="scim", + module_path="litellm.proxy.management_endpoints.scim.scim_v2", + path_prefixes=("/scim",), + register_fn=_include_router("scim_router"), + ), + LazyFeature( + name="cloudzero", + module_path="litellm.proxy.spend_tracking.cloudzero_endpoints", + path_prefixes=("/cloudzero",), + ), + LazyFeature( + name="vantage", + module_path="litellm.proxy.spend_tracking.vantage_endpoints", + path_prefixes=("/vantage",), + ), + LazyFeature( + name="usage_ai", + module_path="litellm.proxy.management_endpoints.usage_endpoints", + path_prefixes=("/usage/ai",), + ), + LazyFeature( + name="prompts", + module_path="litellm.proxy.prompts.prompt_endpoints", + path_prefixes=("/prompts", "/utils/dotprompt_json_converter"), + ), + LazyFeature( + name="jwt_mappings", + module_path="litellm.proxy.management_endpoints.jwt_key_mapping_endpoints", + path_prefixes=("/jwt/key/mapping",), + ), + LazyFeature( + name="compliance", + module_path="litellm.proxy.management_endpoints.compliance_endpoints", + path_prefixes=("/compliance",), + ), + LazyFeature( + name="access_groups", + module_path="litellm.proxy.management_endpoints.access_group_endpoints", + path_prefixes=("/access_group", "/v1/access_group", "/v1/unified_access_group"), + ), +) + + +class LazyFeatureMiddleware: + """ASGI middleware that imports + registers a feature router on first + matching request. Idempotent; once loaded, subsequent requests skip.""" + + def __init__( + self, + app, + fastapi_app: "FastAPI", + features: Tuple[LazyFeature, ...] = LAZY_FEATURES, + ): + self.app = app + self._fastapi_app = fastapi_app + self._features = features + self._loaded: set = set() + # Per-feature locks so independent features can load in parallel. + self._locks: dict = {} + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + # Short-circuit once every feature has loaded. + if scope["type"] in ("http", "websocket") and len(self._loaded) < len( + self._features + ): + path = scope.get("path", "") + for feat in self._features: + if feat.module_path in self._loaded: + continue + if any(path.startswith(p) for p in feat.path_prefixes) or any( + path.endswith(s) for s in feat.path_suffixes + ): + await self._load(feat) + await self.app(scope, receive, send) + + async def _load(self, feat: LazyFeature) -> None: + lock = self._locks.setdefault(feat.module_path, asyncio.Lock()) + async with lock: + if feat.module_path in self._loaded: + return + try: + # Import on a thread (heavy modules take 1-3 s). register_fn + # mutates app.router.routes, so it stays on the loop thread. + loop = asyncio.get_running_loop() + module = await loop.run_in_executor( + None, importlib.import_module, feat.module_path + ) + feat.register_fn(self._fastapi_app, module) + self._loaded.add(feat.module_path) + self._fastapi_app.openapi_schema = None + verbose_proxy_logger.info( + "Lazy-loaded optional feature %r (module: %s)", + feat.name, + feat.module_path, + ) + except Exception as exc: + # Mark loaded anyway so we don't retry on every request. + self._loaded.add(feat.module_path) + verbose_proxy_logger.warning( + "Failed to lazy-load optional feature %r (module: %s): %s. " + "This feature's endpoints will return 404 until restart.", + feat.name, + feat.module_path, + exc, + ) + + +def attach_lazy_features(app: "FastAPI") -> None: + app.add_middleware(LazyFeatureMiddleware, fastapi_app=app) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 8f676df04cd..c03a63f211a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -235,37 +235,11 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.vertex_ai.vertex_llm_base import VertexBase -from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import ( - router as mcp_byok_oauth_router, -) -from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - router as mcp_discoverable_endpoints_router, -) -from litellm.proxy._experimental.mcp_server.rest_endpoints import ( - router as mcp_rest_endpoints_router, -) -from litellm.proxy._experimental.mcp_server.server import app as mcp_app -from litellm.proxy._experimental.mcp_server.tool_registry import ( - global_mcp_tool_registry, -) from litellm.proxy._types import * -from litellm.proxy.agent_endpoints.a2a_endpoints import router as a2a_router -from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry -from litellm.proxy.agent_endpoints.endpoints import router as agent_endpoints_router -from litellm.proxy.agent_endpoints.model_list_helpers import ( - append_agents_to_model_group, - append_agents_to_model_info, -) +from litellm.proxy._lazy_features import attach_lazy_features from litellm.proxy.analytics_endpoints.analytics_endpoints import ( router as analytics_router, ) -from litellm.proxy.anthropic_endpoints.claude_code_endpoints import ( - claude_code_marketplace_router, -) -from litellm.proxy.anthropic_endpoints.endpoints import router as anthropic_router -from litellm.proxy.anthropic_endpoints.skills_endpoints import ( - router as anthropic_skills_router, -) from litellm.proxy.auth.auth_checks import ( ExperimentalUIJWTToken, get_team_object, @@ -328,7 +302,6 @@ from litellm.proxy.discovery_endpoints import ui_discovery_endpoints_router from litellm.proxy.fine_tuning_endpoints.endpoints import router as fine_tuning_router from litellm.proxy.fine_tuning_endpoints.endpoints import set_fine_tuning_config from litellm.proxy.google_endpoints.endpoints import router as google_router -from litellm.proxy.guardrails.guardrail_endpoints import router as guardrails_router from litellm.proxy.guardrails.init_guardrails import ( init_guardrails_v2, initialize_guardrails, @@ -344,9 +317,6 @@ from litellm.proxy.hooks.prompt_injection_detection import ( from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger from litellm.proxy.image_endpoints.endpoints import router as image_router from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request -from litellm.proxy.management_endpoints.access_group_endpoints import ( - router as access_group_router, -) from litellm.proxy.management_endpoints.budget_management_endpoints import ( router as budget_management_router, ) @@ -360,12 +330,6 @@ from litellm.proxy.management_endpoints.common_utils import ( _user_has_admin_privileges, admin_can_invite_user, ) -from litellm.proxy.management_endpoints.compliance_endpoints import ( - router as compliance_router, -) -from litellm.proxy.management_endpoints.config_override_endpoints import ( - router as config_override_router, -) from litellm.proxy.management_endpoints.cost_tracking_settings import ( router as cost_tracking_settings_router, ) @@ -379,9 +343,6 @@ from litellm.proxy.management_endpoints.internal_user_endpoints import ( router as internal_user_router, ) from litellm.proxy.management_endpoints.internal_user_endpoints import user_update -from litellm.proxy.management_endpoints.jwt_key_mapping_endpoints import ( - router as jwt_key_mapping_router, -) from litellm.proxy.management_endpoints.key_management_endpoints import ( delete_verification_tokens, duration_in_seconds, @@ -390,9 +351,6 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( from litellm.proxy.management_endpoints.key_management_endpoints import ( router as key_management_router, ) -from litellm.proxy.management_endpoints.mcp_management_endpoints import ( - router as mcp_management_router, -) from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( router as model_access_group_management_router, ) @@ -407,11 +365,9 @@ from litellm.proxy.management_endpoints.model_management_endpoints import ( from litellm.proxy.management_endpoints.organization_endpoints import ( router as organization_router, ) -from litellm.proxy.management_endpoints.policy_endpoints import router as policy_router from litellm.proxy.management_endpoints.router_settings_endpoints import ( router as router_settings_router, ) -from litellm.proxy.management_endpoints.scim.scim_v2 import scim_router from litellm.proxy.management_endpoints.tag_management_endpoints import ( router as tag_management_router, ) @@ -423,15 +379,11 @@ from litellm.proxy.management_endpoints.team_endpoints import ( update_team, validate_membership, ) -from litellm.proxy.management_endpoints.tool_management_endpoints import ( - router as tool_management_router, -) from litellm.proxy.memory.memory_endpoints import router as memory_router from litellm.proxy.management_endpoints.ui_sso import ( get_disabled_non_admin_personal_key_creation, ) from litellm.proxy.management_endpoints.ui_sso import router as ui_sso_router -from litellm.proxy.management_endpoints.usage_endpoints import router as usage_ai_router from litellm.proxy.management_endpoints.user_agent_analytics_endpoints import ( router as user_agent_analytics_router, ) @@ -441,7 +393,6 @@ from litellm.proxy.middleware.in_flight_requests_middleware import ( ) from litellm.proxy.middleware.prometheus_auth_middleware import PrometheusAuthMiddleware from litellm.proxy.ocr_endpoints.endpoints import router as ocr_router -from litellm.proxy.openai_evals_endpoints.endpoints import router as evals_router from litellm.proxy.openai_files_endpoints.files_endpoints import ( router as openai_files_router, ) @@ -461,27 +412,16 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( router as pass_through_router, ) -from litellm.proxy.policy_engine.policy_endpoints import router as policy_crud_router -from litellm.proxy.policy_engine.policy_resolve_endpoints import ( - router as policy_resolve_router, -) -from litellm.proxy.prompts.prompt_endpoints import router as prompts_router from litellm.proxy.public_endpoints import router as public_endpoints_router from litellm.proxy.rag_endpoints.endpoints import router as rag_router -from litellm.proxy.realtime_endpoints.endpoints import router as webrtc_router from litellm.proxy.rerank_endpoints.endpoints import router as rerank_router from litellm.proxy.response_api_endpoints.endpoints import router as response_router from litellm.proxy.route_llm_request import route_request from litellm.proxy.search_endpoints.endpoints import router as search_router -from litellm.proxy.search_endpoints.search_tool_management import ( - router as search_tool_management_router, -) -from litellm.proxy.spend_tracking.cloudzero_endpoints import router as cloudzero_router from litellm.proxy.spend_tracking.spend_management_endpoints import ( router as spend_management_router, ) from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload -from litellm.proxy.spend_tracking.vantage_endpoints import router as vantage_router from litellm.proxy.types_utils.utils import get_instance_fn from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import ( router as ui_crud_endpoints_router, @@ -511,16 +451,6 @@ from litellm.proxy.utils import ( prefetch_config_params, update_spend, ) -from litellm.proxy.vector_store_endpoints.endpoints import router as vector_store_router -from litellm.proxy.vector_store_endpoints.management_endpoints import ( - router as vector_store_management_router, -) -from litellm.proxy.vector_store_files_endpoints.endpoints import ( - router as vector_store_files_router, -) -from litellm.proxy.vertex_ai_endpoints.langfuse_endpoints import ( - router as langfuse_router, -) from litellm.proxy.video_endpoints.endpoints import router as video_router from litellm.router import ( AssistantsTypedDict, @@ -3854,11 +3784,19 @@ class ProxyConfig: ## MCP TOOLS mcp_tools_config = config.get("mcp_tools", None) if mcp_tools_config: + from litellm.proxy._experimental.mcp_server.tool_registry import ( + global_mcp_tool_registry, + ) + global_mcp_tool_registry.load_tools_from_config(mcp_tools_config) ## AGENTS agent_config = config.get("agent_list", None) if agent_config: + from litellm.proxy.agent_endpoints.agent_registry import ( + global_agent_registry, + ) + global_agent_registry.load_agents_from_config(agent_config) # type: ignore mcp_servers_config = config.get("mcp_servers", None) @@ -10576,6 +10514,10 @@ async def model_info_v2( verbose_proxy_logger.debug("all_models: %s", all_models) # Append A2A agents to models list + from litellm.proxy.agent_endpoints.model_list_helpers import ( + append_agents_to_model_info, + ) + all_models = await append_agents_to_model_info( models=all_models, user_api_key_dict=user_api_key_dict, @@ -11425,6 +11367,10 @@ async def model_group_info( ) # Append A2A agents to model groups + from litellm.proxy.agent_endpoints.model_list_helpers import ( + append_agents_to_model_group, + ) + model_groups = await append_agents_to_model_group( model_groups=model_groups, user_api_key_dict=user_api_key_dict, @@ -14230,65 +14176,40 @@ app.include_router(container_router) app.include_router(search_router) app.include_router(image_router) app.include_router(fine_tuning_router) -app.include_router(vector_store_router) -app.include_router(vector_store_management_router) -app.include_router(vector_store_files_router) app.include_router(credential_router) app.include_router(llm_passthrough_router) -app.include_router(webrtc_router) -app.include_router(mcp_management_router) -app.include_router(mcp_byok_oauth_router) -app.include_router(anthropic_router) -app.include_router(anthropic_skills_router) -app.include_router(evals_router) -app.include_router(claude_code_marketplace_router) -app.include_router(google_router) -app.include_router(langfuse_router) app.include_router(pass_through_router) app.include_router(health_router) app.include_router(key_management_router) app.include_router(internal_user_router) app.include_router(team_router) app.include_router(ui_sso_router) -app.include_router(scim_router) app.include_router(organization_router) app.include_router(customer_router) app.include_router(spend_management_router) -app.include_router(cloudzero_router) -app.include_router(vantage_router) app.include_router(caching_router) app.include_router(analytics_router) -app.include_router(guardrails_router) -app.include_router(policy_router) -app.include_router(usage_ai_router) -app.include_router(policy_crud_router) -app.include_router(policy_resolve_router) -app.include_router(search_tool_management_router) -app.include_router(prompts_router) app.include_router(callback_management_endpoints_router) app.include_router(debugging_endpoints_router) app.include_router(ui_crud_endpoints_router) app.include_router(openai_files_router) app.include_router(team_callback_router) -app.include_router(jwt_key_mapping_router) app.include_router(budget_management_router) app.include_router(model_management_router) app.include_router(model_access_group_management_router) app.include_router(tag_management_router) -app.include_router(tool_management_router) app.include_router(memory_router) app.include_router(cost_tracking_settings_router) app.include_router(router_settings_router) app.include_router(fallback_management_router) app.include_router(cache_settings_router) -app.include_router(config_override_router) app.include_router(user_agent_analytics_router) app.include_router(enterprise_router) app.include_router(ui_discovery_endpoints_router) -app.include_router(agent_endpoints_router) -app.include_router(compliance_router) -app.include_router(a2a_router) -app.include_router(access_group_router) +# Eager: /models/{name}:method overlaps with the OpenAI /models endpoint. +app.include_router(google_router) + +attach_lazy_features(app) async def _stream_mcp_asgi_response( @@ -14521,8 +14442,3 @@ async def dynamic_mcp_route(mcp_server_name: str, request: Request): f"Error handling dynamic MCP route for {mcp_server_name}: {str(e)}" ) raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}") - - -app.mount(path=BASE_MCP_ROUTE, app=mcp_app) -app.include_router(mcp_rest_endpoints_router) -app.include_router(mcp_discoverable_endpoints_router) diff --git a/tests/proxy_unit_tests/test_proxy_routes.py b/tests/proxy_unit_tests/test_proxy_routes.py index 812e4e1ac41..67eca5206d4 100644 --- a/tests/proxy_unit_tests/test_proxy_routes.py +++ b/tests/proxy_unit_tests/test_proxy_routes.py @@ -39,6 +39,20 @@ def test_routes_on_litellm_proxy(): this prevents accidentelly deleting /threads, or /batches etc """ + # Force-load lazy features so the test sees the full route set. Continue + # on per-feature import failure — the assertion below still catches + # missing-route regressions. + import importlib + + from litellm.proxy._lazy_features import LAZY_FEATURES + + for feat in LAZY_FEATURES: + try: + module = importlib.import_module(feat.module_path) + feat.register_fn(app, module) + except Exception as exc: + print(f"warning: failed to force-load {feat.name}: {exc}") + _all_routes = [] for route in app.routes: diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 1f4f82a64ef..7a96f6cbd15 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -5471,3 +5471,255 @@ async def test_reseed_warms_cache_even_on_zero_db_spend(): finally: ps.spend_counter_cache = orig_counter ps.prisma_client = orig_prisma + + +# --------------------------------------------------------------------------- +# Lazy feature loading (LazyFeatureMiddleware) — verifies that optional +# routers are NOT imported at module load and ARE imported on first request +# to a matching path prefix. The same module isn't re-imported on subsequent +# requests. +# --------------------------------------------------------------------------- + + +import sys + + +class TestLazyFeatureRegistry: + """Sanity checks on the registry shape — guards against accidental edits.""" + + def test_registry_entries_have_required_fields(self): + from litellm.proxy._lazy_features import LAZY_FEATURES, LazyFeature + + assert len(LAZY_FEATURES) > 0 + for feat in LAZY_FEATURES: + assert isinstance(feat, LazyFeature) + assert feat.name + assert feat.module_path + assert feat.path_prefixes + assert all(p.startswith("/") for p in feat.path_prefixes) + assert callable(feat.register_fn) + + def test_registry_names_unique(self): + from litellm.proxy._lazy_features import LAZY_FEATURES + + names = [f.name for f in LAZY_FEATURES] + assert len(names) == len(set(names)), "duplicate feature names" + + +class TestLazyFeaturesNotImportedAtStartup: + """ + The whole point of the refactor: gated feature modules must NOT be + present in `sys.modules` immediately after `proxy_server` imports. + """ + + def test_heavy_modules_absent_at_startup(self): + # Force a fresh `proxy_server` import in a subprocess so other tests + # in this run (which may have triggered lazy loads via the TestClient) + # don't pollute the result. + import subprocess + + check = ( + "import sys; " + "from litellm.proxy.proxy_server import app; " # noqa: F401 + "heavy = [" + "'litellm.proxy._experimental.mcp_server.rest_endpoints'," + "'litellm.proxy._experimental.mcp_server.server'," + "'litellm.proxy.management_endpoints.config_override_endpoints'," + "'litellm.proxy.guardrails.guardrail_endpoints'," + "'litellm.proxy.openai_evals_endpoints.endpoints'," + "]; " + "still_present = [m for m in heavy if m in sys.modules]; " + "print('PRESENT_AT_STARTUP:', still_present)" + ) + result = subprocess.run( + [sys.executable, "-c", check], + capture_output=True, + text=True, + timeout=120, + ) + # Last non-empty line of stdout (skip warnings printed before) + out_lines = [ + line for line in result.stdout.strip().splitlines() if line.strip() + ] + report = next((line for line in out_lines if "PRESENT_AT_STARTUP" in line), "") + assert report, f"no report emitted (stderr: {result.stderr[-500:]})" + assert ( + "PRESENT_AT_STARTUP: []" in report + ), f"expected no heavy modules at startup, got: {report}" + + +class TestLazyFeatureMiddleware: + """Behavior of the middleware itself, exercised in isolation.""" + + @pytest.mark.asyncio + async def test_first_request_triggers_load_subsequent_does_not(self): + from fastapi import FastAPI + + from litellm.proxy._lazy_features import ( + LazyFeature, + LazyFeatureMiddleware, + ) + + loads = [] + + def fake_register(app, module): + loads.append(getattr(module, "__name__", "?")) + + feat = LazyFeature( + name="dummy", + module_path="json", # any always-importable stdlib module + path_prefixes=("/dummy",), + register_fn=fake_register, + ) + + # Build a minimal ASGI receiver to satisfy the middleware contract + async def downstream(scope, receive, send): + # echo back; no-op handler + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b""}) + + target_app = FastAPI() + mw = LazyFeatureMiddleware(downstream, fastapi_app=target_app, features=(feat,)) + + async def receive(): + return {"type": "http.request", "body": b"", "more_body": False} + + sent: list = [] + + async def send(message): + sent.append(message) + + # First request matching the prefix triggers register + await mw( + {"type": "http", "path": "/dummy/x", "method": "GET", "headers": []}, + receive, + send, + ) + assert loads == ["json"] + + # Second matching request must NOT re-register + sent.clear() + await mw( + {"type": "http", "path": "/dummy/y", "method": "GET", "headers": []}, + receive, + send, + ) + assert loads == ["json"], "register_fn called twice for the same feature" + + # Non-matching path must not trigger anything + await mw( + {"type": "http", "path": "/unrelated", "method": "GET", "headers": []}, + receive, + send, + ) + assert loads == ["json"] + + @pytest.mark.asyncio + async def test_concurrent_first_requests_only_register_once(self): + """ + Two requests to the same prefix arriving in parallel must result in + exactly one `register_fn` invocation — the lock prevents the import + + register from racing with itself. + """ + from fastapi import FastAPI + + from litellm.proxy._lazy_features import ( + LazyFeature, + LazyFeatureMiddleware, + ) + + loads = [] + + def slow_register(app, module): + loads.append(getattr(module, "__name__", "?")) + + feat = LazyFeature( + name="dummy_concurrent", + module_path="json", + path_prefixes=("/dummy_c",), + register_fn=slow_register, + ) + + async def downstream(scope, receive, send): + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b""}) + + target_app = FastAPI() + mw = LazyFeatureMiddleware(downstream, fastapi_app=target_app, features=(feat,)) + + async def receive(): + return {"type": "http.request", "body": b"", "more_body": False} + + sent: list = [] + + async def send(message): + sent.append(message) + + async def hit(): + await mw( + { + "type": "http", + "path": "/dummy_c/x", + "method": "GET", + "headers": [], + }, + receive, + send, + ) + + await asyncio.gather(hit(), hit(), hit(), hit(), hit()) + assert loads == [ + "json" + ], f"expected one registration despite concurrent first hits, got {loads}" + + @pytest.mark.asyncio + async def test_failing_import_does_not_loop(self): + """ + If a feature's module can't be imported, the middleware should mark it + loaded anyway so subsequent requests don't repeatedly retry the failing + import (which would amplify the cost on every request). + """ + from fastapi import FastAPI + + from litellm.proxy._lazy_features import ( + LazyFeature, + LazyFeatureMiddleware, + ) + + attempts = [] + + def fail_register(app, module): + attempts.append("called") + raise RuntimeError("boom") + + feat = LazyFeature( + name="failing", + module_path="json", + path_prefixes=("/fail",), + register_fn=fail_register, + ) + + async def downstream(scope, receive, send): + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b""}) + + target_app = FastAPI() + mw = LazyFeatureMiddleware(downstream, fastapi_app=target_app, features=(feat,)) + + async def receive(): + return {"type": "http.request", "body": b"", "more_body": False} + + sent: list = [] + + async def send(message): + sent.append(message) + + for _ in range(3): + await mw( + {"type": "http", "path": "/fail/x", "method": "GET", "headers": []}, + receive, + send, + ) + assert attempts == [ + "called" + ], f"failing register_fn should be invoked once, not on every request; got {attempts}" diff --git a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py index 44cc5cc4452..1e596aa5675 100644 --- a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py +++ b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py @@ -786,8 +786,23 @@ class TestVectorStoreManagementEndpointsExist: - POST /vector_store/info - POST /vector_store/update """ + import importlib + + from litellm.proxy._lazy_features import LAZY_FEATURES from litellm.proxy.proxy_server import app + # Force-register the lazy vector_store_management routes so the + # assertions can find them. + already_registered = any( + getattr(r, "path", None) == "/vector_store/new" for r in app.routes + ) + if not already_registered: + for feat in LAZY_FEATURES: + if feat.name == "vector_store_management": + module = importlib.import_module(feat.module_path) + feat.register_fn(app, module) + break + # Define expected endpoints expected_endpoints = [ ("POST", "/vector_store/new"),