fix(realtime): settle the budget reservation only for sessions the success log does not own

The blanket finally release from the previous commit also zeroed the reservation
of successful sessions. Success settlement is enqueued on the logging worker, not
awaited, so the endpoint's finally ran first and released the reservation the cost
callback still had to reconcile, dropping the real spend from the key/team/user
counters.

The relay now stamps a synchronous marker (REALTIME_SESSION_SUCCESS_LOGGED_KEY) on
the shared logging object at the single success-dispatch site, and the endpoint
releases the reservation only when that marker is absent. Refused or failed
sessions, which never log success, still release; successful sessions leave the
reservation for the cost callback to settle to actual spend. Exactly one settler
touches each reservation, so the idempotent reconcile never double-adjusts.

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.

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)
 
 
 ######################################################################

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

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),