From 4fbe2276a1ece2f58d047c60660cfe630fdbbede Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 10 Sep 2026 18:15:24 -0700 Subject: [PATCH] fix(logging): finish response metadata before the sync logging thread reads it (#39869) * fix(logging): finish response metadata before the sync logging thread reads it The async and sync client wrappers handed the response to the threaded success handler before computing its cost, call id, and api_base, so that thread inserted into the same metadata dict the request coroutine was still iterating and a finished chat completion turned into a 500 (dictionary changed size during iteration). Metadata is now finalized first, and the merge and header copies snapshot their dicts before iterating. * fix(logging): snapshot metadata with a dict copy and drop redundant comment Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(logging): copy metadata via dict.copy and dedupe Final import Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/litellm_logging.py | 39 +++--- litellm/utils.py | 112 +++++++++++------- .../test_litellm_logging.py | 59 +++++++++ tests/test_litellm/test_utils.py | 46 +++++++ 4 files changed, 193 insertions(+), 63 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 893bfeefff8..6ee68ab21c5 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -295,6 +295,9 @@ def _get_provider_request_id(original_exception: Exception) -> str | None: # Cache custom pricing keys as frozenset for O(1) lookups instead of looping through 49 keys _CUSTOM_PRICING_KEYS: Final[frozenset[str]] = frozenset(CustomPricingLiteLLMParams.model_fields.keys()) _MODEL_INFO_CUSTOM_PRICING_KEYS: Final[frozenset[str]] = _CUSTOM_PRICING_KEYS | DEPLOYMENT_SCOPED_PRICING_FIELDS +_UNSERIALIZABLE_METADATA_KEYS: Final[frozenset[str]] = frozenset( + ("user_api_key_auth", "user_api_key_budget_reservation") +) sentry_sdk_instance = None capture_exception = None @@ -5386,23 +5389,23 @@ class StandardLoggingPayloadSetup: Returns: dict: Merged metadata with user API key fields taking precedence """ - merged_metadata: Final[dict] = {} - - # Start with metadata (user API key fields) - but skip non-serializable objects - if litellm_params.get("metadata") and isinstance(litellm_params.get("metadata"), dict): - for key, value in litellm_params["metadata"].items(): - # Skip non-serializable objects like UserAPIKeyAuth - if key in {"user_api_key_auth", "user_api_key_budget_reservation"}: - continue - merged_metadata[key] = value - - # Then merge litellm_metadata (model-related fields) - this will NOT overwrite existing keys - if litellm_params.get("litellm_metadata") and isinstance(litellm_params.get("litellm_metadata"), dict): - for key, value in litellm_params["litellm_metadata"].items(): - if key not in merged_metadata: # Don't overwrite existing keys from metadata - merged_metadata[key] = value - - return merged_metadata + metadata: Final = litellm_params.get("metadata") + litellm_metadata: Final = litellm_params.get("litellm_metadata") + user_metadata: Final = MappingProxyType( + { + key: value + for key, value in (metadata.copy().items() if isinstance(metadata, dict) else ()) + if key not in _UNSERIALIZABLE_METADATA_KEYS + } + ) + model_metadata: Final = MappingProxyType( + { + key: value + for key, value in (litellm_metadata.copy().items() if isinstance(litellm_metadata, dict) else ()) + if key not in user_metadata + } + ) + return {**user_metadata, **model_metadata} # mutable-ok: function contract returns a plain dict @staticmethod def get_standard_logging_metadata( @@ -5660,7 +5663,7 @@ class StandardLoggingPayloadSetup: additional_logging_headers[key] = additiona_headers[_key] # Preserve all remaining headers verbatim (e.g. llm_provider-x-request-id) - for k, v in additiona_headers.items(): + for k, v in additiona_headers.copy().items(): if k.lower() not in typed_keys: additional_logging_headers[k] = v diff --git a/litellm/utils.py b/litellm/utils.py index b93938a480e..a765e1b1246 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1196,6 +1196,47 @@ def function_setup( raise e +def _dispatch_success_logging( + logging_obj: LiteLLMLoggingObject, + result: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + is_completion_with_fallbacks: bool, + is_litellm_internal_call: bool, +) -> None: + if not is_litellm_internal_call: + if getattr(logging_obj, "_defer_async_logging", False): + + def _enqueue_deferred_logging() -> None: + asyncio.create_task( + _client_async_logging_helper( + logging_obj=logging_obj, + result=result, + start_time=start_time, + end_time=end_time, + is_completion_with_fallbacks=is_completion_with_fallbacks, + ) + ) + + logging_obj._enqueue_deferred_logging = _enqueue_deferred_logging + else: + asyncio.create_task( + _client_async_logging_helper( + logging_obj=logging_obj, + result=result, + start_time=start_time, + end_time=end_time, + is_completion_with_fallbacks=is_completion_with_fallbacks, + ) + ) + + logging_obj.handle_sync_success_callbacks_for_async_calls( + result=result, + start_time=start_time, + end_time=end_time, + ) + + async def _client_async_logging_helper( logging_obj: LiteLLMLoggingObject, result, @@ -1663,6 +1704,16 @@ def client(original_function): kwargs=kwargs, ) + _update_response_metadata: Final = getattr(sys.modules[__name__], "update_response_metadata") + _update_response_metadata( + result=result, + logging_obj=logging_obj, + model=model, + kwargs=kwargs, + start_time=start_time, + end_time=end_time, + ) + # LOG SUCCESS - handle streaming success logging in the _next_ object, remove `handle_success` once it's deprecated verbose_logger.info("Wrapper: Completed Call, calling success_handler") # Copy the current context to propagate it to the background thread @@ -1677,15 +1728,6 @@ def client(original_function): end_time, ) # RETURN RESULT - update_response_metadata = getattr(sys.modules[__name__], "update_response_metadata") - update_response_metadata( - result=result, - logging_obj=logging_obj, - model=model, - kwargs=kwargs, - start_time=start_time, - end_time=end_time, - ) return result except Exception as e: call_type = original_function.__name__ @@ -1945,48 +1987,20 @@ def client(original_function): args=args, ) - # LOG SUCCESS - handle streaming success logging in the _next_ object - # Internal sub-calls (e.g. emulated file-search steps) share the - # parent's logging obj; skip async logging here so only the outer call bills once. - # NOTE: streaming requests return early (before this point) via - # CustomStreamWrapper, so this block is non-streaming only. - if not _is_litellm_internal_call: - if getattr(logging_obj, "_defer_async_logging", False): - - def _enqueue_deferred_logging() -> None: - asyncio.create_task( - _client_async_logging_helper( - logging_obj=logging_obj, - result=result, - start_time=start_time, - end_time=end_time, - is_completion_with_fallbacks=is_completion_with_fallbacks, - ) - ) - - logging_obj._enqueue_deferred_logging = _enqueue_deferred_logging - else: - asyncio.create_task( - _client_async_logging_helper( - logging_obj=logging_obj, - result=result, - start_time=start_time, - end_time=end_time, - is_completion_with_fallbacks=is_completion_with_fallbacks, - ) - ) - - logging_obj.handle_sync_success_callbacks_for_async_calls( - result=result, - start_time=start_time, - end_time=end_time, - ) # REBUILD EMBEDDING CACHING if ( isinstance(result, EmbeddingResponse) and _caching_handler_response is not None and _caching_handler_response.final_embedding_cached_response is not None ): + _dispatch_success_logging( + logging_obj=logging_obj, + result=result, + start_time=start_time, + end_time=end_time, + is_completion_with_fallbacks=is_completion_with_fallbacks, + is_litellm_internal_call=_is_litellm_internal_call, + ) return _llm_caching_handler._combine_cached_embedding_response_with_api_result( _caching_handler_response=_caching_handler_response, embedding_response=result, @@ -2002,6 +2016,14 @@ def client(original_function): start_time=start_time, end_time=end_time, ) + _dispatch_success_logging( + logging_obj=logging_obj, + result=result, + start_time=start_time, + end_time=end_time, + is_completion_with_fallbacks=is_completion_with_fallbacks, + is_litellm_internal_call=_is_litellm_internal_call, + ) return result except Exception as e: diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index aeaa8f18c26..2f6339dcdbb 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -3,6 +3,7 @@ import contextlib import datetime import os import sys +from collections.abc import Callable from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch @@ -6945,3 +6946,61 @@ def test_classifier_audit_is_not_added_to_other_calls(logging_obj, call_type, or logging_obj.model_call_details["litellm_params"] = {"metadata": {"internal_call_origin": origin}} logging_obj.pre_call(input=[], api_key=None, additional_args={"complete_input_dict": {"input": "embedding"}}) assert logging_obj.classifier_input is None + + +def _run_while_a_thread_grows(target: dict, read: Callable[[], None], reads: int) -> None: + import itertools + import threading + + stop: Final = threading.Event() + + def grow() -> None: + for counter in itertools.count(): + if stop.is_set(): + return + key: Final = f"late_{counter % 64}" + if key in target: + del target[key] + else: + target[key] = counter + + writer: Final = threading.Thread(target=grow, daemon=True) + previous_interval: Final = sys.getswitchinterval() + sys.setswitchinterval(1e-6) + writer.start() + try: + for _ in range(reads): + read() + finally: + stop.set() + writer.join(timeout=5) + sys.setswitchinterval(previous_interval) + + +def test_merge_litellm_metadata_survives_a_thread_growing_metadata_mid_merge(): + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + metadata: Final = {f"key_{i}": i for i in range(2000)} + litellm_params: Final = {"metadata": metadata, "litellm_metadata": {"model_group": "gpt"}} + + def read() -> None: + merged: Final = StandardLoggingPayloadSetup.merge_litellm_metadata(litellm_params) + assert merged["key_1999"] == 1999 + assert merged["model_group"] == "gpt" + + _run_while_a_thread_grows(metadata, read, reads=300) + + +def test_get_additional_headers_survives_a_thread_growing_headers_mid_copy(): + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + headers: Final = {f"llm_provider-x-custom-{i}": str(i) for i in range(2000)} + headers["x-ratelimit-remaining-requests"] = "7" + + def read() -> None: + copied: Final = StandardLoggingPayloadSetup.get_additional_headers(headers) + assert copied is not None + assert copied["x_ratelimit_remaining_requests"] == 7 + assert copied["llm_provider-x-custom-1999"] == "1999" + + _run_while_a_thread_grows(headers, read, reads=300) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 2bdde43ff2d..7753cb7d770 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -1,9 +1,12 @@ import asyncio +import contextlib import json import logging import os +import queue import threading from datetime import datetime, timedelta, timezone +from collections.abc import Iterator from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -6298,3 +6301,46 @@ async def test_async_mock_completion_streaming_obj_raises_mock_exception_before_ ) with pytest.raises(litellm.MockException): await _async_mock_stream_snapshots(mock_exception, 51234) + + + +@contextlib.contextmanager +def _recording_hidden_params_at_submit(submit_target: str) -> "Iterator[queue.SimpleQueue[dict[str, object]]]": + seen: Final = queue.SimpleQueue() + + def record_submit(_fn, *args, **_kwargs): + response: Final = next(arg for arg in args if isinstance(arg, litellm.ModelResponse)) + seen.put(dict(response._hidden_params)) + return MagicMock() + + with patch(submit_target, side_effect=record_submit): + yield seen + + +@pytest.mark.asyncio +async def test_acompletion_finishes_response_metadata_before_handing_the_response_to_the_logging_thread(monkeypatch): + monkeypatch.setattr(litellm, "success_callback", [lambda kwargs, response, start_time, end_time: None]) + with _recording_hidden_params_at_submit("litellm.litellm_core_utils.litellm_logging.executor.submit") as seen: + await litellm.acompletion( + model="gpt-5.5", + messages=[{"role": "user", "content": "hi"}], + mock_response="Hello there!", + num_retries=0, + ) + snapshot: Final = seen.get_nowait() + assert snapshot["litellm_call_id"] + assert snapshot["response_cost"] is not None + assert snapshot["api_base"] + + +def test_completion_finishes_response_metadata_before_handing_the_response_to_the_logging_thread(): + with _recording_hidden_params_at_submit("litellm.utils.executor.submit") as seen: + litellm.completion( + model="gpt-5.5", + messages=[{"role": "user", "content": "hi"}], + mock_response="Hello there!", + ) + snapshot: Final = seen.get_nowait() + assert snapshot["litellm_call_id"] + assert snapshot["response_cost"] is not None + assert snapshot["api_base"]