From fbe9c3084cc0238760ce7b5344e35259f9e87c84 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 15 Oct 2026 20:10:43 -0700 Subject: [PATCH] fix(router): count native Responses WebSocket turns toward deployment limits and budgets --- .../proxy/hooks/proxy_track_cost_callback.py | 6 +-- litellm/router.py | 3 +- litellm/router_strategy/budget_limiter.py | 3 +- .../response_api_endpoints/test_endpoints.py | 23 ---------- .../router_strategy/test_budget_limiter.py | 4 +- .../test_router/test_router_helper_utils.py | 44 +++++++++++++++++++ 6 files changed, 50 insertions(+), 33 deletions(-) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 6fbd0b8f1b3..f6e9928ae1d 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/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"), diff --git a/litellm/router.py b/litellm/router.py index a549bb06e7d..130cdfab309 100644 --- a/litellm/router.py +++ b/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 diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index b20f71cd51f..2780455d604 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/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: diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 0eaf619ba70..dd7a0ec7a2e 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/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.""" diff --git a/tests/unit/router_strategy/test_budget_limiter.py b/tests/unit/router_strategy/test_budget_limiter.py index 685ea46547e..335c399df8a 100644 --- a/tests/unit/router_strategy/test_budget_limiter.py +++ b/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( diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index 4ac917d706d..39fa0970aa1 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/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(): """