Merge d35e5681ec into 7e4ad6eb33

litellm/litellm_core_utils/litellm_logging.py

@@ -901,6 +901,7 @@ class Logging(LiteLLMLoggingBaseClass):
         self._async_success_scheduled: bool = False
         self._async_success_enqueued: bool = False
         self._sync_success_scheduled: bool = False
+        self._superseded_success_logs_in_flight: frozenset[str] = frozenset()
         self.on_detached_stream_failure: Callable[[Exception], Awaitable[None]] | None = None
         self._stream_failure_claimed: bool = False
         self._stash_owed_on_success: bool = False
@@ -2904,6 +2905,9 @@ class Logging(LiteLLMLoggingBaseClass):
             Ignore check on stream, as there can be multiple chunks
             """
             return
+        if event_type in self._superseded_success_logs_in_flight:
+            self._superseded_success_logs_in_flight = self._superseded_success_logs_in_flight - frozenset((event_type,))
+            return
         self.model_call_details[f"has_logged_{event_type}"] = True
         if event_type == "async_failure":
             self.model_call_details["async_failure_logged_attempt"] = self.model_call_details.get(
@@ -3511,6 +3515,26 @@ class Logging(LiteLLMLoggingBaseClass):
         self._reset_success_emission_dedupe()
         await self.async_success_handler(result=result)
 
+    def reopen_success_logging_for_next_attempt(self) -> None:
+        """The router discarded an earlier attempt's response, so the next attempt logs and prices its own
+        response. A success log the discarded attempt parked for the proxy to dispatch later is dropped, since
+        it describes a response the caller never receives. One it already dispatched still runs, and its late
+        completion must not mark the next attempt as logged."""
+        self._superseded_success_logs_in_flight = frozenset(
+            event_type
+            for event_type, scheduled in (
+                ("async_success", self._async_success_scheduled),
+                ("sync_success", self._sync_success_scheduled),
+            )
+            if scheduled
+            and self.deferred_success_result is None
+            and self.model_call_details.get(f"has_logged_{event_type}") is not True
+        )
+        self._reset_success_emission_dedupe()
+        self._async_success_scheduled = False
+        self.deferred_success_result = None
+        self.enqueue_deferred_logging = None
+
     def _reset_success_emission_dedupe(self) -> None:
         """
         Success callbacks dedupe per request, because the sync and async

litellm/router_python_core.py

@@ -365,6 +365,7 @@ from litellm.router_utils.fallback_event_handlers import (
     record_disable_fallbacks,
     record_pre_routing_selection,
     record_retry_attempt,
+    reopen_success_logging_for_next_attempt,
     routed_deployment_id,
     run_async_fallback,
 )
@@ -9847,6 +9848,7 @@ class PythonCore(DeprecatedRouterMembers):
                     _metadata["attempted_retries"] = current_attempt + 1
                     _metadata["max_retries"] = num_retries
                     carry_superseded_attempt_usage(kwargs)
+                    reopen_success_logging_for_next_attempt(_REQUEST_KWARGS_ADAPTER.validate_python(kwargs))
                     # if the function call is successful, no exception will be raised and we'll break out of the loop
                     response = await self.make_call(original_function, *args, **kwargs)
                     if coroutine_checker.is_async_callable(response):  # async errors are often returned as coroutines

litellm/router_utils/fallback_event_handlers.py

@@ -466,6 +466,12 @@ def carry_superseded_attempt_usage(kwargs: Mapping[str, object]) -> None:
         logging_obj.carry_failed_attempt_usage()
 
 
+def reopen_success_logging_for_next_attempt(kwargs: Mapping[str, object]) -> None:
+    logging_obj: Final = kwargs.get("litellm_logging_obj")
+    if isinstance(logging_obj, Logging):
+        logging_obj.reopen_success_logging_for_next_attempt()
+
+
 def _failure_callbacks_blocked_logging_object(kwargs: Mapping[str, object]) -> Logging | None:
     logging_obj: Final = kwargs.get("litellm_logging_obj")
     if not isinstance(logging_obj, Logging):
