mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-01 07:03:02 +00:00
fix(proxy): authorize the Vertex Live setup model like a connect-time model and keep main's group rewriting
Re-run full websocket auth for the setup model, so JWT scopes, model budgets and the pricing block apply, and release its budget reservation. Resolve a listed router alias to its target group, and keep sending another provider's group to its Vertex model as main does
This commit is contained in:
mateo-berri
parent
804d6c91b1
commit
f0c8bc293e
5 changed files with 282 additions and 57 deletions
|
|
@ -15,11 +15,11 @@ import os
|
|||
import posixpath
|
||||
import re
|
||||
import sys
|
||||
from collections.abc import AsyncGenerator, Callable, Mapping, Sequence
|
||||
from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from functools import partial
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, cast
|
||||
from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeAlias, cast
|
||||
|
||||
import httpx
|
||||
import openai
|
||||
|
|
@ -90,6 +90,7 @@ from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy modul
|
|||
get_bearer_token,
|
||||
is_no_auth_dev_mode,
|
||||
user_api_key_auth,
|
||||
user_api_key_auth_accepted_websocket_for_model,
|
||||
user_api_key_auth_websocket,
|
||||
user_api_key_auth_websocket_for_model,
|
||||
)
|
||||
|
|
@ -3933,20 +3934,26 @@ def _deployment_named(model_name: str, deployments: Sequence[DeploymentTypedDict
|
|||
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
|
||||
"""
|
||||
def _upstream_model_and_provider(deployment: DeploymentTypedDict) -> tuple[str, str | None] | None:
|
||||
litellm_params: Final = deployment["litellm_params"]
|
||||
if "model" not in litellm_params:
|
||||
return None
|
||||
upstream: Final = litellm_params["model"]
|
||||
declared: Final = declared_provider(upstream, litellm_params.get("custom_llm_provider"))
|
||||
provider: Final = declared or inferred_provider(upstream)
|
||||
if provider is None:
|
||||
return upstream
|
||||
return upstream.removeprefix(f"{provider}/") if provider in VERTEX_LIVE_PROVIDERS else None
|
||||
return (upstream if provider is None else upstream.removeprefix(f"{provider}/"), provider)
|
||||
|
||||
|
||||
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 a
|
||||
group another provider serves grants nothing on the proxy's Vertex credentials
|
||||
"""
|
||||
resolved: Final = _upstream_model_and_provider(deployment)
|
||||
if resolved is None:
|
||||
return None
|
||||
upstream, provider = resolved
|
||||
return upstream if provider is None or provider in VERTEX_LIVE_PROVIDERS else None
|
||||
|
||||
|
||||
def _vertex_deployments(llm_router: Router | None) -> tuple[DeploymentTypedDict, ...]:
|
||||
|
|
@ -3955,33 +3962,35 @@ def _vertex_deployments(llm_router: Router | None) -> tuple[DeploymentTypedDict,
|
|||
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:
|
||||
def _live_deployment(setup_model: str, llm_router: Router | None) -> DeploymentTypedDict | None:
|
||||
"""
|
||||
The Vertex AI deployment whose group a ``setup`` frame's model names, whichever addressing the client used.
|
||||
The deployment whose group a ``setup`` frame's model names, whichever addressing the client used.
|
||||
|
||||
The Live SDK wraps whatever the caller typed as ``models/<name>``, 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
|
||||
"""
|
||||
deployments: Final = _vertex_deployments(llm_router)
|
||||
if setup_model.startswith("projects/") or llm_router is None:
|
||||
return None
|
||||
deployments: Final = tuple(llm_router.get_model_list() or ())
|
||||
exact: Final = _deployment_named(setup_model, deployments)
|
||||
return exact if exact is not None else _deployment_named(setup_model.rsplit("/", 1)[-1], deployments)
|
||||
|
||||
|
||||
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
|
||||
through, else every group serving the publisher model Vertex runs, and that model's id bare or under LiteLLM's
|
||||
``vertex_ai/`` prefix unless a group holds that name for another model.
|
||||
The names a key can hold for the model a ``setup`` frame runs: the Vertex AI group the setup rewriter sends it
|
||||
through, else every Vertex AI group serving the publisher model Vertex runs, and that model's id bare or under
|
||||
LiteLLM's ``vertex_ai/`` prefix unless a group holds that name for another model.
|
||||
|
||||
A full resource path skips the rewriter and reaches Vertex as sent, so only its publisher model counts, and
|
||||
a wildcard is matched against that id alone so addressing in front of it cannot satisfy the pattern
|
||||
"""
|
||||
deployment: Final = (
|
||||
None if setup_model.startswith("projects/") else _vertex_live_deployment(setup_model, llm_router)
|
||||
)
|
||||
if deployment is not None:
|
||||
return (deployment["model_name"],)
|
||||
publisher_model: Final = _vertex_publisher_model_suffix(setup_model)
|
||||
deployment: Final = _live_deployment(setup_model, llm_router)
|
||||
if deployment is not None and _vertex_upstream_model(deployment) is not None:
|
||||
alias_target: Final = None if llm_router is None else llm_router.get_model_from_alias(deployment["model_name"])
|
||||
return (alias_target or deployment["model_name"],)
|
||||
opened_model: Final = _resolve_alias_to_upstream_model(setup_model, llm_router)
|
||||
publisher_model: Final = _vertex_publisher_model_suffix(opened_model)
|
||||
serving_groups: Final = tuple(
|
||||
dict.fromkeys(
|
||||
candidate["model_name"]
|
||||
|
|
@ -3992,7 +4001,7 @@ def _vertex_live_model_names(setup_model: str, llm_router: Router | None) -> tup
|
|||
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]
|
||||
publisher_model_id: Final = opened_model.rsplit("/", 1)[-1]
|
||||
return (
|
||||
*serving_groups,
|
||||
*(name for name in (publisher_model_id, f"vertex_ai/{publisher_model_id}") if name not in group_names),
|
||||
|
|
@ -4008,6 +4017,19 @@ def _is_model_alias(name: str, valid_token: UserAPIKeyAuth, llm_router: Router |
|
|||
)
|
||||
|
||||
|
||||
def _as_model_denial(refusal: ProxyException | HTTPException) -> ProxyException:
|
||||
match refusal:
|
||||
case ProxyException():
|
||||
return refusal
|
||||
case HTTPException():
|
||||
return ProxyException(
|
||||
message=str(refusal.detail),
|
||||
type=ProxyErrorTypes.auth_error,
|
||||
param="model",
|
||||
code=refusal.status_code,
|
||||
)
|
||||
|
||||
|
||||
async def _model_access_denial(
|
||||
model: str, valid_token: UserAPIKeyAuth, llm_router: Router | None
|
||||
) -> ProxyException | None:
|
||||
|
|
@ -4022,23 +4044,47 @@ async def _model_access_denial(
|
|||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
except ProxyException as denial:
|
||||
return denial
|
||||
except HTTPException as refusal:
|
||||
return ProxyException(
|
||||
message=str(refusal.detail),
|
||||
type=ProxyErrorTypes.auth_error,
|
||||
param="model",
|
||||
code=refusal.status_code,
|
||||
)
|
||||
except (ProxyException, HTTPException) as refusal:
|
||||
return _as_model_denial(refusal)
|
||||
return None
|
||||
|
||||
|
||||
VertexLiveModelAuthorizer: TypeAlias = Callable[[str], Awaitable[UserAPIKeyAuth]]
|
||||
|
||||
|
||||
async def _model_reauthorization_denial(
|
||||
model: str, authorize_model: VertexLiveModelAuthorizer
|
||||
) -> ProxyException | None:
|
||||
from litellm.proxy.spend_tracking.budget_reservation import release_or_invalidate_budget_reservation
|
||||
|
||||
try:
|
||||
model_auth: Final = await authorize_model(model)
|
||||
except (ProxyException, HTTPException) as refusal:
|
||||
return _as_model_denial(refusal)
|
||||
await release_or_invalidate_budget_reservation(budget_reservation=model_auth.budget_reservation)
|
||||
return None
|
||||
|
||||
|
||||
async def _first_reauthorization_denial(
|
||||
model_names: tuple[str, ...], authorize_model: VertexLiveModelAuthorizer
|
||||
) -> ProxyException | None:
|
||||
denial: Final = await _model_reauthorization_denial(model_names[0], authorize_model)
|
||||
if denial is None or len(model_names) == 1:
|
||||
return denial
|
||||
later_denial: Final = await _first_reauthorization_denial(model_names[1:], authorize_model)
|
||||
return None if later_denial is None else denial
|
||||
|
||||
|
||||
async def _vertex_live_model_denial(
|
||||
setup_model: str, valid_token: UserAPIKeyAuth, llm_router: Router | None
|
||||
setup_model: str,
|
||||
valid_token: UserAPIKeyAuth,
|
||||
llm_router: Router | None,
|
||||
authorize_model: VertexLiveModelAuthorizer,
|
||||
) -> ProxyException | None:
|
||||
"""
|
||||
The setup rewriter forwards a name as sent, so a name an alias maps to another model grants nothing here
|
||||
The setup rewriter forwards an alias name as sent unless it is a listed router alias, which already resolved to
|
||||
its target group, so an alias name grants nothing here. A name the key may call is then authorized the way a
|
||||
request naming it at connect time is
|
||||
"""
|
||||
model_names: Final = tuple(
|
||||
name
|
||||
|
|
@ -4058,7 +4104,10 @@ async def _vertex_live_model_denial(
|
|||
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]
|
||||
callable_names: Final = tuple(name for name, denial in zip(model_names, denials, strict=True) if denial is None)
|
||||
if not callable_names:
|
||||
return denials[0]
|
||||
return await _first_reauthorization_denial(callable_names, authorize_model)
|
||||
|
||||
|
||||
async def _caller_has_model_limits(valid_token: UserAPIKeyAuth) -> bool:
|
||||
|
|
@ -4073,11 +4122,14 @@ async def _caller_has_model_limits(valid_token: UserAPIKeyAuth) -> bool:
|
|||
|
||||
|
||||
async def _vertex_live_setup_denial(
|
||||
setup_model: str | _UnreadableVertexLiveFrame, valid_token: UserAPIKeyAuth, llm_router: Router | None
|
||||
setup_model: str | _UnreadableVertexLiveFrame,
|
||||
valid_token: UserAPIKeyAuth,
|
||||
llm_router: Router | None,
|
||||
authorize_model: VertexLiveModelAuthorizer,
|
||||
) -> ProxyException | None:
|
||||
match setup_model:
|
||||
case str():
|
||||
denial: Final = await _vertex_live_model_denial(setup_model, valid_token, llm_router)
|
||||
denial: Final = await _vertex_live_model_denial(setup_model, valid_token, llm_router, authorize_model)
|
||||
if denial is not None:
|
||||
log_model_access_denial(denial)
|
||||
return denial
|
||||
|
|
@ -4094,7 +4146,9 @@ async def _vertex_live_setup_denial(
|
|||
)
|
||||
|
||||
|
||||
def _build_vertex_live_client_frame_gate(llm_router: Router | None) -> WebsocketClientFrameGate:
|
||||
def _build_vertex_live_client_frame_gate(
|
||||
llm_router: Router | None, authorize_model: VertexLiveModelAuthorizer
|
||||
) -> WebsocketClientFrameGate:
|
||||
"""
|
||||
Apply the checks a model named at auth time gets to the model a client's ``setup`` frame names.
|
||||
|
||||
|
|
@ -4111,7 +4165,7 @@ def _build_vertex_live_client_frame_gate(llm_router: Router | None) -> Websocket
|
|||
if model is None:
|
||||
return None
|
||||
try:
|
||||
return await _vertex_live_setup_denial(model, valid_token, llm_router)
|
||||
return await _vertex_live_setup_denial(model, valid_token, llm_router, authorize_model)
|
||||
except Exception: # noqa: BLE001 # fail closed: a gate that cannot answer refuses, and says so in band
|
||||
verbose_proxy_logger.exception("Vertex AI Live passthrough: model access check failed for %s", model)
|
||||
return ProxyException(
|
||||
|
|
@ -4182,9 +4236,9 @@ def _resolve_alias_to_upstream_model(setup_model: str, llm_router: Router | None
|
|||
"""
|
||||
The Live SDK wraps whatever the caller typed as ``models/<name>``, so a gateway alias arrives prefixed
|
||||
"""
|
||||
deployment: Final = _vertex_live_deployment(setup_model, llm_router)
|
||||
upstream: Final = None if deployment is None else _vertex_upstream_model(deployment)
|
||||
return setup_model if upstream is None else upstream
|
||||
deployment: Final = _live_deployment(setup_model, llm_router)
|
||||
resolved: Final = None if deployment is None else _upstream_model_and_provider(deployment)
|
||||
return setup_model if resolved is None else resolved[0]
|
||||
|
||||
|
||||
async def vertex_ai_live_websocket_passthrough(
|
||||
|
|
@ -4193,6 +4247,7 @@ async def vertex_ai_live_websocket_passthrough(
|
|||
vertex_project: str | None = None,
|
||||
vertex_location: str | None = None,
|
||||
user_api_key_dict: UserAPIKeyAuth | None = None,
|
||||
authorize_model: VertexLiveModelAuthorizer | None = None,
|
||||
):
|
||||
"""
|
||||
Vertex AI Live API WebSocket Pass-through Function
|
||||
|
|
@ -4291,7 +4346,10 @@ async def vertex_ai_live_websocket_passthrough(
|
|||
vertex_location=resolved_location,
|
||||
llm_router=_get_llm_router(),
|
||||
),
|
||||
client_frame_gate=_build_vertex_live_client_frame_gate(llm_router=_get_llm_router()),
|
||||
client_frame_gate=_build_vertex_live_client_frame_gate(
|
||||
llm_router=_get_llm_router(),
|
||||
authorize_model=authorize_model or partial(user_api_key_auth_accepted_websocket_for_model, websocket),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -2456,9 +2456,6 @@ def _truncated_close_reason(reason: str) -> str:
|
|||
|
||||
|
||||
def _refusal_close_reason(message: str) -> str:
|
||||
"""
|
||||
The proxy's own refusal message as a close reason, ending on the last whole sentence that fits
|
||||
"""
|
||||
clipped: Final = _truncated_close_reason(message)
|
||||
sentence_end: Final = clipped.rfind(". ")
|
||||
return clipped if clipped == message or sentence_end == -1 else clipped[: sentence_end + 1]
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ import litellm
|
|||
from litellm import CustomLLM
|
||||
from litellm.litellm_core_utils import get_llm_provider_logic
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import (
|
||||
declared_provider,
|
||||
get_llm_provider,
|
||||
resolve_model_and_provider_for_metadata,
|
||||
get_provider_affinity_header,
|
||||
|
|
@ -70,6 +71,26 @@ def test_is_registered_custom_provider(registered_custom_provider: str, candidat
|
|||
assert is_registered_custom_provider(candidate) is expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model", "custom_llm_provider", "expected"),
|
||||
[
|
||||
("vertex_ai/gemini-live-2.5-flash", None, "vertex_ai"),
|
||||
("llama-4-scout", "groq", "groq"),
|
||||
("gemini/gemini-2.0-flash-live-001", "vertex_ai", "vertex_ai"),
|
||||
("veniceai/venice-uncensored", None, "veniceai"),
|
||||
(f"{CUSTOM_PROVIDER}/my-model", None, CUSTOM_PROVIDER),
|
||||
("gemini-live-2.5-flash", None, None),
|
||||
("models/gemini-live-2.5-flash", None, None),
|
||||
],
|
||||
ids=["provider-prefix", "custom-llm-provider", "custom-llm-provider-wins", "json-provider", "custom-provider",
|
||||
"bare-name", "unknown-prefix"],
|
||||
)
|
||||
def test_declared_provider_names_only_a_provider_the_pair_spells_out(
|
||||
registered_custom_provider: str, model: str, custom_llm_provider: str | None, expected: str | None
|
||||
) -> None:
|
||||
assert declared_provider(model, custom_llm_provider) == expected
|
||||
|
||||
|
||||
def test_get_llm_provider_leaves_fal_ai_api_base_unset_for_global_fallback() -> None:
|
||||
_, provider, _, api_base = get_llm_provider(model="fal_ai/fal-ai/flux/schnell")
|
||||
assert provider == "fal_ai"
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
# What is this?
|
||||
## Tests if 'get_end_user_object' works as expected
|
||||
|
||||
import sys, os, asyncio, time, random, uuid
|
||||
import sys, os, asyncio, time, random, re, uuid
|
||||
import hashlib
|
||||
import json
|
||||
import traceback
|
||||
|
|
@ -39,6 +39,7 @@ from litellm.proxy.utils import PrismaClient
|
|||
import litellm.proxy.common_utils.auth_cache_invalidation_pubsub as pubsub_module
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.proxy.auth.auth_checks import (
|
||||
can_caller_call_resolved_model,
|
||||
can_team_access_model,
|
||||
delete_cache_key_objects,
|
||||
is_model_cost_zero,
|
||||
|
|
@ -2425,3 +2426,34 @@ async def test_bulk_eviction_clears_the_redis_copies_before_waiting_on_a_wedged_
|
|||
|
||||
assert refilled == []
|
||||
assert redis.entries == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(("team_id", "refused"), [(None, True), ("live-team", False)], ids=["personal-key", "team-key"])
|
||||
async def test_can_caller_call_resolved_model_holds_a_personal_key_to_its_owners_models(
|
||||
team_id: str | None, refused: bool
|
||||
) -> None:
|
||||
owner_id: Final = f"live-owner-{uuid.uuid4()}"
|
||||
cache: Final = UserApiKeyCache()
|
||||
await cache.async_set_cache(key=owner_id, value=LiteLLM_UserTable(user_id=owner_id, models=["gemini-3.8-live"]))
|
||||
router: Final = litellm.Router(
|
||||
model_list=[
|
||||
{"model_name": "gemini-3.8-live", "litellm_params": {"model": "vertex_ai/gemini-3.8-live"}},
|
||||
{"model_name": "gemini-live-2.5-flash", "litellm_params": {"model": "vertex_ai/gemini-live-2.5-flash"}},
|
||||
]
|
||||
)
|
||||
|
||||
call: Final = can_caller_call_resolved_model(
|
||||
model="gemini-live-2.5-flash",
|
||||
valid_token=UserAPIKeyAuth(api_key="sk-live", user_id=owner_id, team_id=team_id),
|
||||
llm_router=router,
|
||||
prisma_client=MagicMock(),
|
||||
user_api_key_cache=cache,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
if refused:
|
||||
with pytest.raises(ProxyException, match=re.escape("gemini-live-2.5-flash")):
|
||||
await call
|
||||
else:
|
||||
await call
|
||||
|
|
|
|||
|
|
@ -7550,10 +7550,18 @@ def _vertex_live_router() -> litellm.Router:
|
|||
)
|
||||
|
||||
|
||||
def _reauthorized_as(valid_token: UserAPIKeyAuth) -> Callable[[str], Awaitable[UserAPIKeyAuth]]:
|
||||
async def authorize_model(model: str) -> UserAPIKeyAuth:
|
||||
return valid_token
|
||||
|
||||
return authorize_model
|
||||
|
||||
|
||||
async def _run_vertex_live_gated_passthrough(
|
||||
client_frames: Sequence[Mapping[str, str | bytes]],
|
||||
valid_token: UserAPIKeyAuth,
|
||||
llm_router: litellm.Router | None = None,
|
||||
authorize_model: Callable[[str], Awaitable[UserAPIKeyAuth]] | None = None,
|
||||
) -> tuple[MagicMock, RecordingUpstreamWebSocket]:
|
||||
"""Drive /vertex_ai/live with the real frame gate, returning the client and upstream sockets"""
|
||||
upstream_ws = RecordingUpstreamWebSocket()
|
||||
|
|
@ -7562,6 +7570,7 @@ async def _run_vertex_live_gated_passthrough(
|
|||
upstream_ws,
|
||||
valid_token,
|
||||
llm_router or _vertex_live_router(),
|
||||
authorize_model,
|
||||
)
|
||||
return websocket, upstream_ws
|
||||
|
||||
|
|
@ -7571,6 +7580,7 @@ async def _run_vertex_live_gated_session(
|
|||
upstream_ws: "RecordingUpstreamWebSocket | ScriptedUpstreamWebSocket",
|
||||
valid_token: UserAPIKeyAuth,
|
||||
llm_router: litellm.Router,
|
||||
authorize_model: Callable[[str], Awaitable[UserAPIKeyAuth]] | None = None,
|
||||
) -> tuple[MagicMock, SimpleNamespace]:
|
||||
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
|
||||
_build_vertex_live_client_frame_gate,
|
||||
|
|
@ -7591,7 +7601,10 @@ async def _run_vertex_live_gated_session(
|
|||
vertex_location="global",
|
||||
llm_router=llm_router,
|
||||
),
|
||||
client_frame_gate=_build_vertex_live_client_frame_gate(llm_router=llm_router),
|
||||
client_frame_gate=_build_vertex_live_client_frame_gate(
|
||||
llm_router=llm_router,
|
||||
authorize_model=authorize_model or _reauthorized_as(valid_token),
|
||||
),
|
||||
)
|
||||
|
||||
return websocket, environment
|
||||
|
|
@ -8173,23 +8186,26 @@ async def test_vertex_live_setup_frame_naming_an_alias_cannot_open_the_model_the
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vertex_live_setup_frame_naming_a_router_alias_cannot_open_the_model_the_alias_shadows():
|
||||
@pytest.mark.parametrize("hidden", [False, True], ids=["listed-alias", "hidden-alias"])
|
||||
async def test_vertex_live_setup_frame_naming_a_router_alias_cannot_open_the_model_the_alias_shadows(hidden):
|
||||
"""A listed router alias runs as its target group, while the setup rewriter cannot see a hidden one and sends it
|
||||
to Vertex as the model whose id it reuses, which a key holding only the target must not open"""
|
||||
llm_router = litellm.Router(
|
||||
model_list=[
|
||||
{"model_name": "gemini-live-2.5-flash", "litellm_params": {"model": "vertex_ai/gemini-live-2.5-flash"}},
|
||||
],
|
||||
model_group_alias={"gemini-3-pro-live": "gemini-live-2.5-flash"},
|
||||
model_group_alias={"gemini-3-pro-live": {"model": "gemini-live-2.5-flash", "hidden": hidden}},
|
||||
)
|
||||
|
||||
websocket, upstream_ws = await _run_vertex_live_gated_passthrough(
|
||||
_, upstream_ws = await _run_vertex_live_gated_passthrough(
|
||||
[_setup_frame("gemini-3-pro-live")],
|
||||
UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"]),
|
||||
llm_router=llm_router,
|
||||
)
|
||||
|
||||
upstream_ws.send.assert_not_awaited()
|
||||
reason = _policy_close_reason(websocket)
|
||||
assert reason is not None and "gemini-3-pro-live" in reason
|
||||
sent_models = [json.loads(call.args[0])["setup"]["model"] for call in upstream_ws.send.await_args_list]
|
||||
assert not any(model.endswith("/gemini-3-pro-live") for model in sent_models)
|
||||
assert sent_models == ([] if hidden else ["projects/proj-db/locations/global/publishers/google/models/gemini-live-2.5-flash"])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -8235,6 +8251,69 @@ async def test_vertex_live_setup_frame_is_checked_against_the_team_in_every_addr
|
|||
assert reason is not None and "gemini-3-pro-live" in reason
|
||||
|
||||
|
||||
def _another_providers_live_router() -> litellm.Router:
|
||||
return litellm.Router(
|
||||
model_list=[
|
||||
{"model_name": "studio-live", "litellm_params": {"model": "gemini/gemini-2.0-flash-live-001"}},
|
||||
{"model_name": "vertex-flash-live", "litellm_params": {"model": "vertex_ai/gemini-2.0-flash-live-001"}},
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vertex_live_setup_frame_naming_another_providers_group_runs_that_groups_model_on_vertex():
|
||||
"""A caller who may call any model kept reaching Vertex through a group another provider serves before the
|
||||
gate existed, so the setup rewriter still sends that group's model id rather than the group name"""
|
||||
_, upstream_ws = await _run_vertex_live_gated_passthrough(
|
||||
[_setup_frame("models/studio-live")],
|
||||
UserAPIKeyAuth(token="hashed"),
|
||||
llm_router=_another_providers_live_router(),
|
||||
)
|
||||
|
||||
sent_setup = json.loads(upstream_ws.send.await_args.args[0])["setup"]
|
||||
assert sent_setup["model"] == "projects/proj-db/locations/global/publishers/google/models/gemini-2.0-flash-live-001"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("models", "reaches_vertex"),
|
||||
[(["studio-live"], False), (["vertex-flash-live"], True), (["vertex_ai/gemini-2.0-flash-live-001"], True)],
|
||||
ids=["other-provider-group", "vertex-group-serving-it", "vertex-model-id"],
|
||||
)
|
||||
async def test_vertex_live_setup_frame_naming_another_providers_group_is_authorized_by_the_vertex_model_it_runs(
|
||||
models, reaches_vertex
|
||||
):
|
||||
websocket, upstream_ws = await _run_vertex_live_gated_passthrough(
|
||||
[_setup_frame("studio-live")],
|
||||
UserAPIKeyAuth(token="hashed", models=models),
|
||||
llm_router=_another_providers_live_router(),
|
||||
)
|
||||
|
||||
assert upstream_ws.send.await_count == (1 if reaches_vertex else 0)
|
||||
reason = _policy_close_reason(websocket)
|
||||
assert (reason is None) is reaches_vertex
|
||||
assert reaches_vertex or ("vertex-flash-live" in reason and "vertex_ai/studio-live" not in reason)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vertex_live_setup_frame_naming_a_router_alias_is_authorized_by_its_target_group():
|
||||
"""The setup rewriter runs a router alias as its target group, so the key holding that group may open it"""
|
||||
llm_router = litellm.Router(
|
||||
model_list=[{"model_name": "vertex-live", "litellm_params": {"model": f"vertex_ai/{VERTEX_LIVE_GROUP}"}}],
|
||||
model_group_alias={"gemini-live-alias": "vertex-live"},
|
||||
)
|
||||
|
||||
websocket, upstream_ws = await _run_vertex_live_gated_passthrough(
|
||||
[_setup_frame("models/gemini-live-alias")],
|
||||
UserAPIKeyAuth(token="hashed", models=["vertex-live"]),
|
||||
llm_router=llm_router,
|
||||
)
|
||||
|
||||
assert _policy_close_reason(websocket) is None
|
||||
sent_setup = json.loads(upstream_ws.send.await_args.args[0])["setup"]
|
||||
assert sent_setup["model"] == f"projects/proj-db/locations/global/publishers/google/models/{VERTEX_LIVE_GROUP}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("addressing", VERTEX_LIVE_CLIENT_ADDRESSINGS)
|
||||
async def test_vertex_live_setup_frame_resolves_an_access_group_grant_in_every_addressing(addressing):
|
||||
|
|
@ -8255,16 +8334,19 @@ async def test_vertex_live_setup_frame_resolves_an_access_group_grant_in_every_a
|
|||
},
|
||||
]
|
||||
)
|
||||
gate = _build_vertex_live_client_frame_gate(llm_router=llm_router)
|
||||
token = UserAPIKeyAuth(token="hashed", models=["live-team"])
|
||||
gate = _build_vertex_live_client_frame_gate(llm_router=llm_router, authorize_model=_reauthorized_as(token))
|
||||
|
||||
frame = json.dumps({"setup": {"model": addressing.format(model=VERTEX_LIVE_GROUP)}})
|
||||
denial = await gate(frame, UserAPIKeyAuth(token="hashed", models=["live-team"]))
|
||||
denial = await gate(frame, token)
|
||||
|
||||
assert denial is None
|
||||
|
||||
|
||||
async def _run_vertex_live_route(
|
||||
client_frames: Sequence[Mapping[str, str | bytes]], valid_token: UserAPIKeyAuth
|
||||
client_frames: Sequence[Mapping[str, str | bytes]],
|
||||
valid_token: UserAPIKeyAuth,
|
||||
reauthorize_through_user_api_key_auth: bool = False,
|
||||
) -> tuple[MagicMock, RecordingUpstreamWebSocket, SimpleNamespace]:
|
||||
"""Drive the /vertex_ai/live route, so the gate under test is whichever one the route itself wires"""
|
||||
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
|
||||
|
|
@ -8276,6 +8358,12 @@ async def _run_vertex_live_route(
|
|||
websocket = _client_websocket(
|
||||
AsyncMock(side_effect=[*client_frames, {"type": "websocket.disconnect"}]),
|
||||
)
|
||||
websocket.headers = {"authorization": "Bearer sk-live-client"}
|
||||
websocket.scope = {
|
||||
"type": "websocket",
|
||||
"path": "/vertex_ai/live",
|
||||
"headers": [(b"authorization", b"Bearer sk-live-client")],
|
||||
}
|
||||
|
||||
with (
|
||||
_patched_websocket_passthrough_environment(upstream_ws) as environment,
|
||||
|
|
@ -8291,6 +8379,7 @@ async def _run_vertex_live_route(
|
|||
vertex_project="proj-db",
|
||||
vertex_location="global",
|
||||
user_api_key_dict=valid_token,
|
||||
authorize_model=None if reauthorize_through_user_api_key_auth else _reauthorized_as(valid_token),
|
||||
)
|
||||
|
||||
return websocket, upstream_ws, environment
|
||||
|
|
@ -8315,6 +8404,34 @@ async def test_vertex_live_route_gates_the_setup_frame_even_when_the_query_model
|
|||
assert refusal.type == ProxyErrorTypes.key_model_access_denied
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vertex_live_route_authorizes_the_setup_model_the_way_connect_time_auth_checks_a_query_model():
|
||||
"""
|
||||
Connect-time auth applies JWT scopes, per-model budgets and the unpriced-model block to ``?model=``, so a model
|
||||
the key's list allows still has to pass them when the setup frame is the first place it is named
|
||||
"""
|
||||
budget_denial = ProxyException(
|
||||
message="Budget has been exceeded for model gemini-live-2.5-flash.",
|
||||
type=ProxyErrorTypes.budget_exceeded,
|
||||
param="model",
|
||||
code=400,
|
||||
)
|
||||
with patch(
|
||||
"litellm.proxy.auth.user_api_key_auth.user_api_key_auth", autospec=True, side_effect=budget_denial
|
||||
) as reauthorization:
|
||||
websocket, upstream_ws, _ = await _run_vertex_live_route(
|
||||
[_setup_frame("models/gemini-live-2.5-flash")],
|
||||
UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"]),
|
||||
reauthorize_through_user_api_key_auth=True,
|
||||
)
|
||||
|
||||
upstream_ws.send.assert_not_awaited()
|
||||
assert _policy_close_reason(websocket) == budget_denial.message
|
||||
reauthorized_request = reauthorization.call_args.kwargs["request"]
|
||||
assert json.loads(await reauthorized_request.body()) == {"model": "gemini-live-2.5-flash"}
|
||||
assert reauthorization.call_args.kwargs["api_key"] == "Bearer sk-live-client"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vertex_live_route_keeps_the_key_allowlist_out_of_the_close_reason_and_in_the_proxy_log(caplog):
|
||||
"""The client learns which model was refused, while only the proxy log names the models the key may reach"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue