From bb14db22c768c86092b4fe3d319503fc0b50c977 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 16:54:09 -0700 Subject: [PATCH] fix(proxy): count only Vertex AI deployments for Vertex AI Live model access A group another provider serves no longer grants a Vertex full path or gets its name rewritten for Vertex, and a caller with no model limit keeps working when every candidate name is a group serving another model. --- .../llm_passthrough_endpoints.py | 81 ++++++++++------- .../test_pass_through_endpoints.py | 87 +++++++++++++++++++ 2 files changed, 135 insertions(+), 33 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index e51b58a6c7b..2a536b9ffd1 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -3914,31 +3914,45 @@ def _deployment_named(model_name: str, deployments: Sequence[DeploymentTypedDict return next((deployment for deployment in deployments if deployment.get("model_name") == model_name), None) +VERTEX_LIVE_PROVIDERS: Final = frozenset({LlmProviders.VERTEX_AI.value, LlmProviders.VERTEX_AI_BETA.value}) + + +def _vertex_upstream_model(deployment: DeploymentTypedDict) -> str | None: + """ + The publisher model a Vertex AI deployment runs, or ``None`` for any other provider's deployment, since this + route runs every session on the proxy's Vertex credentials whatever provider a group was configured with + """ + litellm_params: Final = deployment["litellm_params"] + if "model" not in litellm_params: + return None + upstream: Final = litellm_params["model"] + try: + _, provider, _, _ = litellm.get_llm_provider( + model=upstream, custom_llm_provider=litellm_params.get("custom_llm_provider") + ) + except litellm.exceptions.BadRequestError: + return upstream + return upstream.removeprefix(f"{provider}/") if provider in VERTEX_LIVE_PROVIDERS else None + + +def _vertex_deployments(llm_router: Router | None) -> tuple[DeploymentTypedDict, ...]: + if llm_router is None: + return () + return tuple(deployment for deployment in (llm_router.get_model_list() or ()) if _vertex_upstream_model(deployment)) + + def _vertex_live_deployment(setup_model: str, llm_router: Router | None) -> DeploymentTypedDict | None: """ - The router deployment whose group a ``setup`` frame's model names, whichever addressing the client used. + The Vertex AI deployment whose group a ``setup`` frame's model names, whichever addressing the client used. The Live SDK wraps whatever the caller typed as ``models/``, and clients also send LiteLLM ids and full Vertex resource paths, so the group is read from the last segment when the whole string names none """ - if llm_router is None: - return None - deployments: Final = tuple(llm_router.get_model_list() or ()) + deployments: Final = _vertex_deployments(llm_router) exact: Final = _deployment_named(setup_model, deployments) return exact if exact is not None else _deployment_named(setup_model.rsplit("/", 1)[-1], deployments) -def _deployment_upstream_model(deployment: DeploymentTypedDict) -> str | None: - upstream: Final = deployment["litellm_params"].get("model") - if upstream is None: - return None - try: - _, provider, _, _ = litellm.get_llm_provider(model=upstream) - except litellm.exceptions.BadRequestError: - return upstream - return upstream.removeprefix(f"{provider}/") - - def _vertex_live_model_names(setup_model: str, llm_router: Router | None) -> tuple[str, ...]: """ The names a key can hold for the model a ``setup`` frame runs: the router group the setup rewriter sends it @@ -3953,16 +3967,17 @@ def _vertex_live_model_names(setup_model: str, llm_router: Router | None) -> tup ) if deployment is not None: return (deployment["model_name"],) - deployments: Final = () if llm_router is None else tuple(llm_router.get_model_list() or ()) publisher_model: Final = _vertex_publisher_model_suffix(setup_model) serving_groups: Final = tuple( dict.fromkeys( candidate["model_name"] - for candidate in deployments - if _vertex_publisher_model_suffix(_deployment_upstream_model(candidate) or "") == publisher_model + for candidate in _vertex_deployments(llm_router) + if _vertex_publisher_model_suffix(_vertex_upstream_model(candidate) or "") == publisher_model ) ) - group_names: Final = frozenset(candidate["model_name"] for candidate in deployments) + group_names: Final = frozenset( + candidate["model_name"] for candidate in (() if llm_router is None else llm_router.get_model_list() or ()) + ) publisher_model_id: Final = setup_model.rsplit("/", 1)[-1] return ( *serving_groups, @@ -4016,20 +4031,20 @@ async def _vertex_live_model_denial( for name in _vertex_live_model_names(setup_model, llm_router) if not _is_model_alias(name, valid_token, llm_router) ) - denials: Final = tuple([await _model_access_denial(name, valid_token, llm_router) for name in model_names]) - if None in denials: - return None - return ( - denials[0] - if denials - else ModelAccessDeniedProxyException( - message=model_access_denied_client_message(model=setup_model), - internal_message=f"Vertex AI Live setup model {setup_model} is only reachable through a model alias", - type=ProxyErrorTypes.key_model_access_denied, - param="model", - code=403, + if not model_names: + return ( + ModelAccessDeniedProxyException( + message=model_access_denied_client_message(model=setup_model), + internal_message=f"Vertex AI Live setup model {setup_model} has no name a model limit can grant", + type=ProxyErrorTypes.key_model_access_denied, + param="model", + code=403, + ) + if await _caller_has_model_limits(valid_token) + else None ) - ) + denials: Final = tuple([await _model_access_denial(name, valid_token, llm_router) for name in model_names]) + return None if None in denials else denials[0] async def _caller_has_model_limits(valid_token: UserAPIKeyAuth) -> bool: @@ -4154,7 +4169,7 @@ def _resolve_alias_to_upstream_model(setup_model: str, llm_router: Router | None The Live SDK wraps whatever the caller typed as ``models/``, so a gateway alias arrives prefixed """ deployment: Final = _vertex_live_deployment(setup_model, llm_router) - upstream: Final = None if deployment is None else _deployment_upstream_model(deployment) + upstream: Final = None if deployment is None else _vertex_upstream_model(deployment) return setup_model if upstream is None else upstream 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 4002816f2e3..6452b6c7ccf 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 @@ -7968,6 +7968,93 @@ async def test_vertex_live_full_path_is_not_granted_by_a_group_named_after_it_th assert reason is not None and "gemini-3-pro-live" in reason +@pytest.mark.asyncio +@pytest.mark.parametrize( + "setup_model", + [VERTEX_LIVE_FULL_PATH.format(model=VERTEX_LIVE_GROUP), "openai-native-audio"], + ids=["full-path", "group-name"], +) +async def test_vertex_live_setup_frame_is_not_granted_by_a_group_another_provider_serves(setup_model: str): + """Every Live session runs on the proxy's Vertex credentials, so a group another provider serves under the + same model id must not open the Vertex model""" + llm_router = litellm.Router( + model_list=[ + { + "model_name": "openai-native-audio", + "litellm_params": {"model": f"openai/{VERTEX_LIVE_GROUP}", "api_base": "https://openai-compatible"}, + }, + ] + ) + + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [_setup_frame(setup_model)], + UserAPIKeyAuth(token="hashed", models=["openai-native-audio"]), + llm_router=llm_router, + ) + + upstream_ws.send.assert_not_awaited() + assert _policy_close_reason(websocket) is not None + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("custom_llm_provider", "reaches_vertex"), [("vertex_ai", True), ("openai", False)]) +async def test_vertex_live_full_path_is_authorized_only_by_a_group_whose_custom_provider_is_vertex( + custom_llm_provider: str, reaches_vertex: bool +): + llm_router = litellm.Router( + model_list=[ + { + "model_name": "native-audio", + "litellm_params": { + "model": VERTEX_LIVE_GROUP, + "custom_llm_provider": custom_llm_provider, + "api_base": "https://openai-compatible", + }, + }, + ] + ) + + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [_setup_frame(VERTEX_LIVE_FULL_PATH.format(model=VERTEX_LIVE_GROUP))], + UserAPIKeyAuth(token="hashed", models=["native-audio"]), + llm_router=llm_router, + ) + + assert upstream_ws.send.await_count == (1 if reaches_vertex else 0) + assert (_policy_close_reason(websocket) is None) is reaches_vertex + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("valid_token", "reaches_vertex"), + [ + (UserAPIKeyAuth(token="hashed", models=[]), True), + (UserAPIKeyAuth(token="hashed", models=["gemini-3-pro-live"]), False), + ], + ids=["unrestricted-key", "limited-key"], +) +async def test_vertex_live_full_path_whose_every_name_is_taken_by_another_model_refuses_only_limited_keys( + valid_token: UserAPIKeyAuth, reaches_vertex: bool +): + """Groups holding both the bare and the ``vertex_ai/`` id for another model leave no name a key could hold, + which says nothing about a key that may call any model""" + llm_router = litellm.Router( + model_list=[ + {"model_name": "gemini-3-pro-live", "litellm_params": {"model": "vertex_ai/gemini-live-2.5-flash"}}, + {"model_name": "vertex_ai/gemini-3-pro-live", "litellm_params": {"model": "vertex_ai/gemini-live-2.5-flash"}}, + ] + ) + + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [_setup_frame(VERTEX_LIVE_FULL_PATH.format(model="gemini-3-pro-live"))], + valid_token, + llm_router=llm_router, + ) + + assert upstream_ws.send.await_count == (1 if reaches_vertex else 0) + assert (_policy_close_reason(websocket) is None) is reaches_vertex + + @pytest.mark.asyncio @pytest.mark.parametrize( "valid_token",