[internal copy of #29089] fix: duplicate claude code traces#29311
Conversation
|
bugbot run |
Greptile SummaryThis PR centralizes success-handler dispatch logic into a new
Confidence Score: 5/5Safe to merge — the change consolidates two unconditional handler calls into a single routed dispatch with an explicit dedup flag, removing the root cause of duplicate traces without altering the observable contract for correctly-configured callers. The refactor is well-scoped: every call site that previously fired both async_success_handler and executor.submit(success_handler) unconditionally has been migrated to dispatch_success_handlers, which owns the dedup and routing logic in one place. The new tests directly cover the dedup invariant (async path, sync path), the prefer_async_handlers executor path, and the proxy deferred-guardrail regression. No prior dedup mechanism is removed without a replacement. The _attach_mock_success_dispatch helper in test_deferred_guardrail_logging.py replaces the real dispatch_success_handlers with a simplified stub that skips the dedup check; this is intentional for guardrail-behaviour tests but worth noting if new deferred-guardrail tests are added that care about exactly-once logging.
|
| Filename | Overview |
|---|---|
| litellm/litellm_core_utils/litellm_logging.py | Core change: adds dispatch_success_handlers, _is_assembled_stream_success, and extracts _is_sync_litellm_request; modifies the early-return guard in async_success_handler to bypass should_run_logging for assembled streams. The dedup for assembled streams is now fully delegated to dispatch_success_handlers. |
| litellm/litellm_core_utils/streaming_handler.py | Async anext migrated from paired async_success_handler + executor.submit(success_handler) to dispatch_success_handlers(prefer_async_handlers=True); per-chunk sync path now conditionally skips success_handler for async SDK entrypoints. |
| litellm/proxy/common_request_processing.py | Deferred-guardrail path and orphaned-stream path both migrated from paired async+sync calls to dispatch_success_handlers(prefer_async_handlers=True). |
| litellm/proxy/pass_through_endpoints/streaming_handler.py | Streaming pass-through handler migrated from paired calls to dispatch_success_handlers(prefer_async_handlers=True); removes now-unused executor import. |
| litellm/proxy/pass_through_endpoints/success_handler.py | Non-streaming pass-through _handle_logging migrated from always-firing thread_pool_executor.submit(success_handler) + async_success_handler to dispatch_success_handlers(prefer_async_handlers=True). |
| tests/test_litellm/litellm_core_utils/test_litellm_logging.py | Adds four new focused unit tests covering dedup (async path, sync path), prefer_async_handlers + executor behaviour, and pass-through async routing; no existing assertions weakened. |
| tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py | Existing guardrail tests updated to mock dispatch_success_handlers via _attach_mock_success_dispatch helper; adds a regression test for sync-classified call_type still reaching the async handler in deferred proxy streaming. |
| tests/proxy_unit_tests/test_proxy_reject_logging.py | Fixes the embedding test case to use the router's actual model name (fake-model); registers the test logger on all callback lists to work with the new dispatch path. |
| tests/pass_through_unit_tests/test_unit_test_streaming.py | Adds two new unit tests verifying that the async handler fires (and sync does not) for SDK pass-through streaming and non-streaming pass-through logging paths. |
| tests/test_litellm/litellm_core_utils/test_streaming_handler.py | Minor formatting cleanup; no logic changes. |
Reviews (5): Last reviewed commit: "revert: call_type-based classification i..." | Re-trigger Greptile
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 1 potential issue.
Bugbot Autofix prepared a fix for the issue found in the latest run.
- ✅ Fixed: Method name and logic don't match documented purpose
- Tightened _is_assembled_stream_success to exclude ModelResponseStream (per-chunk) results so the method matches its documented intent and prevents future callers from prematurely tripping the dedup guard.
Preview (cb95e2648b)
diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py
--- a/litellm/litellm_core_utils/litellm_logging.py
+++ b/litellm/litellm_core_utils/litellm_logging.py
@@ -1612,6 +1612,90 @@
) -> Optional[float]:
return self._response_cost_calculator(result=result, cache_hit=cache_hit)
+ @staticmethod
+ def _is_sync_litellm_request(litellm_params: dict) -> bool:
+ """True for sync SDK entrypoints (``completion``), false for async (``acompletion``, etc.)."""
+ return (
+ litellm_params.get(CallTypes.acompletion.value, False) is not True
+ and litellm_params.get(CallTypes.aresponses.value, False) is not True
+ and litellm_params.get(CallTypes.aembedding.value, False) is not True
+ and litellm_params.get(CallTypes.aimage_generation.value, False) is not True
+ and litellm_params.get(CallTypes.atranscription.value, False) is not True
+ )
+
+ def _is_assembled_stream_success(self, result=None) -> bool:
+ """Final assembled stream export (not a per-chunk success call).
+
+ Per-chunk callers pass a ``ModelResponseStream`` (or ``None``); the
+ final assembled response is any other non-``None`` value (typically a
+ ``ModelResponse``). Treating a chunk as the assembled response would
+ prematurely set the ``has_dispatched_final_stream_success`` dedup
+ guard and silently suppress the real final stream log.
+ """
+ if self.stream is not True:
+ return False
+ if result is not None and not isinstance(result, ModelResponseStream):
+ return True
+ return (
+ "async_complete_streaming_response" in self.model_call_details
+ or self.model_call_details.get("complete_streaming_response") is not None
+ )
+
+ async def dispatch_success_handlers(
+ self,
+ result=None,
+ start_time=None,
+ end_time=None,
+ cache_hit=None,
+ prefer_async_handlers: bool = False,
+ **kwargs,
+ ) -> None:
+ """Route success logging to async and/or sync handlers for this request.
+
+ ``prefer_async_handlers`` only bypasses the sync-SDK-only shortcut (e.g.
+ ``async for`` on a stream from ``completion()``). Legacy string callbacks
+ still run via ``executor.submit(success_handler)`` when configured.
+ """
+ from litellm.litellm_core_utils.thread_pool_executor import executor
+
+ if self._is_assembled_stream_success(result):
+ if self.model_call_details.get("has_dispatched_final_stream_success"):
+ return
+ self.model_call_details["has_dispatched_final_stream_success"] = True
+
+ litellm_params = self.model_call_details.get("litellm_params", {}) or {}
+ sync_sdk = self._is_sync_litellm_request(litellm_params)
+ passthrough = self.call_type == CallTypes.pass_through.value
+ if sync_sdk and not prefer_async_handlers and not passthrough:
+ self.success_handler(
+ result,
+ start_time=start_time,
+ end_time=end_time,
+ cache_hit=cache_hit,
+ **kwargs,
+ )
+ return
+
+ await self.async_success_handler(
+ result,
+ start_time=start_time,
+ end_time=end_time,
+ cache_hit=cache_hit,
+ **kwargs,
+ )
+
+ if not self._should_run_sync_callbacks_for_async_calls():
+ return
+
+ executor.submit(
+ self.success_handler,
+ result,
+ start_time=start_time,
+ end_time=end_time,
+ cache_hit=cache_hit,
+ **kwargs,
+ )
+
def should_run_logging(
self,
event_type: Literal[
@@ -2034,13 +2118,7 @@
standard_logging_object=kwargs.get("standard_logging_object", None),
)
litellm_params = self.model_call_details.get("litellm_params", {})
- is_sync_request = (
- litellm_params.get(CallTypes.acompletion.value, False) is not True
- and litellm_params.get(CallTypes.aresponses.value, False) is not True
- and litellm_params.get(CallTypes.aembedding.value, False) is not True
- and litellm_params.get(CallTypes.aimage_generation.value, False) is not True
- and litellm_params.get(CallTypes.atranscription.value, False) is not True
- )
+ is_sync_request = self._is_sync_litellm_request(litellm_params)
try:
## BUILD COMPLETE STREAMED RESPONSE
complete_streaming_response: Optional[
@@ -2496,9 +2574,11 @@
print_verbose(
"Logging Details LiteLLM-Async Success Call, cache_hit={}".format(cache_hit)
)
- if not self.should_run_logging(
+ if not self._is_assembled_stream_success(
+ result
+ ) and not self.should_run_logging(
event_type="async_success"
- ): # prevent double logging
+ ): # prevent double logging (non-streaming)
return
## CALCULATE COST FOR BATCH JOBS
@@ -2948,13 +3028,7 @@
): # prevent double logging
return
litellm_params = self.model_call_details.get("litellm_params", {})
- is_sync_request = (
- litellm_params.get(CallTypes.acompletion.value, False) is not True
- and litellm_params.get(CallTypes.aresponses.value, False) is not True
- and litellm_params.get(CallTypes.aembedding.value, False) is not True
- and litellm_params.get(CallTypes.aimage_generation.value, False) is not True
- and litellm_params.get(CallTypes.atranscription.value, False) is not True
- )
+ is_sync_request = self._is_sync_litellm_request(litellm_params)
try:
start_time, end_time = self._failure_handler_helper_fn(
diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py
--- a/litellm/litellm_core_utils/streaming_handler.py
+++ b/litellm/litellm_core_utils/streaming_handler.py
@@ -1808,8 +1808,10 @@
processed_chunk, None, None, cache_hit
)
)
- ## SYNC LOGGING
- self.logging_obj.success_handler(processed_chunk, None, None, cache_hit)
+ ## SYNC LOGGING — only for sync SDK entrypoints; async proxy paths export via async_success_handler
+ litellm_params = self.logging_obj.model_call_details.get("litellm_params", {})
+ if self.logging_obj._is_sync_litellm_request(litellm_params):
+ self.logging_obj.success_handler(processed_chunk, None, None, cache_hit)
def finish_reason_handler(self):
model_response = self.model_response_creator()
@@ -2206,23 +2208,19 @@
cache_hit,
)
else:
+ # prefer_async_handlers routes CustomLogger to async_success_handler
+ # when consumers use ``async for`` on sync-SDK streams. Legacy string
+ # callbacks still run via executor.submit inside dispatch_success_handlers.
asyncio.create_task(
- self.logging_obj.async_success_handler(
+ self.logging_obj.dispatch_success_handlers(
complete_streaming_response,
cache_hit=cache_hit,
start_time=None,
end_time=None,
+ prefer_async_handlers=True,
)
)
- executor.submit(
- self.logging_obj.success_handler,
- complete_streaming_response,
- cache_hit=cache_hit,
- start_time=None,
- end_time=None,
- )
-
raise StopAsyncIteration # Re-raise StopIteration
else:
self.sent_last_chunk = True
diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py
--- a/litellm/proxy/common_request_processing.py
+++ b/litellm/proxy/common_request_processing.py
@@ -1290,7 +1290,7 @@
# (ProxyLogging._fire_deferred_stream_logging) fires the
# closure after the full streaming pipeline finishes.
# The closure runs non-apply_guardrail hooks on the
- # assembled response, then fires both logging handlers.
+ # assembled response, then fires success logging.
# Only for CustomStreamWrapper — raw async generators from
# passthrough routes bypass CSW and would orphan the closure.
from litellm.litellm_core_utils.streaming_handler import (
@@ -1411,7 +1411,7 @@
logging_obj._on_deferred_stream_complete = None # type: ignore[union-attr]
try:
asyncio.create_task(
- logging_obj.async_success_handler(
+ logging_obj.dispatch_success_handlers(
response,
cache_hit=None,
start_time=None,
@@ -1422,23 +1422,7 @@
verbose_proxy_logger.exception(
"Error in orphaned streaming async logging: %s", e
)
- try:
- from litellm.litellm_core_utils.thread_pool_executor import (
- executor as _exc,
- )
- _exc.submit(
- logging_obj.success_handler,
- response,
- cache_hit=None,
- start_time=None,
- end_time=None,
- )
- except Exception as e:
- verbose_proxy_logger.exception(
- "Error in orphaned streaming sync logging: %s", e
- )
-
# Always return the client-requested model name (not provider-prefixed internal identifiers)
# for OpenAI-compatible responses.
if requested_model_from_client:
@@ -1639,7 +1623,7 @@
) -> None:
"""
Run non-streaming post-call guardrail hooks on an assembled streaming
- response, then fire both async and sync logging handlers.
+ response, then fire success logging via ``dispatch_success_handlers``.
Called by ProxyLogging._fire_deferred_stream_logging after the full
streaming pipeline (including unified_guardrail end-of-stream blocks)
@@ -1655,8 +1639,6 @@
Extracted as a static method so tests can call the production
implementation directly rather than reimplementing the closure.
"""
- from litellm.litellm_core_utils.thread_pool_executor import executor
-
_response = assembled_response
try:
from litellm.proxy.proxy_server import llm_router as _global_llm_router
@@ -1716,7 +1698,7 @@
finally:
try:
asyncio.create_task(
- captured_logging_obj.async_success_handler(
+ captured_logging_obj.dispatch_success_handlers(
_response,
cache_hit=cache_hit,
start_time=None,
@@ -1725,24 +1707,10 @@
)
except Exception as e:
verbose_proxy_logger.exception(
- "Error in deferred streaming async logging: %s",
+ "Error in deferred streaming success logging: %s",
e,
)
- try:
- executor.submit(
- captured_logging_obj.success_handler,
- _response,
- cache_hit=cache_hit,
- start_time=None,
- end_time=None,
- )
- except Exception as e:
- verbose_proxy_logger.exception(
- "Error in deferred streaming sync logging: %s",
- e,
- )
-
async def _handle_llm_api_exception(
self,
e: Exception,
diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py
--- a/litellm/proxy/pass_through_endpoints/streaming_handler.py
+++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py
@@ -7,7 +7,6 @@
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
-from litellm.litellm_core_utils.thread_pool_executor import executor
from litellm.proxy._types import PassThroughEndpointLoggingResultValues
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType
@@ -145,27 +144,13 @@
end_time=end_time,
model=model,
)
- await litellm_logging_obj.async_success_handler(
+ await litellm_logging_obj.dispatch_success_handlers(
result=standard_logging_response_object,
start_time=start_time,
end_time=end_time,
cache_hit=False,
**kwargs,
)
- if (
- litellm_logging_obj._should_run_sync_callbacks_for_async_calls()
- is False
- ):
- return
-
- executor.submit(
- litellm_logging_obj.success_handler,
- result=standard_logging_response_object,
- end_time=end_time,
- cache_hit=False,
- start_time=start_time,
- **kwargs,
- )
except Exception as e:
verbose_proxy_logger.error(
f"Error in _route_streaming_logging_to_handler: {str(e)}"
diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py
--- a/litellm/proxy/pass_through_endpoints/success_handler.py
+++ b/litellm/proxy/pass_through_endpoints/success_handler.py
@@ -11,7 +11,6 @@
PassthroughStandardLoggingPayload,
)
from litellm.types.utils import StandardPassThroughResponseObject
-from litellm.utils import executor as thread_pool_executor
from .llm_provider_handlers.anthropic_passthrough_logging_handler import (
AnthropicPassthroughLoggingHandler,
@@ -94,19 +93,8 @@
cache_hit: bool,
**kwargs,
):
- """Helper function to handle both sync and async logging operations"""
- # Submit to thread pool for sync logging
- thread_pool_executor.submit(
- logging_obj.success_handler,
- standard_logging_response_object,
- start_time,
- end_time,
- cache_hit,
- **kwargs,
- )
-
- # Handle async logging
- await logging_obj.async_success_handler(
+ """Log pass-through success via the shared async dispatch path."""
+ await logging_obj.dispatch_success_handlers(
result=(
json.dumps(result)
if isinstance(result, dict)
diff --git a/tests/proxy_unit_tests/test_proxy_reject_logging.py b/tests/proxy_unit_tests/test_proxy_reject_logging.py
--- a/tests/proxy_unit_tests/test_proxy_reject_logging.py
+++ b/tests/proxy_unit_tests/test_proxy_reject_logging.py
@@ -95,6 +95,21 @@
)
+def _register_proxy_test_logger(callback_logger: testLogger) -> None:
+ """
+ Register the test logger on global callback lists.
+
+ ``function_setup`` dedupes by object identity; each parametrized case
+ constructs a new ``testLogger`` and must replace the global lists, not
+ only ``litellm.callbacks``.
+ """
+ litellm.callbacks = [callback_logger]
+ litellm.success_callback = [callback_logger]
+ litellm.failure_callback = [callback_logger]
+ litellm._async_success_callback = [callback_logger]
+ litellm._async_failure_callback = [callback_logger]
+
+
@pytest.mark.parametrize(
"route, body",
[
@@ -115,7 +130,7 @@
"/v1/embeddings",
{
"input": "The food was delicious and the waiter...",
- "model": "text-embedding-ada-002",
+ "model": "fake-model",
"encoding_format": "float",
},
),
@@ -133,7 +148,7 @@
setattr(proxy_server, "llm_router", router)
_test_logger = testLogger()
- litellm.callbacks = [_test_logger]
+ _register_proxy_test_logger(_test_logger)
litellm.set_verbose = True
# Prepare the query string
diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py
--- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py
+++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py
@@ -1,6 +1,7 @@
import os
import sys
-from unittest.mock import MagicMock, patch
+import asyncio
+from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@@ -786,6 +787,211 @@
dummy_logger.log_stream_event.assert_not_called()
+def test_is_sync_litellm_request():
+ assert LitellmLogging._is_sync_litellm_request({}) is True
+ assert LitellmLogging._is_sync_litellm_request({"acompletion": True}) is False
+
+
+@pytest.mark.asyncio
+async def test_dispatch_success_handlers_invokes_callbacks_once_for_final_stream(
+ logging_obj,
+):
+ """Second final-stream dispatch must not re-export (CSW + deferred guardrail paths)."""
+ import litellm
+ from litellm.integrations.custom_logger import CustomLogger
+
+ class MockCallback(CustomLogger):
+ pass
+
+ mock_callback = MockCallback()
+ original_async_callbacks = list(litellm._async_success_callback or [])
+ litellm._async_success_callback = [mock_callback]
+
+ result = ModelResponse(
+ id="resp-dedupe",
+ model="gpt-4o-mini",
+ choices=[
+ {
+ "message": {"role": "assistant", "content": "hi"},
+ "finish_reason": "stop",
+ "index": 0,
+ }
+ ],
+ usage={"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
+ )
+
+ try:
+ logging_obj.stream = True
+ logging_obj.model_call_details["litellm_params"] = {"acompletion": True}
+
+ with (
+ patch.object(
+ mock_callback, "async_log_success_event", new_callable=AsyncMock
+ ) as mock_async_log,
+ patch.object(mock_callback, "log_success_event") as mock_sync_log,
+ patch.object(
+ logging_obj,
+ "_success_handler_helper_fn",
+ return_value=(time.time(), time.time(), result),
+ ),
+ patch.object(
+ logging_obj,
+ "_get_assembled_streaming_response",
+ return_value=result,
+ ),
+ patch.object(
+ logging_obj,
+ "_should_run_sync_callbacks_for_async_calls",
+ return_value=True,
+ ),
+ ):
+ await logging_obj.dispatch_success_handlers(result=result)
+ await logging_obj.dispatch_success_handlers(result=result)
+
+ mock_async_log.assert_awaited_once()
+ mock_sync_log.assert_not_called()
+ finally:
+ litellm._async_success_callback = original_async_callbacks
+
+
+@pytest.mark.asyncio
+async def test_dispatch_success_handlers_sync_path_invokes_callback_once_for_final_stream(
+ logging_obj,
+):
+ """Sync dispatch path must also dedupe when dispatch is called twice."""
+ import litellm
+ from litellm.integrations.custom_logger import CustomLogger
+
+ class MockCallback(CustomLogger):
+ pass
+
+ mock_callback = MockCallback()
+ original_success_callbacks = list(litellm.success_callback or [])
+ litellm.success_callback = [mock_callback]
+
+ result = ModelResponse(
+ id="resp-sync-dedupe",
+ model="gpt-4o-mini",
+ choices=[
+ {
+ "message": {"role": "assistant", "content": "hi"},
+ "finish_reason": "stop",
+ "index": 0,
+ }
+ ],
+ usage={"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
+ )
+
+ try:
+ logging_obj.stream = True
+ logging_obj.model_call_details["litellm_params"] = {}
+
+ with (
+ patch.object(mock_callback, "log_success_event") as mock_sync_log,
+ patch.object(
+ mock_callback, "async_log_success_event", new_callable=AsyncMock
+ ) as mock_async_log,
+ patch.object(
+ logging_obj,
+ "_success_handler_helper_fn",
+ return_value=(time.time(), time.time(), result),
+ ),
+ patch.object(
+ logging_obj,
+ "_get_assembled_streaming_response",
+ return_value=result,
+ ),
+ ):
+ await logging_obj.dispatch_success_handlers(result=result)
+ await logging_obj.dispatch_success_handlers(result=result)
+
+ mock_sync_log.assert_called_once()
+ mock_async_log.assert_not_awaited()
+ finally:
+ litellm.success_callback = original_success_callbacks
+
+
+@pytest.mark.asyncio
+async def test_dispatch_prefer_async_handlers_runs_legacy_callbacks(
+ logging_obj,
+):
+ """``prefer_async_handlers`` must not skip executor.submit for string callbacks."""
+ result = ModelResponse(
+ id="resp-prefer-async",
+ model="gpt-4o-mini",
+ choices=[
+ {
+ "message": {"role": "assistant", "content": "hi"},
+ "finish_reason": "stop",
+ "index": 0,
+ }
+ ],
+ )
+
+ logging_obj.stream = True
+ logging_obj.model_call_details["litellm_params"] = {}
+
+ with (
+ patch.object(
+ logging_obj, "async_success_handler", new_callable=AsyncMock
+ ) as mock_async,
+ patch.object(
+ logging_obj, "success_handler", new_callable=MagicMock
+ ) as mock_sync,
+ patch.object(
+ logging_obj,
+ "_should_run_sync_callbacks_for_async_calls",
+ return_value=True,
+ ),
+ patch(
+ "litellm.litellm_core_utils.litellm_logging.executor.submit"
+ ) as mock_submit,
+ ):
+ await logging_obj.dispatch_success_handlers(
+ result=result,
+ prefer_async_handlers=True,
+ )
+
+ mock_async.assert_awaited_once()
+ mock_sync.assert_not_called()
+ mock_submit.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_dispatch_success_handlers_invokes_async_callback_for_pass_through(
+ logging_obj,
+):
+ """Pass-through must use async_success_handler (CustomLogger skips sync success_handler)."""
+ import litellm
+ from litellm.integrations.custom_logger import CustomLogger
+ from litellm.types.utils import CallTypes
+
+ class MockCallback(CustomLogger):
+ pass
+
+ mock_callback = MockCallback()
+ original_async_callbacks = list(litellm._async_success_callback or [])
+ litellm._async_success_callback = [mock_callback]
+
+ logging_obj.call_type = CallTypes.pass_through.value
+ logging_obj.stream = False
+ logging_obj.model_call_details["litellm_params"] = {}
+
+ try:
+ with (
+ patch.object(
+ mock_callback, "async_log_success_event", new_callable=AsyncMock
+ ) as mock_async_log,
+ patch.object(mock_callback, "log_success_event") as mock_sync_log,
+ ):
+ await logging_obj.dispatch_success_handlers(result={"id": "pt-1"})
+
+ mock_async_log.assert_awaited_once()
+ mock_sync_log.assert_not_called()
+ finally:
+ litellm._async_success_callback = original_async_callbacks
+
+
def test_success_handler_skips_guardrail_logging_hook_when_disabled(logging_obj):
"""Ensure CustomGuardrail logging_hook is skipped when should_run_guardrail is False."""
import datetime
@@ -1351,7 +1557,7 @@
Test that _generate_cold_storage_object_key uses s3_path from custom logger instance.
"""
from datetime import datetime, timezone
- from unittest.mock import MagicMock, patch
+ from unittest.mock import AsyncMock, MagicMock, patch
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
@@ -1404,7 +1610,7 @@
Test that _generate_cold_storage_object_key falls back to empty s3_path when logger has no s3_path.
"""
from datetime import datetime, timezone
- from unittest.mock import MagicMock, patch
+ from unittest.mock import AsyncMock, MagicMock, patch
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py
--- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py
+++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py
@@ -569,9 +569,7 @@
== final_usage_block
)
- print(mock_log_success_event.call_args.kwargs.keys())
-
def test_streaming_handler_with_stop_chunk(
initialized_custom_stream_wrapper: CustomStreamWrapper,
):
@@ -2036,23 +2034,19 @@
chunks.append(chunk)
# The prompt_filter chunk should be forwarded with choices=[]
- assert len(chunks[0].choices) == 0, (
- f"Expected prompt_filter chunk with choices=[], got {len(chunks[0].choices)} choices"
- )
+ assert (
+ len(chunks[0].choices) == 0
+ ), f"Expected prompt_filter chunk with choices=[], got {len(chunks[0].choices)} choices"
# At least one chunk must have role='assistant' in its delta
has_role = any(
- len(c.choices) > 0
- and getattr(c.choices[0].delta, "role", None) == "assistant"
+ len(c.choices) > 0 and getattr(c.choices[0].delta, "role", None) == "assistant"
for c in chunks
)
assert has_role, (
"No chunk contained role='assistant' in delta (issue #24221). "
"Chunk deltas: "
- + str([
- c.choices[0].delta if c.choices else "no choices"
- for c in chunks
- ])
+ + str([c.choices[0].delta if c.choices else "no choices" for c in chunks])
)
diff --git a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py
--- a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py
+++ b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py
@@ -38,6 +38,24 @@
# ---------------------------------------------------------------------------
+def _attach_mock_success_dispatch(mock_logging_obj, async_success_fn):
+ """Match production entrypoint: ``_run_deferred_stream_guardrails`` uses dispatch."""
+
+ async def dispatch_success_handlers(
+ result=None, start_time=None, end_time=None, cache_hit=None, **kwargs
+ ):
+ await async_success_fn(
+ result,
+ start_time=start_time,
+ end_time=end_time,
+ cache_hit=cache_hit,
+ **kwargs,
+ )
+
+ mock_logging_obj.dispatch_success_handlers = dispatch_success_handlers
+ mock_logging_obj.async_success_handler = async_success_fn
+
+
class PostCallGuardrail(CustomGuardrail):
"""A post-call guardrail."""
@@ -454,7 +472,7 @@
async def track_async_success(*args, **kwargs):
pass
- mock_logging_obj.async_success_handler = track_async_success
+ _attach_mock_success_dispatch(mock_logging_obj, track_async_success)
tracking_guardrail = TrackingGuardrail()
tracking_logger = TrackingLogger()
@@ -511,7 +529,7 @@
nonlocal logged_response
logged_response = args[0] if args else None
- mock_logging_obj.async_success_handler = track_async_success
+ _attach_mock_success_dispatch(mock_logging_obj, track_async_success)
class ModifyingGuardrail(CustomGuardrail):
def __init__(self):
@@ -573,7 +591,7 @@
nonlocal logging_called
logging_called = True
- mock_logging_obj.async_success_handler = track_async_success
+ _attach_mock_success_dispatch(mock_logging_obj, track_async_success)
guardrail = BlockingGuardrail()
@@ -621,7 +639,7 @@
async def track_async_success(*args, **kwargs):
pass
- mock_logging_obj.async_success_handler = track_async_success
+ _attach_mock_success_dispatch(mock_logging_obj, track_async_success)
guardrail = TransientErrorGuardrail()
@@ -656,7 +674,7 @@
nonlocal logged_response
logged_response = args[0] if args else None
- mock_logging_obj.async_success_handler = track_async_success
+ _attach_mock_success_dispatch(mock_logging_obj, track_async_success)
class TestGuardrail(CustomGuardrail):
def __init__(self):
@@ -739,7 +757,7 @@
async def track_async_success(*args, **kwargs):
pass
- mock_logging_obj.async_success_handler = track_async_success
+ _attach_mock_success_dispatch(mock_logging_obj, track_async_success)
guardrail = ApplyGuardrailType()
@@ -792,7 +810,7 @@
async def track_async_success(*args, **kwargs):
pass
- mock_logging_obj.async_success_handler = track_async_success
+ _attach_mock_success_dispatch(mock_logging_obj, track_async_success)
guardrail = IteratorHookGuardrail()
@@ -847,7 +865,7 @@
async def track_async_success(*args, **kwargs):
pass
- mock_logging_obj.async_success_handler = track_async_success
+ _attach_mock_success_dispatch(mock_logging_obj, track_async_success)
guardrail = InspectingGuardrail()
@@ -914,7 +932,7 @@
async def track_async_success(*args, **kwargs):
pass
- mock_logging_obj.async_success_handler = track_async_success
+ _attach_mock_success_dispatch(mock_logging_obj, track_async_success)
guardrail_a = TaggedGuardrail("guardrail-a")
guardrail_b = TaggedGuardrail("guardrail-b")
@@ -962,7 +980,7 @@
nonlocal logging_called
... diff truncated: showing 800 of 816 linesYou can send follow-ups to the cloud agent here.
Reviewed by Cursor Bugbot for commit f770d8e. Configure here.
_is_assembled_stream_success is documented as detecting the final assembled stream export, but the previous logic returned True for any streaming call with a non-None result, including per-chunk ModelResponseStream values. This created a fragile contract where a future caller passing a chunk to dispatch_success_handlers would prematurely set the has_dispatched_final_stream_success dedup guard and silently suppress the real final stream log. Tighten the check to treat only non-ModelResponseStream results as assembled responses.
|
|
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
PR overviewAll previously flagged issues have been addressed. No open security concerns remain on this pull request. Security reviewNo open security issues remain on this pull request. Fixed/addressed: 1 · PR risk: 0/10 |
…aming (#29312) dispatch_success_handlers classified anthropic_messages and google generate_content streaming as sync SDK requests, since their call_type is not pass_through_endpoint and their litellm_params carries no acompletion flag. Only the sync success_handler ran, so async-only loggers (CustomLogger.async_log_success_event) never recorded the assembled stream; this broke test_anthropic_messages_litellm_router_streaming_with_logging and test_async_streaming_with_logging. _route_streaming_logging_to_handler now passes prefer_async_handlers=True, which is always valid because it runs from an async context (anthropic_messages, google_genai, and proxy pass-through stream tasks). This restores async dispatch while keeping the final-stream dedupe guard intact. Adds a regression test asserting the async handler runs and the sync handler does not for an anthropic_messages streaming call.
The deferred proxy streaming logging closures called dispatch_success_handlers without prefer_async_handlers, so routing depended on _is_sync_litellm_request, which only recognizes a subset of async markers stored in litellm_params. An async proxy stream whose litellm_params lacks a recognized marker would take the sync branch and return before async_success_handler runs. Because the proxy DB/spend logger is async-only, that request would not be recorded for spend tracking. Proxy streaming always runs in an async context and there is no sync SDK consumer of a proxy stream, so both deferred dispatch sites in common_request_processing.py now pass prefer_async_handlers=True. This is behavior-neutral for the acompletion path (already async) and guarantees spend callbacks fire for any async proxy call type. The has_dispatched_final_stream_success dedupe guard and the sync-callback gating are unchanged, so no duplicate logs. Adds a regression test driving the real dispatch via _run_deferred_stream_guardrails with a sync-classified call type and asserting the async handler runs.
Make PassThroughEndpointLogging._handle_logging pass prefer_async_handlers=True to dispatch_success_handlers, matching the streaming sibling _route_streaming_logging_to_handler. _handle_logging is only ever reached from pass_through_async_success_handler (an async context) with call_type "pass_through_endpoint", so the passthrough guard in dispatch_success_handlers already forces the async handler to run -- this is defensive hardening, not a live bug fix. Passing the flag explicitly keeps async-only loggers (e.g. the proxy spend logger) firing regardless of how the call-type classification or the passthrough guard evolves, and keeps the streaming and non-streaming pass-through logging paths symmetric. Add a regression test pinning the contract: _handle_logging must reach async_success_handler even for a call_type that _is_sync_litellm_request would classify as sync. https://claude.ai/code/session_018Fawe7GQF6wULjPhAffmyi
get_litellm_params only stores the acompletion/aembedding async flags, so _is_sync_litellm_request's checks for aresponses/aimage_generation/ atranscription (and the missing atext_completion) never fired for real SDK calls -- those requests were classified as sync and would skip async-only success/failure callbacks if dispatched without prefer_async_handlers=True. Close the gap using call_type (the SDK function name, so async entrypoints carry the async-prefixed value) for exactly the four async completion-family call types whose flag is not stored. acompletion/aembedding detection is unchanged (flag-based), and all existing litellm_params flag checks are kept, so behavior for sync and flag-marked requests is unchanged. Convert the helper from a staticmethod to an instance method (all call sites already invoke it on an instance) and rewrite test_is_sync_litellm_request to cover sync types, the four newly-detected async-only types, the "a"-prefixed but sync anthropic_messages case, and the flag-based path. https://claude.ai/code/session_018Fawe7GQF6wULjPhAffmyi
Revert fe60409. Detecting aresponses/atranscription/aimage_generation/ atext_completion as async from call_type broke the proxy reject-logging contract: _is_sync_litellm_request also gates whether the *sync* success/failure callbacks run alongside the async ones, and proxy requests (e.g. a rejected /v1/completions -> atext_completion) must still fire the sync failure callback. test_proxy_reject_logging.py asserts both the async and sync failure events fire; classifying atext_completion as async suppressed the sync one (reaches_sync_failure_event == False), failing the proxy-config and logging-misc CI jobs. The latent gap that change targeted only matters inside dispatch_success_handlers, where every call site already passes prefer_async_handlers=True, so async-only loggers always run regardless of the heuristic. Restore the original helper and its tests; no behavior change for sync, async, or proxy paths. https://claude.ai/code/session_018Fawe7GQF6wULjPhAffmyi
|
@greptileai the latent |
|
You're right, and I was wrong to flag that as a concern. The PR makes Widening I'll withdraw that concern — no change needed here. |
|
bugbot run |
581c30f
into
litellm_internal_staging
* [internal copy of #29089] fix: duplicate claude code traces (#29311) * refactor(proxy/auth): normalize Bearer prefix in safe-hash helper (#29343) * refactor(proxy/auth): normalize Bearer prefix in safe-hash helper UserAPIKeyAuth._safe_hash_litellm_api_key now strips a leading "Bearer "/"bearer " prefix before its existing sk-/JWT classification, so the helper produces the same hashed output regardless of whether the caller stripped the Authorization header prefix or passed the header value through unchanged. * refactor(proxy/auth): make Bearer-prefix strip case-insensitive Per RFC 7235 the HTTP authorization scheme token is case-insensitive. Replace the two-prefix loop with a single case-insensitive check so the helper normalizes "Bearer ", "bearer ", "BEARER ", and any mixed-case variant before classifying the remainder as sk- or JWT. The contract test gains coverage of "BEARER " and "BeArEr ". * test(mcp): align auth-handler test expectations with safe-hash helper The two MCP auth tests asserted that UserAPIKeyAuth(api_key="Bearer ...") retained the raw header bytes on the api_key field. _safe_hash_litellm_api_key now normalizes that input — stripping the Bearer prefix and hashing the resulting sk- key — so the expectations move to the normalized form: the bare token in the parametrize case, and hash_token("sk-...") in the backward-compat assertion. This matches what the real auth flow produces (the builder strips Bearer and the DB stores the hashed token), so the mocks now line up with production rather than with the un-normalized validator output. --------- Co-authored-by: yuneng-jiang <[email protected]>
The #29311 cherry-pick onto stable/1.85.x carried the test test_route_streaming_logging_runs_async_handler_for_sdk_passthrough, which patches PassThroughStreamingHandler._build_passthrough_logging_result to verify the SDK-passthrough dispatch contract. The manual conflict resolution kept the per-endpoint if/elif/elif chain inline in _route_streaming_logging_to_handler (matching v1.84.4's resolution), so the patched attribute did not exist and the test errored at collection with AttributeError. Extract the chain into the static _build_passthrough_logging_result helper as #29089 originally designed it. _route_streaming_logging_to_handler now resolves (standard_logging_response_object, kwargs) through the helper and dispatches via dispatch_success_handlers; the helper itself is synchronous and CPU-bound, suitable for the unit test's patch target. Verified locally: tests/pass_through_unit_tests/test_unit_test_streaming.py passes (5/5) and tests/test_litellm/litellm_core_utils/test_litellm_logging.py passes (83/83). v1.84.4 ships with the same broken test; this strictly improves on that resolution.



Automated copy of #29089 into
litellm_internal_stagingfor pr-babysitter.Original head:
mubashir1osmani/litellm:litellm__duplicate_cc_trce@f770d8e98716Note
Cursor Bugbot is generating a summary for commit f770d8e. Configure here.