mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-09 08:18:12 +00:00
test(proxy): cover per-frame model checks and session pricing on the OpenAI websocket passthrough
This commit is contained in:
mateo-berri
parent
8d6b8dab3d
commit
a0d7377afe
6 changed files with 928 additions and 223 deletions
|
|
@ -10750,139 +10750,6 @@ async def test_key_budget_error_keeps_the_masked_key_name(key_name):
|
|||
assert f"Key=prod-key ({key_name}) Current cost" in message
|
||||
|
||||
|
||||
class _UntouchedPrisma:
|
||||
def __getattr__(self, name: str) -> object:
|
||||
raise AssertionError(f"database reached through {name}")
|
||||
|
||||
|
||||
class _MissingUserPrisma:
|
||||
class db:
|
||||
class litellm_usertable:
|
||||
@staticmethod
|
||||
async def find_unique(where: dict[str, str], include: dict[str, bool]) -> None:
|
||||
return None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_enforced_model_allowlists_treats_a_missing_user_row_as_unrestricted():
|
||||
from litellm.proxy.auth.auth_checks import enforced_model_allowlists
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
||||
cache = UserApiKeyCache()
|
||||
scopes = await enforced_model_allowlists(
|
||||
valid_token=UserAPIKeyAuth(token="hashed-fake", user_id="default_user_id"),
|
||||
prisma_client=_MissingUserPrisma(),
|
||||
user_api_key_cache=cache,
|
||||
proxy_logging_obj=ProxyLogging(user_api_key_cache=cache),
|
||||
)
|
||||
|
||||
assert [list(scope) for scope in scopes] == [[], [], [], [], []]
|
||||
|
||||
|
||||
class _UnreachableUserPrisma:
|
||||
class db:
|
||||
class litellm_usertable:
|
||||
@staticmethod
|
||||
async def find_unique(where: dict[str, str], include: dict[str, bool]) -> None:
|
||||
raise RuntimeError("database gone")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_enforced_model_allowlists_surfaces_a_failed_user_lookup():
|
||||
from litellm.proxy.auth.auth_checks import enforced_model_allowlists
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
||||
cache = UserApiKeyCache()
|
||||
with pytest.raises(ValueError, match="database gone"):
|
||||
await enforced_model_allowlists(
|
||||
valid_token=UserAPIKeyAuth(token="hashed-fake", user_id="user-fake"),
|
||||
prisma_client=_UnreachableUserPrisma(),
|
||||
user_api_key_cache=cache,
|
||||
proxy_logging_obj=ProxyLogging(user_api_key_cache=cache),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_enforced_model_allowlists_reads_every_level_from_cache():
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_BudgetTable,
|
||||
LiteLLM_ProjectTableCachedObj,
|
||||
LiteLLM_TeamMembership,
|
||||
LiteLLM_TeamTableCachedObj,
|
||||
LiteLLM_UserTable,
|
||||
)
|
||||
from litellm.proxy.auth.auth_checks import enforced_model_allowlists
|
||||
from litellm.proxy.common_utils.user_api_key_cache import (
|
||||
UserApiKeyCache,
|
||||
team_membership_reservation_cache_key,
|
||||
)
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
||||
cache = UserApiKeyCache()
|
||||
proxy_logging_obj = ProxyLogging(user_api_key_cache=cache)
|
||||
await cache.async_set_cache(
|
||||
key="team_id:team-fake", value=LiteLLM_TeamTableCachedObj(team_id="team-fake", models=["gpt-4o"])
|
||||
)
|
||||
await cache.async_set_cache(
|
||||
key=team_membership_reservation_cache_key(user_id="user-fake", team_id="team-fake"),
|
||||
value=LiteLLM_TeamMembership(
|
||||
user_id="user-fake",
|
||||
team_id="team-fake",
|
||||
litellm_budget_table=LiteLLM_BudgetTable(allowed_models=["gpt-4o-mini"]),
|
||||
),
|
||||
)
|
||||
await cache.async_set_cache(
|
||||
key="project_id:project-fake",
|
||||
value=LiteLLM_ProjectTableCachedObj(project_id="project-fake", models=["gpt-4.1"]),
|
||||
)
|
||||
await cache.async_set_cache(key="user-fake", value=LiteLLM_UserTable(user_id="user-fake", models=["o3"]))
|
||||
prisma_client = _UntouchedPrisma()
|
||||
|
||||
team_scoped = await enforced_model_allowlists(
|
||||
valid_token=UserAPIKeyAuth(
|
||||
token="hashed-fake",
|
||||
models=["all-team-models"],
|
||||
team_models=["gpt-4o", "gpt-4o-mini"],
|
||||
user_id="user-fake",
|
||||
team_id="team-fake",
|
||||
project_id="project-fake",
|
||||
),
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
personal = await enforced_model_allowlists(
|
||||
valid_token=UserAPIKeyAuth(token="hashed-fake", user_id="user-fake"),
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
without_database = await enforced_model_allowlists(
|
||||
valid_token=UserAPIKeyAuth(
|
||||
token="hashed-fake",
|
||||
models=["gpt-4o"],
|
||||
team_models=["gpt-4o-mini"],
|
||||
user_id="user-fake",
|
||||
team_id="team-fake",
|
||||
),
|
||||
prisma_client=None,
|
||||
user_api_key_cache=cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
assert [list(scope) for scope in team_scoped] == [
|
||||
["gpt-4o", "gpt-4o-mini"],
|
||||
["gpt-4o"],
|
||||
["gpt-4o-mini"],
|
||||
[],
|
||||
["gpt-4.1"],
|
||||
]
|
||||
assert [list(scope) for scope in personal] == [[], [], [], ["o3"], []]
|
||||
assert [list(scope) for scope in without_database] == [["gpt-4o"], ["gpt-4o-mini"]]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("channel", ["team", "key"])
|
||||
async def test_access_group_model_fallback_uses_the_injected_database(channel: str) -> None:
|
||||
|
|
|
|||
|
|
@ -27,6 +27,7 @@ from litellm.proxy.auth.auth_checks import (
|
|||
_model_access_group_max_budget_check,
|
||||
collect_matched_model_access_groups,
|
||||
common_checks,
|
||||
enforce_model_access_group_budget,
|
||||
stamp_matched_model_access_groups,
|
||||
)
|
||||
from litellm.proxy.common_utils.reset_budget_job import _model_access_group_counter_key
|
||||
|
|
@ -551,3 +552,65 @@ async def test_common_checks_blocks_a_request_whose_group_is_over_budget():
|
|||
async def test_free_model_routes_skip_the_model_access_group_budget_check():
|
||||
"""skip_budget_checks is how free models stay free; it has to cover this budget too."""
|
||||
assert await _common_checks_with_over_budget_group(skip_budget_checks=True) is True
|
||||
|
||||
|
||||
async def _enforce_for_a_route_without_a_request_body(
|
||||
*rows: _MagBudgetRow,
|
||||
key_models: list[str],
|
||||
spend_by_counter_key: dict[str, float],
|
||||
) -> UserAPIKeyAuth:
|
||||
"""Run the pair the websocket frame gate runs, which reaches the proxy for what it needs."""
|
||||
valid_token = UserAPIKeyAuth(api_key="hashed", models=key_models, user_id=USER_ID)
|
||||
cache = await _cache()
|
||||
read, _ = _spend_reader(spend_by_counter_key)
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", _RecordingPrismaClient(*rows)), # test-quality-ok: read lazily inside the helper, which takes no client
|
||||
patch("litellm.proxy.proxy_server.user_api_key_cache", cache), # test-quality-ok: same lazy import
|
||||
patch("litellm.proxy.proxy_server.proxy_logging_obj", ProxyLogging(user_api_key_cache=cache)), # test-quality-ok: same lazy import
|
||||
patch("litellm.proxy.proxy_server.get_current_spend", read), # test-quality-ok: get_current_spend is lazily imported inside the budget check and has no injection point
|
||||
):
|
||||
await enforce_model_access_group_budget(
|
||||
model="gpt-4o",
|
||||
valid_token=valid_token,
|
||||
llm_router=Router(model_list=MODEL_LIST),
|
||||
)
|
||||
return valid_token
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_caller_granted_the_model_only_through_an_exhausted_group_is_refused():
|
||||
"""Routes with no request body, the websocket passthrough among them, run this instead of
|
||||
``common_checks``, so the group budget has to stop them the same way."""
|
||||
with pytest.raises(litellm.BudgetExceededError) as exc_info:
|
||||
await _enforce_for_a_route_without_a_request_body(
|
||||
_MagBudgetRow("tier-a", max_budget=1.0),
|
||||
key_models=["tier-a"],
|
||||
spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 1.5},
|
||||
)
|
||||
|
||||
assert exc_info.value.entity_id == "tier-a"
|
||||
assert exc_info.value.current_cost == 1.5
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_group_that_authorized_the_call_is_stamped_for_the_spend_writer():
|
||||
"""Without the stamp the call runs on the group's budget and never pays into it, so the pool
|
||||
can never catch up with what it authorized."""
|
||||
valid_token = await _enforce_for_a_route_without_a_request_body(
|
||||
_MagBudgetRow("tier-a", max_budget=10.0),
|
||||
key_models=["tier-a"],
|
||||
spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 1.0},
|
||||
)
|
||||
|
||||
assert valid_token.matched_model_access_groups == ["tier-a"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_granted_directly_charges_no_group():
|
||||
valid_token = await _enforce_for_a_route_without_a_request_body(
|
||||
_MagBudgetRow("tier-a", max_budget=1.0),
|
||||
key_models=["gpt-4o"],
|
||||
spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 1.5},
|
||||
)
|
||||
|
||||
assert valid_token.matched_model_access_groups is None
|
||||
|
|
|
|||
|
|
@ -0,0 +1,290 @@
|
|||
from datetime import datetime
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_websocket_passthrough_logging_handler import (
|
||||
OpenAIWebsocketPassthroughLoggingHandler,
|
||||
)
|
||||
|
||||
REALTIME_MODEL: Final = "gpt-realtime-2.1"
|
||||
RESPONSES_MODEL: Final = "gpt-5.4-mini"
|
||||
REASONING_MODEL: Final = "gpt-5.4"
|
||||
|
||||
|
||||
def _logging_obj() -> MagicMock:
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.model_call_details = {}
|
||||
return logging_obj
|
||||
|
||||
|
||||
def _session_created(model: str) -> dict:
|
||||
return {"type": "session.created", "session": {"id": "sess_fake", "model": model}}
|
||||
|
||||
|
||||
def _realtime_response_done(text_in: int, audio_in: int, text_out: int, audio_out: int) -> dict:
|
||||
return {
|
||||
"type": "response.done",
|
||||
"response": {
|
||||
"id": "resp_fake",
|
||||
"status": "completed",
|
||||
"usage": {
|
||||
"input_tokens": text_in + audio_in,
|
||||
"output_tokens": text_out + audio_out,
|
||||
"total_tokens": text_in + audio_in + text_out + audio_out,
|
||||
"input_token_details": {"text_tokens": text_in, "audio_tokens": audio_in, "cached_tokens": 0},
|
||||
"output_token_details": {"text_tokens": text_out, "audio_tokens": audio_out},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _responses_completed(model: str, input_tokens: int, output_tokens: int) -> dict:
|
||||
return {
|
||||
"type": "response.completed",
|
||||
"response": {
|
||||
"id": "resp_fake",
|
||||
"model": model,
|
||||
"status": "completed",
|
||||
"usage": {
|
||||
"input_tokens": input_tokens,
|
||||
"output_tokens": output_tokens,
|
||||
"total_tokens": input_tokens + output_tokens,
|
||||
"input_tokens_details": {"cached_tokens": 0},
|
||||
"output_tokens_details": {"reasoning_tokens": 0},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _handle(messages, kwargs=MappingProxyType({}), logging_obj=None, url_route="/openai/v1/realtime"):
|
||||
return OpenAIWebsocketPassthroughLoggingHandler().openai_websocket_passthrough_handler(
|
||||
websocket_messages=messages,
|
||||
logging_obj=logging_obj if logging_obj is not None else _logging_obj(),
|
||||
url_route=url_route,
|
||||
start_time=datetime.now(),
|
||||
kwargs=kwargs,
|
||||
)
|
||||
|
||||
|
||||
def _realtime_cost(model: str, text_in: int, audio_in: int, text_out: int, audio_out: int) -> float:
|
||||
prices = litellm.model_cost[model]
|
||||
return (
|
||||
text_in * prices["input_cost_per_token"]
|
||||
+ audio_in * prices["input_cost_per_audio_token"]
|
||||
+ text_out * prices["output_cost_per_token"]
|
||||
+ audio_out * prices["output_cost_per_audio_token"]
|
||||
)
|
||||
|
||||
|
||||
def test_realtime_session_is_priced_on_the_model_the_session_reported():
|
||||
logging_obj = _logging_obj()
|
||||
messages = [
|
||||
_session_created(REALTIME_MODEL),
|
||||
{"type": "response.created", "response": {"id": "resp_fake"}},
|
||||
_realtime_response_done(text_in=33, audio_in=120, text_out=16, audio_out=240),
|
||||
]
|
||||
|
||||
handled = _handle(messages, logging_obj=logging_obj)
|
||||
|
||||
expected_cost = _realtime_cost(REALTIME_MODEL, 33, 120, 16, 240)
|
||||
assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost)
|
||||
assert handled["kwargs"]["model"] == REALTIME_MODEL
|
||||
assert handled["kwargs"]["custom_llm_provider"] == "openai"
|
||||
assert handled["result"].model == REALTIME_MODEL
|
||||
assert handled["result"].usage.total_tokens == 409
|
||||
assert logging_obj.model_call_details["response_cost"] == pytest.approx(expected_cost)
|
||||
assert logging_obj.model_call_details["model"] == REALTIME_MODEL
|
||||
|
||||
|
||||
def test_realtime_cost_covers_every_turn_of_the_session():
|
||||
one_turn = _handle([_session_created(REALTIME_MODEL), _realtime_response_done(10, 0, 5, 0)])
|
||||
two_turns = _handle(
|
||||
[
|
||||
_session_created(REALTIME_MODEL),
|
||||
_realtime_response_done(10, 0, 5, 0),
|
||||
_realtime_response_done(10, 0, 5, 0),
|
||||
]
|
||||
)
|
||||
|
||||
assert two_turns["kwargs"]["response_cost"] == pytest.approx(2 * one_turn["kwargs"]["response_cost"])
|
||||
assert two_turns["result"].usage.total_tokens == 30
|
||||
|
||||
|
||||
def test_responses_turn_is_priced_on_the_model_the_response_reported():
|
||||
handled = _handle(
|
||||
[_responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5)],
|
||||
url_route="/openai_passthrough/v1/responses",
|
||||
)
|
||||
|
||||
prices = litellm.model_cost[RESPONSES_MODEL]
|
||||
expected_cost = 13 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"]
|
||||
assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost)
|
||||
assert handled["kwargs"]["model"] == RESPONSES_MODEL
|
||||
assert handled["result"].usage.prompt_tokens == 13
|
||||
assert handled["result"].usage.completion_tokens == 5
|
||||
|
||||
|
||||
def test_responses_cost_covers_every_turn_of_the_connection():
|
||||
handled = _handle(
|
||||
[
|
||||
_responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5),
|
||||
_responses_completed(RESPONSES_MODEL, input_tokens=20, output_tokens=7),
|
||||
],
|
||||
url_route="/openai_passthrough/v1/responses",
|
||||
)
|
||||
|
||||
prices = litellm.model_cost[RESPONSES_MODEL]
|
||||
expected_cost = 33 * prices["input_cost_per_token"] + 12 * prices["output_cost_per_token"]
|
||||
assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost)
|
||||
assert handled["result"].usage.total_tokens == 45
|
||||
|
||||
|
||||
def test_a_model_the_caller_asked_for_never_prices_a_session_the_provider_never_named():
|
||||
"""
|
||||
The caller picks what to ask for, and the provider picks what to run. Only
|
||||
the provider's own frames may set the price of a session, so a model the
|
||||
caller put in the request cannot stand in for one the provider never named.
|
||||
"""
|
||||
logging_obj = _logging_obj()
|
||||
|
||||
handled = _handle(
|
||||
[_realtime_response_done(text_in=10, audio_in=0, text_out=5, audio_out=0)],
|
||||
kwargs=MappingProxyType({"model": REALTIME_MODEL}),
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
|
||||
assert handled["result"] is None
|
||||
assert handled["kwargs"] == {"model": REALTIME_MODEL}
|
||||
assert logging_obj.model_call_details == {}
|
||||
|
||||
|
||||
def test_each_responses_turn_is_priced_on_the_model_that_turn_ran():
|
||||
"""
|
||||
One socket can run several models one after another, so a connection that
|
||||
starts on a big model and finishes on a small one has to pay the big
|
||||
model's price for the turn that used it.
|
||||
"""
|
||||
handled = _handle(
|
||||
[
|
||||
_responses_completed(REASONING_MODEL, input_tokens=13, output_tokens=5),
|
||||
_responses_completed(RESPONSES_MODEL, input_tokens=20, output_tokens=7),
|
||||
],
|
||||
url_route="/openai_passthrough/v1/responses",
|
||||
)
|
||||
|
||||
first = litellm.model_cost[REASONING_MODEL]
|
||||
second = litellm.model_cost[RESPONSES_MODEL]
|
||||
expected_cost = (
|
||||
13 * first["input_cost_per_token"]
|
||||
+ 5 * first["output_cost_per_token"]
|
||||
+ 20 * second["input_cost_per_token"]
|
||||
+ 7 * second["output_cost_per_token"]
|
||||
)
|
||||
assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost)
|
||||
assert handled["result"].usage.total_tokens == 45
|
||||
|
||||
|
||||
def test_a_responses_turn_that_names_no_model_is_priced_on_the_session_model():
|
||||
handled = _handle(
|
||||
[
|
||||
_session_created(RESPONSES_MODEL),
|
||||
{
|
||||
"type": "response.completed",
|
||||
"response": {"usage": {"input_tokens": 13, "output_tokens": 5, "total_tokens": 18}},
|
||||
},
|
||||
],
|
||||
url_route="/openai_passthrough/v1/responses",
|
||||
)
|
||||
|
||||
prices = litellm.model_cost[RESPONSES_MODEL]
|
||||
expected_cost = 13 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"]
|
||||
assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost)
|
||||
|
||||
|
||||
def test_a_session_nobody_named_a_model_for_is_left_unpriced():
|
||||
logging_obj = _logging_obj()
|
||||
handled = _handle(
|
||||
[{"type": "response.done", "response": {"usage": {"input_tokens": 10, "output_tokens": 5}}}],
|
||||
kwargs=MappingProxyType({"litellm_call_id": "call-fake"}),
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
|
||||
assert handled["result"] is None
|
||||
assert handled["kwargs"] == {"litellm_call_id": "call-fake"}
|
||||
assert logging_obj.model_call_details == {}
|
||||
|
||||
|
||||
def test_a_session_that_reported_no_usage_still_records_its_model():
|
||||
logging_obj = _logging_obj()
|
||||
|
||||
handled = _handle([_session_created(REALTIME_MODEL)], logging_obj=logging_obj)
|
||||
|
||||
assert handled["result"] is None
|
||||
assert handled["kwargs"] == {"model": REALTIME_MODEL, "custom_llm_provider": "openai"}
|
||||
assert logging_obj.model_call_details["model"] == REALTIME_MODEL
|
||||
assert logging_obj.model_call_details["custom_llm_provider"] == "openai"
|
||||
|
||||
|
||||
def test_a_responses_turn_without_readable_usage_is_not_billed():
|
||||
handled = _handle(
|
||||
[{"type": "response.completed", "response": {"model": RESPONSES_MODEL, "status": "completed"}}],
|
||||
url_route="/openai_passthrough/v1/responses",
|
||||
)
|
||||
|
||||
assert handled["result"] is None
|
||||
assert handled["kwargs"] == {"model": RESPONSES_MODEL, "custom_llm_provider": "openai"}
|
||||
|
||||
|
||||
def test_the_price_of_a_multi_model_connection_reaches_the_spend_row():
|
||||
"""
|
||||
The logging layer prices whatever result it is handed, and one result can
|
||||
only name one model, so a connection that ran two models has to carry the
|
||||
price it already worked out or the cheaper model pays for both turns. The
|
||||
row it lands in names the model the connection opened on.
|
||||
"""
|
||||
handled = _handle(
|
||||
[
|
||||
_responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5),
|
||||
_responses_completed(REASONING_MODEL, input_tokens=20, output_tokens=7),
|
||||
],
|
||||
url_route="/openai_passthrough/v1/responses",
|
||||
)
|
||||
|
||||
logged_cost = litellm.response_cost_calculator(
|
||||
response_object=handled["result"],
|
||||
model=handled["kwargs"]["model"],
|
||||
custom_llm_provider="openai",
|
||||
call_type="completion",
|
||||
optional_params={},
|
||||
)
|
||||
|
||||
session_prices = litellm.model_cost[RESPONSES_MODEL]
|
||||
priced_on_the_session_model = (
|
||||
33 * session_prices["input_cost_per_token"] + 12 * session_prices["output_cost_per_token"]
|
||||
)
|
||||
assert logged_cost == pytest.approx(handled["kwargs"]["response_cost"])
|
||||
assert logged_cost > priced_on_the_session_model
|
||||
assert handled["kwargs"]["model"] == RESPONSES_MODEL
|
||||
|
||||
|
||||
def test_a_turn_the_caller_walked_out_of_still_names_the_model_it_ran():
|
||||
"""
|
||||
A client that closes before response.completed leaves no usage behind to
|
||||
bill, but response.created already named the model, so the row says which
|
||||
model answered instead of falling back to unknown.
|
||||
"""
|
||||
logging_obj = _logging_obj()
|
||||
|
||||
handled = _handle(
|
||||
[{"type": "response.created", "response": {"id": "resp_fake", "model": REASONING_MODEL}}],
|
||||
logging_obj=logging_obj,
|
||||
url_route="/openai_passthrough/v1/responses",
|
||||
)
|
||||
|
||||
assert handled["result"] is None
|
||||
assert handled["kwargs"] == {"model": REASONING_MODEL, "custom_llm_provider": "openai"}
|
||||
assert logging_obj.model_call_details["model"] == REASONING_MODEL
|
||||
|
|
@ -7,20 +7,20 @@ import os
|
|||
import sys
|
||||
import uuid
|
||||
import zlib
|
||||
from collections.abc import AsyncIterator, Callable, Mapping
|
||||
from collections.abc import AsyncIterator, Callable, Mapping, Sequence
|
||||
from contextlib import ExitStack, contextmanager
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from io import BytesIO
|
||||
from types import MappingProxyType, ModuleType, SimpleNamespace
|
||||
from typing import Final, Optional, cast
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
from unittest.mock import AsyncMock, MagicMock, call, patch
|
||||
from uuid import UUID
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
from fastapi import APIRouter, FastAPI, HTTPException, Request, Response, UploadFile
|
||||
from fastapi import APIRouter, FastAPI, HTTPException, Request, Response, UploadFile, WebSocketDisconnect
|
||||
from fastapi.responses import StreamingResponse
|
||||
from fastapi.routing import APIRoute
|
||||
from fastapi.testclient import TestClient
|
||||
|
|
@ -33,6 +33,7 @@ from litellm._logging import verbose_proxy_logger
|
|||
from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.proxy._lazy_features import LazyFeature, attach_lazy_features
|
||||
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
|
||||
|
|
@ -7515,6 +7516,281 @@ async def test_websocket_passthrough_does_not_close_twice_when_success_logging_f
|
|||
websocket.close.assert_awaited_once_with(code=1008, reason=upstream_reason)
|
||||
|
||||
|
||||
SUCCESS_HANDLER: Final = (
|
||||
"litellm.proxy.pass_through_endpoints.pass_through_endpoints."
|
||||
"pass_through_endpoint_logging.pass_through_async_success_handler"
|
||||
)
|
||||
|
||||
|
||||
class RecordingFrameModelGate:
|
||||
def __init__(self, refusals: Mapping[str, str] = MappingProxyType({})):
|
||||
self.refusals = refusals
|
||||
self.checked: list[str] = []
|
||||
|
||||
async def __call__(self, model: str, valid_token: UserAPIKeyAuth, /) -> Optional[str]:
|
||||
self.checked.append(model)
|
||||
return self.refusals.get(model)
|
||||
|
||||
|
||||
def _received_frame(frame):
|
||||
if isinstance(frame, bytes):
|
||||
return {"type": "websocket.receive", "bytes": frame}
|
||||
return {"type": "websocket.receive", "text": frame}
|
||||
|
||||
|
||||
async def _relay_client_frames(frames, gate, user_api_key_dict=None):
|
||||
upstream_ws = RecordingUpstreamWebSocket()
|
||||
websocket = _client_websocket(
|
||||
AsyncMock(side_effect=[_received_frame(frame) for frame in frames] + [{"type": "websocket.disconnect"}])
|
||||
)
|
||||
|
||||
with (
|
||||
_patched_websocket_passthrough_environment(upstream_ws),
|
||||
patch(SUCCESS_HANDLER, new_callable=AsyncMock) as success_handler,
|
||||
):
|
||||
await websocket_passthrough_request(
|
||||
websocket=websocket,
|
||||
target="wss://api.openai.com/v1/realtime?model=gpt-realtime-2.1",
|
||||
custom_headers={"Authorization": "Bearer sk-provider"},
|
||||
user_api_key_dict=user_api_key_dict or UserAPIKeyAuth(),
|
||||
forward_headers=False,
|
||||
endpoint="/openai/v1/realtime",
|
||||
accept_websocket=False,
|
||||
client_frame_model_gate=gate,
|
||||
)
|
||||
|
||||
return SimpleNamespace(upstream=upstream_ws, client=websocket, success_handler=success_handler)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"frame, expected_model",
|
||||
[
|
||||
({"type": "response.create", "model": "gpt-5.4-mini"}, "gpt-5.4-mini"),
|
||||
({"type": "response.create", "response": {"model": "gpt-5.4"}}, "gpt-5.4"),
|
||||
({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}, "gpt-realtime-2.1"),
|
||||
(
|
||||
{"type": "session.update", "session": {"input_audio_transcription": {"model": "gpt-4o-transcribe"}}},
|
||||
"gpt-4o-transcribe",
|
||||
),
|
||||
(
|
||||
{
|
||||
"type": "session.update",
|
||||
"session": {"audio": {"input": {"transcription": {"model": "gpt-transcribe"}}}},
|
||||
},
|
||||
"gpt-transcribe",
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_websocket_passthrough_checks_the_model_a_client_frame_asks_for(frame, expected_model):
|
||||
gate = RecordingFrameModelGate()
|
||||
|
||||
relayed = await _relay_client_frames([json.dumps(frame)], gate)
|
||||
|
||||
assert gate.checked == [expected_model]
|
||||
assert [call.args[0] for call in relayed.upstream.send.await_args_list] == [json.dumps(frame)]
|
||||
assert relayed.client.close.await_args_list == [call()]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_refuses_a_frame_naming_a_model_the_key_cannot_call():
|
||||
refusal = "key not allowed to access model. This key can only access models=['gpt-5.4-mini']."
|
||||
gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": refusal}))
|
||||
frame = json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}})
|
||||
|
||||
relayed = await _relay_client_frames([frame], gate)
|
||||
|
||||
assert relayed.upstream.send.await_args_list == []
|
||||
assert json.loads(relayed.client.send_text.await_args.args[0]) == {
|
||||
"type": "error",
|
||||
"error": {"type": "invalid_request_error", "message": refusal},
|
||||
}
|
||||
assert relayed.client.close.await_args_list == [call(code=1008, reason=refusal)]
|
||||
relayed.upstream.close.assert_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_refuses_the_whole_frame_when_any_named_model_is_denied():
|
||||
refusal = "key not allowed to access model. This key can only access models=['gpt-realtime-2.1']."
|
||||
gate = RecordingFrameModelGate(MappingProxyType({"gpt-4o-transcribe": refusal}))
|
||||
frame = json.dumps(
|
||||
{
|
||||
"type": "session.update",
|
||||
"session": {"model": "gpt-realtime-2.1", "input_audio_transcription": {"model": "gpt-4o-transcribe"}},
|
||||
}
|
||||
)
|
||||
|
||||
relayed = await _relay_client_frames([frame], gate)
|
||||
|
||||
assert gate.checked == ["gpt-realtime-2.1", "gpt-4o-transcribe"]
|
||||
assert relayed.upstream.send.await_args_list == []
|
||||
assert relayed.client.close.await_args_list == [call(code=1008, reason=refusal)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_stops_relaying_after_a_refused_frame():
|
||||
refusal = "key not allowed to access model. This key can only access models=['gpt-5.4-mini']."
|
||||
gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": refusal}))
|
||||
frames = [
|
||||
json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}),
|
||||
json.dumps({"type": "response.create"}),
|
||||
]
|
||||
|
||||
relayed = await _relay_client_frames(frames, gate)
|
||||
|
||||
assert relayed.upstream.send.await_args_list == []
|
||||
assert gate.checked == ["gpt-realtime-2.1"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"frame",
|
||||
[
|
||||
'{"type": "input_audio_buffer.append", "audio": "AAAA"}',
|
||||
'{"type": "response.create", "response": {"modalities": ["audio"]}}',
|
||||
'{"type": "session.update", "session": {"model": 42}}',
|
||||
'not json at all, but it does say "model"',
|
||||
'{"type": "session.update", "session": "model"}',
|
||||
],
|
||||
)
|
||||
async def test_websocket_passthrough_forwards_frames_that_name_no_model(frame):
|
||||
gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": "denied"}))
|
||||
|
||||
relayed = await _relay_client_frames([frame], gate)
|
||||
|
||||
assert gate.checked == []
|
||||
assert [call_.args[0] for call_ in relayed.upstream.send.await_args_list] == [frame]
|
||||
assert relayed.client.close.await_args_list == [call()]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_never_bills_the_model_a_client_frame_named():
|
||||
"""
|
||||
A frame is what the caller asked for, not what the provider ran. A caller
|
||||
who names a cheap model in a frame the provider ignores must not move the
|
||||
session onto that model's price, so only the provider's own frames decide
|
||||
what a session is billed as.
|
||||
"""
|
||||
gate = RecordingFrameModelGate()
|
||||
frames = [
|
||||
json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}),
|
||||
json.dumps({"type": "response.create", "response": {"model": "gpt-realtime-2.1-mini"}}),
|
||||
]
|
||||
|
||||
relayed = await _relay_client_frames(frames, gate)
|
||||
|
||||
assert gate.checked == ["gpt-realtime-2.1", "gpt-realtime-2.1-mini"]
|
||||
assert "model" not in relayed.success_handler.call_args.kwargs
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"frame, expected_model",
|
||||
[
|
||||
('{"type": "response.create", "\\u006dodel": "gpt-5.4"}', "gpt-5.4"),
|
||||
('{"type": "session.update", "session": {"\\u006d\\u006fdel": "gpt-realtime-2.1"}}', "gpt-realtime-2.1"),
|
||||
],
|
||||
)
|
||||
async def test_websocket_passthrough_checks_a_model_named_with_escaped_json(frame, expected_model):
|
||||
"""
|
||||
JSON lets a caller spell the same key several ways, so a frame carrying
|
||||
"\\u006dodel" asks for a model just as plainly as one carrying "model" and
|
||||
has to face the same check.
|
||||
"""
|
||||
refusal = "key not allowed to access model. This key can only access models=['gpt-5.4-mini']."
|
||||
gate = RecordingFrameModelGate(MappingProxyType({expected_model: refusal}))
|
||||
|
||||
relayed = await _relay_client_frames([frame], gate)
|
||||
|
||||
assert gate.checked == [expected_model]
|
||||
assert relayed.upstream.send.await_args_list == []
|
||||
assert relayed.client.close.await_args_list == [call(code=1008, reason=refusal)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_checks_the_model_a_binary_frame_asks_for():
|
||||
refusal = "key not allowed to access model. This key can only access models=['gpt-5.4-mini']."
|
||||
gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": refusal}))
|
||||
frame = json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}).encode()
|
||||
|
||||
relayed = await _relay_client_frames([frame], gate)
|
||||
|
||||
assert gate.checked == ["gpt-realtime-2.1"]
|
||||
assert relayed.upstream.send.await_args_list == []
|
||||
assert relayed.client.close.await_args_list == [call(code=1008, reason=refusal)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_forwards_binary_audio_untouched():
|
||||
gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": "denied"}))
|
||||
audio = bytes(range(256)) * 4
|
||||
|
||||
relayed = await _relay_client_frames([audio], gate)
|
||||
|
||||
assert gate.checked == []
|
||||
assert [call_.args[0] for call_ in relayed.upstream.send.await_args_list] == [audio]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_without_a_frame_gate_forwards_every_frame():
|
||||
frame = json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}})
|
||||
|
||||
relayed = await _relay_client_frames([frame], None)
|
||||
|
||||
assert [call_.args[0] for call_ in relayed.upstream.send.await_args_list] == [frame]
|
||||
assert "model" not in relayed.success_handler.call_args.kwargs
|
||||
|
||||
|
||||
REALTIME_SESSION_CREATED: Final = json.dumps(
|
||||
{"type": "session.created", "session": {"id": "sess_fake", "model": "gpt-realtime-2.1"}}
|
||||
)
|
||||
REALTIME_RESPONSE_DONE: Final = json.dumps(
|
||||
{
|
||||
"type": "response.done",
|
||||
"response": {"status": "completed", "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}},
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_still_bills_a_session_the_client_walked_out_of():
|
||||
"""
|
||||
A realtime client closes the moment it has the answer it asked for, so the
|
||||
next frame the provider sends lands on a socket that is already gone. That
|
||||
is how a session ends, not a gateway failure, and the turns the caller
|
||||
already used still have to reach spend, including the last frame the client
|
||||
never collected.
|
||||
"""
|
||||
upstream_ws = FakeUpstreamWebSocket(
|
||||
REALTIME_SESSION_CREATED,
|
||||
REALTIME_RESPONSE_DONE,
|
||||
json.dumps({"type": "response.output_text.delta", "delta": "too late"}),
|
||||
)
|
||||
websocket = _client_websocket(_pending_receive)
|
||||
websocket.send_text = AsyncMock(side_effect=[None, None, WebSocketDisconnect(code=1006)])
|
||||
|
||||
with (
|
||||
_patched_websocket_passthrough_environment(upstream_ws),
|
||||
patch(SUCCESS_HANDLER, new_callable=AsyncMock) as success_handler,
|
||||
):
|
||||
await websocket_passthrough_request(
|
||||
websocket=websocket,
|
||||
target="wss://api.openai.com/v1/realtime?model=gpt-realtime-2.1",
|
||||
custom_headers={"Authorization": "Bearer sk-provider"},
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
forward_headers=False,
|
||||
endpoint="/openai_passthrough/v1/realtime",
|
||||
accept_websocket=False,
|
||||
)
|
||||
|
||||
relayed_frames: Final = success_handler.call_args.kwargs["response_body"]
|
||||
assert [frame["type"] for frame in relayed_frames] == [
|
||||
"session.created",
|
||||
"response.done",
|
||||
"response.output_text.delta",
|
||||
]
|
||||
|
||||
|
||||
DEEPGRAM_LISTEN_TARGET = "wss://api.deepgram.com/v1/listen?model=nova-3&encoding=linear16&sample_rate=16000"
|
||||
DEEPGRAM_INTERIM_FRAME = json.dumps(
|
||||
{
|
||||
|
|
@ -8442,6 +8718,41 @@ def test_passthrough_client_cannot_forge_session_id_omission(client_metadata_key
|
|||
)
|
||||
|
||||
|
||||
class StampingFrameModelGate:
|
||||
def __init__(self, groups: list[str]):
|
||||
self.groups = groups
|
||||
|
||||
async def __call__(self, model: str, valid_token: UserAPIKeyAuth, /) -> Optional[str]:
|
||||
valid_token.matched_model_access_groups = self.groups
|
||||
return None
|
||||
|
||||
|
||||
async def _billed_metadata(gate) -> dict:
|
||||
served = await _relay_client_frames(
|
||||
[json.dumps({"type": "response.create", "model": "gpt-5.4-mini"})],
|
||||
gate,
|
||||
UserAPIKeyAuth(api_key="hashed-key"),
|
||||
)
|
||||
return served.success_handler.call_args.kwargs["litellm_params"]["metadata"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_session_is_charged_to_the_group_its_frames_were_authorized_through():
|
||||
"""A websocket names its model in a frame, long after the relay built its logging metadata, so the
|
||||
groups the gate matched have to be picked up at the end or the group pays for nothing it allowed."""
|
||||
metadata = await _billed_metadata(StampingFrameModelGate(["tier-a"]))
|
||||
|
||||
assert metadata[MODEL_ACCESS_GROUP_METADATA_KEY] == ["tier-a"]
|
||||
assert metadata["user_api_key"] == "hashed-key"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_session_that_matched_no_group_is_charged_to_none():
|
||||
metadata = await _billed_metadata(RecordingFrameModelGate())
|
||||
|
||||
assert not metadata[MODEL_ACCESS_GROUP_METADATA_KEY]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("client_metadata_key", ["litellm_metadata", "metadata"])
|
||||
def test_passthrough_logs_the_resolved_deployment_model_info_over_the_request_body(client_metadata_key: str):
|
||||
"""A provider route that resolved a router deployment stashes its model_info on request.state. That
|
||||
|
|
|
|||
|
|
@ -648,3 +648,57 @@ def test_vertex_ai_live_route_without_usage_frames_yields_no_logging_response(
|
|||
|
||||
assert normalized["standard_logging_response_object"] is None
|
||||
assert normalized["kwargs"] == {"model": _LIVE_MODEL}
|
||||
|
||||
|
||||
REALTIME_FRAMES: Final = [
|
||||
{"type": "session.created", "session": {"id": "sess_fake", "model": "gpt-realtime-2.1"}},
|
||||
{
|
||||
"type": "response.done",
|
||||
"response": {
|
||||
"usage": {
|
||||
"input_tokens": 33,
|
||||
"output_tokens": 16,
|
||||
"total_tokens": 49,
|
||||
"input_token_details": {"text_tokens": 33, "audio_tokens": 0, "cached_tokens": 0},
|
||||
"output_token_details": {"text_tokens": 16, "audio_tokens": 0},
|
||||
}
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def _normalize(url_route: str, response_body):
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.model_call_details = {}
|
||||
return PassThroughEndpointLogging().normalize_llm_passthrough_logging_payload(
|
||||
httpx_response=MagicMock(spec=httpx.Response),
|
||||
response_body=response_body,
|
||||
request_body={},
|
||||
logging_obj=logging_obj,
|
||||
url_route=url_route,
|
||||
result="websocket_connection_successful",
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
cache_hit=False,
|
||||
litellm_call_id="call-fake",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("url_route", ["/openai/v1/realtime", "/openai_passthrough/v1/realtime"])
|
||||
def test_openai_websocket_frames_reach_the_websocket_cost_handler(url_route):
|
||||
normalized = _normalize(url_route, REALTIME_FRAMES)
|
||||
|
||||
assert normalized["kwargs"]["model"] == "gpt-realtime-2.1"
|
||||
assert normalized["kwargs"]["custom_llm_provider"] == "openai"
|
||||
prices = litellm.model_cost["gpt-realtime-2.1"]
|
||||
assert normalized["kwargs"]["response_cost"] == pytest.approx(
|
||||
33 * prices["input_cost_per_token"] + 16 * prices["output_cost_per_token"]
|
||||
)
|
||||
assert normalized["standard_logging_response_object"].usage.total_tokens == 49
|
||||
|
||||
|
||||
def test_openai_http_passthrough_bodies_are_left_to_the_http_handlers():
|
||||
normalized = _normalize("/openai/v1/responses", {"model": "gpt-5.4-mini", "usage": {"input_tokens": 13}})
|
||||
|
||||
assert normalized["standard_logging_response_object"] is None
|
||||
assert normalized["kwargs"] == {"litellm_call_id": "call-fake"}
|
||||
|
|
|
|||
|
|
@ -1,11 +1,12 @@
|
|||
"""OpenAI passthrough WebSocket route: registration, opt-in gating, and refusals."""
|
||||
"""OpenAI passthrough WebSocket route: registration, opt-in gating, refusals, and per-frame model checks."""
|
||||
|
||||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from collections.abc import Iterator, Mapping
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType, SimpleNamespace
|
||||
from typing import Final
|
||||
from unittest.mock import patch
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
|
@ -14,18 +15,16 @@ from starlette.routing import WebSocketRoute
|
|||
|
||||
import litellm
|
||||
from litellm.llms.openai.workload_identity import _workload_identity_auth
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy._types import Litellm_EntityType, LiteLLM_UserTable, UserAPIKeyAuth
|
||||
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
|
||||
_OPENAI_WS_DISABLED_REFUSAL,
|
||||
_OPENAI_WS_MODEL_RESTRICTED_REFUSAL,
|
||||
_has_model_restrictions,
|
||||
_OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL,
|
||||
_openai_websocket_refusal,
|
||||
_proxy_model_allowlists,
|
||||
_proxy_frame_model_gate,
|
||||
openai_websocket_proxy_route,
|
||||
router,
|
||||
)
|
||||
|
||||
Scopes = tuple[Sequence[str], ...]
|
||||
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import WebsocketFrameModelGate
|
||||
|
||||
ENABLED: Final = MappingProxyType({"enable_openai_websocket_passthrough": True})
|
||||
DISABLED_SETTINGS: Final = (
|
||||
|
|
@ -76,6 +75,7 @@ class _RelayCall:
|
|||
forward_headers: bool
|
||||
endpoint: str
|
||||
accept_websocket: bool
|
||||
client_frame_model_gate: WebsocketFrameModelGate | None
|
||||
|
||||
|
||||
class _FakeRelay:
|
||||
|
|
@ -92,6 +92,7 @@ class _FakeRelay:
|
|||
forward_headers: bool,
|
||||
endpoint: str,
|
||||
accept_websocket: bool,
|
||||
client_frame_model_gate: WebsocketFrameModelGate | None = None,
|
||||
) -> None:
|
||||
self.calls.append(
|
||||
_RelayCall(
|
||||
|
|
@ -100,24 +101,25 @@ class _FakeRelay:
|
|||
forward_headers=forward_headers,
|
||||
endpoint=endpoint,
|
||||
accept_websocket=accept_websocket,
|
||||
client_frame_model_gate=client_frame_model_gate,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class _FakeModelAllowlists:
|
||||
def __init__(self, scopes: Scopes) -> None:
|
||||
self.scopes = scopes
|
||||
self.calls: list[UserAPIKeyAuth] = []
|
||||
class _FakeFrameModelGate:
|
||||
def __init__(self, refusal: str | None = None) -> None:
|
||||
self.refusal = refusal
|
||||
self.checked: list[str] = []
|
||||
|
||||
async def __call__(self, valid_token: UserAPIKeyAuth, /) -> Scopes:
|
||||
self.calls.append(valid_token)
|
||||
return self.scopes
|
||||
async def __call__(self, model: str, valid_token: UserAPIKeyAuth, /) -> str | None:
|
||||
self.checked.append(model)
|
||||
return self.refusal
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Served:
|
||||
relay: _FakeRelay
|
||||
allowlists: _FakeModelAllowlists
|
||||
frame_model_gate: _FakeFrameModelGate
|
||||
|
||||
|
||||
async def _serve(
|
||||
|
|
@ -125,16 +127,15 @@ async def _serve(
|
|||
endpoint: str,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
general_settings: Mapping[str, object],
|
||||
scopes: Scopes = (),
|
||||
) -> _Served:
|
||||
served = _Served(relay=_FakeRelay(), allowlists=_FakeModelAllowlists(scopes))
|
||||
served = _Served(relay=_FakeRelay(), frame_model_gate=_FakeFrameModelGate())
|
||||
await openai_websocket_proxy_route(
|
||||
websocket=websocket,
|
||||
endpoint=endpoint,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
general_settings=general_settings,
|
||||
relay=served.relay,
|
||||
model_allowlists=served.allowlists,
|
||||
frame_model_gate=served.frame_model_gate,
|
||||
)
|
||||
return served
|
||||
|
||||
|
|
@ -143,18 +144,19 @@ async def _serve(
|
|||
@pytest.mark.parametrize("prefix", ["openai", "openai_passthrough"])
|
||||
async def test_openai_websocket_forwards_query_and_keeps_provider_auth(prefix, monkeypatch):
|
||||
monkeypatch.delenv("OPENAI_API_BASE", raising=False)
|
||||
websocket = _FakeWebSocket(f"/{prefix}/v1/realtime", "model=gpt-4o-realtime-preview")
|
||||
websocket = _FakeWebSocket(f"/{prefix}/v1/realtime", "model=gpt-realtime-2.1")
|
||||
|
||||
with patch(GET_CREDENTIALS, return_value="sk-provider"):
|
||||
served = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), ENABLED)
|
||||
|
||||
assert served.relay.calls == [
|
||||
_RelayCall(
|
||||
target="wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview",
|
||||
target="wss://api.openai.com/v1/realtime?model=gpt-realtime-2.1",
|
||||
custom_headers=MappingProxyType({"Authorization": "Bearer sk-provider"}),
|
||||
forward_headers=False,
|
||||
endpoint=f"/{prefix}/v1/realtime",
|
||||
accept_websocket=False,
|
||||
client_frame_model_gate=served.frame_model_gate,
|
||||
)
|
||||
]
|
||||
assert websocket.accepts == [None]
|
||||
|
|
@ -166,7 +168,7 @@ async def test_openai_websocket_forwards_query_and_keeps_provider_auth(prefix, m
|
|||
async def test_openai_websocket_accepts_first_client_subprotocol():
|
||||
websocket = _FakeWebSocket(
|
||||
"/openai/v1/realtime",
|
||||
"model=gpt-4o-realtime-preview",
|
||||
"model=gpt-realtime-2.1",
|
||||
subprotocols="realtime, openai-insecure-api-key.sk-abc, openai-beta.realtime-v1",
|
||||
)
|
||||
|
||||
|
|
@ -239,7 +241,7 @@ async def test_openai_websocket_closes_cleanly_when_workload_identity_exchange_f
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_websocket_closes_cleanly_when_provider_credentials_missing():
|
||||
websocket = _FakeWebSocket("/openai/v1/realtime", "model=gpt-4o-realtime-preview")
|
||||
websocket = _FakeWebSocket("/openai/v1/realtime", "model=gpt-realtime-2.1")
|
||||
|
||||
with patch(GET_CREDENTIALS, return_value=None):
|
||||
served = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), ENABLED)
|
||||
|
|
@ -255,7 +257,7 @@ async def test_openai_websocket_closes_cleanly_when_provider_credentials_missing
|
|||
@pytest.mark.parametrize("prefix", ["openai", "openai_passthrough"])
|
||||
@pytest.mark.parametrize("general_settings", DISABLED_SETTINGS)
|
||||
async def test_openai_websocket_refused_unless_explicitly_enabled(prefix, general_settings):
|
||||
websocket = _FakeWebSocket(f"/{prefix}/v1/realtime", "model=gpt-4o-realtime-preview")
|
||||
websocket = _FakeWebSocket(f"/{prefix}/v1/realtime", "model=gpt-realtime-2.1")
|
||||
|
||||
served = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), general_settings)
|
||||
|
||||
|
|
@ -265,25 +267,21 @@ async def test_openai_websocket_refused_unless_explicitly_enabled(prefix, genera
|
|||
assert served.relay.calls == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("general_settings", DISABLED_SETTINGS)
|
||||
async def test_openai_websocket_refusal_is_disabled_for_falsy_settings(general_settings):
|
||||
refusal = await _openai_websocket_refusal(UserAPIKeyAuth(), general_settings, _FakeModelAllowlists(()))
|
||||
assert refusal is _OPENAI_WS_DISABLED_REFUSAL
|
||||
def test_openai_websocket_refusal_is_disabled_for_falsy_settings(general_settings):
|
||||
assert _openai_websocket_refusal(general_settings) is _OPENAI_WS_DISABLED_REFUSAL
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("value", [True, "true", "True"])
|
||||
async def test_openai_websocket_refusal_is_none_for_truthy_settings(value):
|
||||
settings = MappingProxyType({"enable_openai_websocket_passthrough": value})
|
||||
assert await _openai_websocket_refusal(UserAPIKeyAuth(), settings, _FakeModelAllowlists(())) is None
|
||||
def test_openai_websocket_refusal_is_none_for_truthy_settings(value):
|
||||
assert _openai_websocket_refusal(MappingProxyType({"enable_openai_websocket_passthrough": value})) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_websocket_refusal_echoes_requested_subprotocol():
|
||||
websocket = _FakeWebSocket(
|
||||
"/openai_passthrough/v1/realtime",
|
||||
"model=gpt-4o-realtime-preview",
|
||||
"model=gpt-realtime-2.1",
|
||||
subprotocols="realtime, openai-beta.realtime-v1",
|
||||
)
|
||||
|
||||
|
|
@ -294,68 +292,190 @@ async def test_openai_websocket_refusal_echoes_requested_subprotocol():
|
|||
assert served.relay.calls == []
|
||||
|
||||
|
||||
RESTRICTED_SCOPES: Final[tuple[Scopes, ...]] = (
|
||||
(("gpt-4o",),),
|
||||
((), ("gpt-4o-realtime-preview",)),
|
||||
(("all-team-models",), ("gpt-4o",)),
|
||||
((), ("all-proxy-models",), ("gpt-4o",)),
|
||||
((), (), (), ("gpt-4o",)),
|
||||
(("*",), (), (), (), ("gpt-4o",)),
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"key_models",
|
||||
[["gpt-5.4-mini"], ["gpt-realtime-2.1", "gpt-5.4"], []],
|
||||
)
|
||||
UNRESTRICTED_SCOPES: Final[tuple[Scopes, ...]] = (
|
||||
(),
|
||||
((),),
|
||||
(("all-proxy-models",),),
|
||||
(("*",),),
|
||||
(("all-team-models",), ("all-proxy-models",)),
|
||||
((), (), (), (), ()),
|
||||
(("*",), ("all-proxy-models",), ("all-team-models",), (), ()),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("scopes", RESTRICTED_SCOPES)
|
||||
async def test_openai_websocket_rejects_model_restricted_identities(scopes):
|
||||
websocket = _FakeWebSocket("/openai/v1/realtime", "model=gpt-4o-realtime-preview")
|
||||
user_api_key_dict = UserAPIKeyAuth(token="hashed-fake", user_id="user-fake", team_id="team-fake")
|
||||
|
||||
served = await _serve(websocket, "v1/realtime", user_api_key_dict, ENABLED, scopes)
|
||||
|
||||
assert "model restrictions" in websocket.error_message()
|
||||
assert websocket.closed == (1008, _OPENAI_WS_MODEL_RESTRICTED_REFUSAL.close_reason)
|
||||
assert served.relay.calls == []
|
||||
assert served.allowlists.calls == [user_api_key_dict]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("scopes", RESTRICTED_SCOPES)
|
||||
async def test_openai_websocket_disabled_refusal_skips_allowlist_lookups(scopes):
|
||||
allowlists = _FakeModelAllowlists(scopes)
|
||||
|
||||
refusal = await _openai_websocket_refusal(UserAPIKeyAuth(), MappingProxyType({}), allowlists)
|
||||
|
||||
assert refusal is _OPENAI_WS_DISABLED_REFUSAL
|
||||
assert allowlists.calls == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("scopes", UNRESTRICTED_SCOPES)
|
||||
async def test_openai_websocket_allows_unrestricted_identities(scopes):
|
||||
async def test_openai_websocket_connects_model_restricted_keys(key_models):
|
||||
websocket = _FakeWebSocket("/openai/v1/responses", "")
|
||||
user_api_key_dict = UserAPIKeyAuth(models=key_models, token="hashed-fake", user_id="user-fake")
|
||||
|
||||
with patch(GET_CREDENTIALS, return_value="sk-provider"):
|
||||
served = await _serve(websocket, "v1/responses", UserAPIKeyAuth(), ENABLED, scopes)
|
||||
served = await _serve(websocket, "v1/responses", user_api_key_dict, ENABLED)
|
||||
|
||||
assert len(served.relay.calls) == 1
|
||||
assert served.relay.calls[0].client_frame_model_gate is served.frame_model_gate
|
||||
assert websocket.sent == []
|
||||
assert websocket.closed is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_model_allowlists_reads_the_token_scopes_without_a_database():
|
||||
token: Final = UserAPIKeyAuth(models=[], team_id="team-fake", team_models=["gpt-4o"])
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", None):
|
||||
scopes = await _proxy_model_allowlists()(token)
|
||||
async def test_proxy_frame_model_gate_allows_a_model_the_key_owns():
|
||||
gate: Final = _proxy_frame_model_gate()
|
||||
token: Final = UserAPIKeyAuth(models=["gpt-realtime-2.1", "gpt-5.4-mini"])
|
||||
|
||||
assert tuple(tuple(scope) for scope in scopes) == ((), ("gpt-4o",))
|
||||
assert _has_model_restrictions(scopes)
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", None): # test-quality-ok: the gate needs no database
|
||||
assert await gate("gpt-realtime-2.1", token) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("requested_model", ["gpt-5.4", "gpt-realtime-2.1-mini"])
|
||||
async def test_proxy_frame_model_gate_refuses_a_model_the_key_lacks(requested_model):
|
||||
gate: Final = _proxy_frame_model_gate()
|
||||
token: Final = UserAPIKeyAuth(models=["gpt-realtime-2.1", "gpt-5.4-mini"])
|
||||
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", None): # test-quality-ok: the gate needs no database
|
||||
refusal = await gate(requested_model, token)
|
||||
|
||||
assert refusal is not None
|
||||
assert requested_model in refusal
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"granted_models, refusal_names",
|
||||
[(["gpt-5.4-mini"], "gpt-5.4"), ([], "no model grants")],
|
||||
ids=["model_outside_the_grant", "no_grants"],
|
||||
)
|
||||
async def test_proxy_frame_model_gate_refuses_a_model_the_agent_was_not_granted(granted_models, refusal_names):
|
||||
from litellm.types.agents import AgentResponse
|
||||
|
||||
gate: Final = _proxy_frame_model_gate()
|
||||
token: Final = UserAPIKeyAuth(models=[], agent_id="agent-fake")
|
||||
token.managed_agent_policy = AgentResponse(
|
||||
agent_id="agent-fake",
|
||||
agent_name="Agent",
|
||||
agent_card_params={},
|
||||
object_permission={"object_permission_id": "grant-fake", "models": granted_models},
|
||||
)
|
||||
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", None): # test-quality-ok: the gate needs no database
|
||||
refusal = await gate("gpt-5.4", token)
|
||||
|
||||
assert refusal is not None
|
||||
assert refusal_names in refusal
|
||||
|
||||
|
||||
GET_USER_OBJECT: Final = "litellm.proxy.auth.auth_checks.get_user_object"
|
||||
|
||||
|
||||
def _personal_user(models: list[str]) -> LiteLLM_UserTable:
|
||||
return LiteLLM_UserTable(user_id="user-fake", models=models, max_budget=None, spend=0.0)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _personal_user_allowlist(models: list[str]) -> Iterator[AsyncMock]:
|
||||
"""Stand up the database-backed owner lookup the allowlist check runs, and report every call it made."""
|
||||
lookup: Final = AsyncMock(return_value=_personal_user(models))
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), # test-quality-ok: read at call time
|
||||
patch(GET_USER_OBJECT, lookup),
|
||||
):
|
||||
yield lookup
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_frame_model_gate_refuses_a_model_the_person_behind_the_key_lacks():
|
||||
"""
|
||||
A key with no team of its own carries its owner's model allowlist, the same
|
||||
one every HTTP route applies, so a frame naming a model that person cannot
|
||||
use is refused even when the key itself names no models at all.
|
||||
"""
|
||||
gate: Final = _proxy_frame_model_gate()
|
||||
token: Final = UserAPIKeyAuth(models=[], user_id="user-fake")
|
||||
|
||||
with _personal_user_allowlist(["gpt-5.4-mini"]):
|
||||
refusal = await gate("gpt-5.4", token)
|
||||
|
||||
assert refusal is not None
|
||||
assert "gpt-5.4" in refusal
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_frame_model_gate_allows_a_model_the_person_behind_the_key_owns():
|
||||
gate: Final = _proxy_frame_model_gate()
|
||||
token: Final = UserAPIKeyAuth(models=[], user_id="user-fake")
|
||||
|
||||
with _personal_user_allowlist(["gpt-5.4-mini", "gpt-realtime-2.1"]):
|
||||
assert await gate("gpt-realtime-2.1", token) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_frame_model_gate_leaves_a_team_key_to_the_teams_allowlist():
|
||||
"""
|
||||
A key that belongs to a team is governed by the team's models, not by the
|
||||
personal allowlist of whoever created it, so the owner's own restrictions
|
||||
must not close a session the team is entitled to.
|
||||
"""
|
||||
gate: Final = _proxy_frame_model_gate()
|
||||
token: Final = UserAPIKeyAuth(models=[], user_id="user-fake", team_id="team-fake")
|
||||
|
||||
with _personal_user_allowlist(["gpt-5.4-mini"]) as owner_lookup:
|
||||
assert await gate("gpt-5.4", token) is None
|
||||
|
||||
assert owner_lookup.await_args_list == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("unrestricted_models", [[], ["all-proxy-models"], ["*"]])
|
||||
async def test_proxy_frame_model_gate_allows_every_model_for_unrestricted_keys(unrestricted_models):
|
||||
gate: Final = _proxy_frame_model_gate()
|
||||
token: Final = UserAPIKeyAuth(models=unrestricted_models)
|
||||
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", None): # test-quality-ok: the gate needs no database
|
||||
assert await gate("gpt-realtime-2.1", token) is None
|
||||
|
||||
|
||||
ENFORCE_GROUP_BUDGET: Final = (
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.enforce_model_access_group_budget"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_frame_model_gate_refuses_a_model_whose_access_group_is_out_of_budget():
|
||||
"""
|
||||
A model that sits in a budgeted access group is refused on every other route once
|
||||
that group has spent its budget, so a frame naming it closes the session instead
|
||||
of quietly running on an empty pool. The key is allowed the model outright here,
|
||||
so nothing but the group's budget can produce the refusal.
|
||||
"""
|
||||
exhausted: Final = litellm.BudgetExceededError(
|
||||
current_cost=2.0,
|
||||
max_budget=1.0,
|
||||
message="Budget has been exceeded! Model access group=tier-a Current cost: 2.0, Max budget: 1.0",
|
||||
entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP.value,
|
||||
entity_id="tier-a",
|
||||
)
|
||||
gate: Final = _proxy_frame_model_gate()
|
||||
token: Final = UserAPIKeyAuth(models=["gpt-5.4-mini"])
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", None), # test-quality-ok: the model check needs no database
|
||||
patch(ENFORCE_GROUP_BUDGET, AsyncMock(side_effect=exhausted)),
|
||||
):
|
||||
refusal = await gate("gpt-5.4-mini", token)
|
||||
|
||||
assert refusal is not None
|
||||
assert "Budget has been exceeded" in refusal
|
||||
assert "Model access group=tier-a" in refusal
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_frame_model_gate_refuses_out_loud_when_it_cannot_reach_the_budgets():
|
||||
"""
|
||||
A gate that cannot answer has to close the session with a reason the caller can read.
|
||||
Letting the failure escape drops the socket with no frame at all, which the caller
|
||||
cannot tell apart from the network going away.
|
||||
"""
|
||||
gate: Final = _proxy_frame_model_gate()
|
||||
token: Final = UserAPIKeyAuth(models=["gpt-5.4-mini"])
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", None), # test-quality-ok: the model check needs no database
|
||||
patch(ENFORCE_GROUP_BUDGET, AsyncMock(side_effect=ConnectionError("connection pool is exhausted"))),
|
||||
):
|
||||
refusal = await gate("gpt-5.4-mini", token)
|
||||
|
||||
assert refusal == _OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL
|
||||
assert "connection pool" not in refusal
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue