Merge 3aa4e0b7c0 into 298911d61b

litellm/litellm_core_utils/get_llm_provider_logic.py

@@ -158,6 +158,22 @@ def declared_authenticating_provider(model: str | None, custom_llm_provider: str
     return declared if declared in PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO else None
 
 
+def declared_provider(model: str, custom_llm_provider: str | None = None) -> str | None:
+    """The provider this pair names outright, or None when only get_llm_provider can tell.
+
+    get_llm_provider reads a named provider's API key and base on the way to returning it, and a hosted
+    secret manager answers each read with a network call, so a caller that needs only the name reads it here.
+    """
+    if custom_llm_provider:
+        return custom_llm_provider
+    prefix, slash, _ = model.partition("/")
+    if not slash:
+        return None
+    if prefix in litellm.provider_list or JSONProviderRegistry.exists(prefix) or is_registered_custom_provider(prefix):
+        return prefix
+    return None
+
+
 def resolve_model_and_provider_for_metadata(model: str, custom_llm_provider: str | None = None) -> tuple[str, str]:
     declared: Final = declared_authenticating_provider(model, custom_llm_provider)
     if declared is not None:

litellm/proxy/auth/auth_checks.py

@@ -5744,6 +5744,45 @@ async def can_key_call_resolved_model(
             )
 
 
+async def can_caller_call_resolved_model(
+    model: str,
+    valid_token: UserAPIKeyAuth,
+    llm_router: litellm.Router | None,
+    prisma_client: PrismaClient | None,
+    user_api_key_cache: UserApiKeyCache,
+    proxy_logging_obj: ProxyLogging,
+) -> None:
+    """``can_key_call_resolved_model`` plus the agent grant and a personal key owner's own models, which
+    ``common_checks`` applies to a model named at auth time and a model named after auth needs as well"""
+    await can_key_call_resolved_model(
+        model=model,
+        llm_model_list=None if llm_router is None else llm_router.registry().deployments(),
+        valid_token=valid_token,
+        llm_router=llm_router,
+    )
+    await can_agent_call_model(
+        model=model,
+        valid_token=valid_token,
+        llm_router=llm_router,
+        prisma_client=prisma_client,
+        user_api_key_cache=user_api_key_cache,
+        proxy_logging_obj=proxy_logging_obj,
+    )
+    if valid_token.team_id is not None or prisma_client is None:
+        return
+    await can_user_call_model(
+        model=model,
+        llm_router=llm_router,
+        user_object=await _user_object_or_none(
+            valid_token=valid_token,
+            prisma_client=prisma_client,
+            user_api_key_cache=user_api_key_cache,
+            proxy_logging_obj=proxy_logging_obj,
+        ),
+        key_model_aliases=key_model_aliases_for_auth_check(valid_token),
+    )
+
+
 def can_org_access_model(
     model: str,
     org_object: LiteLLM_OrganizationTable | None,

litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py

@@ -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
@@ -52,6 +52,7 @@ from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException
 from litellm.integrations.custom_guardrail import guardrail_request_data_with_streaming, is_guardrail_intervention
 from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
 from litellm.litellm_core_utils.fallback_generalizations import match_routing_generalization
+from litellm.litellm_core_utils.get_llm_provider_logic import declared_provider, inferred_provider
 from litellm.llms.anthropic.common_utils import AnthropicModelInfo, merge_anthropic_beta_headers
 from litellm.llms.azure.passthrough.transformation import (
     foreign_azure_deployment,
@@ -76,14 +77,23 @@ from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
 from litellm.llms.vllm.common_utils import endpoint_under_api_base
 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_checks import (
+    can_caller_call_resolved_model,
+    enforced_model_allowlists,
+    key_model_aliases_for_auth_check,
+    personal_key_team_model_access_applies,
+    team_model_aliases_for_auth_check,
+)
+from litellm.proxy.auth.auth_utils import log_model_access_denial
 from litellm.proxy.auth.handle_jwt import JWTHandler
+from litellm.proxy.auth.model_access_denied import model_access_denied_client_message
 from litellm.proxy.auth.route_checks import RouteChecks
 from litellm.proxy.auth.user_api_key_auth import (  # noqa: F401  # legacy module exports
     _get_bearer_token,  # pyright: ignore[reportPrivateUsage,reportUnusedImport]  # backwards-compatible package export
     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,
 )
@@ -121,7 +131,10 @@ from litellm.proxy.pass_through_endpoints.injected_stream_usage import (
     without_injected_usage_frames,
 )
 from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
+    ClientFrameRefusal,
     HttpPassThroughEndpointHelpers,
+    LoggedClientFrameRefusal,
+    WebsocketClientFrameGate,
     YieldsToConfiguredPassThroughRoute,
     create_pass_through_route,
     create_websocket_passthrough_route,
@@ -159,6 +172,7 @@ from .passthrough_endpoint_router import PassthroughEndpointRouter
 if TYPE_CHECKING:
     from litellm.proxy.proxy_server import ProxyConfig as _ProxyConfig
     from litellm.router import Router
+    from litellm.types.router import DeploymentTypedDict
 
     ProxyConfig = _ProxyConfig  # rebind-ok: conditional type alias
 else:
@@ -3940,6 +3954,12 @@ 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_UNVERIFIED_MODEL_CLOSE_REASON: Final = (
+    "Could not check your key's access to the model this setup frame names; ask a proxy admin to read the logs"
+)
+
+VERTEX_LIVE_UNREADABLE_SETUP_CLOSE_REASON: Final = "Send setup frames as strict JSON so the proxy can check their model"
+
 VERTEX_PUBLISHER_MODEL_PREFIX: Final = "publishers/google/models/"
 
 VERTEX_PUBLISHERS_SEGMENT: Final = "publishers/"
@@ -3964,6 +3984,339 @@ def _get_llm_router() -> Router | None:
     return llm_router
 
 
+@dataclass(frozen=True, slots=True)
+class _UnreadableVertexLiveFrame:
+    pass
+
+
+class _VertexLiveSetup(TypedDict, total=False):
+    model: ReadOnly[str]
+
+
+class _VertexLiveClientFrame(TypedDict, total=False):
+    setup: ReadOnly[_VertexLiveSetup]
+
+
+_VERTEX_LIVE_CLIENT_FRAME: Final = TypeAdapter(_VertexLiveClientFrame)
+
+
+def _may_carry_vertex_live_setup(frame_data: str | bytes) -> bool:
+    """
+    Whether a frame is worth parsing: audio and video frames stream continuously and are large, and one that
+    cannot spell ``setup`` cannot carry one. Vertex also reads single-quoted and bare keys, so the bare word is
+    the test, and a backslash means the key may be escaped, which only a parse settles
+    """
+    if isinstance(frame_data, bytes):
+        return b"setup" in frame_data or b"\\" in frame_data
+    return "setup" in frame_data or "\\" in frame_data
+
+
+def _object_with_unique_keys(pairs: Sequence[tuple[str, object]]) -> Mapping[str, object]:
+    unique: Final = MappingProxyType(dict(pairs))
+    if len(unique) < len(pairs):
+        raise ValueError("a JSON object repeats a key")
+    return unique
+
+
+def _vertex_live_setup_model(frame_data: str | bytes) -> str | _UnreadableVertexLiveFrame | None:
+    """
+    The model a Vertex AI Live ``setup`` frame names, which is the only place BidiGenerateContent carries one.
+
+    Vertex parses trailing commas, single quotes and bare keys that ``json.loads`` refuses, and keeps the first of
+    two keys when the second is null, so a frame that may carry a setup but fails a strict parse, repeats a key,
+    or holds a setup of the wrong shape is unreadable rather than setup-free
+    """
+    if not _may_carry_vertex_live_setup(frame_data):
+        return None
+    try:
+        frame: Final = _VERTEX_LIVE_CLIENT_FRAME.validate_python(
+            json.loads(frame_data, object_pairs_hook=_object_with_unique_keys)
+        )
+    except (ValueError, RecursionError):
+        return _UnreadableVertexLiveFrame()
+    setup: Final = frame.get("setup")
+    return None if setup is None else setup.get("model") or None
+
+
+def _deployment_named(model_name: str, deployments: Sequence[DeploymentTypedDict]) -> DeploymentTypedDict | None:
+    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 _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)
+    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 _routed_deployments(llm_router: Router) -> tuple[DeploymentTypedDict, ...]:
+    """Every deployment under the name a request reaches it by, which leaves out a group a router alias shares a name
+    with, since the alias wins"""
+    return (
+        *llm_router.registry().deployments(source="alias"),
+        *(
+            deployment
+            for deployment in llm_router.registry().deployments() or ()
+            if deployment["model_name"] not in llm_router.registry().aliases()
+        ),
+    )
+
+
+def _vertex_deployments(llm_router: Router | None) -> tuple[DeploymentTypedDict, ...]:
+    if llm_router is None:
+        return ()
+    return tuple(deployment for deployment in _routed_deployments(llm_router) if _vertex_upstream_model(deployment))
+
+
+def _live_deployment(setup_model: str, llm_router: Router | None) -> DeploymentTypedDict | None:
+    """
+    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
+    """
+    if setup_model.startswith("projects/") or llm_router is None:
+        return None
+    deployments: Final = tuple(llm_router.registry().deployments() 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 Vertex AI group the setup rewriter sends it
+    through, else every Vertex AI group serving the publisher model Vertex runs, and either way that model's id bare
+    or under LiteLLM's ``vertex_ai/`` prefix unless a group holds that name.
+
+    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 = _live_deployment(setup_model, llm_router)
+    if deployment is not None and (upstream := _vertex_upstream_model(deployment)) is not None:
+        return (deployment["model_name"], *_publisher_model_id_names(upstream, llm_router))
+    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"]
+            for candidate in _vertex_deployments(llm_router)
+            if _vertex_publisher_model_suffix(_vertex_upstream_model(candidate) or "") == publisher_model
+        )
+    )
+    return (*serving_groups, *_publisher_model_id_names(opened_model, llm_router))
+
+
+def _publisher_model_id_names(model: str, llm_router: Router | None) -> tuple[str, ...]:
+    group_names: Final = frozenset(
+        candidate["model_name"]
+        for candidate in (() if llm_router is None else llm_router.registry().deployments() or ())
+    )
+    publisher_model_id: Final = model.rsplit("/", 1)[-1]
+    return tuple(name for name in (publisher_model_id, f"vertex_ai/{publisher_model_id}") if name not in group_names)
+
+
+def _is_listed_router_alias(name: str, llm_router: Router) -> bool:
+    return any(deployment["model_name"] == name for deployment in llm_router.registry().deployments(source="alias"))
+
+
+def _is_alias_the_rewriter_ignores(name: str, valid_token: UserAPIKeyAuth, llm_router: Router | None) -> bool:
+    return (
+        name in (key_model_aliases_for_auth_check(valid_token) or ())
+        or name in (team_model_aliases_for_auth_check(valid_token) or ())
+        or name in litellm.model_alias_map
+        or (
+            llm_router is not None
+            and name in llm_router.registry().aliases()
+            and not _is_listed_router_alias(name, llm_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:
+    from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
+
+    try:
+        await can_caller_call_resolved_model(
+            model=model,
+            valid_token=valid_token,
+            llm_router=llm_router,
+            prisma_client=prisma_client,
+            user_api_key_cache=user_api_key_cache,
+            proxy_logging_obj=proxy_logging_obj,
+        )
+    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
+) -> LoggedClientFrameRefusal | 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 LoggedClientFrameRefusal(_as_model_denial(refusal))
+    await release_or_invalidate_budget_reservation(budget_reservation=model_auth.budget_reservation)
+    return None
+
+
+async def _first_model_denial(
+    model_names: tuple[str, ...],
+    valid_token: UserAPIKeyAuth,
+    llm_router: Router | None,
+    authorize_model: VertexLiveModelAuthorizer,
+) -> ClientFrameRefusal | None:
+    """
+    Every name runs the same deployment, so the first name the key may call is re-authorized and that answer
+    stands, budget and pricing checks included. Only a key that may call none of them gets the first access denial
+    """
+    name: Final = model_names[0]
+    access_denial: Final = await _model_access_denial(name, valid_token, llm_router)
+    if access_denial is None:
+        return await _model_reauthorization_denial(name, authorize_model)
+    if len(model_names) == 1:
+        return access_denial
+    later_denial: Final = await _first_model_denial(model_names[1:], valid_token, llm_router, authorize_model)
+    return access_denial if isinstance(later_denial, ProxyException) else later_denial
+
+
+async def _vertex_live_model_denial(
+    setup_model: str,
+    valid_token: UserAPIKeyAuth,
+    llm_router: Router | None,
+    authorize_model: VertexLiveModelAuthorizer,
+) -> ClientFrameRefusal | None:
+    """
+    The setup rewriter runs a listed router alias as its target group, the way a request naming it does, and sends
+    any other alias name to Vertex as written, so only a listed router alias grants anything 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
+        for name in _vertex_live_model_names(setup_model, llm_router)
+        if not _is_alias_the_rewriter_ignores(name, valid_token, llm_router)
+    )
+    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
+        )
+    return await _first_model_denial(model_names, valid_token, llm_router, authorize_model)
+
+
+async def _caller_has_model_limits(valid_token: UserAPIKeyAuth) -> bool:
+    return (
+        valid_token.agent_id is not None
+        or valid_token.managed_agent_policy is not None
+        or valid_token.agent_caller is not None
+        or valid_token.end_user_id is not None
+        or personal_key_team_model_access_applies(valid_token)
+        or _has_model_restrictions(await _proxy_model_allowlists()(valid_token))
+    )
+
+
+async def _vertex_live_setup_denial(
+    setup_model: str | _UnreadableVertexLiveFrame,
+    valid_token: UserAPIKeyAuth,
+    llm_router: Router | None,
+    authorize_model: VertexLiveModelAuthorizer,
+) -> ClientFrameRefusal | None:
+    match setup_model:
+        case str():
+            denial: Final = await _vertex_live_model_denial(setup_model, valid_token, llm_router, authorize_model)
+            if isinstance(denial, ProxyException):
+                log_model_access_denial(denial)
+            return denial
+        case _UnreadableVertexLiveFrame():
+            return (
+                ProxyException(
+                    message=VERTEX_LIVE_UNREADABLE_SETUP_CLOSE_REASON,
+                    type=ProxyErrorTypes.bad_request_error,
+                    param="setup",
+                    code=400,
+                )
+                if await _caller_has_model_limits(valid_token)
+                else None
+            )
+
+
+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.
+
+    ``?model=`` is optional on this route and the Live protocol carries the real model in the first client
+    frame, so connect-time auth has no model to check and every frame that names one has to be authorized
+    here instead.
+
+    A frame the proxy cannot read the way Vertex does only reaches Vertex for a caller who may call any model,
+    which is a caller with no model allowlist and no agent, end user or team grant standing in for one
+    """
+
+    async def gate(frame_data: str | bytes, valid_token: UserAPIKeyAuth, /) -> ClientFrameRefusal | None:
+        model: Final = _vertex_live_setup_model(frame_data)
+        if model is None:
+            return None
+        try:
+            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(
+                message=VERTEX_LIVE_UNVERIFIED_MODEL_CLOSE_REASON,
+                type=ProxyErrorTypes.internal_server_error,
+                param="model",
+                code=500,
+            )
+
+    return gate
+
+
 def _resolve_vertex_live_credentials(
     vertex_project: str | None,
     vertex_location: str | None,
@@ -4022,24 +4375,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
     """
-    if llm_router is None:
-        return setup_model
-    candidates: Final = (setup_model, setup_model.rsplit("/", 1)[-1])
-    upstream: Final = next(
-        (
-            deployment["litellm_params"].get("model")
-            for deployment in (llm_router.registry().deployments() or ())
-            if deployment.get("model_name") in candidates
-        ),
-        None,
-    )
-    if upstream is None:
-        return setup_model
-    try:
-        _, provider, _, _ = litellm.get_llm_provider(model=upstream)
-    except litellm.exceptions.BadRequestError:
-        return upstream
-    return upstream.removeprefix(f"{provider}/")
+    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(
@@ -4048,6 +4386,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
@@ -4153,6 +4492,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(),
+            authorize_model=authorize_model or partial(user_api_key_auth_accepted_websocket_for_model, websocket),
+        ),
     )
 
 

litellm/proxy/pass_through_endpoints/pass_through_endpoints.py

@@ -11,7 +11,7 @@ from dataclasses import dataclass
 from datetime import datetime
 from itertools import chain, count, groupby
 from types import MappingProxyType
-from typing import TYPE_CHECKING, Any, Final, TypedDict, cast
+from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias, TypedDict, cast, overload
 from urllib.parse import urlencode, urlparse
 
 import httpx
@@ -34,7 +34,7 @@ from starlette.routing import BaseRoute, Match, Route
 from starlette.types import Scope
 from starlette.websockets import WebSocketState
 from typing_extensions import ReadOnly
-from websockets.asyncio.client import connect
+from websockets.asyncio.client import ClientConnection, connect
 from websockets.exceptions import (
     ConnectionClosedError,
     ConnectionClosedOK,
@@ -2492,28 +2492,49 @@ def create_websocket_passthrough_route(
     return websocket_endpoint_func
 
 
-def _rewrite_vertex_live_setup_model(text_data: str, setup_model_rewriter: Callable[[str], str] | None) -> str:
+@overload
+def _rewrite_vertex_live_setup_model(frame_data: str, setup_model_rewriter: Callable[[str], str] | None) -> str: ...
+
+
+@overload
+def _rewrite_vertex_live_setup_model(frame_data: bytes, setup_model_rewriter: Callable[[str], str] | None) -> bytes: ...
+
+
+def _rewrite_vertex_live_setup_model(
+    frame_data: str | bytes, setup_model_rewriter: Callable[[str], str] | None
+) -> str | bytes:
     """
     Rewrite the model of a Vertex AI Live ``setup`` frame, leaving every other frame byte-identical
     """
     if setup_model_rewriter is None:
-        return text_data
+        return frame_data
     try:
-        message: Final[object] = _DECODED_JSON.validate_python(json.loads(text_data))
-    except json.JSONDecodeError:
-        return text_data
+        message: Final[object] = _DECODED_JSON.validate_python(json.loads(frame_data))
+    except (json.JSONDecodeError, UnicodeDecodeError, RecursionError):
+        return frame_data
     if not isinstance(message, dict):
-        return text_data
+        return frame_data
     setup: Final = message.get("setup")
     if not isinstance(setup, dict):
-        return text_data
+        return frame_data
     setup_model: Final = setup.get("model")
     if not isinstance(setup_model, str):
-        return text_data
+        return frame_data
     rewritten_model: Final = setup_model_rewriter(setup_model)
     if rewritten_model == setup_model:
-        return text_data
-    return json.dumps({**message, "setup": {**setup, "model": rewritten_model}})
+        return frame_data
+    rewritten_frame: Final = json.dumps({**message, "setup": {**setup, "model": rewritten_model}})
+    return rewritten_frame if isinstance(frame_data, str) else rewritten_frame.encode("utf-8")
+
+
+def _vertex_live_client_setup_model(
+    frame_data: str | bytes, setup_model_rewriter: Callable[[str], str] | None
+) -> str | None:
+    message: Final = _json_object_frame(frame_data)
+    setup: Final = None if message is None else message.get("setup")
+    if not isinstance(setup, dict) or "model" not in setup:
+        return None
+    return _extract_model_from_vertex_ai_setup(_resolved_vertex_live_setup(setup, setup_model_rewriter))
 
 
 def _resolved_vertex_live_setup(
@@ -2537,7 +2558,7 @@ def _json_object_frame(frame: str | bytes) -> dict[str, object] | None:
         decoded: Final[object] = _DECODED_JSON.validate_python(
             json.loads(frame if isinstance(frame, str) else frame.decode("utf-8"))
         )
-    except (json.JSONDecodeError, UnicodeDecodeError):
+    except (json.JSONDecodeError, UnicodeDecodeError, RecursionError):
         return None
     return decoded if isinstance(decoded, dict) else None
 
@@ -2552,6 +2573,12 @@ def _truncated_close_reason(reason: str) -> str:
     return encoded[:WEBSOCKET_CLOSE_REASON_MAX_BYTES].decode("utf-8", errors="ignore")
 
 
+def _refusal_close_reason(message: str) -> str:
+    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]
+
+
 SENDABLE_CLOSE_CODES: Final = frozenset(CloseCode) - frozenset(
     {CloseCode.NO_STATUS_RCVD, CloseCode.ABNORMAL_CLOSURE, CloseCode.TLS_HANDSHAKE}
 )
@@ -2595,6 +2622,53 @@ def _with_trace_context(headers: Mapping[str, str], parent_span: object) -> dict
     return inject_trace_context(headers, parent_span=parent_span)
 
 
+@dataclass(frozen=True, slots=True)
+class LoggedClientFrameRefusal:
+    denial: ProxyException
+
+
+ClientFrameRefusal: TypeAlias = ProxyException | LoggedClientFrameRefusal
+
+
+def _client_frame_denial(refusal: ClientFrameRefusal) -> ProxyException:
+    match refusal:
+        case LoggedClientFrameRefusal(denial=denial):
+            return denial
+        case ProxyException():
+            return refusal
+
+
+class WebsocketClientFrameGate(Protocol):
+    """Authorize one client frame, returning the denial to report, one its check already reported, or ``None`` to
+    let it through.
+
+    A websocket passthrough authenticates before any frame arrives, so a protocol that names its model
+    inside a frame has no model to authorize at connect time. The frame, not the extracted model, is the
+    argument because where a model sits is part of each provider's wire protocol
+    """
+
+    async def __call__(self, frame_data: str | bytes, valid_token: UserAPIKeyAuth, /) -> ClientFrameRefusal | None: ...
+
+
+async def _relay_client_frame(
+    upstream_ws: ClientConnection,
+    checked_frame: str | bytes,
+    forwarded_frame: str | bytes,
+    gate: WebsocketClientFrameGate | None,
+    valid_token: UserAPIKeyAuth,
+) -> ClientFrameRefusal | None:
+    """Forward one client frame upstream, or refuse it and hand back the denial that ended the session.
+
+    The denial travels out so the caller reports a failed session instead of leaving the success hooks to
+    record a refusal as a success, and closing either socket is the caller's job: awaiting a close here lets
+    the upstream reader finish first, which cancels this task mid-close and loses the denial
+    """
+    refusal: Final = None if gate is None else await gate(checked_frame, valid_token)
+    if refusal is None:
+        await upstream_ws.send(forwarded_frame)
+    return refusal
+
+
 async def websocket_passthrough_request(
     websocket: WebSocket,
     target: str,
@@ -2605,6 +2679,7 @@ async def websocket_passthrough_request(
     cost_per_request: float | None = None,
     accept_websocket: bool = True,
     setup_model_rewriter: Callable[[str], str] | None = None,
+    client_frame_gate: WebsocketClientFrameGate | None = None,
     builtin_provider_route: bool = False,
     model: str | None = None,
 ):
@@ -2620,6 +2695,7 @@ async def websocket_passthrough_request(
         endpoint: The endpoint path (for logging purposes)
         cost_per_request: Optional field - cost per request to the target endpoint
         setup_model_rewriter: Optional rewrite of the setup frame's model before it reaches the upstream
+        client_frame_gate: Optional per-frame authorization, applied before a frame reaches the upstream
         builtin_provider_route: Whether a built-in provider route serves the connection, so its URL names the provider
         model: The model the caller named before connecting, logged on a failure before the setup frame arrives
     """
@@ -2717,6 +2793,21 @@ async def websocket_passthrough_request(
     )
     logging_obj.model_call_details["litellm_call_id"] = litellm_call_id
 
+    async def log_failure(exception: Exception, traceback_str: str | None = None) -> None:
+        await proxy_logging_obj.post_call_failure_hook(
+            user_api_key_dict=user_api_key_dict,
+            original_exception=exception,
+            request_data=_build_passthrough_failure_request_payload(
+                parsed_body=None if model is None else {"model": model},
+                kwargs=kwargs,
+                logging_obj=logging_obj,
+                custom_llm_provider=None,
+                url=httpx.URL(target),
+                builtin_provider_route=builtin_provider_route,
+            ),
+            traceback_str=traceback_str,
+        )
+
     # Pre-call logging
     logging_obj.pre_call(
         input=[{"role": "user", "content": "WebSocket connection"}],
@@ -2748,8 +2839,22 @@ async def websocket_passthrough_request(
                 "WebSocket passthrough (%s): Upstream connection established successfully", endpoint
             )
 
-            async def forward_client_to_upstream() -> None:
-                """Forward messages from client to upstream WebSocket"""
+            is_vertex_live: Final = bool(endpoint and "/vertex_ai/live" in endpoint)
+
+            def record_client_setup_model(frame_data: str | bytes) -> None:
+                if not is_vertex_live:
+                    return
+                extracted_model: Final = _vertex_live_client_setup_model(frame_data, setup_model_rewriter)
+                if not extracted_model:
+                    return
+                kwargs["model"] = extracted_model
+                kwargs["custom_llm_provider"] = "vertex_ai-language-models"
+                logging_obj.model = extracted_model
+                logging_obj.model_call_details["model"] = extracted_model
+                logging_obj.model_call_details["custom_llm_provider"] = "vertex_ai"
+
+            async def forward_client_to_upstream() -> ClientFrameRefusal | None:
+                """Forward messages from client to upstream WebSocket, returning any denial that ended it"""
                 try:
                     while True:
                         message = await websocket.receive()
@@ -2762,65 +2867,29 @@ async def websocket_passthrough_request(
                         bytes_data: bytes | None = message.get("bytes")
 
                         if text_data is not None:
-                            # Try to extract model from client setup message for Vertex AI Live
-                            if endpoint and "/vertex_ai/live" in endpoint:
-                                verbose_proxy_logger.debug(
-                                    "WebSocket passthrough (%s): Processing client message for model extraction",
-                                    endpoint,
+                            if (
+                                text_refusal := await _relay_client_frame(
+                                    upstream_ws,
+                                    text_data,
+                                    _rewrite_vertex_live_setup_model(text_data, setup_model_rewriter),
+                                    client_frame_gate,
+                                    user_api_key_dict,
                                 )
-                                try:
-                                    client_message = _DECODED_JSON.validate_python(json.loads(text_data))
-                                    if isinstance(client_message, dict) and "setup" in client_message:
-                                        setup_data = client_message["setup"]
-                                        verbose_proxy_logger.debug(
-                                            "WebSocket passthrough (%s): Found setup data in client message: %s",
-                                            endpoint,
-                                            setup_data,
-                                        )
-                                        if isinstance(setup_data, dict) and "model" in setup_data:
-                                            extracted_model = _extract_model_from_vertex_ai_setup(
-                                                _resolved_vertex_live_setup(setup_data, setup_model_rewriter)
-                                            )
-                                            if extracted_model:
-                                                kwargs["model"] = extracted_model
-                                                kwargs["custom_llm_provider"] = "vertex_ai-language-models"
-                                                # Update logging object with correct model
-                                                logging_obj.model = extracted_model
-                                                logging_obj.model_call_details["model"] = extracted_model
-                                                logging_obj.model_call_details["custom_llm_provider"] = "vertex_ai"
-                                                verbose_proxy_logger.info(
-                                                    "WebSocket passthrough (%s): Successfully extracted model '%s' and set provider to 'vertex_ai' from client setup message",
-                                                    endpoint,
-                                                    extracted_model,
-                                                )
-                                            else:
-                                                verbose_proxy_logger.warning(
-                                                    "WebSocket passthrough (%s): Failed to extract model from client setup data: %s",
-                                                    endpoint,
-                                                    setup_data,
-                                                )
-                                        else:
-                                            verbose_proxy_logger.debug(
-                                                "WebSocket passthrough (%s): Setup data does not contain model field: %s",
-                                                endpoint,
-                                                setup_data,
-                                            )
-                                    else:
-                                        verbose_proxy_logger.debug(
-                                            "WebSocket passthrough (%s): Client message does not contain setup data",
-                                            endpoint,
-                                        )
-                                except (json.JSONDecodeError, KeyError, TypeError) as e:
-                                    verbose_proxy_logger.debug(
-                                        "WebSocket passthrough (%s): Client message is not a valid setup message: %s",
-                                        endpoint,
-                                        e,
-                                    )
-                                    # Not a JSON message or doesn't contain setup data
-
-                            await upstream_ws.send(_rewrite_vertex_live_setup_model(text_data, setup_model_rewriter))
+                            ) is not None:
+                                return text_refusal
+                            record_client_setup_model(text_data)
                         elif bytes_data is not None:
-                            await upstream_ws.send(bytes_data)
+                            if (
+                                bytes_refusal := await _relay_client_frame(
+                                    upstream_ws,
+                                    bytes_data,
+                                    _rewrite_vertex_live_setup_model(bytes_data, setup_model_rewriter),
+                                    client_frame_gate,
+                                    user_api_key_dict,
+                                )
+                            ) is not None:
+                                return bytes_refusal
+                            record_client_setup_model(bytes_data)
                 except asyncio.CancelledError:
                     raise
                 except Exception:
@@ -2844,7 +2913,6 @@ async def websocket_passthrough_request(
                 logging_obj.model_call_details["model"] = extracted_model
                 logging_obj.model_call_details["custom_llm_provider"] = "vertex_ai_language_models"
 
-            is_vertex_live: Final = bool(endpoint and "/vertex_ai/live" in endpoint)
             json_frame_ordinal: Final = count()
 
             async def relay_upstream_frame(upstream_message: str | bytes) -> None:
@@ -2893,13 +2961,33 @@ async def websocket_passthrough_request(
                 except asyncio.CancelledError:
                     pass
 
+            task_results: Final = tuple(task.result() for task in done if task.exception() is None)
+
+            # A refusal outranks any exception it caused: closing the client socket is what makes the
+            # upstream reader's next send fail, so raising that instead would report the wrong cause.
+            # Turns that completed before a refusal were served, so they are logged and billed as a success
+            client_frame_refusal: Final = next(
+                (result for result in task_results if isinstance(result, (ProxyException, LoggedClientFrameRefusal))),
+                None,
+            )
+            if client_frame_refusal is not None:
+                await websocket.close(
+                    code=CloseCode.POLICY_VIOLATION,
+                    reason=_refusal_close_reason(_client_frame_denial(client_frame_refusal).message),
+                )
+                await upstream_ws.close()
+            if isinstance(client_frame_refusal, ProxyException) and not websocket_messages:
+                await log_failure(client_frame_refusal)
+            if client_frame_refusal is not None and not websocket_messages:
+                return
+
             # Check for exceptions in completed tasks
             for task in done:
                 exception = task.exception()
-                if exception is not None:
+                if exception is not None and client_frame_refusal is None:
                     raise exception
 
-            upstream_close: Final = _upstream_close_to_relay(task.result() for task in done)
+            upstream_close: Final = _upstream_close_to_relay(task_results)
             if upstream_close is not None and _client_socket_is_open(websocket):
                 await websocket.close(
                     code=upstream_close.code,
@@ -2974,21 +3062,7 @@ async def websocket_passthrough_request(
     except InvalidStatus as exc:
         verbose_proxy_logger.exception("WebSocket passthrough (%s): upstream rejected WebSocket connection", endpoint)
 
-        await proxy_logging_obj.post_call_failure_hook(
-            user_api_key_dict=user_api_key_dict,
-            original_exception=exc,
-            request_data=_build_passthrough_failure_request_payload(
-                parsed_body=None if model is None else {"model": model},
-                kwargs=kwargs,
-                logging_obj=logging_obj,
-                custom_llm_provider=None,
-                url=httpx.URL(target),
-                builtin_provider_route=builtin_provider_route,
-            ),
-            traceback_str=traceback.format_exc(
-                limit=MAXIMUM_TRACEBACK_LINES_TO_LOG,
-            ),
-        )
+        await log_failure(exc, traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG))
 
         if _client_socket_is_open(websocket):
             await websocket.close(
@@ -3000,21 +3074,7 @@ async def websocket_passthrough_request(
             "WebSocket passthrough (%s): unexpected error while proxying WebSocket", endpoint
         )
 
-        await proxy_logging_obj.post_call_failure_hook(
-            user_api_key_dict=user_api_key_dict,
-            original_exception=e,
-            request_data=_build_passthrough_failure_request_payload(
-                parsed_body=None if model is None else {"model": model},
-                kwargs=kwargs,
-                logging_obj=logging_obj,
-                custom_llm_provider=None,
-                url=httpx.URL(target),
-                builtin_provider_route=builtin_provider_route,
-            ),
-            traceback_str=traceback.format_exc(
-                limit=MAXIMUM_TRACEBACK_LINES_TO_LOG,
-            ),
-        )
+        await log_failure(e, traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG))
 
         if _client_socket_is_open(websocket):
             await websocket.close(code=1011, reason="WebSocket passthrough error")

tests/unit/litellm_core_utils/test_get_llm_provider_logic.py

@@ -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"

tests/unit/proxy/auth/test_auth_checks.py

@@ -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,
@@ -2427,6 +2428,37 @@ async def test_bulk_eviction_clears_the_redis_copies_before_waiting_on_a_wedged_
     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
+
+
 class _SignallingPublishRedisClient(_SlowPublishRedisClient):
     def __init__(self) -> None:
         super().__init__()