From 8b66fd70f6466b5059a0091d42ae9d92ca747689 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 18:09:35 -0700 Subject: [PATCH 01/11] fix(proxy): reject client-supplied Vertex location and project unless the admin opts in A request body could carry vertex_location, vertex_project, or their vertex_ai_* aliases and override the region and project pinned on a Vertex AI deployment, sending the admin's credentials to a caller-chosen project or region. Add the four fields to the banned request params so they need allow_client_side_credentials or configurable_clientside_auth_params. /vertex_ai/live and /rag/ingest name their Vertex project and location by design and have no deployment pin to override, so they keep those two fields while every other banned param is still rejected there --- litellm/proxy/auth/auth_utils.py | 18 +++- litellm/proxy/rag_endpoints/endpoints.py | 3 +- tests/unit/proxy/auth/test_auth_utils.py | 106 +++++++++++++++++++++++ 3 files changed, 125 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 94cf7ff4dbe..3e28eaf8b2a 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -304,6 +304,18 @@ def route_manages_mcp_servers(route: str | None) -> bool: return route is not None and _MCP_SERVER_MANAGEMENT_ROUTE.fullmatch(route) is not None +_VERTEX_TARGET_NAMING_ROUTES: Final[frozenset[str]] = frozenset(("/vertex_ai/live", "/rag/ingest", "/v1/rag/ingest")) +_CALLER_NAMED_VERTEX_TARGET_PARAMS: Final[frozenset[str]] = frozenset(("vertex_project", "vertex_location")) + + +def _without_caller_named_vertex_target(request_body: Mapping[str, object], route: str | None) -> Mapping[str, object]: + if route not in _VERTEX_TARGET_NAMING_ROUTES: + return request_body + return MappingProxyType( + {key: value for key, value in request_body.items() if key not in _CALLER_NAMED_VERTEX_TARGET_PARAMS} + ) + + _NESTED_CONFIG_KEYS: Final[tuple[str, ...]] = ("litellm_embedding_config", "extra_body") # Metadata containers that carry per-request configuration consumed by the @@ -396,6 +408,10 @@ _BANNED_REQUEST_BODY_PARAMS: Final[tuple[str, ...]] = ( "aws_external_id", "aws_session_tags", "vertex_credentials", + "vertex_location", + "vertex_project", + "vertex_ai_location", + "vertex_ai_project", # Azure managed-identity / federated-auth token. The Azure provider # transformer reads ``azure_ad_token`` (top-level or via # ``extra_body``) and resolves it through ``get_secret`` before @@ -650,7 +666,7 @@ def is_request_body_safe( raise ValueError("Rejected Request: model_list is not allowed in the request body.") manages_deployments: Final = route_manages_deployments(route) _check_banned_params( - request_body, + _without_caller_named_vertex_target(request_body, route), general_settings, llm_router, model, diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index a1a7d62899f..38f9b2200cf 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -24,7 +24,7 @@ from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook i LiteLLM_ManagedVectorStore, ) from litellm.proxy._types import * -from litellm.proxy.auth.auth_utils import is_request_body_safe +from litellm.proxy.auth.auth_utils import get_request_route, is_request_body_safe from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, @@ -615,6 +615,7 @@ async def rag_ingest( general_settings=general_settings, llm_router=llm_router, model="", + route=get_request_route(request=request), ) except ValueError as e: raise HTTPException(status_code=400, detail={"error": str(e)}) diff --git a/tests/unit/proxy/auth/test_auth_utils.py b/tests/unit/proxy/auth/test_auth_utils.py index 7de249bfb6e..4cf67028d9b 100644 --- a/tests/unit/proxy/auth/test_auth_utils.py +++ b/tests/unit/proxy/auth/test_auth_utils.py @@ -4351,6 +4351,112 @@ class TestIsRequestBodySafeBlocksAwsIdentitySelectors: ) +_VERTEX_TARGET_PARAMS: Final = ("vertex_location", "vertex_project", "vertex_ai_location", "vertex_ai_project") + + +def _vertex_pinned_router(configurable_clientside_auth_params: list[str]) -> Router: + return Router( + model_list=[ + { + "model_name": "gemini-2.5-flash", + "litellm_params": { + "model": "vertex_ai/gemini-2.5-flash", + "vertex_project": "admin-project", + "vertex_location": "europe-west4", + "configurable_clientside_auth_params": configurable_clientside_auth_params, + }, + } + ] + ) + + +class TestIsRequestBodySafeBlocksVertexTargetOverrides: + @pytest.mark.parametrize("param", _VERTEX_TARGET_PARAMS) + @pytest.mark.parametrize( + "placement", + [ + lambda param: {param: "us-central1"}, + lambda param: {"extra_body": {param: "us-central1"}}, + lambda param: {"metadata": {param: "us-central1"}}, + lambda param: {"litellm_metadata": {param: "us-central1"}}, + ], + ids=["root", "extra_body", "metadata", "litellm_metadata"], + ) + def test_vertex_target_is_rejected_without_admin_opt_in( + self, param: str, placement: Callable[[str], dict[str, object]] + ) -> None: + with pytest.raises(ValueError, match=param): + is_request_body_safe( + request_body={"model": "gemini-2.5-flash", **placement(param)}, + general_settings={}, + llm_router=_vertex_pinned_router([]), + model="gemini-2.5-flash", + route="/v1/chat/completions", + ) + + def test_vertex_target_is_allowed_under_proxy_wide_opt_in(self) -> None: + assert ( + is_request_body_safe( + request_body={"model": "gemini-2.5-flash", **{param: "us-central1" for param in _VERTEX_TARGET_PARAMS}}, + general_settings={"allow_client_side_credentials": True}, + llm_router=None, + model="gemini-2.5-flash", + route="/v1/chat/completions", + ) + is True + ) + + def test_deployment_opt_in_allows_only_the_listed_vertex_target(self) -> None: + router: Final = _vertex_pinned_router(["vertex_location"]) + assert ( + is_request_body_safe( + request_body={"model": "gemini-2.5-flash", "vertex_location": "us-central1"}, + general_settings={}, + llm_router=router, + model="gemini-2.5-flash", + route="/v1/chat/completions", + ) + is True + ) + with pytest.raises(ValueError, match="vertex_project"): + is_request_body_safe( + request_body={ + "model": "gemini-2.5-flash", + "vertex_location": "us-central1", + "vertex_project": "caller-project", + }, + general_settings={}, + llm_router=router, + model="gemini-2.5-flash", + route="/v1/chat/completions", + ) + + @pytest.mark.parametrize("route", ["/vertex_ai/live", "/rag/ingest", "/v1/rag/ingest"]) + def test_routes_that_name_their_vertex_target_keep_project_and_location(self, route: str) -> None: + assert ( + is_request_body_safe( + request_body={"vertex_project": "caller-project", "vertex_location": "us-central1"}, + general_settings={}, + llm_router=None, + model="", + route=route, + ) + is True + ) + + @pytest.mark.parametrize("route", ["/vertex_ai/live", "/rag/ingest"]) + @pytest.mark.parametrize("param", ["api_base", "vertex_ai_project", "vertex_ai_location"]) + def test_routes_that_name_their_vertex_target_reject_every_other_banned_param(self, route: str, param: str) -> None: + with pytest.raises(ValueError, match=param): + is_request_body_safe( + request_body={"vertex_location": "us-central1", param: "caller-chosen"}, + general_settings={}, + llm_router=None, + model="", + route=route, + ) + + @pytest.mark.parametrize( "allowed_param, input_value, should_return_true", [ From 72a7567ef15300fda864d7533c8d3f2e86d16b42 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 20:11:03 -0700 Subject: [PATCH 02/11] fix(proxy): close the remaining Vertex target overrides and keep inert fields working /vertex_ai/live now closes with 1008 when the query names a project or location that differs from a use_in_pass_through deployment's pin, unless that deployment or the proxy opts in A model group's opt-in now counts only when every deployment in it carries it, and Google-style routes look the opt-in up on the model named in the URL Vertex target names inside metadata, the live route's vertex_ai_* aliases, and /health/test_connection inheritance go back to the base behavior, since nothing reads them as a target --- litellm/proxy/auth/auth_utils.py | 141 ++++++++++-------- .../health_endpoints/_health_endpoints.py | 8 +- .../llm_passthrough_endpoints.py | 75 ++++++++-- .../passthrough_endpoint_router.py | 10 +- tests/unit/proxy/auth/test_auth_utils.py | 128 +++++++++++++++- .../health_endpoints/test_health_endpoints.py | 6 + .../test_llm_pass_through_endpoints.py | 109 ++++++++++++++ 7 files changed, 393 insertions(+), 84 deletions(-) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 3e28eaf8b2a..ef4f99420e9 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -36,7 +36,12 @@ from litellm.proxy.common_utils.http_parsing_utils import extract_nested_form_me from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_ENDPOINT_MARKER, ) -from litellm.types.router import CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS, Deployment, server_owned_wif_fields_present +from litellm.types.router import ( + CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS, + Deployment, + LiteLLMParamsTypedDict, + server_owned_wif_fields_present, +) from litellm.types.router import reject_server_owned_wif_params as _reject_server_owned_wif_params from litellm.types.utils import CustomPricingLiteLLMParams, oauth_token_exchange_litellm_params @@ -201,34 +206,34 @@ def _is_param_allowed( return False +def deployment_clientside_auth_params(litellm_params: LiteLLMParamsTypedDict) -> CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS: + if "configurable_clientside_auth_params" not in litellm_params: + return None + return litellm_params["configurable_clientside_auth_params"] + + def _allow_model_level_clientside_configurable_parameters( model: str, param: str, request_body_value: object, llm_router: Router | None ) -> bool: """ - Check if model is allowed to use configurable client-side params - - get matching model - - check if 'clientside_configurable_parameters' is set for model - - + The router, not the caller, picks which deployment serves the request, so an + opt-in counts only when every deployment the model can route to carries it """ if llm_router is None: return False - # check if model is set - model_info = llm_router.get_model_group_info(model_group=model) - if model_info is None: - # check if wildcard model is set - if model.split("/", 1)[0] in provider_list: - model_info = llm_router.get_model_group_info(model_group=model.split("/", 1)[0]) - - if model_info is None: - return False - - if model_info is None or model_info.configurable_clientside_auth_params is None: - return False - - return _is_param_allowed( - param=param, - request_body_value=request_body_value, - configurable_clientside_auth_params=model_info.configurable_clientside_auth_params, + provider_group: Final = model.split("/", 1)[0] + deployments: Final = ( + llm_router.get_model_list(model_name=model) + or (llm_router.get_model_list(model_name=provider_group) if provider_group in provider_list else None) + or () + ) + return bool(deployments) and all( + _is_param_allowed( + param=param, + request_body_value=request_body_value, + configurable_clientside_auth_params=deployment_clientside_auth_params(deployment["litellm_params"]), + ) + for deployment in deployments ) @@ -304,16 +309,24 @@ def route_manages_mcp_servers(route: str | None) -> bool: return route is not None and _MCP_SERVER_MANAGEMENT_ROUTE.fullmatch(route) is not None -_VERTEX_TARGET_NAMING_ROUTES: Final[frozenset[str]] = frozenset(("/vertex_ai/live", "/rag/ingest", "/v1/rag/ingest")) -_CALLER_NAMED_VERTEX_TARGET_PARAMS: Final[frozenset[str]] = frozenset(("vertex_project", "vertex_location")) +VERTEX_TARGET_PARAMS: Final[frozenset[str]] = frozenset( + ("vertex_project", "vertex_location", "vertex_ai_project", "vertex_ai_location") +) +_RAG_INGEST_VERTEX_TARGET_PARAMS: Final[frozenset[str]] = frozenset(("vertex_project", "vertex_location")) +_ROUTE_NAMED_VERTEX_TARGET_PARAMS: Final[Mapping[str, frozenset[str]]] = MappingProxyType( + { + "/vertex_ai/live": VERTEX_TARGET_PARAMS, + "/rag/ingest": _RAG_INGEST_VERTEX_TARGET_PARAMS, + "/v1/rag/ingest": _RAG_INGEST_VERTEX_TARGET_PARAMS, + } +) -def _without_caller_named_vertex_target(request_body: Mapping[str, object], route: str | None) -> Mapping[str, object]: - if route not in _VERTEX_TARGET_NAMING_ROUTES: +def _without_route_named_vertex_target(request_body: Mapping[str, object], route: str | None) -> Mapping[str, object]: + route_named: Final = _ROUTE_NAMED_VERTEX_TARGET_PARAMS.get(route) if route is not None else None + if route_named is None: return request_body - return MappingProxyType( - {key: value for key, value in request_body.items() if key not in _CALLER_NAMED_VERTEX_TARGET_PARAMS} - ) + return MappingProxyType({key: value for key, value in request_body.items() if key not in route_named}) _NESTED_CONFIG_KEYS: Final[tuple[str, ...]] = ("litellm_embedding_config", "extra_body") @@ -477,6 +490,10 @@ _BANNED_REQUEST_BODY_PARAMS: Final[tuple[str, ...]] = ( *sorted(CustomPricingLiteLLMParams.model_fields.keys()), ) +_BANNED_METADATA_PARAMS: Final[tuple[str, ...]] = tuple( + param for param in _BANNED_REQUEST_BODY_PARAMS if param not in VERTEX_TARGET_PARAMS +) + def _check_banned_params( body: Mapping[str, object], @@ -486,6 +503,7 @@ def _check_banned_params( *, manages_deployments: bool = False, mcp_server_fields_allowed: bool = False, + banned_params: tuple[str, ...] = _BANNED_REQUEST_BODY_PARAMS, ) -> None: """Raise ``ValueError`` if ``body`` carries a banned param without admin opt-in. @@ -500,7 +518,7 @@ def _check_banned_params( reject_server_owned_wif_params(wif_body) if not manages_deployments: reject_federated_credential_reference(body) - for param in _BANNED_REQUEST_BODY_PARAMS: + for param in banned_params: if param not in body: continue if general_settings.get("allow_client_side_credentials") is True: @@ -666,7 +684,7 @@ def is_request_body_safe( raise ValueError("Rejected Request: model_list is not allowed in the request body.") manages_deployments: Final = route_manages_deployments(route) _check_banned_params( - _without_caller_named_vertex_target(request_body, route), + _without_route_named_vertex_target(request_body, route), general_settings, llm_router, model, @@ -680,7 +698,14 @@ def is_request_body_safe( for metadata_key in _NESTED_METADATA_KEYS: metadata = _coerce_metadata_to_dict(request_body.get(metadata_key)) if metadata is not None: - _check_banned_params(metadata, general_settings, llm_router, model, manages_deployments=manages_deployments) + _check_banned_params( + metadata, + general_settings, + llm_router, + model, + manages_deployments=manages_deployments, + banned_params=_BANNED_METADATA_PARAMS, + ) if any(isinstance(key, str) and key.startswith(f"{metadata_key}[") for key in request_body): _check_banned_params( extract_nested_form_metadata(form_data=request_body, prefix=f"{metadata_key}["), @@ -688,6 +713,7 @@ def is_request_body_safe( llm_router, model, manages_deployments=manages_deployments, + banned_params=_BANNED_METADATA_PARAMS, ) for target in iter_request_fallback_targets(request_body): if isinstance(target, dict): @@ -710,6 +736,7 @@ def is_request_body_safe( llm_router, model, manages_deployments=manages_deployments, + banned_params=_BANNED_METADATA_PARAMS, ) return True @@ -770,7 +797,7 @@ async def pre_db_read_auth_checks( request_body=request_data, general_settings=general_settings, llm_router=llm_router, - model=request_data.get("model", ""), # [TODO] use model passed in url as well (azure openai routes) + model=request_data.get("model") or _model_from_route_path(route) or "", route=route, ) @@ -2215,37 +2242,8 @@ def get_model_from_request( team_id=team_id, ) model = _format_model_candidates(candidates) - - # If no explicit model was found, try to extract from route if model is None: - # Parse model from route that follows the pattern /openai/deployments/{model}/* - match: Final = re.match(r"/openai/deployments/([^/]+)", route) - if match: - model = match.group(1) - - # If still not found, extract model from Google generateContent-style routes. - # These routes put the model in the path and allow "/" inside the model id. - # Examples: - # - /v1beta/models/gemini-2.0-flash:generateContent - # - /v1beta/models/bedrock/claude-sonnet-3.7:generateContent - # - /models/custom/ns/model:streamGenerateContent - if model is None and not route.lower().startswith("/vertex"): - google_match = re.search(r"/(?:v1beta|beta)/models/([^:]+):", route) - if google_match: - model = google_match.group(1) - - if model is None and not route.lower().startswith("/vertex"): - google_match = re.search(r"^/models/([^:]+):", route) - if google_match: - model = google_match.group(1) - - # If still not found, extract from Vertex AI passthrough route - # Pattern: /vertex_ai/.../models/{model_id}:* - # Example: /vertex_ai/v1/.../models/gemini-1.5-pro:generateContent - if model is None and route.lower().startswith("/vertex"): - vertex_match: Final = re.search(r"/models/([^:]+)", route) - if vertex_match: - model = vertex_match.group(1) + model = _model_from_route_path(route) if route.lower().startswith("/bedrock"): bedrock_model: Final = _model_from_bedrock_route(route) @@ -2264,6 +2262,23 @@ def get_model_from_request( return model +def _model_from_route_path(route: str) -> str | None: + """ + Google generateContent-style routes allow "/" inside the model id, e.g. + /v1beta/models/bedrock/claude-sonnet-3.7:generateContent + """ + azure_deployment: Final = re.match(r"/openai/deployments/([^/]+)", route) + if azure_deployment: + return azure_deployment.group(1) + if route.lower().startswith("/vertex"): + vertex_model: Final = re.search(r"/models/([^:]+)", route) + return vertex_model.group(1) if vertex_model else None + google_model: Final = re.search(r"/(?:v1beta|beta)/models/([^:]+):", route) or re.search( + r"^/models/([^:]+):", route + ) + return google_model.group(1) if google_model else None + + def _router_model_from_azure_route(route: str, llm_router: Router | None) -> str | None: if llm_router is None: return None diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 3c8dd7890fc..08db52c9882 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -42,6 +42,7 @@ from litellm.proxy.auth.auth_checks import ( ) from litellm.proxy.auth.auth_utils import ( _BANNED_REQUEST_BODY_PARAMS, # pyright: ignore[reportPrivateUsage] # one canonical list, shared with the request-body check + VERTEX_TARGET_PARAMS, reject_server_owned_wif_params, ) from litellm.proxy.auth.model_checks import get_key_models @@ -133,6 +134,11 @@ _CONFIG_CONNECTION_FIELDS: Final[frozenset[str]] = frozenset( ) +_BANNED_CONNECTION_PARAMS: Final[tuple[str, ...]] = tuple( + param for param in _BANNED_REQUEST_BODY_PARAMS if param not in VERTEX_TARGET_PARAMS +) + + def _request_inherits_config_credentials( config_params: Mapping[str, object], request_params: Mapping[str, object], @@ -158,7 +164,7 @@ def _request_inherits_config_credentials( return False if allow_client_side_credentials: return True - return not any(param in request_params for param in _BANNED_REQUEST_BODY_PARAMS) + return not any(param in request_params for param in _BANNED_CONNECTION_PARAMS) def _config_base_for_health_check( diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 464cfee270a..96004ab4793 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -72,6 +72,7 @@ from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.passthrough.main import AsyncPassthroughStreamingResponse, is_provider_error from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import enforced_model_allowlists +from litellm.proxy.auth.auth_utils import deployment_clientside_auth_params from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy module exports @@ -3760,11 +3761,25 @@ def _get_llm_router() -> Router | None: return llm_router -def _resolve_vertex_live_credentials( +@dataclass(frozen=True, slots=True) +class _ConfiguredVertexLiveTarget: + credentials: VertexPassThroughCredentials | None + + +@dataclass(frozen=True, slots=True) +class _DeploymentVertexLiveTarget: + credentials: VertexPassThroughCredentials + deployment_params: LiteLLMParamsTypedDict + + +_VertexLiveTarget = _ConfiguredVertexLiveTarget | _DeploymentVertexLiveTarget + + +def _resolve_vertex_live_target( vertex_project: str | None, vertex_location: str | None, model: str | None, -) -> VertexPassThroughCredentials | None: +) -> _VertexLiveTarget: """ Resolution order: an explicit project/location registration, then ``default_vertex_config`` (which the proxy fills from the ``DEFAULT_VERTEXAI_*`` env vars whenever the yaml leaves it out), then any DB model entry @@ -3778,19 +3793,45 @@ def _resolve_vertex_live_credentials( location=vertex_location, ) if keyed is not None and keyed.vertex_project is not None: - return keyed - from_deployments: Final = passthrough_endpoint_router.get_vertex_credentials_from_router_deployments(model=model) + return _ConfiguredVertexLiveTarget(credentials=keyed) + from_deployments: Final = passthrough_endpoint_router.get_vertex_deployment_from_router(model=model) if from_deployments is not None: - return from_deployments + deployment, credentials = from_deployments + return _DeploymentVertexLiveTarget(credentials=credentials, deployment_params=deployment["litellm_params"]) if keyed is not None: - return keyed + return _ConfiguredVertexLiveTarget(credentials=keyed) passthrough_endpoint_router.set_default_vertex_config() - return passthrough_endpoint_router.get_vertex_credentials( - project_id=vertex_project, - location=vertex_location, + return _ConfiguredVertexLiveTarget( + credentials=passthrough_endpoint_router.get_vertex_credentials( + project_id=vertex_project, + location=vertex_location, + ) ) +def _caller_override_of_live_deployment( + target: _VertexLiveTarget, + vertex_project: str | None, + vertex_location: str | None, + allow_client_side_credentials: bool, +) -> str | None: + match target: + case _ConfiguredVertexLiveTarget(): + return None + case _DeploymentVertexLiveTarget(credentials=credentials, deployment_params=deployment_params): + if allow_client_side_credentials: + return None + opted_in: Final = deployment_clientside_auth_params(deployment_params) or () + requested: Final = ( + ("vertex_project", vertex_project, credentials.vertex_project), + ("vertex_location", vertex_location, credentials.vertex_location), + ) + return next( + (name for name, asked, pinned in requested if asked and asked != pinned and name not in opted_in), + None, + ) + + def _build_vertex_live_setup_model_rewriter( vertex_project: str | None, vertex_location: str | None, @@ -3861,11 +3902,25 @@ async def vertex_ai_live_websocket_passthrough( await websocket.accept() incoming_headers: Final = dict(websocket.headers) - vertex_credentials_config: Final = _resolve_vertex_live_credentials( + live_target: Final = _resolve_vertex_live_target( vertex_project=vertex_project, vertex_location=vertex_location, model=model, ) + overridden_field: Final = _caller_override_of_live_deployment( + live_target, + vertex_project=vertex_project, + vertex_location=vertex_location, + allow_client_side_credentials=_proxy_general_settings().get("allow_client_side_credentials") is True, + ) + if overridden_field is not None: + await websocket.close( + code=1008, + reason=f"Rejected Request: {overridden_field} is pinned by the deployment, " + "see configurable_clientside_auth_params", + ) + return + vertex_credentials_config: Final = live_target.credentials configured_project: Final = vertex_project or ( vertex_credentials_config.vertex_project if vertex_credentials_config is not None else None diff --git a/litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py b/litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py index 3d1be882779..b19fdc5c971 100644 --- a/litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py +++ b/litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py @@ -151,6 +151,12 @@ class PassthroughEndpointRouter: return provider def get_vertex_credentials_from_router_deployments(self, model: str | None) -> VertexPassThroughCredentials | None: + found: Final = self.get_vertex_deployment_from_router(model) + return found[1] if found is not None else None + + def get_vertex_deployment_from_router( + self, model: str | None + ) -> tuple[DeploymentTypedDict, VertexPassThroughCredentials] | None: """ Resolve vertex pass-through credentials from the live router deployments flagged ``use_in_pass_through``. @@ -170,7 +176,7 @@ class PassthroughEndpointRouter: ) matched: Final = next( ( - credentials + (deployment, credentials) for deployment, credentials in resolved if model is not None and self._deployment_matches_model(deployment, model) ), @@ -188,7 +194,7 @@ class PassthroughEndpointRouter: ) if len(targets) != 1: return None - return resolved[0][1] + return resolved[0] def _resolve_vertex_deployment_credentials( self, litellm_params: LiteLLMParamsTypedDict diff --git a/tests/unit/proxy/auth/test_auth_utils.py b/tests/unit/proxy/auth/test_auth_utils.py index 4cf67028d9b..b6fbe779716 100644 --- a/tests/unit/proxy/auth/test_auth_utils.py +++ b/tests/unit/proxy/auth/test_auth_utils.py @@ -33,6 +33,7 @@ from litellm.proxy.auth.auth_utils import ( get_project_model_tpm_limit, get_request_route_template, is_request_body_safe, + pre_db_read_auth_checks, ) from litellm.router import Router from litellm.types.utils import oauth_token_exchange_litellm_params, server_owned_wif_litellm_params @@ -4377,10 +4378,8 @@ class TestIsRequestBodySafeBlocksVertexTargetOverrides: [ lambda param: {param: "us-central1"}, lambda param: {"extra_body": {param: "us-central1"}}, - lambda param: {"metadata": {param: "us-central1"}}, - lambda param: {"litellm_metadata": {param: "us-central1"}}, ], - ids=["root", "extra_body", "metadata", "litellm_metadata"], + ids=["root", "extra_body"], ) def test_vertex_target_is_rejected_without_admin_opt_in( self, param: str, placement: Callable[[str], dict[str, object]] @@ -4394,6 +4393,31 @@ class TestIsRequestBodySafeBlocksVertexTargetOverrides: route="/v1/chat/completions", ) + @pytest.mark.parametrize("param", _VERTEX_TARGET_PARAMS) + @pytest.mark.parametrize( + "placement", + [ + lambda param: {"metadata": {param: "us-central1"}}, + lambda param: {"litellm_metadata": {param: "us-central1"}}, + lambda param: {f"metadata[{param}]": "us-central1"}, + lambda param: {"litellm_params": {"metadata": {param: "us-central1"}}}, + ], + ids=["metadata", "litellm_metadata", "form_metadata", "litellm_params_metadata"], + ) + def test_vertex_target_in_metadata_never_reaches_the_provider_and_is_allowed( + self, param: str, placement: Callable[[str], dict[str, object]] + ) -> None: + assert ( + is_request_body_safe( + request_body={"model": "gemini-2.5-flash", **placement(param)}, + general_settings={}, + llm_router=_vertex_pinned_router([]), + model="gemini-2.5-flash", + route="/v1/chat/completions", + ) + is True + ) + def test_vertex_target_is_allowed_under_proxy_wide_opt_in(self) -> None: assert ( is_request_body_safe( @@ -4431,11 +4455,92 @@ class TestIsRequestBodySafeBlocksVertexTargetOverrides: route="/v1/chat/completions", ) - @pytest.mark.parametrize("route", ["/vertex_ai/live", "/rag/ingest", "/v1/rag/ingest"]) - def test_routes_that_name_their_vertex_target_keep_project_and_location(self, route: str) -> None: + @pytest.mark.parametrize("opted_in_first", [True, False]) + def test_group_opt_in_does_not_count_while_any_deployment_lacks_it(self, opted_in_first: bool) -> None: + opted_in: Final = { + "model_name": "gemini-group", + "litellm_params": { + "model": "vertex_ai/gemini-2.5-flash", + "vertex_location": "us-central1", + "configurable_clientside_auth_params": ["vertex_location"], + }, + } + pinned: Final = { + "model_name": "gemini-group", + "litellm_params": {"model": "vertex_ai/gemini-2.5-flash", "vertex_location": "europe-west4"}, + } + router: Final = Router(model_list=[opted_in, pinned] if opted_in_first else [pinned, opted_in]) + with pytest.raises(ValueError, match="vertex_location"): + is_request_body_safe( + request_body={"model": "gemini-group", "vertex_location": "us-central1"}, + general_settings={}, + llm_router=router, + model="gemini-group", + route="/v1/chat/completions", + ) + + def test_model_with_no_deployment_has_no_opt_in(self) -> None: + with pytest.raises(ValueError, match="vertex_location"): + is_request_body_safe( + request_body={"model": "not-configured", "vertex_location": "us-central1"}, + general_settings={}, + llm_router=_vertex_pinned_router(["vertex_location"]), + model="not-configured", + route="/v1/chat/completions", + ) + + def test_group_opt_in_counts_when_every_deployment_carries_it(self) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "gemini-group", + "litellm_params": { + "model": "vertex_ai/gemini-2.5-flash", + "vertex_location": location, + "configurable_clientside_auth_params": ["vertex_location"], + }, + } + for location in ("us-central1", "europe-west4") + ] + ) assert ( is_request_body_safe( - request_body={"vertex_project": "caller-project", "vertex_location": "us-central1"}, + request_body={"model": "gemini-group", "vertex_location": "us-central1"}, + general_settings={}, + llm_router=router, + model="gemini-group", + route="/v1/chat/completions", + ) + is True + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize("opted_in", [True, False]) + async def test_opt_in_of_the_model_named_in_a_google_route_decides(self, monkeypatch, opted_in: bool) -> None: + route: Final = "/v1beta/models/gemini-2.5-flash:generateContent" + monkeypatch.setattr( + "litellm.proxy.proxy_server.llm_router", _vertex_pinned_router(["vertex_location"] if opted_in else []) + ) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + request: Final = Request({"type": "http", "method": "POST", "path": route, "headers": [], "query_string": b""}) + if opted_in: + await pre_db_read_auth_checks(request=request, request_data={"vertex_location": "us-central1"}, route=route) + return + with pytest.raises(ValueError, match="vertex_location"): + await pre_db_read_auth_checks(request=request, request_data={"vertex_location": "us-central1"}, route=route) + + @pytest.mark.parametrize( + "route, kept_params", + [ + ("/vertex_ai/live", _VERTEX_TARGET_PARAMS), + ("/rag/ingest", ("vertex_project", "vertex_location")), + ("/v1/rag/ingest", ("vertex_project", "vertex_location")), + ], + ) + def test_routes_that_name_their_vertex_target_keep_it(self, route: str, kept_params: tuple[str, ...]) -> None: + assert ( + is_request_body_safe( + request_body={param: "caller-chosen" for param in kept_params}, general_settings={}, llm_router=None, model="", @@ -4444,8 +4549,15 @@ class TestIsRequestBodySafeBlocksVertexTargetOverrides: is True ) - @pytest.mark.parametrize("route", ["/vertex_ai/live", "/rag/ingest"]) - @pytest.mark.parametrize("param", ["api_base", "vertex_ai_project", "vertex_ai_location"]) + @pytest.mark.parametrize( + "route, param", + [ + ("/vertex_ai/live", "api_base"), + ("/rag/ingest", "api_base"), + ("/rag/ingest", "vertex_ai_project"), + ("/rag/ingest", "vertex_ai_location"), + ], + ) def test_routes_that_name_their_vertex_target_reject_every_other_banned_param(self, route: str, param: str) -> None: with pytest.raises(ValueError, match=param): is_request_body_safe( diff --git a/tests/unit/proxy/health_endpoints/test_health_endpoints.py b/tests/unit/proxy/health_endpoints/test_health_endpoints.py index 441c52b6a5c..28e89fe1ffb 100644 --- a/tests/unit/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/unit/proxy/health_endpoints/test_health_endpoints.py @@ -4458,6 +4458,12 @@ class TestConfigBaseForHealthCheck: assert "vertex_credentials" not in base assert base["rpm"] == 100 + @pytest.mark.parametrize("param", ["vertex_project", "vertex_location", "vertex_ai_project", "vertex_ai_location"]) + def test_request_naming_a_vertex_target_still_inherits_config_credentials(self, param: str) -> None: + base = self._base(self.CONFIG, {"model": "openai/gpt-4o", param: "us-central1"}, selected_by_id=False) + assert base["vertex_credentials"] == "configured-creds" + assert base["api_key"] == "sk-configured" + def test_add_model_flow_keeps_its_own_credentials(self): """Adding a second deployment for an already-configured name sends a complete connection; it is tested as sent, not as configured.""" diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 21170abaeb1..8d221e0009f 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -6246,6 +6246,115 @@ class TestVertexAILiveWebsocketPassthrough: assert "default_vertex_config" in close_kwargs["reason"] assert len(close_kwargs["reason"].encode("utf-8")) <= 123 + def _pinned_live_router(self, extra_params: Mapping[str, object]) -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": "gemini-live", + "litellm_params": { + "model": "vertex_ai/gemini-live-2.5-flash", + "use_in_pass_through": True, + "vertex_project": "proj-db", + "vertex_location": "global", + "vertex_credentials": '{"type": "service_account"}', + **extra_params, + }, + } + ] + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "query, field", + [ + ({"vertex_location": "us-central1"}, "vertex_location"), + ({"vertex_project": "caller-project"}, "vertex_project"), + ({"vertex_project": "proj-db", "vertex_location": "us-central1"}, "vertex_location"), + ], + ) + async def test_query_overriding_a_pinned_deployment_closes_before_minting_a_token( + self, monkeypatch, query: Mapping[str, str], field: str + ): + from litellm.proxy.pass_through_endpoints import ( + llm_passthrough_endpoints as passthrough_module, + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", self._pinned_live_router({})) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + monkeypatch.setattr(passthrough_module.passthrough_endpoint_router, "default_vertex_config", None) + self._clear_vertex_env(monkeypatch) + websocket = self._websocket() + ensure_token = AsyncMock(return_value=("token-abc", "proj-db")) + ws_passthrough = AsyncMock() + + with ( + patch.object(passthrough_module.vertex_llm_base, "ensure_access_token_async", ensure_token), + patch.object(passthrough_module, "websocket_passthrough_request", ws_passthrough), + ): + await passthrough_module.vertex_ai_live_websocket_passthrough( + websocket=websocket, + user_api_key_dict=UserAPIKeyAuth(), + **query, + ) + + close_kwargs = websocket.close.await_args.kwargs + assert close_kwargs["code"] == 1008 + assert field in close_kwargs["reason"] + assert len(close_kwargs["reason"].encode("utf-8")) <= 123 + ensure_token.assert_not_awaited() + ws_passthrough.assert_not_awaited() + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "extra_params, general_settings, query, expected_location", + [ + ({}, {}, {"vertex_project": "proj-db", "vertex_location": "global"}, "global"), + ( + {"configurable_clientside_auth_params": ["vertex_location"]}, + {}, + {"vertex_location": "us-central1"}, + "us-central1", + ), + ({}, {"allow_client_side_credentials": True}, {"vertex_location": "us-central1"}, "us-central1"), + ], + ids=["matches_the_pin", "deployment_opt_in", "proxy_wide_opt_in"], + ) + async def test_query_the_admin_allows_reaches_vertex( + self, + monkeypatch, + extra_params: Mapping[str, object], + general_settings: Mapping[str, object], + query: Mapping[str, str], + expected_location: str, + ): + from litellm.proxy.pass_through_endpoints import ( + llm_passthrough_endpoints as passthrough_module, + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", self._pinned_live_router(extra_params)) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", dict(general_settings)) + monkeypatch.setattr(passthrough_module.passthrough_endpoint_router, "default_vertex_config", None) + self._clear_vertex_env(monkeypatch) + websocket = self._websocket() + ensure_token = AsyncMock(return_value=("token-abc", "proj-db")) + ws_passthrough = AsyncMock() + + with ( + patch.object(passthrough_module.vertex_llm_base, "ensure_access_token_async", ensure_token), + patch.object(passthrough_module, "websocket_passthrough_request", ws_passthrough), + ): + await passthrough_module.vertex_ai_live_websocket_passthrough( + websocket=websocket, + user_api_key_dict=UserAPIKeyAuth(), + **query, + ) + + websocket.close.assert_not_awaited() + rewriter = ws_passthrough.await_args.kwargs["setup_model_rewriter"] + assert rewriter("gemini-live") == ( + f"projects/proj-db/locations/{expected_location}/publishers/google/models/gemini-live-2.5-flash" + ) + class TestAnthropicProxyRoute: """The /anthropic passthrough route: custom auth headers must not clobber the From f035ca0fb51788df9291cf194438055da6c3090e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 20:48:14 -0700 Subject: [PATCH 03/11] test(batches): give the batch FakeRouter the model-list lookup the opt-in check reads --- .../proxy/batches_endpoints/test_litellm_executed_batches.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/proxy/batches_endpoints/test_litellm_executed_batches.py b/tests/unit/proxy/batches_endpoints/test_litellm_executed_batches.py index 00f84e75664..993e54304d0 100644 --- a/tests/unit/proxy/batches_endpoints/test_litellm_executed_batches.py +++ b/tests/unit/proxy/batches_endpoints/test_litellm_executed_batches.py @@ -251,7 +251,7 @@ class FakeRouter: def get_model_ids(self, model_name: str) -> list[str]: return [DEPLOYMENT_ID] if model_name == BATCH_MODEL else [] - def get_model_group_info(self, model_group: str) -> None: + def get_model_list(self, model_name: str | None = None) -> None: return None From 5f984fb6ab9fcc27e39f7b34972ed28d19ddc6c0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 05:16:40 -0700 Subject: [PATCH 04/11] fix(proxy): pin default_vertex_config on /vertex_ai/live and the proxy's Vertex project on /rag/ingest --- .../llm_passthrough_endpoints.py | 79 ++++++-- litellm/proxy/rag_endpoints/endpoints.py | 45 +++++ litellm/rag/ingestion/vertex_ai_ingestion.py | 4 +- .../test_llm_pass_through_endpoints.py | 173 +++++++++++++++++- .../proxy/rag_endpoints/test_rag_endpoints.py | 117 ++++++++++++ 5 files changed, 397 insertions(+), 21 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 6b6069e7cae..b092bf4fd37 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -3904,46 +3904,89 @@ def _resolve_vertex_live_target( ) -def _caller_override_of_live_deployment( +def _requested_vertex_live_target( + credentials: VertexPassThroughCredentials, vertex_project: str | None, vertex_location: str | None +) -> tuple[tuple[str, str | None, str | None], ...]: + return ( + ("vertex_project", vertex_project, credentials.vertex_project), + ("vertex_location", vertex_location, credentials.vertex_location), + ) + + +def _live_target_override_reason( target: _VertexLiveTarget, vertex_project: str | None, vertex_location: str | None, allow_client_side_credentials: bool, ) -> str | None: + if allow_client_side_credentials: + return None match target: + case _ConfiguredVertexLiveTarget(credentials=VertexPassThroughCredentials() as credentials): + configured_field: Final = next( + ( + name + for name, asked, pinned in _requested_vertex_live_target( + credentials, vertex_project, vertex_location + ) + if asked and pinned and asked != pinned + ), + None, + ) + if configured_field is None: + return None + return ( + f"Rejected Request: {configured_field} is pinned by default_vertex_config, " + "see allow_client_side_credentials" + ) case _ConfiguredVertexLiveTarget(): return None case _DeploymentVertexLiveTarget(credentials=credentials, deployment_params=deployment_params): - if allow_client_side_credentials: - return None opted_in: Final = deployment_clientside_auth_params(deployment_params) or () - requested: Final = ( - ("vertex_project", vertex_project, credentials.vertex_project), - ("vertex_location", vertex_location, credentials.vertex_location), - ) - return next( - (name for name, asked, pinned in requested if asked and asked != pinned and name not in opted_in), + deployment_field: Final = next( + ( + name + for name, asked, pinned in _requested_vertex_live_target( + credentials, vertex_project, vertex_location + ) + if asked and asked != pinned and name not in opted_in + ), None, ) + if deployment_field is None: + return None + return ( + f"Rejected Request: {deployment_field} is pinned by the deployment, " + "see configurable_clientside_auth_params" + ) + + +_VERTEX_RESOURCE_PROJECT_AND_LOCATION: Final = re.compile(r"projects/[^/]+/locations/[^/]+/") def _build_vertex_live_setup_model_rewriter( vertex_project: str | None, vertex_location: str | None, llm_router: Router | None, + keeps_caller_resource_path: bool, ) -> Callable[[str], str] | None: """ Rewrite the ``setup`` frame's model into the full Vertex resource path the Live API requires. Clients address the gateway the way they address LiteLLM (bare id or model alias); Vertex reads anything - that is not a ``projects/...`` path as a project name and closes the socket + that is not a ``projects/...`` path as a project name and closes the socket. A ``projects/...`` path + picks the project the proxy's token is spent on, so it is moved onto the connection's project and region + unless ``allow_client_side_credentials`` lets callers choose them """ if vertex_project is None or vertex_location is None: return None def rewrite(setup_model: str) -> str: if setup_model.startswith("projects/"): - return setup_model + resource_prefix: Final = _VERTEX_RESOURCE_PROJECT_AND_LOCATION.match(setup_model) + if keeps_caller_resource_path or resource_prefix is None: + return setup_model + return f"projects/{vertex_project}/locations/{vertex_location}/{setup_model[resource_prefix.end() :]}" aliased: Final = _resolve_alias_to_upstream_model(setup_model, llm_router) return f"projects/{vertex_project}/locations/{vertex_location}/{_vertex_publisher_model_suffix(aliased)}" @@ -4002,18 +4045,15 @@ async def vertex_ai_live_websocket_passthrough( vertex_location=vertex_location, model=model, ) - overridden_field: Final = _caller_override_of_live_deployment( + allow_client_side_credentials: Final = _proxy_general_settings().get("allow_client_side_credentials") is True + override_reason: Final = _live_target_override_reason( live_target, vertex_project=vertex_project, vertex_location=vertex_location, - allow_client_side_credentials=_proxy_general_settings().get("allow_client_side_credentials") is True, + allow_client_side_credentials=allow_client_side_credentials, ) - if overridden_field is not None: - await websocket.close( - code=1008, - reason=f"Rejected Request: {overridden_field} is pinned by the deployment, " - "see configurable_clientside_auth_params", - ) + if override_reason is not None: + await websocket.close(code=1008, reason=override_reason) return vertex_credentials_config: Final = live_target.credentials @@ -4091,6 +4131,7 @@ async def vertex_ai_live_websocket_passthrough( vertex_project=resolved_project, vertex_location=resolved_location, llm_router=_get_llm_router(), + keeps_caller_resource_path=allow_client_side_credentials, ), ) diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index 7926d4d77a2..99d2784b82d 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -8,6 +8,7 @@ Provides: import base64 import json +import re from collections.abc import Mapping from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final @@ -24,6 +25,7 @@ from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( LiteLLM_ManagedVectorStore, ) +from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.proxy._types import * from litellm.proxy.auth.auth_utils import get_request_route, is_request_body_safe from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth @@ -58,8 +60,10 @@ from litellm.proxy.vector_store_endpoints.endpoints import ( from litellm.proxy.vector_store_endpoints.utils import ( assert_user_can_access_vector_store_id, ) +from litellm.rag.ingestion.vertex_ai_ingestion import DEFAULT_VERTEX_RAG_LOCATION from litellm.rag.main import get_ingestion_class from litellm.repositories.table_repositories import ManagedVectorStoresRepository +from litellm.secret_managers.main import get_secret_str from litellm.types.utils import ModelResponse if TYPE_CHECKING: @@ -178,6 +182,41 @@ def _ingest_provider_error(vector_store_config: Mapping[str, object]) -> str | N return None +_VERTEX_RAG_CORPUS_NAME: Final = re.compile(r"projects/(?P[^/]+)/locations/(?P[^/]+)/ragCorpora/") + + +def _vertex_target_error( + vector_store_config: Mapping[str, object], + managed_store: LiteLLM_ManagedVectorStore | None, + user_api_key_dict: UserAPIKeyAuth, + general_settings: Mapping[str, object], +) -> str | None: + if ( + managed_store is not None + or vector_store_config.get("custom_llm_provider") != "vertex_ai" + or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN + or general_settings.get("allow_client_side_credentials") is True + ): + return None + configured_project: Final = litellm.vertex_project or get_secret_str("VERTEXAI_PROJECT") + configured_location: Final = VertexBase.safe_get_vertex_ai_location({}) or DEFAULT_VERTEX_RAG_LOCATION + corpus_name: Final = vector_store_config.get("vector_store_id") + corpus: Final = _VERTEX_RAG_CORPUS_NAME.match(corpus_name) if isinstance(corpus_name, str) else None + requested: Final = ( + ("vertex_project", vector_store_config.get("vertex_project"), configured_project), + ("vertex_location", vector_store_config.get("vertex_location"), configured_location), + ("vector_store_id", corpus["project"] if corpus is not None else None, configured_project), + ("vector_store_id", corpus["location"] if corpus is not None else None, configured_location), + ) + field: Final = next((name for name, asked, configured in requested if asked and asked != configured), None) + if field is None: + return None + return ( + f"Rejected Request: {field} picks a Vertex project or region other than the proxy's. " + "Use a vector store a proxy admin created, or set general_settings.allow_client_side_credentials" + ) + + _MANAGED_STORE_CALLER_OPTIONS: Final = frozenset( { "vector_store_id", @@ -659,6 +698,7 @@ async def rag_ingest( from litellm.proxy.proxy_server import ( add_litellm_data_to_request, general_settings, + general_settings_view, llm_router, prisma_client, proxy_config, @@ -689,6 +729,11 @@ async def rag_ingest( raise HTTPException(status_code=400, detail={"error": str(e)}) managed_store: Final = resolved_stores.get(request_vector_store_config.get("vector_store_id")) + vertex_target_error: Final = _vertex_target_error( + request_vector_store_config, managed_store, user_api_key_dict, general_settings_view() + ) + if vertex_target_error is not None: + raise HTTPException(status_code=400, detail={"error": vertex_target_error}) base: Final = { **_caller_vector_store_options(request_vector_store_config, managed_store), **_managed_store_overrides(managed_store), diff --git a/litellm/rag/ingestion/vertex_ai_ingestion.py b/litellm/rag/ingestion/vertex_ai_ingestion.py index cceddf31117..8157c0c6405 100644 --- a/litellm/rag/ingestion/vertex_ai_ingestion.py +++ b/litellm/rag/ingestion/vertex_ai_ingestion.py @@ -28,6 +28,8 @@ if TYPE_CHECKING: from litellm import Router from litellm.types.rag import RAGIngestOptions +DEFAULT_VERTEX_RAG_LOCATION: Final = "us-central1" + def _present_fields(fields: tuple[tuple[str, object], ...]) -> Mapping[str, object]: return {name: value for name, value in fields if value} @@ -90,7 +92,7 @@ class VertexAIRAGIngestion(BaseRAGIngestion, VertexBase): # Get project, location, and credentials using VertexBase methods self.project_id = self.safe_get_vertex_ai_project(litellm_params) - self.location = self.get_vertex_ai_location(litellm_params) or "us-central1" + self.location = self.get_vertex_ai_location(litellm_params) or DEFAULT_VERTEX_RAG_LOCATION self.vertex_credentials = self.safe_get_vertex_ai_credentials(litellm_params) async def embed( diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 15d581108a0..2cf741620f3 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -6217,7 +6217,11 @@ class TestVertexAILiveWebsocketPassthrough: ), ( "projects/other/locations/us-central1/publishers/google/models/gemini-2.0-flash", - "projects/other/locations/us-central1/publishers/google/models/gemini-2.0-flash", + "projects/proj-db/locations/global/publishers/google/models/gemini-2.0-flash", + ), + ( + "projects/other/locations/us-central1/endpoints/1234", + "projects/proj-db/locations/global/endpoints/1234", ), ], ) @@ -6244,6 +6248,7 @@ class TestVertexAILiveWebsocketPassthrough: vertex_project="proj-db", vertex_location="global", llm_router=llm_router, + keeps_caller_resource_path=False, ) assert rewriter is not None @@ -6444,6 +6449,172 @@ class TestVertexAILiveWebsocketPassthrough: f"projects/proj-db/locations/{expected_location}/publishers/google/models/gemini-live-2.5-flash" ) + def _use_default_vertex_config(self, monkeypatch, general_settings: Mapping[str, object], vertex_location): + from litellm.proxy.pass_through_endpoints import ( + llm_passthrough_endpoints as passthrough_module, + ) + from litellm.types.passthrough_endpoints.vertex_ai import ( + VertexPassThroughCredentials, + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", dict(general_settings)) + monkeypatch.setattr( + passthrough_module.passthrough_endpoint_router, + "default_vertex_config", + VertexPassThroughCredentials( + vertex_project="proj-env", + vertex_location=vertex_location, + vertex_credentials='{"type": "env_account"}', + ), + ) + monkeypatch.setattr(passthrough_module.passthrough_endpoint_router, "deployment_key_to_vertex_credentials", {}) + self._clear_vertex_env(monkeypatch) + + async def _open_live(self, query: Mapping[str, str]): + from litellm.proxy.pass_through_endpoints import ( + llm_passthrough_endpoints as passthrough_module, + ) + + websocket = self._websocket() + ensure_token = AsyncMock( + side_effect=lambda credentials, project_id, custom_llm_provider: ("token-abc", project_id or "proj-env") + ) + ws_passthrough = AsyncMock() + with ( + patch.object(passthrough_module.vertex_llm_base, "ensure_access_token_async", ensure_token), + patch.object(passthrough_module, "websocket_passthrough_request", ws_passthrough), + ): + await passthrough_module.vertex_ai_live_websocket_passthrough( + websocket=websocket, + user_api_key_dict=UserAPIKeyAuth(), + **query, + ) + return websocket, ensure_token, ws_passthrough + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "query, field", + [ + ({"vertex_project": "caller-project"}, "vertex_project"), + ({"vertex_location": "us-central1"}, "vertex_location"), + ({"vertex_project": "proj-env", "vertex_location": "us-central1"}, "vertex_location"), + ], + ) + async def test_query_moving_default_vertex_config_closes_before_minting_a_token( + self, monkeypatch, query: Mapping[str, str], field: str + ): + self._use_default_vertex_config(monkeypatch, {}, vertex_location="europe-west4") + + websocket, ensure_token, ws_passthrough = await self._open_live(query) + + close_kwargs = websocket.close.await_args.kwargs + assert close_kwargs["code"] == 1008 + assert f"{field} is pinned by default_vertex_config" in close_kwargs["reason"] + assert len(close_kwargs["reason"].encode("utf-8")) <= 123 + ensure_token.assert_not_awaited() + ws_passthrough.assert_not_awaited() + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "general_settings, default_location, query, expected_project, expected_location", + [ + ( + {}, + "europe-west4", + {"vertex_project": "proj-env", "vertex_location": "europe-west4"}, + "proj-env", + "europe-west4", + ), + ({}, None, {"vertex_location": "us-central1"}, "proj-env", "us-central1"), + ( + {"allow_client_side_credentials": True}, + "europe-west4", + {"vertex_project": "caller-project", "vertex_location": "us-central1"}, + "caller-project", + "us-central1", + ), + ], + ids=["matches_the_pin", "location_left_unset", "proxy_wide_opt_in"], + ) + async def test_default_vertex_config_serves_the_queries_it_allows( + self, + monkeypatch, + general_settings: Mapping[str, object], + default_location, + query: Mapping[str, str], + expected_project: str, + expected_location: str, + ): + self._use_default_vertex_config(monkeypatch, general_settings, vertex_location=default_location) + + websocket, ensure_token, ws_passthrough = await self._open_live(query) + + websocket.close.assert_not_awaited() + assert ensure_token.await_args.kwargs["project_id"] == expected_project + rewriter = ws_passthrough.await_args.kwargs["setup_model_rewriter"] + assert rewriter("gemini-live-2.5-flash") == ( + f"projects/{expected_project}/locations/{expected_location}/publishers/google/models/gemini-live-2.5-flash" + ) + + @pytest.mark.asyncio + async def test_a_registered_project_and_location_pair_still_serves_its_query(self, monkeypatch): + from litellm.proxy.pass_through_endpoints import ( + llm_passthrough_endpoints as passthrough_module, + ) + from litellm.types.passthrough_endpoints.vertex_ai import ( + VertexPassThroughCredentials, + ) + + self._use_default_vertex_config(monkeypatch, {}, vertex_location="europe-west4") + monkeypatch.setattr( + passthrough_module.passthrough_endpoint_router, + "deployment_key_to_vertex_credentials", + { + "caller-project-us-central1": VertexPassThroughCredentials( + vertex_project="caller-project", + vertex_location="us-central1", + vertex_credentials='{"type": "registered_account"}', + ) + }, + ) + + websocket, ensure_token, _ = await self._open_live( + {"vertex_project": "caller-project", "vertex_location": "us-central1"} + ) + + websocket.close.assert_not_awaited() + ensure_token.assert_awaited_once_with( + credentials='{"type": "registered_account"}', + project_id="caller-project", + custom_llm_provider="vertex_ai_beta", + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "general_settings, expected", + [ + ({}, "projects/proj-env/locations/europe-west4/publishers/google/models/gemini-2.0-flash"), + ( + {"allow_client_side_credentials": True}, + "projects/caller-project/locations/us-central1/publishers/google/models/gemini-2.0-flash", + ), + ], + ids=["pinned", "proxy_wide_opt_in"], + ) + async def test_setup_frame_resource_path_cannot_move_the_session_to_another_project( + self, monkeypatch, general_settings: Mapping[str, object], expected: str + ): + self._use_default_vertex_config(monkeypatch, general_settings, vertex_location="europe-west4") + + _, _, ws_passthrough = await self._open_live({}) + + rewriter = ws_passthrough.await_args.kwargs["setup_model_rewriter"] + assert ( + rewriter("projects/caller-project/locations/us-central1/publishers/google/models/gemini-2.0-flash") + == expected + ) + class TestAnthropicProxyRoute: """The /anthropic passthrough route: custom auth headers must not clobber the diff --git a/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py b/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py index d9f3ceda1ed..5f88aec692a 100644 --- a/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py @@ -990,6 +990,123 @@ def test_rag_ingest_unmanaged_store_keeps_the_callers_full_config(client_interna assert mock_aingest.await_args.kwargs["ingest_options"]["vector_store"] == caller_config +_PROXY_VERTEX_PROJECT: Final = "proxy-project" +_PROXY_VERTEX_LOCATION: Final = "europe-west4" + + +@pytest.fixture +def proxy_vertex_target(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "vertex_project", None) + monkeypatch.setattr(litellm, "vertex_location", None) + monkeypatch.setenv("VERTEXAI_PROJECT", _PROXY_VERTEX_PROJECT) + monkeypatch.setenv("VERTEXAI_LOCATION", _PROXY_VERTEX_LOCATION) + + +def _vertex_ingest(client: TestClient, vector_store: Mapping[str, str]) -> tuple[httpx.Response, AsyncMock]: + aingest_patch, registry_patch = _patched_ingest_boundary( + None, {"vector_store_id": "projects/p/locations/l/ragCorpora/7", "file_id": "file_123"} + ) + with aingest_patch as mock_aingest, registry_patch, _patched_prisma_client(None): + response = client.post("/v1/rag/ingest", **_ingest_form({"custom_llm_provider": "vertex_ai", **vector_store})) + return response, mock_aingest + + +@pytest.mark.parametrize( + "vector_store, field", + [ + ({"vertex_project": "other-project"}, "vertex_project"), + ({"vertex_location": "us-east5"}, "vertex_location"), + ( + {"vector_store_id": f"projects/other-project/locations/{_PROXY_VERTEX_LOCATION}/ragCorpora/7"}, + "vector_store_id", + ), + ( + { + "vertex_project": _PROXY_VERTEX_PROJECT, + "vector_store_id": f"projects/other-project/locations/{_PROXY_VERTEX_LOCATION}/ragCorpora/7", + }, + "vector_store_id", + ), + ({"vector_store_id": f"projects/{_PROXY_VERTEX_PROJECT}/locations/us-east5/ragCorpora/7"}, "vector_store_id"), + ], + ids=["project_field", "location_field", "corpus_project", "corpus_project_behind_matching_field", "corpus_region"], +) +@pytest.mark.usefixtures("proxy_vertex_target") +def test_rag_ingest_refuses_a_key_pointing_the_proxys_vertex_credentials_elsewhere( + client_internal_user, vector_store: Mapping[str, str], field: str +): + response, mock_aingest = _vertex_ingest(client_internal_user, vector_store) + + assert response.status_code == 400, response.json() + assert response.json()["detail"]["error"].startswith(f"Rejected Request: {field} picks a Vertex project") + mock_aingest.assert_not_awaited() + + +def test_rag_ingest_refuses_a_key_naming_a_vertex_project_when_the_proxy_names_none( + client_internal_user, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "vertex_project", None) + monkeypatch.delenv("VERTEXAI_PROJECT", raising=False) + + response, mock_aingest = _vertex_ingest(client_internal_user, {"vertex_project": "other-project"}) + + assert response.status_code == 400, response.json() + mock_aingest.assert_not_awaited() + + +@pytest.mark.parametrize( + "vector_store", + [ + {}, + {"vertex_project": _PROXY_VERTEX_PROJECT, "vertex_location": _PROXY_VERTEX_LOCATION}, + {"vector_store_id": f"projects/{_PROXY_VERTEX_PROJECT}/locations/{_PROXY_VERTEX_LOCATION}/ragCorpora/7"}, + ], + ids=["proxy_defaults", "same_fields", "corpus_in_the_proxy_project"], +) +@pytest.mark.usefixtures("proxy_vertex_target") +def test_rag_ingest_lets_a_key_use_the_proxys_own_vertex_target(client_internal_user, vector_store: Mapping[str, str]): + response, mock_aingest = _vertex_ingest(client_internal_user, vector_store) + + assert response.status_code == 200, response.json() + assert mock_aingest.await_args.kwargs["ingest_options"]["vector_store"] == { + "custom_llm_provider": "vertex_ai", + **vector_store, + } + + +@pytest.mark.parametrize( + "user_role, general_settings", + [ + (LitellmUserRoles.PROXY_ADMIN, {}), + (LitellmUserRoles.INTERNAL_USER, {"allow_client_side_credentials": True}), + ], + ids=["proxy_admin", "proxy_wide_opt_in"], +) +@pytest.mark.usefixtures("proxy_vertex_target") +def test_rag_ingest_lets_the_admin_choose_another_vertex_target( + user_role: LitellmUserRoles, general_settings: Mapping[str, object] +): + caller: Final = UserAPIKeyAuth(user_id="vertex-caller", user_role=user_role) + original_overrides: Final = app.dependency_overrides.copy() + app.dependency_overrides[user_api_key_auth] = lambda: caller + vector_store: Final = { + "vertex_project": "other-project", + "vertex_location": "us-east5", + "vector_store_id": "projects/other-project/locations/us-east5/ragCorpora/7", + } + try: + with patch.dict("litellm.proxy.proxy_server.general_settings", general_settings): + response, mock_aingest = _vertex_ingest(TestClient(app), vector_store) + finally: + app.dependency_overrides = original_overrides + + assert response.status_code == 200, response.json() + assert mock_aingest.await_args.kwargs["ingest_options"]["vector_store"] == { + "custom_llm_provider": "vertex_ai", + **vector_store, + } + + def test_rag_ingest_db_managed_store_drops_the_callers_credential_name(client_internal_user): aingest_patch, registry_patch = _patched_ingest_boundary( DB_MANAGED_STORE, {"vector_store_id": "db-store", "file_id": "file_123"} From a486146c9f2ece9aa3937e5e00b3fa49461c4514 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 10:10:45 -0700 Subject: [PATCH 05/11] fix(proxy): require the Vertex opt-in on alias targets and fallbacks, and pass the live rewriter flag in main's tests --- litellm/proxy/auth/auth_utils.py | 47 +++++++++-- tests/unit/proxy/auth/test_auth_utils.py | 82 +++++++++++++++++++ .../test_pass_through_endpoints.py | 30 +++++-- 3 files changed, 144 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index a979eba7b48..83a518ae6e6 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -226,9 +226,10 @@ def _allow_model_level_clientside_configurable_parameters( """ if llm_router is None: return False - provider_group: Final = model.split("/", 1)[0] + routed_group: Final = llm_router.get_model_from_alias(model) or model + provider_group: Final = routed_group.split("/", 1)[0] deployments: Final = ( - llm_router.get_model_list(model_name=model) + llm_router.get_model_list_of_routed_group(routed_group) or (llm_router.get_model_list(model_name=provider_group) if provider_group in provider_list else None) or () ) @@ -506,6 +507,16 @@ _BANNED_METADATA_PARAMS: Final[tuple[str, ...]] = tuple( ) +def _opt_in_model_groups( + model: str, request_fallback_models: tuple[str, ...], llm_router: Router | None +) -> frozenset[str]: + requested: Final = frozenset((model, *request_fallback_models)) + if llm_router is None: + return requested + routed: Final = frozenset(target for name in requested if (target := llm_router.get_model_from_alias(name))) + return (requested | routed).union(*(router_fallback_model_names(name, llm_router) for name in requested | routed)) + + def _check_banned_params( body: Mapping[str, object], general_settings: dict, @@ -515,6 +526,7 @@ def _check_banned_params( manages_deployments: bool = False, mcp_server_fields_allowed: bool = False, banned_params: tuple[str, ...] = _BANNED_REQUEST_BODY_PARAMS, + request_fallback_models: tuple[str, ...] = (), ) -> None: """Raise ``ValueError`` if ``body`` carries a banned param without admin opt-in. @@ -536,14 +548,15 @@ def _check_banned_params( # Proxy-wide opt-in: every banned param is permitted, exit # entirely so the rest of the loop doesn't waste work. return - if ( + if all( _allow_model_level_clientside_configurable_parameters( - model=model, + model=model_group, param=param, request_body_value=body[param], llm_router=llm_router, ) is True + for model_group in _opt_in_model_groups(model, request_fallback_models, llm_router) ): # Per-param opt-in: only THIS param is permitted by the # deployment's ``configurable_clientside_auth_params``. Skip @@ -719,18 +732,28 @@ def is_request_body_safe( if "model_list" in request_body: raise ValueError("Rejected Request: model_list is not allowed in the request body.") manages_deployments: Final = route_manages_deployments(route) + root_body: Final = _without_route_named_vertex_target(request_body, route) + request_fallback_models: Final = request_fallback_model_names(root_body) _check_banned_params( - _without_route_named_vertex_target(request_body, route), + root_body, general_settings, llm_router, model, manages_deployments=manages_deployments, mcp_server_fields_allowed=route_manages_mcp_servers(route), + request_fallback_models=request_fallback_models, ) for nested_key in _NESTED_CONFIG_KEYS: nested = _coerce_metadata_to_dict(request_body.get(nested_key)) if nested is not None: - _check_banned_params(nested, general_settings, llm_router, model, manages_deployments=manages_deployments) + _check_banned_params( + nested, + general_settings, + llm_router, + model, + manages_deployments=manages_deployments, + request_fallback_models=request_fallback_models, + ) for metadata_key in _NESTED_METADATA_KEYS: metadata = _coerce_metadata_to_dict(request_body.get(metadata_key)) if metadata is not None: @@ -741,6 +764,7 @@ def is_request_body_safe( model, manages_deployments=manages_deployments, banned_params=_BANNED_METADATA_PARAMS, + request_fallback_models=request_fallback_models, ) if any(isinstance(key, str) and key.startswith(f"{metadata_key}[") for key in request_body): _check_banned_params( @@ -750,10 +774,18 @@ def is_request_body_safe( model, manages_deployments=manages_deployments, banned_params=_BANNED_METADATA_PARAMS, + request_fallback_models=request_fallback_models, ) for target in iter_request_fallback_targets(request_body): if isinstance(target, dict): - _check_banned_params(target, general_settings, llm_router, model, manages_deployments=manages_deployments) + _check_banned_params( + target, + general_settings, + llm_router, + model, + manages_deployments=manages_deployments, + request_fallback_models=request_fallback_models, + ) target_model = target.get("model") if isinstance(target_model, str): _reject_url_valued_fallback_target(target_model) @@ -773,6 +805,7 @@ def is_request_body_safe( model, manages_deployments=manages_deployments, banned_params=_BANNED_METADATA_PARAMS, + request_fallback_models=request_fallback_models, ) for user_config_params in _user_config_litellm_params(request_body): # pyright: ignore[reportUnknownArgumentType] # request_body is an untyped dict reject_api_key_secret_reference(user_config_params) diff --git a/tests/unit/proxy/auth/test_auth_utils.py b/tests/unit/proxy/auth/test_auth_utils.py index 0e969717782..e69a0b2972c 100644 --- a/tests/unit/proxy/auth/test_auth_utils.py +++ b/tests/unit/proxy/auth/test_auth_utils.py @@ -5107,6 +5107,88 @@ class TestIsRequestBodySafeBlocksVertexTargetOverrides: is True ) + @pytest.mark.parametrize( + "model, is_allowed", [("public-alias", False), ("public-gemini-2.5-flash", True)], ids=["alias", "direct"] + ) + def test_alias_opt_in_comes_from_the_group_the_alias_routes_to(self, model: str, is_allowed: bool) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "gemini-eu/*", + "litellm_params": {"model": "vertex_ai/*", "vertex_location": "europe-west4"}, + }, + { + "model_name": "public-*", + "litellm_params": { + "model": "vertex_ai/*", + "vertex_location": "europe-west4", + "configurable_clientside_auth_params": ["vertex_location"], + }, + }, + ], + model_group_alias={"public-alias": "gemini-eu/gemini-2.5-flash"}, + ) + request_body: Final = {"model": model, "vertex_location": "us-central1"} + if is_allowed: + assert is_request_body_safe(request_body, {}, router, model, route="/v1/chat/completions") is True + return + with pytest.raises(ValueError, match="vertex_location"): + is_request_body_safe(request_body, {}, router, model, route="/v1/chat/completions") + + @staticmethod + def _opted_in_and_pinned_router(router_fallbacks: list[dict[str, list[str]]]) -> Router: + return Router( + model_list=[ + { + "model_name": name, + "litellm_params": { + "model": "vertex_ai/gemini-2.5-flash", + "vertex_location": "europe-west4", + "configurable_clientside_auth_params": opt_in, + }, + } + for name, opt_in in ( + ("gemini-optin", ["vertex_location"]), + ("gemini-optin-other", ["vertex_location"]), + ("gemini-pinned", []), + ) + ], + fallbacks=router_fallbacks, + ) + + @pytest.mark.parametrize( + "fields, is_allowed", + [ + ({"vertex_location": "us-central1", "fallbacks": ["gemini-pinned"]}, False), + ({"fallbacks": [{"model": "gemini-pinned", "vertex_location": "us-central1"}]}, False), + ({"vertex_location": "us-central1", "fallbacks": ["gemini-optin-other"]}, True), + ({"fallbacks": [{"model": "gemini-optin-other", "vertex_location": "us-central1"}]}, True), + ], + ids=["pinned", "pinned_entry_with_field", "opted_in", "opted_in_entry_with_field"], + ) + def test_request_fallback_needs_the_same_opt_in(self, fields: dict[str, object], is_allowed: bool) -> None: + router: Final = self._opted_in_and_pinned_router([]) + request_body: Final = {"model": "gemini-optin", **fields} + if is_allowed: + assert is_request_body_safe(request_body, {}, router, "gemini-optin", route="/v1/chat/completions") is True + return + with pytest.raises(ValueError, match="vertex_location"): + is_request_body_safe(request_body, {}, router, "gemini-optin", route="/v1/chat/completions") + + @pytest.mark.parametrize( + "fallback_group, is_allowed", + [("gemini-pinned", False), ("gemini-optin-other", True)], + ids=["pinned", "opted_in"], + ) + def test_router_fallback_needs_the_same_opt_in(self, fallback_group: str, is_allowed: bool) -> None: + router: Final = self._opted_in_and_pinned_router([{"gemini-optin": [fallback_group]}]) + request_body: Final = {"model": "gemini-optin", "vertex_location": "us-central1"} + if is_allowed: + assert is_request_body_safe(request_body, {}, router, "gemini-optin", route="/v1/chat/completions") is True + return + with pytest.raises(ValueError, match="vertex_location"): + is_request_body_safe(request_body, {}, router, "gemini-optin", route="/v1/chat/completions") + @pytest.mark.asyncio @pytest.mark.parametrize("opted_in", [True, False]) async def test_opt_in_of_the_model_named_in_a_google_route_decides(self, monkeypatch, opted_in: bool) -> None: diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index b0504af1b5e..e1cb79ff5b2 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -7368,7 +7368,7 @@ async def test_websocket_passthrough_keeps_normal_upstream_close_normal(): websocket.close.assert_awaited_once_with() -async def _run_setup_rewrite_passthrough(setup_model: str, llm_router) -> str: +async def _run_setup_rewrite_passthrough(setup_model: str, llm_router, keeps_caller_resource_path: bool = False) -> str: from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( _build_vertex_live_setup_model_rewriter, ) @@ -7397,6 +7397,7 @@ async def _run_setup_rewrite_passthrough(setup_model: str, llm_router) -> str: vertex_project="proj-db", vertex_location="global", llm_router=llm_router, + keeps_caller_resource_path=keeps_caller_resource_path, ), ) @@ -7455,7 +7456,7 @@ def test_vertex_live_setup_model_resolves_before_extraction(setup_model): ) rewriter = _build_vertex_live_setup_model_rewriter( - vertex_project="proj-db", vertex_location="global", llm_router=None + vertex_project="proj-db", vertex_location="global", llm_router=None, keeps_caller_resource_path=False ) setup_data = {"model": setup_model} @@ -7505,7 +7506,10 @@ async def test_websocket_passthrough_logs_a_bare_alias_setup_model(): endpoint="/vertex_ai/live", accept_websocket=False, setup_model_rewriter=_build_vertex_live_setup_model_rewriter( - vertex_project="proj-db", vertex_location="global", llm_router=None + vertex_project="proj-db", + vertex_location="global", + llm_router=None, + keeps_caller_resource_path=False, ), ) @@ -7570,14 +7574,24 @@ async def test_websocket_passthrough_rewrites_alias_of_unrecognised_upstream_mod @pytest.mark.asyncio -async def test_websocket_passthrough_leaves_full_resource_setup_model_untouched(): - full_resource = "projects/other/locations/us-central1/publishers/google/models/gemini-2.0-flash-live-preview-04-09" +@pytest.mark.parametrize( + "keeps_caller_resource_path, expected_prefix", + [(True, "projects/other/locations/us-central1/"), (False, "projects/proj-db/locations/global/")], + ids=["caller_path_kept", "moved_onto_session_target"], +) +async def test_websocket_passthrough_full_resource_setup_model( + keeps_caller_resource_path: bool, expected_prefix: str +) -> None: + model_suffix = "publishers/google/models/gemini-2.0-flash-live-preview-04-09" - sent_frame = await _run_setup_rewrite_passthrough(full_resource, llm_router=None) + sent_frame = await _run_setup_rewrite_passthrough( + f"projects/other/locations/us-central1/{model_suffix}", + llm_router=None, + keeps_caller_resource_path=keeps_caller_resource_path, + ) - assert json.loads(sent_frame)["setup"]["model"] == full_resource assert sent_frame == json.dumps( - {"setup": {"model": full_resource, "generationConfig": {"responseModalities": ["TEXT"]}}} + {"setup": {"model": f"{expected_prefix}{model_suffix}", "generationConfig": {"responseModalities": ["TEXT"]}}} ) From e49aa9f0c416727d2a5a542a8e86f70c68cf9584 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 11:11:43 -0700 Subject: [PATCH 06/11] test(proxy): give the batch tests' fake router the alias and routed-group lookups --- .../batches_endpoints/test_litellm_executed_batches.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/unit/proxy/batches_endpoints/test_litellm_executed_batches.py b/tests/unit/proxy/batches_endpoints/test_litellm_executed_batches.py index dd590ac1bd3..334d636ec91 100644 --- a/tests/unit/proxy/batches_endpoints/test_litellm_executed_batches.py +++ b/tests/unit/proxy/batches_endpoints/test_litellm_executed_batches.py @@ -255,6 +255,12 @@ class FakeRouter: def get_model_list(self, model_name: str | None = None) -> None: return None + def get_model_from_alias(self, model: str) -> None: + return None + + def get_model_list_of_routed_group(self, model_group: str) -> list[object]: + return [] + class FakeStorageBackend: def __init__(self, contents: Mapping[str, bytes]) -> None: From 36665d16fdc49f5b6c400e46a203cccb990204ea Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 08:31:42 -0700 Subject: [PATCH 07/11] fix(proxy): close the Vertex opt-in, RAG corpus, and live setup-frame gaps found in review A Vertex opt-in now counts only when every group the request can reach opts in, router, context window, and content policy fallbacks and one alias hop included, while a group with no Google deployment never needs the fields. A model named in the route path is checked next to the body model. RAG ingest from a non-admin key refuses a bare or dot-segment corpus id, a vertex_ai_project or vertex_ai_location naming another target, and a store a non-admin saved before the upgrade, and reads a corpus named by project number on the proxy's own project. /vertex_ai/live closes 1008 on any caller-named project when the proxy pins none, moves a setup frame sent as a binary message onto the session's target, and resolves the region from default_vertex_config before the model path. --- litellm/proxy/auth/auth_utils.py | 125 ++++++++++----- .../llm_passthrough_endpoints.py | 19 ++- .../pass_through_endpoints.py | 22 ++- litellm/proxy/rag_endpoints/endpoints.py | 123 +++++++++++---- .../test_vertex_image_regional_uplift.py | 10 +- tests/unit/proxy/auth/test_auth_utils.py | 84 +++++++++- .../test_llm_pass_through_endpoints.py | 98 ++++++++++++ .../test_pass_through_endpoints.py | 69 ++++++++- .../proxy/rag_endpoints/test_rag_endpoints.py | 145 +++++++++++++++++- 9 files changed, 600 insertions(+), 95 deletions(-) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index bc9d03a5f03..53eb13ed502 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -2,7 +2,7 @@ import importlib.util import os import re import sys -from collections.abc import Collection, Iterator, Mapping +from collections.abc import Collection, Iterator, Mapping, Sequence from functools import lru_cache from itertools import accumulate, chain, repeat from logging import Logger @@ -23,6 +23,7 @@ from litellm.constants import ( STANDARD_CUSTOMER_ID_HEADERS, ) from litellm.litellm_core_utils.credential_accessor import CredentialAccessor, reject_api_key_secret_reference +from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.url_utils import ( SSRFError, @@ -40,12 +41,13 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( from litellm.types.router import ( CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS, Deployment, + DeploymentTypedDict, LiteLLMParamsTypedDict, reject_server_owned_wif_fields, server_owned_wif_fields_present, ) from litellm.types.router import reject_server_owned_wif_params as _reject_server_owned_wif_params -from litellm.types.utils import CustomPricingLiteLLMParams, oauth_token_exchange_litellm_params +from litellm.types.utils import CustomPricingLiteLLMParams, LlmProviders, oauth_token_exchange_litellm_params from litellm.types.videos.main import DecodedVideoId @@ -217,23 +219,50 @@ def deployment_clientside_auth_params(litellm_params: LiteLLMParamsTypedDict) -> return litellm_params["configurable_clientside_auth_params"] +def _model_group_deployments(model: str, llm_router: Router) -> Sequence[DeploymentTypedDict]: + registry: Final = llm_router.registry() + provider_group: Final = model.split("/", 1)[0] + return ( + registry.deployments(model, source="routed_group") + or (registry.deployments(provider_group) if provider_group in provider_list else None) + or () + ) + + +_VERTEX_TARGET_READING_PROVIDERS: Final[frozenset[str]] = frozenset( + (LlmProviders.VERTEX_AI.value, LlmProviders.VERTEX_AI_BETA.value, LlmProviders.GEMINI.value, LlmProviders.GDC.value) +) + + +def _deployment_provider(litellm_params: LiteLLMParamsTypedDict) -> str | None: + match litellm_params: + case {"custom_llm_provider": str(declared)} if declared: + return declared + case {"model": str(model)}: + return inferred_provider(model) + case _: + return None + + +def _model_group_reads_vertex_target(model: str, llm_router: Router | None) -> bool: + if llm_router is None: + return True + deployments: Final = _model_group_deployments(model, llm_router) + providers: Final = frozenset(_deployment_provider(deployment["litellm_params"]) for deployment in deployments) + return not deployments or None in providers or bool(providers & _VERTEX_TARGET_READING_PROVIDERS) + + def _allow_model_level_clientside_configurable_parameters( model: str, param: str, request_body_value: object, llm_router: Router | None ) -> bool: """ The router, not the caller, picks which deployment serves the request, so an - opt-in counts only when every deployment the model can route to carries it + opt-in counts only when every deployment of the model group, already alias + resolved, carries it """ if llm_router is None: return False - registry: Final = llm_router.registry() - routed_group: Final = registry.resolve(model, kind="alias") or model - provider_group: Final = routed_group.split("/", 1)[0] - deployments: Final = ( - registry.deployments(routed_group, source="routed_group") - or (registry.deployments(provider_group) if provider_group in provider_list else None) - or () - ) + deployments: Final = _model_group_deployments(model, llm_router) return bool(deployments) and all( _is_param_allowed( param=param, @@ -322,21 +351,15 @@ def route_manages_mcp_servers(route: str | None) -> bool: VERTEX_TARGET_PARAMS: Final[frozenset[str]] = frozenset( ("vertex_project", "vertex_location", "vertex_ai_project", "vertex_ai_location") ) -_RAG_INGEST_VERTEX_TARGET_PARAMS: Final[frozenset[str]] = frozenset(("vertex_project", "vertex_location")) -_ROUTE_NAMED_VERTEX_TARGET_PARAMS: Final[Mapping[str, frozenset[str]]] = MappingProxyType( - { - "/vertex_ai/live": VERTEX_TARGET_PARAMS, - "/rag/ingest": _RAG_INGEST_VERTEX_TARGET_PARAMS, - "/v1/rag/ingest": _RAG_INGEST_VERTEX_TARGET_PARAMS, - } +_ROUTES_THAT_NAME_THEIR_VERTEX_TARGET: Final[frozenset[str]] = frozenset( + ("/vertex_ai/live", "/rag/ingest", "/v1/rag/ingest") ) def _without_route_named_vertex_target(request_body: Mapping[str, object], route: str | None) -> Mapping[str, object]: - route_named: Final = _ROUTE_NAMED_VERTEX_TARGET_PARAMS.get(route) if route is not None else None - if route_named is None: + if route not in _ROUTES_THAT_NAME_THEIR_VERTEX_TARGET: return request_body - return MappingProxyType({key: value for key, value in request_body.items() if key not in route_named}) + return MappingProxyType({key: value for key, value in request_body.items() if key not in VERTEX_TARGET_PARAMS}) _NESTED_CONFIG_KEYS: Final[tuple[str, ...]] = ("litellm_embedding_config", "extra_body") @@ -509,14 +532,17 @@ _BANNED_METADATA_PARAMS: Final[tuple[str, ...]] = tuple( def _opt_in_model_groups( - model: str, request_fallback_models: tuple[str, ...], llm_router: Router | None + model: str, also_requested_models: tuple[str, ...], llm_router: Router | None ) -> frozenset[str]: - requested: Final = frozenset((model, *request_fallback_models)) + requested: Final = frozenset((model, *also_requested_models)) if llm_router is None: return requested registry: Final = llm_router.registry() - routed: Final = frozenset(target for name in requested if (target := registry.resolve(name, kind="alias"))) - return (requested | routed).union(*(router_fallback_model_names(name, llm_router) for name in requested | routed)) + routed: Final = frozenset(registry.resolve(name, kind="alias") or name for name in requested) + fallbacks: Final = frozenset[str]().union( + *(router_fallback_model_names(name, llm_router, _EVERY_FALLBACK_TYPE) for name in requested | routed) + ) + return routed | frozenset(registry.resolve(name, kind="alias") or name for name in fallbacks) def _check_banned_params( @@ -528,7 +554,7 @@ def _check_banned_params( manages_deployments: bool = False, mcp_server_fields_allowed: bool = False, banned_params: tuple[str, ...] = _BANNED_REQUEST_BODY_PARAMS, - request_fallback_models: tuple[str, ...] = (), + also_requested_models: tuple[str, ...] = (), ) -> None: """Raise ``ValueError`` if ``body`` carries a banned param without admin opt-in. @@ -558,7 +584,8 @@ def _check_banned_params( llm_router=llm_router, ) is True - for model_group in _opt_in_model_groups(model, request_fallback_models, llm_router) + for model_group in _opt_in_model_groups(model, also_requested_models, llm_router) + if param not in VERTEX_TARGET_PARAMS or _model_group_reads_vertex_target(model_group, llm_router) ): # Per-param opt-in: only THIS param is permitted by the # deployment's ``configurable_clientside_auth_params``. Skip @@ -649,27 +676,39 @@ def fallback_target_model_name(target: object) -> str | None: return None -def _direct_fallback_model_names(model: str, llm_router: Router) -> frozenset[str]: +_EVERY_FALLBACK_TYPE: Final[tuple[str, ...]] = ("general", "context_window", "content_policy") + + +def _direct_fallback_model_names(model: str, llm_router: Router, fallback_types: tuple[str, ...]) -> frozenset[str]: from litellm.proxy.auth.model_checks import get_all_fallbacks - return frozenset( - name for target in get_all_fallbacks(model, llm_router) if (name := fallback_target_model_name(target)) + targets: Final = chain.from_iterable( + get_all_fallbacks(model, llm_router, fallback_type) for fallback_type in fallback_types ) + return frozenset(name for target in targets if (name := fallback_target_model_name(target))) def _fallback_hop( - frontier_and_reached: tuple[frozenset[str], frozenset[str]], llm_router: Router + frontier_and_reached: tuple[frozenset[str], frozenset[str]], router_and_types: tuple[Router, tuple[str, ...]] ) -> tuple[frozenset[str], frozenset[str]]: frontier, reached = frontier_and_reached + llm_router, fallback_types = router_and_types next_hop: Final = ( - frozenset(chain.from_iterable(_direct_fallback_model_names(model, llm_router) for model in frontier)) - reached + frozenset( + chain.from_iterable(_direct_fallback_model_names(model, llm_router, fallback_types) for model in frontier) + ) + - reached ) return next_hop, reached | next_hop -def router_fallback_model_names(model: str, llm_router: Router) -> frozenset[str]: +def router_fallback_model_names( + model: str, llm_router: Router, fallback_types: tuple[str, ...] = ("general",) +) -> frozenset[str]: """The groups a failure on `model` can reach through the router's fallbacks, chained and generic (`*`) ones.""" - hops: Final = accumulate(repeat(llm_router), _fallback_hop, initial=(frozenset({model}), frozenset({model}))) + hops: Final = accumulate( + repeat((llm_router, fallback_types)), _fallback_hop, initial=(frozenset({model}), frozenset({model})) + ) _, reached = next(hop for hop in hops if not hop[0]) return reached - {model} @@ -735,7 +774,11 @@ def is_request_body_safe( raise ValueError("Rejected Request: model_list is not allowed in the request body.") manages_deployments: Final = route_manages_deployments(route) root_body: Final = _without_route_named_vertex_target(request_body, route) - request_fallback_models: Final = request_fallback_model_names(root_body) + route_model: Final = _model_from_route_path(route) if route is not None else None + also_requested_models: Final = ( + *request_fallback_model_names(root_body), + *((route_model,) if route_model else ()), + ) _check_banned_params( root_body, general_settings, @@ -743,7 +786,7 @@ def is_request_body_safe( model, manages_deployments=manages_deployments, mcp_server_fields_allowed=route_manages_mcp_servers(route), - request_fallback_models=request_fallback_models, + also_requested_models=also_requested_models, ) for nested_key in _NESTED_CONFIG_KEYS: nested = _coerce_metadata_to_dict(request_body.get(nested_key)) @@ -754,7 +797,7 @@ def is_request_body_safe( llm_router, model, manages_deployments=manages_deployments, - request_fallback_models=request_fallback_models, + also_requested_models=also_requested_models, ) for metadata_key in _NESTED_METADATA_KEYS: metadata = _coerce_metadata_to_dict(request_body.get(metadata_key)) @@ -766,7 +809,7 @@ def is_request_body_safe( model, manages_deployments=manages_deployments, banned_params=_BANNED_METADATA_PARAMS, - request_fallback_models=request_fallback_models, + also_requested_models=also_requested_models, ) if any(isinstance(key, str) and key.startswith(f"{metadata_key}[") for key in request_body): _check_banned_params( @@ -776,7 +819,7 @@ def is_request_body_safe( model, manages_deployments=manages_deployments, banned_params=_BANNED_METADATA_PARAMS, - request_fallback_models=request_fallback_models, + also_requested_models=also_requested_models, ) for target in iter_request_fallback_targets(request_body): if isinstance(target, dict): @@ -786,7 +829,7 @@ def is_request_body_safe( llm_router, model, manages_deployments=manages_deployments, - request_fallback_models=request_fallback_models, + also_requested_models=also_requested_models, ) target_model = target.get("model") if isinstance(target_model, str): @@ -807,7 +850,7 @@ def is_request_body_safe( model, manages_deployments=manages_deployments, banned_params=_BANNED_METADATA_PARAMS, - request_fallback_models=request_fallback_models, + also_requested_models=also_requested_models, ) for user_config_params in _user_config_litellm_params(request_body): # pyright: ignore[reportUnknownArgumentType] # request_body is an untyped dict reject_api_key_secret_reference(user_config_params) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index df374c1a154..302f41491f5 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -3857,6 +3857,10 @@ VERTEX_LIVE_UNCONFIGURED_CLOSE_REASON: Final = ( "Vertex AI auth failed: set a use_in_pass_through vertex model, default_vertex_config, or DEFAULT_VERTEXAI_* env" ) +VERTEX_LIVE_UNPINNED_PROJECT_CLOSE_REASON: Final = ( + "Rejected Request: vertex_project must come from the proxy's configuration, see allow_client_side_credentials" +) + VERTEX_PUBLISHER_MODEL_PREFIX: Final = "publishers/google/models/" VERTEX_PUBLISHERS_SEGMENT: Final = "publishers/" @@ -3947,6 +3951,10 @@ def _live_target_override_reason( if allow_client_side_credentials: return None match target: + case _ConfiguredVertexLiveTarget(credentials=credentials) if vertex_project and ( + credentials is None or credentials.vertex_project is None + ): + return VERTEX_LIVE_UNPINNED_PROJECT_CLOSE_REASON case _ConfiguredVertexLiveTarget(credentials=VertexPassThroughCredentials() as credentials): configured_field: Final = next( ( @@ -3986,12 +3994,13 @@ def _live_target_override_reason( ) -_VERTEX_RESOURCE_PROJECT_AND_LOCATION: Final = re.compile(r"projects/[^/]+/locations/[^/]+/") +_VERTEX_RESOURCE_PROJECT_AND_LOCATION: Final = re.compile(r"projects/[^/]+/locations/(?P[^/]+)/") def _build_vertex_live_setup_model_rewriter( vertex_project: str | None, vertex_location: str | None, + configured_location: str | None, llm_router: Router | None, keeps_caller_resource_path: bool, ) -> Callable[[str], str] | None: @@ -4000,8 +4009,8 @@ def _build_vertex_live_setup_model_rewriter( Clients address the gateway the way they address LiteLLM (bare id or model alias); Vertex reads anything that is not a ``projects/...`` path as a project name and closes the socket. A ``projects/...`` path - picks the project the proxy's token is spent on, so it is moved onto the connection's project and region - unless ``allow_client_side_credentials`` lets callers choose them + picks the project the proxy's token is spent on, so it is moved onto the connection's project, and onto its + region when one is configured, unless ``allow_client_side_credentials`` lets callers choose them """ if vertex_project is None or vertex_location is None: return None @@ -4011,7 +4020,8 @@ def _build_vertex_live_setup_model_rewriter( resource_prefix: Final = _VERTEX_RESOURCE_PROJECT_AND_LOCATION.match(setup_model) if keeps_caller_resource_path or resource_prefix is None: return setup_model - return f"projects/{vertex_project}/locations/{vertex_location}/{setup_model[resource_prefix.end() :]}" + location: Final = configured_location or resource_prefix.group("location") + return f"projects/{vertex_project}/locations/{location}/{setup_model[resource_prefix.end() :]}" aliased: Final = _resolve_alias_to_upstream_model(setup_model, llm_router) return f"projects/{vertex_project}/locations/{vertex_location}/{_vertex_publisher_model_suffix(aliased)}" @@ -4155,6 +4165,7 @@ async def vertex_ai_live_websocket_passthrough( setup_model_rewriter=_build_vertex_live_setup_model_rewriter( vertex_project=resolved_project, vertex_location=resolved_location, + configured_location=configured_location, llm_router=_get_llm_router(), keeps_caller_resource_path=allow_client_side_credentials, ), diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 7b97bbb670e..9534d33bd91 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -2426,6 +2426,21 @@ def _rewrite_vertex_live_setup_model(text_data: str, setup_model_rewriter: Calla return json.dumps({**message, "setup": {**setup, "model": rewritten_model}}) +VERTEX_LIVE_UNDECODABLE_BINARY_CLOSE_REASON: Final = "Rejected Request: binary messages must be UTF-8 encoded JSON" + + +def _rewrite_vertex_live_binary_setup_model( + bytes_data: bytes, setup_model_rewriter: Callable[[str], str] | None +) -> bytes | None: + if setup_model_rewriter is None: + return bytes_data + try: + text_data: Final = bytes_data.decode("utf-8") + except UnicodeDecodeError: + return None + return _rewrite_vertex_live_setup_model(text_data, setup_model_rewriter).encode("utf-8") + + def _resolved_vertex_live_setup( setup_data: Mapping[str, object], setup_model_rewriter: Callable[[str], str] | None ) -> Mapping[str, object]: @@ -2726,7 +2741,12 @@ async def websocket_passthrough_request( await upstream_ws.send(_rewrite_vertex_live_setup_model(text_data, setup_model_rewriter)) elif bytes_data is not None: - await upstream_ws.send(bytes_data) + rewritten_bytes = _rewrite_vertex_live_binary_setup_model(bytes_data, setup_model_rewriter) + if rewritten_bytes is None: + await websocket.close(code=1008, reason=VERTEX_LIVE_UNDECODABLE_BINARY_CLOSE_REASON) + await upstream_ws.close() + break + await upstream_ws.send(rewritten_bytes) except asyncio.CancelledError: raise except Exception: diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index 19c4bf5ff17..35e7ae71fb7 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -21,7 +21,7 @@ from starlette.datastructures import UploadFile import litellm from litellm._logging import verbose_proxy_logger -from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH +from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH, LITELLM_PROXY_ADMIN_NAME from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( LiteLLM_ManagedVectorStore, ) @@ -64,6 +64,7 @@ from litellm.proxy.vector_store_endpoints.utils import ( from litellm.rag.ingestion.vertex_ai_ingestion import DEFAULT_VERTEX_RAG_LOCATION from litellm.rag.main import get_ingestion_class from litellm.repositories.table_repositories import ManagedVectorStoresRepository +from litellm.repositories.user_repository import UserRepository from litellm.secret_managers.main import get_secret_str from litellm.types.utils import ModelResponse @@ -183,41 +184,95 @@ def _ingest_provider_error(vector_store_config: Mapping[str, object]) -> str | N return None -_VERTEX_RAG_CORPUS_NAME: Final = re.compile(r"projects/(?P[^/]+)/locations/(?P[^/]+)/ragCorpora/") +_VERTEX_RAG_CORPUS_NAME: Final = re.compile( + r"projects/(?P[^/%\\]+)/locations/(?P[^/%\\]+)/ragCorpora/(?P[^/%\\]+)" +) +_DOT_SEGMENTS: Final = frozenset({".", ".."}) -def _vertex_target_error( - vector_store_config: Mapping[str, object], - managed_store: LiteLLM_ManagedVectorStore | None, - user_api_key_dict: UserAPIKeyAuth, - general_settings: Mapping[str, object], -) -> str | None: - if ( - managed_store is not None - or vector_store_config.get("custom_llm_provider") != "vertex_ai" - or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN - or general_settings.get("allow_client_side_credentials") is True - ): - return None - configured_project: Final = litellm.vertex_project or get_secret_str("VERTEXAI_PROJECT") - configured_location: Final = VertexBase.safe_get_vertex_ai_location({}) or DEFAULT_VERTEX_RAG_LOCATION - corpus_name: Final = vector_store_config.get("vector_store_id") - corpus: Final = _VERTEX_RAG_CORPUS_NAME.match(corpus_name) if isinstance(corpus_name, str) else None - requested: Final = ( - ("vertex_project", vector_store_config.get("vertex_project"), configured_project), - ("vertex_location", vector_store_config.get("vertex_location"), configured_location), - ("vector_store_id", corpus["project"] if corpus is not None else None, configured_project), - ("vector_store_id", corpus["location"] if corpus is not None else None, configured_location), - ) - field: Final = next((name for name, asked, configured in requested if asked and asked != configured), None) - if field is None: +def _clean_vertex_corpus_name(corpus_name: object) -> re.Match[str] | None: + if not isinstance(corpus_name, str) or not _DOT_SEGMENTS.isdisjoint(corpus_name.split("/")): return None + return _VERTEX_RAG_CORPUS_NAME.fullmatch(corpus_name) + + +def _vertex_target_message(field: str) -> str: return ( f"Rejected Request: {field} picks a Vertex project or region other than the proxy's. " "Use a vector store a proxy admin created, or set general_settings.allow_client_side_credentials" ) +def _vertex_target_is_checked( + vector_store_config: Mapping[str, object], + user_api_key_dict: UserAPIKeyAuth, + general_settings: Mapping[str, object], +) -> bool: + return ( + vector_store_config.get("custom_llm_provider") == "vertex_ai" + and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN + and general_settings.get("allow_client_side_credentials") is not True + ) + + +def _configured_vertex_project() -> str | None: + return litellm.vertex_project or get_secret_str("VERTEXAI_PROJECT") + + +def _corpus_on_configured_project(vector_store_config: Mapping[str, object]) -> Mapping[str, object]: + corpus: Final = _clean_vertex_corpus_name(vector_store_config.get("vector_store_id")) + configured_project: Final = _configured_vertex_project() + if corpus is None or configured_project is None or not corpus["project"].isdigit(): + return vector_store_config + return MappingProxyType( + { + **vector_store_config, + "vector_store_id": ( + f"projects/{configured_project}/locations/{corpus['location']}/ragCorpora/{corpus['corpus']}" + ), + } + ) + + +def _vertex_target_error(vector_store_config: Mapping[str, object]) -> str | None: + corpus_name: Final = vector_store_config.get("vector_store_id") + corpus: Final = _clean_vertex_corpus_name(corpus_name) + if corpus_name and corpus is None: + return _vertex_target_message("vector_store_id") + configured_project: Final = _configured_vertex_project() + configured_location: Final = VertexBase.safe_get_vertex_ai_location({}) or DEFAULT_VERTEX_RAG_LOCATION + requested: Final = ( + *( + (name, vector_store_config.get(name), configured_project) + for name in ("vertex_project", "vertex_ai_project") + ), + *( + (name, vector_store_config.get(name), configured_location) + for name in ("vertex_location", "vertex_ai_location") + ), + ("vector_store_id", corpus["project"] if corpus is not None else None, configured_project), + ("vector_store_id", corpus["location"] if corpus is not None else None, configured_location), + ) + field: Final = next((name for name, asked, configured in requested if asked and asked != configured), None) + return None if field is None else _vertex_target_message(field) + + +async def _admin_set_store_target( + managed_store: LiteLLM_ManagedVectorStore | None, prisma_client: "PrismaClient | None" +) -> bool: + if managed_store is None: + return False + if managed_store.get("is_config", False): + return True + creator: Final = managed_store.get("user_id") + if creator == LITELLM_PROXY_ADMIN_NAME: + return True + if creator is None or prisma_client is None: + return False + creator_row: Final = await UserRepository(prisma_client).find_by_id(creator) + return creator_row is not None and creator_row.user_role == LitellmUserRoles.PROXY_ADMIN + + _MANAGED_STORE_CALLER_OPTIONS: Final = frozenset( { "vector_store_id", @@ -730,15 +785,17 @@ async def rag_ingest( raise HTTPException(status_code=400, detail={"error": str(e)}) managed_store: Final = resolved_stores.get(request_vector_store_config.get("vector_store_id")) - vertex_target_error: Final = _vertex_target_error( - request_vector_store_config, managed_store, user_api_key_dict, general_settings_view() - ) - if vertex_target_error is not None: - raise HTTPException(status_code=400, detail={"error": vertex_target_error}) - base: Final = { + requested_base: Final = { **_caller_vector_store_options(request_vector_store_config, managed_store), **_managed_store_overrides(managed_store), } + vertex_target_checked: Final = _vertex_target_is_checked( + requested_base, user_api_key_dict, general_settings_view() + ) + base: Final = _corpus_on_configured_project(requested_base) if vertex_target_checked else requested_base + vertex_target_error: Final = _vertex_target_error(base) if vertex_target_checked else None + if vertex_target_error is not None and not await _admin_set_store_target(managed_store, prisma_client): + raise HTTPException(status_code=400, detail={"error": vertex_target_error}) request_data: Final[dict[str, Any]] = await add_litellm_data_to_request( data={}, request=request, diff --git a/tests/integration/spend/test_vertex_image_regional_uplift.py b/tests/integration/spend/test_vertex_image_regional_uplift.py index 6900a6d489a..dfd52d28f1f 100644 --- a/tests/integration/spend/test_vertex_image_regional_uplift.py +++ b/tests/integration/spend/test_vertex_image_regional_uplift.py @@ -100,6 +100,7 @@ def _deployment( location: Mapping[str, JsonValue], multiplier: JsonValue = 1.1, backend: str = _IMAGE_BACKEND, + clientside_params: tuple[str, ...] = (), ) -> str: return scenario.model( model=f"vertex_ai/{backend}", @@ -111,6 +112,7 @@ def _deployment( output_cost_per_token=4e-5, output_cost_per_image=0.05, **({} if multiplier is None else {"regional_endpoint_uplift_multiplier": multiplier}), + **({"configurable_clientside_auth_params": list(clientside_params)} if clientside_params else {}), **location, ) @@ -380,7 +382,13 @@ def test_request_body_location_overrides_the_deployment( gateway: Gateway, deployment_location: str, extra: Mapping[str, JsonValue], expected: float ) -> None: with wire_server(_scripted()) as wire, gateway.scenario() as scenario: - model: Final = _deployment(gateway, scenario, wire, location={"vertex_location": deployment_location}) + model: Final = _deployment( + gateway, + scenario, + wire, + location={"vertex_location": deployment_location}, + clientside_params=("vertex_location", "vertex_ai_location"), + ) call_id, response = _generate(gateway, model, extra=extra) _assert_images(response, call_id, expected) _assert_one_upstream_image_call(wire) diff --git a/tests/unit/proxy/auth/test_auth_utils.py b/tests/unit/proxy/auth/test_auth_utils.py index f5d2ebf234f..54bdc8e5959 100644 --- a/tests/unit/proxy/auth/test_auth_utils.py +++ b/tests/unit/proxy/auth/test_auth_utils.py @@ -5136,7 +5136,11 @@ class TestIsRequestBodySafeBlocksVertexTargetOverrides: is_request_body_safe(request_body, {}, router, model, route="/v1/chat/completions") @staticmethod - def _opted_in_and_pinned_router(router_fallbacks: list[dict[str, list[str]]]) -> Router: + def _opted_in_and_pinned_router( + router_fallbacks: list[dict[str, list[str]]], + fallback_setting: str = "fallbacks", + model_group_alias: dict[str, str] | None = None, + ) -> Router: return Router( model_list=[ { @@ -5153,7 +5157,8 @@ class TestIsRequestBodySafeBlocksVertexTargetOverrides: ("gemini-pinned", []), ) ], - fallbacks=router_fallbacks, + model_group_alias=model_group_alias or {}, + **{fallback_setting: router_fallbacks}, ) @pytest.mark.parametrize( @@ -5175,13 +5180,16 @@ class TestIsRequestBodySafeBlocksVertexTargetOverrides: with pytest.raises(ValueError, match="vertex_location"): is_request_body_safe(request_body, {}, router, "gemini-optin", route="/v1/chat/completions") + @pytest.mark.parametrize("fallback_setting", ["fallbacks", "context_window_fallbacks", "content_policy_fallbacks"]) @pytest.mark.parametrize( "fallback_group, is_allowed", [("gemini-pinned", False), ("gemini-optin-other", True)], ids=["pinned", "opted_in"], ) - def test_router_fallback_needs_the_same_opt_in(self, fallback_group: str, is_allowed: bool) -> None: - router: Final = self._opted_in_and_pinned_router([{"gemini-optin": [fallback_group]}]) + def test_router_fallback_needs_the_same_opt_in( + self, fallback_setting: str, fallback_group: str, is_allowed: bool + ) -> None: + router: Final = self._opted_in_and_pinned_router([{"gemini-optin": [fallback_group]}], fallback_setting) request_body: Final = {"model": "gemini-optin", "vertex_location": "us-central1"} if is_allowed: assert is_request_body_safe(request_body, {}, router, "gemini-optin", route="/v1/chat/completions") is True @@ -5189,6 +5197,67 @@ class TestIsRequestBodySafeBlocksVertexTargetOverrides: with pytest.raises(ValueError, match="vertex_location"): is_request_body_safe(request_body, {}, router, "gemini-optin", route="/v1/chat/completions") + @pytest.mark.parametrize( + "model, fields, is_allowed", + [ + ("gemini-optin", {"vertex_location": "us-central1"}, True), + ("openai-group", {"vertex_location": "us-central1"}, True), + ("gemini-optin", {"api_base": "https://eu.example.com"}, False), + ("studio-group", {"vertex_location": "us-central1"}, False), + ], + ids=["vertex_field_to_openai_fallback", "vertex_field_on_openai", "api_base_to_openai_fallback", "gemini"], + ) + def test_only_google_groups_need_the_vertex_opt_in( + self, model: str, fields: dict[str, object], is_allowed: bool + ) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "gemini-optin", + "litellm_params": { + "model": "vertex_ai/gemini-2.5-flash", + "configurable_clientside_auth_params": ["vertex_location", {"api_base": r"https://eu\..*"}], + }, + }, + {"model_name": "openai-group", "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}}, + {"model_name": "studio-group", "litellm_params": {"model": "gemini/gemini-2.5-flash", "api_key": "k"}}, + ], + fallbacks=[{"gemini-optin": ["openai-group"]}, {"studio-group": ["openai-group"]}], + ) + request_body: Final = {"model": model, **fields} + if is_allowed: + assert is_request_body_safe(request_body, {}, router, model, route="/v1/chat/completions") is True + return + with pytest.raises(ValueError, match=r"vertex_location|api_base"): + is_request_body_safe(request_body, {}, router, model, route="/v1/chat/completions") + + @pytest.mark.parametrize( + "model, is_allowed", [("public-alias", True), ("gemini-optin", False)], ids=["alias", "aliased_target"] + ) + def test_alias_chain_is_checked_one_hop_like_the_router_serves_it(self, model: str, is_allowed: bool) -> None: + router: Final = self._opted_in_and_pinned_router( + [], model_group_alias={"public-alias": "gemini-optin", "gemini-optin": "gemini-pinned"} + ) + request_body: Final = {"model": model, "vertex_location": "us-central1"} + if is_allowed: + assert is_request_body_safe(request_body, {}, router, model, route="/v1/chat/completions") is True + return + with pytest.raises(ValueError, match="vertex_location"): + is_request_body_safe(request_body, {}, router, model, route="/v1/chat/completions") + + @pytest.mark.parametrize( + "route", + [ + "/openai/deployments/gemini-pinned/chat/completions", + "/v1beta/models/gemini-pinned:generateContent", + ], + ) + def test_model_named_in_the_route_needs_the_opt_in_too(self, route: str) -> None: + router: Final = self._opted_in_and_pinned_router([]) + request_body: Final = {"model": "gemini-optin", "vertex_location": "us-central1"} + with pytest.raises(ValueError, match="vertex_location"): + is_request_body_safe(request_body, {}, router, "gemini-optin", route=route) + @pytest.mark.asyncio @pytest.mark.parametrize("opted_in", [True, False]) async def test_opt_in_of_the_model_named_in_a_google_route_decides(self, monkeypatch, opted_in: bool) -> None: @@ -5208,8 +5277,8 @@ class TestIsRequestBodySafeBlocksVertexTargetOverrides: "route, kept_params", [ ("/vertex_ai/live", _VERTEX_TARGET_PARAMS), - ("/rag/ingest", ("vertex_project", "vertex_location")), - ("/v1/rag/ingest", ("vertex_project", "vertex_location")), + ("/rag/ingest", _VERTEX_TARGET_PARAMS), + ("/v1/rag/ingest", _VERTEX_TARGET_PARAMS), ], ) def test_routes_that_name_their_vertex_target_keep_it(self, route: str, kept_params: tuple[str, ...]) -> None: @@ -5229,8 +5298,7 @@ class TestIsRequestBodySafeBlocksVertexTargetOverrides: [ ("/vertex_ai/live", "api_base"), ("/rag/ingest", "api_base"), - ("/rag/ingest", "vertex_ai_project"), - ("/rag/ingest", "vertex_ai_location"), + ("/v1/rag/ingest", "api_base"), ], ) def test_routes_that_name_their_vertex_target_reject_every_other_banned_param(self, route: str, param: str) -> None: diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index beaa7bb380d..75efce38a07 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -6338,6 +6338,7 @@ class TestVertexAILiveWebsocketPassthrough: rewriter = passthrough_module._build_vertex_live_setup_model_rewriter( vertex_project="proj-db", vertex_location="global", + configured_location="global", llm_router=llm_router, keeps_caller_resource_path=False, ) @@ -6706,6 +6707,103 @@ class TestVertexAILiveWebsocketPassthrough: == expected ) + def _pin_no_vertex_project( + self, + monkeypatch, + general_settings: Mapping[str, object], + default_vertex_config: VertexPassThroughCredentials | None, + ): + from litellm.proxy.pass_through_endpoints import ( + llm_passthrough_endpoints as passthrough_module, + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", dict(general_settings)) + monkeypatch.setattr( + passthrough_module.passthrough_endpoint_router, "default_vertex_config", default_vertex_config + ) + monkeypatch.setattr(passthrough_module.passthrough_endpoint_router, "deployment_key_to_vertex_credentials", {}) + self._clear_vertex_env(monkeypatch) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "default_vertex_config", + [None, VertexPassThroughCredentials(vertex_location="europe-west4")], + ids=["no_default_vertex_config", "default_vertex_config_without_a_project"], + ) + async def test_caller_named_project_closes_before_any_upstream_connect_when_nothing_pins_one( + self, monkeypatch, default_vertex_config: VertexPassThroughCredentials | None + ): + self._pin_no_vertex_project(monkeypatch, {}, default_vertex_config) + + websocket, ensure_token, ws_passthrough = await self._open_live({"vertex_project": "other-project"}) + + ws_passthrough.assert_not_awaited() + ensure_token.assert_not_awaited() + close_kwargs = websocket.close.await_args.kwargs + assert close_kwargs["code"] == 1008 + assert "vertex_project" in close_kwargs["reason"] + assert len(close_kwargs["reason"].encode("utf-8")) <= 123 + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "general_settings, query, expected_project, expected_location", + [ + ({}, {"vertex_location": "europe-west4"}, "proj-env", "europe-west4"), + ( + {"allow_client_side_credentials": True}, + {"vertex_project": "other-project", "vertex_location": "europe-west4"}, + "other-project", + "europe-west4", + ), + ], + ids=["caller_region", "proxy_wide_opt_in"], + ) + async def test_with_nothing_pinned_the_caller_keeps_the_region_and_the_opt_in_keeps_the_project( + self, + monkeypatch, + general_settings: Mapping[str, object], + query: Mapping[str, str], + expected_project: str, + expected_location: str, + ): + self._pin_no_vertex_project(monkeypatch, general_settings, None) + + websocket, _, ws_passthrough = await self._open_live(query) + + websocket.close.assert_not_awaited() + passthrough_kwargs = ws_passthrough.await_args.kwargs + assert passthrough_kwargs["custom_headers"]["x-goog-user-project"] == expected_project + assert passthrough_kwargs["setup_model_rewriter"]("gemini-2.0-flash") == ( + f"projects/{expected_project}/locations/{expected_location}/publishers/google/models/gemini-2.0-flash" + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "setup_model, expected", + [ + ( + "projects/proj-env/locations/europe-west4/publishers/google/models/gemini-2.0-flash", + "projects/proj-env/locations/europe-west4/publishers/google/models/gemini-2.0-flash", + ), + ( + "projects/caller-project/locations/europe-west4/publishers/google/models/gemini-2.0-flash", + "projects/proj-env/locations/europe-west4/publishers/google/models/gemini-2.0-flash", + ), + ], + ids=["own_project", "other_project"], + ) + async def test_setup_frame_keeps_its_region_when_no_region_is_pinned( + self, monkeypatch, setup_model: str, expected: str + ): + self._use_default_vertex_config(monkeypatch, {}, vertex_location=None) + + websocket, _, ws_passthrough = await self._open_live({}) + + websocket.close.assert_not_awaited() + rewriter = ws_passthrough.await_args.kwargs["setup_model_rewriter"] + assert rewriter(setup_model) == expected + class TestAnthropicProxyRoute: """The /anthropic passthrough route: custom auth headers must not clobber the diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 228a3934596..53feb88eac8 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -7396,6 +7396,7 @@ async def _run_setup_rewrite_passthrough(setup_model: str, llm_router, keeps_cal setup_model_rewriter=_build_vertex_live_setup_model_rewriter( vertex_project="proj-db", vertex_location="global", + configured_location="global", llm_router=llm_router, keeps_caller_resource_path=keeps_caller_resource_path, ), @@ -7435,6 +7436,67 @@ async def test_websocket_passthrough_rewrites_gateway_alias_setup_model(): assert sent_setup["model"] == "projects/proj-db/locations/global/publishers/google/models/gemini-live-2.5-flash" +async def _run_binary_client_frame_passthrough(frame: bytes): + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + _build_vertex_live_setup_model_rewriter, + ) + + upstream_ws = RecordingUpstreamWebSocket() + websocket = _client_websocket( + AsyncMock( + side_effect=[ + {"type": "websocket.receive", "bytes": frame}, + {"type": "websocket.disconnect"}, + ] + ) + ) + + with _patched_websocket_passthrough_environment(upstream_ws): + await websocket_passthrough_request( + websocket=websocket, + target="wss://aiplatform.googleapis.com/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent", + custom_headers={"Authorization": "Bearer token"}, + user_api_key_dict=UserAPIKeyAuth(), + forward_headers=False, + endpoint="/vertex_ai/live", + accept_websocket=False, + setup_model_rewriter=_build_vertex_live_setup_model_rewriter( + vertex_project="proj-a", + vertex_location="europe-west4", + configured_location="europe-west4", + llm_router=None, + keeps_caller_resource_path=False, + ), + ) + + return websocket, upstream_ws + + +@pytest.mark.asyncio +async def test_websocket_passthrough_moves_a_binary_setup_frame_onto_the_pinned_project(): + setup_frame = json.dumps( + {"setup": {"model": "projects/proj-b/locations/europe-west4/publishers/google/models/gemini-2.0-flash"}} + ).encode("utf-8") + + _, upstream_ws = await _run_binary_client_frame_passthrough(setup_frame) + + upstream_ws.send.assert_awaited_once() + sent_frame = upstream_ws.send.await_args.args[0] + assert isinstance(sent_frame, bytes), "a binary client message must stay binary upstream" + assert json.loads(sent_frame)["setup"]["model"] == ( + "projects/proj-a/locations/europe-west4/publishers/google/models/gemini-2.0-flash" + ) + + +@pytest.mark.asyncio +async def test_websocket_passthrough_closes_on_a_binary_frame_that_is_not_utf8(): + websocket, upstream_ws = await _run_binary_client_frame_passthrough(b"\xff\xfe\x00projects/proj-b") + + upstream_ws.send.assert_not_awaited() + websocket.close.assert_awaited_once() + assert websocket.close.await_args.kwargs["code"] == 1008 + + @pytest.mark.parametrize( "setup_model", ["gemini-live-2.5-flash", "models/gemini-live-2.5-flash", "publishers/google/models/gemini-live-2.5-flash"], @@ -7456,7 +7518,11 @@ def test_vertex_live_setup_model_resolves_before_extraction(setup_model): ) rewriter = _build_vertex_live_setup_model_rewriter( - vertex_project="proj-db", vertex_location="global", llm_router=None, keeps_caller_resource_path=False + vertex_project="proj-db", + vertex_location="global", + configured_location="global", + llm_router=None, + keeps_caller_resource_path=False, ) setup_data = {"model": setup_model} @@ -7508,6 +7574,7 @@ async def test_websocket_passthrough_logs_a_bare_alias_setup_model(): setup_model_rewriter=_build_vertex_live_setup_model_rewriter( vertex_project="proj-db", vertex_location="global", + configured_location="global", llm_router=None, keeps_caller_resource_path=False, ), diff --git a/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py b/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py index 664556a8a4f..098b2065ae1 100644 --- a/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py @@ -15,7 +15,8 @@ from fastapi.testclient import TestClient import litellm from litellm import Router -from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.constants import LITELLM_PROXY_ADMIN_NAME +from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.proxy_server import app from litellm.types.vector_stores import LiteLLM_ManagedVectorStore @@ -284,6 +285,7 @@ VERTEX_REGISTRY_STORE = { "vector_store_id": "projects/registry-project/locations/us-central1/ragCorpora/42", "custom_llm_provider": "vertex_ai", "litellm_params": {"vertex_project": "registry-project", "vertex_location": "us-central1"}, + "is_config": True, } UNSUPPORTED_INGEST_PROVIDER_ERROR = ( "Provider '{provider}' is not supported for RAG ingestion. " @@ -1028,8 +1030,20 @@ def _vertex_ingest(client: TestClient, vector_store: Mapping[str, str]) -> tuple "vector_store_id", ), ({"vector_store_id": f"projects/{_PROXY_VERTEX_PROJECT}/locations/us-east5/ragCorpora/7"}, "vector_store_id"), + ({"vertex_ai_project": "other-project"}, "vertex_ai_project"), + ({"vertex_ai_location": "us-east5"}, "vertex_ai_location"), + ({"vector_store_id": "projects/123456789/locations/us-east5/ragCorpora/7"}, "vector_store_id"), + ], + ids=[ + "project_field", + "location_field", + "corpus_project", + "corpus_project_behind_matching_field", + "corpus_region", + "project_alias_field", + "location_alias_field", + "corpus_region_by_project_number", ], - ids=["project_field", "location_field", "corpus_project", "corpus_project_behind_matching_field", "corpus_region"], ) @pytest.mark.usefixtures("proxy_vertex_target") def test_rag_ingest_refuses_a_key_pointing_the_proxys_vertex_credentials_elsewhere( @@ -1060,8 +1074,9 @@ def test_rag_ingest_refuses_a_key_naming_a_vertex_project_when_the_proxy_names_n {}, {"vertex_project": _PROXY_VERTEX_PROJECT, "vertex_location": _PROXY_VERTEX_LOCATION}, {"vector_store_id": f"projects/{_PROXY_VERTEX_PROJECT}/locations/{_PROXY_VERTEX_LOCATION}/ragCorpora/7"}, + {"vertex_ai_project": _PROXY_VERTEX_PROJECT, "vertex_ai_location": _PROXY_VERTEX_LOCATION}, ], - ids=["proxy_defaults", "same_fields", "corpus_in_the_proxy_project"], + ids=["proxy_defaults", "same_fields", "corpus_in_the_proxy_project", "same_alias_fields"], ) @pytest.mark.usefixtures("proxy_vertex_target") def test_rag_ingest_lets_a_key_use_the_proxys_own_vertex_target(client_internal_user, vector_store: Mapping[str, str]): @@ -1074,6 +1089,11 @@ def test_rag_ingest_lets_a_key_use_the_proxys_own_vertex_target(client_internal_ } +@pytest.mark.parametrize( + "project_field, location_field", + [("vertex_project", "vertex_location"), ("vertex_ai_project", "vertex_ai_location")], + ids=["fields", "alias_fields"], +) @pytest.mark.parametrize( "user_role, general_settings", [ @@ -1084,14 +1104,14 @@ def test_rag_ingest_lets_a_key_use_the_proxys_own_vertex_target(client_internal_ ) @pytest.mark.usefixtures("proxy_vertex_target") def test_rag_ingest_lets_the_admin_choose_another_vertex_target( - user_role: LitellmUserRoles, general_settings: Mapping[str, object] + user_role: LitellmUserRoles, general_settings: Mapping[str, object], project_field: str, location_field: str ): caller: Final = UserAPIKeyAuth(user_id="vertex-caller", user_role=user_role) original_overrides: Final = app.dependency_overrides.copy() app.dependency_overrides[user_api_key_auth] = lambda: caller vector_store: Final = { - "vertex_project": "other-project", - "vertex_location": "us-east5", + project_field: "other-project", + location_field: "us-east5", "vector_store_id": "projects/other-project/locations/us-east5/ragCorpora/7", } try: @@ -1107,6 +1127,119 @@ def test_rag_ingest_lets_the_admin_choose_another_vertex_target( } +@pytest.mark.usefixtures("proxy_vertex_target") +def test_rag_ingest_reads_a_corpus_named_by_project_number_on_the_proxys_own_project(client_internal_user): + response, mock_aingest = _vertex_ingest( + client_internal_user, {"vector_store_id": f"projects/123456789/locations/{_PROXY_VERTEX_LOCATION}/ragCorpora/7"} + ) + + assert response.status_code == 200, response.json() + assert mock_aingest.await_args.kwargs["ingest_options"]["vector_store"]["vector_store_id"] == ( + f"projects/{_PROXY_VERTEX_PROJECT}/locations/{_PROXY_VERTEX_LOCATION}/ragCorpora/7" + ) + + +_PROXY_CORPUS_PREFIX: Final = f"projects/{_PROXY_VERTEX_PROJECT}/locations/{_PROXY_VERTEX_LOCATION}/ragCorpora" +_FOREIGN_CORPUS: Final = f"projects/other-project/locations/{_PROXY_VERTEX_LOCATION}/ragCorpora/1" + + +@pytest.mark.parametrize( + "vector_store_id", + [ + f"./{_FOREIGN_CORPUS}", + f"x/../{_FOREIGN_CORPUS}", + f"{_PROXY_CORPUS_PREFIX}/1{'/..' * 6}/{_FOREIGN_CORPUS}", + f"{_PROXY_CORPUS_PREFIX}/..", + f"{_PROXY_CORPUS_PREFIX}/1%2F..%2F..", + f"{_PROXY_CORPUS_PREFIX}/1\\..\\..", + ], + ids=["dot_prefix", "dot_dot_prefix", "own_corpus_then_dot_dots", "dot_dot_corpus", "percent", "backslash"], +) +@pytest.mark.usefixtures("proxy_vertex_target") +def test_rag_ingest_refuses_a_vertex_corpus_name_with_path_tricks(client_internal_user, vector_store_id: str): + response, mock_aingest = _vertex_ingest(client_internal_user, {"vector_store_id": vector_store_id}) + + assert response.status_code == 400, response.json() + assert response.json()["detail"]["error"].startswith("Rejected Request: vector_store_id picks a Vertex project") + mock_aingest.assert_not_awaited() + + +def _foreign_vertex_store(created_by: str | None) -> LiteLLM_ManagedVectorStore: + return LiteLLM_ManagedVectorStore( + vector_store_id=_FOREIGN_CORPUS, custom_llm_provider="vertex_ai", litellm_params={}, user_id=created_by + ) + + +def _prisma_with_user_row(user_row: LiteLLM_UserTable | None) -> MagicMock: + prisma_client: Final = MagicMock() + prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=user_row) + prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock( + return_value=MagicMock(vector_store_metadata={}) + ) + prisma_client.db.litellm_managedvectorstorestable.update = AsyncMock() + return prisma_client + + +def _ingest_into_saved_store( + store: LiteLLM_ManagedVectorStore, creator_row: LiteLLM_UserTable | None +) -> tuple[httpx.Response, AsyncMock]: + caller: Final = UserAPIKeyAuth(user_id="vertex-caller", user_role=LitellmUserRoles.INTERNAL_USER) + original_overrides: Final = app.dependency_overrides.copy() + app.dependency_overrides[user_api_key_auth] = lambda: caller + aingest_patch, registry_patch = _patched_ingest_boundary( + store, {"vector_store_id": store["vector_store_id"], "file_id": "file_123"} + ) + try: + with aingest_patch as mock_aingest, registry_patch, _patched_prisma_client(_prisma_with_user_row(creator_row)): + response = TestClient(app).post( + "/v1/rag/ingest", **_ingest_form({"vector_store_id": store["vector_store_id"]}) + ) + finally: + app.dependency_overrides = original_overrides + return response, mock_aingest + + +@pytest.mark.parametrize( + "creator_row", + [LiteLLM_UserTable(user_id="vertex-caller", user_role=LitellmUserRoles.INTERNAL_USER.value), None], + ids=["saved_by_an_internal_user", "saved_by_a_key_with_no_user"], +) +@pytest.mark.usefixtures("proxy_vertex_target") +def test_rag_ingest_refuses_a_saved_vertex_store_a_non_admin_pointed_elsewhere(creator_row: LiteLLM_UserTable | None): + store: Final = _foreign_vertex_store(creator_row.user_id if creator_row is not None else None) + + response, mock_aingest = _ingest_into_saved_store(store, creator_row) + + assert response.status_code == 400, response.json() + assert response.json()["detail"]["error"].startswith("Rejected Request: vector_store_id picks a Vertex project") + mock_aingest.assert_not_awaited() + + +@pytest.mark.parametrize( + "store, creator_row", + [ + (_foreign_vertex_store(LITELLM_PROXY_ADMIN_NAME), None), + ( + _foreign_vertex_store("admin-user"), + LiteLLM_UserTable(user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN.value), + ), + (LiteLLM_ManagedVectorStore(**_foreign_vertex_store(None), is_config=True), None), + ], + ids=["master_key_created", "admin_user_created", "config_store"], +) +@pytest.mark.usefixtures("proxy_vertex_target") +def test_rag_ingest_lets_a_key_use_a_vertex_store_an_admin_pointed_elsewhere( + store: LiteLLM_ManagedVectorStore, creator_row: LiteLLM_UserTable | None +): + response, mock_aingest = _ingest_into_saved_store(store, creator_row) + + assert response.status_code == 200, response.json() + assert mock_aingest.await_args.kwargs["ingest_options"]["vector_store"] == { + "vector_store_id": _FOREIGN_CORPUS, + "custom_llm_provider": "vertex_ai", + } + + def test_rag_ingest_db_managed_store_drops_the_callers_credential_name(client_internal_user): aingest_patch, registry_patch = _patched_ingest_boundary( DB_MANAGED_STORE, {"vector_store_id": "db-store", "file_id": "file_123"} From 53c743eaaabab552162e9a5fac3db6f425889321 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 12:22:53 -0700 Subject: [PATCH 08/11] fix(proxy): let a live caller name the proxy's configured Vertex project when nothing pins one With default_vertex_config set but no project pinned, /vertex_ai/live closed every query that named a project, including the one the admin set in VERTEXAI_PROJECT or DEFAULT_VERTEXAI_PROJECT. A query naming one of those projects now connects, and any other project still closes with 1008 before a token is minted. --- .../llm_passthrough_endpoints.py | 26 +++++++++++++++++-- .../test_llm_pass_through_endpoints.py | 24 +++++++++++++++++ 2 files changed, 48 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 4dffa42f9cd..a91dc62cf4b 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -4026,6 +4026,28 @@ def _requested_vertex_live_target( ) +def _proxy_configured_vertex_projects() -> frozenset[str]: + return frozenset( + project + for project in ( + litellm.vertex_project, + get_secret_str("VERTEXAI_PROJECT"), + get_secret_str("DEFAULT_VERTEXAI_PROJECT"), + ) + if project + ) + + +def _names_an_unconfigured_project( + credentials: VertexPassThroughCredentials | None, vertex_project: str | None +) -> bool: + return ( + bool(vertex_project) + and (credentials is None or credentials.vertex_project is None) + and vertex_project not in _proxy_configured_vertex_projects() + ) + + def _live_target_override_reason( target: _VertexLiveTarget, vertex_project: str | None, @@ -4035,8 +4057,8 @@ def _live_target_override_reason( if allow_client_side_credentials: return None match target: - case _ConfiguredVertexLiveTarget(credentials=credentials) if vertex_project and ( - credentials is None or credentials.vertex_project is None + case _ConfiguredVertexLiveTarget(credentials=credentials) if _names_an_unconfigured_project( + credentials, vertex_project ): return VERTEX_LIVE_UNPINNED_PROJECT_CLOSE_REASON case _ConfiguredVertexLiveTarget(credentials=VertexPassThroughCredentials() as credentials): diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index afdfaa82e34..0475fc1ba5d 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -6811,6 +6811,30 @@ class TestVertexAILiveWebsocketPassthrough: assert "vertex_project" in close_kwargs["reason"] assert len(close_kwargs["reason"].encode("utf-8")) <= 123 + @pytest.mark.asyncio + @pytest.mark.parametrize( + "configured_by, default_vertex_config", + [ + ("VERTEXAI_PROJECT", None), + ("DEFAULT_VERTEXAI_PROJECT", VertexPassThroughCredentials(vertex_location="europe-west4")), + ], + ids=["vertexai_project_env", "default_vertexai_project_env"], + ) + async def test_with_nothing_pinned_the_caller_may_name_the_project_the_admin_configured( + self, monkeypatch, configured_by: str, default_vertex_config: VertexPassThroughCredentials | None + ): + self._pin_no_vertex_project(monkeypatch, {}, default_vertex_config) + monkeypatch.setattr(litellm, "vertex_project", None) + monkeypatch.delenv("VERTEXAI_PROJECT", raising=False) + monkeypatch.setenv(configured_by, "proj-configured") + + websocket, _, ws_passthrough = await self._open_live( + {"vertex_project": "proj-configured", "vertex_location": "europe-west4"} + ) + + websocket.close.assert_not_awaited() + assert ws_passthrough.await_args.kwargs["custom_headers"]["x-goog-user-project"] == "proj-configured" + @pytest.mark.asyncio @pytest.mark.parametrize( "general_settings, query, expected_project, expected_location", From 29e241c6874ecf51ad87cd9a89196b77dcb9ad8c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 21:18:17 -0700 Subject: [PATCH 09/11] fix(proxy): keep an admin-created vector store's corpus as saved on RAG ingest A non-admin ingest into a store an admin created named by project number had its corpus moved onto the proxy's own project before the admin-store exemption ran, so the upload went to a different corpus than the admin saved. The exemption now runs first, and a non-admin store still gets the project-number rewrite and the target check. --- litellm/proxy/rag_endpoints/endpoints.py | 4 +-- .../proxy/rag_endpoints/test_rag_endpoints.py | 25 +++++++++++-------- 2 files changed, 16 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index 35e7ae71fb7..aee958fae3d 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -791,10 +791,10 @@ async def rag_ingest( } vertex_target_checked: Final = _vertex_target_is_checked( requested_base, user_api_key_dict, general_settings_view() - ) + ) and not await _admin_set_store_target(managed_store, prisma_client) base: Final = _corpus_on_configured_project(requested_base) if vertex_target_checked else requested_base vertex_target_error: Final = _vertex_target_error(base) if vertex_target_checked else None - if vertex_target_error is not None and not await _admin_set_store_target(managed_store, prisma_client): + if vertex_target_error is not None: raise HTTPException(status_code=400, detail={"error": vertex_target_error}) request_data: Final[dict[str, Any]] = await add_litellm_data_to_request( data={}, diff --git a/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py b/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py index 098b2065ae1..620dc13aacb 100644 --- a/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py @@ -1164,9 +1164,12 @@ def test_rag_ingest_refuses_a_vertex_corpus_name_with_path_tricks(client_interna mock_aingest.assert_not_awaited() -def _foreign_vertex_store(created_by: str | None) -> LiteLLM_ManagedVectorStore: +_FOREIGN_CORPUS_BY_NUMBER: Final = f"projects/987654321/locations/{_PROXY_VERTEX_LOCATION}/ragCorpora/1" + + +def _foreign_vertex_store(created_by: str | None, corpus: str = _FOREIGN_CORPUS) -> LiteLLM_ManagedVectorStore: return LiteLLM_ManagedVectorStore( - vector_store_id=_FOREIGN_CORPUS, custom_llm_provider="vertex_ai", litellm_params={}, user_id=created_by + vector_store_id=corpus, custom_llm_provider="vertex_ai", litellm_params={}, user_id=created_by ) @@ -1215,27 +1218,27 @@ def test_rag_ingest_refuses_a_saved_vertex_store_a_non_admin_pointed_elsewhere(c mock_aingest.assert_not_awaited() +@pytest.mark.parametrize("corpus", [_FOREIGN_CORPUS, _FOREIGN_CORPUS_BY_NUMBER], ids=["by_id", "by_number"]) @pytest.mark.parametrize( - "store, creator_row", + "created_by, is_config, creator_row", [ - (_foreign_vertex_store(LITELLM_PROXY_ADMIN_NAME), None), - ( - _foreign_vertex_store("admin-user"), - LiteLLM_UserTable(user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN.value), - ), - (LiteLLM_ManagedVectorStore(**_foreign_vertex_store(None), is_config=True), None), + (LITELLM_PROXY_ADMIN_NAME, False, None), + ("admin-user", False, LiteLLM_UserTable(user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN.value)), + (None, True, None), ], ids=["master_key_created", "admin_user_created", "config_store"], ) @pytest.mark.usefixtures("proxy_vertex_target") def test_rag_ingest_lets_a_key_use_a_vertex_store_an_admin_pointed_elsewhere( - store: LiteLLM_ManagedVectorStore, creator_row: LiteLLM_UserTable | None + corpus: str, created_by: str | None, is_config: bool, creator_row: LiteLLM_UserTable | None ): + store: Final = LiteLLM_ManagedVectorStore(**_foreign_vertex_store(created_by, corpus), is_config=is_config) + response, mock_aingest = _ingest_into_saved_store(store, creator_row) assert response.status_code == 200, response.json() assert mock_aingest.await_args.kwargs["ingest_options"]["vector_store"] == { - "vector_store_id": _FOREIGN_CORPUS, + "vector_store_id": corpus, "custom_llm_provider": "vertex_ai", } From 31e7af423da45568949d9c17c207d088f1baefc7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 31 Oct 2026 01:09:55 -0700 Subject: [PATCH 10/11] fix(proxy): authorize a RAG corpus moved onto the proxy's project and check client-side opt-ins once per request --- litellm/proxy/auth/auth_utils.py | 80 ++++++++++++++----- litellm/proxy/rag_endpoints/endpoints.py | 9 +++ tests/unit/proxy/auth/test_auth_utils.py | 34 ++++++++ .../proxy/rag_endpoints/test_rag_endpoints.py | 35 ++++++++ 4 files changed, 140 insertions(+), 18 deletions(-) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 8270314faf1..f87400371b3 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -2,8 +2,9 @@ import importlib.util import os import re import sys -from collections.abc import Collection, Iterator, Mapping, Sequence -from functools import lru_cache +from collections.abc import Callable, Collection, Iterator, Mapping, Sequence +from dataclasses import dataclass +from functools import cache, lru_cache from itertools import accumulate, chain, repeat from logging import Logger from types import MappingProxyType @@ -546,16 +547,63 @@ def _opt_in_model_groups( return routed | frozenset(registry.resolve(name, kind="alias") or name for name in fallbacks) +_MAX_OPT_IN_FALLBACK_MODELS: Final = 100 + + +@dataclass(frozen=True, slots=True) +class _OptInGroups: + every: tuple[str, ...] + vertex_target: tuple[str, ...] + + +def _deployment_id(deployment: Mapping[str, object]) -> str | None: + match deployment.get("model_info"): + case {"id": str(deployment_id)} if deployment_id: + return deployment_id + case _: + return None + + +def _same_deployments_key(model_group: str, llm_router: Router) -> frozenset[tuple[str | None, str | None]] | str: + deployments: Final = _model_group_deployments(model_group, llm_router) + key: Final = tuple( + (_deployment_id(deployment), _deployment_provider(deployment["litellm_params"])) for deployment in deployments + ) + return frozenset(key) if key and all(deployment_id for deployment_id, _ in key) else model_group + + +def _distinct_opt_in_groups(model_groups: frozenset[str], llm_router: Router | None) -> tuple[str, ...]: + if llm_router is None: + return tuple(sorted(model_groups)) + return tuple( + {_same_deployments_key(model_group, llm_router): model_group for model_group in sorted(model_groups)}.values() + ) + + +def _resolve_opt_in_groups( + model: str, also_requested_models: tuple[str, ...], llm_router: Router | None +) -> _OptInGroups: + if len(frozenset(also_requested_models)) > _MAX_OPT_IN_FALLBACK_MODELS: + raise ValueError( + f"Rejected Request: the request names more than {_MAX_OPT_IN_FALLBACK_MODELS} fallback models, " + "too many to check against configurable_clientside_auth_params." + ) + every: Final = _distinct_opt_in_groups(_opt_in_model_groups(model, also_requested_models, llm_router), llm_router) + return _OptInGroups( + every=every, + vertex_target=tuple(group for group in every if _model_group_reads_vertex_target(group, llm_router)), + ) + + def _check_banned_params( body: Mapping[str, object], general_settings: dict, llm_router: Router | None, - model: str, *, manages_deployments: bool = False, mcp_server_fields_allowed: bool = False, banned_params: tuple[str, ...] = _BANNED_REQUEST_BODY_PARAMS, - also_requested_models: tuple[str, ...] = (), + opt_in_groups: Callable[[], _OptInGroups], ) -> None: """Raise ``ValueError`` if ``body`` carries a banned param without admin opt-in. @@ -585,8 +633,9 @@ def _check_banned_params( llm_router=llm_router, ) is True - for model_group in _opt_in_model_groups(model, also_requested_models, llm_router) - if param not in VERTEX_TARGET_PARAMS or _model_group_reads_vertex_target(model_group, llm_router) + for model_group in ( + opt_in_groups().vertex_target if param in VERTEX_TARGET_PARAMS else opt_in_groups().every + ) ): # Per-param opt-in: only THIS param is permitted by the # deployment's ``configurable_clientside_auth_params``. Skip @@ -780,14 +829,14 @@ def is_request_body_safe( *request_fallback_model_names(root_body), *((route_model,) if route_model else ()), ) + opt_in_groups: Final = cache(lambda: _resolve_opt_in_groups(model, also_requested_models, llm_router)) _check_banned_params( root_body, general_settings, llm_router, - model, manages_deployments=manages_deployments, mcp_server_fields_allowed=route_manages_mcp_servers(route), - also_requested_models=also_requested_models, + opt_in_groups=opt_in_groups, ) for nested_key in _NESTED_CONFIG_KEYS: nested = _coerce_metadata_to_dict(request_body.get(nested_key)) @@ -796,9 +845,8 @@ def is_request_body_safe( nested, general_settings, llm_router, - model, manages_deployments=manages_deployments, - also_requested_models=also_requested_models, + opt_in_groups=opt_in_groups, ) for metadata_key in _NESTED_METADATA_KEYS: metadata = _coerce_metadata_to_dict(request_body.get(metadata_key)) @@ -807,20 +855,18 @@ def is_request_body_safe( metadata, general_settings, llm_router, - model, manages_deployments=manages_deployments, banned_params=_BANNED_METADATA_PARAMS, - also_requested_models=also_requested_models, + opt_in_groups=opt_in_groups, ) if any(isinstance(key, str) and key.startswith(f"{metadata_key}[") for key in request_body): _check_banned_params( extract_nested_form_metadata(form_data=request_body, prefix=f"{metadata_key}["), general_settings, llm_router, - model, manages_deployments=manages_deployments, banned_params=_BANNED_METADATA_PARAMS, - also_requested_models=also_requested_models, + opt_in_groups=opt_in_groups, ) for target in iter_request_fallback_targets(request_body): if isinstance(target, dict): @@ -828,9 +874,8 @@ def is_request_body_safe( target, general_settings, llm_router, - model, manages_deployments=manages_deployments, - also_requested_models=also_requested_models, + opt_in_groups=opt_in_groups, ) target_model = target.get("model") if isinstance(target_model, str): @@ -848,10 +893,9 @@ def is_request_body_safe( litellm_params_metadata, general_settings, llm_router, - model, manages_deployments=manages_deployments, banned_params=_BANNED_METADATA_PARAMS, - also_requested_models=also_requested_models, + opt_in_groups=opt_in_groups, ) for user_config_params in _user_config_litellm_params(request_body): # pyright: ignore[reportUnknownArgumentType] # request_body is an untyped dict reject_api_key_secret_reference(user_config_params) diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index aee958fae3d..e105afda33c 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -234,6 +234,14 @@ def _corpus_on_configured_project(vector_store_config: Mapping[str, object]) -> ) +async def _authorize_rewritten_corpus( + requested: Mapping[str, object], rewritten: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth +) -> None: + rewritten_id: Final = rewritten.get("vector_store_id") + if isinstance(rewritten_id, str) and rewritten_id != requested.get("vector_store_id"): + await assert_user_can_access_vector_store_id(vector_store_id=rewritten_id, user_api_key_dict=user_api_key_dict) + + def _vertex_target_error(vector_store_config: Mapping[str, object]) -> str | None: corpus_name: Final = vector_store_config.get("vector_store_id") corpus: Final = _clean_vertex_corpus_name(corpus_name) @@ -793,6 +801,7 @@ async def rag_ingest( requested_base, user_api_key_dict, general_settings_view() ) and not await _admin_set_store_target(managed_store, prisma_client) base: Final = _corpus_on_configured_project(requested_base) if vertex_target_checked else requested_base + await _authorize_rewritten_corpus(requested_base, base, user_api_key_dict) vertex_target_error: Final = _vertex_target_error(base) if vertex_target_checked else None if vertex_target_error is not None: raise HTTPException(status_code=400, detail={"error": vertex_target_error}) diff --git a/tests/unit/proxy/auth/test_auth_utils.py b/tests/unit/proxy/auth/test_auth_utils.py index 68b2326ac64..4db586645f1 100644 --- a/tests/unit/proxy/auth/test_auth_utils.py +++ b/tests/unit/proxy/auth/test_auth_utils.py @@ -4,6 +4,7 @@ Unit tests for auth_utils functions related to rate limiting and customer ID ext import base64 import logging +import time from collections.abc import Callable from typing import Final, Optional from unittest.mock import MagicMock, patch @@ -5265,6 +5266,39 @@ class TestIsRequestBodySafeBlocksVertexTargetOverrides: with pytest.raises(ValueError, match=r"vertex_location|api_base"): is_request_body_safe(request_body, {}, router, model, route="/v1/chat/completions") + def test_many_fallback_entries_on_one_wildcard_are_checked_in_linear_time(self) -> None: + router: Final = Router(model_list=[{"model_name": "*", "litellm_params": {"model": "openai/*", "api_key": "k"}}]) + request_body: Final = { + "model": "gpt-4o-mini", + "fallbacks": [{"model": f"model-{index % 100}", "vertex_location": f"r{index}"} for index in range(5000)], + } + + started: Final = time.perf_counter() + assert is_request_body_safe(request_body, {}, router, "gpt-4o-mini", route="/v1/chat/completions") is True + assert time.perf_counter() - started < 10 + + def test_too_many_fallback_models_for_an_opt_in_check_are_rejected(self) -> None: + router: Final = Router(model_list=[{"model_name": "*", "litellm_params": {"model": "openai/*", "api_key": "k"}}]) + request_body: Final = { + "model": "gpt-4o-mini", + "vertex_location": "us-central1", + "fallbacks": [f"model-{index}" for index in range(101)], + } + + with pytest.raises(ValueError, match="more than 100 fallback models"): + is_request_body_safe(request_body, {}, router, "gpt-4o-mini", route="/v1/chat/completions") + + def test_one_wildcard_serving_two_providers_checks_the_google_one(self) -> None: + router: Final = Router(model_list=[{"model_name": "*", "litellm_params": {"model": "*"}}]) + request_body: Final = { + "model": "vertex_ai/gemini-2.5-flash", + "vertex_location": "us-central1", + "fallbacks": ["xai/grok-3"], + } + + with pytest.raises(ValueError, match="vertex_location"): + is_request_body_safe(request_body, {}, router, "vertex_ai/gemini-2.5-flash", route="/v1/chat/completions") + @pytest.mark.parametrize( "model, is_allowed", [("public-alias", True), ("gemini-optin", False)], ids=["alias", "aliased_target"] ) diff --git a/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py b/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py index 620dc13aacb..3d92a812ef4 100644 --- a/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py @@ -1143,6 +1143,41 @@ _PROXY_CORPUS_PREFIX: Final = f"projects/{_PROXY_VERTEX_PROJECT}/locations/{_PRO _FOREIGN_CORPUS: Final = f"projects/other-project/locations/{_PROXY_VERTEX_LOCATION}/ragCorpora/1" +@pytest.mark.usefixtures("proxy_vertex_target") +def test_rag_ingest_refuses_another_teams_store_named_by_a_project_number(client_internal_user): + other_teams_store: Final = LiteLLM_ManagedVectorStore( + vector_store_id=f"{_PROXY_CORPUS_PREFIX}/1", + custom_llm_provider="vertex_ai", + litellm_params={}, + user_id="other-user", + team_id="other-team", + ) + registry: Final = MagicMock() + registry.get_litellm_managed_vector_store_from_registry.side_effect = lambda vector_store_id: ( + other_teams_store if vector_store_id == other_teams_store["vector_store_id"] else None + ) + with ( + patch( # test-quality-ok: aingest is the endpoint's downstream boundary; the test asserts it is never called + "litellm.proxy.rag_endpoints.endpoints.litellm.aingest", + new=AsyncMock(return_value={"vector_store_id": f"{_PROXY_CORPUS_PREFIX}/1", "file_id": "file_123"}), + ) as mock_aingest, + patch.object(litellm, "vector_store_registry", registry), # test-quality-ok: seeds the saved stores by id + _patched_prisma_client(None), + ): + response = client_internal_user.post( + "/v1/rag/ingest", + **_ingest_form( + { + "custom_llm_provider": "vertex_ai", + "vector_store_id": f"projects/0/locations/{_PROXY_VERTEX_LOCATION}/ragCorpora/1", + } + ), + ) + + assert response.status_code == 403, response.json() + mock_aingest.assert_not_awaited() + + @pytest.mark.parametrize( "vector_store_id", [ From 19c0adde5482231c666da5a329cb1a601225d763 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 31 Oct 2026 03:15:20 -0700 Subject: [PATCH 11/11] fix(proxy): merge opt-in fallback groups only when their deployments share provider and opt-in list --- litellm/proxy/auth/auth_utils.py | 27 +++++++++++++----------- tests/unit/proxy/auth/test_auth_utils.py | 20 ++++++++++++++++++ 2 files changed, 35 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index f87400371b3..50b96f6eeb9 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -556,27 +556,30 @@ class _OptInGroups: vertex_target: tuple[str, ...] -def _deployment_id(deployment: Mapping[str, object]) -> str | None: - match deployment.get("model_info"): - case {"id": str(deployment_id)} if deployment_id: - return deployment_id - case _: - return None +_OptInKey = tuple[str | tuple[str, str | None], ...] -def _same_deployments_key(model_group: str, llm_router: Router) -> frozenset[tuple[str | None, str | None]] | str: - deployments: Final = _model_group_deployments(model_group, llm_router) - key: Final = tuple( - (_deployment_id(deployment), _deployment_provider(deployment["litellm_params"])) for deployment in deployments +def _opt_in_key(litellm_params: LiteLLMParamsTypedDict) -> _OptInKey: + return tuple( + item if isinstance(item, str) else ("api_base", item.get("api_base")) + for item in deployment_clientside_auth_params(litellm_params) or () ) - return frozenset(key) if key and all(deployment_id for deployment_id, _ in key) else model_group + + +def _same_answer_key(model_group: str, llm_router: Router) -> frozenset[tuple[str | None, _OptInKey]] | str: + deployments: Final = _model_group_deployments(model_group, llm_router) + key: Final = frozenset( + (_deployment_provider(deployment["litellm_params"]), _opt_in_key(deployment["litellm_params"])) + for deployment in deployments + ) + return key or model_group def _distinct_opt_in_groups(model_groups: frozenset[str], llm_router: Router | None) -> tuple[str, ...]: if llm_router is None: return tuple(sorted(model_groups)) return tuple( - {_same_deployments_key(model_group, llm_router): model_group for model_group in sorted(model_groups)}.values() + {_same_answer_key(model_group, llm_router): model_group for model_group in sorted(model_groups)}.values() ) diff --git a/tests/unit/proxy/auth/test_auth_utils.py b/tests/unit/proxy/auth/test_auth_utils.py index 4db586645f1..7310b820472 100644 --- a/tests/unit/proxy/auth/test_auth_utils.py +++ b/tests/unit/proxy/auth/test_auth_utils.py @@ -5299,6 +5299,26 @@ class TestIsRequestBodySafeBlocksVertexTargetOverrides: with pytest.raises(ValueError, match="vertex_location"): is_request_body_safe(request_body, {}, router, "vertex_ai/gemini-2.5-flash", route="/v1/chat/completions") + def test_a_pinned_group_sharing_a_deployment_id_with_an_opted_in_group_is_still_checked(self) -> None: + router: Final = Router( + model_list=[ + { + "model_name": name, + "litellm_params": { + "model": "vertex_ai/gemini-2.5-flash", + "vertex_location": "europe-west4", + "configurable_clientside_auth_params": opt_in, + }, + "model_info": {"id": "shared-id"}, + } + for name, opt_in in (("a-pinned", []), ("z-optin", ["vertex_location"])) + ] + ) + request_body: Final = {"model": "a-pinned", "vertex_location": "us-central1", "fallbacks": ["z-optin"]} + + with pytest.raises(ValueError, match="vertex_location"): + is_request_body_safe(request_body, {}, router, "a-pinned", route="/v1/chat/completions") + @pytest.mark.parametrize( "model, is_allowed", [("public-alias", True), ("gemini-optin", False)], ids=["alias", "aliased_target"] )