diff --git a/litellm/constants.py b/litellm/constants.py index 1dbb8a842fb..745a4d9294e 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -317,6 +317,8 @@ WEBSOCKET_CLOSE_REASON_MAX_BYTES: Final = 123 BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY: Final = "litellm.bedrock_realtime.pending_session_update" BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY: Final = "litellm.bedrock_realtime.session_committed" BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY: Final = "litellm.bedrock_realtime.committed_failure" +CLIENT_REQUESTED_MODEL_SCOPE_KEY: Final = "litellm.client_requested_model" +MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY: Final = "litellm.model_group_alias_resolved" REALTIME_SESSION_SUCCESS_LOGGED_KEY: Final = "realtime_session_success_logged" REALTIME_SESSION_FAILURE_LOGGED_KEY: Final = "realtime_session_failure_logged" diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index dc304a156cf..b4c123af762 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -1967,6 +1967,11 @@ def request_dispatched_to_pass_through_endpoint(request: Request | None) -> bool return getattr(endpoint, LITELLM_PASS_THROUGH_ENDPOINT_MARKER, False) is True +def request_dispatched_to_provider_pass_through(request: Request) -> bool: + """Built-in provider pass-through handlers (``/anthropic/{endpoint:path}``, ...) bind ``endpoint``.""" + return "endpoint" in request.path_params + + def get_model_from_request( request_data: dict, route: str, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 5958a68f975..c5297ac83dc 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -26,11 +26,13 @@ from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._service_logger import ServiceLogging from litellm.caching.redis_cache import RedisCache from litellm.constants import ( + CLIENT_REQUESTED_MODEL_SCOPE_KEY, GLOBAL_PROXY_SPEND_CACHE_KEY, INVALID_VIRTUAL_KEY_ERROR_MARKER, INVALID_VIRTUAL_KEY_ERROR_MESSAGE, LITELLM_PROXY_BUDGET_NAME, LITELLM_PROXY_MASTER_KEY_ALIAS, + MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY, ) from litellm.integrations.otel.model.config import is_otel_v2_enabled from litellm.integrations.otel.runtime import phase_span, seed_request_identity @@ -76,6 +78,8 @@ from litellm.proxy.auth.auth_utils import ( iter_request_fallback_targets, normalize_request_route, pre_db_read_auth_checks, + request_dispatched_to_pass_through_endpoint, + request_dispatched_to_provider_pass_through, route_in_additonal_public_routes, ) from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler @@ -103,6 +107,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( _safe_set_request_parsed_body, populate_request_with_path_params, read_raw_json_body, + rewrite_request_model, ) from litellm.proxy.common_utils.model_listing_utils import claude_code_requested_group from litellm.proxy.common_utils.realtime_utils import _realtime_request_body @@ -124,6 +129,7 @@ from litellm.proxy.utils import ( normalize_route_for_root_path, ) from litellm.repositories.table_repositories import TeamMembershipRepository +from litellm.router_utils.common_utils import resolve_model_group_alias from litellm.secret_managers.main import get_secret_bool from litellm.types.services import ServiceTypes @@ -235,11 +241,45 @@ async def _normalize_claude_model( request.scope[_CLAUDE_MODEL_NORMALIZED] = True if source is None: return - request_data["model"] = source - _safe_set_request_parsed_body(request=request, parsed_body=request_data) - if request is not None: - request._json = request_data - request._body = orjson.dumps(request_data) + rewrite_request_model(request_data, request, source) + + +async def _resolve_router_settings_model_group_alias( + request_data: dict[str, object], # mutable-ok: the request body is rewritten in place for every downstream reader + valid_token: UserAPIKeyAuth, + request: Request | None, + route: str, +) -> None: + """Rewrite the requested model through the key's or team's ``router_settings.model_group_alias`` + before the allowlist checks, so they authorize the model group the request is routed to. + """ + from litellm.proxy.proxy_server import llm_router, prisma_client, proxy_config, proxy_logging_obj + + if request is None or llm_router is None or not RouteChecks.is_llm_api_route(route=route): + return + if request.scope.get(MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY) is True: + return + request.scope[MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY] = True + if request_dispatched_to_pass_through_endpoint(request) or request_dispatched_to_provider_pass_through(request): + return + requested: Final = request_data.get("model") + if not isinstance(requested, str) or await read_raw_json_body(request=request) is None: + return + settings: Final = await proxy_config.get_hierarchical_router_settings( + user_api_key_dict=valid_token, prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj + ) + if not isinstance(settings, Mapping): + return + target: Final = resolve_model_group_alias(settings.get("model_group_alias"), requested) + if target is None or target == requested: + return + verbose_proxy_logger.debug( + "router_settings.model_group_alias resolved %s -> %s before auth", + requested.replace("\r", "").replace("\n", ""), + target.replace("\r", "").replace("\n", ""), + ) + request.scope.setdefault(CLIENT_REQUESTED_MODEL_SCOPE_KEY, requested) + rewrite_request_model(request_data, request, target) def _get_model_names_for_budget_checks( @@ -2990,6 +3030,7 @@ async def _authorize_authenticated_request( ## ENSURE DISABLE ROUTE WORKS ACROSS ALL USER AUTH FLOWS ## RouteChecks.should_call_route(route=route, valid_token=user_api_key_auth_obj, request=request) await _normalize_claude_model(request_data, user_api_key_auth_obj, request, route) + await _resolve_router_settings_model_group_alias(request_data, user_api_key_auth_obj, request, route) # Single authorization point. Builder paths MUST NOT call common_checks. # Route through the same exception handler the builder uses so @@ -3376,6 +3417,7 @@ async def _enforce_key_and_fallback_model_access( Not included in common_checks — common_checks enforces team/user/project model access only. """ await _normalize_claude_model(request_data, valid_token, request, route) + await _resolve_router_settings_model_group_alias(request_data, valid_token, request, route) config: Final = valid_token.config if config != {}: diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 37c1ab39632..46b222a4fc9 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -56,6 +56,7 @@ from litellm.proxy.common_utils.callback_utils import ( get_logging_caching_headers, get_remaining_tokens_and_requests_from_request_data, ) +from litellm.proxy.common_utils.http_parsing_utils import get_client_requested_model from litellm.proxy.common_utils.openai_error_payload import ( attribute_of, error_status_code, @@ -622,9 +623,9 @@ async def _resolve_per_request_model_group_alias( holds the global config map and is shared across requests, so a per-request map has to be applied here instead of being forwarded to the Router. - Model access was authorized against the requested group, so the target is - authorized in its own right before the rewrite; a key that may not call the - target gets the usual 403 rather than being quietly served it. + Auth already rewrote the body through this map for LLM API routes, so this is + a fallback for callers that skipped it; the target is authorized in its own + right before the rewrite, so a key that may not call it gets the usual 403. Returns the target model group, or None when no alias applies. """ @@ -2346,9 +2347,8 @@ class ProxyBaseLLMRequestProcessing: """ Common request processing logic for both chat completions and responses API endpoints """ - requested_model_from_client: Final[str | None] = ( - self.data.get("model") if isinstance(self.data.get("model"), str) else None - ) + client_model: Final = get_client_requested_model(request) or self.data.get("model") + requested_model_from_client: Final[str | None] = client_model if isinstance(client_model, str) else None self._debug_log_request_payload() if skip_pre_call_logic: diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index 9c2767c7771..f5b6a0a766d 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -9,7 +9,7 @@ from fastapi import Request, UploadFile, status from typing_extensions import NotRequired, ReadOnly, Required from litellm._logging import verbose_proxy_logger -from litellm.constants import MAX_REQUEST_BODY_SIZE_TO_REPAIR_MB +from litellm.constants import CLIENT_REQUESTED_MODEL_SCOPE_KEY, MAX_REQUEST_BODY_SIZE_TO_REPAIR_MB from litellm.proxy._types import ProxyException from litellm.proxy.common_utils.callback_utils import ( get_metadata_variable_name_from_kwargs, @@ -235,6 +235,13 @@ def _safe_get_request_parsed_body(request: Request | None) -> dict | None: return None +def get_client_requested_model(request: Request | None) -> str | None: + if request is None or not hasattr(request, "scope"): + return None + model: Final = request.scope.get(CLIENT_REQUESTED_MODEL_SCOPE_KEY) + return model if isinstance(model, str) else None + + def _safe_get_request_query_params(request: Request | None) -> dict: if request is None: return {} @@ -259,6 +266,24 @@ def _safe_set_request_parsed_body( verbose_proxy_logger.debug("Unexpected error setting request parsed body - %s", e) +def rewrite_request_model( + request_data: dict[str, object], # mutable-ok: the request body is rewritten in place for every downstream reader + request: Request | None, + model: str, +) -> None: + """Point the auth-time payload, the parsed-body cache, ``request.json()`` and ``request.body()`` at ``model``. + The cache and raw body keep only the keys the client sent, not params auth merged into ``request_data``. + """ + request_data["model"] = model + if request is None: + return + cached_body: Final = _safe_get_request_parsed_body(request=request) + body: Final = {**cached_body, "model": model} if cached_body is not None else request_data + _safe_set_request_parsed_body(request=request, parsed_body=body) + request._json = body + request._body = orjson.dumps(body) + + def _safe_get_request_headers(request: Request | None) -> dict: """ [Non-Blocking] Safely get the request headers. diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index bd7ff62ac8b..896acc5fcef 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -6,6 +6,7 @@ import subprocess import sys from contextlib import contextmanager from datetime import datetime, timedelta, timezone +from functools import partial from pathlib import Path from textwrap import dedent from types import SimpleNamespace @@ -40,6 +41,7 @@ from litellm.proxy.auth.auth_checks import ( jwt_key_mapping_cache_key, ) from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.common_utils.http_parsing_utils import get_client_requested_model from litellm.proxy.auth.user_api_key_auth import ( _check_key_model_budget_with_fallback, _ensure_litellm_received_at_on_request_state, @@ -8298,3 +8300,189 @@ async def test_auth_flow_enters_virtual_key_mapping_when_only_an_issuer_configur assert resolve_mock.await_args.kwargs["jwt_claims"][JWTHandler.LITELLM_JWT_ISSUER_CLAIM] == ISSUER_TWO assert result.api_key == "hashed-mapped-key" assert result.team_id == "svc-team" + + +def _alias_router() -> litellm.Router: + return litellm.Router( + model_list=[ + {"model_name": name, "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}} + for name in ("claude-haiku", "claude-sonnet") + ] + ) + + +def _alias_request(route: str, data: dict, content_type: str = "application/json", path_params: dict | None = None): + """A request as auth sees it: the body already read once and cached alongside its parsed form.""" + from starlette.requests import Request + + scope = { + "type": "http", + "method": "POST", + "path": route, + "headers": [(b"content-type", content_type.encode())], + "query_string": b"", + "path_params": path_params or {}, + "parsed_body": (tuple(data), data), + } + request = Request(scope) + request._body = json.dumps(data).encode() + return request + + +async def _enforce_alias_access(token: UserAPIKeyAuth, data: dict, route: str, request, router: litellm.Router): + from litellm.proxy.auth.user_api_key_auth import _enforce_key_and_fallback_model_access + + await _enforce_key_and_fallback_model_access( + valid_token=token, + request_data=data, + route=route, + request=request, + llm_model_list=router.model_list, + llm_router=router, + ) + + +def _alias_token(monkeypatch, level: str, alias: dict, models: list) -> UserAPIKeyAuth: + """A key whose ``router_settings.model_group_alias`` lives on the key itself or on its cached team row.""" + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + if level == "key": + return UserAPIKeyAuth(models=models, router_settings={"model_group_alias": alias}) + cache = UserApiKeyCache() + team = LiteLLM_TeamTableCachedObj(team_id="team-alias", models=models, router_settings={"model_group_alias": alias}) + cache.set_cache(key="team_id:team-alias", value=team) + monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", cache) + monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", MagicMock()) + return UserAPIKeyAuth(team_id="team-alias", models=models) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("level", ["key", "team"]) +@pytest.mark.parametrize( + "route", + ["/v1/chat/completions", "/v1/messages", "/v1/embeddings", "/openai/v1/responses", "/cursor/chat/completions"], +) +async def test_router_settings_model_group_alias_authorizes_target_for_key(monkeypatch, level, route): + """LIT-3054: a key allowed only the alias target must be able to call the alias, and a key not + allowed the target must still be denied even when the alias itself is what it requested.""" + router = _alias_router() + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + data = {"model": "AgentX-LLM", "messages": [{"role": "user", "content": "hi"}]} + request = _alias_request(route, data) + token = _alias_token(monkeypatch, level, {"AgentX-LLM": "claude-haiku"}, ["claude-haiku"]) + await _enforce_alias_access(token, data, route, request, router) + assert data["model"] == "claude-haiku" + assert (await request.json())["model"] == "claude-haiku" + assert json.loads(await request.body())["model"] == "claude-haiku" + assert request.scope["parsed_body"][1]["model"] == "claude-haiku" + assert get_client_requested_model(request) == "AgentX-LLM" + + denied = _alias_token(monkeypatch, level, {"AgentX-LLM": "claude-sonnet"}, ["claude-haiku"]) + denied_data = {"model": "AgentX-LLM"} + with pytest.raises(ProxyException) as exc: + await _enforce_alias_access(denied, denied_data, route, _alias_request(route, denied_data), router) + assert "claude-sonnet" in exc.value.message + + +@pytest.mark.asyncio +async def test_router_settings_model_group_alias_leaves_form_bodies_alone(monkeypatch): + """LIT-3054: a multipart body cannot be re-serialized as JSON, so auth must not rewrite it.""" + router = _alias_router() + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + data = {"model": "AgentX-LLM"} + route = "/v1/audio/transcriptions" + request = _alias_request(route, data, content_type="multipart/form-data; boundary=x") + token = _alias_token(monkeypatch, "key", {"AgentX-LLM": "claude-haiku"}, ["claude-haiku", "AgentX-LLM"]) + await _enforce_alias_access(token, data, route, request, router) + assert data["model"] == "AgentX-LLM" + assert get_client_requested_model(request) is None + + +@pytest.mark.asyncio +async def test_router_settings_model_group_alias_rewrite_keeps_query_params_out_of_body(monkeypatch): + """LIT-3054: auth merges query params into its own copy of the body; the rewrite must not forward them.""" + from litellm.proxy.common_utils.http_parsing_utils import _read_request_body, populate_request_with_path_params + + router = _alias_router() + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + body = {"model": "AgentX-LLM", "messages": [{"role": "user", "content": "hi"}]} + request = _alias_request("/v1/chat/completions", body) + request.scope["query_string"] = b"api-version=2024-10-21&stream=true" + data = populate_request_with_path_params(request_data=await _read_request_body(request), request=request) + assert data["api-version"] == "2024-10-21" + token = _alias_token(monkeypatch, "key", {"AgentX-LLM": "claude-haiku"}, ["claude-haiku"]) + await _enforce_alias_access(token, data, "/v1/chat/completions", request, router) + downstream = await _read_request_body(request) + assert downstream == {**body, "model": "claude-haiku"} + assert json.loads(await request.body()) == downstream + assert await request.json() == downstream + + +def _user_defined_pass_through_endpoint(): + from litellm.types.passthrough_endpoints.pass_through_endpoints import LITELLM_PASS_THROUGH_ENDPOINT_MARKER + + async def endpoint(): + return None + + setattr(endpoint, LITELLM_PASS_THROUGH_ENDPOINT_MARKER, True) + return endpoint + + +@pytest.mark.asyncio +@pytest.mark.parametrize("user_defined", [False, True]) +async def test_router_settings_model_group_alias_leaves_pass_through_bodies_alone(monkeypatch, user_defined): + """LIT-3054: pass-through handlers forward the body verbatim to the provider, so auth must not rewrite it. + Built-in provider handlers bind ``{endpoint:path}``; user-defined ones carry the pass-through marker.""" + router = _alias_router() + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + data = {"model": "AgentX-LLM", "messages": [{"role": "user", "content": "hi"}]} + route = "/custom-upstream/chat" if user_defined else "/anthropic/v1/messages" + request = _alias_request(route, data, path_params={} if user_defined else {"endpoint": "v1/messages"}) + if user_defined: + request.scope["endpoint"] = _user_defined_pass_through_endpoint() + LiteLLMRoutes.openai_routes.value.append(route) + token = _alias_token(monkeypatch, "key", {"AgentX-LLM": "claude-haiku"}, ["claude-haiku", "AgentX-LLM"]) + try: + await _enforce_alias_access(token, data, route, request, router) + finally: + if user_defined: + LiteLLMRoutes.openai_routes.value.remove(route) + assert data["model"] == "AgentX-LLM" + assert (await request.json())["model"] == "AgentX-LLM" + assert get_client_requested_model(request) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("target, expect_denied", [("claude-haiku", False), ("claude-sonnet", True)]) +async def test_router_settings_model_group_alias_authorizes_target_for_team(monkeypatch, target, expect_denied): + """LIT-3054: the team allowlist check in common_checks must judge the alias target, not the alias.""" + import litellm.proxy.proxy_server as _proxy_server_mod + from litellm.proxy.auth.user_api_key_auth import _authorize_authenticated_request + + router = _alias_router() + logging_obj = MagicMock(post_call_failure_hook=AsyncMock(return_value=None)) + attrs = {**_proxy_attrs_for_centralized_checks(), "llm_router": router, "proxy_logging_obj": logging_obj} + for k, v in attrs.items(): + monkeypatch.setattr(_proxy_server_mod, k, v) + token = _alias_token(monkeypatch, "team", {"AgentX-LLM": target}, ["claude-haiku"]) + token.team_models = ["claude-haiku"] + data = {"model": "AgentX-LLM", "messages": [{"role": "user", "content": "hi"}]} + route = "/v1/chat/completions" + request = _alias_request(route, data) + authorize = partial( + _authorize_authenticated_request, + user_api_key_auth_obj=token, + request=request, + request_data=data, + route=route, + api_key="sk-test", + ) + if expect_denied: + with pytest.raises(ProxyException) as exc: + await authorize() + assert exc.value.type == ProxyErrorTypes.team_model_access_denied + assert target in exc.value.message + return + await authorize() + assert (await request.json())["model"] == target + assert get_client_requested_model(request) == "AgentX-LLM" diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 735b0dee3dc..099204cd6c6 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -13,7 +13,11 @@ from fastapi.responses import JSONResponse, StreamingResponse import litellm from litellm._uuid import uuid -from litellm.constants import MAX_LITELLM_CALL_ID_LENGTH, RETURN_RAW_MODEL_NAME_METADATA_KEY +from litellm.constants import ( + CLIENT_REQUESTED_MODEL_SCOPE_KEY, + MAX_LITELLM_CALL_ID_LENGTH, + RETURN_RAW_MODEL_NAME_METADATA_KEY, +) from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.opentelemetry import UserAPIKeyAuth from litellm.proxy.common_request_processing import ( @@ -4395,6 +4399,53 @@ class TestDisconnectGatherCleanup: ) +@pytest.mark.asyncio +@pytest.mark.parametrize("client_model, expected", [("AgentX-LLM", "AgentX-LLM"), (None, "gpt-mini")]) +async def test_response_model_echoes_the_name_the_client_sent_before_auth_rewrote_it( + monkeypatch, client_model, expected +): + """LIT-3054: auth resolves router_settings.model_group_alias in the body, so the alias the + client sent only survives in the request scope. The response must still echo it.""" + import litellm.proxy.common_request_processing as cpr + + async def llm(): + return litellm.ModelResponse( + model="gpt-4o-mini", choices=[{"message": {"role": "assistant", "content": "pong"}}] + ) + + async def fake_route_request(**_kwargs): + return llm() + + logging_obj = MagicMock(litellm_call_id="call-id", _defer_async_logging=False) + proxy_logging = MagicMock(spec=ProxyLogging) + proxy_logging.during_call_hook = AsyncMock(return_value=None) + proxy_logging.post_call_success_hook = AsyncMock(side_effect=lambda data, user_api_key_dict, response: response) + proxy_logging.post_call_response_headers_hook = AsyncMock(return_value={}) + proxy_logging._callback_capabilities_cache = {} + monkeypatch.setattr(cpr, "route_request", fake_route_request) + + processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-mini", "messages": []}) + monkeypatch.setattr( + processor, "common_processing_pre_call_logic", AsyncMock(return_value=({"model": "gpt-mini"}, logging_obj)) + ) + monkeypatch.setattr(processor, "_has_post_call_guardrails", MagicMock(return_value=False)) + scope = {"type": "http", "method": "POST", "path": "/v1/chat/completions", "headers": [], "query_string": b""} + request = Request({**scope, CLIENT_REQUESTED_MODEL_SCOPE_KEY: client_model} if client_model else scope) + + response = await processor.base_process_llm_request( + request=request, + fastapi_response=Response(), + user_api_key_dict=ProxyUserAPIKeyAuth(), + proxy_logging_obj=proxy_logging, + general_settings={}, + proxy_config=MagicMock(spec=ProxyConfig), + route_type="acompletion", + version=None, + ) + + assert response.model == expected + + class TestStreamingClientDisconnectLogging: @pytest.mark.asyncio async def test_record_streaming_client_disconnect_sets_error_information(self):