scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 88

Cherry-pick #29311 and #29343 onto patch/v1.84.3 (#29352)

Browse source 
* [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 <yuneng@berri.ai>
This commit is contained in:
Mateo Wang 2026-05-30 17:52:11 -07:00 • committed by GitHub
parent aab7ff639e
commit a0bb8e6e51
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 596 additions and 129 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

106
litellm/litellm_core_utils/litellm_logging.py
View file

@ -1603,6 +1603,90 @@ class Logging(LiteLLMLoggingBaseClass):
) -> 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[
@ -2022,13 +2106,7 @@ class Logging(LiteLLMLoggingBaseClass):
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[
@ -2484,9 +2562,11 @@ class Logging(LiteLLMLoggingBaseClass):
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
@ -2936,13 +3016,7 @@ class Logging(LiteLLMLoggingBaseClass):
): # 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(

20
litellm/litellm_core_utils/streaming_handler.py
View file

@ -1808,8 +1808,10 @@ class CustomStreamWrapper:
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 @@ class CustomStreamWrapper:
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

13
litellm/proxy/_types.py
View file

@ -2676,13 +2676,16 @@ class UserAPIKeyAuth(
1. Regular API keys from LiteLLM DB
2. JWT tokens used for connecting to LiteLLM API
"""
if api_key.startswith("sk-"):
return hash_token(api_key)
normalized = api_key
if normalized[:7].lower() == "bearer ":
normalized = normalized[7:]
if normalized.startswith("sk-"):
return hash_token(normalized)
from litellm.proxy.auth.handle_jwt import JWTHandler
if JWTHandler.is_jwt(token=api_key):
return f"hashed-jwt-{hash_token(token=api_key)}"
return api_key
if JWTHandler.is_jwt(token=normalized):
return f"hashed-jwt-{hash_token(token=normalized)}"
return normalized
@classmethod
def get_litellm_internal_health_check_user_api_key_auth(cls) -> "UserAPIKeyAuth":

49
litellm/proxy/common_request_processing.py
View file

@ -1266,7 +1266,7 @@ class ProxyBaseLLMRequestProcessing:
# (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 (
@ -1387,33 +1387,18 @@ class ProxyBaseLLMRequestProcessing:
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,
end_time=None,
prefer_async_handlers=True,
)
)
except Exception as e:
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.
@ -1615,7 +1600,7 @@ class ProxyBaseLLMRequestProcessing:
) -> 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)
@ -1631,8 +1616,6 @@ class ProxyBaseLLMRequestProcessing:
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
@ -1691,31 +1674,23 @@ class ProxyBaseLLMRequestProcessing:
)
finally:
try:
# Proxy streaming always runs in async context and proxy spend
# logging is async-only; force async dispatch so DB/spend
# callbacks fire regardless of the call-type heuristic in
# _is_sync_litellm_request (which only recognizes a subset of
# async markers stored in litellm_params).
asyncio.create_task(
captured_logging_obj.async_success_handler(
captured_logging_obj.dispatch_success_handlers(
_response,
cache_hit=cache_hit,
start_time=None,
end_time=None,
prefer_async_handlers=True,
)
)
except Exception as e:
verbose_proxy_logger.exception(
"Error in deferred streaming async 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",
"Error in deferred streaming success logging: %s",
e,
)

22
litellm/proxy/pass_through_endpoints/streaming_handler.py
View file

@ -7,7 +7,6 @@ import httpx
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
@ -173,25 +172,16 @@ class PassThroughStreamingHandler:
standard_logging_response_object = StandardPassThroughResponseObject(
response=f"cannot parse chunks to standard response object. Chunks={all_chunks}"
)
await litellm_logging_obj.async_success_handler(
# Always reached from an async context (anthropic_messages,
# google_genai, and proxy pass-through stream tasks). prefer_async_handlers
# keeps async-only loggers running even when call_type isn't pass_through
# and litellm_params lacks an async flag (e.g. aanthropic_messages).
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,
prefer_async_handlers=True,
**kwargs,
)
except Exception as e:

24
litellm/proxy/pass_through_endpoints/success_handler.py
View file

@ -11,7 +11,6 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import (
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,15 @@ class PassThroughEndpointLogging:
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."""
# Always reached from pass_through_async_success_handler, which runs in
# an async context. call_type is "pass_through_endpoint" here, so the
# passthrough guard in dispatch_success_handlers already forces the
# async handler to run; pass prefer_async_handlers explicitly to match
# the streaming sibling (_route_streaming_logging_to_handler) and keep
# async-only loggers (e.g. the proxy spend logger) firing regardless of
# how the call-type classification evolves.
await logging_obj.dispatch_success_handlers(
result=(
json.dumps(result)
if isinstance(result, dict)
@ -115,6 +110,7 @@ class PassThroughEndpointLogging:
start_time=start_time,
end_time=end_time,
cache_hit=False,
prefer_async_handlers=True,
**kwargs,
)

117
tests/pass_through_unit_tests/test_unit_test_streaming.py
View file

@ -97,6 +97,123 @@ async def test_chunk_processor_yields_raw_bytes(endpoint_type, url_route):
), "Collected chunks do not match raw chunks"
@pytest.mark.asyncio
async def test_route_streaming_logging_runs_async_handler_for_sdk_passthrough():
"""
SDK pass-through streaming (anthropic_messages, google generate_content) must run
the async success handler so async-only loggers record the assembled stream.
Regression for duplicate-trace dedupe: dispatch_success_handlers treated these as
sync SDK requests because call_type is not ``pass_through_endpoint`` and
litellm_params carries no ``acompletion`` flag, so only the sync success_handler
ran and CustomLogger.async_log_success_event never fired.
"""
import time
from litellm.types.utils import CallTypes
logging_obj = LiteLLMLoggingObj(
model="claude-sonnet-4-5",
messages=[{"role": "user", "content": "hi"}],
stream=True,
call_type=CallTypes.anthropic_messages.value,
start_time=time.time(),
litellm_call_id="test-id",
function_id="fn",
)
logging_obj.model_call_details["litellm_params"] = {"anthropic_messages": True}
with (
patch.object(
PassThroughStreamingHandler,
"_build_passthrough_logging_result",
return_value=({"id": "slp"}, {}),
),
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=False,
),
):
await PassThroughStreamingHandler._route_streaming_logging_to_handler(
litellm_logging_obj=logging_obj,
passthrough_success_handler_obj=MagicMock(),
url_route="/v1/messages",
request_body={},
endpoint_type=EndpointType.ANTHROPIC,
start_time=datetime.now(),
raw_bytes=[],
end_time=datetime.now(),
)
mock_async.assert_awaited_once()
mock_sync.assert_not_called()
@pytest.mark.asyncio
async def test_handle_logging_runs_async_handler_for_passthrough():
"""
Non-streaming pass-through logging (_handle_logging) must always run the
async success handler so async-only loggers (e.g. the proxy spend logger)
record the request.
_handle_logging is only ever reached from pass_through_async_success_handler
(an async context), so it forces async dispatch via prefer_async_handlers.
This pins that contract independent of the call-type classification: even a
call_type that _is_sync_litellm_request would classify as sync (here
"completion" with no async marker in litellm_params) must still reach
async_success_handler. Without prefer_async_handlers=True the sync-only
branch would return early and async_log_success_event would never fire.
"""
import time
from litellm.types.utils import CallTypes
logging_obj = LiteLLMLoggingObj(
model="claude-sonnet-4-5",
messages=[{"role": "user", "content": "hi"}],
stream=False,
call_type=CallTypes.completion.value,
start_time=time.time(),
litellm_call_id="test-id",
function_id="fn",
)
logging_obj.model_call_details["litellm_params"] = {}
handler = PassThroughEndpointLogging()
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=False,
),
):
await handler._handle_logging(
logging_obj=logging_obj,
standard_logging_response_object={"id": "slp"},
result="",
start_time=datetime.now(),
end_time=datetime.now(),
cache_hit=False,
)
mock_async.assert_awaited_once()
mock_sync.assert_not_called()
def test_convert_raw_bytes_to_str_lines():
"""
Test that the _convert_raw_bytes_to_str_lines method correctly converts raw bytes to a list of strings

19
tests/proxy_unit_tests/test_proxy_reject_logging.py
View file

@ -95,6 +95,21 @@ router = Router(
)
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 @@ router = Router(
"/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 @@ async def test_chat_completion_request_with_redaction(route, body):
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

212
tests/test_litellm/litellm_core_utils/test_litellm_logging.py
View file

@ -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
@ -773,6 +774,211 @@ def test_success_handler_runs_sync_callbacks_for_sync_requests(logging_obj, call
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
@ -1338,7 +1544,7 @@ async def test_e2e_generate_cold_storage_object_key_with_custom_logger_s3_path()
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
@ -1391,7 +1597,7 @@ async def test_e2e_generate_cold_storage_object_key_with_logger_no_s3_path():
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

16
tests/test_litellm/litellm_core_utils/test_streaming_handler.py
View file

@ -569,8 +569,6 @@ async def test_streaming_with_usage_and_logging(sync_mode: bool):
== 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 @@ async def test_azure_streaming_role_preserved_with_include_usage(sync_mode: bool
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])
)

6
tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
View file

@ -213,7 +213,7 @@ class TestMCPRequestHandler:
# Test case 2: Authorization header present (fallback)
(
[(b"authorization", b"Bearer test-auth-token")],
"Bearer test-auth-token",
"test-auth-token",
None,
{},
),
@ -674,7 +674,9 @@ class TestMCPOAuth2AuthFlow:
) = await MCPRequestHandler.process_mcp_request(scope)
# Should succeed with the LiteLLM key from Authorization header
assert auth_result.api_key == "Bearer sk-litellm-valid-key"
from litellm.proxy.utils import hash_token
assert auth_result.api_key == hash_token("sk-litellm-valid-key")
mock_auth.assert_called_once()
async def test_non_auth_http_exception_still_raises(self):

103
tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py
View file

@ -18,7 +18,7 @@ import asyncio
import os
import sys
from typing import Any
from unittest.mock import MagicMock, patch
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -38,6 +38,24 @@ from litellm.types.guardrails import GuardrailEventHooks
# ---------------------------------------------------------------------------
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 @@ class TestDeferredStreamingClosure:
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 @@ class TestDeferredStreamingClosure:
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 @@ class TestDeferredStreamingClosure:
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 @@ class TestDeferredStreamingClosure:
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 @@ class TestDeferredStreamingClosure:
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 @@ class TestDeferredStreamingClosure:
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 @@ class TestDeferredStreamingClosure:
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 @@ class TestDeferredStreamingClosure:
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 @@ class TestDeferredStreamingClosure:
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 @@ class TestDeferredStreamingClosure:
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)
def exploding_merge(data, llm_router):
raise RuntimeError("Simulated init failure")
@ -986,6 +1004,67 @@ class TestDeferredStreamingClosure:
logging_called is True
), "Logging must fire even when guardrail initialization raises"
@pytest.mark.asyncio
async def test_deferred_logging_forces_async_for_sync_classified_call_type(self):
"""
Regression: proxy deferred streaming logging must reach the async success
handler (which runs the async-only DB/spend logger) even when the call
type is classified as a sync SDK request by _is_sync_litellm_request.
Without prefer_async_handlers=True, an async proxy stream whose
litellm_params lacks a recognized async marker would enter the sync
branch of dispatch_success_handlers and silently skip spend tracking.
Uses the real dispatch_success_handlers via the production
_run_deferred_stream_guardrails entrypoint.
"""
import time
from litellm.litellm_core_utils.litellm_logging import (
Logging as LiteLLMLoggingObj,
)
logging_obj = LiteLLMLoggingObj(
model="gpt-4o-mini",
messages=[{"role": "user", "content": "hi"}],
stream=True,
call_type="completion", # not pass_through_endpoint
start_time=time.time(),
litellm_call_id="test-id",
function_id="fn",
)
# litellm_params with no recognized async marker -> classified sync.
logging_obj.model_call_details["litellm_params"] = {}
assert LiteLLMLoggingObj._is_sync_litellm_request({}) is True
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=False,
),
patch("litellm.callbacks", [PostCallGuardrail()]),
):
await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails(
captured_data={"model": "gpt-4o-mini", "metadata": {}},
captured_user_api_key_dict=UserAPIKeyAuth(api_key="test"),
captured_logging_obj=logging_obj,
assembled_response=MagicMock(),
cache_hit=False,
)
await asyncio.sleep(0)
await asyncio.sleep(0)
mock_async.assert_awaited_once()
mock_sync.assert_not_called()
# ---------------------------------------------------------------------------
# 7. _fire_deferred_stream_logging
@ -1054,7 +1133,7 @@ class TestFireDeferredStreamLogging:
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 InfoWritingGuardrail(CustomGuardrail):
def __init__(self):

18
tests/test_litellm/proxy/test_proxy_types.py
View file

@ -69,3 +69,21 @@ def test_internal_jobs_user_has_proxy_admin_role():
assert system_user.user_id == "system"
assert system_user.team_id == "system"
assert system_user.team_alias == "system"
def test_user_api_key_auth_hashes_authorization_header_form_of_key():
from litellm.proxy._types import UserAPIKeyAuth
raw_key = "sk-AbCdEfGhIjKlMnOpQrStUvWxYz0123456789"
baseline = UserAPIKeyAuth(api_key=raw_key)
for header_form in (
f"Bearer {raw_key}",
f"bearer {raw_key}",
f"BEARER {raw_key}",
f"BeArEr {raw_key}",
):
from_header = UserAPIKeyAuth(api_key=header_form)
assert from_header.api_key == baseline.api_key
assert from_header.token == baseline.token
assert not from_header.api_key.lower().startswith("bearer")