From 00fcc96002b5c43b91081d88b8d168ef4240cd3b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 06:11:52 -0700 Subject: [PATCH] fix(proxy): keep a Vertex Live setup group's budget and pricing answer when the key may call it Every name a setup frame matches runs the same deployment, so the first name the key may call is re-authorized and its denial stands instead of falling through to the bare model id, which no group's budget or pricing check covers --- .../llm_passthrough_endpoints.py | 22 ++++++--------- .../test_pass_through_endpoints.py | 27 +++++++++++++++++++ 2 files changed, 35 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 23650ded311..8ffc6f04c1b 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -4095,23 +4095,17 @@ async def _first_model_denial( authorize_model: VertexLiveModelAuthorizer, ) -> ClientFrameRefusal | None: """ - The first name the key may call that also passes re-authorization opens the session, and no later name is - checked. Otherwise the refusal is the first re-authorization denial, which says more than an access denial + 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] - denial: Final = await _model_access_denial(name, valid_token, llm_router) or await _model_reauthorization_denial( - name, authorize_model - ) - if denial is None or len(model_names) == 1: - return denial + 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) - if later_denial is None: - return None - return ( - later_denial - if isinstance(later_denial, LoggedClientFrameRefusal) and isinstance(denial, ProxyException) - else denial - ) + return access_denial if isinstance(later_denial, ProxyException) else later_denial async def _vertex_live_model_denial( 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 2b71e6dfc96..0db2500bec0 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 @@ -8541,6 +8541,33 @@ async def test_vertex_live_route_releases_the_budget_it_reserved_to_authorize_an assert valid_token.budget_reservation is not None and valid_token.budget_reservation["finalized"] is True +@pytest.mark.asyncio +async def test_vertex_live_setup_group_the_key_may_call_but_fails_budget_is_not_opened_under_its_bare_model_id(): + """ + A setup frame naming a group runs that group's deployment whichever name the key holds, so when the key may call + the group its budget and pricing checks decide, and the group's bare model id cannot open the session instead + """ + group_denial: Final = ProxyException( + message=f"Budget has been exceeded for model {VERTEX_LIVE_GROUP}.", + type=ProxyErrorTypes.budget_exceeded, + param="model", + code=400, + ) + valid_token: Final = UserAPIKeyAuth(token="hashed") + + async def authorize_model(model: str) -> UserAPIKeyAuth: + if model == VERTEX_LIVE_GROUP: + raise group_denial + return valid_token + + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [_setup_frame(VERTEX_LIVE_GROUP)], valid_token, authorize_model=authorize_model + ) + + upstream_ws.send.assert_not_awaited() + assert _policy_close_reason(websocket) == group_denial.message + + @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"""