diff --git a/.github/workflows/_test-unit-base.yml b/.github/workflows/_test-unit-base.yml index 62790e23143..4d939b3e838 100644 --- a/.github/workflows/_test-unit-base.yml +++ b/.github/workflows/_test-unit-base.yml @@ -68,6 +68,16 @@ jobs: pull-requests: read outputs: decision: ${{ steps.changes.outputs.decision }} + services: + redis: + image: ${{ inputs.artifact-name == 'proxy-auth' && 'redis:8.2.9-alpine@sha256:30abb90e62f14b737010746def3ba99cc79fe19dcdb3d37b41f21fc62e7da19d' || '' }} + ports: + - '127.0.0.1::6379' + options: >- + --health-cmd "redis-cli ping" + --health-interval 2s + --health-timeout 2s + --health-retries 15 steps: - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 @@ -137,6 +147,7 @@ jobs: RERUNS: ${{ inputs.reruns }} DIST: ${{ inputs.dist }} COVERAGE_CORE: sysmon + LITELLM_TEST_REDIS_PORT: ${{ job.services.redis.ports['6379'] }} run: | if [ "${WORKERS}" = "0" ]; then uv run --no-sync pytest ${TEST_PATH:?} \ diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index f55c87c2ae5..cdbbe28f1cd 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -133,6 +133,7 @@ jobs: tests/test_litellm/proxy/hooks tests/test_litellm/proxy/policy_engine tests/test_litellm/proxy/client + tests/local_testing/test_realtime_call_redis.py workers: 2 reruns: 2 timeout-minutes: 20 diff --git a/gateway/routes/allowlist.py b/gateway/routes/allowlist.py index 3733072a948..02b4664bf65 100644 --- a/gateway/routes/allowlist.py +++ b/gateway/routes/allowlist.py @@ -105,6 +105,8 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = ( "/{provider}/", "/toolset/", # Realtime / streaming + "/v1/live", + "/live", "/v1/realtime", "/realtime", # Health & ops diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 814eaaf76f7..15060004945 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -8,7 +8,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, cast from httpx import Response -from pydantic import BaseModel +from pydantic import BaseModel, Field, ValidationError import litellm import litellm._logging @@ -2563,7 +2563,12 @@ def handle_realtime_stream_cost_calculation( if any(r.get("type") == _TRANSCRIPTION_COMPLETED_EVENT_TYPE for r in results) else 0.0 ) - total_cost: Final = input_cost_per_token + output_cost_per_token + transcription_cost + live_audio_cost: Final = handle_live_session_duration_cost( + results=results, + custom_llm_provider=custom_llm_provider, + litellm_model_name=litellm_model_name, + ) + total_cost: Final = input_cost_per_token + output_cost_per_token + transcription_cost + live_audio_cost _store_cost_breakdown_in_logging_obj( litellm_logging_obj=litellm_logging_obj, @@ -2571,13 +2576,45 @@ def handle_realtime_stream_cost_calculation( completion_tokens_cost_usd_dollar=output_cost_per_token, cost_for_built_in_tools_cost_usd_dollar=0.0, total_cost_usd_dollar=total_cost, - additional_costs={"transcription_cost": transcription_cost} if transcription_cost > 0 else None, + additional_costs={ # mutable-ok: logging cost breakdown requires a concrete dict + name: cost + for name, cost in (("transcription_cost", transcription_cost), ("live_audio_cost", live_audio_cost)) + if cost > 0 + } + or None, data_residency=data_residency, ) return total_cost +class _LiveSessionDurationUsage(BaseModel): + audio_duration_ms: float = Field(strict=True, ge=0, allow_inf_nan=False) + + +class _LiveSessionClosedEvent(BaseModel): + usage: _LiveSessionDurationUsage + + +def handle_live_session_duration_cost( + results: OpenAIRealtimeStreamList, + custom_llm_provider: str, + litellm_model_name: str, +) -> float: + terminal: Final = next((event for event in reversed(results) if event.get("type") == "session.closed"), None) + if terminal is None: + return 0.0 + try: + usage: Final = _LiveSessionClosedEvent.model_validate(terminal).usage + except ValidationError: + return 0.0 + try: + model_info: Final = litellm.get_model_info(model=litellm_model_name, custom_llm_provider=custom_llm_provider) + except Exception: + return 0.0 + return usage.audio_duration_ms / 1000 * (model_info.get("input_cost_per_second") or 0.0) + + def handle_realtime_transcription_cost_calculation( results: OpenAIRealtimeStreamList, custom_llm_provider: str, diff --git a/litellm/images/main.py b/litellm/images/main.py index 6a94e7c8df2..11df9728ede 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -377,6 +377,7 @@ def image_generation( # Providers using llm_http_handler ######################################################### elif custom_llm_provider in ( + litellm.LlmProviders.CHATGPT, litellm.LlmProviders.RECRAFT, litellm.LlmProviders.AIML, litellm.LlmProviders.GEMINI, @@ -401,6 +402,7 @@ def image_generation( model=model, prompt=prompt, image_generation_provider_config=image_generation_config, + extra_headers=extra_headers, image_generation_optional_request_params=optional_params, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params_dict, @@ -965,9 +967,9 @@ def image_edit( @client async def aimage_edit( - image: FileTypes | list[FileTypes], - model: str, - prompt: str, + image: FileTypes | list[FileTypes] | None = None, + model: str = "", + prompt: str = "", mask: str | None = None, n: int | None = None, quality: str | ImageGenerationRequestQuality | None = None, @@ -1005,11 +1007,9 @@ async def aimage_edit( model=model, api_base=local_vars.get("base_url", None) ) - images: Final = image if isinstance(image, list) else [image] - func: Final = partial( image_edit, - image=images, + image=image, prompt=prompt, mask=mask, model=model, diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 4923bdda305..09b646b7c70 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -1,11 +1,13 @@ import asyncio import json import traceback -from collections.abc import Coroutine, Mapping, Sequence +from collections.abc import Awaitable, Callable, Coroutine, Mapping, Sequence +from contextvars import ContextVar from dataclasses import dataclass from enum import Enum, auto from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypedDict, cast +from pydantic import TypeAdapter from typing_extensions import ReadOnly import litellm @@ -16,6 +18,7 @@ from litellm.types.llms.openai import ( OpenAIRealtimeEvents, OpenAIRealtimeOutputItemDone, OpenAIRealtimeResponseDelta, + OpenAIRealtimeSessionClosed, OpenAIRealtimeStreamResponseBaseObject, OpenAIRealtimeStreamSessionEvents, ) @@ -24,6 +27,10 @@ from litellm.types.realtime import ALL_DELTA_TYPES from .litellm_logging import Logging as LiteLLMLogging from .realtime_errors import client_close_code, realtime_error_event, websocket_close_reason +realtime_attachment_cleanup: Final[ContextVar[Callable[[], Awaitable[None]] | None]] = ContextVar( + "realtime_attachment_cleanup", default=None +) + if TYPE_CHECKING: from websockets.asyncio.client import ClientConnection from websockets.exceptions import ConnectionClosed @@ -139,11 +146,14 @@ class RealTimeStreaming: force_transcription_model: str | None = None, event_normalizer: RealtimeEventNormalizer | None = None, logging_worker: _LoggingWorker = GLOBAL_LOGGING_WORKER, + *, + account_usage: bool = True, ): self.websocket: _ClientWebSocket = websocket self.backend_ws = backend_ws self.logging_obj = logging_obj self._logging_worker = logging_worker + self._account_usage = account_usage self.messages: list[OpenAIRealtimeEvents] = [] self._backend_sent_frames: bool = False self.input_message: dict = {} @@ -256,6 +266,9 @@ class RealTimeStreaming: else: message_obj = cast(dict[str, Any], json.loads(cast(str, message))) self._collect_tool_calls_from_response_done(cast(dict, message_obj)) + if message_obj.get("type") == "session.closed" and isinstance(message_obj.get("usage"), dict): + self.messages.append(TypeAdapter(OpenAIRealtimeSessionClosed).validate_python(message_obj)) + return if not self._should_store_message(message_obj): return try: @@ -410,8 +423,10 @@ class RealTimeStreaming: if self.logging_obj: self.logging_obj.pre_call(input=message, api_key="") - async def log_messages(self): + async def log_messages(self, *, wait_for_dispatch: bool = False): """Log messages in list""" + if not self._account_usage: + return if self.logging_obj: if self.input_messages: self.logging_obj.model_call_details["messages"] = self.input_messages @@ -421,9 +436,12 @@ class RealTimeStreaming: # Route through the bounded logging worker (per-coroutine timeout + # concurrency cap) instead of a bare create_task, so a slow callback # can't leave suspended tasks pinning each call's response in memory. - self._logging_worker.ensure_initialized_and_enqueue( - self.logging_obj.dispatch_success_handlers(self.messages, prefer_async_handlers=True) - ) + if wait_for_dispatch: + await self.logging_obj.dispatch_success_handlers(self.messages, prefer_async_handlers=True) + else: + self._logging_worker.ensure_initialized_and_enqueue( + self.logging_obj.dispatch_success_handlers(self.messages, prefer_async_handlers=True) + ) self.logging_obj.model_call_details[REALTIME_SESSION_SUCCESS_LOGGED_KEY] = True async def _send_to_backend(self, message: str) -> bool: @@ -1568,7 +1586,12 @@ class RealTimeStreaming: finally: forward_task.cancel() client_task.cancel() - await asyncio.gather(forward_task, client_task, return_exceptions=True) + try: + await asyncio.gather(forward_task, client_task, return_exceptions=True) + finally: + cleanup: Final = realtime_attachment_cleanup.get() + if not self._account_usage and cleanup is not None: + await cleanup() async def _close_client(self, close: BackendClose) -> None: redacted_message: Final = redact_internal_details_from_client_message(close.message) diff --git a/litellm/llms/base_llm/realtime/http_transformation.py b/litellm/llms/base_llm/realtime/http_transformation.py index 43a80edb493..80daebd88b2 100644 --- a/litellm/llms/base_llm/realtime/http_transformation.py +++ b/litellm/llms/base_llm/realtime/http_transformation.py @@ -7,6 +7,7 @@ These are HTTP (not WebSocket) endpoints used by the WebRTC flow: """ from abc import ABC, abstractmethod +from collections.abc import Mapping from typing import Final import httpx @@ -36,6 +37,14 @@ class BaseRealtimeHTTPConfig(ABC): explicit api_base → litellm.api_base → env var → hard-coded default """ + def resolve_api_base(self, api_base: str | None, dynamic_api_base: str | None) -> str: + return self.get_api_base(dynamic_api_base or api_base) + + def get_realtime_calls_extra_headers( + self, headers: dict[str, object] | None + ) -> dict[str, object] | None: # mutable-ok: shared HTTP handler accepts a mutable header dictionary + return headers + @abstractmethod def get_api_key( self, @@ -97,6 +106,11 @@ class BaseRealtimeHTTPConfig(ABC): "Authorization": f"Bearer {ephemeral_key}", } + def transform_realtime_calls_response( + self, response: httpx.Response, model: str, model_id: str | None, headers: Mapping[str, object] | None + ) -> httpx.Response: + return response + # ------------------------------------------------------------------ # # Error handling # # ------------------------------------------------------------------ # diff --git a/litellm/llms/chatgpt/authenticator.py b/litellm/llms/chatgpt/authenticator.py index 563826c2b93..4a5045e02e7 100644 --- a/litellm/llms/chatgpt/authenticator.py +++ b/litellm/llms/chatgpt/authenticator.py @@ -49,8 +49,9 @@ class Authenticator: self.auth_file = os.path.join(self.token_dir, os.getenv("CHATGPT_AUTH_FILE", "auth.json")) self._ensure_token_dir() - def get_api_base(self) -> str: - return os.getenv("CHATGPT_API_BASE") or os.getenv("OPENAI_CHATGPT_API_BASE") or CHATGPT_API_BASE + @staticmethod + def get_api_base(default_base: str = CHATGPT_API_BASE) -> str: + return os.getenv("CHATGPT_API_BASE") or os.getenv("OPENAI_CHATGPT_API_BASE") or default_base def get_access_token(self) -> str: auth_data: Final = self._read_auth_file() diff --git a/litellm/llms/chatgpt/chat/transformation.py b/litellm/llms/chatgpt/chat/transformation.py index e35408b0829..1926283e85e 100644 --- a/litellm/llms/chatgpt/chat/transformation.py +++ b/litellm/llms/chatgpt/chat/transformation.py @@ -30,7 +30,7 @@ class ChatGPTConfig(OpenAIConfig): api_key: str | None, custom_llm_provider: str, ) -> tuple[str | None, str | None, str]: - dynamic_api_base: Final = self.authenticator.get_api_base() + dynamic_api_base: Final = api_base or self.authenticator.get_api_base() try: dynamic_api_key: Final = self.authenticator.get_access_token() except GetAccessTokenError as e: diff --git a/litellm/llms/chatgpt/codex.py b/litellm/llms/chatgpt/codex.py new file mode 100644 index 00000000000..ea736c97a0d --- /dev/null +++ b/litellm/llms/chatgpt/codex.py @@ -0,0 +1,94 @@ +from collections.abc import Mapping +from typing import Final +from urllib.parse import urlsplit + +import httpx +from pydantic import BaseModel, Field +from typing_extensions import ReadOnly, TypedDict + +from litellm.types.realtime import RealtimeQueryParams, RealtimeSessionConfig + + +class CodexRealtimeOffer(BaseModel): + sdp: str = Field(min_length=1) + session: RealtimeSessionConfig + + +class CodexRealtimeCall(BaseModel): + call_id: str = Field(pattern=r"^rtc_[A-Za-z0-9_-]+$") + model: str + model_id: str | None = None + alias: str + api_base: str | None = None + extra_headers: Mapping[str, str] | None = None + extra_query: Mapping[str, str | tuple[str, ...]] | None = None + usage_supervised: bool = False + parallel_reserved: bool = False + owner: str + expires_at: float + + +class ChatGPTCallRouting(BaseModel): + model: str + model_id: str | None = None + api_base: str | None = None + extra_headers: Mapping[str, str] | None = None + extra_query: Mapping[str, str | tuple[str, ...]] | None = None + + +class CodexSidebandRequest(TypedDict): + api_base: ReadOnly[str | None] + model: ReadOnly[str] + chatgpt_realtime_call_id: ReadOnly[str] + query_params: ReadOnly[RealtimeQueryParams] + extra_headers: ReadOnly[Mapping[str, str] | None] + extra_query: ReadOnly[Mapping[str, str | tuple[str, ...]] | None] + + +def build_call_request( + offer: CodexRealtimeOffer, query: Mapping[str, str], headers: Mapping[str, str] +) -> dict[str, object]: # mutable-ok: proxy processor enriches the request dictionary + return { # mutable-ok: proxy processor enriches the request dictionary + "model": offer.session.model, + "sdp_body": offer.sdp.encode(), + "session": offer.session.model_dump(exclude_none=True), + "openai_ephemeral_key": "", + "chatgpt_realtime_client_query": { # mutable-ok: router request parameters + key: value for key, value in query.items() if key in ("intent", "architecture") + }, + "chatgpt_realtime_client_headers": { # mutable-ok: router request headers + key: value + for key, value in headers.items() + if key in ("openai-alpha", "openai-beta", "x-session-id", "x-oai-attestation") + }, + } + + +def parse_call_response(response: httpx.Response, alias: str, owner: str, expires_at: float) -> CodexRealtimeCall: + routing_data: Final = response.extensions.get("chatgpt_realtime") + if not routing_data: + raise ValueError("Direct call signaling requires a ChatGPT deployment") + routing: Final = ChatGPTCallRouting.model_validate(routing_data) + call_id: Final = urlsplit(response.headers.get("location", "")).path.rstrip("/").rsplit("/", 1)[-1] + return CodexRealtimeCall( + call_id=call_id, + model=routing.model, + model_id=routing.model_id, + alias=alias, + owner=owner, + expires_at=expires_at, + api_base=routing.api_base, + extra_headers=routing.extra_headers, + extra_query=routing.extra_query, + ) + + +def build_sideband_request(call: CodexRealtimeCall) -> CodexSidebandRequest: + return CodexSidebandRequest( + api_base=call.api_base, + model=f"chatgpt/{call.model}", + chatgpt_realtime_call_id=call.call_id, + query_params=RealtimeQueryParams(model=call.model), + extra_headers=call.extra_headers, + extra_query=call.extra_query, + ) diff --git a/litellm/llms/chatgpt/common_utils.py b/litellm/llms/chatgpt/common_utils.py index 35e32e4172f..27fc28d2cd7 100644 --- a/litellm/llms/chatgpt/common_utils.py +++ b/litellm/llms/chatgpt/common_utils.py @@ -4,6 +4,8 @@ Constants and helpers for ChatGPT subscription OAuth. import os import platform +from collections.abc import Mapping +from types import MappingProxyType from typing import Any, Final from uuid import uuid4 @@ -105,6 +107,12 @@ You are producing plain text that will later be styled by the CLI. Follow these """ +def without_oauth_identity_headers(headers: Mapping[str, object]) -> Mapping[str, object]: + return MappingProxyType( + {key: value for key, value in headers.items() if key.lower() not in ("authorization", "chatgpt-account-id")} + ) + + class ChatGPTAuthError(BaseLLMException): def __init__( self, diff --git a/litellm/llms/chatgpt/images.py b/litellm/llms/chatgpt/images.py new file mode 100644 index 00000000000..8366e5016aa --- /dev/null +++ b/litellm/llms/chatgpt/images.py @@ -0,0 +1,154 @@ +import base64 +import os +from collections.abc import Mapping, Sequence +from pathlib import Path +from types import MappingProxyType +from typing import Final + +from httpx._types import FileTypes as HTTPFileTypes +from httpx._types import RequestFiles +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter + +from litellm.images.utils import ImageEditRequestUtils +from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig +from litellm.llms.openai.image_generation.gpt_transformation import GPTImageGenerationConfig +from litellm.types.llms.openai import AllMessageValues, FileTypes +from litellm.types.router import GenericLiteLLMParams + +from .authenticator import Authenticator +from .common_utils import without_oauth_identity_headers +from .responses.transformation import ChatGPTResponsesAPIConfig + + +class ReferenceImage(BaseModel): + model_config = ConfigDict(extra="forbid") + image_url: str = Field(pattern=r"^(data:image/(png|jpeg|webp);base64,|https://)") + + +def encode_reference( + file: HTTPFileTypes | FileTypes, +) -> dict[str, str]: # mutable-ok: image handler requires dictionaries + content: Final = file[1] if isinstance(file, tuple) else file + raw: Final = ( + Path(os.fsdecode(content)).read_bytes() + if isinstance(content, os.PathLike) + else content.encode() + if isinstance(content, str) + else content + if isinstance(content, bytes) + else content.read() + ) + content_type: Final = ( + file[2] + if isinstance(file, tuple) and len(file) >= 3 and file[2] + else ImageEditRequestUtils.get_image_content_type(raw) + ) + if content_type not in ("image/png", "image/jpeg", "image/webp"): + raise ValueError("Reference images must be PNG, JPEG, or WEBP") + return { # mutable-ok: JSON request serialization + "image_url": f"data:{content_type};base64," + base64.b64encode(raw).decode("ascii") + } + + +def image_headers( + headers: Mapping[str, object], model: str, params: Mapping[str, object] +) -> dict[str, object]: # mutable-ok: image handler requires dictionaries + auth_headers: Final = ChatGPTResponsesAPIConfig().validate_environment( + headers={}, # mutable-ok: Responses adapter header contract + model=model, + litellm_params=GenericLiteLLMParams.model_validate(params), + ) + return { # mutable-ok: image handler requires dictionaries + **without_oauth_identity_headers(headers), + **auth_headers, + "accept": "application/json", + } + + +class ChatGPTImageGenerationConfig(GPTImageGenerationConfig): + def validate_environment( + self, + headers: Mapping[str, object], + model: str, + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + api_key: str | None = None, + api_base: str | None = None, + ) -> dict[str, object]: # mutable-ok: image handler requires dictionaries + return image_headers(headers, model, litellm_params) + + def get_complete_url( + self, + api_base: str | None, + api_key: str | None, + model: str, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + stream: bool | None = None, + ) -> str: + return f"{(api_base or Authenticator.get_api_base()).rstrip('/')}/images/generations" + + def transform_image_generation_request( + self, + model: str, + prompt: str, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + headers: Mapping[str, object], + ) -> dict[str, object]: # mutable-ok: image handler requires dictionaries + return {"model": model, "prompt": prompt, **optional_params} # mutable-ok: JSON request serialization + + +class ChatGPTImageEditConfig(OpenAIImageEditConfig): + def validate_environment( + self, + headers: Mapping[str, object], + model: str, + api_key: str | None = None, + litellm_params: Mapping[str, object] | None = None, + api_base: str | None = None, + ) -> dict[str, object]: # mutable-ok: image handler requires dictionaries + return image_headers(headers, model, litellm_params or MappingProxyType({})) + + def get_complete_url(self, model: str, api_base: str | None, litellm_params: Mapping[str, object]) -> str: + return f"{(api_base or Authenticator.get_api_base()).rstrip('/')}/images/edits" + + def use_multipart_form_data(self) -> bool: + return False + + def transform_image_edit_request( + self, + model: str, + prompt: str | None, + image: FileTypes | Sequence[FileTypes] | None, + image_edit_optional_request_params: Mapping[str, object], + litellm_params: GenericLiteLLMParams, + headers: Mapping[str, object], + ) -> tuple[dict[str, object], RequestFiles]: # mutable-ok: image handler requires dictionaries + if image_edit_optional_request_params.get("mask") is not None: + raise ValueError("ChatGPT image editing does not support masks") + references: Final = getattr(litellm_params, "images", None) + if references is not None: + if image: + raise ValueError("Specify only one of image or images") + validated: Final = TypeAdapter(tuple[ReferenceImage, ...]).validate_python(references) + if not 1 <= len(validated) <= 5: + raise ValueError("images must contain between 1 and 5 reference images") + return { # mutable-ok: JSON request serialization + "model": model, + "prompt": prompt, + **image_edit_optional_request_params, + "images": tuple(item.model_dump() for item in validated), + }, () + + inputs: Final = tuple(image) if isinstance(image, list) else (image,) if image is not None else () + encoded: Final = tuple(encode_reference(file) for file in inputs) + if not 1 <= len(encoded) <= 5: + raise ValueError("images must contain between 1 and 5 reference images") + return { # mutable-ok: JSON request serialization + "model": model, + "prompt": prompt, + **image_edit_optional_request_params, + "images": encoded, + }, () diff --git a/litellm/llms/chatgpt/realtime.py b/litellm/llms/chatgpt/realtime.py new file mode 100644 index 00000000000..6c275a1dec7 --- /dev/null +++ b/litellm/llms/chatgpt/realtime.py @@ -0,0 +1,285 @@ +from collections.abc import Mapping +from enum import Enum, auto +from types import MappingProxyType +from typing import TYPE_CHECKING, Final + +from httpx import URL, QueryParams, Response +from pydantic import TypeAdapter + +from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES +from litellm.llms.openai.realtime.handler import OpenAIRealtime +from litellm.llms.openai.realtime.http_transformation import OpenAIRealtimeHTTPConfig +from litellm.types.realtime import RealtimeQueryParams +from litellm.types.router import GenericLiteLLMParams +from litellm.utils import get_model_info + +from .authenticator import Authenticator +from .common_utils import without_oauth_identity_headers +from .responses.transformation import ChatGPTResponsesAPIConfig + +if TYPE_CHECKING: + from websockets.asyncio.client import ClientConnection + + +class CallAccounting(Enum): + SUPERVISED = auto() + + +def accounts_for_call_usage(params: GenericLiteLLMParams) -> bool: + return getattr(params, "chatgpt_call_accounting", None) is not CallAccounting.SUPERVISED + + +def configured_realtime_headers(headers: Mapping[str, object] | None) -> Mapping[str, str]: + validated: Final = TypeAdapter(Mapping[str, str]).validate_python( + without_oauth_identity_headers(headers or MappingProxyType({})) + ) + return MappingProxyType({key.lower(): value for key, value in validated.items()}) + + +def configured_realtime_query(params: GenericLiteLLMParams) -> Mapping[str, str | tuple[str, ...]]: + inbound: Final = TypeAdapter(Mapping[str, str]).validate_python( + getattr(params, "chatgpt_realtime_client_query", None) or MappingProxyType({}) + ) + configured: Final = TypeAdapter( + Mapping[str, str | int | float | bool | None | tuple[str | int | float | bool | None, ...]] + ).validate_python(getattr(params, "extra_query", None) or MappingProxyType({})) + merged: Final = QueryParams( + tuple((key, value) for key, value in inbound.items() if key in ("intent", "architecture")) + ).merge(configured) + return MappingProxyType( + {key: merged[key] if len(merged.get_list(key)) == 1 else tuple(merged.get_list(key)) for key in merged} + ) + + +def realtime_call_headers(params: GenericLiteLLMParams) -> dict[str, str]: # mutable-ok: HTTP handler header contract + inbound: Final = TypeAdapter(Mapping[str, str]).validate_python( + getattr(params, "chatgpt_realtime_client_headers", None) or MappingProxyType({}) + ) + configured: Final = TypeAdapter(Mapping[str, object]).validate_python( + getattr(params, "extra_headers", None) or MappingProxyType({}) + ) + return { # mutable-ok: HTTP handler header contract + **MappingProxyType( + { + key.lower(): value + for key, value in inbound.items() + if key.lower() in ("openai-alpha", "openai-beta", "x-session-id", "x-oai-attestation") + } + ), + **configured_realtime_headers(configured), + } + + +def realtime_headers( + params: GenericLiteLLMParams, headers: Mapping[str, str], extra_headers: Mapping[str, object] | None = None +) -> dict[str, str]: # mutable-ok: HTTP handler header contract + forwarded: Final = MappingProxyType( + { + key.lower(): value + for key, value in headers.items() + if key.lower() in ("openai-alpha", "openai-beta", "x-session-id", "x-oai-attestation") + } + ) + return { # mutable-ok: HTTP handler updates headers + **ChatGPTResponsesAPIConfig().validate_environment( + headers={}, # mutable-ok: Responses adapter header contract + model="", + litellm_params=params, + ), + **forwarded, + **configured_realtime_headers(extra_headers), + } + + +def realtime_endpoint(model: str) -> str: + try: + model_info: Final = get_model_info(model, custom_llm_provider="chatgpt") + except Exception: # noqa: BLE001 # get_model_info raises bare Exception for unmapped models + return "realtime" + return "live" if "/v1/live" in (model_info.get("supported_endpoints") or ()) else "realtime" + + +class ChatGPTRealtime(OpenAIRealtime): + async def open_call_connection(self, model: str, api_base: str) -> "ClientConnection": + import websockets + + url: Final = self._construct_url(api_base, RealtimeQueryParams(model=model)) + return await websockets.connect( + url, + additional_headers=self._profile_headers, + max_size=REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES, + ssl=self._get_ssl_config(url), + open_timeout=20, + ) + + async def close_call(self, connection: "ClientConnection", model: str, api_base: str) -> None: + from websockets.exceptions import ConnectionClosed + + if realtime_endpoint(model) == "live": + try: + await connection.send('{"type":"session.close"}') + return + except (ConnectionClosed, OSError): + await self.hangup_call(api_base) + return + await self.hangup_call(api_base) + + async def hangup_call(self, api_base: str) -> None: + from litellm.llms.custom_httpx.http_handler import get_async_httpx_client + from litellm.types.utils import LlmProviders + + base: Final = URL(api_base) + url: Final = base.copy_with( + scheme="https" if base.scheme in ("https", "wss") else "http", + path=f"{base.path.rstrip('/')}/realtime/calls/{self._call_id}/hangup", + params=tuple( + (key, value) + for key, value in QueryParams(self._extra_query).multi_items() + if key not in ("model", "call_id") + ), + ) + client: Final = get_async_httpx_client(llm_provider=LlmProviders.CHATGPT) + response: Final = await client.post(str(url), headers=self._profile_headers, data=b"", timeout=10) + response.raise_for_status() + + @staticmethod + def get_api_base(api_base: str | None = None) -> str: + return api_base or Authenticator.get_api_base(default_base="https://api.openai.com/v1") + + def __init__( + self, + params: GenericLiteLLMParams, + headers: Mapping[str, str], + extra_headers: Mapping[str, object] | None = None, + ) -> None: + super().__init__() + self._profile_headers = realtime_headers(params, headers, extra_headers) + self._call_id = TypeAdapter(str | None).validate_python(getattr(params, "chatgpt_realtime_call_id", None)) + self._extra_query = configured_realtime_query(params) + self._account_usage = accounts_for_call_usage(params) + + def _get_default_api_base(self) -> str: + return self.get_api_base() + + def _resolve_api_key(self, api_key: str | None) -> str: + return "chatgpt-oauth" + + def _accounts_for_call_usage(self) -> bool: + return self._account_usage + + def _get_additional_headers( + self, api_key: str, *, openai_beta_realtime: bool = False + ) -> dict[str, str]: # mutable-ok: HTTP handler header contract + return { # mutable-ok: HTTP handler updates headers + **(MappingProxyType({"OpenAI-Beta": "realtime=v1"}) if openai_beta_realtime else MappingProxyType({})), + **self._profile_headers, + } + + def _construct_url(self, api_base: str, query_params: RealtimeQueryParams) -> str: + base: Final = URL(api_base) + endpoint: Final = realtime_endpoint(query_params.get("model", "")) + if self._call_id: + gateway_query: Final = tuple( + (key, value) + for key, value in QueryParams(self._extra_query).multi_items() + if key not in ("model", "call_id") + ) + return str( + base.copy_with( + scheme="wss" if base.scheme in ("https", "wss") else "ws", + path=f"{base.path.rstrip('/')}/{endpoint}/{self._call_id}" + if endpoint == "live" + else f"{base.path.rstrip('/')}/realtime", + params=gateway_query + (() if endpoint == "live" else (("call_id", self._call_id),)), + ) + ) + return str( + base.copy_with( + scheme="wss" if base.scheme in ("https", "wss") else "ws", + path=f"{base.path.rstrip('/')}/{endpoint}", + params=QueryParams(TypeAdapter(Mapping[str, str | None]).validate_python(query_params)).merge( + tuple( + (key, value) + for key, value in QueryParams(self._extra_query).multi_items() + if key not in ("model", "call_id") + ) + ), + ) + ) + + +class ChatGPTRealtimeHTTPConfig(OpenAIRealtimeHTTPConfig): + realtime_calls_json: Final = True + + def __init__(self, params: GenericLiteLLMParams, use_codex_backend: bool = True) -> None: + self._params = params + self._use_codex_backend = use_codex_backend + + def get_api_base( + self, + api_base: str | None, + **kwargs: object, # kwargs-ok: provider interface accepts optional credentials + ) -> str: + return api_base or (Authenticator.get_api_base() if self._use_codex_backend else ChatGPTRealtime.get_api_base()) + + def resolve_api_base(self, api_base: str | None, dynamic_api_base: str | None) -> str: + return self.get_api_base(api_base) + + def get_realtime_calls_extra_headers( + self, headers: dict[str, object] | None + ) -> dict[str, object]: # mutable-ok: shared HTTP handler accepts a mutable header dictionary + return {**realtime_call_headers(self._params)} # mutable-ok: shared HTTP header contract + + def get_api_key( + self, + api_key: str | None, + **kwargs: object, # kwargs-ok: provider interface accepts optional credentials + ) -> str: + return "chatgpt-oauth" + + def get_realtime_calls_url(self, api_base: str | None, model: str, api_version: str | None = None) -> str: + query: Final = configured_realtime_query(self._params) + return str(URL(f"{self.get_api_base(api_base).rstrip('/')}/realtime/calls", params=query)) + + def transform_realtime_calls_response( + self, response: Response, model: str, model_id: str | None, headers: Mapping[str, object] | None + ) -> Response: + response.extensions["chatgpt_realtime"] = ( + MappingProxyType( # rebind-ok: HTTPX response extensions carry provider routing metadata + { + "model": model, + "model_id": model_id, + "api_base": ChatGPTRealtime.get_api_base(self._params.api_base), + "extra_headers": configured_realtime_headers(headers), + "extra_query": configured_realtime_query(self._params), + } + ) + ) + return response + + def get_realtime_calls_headers( + self, ephemeral_key: str + ) -> dict[str, str]: # mutable-ok: HTTP handler header contract + return realtime_headers(self._params, MappingProxyType({})) + + def validate_environment( + self, + headers: Mapping[str, str], + model: str, + api_key: str | None = None, + ) -> dict[str, str]: # mutable-ok: HTTP handler header contract + return { # mutable-ok: HTTP handler updates headers + **realtime_headers(self._params, headers), + "Content-Type": "application/json", + } + + def get_complete_url(self, api_base: str | None, model: str, api_version: str | None = None) -> str: + return f"{self.get_api_base(api_base).rstrip('/')}/realtime/client_secrets" + + def get_transcription_session_url( + self, + api_base: str | None, + model: str, + api_version: str | None = None, + ) -> str: + return f"{self.get_api_base(api_base).rstrip('/')}/realtime/transcription_sessions" diff --git a/litellm/llms/chatgpt/responses/transformation.py b/litellm/llms/chatgpt/responses/transformation.py index b96e06be3d8..1a5c6302a01 100644 --- a/litellm/llms/chatgpt/responses/transformation.py +++ b/litellm/llms/chatgpt/responses/transformation.py @@ -102,6 +102,7 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig): "reasoning", "previous_response_id", "truncation", + "text", } return {k: v for k, v in request.items() if k in allowed_keys} diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 7109e6942d1..b118e4c17e7 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -6336,6 +6336,9 @@ class BaseLLMHTTPHandler: Uses provider_config (BaseRealtimeHTTPConfig) for URL construction and header auth when available; falls back to the legacy OpenAI-style defaults. """ + from litellm.llms.chatgpt.common_utils import without_oauth_identity_headers + from litellm.llms.chatgpt.realtime import ChatGPTRealtimeHTTPConfig + if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( llm_provider=litellm.LlmProviders.OPENAI, @@ -6361,7 +6364,11 @@ class BaseLLMHTTPHandler: } if extra_headers: - headers.update(extra_headers) + headers.update( + without_oauth_identity_headers(extra_headers) + if isinstance(provider_config, ChatGPTRealtimeHTTPConfig) + else extra_headers + ) logging_obj.pre_call( input=request_data, @@ -6412,6 +6419,9 @@ class BaseLLMHTTPHandler: - sdp: the SDP offer (text) - session: JSON string with {"type": "realtime", "model": "...", ...} """ + from litellm.llms.chatgpt.common_utils import without_oauth_identity_headers + from litellm.llms.chatgpt.realtime import ChatGPTRealtimeHTTPConfig + if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( llm_provider=litellm.LlmProviders.OPENAI, @@ -6429,11 +6439,15 @@ class BaseLLMHTTPHandler: } if extra_headers: - headers.update(extra_headers) + headers.update( + without_oauth_identity_headers(extra_headers) + if isinstance(provider_config, ChatGPTRealtimeHTTPConfig) + else extra_headers + ) # Build multipart form data: sdp + session JSON session_data: Final = session_config or {} - if "type" not in session_data: + if "type" not in session_data and not getattr(provider_config, "realtime_calls_json", False): session_data["type"] = "realtime" if "model" not in session_data and model: session_data["model"] = model @@ -6456,6 +6470,13 @@ class BaseLLMHTTPHandler: ) try: + if getattr(provider_config, "realtime_calls_json", False): + return await async_httpx_client.post( + url=url, + headers=headers, + json={"sdp": sdp_text, "session": session_data}, # mutable-ok: JSON signaling payload + timeout=timeout, + ) return await async_httpx_client.post( url=url, headers=headers, @@ -6653,6 +6674,14 @@ class BaseLLMHTTPHandler: else: raise Exception(f"Unexpected error while closing WebSocket: {close_error}") + @staticmethod + def _image_extra_headers(custom_llm_provider: str, headers: Mapping[str, object]) -> Mapping[str, object]: + if custom_llm_provider == "chatgpt": + from litellm.llms.chatgpt.common_utils import without_oauth_identity_headers + + return without_oauth_identity_headers(headers) + return headers + def image_edit_handler( self, model: str, @@ -6709,7 +6738,7 @@ class BaseLLMHTTPHandler: ) if extra_headers: - headers.update(extra_headers) + headers.update(self._image_extra_headers(custom_llm_provider, extra_headers)) api_base: Final = image_edit_provider_config.get_complete_url( model=model, @@ -6808,7 +6837,7 @@ class BaseLLMHTTPHandler: ) if extra_headers: - headers.update(extra_headers) + headers.update(self._image_extra_headers(custom_llm_provider, extra_headers)) api_base: Final = image_edit_provider_config.get_complete_url( model=model, @@ -6925,7 +6954,7 @@ class BaseLLMHTTPHandler: ) if extra_headers: - headers.update(extra_headers) + headers.update(self._image_extra_headers(custom_llm_provider, extra_headers)) api_base: Final = image_generation_provider_config.get_complete_url( model=model, @@ -7032,7 +7061,7 @@ class BaseLLMHTTPHandler: ) if extra_headers: - headers.update(extra_headers) + headers.update(self._image_extra_headers(custom_llm_provider, extra_headers)) api_base: Final = image_generation_provider_config.get_complete_url( model=model, diff --git a/litellm/llms/openai/realtime/handler.py b/litellm/llms/openai/realtime/handler.py index e3ecbac1a53..a896594db41 100644 --- a/litellm/llms/openai/realtime/handler.py +++ b/litellm/llms/openai/realtime/handler.py @@ -38,6 +38,14 @@ class OpenAIRealtime(OpenAIChatCompletion): """ return "https://api.openai.com/" + def _resolve_api_key(self, api_key: str | None) -> str: + if api_key is None: + raise ValueError("api_key is required for OpenAI realtime calls") + return api_key + + def _accounts_for_call_usage(self) -> bool: + return True + def _get_additional_headers( self, api_key: str, @@ -117,6 +125,7 @@ class OpenAIRealtime(OpenAIChatCompletion): query_params: RealtimeQueryParams | None = None, user_api_key_dict: object | None = None, litellm_metadata: dict | None = None, + account_usage: bool = True, **kwargs: object, ): import websockets @@ -124,8 +133,7 @@ class OpenAIRealtime(OpenAIChatCompletion): if api_base is None: api_base = self._get_default_api_base() - if api_key is None: - raise ValueError("api_key is required for OpenAI realtime calls") + resolved_api_key: Final = self._resolve_api_key(api_key) # Use all query params if provided, else fallback to just model if query_params is None: @@ -143,12 +151,12 @@ class OpenAIRealtime(OpenAIChatCompletion): "If your client expects beta event names, add 'OpenAI-Beta: realtime=v1' " "to the WebSocket headers sent to the LiteLLM proxy." ) - headers: Final = self._get_additional_headers(api_key, openai_beta_realtime=openai_beta_realtime) + headers: Final = self._get_additional_headers(resolved_api_key, openai_beta_realtime=openai_beta_realtime) # Log a masked request preview consistent with other endpoints. logging_obj.pre_call( input=None, - api_key=api_key, + api_key=resolved_api_key, additional_args={ "api_base": url, "headers": headers, @@ -172,6 +180,7 @@ class OpenAIRealtime(OpenAIChatCompletion): model if (query_params or {}).get("intent") == "transcription" else None ), event_normalizer=self._make_event_normalizer(), + account_usage=account_usage and self._accounts_for_call_usage(), ) await realtime_streaming.bidirectional_forward() diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 7fa09951eae..4b25dae8aeb 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -27932,6 +27932,16 @@ "max_tokens": 8191, "mode": "embedding" }, + "chatgpt/gpt-live-1-codex": { + "litellm_provider": "chatgpt", + "mode": "realtime", + "supported_endpoints": [ + "/v1/realtime/calls", + "/v1/live" + ], + "supports_audio_input": true, + "supports_audio_output": true + }, "chatgpt/gpt-5.5": { "litellm_provider": "chatgpt", "source": "https://platform.openai.com/docs/models/gpt-5.5", diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 85677f4eb3c..890b1bf2b1f 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -102,6 +102,10 @@ class ReconcileOutcome(NamedTuple): live_after: frozenset[str] | None +class InternalRequestOrigin(enum.Enum): + REALTIME_OBSERVER = enum.auto() + + class SupportedDBObjectType(str, enum.Enum): """ Supported database object types for fine-grained DB storage control. @@ -405,6 +409,9 @@ class LiteLLMRoutes(enum.Enum): "/realtime?{model}", "/v1/realtime?{model}", "/openai/v1/realtime?{model}", + "/live", + "/v1/live", + "/v1/live/{call_id}", # realtime (GA WebRTC HTTP routes) "/realtime/client_secrets", "/v1/realtime/client_secrets", diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index be65c3b39ec..2cb739878bd 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -1866,15 +1866,27 @@ def _extract_model_candidates_from_request( uses_completion_model_sources: Final = _route_matches_any_marker( route=route, markers=_MODEL_ROUTING_COMPLETION_MODEL_ROUTE_MARKERS ) + session: Final[object] = ( + request_data.get("session") + if _route_matches_any_marker(route=route, markers=_MODEL_ROUTING_SESSION_MODEL_ROUTE_MARKERS) + else None + ) + parsed_session: Final[object] = safe_json_loads(session) if isinstance(session, str) else session + session_model: Final[object] = parsed_session.get("model") if isinstance(parsed_session, dict) else None + if ( + _route_matches_any_marker(route=route, markers=("/realtime/calls",)) + and isinstance(session_model, str) + and session_model + ): + candidates.append(session_model) + return candidates body_model: Final = request_data.get("model") _append_model_candidates(candidates, body_model) if uses_body_target_model_sources or not body_model: _append_model_candidates(candidates, request_data.get("target_model_names")) if _route_matches_any_marker(route=route, markers=_MODEL_ROUTING_SESSION_MODEL_ROUTE_MARKERS): - session: Final = request_data.get("session") - if isinstance(session, dict): - _append_model_candidates(candidates, session.get("model")) + _append_model_candidates(candidates, session_model) if uses_completion_model_sources and isinstance(request_data.get("completion"), dict): _append_model_candidates(candidates, request_data["completion"].get("model")) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 9828311112e..748344dd33e 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -350,7 +350,7 @@ async def _check_key_model_budget_with_fallback( model=model_name, ) except litellm.BudgetExceededError as e: - if request_data.get("model") != model_name: + if request_data.get("model") != model_name or request.scope.get("litellm_pinned_realtime_model") == model_name: raise e fallback_model: Final = await model_max_budget_limiter.get_fallback_model_within_budget( user_api_key_dict=valid_token, @@ -535,6 +535,38 @@ def _apply_budget_limits_to_end_user_params( verbose_proxy_logger.debug("Applied budget limits to end user %s", end_user_id) +def get_websocket_api_key(websocket: WebSocket) -> str | None: + from litellm.proxy.proxy_server import general_settings + + custom_header: Final = general_settings.get("litellm_key_header_name") + if isinstance(custom_header, str): + if not websocket.headers.get(custom_header): + return None + request: Final = Request( + {"type": "http", "headers": websocket.scope.get("headers", [])} # mutable-ok: ASGI request scope + ) + return get_api_key_from_custom_header(request, custom_header) + custom_key: Final = websocket.headers.get("x-litellm-api-key") + if custom_key is not None: + return _get_bearer_token_or_received_api_key(custom_key) + authorization: Final = websocket.headers.get("authorization") + if authorization: + if not authorization.startswith("Bearer "): + raise HTTPException(status_code=403, detail="Invalid Authorization header format") + return authorization[len("Bearer ") :].strip() + api_key: Final = websocket.headers.get("api-key") + if api_key: + return api_key + return next( + ( + protocol.strip().removeprefix("openai-insecure-api-key.") + for protocol in websocket.headers.get("sec-websocket-protocol", "").split(",") + if protocol.strip().startswith("openai-insecure-api-key.") + ), + None, + ) + + async def user_api_key_auth_websocket(websocket: WebSocket): # Accept the WebSocket connection @@ -557,40 +589,33 @@ async def user_api_key_auth_websocket(websocket: WebSocket): request._url = websocket.url - query_params: Final = websocket.query_params - - model: Final = query_params.get("model") - - async def return_body(): - return _realtime_request_body(model) - - request.body = return_body - - authorization: Final = websocket.headers.get("authorization") - # If no Authorization header, try the api-key header - if not authorization: - api_key = websocket.headers.get("api-key") - if not api_key: - # Try extracting from WebSocket subprotocol (browser clients) - for protocol in websocket.headers.get("sec-websocket-protocol", "").split(","): - protocol = protocol.strip() - if protocol.startswith("openai-insecure-api-key."): - api_key = protocol[len("openai-insecure-api-key.") :] - break - if not api_key: - await websocket.close(code=status.WS_1008_POLICY_VIOLATION) - raise HTTPException(status_code=403, detail="No API key provided") - else: - # Extract the API key from the Bearer token - if not authorization.startswith("Bearer "): - await websocket.close(code=status.WS_1008_POLICY_VIOLATION) - raise HTTPException(status_code=403, detail="Invalid Authorization header format") - - api_key = authorization[len("Bearer ") :].strip() + try: + api_key: Final = get_websocket_api_key(websocket) + except HTTPException: + await websocket.close(code=status.WS_1008_POLICY_VIOLATION) + raise + if not api_key: + await websocket.close(code=status.WS_1008_POLICY_VIOLATION) + raise HTTPException(status_code=403, detail="No API key provided") # Call user_api_key_auth with the extracted API key # Note: You'll need to modify this to work with WebSocket context if needed try: + from litellm.proxy.realtime_endpoints.call_sessions import decode_call + + call_token: Final = websocket.path_params.get("call_id") or websocket.query_params.get("call_id") + model: Final = ( + decode_call(call_token, f"Bearer {api_key}").alias + if call_token is not None + else websocket.query_params.get("model") + ) + if call_token is not None: + request.scope["litellm_pinned_realtime_model"] = model + + async def return_body(): + return _realtime_request_body(model) + + request.body = return_body return await user_api_key_auth(request=request, api_key=f"Bearer {api_key}") except Exception as e: if is_invalid_virtual_key_error(e): diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 9f58aaf24f1..c06fab872c7 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1832,6 +1832,8 @@ class ProxyBaseLLMRequestProcessing: user_api_base: str | None = None, model: str | None = None, llm_router: Router | None = None, + *, + internal_realtime_observer: bool = False, ) -> tuple[dict, LiteLLMLoggingObj]: start_time: Final = datetime.now() # start before calling guardrail hooks @@ -1996,6 +1998,11 @@ class ProxyBaseLLMRequestProcessing: user_api_key_dict=user_api_key_dict, data=self.data, call_type=route_type, + **( + MappingProxyType({"internal_realtime_observer": True}) + if internal_realtime_observer + else MappingProxyType({}) + ), ) if route_type == "aget_responses": attach_post_call_pipelines_to_retrieval( diff --git a/litellm/proxy/hooks/parallel_request_limiter.py b/litellm/proxy/hooks/parallel_request_limiter.py index b313cb64c3f..d8dc6bcb5e1 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -1,9 +1,10 @@ import asyncio import sys +from collections.abc import Mapping from datetime import datetime, timedelta from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter from typing_extensions import TypedDict import litellm @@ -12,7 +13,7 @@ from litellm._logging import verbose_proxy_logger from litellm.exceptions import RateLimitType from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs -from litellm.proxy._types import CommonProxyErrors, CurrentItemRateLimit, UserAPIKeyAuth +from litellm.proxy._types import CommonProxyErrors, CurrentItemRateLimit, InternalRequestOrigin, UserAPIKeyAuth from litellm.proxy.auth.auth_utils import ( get_key_model_rpm_limit, get_key_model_tpm_limit, @@ -50,11 +51,81 @@ class CacheObject(TypedDict): request_count_end_user_id: dict | None +class _RealtimeAttachmentReservations(BaseModel): + cache_keys: tuple[str, ...] = () + global_acquired: bool = False + + def acquire(self, key: str) -> None: + self.cache_keys = tuple(dict.fromkeys((*self.cache_keys, key))) + + def acquire_global(self) -> None: + self.global_acquired = True + + def take(self) -> tuple[tuple[str, ...], bool]: + owned: Final = (self.cache_keys, self.global_acquired) + self.cache_keys = () + self.global_acquired = False + return owned + + +_RELEASE_REALTIME_COUNTER_LUA: Final = """ +local raw = redis.call('GET', KEYS[1]) +if not raw then return 0 end +local value = cjson.decode(raw) +value.current_requests = math.max(value.current_requests - 1, 0) +redis.call('SET', KEYS[1], cjson.encode(value), 'KEEPTTL') +return 1 +""" + + class _PROXY_MaxParallelRequestsHandler(CustomLogger): # Class variables or attributes def __init__(self, internal_usage_cache: InternalUsageCache): self.internal_usage_cache = internal_usage_cache + def begin_realtime_attachment(self, request_data: dict[str, object]) -> None: + request_data["_legacy_realtime_attachment_reservations"] = ( # rebind-ok: request-scoped cleanup receipt + _RealtimeAttachmentReservations() + ) + + async def async_release_realtime_attachment( + self, request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth + ) -> None: + receipt: Final = request_data.get("_legacy_realtime_attachment_reservations") + if not isinstance(receipt, _RealtimeAttachmentReservations): + return + keys, global_acquired = receipt.take() + if global_acquired: + await self.internal_usage_cache.async_increment_cache( + key="global_max_parallel_requests", + value=-1, + local_only=True, + litellm_parent_otel_span=user_api_key_dict.parent_otel_span, + ) + for key in keys: + await self._release_realtime_counter(key) + + async def _release_realtime_counter(self, key: str) -> None: + local: Final = self.internal_usage_cache.dual_cache.in_memory_cache + remote: Final = self.internal_usage_cache.dual_cache.redis_cache + raw: Final[object] = local.get_cache(key) + current: Final = TypeAdapter(Mapping[str, int] | None).validate_python(raw) + updated: Final = ( + { # mutable-ok: shared cache counter dict + **current, + "current_requests": max(current["current_requests"] - 1, 0), + } + if current is not None + else None + ) + if updated is not None: + local.set_cache(key, updated, ttl=60) + if remote is not None: + release: Final = remote.async_register_script(_RELEASE_REALTIME_COUNTER_LUA) + await release(keys=(key,), args=()) + if local.get_cache(key) is updated: + local.delete_cache(key) + def print_verbose(self, print_statement): try: verbose_proxy_logger.debug(print_statement) @@ -142,6 +213,9 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): litellm_parent_otel_span=user_api_key_dict.parent_otel_span, local_only=True, ) + receipt: Final = data.get("_legacy_realtime_attachment_reservations") + if isinstance(receipt, _RealtimeAttachmentReservations): + receipt.acquire(request_count_api_key) return new_val def time_to_next_minute(self) -> float: @@ -299,6 +373,9 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): local_only=True, litellm_parent_otel_span=user_api_key_dict.parent_otel_span, ) + receipt: Final = data.get("_legacy_realtime_attachment_reservations") + if isinstance(receipt, _RealtimeAttachmentReservations): + receipt.acquire_global() _model = data.get("model", None) current_date: Final = datetime.now().strftime("%Y-%m-%d") @@ -480,6 +557,13 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): values_to_update_in_cache=values_to_update_in_cache, ) + if isinstance(data.get("_legacy_realtime_attachment_reservations"), _RealtimeAttachmentReservations): + await self.internal_usage_cache.async_batch_set_cache( + cache_list=values_to_update_in_cache, + ttl=60, + litellm_parent_otel_span=user_api_key_dict.parent_otel_span, + ) + return asyncio.create_task( self.internal_usage_cache.async_batch_set_cache( cache_list=values_to_update_in_cache, @@ -489,6 +573,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): ) async def async_log_success_event(self, kwargs, response_obj: object, start_time, end_time): + releases_slot: Final = kwargs.get("internal_request_origin") is not InternalRequestOrigin.REALTIME_OBSERVER from litellm.proxy.common_utils.callback_utils import ( get_model_group_from_litellm_kwargs, ) @@ -521,7 +606,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): # Setup values # ------------ - if global_max_parallel_requests is not None: + if releases_slot and global_max_parallel_requests is not None: # get value from cache _key: Final = "global_max_parallel_requests" # decrement @@ -552,13 +637,13 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): key=request_count_api_key, litellm_parent_otel_span=litellm_parent_otel_span, ) or { - "current_requests": 1, + "current_requests": int(releases_slot), "current_tpm": 0, "current_rpm": 0, } new_val = { - "current_requests": max(current["current_requests"] - 1, 0), + "current_requests": max(current["current_requests"] - int(releases_slot), 0), "current_tpm": current["current_tpm"] + total_tokens, "current_rpm": current["current_rpm"], } @@ -593,13 +678,13 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): key=request_count_api_key, litellm_parent_otel_span=litellm_parent_otel_span, ) or { - "current_requests": 1, + "current_requests": int(releases_slot), "current_tpm": 0, "current_rpm": 0, } new_val = { - "current_requests": max(current["current_requests"] - 1, 0), + "current_requests": max(current["current_requests"] - int(releases_slot), 0), "current_tpm": current["current_tpm"] + total_tokens, "current_rpm": current["current_rpm"], } @@ -619,13 +704,13 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): key=request_count_api_key, litellm_parent_otel_span=litellm_parent_otel_span, ) or { - "current_requests": 1, - "current_tpm": total_tokens, - "current_rpm": 1, + "current_requests": int(releases_slot), + "current_tpm": total_tokens if releases_slot else 0, + "current_rpm": int(releases_slot), } new_val = { - "current_requests": max(current["current_requests"] - 1, 0), + "current_requests": max(current["current_requests"] - int(releases_slot), 0), "current_tpm": current["current_tpm"] + total_tokens, "current_rpm": current["current_rpm"], } @@ -645,13 +730,13 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): key=request_count_api_key, litellm_parent_otel_span=litellm_parent_otel_span, ) or { - "current_requests": 1, - "current_tpm": total_tokens, - "current_rpm": 1, + "current_requests": int(releases_slot), + "current_tpm": total_tokens if releases_slot else 0, + "current_rpm": int(releases_slot), } new_val = { - "current_requests": max(current["current_requests"] - 1, 0), + "current_requests": max(current["current_requests"] - int(releases_slot), 0), "current_tpm": current["current_tpm"] + total_tokens, "current_rpm": current["current_rpm"], } @@ -671,13 +756,13 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): key=request_count_api_key, litellm_parent_otel_span=litellm_parent_otel_span, ) or { - "current_requests": 1, - "current_tpm": total_tokens, - "current_rpm": 1, + "current_requests": int(releases_slot), + "current_tpm": total_tokens if releases_slot else 0, + "current_rpm": int(releases_slot), } new_val = { - "current_requests": max(current["current_requests"] - 1, 0), + "current_requests": max(current["current_requests"] - int(releases_slot), 0), "current_tpm": current["current_tpm"] + total_tokens, "current_rpm": current["current_rpm"], } @@ -694,6 +779,8 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): self.print_verbose(e) async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + if kwargs.get("internal_request_origin") is InternalRequestOrigin.REALTIME_OBSERVER: + return try: self.print_verbose("Inside Max Parallel Request Failure Hook") litellm_parent_otel_span: Final[Span | None] = _get_parent_otel_span_from_kwargs(kwargs=kwargs) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index c398abff099..f17e920802b 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -9,10 +9,12 @@ import binascii import logging import os import uuid -from collections.abc import Awaitable, Callable, Mapping, Sequence, Set +from collections.abc import Awaitable, Callable, Generator, Mapping, Sequence, Set +from contextlib import contextmanager from contextvars import ContextVar from dataclasses import dataclass, field from datetime import datetime +from types import MappingProxyType from typing import ( TYPE_CHECKING, Any, @@ -55,6 +57,7 @@ from litellm.proxy.hooks.batch_enqueued_tokens import ( canonical_provider_batch_id, ) from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit +from litellm.proxy.hooks.realtime_call_lease import RealtimeCallLease, is_realtime_call_attachment from litellm.router_utils.add_retry_fallback_headers import ( ensure_response_additional_headers, response_has_hidden_params, @@ -306,6 +309,23 @@ end return results """ +PARALLEL_RENEW_SCRIPT: Final = """ +local clock = redis.call('TIME') +local now = tonumber(clock[1]) +local ttl = tonumber(ARGV[2]) +for i = 1, #KEYS do + local score = redis.call('ZSCORE', KEYS[i], ARGV[1]) + if not score or tonumber(score) <= now - ttl then + return {0} + end +end +for i = 1, #KEYS do + redis.call('ZADD', KEYS[i], 'XX', now, ARGV[1]) + redis.call('EXPIRE', KEYS[i], ttl) +end +return {1} +""" + TOKEN_INCREMENT_SCRIPT: Final = """ local results = {} @@ -418,9 +438,17 @@ class ParallelRequestGauge(TypedDict): descriptor_key: str +def _without_parallel_limit(descriptor: RateLimitDescriptor) -> RateLimitDescriptor: + rate_limit: Final[RateLimitDescriptorRateLimitObject] = { + **(descriptor.get("rate_limit") or MappingProxyType({})), + "max_parallel_requests": None, + } + return RateLimitDescriptor(key=descriptor["key"], value=descriptor["value"], rate_limit=rate_limit) + + class ParallelSlotAcquisition(TypedDict): slot_id: str - counter_keys: list[str] + counter_keys: Sequence[str] class RateLimitStatus(TypedDict): @@ -548,6 +576,15 @@ def get_request_stash() -> RequestRateLimiterStash | None: return _request_stash.get() +@contextmanager +def isolated_request_stash() -> Generator[None]: + token: Final = _request_stash.set(None) + try: + yield + finally: + _request_stash.reset(token) + + def get_or_create_request_stash() -> RequestRateLimiterStash: stash = _request_stash.get() if stash is None: @@ -597,6 +634,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): parallel_acquire_script: _AsyncLuaScript | None parallel_release_script: _AsyncLuaScript | None parallel_count_script: _AsyncLuaScript | None + parallel_renew_script: _AsyncLuaScript | None def __init__( self, @@ -629,6 +667,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): self.parallel_count_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script( PARALLEL_COUNT_SCRIPT ) + self.parallel_renew_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script( + PARALLEL_RENEW_SCRIPT + ) else: self.batch_rate_limiter_script = None self.token_increment_script = None @@ -637,6 +678,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): self.parallel_acquire_script = None self.parallel_release_script = None self.parallel_count_script = None + self.parallel_renew_script = None self.window_size = int(os.getenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", 60)) @@ -998,7 +1040,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): async def in_memory_cache_sliding_window( self, - keys: list[str], + keys: Sequence[str], now_int: int, window_size: int, ) -> CacheCounterValues: @@ -1152,7 +1194,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): crc: Final = binascii.crc_hqx(key.encode("utf-8"), 0) return crc % REDIS_CLUSTER_SLOTS - def _group_keys_by_hash_tag(self, keys: list[str]) -> dict[str, list[str]]: + def _group_keys_by_hash_tag(self, keys: Sequence[str]) -> Mapping[str, Sequence[str]]: """ Group keys by their Redis hash tag to ensure cluster compatibility. @@ -1172,7 +1214,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): groups[slot_key].append(key) else: # For regular Redis, no grouping needed - process all keys together - groups[REDIS_NODE_HASHTAG_NAME] = keys + return MappingProxyType({REDIS_NODE_HASHTAG_NAME: keys}) return groups @@ -1470,12 +1512,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ gauge_keys: Final = [gauge["counter_key"] for gauge in gauges] + if self._is_redis_cluster() and self.parallel_acquire_script is not None: + return await self._check_cluster_parallel_gauges(gauges, slot_id, parent_otel_span, read_only) + if read_only: if self.parallel_count_script is not None: try: raw_counts: Final[list[CacheCounterValue]] = await self.parallel_count_script( keys=gauge_keys, - args=[PARALLEL_REQUEST_SLOT_TTL_SECONDS for _ in gauges], + args=tuple(PARALLEL_REQUEST_SLOT_TTL_SECONDS for _ in gauges), ) counts = [max(0, int(value)) for value in raw_counts] except Exception as e: # noqa: BLE001 - any Redis/Lua failure degrades to the local mirror, never a 500 @@ -1540,6 +1585,135 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): async with self._check_and_increment_lock: return await self._acquire_parallel_slots_in_memory(gauges, slot_id, parent_otel_span) + async def _check_cluster_parallel_gauges( + self, + gauges: Sequence[ParallelRequestGauge], + slot_id: str, + parent_otel_span: Span | None, + read_only: bool, + ) -> RateLimitResponse: + by_key: Final = MappingProxyType( + { + gauge["counter_key"]: min( + (candidate for candidate in gauges if candidate["counter_key"] == gauge["counter_key"]), + key=lambda candidate: candidate["limit"], + ) + for gauge in gauges + } + ) + groups: Final = self._group_keys_by_hash_tag(tuple(by_key)) + counts: Final[dict[str, int]] = {} # mutable-ok: gather independent Redis-slot results + attempted: Final[list[str]] = [] # mutable-ok: rollback includes requests whose responses were lost + try: + for keys in groups.values(): + if read_only: + if self.parallel_count_script is None: + raise RuntimeError("Redis cluster parallel count script is unavailable") + counts.update( + (key, max(0, int(count))) + for key, count in zip( + keys, + await self.parallel_count_script( + keys=keys, args=tuple(PARALLEL_REQUEST_SLOT_TTL_SECONDS for _ in keys) + ), + strict=True, + ) + ) + continue + if self.parallel_acquire_script is None: + raise RuntimeError("Redis cluster parallel acquire script is unavailable") + attempted.extend(keys) + (raw,) = ( + await self.parallel_acquire_script( + keys=keys, + args=tuple( + arg + for key in keys + for arg in (by_key[key]["limit"], PARALLEL_REQUEST_SLOT_TTL_SECONDS, slot_id) + ), + ), + ) + if int(raw[0]) == 1: + await self._rollback_cluster_parallel_slots(tuple(attempted), slot_id, parent_otel_span) + return RateLimitResponse( + overall_code="OVER_LIMIT", + statuses=[ # mutable-ok: response contract requires a list + self._gauge_status(by_key[keys[int(raw[1]) - 1]], int(raw[2]), "OVER_LIMIT") + ], + ) + counts.update((key, int(count)) for key, count in zip(keys, raw[1:], strict=True)) + for key in keys: + await self.internal_usage_cache.async_set_cache( + key=key, + value=counts[key], + ttl=PARALLEL_REQUEST_SLOT_TTL_SECONDS, + litellm_parent_otel_span=parent_otel_span, + local_only=True, + ) + except BaseException: + if attempted: + await self._rollback_cluster_parallel_slots(tuple(attempted), slot_id, parent_otel_span) + raise + statuses: Final = tuple( + self._gauge_status( + gauge, + counts[gauge["counter_key"]], + "OVER_LIMIT" if read_only and counts[gauge["counter_key"]] >= gauge["limit"] else "OK", + ) + for gauge in gauges + ) + return RateLimitResponse( + overall_code="OVER_LIMIT" if any(item["code"] == "OVER_LIMIT" for item in statuses) else "OK", + statuses=list(statuses), # mutable-ok: RateLimitResponse contract requires a list + ) + + async def _rollback_cluster_parallel_slots( + self, counter_keys: tuple[str, ...], slot_id: str, parent_otel_span: Span | None + ) -> None: + rollback: Final = asyncio.create_task( + self._release_cluster_parallel_slots(counter_keys, slot_id, parent_otel_span) + ) + cancelled = False # rebind-ok: defer repeated caller cancellation until compensation finishes + while not rollback.done(): + try: + await asyncio.shield(rollback) + except asyncio.CancelledError: + cancelled = True + except Exception: # noqa: BLE001 # retrieve and report the completed task's exception below + break + try: + rollback.result() + except Exception: # noqa: BLE001 # preserve admission failure; unreachable Redis slots expire by TTL + verbose_proxy_logger.error("Could not roll back all Redis cluster parallel request slots") + if cancelled: + raise asyncio.CancelledError + + async def _release_cluster_parallel_slots( + self, counter_keys: tuple[str, ...], slot_id: str, parent_otel_span: Span | None + ) -> None: + first_error: Exception | None = None # rebind-ok: finish every shard before reporting the first failure + for keys in self._group_keys_by_hash_tag(counter_keys).values(): + try: + if self.parallel_release_script is None: + raise RuntimeError("Redis cluster parallel release script is unavailable") + for key, count in zip( + keys, + await self.parallel_release_script(keys=keys, args=tuple(slot_id for _ in keys)), + strict=True, + ): + await self.internal_usage_cache.async_set_cache( + key=key, + value=max(0, int(count)), + ttl=PARALLEL_REQUEST_SLOT_TTL_SECONDS, + litellm_parent_otel_span=parent_otel_span, + local_only=True, + ) + except Exception as exc: # noqa: BLE001 # one unreachable shard must not strand the other shards + if first_error is None: + first_error = exc + if first_error is not None: + raise first_error + async def _read_local_gauge_counts( self, gauge_keys: list[str], @@ -1609,6 +1783,68 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): statuses.append(self._gauge_status(gauge, in_flight + 1, "OK")) return RateLimitResponse(overall_code="OK", statuses=statuses) + def transfer_realtime_call_slot(self, request_data: Mapping[str, object]) -> RealtimeCallLease | None: + call_id: Final = request_data.get("litellm_call_id") + if not isinstance(call_id, str): + return None + stash: Final = get_request_stash_for_call(call_id) + if stash is None or stash.parallel_slot is None: + return None + slot_id: Final = stash.parallel_slot["slot_id"] + counter_keys: Final = tuple(stash.parallel_slot["counter_keys"]) + stash.parallel_slot = None + + async def renew() -> bool: + return await self._renew_realtime_call_slot(slot_id, counter_keys) + + async def release() -> None: + await self._release_parallel_request_slots( + ParallelSlotAcquisition(slot_id=slot_id, counter_keys=counter_keys) + ) + + return RealtimeCallLease(renew=renew, release=release) + + async def _renew_realtime_call_slot(self, slot_id: str, counter_keys: tuple[str, ...]) -> bool: + if self.parallel_renew_script is not None: + try: + for keys in self._group_keys_by_hash_tag(counter_keys).values(): + if tuple( + await self.parallel_renew_script(keys=keys, args=(slot_id, PARALLEL_REQUEST_SLOT_TTL_SECONDS)) + ) != (1,): + return False + return True + except Exception: # noqa: BLE001 # Redis ownership cannot be established by a local count mirror + return False + async with self._check_and_increment_lock: + now: Final = self._get_current_time().timestamp() + cutoff: Final = now - PARALLEL_REQUEST_SLOT_TTL_SECONDS + values: Final[tuple[ParallelGaugeCacheValue | None, ...]] = tuple( + [ + await self.internal_usage_cache.async_get_cache( + key=counter_key, local_only=True, litellm_parent_otel_span=None + ) + for counter_key in counter_keys + ] + ) + if any( + not isinstance(value, dict) + or not isinstance(score := value.get(slot_id), (int, float)) + or score <= cutoff + for value in values + ): + return False + for counter_key, value in zip(counter_keys, values): + if not isinstance(value, dict): + return False + await self.internal_usage_cache.async_set_cache( + key=counter_key, + value={**value, slot_id: now}, # mutable-ok: slot registry readers require a concrete dict + ttl=PARALLEL_REQUEST_SLOT_TTL_SECONDS, + local_only=True, + litellm_parent_otel_span=None, + ) + return True + async def _release_parallel_request_slots( self, acquisition: ParallelSlotAcquisition, @@ -1626,11 +1862,14 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): slot_id: Final = acquisition["slot_id"] if not counter_keys or not slot_id: return + if self._is_redis_cluster() and self.parallel_release_script is not None: + await self._release_cluster_parallel_slots(tuple(counter_keys), slot_id, parent_otel_span) + return if self.parallel_release_script is not None: try: raw: Final[list[CacheCounterValue]] = await self.parallel_release_script( keys=counter_keys, - args=[slot_id for _ in counter_keys], + args=tuple(slot_id for _ in counter_keys), ) for counter_key, remaining in zip(counter_keys, raw): await self.internal_usage_cache.async_set_cache( @@ -3497,6 +3736,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # Org Level Rate Limits descriptors.extend(self.create_organization_rate_limit_descriptor(user_api_key_dict, requested_model)) + effective_descriptors: Final = ( + tuple(_without_parallel_limit(descriptor) for descriptor in descriptors) + if call_type == "_arealtime" and is_realtime_call_attachment(data.get("websocket")) + else descriptors + ) + # Only check rate limits if we have descriptors with actual limits if descriptors: # First pass: RPM and max_parallel_requests sliding-window check. @@ -3515,16 +3760,18 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # double-charge every request. parallel_counter_keys: Final = [ self.create_rate_limit_keys(d["key"], d["value"], "max_parallel_requests") - for d in descriptors + for d in effective_descriptors if (d.get("rate_limit") or {}).get("max_parallel_requests") is not None ] parallel_slot_id: Final = uuid.uuid4().hex if parallel_counter_keys else None first_pass_descriptors: Final = ( - descriptors + effective_descriptors if self.tpm_reservation_enabled else tuple( - d for d in descriptors if d["key"] not in (PROJECT_ITPM_DESCRIPTOR_KEY, PROJECT_OTPM_DESCRIPTOR_KEY) + d + for d in effective_descriptors + if d["key"] not in (PROJECT_ITPM_DESCRIPTOR_KEY, PROJECT_OTPM_DESCRIPTOR_KEY) ) ) response: Final = await self.should_rate_limit( @@ -4700,6 +4947,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) stash.parallel_slot = None + async def async_release_realtime_attachment( + self, request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth + ) -> None: + await self.async_post_call_failure_hook( + request_data={}, # mutable-ok: existing failure hook requires dict; attachment has no billable usage + original_exception=Exception("Realtime attachment completed"), + user_api_key_dict=user_api_key_dict, + ) + async def async_post_call_success_hook(self, data: dict, user_api_key_dict: UserAPIKeyAuth, response): """ Post-call hook to update rate limit headers in the response. diff --git a/litellm/proxy/hooks/realtime_call_lease.py b/litellm/proxy/hooks/realtime_call_lease.py new file mode 100644 index 00000000000..c33d81f92d3 --- /dev/null +++ b/litellm/proxy/hooks/realtime_call_lease.py @@ -0,0 +1,74 @@ +import asyncio +from collections.abc import Awaitable, Callable, Generator +from contextlib import contextmanager +from contextvars import ContextVar +from typing import Final + +_realtime_call_attachment: Final[ContextVar[object | None]] = ContextVar("realtime_call_attachment", default=None) + + +@contextmanager +def realtime_call_attachment(websocket: object) -> Generator[None]: + token: Final = _realtime_call_attachment.set(websocket) + try: + yield + finally: + _realtime_call_attachment.reset(token) + + +def is_realtime_call_attachment(websocket: object) -> bool: + bound: Final = _realtime_call_attachment.get() + return bound is not None and bound is websocket + + +class RealtimeCallLease: + def __init__( + self, + *, + renew: Callable[[], Awaitable[bool]], + release: Callable[[], Awaitable[None]], + interval: float = 300, + renewal_timeout: float = 10, + ) -> None: + self._renew = renew + self._release = release + self._interval = interval + self._renewal_timeout = renewal_timeout + self._failed = asyncio.Event() + self._heartbeat: asyncio.Task[None] | None = None + self._closing: asyncio.Task[None] | None = None + + def start(self) -> None: + if self._heartbeat is None and self._closing is None: + self._heartbeat = asyncio.create_task(self._run()) + + async def renew(self) -> bool: + if self._closing is not None or self._failed.is_set(): + return False + try: + renewed: Final = await asyncio.wait_for(self._renew(), timeout=self._renewal_timeout) + except Exception: # noqa: BLE001 # fail closed without exposing cache credentials + self._failed.set() + return False + if not renewed: + self._failed.set() + return renewed and not self._failed.is_set() and self._closing is None + + async def wait_failed(self) -> None: + await self._failed.wait() + + async def _run(self) -> None: + while await self.renew(): + await asyncio.sleep(self._interval) + self._failed.set() + + async def close(self) -> None: + if self._closing is None: + self._closing = asyncio.create_task(self._close()) + await asyncio.shield(self._closing) + + async def _close(self) -> None: + if self._heartbeat is not None: + self._heartbeat.cancel() + await asyncio.gather(self._heartbeat, return_exceptions=True) + await self._release() diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 606a590c24b..f6c5e2d9ed5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1399,6 +1399,10 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: except Exception as e: verbose_proxy_logger.error("Error stopping DB health watchdog task: %s", e) + from litellm.proxy.realtime_endpoints.call_supervision import CALL_SUPERVISORS + + await CALL_SUPERVISORS.shutdown() + await _drain_spend_event_producer_on_shutdown() await flush_spend_counters_on_shutdown() @@ -11917,6 +11921,19 @@ async def _reject_realtime_session( await _release_realtime_budget_reservation(user_api_key_dict) +@app.websocket("/v1/live/{call_id}") +async def codex_live_sideband_endpoint( + websocket: WebSocket, + call_id: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth_websocket), +) -> None: + from litellm.proxy.realtime_endpoints.call_sessions import codex_realtime_sideband + + await codex_realtime_sideband(websocket, call_id, user_api_key_dict) + + +@app.websocket("/v1/live") +@app.websocket("/live") @app.websocket("/openai/v1/realtime") @app.websocket("/v1/realtime") @app.websocket("/realtime") @@ -11924,12 +11941,18 @@ async def realtime_websocket_endpoint( websocket: WebSocket, model: str | None = fastapi.Query(None, description="The model to use for the websocket connection."), intent: str | None = fastapi.Query(None, description="The intent of the websocket connection."), + call_id: str | None = None, guardrails: str | None = fastapi.Query( None, description="Comma-separated list of guardrail names to apply to this request.", ), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth_websocket), ): + if call_id is not None: + from litellm.proxy.realtime_endpoints.call_sessions import codex_realtime_sideband + + await codex_realtime_sideband(websocket, call_id, user_api_key_dict) + return requested_protocols: Final = [ p.strip() for p in (websocket.headers.get("sec-websocket-protocol") or "").split(",") if p.strip() ] @@ -11960,7 +11983,10 @@ async def realtime_websocket_endpoint( await websocket.accept(**accept_kwargs) # Only use explicit parameters, not all query params - query_params: Final = cast(RealtimeQueryParams, dict(_realtime_query_params_template(model, intent))) + query_params: Final = cast( + RealtimeQueryParams, + dict(_realtime_query_params_template(model, intent) + ((("call_id", call_id),) if call_id is not None else ())), + ) data: dict[str, object] = { "model": route_model, diff --git a/litellm/proxy/realtime_endpoints/call_sessions.py b/litellm/proxy/realtime_endpoints/call_sessions.py new file mode 100644 index 00000000000..8989fadfead --- /dev/null +++ b/litellm/proxy/realtime_endpoints/call_sessions.py @@ -0,0 +1,542 @@ +import asyncio +import base64 +import hashlib +import json +import time +from collections.abc import Awaitable, Callable, Mapping +from contextlib import AsyncExitStack, nullcontext +from contextvars import Token +from types import MappingProxyType +from typing import Final, Literal + +import httpx +from fastapi import HTTPException, Request, Response, WebSocket +from pydantic import TypeAdapter +from starlette.formparsers import MultiPartException, MultiPartParser +from starlette.types import Message, Scope + +from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.realtime_streaming import ( + REALTIME_SESSION_SUCCESS_LOGGED_KEY, + RealTimeStreaming, + realtime_attachment_cleanup, +) +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.chatgpt.codex import ( + CodexRealtimeCall, + CodexRealtimeOffer, + build_call_request, + build_sideband_request, + parse_call_response, +) +from litellm.llms.chatgpt.realtime import ( + CallAccounting, + ChatGPTRealtime, + configured_realtime_headers, + realtime_endpoint, +) +from litellm.proxy._types import InternalRequestOrigin, ProxyException, UserAPIKeyAuth +from litellm.proxy.auth.auth_checks import can_key_call_resolved_model +from litellm.proxy.auth.user_api_key_auth import ( + get_api_key, + get_api_key_from_custom_header, + get_websocket_api_key, + user_api_key_auth, +) +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper, encrypt_value_helper +from litellm.proxy.common_utils.http_parsing_utils import ( + _normalize_media_type, # pyright: ignore[reportPrivateUsage] # reuse the shared HTTP media-type normalization contract +) +from litellm.proxy.hooks.parallel_request_limiter import ( + _PROXY_MaxParallelRequestsHandler, # pyright: ignore[reportPrivateUsage] # existing built-in limiter has no public alias +) +from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + _PROXY_MaxParallelRequestsHandler_v3, # pyright: ignore[reportPrivateUsage] # existing built-in limiter has no public alias + isolated_request_stash, +) +from litellm.proxy.hooks.realtime_call_lease import RealtimeCallLease, realtime_call_attachment +from litellm.proxy.spend_tracking.budget_reservation import ( + invalidate_budget_reservation_counters, + release_or_invalidate_budget_reservation, +) +from litellm.types.realtime import RealtimeQueryParams +from litellm.types.router import GenericLiteLLMParams + + +async def supervise_codex_call( + request: Request, call: CodexRealtimeCall, auth: UserAPIKeyAuth, lease: RealtimeCallLease | None = None +) -> None: + with isolated_request_stash(): + await _start_codex_supervisor(request, call, auth, lease) + + +async def _start_codex_supervisor( + request: Request, call: CodexRealtimeCall, auth: UserAPIKeyAuth, lease: RealtimeCallLease | None +) -> None: + import litellm + from litellm.proxy.realtime_endpoints.call_supervision import CALL_SUPERVISORS, CallSupervisor + + async def receive() -> Message: + body: Final[RealtimeQueryParams] = {"model": call.alias} + message: Final[Message] = { + "type": "http.request", + "body": json.dumps(body).encode(), + "more_body": False, + } + return message + + async def send(_message: Message) -> None: + return None + + supervision_owned = False # rebind-ok: supervisor owns cleanup after construction + effective_handler: ChatGPTRealtime | None = None # rebind-ok: reuse hook-enriched credentials for cleanup + sockets: Final = AsyncExitStack() + try: + observer_scope: Final[Scope] = {**request.scope} + observer_request: Final = Request(observer_scope, receive=receive) + processed, logger = await process_codex_request( + observer_request, + { # mutable-ok: common request processing enriches metadata + **build_sideband_request(call), + "model": call.alias, + }, + auth, + call.alias, + "_arealtime", + internal_realtime_observer=True, + ) + pinned: Final = { # mutable-ok: logging and provider parameter contract + **processed, + **build_sideband_request(call), + "extra_headers": MappingProxyType( + { + **configured_realtime_headers( + TypeAdapter(Mapping[str, object] | None).validate_python(processed.get("extra_headers")) + ), + **configured_realtime_headers(call.extra_headers), + } + ), + "litellm_metadata": { # mutable-ok: Logging.update_from_kwargs requires a dict to retain ownership metadata + **TypeAdapter(Mapping[str, object]).validate_python( + processed.get("litellm_metadata") or MappingProxyType({}) + ), + **( + MappingProxyType( + { + "model_info": { # mutable-ok: logging and cost callbacks require a concrete model-info dict + **litellm.get_model_info(model=call.model_id), + "id": call.model_id, + } + } + ) + if call.model_id is not None + else MappingProxyType({}) + ), + }, + } + logger.update_from_kwargs( + kwargs=pinned, + model=call.model, + user=None, + optional_params={}, # mutable-ok: logging contract + litellm_params={ # mutable-ok: Logging.update_from_kwargs pops metadata from its argument + **logger.litellm_params, + "litellm_metadata": pinned["litellm_metadata"], + "arealtime": True, + }, + custom_llm_provider="chatgpt", + ) + params: Final = GenericLiteLLMParams.model_validate(pinned) + handler: Final = ChatGPTRealtime( + params, request.headers, TypeAdapter(Mapping[str, object]).validate_python(pinned["extra_headers"]) + ) + effective_handler = handler + api_base: Final = ChatGPTRealtime.get_api_base(call.api_base) + connection: Final = await handler.open_call_connection(call.model, api_base) + sockets.push_async_callback(connection.close) + + async def close_call() -> None: + await handler.close_call(connection, call.model, api_base) + + async def force_close_call() -> None: + await handler.hangup_call(api_base) + + frontend_scope: Final[Scope] = {**request.scope, "type": "websocket"} + frontend: Final = WebSocket(frontend_scope, receive=receive, send=send) + stream: Final = RealTimeStreaming(frontend, connection, logger, model=call.model, user_api_key_dict=auth) + supervisor: Final = CallSupervisor( + connection, + stream, + logger, + auth, + close_call, + force_close_call=force_close_call, + terminal_usage_required=realtime_endpoint(call.model) == "live", + lease=lease, + ) + supervision_owned = True + sockets.pop_all() + await CALL_SUPERVISORS.start(supervisor) + except BaseException: + if not supervision_owned: + try: + fallback_handler: Final = effective_handler or ChatGPTRealtime( + GenericLiteLLMParams.model_validate(build_sideband_request(call)), + request.headers, + call.extra_headers, + ) + await fallback_handler.hangup_call(ChatGPTRealtime.get_api_base(call.api_base)) + except Exception: # noqa: BLE001 # preserve original failure without logging provider credentials + verbose_proxy_logger.error("Realtime startup cleanup could not confirm upstream termination") + try: + await invalidate_budget_reservation_counters(budget_reservation=auth.budget_reservation) + except Exception: # noqa: BLE001 # cleanup errors must not replace the original startup failure + verbose_proxy_logger.error("Realtime startup cleanup could not invalidate budget counters") + else: + await release_or_invalidate_budget_reservation(budget_reservation=auth.budget_reservation) + finally: + try: + await sockets.aclose() + except Exception: # noqa: BLE001 # socket cleanup must preserve the original startup failure + verbose_proxy_logger.error("Realtime startup cleanup could not close observer socket") + raise + + +def encode_call(call: CodexRealtimeCall) -> str: + encrypted: Final = encrypt_value_helper(call.model_dump_json()) + return "rtc_litellm_" + base64.urlsafe_b64encode(encrypted.encode()).decode().rstrip("=") + + +def decode_call(token: str, authorization: str) -> CodexRealtimeCall: + try: + if not token.startswith("rtc_litellm_"): + raise ValueError("Invalid call prefix") + encoded: Final = token.removeprefix("rtc_litellm_") + encrypted: Final = base64.b64decode(encoded + "=" * (-len(encoded) % 4), altchars=b"-_", validate=True) + plaintext: Final = decrypt_value_helper(encrypted.decode(), key="codex_realtime_call") + call: Final = CodexRealtimeCall.model_validate_json(plaintext or "") + except (ValueError, TypeError, UnicodeError) as exc: + raise HTTPException(403, "Invalid realtime call") from exc + if call.expires_at < time.time() or call.owner != hashlib.sha256(authorization.encode()).hexdigest(): + raise HTTPException(403, "Invalid or expired realtime call") + return call + + +MAX_REALTIME_OFFER_BYTES: Final = 8 * 1024 * 1024 + + +async def _cache_bounded_offer_body(request: Request) -> None: + try: + if int(request.headers.get("content-length", "")) > MAX_REALTIME_OFFER_BYTES: + raise HTTPException(413, "Realtime offer exceeds the 8 MiB limit") + except ValueError: + pass + if hasattr(request, "_body"): + if len(request._body) > MAX_REALTIME_OFFER_BYTES: # pyright: ignore[reportPrivateUsage] # validate Starlette's cached body without consuming it again + raise HTTPException(413, "Realtime offer exceeds the 8 MiB limit") + return + if request._form is not None and request._stream_consumed: # pyright: ignore[reportPrivateUsage] # a mixed-case empty form cache may leave the stream unread + return + body: Final = bytearray() + async for chunk in request.stream(): + if len(body) + len(chunk) > MAX_REALTIME_OFFER_BYTES: + raise HTTPException(413, "Realtime offer exceeds the 8 MiB limit") + body.extend(chunk) + request._body = bytes(body) # pyright: ignore[reportPrivateUsage] # Starlette has no public setter for its shared body cache + + +async def read_codex_offer(request: Request) -> CodexRealtimeOffer: + await _cache_bounded_offer_body(request) + content_type: Final = request.headers.get("content-type", "") + if _normalize_media_type(content_type) == "multipart/form-data": + if content_type.split(";", 1)[0] != "multipart/form-data" and not await request.form(): + try: + request._form = await MultiPartParser(request.headers, request.stream()).parse() # pyright: ignore[reportPrivateUsage] # Starlette exposes no setter for its shared form cache; # rebind-ok: Request.close must own and close uploaded files + request.scope.pop("parsed_body", None) + except MultiPartException as exc: + raise HTTPException(400, "Invalid realtime multipart offer") from exc + form: Final = await request.form() + return CodexRealtimeOffer.model_validate( + MappingProxyType({"sdp": form.get("sdp"), "session": json.loads(str(form.get("session", "{}")))}) + ) + return CodexRealtimeOffer.model_validate(await request.json()) + + +async def process_codex_request( + request: Request, + data: dict[str, object], # mutable-ok: common request processor enriches this dictionary + auth: UserAPIKeyAuth, + model: str, + route_type: Literal["arealtime_calls", "_arealtime"], + *, + internal_realtime_observer: bool = False, +) -> tuple[dict[str, object], Logging]: # mutable-ok: common request processor returns enriched routing arguments + from litellm.proxy import proxy_server as server + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + processor: Final = ProxyBaseLLMRequestProcessing(data=data) + processed, logging_obj = await processor.common_processing_pre_call_logic( + request=request, + general_settings=server.general_settings, + user_api_key_dict=auth, + version=server.version, + proxy_logging_obj=server.proxy_logging_obj, + proxy_config=server.proxy_config, + llm_router=server.llm_router, + user_model=server.user_model, + user_temperature=server.user_temperature, + user_request_timeout=server.user_request_timeout, + user_max_tokens=server.user_max_tokens, + user_api_base=server.user_api_base, + model=model, + route_type=route_type, + **( + MappingProxyType({"internal_realtime_observer": True}) + if internal_realtime_observer + else MappingProxyType({}) + ), + ) + if internal_realtime_observer: + logging_obj.model_call_details["internal_request_origin"] = InternalRequestOrigin.REALTIME_OBSERVER + return processed, logging_obj + + +async def create_codex_realtime_call(request: Request) -> Response: + try: + with isolated_request_stash(): + return await _create_codex_realtime_call(request) + finally: + await request.close() + + +async def _create_codex_realtime_call(request: Request) -> Response: + from litellm.proxy import proxy_server as server + + try: + offer: Final = await read_codex_offer(request) + except ValueError as exc: + raise HTTPException(400, "Invalid realtime offer: expected sdp and session") from exc + model: Final = offer.session.model + if not model: + raise HTTPException(400, "session.model is required") + auth: Final = await user_api_key_auth( + request=request, + api_key=request.headers.get("authorization", ""), + azure_api_key_header=request.headers.get("api-key", ""), + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + custom_litellm_key_header=request.headers.get("x-litellm-api-key"), + ) + selected_key, _ = get_api_key( + request=request, + api_key=request.headers.get("authorization", ""), + azure_api_key_header=request.headers.get("api-key", ""), + custom_litellm_key_header=request.headers.get("x-litellm-api-key"), + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + pass_through_endpoints=None, + route="/v1/realtime/calls", + ) + custom_header: Final = server.general_settings.get("litellm_key_header_name") + owner_key: Final = ( + get_api_key_from_custom_header(request, custom_header) if isinstance(custom_header, str) else selected_key + ) + supervision_started = False # rebind-ok: transfer reservation ownership only after supervision is established + call_lease: RealtimeCallLease | None = None + lease_transferred = False # rebind-ok: failed startup leaves the signaling task responsible for its lease + preprocessing_started = False # rebind-ok: only refund reservations belonging to this signaling request + limiter: Final = server.proxy_logging_obj.get_proxy_hook("parallel_request_limiter") + try: + await can_key_call_resolved_model( + model=model, + llm_model_list=server.llm_model_list, + valid_token=auth, + llm_router=server.llm_router, + ) + data: Final = build_call_request(offer, request.query_params, request.headers) + signaling_auth: Final = auth.model_copy(update=MappingProxyType({"budget_reservation": None})) + if isinstance(limiter, _PROXY_MaxParallelRequestsHandler) and ( + auth.max_parallel_requests is not None + or server.general_settings.get("global_max_parallel_requests") is not None + ): + raise HTTPException(400, "Realtime calls with parallel limits require the V3 rate limiter") + preprocessing_started = True + processed, _ = await process_codex_request(request, data, signaling_auth, model, "arealtime_calls") + if isinstance(limiter, _PROXY_MaxParallelRequestsHandler_v3): + call_lease = limiter.transfer_realtime_call_slot(processed) + if call_lease is not None: + call_lease.start() + if not await call_lease.renew(): + raise HTTPException(503, "Realtime call quota reservation was lost") + with isolated_request_stash(): + result: Final = await server.route_request( + data=processed, + route_type="arealtime_calls", + llm_router=server.llm_router, + user_model=server.user_model, + ) + try: + response: Final = await result + except BaseLLMException as exc: + raise HTTPException(exc.status_code, str(exc)) from exc + if not isinstance(response, httpx.Response): + raise HTTPException(502, "Invalid realtime signaling response") + if response.is_error: + return Response(response.content, status_code=response.status_code, media_type="application/json") + try: + call: Final = parse_call_response( + response, + alias=model, + owner=hashlib.sha256(f"Bearer {owner_key}".encode()).hexdigest(), + expires_at=time.time() + 3600, + ) + except ValueError as exc: + raise HTTPException(400, str(exc)) from exc + supervised_call: Final = call.model_copy( + update=MappingProxyType({"usage_supervised": True, "parallel_reserved": call_lease is not None}) + ) + token: Final = encode_call(supervised_call) + supervision_started = True + if call_lease is None: + await supervise_codex_call(request, supervised_call, auth) + else: + await supervise_codex_call(request, supervised_call, auth, call_lease) + lease_transferred = True + return Response( + response.content, + status_code=response.status_code, + media_type="application/sdp", + headers=MappingProxyType({"Location": f"/v1/realtime/calls/{token}"}), + ) + finally: + try: + if call_lease is not None and not lease_transferred: + await call_lease.close() + finally: + try: + if preprocessing_started and isinstance(limiter, _PROXY_MaxParallelRequestsHandler_v3): + await asyncio.shield( + limiter.async_post_call_failure_hook( + request_data={}, # mutable-ok: existing failure-hook contract + original_exception=Exception("Realtime signaling completed without token usage"), + user_api_key_dict=auth, + ) + ) + finally: + if not supervision_started: + await release_or_invalidate_budget_reservation(budget_reservation=auth.budget_reservation) + + +async def codex_realtime_sideband(websocket: WebSocket, token: str, auth: UserAPIKeyAuth) -> None: + import litellm + from litellm.proxy import proxy_server as server + + protocols: Final = tuple( + p.strip() for p in websocket.headers.get("sec-websocket-protocol", "").split(",") if p.strip() + ) + logging_obj: Logging | None = None # rebind-ok: cleanup needs the logger only after pre-call succeeds + attachment_limiter: _PROXY_MaxParallelRequestsHandler | _PROXY_MaxParallelRequestsHandler_v3 | None = None + cleanup_token: Token[Callable[[], Awaitable[None]] | None] | None = None + try: + try: + api_key: Final = get_websocket_api_key(websocket) + if not api_key: + raise HTTPException(403, "No API key provided") + call: Final = decode_call(token, f"Bearer {api_key}") + await can_key_call_resolved_model( + model=call.alias, + llm_model_list=server.llm_model_list, + valid_token=auth, + llm_router=server.llm_router, + ) + except (HTTPException, ProxyException): + await websocket.close(code=1008, reason="Invalid realtime call") + return + + async def receive() -> Message: + return { # mutable-ok: ASGI receive message + "type": "http.request", + "body": json.dumps({"model": call.alias}).encode(), # mutable-ok: JSON request serialization + "more_body": False, + } + + request: Final = Request( + { # mutable-ok: Starlette stores request state in the ASGI scope + **websocket.scope, + "type": "http", + "method": "POST", + "path": websocket.scope.get("path", "/v1/realtime"), + }, + receive=receive, + ) + data: Final = { # mutable-ok: common request processor enriches routing arguments + **build_sideband_request(call), + "model": call.alias, + "websocket": websocket, + "guardrails": [ # mutable-ok: guardrail processing expects a list + name.strip() for name in websocket.query_params.get("guardrails", "").split(",") if name.strip() + ], + } + limiter: Final = server.proxy_logging_obj.get_proxy_hook("parallel_request_limiter") + if call.usage_supervised and isinstance( + limiter, (_PROXY_MaxParallelRequestsHandler, _PROXY_MaxParallelRequestsHandler_v3) + ): + attachment_limiter = limiter + if isinstance(limiter, _PROXY_MaxParallelRequestsHandler): + limiter.begin_realtime_attachment(data) + try: + with realtime_call_attachment(websocket) if call.parallel_reserved else nullcontext(): + processed, logging_obj = await process_codex_request(request, data, auth, call.alias, "_arealtime") + except Exception: # noqa: BLE001 # custom hook exceptions must reject the connection + verbose_proxy_logger.exception("Realtime sideband pre-call rejected") + await websocket.close(code=1008, reason="Realtime pre-call rejected") + return + await websocket.accept( + subprotocol=next((p for p in protocols if not p.startswith("openai-insecure-api-key.")), None) + ) + if attachment_limiter is not None: + selected_limiter: Final = attachment_limiter + + async def release_attachment() -> None: + await selected_limiter.async_release_realtime_attachment(data, auth) + + cleanup_token = realtime_attachment_cleanup.set(release_attachment) + await litellm._arealtime( # pyright: ignore[reportPrivateUsage] # dispatch for an already authorized call + model=f"chatgpt/{call.model}", + websocket=websocket, + **MappingProxyType( + { + key: value + for key, value in { # mutable-ok: retain processed metadata with pinned routing + **processed, + **build_sideband_request(call), + "extra_headers": MappingProxyType( + { + **configured_realtime_headers( + TypeAdapter(Mapping[str, object] | None).validate_python( + processed.get("extra_headers") + ) + ), + **configured_realtime_headers(call.extra_headers), + } + ), + "websocket": websocket, + "user_api_key_dict": auth, + "chatgpt_call_accounting": CallAccounting.SUPERVISED if call.usage_supervised else None, + }.items() + if key not in ("model", "websocket") + } + ), + ) + finally: + try: + if attachment_limiter is not None: + await attachment_limiter.async_release_realtime_attachment(data, auth) + finally: + if cleanup_token is not None: + realtime_attachment_cleanup.reset(cleanup_token) + if logging_obj is None or not logging_obj.model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY): + await release_or_invalidate_budget_reservation(budget_reservation=auth.budget_reservation) diff --git a/litellm/proxy/realtime_endpoints/call_supervision.py b/litellm/proxy/realtime_endpoints/call_supervision.py new file mode 100644 index 00000000000..f511d943c3a --- /dev/null +++ b/litellm/proxy/realtime_endpoints/call_supervision.py @@ -0,0 +1,270 @@ +import asyncio +from collections.abc import AsyncIterator, Awaitable, Callable +from contextlib import suppress +from typing import Final, Protocol + +from pydantic import BaseModel, Field, ValidationError +from websockets.exceptions import ConnectionClosedOK + +from litellm._logging import verbose_proxy_logger +from litellm.constants import LOGGING_WORKER_MAX_TIME_PER_COROUTINE +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.realtime_streaming import REALTIME_SESSION_SUCCESS_LOGGED_KEY +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.hooks.realtime_call_lease import RealtimeCallLease +from litellm.proxy.spend_tracking.budget_reservation import ( + invalidate_budget_reservation_counters, + release_or_invalidate_budget_reservation, +) + + +class ObserverSocket(Protocol): + def __aiter__(self) -> AsyncIterator[str | bytes]: ... + + async def close(self) -> None: ... + + +class UsageSink(Protocol): + def store_message(self, message: str) -> None: ... + + async def log_messages(self, *, wait_for_dispatch: bool = False) -> None: ... + + +class _ObserverEvent(BaseModel): + type: str + + +class _LiveDurationUsage(BaseModel): + audio_duration_ms: float = Field(strict=True, ge=0, allow_inf_nan=False) + + +class _LiveTerminalEvent(BaseModel): + usage: _LiveDurationUsage + + +class CallSupervisor: + def __init__( + self, + upstream: ObserverSocket, + stream: UsageSink, + logging_obj: Logging, + auth: UserAPIKeyAuth, + close_call: Callable[[], Awaitable[None]], + *, + ready_timeout: float = 20, + lifetime: float = 3600, + drain_timeout: float = 5, + termination_timeout: float = 60, + logging_timeout: float = LOGGING_WORKER_MAX_TIME_PER_COROUTINE, + terminal_usage_required: bool = True, + force_close_call: Callable[[], Awaitable[None]] | None = None, + lease: RealtimeCallLease | None = None, + ) -> None: + self._upstream = upstream + self._stream = stream + self._logging = logging_obj + self._auth = auth + self._close_call = close_call + self._force_close_call = force_close_call + self._lease = lease + self._ready_timeout = ready_timeout + self._lifetime = lifetime + self._drain_timeout = drain_timeout + self._termination_timeout = termination_timeout + self._logging_timeout = logging_timeout + self._terminal_usage_required = terminal_usage_required + self._ready = asyncio.Event() + self._stop = asyncio.Event() + self._started = False + self._terminal = False + self._terminal_usage_valid = False + self._close_confirmed = False + self._accounting_complete = False + self._task: asyncio.Task[None] | None = None + + async def start(self) -> None: + if self._task is not None: + raise RuntimeError("Call observer already started") + self._task = asyncio.create_task(self._run()) + try: + await asyncio.wait_for(self._ready.wait(), timeout=self._ready_timeout) + if self._lease is not None and not await self._lease.renew(): + raise RuntimeError("Call observer lost its quota reservation during startup") + if not self._started or self._terminal or self._task.done(): + raise RuntimeError("Call observer ended before session became available") + except BaseException: + await self._close_after_failed_start() + raise + + async def _close_after_failed_start(self) -> None: + cleanup: Final = asyncio.create_task(self.close()) + while not cleanup.done(): + with suppress(asyncio.CancelledError): + await asyncio.shield(cleanup) + cleanup.result() + + async def close(self) -> None: + self._stop.set() + await self.wait() + + async def wait(self) -> None: + if self._task is not None: + await asyncio.shield(self._task) + + async def _read(self) -> None: + try: + await self._read_events() + except ConnectionClosedOK: + return + + async def _read_events(self) -> None: + event: _ObserverEvent + async for message in self._upstream: + self._stream.store_message(message.decode("utf-8") if isinstance(message, bytes) else message) + event = _ObserverEvent.model_validate_json(message) + if event.type in ("session.started", "session.created"): + self._started = True + self._ready.set() + if event.type == "session.closed": + self._terminal = True + try: + _LiveTerminalEvent.model_validate_json(message) + except ValidationError: + self._terminal_usage_valid = False + else: + self._terminal_usage_valid = True + return + + def _usage_complete(self) -> bool: + if self._terminal_usage_required: + return self._terminal and self._terminal_usage_valid + return self._terminal or self._close_confirmed + + async def _run(self) -> None: + try: + await self._observe() + finally: + try: + if self._lease is not None: + await self._lease.close() + finally: + self._ready.set() + + async def _observe(self) -> None: + reader: Final = asyncio.create_task(self._read()) + stopped: Final = asyncio.create_task(self._stop.wait()) + lease_failed: Final = asyncio.create_task(self._lease.wait_failed()) if self._lease is not None else None + try: + await asyncio.wait( + (reader, stopped, lease_failed) if lease_failed is not None else (reader, stopped), + timeout=self._lifetime, + return_when=asyncio.FIRST_COMPLETED, + ) + finally: + try: + if not self._terminal: + deadline: Final = asyncio.get_running_loop().time() + self._termination_timeout + primary_deadline: Final = ( + deadline - self._termination_timeout / 2 + if self._terminal_usage_required and self._force_close_call is not None + else deadline + ) + try: + await asyncio.wait_for( + self._close_call(), timeout=max(0.0, primary_deadline - asyncio.get_running_loop().time()) + ) + self._close_confirmed = True + except Exception: # noqa: BLE001 # provider exceptions can contain credentials + verbose_proxy_logger.error("Realtime observer could not terminate upstream call") + await self._drain(reader, timeout=max(0.0, primary_deadline - asyncio.get_running_loop().time())) + if self._terminal_usage_required and not self._terminal and self._force_close_call is not None: + remaining: Final = max(0.0, deadline - asyncio.get_running_loop().time()) + try: + await asyncio.wait_for(self._force_close_call(), timeout=remaining) + self._close_confirmed = True + except Exception: # noqa: BLE001 # provider exceptions can contain credentials + verbose_proxy_logger.error("Realtime observer independent hangup failed") + await self._drain(reader, timeout=max(0.0, deadline - asyncio.get_running_loop().time())) + finally: + stopped.cancel() + reader.cancel() + if lease_failed is not None: + lease_failed.cancel() + await asyncio.gather( + *((reader, stopped, lease_failed) if lease_failed is not None else (reader, stopped)), + return_exceptions=True, + ) + with suppress(Exception): + await self._upstream.close() + if not self._usage_complete(): + self._logging.model_call_details["realtime_usage_incomplete"] = True + verbose_proxy_logger.error( + "Realtime observer ended without terminal usage; recorded usage is partial" + ) + try: + try: + await asyncio.wait_for( + self._stream.log_messages(wait_for_dispatch=True), timeout=self._logging_timeout + ) + self._accounting_complete = True + except asyncio.TimeoutError: + verbose_proxy_logger.error("Realtime observer timed out dispatching usage accounting") + finally: + if not self._accounting_complete: + self._logging.model_call_details["realtime_accounting_incomplete"] = True + if self._started and (not self._usage_complete() or not self._accounting_complete): + await invalidate_budget_reservation_counters( + budget_reservation=self._auth.budget_reservation + ) + elif not self._logging.model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY): + await release_or_invalidate_budget_reservation( + budget_reservation=self._auth.budget_reservation + ) + finally: + self._ready.set() + + async def _drain(self, reader: asyncio.Task[None], *, timeout: float | None = None) -> None: + try: + await asyncio.wait_for( + asyncio.shield(reader), + timeout=self._drain_timeout if timeout is None else min(self._drain_timeout, timeout), + ) + except asyncio.TimeoutError: + if not self._usage_complete(): + verbose_proxy_logger.error("Realtime observer timed out draining terminal usage") + except Exception: # noqa: BLE001 # cleanup must settle the socket even when reading or closing fails + verbose_proxy_logger.error("Realtime observer could not drain terminal usage") + return + + +class CallSupervisors: + def __init__(self) -> None: + self._tasks: tuple[asyncio.Task[None], ...] = () + self._calls: tuple[CallSupervisor, ...] = () + + async def start(self, supervisor: CallSupervisor) -> None: + self._calls = (*self._calls, supervisor) + try: + await supervisor.start() + except BaseException: + self._calls = tuple(call for call in self._calls if call is not supervisor) + raise + task: Final = asyncio.create_task(self._watch(supervisor)) + self._tasks = (*self._tasks, task) + + async def _watch(self, supervisor: CallSupervisor) -> None: + try: + try: + await supervisor.wait() + except Exception: # noqa: BLE001 # task must be consumed without exposing provider exception payloads + verbose_proxy_logger.error("Realtime observer accounting failed") + finally: + self._calls = tuple(call for call in self._calls if call is not supervisor) + self._tasks = tuple(task for task in self._tasks if task is not asyncio.current_task()) + + async def shutdown(self) -> None: + await asyncio.gather(*(call.close() for call in self._calls), return_exceptions=True) + await asyncio.gather(*self._tasks, return_exceptions=True) + + +CALL_SUPERVISORS: Final = CallSupervisors() diff --git a/litellm/proxy/realtime_endpoints/endpoints.py b/litellm/proxy/realtime_endpoints/endpoints.py index be2ac2ff33e..82041b4274f 100644 --- a/litellm/proxy/realtime_endpoints/endpoints.py +++ b/litellm/proxy/realtime_endpoints/endpoints.py @@ -16,7 +16,10 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_value_helper, ) -from litellm.proxy.common_utils.http_parsing_utils import _read_request_body +from litellm.proxy.common_utils.http_parsing_utils import ( + _normalize_media_type, # pyright: ignore[reportPrivateUsage] # reuse the shared HTTP media-type normalization contract + _read_request_body, +) from litellm.proxy.common_utils.openai_error_payload import ( error_status_code, openai_error_param, @@ -375,6 +378,11 @@ async def proxy_realtime_calls( request: Request, fastapi_response: Response, ) -> Response: + if _normalize_media_type(request.headers.get("content-type", "")) in ("application/json", "multipart/form-data"): + from litellm.proxy.realtime_endpoints.call_sessions import create_codex_realtime_call + + return await create_codex_realtime_call(request) + from litellm.proxy.proxy_server import ( add_litellm_data_to_request, general_settings, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 89625021e37..c24a2627562 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2061,6 +2061,8 @@ class ProxyLogging: data: None, call_type: CallTypesLiteral, guardrails_only: bool = False, + *, + internal_realtime_observer: bool = False, ) -> None: pass @@ -2071,6 +2073,8 @@ class ProxyLogging: data: dict, call_type: CallTypesLiteral, guardrails_only: bool = False, + *, + internal_realtime_observer: bool = False, ) -> dict: pass @@ -2080,6 +2084,8 @@ class ProxyLogging: data: dict | None, call_type: CallTypesLiteral, guardrails_only: bool = False, + *, + internal_realtime_observer: bool = False, ) -> dict | None: """ Allows users to modify/reject the incoming request to the proxy, without having to deal with parsing Request body. @@ -2178,6 +2184,10 @@ class ProxyLogging: deferred_route_exc: SensitiveDataRouteException | None = None for _callback in caps.resolved_callbacks: + if internal_realtime_observer and isinstance( + _callback, (_PROXY_MaxParallelRequestsHandler, _PROXY_MaxParallelRequestsHandler_v3) + ): + continue start_time = time.time() try: if isinstance(_callback, CustomGuardrail) and data is not None: diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index 44c47af57f4..7ed78dac899 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -76,6 +76,7 @@ def _get_realtime_http_provider_config( dynamic_api_base: str | None, dynamic_api_key: str | None, litellm_params: GenericLiteLLMParams, + is_call: bool = False, ) -> tuple["BaseRealtimeHTTPConfig | None", str, str]: """ Return (provider_config, resolved_api_base, resolved_api_key) for the @@ -93,13 +94,15 @@ def _get_realtime_http_provider_config( provider_config = ProviderConfigManager.get_provider_realtime_http_config( model="", provider=LlmProviders(custom_llm_provider), + params=litellm_params, + is_call=is_call, ) raw_api_base: Final = dynamic_api_base or litellm_params.api_base raw_api_key: Final = dynamic_api_key or litellm_params.api_key if provider_config is not None: - resolved_api_base = provider_config.get_api_base(api_base=raw_api_base) + resolved_api_base = provider_config.resolve_api_base(litellm_params.api_base, dynamic_api_base) resolved_api_key = provider_config.get_api_key(api_key=raw_api_key) else: # Fallback for providers without a dedicated HTTP config (treated as OpenAI-compatible). @@ -272,9 +275,15 @@ async def arealtime_calls( dynamic_api_base=dynamic_api_base, dynamic_api_key=dynamic_api_key, litellm_params=litellm_params, + is_call=True, ) if session is not None: session = _with_resolved_session_model(session, model_name) + call_headers: Final = ( + provider_config.get_realtime_calls_extra_headers(kwargs.get("extra_headers")) + if provider_config is not None + else kwargs.get("extra_headers") + ) litellm_logging_obj.update_from_kwargs( kwargs=kwargs, model=model_name, @@ -282,7 +291,7 @@ async def arealtime_calls( litellm_params={"api_base": resolved_api_base}, custom_llm_provider=custom_llm_provider, ) - return await base_llm_http_handler.async_realtime_calls_handler( + response: Final = await base_llm_http_handler.async_realtime_calls_handler( api_base=resolved_api_base, openai_ephemeral_key=openai_ephemeral_key, sdp_body=sdp_body, @@ -291,10 +300,17 @@ async def arealtime_calls( provider_config=provider_config, model=model_name, session_config=session, - extra_headers=kwargs.get("extra_headers"), + extra_headers=call_headers, client=kwargs.get("client"), api_version=litellm_params.api_version, ) + return ( + provider_config.transform_realtime_calls_response( + response, model_name, litellm_logging_obj.get_router_model_id(), call_headers + ) + if provider_config is not None + else response + ) async def vertex_access_token_resolver( @@ -391,7 +407,26 @@ async def _arealtime( model=model, provider=LlmProviders(_custom_llm_provider), ) - if provider_config is not None: + provider_handler: Final = ( + ProviderConfigManager.get_provider_realtime_handler( + LlmProviders(_custom_llm_provider), litellm_params, lambda: websocket.headers, headers + ) + if _custom_llm_provider in LlmProviders._member_map_.values() + else None + ) + if provider_handler is not None: + await provider_handler.async_realtime( + model=model, + websocket=websocket, + logging_obj=litellm_logging_obj, + api_base=api_base or None, + api_key=api_key, + timeout=timeout, + query_params=query_params, + user_api_key_dict=kwargs.get("user_api_key_dict"), + litellm_metadata=_build_litellm_metadata(kwargs), + ) + elif provider_config is not None: await base_llm_http_handler.async_realtime( model=model, websocket=websocket, diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index b7c4371f32f..83eb3c4aa2a 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -2034,6 +2034,11 @@ class OpenAIRealtimeStreamResponseBaseObject(TypedDict): type: str +class OpenAIRealtimeSessionClosed(TypedDict): + type: ReadOnly[Literal["session.closed"]] + usage: ReadOnly[Mapping[str, object]] + + class OpenAIRealtimeConversationObject(TypedDict, total=False): id: str object: Required[Literal["realtime.conversation"]] @@ -2239,6 +2244,7 @@ class OpenAIRealtimeEventTypes(Enum): OpenAIRealtimeEvents = ( OpenAIRealtimeStreamResponseBaseObject + | OpenAIRealtimeSessionClosed | OpenAIRealtimeStreamSessionEvents | OpenAIRealtimeStreamResponseOutputItemAdded | OpenAIRealtimeResponseContentPartAdded diff --git a/litellm/types/realtime.py b/litellm/types/realtime.py index 17dc70126f3..ae3a3c6ef2f 100644 --- a/litellm/types/realtime.py +++ b/litellm/types/realtime.py @@ -49,6 +49,7 @@ class RealtimeModalityResponseTransformOutput(TypedDict): class RealtimeQueryParams(TypedDict, total=False): model: str intent: str | None + call_id: ReadOnly[str] # Add more fields as needed diff --git a/litellm/utils.py b/litellm/utils.py index 1a77655a5a4..44899234a11 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -420,6 +420,7 @@ if TYPE_CHECKING: from litellm.llms.cohere.common_utils import CohereModelInfo from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.mistral.ocr.transformation import MistralOCRConfig + from litellm.llms.openai.realtime.handler import OpenAIRealtime from litellm.proxy._types import AllowedModelRegion from litellm.router_utils.get_retry_from_policy import ( get_num_retries_from_retry_policy, @@ -435,7 +436,7 @@ if TYPE_CHECKING: ChatCompletionToolCallFunctionChunk, ) from litellm.types.rerank import RerankResponse - from litellm.types.router import LiteLLM_Params + from litellm.types.router import GenericLiteLLMParams, LiteLLM_Params from litellm.llms.base_llm.chat.transformation import BaseConfig from litellm.llms.base_llm.completion.transformation import BaseTextCompletionConfig @@ -9119,6 +9120,10 @@ class ProviderConfigManager: model: str, provider: LlmProviders, ) -> BaseImageGenerationConfig | None: + if LlmProviders.CHATGPT == provider: + from litellm.llms.chatgpt.images import ChatGPTImageGenerationConfig + + return ChatGPTImageGenerationConfig() if LlmProviders.OPENAI == provider: from litellm.llms.openai.image_generation import ( get_openai_image_generation_config, @@ -9287,16 +9292,36 @@ class ProviderConfigManager: return GeminiRealtimeConfig() return None + @staticmethod + def get_provider_realtime_handler( + provider: LlmProviders, + params: GenericLiteLLMParams, + get_headers: Callable[[], Mapping[str, str]], + extra_headers: Mapping[str, object] | None = None, + ) -> OpenAIRealtime | None: + if provider == LlmProviders.CHATGPT: + from litellm.llms.chatgpt.realtime import ChatGPTRealtime + + return ChatGPTRealtime(params, get_headers(), extra_headers) + return None + @staticmethod def get_provider_realtime_http_config( model: str, provider: LlmProviders, + params: GenericLiteLLMParams | None = None, + is_call: bool = False, ) -> BaseRealtimeHTTPConfig | None: """ Return the HTTP transformation config for realtime HTTP endpoints (POST /realtime/client_secrets and POST /realtime/calls). """ + if LlmProviders.CHATGPT == provider: + from litellm.llms.chatgpt.realtime import ChatGPTRealtimeHTTPConfig + from litellm.types.router import GenericLiteLLMParams + + return ChatGPTRealtimeHTTPConfig(params or GenericLiteLLMParams(), use_codex_backend=is_call) if LlmProviders.OPENAI == provider: from litellm.llms.openai.realtime.http_transformation import ( OpenAIRealtimeHTTPConfig, @@ -9316,6 +9341,10 @@ class ProviderConfigManager: model: str, provider: LlmProviders, ) -> BaseImageEditConfig | None: + if LlmProviders.CHATGPT == provider: + from litellm.llms.chatgpt.images import ChatGPTImageEditConfig + + return ChatGPTImageEditConfig() if LlmProviders.OPENAI == provider: from litellm.llms.openai.image_edit import get_openai_image_edit_config diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 7fa09951eae..4b25dae8aeb 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -27932,6 +27932,16 @@ "max_tokens": 8191, "mode": "embedding" }, + "chatgpt/gpt-live-1-codex": { + "litellm_provider": "chatgpt", + "mode": "realtime", + "supported_endpoints": [ + "/v1/realtime/calls", + "/v1/live" + ], + "supports_audio_input": true, + "supports_audio_output": true + }, "chatgpt/gpt-5.5": { "litellm_provider": "chatgpt", "source": "https://platform.openai.com/docs/models/gpt-5.5", diff --git a/tests/local_testing/test_realtime_call_redis.py b/tests/local_testing/test_realtime_call_redis.py new file mode 100644 index 00000000000..d55f3488409 --- /dev/null +++ b/tests/local_testing/test_realtime_call_redis.py @@ -0,0 +1,102 @@ +import asyncio +import os +from contextlib import AsyncExitStack +from datetime import datetime +from uuid import uuid4 + +import pytest +import pytest_asyncio + +import litellm +from litellm.caching.redis_cache import RedisCache +from litellm.caching.caching import DualCache +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.hooks.parallel_request_limiter import _PROXY_MaxParallelRequestsHandler +from litellm.proxy.hooks.parallel_request_limiter_v3 import PARALLEL_REQUEST_SLOT_TTL_SECONDS +from litellm.proxy.utils import InternalUsageCache + + +@pytest_asyncio.fixture(loop_scope="function") +async def isolated_test_redis(monkeypatch): + raw_port = os.environ.get("LITELLM_TEST_REDIS_PORT", "") + if not raw_port.isdecimal() or not 1 <= int(raw_port) <= 65535: + pytest.fail("Set LITELLM_TEST_REDIS_PORT to an isolated Redis server's loopback port") + for name in tuple(os.environ): + if name.startswith("REDIS_"): + monkeypatch.delenv(name) + namespace = f"litellm-lua-test-{uuid4().hex}" + cache = RedisCache( + host="127.0.0.1", + port=int(raw_port), + namespace=namespace, + client_name=namespace, + socket_timeout=2, + socket_connect_timeout=2, + ) + async with AsyncExitStack() as cleanup: + cleanup.callback(cache.redis_client.close) + cleanup.push_async_callback(cache.async_redis_conn_pool.disconnect) + client = cache.init_async_client() + cleanup.push_async_callback(cache.async_redis_conn_pool.disconnect) + cleanup.push_async_callback(client.aclose) + cleanup.callback(litellm.in_memory_llm_clients_cache.delete_cache, cache._get_async_client_cache_key()) + try: + await client.ping() + yield cache + finally: + async for key in client.scan_iter(match=f"{namespace}:*"): + await client.delete(key) + + +@pytest.mark.asyncio +async def test_concurrent_realtime_releases_update_redis_without_lost_decrement(isolated_test_redis): + remote = isolated_test_redis + first_cache, second_cache = DualCache(redis_cache=remote), DualCache(redis_cache=remote) + first, second = (_PROXY_MaxParallelRequestsHandler(InternalUsageCache(c)) for c in (first_cache, second_cache)) + auth = UserAPIKeyAuth(api_key="concurrent-key", max_parallel_requests=2) + first_data, second_data = {"model": "test"}, {"model": "test"} + first.begin_realtime_attachment(first_data) + second.begin_realtime_attachment(second_data) + await first.async_pre_call_hook(auth, first_cache, first_data, "_arealtime") + await second.async_pre_call_hook(auth, second_cache, second_data, "_arealtime") + key = f"concurrent-key::{datetime.now().strftime('%Y-%m-%d-%H-%M')}::request_count" + counter = {"current_requests": 2, "current_rpm": 2, "current_tpm": 17} + await first_cache.async_set_cache(key, counter) + await second_cache.async_set_cache(key, counter, local_only=True) + remote.redis_client.pexpire(remote.check_and_fix_namespace(key), 15000) + await asyncio.gather( + first.async_release_realtime_attachment(first_data, auth), + second.async_release_realtime_attachment(second_data, auth), + ) + expected = {"current_requests": 0, "current_rpm": 2, "current_tpm": 17} + assert await remote.async_get_cache(key) == expected + assert 0 < remote.redis_client.pttl(remote.check_and_fix_namespace(key)) <= 15000 + assert await first_cache.async_get_cache(key) == expected + assert await second_cache.async_get_cache(key) == expected + await first_cache.async_set_cache("missing", counter, local_only=True) + await first._release_realtime_counter("missing") + assert await remote.async_get_cache("missing") is None + assert await first_cache.async_get_cache("missing", local_only=True) is None + + +@pytest.mark.asyncio +async def test_realtime_lease_redis_renewal_is_atomic_and_does_not_resurrect(isolated_test_redis): + from litellm.proxy.hooks.parallel_request_limiter_v3 import PARALLEL_RENEW_SCRIPT + + client = isolated_test_redis.init_async_client() + first_key = isolated_test_redis.check_and_fix_namespace("first") + second_key = isolated_test_redis.check_and_fix_namespace("second") + now = (await client.time())[0] + await client.zadd(first_key, {"owner": now - 10, "other": now}) + await client.zadd(second_key, {"owner": now - PARALLEL_REQUEST_SLOT_TTL_SECONDS}) + renew = client.register_script(PARALLEL_RENEW_SCRIPT) + assert await renew(keys=[first_key, second_key], args=["owner", PARALLEL_REQUEST_SLOT_TTL_SECONDS]) == [0] + assert await client.zscore(first_key, "owner") == now - 10 + await client.zadd(second_key, {"owner": now - 10}) + assert await renew(keys=[first_key, second_key], args=["owner", PARALLEL_REQUEST_SLOT_TTL_SECONDS]) == [1] + assert await client.zscore(first_key, "owner") >= now + assert await client.ttl(first_key) > PARALLEL_REQUEST_SLOT_TTL_SECONDS - 10 + await client.zrem(second_key, "owner") + assert await renew(keys=[first_key, second_key], args=["owner", PARALLEL_REQUEST_SLOT_TTL_SECONDS]) == [0] + assert await client.zscore(second_key, "owner") is None + assert await client.zscore(first_key, "other") == now diff --git a/tests/proxy_unit_tests/test_user_api_key_auth.py b/tests/proxy_unit_tests/test_user_api_key_auth.py index 0cdf3500d50..7b8d20543a5 100644 --- a/tests/proxy_unit_tests/test_user_api_key_auth.py +++ b/tests/proxy_unit_tests/test_user_api_key_auth.py @@ -849,6 +849,7 @@ async def test_user_api_key_auth_websocket(): # Prepare a mock WebSocket object mock_websocket = MagicMock(spec=WebSocket) mock_websocket.query_params = {"model": "some_model"} + mock_websocket.path_params = {} mock_websocket.headers = {"authorization": "Bearer some_api_key"} # Mock the scope attribute that user_api_key_auth_websocket accesses mock_websocket.scope = {"headers": [(b"authorization", b"Bearer some_api_key")]} @@ -872,6 +873,7 @@ async def test_user_api_key_auth_websocket(): assert request_arg.headers["authorization"] == "Bearer some_api_key" assert mock_user_api_key_auth.call_args.kwargs["api_key"] == "Bearer some_api_key" + assert await request_arg.json() == {"model": "some_model"} @pytest.mark.asyncio @@ -885,6 +887,7 @@ async def test_user_api_key_auth_websocket_carries_asgi_path(): mock_websocket = MagicMock(spec=WebSocket) mock_websocket.query_params = {"model": "some_model"} + mock_websocket.path_params = {} mock_websocket.headers = {"authorization": "Bearer some_api_key"} mock_websocket.scope = { "type": "websocket", diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py index 9c0f6f59463..5b263f9d557 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -3412,3 +3412,63 @@ async def test_refused_session_does_not_stamp_the_reservation_ownership_marker() assert session.logging.logged_failures == (upstream_close,) assert REALTIME_SESSION_SUCCESS_LOGGED_KEY not in session.logging.model_call_details + + +def test_live_terminal_usage_survives_filtered_event_logging(monkeypatch): + from litellm.cost_calculator import RealtimeAPITokenUsageProcessor + + def terminal(): + return {"type": "session.closed", "usage": {"audio_duration_ms": 4000, "backend_model_usage": []}} + + monkeypatch.setattr(litellm, "logged_real_time_event_types", []) + stream = RealTimeStreaming(MagicMock(), MagicMock(), MagicMock()) + event = {**terminal(), "private_transcript": "Do not retain this text"} + stream.store_message(event) + assert stream.messages == [terminal()] + usage = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results(stream.messages) + assert usage.total_tokens == 0 + + +@pytest.mark.asyncio +async def test_live_attachment_does_not_dispatch_duplicate_usage(): + worker = MagicMock() + logger = MagicMock() + stream = RealTimeStreaming(MagicMock(), MagicMock(), logger, logging_worker=worker, account_usage=False) + stream.store_message({"type": "session.closed", "usage": {"audio_duration_ms": 4000}}) + await stream.log_messages() + worker.ensure_initialized_and_enqueue.assert_not_called() + logger.dispatch_success_handlers.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("account_usage", [False, True]) +async def test_attachment_cleanup_runs_in_owning_context_only(account_usage): + from litellm.litellm_core_utils.realtime_streaming import realtime_attachment_cleanup + + contexts = [] + + async def one(name): + task = asyncio.current_task() + callback = AsyncMock(side_effect=lambda: contexts.append((name, asyncio.current_task() is task))) + token = realtime_attachment_cleanup.set(callback) + try: + websocket = MagicMock() + websocket.receive_text = AsyncMock(side_effect=RuntimeError("disconnected")) + backend = MagicMock() + + async def recv(**kwargs): + await asyncio.Event().wait() + + backend.recv = recv + stream = RealTimeStreaming(websocket, backend, MagicMock(), account_usage=account_usage) + await stream.bidirectional_forward() + if account_usage: + callback.assert_not_awaited() + else: + callback.assert_awaited_once() + finally: + realtime_attachment_cleanup.reset(token) + + await asyncio.gather(one("first"), one("second")) + assert sorted(contexts) == ([] if account_usage else [("first", True), ("second", True)]) + assert realtime_attachment_cleanup.get() is None diff --git a/tests/test_litellm/llms/chatgpt/conftest.py b/tests/test_litellm/llms/chatgpt/conftest.py new file mode 100644 index 00000000000..fb5b4cd5a1c --- /dev/null +++ b/tests/test_litellm/llms/chatgpt/conftest.py @@ -0,0 +1,22 @@ +import json +import time + +import pytest + + +@pytest.fixture +def chatgpt_tokens(tmp_path, monkeypatch): + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + monkeypatch.setenv("CHATGPT_AUTH_FILE", "auth.json") + for profile in ("default", "account2", "account3"): + name = "auth.json" if profile == "default" else profile + ".json" + (tmp_path / name).write_text( + json.dumps( + { + "access_token": "test-token-" + profile, + "account_id": "test-account-" + profile, + "expires_at": time.time() + 3600, + } + ) + ) + return str(tmp_path) diff --git a/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py b/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py index a7520bd5955..bd8076dde97 100644 --- a/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py +++ b/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py @@ -20,6 +20,29 @@ from litellm.utils import ProviderConfigManager class TestChatGPTResponsesAPITransformation: + def test_guardian_preserves_strict_output_schema(self): + text = { + "format": { + "type": "json_schema", + "name": "review", + "strict": True, + "schema": { + "type": "object", + "properties": {"allowed": {"type": "boolean"}}, + "required": ["allowed"], + "additionalProperties": False, + }, + } + } + request = ChatGPTResponsesAPIConfig().transform_responses_api_request( + model="codex-auto-review", + input="Review the command pwd", + response_api_optional_request_params={"text": text}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert request["text"] == text + @pytest.mark.parametrize( "model_name", [ diff --git a/tests/test_litellm/llms/chatgpt/test_codex.py b/tests/test_litellm/llms/chatgpt/test_codex.py new file mode 100644 index 00000000000..35606931e24 --- /dev/null +++ b/tests/test_litellm/llms/chatgpt/test_codex.py @@ -0,0 +1,57 @@ +import hashlib +import time + +import httpx +import pytest + +from litellm.llms.chatgpt.codex import CodexRealtimeCall, build_sideband_request, parse_call_response + + +def test_encrypted_call_preserves_repeated_gateway_query(monkeypatch): + from litellm.proxy.realtime_endpoints.call_sessions import decode_call, encode_call + + monkeypatch.setenv("LITELLM_SALT_KEY", "test-only-repeated-query") + authorization = "Bearer test-owner" + call = CodexRealtimeCall( + call_id="rtc_repeated", + model="gpt-live-1-codex", + alias="voice", + owner=hashlib.sha256(authorization.encode()).hexdigest(), + expires_at=time.time() + 60, + extra_query={"tag": ["alpha +/&", "beta"], "gateway": "tenant"}, + ) + restored = decode_call(encode_call(call), authorization) + assert restored.extra_query == {"tag": ("alpha +/&", "beta"), "gateway": "tenant"} + assert build_sideband_request(restored)["extra_query"] == restored.extra_query + + +@pytest.mark.parametrize("location", ["", "/v1/realtime/calls/foreign-id"]) +def test_signaling_rejects_invalid_upstream_call_id(location): + response = httpx.Response(201, headers={"Location": location}, + extensions={"chatgpt_realtime": {"model": "gpt-live-1-codex"}}) + with pytest.raises(ValueError, match="String should match pattern"): + parse_call_response(response, "voice", "owner", 1000) + + +@pytest.mark.parametrize("extra_query", [None, {"gateway_token": "opaque +/& value"}]) +def test_signaling_preserves_selected_model_for_sideband(extra_query): + response = httpx.Response( + 201, + headers={"Location": "/v1/realtime/calls/rtc_provider"}, + extensions={ + "chatgpt_realtime": { + "model": "gpt-live-1-codex", + "api_base": "https://voice.example/codex", + "extra_headers": {"x-gateway-route": "voice"}, + **({"extra_query": extra_query} if extra_query is not None else {}), + } + }, + ) + call = parse_call_response(response, "voice", "owner", 1000) + request = build_sideband_request(CodexRealtimeCall.model_validate_json(call.model_dump_json(exclude_none=True))) + assert request["api_base"] == "https://voice.example/codex" + assert request["model"] == "chatgpt/gpt-live-1-codex" + assert request["chatgpt_realtime_call_id"] == "rtc_provider" + assert request["query_params"] == {"model": "gpt-live-1-codex"} + assert request["extra_headers"] == {"x-gateway-route": "voice"} + assert request["extra_query"] == extra_query diff --git a/tests/test_litellm/llms/chatgpt/test_images.py b/tests/test_litellm/llms/chatgpt/test_images.py new file mode 100644 index 00000000000..2c51b932e3e --- /dev/null +++ b/tests/test_litellm/llms/chatgpt/test_images.py @@ -0,0 +1,149 @@ +import base64 + +import httpx +import pytest + +import litellm +from litellm.llms.chatgpt.images import ChatGPTImageEditConfig, ChatGPTImageGenerationConfig +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.types.router import GenericLiteLLMParams + + +@pytest.mark.parametrize("api_base", [None, "https://image-gateway.test"]) +def test_generation_routes_with_chatgpt_oauth(chatgpt_tokens, api_base): + requests = [] + + def respond(request): + requests.append(request) + return httpx.Response(200, json={"created": 1, "data": [{"b64_json": "aGVsbG8="}]}) + + client = HTTPHandler() + client.client = httpx.Client(transport=httpx.MockTransport(respond)) + result = litellm.image_generation( + model="chatgpt/gpt-image-2", + prompt="blue circle", + api_base=api_base, + client=client, + quality="auto", + size="auto", + background="auto", + extra_headers={"x-gateway-route": "images", "aUtHoRiZaTiOn": "Bearer wrong", "CHATGPT-ACCOUNT-ID": "wrong"}, + ) + assert requests[0].headers["x-gateway-route"] == "images" + assert result.data[0].b64_json == "aGVsbG8=" + assert str(requests[0].url) == (api_base or "https://chatgpt.com/backend-api/codex") + "/images/generations" + assert requests[0].headers["authorization"] == "Bearer test-token-" + "default" + assert requests[0].headers["chatgpt-account-id"] == "test-account-" + "default" + assert b'"model":"gpt-image-2"' in requests[0].content + + +def test_codex_json_edit_survives_sdk_dispatch(chatgpt_tokens): + requests = [] + + def respond(request): + requests.append(request) + return httpx.Response(200, json={"created": 1, "data": [{"b64_json": "aGVsbG8="}]}) + + client = HTTPHandler() + client.client = httpx.Client(transport=httpx.MockTransport(respond)) + references = [{"image_url": "data:image/png;base64,aGVsbG8="}] + result = litellm.image_edit( + model="chatgpt/gpt-image-2", + prompt="red circle", + extra_headers={"x-gateway-route": "images", "authorization": "Bearer wrong", "CHATGPT-ACCOUNT-ID": "wrong"}, + images=references, + client=client, + quality="auto", + size="auto", + ) + assert result.data[0].b64_json == "aGVsbG8=" + assert str(requests[0].url) == "https://chatgpt.com/backend-api/codex/images/edits" + import json + + assert json.loads(requests[0].content)["images"] == references + + assert requests[0].headers["authorization"] == "Bearer test-token-default" + assert requests[0].headers["chatgpt-account-id"] == "test-account-default" + assert requests[0].headers["x-gateway-route"] == "images" + + +@pytest.mark.parametrize( + "references", [[], [{"image_url": "file:///etc/passwd"}], [{}], [{"image_url": "https://example.com/a.png"}] * 6] +) +def test_edit_rejects_invalid_references(references): + with pytest.raises(ValueError, match=r"images must contain|validation error"): + ChatGPTImageEditConfig().transform_image_edit_request( + "gpt-image-2", "edit", None, {}, GenericLiteLLMParams(images=references), {} + ) + + +def test_edit_converts_multipart_image_bytes(): + data, files = ChatGPTImageEditConfig().transform_image_edit_request( + "gpt-image-2", "edit", b"example", {}, GenericLiteLLMParams(), {} + ) + assert not files + assert base64.b64decode(data["images"][0]["image_url"].split(",", 1)[1]) == b"example" + + +def test_image_auth_does_not_accept_inbound_override(chatgpt_tokens): + headers = ChatGPTImageGenerationConfig().validate_environment( + {"authorization": "Bearer wrong", "CHATGPT-ACCOUNT-ID": "wrong"}, "gpt-image-2", [], {}, {"chatgpt_token_dir": chatgpt_tokens} + ) + assert httpx.Headers(headers)["authorization"] == "Bearer test-token-default" + assert httpx.Headers(headers)["chatgpt-account-id"] == "test-account-default" + + +@pytest.mark.asyncio +async def test_async_codex_edit_without_multipart_image(chatgpt_tokens): + requests = [] + + def respond(request): + requests.append(request) + return httpx.Response(200, json={"created": 1, "data": [{"b64_json": "aGVsbG8="}]}) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(respond)) + response = await litellm.aimage_edit( + model="chatgpt/gpt-image-2", + prompt="red circle", + extra_headers={"x-gateway-route": "images", "authorization": "Bearer wrong", "CHATGPT-ACCOUNT-ID": "wrong"}, + client=client, + images=[{"image_url": "data:image/png;base64,aGVsbG8="}], + chatgpt_auth_profile="account3", + ) + assert response.data[0].b64_json == "aGVsbG8=" + assert str(requests[0].url).endswith("/codex/images/edits") + assert requests[0].headers["content-type"] == "application/json" + await client.client.aclose() + + assert requests[0].headers["authorization"] == "Bearer test-token-default" + assert requests[0].headers["chatgpt-account-id"] == "test-account-default" + assert requests[0].headers["x-gateway-route"] == "images" + + +@pytest.mark.parametrize("as_tuple", [False, True]) +def test_edit_accepts_filesystem_path(tmp_path, as_tuple): + image = tmp_path / "reference.png" + image.write_bytes(b"reference image bytes") + data, files = ChatGPTImageEditConfig().transform_image_edit_request( + "gpt-image-2", "edit", ("reference.png", image, "image/png") if as_tuple else image, + {}, GenericLiteLLMParams(), {} + ) + assert not files + assert data["images"] == ({"image_url": "data:image/png;base64," + base64.b64encode(image.read_bytes()).decode()},) + + +@pytest.mark.parametrize("env_name", ["CHATGPT_API_BASE", "OPENAI_CHATGPT_API_BASE"]) +@pytest.mark.parametrize("api_base", [None, "https://deployment.example/codex"]) +def test_image_routes_use_configured_gateway(monkeypatch, env_name, api_base, tmp_path): + token_path = tmp_path / "unavailable-token-directory" + token_path.write_text("not a directory") + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(token_path)) + monkeypatch.delenv("CHATGPT_API_BASE", raising=False) + monkeypatch.delenv("OPENAI_CHATGPT_API_BASE", raising=False) + monkeypatch.setenv(env_name, "https://gateway.example/codex/") + expected = api_base or "https://gateway.example/codex" + assert ChatGPTImageGenerationConfig().get_complete_url(api_base, None, "gpt-image-2", {}, {}) == ( + expected + "/images/generations" + ) + assert ChatGPTImageEditConfig().get_complete_url("gpt-image-2", api_base, {}) == expected + "/images/edits" diff --git a/tests/test_litellm/llms/chatgpt/test_realtime.py b/tests/test_litellm/llms/chatgpt/test_realtime.py new file mode 100644 index 00000000000..2e3d091c088 --- /dev/null +++ b/tests/test_litellm/llms/chatgpt/test_realtime.py @@ -0,0 +1,477 @@ +import json +import sys +from contextlib import nullcontext +from types import SimpleNamespace +from unittest.mock import AsyncMock, patch + +import httpx +import pytest + +import litellm +from litellm.llms.chatgpt.realtime import ChatGPTRealtime +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.types.router import GenericLiteLLMParams + +pytestmark = pytest.mark.usefixtures("local_model_cost_map") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", ["closed", "network"]) +@pytest.mark.parametrize("hangup_status", [200, 503]) +async def test_live_closed_observer_uses_independent_hangup(failure, hangup_status, chatgpt_tokens, monkeypatch): + from websockets.exceptions import ConnectionClosedOK + from websockets.frames import Close + + from litellm.caching.llm_caching_handler import LLMClientCache + + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) + handler = ChatGPTRealtime( + GenericLiteLLMParams( + chatgpt_realtime_call_id="rtc_live_closed", + chatgpt_token_dir=chatgpt_tokens, + extra_query={"gateway": "tenant", "tag": ["alpha +/&", "beta"]}, + ), + {}, + {"x-gateway-token": "test-only"}, + ) + connection = SimpleNamespace( + send=AsyncMock( + side_effect=( + ConnectionClosedOK(Close(1000, ""), Close(1000, ""), True) + if failure == "closed" + else OSError("socket unavailable") + ) + ) + ) + requests = [] + + def respond(request): + requests.append(request) + return httpx.Response(hangup_status) + + client = httpx.AsyncClient(transport=httpx.MockTransport(respond)) + try: + with patch("httpx.AsyncClient", return_value=client) as create_client: + for _ in range(2): + with pytest.raises(httpx.HTTPStatusError) if hangup_status == 503 else nullcontext(): + await handler.close_call(connection, "gpt-live-1-codex", "https://gateway.example/v1") + assert not client.is_closed + create_client.assert_called_once() + finally: + await client.aclose() + assert len(requests) == 2 + assert requests[0].method == "POST" + assert requests[0].url.path == "/v1/realtime/calls/rtc_live_closed/hangup" + assert requests[0].url.params.get_list("tag") == ["alpha +/&", "beta"] + assert requests[0].url.params["gateway"] == "tenant" + assert requests[0].headers["x-gateway-token"] == "test-only" + assert requests[0].headers["Authorization"] == "Bearer test-token-default" + assert requests[0].extensions["timeout"]["read"] == 10 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("endpoint", ["client_secrets", "transcription_sessions"]) +@pytest.mark.parametrize("source", ["default", "explicit", "CHATGPT_API_BASE", "OPENAI_CHATGPT_API_BASE"]) +async def test_realtime_session_urls_honor_gateway(endpoint, source, chatgpt_tokens, monkeypatch): + monkeypatch.setenv("CHATGPT_TOKEN_DIR", chatgpt_tokens) + monkeypatch.delenv("CHATGPT_API_BASE", raising=False) + monkeypatch.delenv("OPENAI_CHATGPT_API_BASE", raising=False) + gateway = "https://voice.example/custom/v1/" + if source in ("CHATGPT_API_BASE", "OPENAI_CHATGPT_API_BASE"): + monkeypatch.setenv(source, gateway) + requests = [] + + def respond(request): + requests.append(request) + return httpx.Response(200, json={"client_secret": {"value": "test-secret"}}) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(respond)) + kwargs = {"model": "chatgpt/gpt-realtime-1.5", "client": client} + if source == "explicit": + kwargs["api_base"] = gateway + try: + if endpoint == "client_secrets": + await litellm.acreate_realtime_client_secret(**kwargs) + else: + await litellm.acreate_realtime_transcription_session(**kwargs) + finally: + await client.client.aclose() + base = "https://api.openai.com/v1" if source == "default" else gateway.rstrip("/") + assert len(requests) == 1 + assert str(requests[0].url) == f"{base}/realtime/{endpoint}" + assert requests[0].headers["authorization"] == "Bearer test-token-default" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("inbound_headers", [{}, {"openai-alpha": "quicksilver=v2"}]) +@pytest.mark.parametrize("model, endpoint", [("gpt-live-1-codex", "live"), ("gpt-realtime-1.5", "realtime")]) +async def test_routed_call_preserves_deployment_gateway_headers( + inbound_headers, model, endpoint, chatgpt_tokens, monkeypatch +): + from litellm.llms.chatgpt.codex import ( + CodexRealtimeCall, + CodexRealtimeOffer, + build_call_request, + build_sideband_request, + parse_call_response, + ) + + monkeypatch.setenv("CHATGPT_TOKEN_DIR", chatgpt_tokens) + requests = [] + + def respond(request): + requests.append(request) + return httpx.Response(201, text="v=0\r\n", headers={"location": "/v1/realtime/calls/rtc_test"}) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(respond)) + router = litellm.Router( + model_list=[ + { + "model_name": "voice-gateway", + "litellm_params": { + "model": f"chatgpt/{model}", + "api_base": "https://voice.example/backend-api/codex", + "extra_headers": {"x-gateway-route": "configured"}, + "extra_query": { + "gateway_token": "configured", + "intent": "pinned-intent", + "count": 7, + "fraction": 1.5, + "enabled": True, + "disabled": False, + "blank": None, + "tag": ["alpha +/&", "beta"], + "empty": [], + "model": "other-model", + "call_id": "rtc_wrong", + }, + }, + "model_info": {"id": "selected-gateway-deployment"}, + } + ], + num_retries=0, + ) + offer = CodexRealtimeOffer(sdp="v=0\r\n", session={"model": "voice-gateway"}) + try: + response = await router.arealtime_calls( + **build_call_request(offer, {"intent": "quicksilver", "architecture": "avas"}, inbound_headers), + client=client, + ) + assert requests[0].headers.get("x-gateway-route") == "configured" + assert dict(requests[0].url.params) == { + "gateway_token": "configured", + "intent": "pinned-intent", + "architecture": "avas", + "count": "7", + "fraction": "1.5", + "enabled": "true", + "disabled": "false", + "blank": "", + "tag": "alpha +/&", + "model": "other-model", + "call_id": "rtc_wrong", + } + assert requests[0].url.params.get_list("tag") == ["alpha +/&", "beta"] + assert response.extensions["chatgpt_realtime"]["extra_query"] == { + **dict(requests[0].url.params), + "tag": ("alpha +/&", "beta"), + "empty": (), + } + assert response.extensions["chatgpt_realtime"]["extra_headers"]["x-gateway-route"] == "configured" + for name, value in inbound_headers.items(): + assert requests[0].headers[name] == value + call = parse_call_response(response, alias="voice-gateway", owner="test-owner", expires_at=1) + restored = CodexRealtimeCall.model_validate_json(call.model_dump_json()) + assert restored.model_id == "selected-gateway-deployment" + assert restored.model == model + handler = ChatGPTRealtime(GenericLiteLLMParams.model_validate(build_sideband_request(restored)), {}) + sideband_url = httpx.URL(handler._construct_url(restored.api_base, {"model": restored.model})) + assert {key: value for key, value in sideband_url.params.items() if key != "call_id"} == { + key: value for key, value in requests[0].url.params.items() if key not in ("model", "call_id") + } + assert sideband_url.params.get("call_id") == ("rtc_test" if endpoint == "realtime" else None) + assert sideband_url.params.get_list("tag") == ["alpha +/&", "beta"] + assert sideband_url.path.endswith("/realtime" if endpoint == "realtime" else "/live/rtc_test") + finally: + await client.client.aclose() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", ["gpt-realtime-1.5", "gpt-live-1-codex"]) +@pytest.mark.parametrize("call_id", [None, "rtc_existing"]) +async def test_websocket_forwards_configured_headers_without_client_identity(model, call_id, chatgpt_tokens): + websocket = SimpleNamespace( + headers={"authorization": "Bearer client", "cookie": "private-cookie", "openai-alpha": "client-value"}, + scope={}, + receive_text=AsyncMock(side_effect=RuntimeError("client disconnected")), + send_text=AsyncMock(), + close=AsyncMock(), + ) + with patch("websockets.connect") as connect: + connect.return_value.__aenter__ = AsyncMock(side_effect=RuntimeError("stop before streaming")) + await litellm._arealtime( + model=f"chatgpt/{model}", + websocket=websocket, + api_base="https://voice.example/codex", + chatgpt_realtime_call_id=call_id, + query_params={"model": model, "intent": "client-intent"}, + extra_query={"intent": "configured-intent", "tag": ["alpha +/&", "beta"]}, + headers={"x-deployment-header": "configured"}, + extra_headers={ + "X-Gateway-Route": "voice", + "OpenAI-Alpha": "configured-value", + "aUtHoRiZaTiOn": "Bearer wrong", + "CHATGPT-ACCOUNT-ID": "wrong", + }, + ) + connect.assert_called_once() + headers = httpx.Headers(connect.call_args.kwargs["additional_headers"]) + upstream_url = httpx.URL(connect.call_args.args[0]) + assert upstream_url.params.get_list("intent") == ["configured-intent"] + assert upstream_url.params.get_list("tag") == ["alpha +/&", "beta"] + assert headers["x-deployment-header"] == "configured" + assert headers["x-gateway-route"] == "voice" + assert headers["openai-alpha"] == "configured-value" + assert headers["authorization"] == "Bearer test-token-default" + assert headers["chatgpt-account-id"] == "test-account-default" + assert "cookie" not in headers + + +@pytest.mark.asyncio +@pytest.mark.parametrize("api_base", [None, "https://voice.example/backend-api/codex"]) +async def test_chatgpt_call_keeps_oauth_and_frameless_session(chatgpt_tokens, api_base): + requests = [] + + def respond(request): + requests.append(request) + return httpx.Response(201, text="v=0\r\n", headers={"location": "/v1/realtime/calls/rtc_test"}) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(respond)) + response = await litellm.arealtime_calls( + model="chatgpt/gpt-live-1-codex", + api_base=api_base, + openai_ephemeral_key="", + sdp_body=b"v=0\r\n", + session={"model": "chatgpt/gpt-live-1-codex", "audio": {"output": {"voice": "sol"}}}, + extra_query={"intent": "quicksilver", "architecture": "avas"}, + chatgpt_realtime_client_query={"intent": "untrusted-override", "architecture": "avas", "untrusted": "bad"}, + extra_headers={ + "openai-alpha": "quicksilver=v2", + "x-gateway-route": "voice", + "aUtHoRiZaTiOn": "Bearer wrong", + "CHATGPT-ACCOUNT-ID": "wrong", + }, + client=client, + ) + assert response.extensions["chatgpt_realtime"]["api_base"] == (api_base or "https://api.openai.com/v1") + assert response.extensions["chatgpt_realtime"]["extra_headers"] == { + "openai-alpha": "quicksilver=v2", + "x-gateway-route": "voice", + } + assert requests[0].url.host == ("voice.example" if api_base else "chatgpt.com") + assert response.status_code == 201 + assert response.extensions["chatgpt_realtime"]["extra_query"] == {"intent": "quicksilver", "architecture": "avas"} + assert requests[0].url.path == "/backend-api/codex/realtime/calls" + assert requests[0].url.params["architecture"] == "avas" + assert requests[0].headers["authorization"] == "Bearer test-token-" + "default" + assert requests[0].headers["chatgpt-account-id"] == "test-account-default" + assert requests[0].headers["openai-alpha"] == "quicksilver=v2" + assert requests[0].headers["x-gateway-route"] == "voice" + assert json.loads(requests[0].content) == { + "sdp": "v=0\r\n", + "session": {"model": "gpt-live-1-codex", "audio": {"output": {"voice": "sol"}}}, + } + await client.client.aclose() + + +@pytest.mark.asyncio +async def test_openai_call_preserves_explicit_identity_headers(): + requests = [] + + def respond(request): + requests.append(request) + return httpx.Response(201, text="v=0\r\n") + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(respond)) + try: + response = await litellm.arealtime_calls( + model="openai/gpt-realtime-1.5", + openai_ephemeral_key="original-key", + sdp_body=b"v=0\r\n", + extra_headers={"Authorization": "Bearer explicit-key", "chatgpt-account-id": "custom-account"}, + client=client, + ) + assert response.status_code == 201 + assert requests[0].headers["authorization"] == "Bearer explicit-key" + assert requests[0].headers["chatgpt-account-id"] == "custom-account" + assert requests[0].headers["content-type"].startswith("multipart/form-data") + finally: + await client.client.aclose() + + +@pytest.mark.parametrize("model,endpoint", [("gpt-realtime-1.5", "realtime"), ("gpt-live-1-codex", "live")]) +def test_realtime_uses_platform_endpoint_with_oauth_headers(model, endpoint, chatgpt_tokens, local_model_cost_map): + handler = ChatGPTRealtime( + GenericLiteLLMParams(), + { + "authorization": "Bearer proxy-key", + "openai-alpha": "quicksilver=v2", + }, + ) + assert handler._construct_url("https://api.openai.com/v1", {"model": model}) == ( + f"wss://api.openai.com/v1/{endpoint}?model={model}" + ) + headers = handler._get_additional_headers("unused") + assert headers["Authorization"] == "Bearer test-token-default" + assert "authorization" not in headers + assert headers["openai-alpha"] == "quicksilver=v2" + + +@pytest.mark.parametrize("endpoint", ["live", "realtime"]) +def test_new_realtime_session_preserves_gateway_query(endpoint, chatgpt_tokens, local_model_cost_map): + model = "gpt-live-1-codex" if endpoint == "live" else "gpt-realtime-1.5" + handler = ChatGPTRealtime( + GenericLiteLLMParams( + chatgpt_token_dir=chatgpt_tokens, + chatgpt_realtime_client_query={"intent": "conversation", "architecture": "client-architecture"}, + extra_query={ + "gateway_token": "opaque +/& value", + "intent": "gateway-intent", + "architecture": "gateway-architecture", + "model": "other-model", + "call_id": "rtc_other", + }, + ), + {}, + ) + url = httpx.URL(handler._construct_url("https://gateway.example/v1", {"model": model, "intent": "query-intent"})) + assert url.path == f"/v1/{endpoint}" + assert dict(url.params) == { + "model": model, + "gateway_token": "opaque +/& value", + "intent": "gateway-intent", + "architecture": "gateway-architecture", + } + + +@pytest.mark.asyncio +async def test_openai_http_call_does_not_require_websockets(monkeypatch): + monkeypatch.delitem(sys.modules, "litellm.llms.chatgpt.realtime", raising=False) + for name in tuple(sys.modules): + if name == "websockets" or name.startswith("websockets."): + monkeypatch.delitem(sys.modules, name) + monkeypatch.setitem(sys.modules, "websockets", None) + requests = [] + + def respond(request): + requests.append(request) + return httpx.Response(201, text="v=0\r\n") + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(respond)) + try: + response = await litellm.arealtime_calls( + model="openai/gpt-realtime-1.5", + openai_ephemeral_key="test-only", + sdp_body=b"v=0\r\n", + api_key="test-only", + client=client, + ) + assert response.status_code == 201 + assert len(requests) == 1 + assert requests[0].url.path == "/v1/realtime/calls" + finally: + await client.client.aclose() + + +@pytest.mark.parametrize("endpoint", ["live", "realtime"]) +@pytest.mark.parametrize("call_id", [None, "rtc_metadata"]) +def test_realtime_routes_new_models_using_registered_metadata(endpoint, call_id, chatgpt_tokens, local_model_cost_map): + model = "metadata-voice-model" + litellm.register_model({f"chatgpt/{model}": { + "litellm_provider": "chatgpt", "mode": "realtime", "supported_endpoints": [f"/v1/{endpoint}"] + }}) + handler = ChatGPTRealtime(GenericLiteLLMParams(chatgpt_realtime_call_id=call_id), {}) + expected = ( + f"wss://api.openai.com/v1/{endpoint}?model={model}" if call_id is None + else f"wss://api.openai.com/v1/live/{call_id}" if endpoint == "live" + else f"wss://api.openai.com/v1/realtime?call_id={call_id}" + ) + assert handler._construct_url("https://api.openai.com/v1", {"model": model}) == expected + + +def test_realtime_unknown_model_keeps_standard_endpoint(chatgpt_tokens, local_model_cost_map): + handler = ChatGPTRealtime(GenericLiteLLMParams(), {}) + assert handler._construct_url("https://api.openai.com/v1", {"model": "unknown-voice-model"}) == ( + "wss://api.openai.com/v1/realtime?model=unknown-voice-model" + ) + + +@pytest.mark.parametrize("env_name", ["CHATGPT_API_BASE", "OPENAI_CHATGPT_API_BASE"]) +@pytest.mark.parametrize("api_base", [None, "https://deployment.example/codex"]) +def test_realtime_routes_use_configured_gateway(monkeypatch, env_name, api_base, chatgpt_tokens): + from litellm.llms.chatgpt.realtime import ChatGPTRealtimeHTTPConfig + + monkeypatch.delenv("CHATGPT_API_BASE", raising=False) + monkeypatch.delenv("OPENAI_CHATGPT_API_BASE", raising=False) + monkeypatch.setenv(env_name, "https://gateway.example/codex/") + expected = api_base or "https://gateway.example/codex" + config = ChatGPTRealtimeHTTPConfig(GenericLiteLLMParams()) + assert config.get_realtime_calls_url(api_base, "gpt-live-1-codex") == expected + "/realtime/calls" + handler = ChatGPTRealtime(GenericLiteLLMParams(), {}) + assert handler._construct_url(handler.get_api_base(api_base), {"model": "gpt-realtime-1.5"}) == ( + expected.replace("https://", "wss://") + "/realtime?model=gpt-realtime-1.5" + ) + + +@pytest.mark.parametrize("model,endpoint", [("gpt-live-1-codex", "live"), ("gpt-realtime-1.5", "realtime")]) +def test_sideband_restores_gateway_query_without_overriding_call(model, endpoint, chatgpt_tokens): + handler = ChatGPTRealtime( + GenericLiteLLMParams( + chatgpt_realtime_call_id="rtc_selected", + extra_query={"gateway_token": "opaque +/& value", "model": "other", "call_id": "rtc_other"}, + ), + {}, + ) + url = httpx.URL(handler._construct_url("https://gateway.example/v1", {"model": model})) + assert url.params["gateway_token"] == "opaque +/& value" + assert "model" not in url.params + if endpoint == "live": + assert url.path == "/v1/live/rtc_selected" + assert "call_id" not in url.params + else: + assert url.path == "/v1/realtime" + assert url.params["call_id"] == "rtc_selected" + + +def test_client_cannot_forge_supervised_call_accounting(chatgpt_tokens): + from litellm.llms.chatgpt.realtime import CallAccounting, accounts_for_call_usage + + assert accounts_for_call_usage(GenericLiteLLMParams(chatgpt_call_accounting={"supervised": True})) + assert accounts_for_call_usage(GenericLiteLLMParams(chatgpt_call_accounting="supervised")) + assert not accounts_for_call_usage(GenericLiteLLMParams(chatgpt_call_accounting=CallAccounting.SUPERVISED)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", ["gpt-live-1-codex", "gpt-realtime-1.5"]) +async def test_supervisor_connection_preserves_call_routing(model, chatgpt_tokens): + handler = ChatGPTRealtime( + GenericLiteLLMParams( + chatgpt_token_dir=chatgpt_tokens, + chatgpt_realtime_call_id="rtc_owner", + extra_query={"gateway_token": "a+b&c"}, + ), + {"openai-alpha": "quicksilver=v2"}, + {"x-gateway-token": "configured"}, + ) + connection = AsyncMock() + with patch("websockets.connect", AsyncMock(return_value=connection)) as connect: + assert await handler.open_call_connection(model, "https://gateway.example/v1") is connection + url = httpx.URL(connect.call_args.args[0]) + assert url.params["gateway_token"] == "a+b&c" + assert connect.call_args.kwargs["additional_headers"]["x-gateway-token"] == "configured" + assert url.path.endswith("/rtc_owner") if model == "gpt-live-1-codex" else url.params["call_id"] == "rtc_owner" diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index c39779972c0..bed15399fbb 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -3713,3 +3713,62 @@ def test_image_edit_handler_keeps_the_sync_transform(): assert config.transform_calls == ["sync"] assert captured["body"] == {"transformed_by": "sync"} assert response.data[0].b64_json == "sync" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("endpoint", ["client_secrets", "transcription_sessions"]) +@pytest.mark.parametrize("provider", ["chatgpt", "openai"]) +@pytest.mark.parametrize("authorization_header", ["Authorization", "aUtHoRiZaTiOn"]) +async def test_realtime_http_sessions_preserve_provider_identity( + endpoint, provider, authorization_header, tmp_path, monkeypatch +): + import time + + from litellm.llms.chatgpt.realtime import ChatGPTRealtimeHTTPConfig + from litellm.llms.openai.realtime.http_transformation import OpenAIRealtimeHTTPConfig + from litellm.types.router import GenericLiteLLMParams + + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + monkeypatch.setenv("CHATGPT_AUTH_FILE", "auth.json") + (tmp_path / "auth.json").write_text( + json.dumps({"access_token": "test-resolved", "account_id": "test-selected", "expires_at": time.time() + 3600}) + ) + config = ChatGPTRealtimeHTTPConfig(GenericLiteLLMParams()) if provider == "chatgpt" else OpenAIRealtimeHTTPConfig() + requests = [] + + def respond(request): + requests.append(request) + return httpx.Response(200, json={"id": "session-test"}) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(respond)) + try: + response = await BaseLLMHTTPHandler()._async_realtime_session_post( + endpoint=endpoint, + api_base="https://gateway.example/v1", + api_key="test-openai", + request_data={"session": {"model": "gpt-realtime-1.5"}}, + logging_obj=Mock(), + timeout=5, + provider_config=config, + model="gpt-realtime-1.5", + extra_headers={ + authorization_header: "Bearer test-override", + "CHATGPT-ACCOUNT-ID": "test-other-account", + "x-gateway-route": "required", + }, + client=client, + ) + assert response.status_code == 200 + assert not client.client.is_closed + finally: + await client.client.aclose() + assert len(requests) == 1 + assert requests[0].url.path == f"/v1/realtime/{endpoint}" + assert requests[0].headers["x-gateway-route"] == "required" + if provider == "chatgpt": + assert requests[0].headers.get_list("authorization") == ["Bearer test-resolved"] + assert requests[0].headers.get_list("chatgpt-account-id") == ["test-selected"] + else: + assert requests[0].headers.get_list("authorization")[-1] == "Bearer test-override" + assert requests[0].headers["chatgpt-account-id"] == "test-other-account" diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index cdf1f897707..8f41fc22d71 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -947,7 +947,8 @@ def test_get_model_from_request_handles_managed_id_decoder_failures(): "/openai/v1/realtime/calls", ], ) -def test_get_model_from_request_extracts_realtime_session_model(route): +@pytest.mark.parametrize("encoded", [False, True]) +def test_get_model_from_request_extracts_realtime_session_model(route, encoded): """The effective realtime model lives in ``session.model`` (not the top-level ``model``). It must be surfaced so can_key_call_model() can validate the model a restricted key is actually requesting. @@ -957,13 +958,36 @@ def test_get_model_from_request_extracts_realtime_session_model(route): """ assert ( get_model_from_request( - request_data={"session": {"type": "realtime", "model": "gpt-realtime"}}, + request_data={"session": '{"model":"gpt-realtime"}' if encoded else {"model": "gpt-realtime"}}, route=route, ) == "gpt-realtime" ) +@pytest.mark.parametrize("session", ['{"model":"actual-voice"}', {"model": "actual-voice"}]) +def test_realtime_calls_auth_uses_executed_session_model_despite_decoys(session): + assert ( + get_model_from_request( + request_data={"model": "body-decoy", "session": session}, + route="/v1/realtime/calls", + request_query_params={"model": "query-decoy"}, + request_headers={"x-litellm-model": "header-decoy"}, + ) + == "actual-voice" + ) + + +@pytest.mark.parametrize("model", ["voice,alias", " voice "]) +def test_realtime_calls_auth_preserves_exact_session_model(model): + assert get_model_from_request(request_data={"session": {"model": model}}, route="/v1/realtime/calls") == model + + +@pytest.mark.parametrize("session", ["invalid", "null", "[]", "12", '"text"', "{}"]) +def test_realtime_model_extraction_ignores_invalid_serialized_session(session): + assert get_model_from_request(request_data={"session": session}, route="/v1/realtime/calls") is None + + def test_get_model_from_request_realtime_includes_top_level_and_session_model(): """When both top-level and session model are present, both are returned so neither path can smuggle a disallowed model past the model-access check.""" diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index c8b3d789665..c5a34daa50a 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -17,6 +17,24 @@ from litellm.proxy.auth.auth_checks_organization import _user_is_org_admin from litellm.proxy.auth.route_checks import RouteChecks +@pytest.mark.parametrize("route", ["/live", "/v1/live", "/v1/live/rtc_litellm_test"]) +def test_codex_live_routes_allow_inference_keys(route: str): + from litellm.proxy.auth.auth_checks import _allowed_routes_check + + assert RouteChecks.is_llm_api_route(route) + assert _allowed_routes_check(user_route=route, allowed_routes=["openai_routes"]) + token = UserAPIKeyAuth(allowed_routes=["llm_api_routes"]) + RouteChecks.is_virtual_key_allowed_to_call_route(route=route, valid_token=token) + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=None, + _user_role=LitellmUserRoles.INTERNAL_USER.value, + route=route, + request=Request({"type": "http", "path": route, "query_string": b"", "headers": []}), + valid_token=token, + request_data={}, + ) + + def test_non_admin_config_update_route_rejected(): """Test that non-admin users are rejected when trying to call /config/update""" diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 0cdbcde6abc..8cfb8d0905d 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -7285,6 +7285,52 @@ def test_user_api_key_auth_opens_a_datadog_span_for_accepted_and_rejected_keys(t assert [span for span in report["spans"] if span == auth_span] == [auth_span, auth_span] +@pytest.mark.asyncio +@pytest.mark.parametrize("attachment", ["path", "query"]) +@pytest.mark.parametrize("credential", ["authorization", "api-key", "subprotocol", "x-litellm-api-key", "custom", "custom-mixed"]) +@pytest.mark.parametrize("query_model", [b"", b"model=unbudgeted"]) +async def test_sideband_auth_uses_encrypted_model_for_budget_checks(monkeypatch, attachment, credential, query_model): + import hashlib + import importlib + import time + from unittest.mock import AsyncMock + from fastapi import WebSocket + from litellm.llms.chatgpt.codex import CodexRealtimeCall + from litellm.proxy.realtime_endpoints.call_sessions import encode_call + + auth_module = importlib.import_module("litellm.proxy.auth.user_api_key_auth") + monkeypatch.setenv("LITELLM_SALT_KEY", "test-only-sideband-budget-salt") + token = encode_call(CodexRealtimeCall( + call_id="rtc_test", model="gpt-live-1-codex", alias="budgeted-voice", + owner=hashlib.sha256(b"Bearer owner").hexdigest(), expires_at=time.time() + 300, + )) + from litellm.proxy import proxy_server + monkeypatch.setattr(proxy_server, "general_settings", {"litellm_key_header_name": "x-proxy-key"} if credential.startswith("custom") else {}) + seen = [] + + async def authenticate(request, api_key): + seen.append((await request.json(), api_key)) + return "authenticated-with-model" + + monkeypatch.setattr(auth_module, "user_api_key_auth", authenticate) + websocket = WebSocket({ + "type": "websocket", "scheme": "ws", "server": ("localhost", 4000), + "path": "/v1/live/" + token if attachment == "path" else "/v1/realtime", + "path_params": {"call_id": token} if attachment == "path" else {}, + "query_string": query_model + (b"&call_id=" + token.encode() if attachment == "query" else b""), + "headers": { + "authorization": [(b"authorization", b"Bearer owner")], + "api-key": [(b"api-key", b"owner")], + "x-litellm-api-key": [(b"x-litellm-api-key", b"owner")], + "custom": [(b"x-proxy-key", b"Bearer owner")], + "custom-mixed": [(b"x-proxy-key", b"Bearer owner"), (b"authorization", b"Bearer other-owner")], + "subprotocol": [(b"sec-websocket-protocol", b"realtime, openai-insecure-api-key.owner")], + }[credential], + }, AsyncMock(), AsyncMock()) + assert await auth_module.user_api_key_auth_websocket(websocket) == "authenticated-with-model" + assert seen == [({"model": "budgeted-voice"}, "Bearer owner")] + + @pytest.mark.asyncio @pytest.mark.parametrize("is_proxy_admin", [False, True], ids=["standard-return", "proxy-admin-return"]) async def test_jwt_builder_returns_every_team_grant_the_key_path_gets(is_proxy_admin): @@ -7401,6 +7447,87 @@ async def test_jwt_builder_returns_every_team_grant_the_key_path_gets(is_proxy_a assert token.jwt_claims == {"sub": "jwt-user"} +@pytest.mark.asyncio +@pytest.mark.parametrize("attachment", ["path", "query"]) +async def test_sideband_rejects_budget_fallback_before_rerouting(monkeypatch, attachment): + import hashlib + import importlib + import time + from types import SimpleNamespace + from unittest.mock import AsyncMock + from fastapi import HTTPException, WebSocket + from litellm.llms.chatgpt.codex import CodexRealtimeCall + from litellm.proxy.realtime_endpoints.call_sessions import encode_call + + auth_module = importlib.import_module("litellm.proxy.auth.user_api_key_auth") + monkeypatch.setenv("LITELLM_SALT_KEY", "test-only-sideband-budget-salt") + token = encode_call(CodexRealtimeCall( + call_id="rtc_test", model="gpt-live-1-codex", alias="budgeted-voice", + owner=hashlib.sha256(b"Bearer owner").hexdigest(), expires_at=time.time() + 300, + )) + limiter = SimpleNamespace( + is_key_within_model_budget=AsyncMock(side_effect=litellm.BudgetExceededError(current_cost=2, max_budget=1)), + get_fallback_model_within_budget=AsyncMock(return_value="cheap-voice"), + ) + auth = UserAPIKeyAuth(models=["budgeted-voice", "cheap-voice"]) + + async def authenticate(request, api_key): + data = await request.json() + await auth_module._check_key_model_budget_with_fallback(auth, limiter, data["model"], data, request) + return auth + + monkeypatch.setattr(auth_module, "user_api_key_auth", authenticate) + monkeypatch.setattr(auth_module, "can_key_call_model", AsyncMock()) + send = AsyncMock() + websocket = WebSocket({ + "type": "websocket", "scheme": "ws", "server": ("localhost", 4000), + "path": "/v1/live/" + token if attachment == "path" else "/v1/realtime", + "path_params": {"call_id": token} if attachment == "path" else {}, + "query_string": b"call_id=" + token.encode() if attachment == "query" else b"", + "headers": [(b"authorization", b"Bearer owner")], + }, AsyncMock(), send) + with pytest.raises(HTTPException) as error: + await auth_module.user_api_key_auth_websocket(websocket) + assert error.value.status_code == 403 + limiter.get_fallback_model_within_budget.assert_not_awaited() + send.assert_awaited_once_with({"type": "websocket.close", "code": 1008, "reason": ""}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("custom_value", [None, b"Bearer different-owner"]) +async def test_sideband_custom_header_cannot_fall_back_to_other_credentials(monkeypatch, custom_value): + import hashlib + import importlib + import time + from unittest.mock import AsyncMock + from fastapi import HTTPException, WebSocket + from litellm.proxy import proxy_server + from litellm.llms.chatgpt.codex import CodexRealtimeCall + from litellm.proxy.realtime_endpoints.call_sessions import encode_call + + auth_module = importlib.import_module("litellm.proxy.auth.user_api_key_auth") + monkeypatch.setenv("LITELLM_SALT_KEY", "test-only-custom-header-salt") + monkeypatch.setattr(proxy_server, "general_settings", {"litellm_key_header_name": "x-proxy-key"}) + token = encode_call(CodexRealtimeCall( + call_id="rtc_test", model="gpt-live-1-codex", alias="voice", + owner=hashlib.sha256(b"Bearer owner").hexdigest(), expires_at=time.time() + 300, + )) + authenticate = AsyncMock() + monkeypatch.setattr(auth_module, "user_api_key_auth", authenticate) + send = AsyncMock() + websocket = WebSocket({ + "type": "websocket", "scheme": "ws", "server": ("localhost", 4000), + "path": "/v1/live/" + token, "path_params": {"call_id": token}, "query_string": b"", + "headers": [(b"authorization", b"Bearer owner")] + + ([(b"x-proxy-key", custom_value)] if custom_value is not None else []), + }, AsyncMock(), send) + with pytest.raises(HTTPException) as error: + await auth_module.user_api_key_auth_websocket(websocket) + assert error.value.status_code == 403 + authenticate.assert_not_awaited() + send.assert_awaited_once_with({"type": "websocket.close", "code": 1008, "reason": ""}) + + @pytest.mark.asyncio @pytest.mark.parametrize("route", ["/v1/messages", "/messages", "/v1/chat/completions", "/chat/completions", "/v1/responses", "/responses"]) async def test_claude_view_normalizes_before_model_access(monkeypatch, route): diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter.py index 0e2683dcbfd..589c8d1bd06 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter.py @@ -2,11 +2,16 @@ Unit Tests for the max parallel request limiter v1 for the proxy """ +import asyncio from datetime import datetime +from unittest.mock import AsyncMock, MagicMock, patch import pytest from litellm.caching.caching import DualCache +from litellm.caching.redis_cache import RedisCache +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.hooks.parallel_request_limiter import ( _PROXY_MaxParallelRequestsHandler, ) @@ -14,6 +19,100 @@ from litellm.proxy.utils import InternalUsageCache, hash_token from litellm.types.utils import EmbeddingResponse, TextCompletionResponse, Usage +@pytest.mark.asyncio +async def test_realtime_release_preserves_newer_local_admission_while_redis_finishes(): + started, finish = asyncio.Event(), asyncio.Event() + + async def release(**kwargs): + started.set() + await finish.wait() + + remote = MagicMock(spec=RedisCache) + remote.async_register_script.return_value = AsyncMock(side_effect=release) + cache = DualCache(redis_cache=remote) + handler = _PROXY_MaxParallelRequestsHandler(InternalUsageCache(cache)) + await cache.async_set_cache("key", {"current_requests": 1, "current_rpm": 1, "current_tpm": 7}, local_only=True) + task = asyncio.create_task(handler._release_realtime_counter("key")) + await started.wait() + next_admission = {"current_requests": 1, "current_rpm": 2, "current_tpm": 7} + await cache.async_set_cache("key", next_admission, local_only=True) + finish.set() + await task + assert await cache.async_get_cache("key", local_only=True) == next_admission + + +@pytest.mark.asyncio +@pytest.mark.parametrize("reject_team", [False, True]) +async def test_realtime_attachment_releases_only_acquired_legacy_slots(reject_team): + cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler(InternalUsageCache(cache)) + auth = UserAPIKeyAuth( + api_key="attachment-key", + user_id="attachment-user", + team_id="attachment-team", + team_rpm_limit=0 if reject_team else 100, + max_parallel_requests=1, + end_user_id="attachment-end-user", + metadata={"model_rpm_limit": {"test-model": 100}}, + ) + data = {"model": "test-model", "metadata": {"global_max_parallel_requests": 10}} + minute = datetime.now().strftime("%Y-%m-%d-%H-%M") + team_key = f"attachment-team::{minute}::request_count" + await cache.async_set_cache(team_key, {"current_requests": 3, "current_tpm": 7, "current_rpm": 4}) + handler.begin_realtime_attachment(data) + if reject_team: + with pytest.raises(ProxyRateLimitError, match="Rate Limit Handler"): + await handler.async_pre_call_hook(auth, cache, data, "_arealtime") + else: + await handler.async_pre_call_hook(auth, cache, data, "_arealtime") + await handler.async_release_realtime_attachment(data, auth) + await handler.async_release_realtime_attachment(data, auth) + assert await cache.async_get_cache("global_max_parallel_requests") == 0 + assert await cache.async_get_cache(f"attachment-key::{minute}::request_count") == { + "current_requests": 0, + "current_tpm": 0, + "current_rpm": 1, + } + assert await cache.async_get_cache(f"attachment-user::{minute}::request_count") == { + "current_requests": 0, + "current_tpm": 0, + "current_rpm": 1, + } + assert await cache.async_get_cache(team_key) == { + "current_requests": 3, + "current_tpm": 7, + "current_rpm": 4 if reject_team else 5, + } + assert await cache.async_get_cache(f"attachment-key::test-model::{minute}::request_count") == { + "current_requests": 0, + "current_tpm": 0, + "current_rpm": 1, + } + end_user = await cache.async_get_cache(f"attachment-end-user::{minute}::request_count") + assert end_user == (None if reject_team else {"current_requests": 0, "current_tpm": 0, "current_rpm": 1}) + if not reject_team: + handler.begin_realtime_attachment(data) + await handler.async_pre_call_hook(auth, cache, data, "_arealtime") + await handler.async_release_realtime_attachment(data, auth) + + +@pytest.mark.asyncio +async def test_realtime_attachment_rejected_before_acquisition_preserves_other_slot(): + cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler(InternalUsageCache(cache)) + auth = UserAPIKeyAuth(api_key="busy-key", max_parallel_requests=1) + minute = datetime.now().strftime("%Y-%m-%d-%H-%M") + key = f"busy-key::{minute}::request_count" + current = {"current_requests": 1, "current_tpm": 13, "current_rpm": 2} + await cache.async_set_cache(key, current) + data = {"model": "test-model"} + handler.begin_realtime_attachment(data) + with pytest.raises(ProxyRateLimitError, match="Rate Limit Handler"): + await handler.async_pre_call_hook(auth, cache, data, "_arealtime") + await handler.async_release_realtime_attachment(data, auth) + assert await cache.async_get_cache(key) == current + + @pytest.mark.parametrize( "response_obj", [ @@ -39,9 +138,7 @@ async def test_async_log_success_event_counts_non_chat_response_tokens(response_ team_id = "litellm-team" end_user_id = "customer-1" - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) current_date = datetime.now().strftime("%Y-%m-%d") current_hour = datetime.now().strftime("%H") @@ -80,7 +177,4 @@ async def test_async_log_success_event_counts_non_chat_response_tokens(response_ key=f"{scope_id}::{precise_minute}::request_count", litellm_parent_otel_span=None, ) - assert current["current_tpm"] == 50, ( - f"expected 50 tokens counted for {scope_id}, " - f"got {current['current_tpm']}" - ) + assert current["current_tpm"] == 50, f"expected 50 tokens counted for {scope_id}, got {current['current_tpm']}" diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 10c0bb88a82..14d078820af 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -31,6 +31,8 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import ( _PROXY_MaxParallelRequestsHandler_v3 as _PROXY_MaxParallelRequestsHandler, ) from litellm.proxy.utils import InternalUsageCache, ProxyLogging, hash_token +from litellm.proxy.hooks.parallel_request_limiter_v3 import isolated_request_stash +from litellm.proxy.hooks.realtime_call_lease import realtime_call_attachment from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.utils import ( EmbeddingResponse, @@ -51,6 +53,94 @@ class TimeController: self._current += timedelta(seconds=seconds) +@pytest.mark.asyncio +async def test_realtime_lease_retains_quota_across_signaling_and_three_attachments(): + cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler(InternalUsageCache(cache)) + auth = UserAPIKeyAuth(api_key="logical-owner", max_parallel_requests=1, rpm_limit=4, tpm_limit=100000) + data = {"model": "gpt-3.5-turbo", "litellm_call_id": "signaling"} + await handler.async_pre_call_hook(auth, cache, data, "arealtime_calls") + stash = get_request_stash() + assert stash.reserved_tokens > 0 + assert handler.transfer_realtime_call_slot({"litellm_call_id": "other-call"}) is None + assert stash.parallel_slot is not None + lease = handler.transfer_realtime_call_slot(data) + assert lease is not None + assert stash.parallel_slot is None + assert stash.reserved_tokens > 0 + assert not stash.reservation_released + assert handler.transfer_realtime_call_slot(data) is None + await handler.async_log_success_event( + kwargs={ + "litellm_call_id": "signaling", + "standard_logging_object": {"metadata": {"user_api_key_hash": auth.api_key}}, + }, + response_obj=ModelResponse(usage=Usage()), + start_time=datetime.now(), + end_time=datetime.now(), + ) + assert await cache.async_get_cache("{api_key:logical-owner}:tokens") == 0 + socket = object() + for attachment in range(3): + with isolated_request_stash(), realtime_call_attachment(socket): + attachment_data = { + "model": "gpt-3.5-turbo", + "litellm_call_id": f"attachment-{attachment}", + "websocket": socket, + } + await handler.async_pre_call_hook(auth, cache, attachment_data, "_arealtime") + assert get_request_stash().parallel_slot is None + await handler.async_release_realtime_attachment(attachment_data, auth) + with isolated_request_stash(), realtime_call_attachment(socket), pytest.raises(HTTPException) as error: + await handler.async_pre_call_hook(auth, cache, {"model": "gpt-3.5-turbo", "websocket": socket}, "_arealtime") + assert error.value.status_code == 429 + assert "requests" in str(error.value.detail) + quota_only = auth.model_copy(update={"rpm_limit": None}) + with isolated_request_stash(), realtime_call_attachment(object()), pytest.raises(HTTPException) as error: + await handler.async_pre_call_hook( + quota_only, cache, {"model": "gpt-3.5-turbo", "websocket": socket}, "_arealtime" + ) + assert "max_parallel_requests" in str(error.value.detail) + with isolated_request_stash(), realtime_call_attachment(socket), pytest.raises(HTTPException) as error: + await handler.async_pre_call_hook( + quota_only, cache, {"model": "gpt-3.5-turbo", "websocket": socket}, "acompletion" + ) + assert "max_parallel_requests" in str(error.value.detail) + with isolated_request_stash(), pytest.raises(HTTPException) as error: + await handler.async_pre_call_hook(quota_only, cache, {"model": "gpt-3.5-turbo"}, "arealtime_calls") + assert "max_parallel_requests" in str(error.value.detail) + assert get_request_stash() is stash + await lease.close() + with isolated_request_stash(): + await handler.async_pre_call_hook(quota_only, cache, {"model": "gpt-3.5-turbo"}, "arealtime_calls") + + +@pytest.mark.asyncio +async def test_realtime_lease_renewal_preserves_quota_past_ttl_and_does_not_resurrect_expiry(): + cache = DualCache() + clock = TimeController() + handler = _PROXY_MaxParallelRequestsHandler(InternalUsageCache(cache), time_provider=clock.now) + auth = UserAPIKeyAuth(api_key="long-call", max_parallel_requests=1) + data = {"model": "gpt-3.5-turbo", "litellm_call_id": "long-call"} + await handler.async_pre_call_hook(auth, cache, data, "arealtime_calls") + lease = handler.transfer_realtime_call_slot(data) + assert lease is not None + clock.advance(PARALLEL_REQUEST_SLOT_TTL_SECONDS - 1) + assert await lease.renew() + clock.advance(2) + with isolated_request_stash(), pytest.raises(HTTPException): + await handler.async_pre_call_hook(auth, cache, {"model": "gpt-3.5-turbo"}, "arealtime_calls") + + clock.advance(PARALLEL_REQUEST_SLOT_TTL_SECONDS) + assert not await lease.renew() + await asyncio.wait_for(lease.wait_failed(), 1) + with isolated_request_stash(): + await handler.async_pre_call_hook(auth, cache, {"model": "gpt-3.5-turbo"}, "arealtime_calls") + await lease.close() + with isolated_request_stash(), pytest.raises(HTTPException): + await handler.async_pre_call_hook(auth, cache, {"model": "gpt-3.5-turbo"}, "arealtime_calls") + + @pytest.fixture def time_controller(monkeypatch): controller = TimeController() @@ -73,14 +163,10 @@ def _isolated_request_stash(): (0.5, 50, 500), ], ) -def test_api_key_descriptor_applies_budget_throttle( - throttle_pct, expected_rpm, expected_tpm -): +def test_api_key_descriptor_applies_budget_throttle(throttle_pct, expected_rpm, expected_tpm): """The api_key rate-limit descriptor scales the key's configured TPM/RPM by the request-scoped budget_throttle_pct, leaving the configured limits intact.""" - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) user_api_key_dict = UserAPIKeyAuth( api_key=hash_token("sk-throttle"), rpm_limit=100, @@ -138,18 +224,12 @@ async def test_sliding_window_rate_limit_v3(monkeypatch, time_controller): window_starts[window_key] = now new_counter = 1 request_counts[counter_key] = new_counter - await local_cache.async_set_cache( - key=window_key, value=now, ttl=window_size - ) - await local_cache.async_set_cache( - key=counter_key, value=new_counter, ttl=window_size - ) + await local_cache.async_set_cache(key=window_key, value=now, ttl=window_size) + await local_cache.async_set_cache(key=counter_key, value=new_counter, ttl=window_size) else: new_counter = prev_counter + 1 request_counts[counter_key] = new_counter - await local_cache.async_set_cache( - key=counter_key, value=new_counter, ttl=window_size - ) + await local_cache.async_set_cache(key=counter_key, value=new_counter, ttl=window_size) results.append(now) results.append(new_counter) return results @@ -231,18 +311,12 @@ async def test_rate_limiter_script_return_values_v3(monkeypatch, time_controller window_starts[window_key] = now new_counter = 1 request_counts[counter_key] = new_counter - await local_cache.async_set_cache( - key=window_key, value=now, ttl=window_size - ) - await local_cache.async_set_cache( - key=counter_key, value=new_counter, ttl=window_size - ) + await local_cache.async_set_cache(key=window_key, value=now, ttl=window_size) + await local_cache.async_set_cache(key=counter_key, value=new_counter, ttl=window_size) else: new_counter = prev_counter + 1 request_counts[counter_key] = new_counter - await local_cache.async_set_cache( - key=counter_key, value=new_counter, ttl=window_size - ) + await local_cache.async_set_cache(key=counter_key, value=new_counter, ttl=window_size) results.append(now) results.append(new_counter) return results @@ -274,9 +348,7 @@ async def test_rate_limiter_script_return_values_v3(monkeypatch, time_controller new_window_value = await local_cache.async_get_cache(key=window_key) new_counter_value = await local_cache.async_get_cache(key=counter_key) - assert ( - new_window_value == window_value - ), "Window value should not change within window" + assert new_window_value == window_value, "Window value should not change within window" assert new_counter_value == 2, "Counter should be 2 after second request" # Wait for window to expire @@ -307,9 +379,7 @@ async def test_rate_limiter_script_return_values_v3(monkeypatch, time_controller ) @pytest.mark.flaky(reruns=3) @pytest.mark.asyncio -async def test_normal_router_call_tpm_v3( - monkeypatch, rate_limit_object, time_controller -): +async def test_normal_router_call_tpm_v3(monkeypatch, rate_limit_object, time_controller): """ Test normal router call with parallel request limiter v3 for TPM rate limiting """ @@ -385,18 +455,12 @@ async def test_normal_router_call_tpm_v3( window_starts[window_key] = now new_counter = 1 request_counts[counter_key] = new_counter - await local_cache.async_set_cache( - key=window_key, value=now, ttl=window_size - ) - await local_cache.async_set_cache( - key=counter_key, value=new_counter, ttl=window_size - ) + await local_cache.async_set_cache(key=window_key, value=now, ttl=window_size) + await local_cache.async_set_cache(key=counter_key, value=new_counter, ttl=window_size) else: new_counter = prev_counter + 1 request_counts[counter_key] = new_counter - await local_cache.async_set_cache( - key=counter_key, value=new_counter, ttl=window_size - ) + await local_cache.async_set_cache(key=counter_key, value=new_counter, ttl=window_size) results.append(now) results.append(new_counter) return results @@ -419,9 +483,7 @@ async def test_normal_router_call_tpm_v3( return None value = get_value_for_key(rate_limit_object, user_api_key_dict, "azure-model") - counter_key = parallel_request_handler.create_rate_limit_keys( - rate_limit_object, value, "tokens" - ) + counter_key = parallel_request_handler.create_rate_limit_keys(rate_limit_object, value, "tokens") # First request should succeed. Include messages + a tight max_tokens so # the atomic reserve_tpm_tokens path populates the :tokens counter with a @@ -434,12 +496,8 @@ async def test_normal_router_call_tpm_v3( "messages": [{"role": "user", "content": "hi"}], "max_tokens": 5, } - expected_reservation = parallel_request_handler._estimate_tokens_for_request( - data=pre_call_data - ) - assert ( - expected_reservation < 10 - ), "Test premise: reservation must fit under tpm_limit=10" + expected_reservation = parallel_request_handler._estimate_tokens_for_request(data=pre_call_data) + assert expected_reservation < 10, "Test premise: reservation must fit under tpm_limit=10" await parallel_request_handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, @@ -467,15 +525,11 @@ async def test_normal_router_call_tpm_v3( counter_value = await local_cache.async_get_cache(key=counter_key) print(f"local_cache: {local_cache.in_memory_cache.cache_dict}") - assert ( - counter_value is not None - ), f"Counter value should be stored in cache for {counter_key}" + assert counter_value is not None, f"Counter value should be stored in cache for {counter_key}" # Manually increment the token counter to simulate token usage from previous call # This simulates what would happen after a successful call - await local_cache.async_increment_cache( - key=counter_key, value=15, ttl=2 - ) # Use up most of our 10 token limit + await local_cache.async_increment_cache(key=counter_key, value=15, ttl=2) # Use up most of our 10 token limit # Make another request to test rate limiting - this should fail as we've consumed tokens with pytest.raises(HTTPException) as exc_info: @@ -523,17 +577,13 @@ async def test_token_rate_limit_type_respected_v3(monkeypatch, token_rate_limit_ _api_key = hash_token(_api_key) user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, tpm_limit=100) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Mock the get_rate_limit_type method directly since it imports general_settings internally def mock_get_rate_limit_type(): return token_rate_limit_type - monkeypatch.setattr( - parallel_request_handler, "get_rate_limit_type", mock_get_rate_limit_type - ) + monkeypatch.setattr(parallel_request_handler, "get_rate_limit_type", mock_get_rate_limit_type) # Create a mock response with different token counts mock_usage = Usage(prompt_tokens=20, completion_tokens=30, total_tokens=50) @@ -582,9 +632,9 @@ async def test_token_rate_limit_type_respected_v3(monkeypatch, token_rate_limit_ ) # Verify that the correct token count was used based on the rate limit type - assert ( - len(captured_operations) == 1 - ), "Should have 1 operation: the TPM increment (parallel slots are released via the gauge, not the pipeline)" + assert len(captured_operations) == 1, ( + "Should have 1 operation: the TPM increment (parallel slots are released via the gauge, not the pipeline)" + ) tpm_operation = None for op in captured_operations: @@ -601,9 +651,9 @@ async def test_token_rate_limit_type_respected_v3(monkeypatch, token_rate_limit_ "total": mock_usage.total_tokens, # 50 } - assert ( - tpm_operation["increment_value"] == expected_tokens[token_rate_limit_type] - ), f"Expected {expected_tokens[token_rate_limit_type]} tokens for type '{token_rate_limit_type}', got {tpm_operation['increment_value']}" + assert tpm_operation["increment_value"] == expected_tokens[token_rate_limit_type], ( + f"Expected {expected_tokens[token_rate_limit_type]} tokens for type '{token_rate_limit_type}', got {tpm_operation['increment_value']}" + ) @pytest.mark.parametrize( @@ -620,9 +670,7 @@ async def test_token_rate_limit_type_respected_v3(monkeypatch, token_rate_limit_ ], ) @pytest.mark.asyncio -async def test_async_log_success_event_counts_non_chat_response_tokens( - monkeypatch, response_obj -): +async def test_async_log_success_event_counts_non_chat_response_tokens(monkeypatch, response_obj): """ Embedding and text completion responses must increment the TPM counter, not just chat completion ModelResponse objects. @@ -630,12 +678,8 @@ async def test_async_log_success_event_counts_non_chat_response_tokens( monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60") _api_key = hash_token("sk-12345") - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) - monkeypatch.setattr( - parallel_request_handler, "get_rate_limit_type", lambda: "total" - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + monkeypatch.setattr(parallel_request_handler, "get_rate_limit_type", lambda: "total") mock_kwargs = { "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}, @@ -661,9 +705,7 @@ async def test_async_log_success_event_counts_non_chat_response_tokens( end_time=datetime.now(), ) - tpm_operation = next( - (op for op in captured_operations if op["key"].endswith(":tokens")), None - ) + tpm_operation = next((op for op in captured_operations if op["key"].endswith(":tokens")), None) assert tpm_operation is not None, "Should have a TPM increment operation" assert tpm_operation["increment_value"] == 50 @@ -679,9 +721,7 @@ async def test_async_log_failure_event_v3(): _api_key = "sk-12345" _api_key = hash_token(_api_key) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" await _seed_max_parallel_requests_slots(local_cache, counter_key, ["slot-a", "slot-b"]) @@ -731,29 +771,18 @@ async def test_failure_event_without_acquired_slot_does_not_release_v3(): """ _api_key = hash_token("sk-12345") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" - await _seed_max_parallel_requests_slots( - local_cache, counter_key, ["slot-a", "slot-b", "slot-c"] - ) + await _seed_max_parallel_requests_slots(local_cache, counter_key, ["slot-a", "slot-b", "slot-c"]) await handler.async_log_failure_event( - kwargs={ - "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}} - }, + kwargs={"standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}}, response_obj=None, start_time=None, end_time=None, ) - assert ( - handler._gauge_in_flight_from_cache_value( - await local_cache.async_get_cache(key=counter_key) - ) - == 3 - ) + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 3 @pytest.mark.asyncio @@ -804,9 +833,7 @@ async def test_rejected_request_does_not_consume_parallel_slot_v3(): rejected requests that should have been admitted after a release. """ local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) _api_key = hash_token("sk-12345") user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1) @@ -859,9 +886,7 @@ async def test_parallel_gauge_uses_atomic_redis_script_v3(): and an over-limit script result maps to a 429 without occupying a slot. """ local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) _api_key = hash_token("sk-12345") user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=5) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" @@ -886,9 +911,7 @@ async def test_parallel_gauge_uses_atomic_redis_script_v3(): stashed_slot_id = stashed_acquisition["slot_id"] assert isinstance(stashed_slot_id, str) and stashed_slot_id assert stashed_acquisition["counter_keys"] == [counter_key] - assert captured_calls == [ - ([counter_key], [5, PARALLEL_REQUEST_SLOT_TTL_SECONDS, stashed_slot_id]) - ] + assert captured_calls == [([counter_key], [5, PARALLEL_REQUEST_SLOT_TTL_SECONDS, stashed_slot_id])] assert ( await handler.internal_usage_cache.async_get_cache( key=counter_key, litellm_parent_otel_span=None, local_only=True @@ -935,9 +958,7 @@ async def test_should_rate_limit_only_called_when_limits_exist_v3(): _api_key = "sk-12345" _api_key = hash_token(_api_key) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Mock should_rate_limit to track if it's called should_rate_limit_called = False @@ -966,9 +987,7 @@ async def test_should_rate_limit_only_called_when_limits_exist_v3(): call_type="", ) - assert ( - not should_rate_limit_called - ), "should_rate_limit should not be called when no rate limits are configured" + assert not should_rate_limit_called, "should_rate_limit should not be called when no rate limits are configured" # Test 2: API key rate limits configured - should_rate_limit SHOULD be called should_rate_limit_called = False @@ -984,9 +1003,7 @@ async def test_should_rate_limit_only_called_when_limits_exist_v3(): call_type="", ) - assert ( - should_rate_limit_called - ), "should_rate_limit should be called when API key rate limits are configured" + assert should_rate_limit_called, "should_rate_limit should be called when API key rate limits are configured" # Test 3: User rate limits configured - should_rate_limit SHOULD be called should_rate_limit_called = False @@ -1003,9 +1020,7 @@ async def test_should_rate_limit_only_called_when_limits_exist_v3(): call_type="", ) - assert ( - should_rate_limit_called - ), "should_rate_limit should be called when user rate limits are configured" + assert should_rate_limit_called, "should_rate_limit should be called when user rate limits are configured" # Test 4: Team rate limits configured - should_rate_limit SHOULD be called should_rate_limit_called = False @@ -1022,9 +1037,7 @@ async def test_should_rate_limit_only_called_when_limits_exist_v3(): call_type="", ) - assert ( - should_rate_limit_called - ), "should_rate_limit should be called when team rate limits are configured" + assert should_rate_limit_called, "should_rate_limit should be called when team rate limits are configured" # Test 5: End user rate limits configured - should_rate_limit SHOULD be called should_rate_limit_called = False @@ -1041,9 +1054,7 @@ async def test_should_rate_limit_only_called_when_limits_exist_v3(): call_type="", ) - assert ( - should_rate_limit_called - ), "should_rate_limit should be called when end user rate limits are configured" + assert should_rate_limit_called, "should_rate_limit should be called when end user rate limits are configured" # Test 6: Max parallel requests configured - should_rate_limit SHOULD be called should_rate_limit_called = False @@ -1059,9 +1070,7 @@ async def test_should_rate_limit_only_called_when_limits_exist_v3(): call_type="", ) - assert ( - should_rate_limit_called - ), "should_rate_limit should be called when max parallel requests are configured" + assert should_rate_limit_called, "should_rate_limit should be called when max parallel requests are configured" @pytest.mark.asyncio @@ -1077,9 +1086,7 @@ async def test_model_specific_rate_limits_only_called_when_configured_v3(): _api_key = "sk-12345" _api_key = hash_token(_api_key) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Mock should_rate_limit to track if it's called should_rate_limit_called = False @@ -1095,9 +1102,7 @@ async def test_model_specific_rate_limits_only_called_when_configured_v3(): should_rate_limit_called = False user_api_key_dict_with_model_limits = UserAPIKeyAuth( api_key=_api_key, - metadata={ - "model_tpm_limit": {"gpt-4": 1000} - }, # Rate limit for gpt-4, not gpt-3.5-turbo + metadata={"model_tpm_limit": {"gpt-4": 1000}}, # Rate limit for gpt-4, not gpt-3.5-turbo ) await parallel_request_handler.async_pre_call_hook( @@ -1107,17 +1112,15 @@ async def test_model_specific_rate_limits_only_called_when_configured_v3(): call_type="", ) - assert ( - not should_rate_limit_called - ), "should_rate_limit should not be called when model-specific limits don't match requested model" + assert not should_rate_limit_called, ( + "should_rate_limit should not be called when model-specific limits don't match requested model" + ) # Test 2: Model-specific rate limits configured for requested model - SHOULD be called should_rate_limit_called = False user_api_key_dict_with_matching_model_limits = UserAPIKeyAuth( api_key=_api_key, - metadata={ - "model_tpm_limit": {"gpt-3.5-turbo": 1000} - }, # Rate limit for requested model + metadata={"model_tpm_limit": {"gpt-3.5-turbo": 1000}}, # Rate limit for requested model ) await parallel_request_handler.async_pre_call_hook( @@ -1127,9 +1130,9 @@ async def test_model_specific_rate_limits_only_called_when_configured_v3(): call_type="", ) - assert ( - should_rate_limit_called - ), "should_rate_limit should be called when model-specific limits match requested model" + assert should_rate_limit_called, ( + "should_rate_limit should be called when model-specific limits match requested model" + ) @pytest.mark.asyncio @@ -1156,9 +1159,7 @@ async def test_tpm_api_key_rate_limits_v3(): user_api_key_dict.metadata["model_rpm_limit"] = rpms local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Mock should_rate_limit to capture the descriptors captured_descriptors = None @@ -1216,15 +1217,11 @@ async def test_tpm_api_key_rate_limits_v3(): break assert model_per_key_descriptor is not None, "Api-Key descriptor should be present" - assert ( - model_per_key_descriptor["value"] == f"{_api_key_hash}:{model}" - ), "Api-Key value should combine api_key and model" - assert ( - model_per_key_descriptor["rate_limit"]["requests_per_unit"] == rpm_limit - ), "Api-Key RPM limit should be set" - assert ( - model_per_key_descriptor["rate_limit"]["tokens_per_unit"] == tpm_limit - ), "Api-Key TPM limit should be set" + assert model_per_key_descriptor["value"] == f"{_api_key_hash}:{model}", ( + "Api-Key value should combine api_key and model" + ) + assert model_per_key_descriptor["rate_limit"]["requests_per_unit"] == rpm_limit, "Api-Key RPM limit should be set" + assert model_per_key_descriptor["rate_limit"]["tokens_per_unit"] == tpm_limit, "Api-Key TPM limit should be set" @pytest.mark.asyncio @@ -1251,9 +1248,7 @@ async def test_rpm_api_key_rate_limits_v3(): user_api_key_dict.metadata["model_rpm_limit"] = rpms local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Mock should_rate_limit to capture the descriptors captured_descriptors = None @@ -1311,15 +1306,11 @@ async def test_rpm_api_key_rate_limits_v3(): break assert model_per_key_descriptor is not None, "Api-Key descriptor should be present" - assert ( - model_per_key_descriptor["value"] == f"{_api_key_hash}:{model}" - ), "Api-Key value should combine api_key and model" - assert ( - model_per_key_descriptor["rate_limit"]["requests_per_unit"] == rpm_limit - ), "Api-Key RPM limit should be set" - assert ( - model_per_key_descriptor["rate_limit"]["tokens_per_unit"] == tpm_limit - ), "Api-Key TPM limit should be set" + assert model_per_key_descriptor["value"] == f"{_api_key_hash}:{model}", ( + "Api-Key value should combine api_key and model" + ) + assert model_per_key_descriptor["rate_limit"]["requests_per_unit"] == rpm_limit, "Api-Key RPM limit should be set" + assert model_per_key_descriptor["rate_limit"]["tokens_per_unit"] == tpm_limit, "Api-Key TPM limit should be set" @pytest.mark.asyncio @@ -1341,9 +1332,7 @@ async def test_team_member_rate_limits_v3(): ) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Mock should_rate_limit to capture the descriptors captured_descriptors = None @@ -1375,18 +1364,12 @@ async def test_team_member_rate_limits_v3(): team_member_descriptor = descriptor break - assert ( - team_member_descriptor is not None - ), "Team member descriptor should be present" - assert ( - team_member_descriptor["value"] == f"{_team_id}:{_user_id}" - ), "Team member value should combine team_id and user_id" - assert ( - team_member_descriptor["rate_limit"]["requests_per_unit"] == 10 - ), "Team member RPM limit should be set" - assert ( - team_member_descriptor["rate_limit"]["tokens_per_unit"] == 1000 - ), "Team member TPM limit should be set" + assert team_member_descriptor is not None, "Team member descriptor should be present" + assert team_member_descriptor["value"] == f"{_team_id}:{_user_id}", ( + "Team member value should combine team_id and user_id" + ) + assert team_member_descriptor["rate_limit"]["requests_per_unit"] == 10, "Team member RPM limit should be set" + assert team_member_descriptor["rate_limit"]["tokens_per_unit"] == 1000, "Team member TPM limit should be set" @pytest.mark.asyncio @@ -1409,9 +1392,7 @@ async def test_team_member_rate_limits_v3_raises_429_when_over_limit(): ) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) captured_descriptors = None @@ -1487,9 +1468,7 @@ async def test_dynamic_rate_limiting_v3(): ) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Mock should_rate_limit to track if limits are enforced captured_descriptors = [] @@ -1522,9 +1501,9 @@ async def test_dynamic_rate_limiting_v3(): break assert api_key_descriptor is not None, "API key descriptor should be present" - assert ( - api_key_descriptor["rate_limit"]["requests_per_unit"] is None - ), "RPM limit should be None when dynamic mode and no failures" + assert api_key_descriptor["rate_limit"]["requests_per_unit"] is None, ( + "RPM limit should be None when dynamic mode and no failures" + ) # Test 2: With failures - rate limits SHOULD be enforced (rpm_limit should be set) async def mock_check_with_failures(*args, **kwargs): @@ -1548,9 +1527,9 @@ async def test_dynamic_rate_limiting_v3(): break assert api_key_descriptor is not None, "API key descriptor should be present" - assert ( - api_key_descriptor["rate_limit"]["requests_per_unit"] == 2 - ), "RPM limit should be enforced when dynamic mode and failures detected" + assert api_key_descriptor["rate_limit"]["requests_per_unit"] == 2, ( + "RPM limit should be enforced when dynamic mode and failures detected" + ) @pytest.mark.flaky(retries=3, delay=2) @@ -1596,9 +1575,7 @@ async def test_async_increment_tokens_with_ttl_preservation(): ) local_cache = DualCache(redis_cache=redis_cache) - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Verify Redis connection is working try: @@ -1608,9 +1585,7 @@ async def test_async_increment_tokens_with_ttl_preservation(): # Verify the TTL preservation script is registered if parallel_request_handler.token_increment_script is None: - pytest.skip( - "Token increment script not available - Redis Lua scripting may not be supported" - ) + pytest.skip("Token increment script not available - Redis Lua scripting may not be supported") # Test keys - use hash tags to ensure they map to same Redis cluster slot # Use a unique suffix per test run to avoid stale state from prior runs @@ -1631,11 +1606,11 @@ async def test_async_increment_tokens_with_ttl_preservation(): # First increment: Create operations with mixed TTL scenarios pipeline_operations_first = [ + RedisPipelineIncrementOperation(key=test_key_with_ttl, increment_value=10.0, ttl=60), RedisPipelineIncrementOperation( - key=test_key_with_ttl, increment_value=10.0, ttl=60 - ), - RedisPipelineIncrementOperation( - key=test_key_without_ttl, increment_value=5.0, ttl=None # No TTL + key=test_key_without_ttl, + increment_value=5.0, + ttl=None, # No TTL ), ] @@ -1649,29 +1624,21 @@ async def test_async_increment_tokens_with_ttl_preservation(): # Verify keys exist and check initial TTL ttl_after_first = await redis_cache.async_get_ttl(test_key_with_ttl) - value_after_first_with_ttl = await redis_cache.async_get_cache( - test_key_with_ttl - ) - value_after_first_without_ttl = await redis_cache.async_get_cache( - test_key_without_ttl - ) + value_after_first_with_ttl = await redis_cache.async_get_cache(test_key_with_ttl) + value_after_first_without_ttl = await redis_cache.async_get_cache(test_key_without_ttl) - assert ( - value_after_first_with_ttl == 10.0 - ), f"First increment should set value to 10.0, got {value_after_first_with_ttl}" - assert ( - value_after_first_without_ttl == 5.0 - ), "First increment should set value to 5.0" - assert ( - ttl_after_first is not None and ttl_after_first > 0 - ), "Key with TTL should have positive TTL after first increment" + assert value_after_first_with_ttl == 10.0, ( + f"First increment should set value to 10.0, got {value_after_first_with_ttl}" + ) + assert value_after_first_without_ttl == 5.0, "First increment should set value to 5.0" + assert ttl_after_first is not None and ttl_after_first > 0, ( + "Key with TTL should have positive TTL after first increment" + ) assert ttl_after_first <= 60, "TTL should not exceed the set value" # Check TTL for key without TTL (should be None, meaning no expiry) ttl_no_ttl_key = await redis_cache.async_get_ttl(test_key_without_ttl) - assert ( - ttl_no_ttl_key is None - ), "Key without TTL should have no expiry (None from async_get_ttl)" + assert ttl_no_ttl_key is None, "Key without TTL should have no expiry (None from async_get_ttl)" # Wait a moment to ensure TTL decreases await asyncio.sleep(2) @@ -1679,10 +1646,14 @@ async def test_async_increment_tokens_with_ttl_preservation(): # Second increment: Same operations to test TTL preservation pipeline_operations_second = [ RedisPipelineIncrementOperation( - key=test_key_with_ttl, increment_value=15.0, ttl=60 # Same TTL value + key=test_key_with_ttl, + increment_value=15.0, + ttl=60, # Same TTL value ), RedisPipelineIncrementOperation( - key=test_key_without_ttl, increment_value=7.0, ttl=None # No TTL + key=test_key_without_ttl, + increment_value=7.0, + ttl=None, # No TTL ), ] @@ -1696,39 +1667,23 @@ async def test_async_increment_tokens_with_ttl_preservation(): # Verify TTL preservation and value updates ttl_after_second = await redis_cache.async_get_ttl(test_key_with_ttl) - value_after_second_with_ttl = await redis_cache.async_get_cache( - test_key_with_ttl - ) - value_after_second_without_ttl = await redis_cache.async_get_cache( - test_key_without_ttl - ) + value_after_second_with_ttl = await redis_cache.async_get_cache(test_key_with_ttl) + value_after_second_without_ttl = await redis_cache.async_get_cache(test_key_without_ttl) - assert ( - value_after_second_with_ttl == 25.0 - ), "Second increment should update value to 25.0" - assert ( - value_after_second_without_ttl == 12.0 - ), "Second increment should update value to 12.0" + assert value_after_second_with_ttl == 25.0, "Second increment should update value to 25.0" + assert value_after_second_without_ttl == 12.0, "Second increment should update value to 12.0" # Critical test: TTL should be preserved (not reset to 60) assert ttl_after_second is not None, "TTL should still exist" - assert ( - ttl_after_second < ttl_after_first - ), "TTL should have decreased (not been reset)" + assert ttl_after_second < ttl_after_first, "TTL should have decreased (not been reset)" assert ttl_after_second > 0, "TTL should still be positive" # TTL should not be close to the original 60 seconds (proving it wasn't reset) - assert ( - ttl_after_second < 59 - ), "TTL should be significantly less than original, proving preservation" + assert ttl_after_second < 59, "TTL should be significantly less than original, proving preservation" # Key without TTL should still have no expiry - ttl_no_ttl_key_after_second = await redis_cache.async_get_ttl( - test_key_without_ttl - ) - assert ( - ttl_no_ttl_key_after_second is None - ), "Key without TTL should still have no expiry" + ttl_no_ttl_key_after_second = await redis_cache.async_get_ttl(test_key_without_ttl) + assert ttl_no_ttl_key_after_second is None, "Key without TTL should still have no expiry" finally: # Clean up test keys @@ -1755,44 +1710,30 @@ async def test_async_increment_tokens_fallback_behavior(): from litellm.types.caching import RedisPipelineIncrementOperation local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Mock the token_increment_script to None to simulate unavailable script parallel_request_handler.token_increment_script = None # Mock the fallback method fallback_called = False - original_method = ( - parallel_request_handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline - ) + original_method = parallel_request_handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline async def mock_fallback(*args, **kwargs): nonlocal fallback_called fallback_called = True return await original_method(*args, **kwargs) - parallel_request_handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = ( - mock_fallback - ) + parallel_request_handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = mock_fallback # Test operations - pipeline_operations = [ - RedisPipelineIncrementOperation( - key="test_fallback_key", increment_value=10.0, ttl=60 - ) - ] + pipeline_operations = [RedisPipelineIncrementOperation(key="test_fallback_key", increment_value=10.0, ttl=60)] # Execute increment - await parallel_request_handler.async_increment_tokens_with_ttl_preservation( - pipeline_operations=pipeline_operations - ) + await parallel_request_handler.async_increment_tokens_with_ttl_preservation(pipeline_operations=pipeline_operations) # Verify fallback was called - assert ( - fallback_called - ), "Fallback method should be called when Lua script is not available" + assert fallback_called, "Fallback method should be called when Lua script is not available" # Redis Cluster Compatibility Tests @@ -1803,9 +1744,7 @@ def test_group_keys_by_hash_tag_regular_redis(): For regular Redis, all keys should be grouped together under a single group. """ local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Test keys with different hash tags test_keys = [ @@ -1825,9 +1764,7 @@ def test_group_keys_by_hash_tag_regular_redis(): # Verify all keys are in single group for regular Redis assert len(groups) == 1, f"Expected 1 group for regular Redis, got {len(groups)}" assert "all_keys" in groups, "Expected 'all_keys' group for regular Redis" - assert set(groups["all_keys"]) == set( - test_keys - ), "All keys should be in single group" + assert set(groups["all_keys"]) == set(test_keys), "All keys should be in single group" def test_group_keys_by_hash_tag_redis_cluster(): @@ -1839,9 +1776,7 @@ def test_group_keys_by_hash_tag_redis_cluster(): from unittest.mock import patch local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Mock _is_redis_cluster to return True with patch.object(handler, "_is_redis_cluster", return_value=True): @@ -1861,17 +1796,13 @@ def test_group_keys_by_hash_tag_redis_cluster(): # All group keys should start with "slot_" for group_key in groups.keys(): - assert group_key.startswith( - "slot_" - ), f"Group key {group_key} should start with 'slot_'" + assert group_key.startswith("slot_"), f"Group key {group_key} should start with 'slot_'" # Verify all original keys are present across groups all_grouped_keys = [] for group_keys in groups.values(): all_grouped_keys.extend(group_keys) - assert set(all_grouped_keys) == set( - test_keys - ), "All keys should be present in groups" + assert set(all_grouped_keys) == set(test_keys), "All keys should be present in groups" def test_keyslot_for_redis_cluster(): @@ -1879,9 +1810,7 @@ def test_keyslot_for_redis_cluster(): Test the keyslot calculation for Redis cluster. """ local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Test basic key slot1 = handler.keyslot_for_redis_cluster("user:1000") @@ -1909,18 +1838,14 @@ async def test_execute_redis_batch_rate_limiter_script_cluster_compatibility(): from unittest.mock import AsyncMock, patch local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Mock _is_redis_cluster to return True for this test with patch.object(handler, "_is_redis_cluster", return_value=True): # Mock script that simulates Redis cluster slot conflict mock_script = AsyncMock() mock_script.side_effect = [ - Exception( - "EVALSHA - all keys must map to the same key slot" - ), # First group fails + Exception("EVALSHA - all keys must map to the same key slot"), # First group fails [1234, 1, 1234, 2], # Second group succeeds ] handler.batch_rate_limiter_script = mock_script @@ -1937,9 +1862,7 @@ async def test_execute_redis_batch_rate_limiter_script_cluster_compatibility(): ] # Execute the method - results = await handler._execute_redis_batch_rate_limiter_script( - keys_to_fetch=test_keys, now_int=1234 - ) + results = await handler._execute_redis_batch_rate_limiter_script(keys_to_fetch=test_keys, now_int=1234) # Verify results: 2 from fallback + 4 from successful script = 6 total assert len(results) == 6, f"Expected 6 results, got {len(results)}" @@ -1964,9 +1887,7 @@ async def test_execute_redis_batch_rate_limiter_script_cluster_compatibility(): # Should have processed all keys (some might be duplicated due to fallback) unique_processed_keys = set(all_processed_keys) - assert ( - len(unique_processed_keys) >= 2 - ), "Should have processed at least some keys" + assert len(unique_processed_keys) >= 2, "Should have processed at least some keys" @pytest.mark.asyncio @@ -1992,9 +1913,7 @@ async def test_multiple_rate_limits_per_descriptor(): ) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Mock should_rate_limit to return a response with multiple statuses where one hits the limit # This simulates the case where we have more statuses than descriptors due to multiple rate limit types @@ -2072,9 +1991,7 @@ async def test_missing_descriptor_fallback(): ) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Mock should_rate_limit to return a status with descriptor_key that doesn't match descriptors async def mock_should_rate_limit(descriptors, **kwargs): @@ -2118,9 +2035,7 @@ async def test_get_rate_limit_type_default_is_total(monkeypatch): This verifies the change from 'output' to 'total' as the default value. """ local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Mock general_settings to return empty dict (no token_rate_limit_type set) import litellm.proxy.proxy_server as proxy_server @@ -2130,9 +2045,7 @@ async def test_get_rate_limit_type_default_is_total(monkeypatch): try: result = parallel_request_handler.get_rate_limit_type() - assert ( - result == "total" - ), f"Default rate limit type should be 'total', got '{result}'" + assert result == "total", f"Default rate limit type should be 'total', got '{result}'" finally: monkeypatch.setattr(proxy_server, "general_settings", original_settings) @@ -2143,23 +2056,17 @@ async def test_get_rate_limit_type_invalid_falls_back_to_total(monkeypatch): Test that get_rate_limit_type falls back to 'total' when an invalid value is specified. """ local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Mock general_settings to return an invalid token_rate_limit_type import litellm.proxy.proxy_server as proxy_server original_settings = getattr(proxy_server, "general_settings", {}) - monkeypatch.setattr( - proxy_server, "general_settings", {"token_rate_limit_type": "invalid_type"} - ) + monkeypatch.setattr(proxy_server, "general_settings", {"token_rate_limit_type": "invalid_type"}) try: result = parallel_request_handler.get_rate_limit_type() - assert ( - result == "total" - ), f"Invalid rate limit type should fall back to 'total', got '{result}'" + assert result == "total", f"Invalid rate limit type should fall back to 'total', got '{result}'" finally: monkeypatch.setattr(proxy_server, "general_settings", original_settings) @@ -2173,9 +2080,7 @@ async def test_get_rate_limit_type_invalid_falls_back_to_total(monkeypatch): ], ) @pytest.mark.asyncio -async def test_async_log_success_event_with_dict_usage( - monkeypatch, token_rate_limit_type, expected_field -): +async def test_async_log_success_event_with_dict_usage(monkeypatch, token_rate_limit_type, expected_field): """ Test that async_log_success_event correctly handles usage as a dict (Responses API format). @@ -2187,17 +2092,13 @@ async def test_async_log_success_event_with_dict_usage( _api_key = "sk-12345" _api_key = hash_token(_api_key) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Mock the get_rate_limit_type method def mock_get_rate_limit_type(): return token_rate_limit_type - monkeypatch.setattr( - parallel_request_handler, "get_rate_limit_type", mock_get_rate_limit_type - ) + monkeypatch.setattr(parallel_request_handler, "get_rate_limit_type", mock_get_rate_limit_type) # Create a mock response object with usage as a dict (Responses API format) from litellm.types.utils import BaseLiteLLMOpenAIResponseObject @@ -2260,9 +2161,9 @@ async def test_async_log_success_event_with_dict_usage( "total": 60, # total_tokens } - assert ( - tpm_operation["increment_value"] == expected_tokens[token_rate_limit_type] - ), f"Expected {expected_tokens[token_rate_limit_type]} tokens for type '{token_rate_limit_type}', got {tpm_operation['increment_value']}" + assert tpm_operation["increment_value"] == expected_tokens[token_rate_limit_type], ( + f"Expected {expected_tokens[token_rate_limit_type]} tokens for type '{token_rate_limit_type}', got {tpm_operation['increment_value']}" + ) @pytest.mark.asyncio @@ -2277,17 +2178,13 @@ async def test_async_log_success_event_with_dict_usage_missing_fields(monkeypatc _api_key = "sk-12345" _api_key = hash_token(_api_key) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Mock the get_rate_limit_type method def mock_get_rate_limit_type(): return "output" - monkeypatch.setattr( - parallel_request_handler, "get_rate_limit_type", mock_get_rate_limit_type - ) + monkeypatch.setattr(parallel_request_handler, "get_rate_limit_type", mock_get_rate_limit_type) # Create a mock response object with usage as a dict missing some fields mock_response = MagicMock() @@ -2298,9 +2195,7 @@ async def test_async_log_success_event_with_dict_usage_missing_fields(monkeypatc } from litellm.types.utils import BaseLiteLLMOpenAIResponseObject - mock_response.__class__ = type( - "MockResponse", (BaseLiteLLMOpenAIResponseObject,), {} - ) + mock_response.__class__ = type("MockResponse", (BaseLiteLLMOpenAIResponseObject,), {}) # Create mock kwargs for the success event mock_kwargs = { @@ -2356,9 +2251,7 @@ async def test_execute_token_increment_script_cluster_compatibility(): from litellm.types.caching import RedisPipelineIncrementOperation local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Mock _is_redis_cluster to return True for this test with patch.object(handler, "_is_redis_cluster", return_value=True): @@ -2396,18 +2289,14 @@ async def test_execute_token_increment_script_cluster_compatibility(): "{api_key:sk-123}:max_parallel_requests", "{user:user-456}:tokens", } - assert ( - set(all_processed_keys) == expected_keys - ), "All operation keys should be processed" + assert set(all_processed_keys) == expected_keys, "All operation keys should be processed" # Verify args structure is correct for each call for call_args in call_args_list: keys = call_args[1]["keys"] args = call_args[1]["args"] # Each key should have 2 args (increment_value, ttl) - assert ( - len(args) == len(keys) * 2 - ), f"Each key should have 2 args, got {len(args)} args for {len(keys)} keys" + assert len(args) == len(keys) * 2, f"Each key should have 2 args, got {len(args)} args for {len(keys)} keys" @pytest.mark.asyncio @@ -2430,9 +2319,7 @@ async def test_agent_level_rate_limit_descriptors(): ) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) mock_agent = AgentResponse( agent_id=_agent_id, @@ -2497,9 +2384,7 @@ async def test_agent_session_rate_limit_descriptors(): ) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) mock_agent = AgentResponse( agent_id=_agent_id, @@ -2566,9 +2451,7 @@ async def test_agent_session_rate_limit_skipped_without_session_id(): ) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) mock_agent = AgentResponse( agent_id=_agent_id, @@ -2601,8 +2484,7 @@ async def test_agent_session_rate_limit_skipped_without_session_id(): # should_rate_limit should not have been called (no agent-level limits, only session limits # but no session_id) assert captured_descriptors is None, ( - "No descriptors should be created when agent has only session limits " - "but no session_id in request" + "No descriptors should be created when agent has only session limits but no session_id in request" ) @@ -2625,9 +2507,7 @@ async def test_agent_rate_limit_from_metadata_agent_id(): ) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) mock_agent = AgentResponse( agent_id=_agent_id, @@ -2668,9 +2548,7 @@ async def test_agent_rate_limit_from_metadata_agent_id(): agent_descriptor = d break - assert ( - agent_descriptor is not None - ), "Agent descriptor should be created from metadata agent_id" + assert agent_descriptor is not None, "Agent descriptor should be created from metadata agent_id" assert agent_descriptor["value"] == _agent_id assert agent_descriptor["rate_limit"]["requests_per_unit"] == 25 @@ -2696,9 +2574,7 @@ async def test_agent_both_agent_and_session_rate_limits(): ) local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) mock_agent = AgentResponse( agent_id=_agent_id, @@ -2765,16 +2641,12 @@ async def test_agent_rate_limit_tpm_increment_on_success(monkeypatch): _session_id = "sess_tpm_test" local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) def mock_get_rate_limit_type(): return "total" - monkeypatch.setattr( - parallel_request_handler, "get_rate_limit_type", mock_get_rate_limit_type - ) + monkeypatch.setattr(parallel_request_handler, "get_rate_limit_type", mock_get_rate_limit_type) mock_usage = Usage(prompt_tokens=20, completion_tokens=30, total_tokens=50) mock_response = ModelResponse( @@ -2885,18 +2757,12 @@ async def test_agent_rate_limit_429_on_over_limit(monkeypatch, time_controller): window_starts[window_key] = now new_counter = 1 request_counts[counter_key] = new_counter - await local_cache.async_set_cache( - key=window_key, value=now, ttl=window_size - ) - await local_cache.async_set_cache( - key=counter_key, value=new_counter, ttl=window_size - ) + await local_cache.async_set_cache(key=window_key, value=now, ttl=window_size) + await local_cache.async_set_cache(key=counter_key, value=new_counter, ttl=window_size) else: new_counter = prev_counter + 1 request_counts[counter_key] = new_counter - await local_cache.async_set_cache( - key=counter_key, value=new_counter, ttl=window_size - ) + await local_cache.async_set_cache(key=counter_key, value=new_counter, ttl=window_size) results.append(now) results.append(new_counter) return results @@ -3034,9 +2900,7 @@ async def test_project_model_rate_limits_enforced_v3(): """ _api_key = hash_token("sk-project-test") local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) captured_descriptors = [] @@ -3064,13 +2928,9 @@ async def test_project_model_rate_limits_enforced_v3(): ) descriptor_keys = [d["key"] for d in captured_descriptors] - assert ( - "model_per_project" in descriptor_keys - ), f"Expected model_per_project descriptor, got: {descriptor_keys}" + assert "model_per_project" in descriptor_keys, f"Expected model_per_project descriptor, got: {descriptor_keys}" - model_per_project = next( - d for d in captured_descriptors if d["key"] == "model_per_project" - ) + model_per_project = next(d for d in captured_descriptors if d["key"] == "model_per_project") assert model_per_project["value"] == "proj-abc123:gpt-4" assert model_per_project["rate_limit"]["requests_per_unit"] == 5 assert model_per_project["rate_limit"]["tokens_per_unit"] == 1000 @@ -3081,9 +2941,7 @@ async def test_project_model_rate_limits_not_triggered_for_other_model_v3(): """Project model limits should not trigger for a model not in project_metadata.""" _api_key = hash_token("sk-project-test-2") local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) captured_descriptors = [] @@ -3110,9 +2968,9 @@ async def test_project_model_rate_limits_not_triggered_for_other_model_v3(): ) descriptor_keys = [d["key"] for d in captured_descriptors] - assert ( - "model_per_project" not in descriptor_keys - ), f"model_per_project should not be added for unrelated model, got: {descriptor_keys}" + assert "model_per_project" not in descriptor_keys, ( + f"model_per_project should not be added for unrelated model, got: {descriptor_keys}" + ) @pytest.mark.asyncio @@ -3123,9 +2981,7 @@ async def test_project_model_itpm_otpm_limits_enforced_v3(): """ _api_key = hash_token("sk-project-io-test") local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) captured_descriptors = [] @@ -3156,12 +3012,8 @@ async def test_project_model_itpm_otpm_limits_enforced_v3(): assert "model_per_project_otpm" in descriptor_keys assert "model_per_project" not in descriptor_keys - itpm_descriptor = next( - d for d in captured_descriptors if d["key"] == "model_per_project_itpm" - ) - otpm_descriptor = next( - d for d in captured_descriptors if d["key"] == "model_per_project_otpm" - ) + itpm_descriptor = next(d for d in captured_descriptors if d["key"] == "model_per_project_itpm") + otpm_descriptor = next(d for d in captured_descriptors if d["key"] == "model_per_project_otpm") assert itpm_descriptor["value"] == "proj-mantle:bedrock_mantle/claude-opus" assert itpm_descriptor["rate_limit"]["tokens_per_unit"] == 20000000 assert otpm_descriptor["value"] == "proj-mantle:bedrock_mantle/claude-opus" @@ -3173,9 +3025,7 @@ async def test_project_model_itpm_otpm_limits_not_triggered_for_other_model_v3() """Split project limits must not apply to an unrelated model.""" _api_key = hash_token("sk-project-io-test-2") local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) captured_descriptors = [] @@ -3210,9 +3060,7 @@ async def test_project_model_itpm_and_tpm_limits_coexist_v3(): """Combined project TPM and split ITPM/OTPM limits are enforced together.""" _api_key = hash_token("sk-project-io-test-3") local_cache = DualCache() - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) captured_descriptors = [] @@ -3254,9 +3102,7 @@ async def test_enforce_project_io_token_quota_for_frame_blocks_over_limit_otpm() limit and reject once a frame's estimated output tokens exceed it.""" _api_key = hash_token("sk-ws-frame-otpm") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) user_api_key_dict = UserAPIKeyAuth( api_key=_api_key, project_id="proj-mantle-ws", @@ -3288,9 +3134,7 @@ async def test_enforce_project_io_token_quota_for_frame_noop_without_project_lim the per-frame check (no descriptors to reserve against).""" _api_key = hash_token("sk-ws-frame-no-limits") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) user_api_key_dict = UserAPIKeyAuth(api_key=_api_key) await handler.enforce_project_io_token_quota_for_frame( @@ -3425,17 +3269,13 @@ async def test_chat_tpm_refund_and_slot_release_via_context_stash(monkeypatch): monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) _api_key = hash_token("sk-refund-lifecycle") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) user_api_key_dict = UserAPIKeyAuth( api_key=_api_key, tpm_limit=10_000, max_parallel_requests=2, ) - tokens_key = handler.create_rate_limit_keys( - key="api_key", value=_api_key, rate_limit_type="tokens" - ) + tokens_key = handler.create_rate_limit_keys(key="api_key", value=_api_key, rate_limit_type="tokens") parallel_key = f"{{api_key:{_api_key}}}:max_parallel_requests" await handler.async_pre_call_hook( @@ -3452,28 +3292,18 @@ async def test_chat_tpm_refund_and_slot_release_via_context_stash(monkeypatch): reserved = get_request_stash().reserved_tokens assert reserved > 0 assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == reserved - assert handler._gauge_in_flight_from_cache_value( - await local_cache.async_get_cache(key=parallel_key) - ) == 1 + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=parallel_key)) == 1 kwargs = {"standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}} - await handler.async_log_failure_event( - kwargs=kwargs, response_obj=None, start_time=None, end_time=None - ) + await handler.async_log_failure_event(kwargs=kwargs, response_obj=None, start_time=None, end_time=None) assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == 0 - assert handler._gauge_in_flight_from_cache_value( - await local_cache.async_get_cache(key=parallel_key) - ) == 0 + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=parallel_key)) == 0 assert get_request_stash().reservation_released is True - await handler.async_log_failure_event( - kwargs=kwargs, response_obj=None, start_time=None, end_time=None - ) + await handler.async_log_failure_event(kwargs=kwargs, response_obj=None, start_time=None, end_time=None) assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == 0 - assert handler._gauge_in_flight_from_cache_value( - await local_cache.async_get_cache(key=parallel_key) - ) == 0 + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=parallel_key)) == 0 @pytest.mark.asyncio @@ -3517,9 +3347,7 @@ async def test_pre_call_hook_ignores_caller_supplied_stash_values(): async def spy_increment_pipeline(increment_list, **kwargs): refund_calls.append(increment_list) - handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = ( - spy_increment_pipeline - ) + handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = spy_increment_pipeline await handler.async_post_call_failure_hook( request_data=data, @@ -3545,17 +3373,13 @@ async def test_log_events_from_nested_calls_leave_owner_stash_alone(monkeypatch) monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) _api_key = hash_token("sk-nested-guard") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) user_api_key_dict = UserAPIKeyAuth( api_key=_api_key, tpm_limit=10_000, max_parallel_requests=2, ) - tokens_key = handler.create_rate_limit_keys( - key="api_key", value=_api_key, rate_limit_type="tokens" - ) + tokens_key = handler.create_rate_limit_keys(key="api_key", value=_api_key, rate_limit_type="tokens") parallel_key = f"{{api_key:{_api_key}}}:max_parallel_requests" await handler.async_pre_call_hook( @@ -3580,33 +3404,23 @@ async def test_log_events_from_nested_calls_leave_owner_stash_alone(monkeypatch) "litellm_call_id": "nested-guardrail-call", "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}, } - await handler.async_log_success_event( - kwargs=nested_kwargs, response_obj=None, start_time=None, end_time=None - ) - await handler.async_log_failure_event( - kwargs=nested_kwargs, response_obj=None, start_time=None, end_time=None - ) + await handler.async_log_success_event(kwargs=nested_kwargs, response_obj=None, start_time=None, end_time=None) + await handler.async_log_failure_event(kwargs=nested_kwargs, response_obj=None, start_time=None, end_time=None) assert stash.parallel_slot is not None assert stash.reservation_released is False - assert handler._gauge_in_flight_from_cache_value( - await local_cache.async_get_cache(key=parallel_key) - ) == 1 + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=parallel_key)) == 1 assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == reserved owner_kwargs = { "litellm_call_id": "owner-call-id", "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}, } - await handler.async_log_failure_event( - kwargs=owner_kwargs, response_obj=None, start_time=None, end_time=None - ) + await handler.async_log_failure_event(kwargs=owner_kwargs, response_obj=None, start_time=None, end_time=None) assert stash.parallel_slot is None assert stash.reservation_released is True - assert handler._gauge_in_flight_from_cache_value( - await local_cache.async_get_cache(key=parallel_key) - ) == 0 + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=parallel_key)) == 0 assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == 0 @@ -3620,9 +3434,7 @@ async def test_stash_applies_when_owner_or_callback_call_id_missing(): do not thread ``litellm_call_id`` into their logging kwargs. """ local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) unclaimed = get_or_create_request_stash() unclaimed.reserved_tokens = 42 @@ -3634,9 +3446,7 @@ async def test_stash_applies_when_owner_or_callback_call_id_missing(): ) assert unclaimed.reservation_released is True - claimed = RequestRateLimiterStash( - owner_litellm_call_id="owner-1", reserved_tokens=42 - ) + claimed = RequestRateLimiterStash(owner_litellm_call_id="owner-1", reserved_tokens=42) _request_stash.set(claimed) await handler.async_log_failure_event( kwargs={"standard_logging_object": {}}, @@ -3819,9 +3629,7 @@ async def test_failure_event_settles_project_itpm_otpm_at_recovered_partial_usag def _make_mcp_handler(): local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) return handler, local_cache @@ -3848,9 +3656,7 @@ def test_mcp_per_key_descriptor_created_for_matching_server_v3(): metadata={"mcp_rpm_limit": {"github": 5}}, ) - descriptors = _build_mcp_descriptors( - handler, user_api_key_dict, {"mcp_server_name": "github"} - ) + descriptors = _build_mcp_descriptors(handler, user_api_key_dict, {"mcp_server_name": "github"}) descriptor = _find_descriptor(descriptors, "mcp_per_key") assert descriptor is not None @@ -3868,9 +3674,7 @@ def test_mcp_per_key_descriptor_skipped_for_non_matching_server_v3(): metadata={"mcp_rpm_limit": {"github": 5}}, ) - descriptors = _build_mcp_descriptors( - handler, user_api_key_dict, {"mcp_server_name": "slack"} - ) + descriptors = _build_mcp_descriptors(handler, user_api_key_dict, {"mcp_server_name": "slack"}) assert _find_descriptor(descriptors, "mcp_per_key") is None @@ -3927,9 +3731,7 @@ def test_mcp_per_team_descriptor_created_from_team_metadata_v3(): team_metadata={"mcp_rpm_limit": {"github": 3}}, ) - descriptors = _build_mcp_descriptors( - handler, user_api_key_dict, {"mcp_server_name": "github"} - ) + descriptors = _build_mcp_descriptors(handler, user_api_key_dict, {"mcp_server_name": "github"}) descriptor = _find_descriptor(descriptors, "mcp_per_team") assert descriptor is not None @@ -3948,9 +3750,7 @@ async def test_mcp_per_key_rpm_enforced_v3(monkeypatch): monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60") api_key = hash_token("sk-mcp-enforce") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) window_starts: Dict[str, int] = {} request_counts: Dict[str, int] = {} @@ -4044,9 +3844,7 @@ def test_get_key_mcp_rpm_limit_precedence(): _TEST_SLOT_ID = "slot-disconnect-test" -async def _seed_max_parallel_requests_slots( - dual_cache: DualCache, counter_key: str, slot_ids: List[str] -) -> None: +async def _seed_max_parallel_requests_slots(dual_cache: DualCache, counter_key: str, slot_ids: List[str]) -> None: await dual_cache.async_set_cache( key=counter_key, value={slot_id: time.time() for slot_id in slot_ids}, @@ -4058,9 +3856,7 @@ async def _build_seeded_limiter(): """Build a v3 limiter whose api-key slot registry already holds the pre-call slot.""" api_key = hash_token("sk-disconnect") cache = DualCache() - limiter = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(cache) - ) + limiter = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache)) counter_key = f"{{api_key:{api_key}}}:max_parallel_requests" await _seed_max_parallel_requests_slots(cache, counter_key, [_TEST_SLOT_ID]) user_api_key_dict = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=2) @@ -4095,16 +3891,12 @@ async def test_release_max_parallel_requests_on_disconnect_v3(): """ _api_key = hash_token("sk-12345") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=2) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" await _seed_max_parallel_requests_slots(local_cache, counter_key, [_TEST_SLOT_ID]) - assert handler._gauge_in_flight_from_cache_value( - await local_cache.async_get_cache(key=counter_key) - ) == 1 + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 1 get_or_create_request_stash().parallel_slot = ParallelSlotAcquisition( slot_id=_TEST_SLOT_ID, @@ -4113,9 +3905,7 @@ async def test_release_max_parallel_requests_on_disconnect_v3(): await handler.async_release_max_parallel_requests_on_disconnect(user_api_key_dict) assert get_request_stash().parallel_slot is None - assert handler._gauge_in_flight_from_cache_value( - await local_cache.async_get_cache(key=counter_key) - ) == 0 + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 0 @pytest.mark.asyncio @@ -4128,9 +3918,7 @@ async def test_release_on_disconnect_works_when_key_config_changed_v3(): """ _api_key = hash_token("sk-12345") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" await _seed_max_parallel_requests_slots(local_cache, counter_key, [_TEST_SLOT_ID]) @@ -4141,9 +3929,7 @@ async def test_release_on_disconnect_works_when_key_config_changed_v3(): await handler.async_release_max_parallel_requests_on_disconnect( UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=None) ) - assert handler._gauge_in_flight_from_cache_value( - await local_cache.async_get_cache(key=counter_key) - ) == 0 + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 0 @pytest.mark.asyncio @@ -4159,9 +3945,7 @@ async def test_post_call_failure_hook_releases_parallel_slot_v3(): """ _api_key = hash_token("sk-12345") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" @@ -4172,18 +3956,14 @@ async def test_post_call_failure_hook_releases_parallel_slot_v3(): data=admitted_data, call_type="", ) - assert handler._gauge_in_flight_from_cache_value( - await local_cache.async_get_cache(key=counter_key) - ) == 1 + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 1 await handler.async_post_call_failure_hook( request_data=admitted_data, original_exception=Exception("guardrail rejected the request"), user_api_key_dict=user_api_key_dict, ) - assert handler._gauge_in_flight_from_cache_value( - await local_cache.async_get_cache(key=counter_key) - ) == 0 + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 0 await handler.async_log_failure_event( kwargs={ @@ -4193,9 +3973,7 @@ async def test_post_call_failure_hook_releases_parallel_slot_v3(): start_time=None, end_time=None, ) - assert handler._gauge_in_flight_from_cache_value( - await local_cache.async_get_cache(key=counter_key) - ) == 0 + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 0 await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, @@ -4214,9 +3992,7 @@ async def test_success_event_releases_parallel_slot_v3(monkeypatch): """ _api_key = hash_token("sk-12345") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) monkeypatch.setattr(handler, "get_rate_limit_type", lambda: "total") user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" @@ -4228,23 +4004,17 @@ async def test_success_event_releases_parallel_slot_v3(monkeypatch): data=admitted_data, call_type="", ) - assert handler._gauge_in_flight_from_cache_value( - await local_cache.async_get_cache(key=counter_key) - ) == 1 + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 1 await handler.async_log_success_event( kwargs={ "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}, }, - response_obj=ModelResponse( - usage=Usage(prompt_tokens=5, completion_tokens=5, total_tokens=10) - ), + response_obj=ModelResponse(usage=Usage(prompt_tokens=5, completion_tokens=5, total_tokens=10)), start_time=datetime.now(), end_time=datetime.now(), ) - assert handler._gauge_in_flight_from_cache_value( - await local_cache.async_get_cache(key=counter_key) - ) == 0 + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 0 await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, @@ -4264,9 +4034,7 @@ async def test_read_only_gauge_check_counts_without_acquiring_v3(): """ _api_key = hash_token("sk-12345") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" descriptors = [ { @@ -4285,9 +4053,7 @@ async def test_read_only_gauge_check_counts_without_acquiring_v3(): handler.parallel_count_script = fake_count response = await handler.should_rate_limit(descriptors=descriptors, read_only=True) - assert captured_calls == [ - ([counter_key], [PARALLEL_REQUEST_SLOT_TTL_SECONDS]) - ] + assert captured_calls == [([counter_key], [PARALLEL_REQUEST_SLOT_TTL_SECONDS])] assert response["overall_code"] == "OK" assert response["statuses"] == [ { @@ -4304,9 +4070,7 @@ async def test_read_only_gauge_check_counts_without_acquiring_v3(): raise ConnectionError("redis unavailable") handler.parallel_count_script = failing_count - await _seed_max_parallel_requests_slots( - local_cache, counter_key, ["s1", "s2", "s3", "s4", "s5"] - ) + await _seed_max_parallel_requests_slots(local_cache, counter_key, ["s1", "s2", "s3", "s4", "s5"]) response = await handler.should_rate_limit(descriptors=descriptors, read_only=True) assert response["overall_code"] == "OVER_LIMIT" assert response["statuses"][0]["rate_limit_type"] == "max_parallel_requests" @@ -4321,9 +4085,7 @@ async def test_redis_release_script_updates_local_mirror_v3(): """ _api_key = hash_token("sk-12345") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" captured_calls = [] @@ -4361,12 +4123,8 @@ async def test_tpm_over_limit_rejection_releases_parallel_slot_v3(monkeypatch): monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) _api_key = hash_token("sk-12345") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) - user_api_key_dict = UserAPIKeyAuth( - api_key=_api_key, max_parallel_requests=5, tpm_limit=100 - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=5, tpm_limit=100) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" async def over_limit_reservation(descriptors, estimated_tokens, parent_otel_span=None): @@ -4393,9 +4151,7 @@ async def test_tpm_over_limit_rejection_releases_parallel_slot_v3(monkeypatch): call_type="", ) assert exc_info.value.status_code == 429 - assert handler._gauge_in_flight_from_cache_value( - await local_cache.async_get_cache(key=counter_key) - ) == 0 + assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 0 @pytest.mark.asyncio @@ -4410,9 +4166,7 @@ async def test_in_memory_fallback_respects_mirrored_redis_count_v3(): """ _api_key = hash_token("sk-12345") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=5) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" @@ -4462,9 +4216,7 @@ async def test_release_max_parallel_requests_on_disconnect_noop_v3(): """ _api_key = hash_token("sk-12345") local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" await handler.async_release_max_parallel_requests_on_disconnect( @@ -4495,9 +4247,7 @@ async def test_async_streaming_data_generator_releases_counter_on_disconnect_v3( from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing limiter, cache, counter_key, user_api_key_dict = await _build_seeded_limiter() - assert limiter._gauge_in_flight_from_cache_value( - await cache.async_get_cache(key=counter_key) - ) == 1 + assert limiter._gauge_in_flight_from_cache_value(await cache.async_get_cache(key=counter_key)) == 1 proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) proxy_logging_obj.proxy_hook_mapping["parallel_request_limiter"] = limiter @@ -4528,9 +4278,7 @@ async def test_async_streaming_data_generator_releases_counter_on_disconnect_v3( await gen.aclose() await _drain_release_task() - assert limiter._gauge_in_flight_from_cache_value( - await cache.async_get_cache(key=counter_key) - ) == 0 + assert limiter._gauge_in_flight_from_cache_value(await cache.async_get_cache(key=counter_key)) == 0 @pytest.mark.parametrize("disconnect", ["cancel", "aclose"]) @@ -4577,14 +4325,10 @@ async def test_async_data_generator_releases_counter_on_disconnect_v3(disconnect else: await gen.aclose() await _drain_release_task() - assert limiter._gauge_in_flight_from_cache_value( - await cache.async_get_cache(key=counter_key) - ) == 0 + assert limiter._gauge_in_flight_from_cache_value(await cache.async_get_cache(key=counter_key)) == 0 finally: if saved_hook is not None: - proxy_logging_obj.proxy_hook_mapping["parallel_request_limiter"] = ( - saved_hook - ) + proxy_logging_obj.proxy_hook_mapping["parallel_request_limiter"] = saved_hook else: proxy_logging_obj.proxy_hook_mapping.pop("parallel_request_limiter", None) @@ -4602,9 +4346,7 @@ async def test_async_data_generator_releases_counter_when_wrapped_v3(): import litellm.proxy.proxy_server as proxy_server class _PassthroughIteratorOverride(CustomLogger): - async def async_post_call_streaming_iterator_hook( - self, user_api_key_dict, response, request_data - ): + async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data): async for chunk in response: yield chunk @@ -4632,14 +4374,10 @@ async def test_async_data_generator_releases_counter_when_wrapped_v3(): await gen.__anext__() await gen.aclose() await _drain_release_task() - assert limiter._gauge_in_flight_from_cache_value( - await cache.async_get_cache(key=counter_key) - ) == 0 + assert limiter._gauge_in_flight_from_cache_value(await cache.async_get_cache(key=counter_key)) == 0 finally: if saved_hook is not None: - proxy_logging_obj.proxy_hook_mapping["parallel_request_limiter"] = ( - saved_hook - ) + proxy_logging_obj.proxy_hook_mapping["parallel_request_limiter"] = saved_hook else: proxy_logging_obj.proxy_hook_mapping.pop("parallel_request_limiter", None) @@ -4647,18 +4385,14 @@ async def test_async_data_generator_releases_counter_when_wrapped_v3(): def test_tpm_reservation_enabled_by_default(monkeypatch): """Upfront TPM reservation is on unless explicitly disabled via env.""" monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) assert handler.tpm_reservation_enabled is True @pytest.mark.parametrize("value", ["false", "False", "FALSE"]) def test_tpm_reservation_disabled_via_env(monkeypatch, value): monkeypatch.setenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", value) - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) assert handler.tpm_reservation_enabled is False @@ -4670,9 +4404,7 @@ async def test_pre_call_hook_reserves_tpm_when_enabled(monkeypatch): only the reservation path owns it. """ monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) user_api_key_dict = UserAPIKeyAuth(api_key=hash_token("sk-tpm"), tpm_limit=10_000) @@ -4712,9 +4444,7 @@ async def test_pre_call_hook_skips_reservation_when_disabled(monkeypatch): pre-v1.82 post-call accounting behavior. """ monkeypatch.setenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", "false") - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) user_api_key_dict = UserAPIKeyAuth(api_key=hash_token("sk-tpm"), tpm_limit=10_000) @@ -4762,9 +4492,7 @@ async def test_per_tag_rate_limit_independent_counters_v3(monkeypatch): metadata={"tag_rpm_limit": {"cell-1": 2}}, ) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) async def call(tag: str) -> None: await handler.async_pre_call_hook( @@ -4798,9 +4526,7 @@ async def test_per_tag_descriptor_creation_v3(): api_key=_api_key, metadata={"tag_rpm_limit": {"cell-1": 5}}, ) - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) descriptors = handler._create_rate_limit_descriptors( user_api_key_dict=user_api_key_dict, @@ -4824,9 +4550,7 @@ async def test_per_tag_descriptor_absent_without_config_v3(): api_key=hash_token("sk-no-tag"), rpm_limit=10, ) - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) descriptors = handler._create_rate_limit_descriptors( user_api_key_dict=user_api_key_dict, @@ -4857,9 +4581,7 @@ async def test_per_tag_untagged_request_governed_by_key_limit_v3(monkeypatch): metadata={"tag_rpm_limit": {"cell-1": 2}}, ) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) async def call(metadata: dict) -> None: await handler.async_pre_call_hook( @@ -4904,9 +4626,7 @@ async def test_streaming_end_to_end_populates_slp_ratelimit_headers(monkeypatch) tpm_limit=10000, ) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Real pre-call: populates data and stashes the response into metadata # so the success callback can find it via litellm_params.metadata. @@ -4963,27 +4683,16 @@ async def test_streaming_end_to_end_populates_slp_ratelimit_headers(monkeypatch) call_type="acompletion", ) - additional_headers = ( - mock_kwargs["standard_logging_object"] - .get("hidden_params", {}) - .get("additional_headers", {}) - ) + additional_headers = mock_kwargs["standard_logging_object"].get("hidden_params", {}).get("additional_headers", {}) # api_key-scoped remaining/limit values are the baseline every request # emits and must always reach the SLP. - remaining_keys = [ - k for k in additional_headers if "-remaining-" in k - ] - assert ( - remaining_keys - ), f"streaming success must populate remaining values, got {additional_headers!r}" + remaining_keys = [k for k in additional_headers if "-remaining-" in k] + assert remaining_keys, f"streaming success must populate remaining values, got {additional_headers!r}" limit_keys = [k for k in additional_headers if "-limit-" in k] assert limit_keys, "streaming success must also populate limit values" - assert ( - additional_headers.get("x-ratelimit-api_key-remaining-requests") == 99 - ), ( - "api_key remaining requests should reflect the just-consumed slot;" - f" got {additional_headers!r}" + assert additional_headers.get("x-ratelimit-api_key-remaining-requests") == 99, ( + f"api_key remaining requests should reflect the just-consumed slot; got {additional_headers!r}" ) @@ -5003,9 +4712,7 @@ async def test_streaming_populates_model_per_key_ratelimit_headers(monkeypatch): }, ) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) async def _noop_increment(increment_list, **_): return True @@ -5055,9 +4762,7 @@ async def test_streaming_populates_model_per_key_ratelimit_headers(monkeypatch): hidden_params = mock_kwargs["standard_logging_object"].get("hidden_params") or {} additional_headers = hidden_params.get("additional_headers") or {} - assert ( - additional_headers.get("x-ratelimit-model_per_key-remaining-requests") == 99 - ), f"got {additional_headers!r}" + assert additional_headers.get("x-ratelimit-model_per_key-remaining-requests") == 99, f"got {additional_headers!r}" assert additional_headers.get("x-ratelimit-model_per_key-limit-requests") == 100 # response._hidden_params is also updated for late readers. @@ -5074,9 +4779,7 @@ async def test_async_log_success_event_no_mirror_when_no_snapshot(monkeypatch): """ monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60") _api_key = hash_token("sk-stream-no-mirror") - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) async def _noop_increment(increment_list, **_): return True @@ -5118,9 +4821,7 @@ async def test_async_log_success_event_no_mirror_when_no_snapshot(monkeypatch): hidden_params = mock_kwargs["standard_logging_object"].get("hidden_params") or {} additional_headers = hidden_params.get("additional_headers") or {} ratelimit_keys = [k for k in additional_headers if k.startswith("x-ratelimit-")] - assert ( - not ratelimit_keys - ), f"no snapshot must produce no rate-limit headers, got {ratelimit_keys}" + assert not ratelimit_keys, f"no snapshot must produce no rate-limit headers, got {ratelimit_keys}" @pytest.mark.asyncio @@ -5140,9 +4841,7 @@ async def test_streaming_mirror_matches_non_streaming_header_shape(monkeypatch): }, ) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) async def _noop_increment(increment_list, **_): return True @@ -5177,15 +4876,11 @@ async def test_streaming_mirror_matches_non_streaming_header_shape(monkeypatch): user_api_key_dict=user_api_key_dict, response=non_stream_response, ) - non_stream_headers = non_stream_response._hidden_params.get( - "additional_headers", {} - ) + non_stream_headers = non_stream_response._hidden_params.get("additional_headers", {}) # Streaming path: async_logging_hook mirrors into standard_logging_object. stream_kwargs: Dict[str, Any] = { - "standard_logging_object": { - "metadata": {"user_api_key_hash": _api_key} - }, + "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}, "litellm_params": {"metadata": data["metadata"]}, "model": "gpt-4o-mini", } @@ -5202,18 +4897,13 @@ async def test_streaming_mirror_matches_non_streaming_header_shape(monkeypatch): result=stream_response, call_type="acompletion", ) - stream_slp_headers = ( - stream_kwargs["standard_logging_object"] - .get("hidden_params", {}) - .get("additional_headers", {}) - ) + stream_slp_headers = stream_kwargs["standard_logging_object"].get("hidden_params", {}).get("additional_headers", {}) def _rl_only(headers: Dict[str, Any]) -> Dict[str, Any]: return {k: v for k, v in headers.items() if k.startswith("x-ratelimit-")} assert _rl_only(stream_slp_headers) == _rl_only(non_stream_headers), ( - f"streaming={_rl_only(stream_slp_headers)}" - f" non_streaming={_rl_only(non_stream_headers)}" + f"streaming={_rl_only(stream_slp_headers)} non_streaming={_rl_only(non_stream_headers)}" ) assert "x-ratelimit-model_per_key-remaining-requests" in stream_slp_headers @@ -5229,12 +4919,8 @@ async def test_async_log_success_event_counts_passthrough_reported_tokens(monkey monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60") _api_key = hash_token("sk-passthrough") - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) - monkeypatch.setattr( - parallel_request_handler, "get_rate_limit_type", lambda: "total" - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + monkeypatch.setattr(parallel_request_handler, "get_rate_limit_type", lambda: "total") captured_operations = [] @@ -5270,9 +4956,7 @@ async def test_async_log_success_event_counts_passthrough_reported_tokens(monkey @pytest.mark.parametrize("rate_limit_type", ["input", "output", "total"]) @pytest.mark.asyncio -async def test_aggregate_only_usage_charges_tpm_under_every_limit_type( - monkeypatch, rate_limit_type -): +async def test_aggregate_only_usage_charges_tpm_under_every_limit_type(monkeypatch, rate_limit_type): """ A pass-through target reports one total for the whole request and cannot split it into prompt/completion. Reading a split out of it yields 0, which @@ -5282,12 +4966,8 @@ async def test_aggregate_only_usage_charges_tpm_under_every_limit_type( monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60") _api_key = hash_token("sk-aggregate-only") - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) - monkeypatch.setattr( - parallel_request_handler, "get_rate_limit_type", lambda: rate_limit_type - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + monkeypatch.setattr(parallel_request_handler, "get_rate_limit_type", lambda: rate_limit_type) captured_operations = [] @@ -5322,9 +5002,7 @@ async def test_split_usage_still_respects_the_configured_limit_type(monkeypatch) monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60") _api_key = hash_token("sk-split-usage") - parallel_request_handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) + parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) monkeypatch.setattr(parallel_request_handler, "get_rate_limit_type", lambda: "output") captured_operations = [] @@ -5366,9 +5044,7 @@ async def test_atomic_check_with_zero_increment_still_enforces_token_limit(): """ from litellm.proxy.hooks.parallel_request_limiter_v3 import RateLimitDescriptor - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) descriptor = RateLimitDescriptor( key="model_saturation_check", value="tpm-only-model", @@ -5383,15 +5059,9 @@ async def test_atomic_check_with_zero_increment_still_enforces_token_limit(): assert under_limit["overall_code"] == "OK" assert [s["rate_limit_type"] for s in under_limit["statuses"]] == ["tokens"] - counter_key = handler.create_rate_limit_keys( - "model_saturation_check", "tpm-only-model", "tokens" - ) + counter_key = handler.create_rate_limit_keys("model_saturation_check", "tpm-only-model", "tokens") await handler.async_increment_tokens_with_ttl_preservation( - pipeline_operations=[ - RedisPipelineIncrementOperation( - key=counter_key, increment_value=100, ttl=60 - ) - ], + pipeline_operations=[RedisPipelineIncrementOperation(key=counter_key, increment_value=100, ttl=60)], ) at_limit = await handler.atomic_check_and_increment_by_n( @@ -5436,9 +5106,7 @@ async def test_reserve_tpm_tokens_never_evaluates_the_requests_dimension(): its descriptors or an exhausted RPM budget would double-enforce here.""" from litellm.proxy.hooks.parallel_request_limiter_v3 import RateLimitDescriptor - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) descriptor = RateLimitDescriptor( key="api_key", value="reserve-test-key", @@ -5450,8 +5118,7 @@ async def test_reserve_tpm_tokens_never_evaluates_the_requests_dimension(): estimated_tokens=10, ) assert response["overall_code"] == "OK", ( - "an exhausted requests budget (limit 0) must not block the token " - f"reservation pass, got: {response}" + f"an exhausted requests budget (limit 0) must not block the token reservation pass, got: {response}" ) assert [s["rate_limit_type"] for s in response["statuses"]] == ["tokens"] @@ -5553,9 +5220,7 @@ async def test_estimated_output_tokens_resolution_precedence( """ monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) user_api_key_dict = UserAPIKeyAuth( api_key=hash_token(f"sk-estimate-{expected_output_estimate}-{tier}"), tpm_limit=1_000_000, @@ -5578,9 +5243,7 @@ async def test_request_max_tokens_outranks_configured_estimate(monkeypatch): """An explicit request-level max_tokens stays the top of the precedence order.""" monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) user_api_key_dict = UserAPIKeyAuth( api_key=hash_token("sk-estimate-explicit-max-tokens"), tpm_limit=1_000_000, @@ -5602,9 +5265,7 @@ async def test_configured_estimate_does_not_apply_to_embeddings(monkeypatch): """Embeddings generate no output, so a declared output estimate must not be reserved.""" monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) user_api_key_dict = UserAPIKeyAuth( api_key=hash_token("sk-estimate-embeddings"), tpm_limit=1_000_000, @@ -5631,9 +5292,7 @@ async def test_configured_estimate_applies_to_contentless_requests(monkeypatch): """ monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) configured = UserAPIKeyAuth( api_key=hash_token("sk-estimate-contentless-configured"), tpm_limit=1_000_000, @@ -5645,17 +5304,9 @@ async def test_configured_estimate_applies_to_contentless_requests(monkeypatch): ) assert ( - await _reserved_tokens_for( - handler, local_cache, configured, {"model": "gpt-4o-mini", "messages": []} - ) - == 2002 - ) - assert ( - await _reserved_tokens_for( - handler, local_cache, unconfigured, {"model": "gpt-4o-mini", "messages": []} - ) - == 1 + await _reserved_tokens_for(handler, local_cache, configured, {"model": "gpt-4o-mini", "messages": []}) == 2002 ) + assert await _reserved_tokens_for(handler, local_cache, unconfigured, {"model": "gpt-4o-mini", "messages": []}) == 1 @pytest.mark.asyncio @@ -5672,9 +5323,7 @@ async def test_declared_estimate_never_tightens_the_small_tpm_clamp(monkeypatch) """ monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) raised_data: Dict[str, Any] = {"model": "gpt-4o-mini", "messages": ONE_TOKEN_PROMPT} raised_reserved = await _reserved_tokens_for( handler, @@ -5726,9 +5375,7 @@ async def test_one_malformed_estimate_field_does_not_discard_the_other(monkeypat """ monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) broken_map = await _reserved_tokens_for( handler, @@ -5777,9 +5424,7 @@ async def test_declared_estimate_over_the_tpm_budget_is_honored_and_explained(mo """ monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) user_api_key_dict = UserAPIKeyAuth( api_key=hash_token(f"sk-estimate-over-budget-{declared}"), tpm_limit=5000, @@ -5818,9 +5463,7 @@ async def test_a_key_that_declared_nothing_is_never_blamed_for_a_declaration(mon """ monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): with pytest.raises(HTTPException) as exc_info: @@ -5848,9 +5491,7 @@ async def test_declared_estimate_inside_the_tpm_budget_is_not_explained(monkeypa """ monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): await handler.async_pre_call_hook( @@ -5883,9 +5524,7 @@ async def test_configured_estimate_blocks_the_overrun_the_static_floor_admits(mo async def admitted(metadata): local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) user_api_key_dict = UserAPIKeyAuth( api_key=hash_token(f"sk-overrun-{metadata}"), tpm_limit=8000, @@ -5913,9 +5552,7 @@ def test_internal_call_origin_success_ops_are_skipped(): """Internal sub-calls (auto-router classifier, shadow eval shadow/judge) bill spend to the caller's key but must not consume its TPM counters: the same kwargs charge ops without the origin stamp and none with it.""" - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) response = ModelResponse( id="internal-origin-tpm", object="chat.completion", @@ -5927,9 +5564,7 @@ def test_internal_call_origin_success_ops_are_skipped(): def _kwargs(metadata: Dict[str, Any]) -> Dict[str, Any]: return { - "standard_logging_object": { - "metadata": {"user_api_key_hash": hash_token("sk-internal-origin")} - }, + "standard_logging_object": {"metadata": {"user_api_key_hash": hash_token("sk-internal-origin")}}, "litellm_params": {"metadata": metadata}, "model": "gpt-4o-mini", } @@ -5963,15 +5598,10 @@ def test_conflicting_token_limits_reserve_the_larger_declared_budget(): A request declaring max_tokens=1 alongside max_completion_tokens=10000 previously reserved one output token while the provider stayed free to emit ten thousand. """ - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) bodies = _conflicting_budget_bodies() - reserved = { - label: handler._estimate_tokens_for_request(data=body) - for label, body in bodies.items() - } + reserved = {label: handler._estimate_tokens_for_request(data=body) for label, body in bodies.items()} assert reserved["both"] == reserved["only_large"] assert reserved["both"] > reserved["only_small"] @@ -5985,9 +5615,7 @@ def test_non_integer_output_budgets_still_reserve_their_declared_size(declared): output floor, so dropping it from the estimate under-reserves and reopens the same TPM bypass that reading both spellings was meant to close. """ - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) base = {"model": "gpt-5-chat", "messages": [{"role": "user", "content": "hi"}]} reserved = handler._estimate_tokens_for_request(data={**base, "max_tokens": declared}) @@ -6000,12 +5628,8 @@ def test_non_integer_output_budgets_still_reserve_their_declared_size(declared): async def test_conflicting_token_limits_cannot_bypass_tpm_reservation(): """The pre-call hook must refuse a request whose larger declared budget exceeds the TPM limit.""" local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(local_cache) - ) - user_api_key_dict = UserAPIKeyAuth( - api_key=hash_token("sk-conflicting-budgets"), tpm_limit=100, models=[] - ) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + user_api_key_dict = UserAPIKeyAuth(api_key=hash_token("sk-conflicting-budgets"), tpm_limit=100, models=[]) bodies = _conflicting_budget_bodies() await handler.async_pre_call_hook( @@ -6032,7 +5656,9 @@ async def test_conflicting_token_limits_cannot_bypass_tpm_reservation(): def _enqueued_test_handler() -> _PROXY_MaxParallelRequestsHandler: - return _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache(default_in_memory_ttl=60))) + return _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache(default_in_memory_ttl=60)) + ) def _batch_response(batch_id: str, status: str): @@ -6238,6 +5864,201 @@ async def test_post_call_success_hook_leaves_raw_provider_dict_untouched(): assert response == {"id": "msg_123", "type": "message", "role": "assistant", "content": []} +class _ClusterParallelTransport: + def __init__(self, handler): + self.handler = handler + self.members = {} + self.now = 10000 + self.calls = [] + self.fail = None + self.lose_acquire_response = None + self.pause_acquire = None + self.entered = asyncio.Event() + + def script(self, operation): + async def run(*, keys, args): + self.calls.append((operation, tuple(keys))) + if len({self.handler.keyslot_for_redis_cluster(key) for key in keys}) > 1: + raise RuntimeError("CROSSSLOT Keys in request do not hash to the same slot") + if self.fail is not None and (operation, keys[0]) == self.fail: + raise RuntimeError("Shard unavailable") + if operation in ("acquire", "count"): + for key in keys: + self.members[key] = { + slot: score + for slot, score in self.members.get(key, {}).items() + if score > self.now - PARALLEL_REQUEST_SLOT_TTL_SECONDS + } + if operation == "acquire": + for index, key in enumerate(keys): + if len(self.members[key]) >= args[index * 3]: + return [1, index + 1, len(self.members[key])] + for index, key in enumerate(keys): + self.members[key][args[index * 3 + 2]] = self.now + if keys[0] == self.pause_acquire: + self.entered.set() + await asyncio.Event().wait() + if keys[0] == self.lose_acquire_response: + raise RuntimeError("Reply lost after Redis admitted slot") + return [0, *(len(self.members[key]) for key in keys)] + if operation == "count": + return [len(self.members.get(key, {})) for key in keys] + if operation == "renew": + if any(self.members.get(key, {}).get(args[0], 0) <= self.now - args[1] for key in keys): + return [0] + for key in keys: + self.members[key][args[0]] = self.now + return [1] + assert operation == "release" + for index, key in enumerate(keys): + self.members.get(key, {}).pop(args[index], None) + return [len(self.members.get(key, {})) for key in keys] + + return run + + +def _cluster_parallel_fixture(monkeypatch): + handler = _PROXY_MaxParallelRequestsHandler(InternalUsageCache(DualCache())) + monkeypatch.setattr(handler, "_is_redis_cluster", lambda: True) + transport = _ClusterParallelTransport(handler) + for operation in ("acquire", "count", "renew", "release"): + monkeypatch.setattr(handler, f"parallel_{operation}_script", transport.script(operation)) + gauges = [ + {"counter_key": "{api_key:owner}:max_parallel_requests", "limit": 1, "descriptor_key": "api_key"}, + {"counter_key": "{team:group}:max_parallel_requests", "limit": 2, "descriptor_key": "team"}, + {"counter_key": "{api_key:owner}:another-parallel-scope", "limit": 1, "descriptor_key": "extra"}, + ] + return handler, transport, gauges + + +@pytest.mark.asyncio +async def test_cluster_parallel_slots_admit_count_renew_release_across_hash_slots(monkeypatch): + handler, transport, gauges = _cluster_parallel_fixture(monkeypatch) + keys = tuple(gauge["counter_key"] for gauge in gauges) + transport.members[keys[1]] = {"unrelated": transport.now} + result = await handler._check_parallel_request_gauges(gauges, "owner") + assert result["overall_code"] == "OK" + assert len([call for call in transport.calls if call[0] == "acquire"]) == 2 + assert all("owner" in transport.members[key] for key in keys) + result = await handler._check_parallel_request_gauges(gauges, "reader", read_only=True) + assert result["overall_code"] == "OVER_LIMIT" + assert [status["descriptor_key"] for status in result["statuses"]] == ["api_key", "team", "extra"] + transport.now += PARALLEL_REQUEST_SLOT_TTL_SECONDS - 1 + assert await handler._renew_realtime_call_slot("owner", keys) + transport.now += 2 + result = await handler._check_parallel_request_gauges(gauges, "second") + assert result["overall_code"] == "OVER_LIMIT" + acquisition = ParallelSlotAcquisition(slot_id="owner", counter_keys=list(keys)) + await handler._release_parallel_request_slots(acquisition) + await handler._release_parallel_request_slots(acquisition) + assert all("owner" not in transport.members[key] for key in keys) + assert not await handler._renew_realtime_call_slot("owner", keys) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", ["limit", "unreachable", "lost_reply", "cancel"]) +async def test_cluster_parallel_acquire_rolls_back_attempted_shards_without_releasing_others(monkeypatch, failure): + handler, transport, gauges = _cluster_parallel_fixture(monkeypatch) + keys = tuple(gauge["counter_key"] for gauge in gauges) + transport.members[keys[1]] = {"unrelated": transport.now} + if failure == "limit": + gauges[1]["limit"] = 1 + elif failure == "unreachable": + transport.fail = ("acquire", keys[1]) + elif failure == "lost_reply": + transport.lose_acquire_response = keys[1] + else: + transport.pause_acquire = keys[1] + task = asyncio.create_task(handler._check_parallel_request_gauges(gauges, "owner")) + if failure == "cancel": + await asyncio.wait_for(transport.entered.wait(), timeout=1) + task.cancel() + if failure == "limit": + assert (await task)["overall_code"] == "OVER_LIMIT" + else: + with pytest.raises(asyncio.CancelledError if failure == "cancel" else RuntimeError): + await task + assert all("owner" not in transport.members.get(key, {}) for key in keys) + assert transport.members[keys[1]] == {"unrelated": transport.now} + released_keys = {key for operation, group in transport.calls if operation == "release" for key in group} + assert released_keys == set(keys) + + +@pytest.mark.asyncio +async def test_cluster_parallel_release_continues_after_shard_failure_and_renewal_fails_closed(monkeypatch): + handler, transport, gauges = _cluster_parallel_fixture(monkeypatch) + keys = tuple(gauge["counter_key"] for gauge in gauges) + assert (await handler._check_parallel_request_gauges(gauges, "owner"))["overall_code"] == "OK" + transport.fail = ("renew", keys[1]) + assert not await handler._renew_realtime_call_slot("owner", keys) + transport.fail = None + transport.members[keys[1]].pop("owner") + assert not await handler._renew_realtime_call_slot("owner", keys) + assert "owner" not in transport.members[keys[1]] + transport.fail = ("count", keys[1]) + with pytest.raises(RuntimeError, match="Shard unavailable"): + await handler._check_parallel_request_gauges(gauges, "reader", read_only=True) + transport.fail = ("release", keys[0]) + receipt = ParallelSlotAcquisition(slot_id="owner", counter_keys=list(keys)) + with pytest.raises(RuntimeError, match="Shard unavailable"): + await handler._release_parallel_request_slots(receipt) + assert "owner" not in transport.members[keys[1]] + transport.fail = None + await handler._release_parallel_request_slots(receipt) + assert all("owner" not in transport.members.get(key, {}) for key in keys) + + +@pytest.mark.asyncio +async def test_cluster_parallel_duplicate_scope_keeps_strictest_limit(monkeypatch): + handler, transport, gauges = _cluster_parallel_fixture(monkeypatch) + gauges.append({**gauges[0], "limit": 100}) + transport.members[gauges[0]["counter_key"]] = {"unrelated": transport.now} + assert (await handler._check_parallel_request_gauges(gauges, "owner"))["overall_code"] == "OVER_LIMIT" + assert transport.members[gauges[0]["counter_key"]] == {"unrelated": transport.now} + + +@pytest.mark.asyncio +async def test_cluster_rollback_waits_through_repeated_cancellation(monkeypatch): + handler, transport, gauges = _cluster_parallel_fixture(monkeypatch) + keys = tuple(gauge["counter_key"] for gauge in gauges) + transport.lose_acquire_response = keys[1] + release_entered, finish_release = asyncio.Event(), asyncio.Event() + release = handler.parallel_release_script + + async def blocked_release(*, keys, args): + release_entered.set() + await finish_release.wait() + return await release(keys=keys, args=args) + + monkeypatch.setattr(handler, "parallel_release_script", blocked_release) + task = asyncio.create_task(handler._check_parallel_request_gauges(gauges, "owner")) + await asyncio.wait_for(release_entered.wait(), timeout=1) + try: + for _ in range(3): + task.cancel() + await asyncio.sleep(0) + assert not task.done(), "admission returned while its Redis compensation was still running" + finally: + finish_release.set() + await asyncio.gather(task, return_exceptions=True) + await asyncio.sleep(0) + assert task.cancelled() + assert all("owner" not in transport.members.get(key, {}) for key in keys) + + +@pytest.mark.asyncio +async def test_standalone_realtime_renewal_keeps_single_atomic_batch(monkeypatch): + from unittest.mock import AsyncMock + + handler = _PROXY_MaxParallelRequestsHandler(InternalUsageCache(DualCache())) + monkeypatch.setattr(handler, "_is_redis_cluster", lambda: False) + renew = AsyncMock(return_value=[1]) + monkeypatch.setattr(handler, "parallel_renew_script", renew) + keys = ("{api_key:owner}:max_parallel_requests", "{team:group}:max_parallel_requests") + assert await handler._renew_realtime_call_slot("owner", keys) + renew.assert_awaited_once_with(keys=keys, args=("owner", PARALLEL_REQUEST_SLOT_TTL_SECONDS)) + + class _OpenBreakerRedis: def async_register_script(self, script: str): async def refused(keys, args): diff --git a/tests/test_litellm/proxy/hooks/test_realtime_call_lease.py b/tests/test_litellm/proxy/hooks/test_realtime_call_lease.py new file mode 100644 index 00000000000..cb8ee7fc6b9 --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_realtime_call_lease.py @@ -0,0 +1,85 @@ +import asyncio +from unittest.mock import AsyncMock + +import pytest + +from litellm.proxy.hooks.realtime_call_lease import RealtimeCallLease + + +@pytest.mark.asyncio +async def test_failed_renewal_signals_owner_and_close_releases_once(): + renew = AsyncMock(side_effect=[True, False, True]) + release = AsyncMock() + lease = RealtimeCallLease(renew=renew, release=release, interval=0.001) + lease.start() + await asyncio.wait_for(lease.wait_failed(), timeout=1) + assert renew.await_count == 2 + assert not await lease.renew() + assert renew.await_count == 2 + await asyncio.gather(lease.close(), lease.close()) + assert release.await_count == 1 + + +@pytest.mark.asyncio +async def test_renewal_exception_and_close_before_start(): + release = AsyncMock() + lease = RealtimeCallLease(renew=AsyncMock(side_effect=RuntimeError("backend")), release=release, interval=0.001) + lease.start() + await asyncio.wait_for(lease.wait_failed(), timeout=1) + await lease.close() + assert release.await_count == 1 + unused = RealtimeCallLease(renew=AsyncMock(), release=release) + await unused.close() + assert release.await_count == 2 + + +@pytest.mark.asyncio +async def test_renewal_timeout_signals_failure_without_start(): + lease = RealtimeCallLease(renew=asyncio.Event().wait, release=AsyncMock(), renewal_timeout=0.001) + assert not await lease.renew() + await asyncio.wait_for(lease.wait_failed(), timeout=1) + await lease.close() + + +@pytest.mark.asyncio +async def test_cancelled_close_still_releases_exactly_once(): + entered = asyncio.Event() + finish = asyncio.Event() + + async def release(): + entered.set() + await finish.wait() + + cleanup = AsyncMock(side_effect=release) + lease = RealtimeCallLease(renew=AsyncMock(return_value=True), release=cleanup) + lease.start() + closing = asyncio.create_task(lease.close()) + await asyncio.wait_for(entered.wait(), timeout=1) + closing.cancel() + with pytest.raises(asyncio.CancelledError): + await closing + finish.set() + await lease.close() + assert cleanup.await_count == 1 + + +@pytest.mark.asyncio +async def test_concurrent_renewal_cannot_restore_a_failed_lease(): + pending = asyncio.Event() + entered = asyncio.Event() + + async def delayed_success(): + entered.set() + await pending.wait() + return True + + renew = AsyncMock(side_effect=delayed_success) + lease = RealtimeCallLease(renew=renew, release=AsyncMock()) + first = asyncio.create_task(lease.renew()) + await asyncio.wait_for(entered.wait(), timeout=1) + renew.side_effect = None + renew.return_value = False + assert not await lease.renew() + pending.set() + assert not await first + await lease.close() diff --git a/tests/test_litellm/proxy/realtime_endpoints/test_call_sessions.py b/tests/test_litellm/proxy/realtime_endpoints/test_call_sessions.py new file mode 100644 index 00000000000..dffd4500ed7 --- /dev/null +++ b/tests/test_litellm/proxy/realtime_endpoints/test_call_sessions.py @@ -0,0 +1,1307 @@ +import hashlib +import time +from types import SimpleNamespace + +import pytest +from fastapi import HTTPException, WebSocket + +from litellm.proxy._types import ProxyException, UserAPIKeyAuth +from litellm.proxy.realtime_endpoints import call_sessions as codex +from litellm.llms.chatgpt.codex import CodexRealtimeCall +from litellm.proxy.realtime_endpoints.call_sessions import decode_call, encode_call + + +@pytest.mark.asyncio +@pytest.mark.parametrize("multipart", [False, True]) +@pytest.mark.parametrize("content_length", [None, "1", "999999999"]) +async def test_oversized_offer_stops_before_auth_or_multipart_files(monkeypatch, multipart, content_length): + import json + from unittest.mock import AsyncMock, Mock + + from fastapi import Request + + monkeypatch.setattr(codex, "MAX_REALTIME_OFFER_BYTES", 1024) + if multipart: + body = ( + b'--Boundary\r\nContent-Disposition: form-data; name="extra"; filename="large.bin"\r\n\r\n' + + b"x" * 2048 + + b"\r\n--Boundary--\r\n" + ) + media_type = b"Multipart/Form-Data; boundary=Boundary" + else: + body = json.dumps({"sdp": "x" * 2048, "session": {"model": "voice"}}).encode() + media_type = b"application/json" + chunks = [body[offset : offset + 256] for offset in range(0, len(body), 256)] + received = [] + + async def receive(): + chunk = chunks.pop(0) + received.append(len(chunk)) + return {"type": "http.request", "body": chunk, "more_body": bool(chunks)} + + headers = [(b"content-type", media_type)] + if content_length is not None: + headers.append((b"content-length", content_length.encode())) + request = Request({"type": "http", "headers": headers}, receive) + if multipart: + from litellm.proxy.common_utils.http_parsing_utils import _read_request_body + + assert await _read_request_body(request) == {} + authenticate = AsyncMock() + create_file = Mock(side_effect=AssertionError("Oversized offers must not create temporary files")) + monkeypatch.setattr(codex, "user_api_key_auth", authenticate) + monkeypatch.setattr("starlette.formparsers.SpooledTemporaryFile", create_file) + with pytest.raises(HTTPException) as rejected: + await codex.create_codex_realtime_call(request) + assert rejected.value.status_code == 413 + assert sum(received) <= 1280 + assert chunks + authenticate.assert_not_awaited() + create_file.assert_not_called() + + +@pytest.mark.asyncio +async def test_offer_at_size_limit_keeps_body_available_for_custom_auth(monkeypatch): + import json + + from fastapi import Request + + monkeypatch.setattr(codex, "MAX_REALTIME_OFFER_BYTES", 1024) + empty = {"sdp": "", "session": {"model": "voice"}} + sdp = "x" * (1024 - len(json.dumps(empty).encode())) + body = json.dumps({"sdp": sdp, "session": {"model": "voice"}}).encode() + chunks = [body[:512], body[512:]] + + async def receive(): + return {"type": "http.request", "body": chunks.pop(0), "more_body": bool(chunks)} + + request = Request({"type": "http", "headers": [(b"content-type", b"application/json")]}, receive) + offer = await codex.read_codex_offer(request) + assert offer.sdp == sdp + assert offer.session.model == "voice" + assert await request.body() == body + assert not chunks + + +@pytest.mark.asyncio +async def test_empty_pre_read_multipart_offer_returns_invalid_offer(): + from fastapi import Request + + async def receive(): + return {"type": "http.request", "body": b"--Boundary--\r\n", "more_body": False} + + request = Request( + {"type": "http", "headers": [(b"content-type", b"multipart/form-data; boundary=Boundary")]}, receive + ) + assert not await request.form() + with pytest.raises(HTTPException) as rejected: + await codex.create_codex_realtime_call(request) + assert rejected.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_oversized_pre_read_offer_is_rejected_before_decoding(monkeypatch): + from fastapi import Request + + monkeypatch.setattr(codex, "MAX_REALTIME_OFFER_BYTES", 1024) + + async def receive(): + return {"type": "http.request", "body": b"x" * 2048, "more_body": False} + + request = Request({"type": "http", "headers": [(b"content-type", b"application/json")]}, receive) + await request.body() + with pytest.raises(HTTPException) as rejected: + await codex.read_codex_offer(request) + assert rejected.value.status_code == 413 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("malformed", [False, True]) +async def test_mixed_case_offer_preserves_boundary_metadata_and_closes_extra_files(monkeypatch, malformed): + from fastapi import Request + from litellm.proxy.common_utils.http_parsing_utils import _read_request_body + + boundary = "AbCdEf123" + fields = { + "sdp": "v=0", + "session": "invalid" if malformed else '{"model":"voice"}', + "metadata": '{"policy":"keep"}', + "extra_policy": "keep", + } + body = ( + "".join( + f'--{boundary}\r\nContent-Disposition: form-data; name="{name}"\r\n\r\n{value}\r\n' + for name, value in fields.items() + ) + + f'--{boundary}\r\nContent-Disposition: form-data; name="extra_file"; filename="test.txt"\r\n\r\nextra\r\n--{boundary}--\r\n' + ).encode() + + async def receive(): + return {"type": "http.request", "body": body, "more_body": False} + + request = Request( + {"type": "http", "headers": [(b"content-type", f'Multipart/Form-Data; boundary="{boundary}"'.encode())]}, + receive, + ) + assert not await request.form() + if malformed: + with pytest.raises(HTTPException) as error: + await codex.create_codex_realtime_call(request) + assert error.value.status_code == 400 + assert (await request.form())["extra_file"].file.closed + return + first = await codex.read_codex_offer(request) + second = await codex.read_codex_offer(request) + assert first == second + parsed = await _read_request_body(request) + assert parsed["metadata"] == {"policy": "keep"} + assert parsed["extra_policy"] == "keep" + assert not parsed["extra_file"].file.closed + + async def deny_auth(**kwargs): + assert kwargs["request"] is request + auth_form = await request.form() + assert auth_form["extra_policy"] == "keep" + assert await auth_form["extra_file"].read() == b"extra" + assert not auth_form["extra_file"].file.closed + raise HTTPException(403, "policy denied") + + monkeypatch.setattr(codex, "user_api_key_auth", deny_auth) + with pytest.raises(HTTPException, match="policy denied"): + await codex.create_codex_realtime_call(request) + assert parsed["extra_file"].file.closed + assert request.headers["content-type"] == f'Multipart/Form-Data; boundary="{boundary}"' + + +@pytest.mark.asyncio +@pytest.mark.parametrize("multipart", [False, True]) +@pytest.mark.parametrize("policy", ["budget", "personal_models"]) +@pytest.mark.parametrize("mixed_case", [False, True]) +@pytest.mark.parametrize("pre_read", [False, True]) +async def test_offer_auth_enforces_session_model_policy_before_upstream(monkeypatch, multipart, policy, mixed_case, pre_read): + import json + from unittest.mock import AsyncMock, MagicMock + + import httpx + from fastapi import Request, Response + + import litellm + from litellm.exceptions import BudgetExceededError + from litellm.proxy import proxy_server as server + from litellm.proxy._types import LiteLLM_UserTable + from litellm.proxy.auth.auth_checks import common_checks + from litellm.proxy.common_utils.http_parsing_utils import _read_request_body + from litellm.proxy.realtime_endpoints.endpoints import proxy_realtime_calls + + session = {"model": "forbidden-voice"} + payload = ( + {"files": {"sdp": (None, "v=0"), "session": (None, json.dumps(session)), "model": (None, "body-decoy")}} + if multipart + else {"json": {"sdp": "v=0", "session": session, "model": "body-decoy"}} + ) + outbound = httpx.Request("POST", "http://localhost/v1/realtime/calls", **payload) + body = outbound.read() + content_type = outbound.headers["content-type"] + if mixed_case: + content_type = content_type.replace("multipart/form-data", "Multipart/Form-Data").replace( + "application/json", "Application/JSON" + ) + receives = [] + + async def receive(): + receives.append(True) + assert len(receives) == 1 + return {"type": "http.request", "body": body, "more_body": False} + + request = Request( + { + "type": "http", + "method": "POST", + "path": "/v1/realtime/calls", + "query_string": b"model=query-decoy&policy=keep", + "client": ("127.0.0.7", 1234), + "headers": [ + (b"content-type", content_type.encode()), + (b"x-policy-key", b"Bearer test-key"), + (b"x-custom-policy", b"preserved"), + (b"x-litellm-model", b"header-decoy"), + ], + }, + receive, + ) + token = UserAPIKeyAuth(token="test-key", user_id="personal-user", model_max_budget={"forbidden-voice": 0}) + budget = AsyncMock(side_effect=BudgetExceededError(current_cost=1, max_budget=0)) + upstream = AsyncMock() + original_request = request + + async def custom_auth(request: Request, api_key: str): + assert request is original_request + assert request.headers["content-type"] == content_type + assert api_key == "test-key" + assert request.headers["x-custom-policy"] == "preserved" + assert request.query_params["policy"] == "keep" + assert request.client.host == "127.0.0.7" + if multipart: + assert (await request.form())["model"] == "body-decoy" + parsed = await _read_request_body(request) + assert parsed["model"] == "body-decoy" + assert isinstance(parsed["session"], str) is multipart + if policy == "personal_models": + await common_checks( + request_body=parsed, + team_object=None, + user_object=LiteLLM_UserTable( + user_id="personal-user", models=["allowed-voice", "body-decoy", "query-decoy", "header-decoy"] + ), + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/v1/realtime/calls", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=token, + request=request, + skip_budget_checks=True, + ) + return token + + custom = AsyncMock(side_effect=custom_auth) + monkeypatch.setattr(server, "general_settings", {"litellm_key_header_name": "x-policy-key"}) + monkeypatch.setattr(server, "user_custom_auth", custom) + monkeypatch.setattr(server, "llm_router", None) + monkeypatch.setattr(server, "llm_model_list", []) + monkeypatch.setattr(server, "model_max_budget_limiter", SimpleNamespace(is_key_within_model_budget=budget)) + monkeypatch.setattr(server, "route_request", upstream) + monkeypatch.setattr(litellm, "enable_post_custom_auth_checks", True, raising=False) + if pre_read: + await _read_request_body(request) + with pytest.raises(ProxyException) as denied: + await proxy_realtime_calls(request, Response()) + if policy == "personal_models": + assert "user not allowed to access model" in str(denied.value) + assert "forbidden-voice" in str(denied.value) + custom.assert_awaited_once() + upstream.assert_not_awaited() + if policy == "budget": + budget.assert_awaited_once() + assert budget.await_args.kwargs["model"] == "forbidden-voice" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route_type", ["arealtime_calls", "_arealtime"]) +@pytest.mark.parametrize("observer", [False, True]) +async def test_codex_processing_merges_model_guardrails(monkeypatch, route_type, observer): + from fastapi import Request + from litellm import Router + from litellm.proxy import proxy_server as server + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.realtime_endpoints.call_sessions import process_codex_request + + class PolicyHook: + async def pre_call_hook(self, user_api_key_dict, data, call_type, *, internal_realtime_observer=False): + assert internal_realtime_observer is observer + if "model-policy" in data.get("metadata", {}).get("guardrails", []): + raise HTTPException(403, "Model policy rejected request") + return data + + router = Router(model_list=[{ + "model_name": "voice-policy", + "litellm_params": {"model": "openai/gpt-realtime-1.5", "api_key": "test", "guardrails": ["model-policy"]}, + }]) + monkeypatch.setattr(server, "llm_router", router) + monkeypatch.setattr(server, "proxy_logging_obj", PolicyHook()) + request = Request({"type": "http", "method": "POST", "path": "/v1/realtime/calls", "headers": [], "query_string": b"", "scheme": "http", "server": ("localhost", 80)}) + with pytest.raises(HTTPException) as error: + await process_codex_request( + request, + {"model": "voice-policy"}, + UserAPIKeyAuth(), + "voice-policy", + route_type, + internal_realtime_observer=observer, + ) + assert error.value.status_code == 403 + assert error.value.detail == "Model policy rejected request" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("logged_success", [False, True]) +@pytest.mark.parametrize("disconnect_error", [False, True]) +async def test_sideband_preserves_pending_cost_reconciliation(monkeypatch, logged_success, disconnect_error): + import litellm + from unittest.mock import AsyncMock + from litellm.litellm_core_utils.realtime_streaming import REALTIME_SESSION_SUCCESS_LOGGED_KEY + + monkeypatch.setenv("LITELLM_SALT_KEY", "test-only-salt-for-codex-realtime") + call = CodexRealtimeCall(call_id="rtc_test", model="gpt-live-1-codex", alias="voice", + owner=hashlib.sha256(b"Bearer owner").hexdigest(), expires_at=time.time()+300) + auth = UserAPIKeyAuth() + auth.budget_reservation = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} + logger = SimpleNamespace(model_call_details={}) + monkeypatch.setattr(codex, "can_key_call_resolved_model", AsyncMock()) + monkeypatch.setattr(codex, "process_codex_request", AsyncMock(return_value=({}, logger))) + + async def forward(**kwargs): + if logged_success: + logger.model_call_details[REALTIME_SESSION_SUCCESS_LOGGED_KEY] = True + if disconnect_error: + raise RuntimeError("Backend disconnected") + + monkeypatch.setattr(litellm, "_arealtime", forward) + websocket = WebSocket({"type": "websocket", "path": "/v1/live/opaque", "query_string": b"", + "headers": [(b"authorization", b"Bearer owner")]}, + AsyncMock(return_value={"type": "websocket.connect"}), AsyncMock()) + if disconnect_error: + with pytest.raises(RuntimeError, match="Backend disconnected"): + await codex.codex_realtime_sideband(websocket, encode_call(call), auth) + else: + await codex.codex_realtime_sideband(websocket, encode_call(call), auth) + assert auth.budget_reservation["finalized"] is not logged_success + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ending", ["normal", "disconnect", "pre_call", "admission"]) +async def test_supervised_attachments_release_real_limiter_before_reconnect(monkeypatch, ending): + import asyncio + from unittest.mock import AsyncMock + + import litellm + from litellm.caching.caching import DualCache + from litellm.proxy import proxy_server as server + from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + _PROXY_MaxParallelRequestsHandler_v3, + _request_stash, + get_request_stash, + ) + from litellm.proxy.utils import InternalUsageCache + + cache = DualCache() + limiter = _PROXY_MaxParallelRequestsHandler_v3(InternalUsageCache(cache)) + auth = UserAPIKeyAuth(api_key="attachment-owner", max_parallel_requests=1, tpm_limit=10000) + token_key = limiter.create_rate_limit_keys(key="api_key", value=auth.api_key, rate_limit_type="tokens") + parallel_key = f"{{api_key:{auth.api_key}}}:max_parallel_requests" + call = CodexRealtimeCall( + call_id="rtc_test", + model="gpt-live-1-codex", + alias="voice", + usage_supervised=True, + owner=hashlib.sha256(b"Bearer owner").hexdigest(), + expires_at=time.time() + 300, + ) + monkeypatch.setenv("LITELLM_SALT_KEY", "attachment-cleanup-test") + monkeypatch.setattr(codex, "can_key_call_resolved_model", AsyncMock()) + monkeypatch.setattr(server, "proxy_logging_obj", SimpleNamespace(get_proxy_hook=lambda name: limiter)) + + async def process(request, data, selected_auth, model, call_type): + await limiter.async_pre_call_hook( + user_api_key_dict=selected_auth, + cache=cache, + data={"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hello"}], "max_tokens": 50}, + call_type="completion", + ) + assert get_request_stash().reserved_tokens > 0 + if ending == "pre_call": + raise RuntimeError("Later policy rejected attachment") + return data, SimpleNamespace(model_call_details={}) + + async def forward(**kwargs): + if ending == "disconnect": + raise asyncio.CancelledError() + + monkeypatch.setattr(codex, "process_codex_request", process) + monkeypatch.setattr(litellm, "_arealtime", forward) + blocker_stash = None + blocker_reserved = 0 + if ending == "admission": + setup_token = _request_stash.set(None) + try: + await limiter.async_pre_call_hook( + user_api_key_dict=auth, + cache=cache, + data={"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hello"}], "max_tokens": 50}, + call_type="completion", + ) + blocker_stash = get_request_stash() + blocker_reserved = blocker_stash.reserved_tokens + finally: + _request_stash.reset(setup_token) + for _ in range(3): + stash_token = _request_stash.set(None) + try: + websocket = WebSocket( + { + "type": "websocket", + "path": "/v1/live/opaque", + "query_string": b"", + "headers": [(b"authorization", b"Bearer owner")], + }, + AsyncMock(return_value={"type": "websocket.connect"}), + AsyncMock(), + ) + if ending == "disconnect": + with pytest.raises(asyncio.CancelledError): + await codex.codex_realtime_sideband(websocket, encode_call(call), auth) + else: + await codex.codex_realtime_sideband(websocket, encode_call(call), auth) + assert limiter._gauge_in_flight_from_cache_value(await cache.async_get_cache(parallel_key)) == int( + ending == "admission" + ) + assert int(await cache.async_get_cache(token_key) or 0) == blocker_reserved + finally: + _request_stash.reset(stash_token) + if blocker_stash is not None: + cleanup_token = _request_stash.set(blocker_stash) + try: + await limiter.async_release_realtime_attachment({}, auth) + finally: + _request_stash.reset(cleanup_token) + + +def test_sideband_token_binds_owner_and_model(monkeypatch): + monkeypatch.setenv("LITELLM_SALT_KEY", "test-only-salt-for-codex-realtime") + call = CodexRealtimeCall( + call_id="rtc_test", + model="gpt-live-1-codex", + alias="gpt-live-1-codex", + extra_headers={"x-gateway-secret": "configured-secret"}, + owner=hashlib.sha256(b"Bearer test-owner").hexdigest(), + expires_at=time.time() + 300, + ) + token = encode_call(call) + assert "/" not in token + assert "configured-secret" not in token + assert decode_call(token, "Bearer test-owner") == call + with pytest.raises(HTTPException) as error: + decode_call(token, "Bearer different-owner") + assert error.value.status_code == 403 + with pytest.raises(HTTPException): + decode_call(token[:30] + "tampered" + token[30:], "Bearer test-owner") + + +def test_sideband_rejects_expired_token(monkeypatch): + monkeypatch.setenv("LITELLM_SALT_KEY", "test-only-salt-for-codex-realtime") + call = CodexRealtimeCall( + call_id="rtc_test", + model="gpt-realtime-1.5", + alias="gpt-realtime-1.5", + owner=hashlib.sha256(b"Bearer test-owner").hexdigest(), + expires_at=time.time() - 1, + ) + with pytest.raises(HTTPException): + decode_call(encode_call(call), "Bearer test-owner") + + +@pytest.mark.parametrize("token", ["", "rtc_other", "rtc_litellm_%%%%", "rtc_litellm_a"]) +def test_sideband_rejects_malformed_tokens(token): + with pytest.raises(HTTPException): + decode_call(token, "Bearer test-owner") + + +@pytest.mark.asyncio +async def test_sideband_rejects_revoked_model_access(monkeypatch): + monkeypatch.setenv("LITELLM_SALT_KEY", "test-only-salt-for-codex-realtime") + call = CodexRealtimeCall( + call_id="rtc_test", + model="gpt-live-1-codex", + alias="voice", + owner=hashlib.sha256(b"Bearer test-owner").hexdigest(), + expires_at=time.time() + 300, + ) + sent = [] + + async def receive(): + return {"type": "websocket.connect"} + + async def send(message): + sent.append(message) + + async def deny_model(**kwargs): + raise ProxyException("Model access revoked", "auth_error", "model", 403) + + monkeypatch.setattr(codex, "can_key_call_resolved_model", deny_model) + websocket = WebSocket( + {"type": "websocket", "headers": [(b"authorization", b"Bearer test-owner")]}, receive, send + ) + await codex.codex_realtime_sideband(websocket, encode_call(call), UserAPIKeyAuth()) + assert sent == [{"type": "websocket.close", "code": 1008, "reason": "Invalid realtime call"}] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_id", ["rtc_raw", "", "rtc_litellm_invalid"]) +async def test_realtime_endpoint_rejects_untrusted_call_ids(monkeypatch, call_id): + from unittest.mock import AsyncMock + from fastapi import WebSocket + from litellm.proxy import proxy_server as server + from litellm.proxy._types import UserAPIKeyAuth + + sent = [] + + async def receive(): + return {"type": "websocket.connect"} + + async def send(message): + sent.append(message) + + route = AsyncMock() + monkeypatch.setattr(server, "route_request", route) + websocket = WebSocket({"type": "websocket", "headers": [], "query_string": b""}, receive, send) + await server.realtime_websocket_endpoint( + websocket, model="gpt-realtime-1.5", call_id=call_id, + intent=None, guardrails=None, user_api_key_dict=UserAPIKeyAuth() + ) + assert sent == [{"type": "websocket.close", "code": 1008, "reason": "Invalid realtime call"}] + route.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("multipart", [False, True]) +@pytest.mark.parametrize("credential", ["authorization", "api-key", "subprotocol", "x-litellm-api-key", "custom"]) +@pytest.mark.parametrize("signaling_credential", ["authorization", "api-key", "x-litellm-api-key", "mixed"]) +async def test_offer_exchange_wraps_call_and_filters_client_headers( + monkeypatch, multipart, credential, signaling_credential +): + import json + from unittest.mock import AsyncMock + + import httpx + from fastapi import Request, WebSocket + from litellm.proxy import common_request_processing, proxy_server + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.realtime_endpoints import call_sessions as codex + import litellm + + monkeypatch.setenv("LITELLM_SALT_KEY", "test-only-salt-for-codex-realtime") + session = {"model": "voice-alias", "audio": {"output": {"voice": "sol"}}} + if multipart: + body_request = httpx.Request( + "POST", + "http://test/v1/realtime/calls", + files={"sdp": (None, "v=0\r\n"), "session": (None, json.dumps(session))}, + ) + else: + body_request = httpx.Request( + "POST", "http://test/v1/realtime/calls", json={"sdp": "v=0\r\n", "session": session} + ) + body = body_request.read() + + async def receive(): + return {"type": "http.request", "body": body, "more_body": False} + + signaling_headers = ( + [(b"authorization", b"Bearer other-owner"), (b"x-litellm-api-key", b"owner")] + if signaling_credential == "mixed" + else [(signaling_credential.encode(), b"Bearer owner" if signaling_credential == "authorization" else b"owner")] + ) + request = Request( + { + "type": "http", + "method": "POST", + "path": "/v1/realtime/calls", + "scheme": "http", + "server": ("localhost", 80), + "query_string": b"intent=quicksilver&architecture=avas&untrusted=bad", + "headers": [ + (b"content-type", body_request.headers["content-type"].encode()), + *signaling_headers, + *([(b"x-proxy-key", b"Bearer owner")] if credential == "custom" else []), + (b"openai-alpha", b"quicksilver=v2"), + (b"x-untrusted", b"bad"), + ], + }, + receive, + ) + auth = UserAPIKeyAuth() + authorize = AsyncMock() + monkeypatch.setattr(proxy_server, "master_key", "owner") + monkeypatch.setattr( + proxy_server, "general_settings", {"litellm_key_header_name": "x-proxy-key"} if credential == "custom" else {} + ) + monkeypatch.setattr(codex, "can_key_call_resolved_model", authorize) + + class Processor: + def __init__(self, data): + self.data = data + + async def common_processing_pre_call_logic(self, **kwargs): + assert isinstance(kwargs["user_api_key_dict"], UserAPIKeyAuth) + if kwargs["route_type"] == "_arealtime": + assert self.data["model"] == "voice-alias" + assert self.data["guardrails"] == ["query-guardrail"] + assert await kwargs["request"].json() == {"model": "voice-alias"} + return { + **self.data, + "extra_headers": { + "X-Hook-Required": "policy-value", + "x-gateway-token": "untrusted-override", + "Authorization": "Bearer untrusted", + }, + "extra_query": {"gateway_token": "untrusted-override"}, + "metadata": {"guardrails": ["policy-guardrail"], "user_api_key_team_id": "team"}, + }, None + return self.data, None + + monkeypatch.setattr(common_request_processing, "ProxyBaseLLMRequestProcessing", Processor) + + async def route(**kwargs): + data = kwargs["data"] + assert data["sdp_body"] == b"v=0\r\n" + assert data["session"] == session + assert data["chatgpt_realtime_client_headers"] == {"openai-alpha": "quicksilver=v2"} + assert "extra_headers" not in data + assert data["chatgpt_realtime_client_query"] == {"intent": "quicksilver", "architecture": "avas"} + + async def respond(): + return httpx.Response( + 201, + content=b"v=0\r\nanswer", + headers={"Location": "/v1/realtime/calls/rtc_private"}, + extensions={ + "chatgpt_realtime": { + "model": "gpt-live-1-codex", + "api_base": "https://voice.example/codex", + "extra_headers": {"X-Gateway-Token": "pinned-value"}, + "extra_query": {"gateway_token": "pinned-query-value"}, + } + }, + ) + + return respond() + + monkeypatch.setattr(proxy_server, "route_request", route) + supervise = AsyncMock() + monkeypatch.setattr(codex, "supervise_codex_call", supervise) + response = await codex.create_codex_realtime_call(request) + assert response.status_code == 201 + assert response.body == b"v=0\r\nanswer" + token = response.headers["location"].rsplit("/", 1)[-1] + call = codex.decode_call(token, "Bearer owner") + assert call.call_id == "rtc_private" + assert call.alias == "voice-alias" + assert call.model == "gpt-live-1-codex" + assert call.usage_supervised + supervise.assert_awaited_once() + assert "rtc_private" not in token + assert "pinned-query-value" not in token + assert call.extra_query == {"gateway_token": "pinned-query-value"} + assert time.time() < call.expires_at < time.time() + 3601 + authorize.assert_awaited_once() + + sent = [] + + async def send(message): + sent.append(message) + + async def receive_ws(): + return {"type": "websocket.connect"} + + credential_headers = { + "authorization": [(b"authorization", b"Bearer owner")], + "api-key": [(b"api-key", b"owner")], + "x-litellm-api-key": [(b"x-litellm-api-key", b"owner")], + "custom": [(b"x-proxy-key", b"Bearer owner")], + "subprotocol": [(b"sec-websocket-protocol", b"realtime, openai-insecure-api-key.owner")], + } + websocket = WebSocket( + { + "type": "websocket", + "path": "/v1/live/opaque", + "query_string": b"guardrails=query-guardrail", + "headers": credential_headers[credential], + }, + receive_ws, + send, + ) + forward = AsyncMock() + monkeypatch.setattr(litellm, "_arealtime", forward) + await codex.codex_realtime_sideband(websocket, token, auth) + assert sent[0]["type"] == "websocket.accept" + if credential == "subprotocol": + assert sent[0]["subprotocol"] == "realtime" + assert forward.await_args.kwargs["extra_headers"] == { + "x-hook-required": "policy-value", + "x-gateway-token": "pinned-value", + } + assert forward.await_args.kwargs["metadata"] == {"guardrails": ["policy-guardrail"], "user_api_key_team_id": "team"} + assert forward.await_args.kwargs["extra_query"] == {"gateway_token": "pinned-query-value"} + assert forward.await_args.kwargs["chatgpt_realtime_call_id"] == "rtc_private" + assert forward.await_args.kwargs["model"] == "chatgpt/gpt-live-1-codex" + assert forward.await_args.kwargs["api_base"] == "https://voice.example/codex" + assert authorize.await_count == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("body", [b"not json", b'{}', b'{"sdp":"v=0","session":{}}']) +async def test_invalid_offers_fail_before_authentication(monkeypatch, body): + from unittest.mock import AsyncMock + from fastapi import Request + from litellm.proxy.realtime_endpoints import call_sessions as codex + + async def receive(): + return {"type": "http.request", "body": body} + + request = Request({"type": "http", "headers": [(b"content-type", b"application/json")]}, receive) + authenticate = AsyncMock() + monkeypatch.setattr(codex, "user_api_key_auth", authenticate) + with pytest.raises(HTTPException) as error: + await codex.create_codex_realtime_call(request) + assert error.value.status_code == 400 + authenticate.assert_not_called() + + +@pytest.mark.asyncio +async def test_sideband_pre_call_block_prevents_upstream_connection(monkeypatch): + from unittest.mock import AsyncMock + from fastapi import WebSocket + import litellm + from litellm.proxy import common_request_processing + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.realtime_endpoints import call_sessions as codex + + monkeypatch.setenv("LITELLM_SALT_KEY", "test-only-salt-for-codex-realtime") + call = CodexRealtimeCall(call_id="rtc_test", model="gpt-live-1-codex", alias="voice", + owner=hashlib.sha256(b"Bearer owner").hexdigest(), expires_at=time.time()+300) + token = encode_call(call) + sent = [] + + async def receive(): + return {"type": "websocket.connect"} + + async def send(message): + sent.append(message) + + class BlockingProcessor: + def __init__(self, data): + assert data["model"] == "voice" + + async def common_processing_pre_call_logic(self, **kwargs): + assert kwargs["route_type"] == "_arealtime" + raise HTTPException(403, "Policy blocked this call") + + forward = AsyncMock() + monkeypatch.setattr(litellm, "_arealtime", forward) + monkeypatch.setattr(codex, "can_key_call_resolved_model", AsyncMock()) + monkeypatch.setattr(common_request_processing, "ProxyBaseLLMRequestProcessing", BlockingProcessor) + websocket = WebSocket({"type": "websocket", "path": "/v1/live/opaque", "query_string": b"", + "headers": [(b"authorization", b"Bearer owner")]}, receive, send) + await codex.codex_realtime_sideband(websocket, token, UserAPIKeyAuth()) + forward.assert_not_called() + assert sent == [{"type": "websocket.close", "code": 1008, "reason": "Realtime pre-call rejected"}] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("observer_fails", [False, True]) +async def test_signaling_transfers_reservation_only_to_ready_observer(monkeypatch, observer_fails): + import json + from unittest.mock import AsyncMock + + import httpx + from fastapi import Request + + from litellm.proxy import proxy_server + + monkeypatch.setenv("LITELLM_SALT_KEY", "test-only-reservation-transfer") + reservation = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []} + auth = UserAPIKeyAuth(budget_reservation=reservation) + monkeypatch.setattr(codex, "user_api_key_auth", AsyncMock(return_value=auth)) + monkeypatch.setattr(codex, "can_key_call_resolved_model", AsyncMock()) + monkeypatch.setattr(proxy_server, "general_settings", {}) + process = AsyncMock(return_value=({}, None)) + monkeypatch.setattr(codex, "process_codex_request", process) + + async def response(): + return httpx.Response( + 201, + text="v=0\r\n", + headers={"Location": "/v1/realtime/calls/rtc_ready"}, + extensions={"chatgpt_realtime": {"model": "gpt-live-1-codex"}}, + ) + + async def route(**kwargs): + return response() + + monkeypatch.setattr(proxy_server, "route_request", route) + + async def supervise(request, call, owner): + assert owner is auth + assert not owner.budget_reservation["finalized"] + assert call.usage_supervised + if observer_fails: + await codex.release_or_invalidate_budget_reservation(budget_reservation=owner.budget_reservation) + raise RuntimeError("Observer unavailable") + + monkeypatch.setattr(codex, "supervise_codex_call", supervise) + + async def receive(): + return {"type": "http.request", "body": json.dumps({"sdp": "v=0", "session": {"model": "voice"}}).encode()} + + request = Request( + { + "type": "http", + "method": "POST", + "path": "/v1/realtime/calls", + "query_string": b"", + "headers": [(b"content-type", b"application/json"), (b"authorization", b"Bearer owner")], + }, + receive, + ) + if observer_fails: + with pytest.raises(RuntimeError, match="Observer unavailable"): + await codex.create_codex_realtime_call(request) + else: + assert (await codex.create_codex_realtime_call(request)).status_code == 201 + assert process.await_args.args[2].budget_reservation is None + assert auth.budget_reservation["finalized"] is observer_fails + + +@pytest.mark.asyncio +async def test_supervisor_policy_failure_hangs_up_before_releasing(monkeypatch): + from unittest.mock import AsyncMock + + from fastapi import Request + + call = CodexRealtimeCall( + call_id="rtc_open", model="gpt-live-1-codex", alias="voice", owner="owner", expires_at=time.time() + 60 + ) + auth = UserAPIKeyAuth(budget_reservation={"reserved_cost": 0.5, "finalized": False, "entries": []}) + monkeypatch.setattr(codex, "process_codex_request", AsyncMock(side_effect=HTTPException(403, "Policy rejected"))) + closed = [] + + class Handler: + def __init__(self, *args): + pass + + @staticmethod + def get_api_base(base): + return "https://gateway.test/v1" + + async def hangup_call(self, base): + assert not auth.budget_reservation["finalized"] + closed.append(base) + + monkeypatch.setattr(codex, "ChatGPTRealtime", Handler) + request = Request({"type": "http", "headers": [], "method": "POST", "path": "/v1/realtime/calls"}) + with pytest.raises(HTTPException) as error: + await codex.supervise_codex_call(request, call, auth) + assert error.value.status_code == 403 + assert closed == ["https://gateway.test/v1"] + assert auth.budget_reservation["finalized"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("hangup_fails", [False, True]) +async def test_supervisor_constructor_failure_closes_effective_connection(monkeypatch, hangup_fails, caplog): + from unittest.mock import AsyncMock, MagicMock + + from fastapi import Request + + call = CodexRealtimeCall( + call_id="rtc_open", model="gpt-live-1-codex", alias="voice", owner="owner", expires_at=time.time() + 60 + ) + auth = UserAPIKeyAuth(budget_reservation={"reserved_cost": 0.5, "finalized": False, "entries": []}) + logger = MagicMock() + logger.litellm_params = {} + connection = AsyncMock() + handlers = [] + invalidate = AsyncMock() + release = AsyncMock() + monkeypatch.setattr(codex, "invalidate_budget_reservation_counters", invalidate, raising=False) + monkeypatch.setattr(codex, "release_or_invalidate_budget_reservation", release) + monkeypatch.setattr( + codex, "process_codex_request", AsyncMock(return_value=({"extra_headers": {"x-hook": "effective"}}, logger)) + ) + + class Handler: + def __init__(self, params, headers, extra_headers): + self.headers = extra_headers + handlers.append(self) + + @staticmethod + def get_api_base(base): + return "https://gateway.test/v1" + + async def open_call_connection(self, model, base): + return connection + + async def hangup_call(self, base): + assert self.headers["x-hook"] == "effective" + if hangup_fails: + raise RuntimeError("private-cleanup-credential") + + monkeypatch.setattr(codex, "ChatGPTRealtime", Handler) + monkeypatch.setattr(codex, "RealTimeStreaming", MagicMock(side_effect=ValueError("original constructor failure"))) + request = Request({"type": "http", "headers": [], "method": "POST", "path": "/v1/realtime/calls"}) + with pytest.raises(ValueError, match="original constructor failure"): + await codex.supervise_codex_call(request, call, auth) + connection.close.assert_awaited_once() + assert len(handlers) == 1 + if hangup_fails: + invalidate.assert_awaited_once_with(budget_reservation=auth.budget_reservation) + release.assert_not_awaited() + else: + release.assert_awaited_once_with(budget_reservation=auth.budget_reservation) + invalidate.assert_not_awaited() + assert "private-cleanup-credential" not in caplog.text + + +@pytest.mark.asyncio +async def test_attachment_releases_quota_before_upstream_close_handshake(monkeypatch): + import asyncio + from unittest.mock import AsyncMock + + import websockets + + import litellm + from litellm.proxy import proxy_server as server + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.hooks.parallel_request_limiter_v3 import _request_stash + from litellm.proxy.utils import ProxyLogging + + proxy = ProxyLogging(UserApiKeyCache()) + monkeypatch.setattr(litellm, "callbacks", []) + proxy._add_proxy_hooks() + monkeypatch.setattr(server, "proxy_logging_obj", proxy) + monkeypatch.setattr( + server, + "llm_router", + litellm.Router( + model_list=[ + {"model_name": "voice", "litellm_params": {"model": "openai/gpt-realtime-1.5", "api_key": "test"}} + ] + ), + ) + monkeypatch.setattr(codex, "can_key_call_resolved_model", AsyncMock()) + monkeypatch.setenv("LITELLM_SALT_KEY", "attachment-close-order-test") + from litellm.llms.chatgpt.authenticator import Authenticator + + monkeypatch.setattr(Authenticator, "get_access_token", lambda self: "test-token") + monkeypatch.setattr(Authenticator, "get_account_id", lambda self: "test-account") + closing = asyncio.Event() + finish_close = asyncio.Event() + + class Backend: + async def recv(self, **kwargs): + await asyncio.Event().wait() + + async def send(self, value): + return None + + class Connection: + async def __aenter__(self): + return Backend() + + async def __aexit__(self, *args): + closing.set() + await finish_close.wait() + + monkeypatch.setattr(websockets, "connect", lambda *args, **kwargs: Connection()) + auth = UserAPIKeyAuth(api_key="close-order-owner", max_parallel_requests=1) + call = CodexRealtimeCall( + call_id="rtc_test", + model="gpt-live-1-codex", + alias="voice", + usage_supervised=True, + owner=hashlib.sha256(b"Bearer owner").hexdigest(), + expires_at=time.time() + 300, + ) + ws = WebSocket( + { + "type": "websocket", + "path": "/v1/live/test", + "query_string": b"", + "headers": [(b"authorization", b"Bearer owner")], + "scheme": "ws", + "server": ("localhost", 80), + }, + AsyncMock(side_effect=[{"type": "websocket.connect"}, {"type": "websocket.disconnect", "code": 1000}]), + AsyncMock(), + ) + token = _request_stash.set(None) + request = asyncio.create_task(codex.codex_realtime_sideband(ws, encode_call(call), auth)) + try: + await asyncio.wait_for(closing.wait(), timeout=5) + limiter = proxy.get_proxy_hook("parallel_request_limiter") + value = await proxy.internal_usage_cache.async_get_cache( + "{api_key:close-order-owner}:max_parallel_requests", litellm_parent_otel_span=None, local_only=True + ) + assert limiter._gauge_in_flight_from_cache_value(value) == 0 + finally: + finish_close.set() + await asyncio.wait_for(request, timeout=5) + _request_stash.reset(token) + value = await proxy.internal_usage_cache.async_get_cache( + "{api_key:close-order-owner}:max_parallel_requests", litellm_parent_otel_span=None, local_only=True + ) + assert limiter._gauge_in_flight_from_cache_value(value) == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", [None, "provider", "observer", "renewal", "legacy_key", "legacy_global"]) +async def test_signaling_keeps_or_releases_owned_call_lease(monkeypatch, failure): + import json + from unittest.mock import AsyncMock, MagicMock + + import httpx + from fastapi import Request + from litellm.proxy.hooks.realtime_call_lease import RealtimeCallLease + + from litellm.proxy import proxy_server as server + from litellm.proxy.hooks.parallel_request_limiter import _PROXY_MaxParallelRequestsHandler + from litellm.proxy.hooks.parallel_request_limiter_v3 import _PROXY_MaxParallelRequestsHandler_v3 + + auth = UserAPIKeyAuth(max_parallel_requests=None if failure == "legacy_global" else 1) + lease = MagicMock(spec=RealtimeCallLease) + lease.renew = AsyncMock(return_value=failure != "renewal") + lease.close = AsyncMock() + legacy = failure in ("legacy_key", "legacy_global") + limiter = MagicMock(spec=_PROXY_MaxParallelRequestsHandler if legacy else _PROXY_MaxParallelRequestsHandler_v3) + if not legacy: + limiter.transfer_realtime_call_slot.return_value = lease + proxy = MagicMock() + proxy.get_proxy_hook.return_value = limiter + monkeypatch.setattr(server, "proxy_logging_obj", proxy) + monkeypatch.setattr( + server, "general_settings", {"global_max_parallel_requests": 1} if failure == "legacy_global" else {} + ) + monkeypatch.setattr(codex, "user_api_key_auth", AsyncMock(return_value=auth)) + monkeypatch.setattr(codex, "can_key_call_resolved_model", AsyncMock()) + process = AsyncMock(return_value=({}, None)) + monkeypatch.setattr(codex, "process_codex_request", process) + monkeypatch.setenv("LITELLM_SALT_KEY", "lease-transfer-test") + + async def route(**kwargs): + lease.start.assert_called_once() + if failure == "provider": + raise RuntimeError("Provider unavailable") + + async def respond(): + return httpx.Response( + 201, + text="v=0\r\n", + headers={"Location": "/v1/realtime/calls/rtc_lease"}, + extensions={"chatgpt_realtime": {"model": "gpt-live-1-codex"}}, + ) + + return respond() + + async def supervise(request, call, owner, selected_lease): + assert owner is auth + assert selected_lease is lease + assert call.parallel_reserved + if failure == "observer": + raise RuntimeError("Observer unavailable") + + monkeypatch.setattr(server, "route_request", route) + monkeypatch.setattr(codex, "supervise_codex_call", supervise) + request = Request( + { + "type": "http", + "method": "POST", + "path": "/v1/realtime/calls", + "query_string": b"", + "headers": [(b"content-type", b"application/json"), (b"authorization", b"Bearer owner")], + }, + AsyncMock( + return_value={ + "type": "http.request", + "body": json.dumps({"sdp": "v=0", "session": {"model": "voice"}}).encode(), + } + ), + ) + if failure is None: + response = await codex.create_codex_realtime_call(request) + token = response.headers["location"].rsplit("/", 1)[-1] + assert codex.decode_call(token, "Bearer owner").parallel_reserved + lease.close.assert_not_awaited() + else: + with pytest.raises((RuntimeError, HTTPException)) as raised: + await codex.create_codex_realtime_call(request) + if legacy: + assert raised.value.status_code == 400 + assert "V3 rate limiter" in raised.value.detail + process.assert_not_awaited() + lease.close.assert_not_awaited() + else: + if failure == "renewal": + assert raised.value.status_code == 503 + lease.close.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_signaling_rejection_after_admission_refunds_parallel_slot(monkeypatch): + import json + from unittest.mock import AsyncMock + + from fastapi import Request + + import litellm + from litellm.integrations.custom_logger import CustomLogger + from litellm.proxy import proxy_server as server + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.utils import ProxyLogging + + proxy = ProxyLogging(UserApiKeyCache()) + monkeypatch.setattr(litellm, "callbacks", []) + proxy._add_proxy_hooks() + limiter = proxy.get_proxy_hook("parallel_request_limiter") + key = "{api_key:rejected-signaling-owner}:max_parallel_requests" + + class Reject(CustomLogger): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + current = await proxy.internal_usage_cache.async_get_cache(key, litellm_parent_otel_span=None, local_only=True) + assert limiter._gauge_in_flight_from_cache_value(current) == 1 + raise RuntimeError("Policy rejected after admission") + + litellm.callbacks.append(Reject()) + monkeypatch.setattr(server, "proxy_logging_obj", proxy) + monkeypatch.setattr(server, "general_settings", {}) + monkeypatch.setattr( + server, + "llm_router", + litellm.Router( + model_list=[ + {"model_name": "voice", "litellm_params": {"model": "openai/gpt-realtime-1.5", "api_key": "test"}} + ] + ), + ) + auth = UserAPIKeyAuth(api_key="rejected-signaling-owner", max_parallel_requests=1) + monkeypatch.setattr(codex, "user_api_key_auth", AsyncMock(return_value=auth)) + monkeypatch.setattr(codex, "can_key_call_resolved_model", AsyncMock()) + route = AsyncMock() + monkeypatch.setattr(server, "route_request", route) + for _ in range(2): + request = Request( + { + "type": "http", + "method": "POST", + "path": "/v1/realtime/calls", + "query_string": b"", + "headers": [(b"content-type", b"application/json")], + }, + AsyncMock( + return_value={ + "type": "http.request", + "body": json.dumps({"sdp": "v=0", "session": {"model": "voice"}}).encode(), + } + ), + ) + with pytest.raises(RuntimeError, match="Policy rejected after admission"): + await codex.create_codex_realtime_call(request) + current = await proxy.internal_usage_cache.async_get_cache(key, litellm_parent_otel_span=None, local_only=True) + assert limiter._gauge_in_flight_from_cache_value(current) == 0 + route.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("callback_order", ["before", "after", "cancel"]) +async def test_signaling_settles_tokens_once_with_isolated_sdk_callbacks(monkeypatch, callback_order): + import asyncio + import json + from datetime import datetime + from unittest.mock import AsyncMock + + import httpx + import litellm + from fastapi import Request + from litellm.proxy import proxy_server as server + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.hooks.parallel_request_limiter_v3 import get_request_stash, isolated_request_stash + from litellm.proxy.utils import ProxyLogging + + proxy = ProxyLogging(UserApiKeyCache()) + monkeypatch.setattr(litellm, "callbacks", []) + proxy._add_proxy_hooks() + limiter = proxy.get_proxy_hook("parallel_request_limiter") + monkeypatch.setattr(server, "proxy_logging_obj", proxy) + monkeypatch.setattr(server, "general_settings", {}) + monkeypatch.setattr( + server, + "llm_router", + litellm.Router( + model_list=[ + {"model_name": "voice", "litellm_params": {"model": "openai/gpt-realtime-1.5", "api_key": "test"}} + ] + ), + ) + auth = UserAPIKeyAuth(api_key="signaling-settlement-owner", max_parallel_requests=1, tpm_limit=10000) + monkeypatch.setattr(codex, "user_api_key_auth", AsyncMock(return_value=auth)) + monkeypatch.setattr(codex, "can_key_call_resolved_model", AsyncMock()) + supervisor = AsyncMock() + monkeypatch.setattr(codex, "supervise_codex_call", supervisor) + monkeypatch.setenv("LITELLM_SALT_KEY", "signaling-settlement-test") + ready, release_callback = asyncio.Event(), asyncio.Event() + callbacks = [] + + async def counter(kind): + return await proxy.internal_usage_cache.async_get_cache( + f"{{api_key:{auth.api_key}}}:{kind}", litellm_parent_otel_span=None, local_only=True + ) + + async def route(**kwargs): + assert get_request_stash() is None + assert await counter("tokens") > 0 + + async def callback(): + await release_callback.wait() + assert get_request_stash() is None + await limiter.async_log_success_event( + kwargs={ + "litellm_call_id": kwargs["data"]["litellm_call_id"], + "standard_logging_object": {"metadata": {"user_api_key_hash": auth.api_key}}, + }, + response_obj=litellm.ModelResponse(usage=litellm.Usage()), + start_time=datetime.now(), + end_time=datetime.now(), + ) + + async def respond(): + assert get_request_stash() is None + ready.set() + if callback_order == "cancel": + await asyncio.Event().wait() + callbacks.append(asyncio.create_task(callback())) + if callback_order == "before": + release_callback.set() + await callbacks[0] + assert await counter("tokens") > 0 + return httpx.Response( + 201, + text="v=0\r\n", + headers={"Location": "/v1/realtime/calls/rtc_settlement"}, + extensions={"chatgpt_realtime": {"model": "gpt-live-1-codex"}}, + ) + + return respond() + + monkeypatch.setattr(server, "route_request", route) + request = Request( + { + "type": "http", + "method": "POST", + "path": "/v1/realtime/calls", + "query_string": b"", + "headers": [(b"content-type", b"application/json"), (b"authorization", b"Bearer owner")], + }, + AsyncMock( + return_value={ + "type": "http.request", + "body": json.dumps({"sdp": "v=0", "session": {"model": "voice"}}).encode(), + } + ), + ) + with isolated_request_stash(): + signaling = asyncio.create_task(codex.create_codex_realtime_call(request)) + await asyncio.wait_for(ready.wait(), timeout=2) + if callback_order == "cancel": + signaling.cancel() + with pytest.raises(asyncio.CancelledError): + await signaling + supervisor.assert_not_awaited() + else: + assert (await signaling).status_code == 201 + assert limiter._gauge_in_flight_from_cache_value(await counter("max_parallel_requests")) == 1 + await supervisor.call_args.args[3].close() + assert await counter("tokens") == 0 + release_callback.set() + await asyncio.gather(*callbacks) + assert await counter("tokens") == 0 + assert limiter._gauge_in_flight_from_cache_value(await counter("max_parallel_requests")) == 0 diff --git a/tests/test_litellm/proxy/realtime_endpoints/test_call_supervision.py b/tests/test_litellm/proxy/realtime_endpoints/test_call_supervision.py new file mode 100644 index 00000000000..6603b4a0cdb --- /dev/null +++ b/tests/test_litellm/proxy/realtime_endpoints/test_call_supervision.py @@ -0,0 +1,697 @@ +import asyncio +import json +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.realtime_streaming import REALTIME_SESSION_SUCCESS_LOGGED_KEY +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.realtime_endpoints.call_supervision import CallSupervisor, CallSupervisors + + +class Socket: + def __init__(self): + self.messages = asyncio.Queue() + self.closed = False + + def __aiter__(self): + return self + + async def __anext__(self): + message = await self.messages.get() + if message is None: + raise StopAsyncIteration + if isinstance(message, Exception): + raise message + return json.dumps(message) + + async def close(self): + self.closed = True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("lease_lost", [False, True]) +async def test_supervisor_holds_call_lease_until_terminal_accounting(lease_lost): + from litellm.proxy.hooks.realtime_call_lease import RealtimeCallLease + + socket = Socket() + logger = MagicMock(spec=Logging) + logger.model_call_details = {} + sink = Sink(logger) + lost = asyncio.Event() + lease = MagicMock(spec=RealtimeCallLease) + lease.wait_failed = lost.wait + + async def release(): + assert socket.closed + assert sink.logs == 1 + + lease.close = AsyncMock(side_effect=release) + + async def close(): + await socket.messages.put({"type": "session.closed", "usage": {"audio_duration_ms": 1000}}) + + terminate = AsyncMock(side_effect=close) + supervisor = CallSupervisor(socket, sink, logger, UserAPIKeyAuth(), terminate, lease=lease) + await socket.messages.put({"type": "session.started"}) + await supervisor.start() + lease.close.assert_not_awaited() + if lease_lost: + lost.set() + else: + await close() + await asyncio.wait_for(supervisor.wait(), 1) + assert terminate.await_count == int(lease_lost) + lease.close.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stalled_step", ["close", "drain"]) +async def test_live_initial_close_reserves_time_for_independent_hangup(stalled_step): + socket = Socket() + logger = MagicMock(spec=Logging) + logger.model_call_details = {} + close_cancelled = asyncio.Event() + + async def close(): + if stalled_step == "close": + try: + await asyncio.Event().wait() + finally: + close_cancelled.set() + + async def force_close(): + await socket.messages.put({"type": "session.closed", "usage": {"audio_duration_ms": 1000}}) + + force = AsyncMock(side_effect=force_close) + sink = Sink(logger) + supervisor = CallSupervisor( + socket, + sink, + logger, + UserAPIKeyAuth(), + close, + force_close_call=force, + drain_timeout=1, + termination_timeout=0.08, + ) + await socket.messages.put({"type": "session.started"}) + await supervisor.start() + await asyncio.wait_for(supervisor.close(), timeout=0.5) + force.assert_awaited_once() + assert close_cancelled.is_set() == (stalled_step == "close") + assert any(event["type"] == "session.closed" for event in sink.events) + assert not logger.model_call_details.get("realtime_usage_incomplete") + assert sink.logs == 1 + assert socket.closed + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fallback", ["terminal", "no_terminal", "timeout"]) +async def test_live_unacknowledged_close_uses_bounded_independent_hangup(monkeypatch, fallback): + from litellm.proxy.realtime_endpoints import call_supervision + + socket = Socket() + logger = MagicMock(spec=Logging) + logger.model_call_details = {} + sink = Sink(logger) + invalidate = AsyncMock() + monkeypatch.setattr(call_supervision, "invalidate_budget_reservation_counters", invalidate) + + async def force_close(): + if fallback == "terminal": + await socket.messages.put({"type": "session.closed", "usage": {"audio_duration_ms": 1000}}) + elif fallback == "timeout": + await asyncio.Event().wait() + + force = AsyncMock(side_effect=force_close) + close = AsyncMock() + supervisor = CallSupervisor( + socket, + sink, + logger, + UserAPIKeyAuth(), + close, + force_close_call=force, + drain_timeout=0.01, + termination_timeout=0.08, + ) + await socket.messages.put({"type": "session.started"}) + await supervisor.start() + await asyncio.wait_for(supervisor.close(), timeout=0.5) + close.assert_awaited_once() + force.assert_awaited_once() + assert socket.closed + if fallback == "terminal": + invalidate.assert_not_awaited() + assert not logger.model_call_details.get("realtime_usage_incomplete") + else: + invalidate.assert_awaited_once() + assert logger.model_call_details["realtime_usage_incomplete"] is True + + +@pytest.mark.asyncio +async def test_live_confirmed_terminal_does_not_force_hangup(): + socket = Socket() + logger = MagicMock(spec=Logging) + logger.model_call_details = {} + + async def close(): + await socket.messages.put({"type": "session.closed"}) + + force = AsyncMock() + supervisor = CallSupervisor(socket, Sink(logger), logger, UserAPIKeyAuth(), close, force_close_call=force) + await socket.messages.put({"type": "session.started"}) + await supervisor.start() + await supervisor.close() + force.assert_not_awaited() + + +class Sink: + def __init__(self, logger): + self.logger = logger + self.events = [] + self.logs = 0 + + def store_message(self, message): + self.events.append(json.loads(message)) + + async def log_messages(self, *, wait_for_dispatch=False): + assert wait_for_dispatch + self.logs += 1 + self.logger.model_call_details[REALTIME_SESSION_SUCCESS_LOGGED_KEY] = True + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "duration,valid", [(0, True), (1000, True), (None, False), (-1, False), (True, False), ("1000", False)] +) +async def test_live_terminal_requires_valid_duration_for_accounting(monkeypatch, duration, valid): + from litellm.proxy.realtime_endpoints import call_supervision + + socket = Socket() + logger = MagicMock(spec=Logging) + logger.model_call_details = {} + invalidate = AsyncMock() + monkeypatch.setattr(call_supervision, "invalidate_budget_reservation_counters", invalidate) + close = AsyncMock() + force = AsyncMock() + supervisor = CallSupervisor(socket, Sink(logger), logger, UserAPIKeyAuth(), close, force_close_call=force) + await socket.messages.put({"type": "session.started"}) + await supervisor.start() + await socket.messages.put( + {"type": "session.closed", **({"usage": {"audio_duration_ms": duration}} if duration is not None else {})} + ) + await supervisor.wait() + close.assert_not_awaited() + force.assert_not_awaited() + assert socket.closed + assert bool(logger.model_call_details.get("realtime_usage_incomplete")) is not valid + assert invalidate.await_count == (0 if valid else 1) + + +def fixture(*, ready_timeout=1, lifetime=1): + socket = Socket() + logger = MagicMock(spec=Logging) + logger.model_call_details = {} + sink = Sink(logger) + + async def hangup(): + assert not socket.closed + await socket.messages.put({"type": "session.closed", "usage": {"total_tokens": 42}}) + + close_call = AsyncMock(side_effect=hangup) + supervisor = CallSupervisor( + socket, + sink, + logger, + UserAPIKeyAuth(), + close_call, + ready_timeout=ready_timeout, + lifetime=lifetime, + drain_timeout=0.05, + ) + return socket, sink, close_call, supervisor + + +@pytest.mark.asyncio +async def test_observer_logs_webrtc_usage_without_client_sideband(): + socket, sink, close_call, supervisor = fixture() + await socket.messages.put({"type": "session.started"}) + await supervisor.start() + await socket.messages.put({"type": "response.done", "response": {"usage": {"total_tokens": 15}}}) + await socket.messages.put({"type": "session.closed", "usage": {"total_tokens": 19}}) + await supervisor.wait() + await supervisor.close() + assert sink.logs == 1 + assert sink.events[-1]["usage"]["total_tokens"] == 19 + assert sink.events[1]["response"]["usage"]["total_tokens"] == 15 + assert socket.closed + close_call.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_early_upstream_eof_rejects_start(): + socket, sink, close_call, supervisor = fixture() + await socket.messages.put(None) + with pytest.raises(RuntimeError, match="ended before"): + await supervisor.start() + assert socket.closed + assert sink.logs == 1 + close_call.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_cancelled_start_hangs_up_and_drains_terminal_usage(): + socket, sink, close_call, supervisor = fixture() + started = asyncio.create_task(supervisor.start()) + await asyncio.sleep(0) + started.cancel() + with pytest.raises(asyncio.CancelledError): + await started + close_call.assert_awaited_once() + assert socket.closed + assert sink.logs == 1 + assert sink.events[-1]["usage"]["total_tokens"] == 42 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cancel_count", [1, 2, 3]) +async def test_repeated_start_cancellation_keeps_lease_until_shutdown_finishes(cancel_count): + from litellm.proxy.hooks.realtime_call_lease import RealtimeCallLease + + reading = asyncio.Event() + close_entered = asyncio.Event() + allow_close = asyncio.Event() + released = asyncio.Event() + + class ObservedSocket(Socket): + async def __anext__(self): + reading.set() + return await super().__anext__() + + socket = ObservedSocket() + logger = MagicMock(spec=Logging) + logger.model_call_details = {} + sink = Sink(logger) + + async def close_call(): + close_entered.set() + await allow_close.wait() + await socket.messages.put({"type": "session.closed", "usage": {"audio_duration_ms": 1000}}) + + async def release(): + released.set() + + lease = RealtimeCallLease(renew=AsyncMock(return_value=True), release=release) + lease.start() + supervisor = CallSupervisor( + socket, sink, logger, UserAPIKeyAuth(), close_call, lease=lease, ready_timeout=10, termination_timeout=10 + ) + registry = CallSupervisors() + + async def signaling(): + transferred = False + try: + await registry.start(supervisor) + transferred = True + finally: + # The signaling endpoint retains lease ownership until registry startup succeeds. + if not transferred: + await lease.close() + + started = asyncio.create_task(signaling()) + shutdown = None + try: + await asyncio.wait_for(reading.wait(), timeout=1) + started.cancel() + await asyncio.wait_for(close_entered.wait(), timeout=1) + for _ in range(cancel_count - 1): + started.cancel() + done, _ = await asyncio.wait({started}, timeout=0.02) + assert not done + assert not released.is_set() + shutdown = asyncio.create_task(registry.shutdown()) + done, _ = await asyncio.wait({started, shutdown}, timeout=0.02) + assert not done + assert not released.is_set() + assert not socket.closed + assert sink.logs == 0 + allow_close.set() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(started, timeout=1) + await asyncio.wait_for(shutdown, timeout=1) + assert socket.closed + assert sink.logs == 1 + assert released.is_set() + assert sink.events[-1]["usage"]["audio_duration_ms"] == 1000 + finally: + allow_close.set() + await asyncio.wait_for(supervisor.wait(), timeout=1) + await asyncio.gather(started, return_exceptions=True) + if shutdown is not None: + await shutdown + await registry.shutdown() + await lease.close() + + +@pytest.mark.asyncio +async def test_worker_shutdown_drains_all_calls(): + registry = CallSupervisors() + socket, sink, close_call, supervisor = fixture() + await socket.messages.put({"type": "session.created"}) + await registry.start(supervisor) + await registry.shutdown() + await registry.shutdown() + close_call.assert_awaited_once() + assert socket.closed + assert sink.logs == 1 + assert sink.events[-1]["usage"]["total_tokens"] == 42 + + +@pytest.mark.asyncio +async def test_ready_timeout_hangs_up_before_returning_error(): + socket, sink, close_call, supervisor = fixture(ready_timeout=0.01) + with pytest.raises(asyncio.TimeoutError): + await supervisor.start() + close_call.assert_awaited_once() + assert socket.closed + assert sink.logs == 1 + + +@pytest.mark.asyncio +async def test_lifetime_limit_closes_call_and_collects_final_usage(): + socket, sink, close_call, supervisor = fixture(lifetime=0.01) + await socket.messages.put({"type": "session.started"}) + await supervisor.start() + await supervisor.wait() + close_call.assert_awaited_once() + assert socket.closed + assert sink.events[-1]["usage"]["total_tokens"] == 42 + + +@pytest.mark.asyncio +async def test_socket_eof_after_ready_still_hangs_up_provider_call(monkeypatch): + from litellm.proxy.realtime_endpoints import call_supervision + + invalidate = AsyncMock() + release = AsyncMock() + monkeypatch.setattr(call_supervision, "invalidate_budget_reservation_counters", invalidate) + monkeypatch.setattr(call_supervision, "release_or_invalidate_budget_reservation", release) + socket, sink, close_call, supervisor = fixture() + await socket.messages.put({"type": "session.started"}) + await supervisor.start() + await socket.messages.put(None) + await supervisor.wait() + close_call.assert_awaited_once() + assert socket.closed + assert sink.logs == 1 + assert sink.logger.model_call_details["realtime_usage_incomplete"] is True + invalidate.assert_awaited_once() + release.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_observer_error_rejects_start(caplog): + socket, sink, close_call, supervisor = fixture() + await socket.messages.put(RuntimeError("private-provider-credential")) + with pytest.raises(RuntimeError, match="ended before"): + await supervisor.start() + assert socket.closed + close_call.assert_awaited_once() + assert "private-provider-credential" not in caplog.text + + +@pytest.mark.asyncio +async def test_failed_logging_invalidates_reservation_without_zeroing_spend(monkeypatch): + from litellm.proxy.realtime_endpoints import call_supervision + + socket = Socket() + logger = MagicMock(spec=Logging) + logger.model_call_details = {} + sink = MagicMock() + sink.log_messages = AsyncMock(side_effect=RuntimeError("logging unavailable")) + release = AsyncMock() + invalidate = AsyncMock() + monkeypatch.setattr(call_supervision, "release_or_invalidate_budget_reservation", release) + monkeypatch.setattr(call_supervision, "invalidate_budget_reservation_counters", invalidate) + supervisor = CallSupervisor(socket, sink, logger, UserAPIKeyAuth(), AsyncMock()) + await socket.messages.put({"type": "session.started"}) + await supervisor.start() + await socket.messages.put({"type": "session.closed"}) + with pytest.raises(RuntimeError, match="logging unavailable"): + await supervisor.wait() + release.assert_not_awaited() + invalidate.assert_awaited_once_with(budget_reservation=None) + assert logger.model_call_details["realtime_accounting_incomplete"] is True + assert socket.closed + sink.log_messages.assert_awaited_once_with(wait_for_dispatch=True) + + +@pytest.mark.asyncio +async def test_start_rejects_terminal_session_while_accounting_is_pending(): + socket = Socket() + logger = MagicMock(spec=Logging) + logger.model_call_details = {} + dispatch_started = asyncio.Event() + allow_dispatch = asyncio.Event() + + async def log_messages(*, wait_for_dispatch=False): + dispatch_started.set() + await allow_dispatch.wait() + logger.model_call_details[REALTIME_SESSION_SUCCESS_LOGGED_KEY] = True + + sink = MagicMock() + sink.log_messages = AsyncMock(side_effect=log_messages) + supervisor = CallSupervisor(socket, sink, logger, UserAPIKeyAuth(), AsyncMock()) + await socket.messages.put({"type": "session.created"}) + await socket.messages.put({"type": "session.closed"}) + startup = asyncio.create_task(supervisor.start()) + try: + await asyncio.wait_for(dispatch_started.wait(), timeout=1) + finally: + allow_dispatch.set() + with pytest.raises(RuntimeError, match="ended before"): + await asyncio.wait_for(startup, timeout=1) + assert socket.closed + sink.log_messages.assert_awaited_once_with(wait_for_dispatch=True) + + +@pytest.mark.asyncio +async def test_shutdown_bounds_accounting_and_invalidates_partial_dispatch(monkeypatch): + from litellm.proxy.realtime_endpoints import call_supervision + + socket = Socket() + logger = MagicMock(spec=Logging) + logger.model_call_details = {} + dispatch_cancelled = asyncio.Event() + invalidate = AsyncMock() + release = AsyncMock() + monkeypatch.setattr(call_supervision, "invalidate_budget_reservation_counters", invalidate) + monkeypatch.setattr(call_supervision, "release_or_invalidate_budget_reservation", release) + + async def log_messages(*, wait_for_dispatch=False): + try: + await asyncio.Event().wait() + finally: + dispatch_cancelled.set() + + async def hangup(): + await socket.messages.put({"type": "session.closed", "usage": {"total_tokens": 42}}) + + sink = MagicMock() + sink.log_messages = AsyncMock(side_effect=log_messages) + supervisor = CallSupervisor(socket, sink, logger, UserAPIKeyAuth(), hangup, logging_timeout=0.01) + registry = CallSupervisors() + await socket.messages.put({"type": "session.created"}) + await registry.start(supervisor) + await asyncio.wait_for(registry.shutdown(), timeout=1) + assert dispatch_cancelled.is_set() + assert socket.closed + assert logger.model_call_details["realtime_accounting_incomplete"] is True + assert not logger.model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY) + invalidate.assert_awaited_once_with(budget_reservation=None) + release.assert_not_awaited() + sink.log_messages.assert_awaited_once_with(wait_for_dispatch=True) + + +@pytest.mark.asyncio +async def test_shutdown_waits_for_usage_dispatch_completion(): + socket = Socket() + logger = MagicMock(spec=Logging) + logger.model_call_details = {} + dispatch_started = asyncio.Event() + dispatch_complete = asyncio.Event() + dispatch_finished = asyncio.Event() + + async def log_messages(*, wait_for_dispatch=False): + assert wait_for_dispatch + dispatch_started.set() + await dispatch_complete.wait() + logger.model_call_details[REALTIME_SESSION_SUCCESS_LOGGED_KEY] = True + dispatch_finished.set() + + async def hangup(): + await socket.messages.put({"type": "session.closed", "usage": {"total_tokens": 42}}) + + sink = MagicMock() + sink.log_messages = AsyncMock(side_effect=log_messages) + registry = CallSupervisors() + supervisor = CallSupervisor(socket, sink, logger, UserAPIKeyAuth(), hangup) + await socket.messages.put({"type": "session.created"}) + await registry.start(supervisor) + shutdown = asyncio.create_task(registry.shutdown()) + try: + await asyncio.wait_for(dispatch_started.wait(), timeout=1) + assert not shutdown.done() + assert not dispatch_finished.is_set() + finally: + dispatch_complete.set() + await asyncio.wait_for(shutdown, timeout=1) + assert dispatch_finished.is_set() + sink.log_messages.assert_awaited_once_with(wait_for_dispatch=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("terminal_usage_required", [True, False]) +async def test_confirmed_hangup_without_terminal_usage_matches_protocol(terminal_usage_required): + socket = Socket() + logger = MagicMock(spec=Logging) + logger.model_call_details = {} + sink = Sink(logger) + close_call = AsyncMock() + supervisor = CallSupervisor( + socket, + sink, + logger, + UserAPIKeyAuth(), + close_call, + terminal_usage_required=terminal_usage_required, + drain_timeout=0.01, + ) + await socket.messages.put({"type": "session.created"}) + await supervisor.start() + await socket.messages.put({"type": "response.done", "response": {"usage": {"total_tokens": 17}}}) + await supervisor.close() + assert bool(logger.model_call_details.get("realtime_usage_incomplete")) == terminal_usage_required + assert sink.events[-1]["response"]["usage"]["total_tokens"] == 17 + assert sink.logs == 1 + assert socket.closed + close_call.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_shutdown_allows_hangup_longer_than_usage_drain_timeout(): + socket = Socket() + logger = MagicMock(spec=Logging) + logger.model_call_details = {} + sink = Sink(logger) + hangup_started = asyncio.Event() + allow_hangup = asyncio.Event() + hangup_finished = asyncio.Event() + + async def hangup(): + hangup_started.set() + await allow_hangup.wait() + await socket.messages.put({"type": "session.closed", "usage": {"total_tokens": 42}}) + hangup_finished.set() + + supervisor = CallSupervisor( + socket, + sink, + logger, + UserAPIKeyAuth(), + hangup, + drain_timeout=0.01, + termination_timeout=1, + terminal_usage_required=False, + ) + registry = CallSupervisors() + await socket.messages.put({"type": "session.created"}) + await registry.start(supervisor) + shutdown = asyncio.create_task(registry.shutdown()) + try: + await asyncio.wait_for(hangup_started.wait(), timeout=1) + await asyncio.sleep(0.04) + assert not shutdown.done() + assert not socket.closed + assert not hangup_finished.is_set() + finally: + allow_hangup.set() + await asyncio.wait_for(shutdown, timeout=1) + assert hangup_finished.is_set() + assert socket.closed + assert sink.logs == 1 + assert sink.events[-1]["usage"]["total_tokens"] == 42 + assert not logger.model_call_details.get("realtime_usage_incomplete") + + +@pytest.mark.asyncio +async def test_termination_timeout_cancels_hangup_and_finishes_cleanup(): + socket = Socket() + logger = MagicMock(spec=Logging) + logger.model_call_details = {} + sink = Sink(logger) + hangup_cancelled = asyncio.Event() + + async def hangup(): + try: + await asyncio.Event().wait() + finally: + hangup_cancelled.set() + + supervisor = CallSupervisor( + socket, + sink, + logger, + UserAPIKeyAuth(), + hangup, + drain_timeout=0.01, + termination_timeout=0.02, + terminal_usage_required=False, + ) + await socket.messages.put({"type": "session.created"}) + await supervisor.start() + await asyncio.wait_for(supervisor.close(), timeout=1) + assert hangup_cancelled.is_set() + assert socket.closed + assert sink.logs == 1 + assert logger.model_call_details["realtime_usage_incomplete"] is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("closure", ["eof", "normal_close", "error"]) +@pytest.mark.parametrize("hangup_succeeds", [True, False]) +async def test_ga_observer_disconnect_requires_confirmed_hangup(closure, hangup_succeeds): + from websockets.exceptions import ConnectionClosedOK + from websockets.frames import Close + + socket = Socket() + logger = MagicMock(spec=Logging) + logger.model_call_details = {} + sink = Sink(logger) + close_call = AsyncMock(side_effect=None if hangup_succeeds else RuntimeError("unconfirmed hangup")) + supervisor = CallSupervisor( + socket, + sink, + logger, + UserAPIKeyAuth(), + close_call, + terminal_usage_required=False, + drain_timeout=0.01, + ) + await socket.messages.put({"type": "session.created"}) + await supervisor.start() + await socket.messages.put( + None + if closure == "eof" + else ConnectionClosedOK(Close(1000, ""), Close(1000, ""), True) + if closure == "normal_close" + else RuntimeError("observer failed") + ) + await supervisor.wait() + close_call.assert_awaited_once() + assert bool(logger.model_call_details.get("realtime_usage_incomplete")) == (not hangup_succeeds) + assert sink.logs == 1 + assert socket.closed diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index 9def21c0573..e5bc60139fc 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -1,6 +1,7 @@ import datetime as real_datetime import smtplib from typing import Final +from unittest.mock import MagicMock, patch import pytest from fastapi import HTTPException @@ -9,15 +10,10 @@ from litellm.caching.caching import DualCache from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import ProxyErrorTypes, UserAPIKeyAuth -from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.proxy.utils import PrismaClient, ProxyLogging, get_custom_url, join_paths from litellm.types.guardrails import GuardrailEventHooks -from unittest.mock import MagicMock, patch - -from litellm.proxy.utils import get_custom_url, join_paths - - def test_get_custom_url(monkeypatch): monkeypatch.setenv("SERVER_ROOT_PATH", "/litellm") custom_url = get_custom_url(request_base_url="http://0.0.0.0:4000", route="ui/") @@ -2110,9 +2106,7 @@ def test_create_model_info_response_resolves_alias_to_deployment_model(): ] ) - response = create_model_info_response( - model_id="bedrock-claude-opus-5", provider="openai", llm_router=router - ) + response = create_model_info_response(model_id="bedrock-claude-opus-5", provider="openai", llm_router=router) finally: litellm.model_cost.clear() litellm.model_cost.update(saved_model_cost) @@ -2141,9 +2135,7 @@ def test_create_model_info_response_keeps_exact_alias_over_generalized_deploymen ] ) - response = create_model_info_response( - model_id="claude-opus-5", provider="openai", llm_router=router - ) + response = create_model_info_response(model_id="claude-opus-5", provider="openai", llm_router=router) finally: litellm.model_cost.clear() litellm.model_cost.update(saved_model_cost) @@ -2167,9 +2159,7 @@ def test_create_model_info_response_falls_back_to_alias_for_opaque_deployment_na ] ) - response = create_model_info_response( - model_id="gpt-4o", provider="openai", llm_router=router - ) + response = create_model_info_response(model_id="gpt-4o", provider="openai", llm_router=router) finally: litellm.model_cost.clear() litellm.model_cost.update(saved_model_cost) @@ -2194,9 +2184,7 @@ def test_create_model_info_response_resolves_mode_through_deployment_model(): ] ) - response = create_model_info_response( - model_id="my-embeddings", provider="openai", llm_router=router - ) + response = create_model_info_response(model_id="my-embeddings", provider="openai", llm_router=router) finally: litellm.model_cost.clear() litellm.model_cost.update(saved_model_cost) @@ -2274,7 +2262,9 @@ async def test_post_call_failure_hook_redacts_traceback_before_callbacks(monkeyp with patch.object(proxy_logging_obj, "update_request_status", new=AsyncMock()): await proxy_logging_obj.post_call_failure_hook( request_data={"metadata": {}}, - original_exception=HTTPException(status_code=400, detail="Upstream passthrough request failed with status 400"), + original_exception=HTTPException( + status_code=400, detail="Upstream passthrough request failed with status 400" + ), user_api_key_dict=UserAPIKeyAuth(), traceback_str=upstream_traceback, ) @@ -2284,6 +2274,130 @@ async def test_post_call_failure_hook_redacts_traceback_before_callbacks(monkeyp assert "REDACTED" in recorder.received_traceback +@pytest.mark.asyncio +@pytest.mark.parametrize("limiter_version", [1, 3]) +@pytest.mark.parametrize("limit", ["rpm_limit", "max_parallel_requests"]) +async def test_internal_realtime_observer_preserves_quota_and_custom_hooks(monkeypatch, limiter_version, limit): + import asyncio + from datetime import datetime + + from litellm.caching.caching import DualCache + from litellm.integrations.custom_logger import CustomLogger + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.hooks.parallel_request_limiter import _PROXY_MaxParallelRequestsHandler + from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + _PROXY_MaxParallelRequestsHandler_v3, + _request_stash, + get_request_stash, + ) + from litellm.proxy.utils import InternalUsageCache, ProxyLogging + + observed = [] + + class Hook(CustomLogger): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + observed.append(call_type) + return {**data, "extra_headers": {"x-hook": "required"}} + + cache = DualCache() + limiter_type = _PROXY_MaxParallelRequestsHandler if limiter_version == 1 else _PROXY_MaxParallelRequestsHandler_v3 + limiter = limiter_type(InternalUsageCache(dual_cache=cache)) + proxy = ProxyLogging(UserApiKeyCache()) + monkeypatch.setattr(litellm, "callbacks", [limiter, Hook()]) + token = _request_stash.set(None) + try: + auth = UserAPIKeyAuth(api_key="observer-quota-test", **{limit: 1}) + await proxy.pre_call_hook( + auth, {"model": "voice", "litellm_call_id": "signaling", "metadata": {}}, "arealtime_calls" + ) + await asyncio.sleep(0) + initial_stash = get_request_stash() + result = await proxy.pre_call_hook( + auth, + {"model": "voice", "litellm_call_id": "observer", "metadata": {}}, + "_arealtime", + internal_realtime_observer=True, + ) + assert result["extra_headers"] == {"x-hook": "required"} + assert observed == ["arealtime_calls", "_arealtime"] + if limiter_version == 3: + assert get_request_stash() is initial_stash + assert initial_stash.owner_litellm_call_id == "signaling" + if limit == "max_parallel_requests": + await limiter.async_log_success_event( + { + "litellm_call_id": "signaling", + "litellm_params": {"metadata": {"user_api_key": auth.api_key, "user_api_key_model_max_budget": {}}}, + }, + litellm.ModelResponse(usage=litellm.Usage(total_tokens=0)), + datetime.now(), + datetime.now(), + ) + if limiter_version == 3: + assert initial_stash.parallel_slot is None + await proxy.pre_call_hook( + auth, {"model": "voice", "litellm_call_id": "next", "metadata": {}}, "arealtime_calls" + ) + if limiter_version == 1 and limit == "max_parallel_requests": + from litellm.proxy._types import InternalRequestOrigin + + await asyncio.sleep(0) + observer_kwargs = { + "internal_request_origin": InternalRequestOrigin.REALTIME_OBSERVER, + "litellm_call_id": "observer", + "litellm_params": {"metadata": {"user_api_key": auth.api_key, "user_api_key_model_max_budget": {}}}, + } + await limiter.async_log_success_event( + observer_kwargs, + litellm.ModelResponse(usage=litellm.Usage(total_tokens=17)), + datetime.now(), + datetime.now(), + ) + current = await limiter.internal_usage_cache.async_get_cache( + key=f"{auth.api_key}::{datetime.now():%Y-%m-%d-%H-%M}::request_count", litellm_parent_otel_span=None + ) + assert current["current_requests"] == 1 + assert current["current_tpm"] == 17 + with pytest.raises(HTTPException) as error: + await proxy.pre_call_hook( + auth, + {"model": "voice", "litellm_call_id": "forged", "metadata": {}, "internal_realtime_observer": True}, + "_arealtime", + ) + assert error.value.status_code == 429 + finally: + _request_stash.reset(token) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("scope", ["key", "user", "team", "end_user"]) +async def test_internal_observer_missing_legacy_counter_only_adds_usage(scope): + from datetime import datetime + + from litellm.proxy._types import InternalRequestOrigin + from litellm.proxy.hooks.parallel_request_limiter import _PROXY_MaxParallelRequestsHandler + from litellm.proxy.utils import InternalUsageCache + + limiter = _PROXY_MaxParallelRequestsHandler(InternalUsageCache(dual_cache=DualCache())) + metadata = {"user_api_key": "expired-key", "user_api_key_model_max_budget": {}} + if scope in ("user", "team"): + metadata[f"user_api_key_{scope}_id"] = "expired-scope" + kwargs = { + "internal_request_origin": InternalRequestOrigin.REALTIME_OBSERVER, + "litellm_params": {"metadata": metadata}, + **({"user": "expired-scope"} if scope == "end_user" else {}), + } + await limiter.async_log_success_event( + kwargs, litellm.ModelResponse(usage=litellm.Usage(total_tokens=23)), datetime.now(), datetime.now() + ) + identity = "expired-key" if scope == "key" else "expired-scope" + current = await limiter.internal_usage_cache.async_get_cache( + key=f"{identity}::{datetime.now():%Y-%m-%d-%H-%M}::request_count", litellm_parent_otel_span=None + ) + assert current == {"current_requests": 0, "current_tpm": 23, "current_rpm": 0} + + class TestPrismaClientTokenAuthBehindThePool: """Behind the in-container pool the supervisor renews the writer's database token and hands the workers a loopback URL with a static password, so the diff --git a/tests/test_litellm/realtime_api/test_main.py b/tests/test_litellm/realtime_api/test_main.py index 0827bbcdc38..2b6bb8d123c 100644 --- a/tests/test_litellm/realtime_api/test_main.py +++ b/tests/test_litellm/realtime_api/test_main.py @@ -1,4 +1,5 @@ import asyncio +import json import time from types import TracebackType from typing import Final @@ -17,6 +18,42 @@ class FakeLogging: pass +@pytest.mark.parametrize("provider", [litellm.LlmProviders.XAI, litellm.LlmProviders.OPENAI, litellm.LlmProviders.GEMINI]) +def test_realtime_handler_factory_does_not_read_headers_without_a_handler(provider): + from litellm.types.router import GenericLiteLLMParams + + read_headers = MagicMock(side_effect=AssertionError("Headers must not be read")) + assert realtime_main.ProviderConfigManager.get_provider_realtime_handler( + provider, GenericLiteLLMParams(), read_headers + ) is None + read_headers.assert_not_called() + + +def test_realtime_handler_factory_passes_actual_chatgpt_headers(tmp_path, monkeypatch): + from litellm.llms.chatgpt.realtime import ChatGPTRealtime + from litellm.types.router import GenericLiteLLMParams + + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + monkeypatch.setenv("CHATGPT_AUTH_FILE", "auth.json") + (tmp_path / "auth.json").write_text( + json.dumps({"access_token": "factory-test-token", "account_id": "factory-account", "expires_at": time.time() + 3600}) + ) + params = GenericLiteLLMParams(litellm_session_id="factory-session") + headers = {"openai-alpha": "quicksilver=v2"} + extra_headers = {"x-gateway-route": "required"} + read_headers = MagicMock(return_value=headers) + result = realtime_main.ProviderConfigManager.get_provider_realtime_handler( + litellm.LlmProviders.CHATGPT, params, read_headers, extra_headers + ) + assert isinstance(result, ChatGPTRealtime) + read_headers.assert_called_once_with() + outgoing_headers = result._get_additional_headers("unused") + assert outgoing_headers["openai-alpha"] == headers["openai-alpha"] + assert outgoing_headers["x-gateway-route"] == extra_headers["x-gateway-route"] + assert outgoing_headers["session_id"] == "factory-session" + assert outgoing_headers["Authorization"] == "Bearer factory-test-token" + + def test_resolves_top_level_session_model(): resolved = _with_resolved_session_model({"model": "alias/gpt-realtime"}, "gpt-realtime") assert resolved == {"model": "gpt-realtime"} @@ -404,3 +441,29 @@ async def test_arealtime_azure_env_beta_protocol_wins_over_a_ga_client(monkeypat assert await _azure_backend_url_dialed_for(_GA_CLIENT) == ( "wss://my-endpoint.openai.azure.com/openai/realtime?api-version=2024-10-01-preview&deployment=gpt-realtime" ) + + +@pytest.mark.parametrize("is_call", [False, True]) +@pytest.mark.parametrize("provider", ["chatgpt", "openai", "azure"]) +def test_realtime_http_provider_controls_dynamic_base_precedence(provider, is_call, monkeypatch): + from litellm.types.router import GenericLiteLLMParams + + monkeypatch.delenv("CHATGPT_API_BASE", raising=False) + monkeypatch.delenv("OPENAI_CHATGPT_API_BASE", raising=False) + config, base, key = realtime_main._get_realtime_http_provider_config( + custom_llm_provider=provider, + dynamic_api_base="https://dynamic.example/v1", + dynamic_api_key="dynamic-key", + litellm_params=GenericLiteLLMParams(api_base="https://configured.example/v1"), + is_call=is_call, + ) + expected_base = "https://configured.example/v1" if provider == "chatgpt" else "https://dynamic.example/v1" + assert base == expected_base + assert key == ("chatgpt-oauth" if provider == "chatgpt" else "dynamic-key") + assert config is not None + if provider == "chatgpt": + assert config.get_realtime_calls_url(base, "gpt-realtime-1.5") == expected_base + "/realtime/calls" + else: + assert config.get_realtime_calls_extra_headers({"x-gateway-route": "required"}) == { + "x-gateway-route": "required" + } diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index f610821e06a..9f9db7588c6 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -4768,3 +4768,102 @@ def test_collect_and_combine_realtime_usage_stores_partitioned_text_tokens() -> assert combined.completion_tokens_details.reasoning_tokens == 95 assert combined.completion_tokens_details.text_tokens == 38 assert combined.completion_tokens_details.audio_tokens == 0 + + +def _live_terminal_event(duration=4000): + return {"type": "session.closed", "usage": {"audio_duration_ms": duration, "backend_model_usage": []}} + + +@pytest.mark.parametrize("rate,expected", [(0.025, 0.1), (0, 0), (None, 0)]) +def test_live_terminal_duration_uses_configured_second_price(monkeypatch, rate, expected): + monkeypatch.setitem( + litellm.model_cost, + "live-priced-test", + {"litellm_provider": "chatgpt", "mode": "realtime", "input_cost_per_second": rate}, + ) + assert handle_realtime_stream_cost_calculation( + [_live_terminal_event()], Usage(), "chatgpt", "live-priced-test" + ) == pytest.approx(expected) + + +def test_live_terminal_duration_honors_deployment_override(monkeypatch): + monkeypatch.setitem( + litellm.model_cost, + "live-deployment-test", + {"litellm_provider": "chatgpt", "mode": "realtime", "input_cost_per_second": 0.025}, + ) + result = RealtimeAPITokenUsageProcessor.create_logging_realtime_object(Usage(), [_live_terminal_event()]) + assert completion_cost( + completion_response=result, + model="gpt-live-1", + custom_llm_provider="chatgpt", + call_type="_arealtime", + custom_pricing=True, + router_model_id="live-deployment-test", + ) == pytest.approx(0.1) + + +@pytest.mark.parametrize("duration", [-1, True, "4000", float("inf"), float("nan"), None]) +def test_live_terminal_invalid_duration_does_not_create_spend(monkeypatch, duration): + monkeypatch.setitem( + litellm.model_cost, + "live-priced-test", + {"litellm_provider": "chatgpt", "mode": "realtime", "input_cost_per_second": 0.025}, + ) + assert ( + handle_realtime_stream_cost_calculation( + [_live_terminal_event(duration)], Usage(), "chatgpt", "live-priced-test" + ) + == 0 + ) + + +def test_live_terminal_is_not_counted_twice(monkeypatch): + monkeypatch.setitem( + litellm.model_cost, + "live-priced-test", + {"litellm_provider": "chatgpt", "mode": "realtime", "input_cost_per_second": 0.025}, + ) + assert handle_realtime_stream_cost_calculation( + [_live_terminal_event(), _live_terminal_event()], Usage(), "chatgpt", "live-priced-test" + ) == pytest.approx(0.1) + + +@pytest.mark.parametrize("with_tokens", [False, True]) +@pytest.mark.parametrize("terminal_count", [1, 2]) +@pytest.mark.parametrize("duration_priced", [False, True]) +def test_live_terminal_with_response_done_preserves_configured_billing( + monkeypatch, with_tokens, terminal_count, duration_priced +): + monkeypatch.setitem( + litellm.model_cost, + "realtime-deployment-test", + { + "litellm_provider": "chatgpt", + "mode": "realtime", + **( + {"input_cost_per_second": 0.025} + if duration_priced + else {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002} + ), + }, + ) + events = [ + { + "type": "response.done", + "response": { + "usage": ({"input_tokens": 10, "output_tokens": 5, "total_tokens": 15} if with_tokens else {}) + }, + }, + *(_live_terminal_event() for _ in range(terminal_count)), + ] + usage = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results(events) + result = RealtimeAPITokenUsageProcessor.create_logging_realtime_object(usage, events) + assert completion_cost( + completion_response=result, + model="gpt-live-1-codex" if duration_priced else "gpt-realtime-1.5", + custom_llm_provider="chatgpt", + call_type="_arealtime", + custom_pricing=True, + router_model_id="realtime-deployment-test", + ) == pytest.approx(0.1 if duration_priced else (0.02 if with_tokens else 0)) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index c7e46829aba..3342892c017 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -1153,6 +1153,8 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "/v1/messages", "/v1/images/generations", "/v1/realtime", + "/v1/realtime/calls", + "/v1/live", "/v1/realtime/transcription_sessions", "/v1/images/variations", "/v1/images/edits", diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 7ed8df6815d..38279183ec2 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -8320,6 +8320,26 @@ export interface paths { patch?: never; trace?: never; }; + "/live": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * WebSocket: realtime_websocket_endpoint + * @description WebSocket connection endpoint + */ + get: operations["websocket_realtime_websocket_endpoint_get_4"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/login": { parameters: { query?: never; @@ -18340,6 +18360,46 @@ export interface paths { patch?: never; trace?: never; }; + "/v1/live": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * WebSocket: realtime_websocket_endpoint + * @description WebSocket connection endpoint + */ + get: operations["websocket_realtime_websocket_endpoint_get_5"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/v1/live/": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * WebSocket: codex_live_sideband_endpoint + * @description WebSocket connection endpoint + */ + get: operations["websocket_codex_live_sideband_endpoint"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/v1/mcp/access_groups": { parameters: { query?: never; @@ -50839,6 +50899,24 @@ export interface operations { }; }; }; + websocket_realtime_websocket_endpoint_get_4: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description WebSocket Protocol Switched */ + 101: { + headers: { + [name: string]: unknown; + }; + content?: never; + }; + }; + }; login_login_post: { parameters: { query?: never; @@ -63159,6 +63237,42 @@ export interface operations { }; }; }; + websocket_realtime_websocket_endpoint_get_5: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description WebSocket Protocol Switched */ + 101: { + headers: { + [name: string]: unknown; + }; + content?: never; + }; + }; + }; + websocket_codex_live_sideband_endpoint: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description WebSocket Protocol Switched */ + 101: { + headers: { + [name: string]: unknown; + }; + content?: never; + }; + }; + }; get_mcp_access_groups_v1_mcp_access_groups_get: { parameters: { query?: never;