@@ -900,6 +906,7 @@ async def run_async_fallback(
             if include_fallback_errors:
                 kwargs["include_fallback_errors"] = include_fallback_errors
             carry_superseded_attempt_usage(_FALLBACK_KWARGS_ADAPTER.validate_python(kwargs))
+            reopen_success_logging_for_next_attempt(_FALLBACK_KWARGS_ADAPTER.validate_python(kwargs))
             response = await ctx.with_fallbacks(*args, **kwargs)
             verbose_router_logger.info("Successful fallback b/w models.")
             response = add_fallback_headers_to_response(

tests/unit/test_router/test_router_retries.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
 import asyncio
+from collections.abc import Callable
+from datetime import datetime
 from typing import Final
 from unittest.mock import AsyncMock
 
@@ -13,6 +15,7 @@ from pydantic import TypeAdapter
 import litellm
 from litellm import Router
 from litellm.integrations.custom_logger import CustomLogger
+from litellm.litellm_core_utils.litellm_logging import Logging
 from litellm.router import AllowedFailsPolicy, RetryPolicy
 import os
 
@@ -336,6 +339,100 @@ async def test_authentication_errors_are_not_retried(
     assert len(route.calls) == 1
 
 
+class _SuccessLogRecorder(CustomLogger):
+    def __init__(self) -> None:
+        super().__init__()
+        self.logged: tuple[tuple[object, object], ...] = ()
+
+    async def async_log_success_event(
+        self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime
+    ) -> None:
+        payload: Final = TypeAdapter(dict[str, object]).validate_python(kwargs["standard_logging_object"])
+        self.logged = (*self.logged, (payload["id"], payload["response_cost"]))
+
+
+def _chat_completion(response_id: str, finish_reason: str, completion_tokens: int) -> httpx.Response:
+    return httpx.Response(
+        200,
+        json={
+            "id": response_id,
+            "object": "chat.completion",
+            "created": 1,
+            "model": "gpt-4o-mini",
+            "choices": [
+                {"index": 0, "finish_reason": finish_reason, "message": {"role": "assistant", "content": "hi"}}
+            ],
+            "usage": {
+                "prompt_tokens": 10,
+                "completion_tokens": completion_tokens,
+                "total_tokens": 10 + completion_tokens,
+            },
+        },
+    )
+
+
+async def _wait_until(condition: Callable[[], bool]) -> None:
+    for _ in range(200):
+        if condition():
+            return
+        await asyncio.sleep(0.01)
+
+
+@pytest.mark.parametrize("deferred", [False, True])
+@pytest.mark.parametrize("redispatch", ["retry", "content_policy_fallback"])
+@pytest.mark.asyncio
+async def test_attempt_after_a_discarded_response_logs_its_own_cost(
+    deferred: bool,
+    redispatch: str,
+    respx_mock: respx.MockRouter,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    respx_mock.post("https://api.openai.com/v1/chat/completions").mock(
+        side_effect=[
+            _chat_completion("chatcmpl-filtered", "content_filter", completion_tokens=40),
+            _chat_completion("chatcmpl-served", "stop", completion_tokens=4),
+        ]
+    )
+    recorder: Final = _SuccessLogRecorder()
+    messages: Final = [{"role": "user", "content": "hi"}]
+    logging_obj: Final = Logging(
+        model="retry-test-model",
+        messages=messages,
+        stream=False,
+        call_type="acompletion",
+        start_time=datetime(2026, 1, 1),
+        litellm_call_id="discarded-attempt",
+        function_id="discarded-attempt",
+        dynamic_async_success_callbacks=[recorder],
+    )
+    logging_obj.defer_async_logging = deferred
+    retry_policy: Final = (
+        {"retry_policy": RetryPolicy(ContentPolicyViolationErrorRetries=1)} if redispatch == "retry" else {}
+    )
+    router: Final = Router(
+        model_list=[
+            {"model_name": group, "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "test-key"}}
+            for group in ("retry-test-model", "fallback-model")
+        ],
+        content_policy_fallbacks=[{"retry-test-model": ["fallback-model"]}],
+        **retry_policy,
+    )
+
+    response: Final = await router.acompletion(
+        model="retry-test-model", messages=messages, litellm_logging_obj=logging_obj
+    )
+    if deferred:
+        assert logging_obj.enqueue_deferred_logging is not None
+        logging_obj.enqueue_deferred_logging()
+    expected_ids: Final = ("chatcmpl-served",) if deferred else ("chatcmpl-filtered", "chatcmpl-served")
+    await _wait_until(lambda: len(recorder.logged) >= len(expected_ids))
+
+    assert response.id == "chatcmpl-served"
+    assert tuple(logged_id for logged_id, _ in recorder.logged) == expected_ids
+    assert recorder.logged[-1][1] == response._hidden_params["response_cost"]
+
+
 def test_retry_rate_limit_error_with_healthy_deployments():
     """
     Test 1. It SHOULD retry when there is a rate limit error and len(healthy_deployments) > 0