mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
Merge pull request #39851 from BerriAI/litellm_fix_realtime_backend_close_hang
fix(realtime): relay the upstream websocket close to the client instead of hanging
This commit is contained in:
Mateo Wang
GitHub
commit
74ebe2974c
10 changed files with 800 additions and 95 deletions
|
|
@ -1820,6 +1820,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
and litellm_params.get(CallTypes.aanthropic_messages.value, False) is not True
|
||||
and litellm_params.get(CallTypes.agenerate_content.value, False) is not True
|
||||
and litellm_params.get(CallTypes.agenerate_content_stream.value, False) is not True
|
||||
and litellm_params.get(CallTypes.arealtime.value, False) is not True
|
||||
)
|
||||
|
||||
def _is_assembled_stream_success(self, result=None) -> bool:
|
||||
|
|
|
|||
|
|
@ -29,3 +29,11 @@ def websocket_close_reason(message: str, fallback: str) -> str:
|
|||
if len(encoded) <= WEBSOCKET_CLOSE_REASON_MAX_BYTES:
|
||||
return message
|
||||
return encoded[:WEBSOCKET_CLOSE_REASON_MAX_BYTES].decode("utf-8", errors="ignore")
|
||||
|
||||
|
||||
def client_close_code(upstream_code: int) -> int:
|
||||
from websockets.frames import EXTERNAL_CLOSE_CODES, CloseCode
|
||||
|
||||
if upstream_code in EXTERNAL_CLOSE_CODES or 3000 <= upstream_code < 5000:
|
||||
return upstream_code
|
||||
return int(CloseCode.INTERNAL_ERROR)
|
||||
|
|
|
|||
|
|
@ -1,12 +1,15 @@
|
|||
import asyncio
|
||||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict, cast
|
||||
import traceback
|
||||
from collections.abc import Coroutine, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum, auto
|
||||
from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypedDict, cast
|
||||
|
||||
from typing_extensions import ReadOnly
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm._logging import redact_internal_details_from_client_message, verbose_logger
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig
|
||||
from litellm.types.llms.openai import (
|
||||
|
|
@ -19,9 +22,11 @@ from litellm.types.llms.openai import (
|
|||
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
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from websockets.asyncio.client import ClientConnection
|
||||
from websockets.exceptions import ConnectionClosed
|
||||
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
||||
|
|
@ -30,8 +35,30 @@ else:
|
|||
CLIENT_CONNECTION_CLASS = Any
|
||||
|
||||
|
||||
class _ClientWebSocketExceptions(Protocol):
|
||||
ConnectionClosed: type[Exception]
|
||||
REALTIME_SESSION_SUCCESS_LOGGED_KEY: Final = "realtime_session_success_logged"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class BackendClose:
|
||||
code: int
|
||||
reason: str
|
||||
|
||||
@property
|
||||
def message(self) -> str:
|
||||
if not self.reason:
|
||||
return f"upstream websocket closed with code {self.code}"
|
||||
return f"upstream websocket closed with code {self.code}: {self.reason}"
|
||||
|
||||
|
||||
class ClientLoopExit(Enum):
|
||||
CLIENT_DISCONNECTED = auto()
|
||||
BACKEND_CLOSED = auto()
|
||||
|
||||
|
||||
def backend_close_from(error: "ConnectionClosed") -> BackendClose:
|
||||
if error.rcvd is None:
|
||||
return BackendClose(code=1006, reason=str(error))
|
||||
return BackendClose(code=error.rcvd.code, reason=error.rcvd.reason)
|
||||
|
||||
|
||||
class _ASGIScope(TypedDict, total=False):
|
||||
|
|
@ -69,10 +96,13 @@ class _ScopedWebSocket(Protocol):
|
|||
|
||||
|
||||
class _ClientWebSocket(_ScopedWebSocket, Protocol):
|
||||
exceptions: _ClientWebSocketExceptions
|
||||
|
||||
async def send_text(self, data: str) -> None: ...
|
||||
async def receive_text(self) -> str: ...
|
||||
async def close(self, code: int = 1000, reason: str | None = None) -> None: ...
|
||||
|
||||
|
||||
class _LoggingWorker(Protocol):
|
||||
def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine[object, object, None]) -> None: ...
|
||||
|
||||
|
||||
def _decode_json_object(payload: str) -> Mapping[str, object]:
|
||||
|
|
@ -108,11 +138,14 @@ class RealTimeStreaming:
|
|||
backend_uses_beta_protocol: bool | None = None,
|
||||
force_transcription_model: str | None = None,
|
||||
event_normalizer: RealtimeEventNormalizer | None = None,
|
||||
logging_worker: _LoggingWorker = GLOBAL_LOGGING_WORKER,
|
||||
):
|
||||
self.websocket: _ClientWebSocket = websocket
|
||||
self.backend_ws = backend_ws
|
||||
self.logging_obj = logging_obj
|
||||
self._logging_worker = logging_worker
|
||||
self.messages: list[OpenAIRealtimeEvents] = []
|
||||
self._backend_sent_frames: bool = False
|
||||
self.input_message: dict = {}
|
||||
self.input_messages: list[dict[str, str]] = []
|
||||
self.session_tools: list[dict] = []
|
||||
|
|
@ -388,9 +421,10 @@ 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.
|
||||
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
|
||||
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:
|
||||
"""Send a message to the backend WebSocket.
|
||||
|
|
@ -1035,60 +1069,84 @@ class RealTimeStreaming:
|
|||
return True
|
||||
return False
|
||||
|
||||
async def backend_to_client_send_messages(self):
|
||||
async def _relay_backend_messages(self) -> NoReturn:
|
||||
while True:
|
||||
try:
|
||||
raw_response = await self.backend_ws.recv(decode=False)
|
||||
except TypeError:
|
||||
raw_response = await self.backend_ws.recv()
|
||||
self._backend_sent_frames = True
|
||||
|
||||
if isinstance(raw_response, bytes):
|
||||
try:
|
||||
raw_response = raw_response.decode("utf-8")
|
||||
except UnicodeDecodeError:
|
||||
verbose_logger.warning("Received non-UTF-8 binary frame from backend, skipping.")
|
||||
continue
|
||||
|
||||
if self.provider_config:
|
||||
try:
|
||||
await self._handle_provider_config_message(raw_response)
|
||||
except Exception as e:
|
||||
verbose_logger.exception("Error processing backend message, skipping: %s", e)
|
||||
continue
|
||||
else:
|
||||
event = self._parse_backend_event(raw_response)
|
||||
if event is None:
|
||||
await self.websocket.send_text(raw_response)
|
||||
continue
|
||||
|
||||
if self._should_drop_event_from_client(event):
|
||||
continue
|
||||
|
||||
if await self._handle_raw_backend_message(event, raw_response):
|
||||
continue
|
||||
|
||||
event = self._normalize_event_for_ga_client(event)
|
||||
self.store_message(event)
|
||||
|
||||
if not self._client_wants_beta:
|
||||
await self.websocket.send_text(json.dumps(event))
|
||||
continue
|
||||
|
||||
translated = self._translate_event_to_beta(event)
|
||||
if translated is None:
|
||||
continue
|
||||
await self.websocket.send_text(json.dumps(translated))
|
||||
|
||||
async def backend_to_client_send_messages(self) -> BackendClose:
|
||||
import websockets
|
||||
|
||||
try:
|
||||
while True:
|
||||
try:
|
||||
raw_response = await self.backend_ws.recv(decode=False)
|
||||
except TypeError:
|
||||
raw_response = await self.backend_ws.recv()
|
||||
|
||||
if isinstance(raw_response, bytes):
|
||||
try:
|
||||
raw_response = raw_response.decode("utf-8")
|
||||
except UnicodeDecodeError:
|
||||
verbose_logger.warning("Received non-UTF-8 binary frame from backend, skipping.")
|
||||
continue
|
||||
|
||||
if self.provider_config:
|
||||
try:
|
||||
await self._handle_provider_config_message(raw_response)
|
||||
except Exception as e:
|
||||
verbose_logger.exception("Error processing backend message, skipping: %s", e)
|
||||
continue
|
||||
else:
|
||||
event = self._parse_backend_event(raw_response)
|
||||
if event is None:
|
||||
await self.websocket.send_text(raw_response)
|
||||
continue
|
||||
|
||||
if self._should_drop_event_from_client(event):
|
||||
continue
|
||||
|
||||
if await self._handle_raw_backend_message(event, raw_response):
|
||||
continue
|
||||
|
||||
event = self._normalize_event_for_ga_client(event)
|
||||
self.store_message(event)
|
||||
|
||||
if not self._client_wants_beta:
|
||||
await self.websocket.send_text(json.dumps(event))
|
||||
continue
|
||||
|
||||
translated = self._translate_event_to_beta(event)
|
||||
if translated is None:
|
||||
continue
|
||||
await self.websocket.send_text(json.dumps(translated))
|
||||
|
||||
await self._relay_backend_messages()
|
||||
except websockets.exceptions.ConnectionClosed as e:
|
||||
verbose_logger.exception("Connection closed in backend to client send messages - %s", e)
|
||||
except Exception as e:
|
||||
verbose_logger.exception("Error in backend to client send messages: %s", e)
|
||||
finally:
|
||||
close: Final = backend_close_from(e)
|
||||
self._flush_unbilled_transcription_usage()
|
||||
if self._backend_refused_session(close):
|
||||
await self.log_backend_refusal(e)
|
||||
else:
|
||||
await self.log_messages()
|
||||
return close
|
||||
except asyncio.CancelledError:
|
||||
self._flush_unbilled_transcription_usage()
|
||||
await self.log_messages()
|
||||
raise
|
||||
except Exception as e:
|
||||
verbose_logger.exception("Error in backend to client send messages: %s", e)
|
||||
self._flush_unbilled_transcription_usage()
|
||||
await self.log_messages()
|
||||
return BackendClose(code=1011, reason="proxy failed while relaying the upstream websocket")
|
||||
|
||||
def _backend_refused_session(self, close: BackendClose) -> bool:
|
||||
return close.code != 1000 and not self._backend_sent_frames
|
||||
|
||||
async def log_backend_refusal(self, error: Exception) -> None:
|
||||
if not self.logging_obj:
|
||||
return
|
||||
self._logging_worker.ensure_initialized_and_enqueue(
|
||||
self.logging_obj.dispatch_failure_handlers(error, traceback.format_exc(), prefer_async_handlers=True)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _detect_beta_header(websocket: _ScopedWebSocket) -> bool:
|
||||
|
|
@ -1243,11 +1301,22 @@ class RealTimeStreaming:
|
|||
item["content"] = new_content
|
||||
return item
|
||||
|
||||
async def client_ack_messages(self):
|
||||
async def _receive_client_message(self) -> str | None:
|
||||
try:
|
||||
return await self.websocket.receive_text()
|
||||
except Exception as e: # noqa: BLE001 # whatever the client socket raises, the client is gone
|
||||
verbose_logger.debug("Client disconnected: %s", e)
|
||||
return None
|
||||
|
||||
async def client_ack_messages(self) -> ClientLoopExit:
|
||||
import websockets
|
||||
|
||||
client_event: _ClientEventFrame
|
||||
try:
|
||||
while True:
|
||||
message = await self.websocket.receive_text()
|
||||
message = await self._receive_client_message()
|
||||
if message is None:
|
||||
return ClientLoopExit.CLIENT_DISCONNECTED
|
||||
|
||||
## GUARDRAIL: intercept conversation.item.create for text-based injection.
|
||||
guardrail_turn_detection_injected = False
|
||||
|
|
@ -1481,23 +1550,38 @@ class RealTimeStreaming:
|
|||
if guardrail_turn_detection_injected and sent:
|
||||
self._guardrail_turn_detection_update_sent = True
|
||||
|
||||
except websockets.exceptions.ConnectionClosed as e:
|
||||
verbose_logger.debug("Backend closed while forwarding a client message: %s", e)
|
||||
return ClientLoopExit.BACKEND_CLOSED
|
||||
except Exception as e:
|
||||
verbose_logger.debug("Error in client ack messages: %s", e)
|
||||
return ClientLoopExit.CLIENT_DISCONNECTED
|
||||
|
||||
async def bidirectional_forward(self):
|
||||
async def bidirectional_forward(self) -> None:
|
||||
forward_task: Final = asyncio.create_task(self.backend_to_client_send_messages())
|
||||
client_task: Final = asyncio.create_task(self.client_ack_messages())
|
||||
try:
|
||||
await self.client_ack_messages()
|
||||
except self.websocket.exceptions.ConnectionClosed:
|
||||
verbose_logger.debug("Connection closed")
|
||||
forward_task.cancel()
|
||||
await asyncio.wait((forward_task, client_task), return_when=asyncio.FIRST_COMPLETED)
|
||||
if client_task.done() and client_task.result() is ClientLoopExit.CLIENT_DISCONNECTED:
|
||||
return
|
||||
await self._close_client(await forward_task)
|
||||
finally:
|
||||
if not forward_task.done():
|
||||
forward_task.cancel()
|
||||
try:
|
||||
await forward_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
forward_task.cancel()
|
||||
client_task.cancel()
|
||||
await asyncio.gather(forward_task, client_task, return_exceptions=True)
|
||||
|
||||
async def _close_client(self, close: BackendClose) -> None:
|
||||
redacted_message: Final = redact_internal_details_from_client_message(close.message)
|
||||
redacted_reason: Final = redact_internal_details_from_client_message(close.reason)
|
||||
try:
|
||||
if close.code != 1000:
|
||||
await self.websocket.send_text(realtime_error_event(redacted_message, error_type="server_error"))
|
||||
await self.websocket.close(
|
||||
code=client_close_code(close.code),
|
||||
reason=websocket_close_reason(redacted_reason, fallback=redacted_message),
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # the client may already be gone; the session is over either way
|
||||
verbose_logger.debug("Could not relay the upstream close to the client: %s", e)
|
||||
|
||||
|
||||
def client_sent_openai_beta_realtime_header(websocket: _ScopedWebSocket) -> bool:
|
||||
|
|
|
|||
|
|
@ -11458,6 +11458,37 @@ def _realtime_query_params_template(model: str | None, intent: str | None) -> tu
|
|||
return tuple(params)
|
||||
|
||||
|
||||
async def _release_realtime_budget_reservation(user_api_key_dict: UserAPIKeyAuth) -> None:
|
||||
from litellm.proxy.spend_tracking.budget_reservation import (
|
||||
release_or_invalidate_budget_reservation,
|
||||
)
|
||||
|
||||
await release_or_invalidate_budget_reservation(
|
||||
budget_reservation=user_api_key_dict.budget_reservation,
|
||||
)
|
||||
|
||||
|
||||
async def _reject_realtime_session(
|
||||
websocket: WebSocket,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
*,
|
||||
code: int,
|
||||
reason: str,
|
||||
error_message: str | None = None,
|
||||
) -> None:
|
||||
try:
|
||||
if error_message is not None:
|
||||
try:
|
||||
await websocket.send_text(
|
||||
json.dumps({"type": "error", "error": {"type": "guardrail_error", "message": error_message}})
|
||||
)
|
||||
except Exception: # noqa: BLE001 # best-effort notice: a dead client socket must not skip the close below
|
||||
verbose_proxy_logger.debug("Could not send realtime pre-call error event to client; closing anyway")
|
||||
await websocket.close(code=code, reason=reason)
|
||||
finally:
|
||||
await _release_realtime_budget_reservation(user_api_key_dict)
|
||||
|
||||
|
||||
@app.websocket("/openai/v1/realtime")
|
||||
@app.websocket("/v1/realtime")
|
||||
@app.websocket("/realtime")
|
||||
|
|
@ -11483,7 +11514,9 @@ async def realtime_websocket_endpoint(
|
|||
if intent == "transcription":
|
||||
route_model = "gpt-realtime-whisper"
|
||||
else:
|
||||
await websocket.close(code=1008, reason="model query parameter is required")
|
||||
await _reject_realtime_session(
|
||||
websocket, user_api_key_dict, code=1008, reason="model query parameter is required"
|
||||
)
|
||||
return
|
||||
assert route_model is not None
|
||||
try:
|
||||
|
|
@ -11494,7 +11527,7 @@ async def realtime_websocket_endpoint(
|
|||
llm_router=llm_router,
|
||||
)
|
||||
except ProxyException as e:
|
||||
await websocket.close(code=1008, reason=e.message[:120])
|
||||
await _reject_realtime_session(websocket, user_api_key_dict, code=1008, reason=e.message[:120])
|
||||
return
|
||||
await websocket.accept(**accept_kwargs)
|
||||
|
||||
|
|
@ -11553,21 +11586,9 @@ async def realtime_websocket_endpoint(
|
|||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception("Realtime pre-call error")
|
||||
try:
|
||||
await websocket.send_text(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "error",
|
||||
"error": {
|
||||
"type": "guardrail_error",
|
||||
"message": str(e),
|
||||
},
|
||||
}
|
||||
)
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
await websocket.close(code=1011, reason="Pre-call error")
|
||||
await _reject_realtime_session(
|
||||
websocket, user_api_key_dict, code=1011, reason="Pre-call error", error_message=str(e)
|
||||
)
|
||||
return
|
||||
|
||||
# Phase 2: route to upstream LLM.
|
||||
|
|
@ -11597,6 +11618,13 @@ async def realtime_websocket_endpoint(
|
|||
)
|
||||
except Exception: # noqa: BLE001 # the lower layer may have closed the socket already; closing twice is not an error
|
||||
verbose_proxy_logger.debug("Could not close realtime client websocket; it is already gone")
|
||||
finally:
|
||||
from litellm.litellm_core_utils.realtime_streaming import (
|
||||
REALTIME_SESSION_SUCCESS_LOGGED_KEY,
|
||||
)
|
||||
|
||||
if not litellm_logging_obj.model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY):
|
||||
await _release_realtime_budget_reservation(user_api_key_dict)
|
||||
|
||||
|
||||
######################################################################
|
||||
|
|
|
|||
|
|
@ -373,6 +373,33 @@ async def invalidate_budget_reservation_counters(
|
|||
await _invalidate_spend_counter(counter_key=counter_key)
|
||||
|
||||
|
||||
async def release_or_invalidate_budget_reservation(
|
||||
budget_reservation: dict | None, # mutable-ok: stamps finalized on the caller's shared reservation dict
|
||||
) -> None:
|
||||
"""Reconcile a still-open reservation on a terminal path that settles no cost.
|
||||
|
||||
A failed or upstream-refused request never runs the success cost callback, so
|
||||
its pre-call reservation stays open and keeps the spend counter pinned above
|
||||
real spend until the counter's TTL expires, 429ing later requests on the same
|
||||
key. Release it to zero; if the release itself fails (e.g. the counter store is
|
||||
unreachable) drop the reserved counters directly and mark the reservation
|
||||
finalized so nothing reprocesses it. Idempotent: the finalized guard makes a
|
||||
second call a no-op once success or failure handling already reconciled.
|
||||
"""
|
||||
if budget_reservation is None or budget_reservation.get("finalized") is True:
|
||||
return
|
||||
try:
|
||||
await asyncio.shield(release_budget_reservation(budget_reservation=budget_reservation))
|
||||
except Exception: # noqa: BLE001 # a cleanup failure must not pin the counter; drop it directly instead
|
||||
verbose_proxy_logger.exception("Failed to release budget reservation; invalidating counters")
|
||||
try:
|
||||
await invalidate_budget_reservation_counters(budget_reservation=budget_reservation)
|
||||
except Exception: # noqa: BLE001 # nothing left to try; the finalized stamp below keeps it from being reprocessed
|
||||
verbose_proxy_logger.exception("Failed to invalidate budget reservation counters after release failed")
|
||||
finally:
|
||||
budget_reservation["finalized"] = True
|
||||
|
||||
|
||||
async def _get_budget_counters(
|
||||
request_body: dict,
|
||||
valid_token: UserAPIKeyAuth,
|
||||
|
|
|
|||
|
|
@ -27,7 +27,7 @@ from litellm.types.realtime import (
|
|||
RealtimeTranscriptionSessionRequest,
|
||||
)
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.utils import LlmProviders
|
||||
from litellm.types.utils import CallTypes, LlmProviders
|
||||
from litellm.utils import ProviderConfigManager
|
||||
|
||||
from ..litellm_core_utils.get_litellm_params import get_litellm_params
|
||||
|
|
@ -360,7 +360,7 @@ async def _arealtime(
|
|||
user: Final = kwargs.get("user", None)
|
||||
litellm_params: Final = GenericLiteLLMParams(**kwargs)
|
||||
|
||||
litellm_params_dict: Final = get_litellm_params(**kwargs)
|
||||
litellm_params_dict: Final = {**get_litellm_params(**kwargs), CallTypes.arealtime.value: True}
|
||||
|
||||
model, _custom_llm_provider, dynamic_api_key, dynamic_api_base = get_llm_provider(
|
||||
model=model,
|
||||
|
|
|
|||
|
|
@ -995,6 +995,35 @@ async def test_anthropic_messages_marks_litellm_params_async():
|
|||
litellm.callbacks = original_callbacks
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_arealtime_marks_litellm_params_async(monkeypatch):
|
||||
"""LIT-6973: ``_arealtime`` must plant ``_arealtime`` in ``litellm_params`` so
|
||||
``_is_sync_litellm_request`` classifies the session async and a failed session
|
||||
reaches a CustomLogger's failure hook once, through the async path only, even
|
||||
though the sync ``failure_handler`` still runs ahead of the async one."""
|
||||
captured = {}
|
||||
async_logged = asyncio.Event()
|
||||
|
||||
class CaptureLogger(CustomLogger):
|
||||
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
||||
captured["litellm_params"] = kwargs.get("litellm_params", {})
|
||||
async_logged.set()
|
||||
|
||||
logger = CaptureLogger()
|
||||
logger.log_failure_event = MagicMock()
|
||||
monkeypatch.setattr(litellm, "callbacks", [logger])
|
||||
monkeypatch.setattr(litellm, "failure_callback", [])
|
||||
monkeypatch.setattr(litellm, "_async_failure_callback", [])
|
||||
monkeypatch.setattr(litellm, "success_callback", [])
|
||||
monkeypatch.setattr(litellm, "_async_success_callback", [])
|
||||
with pytest.raises(ValueError, match="Unsupported model"):
|
||||
await litellm._arealtime(model="anthropic/claude-x", websocket=MagicMock())
|
||||
await asyncio.wait_for(async_logged.wait(), timeout=10)
|
||||
logger.log_failure_event.assert_not_called()
|
||||
assert captured["litellm_params"].get("_arealtime") is True
|
||||
assert LitellmLogging._is_sync_litellm_request(captured["litellm_params"]) is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_agenerate_content_marks_litellm_params_async():
|
||||
"""LIT-4475: the async ``agenerate_content`` entrypoint must plant
|
||||
|
|
@ -1180,6 +1209,7 @@ def test_is_sync_litellm_request():
|
|||
assert LitellmLogging._is_sync_litellm_request({}) is True
|
||||
assert LitellmLogging._is_sync_litellm_request({"acompletion": True}) is False
|
||||
assert LitellmLogging._is_sync_litellm_request({"allm_passthrough_route": True}) is False
|
||||
assert LitellmLogging._is_sync_litellm_request({"_arealtime": True}) is False
|
||||
assert LitellmLogging._is_sync_litellm_request({"aanthropic_messages": True}) is False
|
||||
assert LitellmLogging._is_sync_litellm_request({"agenerate_content": True}) is False
|
||||
assert LitellmLogging._is_sync_litellm_request({"agenerate_content_stream": True}) is False
|
||||
|
|
|
|||
|
|
@ -1,8 +1,10 @@
|
|||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.litellm_core_utils.realtime_errors import (
|
||||
WEBSOCKET_CLOSE_REASON_MAX_BYTES,
|
||||
client_close_code,
|
||||
realtime_error_event,
|
||||
websocket_close_reason,
|
||||
)
|
||||
|
|
@ -42,3 +44,11 @@ def test_websocket_close_reason_truncates_multibyte_message_by_bytes():
|
|||
assert len(reason.encode("utf-8")) <= WEBSOCKET_CLOSE_REASON_MAX_BYTES
|
||||
assert reason == "あ" * (WEBSOCKET_CLOSE_REASON_MAX_BYTES // 3)
|
||||
assert "<EFBFBD>" not in reason
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("upstream_code", "expected"),
|
||||
[(1000, 1000), (1008, 1008), (1011, 1011), (4001, 4001), (1005, 1011), (1006, 1011), (1015, 1011), (2999, 1011)],
|
||||
)
|
||||
def test_client_close_code_only_forwards_codes_a_server_may_send(upstream_code, expected):
|
||||
assert client_close_code(upstream_code) == expected
|
||||
|
|
|
|||
|
|
@ -1,14 +1,20 @@
|
|||
import asyncio
|
||||
import json
|
||||
from collections.abc import Coroutine
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from websockets.exceptions import ConnectionClosed
|
||||
from websockets.frames import Close
|
||||
|
||||
import litellm
|
||||
|
||||
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.litellm_core_utils.realtime_streaming import (
|
||||
REALTIME_SESSION_SUCCESS_LOGGED_KEY,
|
||||
RealTimeStreaming,
|
||||
client_sent_openai_beta_realtime_header,
|
||||
)
|
||||
|
|
@ -2941,13 +2947,11 @@ async def test_log_messages_routes_async_logging_through_bounded_worker():
|
|||
realtime turn leaves a suspended task pinning its response in memory -> an
|
||||
unbounded leak. Regression for that fix."""
|
||||
logging_obj = MagicMock()
|
||||
streaming = RealTimeStreaming(MagicMock(), MagicMock(), logging_obj)
|
||||
mock_worker = MagicMock()
|
||||
streaming = RealTimeStreaming(MagicMock(), MagicMock(), logging_obj, logging_worker=mock_worker)
|
||||
streaming.messages = [{"type": "session.created"}]
|
||||
|
||||
with (
|
||||
patch("litellm.litellm_core_utils.realtime_streaming.GLOBAL_LOGGING_WORKER") as mock_worker,
|
||||
patch("litellm.litellm_core_utils.realtime_streaming.asyncio.create_task") as mock_create_task,
|
||||
):
|
||||
with patch("litellm.litellm_core_utils.realtime_streaming.asyncio.create_task") as mock_create_task:
|
||||
await streaming.log_messages()
|
||||
|
||||
mock_worker.ensure_initialized_and_enqueue.assert_called_once()
|
||||
|
|
@ -3028,12 +3032,12 @@ async def test_session_close_flushes_unbilled_transcription_usage():
|
|||
messages before log_messages runs, and never forwarded to the client."""
|
||||
from typing import Final
|
||||
|
||||
from litellm.types.realtime import RealtimeInputAudioTranscriptionUsage
|
||||
from litellm.types.realtime import RealtimeInputAudioTranscriptionUsage, RealtimeResponseTypedDict
|
||||
|
||||
client_ws: Final = MagicMock()
|
||||
client_ws.send_text = AsyncMock()
|
||||
backend_ws: Final = MagicMock()
|
||||
backend_ws.recv = AsyncMock(side_effect=ConnectionClosed(None, None))
|
||||
backend_ws.recv = AsyncMock(side_effect=[b'{"serverContent": {}}', ConnectionClosed(None, None)])
|
||||
logging_obj: Final = MagicMock()
|
||||
logging_obj.async_success_handler = AsyncMock()
|
||||
logging_obj.success_handler = MagicMock()
|
||||
|
|
@ -3045,7 +3049,24 @@ async def test_session_close_flushes_unbilled_transcription_usage():
|
|||
"total_tokens": 171,
|
||||
"input_token_details": {"text_tokens": 0, "audio_tokens": 153},
|
||||
}
|
||||
transcript_frame: Final[RealtimeResponseTypedDict] = {
|
||||
"response": {
|
||||
"type": "conversation.item.input_audio_transcription.completed",
|
||||
"event_id": "event_1",
|
||||
"transcript": "ahoy",
|
||||
"item_id": "item_1",
|
||||
"content_index": 0,
|
||||
},
|
||||
"current_output_item_id": None,
|
||||
"current_response_id": None,
|
||||
"current_delta_chunks": None,
|
||||
"current_conversation_id": None,
|
||||
"current_item_chunks": None,
|
||||
"current_delta_type": None,
|
||||
"session_configuration_request": None,
|
||||
}
|
||||
provider_config: Final = MagicMock()
|
||||
provider_config.transform_realtime_response = MagicMock(return_value=transcript_frame)
|
||||
provider_config.unbilled_usage_on_session_close = MagicMock(return_value=usage)
|
||||
|
||||
streaming: Final = RealTimeStreaming(
|
||||
|
|
@ -3077,7 +3098,9 @@ async def test_session_close_flushes_unbilled_transcription_usage():
|
|||
)
|
||||
assert len(flushed) == 1
|
||||
assert flushed[0] in logged_snapshots[0]
|
||||
assert not client_ws.send_text.called
|
||||
forwarded: Final = tuple(json.loads(call.args[0]) for call in client_ws.send_text.await_args_list)
|
||||
assert [event.get("transcript") for event in forwarded] == ["ahoy"]
|
||||
assert all("usage" not in event for event in forwarded)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -3111,3 +3134,281 @@ async def test_session_close_flush_noop_without_unbilled_usage():
|
|||
isinstance(message, dict) and message.get("type") == "conversation.item.input_audio_transcription.completed"
|
||||
for message in streaming.messages
|
||||
)
|
||||
|
||||
|
||||
|
||||
_UPSTREAM_REFUSAL: Final = "Publisher model `publishers/google/models/gemini-live-2.5-flash` was not found"
|
||||
|
||||
|
||||
class _InlineLoggingWorker:
|
||||
def __init__(self) -> None:
|
||||
self.enqueued: tuple[Coroutine[object, object, None], ...] = ()
|
||||
|
||||
def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine[object, object, None]) -> None:
|
||||
self.enqueued = (*self.enqueued, async_coroutine)
|
||||
|
||||
async def drain(self) -> None:
|
||||
for coroutine in self.enqueued:
|
||||
await coroutine
|
||||
|
||||
|
||||
class _RecordingLogging:
|
||||
def __init__(self) -> None:
|
||||
self.model_call_details: dict[str, object] = {}
|
||||
self.logged_sessions: tuple[tuple[dict, ...], ...] = ()
|
||||
self.logged_failures: tuple[Exception, ...] = ()
|
||||
|
||||
def pre_call(self, input: str | dict, api_key: str) -> None:
|
||||
return None
|
||||
|
||||
async def dispatch_success_handlers(self, result: list[dict], prefer_async_handlers: bool = False) -> None:
|
||||
self.logged_sessions = (*self.logged_sessions, tuple(result))
|
||||
|
||||
async def dispatch_failure_handlers(
|
||||
self, exception: Exception, traceback_exception: str, prefer_async_handlers: bool = False
|
||||
) -> None:
|
||||
self.logged_failures = (*self.logged_failures, exception)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _RelaySession:
|
||||
streaming: RealTimeStreaming
|
||||
logging: _RecordingLogging
|
||||
worker: _InlineLoggingWorker
|
||||
|
||||
async def run(self) -> None:
|
||||
await asyncio.wait_for(self.streaming.bidirectional_forward(), timeout=2)
|
||||
await self.worker.drain()
|
||||
|
||||
|
||||
async def _wait_forever() -> str:
|
||||
await asyncio.Event().wait()
|
||||
raise AssertionError("unreachable")
|
||||
|
||||
|
||||
def _client_ws_that_never_sends() -> MagicMock:
|
||||
client_ws: Final = MagicMock()
|
||||
client_ws.headers = {}
|
||||
client_ws.receive_text = AsyncMock(side_effect=_wait_forever)
|
||||
client_ws.send_text = AsyncMock()
|
||||
client_ws.close = AsyncMock()
|
||||
return client_ws
|
||||
|
||||
|
||||
def _backend_ws_closing_with(*frames: bytes | Exception) -> MagicMock:
|
||||
backend_ws: Final = MagicMock()
|
||||
backend_ws.recv = AsyncMock(side_effect=list(frames))
|
||||
return backend_ws
|
||||
|
||||
|
||||
def _relay_session(client_ws: MagicMock, backend_ws: MagicMock) -> _RelaySession:
|
||||
logging: Final = _RecordingLogging()
|
||||
worker: Final = _InlineLoggingWorker()
|
||||
streaming: Final = RealTimeStreaming(
|
||||
client_ws, backend_ws, logging, model="gpt-realtime", logging_worker=worker
|
||||
)
|
||||
return _RelaySession(streaming=streaming, logging=logging, worker=worker)
|
||||
|
||||
|
||||
def _error_events_sent_to(client_ws: MagicMock) -> list[dict]:
|
||||
events: Final = (json.loads(call.args[0]) for call in client_ws.send_text.await_args_list)
|
||||
return [event for event in events if event.get("type") == "error"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bidirectional_forward_relays_upstream_policy_close_to_client():
|
||||
client_ws: Final = _client_ws_that_never_sends()
|
||||
upstream_close: Final = ConnectionClosed(Close(1008, _UPSTREAM_REFUSAL), None)
|
||||
session: Final = _relay_session(client_ws, _backend_ws_closing_with(upstream_close))
|
||||
|
||||
await session.run()
|
||||
|
||||
(error_event,) = _error_events_sent_to(client_ws)
|
||||
assert error_event["error"]["type"] == "server_error"
|
||||
assert "1008" in error_event["error"]["message"]
|
||||
assert _UPSTREAM_REFUSAL in error_event["error"]["message"]
|
||||
client_ws.close.assert_awaited_once_with(code=1008, reason=_UPSTREAM_REFUSAL)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"leaked_detail",
|
||||
(
|
||||
pytest.param("sk-live-abcdef0123456789abcdef0123", id="credential"),
|
||||
pytest.param("vertex-int.svc.cluster.local", id="internal-hostname"),
|
||||
pytest.param("/etc/litellm/service-account.json", id="filesystem-path"),
|
||||
),
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_upstream_close_details_are_scrubbed_before_reaching_the_client(leaked_detail: str):
|
||||
"""LIT-6973: the relayed close goes through the proxy's client-facing redaction, so an upstream
|
||||
error echoing a credential, an internal host, or a server path never reaches the client verbatim."""
|
||||
client_ws: Final = _client_ws_that_never_sends()
|
||||
upstream_close: Final = ConnectionClosed(Close(1008, f"upstream rejected: {leaked_detail}"), None)
|
||||
session: Final = _relay_session(client_ws, _backend_ws_closing_with(upstream_close))
|
||||
|
||||
await session.run()
|
||||
|
||||
(error_event,) = _error_events_sent_to(client_ws)
|
||||
assert leaked_detail not in error_event["error"]["message"]
|
||||
relayed_reason: Final = client_ws.close.await_args.kwargs["reason"]
|
||||
assert leaked_detail not in relayed_reason
|
||||
assert "REDACTED" in relayed_reason
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bidirectional_forward_maps_abnormal_upstream_close_to_internal_error():
|
||||
client_ws: Final = _client_ws_that_never_sends()
|
||||
session: Final = _relay_session(client_ws, _backend_ws_closing_with(ConnectionClosed(None, None)))
|
||||
|
||||
await session.run()
|
||||
|
||||
(error_event,) = _error_events_sent_to(client_ws)
|
||||
assert "1006" in error_event["error"]["message"]
|
||||
client_ws.close.assert_awaited_once()
|
||||
assert client_ws.close.await_args.kwargs["code"] == 1011
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bidirectional_forward_relays_normal_upstream_close_without_error_event():
|
||||
client_ws: Final = _client_ws_that_never_sends()
|
||||
session: Final = _relay_session(client_ws, _backend_ws_closing_with(ConnectionClosed(Close(1000, ""), None)))
|
||||
|
||||
await session.run()
|
||||
|
||||
assert _error_events_sent_to(client_ws) == []
|
||||
client_ws.close.assert_awaited_once()
|
||||
assert client_ws.close.await_args.kwargs["code"] == 1000
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upstream_refusal_before_any_frame_logs_a_failure_not_a_success():
|
||||
upstream_close: Final = ConnectionClosed(Close(1008, _UPSTREAM_REFUSAL), None)
|
||||
session: Final = _relay_session(_client_ws_that_never_sends(), _backend_ws_closing_with(upstream_close))
|
||||
|
||||
await session.run()
|
||||
|
||||
assert session.logging.logged_failures == (upstream_close,)
|
||||
assert session.logging.logged_sessions == ()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upstream_refusal_after_a_synthetic_session_created_still_logs_a_failure():
|
||||
"""LIT-6973: deferred Gemini Live setup stores a synthetic ``session.created`` before
|
||||
the relay starts. It is not an upstream frame, so a refusal after it is still a refusal."""
|
||||
upstream_close: Final = ConnectionClosed(Close(1008, _UPSTREAM_REFUSAL), None)
|
||||
session: Final = _relay_session(_client_ws_that_never_sends(), _backend_ws_closing_with(upstream_close))
|
||||
session.streaming.store_message(json.dumps({"type": "session.created", "session": {"id": "sess_synthetic"}}))
|
||||
|
||||
await session.run()
|
||||
|
||||
assert session.logging.logged_failures == (upstream_close,)
|
||||
assert session.logging.logged_sessions == ()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upstream_close_after_relayed_events_still_logs_the_session_as_success():
|
||||
client_ws: Final = _client_ws_that_never_sends()
|
||||
session_created: Final = json.dumps({"type": "session.created", "session": {"id": "sess_1"}}).encode()
|
||||
upstream_close: Final = ConnectionClosed(Close(1008, _UPSTREAM_REFUSAL), None)
|
||||
session: Final = _relay_session(client_ws, _backend_ws_closing_with(session_created, upstream_close))
|
||||
|
||||
await session.run()
|
||||
|
||||
(logged_session,) = session.logging.logged_sessions
|
||||
assert [event["type"] for event in logged_session] == ["session.created"]
|
||||
assert session.logging.logged_failures == ()
|
||||
client_ws.close.assert_awaited_once_with(code=1008, reason=_UPSTREAM_REFUSAL)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upstream_closing_while_a_client_message_is_forwarded_still_reaches_the_client():
|
||||
upstream_close: Final = ConnectionClosed(Close(1008, _UPSTREAM_REFUSAL), None)
|
||||
backend_closed: Final = asyncio.Event()
|
||||
client_messages: Final = iter((json.dumps({"type": "response.create"}),))
|
||||
|
||||
async def receive_text() -> str:
|
||||
message = next(client_messages, None)
|
||||
return message if message is not None else await _wait_forever()
|
||||
|
||||
async def send_to_backend(_message: str) -> None:
|
||||
backend_closed.set()
|
||||
raise upstream_close
|
||||
|
||||
async def recv_from_backend() -> bytes:
|
||||
await backend_closed.wait()
|
||||
raise upstream_close
|
||||
|
||||
client_ws: Final = _client_ws_that_never_sends()
|
||||
client_ws.receive_text = receive_text
|
||||
backend_ws: Final = MagicMock()
|
||||
backend_ws.send = send_to_backend
|
||||
backend_ws.recv = recv_from_backend
|
||||
session: Final = _relay_session(client_ws, backend_ws)
|
||||
|
||||
await session.run()
|
||||
|
||||
(error_event,) = _error_events_sent_to(client_ws)
|
||||
assert _UPSTREAM_REFUSAL in error_event["error"]["message"]
|
||||
client_ws.close.assert_awaited_once_with(code=1008, reason=_UPSTREAM_REFUSAL)
|
||||
assert session.logging.logged_failures == (upstream_close,)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_hanging_up_first_ends_the_session_without_a_relayed_close():
|
||||
client_ws: Final = _client_ws_that_never_sends()
|
||||
client_ws.receive_text = AsyncMock(side_effect=RuntimeError("client went away"))
|
||||
backend_ws: Final = MagicMock()
|
||||
backend_ws.recv = AsyncMock(side_effect=_wait_forever)
|
||||
session: Final = _relay_session(client_ws, backend_ws)
|
||||
|
||||
await session.run()
|
||||
|
||||
assert session.logging.logged_sessions == ((),)
|
||||
assert session.logging.logged_failures == ()
|
||||
client_ws.close.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_hanging_up_with_a_websockets_close_is_not_mistaken_for_the_backend_closing():
|
||||
client_ws: Final = _client_ws_that_never_sends()
|
||||
client_ws.receive_text = AsyncMock(side_effect=ConnectionClosed(None, None))
|
||||
backend_ws: Final = MagicMock()
|
||||
backend_ws.recv = AsyncMock(side_effect=_wait_forever)
|
||||
session: Final = _relay_session(client_ws, backend_ws)
|
||||
|
||||
await session.run()
|
||||
|
||||
assert session.logging.logged_sessions == ((),)
|
||||
assert session.logging.logged_failures == ()
|
||||
client_ws.close.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_success_logging_stamps_the_reservation_ownership_marker():
|
||||
"""LIT-6973: only the success path enqueues the cost callback that settles the
|
||||
session's budget reservation, so it stamps REALTIME_SESSION_SUCCESS_LOGGED_KEY on
|
||||
the shared logging object. The proxy endpoint reads that stamp to decide whether to
|
||||
release the reservation itself, so a logged-as-success session must carry it."""
|
||||
client_ws: Final = _client_ws_that_never_sends()
|
||||
session_created: Final = json.dumps({"type": "session.created", "session": {"id": "sess_1"}}).encode()
|
||||
upstream_close: Final = ConnectionClosed(Close(1000, ""), None)
|
||||
session: Final = _relay_session(client_ws, _backend_ws_closing_with(session_created, upstream_close))
|
||||
|
||||
await session.run()
|
||||
|
||||
assert session.logging.logged_sessions != ()
|
||||
assert session.logging.model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY) is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refused_session_does_not_stamp_the_reservation_ownership_marker():
|
||||
"""A refused session logs a failure, not a success, so it must not stamp
|
||||
REALTIME_SESSION_SUCCESS_LOGGED_KEY. If it did, the proxy endpoint would skip its
|
||||
own reservation release and the refused session's reservation would stay pinned."""
|
||||
upstream_close: Final = ConnectionClosed(Close(1008, _UPSTREAM_REFUSAL), None)
|
||||
session: Final = _relay_session(_client_ws_that_never_sends(), _backend_ws_closing_with(upstream_close))
|
||||
|
||||
await session.run()
|
||||
|
||||
assert session.logging.logged_failures == (upstream_close,)
|
||||
assert REALTIME_SESSION_SUCCESS_LOGGED_KEY not in session.logging.model_call_details
|
||||
|
|
|
|||
|
|
@ -9521,6 +9521,222 @@ def test_realtime_websocket_route_aliases_registered():
|
|||
)
|
||||
|
||||
|
||||
def _lit6973_fake_realtime_ws() -> MagicMock:
|
||||
ws = MagicMock()
|
||||
ws.headers = {}
|
||||
ws.scope = {"headers": [], "type": "websocket"}
|
||||
ws.url = "ws://testserver/v1/realtime"
|
||||
ws.accept = AsyncMock()
|
||||
ws.send_text = AsyncMock()
|
||||
ws.close = AsyncMock()
|
||||
return ws
|
||||
|
||||
|
||||
async def _lit6973_drive_realtime_session(
|
||||
reservation: dict,
|
||||
*,
|
||||
backend_logged_success: bool,
|
||||
phase_one_exit: str | None = None,
|
||||
websocket: MagicMock | None = None,
|
||||
) -> MagicMock:
|
||||
"""Drive realtime_websocket_endpoint through one of its reservation-settling exits.
|
||||
|
||||
phase_one_exit picks a rejection before the relay: "model_access" makes the
|
||||
key/model check raise ProxyException, "pre_call" makes pre-call processing
|
||||
(rate limits, guardrails) raise. Neither reaches route_request, so no success
|
||||
log can own the reservation and the endpoint has to release it on that exit.
|
||||
|
||||
route_request resolves normally in both cases: the relay owns the session
|
||||
once route_request returns. A successful session enqueues its success cost
|
||||
callback and stamps REALTIME_SESSION_SUCCESS_LOGGED_KEY on the shared logging
|
||||
object; a refused one does neither. The endpoint keys its reservation cleanup
|
||||
off that stamp, so backend_logged_success reproduces both branches. The fake
|
||||
logging object carries a real model_call_details dict so the stamp is
|
||||
observable, and the reservation has empty entries so the real release touches
|
||||
no counter store."""
|
||||
from litellm.litellm_core_utils.realtime_streaming import REALTIME_SESSION_SUCCESS_LOGGED_KEY
|
||||
from litellm.proxy import proxy_server as ps
|
||||
|
||||
user_api_key_dict: Final = UserAPIKeyAuth(api_key="sk-test", token="hashed-token")
|
||||
user_api_key_dict.budget_reservation = reservation
|
||||
|
||||
logging_obj: Final = MagicMock()
|
||||
logging_obj.model_call_details = {}
|
||||
|
||||
async def fake_llm_call() -> None:
|
||||
if backend_logged_success:
|
||||
logging_obj.model_call_details[REALTIME_SESSION_SUCCESS_LOGGED_KEY] = True
|
||||
|
||||
from litellm.proxy._types import ProxyException
|
||||
|
||||
model_access_error: Final = (
|
||||
ProxyException(message="key cannot access model", type="auth_error", param="model", code=401)
|
||||
if phase_one_exit == "model_access"
|
||||
else None
|
||||
)
|
||||
pre_call_error: Final = Exception("Rate limit exceeded") if phase_one_exit == "pre_call" else None
|
||||
pre_call: Final = AsyncMock(
|
||||
side_effect=pre_call_error, return_value=({"model": "vertex_ai/gemini-live-2.5-flash"}, logging_obj)
|
||||
)
|
||||
ws: Final = websocket if websocket is not None else _lit6973_fake_realtime_ws()
|
||||
can_call = patch.object(ps, "can_key_call_resolved_model", new=AsyncMock(side_effect=model_access_error)) # test-quality-ok: no HTTP boundary; fakes in-process auth to reach the exit under test
|
||||
pre = patch.object(ps.ProxyBaseLLMRequestProcessing, "common_processing_pre_call_logic", new=pre_call) # test-quality-ok: fakes phase-1 wiring; assertion checks observable reservation state
|
||||
route = patch.object(ps, "route_request", new=AsyncMock(return_value=fake_llm_call())) # test-quality-ok: fakes the relay whose success/refusal outcome the endpoint reads off the logging object
|
||||
with can_call, pre, route:
|
||||
await ps.realtime_websocket_endpoint(
|
||||
websocket=ws,
|
||||
model="vertex_ai/gemini-live-2.5-flash",
|
||||
intent=None,
|
||||
guardrails=None,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
return ws
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refused_realtime_session_releases_the_budget_reservation():
|
||||
"""LIT-6973: a refused realtime session enqueues no success cost callback, so
|
||||
the pre-call reservation would stay open and pin the key/team/user spend
|
||||
counters, locking the key after a couple of refusals. The endpoint sees no
|
||||
success stamp and reconciles it: the reservation ends up finalized."""
|
||||
reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []}
|
||||
|
||||
await _lit6973_drive_realtime_session(reservation, backend_logged_success=False)
|
||||
|
||||
assert reservation["finalized"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_realtime_session_rejected_in_pre_call_releases_the_budget_reservation():
|
||||
"""A rate-limit or guardrail rejection happens before route_request, so the
|
||||
relay never runs and no success log can own the reservation. The endpoint
|
||||
must release it on that exit too, or the key stays pinned at the reserved
|
||||
amount and its next requests 429 with budget_exceeded while /key/info shows
|
||||
spend 0 (reproduced live with rpm_limit=1). The client still gets the
|
||||
pre-call error event and the 1011 close it got before."""
|
||||
reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []}
|
||||
|
||||
ws: Final = await _lit6973_drive_realtime_session(
|
||||
reservation, backend_logged_success=False, phase_one_exit="pre_call"
|
||||
)
|
||||
|
||||
assert reservation["finalized"] is True
|
||||
assert json.loads(ws.send_text.await_args.args[0])["error"]["message"] == "Rate limit exceeded"
|
||||
ws.close.assert_awaited_once_with(code=1011, reason="Pre-call error")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_realtime_session_denied_model_access_releases_the_budget_reservation():
|
||||
"""The key/model access check rejects before the socket is even accepted;
|
||||
that exit skipped the release as well, pinning the reservation."""
|
||||
reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []}
|
||||
|
||||
ws: Final = await _lit6973_drive_realtime_session(
|
||||
reservation, backend_logged_success=False, phase_one_exit="model_access"
|
||||
)
|
||||
|
||||
assert reservation["finalized"] is True
|
||||
ws.close.assert_awaited_once_with(code=1008, reason="key cannot access model")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rejected_realtime_session_closes_the_client_before_releasing_the_reservation():
|
||||
"""The counter release can block on a slow or unreachable store, and a
|
||||
rejected client must not sit behind it: the relay's own failure path closes
|
||||
the client first and releases in its finally, so the pre-relay rejection
|
||||
has to close first as well. The fake close checks the reservation is still
|
||||
open when the client is closed."""
|
||||
reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []}
|
||||
ws: Final = _lit6973_fake_realtime_ws()
|
||||
|
||||
async def close_while_reservation_is_still_open(**_: object) -> None:
|
||||
assert reservation["finalized"] is False, "client was closed only after the reservation release"
|
||||
|
||||
ws.close = AsyncMock(side_effect=close_while_reservation_is_still_open)
|
||||
|
||||
await _lit6973_drive_realtime_session(
|
||||
reservation, backend_logged_success=False, phase_one_exit="pre_call", websocket=ws
|
||||
)
|
||||
|
||||
ws.close.assert_awaited_once_with(code=1011, reason="Pre-call error")
|
||||
assert reservation["finalized"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rejected_realtime_session_releases_the_reservation_when_the_client_is_already_gone():
|
||||
"""A client that hung up before the rejection makes the close raise; the
|
||||
reservation must still be released, or the key stays pinned."""
|
||||
reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []}
|
||||
ws: Final = _lit6973_fake_realtime_ws()
|
||||
ws.close = AsyncMock(side_effect=RuntimeError("client already disconnected"))
|
||||
|
||||
with pytest.raises(RuntimeError, match="client already disconnected"):
|
||||
await _lit6973_drive_realtime_session(
|
||||
reservation, backend_logged_success=False, phase_one_exit="model_access", websocket=ws
|
||||
)
|
||||
|
||||
assert reservation["finalized"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_successful_realtime_session_leaves_the_reservation_for_the_cost_callback():
|
||||
"""A billable realtime session settles its reservation through the enqueued
|
||||
success cost callback, not the endpoint. The endpoint must not finalize it in
|
||||
its finally, or it would reconcile the reservation to zero before the cost
|
||||
callback applies real spend, so billable sessions stop counting against budget.
|
||||
With the success stamp present, the endpoint leaves the reservation untouched."""
|
||||
reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []}
|
||||
|
||||
await _lit6973_drive_realtime_session(reservation, backend_logged_success=True)
|
||||
|
||||
assert reservation["finalized"] is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_release_or_invalidate_falls_back_to_invalidating_the_counters():
|
||||
"""If releasing the reservation itself fails (e.g. the counter store is down),
|
||||
the reserved counters must be invalidated directly so the estimate does not
|
||||
stay pinned, and the reservation is finalized so nothing reprocesses it."""
|
||||
from litellm.proxy import proxy_server as ps
|
||||
from litellm.proxy.spend_tracking import budget_reservation as br
|
||||
|
||||
reservation: Final = {
|
||||
"reserved_cost": 0.55,
|
||||
"input_cost": 0.0,
|
||||
"finalized": False,
|
||||
"entries": [{"counter_key": "spend:key:hashed-token"}],
|
||||
}
|
||||
invalidated: Final[list[str]] = []
|
||||
|
||||
async def _record(counter_key: str) -> None:
|
||||
invalidated.append(counter_key)
|
||||
|
||||
failing_release = patch.object(br, "release_budget_reservation", new=AsyncMock(side_effect=RuntimeError("counter store down"))) # test-quality-ok: forces the failure branch; assertion observes which counter key got invalidated
|
||||
sink = patch.object(ps, "_invalidate_spend_counter", new=_record) # test-quality-ok: fakes the counter-store sink so the invalidated key is observable
|
||||
with failing_release, sink:
|
||||
await br.release_or_invalidate_budget_reservation(budget_reservation=reservation)
|
||||
|
||||
assert invalidated == ["spend:key:hashed-token"]
|
||||
assert reservation["finalized"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_release_or_invalidate_finalizes_even_when_the_invalidate_fallback_fails():
|
||||
"""Both counter-store calls failing must not raise out of the realtime
|
||||
endpoint's finally (it would mask the session's own outcome) and must still
|
||||
stamp finalized so nothing retries the same reservation."""
|
||||
from litellm.proxy.spend_tracking import budget_reservation as br
|
||||
|
||||
reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []}
|
||||
failing_release = patch.object(br, "release_budget_reservation", new=AsyncMock(side_effect=RuntimeError("counter store down"))) # test-quality-ok: forces the fallback branch
|
||||
failing_invalidate = patch.object(br, "invalidate_budget_reservation_counters", new=AsyncMock(side_effect=RuntimeError("still down"))) # test-quality-ok: forces the fallback itself to fail
|
||||
|
||||
with failing_release, failing_invalidate:
|
||||
await br.release_or_invalidate_budget_reservation(budget_reservation=reservation)
|
||||
|
||||
assert reservation["finalized"] is True
|
||||
|
||||
|
||||
class TestTransformRequestBannedParams:
|
||||
"""
|
||||
/utils/transform_request applies the same banned-param check as LLM endpoints.
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue