diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index c670278d3fb..75046f2cf87 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -35,6 +35,9 @@ else: CLIENT_CONNECTION_CLASS = Any +REALTIME_SESSION_SUCCESS_LOGGED_KEY: Final = "realtime_session_success_logged" + + @dataclass(frozen=True, slots=True) class BackendClose: code: int @@ -421,6 +424,7 @@ class RealTimeStreaming: 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. diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 06cfee45918..7d59dfa86c4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11603,7 +11603,12 @@ 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: - await _release_realtime_budget_reservation(user_api_key_dict) + 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) ###################################################################### 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 5352c894f87..9c0f6f59463 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -14,6 +14,7 @@ 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, ) @@ -3380,3 +3381,34 @@ async def test_client_hanging_up_with_a_websockets_close_is_not_mistaken_for_the 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 diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 57e2cfc3332..697d4c182c7 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -9532,27 +9532,34 @@ def _lit6973_fake_realtime_ws() -> MagicMock: return ws -async def _lit6973_drive_refused_realtime_session(reservation: dict) -> None: - """Drive realtime_websocket_endpoint through a session the upstream refused. +async def _lit6973_drive_realtime_session(reservation: dict, *, backend_logged_success: bool) -> None: + """Drive realtime_websocket_endpoint to just before its budget-reservation finally. - route_request resolves normally because the relay handles the refusal - internally (sends the error event, closes the client), so neither the - success cost callback nor a failure hook runs on _ProxyDBLogger. The - endpoint itself must reconcile the pre-call budget reservation, so the - real release runs (entries is empty, so it touches no counter store) and - the caller asserts on the observable reservation state afterwards.""" + 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 - completed: Final = asyncio.get_running_loop().create_future() - completed.set_result(None) + logging_obj: Final = MagicMock() + logging_obj.model_call_details = {} - pre_call: Final = AsyncMock(return_value=({"model": "vertex_ai/gemini-live-2.5-flash"}, MagicMock())) + async def fake_llm_call() -> None: + if backend_logged_success: + logging_obj.model_call_details[REALTIME_SESSION_SUCCESS_LOGGED_KEY] = True + + pre_call: Final = AsyncMock(return_value=({"model": "vertex_ai/gemini-live-2.5-flash"}, logging_obj)) can_call = patch.object(ps, "can_key_call_resolved_model", new=AsyncMock()) # test-quality-ok: no HTTP boundary; fakes in-process auth to reach the finally 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=completed)) # test-quality-ok: fakes the relay that already handled the refusal so the session returns normally + 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=_lit6973_fake_realtime_ws(), @@ -9565,17 +9572,31 @@ async def _lit6973_drive_refused_realtime_session(reservation: dict) -> None: @pytest.mark.asyncio async def test_refused_realtime_session_releases_the_budget_reservation(): - """LIT-6973: reclassifying a refused realtime session as a failure removed the - success-path reservation release, so the pre-call reservation stayed open and - pinned the key/team/user spend counters, locking the key after a couple of - refusals. The endpoint must reconcile it: the reservation ends up finalized.""" + """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_refused_realtime_session(reservation) + await _lit6973_drive_realtime_session(reservation, backend_logged_success=False) 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),