From e73d053de3fca84402c40db89ee390f7f2a35f3a Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 1 Oct 2025 14:09:01 -0700 Subject: [PATCH] [Fix] Proxy Auth - Ensure LLM_API_KEYs can access pass through routes (#15115) * test_virtual_key_llm_api_routes_allows_registered_pass_through_endpoints * fix: is_registered_pass_through_route * docs fix --- docs/my-website/docs/providers/lemonade.md | 3 + litellm/proxy/auth/route_checks.py | 12 ++- .../pass_through_endpoints.py | 33 ++++++- .../proxy/auth/test_route_checks.py | 91 +++++++++++++++++++ 4 files changed, 137 insertions(+), 2 deletions(-) diff --git a/docs/my-website/docs/providers/lemonade.md b/docs/my-website/docs/providers/lemonade.md index 8ff7d48b70..fc77b78a76 100644 --- a/docs/my-website/docs/providers/lemonade.md +++ b/docs/my-website/docs/providers/lemonade.md @@ -186,3 +186,6 @@ print("Available models:", [model['id'] for model in models.get('data', [])]) ## Support For more information regarding Lemonade please go to to the [Lemonade website](https://lemonade-server.ai/) or [Lemonade repository](https://github.com/lemonade-sdk/lemonade). + + + diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index 56218b3345..39f11e64bb 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -62,12 +62,22 @@ class RouteChecks: for allowed_route in valid_token.allowed_routes ): for allowed_route in valid_token.allowed_routes: - if allowed_route in LiteLLMRoutes._member_names_: + if allowed_route in LiteLLMRoutes._member_names_: if RouteChecks.check_route_access( route=route, allowed_routes=LiteLLMRoutes._member_map_[allowed_route].value, ): return True + + ################################################ + # For llm_api_routes, also check registered pass-through endpoints + ################################################ + if allowed_route == "llm_api_routes": + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + InitPassThroughEndpointHelpers, + ) + if InitPassThroughEndpointHelpers.is_registered_pass_through_route(route=route): + return True # check if wildcard pattern is allowed for allowed_route in valid_token.allowed_routes: diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index e001b27ad3..53cc3d0ee1 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -1418,7 +1418,7 @@ async def websocket_passthrough_request( # noqa: PLR0915 if websocket.client_state != WebSocketState.DISCONNECTED: await websocket.close( - code=exc.status_code if hasattr(exc, "status_code") else 1011, + code=getattr(exc, "status_code", 1011), reason="Upstream connection rejected", ) except Exception as e: @@ -1603,6 +1603,37 @@ class InitPassThroughEndpointHelpers: del _registered_pass_through_routes[key] verbose_proxy_logger.debug("Removed pass-through route from registry: %s", key) + @staticmethod + def is_registered_pass_through_route(route: str) -> bool: + """ + Check if route is a registered pass-through endpoint from DB + + Uses the in-memory registry to avoid additional DB queries + Optimized for minimal latency + + Args: + route: The route to check + + Returns: + bool: True if route is a registered pass-through endpoint, False otherwise + """ + # Fast path: check if any registered route key contains this path + # Keys are in format: "{endpoint_id}:exact:{path}" or "{endpoint_id}:subpath:{path}" + # Extract unique paths from keys for quick checking + for key in _registered_pass_through_routes.keys(): + parts = key.split(":", 2) # Split into [endpoint_id, type, path] + if len(parts) == 3: + route_type = parts[1] + registered_path = parts[2] + + if route_type == "exact" and route == registered_path: + return True + elif route_type == "subpath": + if route == registered_path or route.startswith(registered_path + "/"): + return True + + return False + async def initialize_pass_through_endpoints( pass_through_endpoints: Union[List[Dict], List[PassThroughGenericEndpoint]], diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 539ee4a9ba..37cff2bd94 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -247,3 +247,94 @@ def test_anthropic_count_tokens_route_accessible_to_internal_users(): # Also test that the regular messages route still works assert RouteChecks.is_llm_api_route("/v1/messages") is True + + +def test_virtual_key_llm_api_routes_allows_registered_pass_through_endpoints(): + """ + Test that virtual keys with llm_api_routes permission can access registered pass-through endpoints. + + This tests the scenario where a pass-through endpoint is registered from the DB + (e.g., /azure-assistant) and a virtual key with llm_api_routes permission should be able to access + both the exact path and subpaths (e.g., /azure-assistant/openai/assistants). + """ + from unittest.mock import patch + + # Mock the registered pass-through routes + mock_registered_routes = { + "test-uuid-1:exact:/azure-assistant": { + "endpoint_id": "test-uuid-1", + "path": "/azure-assistant", + "type": "exact", + }, + "test-uuid-2:subpath:/custom-endpoint": { + "endpoint_id": "test-uuid-2", + "path": "/custom-endpoint", + "type": "subpath", + }, + } + + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes", + mock_registered_routes, + ): + # Create a virtual key with llm_api_routes permission + valid_token = UserAPIKeyAuth( + user_id="test_user", + allowed_routes=["llm_api_routes"], + ) + + # Test exact match for registered pass-through endpoint + result1 = RouteChecks.is_virtual_key_allowed_to_call_route( + route="/azure-assistant", + valid_token=valid_token, + ) + assert result1 is True + + # Test subpath for registered pass-through endpoint with subpath type + result2 = RouteChecks.is_virtual_key_allowed_to_call_route( + route="/custom-endpoint/openai/assistants", + valid_token=valid_token, + ) + assert result2 is True + + # Test exact match for subpath type + result3 = RouteChecks.is_virtual_key_allowed_to_call_route( + route="/custom-endpoint", + valid_token=valid_token, + ) + assert result3 is True + + +def test_virtual_key_without_llm_api_routes_cannot_access_pass_through(): + """ + Test that virtual keys without llm_api_routes permission cannot access registered pass-through endpoints. + """ + from unittest.mock import patch + + # Mock the registered pass-through routes + mock_registered_routes = { + "test-uuid-1:exact:/azure-assistant": { + "endpoint_id": "test-uuid-1", + "path": "/azure-assistant", + "type": "exact", + }, + } + + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes", + mock_registered_routes, + ): + # Create a virtual key without llm_api_routes permission + valid_token = UserAPIKeyAuth( + user_id="test_user", + allowed_routes=["info_routes"], + ) + + # Test that access is denied + with pytest.raises(Exception) as exc_info: + RouteChecks.is_virtual_key_allowed_to_call_route( + route="/azure-assistant", + valid_token=valid_token, + ) + + assert "Virtual key is not allowed to call this route" in str(exc_info.value)