From 0ed960fe7d6f2d13abcaaaa3f6f1c87f0da17ceb Mon Sep 17 00:00:00 2001 From: Michael Riad Zaky Date: Thu, 23 Apr 2026 15:05:35 -0700 Subject: [PATCH] add flags that remove unneeded imports --- litellm/proxy/proxy_server.py | 422 +++++++++++++----- tests/test_litellm/proxy/test_proxy_server.py | 172 +++++++ 2 files changed, 478 insertions(+), 116 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4ca6895c2ce..1fbee780e8c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -235,37 +235,10 @@ 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.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, @@ -326,8 +299,7 @@ from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler 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, @@ -343,9 +315,7 @@ 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, ) @@ -359,12 +329,7 @@ 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, ) @@ -378,9 +343,7 @@ 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, @@ -389,9 +352,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, ) @@ -406,11 +366,10 @@ 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, ) @@ -422,15 +381,12 @@ 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, ) @@ -440,7 +396,7 @@ 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, ) @@ -460,27 +416,20 @@ 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, @@ -506,16 +455,7 @@ from litellm.proxy.utils import ( model_dump_with_preserved_fields, 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, @@ -3892,11 +3832,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) @@ -10525,11 +10473,15 @@ async def model_info_v2( verbose_proxy_logger.debug("all_models: %s", all_models) - # Append A2A agents to models list - all_models = await append_agents_to_model_info( - models=all_models, - user_api_key_dict=user_api_key_dict, - ) + if _should_import_feature("agents"): + 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, + ) # Update total count to include agents search_total_count = len(all_models) @@ -11374,11 +11326,15 @@ async def model_group_info( llm_router=llm_router, all_models_str=all_models_str, model_group=model_group ) - # Append A2A agents to model groups - model_groups = await append_agents_to_model_group( - model_groups=model_groups, - user_api_key_dict=user_api_key_dict, - ) + if _should_import_feature("agents"): + 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, + ) return {"data": model_groups} @@ -14159,6 +14115,271 @@ async def get_routes(): # return {"token": token} +_OPTIONAL_FEATURE_GATES = ( + "guardrails", + "policies", + "agents", + "vector_stores", + "tools", + "search_tools", + "mcp", + "config_overrides", + "realtime", + "anthropic_passthrough", + "google_passthrough", + "langfuse_passthrough", + "evals", + "claude_code_marketplace", + "scim", + "cloudzero", + "vantage", + "usage_ai", + "prompts", + "jwt_mappings", + "compliance", + "access_groups", +) + + +def _parse_enabled_features_env() -> Optional[Set[str]]: + """ + Parse `LITELLM_ENABLED_FEATURES` once at module load. Returns None when + unset or empty (load everything), otherwise the allowlist as a set. + """ + raw = os.environ.get("LITELLM_ENABLED_FEATURES") + if raw is None or not raw.strip(): + return None + allowed = {s.strip() for s in raw.split(",") if s.strip()} + unknown = sorted(allowed - set(_OPTIONAL_FEATURE_GATES)) + if unknown: + verbose_proxy_logger.warning( + "LITELLM_ENABLED_FEATURES contains unknown keys %s — treated as " + "disabled. Valid keys: %s", + unknown, + list(_OPTIONAL_FEATURE_GATES), + ) + return allowed + + +_ENABLED_FEATURES: Optional[Set[str]] = _parse_enabled_features_env() + +_disabled_optional_features: Set[str] = set() + + +def _should_import_feature(name: str) -> bool: + """ + True if the optional feature `name` should be imported at startup. + Controlled by `LITELLM_ENABLED_FEATURES` (CSV of keys from + `_OPTIONAL_FEATURE_GATES`); unset means load everything. + """ + if _ENABLED_FEATURES is None: + return True + if name in _ENABLED_FEATURES: + return True + _disabled_optional_features.add(name) + return False + + +def _register_optional_feature_routers(app: FastAPI) -> None: # noqa: PLR0915 + """ + Mount optional feature routers, gated on `LITELLM_ENABLED_FEATURES`. + Imports inside each branch (not at module top) so skipped features + never pay their import-time memory cost. + """ + _disabled_optional_features.clear() + + if _should_import_feature("guardrails"): + from litellm.proxy.guardrails.guardrail_endpoints import ( + router as guardrails_router, + ) + + app.include_router(guardrails_router) + + if _should_import_feature("policies"): + from litellm.proxy.management_endpoints.policy_endpoints import ( + router as policy_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, + ) + + app.include_router(policy_router) + app.include_router(policy_crud_router) + app.include_router(policy_resolve_router) + + if _should_import_feature("agents"): + from litellm.proxy.agent_endpoints.a2a_endpoints import router as a2a_router + from litellm.proxy.agent_endpoints.endpoints import ( + router as agent_endpoints_router, + ) + + app.include_router(agent_endpoints_router) + app.include_router(a2a_router) + + if _should_import_feature("vector_stores"): + 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, + ) + + app.include_router(vector_store_router) + app.include_router(vector_store_management_router) + app.include_router(vector_store_files_router) + + if _should_import_feature("tools"): + from litellm.proxy.management_endpoints.tool_management_endpoints import ( + router as tool_management_router, + ) + + app.include_router(tool_management_router) + + if _should_import_feature("search_tools"): + from litellm.proxy.search_endpoints.search_tool_management import ( + router as search_tool_management_router, + ) + + app.include_router(search_tool_management_router) + + if _should_import_feature("mcp"): + 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.management_endpoints.mcp_management_endpoints import ( + router as mcp_management_router, + ) + + app.include_router(mcp_management_router) + app.include_router(mcp_byok_oauth_router) + app.include_router(mcp_rest_endpoints_router) + app.include_router(mcp_discoverable_endpoints_router) + app.mount(path=BASE_MCP_ROUTE, app=mcp_app) + + if _should_import_feature("config_overrides"): + from litellm.proxy.management_endpoints.config_override_endpoints import ( + router as config_override_router, + ) + + app.include_router(config_override_router) + + if _should_import_feature("realtime"): + from litellm.proxy.realtime_endpoints.endpoints import router as webrtc_router + + app.include_router(webrtc_router) + + if _should_import_feature("anthropic_passthrough"): + from litellm.proxy.anthropic_endpoints.endpoints import ( + router as anthropic_router, + ) + from litellm.proxy.anthropic_endpoints.skills_endpoints import ( + router as anthropic_skills_router, + ) + + app.include_router(anthropic_router) + app.include_router(anthropic_skills_router) + + if _should_import_feature("google_passthrough"): + from litellm.proxy.google_endpoints.endpoints import router as google_router + + app.include_router(google_router) + + if _should_import_feature("langfuse_passthrough"): + from litellm.proxy.vertex_ai_endpoints.langfuse_endpoints import ( + router as langfuse_router, + ) + + app.include_router(langfuse_router) + + if _should_import_feature("evals"): + from litellm.proxy.openai_evals_endpoints.endpoints import ( + router as evals_router, + ) + + app.include_router(evals_router) + + if _should_import_feature("claude_code_marketplace"): + from litellm.proxy.anthropic_endpoints.claude_code_endpoints import ( + claude_code_marketplace_router, + ) + + app.include_router(claude_code_marketplace_router) + + if _should_import_feature("scim"): + from litellm.proxy.management_endpoints.scim.scim_v2 import ( + scim_router, + ) + + app.include_router(scim_router) + + if _should_import_feature("cloudzero"): + from litellm.proxy.spend_tracking.cloudzero_endpoints import ( + router as cloudzero_router, + ) + + app.include_router(cloudzero_router) + + if _should_import_feature("vantage"): + from litellm.proxy.spend_tracking.vantage_endpoints import ( + router as vantage_router, + ) + + app.include_router(vantage_router) + + if _should_import_feature("usage_ai"): + from litellm.proxy.management_endpoints.usage_endpoints import ( + router as usage_ai_router, + ) + + app.include_router(usage_ai_router) + + if _should_import_feature("prompts"): + from litellm.proxy.prompts.prompt_endpoints import router as prompts_router + + app.include_router(prompts_router) + + if _should_import_feature("jwt_mappings"): + from litellm.proxy.management_endpoints.jwt_key_mapping_endpoints import ( + router as jwt_key_mapping_router, + ) + + app.include_router(jwt_key_mapping_router) + + if _should_import_feature("compliance"): + from litellm.proxy.management_endpoints.compliance_endpoints import ( + router as compliance_router, + ) + + app.include_router(compliance_router) + + if _should_import_feature("access_groups"): + from litellm.proxy.management_endpoints.access_group_endpoints import ( + router as access_group_router, + ) + + app.include_router(access_group_router) + + if _disabled_optional_features: + verbose_proxy_logger.info( + "Optional feature routers disabled via LITELLM_ENABLED_FEATURES " + "(requests to their endpoints will 404): %s", + ", ".join(sorted(_disabled_optional_features)), + ) + + app.include_router(router) app.include_router(response_router) app.include_router(batches_router) @@ -14171,65 +14392,39 @@ 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) + +# Mount optional feature routers (gated on LITELLM_ENABLED_FEATURES). +_register_optional_feature_routers(app) async def _stream_mcp_asgi_response( @@ -14462,8 +14657,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/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index efd1abbb383..250ac343bea 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -5085,3 +5085,175 @@ async def test_reseed_spend_from_db_skips_window_variant_keys(): fake_prisma.db.litellm_teamtable.find_unique.assert_not_awaited() finally: ps.prisma_client = orig_prisma + + +# --------------------------------------------------------------------------- +# Optional-feature gate (LITELLM_ENABLED_FEATURES) — import-time memory gate. +# +# Three testable units in proxy_server.py: +# * _parse_enabled_features_env — reads the env var, returns None (unset / +# empty / whitespace) or a set of allowed feature keys; warns on unknowns. +# * _should_import_feature — checks a name against the parsed allowlist and +# records skips to _disabled_optional_features. +# * _register_optional_feature_routers — uses the above to mount routers; +# the module-load-time call already ran against `app`, so we sample +# `app.routes` to verify the default-unset path registers everything. +# --------------------------------------------------------------------------- + + +import logging + + +class TestParseEnabledFeaturesEnv: + """Unit tests for _parse_enabled_features_env.""" + + def test_unset_returns_none(self, monkeypatch): + from litellm.proxy.proxy_server import _parse_enabled_features_env + + monkeypatch.delenv("LITELLM_ENABLED_FEATURES", raising=False) + assert _parse_enabled_features_env() is None + + def test_empty_string_returns_none(self, monkeypatch): + from litellm.proxy.proxy_server import _parse_enabled_features_env + + monkeypatch.setenv("LITELLM_ENABLED_FEATURES", "") + assert _parse_enabled_features_env() is None + + def test_whitespace_only_returns_none(self, monkeypatch): + from litellm.proxy.proxy_server import _parse_enabled_features_env + + monkeypatch.setenv("LITELLM_ENABLED_FEATURES", " \t ") + assert _parse_enabled_features_env() is None + + def test_single_valid_key(self, monkeypatch): + from litellm.proxy.proxy_server import _parse_enabled_features_env + + monkeypatch.setenv("LITELLM_ENABLED_FEATURES", "mcp") + assert _parse_enabled_features_env() == {"mcp"} + + def test_multiple_valid_keys(self, monkeypatch): + from litellm.proxy.proxy_server import _parse_enabled_features_env + + monkeypatch.setenv("LITELLM_ENABLED_FEATURES", "mcp,guardrails,prompts") + assert _parse_enabled_features_env() == {"mcp", "guardrails", "prompts"} + + def test_whitespace_around_entries_trimmed(self, monkeypatch): + from litellm.proxy.proxy_server import _parse_enabled_features_env + + monkeypatch.setenv("LITELLM_ENABLED_FEATURES", " mcp , guardrails ") + assert _parse_enabled_features_env() == {"mcp", "guardrails"} + + def test_empty_entries_dropped(self, monkeypatch): + from litellm.proxy.proxy_server import _parse_enabled_features_env + + monkeypatch.setenv("LITELLM_ENABLED_FEATURES", "mcp,,guardrails,") + assert _parse_enabled_features_env() == {"mcp", "guardrails"} + + def test_unknown_keys_emit_warning(self, monkeypatch, caplog): + from litellm.proxy.proxy_server import ( + _parse_enabled_features_env, + verbose_proxy_logger, + ) + + monkeypatch.setenv("LITELLM_ENABLED_FEATURES", "mcp,guadrails,prmopts") + with caplog.at_level(logging.WARNING, logger=verbose_proxy_logger.name): + result = _parse_enabled_features_env() + # Unknown keys are still returned (their gate check will simply never + # match a real feature name), but a warning must surface them. + assert result == {"mcp", "guadrails", "prmopts"} + joined = " ".join(r.getMessage() for r in caplog.records) + assert "unknown" in joined.lower() + assert "guadrails" in joined + assert "prmopts" in joined + + +class TestShouldImportFeature: + """Unit tests for _should_import_feature.""" + + def test_none_allowlist_loads_everything(self): + import litellm.proxy.proxy_server as ps + + with patch.object(ps, "_ENABLED_FEATURES", None): + assert ps._should_import_feature("mcp") is True + assert ps._should_import_feature("guardrails") is True + # Even unknown names return True when allowlist is None — that's + # the backward-compat default. + assert ps._should_import_feature("not_a_real_feature") is True + + def test_set_allowlist_gates_unlisted_features(self): + import litellm.proxy.proxy_server as ps + + with patch.object(ps, "_ENABLED_FEATURES", {"mcp"}): + assert ps._should_import_feature("mcp") is True + assert ps._should_import_feature("guardrails") is False + assert ps._should_import_feature("policies") is False + + def test_disabled_features_recorded_once(self): + """Skipped feature names accumulate in the set for the summary log.""" + import litellm.proxy.proxy_server as ps + + with patch.object(ps, "_ENABLED_FEATURES", {"mcp"}): + ps._disabled_optional_features.clear() + ps._should_import_feature("mcp") # enabled → not tracked + ps._should_import_feature("guardrails") + ps._should_import_feature("agents") + ps._should_import_feature("guardrails") # repeated → still one entry + + tracked = ps._disabled_optional_features + assert "mcp" not in tracked + assert tracked >= {"guardrails", "agents"} + assert isinstance(tracked, set) + + +@pytest.mark.skipif( + os.environ.get("LITELLM_ENABLED_FEATURES") is not None, + reason=( + "LITELLM_ENABLED_FEATURES is set in the test runner — " + "_register_optional_feature_routers mounted only the listed " + "features at import time, so assertions about the backward-compat " + "default path cannot be verified. Unset the env var to run these." + ), +) +class TestDefaultUnsetRegistersAllGatedRouters: + """ + Backward-compat contract: when LITELLM_ENABLED_FEATURES is unset at + module load (the default for existing deployments), every optional + feature router must register so no endpoint silently disappears. + + `_register_optional_feature_routers(app)` ran at import time — we + inspect `app.routes` rather than re-invoke the helper (re-invoking + would double-register routes). + """ + + @pytest.fixture(scope="class") + def app_paths(self): + from litellm.proxy.proxy_server import app + + return { + (method, getattr(route, "path", None)) + for route in app.routes + for method in (getattr(route, "methods", None) or set()) + } + + @pytest.mark.parametrize( + "method,path,gate_name", + [ + # One representative endpoint per gate family. If any of these + # is missing in the default-unset case, the gate accidentally + # swallowed the feature on the backward-compat path. + ("GET", "/guardrails/list", "guardrails"), + ("GET", "/v1/agents", "agents"), + ("POST", "/vector_store/new", "vector_stores"), + ("GET", "/vector_store/list", "vector_stores"), + ("POST", "/scim/v2/Users", "scim"), + ("GET", "/access_group/list", "access_groups"), + ], + ) + def test_gated_endpoint_registered_by_default( + self, app_paths, method, path, gate_name + ): + assert (method, path) in app_paths, ( + f"{method} {path} (gate={gate_name!r}) missing from app.routes. " + "The default env-unset path must register every gated router " + "to stay backward compatible." + )