fix(router): count native Responses WebSocket turns toward deployment limits and budgets

litellm/proxy/hooks/proxy_track_cost_callback.py

@@ -471,11 +471,7 @@ class _ProxyDBLogger(CustomLogger):
                 await _release_budget_reservation(budget_reservation=budget_reservation)
                 # Non-model call types (health checks, afile_delete) have no model or standard_logging_object.
                 # Use .get() for "stream" to avoid KeyError on health checks.
-                # WS session wrappers (_aresponses_websocket, _arealtime) also reach here with
-                # result=None; their per-turn costs are tracked on the inner aresponses/realtime calls.
-                if sl_object is None and (
-                    not kwargs.get("model") or kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime")
-                ):
+                if sl_object is None and (not kwargs.get("model") or kwargs.get("call_type") == "_arealtime"):
                     verbose_proxy_logger.warning(
                         "Cost tracking - skipping, no standard_logging_object for call_type=%s",
                         kwargs.get("call_type", "unknown"),

litellm/router.py

@@ -8579,8 +8579,7 @@ class Router:
         Track remaining tpm/rpm quota for model in model_list
         """
         try:
-            # WS session wrappers fire with result=None; per-turn costs tracked by inner calls.
-            if kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime"):
+            if kwargs.get("call_type") == "_arealtime":
                 return
             if is_batch_retrieve_call_type(kwargs.get("call_type")):
                 return

litellm/router_strategy/budget_limiter.py

@@ -502,8 +502,7 @@ class RouterBudgetLimiting(CustomLogger):
     async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
         """Original method now uses helper functions"""
         verbose_router_logger.debug("in RouterBudgetLimiting.async_log_success_event")
-        # WS session wrappers fire with result=None; per-turn costs tracked by inner calls.
-        if kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime"):
+        if kwargs.get("call_type") == "_arealtime":
             return
         standard_logging_payload: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None)
         if standard_logging_payload is None:

tests/unit/proxy/response_api_endpoints/test_endpoints.py

@@ -1642,29 +1642,6 @@ class TestResponsesWSStreamingFirstMessage:
 
 
 class TestWSSessionCostTracking:
-    @pytest.mark.asyncio
-    async def test_router_budget_limiter_skips_aresponses_websocket_call_type(self):
-        """
-        RouterBudgetLimiting.async_log_success_event must not raise when
-        call_type='_aresponses_websocket', even when standard_logging_object is None.
-        Per-turn costs are tracked by individual aresponses calls inside the session;
-        the outer session wrapper fires with result=None.
-        """
-        from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
-
-        limiter = RouterBudgetLimiting.__new__(RouterBudgetLimiting)
-        kwargs = {
-            "call_type": "_aresponses_websocket",
-            "standard_logging_object": None,
-            "litellm_params": {"custom_llm_provider": "vertex_ai"},
-        }
-        await limiter.async_log_success_event(
-            kwargs=kwargs,
-            response_obj=None,
-            start_time=None,
-            end_time=None,
-        )
-
     @pytest.mark.asyncio
     async def test_router_budget_limiter_skips_arealtime_call_type(self):
         """Same guard applies to _arealtime WS session wrappers."""

tests/unit/router_strategy/test_budget_limiter.py

@@ -59,7 +59,9 @@ async def _log_success(limiter: RouterBudgetLimiting, kwargs: dict[str, object])
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize("call_type", ["aresponses", "anthropic_messages", "aembedding", "arerank"])
+@pytest.mark.parametrize(
+    "call_type", ["aresponses", "_aresponses_websocket", "anthropic_messages", "aembedding", "arerank"]
+)
 async def test_provider_spend_tracked_when_litellm_params_omits_provider(disable_budget_sync, call_type):
     """Non-chat surfaces carry the provider only on the standard logging payload."""
     limiter = RouterBudgetLimiting(

tests/unit/test_router/test_router_helper_utils.py

@@ -482,6 +482,50 @@ async def test_deployment_callback_on_success(sync_mode):
     assert tpm_key is not None
 
 
+@pytest.mark.asyncio
+async def test_deployment_callback_on_success_counts_a_responses_websocket_turn():
+    import time
+
+    model_list = [
+        {
+            "model_name": "gpt-5-mini",
+            "litellm_params": {"model": "gpt-5-mini", "api_key": "sk-fake", "rpm": 100},
+            "model_info": {"id": "ws-deployment"},
+        }
+    ]
+    router = Router(model_list=model_list)
+    standard_logging_payload = create_standard_logging_payload()
+    standard_logging_payload["total_tokens"] = 17
+
+    def turn_kwargs(call_type: str) -> dict:
+        return {
+            "call_type": call_type,
+            "litellm_params": {
+                "metadata": {"model_group": "gpt-5-mini"},
+                "model_info": {"id": "ws-deployment"},
+            },
+            "standard_logging_object": standard_logging_payload,
+        }
+
+    tpm_key = await router.deployment_callback_on_success(
+        kwargs=turn_kwargs("_aresponses_websocket"),
+        completion_response=None,
+        start_time=time.time(),
+        end_time=time.time(),
+    )
+    assert tpm_key is not None
+    assert await router.cache.async_get_cache(key=tpm_key) == 17
+
+    realtime_key = await router.deployment_callback_on_success(
+        kwargs=turn_kwargs("_arealtime"),
+        completion_response=None,
+        start_time=time.time(),
+        end_time=time.time(),
+    )
+    assert realtime_key is None
+    assert await router.cache.async_get_cache(key=tpm_key) == 17
+
+
 @pytest.mark.asyncio
 async def test_deployment_callback_on_success_tracks_tpm_for_io_deployment():
     """