fix(vertex-live): resolve a Live setup model before logging reads it

A client that named a bare gateway alias logged the session as "unknown" and billed nothing,
because the model was read off the raw setup frame and the extractor only yields a name when the
string already contains "/models/". The rewriter qualifies that same model a few lines later for
the upstream, so the supported client form, an alias, was the one that went unbilled.

Resolving through the rewriter first means the real model reaches the logging object, and from
there the cost map. A route with no rewriter, which is every non-Live passthrough, hands the frame
over untouched.

(cherry picked from commit 573982803df612fd94144e2e06dd647f8530d4e8)
This commit is contained in:
Marty Sullivan 2026-09-07 23:26:31 -04:00
parent b1262b05f4
commit 9580b89bb1
2 changed files with 116 additions and 1 deletions

View file

@ -2090,6 +2090,22 @@ def _rewrite_vertex_live_setup_model(text_data: str, setup_model_rewriter: Calla
return json.dumps({**message, "setup": {**setup, "model": rewritten_model}}) # mutable-ok: one-shot json payload
def _resolved_vertex_live_setup(
setup_data: Mapping[str, object], setup_model_rewriter: Callable[[str], str] | None
) -> Mapping[str, object]:
"""
Give the model extractor the same fully qualified path the upstream will receive.
Clients may name a bare gateway alias, which the rewriter turns into a ``projects/...`` path before
it reaches Vertex. The extractor only reads a path containing ``/models/``, so running it on the raw
frame logs the session as ``unknown`` at no cost, which is precisely the supported client form
"""
setup_model: Final = setup_data.get("model")
if setup_model_rewriter is None or not isinstance(setup_model, str):
return setup_data
return {**setup_data, "model": setup_model_rewriter(setup_model)}
def _truncated_close_reason(reason: str) -> str:
"""
Fit a close reason inside the byte budget a WebSocket close frame allows, without splitting a character
@ -2314,7 +2330,12 @@ async def websocket_passthrough_request(
setup_data,
)
if isinstance(setup_data, dict) and "model" in setup_data:
extracted_model = _extract_model_from_vertex_ai_setup(setup_data)
# Resolve the alias first: a client may name a bare gateway model,
# which carries no "/models/" for the extractor to read, so reading
# the raw frame leaves the session logged as "unknown" and unbilled.
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"

View file

@ -5030,6 +5030,100 @@ async def test_websocket_passthrough_rewrites_gateway_alias_setup_model():
assert sent_setup["model"] == "projects/proj-db/locations/global/publishers/google/models/gemini-live-2.5-flash"
@pytest.mark.parametrize(
"setup_model",
["gemini-live-2.5-flash", "models/gemini-live-2.5-flash", "publishers/google/models/gemini-live-2.5-flash"],
)
def test_vertex_live_setup_model_resolves_before_extraction(setup_model):
"""A bare gateway alias left the session logged as ``unknown`` at zero cost.
The model was read off the raw client frame, and the extractor only yields a name when the string
already contains ``/models/``. The rewriter qualifies it a few lines later for the upstream, so a
client that addressed the gateway the documented way, by alias, logged no model and therefore
resolved no cost-map entry. Resolving first is what puts the real name on the logging object.
"""
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
_build_vertex_live_setup_model_rewriter,
)
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
_extract_model_from_vertex_ai_setup,
_resolved_vertex_live_setup,
)
rewriter = _build_vertex_live_setup_model_rewriter(
vertex_project="proj-db", vertex_location="global", llm_router=None
)
setup_data = {"model": setup_model}
resolved = _extract_model_from_vertex_ai_setup(_resolved_vertex_live_setup(setup_data, rewriter))
assert resolved == "gemini-live-2.5-flash", "an unresolved setup model logs the session as 'unknown'"
@pytest.mark.asyncio
async def test_websocket_passthrough_logs_a_bare_alias_setup_model():
"""End to end through the relay: a bare alias must reach the logging object as a real model name.
This is the call-site half of the fix. The helper tests above pass even if extraction moves back
before the rewrite, so this one drives the real websocket relay and asserts on what got logged,
which is the name the cost map is looked up by. An unbilled session logs ``unknown``.
"""
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
_build_vertex_live_setup_model_rewriter,
)
upstream_ws = RecordingUpstreamWebSocket()
setup_frame = json.dumps({"setup": {"model": "gemini-live-2.5-flash"}})
websocket = _client_websocket(
AsyncMock(
side_effect=[
{"type": "websocket.receive", "text": setup_frame},
{"type": "websocket.disconnect"},
]
)
)
built = []
real_logging = litellm.litellm_core_utils.litellm_logging.Logging
def _capture(*args, **kwargs):
obj = real_logging(*args, **kwargs)
built.append(obj)
return obj
with _patched_websocket_passthrough_environment(upstream_ws):
with patch("litellm.litellm_core_utils.litellm_logging.Logging", side_effect=_capture):
await websocket_passthrough_request(
websocket=websocket,
target="wss://aiplatform.googleapis.com/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent",
custom_headers={"Authorization": "Bearer token"},
user_api_key_dict=UserAPIKeyAuth(),
forward_headers=False,
endpoint="/vertex_ai/live",
accept_websocket=False,
setup_model_rewriter=_build_vertex_live_setup_model_rewriter(
vertex_project="proj-db", vertex_location="global", llm_router=None
),
)
assert built, "the relay should have built a logging object"
assert built[0].model == "gemini-live-2.5-flash", "a bare alias must not log as 'unknown'"
def test_vertex_live_setup_resolution_is_inert_without_a_rewriter():
"""Non-Live passthrough routes pass no rewriter, so the frame must be handed over untouched."""
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
_extract_model_from_vertex_ai_setup,
_resolved_vertex_live_setup,
)
setup_data = {"model": "projects/p/locations/global/publishers/google/models/gemini-live-2.5-flash"}
assert _resolved_vertex_live_setup(setup_data, None) is setup_data
assert _extract_model_from_vertex_ai_setup(_resolved_vertex_live_setup(setup_data, None)) == (
"gemini-live-2.5-flash"
)
@pytest.mark.asyncio
@pytest.mark.parametrize("rcvd_close", [None, "abnormal", "no_status"])
async def test_websocket_passthrough_does_not_relay_unsendable_upstream_close(rcvd_close):