From cb264540343fd04a9abcf8317e04fe693e9acb38 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 16 Sep 2026 19:28:50 +0000 Subject: [PATCH] test(router): cover success callback racing the pre-header count and drop redundant docstrings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 1 - .../test_router_helper_utils.py | 77 +++++++++++++------ 2 files changed, 53 insertions(+), 25 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 71588109c45..cc11a21d966 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -8202,7 +8202,6 @@ class Router: response: object, request_kwargs: dict[str, object], ) -> None: - """Count the request before the headers are read; the success callback adds only what is still missing""" if response is None: return try: diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index 8646843c0de..664f15889d3 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -11,7 +11,12 @@ import pytest import litellm from unittest.mock import patch, MagicMock, AsyncMock from create_mock_standard_logging_payload import create_standard_logging_payload -from litellm.types.utils import StandardLoggingPayload +from litellm.types.utils import ModelResponse, StandardLoggingPayload +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.types.caching import RedisPipelineIncrementOperation +from litellm.router_utils.router_callbacks.track_deployment_metrics import get_deployment_successes_for_current_minute from litellm.types.router import Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY @@ -930,10 +935,6 @@ async def test_set_response_headers(model_list): @pytest.mark.asyncio async def test_set_response_headers_passes_through_post_increment_counters(model_list): - """ - LIT-3058: `make_call` increments the router's TPM/RPM counter before the headers - are built, so `set_response_headers` writes the remaining values it reads as-is. - """ from pydantic import BaseModel class _Usage(BaseModel): @@ -980,17 +981,12 @@ def _rpm_tpm_router(model_id: str) -> Router: ) -def _ratelimit_headers(response) -> dict: +def _ratelimit_headers(response: ModelResponse | CustomStreamWrapper) -> dict[str, int]: return {k: v for k, v in response._hidden_params["additional_headers"].items() if k.startswith("x-ratelimit-")} @pytest.mark.asyncio async def test_acompletion_headers_read_post_increment_counter_and_count_once(): - """ - LIT-3058 regression: the remaining-* headers on the response must already include - this request, and the request must land in the router counter exactly once even - though `deployment_callback_on_success` still runs after the response returns. - """ router = _rpm_tpm_router("lit-3058-async") response = await router.acompletion( @@ -1010,7 +1006,6 @@ async def test_acompletion_headers_read_post_increment_counter_and_count_once(): @pytest.mark.asyncio async def test_acompletion_wildcard_route_headers_and_counter_use_resolved_deployment_name(): - """The counter key is written under the resolved model name, which is what the usage reader looks up.""" router = Router( model_list=[ { @@ -1034,10 +1029,6 @@ async def test_acompletion_wildcard_route_headers_and_counter_use_resolved_deplo @pytest.mark.asyncio async def test_acompletion_stream_counts_request_before_headers_and_tokens_once_on_completion(): - """ - A stream has no usage when the headers are built: the request is counted before the - headers and the final token usage is added once when the stream completes. - """ router = _rpm_tpm_router("lit-3058-stream") stream = await router.acompletion( @@ -1062,10 +1053,6 @@ async def test_acompletion_stream_counts_request_before_headers_and_tokens_once_ @pytest.mark.asyncio async def test_deployment_callback_on_success_adds_only_uncounted_tokens(): - """ - When the request was already counted before the headers, the success callback adds - only the tokens that were not known at that point and never a second request. - """ import time router = _rpm_tpm_router("lit-3058-callback") @@ -1094,9 +1081,54 @@ async def test_deployment_callback_on_success_adds_only_uncounted_tokens(): assert await router.get_model_group_usage("gpt-5-mini") == (40, None) +class _GatedIncrementCache(DualCache): + def __init__(self) -> None: + super().__init__(in_memory_cache=InMemoryCache()) + self.first_increment_started = asyncio.Event() + self.release_first_increment = asyncio.Event() + self.increment_calls = 0 + + async def async_increment_cache_pipeline( + self, + increment_list: list[RedisPipelineIncrementOperation], + local_only: bool = False, + parent_otel_span: object = None, + **kwargs: object, + ) -> list[float] | None: + self.increment_calls += 1 + if self.increment_calls == 1: + self.first_increment_started.set() + await self.release_first_increment.wait() + return await super().async_increment_cache_pipeline( + increment_list, local_only=local_only, parent_otel_span=parent_otel_span, **kwargs + ) + + +@pytest.mark.asyncio +async def test_success_callback_running_during_pre_header_increment_does_not_double_count(): + router = _rpm_tpm_router("lit-3058-race") + cache = _GatedIncrementCache() + router.cache = cache + + request = asyncio.ensure_future( + router.acompletion(model="gpt-5-mini", messages=[{"role": "user", "content": "hi"}], mock_response="pong") + ) + await asyncio.wait_for(cache.first_increment_started.wait(), timeout=5) + for _ in range(50): + if get_deployment_successes_for_current_minute(router, "lit-3058-race") == 1: + break + await asyncio.sleep(0.1) + assert get_deployment_successes_for_current_minute(router, "lit-3058-race") == 1 + assert cache.increment_calls == 1 + + cache.release_first_increment.set() + response = await request + + assert await router.get_model_group_usage("gpt-5-mini") == (response.usage.total_tokens, 1) + + @pytest.mark.asyncio async def test_increment_deployment_usage_for_response_skips_session_wrappers(): - """WebSocket and realtime session wrappers return None and are not counted as a request.""" router = _rpm_tpm_router("lit-3058-ws") request_kwargs = { "model": "gpt-5-mini", @@ -1179,9 +1211,6 @@ def _shared_redis_stub(store: dict) -> MagicMock: @pytest.mark.asyncio async def test_headers_on_fresh_worker_reflect_shared_redis_usage(): - from litellm.caching.dual_cache import DualCache - from litellm.caching.in_memory_cache import InMemoryCache - store: dict = {} worker_a = _rpm_tpm_router("lit-3058-workers") worker_b = _rpm_tpm_router("lit-3058-workers")