diff --git a/litellm/__init__.py b/litellm/__init__.py index e0c7d56361c..bc8a13ec2cd 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -197,6 +197,7 @@ standard_logging_payload_excluded_fields: Optional[List[str]] = ( None # Fields to exclude from StandardLoggingPayload before callbacks receive it ) log_raw_request_response: bool = False +request_correlation_in_logs: bool = False redact_messages_in_exceptions: Optional[bool] = False redact_user_api_key_info: Optional[bool] = False # When True (default — preserves historical behavior), the Router appends diff --git a/litellm/_logging.py b/litellm/_logging.py index b9e102e2b3c..6add9d79a5b 100644 --- a/litellm/_logging.py +++ b/litellm/_logging.py @@ -1,4 +1,5 @@ import ast +import contextvars import logging import os import sys @@ -6,12 +7,44 @@ from datetime import datetime from logging import Formatter from typing import Any, Final +import litellm from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.secret_redaction import redact_string set_verbose = False +session_id_var: Final[contextvars.ContextVar[str]] = contextvars.ContextVar("session_id", default="") +trace_id_var: Final[contextvars.ContextVar[str]] = contextvars.ContextVar("trace_id", default="") + +_MAX_CORRELATION_ID_LENGTH: Final = 256 + + +def _sanitize_correlation_id(value: str) -> str: + """Strip control characters, bound length, and redact credential-shaped + content before a caller-controlled trace_id/session_id (e.g. + litellm_session_id, x-litellm-trace-id) is stamped into log lines. + + Without the first two, a caller could embed \\r/\\n or terminal escape + sequences to forge fake log entries, or submit an oversized value repeated + across every log line for the request. Without the redaction, a caller + could smuggle a real credential (e.g. an sk-... key) through this field: + CorrelationContextFilter stamps trace_id/session_id onto the record after + SecretRedactionFilter has already run, so those two fields never otherwise + pass through credential redaction. + """ + stripped: Final = "".join(ch for ch in value if ch.isprintable()) + return _redact_string(stripped[:_MAX_CORRELATION_ID_LENGTH]) + + +def set_session_id(session_id: str) -> "contextvars.Token[str]": + return session_id_var.set(_sanitize_correlation_id(session_id)) + + +def set_trace_id(trace_id: str) -> "contextvars.Token[str]": + return trace_id_var.set(_sanitize_correlation_id(trace_id)) + + if set_verbose is True: logging.warning( "`litellm.set_verbose` is deprecated. Please set `os.environ['LITELLM_LOG'] = 'DEBUG'` for debug logs." @@ -77,6 +110,28 @@ class SecretRedactionFilter(logging.Filter): _secret_filter: Final = SecretRedactionFilter() +class CorrelationContextFilter(logging.Filter): + """Stamps each log record with the current request's trace_id and session_id from contextvars. + + Works in tandem with JsonFormatter: the formatter's record.__dict__ loop picks up these + attributes as first-class JSON fields without any formatter-level code. + """ + + def filter(self, record: logging.LogRecord) -> bool: + if not litellm.request_correlation_in_logs: + return True + trace_id: Final = trace_id_var.get() + if trace_id: + record.trace_id = trace_id # rebind-ok: stamping the LogRecord is the Filter interface's contract + session_id: Final = session_id_var.get() + if session_id: + record.session_id = session_id # rebind-ok: stamping the LogRecord is the Filter interface's contract + return True + + +_correlation_filter: Final = CorrelationContextFilter() + + json_logs = bool(os.getenv("JSON_LOGS", False)) # Create a handler for the logger (you may need to adapt this based on your needs) log_level: Final = os.getenv("LITELLM_LOG", "DEBUG") @@ -84,6 +139,7 @@ numeric_level: Final[str] = getattr(logging, log_level.upper()) handler: Final = logging.StreamHandler() handler.setLevel(numeric_level) handler.addFilter(_secret_filter) +handler.addFilter(_correlation_filter) def _try_parse_json_message(message: str) -> dict[str, Any] | None: @@ -146,6 +202,11 @@ def _get_standard_record_attrs() -> frozenset: _STANDARD_RECORD_ATTRS: Final = _get_standard_record_attrs() +# CorrelationContextFilter is the only legitimate source for these two JSON fields; +# see JsonFormatter.format() for why they're excluded from the generic message-content +# and extra-attribute promotion paths. +_RESERVED_CORRELATION_FIELDS: Final = frozenset(("trace_id", "session_id")) + class JsonFormatter(Formatter): def __init__(self): @@ -164,13 +225,18 @@ class JsonFormatter(Formatter): "timestamp": self.formatTime(record), } - # Parse embedded JSON or Python dict repr in message so sub-fields become first-class properties + # Parse embedded JSON or Python dict repr in message so sub-fields become first-class properties. + # trace_id/session_id are excluded here unconditionally (not just "if not already + # set") - CorrelationContextFilter is the only legitimate source for these two + # fields, and a message that merely happens to parse as JSON/dict (e.g. a proxy + # log line dumping raw request headers) must never be able to claim them, even on + # a record the filter hasn't stamped yet (no correlation context active for it). parsed = _try_parse_json_message(message_str) if parsed is None: parsed = _try_parse_embedded_python_dict(message_str) if parsed is not None: for key, value in parsed.items(): - if key not in json_record: + if key not in json_record and key not in _RESERVED_CORRELATION_FIELDS: json_record[key] = value # Include extra attributes passed via logger.debug("msg", extra={...}) @@ -178,6 +244,18 @@ class JsonFormatter(Formatter): if key not in _STANDARD_RECORD_ATTRS and key not in json_record: json_record[key] = value + # trace_id/session_id are reserved: CorrelationContextFilter is the only + # legitimate source for these two fields. Without this, a message string + # that happens to parse as JSON/dict (e.g. a proxy log line dumping raw + # request headers) with a "trace_id"/"session_id" key would have already + # claimed the key at the parsed-message step above, and the extra-attributes + # loop's "key not in json_record" guard would then skip the real value - + # letting a caller-supplied header spoof another request's correlation ids. + for reserved_key in _RESERVED_CORRELATION_FIELDS: + value = getattr(record, reserved_key, None) + if value: + json_record[reserved_key] = value + # Set component/logger only if not already supplied via extra={...} if "component" not in json_record: json_record["component"] = record.name @@ -190,12 +268,34 @@ class JsonFormatter(Formatter): return safe_dumps(json_record) +class CorrelationPlainFormatter(logging.Formatter): + """Appends trace_id/session_id to plain-text log lines stamped by CorrelationContextFilter. + + Mirrors JsonFormatter's handling of these two fields so request_correlation_in_logs + behaves the same whether or not json_logs is enabled. + """ + + def format(self, record: logging.LogRecord) -> str: + formatted: Final = super().format(record) + trace_id: Final = getattr(record, "trace_id", None) + session_id: Final = getattr(record, "session_id", None) + if not trace_id and not session_id: + return formatted + parts: Final = tuple( + p + for p in (f"trace_id={trace_id}" if trace_id else None, f"session_id={session_id}" if session_id else None) + if p + ) + return f"{formatted} [{' '.join(parts)}]" + + # Function to set up exception handlers for JSON logging def _setup_json_exception_handlers(formatter): # Create a handler with JSON formatting for exceptions error_handler: Final = logging.StreamHandler() error_handler.setFormatter(formatter) error_handler.addFilter(_secret_filter) + error_handler.addFilter(_correlation_filter) # Setup excepthook for uncaught exceptions def json_excepthook(exc_type, exc_value, exc_traceback): @@ -243,7 +343,7 @@ if json_logs: handler.setFormatter(JsonFormatter()) _setup_json_exception_handlers(JsonFormatter()) else: - formatter: Final = logging.Formatter( + formatter: Final = CorrelationPlainFormatter( "\033[92m%(asctime)s - %(name)s:%(levelname)s\033[0m: %(filename)s:%(lineno)s - %(message)s", datefmt="%H:%M:%S", ) @@ -346,6 +446,7 @@ def _initialize_loggers_with_handler(handler: logging.Handler): - Prevents bubbling to parent/root (critical to prevent duplicate JSON logs) """ handler.addFilter(_secret_filter) + handler.addFilter(_correlation_filter) for lg in _get_loggers_to_initialize(): lg.handlers.clear() # remove any existing handlers lg.addHandler(handler) # add JSON formatter handler diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 99721c3ffa2..a3ff048e92a 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -10,7 +10,7 @@ import subprocess import sys import time import traceback -from collections.abc import Callable +from collections.abc import Callable, Mapping from datetime import datetime as dt_object from functools import lru_cache from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast @@ -25,7 +25,15 @@ from litellm import ( log_raw_request_response, turn_off_message_logging, ) -from litellm._logging import _is_debugging_on, _redact_string, verbose_logger +from litellm._logging import ( + _is_debugging_on, + _redact_string, + session_id_var, + set_session_id, + set_trace_id, + trace_id_var, + verbose_logger, +) from litellm._uuid import uuid from litellm.batches.batch_utils import _handle_completed_batch from litellm.caching.caching import DualCache, InMemoryCache @@ -313,6 +321,7 @@ class Logging(LiteLLMLoggingBaseClass): applied_guardrails: list[str] | None = None, kwargs: dict | None = None, log_raw_request_response: bool = False, + supports_correlation_logging: bool = True, ): _input: Final[str | None] = messages # save original value of messages if messages is not None: @@ -338,6 +347,36 @@ class Logging(LiteLLMLoggingBaseClass): self.call_type = call_type self.litellm_call_id = litellm_call_id self.litellm_trace_id: str = litellm_trace_id if litellm_trace_id else str(uuid.uuid4()) + + # Capture the pre-call *value* (not a contextvars.Token) so restoration works + # even if this attempt's own logging ends up dispatched onto a different + # asyncio Task/context (e.g. via asyncio.create_task or the logging worker) - + # a Token can only be reset in the exact Context where it was created. + self._pre_call_trace_id: str = trace_id_var.get() + self._pre_call_session_id: str = session_id_var.get() + _sid: Final = kwargs.get("litellm_session_id") if kwargs else None + self.litellm_session_id: str = str(_sid) if _sid else "" + # supports_correlation_logging is False for calls originating from the + # sync client entry point (wrapper() in utils.py): a plain OS thread + # has no per-call context isolation the way an asyncio Task does, and + # a thread pool's worker threads are recycled across unrelated + # requests, so stamping trace_id/session_id there risks one request's + # ids leaking into a different, later request on the same thread. Sync + # support is deferred to a follow-up PR with its own safe-restore + # mechanism; async calls (the proxy's only call path) are unaffected. + if supports_correlation_logging: + set_trace_id(self.litellm_trace_id) + set_session_id(self.litellm_session_id) + # set_trace_id()/set_session_id() sanitize (strip control chars, bound + # length) before storing, so the contextvar's actual value can differ + # from self.litellm_trace_id/litellm_session_id. Capture what was + # really stored - _restore_correlation_context_if_unclaimed() must + # compare against this, not the raw ids, or a caller-supplied id + # containing control characters/oversized input would never match + # and cleanup would be skipped forever. + self._own_trace_id: str = trace_id_var.get() + self._own_session_id: str = session_id_var.get() + self.function_id = function_id self.streaming_chunks: list[Any] = [] # for generating complete stream response self.sync_streaming_chunks: list[Any] = [] # for generating complete stream response @@ -1992,7 +2031,67 @@ class Logging(LiteLLMLoggingBaseClass): if complete_streaming_response is not None: await self.async_success_handler(result=complete_streaming_response) - def success_handler(self, result=None, start_time=None, end_time=None, cache_hit=None, **kwargs): + def _restore_correlation_context(self) -> None: + """Restore trace_id/session_id contextvars to their pre-call value. + + Without this, a nested LiteLLM call sharing the same asyncio Task as an + outer request (e.g. a guardrail's own LLM-as-judge call, an MCP sampling + call) would leave the outer request's subsequent log lines stamped with + the nested call's trace_id/session_id instead of its own. + + Uses a plain set() of the captured pre-call value rather than + contextvars.Token-based reset(), since this can end up called from a + different asyncio Task/context than __init__ ran in (e.g. the request + task's own wrapper() finally block, plus async_success_handler + dispatched separately via asyncio.create_task/the logging worker) - + reset() only works in the exact Context a Token was created in and + raises otherwise. Deliberately NOT idempotent/guarded: each distinct + Task that calls this needs its own restore to actually take effect in + that Task's view of the contextvars, so calling it multiple times + (once per Task involved in this attempt) is required, not just safe. + """ + set_trace_id(self._pre_call_trace_id) + set_session_id(self._pre_call_session_id) + + def _restore_correlation_context_if_unclaimed(self) -> None: + """Guarded variant for __del__-triggered cleanup only. + + __del__ can fire arbitrarily late (delayed by cyclic GC, possibly + after the consuming Task/thread has already moved on to a different, + still-active call). Unconditionally restoring in that case would + stomp the active call's trace_id/session_id with this abandoned + stream's stale pre-call snapshot. Only restore if the contextvars + still hold the ids *this* call set - i.e. nothing has claimed them + since - so an unrelated active call is never overwritten. + """ + if trace_id_var.get() == self._own_trace_id and session_id_var.get() == self._own_session_id: + self._restore_correlation_context() + + def success_handler( + self, + result: Any = None, # heterogeneous response object; varies by call type (ANN401 ignored, see ruff-strict.toml) + start_time: datetime.datetime | None = None, + end_time: datetime.datetime | None = None, + cache_hit: bool | None = None, + **kwargs: Any, # kwargs-ok: forwarded to _success_handler_body + ) -> None: + """Restores trace_id/session_id contextvars once this attempt's own success + logging (including any nested calls its callbacks trigger) is fully done.""" + try: + return self._success_handler_body( + result=result, start_time=start_time, end_time=end_time, cache_hit=cache_hit, **kwargs + ) + finally: + self._restore_correlation_context() + + def _success_handler_body( + self, + result: Any = None, # heterogeneous response object; varies by call type (ANN401 ignored, see ruff-strict.toml) + start_time: datetime.datetime | None = None, + end_time: datetime.datetime | None = None, + cache_hit: bool | None = None, + **kwargs: Any, # kwargs-ok: forwarded from success_handler + ) -> None: verbose_logger.debug("Logging Details LiteLLM-Success Call: Cache_hit=%s", cache_hit) if not self.should_run_logging(event_type="sync_success"): # prevent double logging return @@ -2399,7 +2498,31 @@ class Logging(LiteLLMLoggingBaseClass): e, ) - async def async_success_handler(self, result=None, start_time=None, end_time=None, cache_hit=None, **kwargs): + async def async_success_handler( + self, + result: Any = None, # heterogeneous response object; varies by call type (ANN401 ignored, see ruff-strict.toml) + start_time: datetime.datetime | None = None, + end_time: datetime.datetime | None = None, + cache_hit: bool | None = None, + **kwargs: Any, # kwargs-ok: forwarded to _async_success_handler_body + ) -> None: + """Restores trace_id/session_id contextvars once this attempt's own success + logging (including any nested calls its callbacks trigger) is fully done.""" + try: + return await self._async_success_handler_body( + result=result, start_time=start_time, end_time=end_time, cache_hit=cache_hit, **kwargs + ) + finally: + self._restore_correlation_context() + + async def _async_success_handler_body( + self, + result: Any = None, # heterogeneous response object; varies by call type (ANN401 ignored, see ruff-strict.toml) + start_time: datetime.datetime | None = None, + end_time: datetime.datetime | None = None, + cache_hit: bool | None = None, + **kwargs: Any, # kwargs-ok: forwarded from async_success_handler + ) -> None: """ Implementing async callbacks, to handle asyncio event loop issues when custom integrations need to use async functions. """ @@ -2791,7 +2914,32 @@ class Logging(LiteLLMLoggingBaseClass): kwargs=self.model_call_details, ) - def failure_handler(self, exception, traceback_exception, start_time=None, end_time=None): + def failure_handler( + self, + exception: Exception, + traceback_exception: str, + start_time: datetime.datetime | None = None, + end_time: datetime.datetime | None = None, + ) -> None: + """Restores trace_id/session_id contextvars once this attempt's own failure + logging (including any nested calls its callbacks trigger) is fully done.""" + try: + return self._failure_handler_body( + exception=exception, + traceback_exception=traceback_exception, + start_time=start_time, + end_time=end_time, + ) + finally: + self._restore_correlation_context() + + def _failure_handler_body( + self, + exception: Exception, + traceback_exception: str, + start_time: datetime.datetime | None = None, + end_time: datetime.datetime | None = None, + ) -> None: verbose_logger.debug("Logging Details LiteLLM-Failure Call: %s", litellm.failure_callback) if not self.should_run_logging(event_type="sync_failure"): # prevent double logging return @@ -2960,7 +3108,32 @@ class Logging(LiteLLMLoggingBaseClass): "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while failure logging %s", e ) - async def async_failure_handler(self, exception, traceback_exception, start_time=None, end_time=None): + async def async_failure_handler( + self, + exception: Exception, + traceback_exception: str, + start_time: datetime.datetime | None = None, + end_time: datetime.datetime | None = None, + ) -> None: + """Restores trace_id/session_id contextvars once this attempt's own failure + logging (including any nested calls its callbacks trigger) is fully done.""" + try: + return await self._async_failure_handler_body( + exception=exception, + traceback_exception=traceback_exception, + start_time=start_time, + end_time=end_time, + ) + finally: + self._restore_correlation_context() + + async def _async_failure_handler_body( + self, + exception: Exception, + traceback_exception: str, + start_time: datetime.datetime | None = None, + end_time: datetime.datetime | None = None, + ) -> None: """ Implementing async callbacks, to handle asyncio event loop issues when custom integrations need to use async functions. """ @@ -5061,33 +5234,61 @@ class StandardLoggingPayloadSetup: return end_time_float - start_time_float @staticmethod - def _get_standard_logging_payload_trace_id( + def get_standard_logging_payload_trace_id( logging_obj: Logging, - litellm_params: dict, + litellm_params: Mapping[str, Any], ) -> str: """ Returns the `litellm_trace_id` for this request This helps link sessions when multiple requests are made in a single session + + Gated behind `litellm.request_correlation_in_logs`: + - Off (default): legacy behavior, preserved for backward compatibility - + `litellm_session_id` takes priority over `litellm_trace_id` since historically + this field doubled as the session-grouping field. + - On: `litellm_trace_id` takes priority - trace_id and session_id are independent, + see `get_standard_logging_payload_session_id` for session tracking. """ dynamic_litellm_session_id: Final = litellm_params.get("litellm_session_id") dynamic_litellm_trace_id: Final = litellm_params.get("litellm_trace_id") + metadata: Final = litellm_params.get("metadata") + metadata_session_id: Final = metadata.get("session_id") if metadata else None + metadata_trace_id: Final = metadata.get("trace_id") if metadata else None - # Note: we recommend using `litellm_session_id` for session tracking - # `litellm_trace_id` is an internal litellm param + ordered_candidates: Final[tuple[Any, Any, Any, Any]] = ( + (dynamic_litellm_trace_id, dynamic_litellm_session_id, metadata_trace_id, metadata_session_id) + if litellm.request_correlation_in_logs + else (dynamic_litellm_session_id, dynamic_litellm_trace_id, metadata_session_id, metadata_trace_id) + ) + for candidate in ordered_candidates: + if candidate: + return str(candidate) + return logging_obj.litellm_trace_id + + @staticmethod + def get_standard_logging_payload_session_id( + logging_obj: Logging, + litellm_params: Mapping[str, Any], + ) -> str: + """ + Returns the end-user/conversation `litellm_session_id` for this request, independent of trace_id. + + Only populated when `litellm.request_correlation_in_logs` is enabled - off by default + to avoid changing existing StandardLoggingPayload shape for callers who haven't opted in. + Unlike `get_standard_logging_payload_trace_id`, this never falls back to a generated + per-call trace id: it's empty when the caller never supplied a session id. + """ + if not litellm.request_correlation_in_logs: + return "" + dynamic_litellm_session_id: Final = litellm_params.get("litellm_session_id") if dynamic_litellm_session_id: return str(dynamic_litellm_session_id) - elif dynamic_litellm_trace_id: - return str(dynamic_litellm_trace_id) - # Fallback: use metadata.session_id or metadata.trace_id for call chaining - metadata: Final = litellm_params.get("metadata") or {} - metadata_session_id: Final = metadata.get("session_id") - metadata_trace_id: Final = metadata.get("trace_id") + metadata: Final = litellm_params.get("metadata") + metadata_session_id: Final = metadata.get("session_id") if metadata else None if metadata_session_id: return str(metadata_session_id) - if metadata_trace_id: - return str(metadata_trace_id) - return logging_obj.litellm_trace_id + return logging_obj.litellm_session_id @staticmethod def _get_user_agent_tags(proxy_server_request: dict) -> list[str] | None: @@ -5392,7 +5593,11 @@ def get_standard_logging_object_payload( payload: Final[StandardLoggingPayload] = StandardLoggingPayload( id=str(id), litellm_call_id=kwargs.get("litellm_call_id") or litellm_params.get("litellm_call_id"), - trace_id=StandardLoggingPayloadSetup._get_standard_logging_payload_trace_id( + trace_id=StandardLoggingPayloadSetup.get_standard_logging_payload_trace_id( + logging_obj=logging_obj, + litellm_params=litellm_params, + ), + session_id=StandardLoggingPayloadSetup.get_standard_logging_payload_session_id( logging_obj=logging_obj, litellm_params=litellm_params, ), diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 68465d06b15..2dc71abee3e 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -213,7 +213,75 @@ class CustomStreamWrapper: def __aiter__(self) -> AsyncIterator["ModelResponseStream"]: return self + def _restore_consumer_correlation_context(self, *, guarded: bool = False) -> None: + """Restore trace_id/session_id in the *consuming* thread/task/context. + + wrapper_async() deliberately skips restoring correlation context when + it returns a stream, so log lines emitted while the caller iterates it + still carry this call's ids (see request_correlation_in_logs). + wrapper() (the sync path) never stamps anything in the first place - + see Logging.__init__'s supports_correlation_logging - so this method + is an inert no-op for sync-created streams, harmless to call anyway + since the class is shared between __next__ and __anext__. + But the terminal success/failure handlers this stream dispatches to + finish the job run on a *different* Task/thread (asyncio.create_task, + threading.Thread, or the shared executor) - restoring there fixes up + that detached context, not the one actually running the caller's + `for`/`async for` loop. Call this at every point control genuinely + returns to that consuming context: natural exhaustion (StopIteration/ + StopAsyncIteration), a raised failure, or explicit aclose(). Never let + this raise - it must not break the caller's actual stream handling. + + guarded=True (only __del__ uses this) skips the restore unless the + contextvars still hold the ids this stream's own call set, so a + delayed finalizer never overwrites a different, still-active call + that has since taken over the same Task/thread's context. + """ + try: + logging_obj: Final = getattr(self, "logging_obj", None) + if logging_obj is None: + return + method_name: Final = ( + "_restore_correlation_context_if_unclaimed" if guarded else "_restore_correlation_context" + ) + restore: Final = getattr(logging_obj, method_name, None) + if restore is not None: + restore() + except Exception as restore_error: # noqa: BLE001 # best-effort cleanup; must not raise into the caller + verbose_logger.debug("could not restore correlation context: %s", restore_error) + + def __del__(self) -> None: + """Best-effort correlation-context cleanup for an abandoned async stream. + + Only meaningfully applies to streams created by wrapper_async(): it + leaves contextvars "open" across the caller's iteration, so if the + caller never fully consumes the stream - stops early, drops the + reference, cancels it - none of the exit points + _restore_consumer_correlation_context() is called from ever run. For a + sync stream (wrapper()), this is a no-op in practice: wrapper() never + stamps trace_id/session_id for sync calls in the first place (see + Logging.__init__'s supports_correlation_logging), so there is nothing + for this to clean up. + + This is a best-effort fallback, not a guarantee: __del__ timing is + unpredictable (delayed by cyclic GC, not guaranteed at interpreter + shutdown, and may run on a different thread), so this can only reduce + how long the leak persists, not eliminate it. That's an acceptable + trade specifically because its blast radius is bounded to the one + asyncio Task this stream's own call ran in - each async call has its + own copy of the contextvars, and Tasks (unlike a thread pool's worker + threads) are never recycled across requests, so a delayed or missed + cleanup here can never misattribute a *different* request's logs. + guarded=True additionally ensures it never clobbers a different, + still-active call's context within that same Task if this fires late. + """ + self._restore_consumer_correlation_context(guarded=True) + async def aclose(self): + # Restore the consumer's outer context only after the underlying + # provider stream's own close (and its diagnostic logging below, if + # closing fails) completes - not before - so those log lines still + # carry this closing stream's own trace_id/session_id. if self.completion_stream is not None: stream_to_close: Final = self.completion_stream self.completion_stream = None @@ -233,6 +301,7 @@ class CustomStreamWrapper: "CustomStreamWrapper.aclose: error closing completion_stream: %s", e, ) + self._restore_consumer_correlation_context() def check_send_stream_usage(self, stream_options: dict | None): return stream_options is not None and stream_options.get("include_usage", False) is True @@ -1839,6 +1908,7 @@ class CustomStreamWrapper: if self.sent_stream_usage is False and self.send_stream_usage is True: self.sent_stream_usage = True return response + self._restore_consumer_correlation_context() raise # Re-raise StopIteration else: self.sent_last_chunk = True @@ -1852,6 +1922,19 @@ class CustomStreamWrapper: processed_chunk, cache_hit, ) # log response + # Deliberately do NOT restore context here even though + # completion_stream is already exhausted: this chunk is still + # real data belonging to this call, and the caller's own + # (application-level) log statements processing it run + # immediately after this return, in this same synchronous + # frame - restoring first would make those lines carry the + # wrong ids, which is exactly what leaving context open during + # iteration is meant to prevent (see + # _restore_consumer_correlation_context's docstring). A caller + # that keeps iterating gets cleaned up on its next __next__() + # call (immediate StopIteration, handled above); one that + # stops right here relies on aclose() or the best-effort + # __del__ guard instead. return processed_chunk except Exception as e: traceback_exception: Final = traceback.format_exc() @@ -1879,8 +1962,12 @@ class CustomStreamWrapper: cache_hit = False if self.custom_llm_provider is not None and self.custom_llm_provider == "cached_response": cache_hit = True - self._check_max_streaming_duration() try: + # Inside the try (not before it) so a raised litellm.Timeout flows + # through the same except Exception -> _handle_stream_fallback_error + # path as every other failure, restoring the consumer's correlation + # context - a check before the try would bypass that entirely. + self._check_max_streaming_duration() if self.completion_stream is None: await self.fetch_stream() @@ -2083,10 +2170,17 @@ class CustomStreamWrapper: ) ) + self._restore_consumer_correlation_context() raise StopAsyncIteration # Re-raise StopIteration else: self.sent_last_chunk = True processed_chunk: Final = self.finish_reason_handler() + # see sync __next__'s sibling branch: deliberately do NOT restore + # here - this chunk is still this call's own data, and restoring + # before returning it would corrupt the caller's own log + # statements processing it. A caller that keeps iterating gets + # cleaned up on the next __anext__() call; one that stops here + # relies on aclose() or the best-effort __del__ guard. return processed_chunk def _log_stream_failure_and_raise(self, e: Exception) -> NoReturn: @@ -2138,7 +2232,12 @@ class CustomStreamWrapper: """ from litellm.exceptions import MidStreamFallbackError - # Map to OpenAI exception format + # Map to OpenAI exception format. Some providers' mappers (e.g. + # _map_anthropic_exception, _map_aleph_alpha_exception) synchronously + # log a debug diagnostic (the raw status code) as part of mapping - + # restore the consumer's outer context only after this completes, so + # that diagnostic log line still carries the failing stream's own + # trace_id/session_id instead of the consumer's (or an empty one). if isinstance(e, OpenAIError): mapped_exception: Exception = e else: @@ -2152,6 +2251,7 @@ class CustomStreamWrapper: ) except Exception as mapping_error: mapped_exception = mapping_error + self._restore_consumer_correlation_context() def _normalize_status_code(exc: Exception) -> int | None: """Best-effort status_code extraction.""" diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index c48fee96646..6fd0d52e8da 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -5,6 +5,7 @@ import re import time from collections import OrderedDict from collections.abc import Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final from fastapi import HTTPException, Request @@ -66,6 +67,32 @@ _SESSION_ID_VALUE_RE: Final = re.compile(r"^[a-zA-Z0-9_\-]{8,}$") _SHA256_HEX_RE: Final = re.compile(r"^[0-9a-f]{64}$") +# W3C Trace Context traceparent header: https://www.w3.org/TR/trace-context/ +# e.g. "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01" +_TRACEPARENT_RE: Final = re.compile(r"^[0-9a-f]{2}-([0-9a-f]{32})-[0-9a-f]{16}-[0-9a-f]{2}$", re.IGNORECASE) + + +def _trace_id_from_traceparent(traceparent: str) -> str | None: + """Extract the trace-id from a W3C Trace Context traceparent header, e.g. + "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01" -> the 32-hex + trace-id in the middle. An all-zero trace-id is invalid per spec and is + rejected, matching how the OpenTelemetry SDK itself treats it.""" + match: Final = _TRACEPARENT_RE.match(traceparent.strip()) + if not match: + return None + trace_id: Final = match.group(1).lower() + return trace_id if trace_id != "0" * 32 else None + + +def _session_id_from_baggage(baggage: str) -> str | None: + """Extract a session.id entry from a W3C Baggage header + (https://www.w3.org/TR/baggage/), e.g. "session.id=abc-123,user.id=42".""" + for pair in baggage.split(","): + key, _, value = pair.strip().partition("=") + if key.strip() == "session.id" and value.strip(): + return value.strip() + return None + def _stampable_key_hash(user_api_key_dict: UserAPIKeyAuth) -> str | None: """Only proxy-validated keys are stamped, proven by the unforgeable @@ -1113,6 +1140,33 @@ class LiteLLMProxyRequestSetup: body_metadata["user_id"] = session_id verbose_proxy_logger.debug("Extracted session_id from Anthropic metadata.user_id") + # Last-resort fallback: the W3C standards for trace/session propagation + # (https://www.w3.org/TR/trace-context/, https://www.w3.org/TR/baggage/). + # Lower priority than everything above - only fires when neither the + # explicit litellm headers nor the Anthropic-metadata path found + # anything - but lets a caller's existing traceparent/baggage headers + # (from real OTel instrumentation) correlate with litellm's own logs + # instead of generating an unrelated trace_id. + normalized_headers: Final = MappingProxyType({k.lower(): v for k, v in headers.items() if isinstance(k, str)}) + if "litellm_trace_id" not in data: + traceparent: Final = normalized_headers.get("traceparent") + if isinstance(traceparent, str): + trace_id_from_traceparent: Final = _trace_id_from_traceparent(traceparent) + if trace_id_from_traceparent: + metadata_from_headers["trace_id"] = trace_id_from_traceparent + data["litellm_trace_id"] = trace_id_from_traceparent # rebind-ok: data is an out-param + verbose_proxy_logger.debug( + "Extracted trace_id from W3C traceparent header: %s", trace_id_from_traceparent + ) + if "litellm_session_id" not in data: + baggage: Final = normalized_headers.get("baggage") + if isinstance(baggage, str): + session_id_from_baggage: Final = _session_id_from_baggage(baggage) + if session_id_from_baggage: + metadata_from_headers["session_id"] = session_id_from_baggage + data["litellm_session_id"] = session_id_from_baggage # rebind-ok: data is an out-param + verbose_proxy_logger.debug("Extracted session_id from W3C baggage header") + if isinstance(data[_metadata_variable_name], dict): data[_metadata_variable_name].update(metadata_from_headers) return data diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 18cf9461648..abf9382845b 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3129,6 +3129,7 @@ class StandardAuditLogPayload(TypedDict): class StandardLoggingPayload(TypedDict): id: str trace_id: str # Trace multiple LLM calls belonging to same overall request (e.g. fallbacks/retries) + session_id: str # End-user/conversation session id (litellm_session_id), independent of trace_id litellm_call_id: str | None # UUID returned in x-litellm-call-id response header call_type: str stream: bool | None diff --git a/litellm/utils.py b/litellm/utils.py index 911de83b785..87937c99a0c 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -711,14 +711,71 @@ def _remove_thought_signatures_from_messages(messages: list, thought_signature_s return processed_messages +def _restore_correlation_context_if_supported(logging_obj: object) -> None: + """Call logging_obj._restore_correlation_context() if it's actually there. + + Some call sites (tests, narrow unit paths) inject a minimal stand-in + object as litellm_logging_obj instead of a real Logging instance - this + method is new plumbing specific to request_correlation_in_logs, not part + of any pre-existing stand-in's expected interface. `object` (not `Any`) + is deliberate: the getattr() below is exactly how this stays type-safe + while still tolerating a stand-in that lacks the method. + """ + restore: Final = getattr(logging_obj, "_restore_correlation_context", None) + if restore is not None: + restore() + + +def _is_streaming_response_for_correlation(result: object) -> bool: + """True if `result` is a lazy stream wrapper rather than an already-complete response. + + Only wrapper_async() consults this - it must NOT restore the originating + Task's trace_id/session_id as soon as a streaming call returns this: the + caller is about to iterate it over however many subsequent lines of their + own code, and those log lines should still show this call's ids, not the + pre-call ones. This is safe specifically because each async call already + runs in its own asyncio Task with its own copy of the contextvars, so + leaving it "open" can only affect that one Task, never a different, + unrelated future request - Tasks, unlike a thread pool's worker threads, + are never recycled across requests. The corresponding terminal handler + (async_success_handler, dispatched once the full stream is actually + assembled) is what restores it once streaming genuinely finishes. + + wrapper() (the sync path) does NOT consult this at all: sync calls pass + supports_correlation_logging=False into function_setup()/Logging(), so + they never stamp trace_id/session_id in the first place - a plain OS + thread has no per-call isolation the way an asyncio Task does, and a + thread pool's worker threads *are* recycled across unrelated requests, so + stamping ids there without a safe restore mechanism could permanently + misattribute a later, unrelated request's logs. Full sync support is + deferred to a follow-up PR with its own restore mechanism; see + Logging.__init__'s supports_correlation_logging parameter. + + Genuinely circular otherwise: utils.py -> streaming_handler.py -> + redact_messages.py -> llms/vertex_ai/common_utils.py -> utils.py, which + needs names (supports_response_schema, etc.) this module hasn't finished + defining yet at that point in its own top-to-bottom execution. + """ + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + + return isinstance(result, CustomStreamWrapper) + + +# Runs once per call to check if the user wants to send their data anywhere - PostHog/Sentry/Slack/etc. def function_setup( - original_function: str, rules_obj, start_time, *args, **kwargs -): # just run once to check if user wants to send their data anywhere - PostHog/Sentry/Slack/etc. + original_function: str, + rules_obj: Rules, + start_time: datetime.datetime, + *args: Any, # positional passthrough to the wrapped LLM call (ANN401 ignored, see ruff-strict.toml) + is_async_call: bool = True, + **kwargs: Any, # kwargs-ok: forwarded to Logging()/callbacks, varies per call_type +) -> tuple[LiteLLMLoggingObject, dict[str, Any]]: ### NOTICES ### if litellm.set_verbose is True: verbose_logger.warning( "`litellm.set_verbose` is deprecated. Please set `os.environ['LITELLM_LOG'] = 'DEBUG'` for debug logs." ) + logging_obj: LiteLLMLoggingObject | None = None # rebind-ok: set to the real object further down on success try: global callback_list, add_breadcrumb, user_logger_fn, Logging @@ -1001,7 +1058,8 @@ def function_setup( ): stream = True get_litellm_logging_class: Final = getattr(sys.modules[__name__], "get_litellm_logging_class") - logging_obj: Final = get_litellm_logging_class()( # Victim for object pool + # Victim for object pool + logging_obj = get_litellm_logging_class()( # rebind-ok: 2nd assignment to logging_obj (see initial None above) model=model, messages=messages, stream=stream, @@ -1016,6 +1074,7 @@ def function_setup( dynamic_async_failure_callbacks=dynamic_async_failure_callbacks, kwargs=kwargs, applied_guardrails=applied_guardrails, + supports_correlation_logging=is_async_call, ) ## check if metadata is passed in @@ -1040,6 +1099,15 @@ def function_setup( ) return logging_obj, kwargs except Exception as e: + # If Logging() was constructed above before this failed, its __init__ already + # mutated trace_id_var/session_id_var - restore them *before* logging the + # exception below, since we're about to raise without ever returning + # logging_obj to the caller's wrapper()/wrapper_async() (which would + # otherwise be the one doing this restore). Restoring first means this + # diagnostic log line itself doesn't get stamped with a call's ids when + # that call never actually produced a usable logging object. + if logging_obj is not None: + _restore_correlation_context_if_supported(logging_obj) verbose_logger.exception("litellm.utils.py::function_setup() - [Non-Blocking] Error in function_setup") raise e @@ -1296,7 +1364,9 @@ def client(original_function): try: if logging_obj is None: - logging_obj, kwargs = function_setup(original_function.__name__, rules_obj, start_time, *args, **kwargs) + logging_obj, kwargs = function_setup( + original_function.__name__, rules_obj, start_time, *args, is_async_call=False, **kwargs + ) # Type assertion: logging_obj is guaranteed to be non-None after function_setup assert logging_obj is not None, "logging_obj should not be None after function_setup" @@ -1807,9 +1877,11 @@ def client(original_function): kwargs["retry_strategy"] = "exponential_backoff_retry" elif isinstance(e, openai.APIError): # generic api error kwargs["retry_strategy"] = "constant_retry" - return await litellm.acompletion_with_retries(*args, **kwargs) + result = await litellm.acompletion_with_retries(*args, **kwargs) except Exception: pass + else: + return result elif ( isinstance(e, litellm.exceptions.ContextWindowExceededError) and context_window_fallback_dict @@ -1820,7 +1892,8 @@ def client(original_function): args[0] = context_window_fallback_dict[model] else: kwargs["model"] = context_window_fallback_dict[model] - return await original_function(*args, **kwargs) + result = await original_function(*args, **kwargs) + return result elif call_type == CallTypes.aresponses.value: _is_litellm_router_call = "model_group" in ( kwargs.get("metadata") or {} @@ -1837,9 +1910,11 @@ def client(original_function): kwargs["retry_strategy"] = "exponential_backoff_retry" elif isinstance(e, openai.APIError): # generic api error kwargs["retry_strategy"] = "constant_retry" - return await litellm.aresponses_with_retries(*args, **kwargs) + result = await litellm.aresponses_with_retries(*args, **kwargs) except Exception: pass + else: + return result deployment_num_retries: Final = kwargs.get("num_retries") if deployment_num_retries is not None: @@ -1849,6 +1924,21 @@ def client(original_function): setattr(e, "timeout", timeout) raise e + finally: + # Restore trace_id/session_id contextvars to their pre-call value once + # this call (in this asyncio Task) is fully done - see + # request_correlation_in_logs. Unlike wrapper()'s sync path, it's safe to + # skip restoring when returning a stream: each async call already runs in + # its own Task with its own copy of the contextvars (asyncio.create_task + # copies context at creation), so leaving this Task's own view "open" + # while the caller iterates the stream can only affect that one Task - + # never a different, unrelated future request, since Tasks (unlike a + # thread pool's worker threads) are never recycled across requests. The + # corresponding terminal handler (async_success_handler) restores it once + # streaming genuinely finishes; aclose()/__del__ cover early termination. + if not _is_streaming_response_for_correlation(result): + _restore_correlation_context_if_supported(logging_obj) + get_coroutine_checker: Final = getattr(sys.modules[__name__], "get_coroutine_checker") is_coroutine: Final = get_coroutine_checker().is_async_callable(original_function) diff --git a/ruff-strict.toml b/ruff-strict.toml index 974c49c787b..7afc5da71ee 100644 --- a/ruff-strict.toml +++ b/ruff-strict.toml @@ -16,6 +16,17 @@ external = [ "PLC0415", "E402", "BLE001", "ARG002", "S102", "S324", "S606", "D401", "F403", "F405", ] +[lint.per-file-ignores] +# ANN401 (explicit `Any` disallowed) has no per-line/function-level ignore mechanism +# in ruff, only file-level. These two files each have a handful of parameters that +# are genuinely heterogeneous with no fitting concrete type: a response object that +# varies across every LLM call type (completion/embedding/transcription/etc. each +# return a different shape), and *args/**kwargs forwarded verbatim with no fixed +# shape. Tried the closest existing union (CostResponseTypes) first; basedpyright +# caught a real mismatch, confirming Any is correct here, not a shortcut. +"litellm/litellm_core_utils/litellm_logging.py" = ["ANN401"] +"litellm/utils.py" = ["ANN401"] + [lint.mccabe] max-complexity = 15 diff --git a/tests/logging_callback_tests/test_standard_logging_payload.py b/tests/logging_callback_tests/test_standard_logging_payload.py index f29b245b3be..d13cdf1337a 100644 --- a/tests/logging_callback_tests/test_standard_logging_payload.py +++ b/tests/logging_callback_tests/test_standard_logging_payload.py @@ -471,8 +471,8 @@ def test_get_final_response_obj(): litellm.turn_off_message_logging = False -def test_get_standard_logging_payload_trace_id(): - """Test _get_standard_logging_payload_trace_id with different input scenarios""" +def testget_standard_logging_payload_trace_id(): + """Test get_standard_logging_payload_trace_id with different input scenarios""" # Test case 1: When litellm_trace_id is provided in litellm_params from unittest.mock import MagicMock @@ -482,33 +482,134 @@ def test_get_standard_logging_payload_trace_id(): # Test when litellm_trace_id is in litellm_params litellm_params = {"litellm_trace_id": "dynamic-trace-id"} - result = StandardLoggingPayloadSetup._get_standard_logging_payload_trace_id( + result = StandardLoggingPayloadSetup.get_standard_logging_payload_trace_id( logging_obj=mock_logging_obj, litellm_params=litellm_params ) assert result == "dynamic-trace-id" # Test case 2: When litellm_trace_id is not provided in litellm_params litellm_params = {} - result = StandardLoggingPayloadSetup._get_standard_logging_payload_trace_id( + result = StandardLoggingPayloadSetup.get_standard_logging_payload_trace_id( logging_obj=mock_logging_obj, litellm_params=litellm_params ) assert result == "default-trace-id" # Test case 3: When litellm_params is None - result = StandardLoggingPayloadSetup._get_standard_logging_payload_trace_id( + result = StandardLoggingPayloadSetup.get_standard_logging_payload_trace_id( logging_obj=mock_logging_obj, litellm_params={} ) assert result == "default-trace-id" # Test case 4: When litellm_trace_id in params is not a string litellm_params = {"litellm_trace_id": 12345} - result = StandardLoggingPayloadSetup._get_standard_logging_payload_trace_id( + result = StandardLoggingPayloadSetup.get_standard_logging_payload_trace_id( logging_obj=mock_logging_obj, litellm_params=litellm_params ) assert result == "12345" assert isinstance(result, str) +def testget_standard_logging_payload_trace_id_prioritizes_trace_id_when_flag_on(monkeypatch): + """With request_correlation_in_logs on, an explicit litellm_trace_id wins over litellm_session_id.""" + from unittest.mock import MagicMock + + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + mock_logging_obj = MagicMock() + mock_logging_obj.litellm_trace_id = "default-trace-id" + + litellm_params = {"litellm_trace_id": "the-trace-id", "litellm_session_id": "the-session-id"} + result = StandardLoggingPayloadSetup.get_standard_logging_payload_trace_id( + logging_obj=mock_logging_obj, litellm_params=litellm_params + ) + assert result == "the-trace-id" + + +def testget_standard_logging_payload_trace_id_prioritizes_session_id_when_flag_off(monkeypatch): + """With request_correlation_in_logs off (default), legacy behavior is preserved: + litellm_session_id still wins over litellm_trace_id.""" + from unittest.mock import MagicMock + + monkeypatch.setattr(litellm, "request_correlation_in_logs", False) + mock_logging_obj = MagicMock() + mock_logging_obj.litellm_trace_id = "default-trace-id" + + litellm_params = {"litellm_trace_id": "the-trace-id", "litellm_session_id": "the-session-id"} + result = StandardLoggingPayloadSetup.get_standard_logging_payload_trace_id( + logging_obj=mock_logging_obj, litellm_params=litellm_params + ) + assert result == "the-session-id" + + +def testget_standard_logging_payload_session_id_when_flag_on(monkeypatch): + """Test get_standard_logging_payload_session_id with different input scenarios, flag enabled""" + from unittest.mock import MagicMock + + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + mock_logging_obj = MagicMock() + mock_logging_obj.litellm_session_id = "" + + # Test case 1: litellm_session_id provided directly in litellm_params + litellm_params = {"litellm_session_id": "dynamic-session-id"} + result = StandardLoggingPayloadSetup.get_standard_logging_payload_session_id( + logging_obj=mock_logging_obj, litellm_params=litellm_params + ) + assert result == "dynamic-session-id" + + # Test case 2: falls back to metadata.session_id when not in litellm_params directly + litellm_params = {"metadata": {"session_id": "metadata-session-id"}} + result = StandardLoggingPayloadSetup.get_standard_logging_payload_session_id( + logging_obj=mock_logging_obj, litellm_params=litellm_params + ) + assert result == "metadata-session-id" + + # Test case 3: falls back to logging_obj.litellm_session_id when nothing else is set + mock_logging_obj.litellm_session_id = "obj-session-id" + result = StandardLoggingPayloadSetup.get_standard_logging_payload_session_id( + logging_obj=mock_logging_obj, litellm_params={} + ) + assert result == "obj-session-id" + + # Test case 4: empty string when no session id was supplied anywhere + mock_logging_obj.litellm_session_id = "" + result = StandardLoggingPayloadSetup.get_standard_logging_payload_session_id( + logging_obj=mock_logging_obj, litellm_params={} + ) + assert result == "" + + # Test case 5: non-string session id in params is coerced to str + litellm_params = {"litellm_session_id": 98765} + result = StandardLoggingPayloadSetup.get_standard_logging_payload_session_id( + logging_obj=mock_logging_obj, litellm_params=litellm_params + ) + assert result == "98765" + assert isinstance(result, str) + + # Test case 6: trace_id and session_id are independent - passing only a trace id + # must not populate session_id + litellm_params = {"litellm_trace_id": "some-trace-id"} + result = StandardLoggingPayloadSetup.get_standard_logging_payload_session_id( + logging_obj=mock_logging_obj, litellm_params=litellm_params + ) + assert result == "" + + +def testget_standard_logging_payload_session_id_empty_when_flag_off(monkeypatch): + """When request_correlation_in_logs is off (default), session_id is always empty, + even if litellm_session_id was explicitly supplied - preserves the pre-existing + StandardLoggingPayload shape for callers who haven't opted in.""" + from unittest.mock import MagicMock + + monkeypatch.setattr(litellm, "request_correlation_in_logs", False) + mock_logging_obj = MagicMock() + mock_logging_obj.litellm_session_id = "obj-session-id" + + litellm_params = {"litellm_session_id": "dynamic-session-id"} + result = StandardLoggingPayloadSetup.get_standard_logging_payload_session_id( + logging_obj=mock_logging_obj, litellm_params=litellm_params + ) + assert result == "" + + def test_truncate_standard_logging_payload(): """ 1. original messages, response, and error_str should NOT BE MODIFIED, since these are from kwargs 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 23e0975cd08..9fa3116657b 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -15,6 +15,7 @@ import httpx from openai._legacy_response import HttpxBinaryResponseContent import litellm +from litellm._logging import session_id_var, trace_id_var from litellm.constants import SENTRY_DENYLIST, SENTRY_PII_DENYLIST from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging @@ -3312,6 +3313,51 @@ def test_failure_handler_runs_sync_callbacks_for_non_pass_through_requests( dummy_logger.log_failure_event.assert_called_once() +@pytest.mark.asyncio +async def test_async_failure_handler_runs_callbacks_and_restores_correlation_context(logging_obj): + """await logging_obj.async_failure_handler(...) must dispatch async failure callbacks + and, once its own body completes, restore trace_id/session_id contextvars via + _restore_correlation_context() (the fix for the nested-call context leak).""" + from litellm._logging import session_id_var, trace_id_var + from litellm.integrations.custom_logger import CustomLogger + + class DummyLogger(CustomLogger): + pass + + logging_obj.call_type = "acompletion" + logging_obj.stream = False + logging_obj.model_call_details["litellm_params"] = {} + logging_obj.litellm_params = {} + + dummy_logger = DummyLogger() + dummy_logger.async_log_failure_event = AsyncMock() + + # logging_obj is constructed by the fixture (before this line runs), so it + # already captured whatever was ambient at that point as its own pre-call + # value - assert restoration lands back on THAT captured value, not a + # value set here (which would be too late to affect __init__'s snapshot). + trace_id_var.set("mutated-during-call") + session_id_var.set("mutated-during-call") + try: + with patch.object( + logging_obj, + "get_combined_callback_list", + return_value=[dummy_logger], + ): + await logging_obj.async_failure_handler( + exception=Exception("test error"), + traceback_exception="", + ) + + dummy_logger.async_log_failure_event.assert_called_once() + assert trace_id_var.get() == logging_obj._pre_call_trace_id + assert session_id_var.get() == logging_obj._pre_call_session_id + assert trace_id_var.get() != "mutated-during-call" + finally: + trace_id_var.set("") + session_id_var.set("") + + def test_merge_hidden_params_from_response_into_metadata_populates_metadata(): """Streaming completion path should mirror non-stream: metadata.hidden_params from response.""" from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -4230,3 +4276,199 @@ def test_pre_call_does_not_pin_request_in_module_state(logging_obj): logging_obj.post_call(original_response='{"ok": true}', input=big_input, api_key="sk-test") assert litellm.error_logs == {} + + +def test_logging_init_sets_trace_id(): + """Logging.__init__() must call set_trace_id with self.litellm_trace_id.""" + from litellm.litellm_core_utils.litellm_logging import Logging + + trace_id_var.set("") + + log_obj = Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="completion", + start_time=None, + litellm_call_id="call-001", + function_id="fn-001", + kwargs={}, + ) + assert trace_id_var.get() == log_obj.litellm_trace_id + + +def test_logging_init_skips_stamping_when_correlation_logging_unsupported(): + """supports_correlation_logging=False (what wrapper(), the sync entry + point, always passes) must leave trace_id_var/session_id_var completely + untouched, even though self.litellm_trace_id/litellm_session_id (the + plain attributes used by StandardLoggingPayload) are still populated as + usual - only the ambient contextvar stamping is gated.""" + from litellm.litellm_core_utils.litellm_logging import Logging + + trace_id_var.set("") + session_id_var.set("") + + log_obj = Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="completion", + start_time=None, + litellm_call_id="call-sync-excluded", + function_id="fn-sync-excluded", + kwargs={"litellm_session_id": "should-not-be-stamped"}, + litellm_trace_id="should-not-be-stamped-either", + supports_correlation_logging=False, + ) + + assert trace_id_var.get() == "" + assert session_id_var.get() == "" + # The plain attributes are unaffected - only the contextvar stamping is gated. + assert log_obj.litellm_trace_id == "should-not-be-stamped-either" + assert log_obj.litellm_session_id == "should-not-be-stamped" + + +def test_logging_init_sets_session_id_when_provided(): + """Logging.__init__() must call set_session_id when litellm_session_id is in kwargs.""" + from litellm.litellm_core_utils.litellm_logging import Logging + + session_id_var.set("") + + Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="completion", + start_time=None, + litellm_call_id="call-002", + function_id="fn-002", + kwargs={"litellm_session_id": "my-session-99"}, + ) + assert session_id_var.get() == "my-session-99" + + +def test_logging_init_resets_session_id_to_empty_when_absent(): + """When no session_id is in kwargs, Logging.__init__() must reset session_id_var to "" + so a prior request's session_id does not leak into subsequent log records.""" + from litellm.litellm_core_utils.litellm_logging import Logging + + session_id_var.set("preexisting-sid") + + Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="completion", + start_time=None, + litellm_call_id="call-003", + function_id="fn-003", + kwargs={}, + ) + assert session_id_var.get() == "" + + +def test_restore_correlation_context_resets_to_pre_call_value(): + """_restore_correlation_context() must put trace_id_var/session_id_var back to + whatever they were immediately before this Logging instance was constructed. + This is the mechanism that prevents a nested call (e.g. a guardrail's own + LLM-as-judge call sharing the same asyncio Task) from leaking its trace_id/ + session_id into the outer call's subsequent log lines.""" + from litellm.litellm_core_utils.litellm_logging import Logging + + trace_id_var.set("outer-trace") + session_id_var.set("outer-session") + try: + inner = Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="completion", + start_time=None, + litellm_call_id="inner-call", + function_id="fn-inner", + kwargs={"litellm_session_id": "inner-session"}, + ) + assert trace_id_var.get() == inner.litellm_trace_id + assert session_id_var.get() == "inner-session" + + inner._restore_correlation_context() + + assert trace_id_var.get() == "outer-trace" + assert session_id_var.get() == "outer-session" + finally: + trace_id_var.set("") + session_id_var.set("") + + +def test_restore_correlation_context_safe_to_call_repeatedly(): + """Calling _restore_correlation_context() more than once must not raise. + + It's deliberately NOT guarded against repeat calls: wrapper()'s finally + block and a terminal handler (success_handler/failure_handler) can both + end up calling it for the same instance, potentially from different + asyncio Tasks - each call needs to take effect in its own Task's view of + the contextvars, so repeat calls are expected, not just tolerated.""" + from litellm.litellm_core_utils.litellm_logging import Logging + + log_obj = Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="completion", + start_time=None, + litellm_call_id="call-idempotent", + function_id="fn-idempotent", + kwargs={}, + ) + log_obj._restore_correlation_context() + log_obj._restore_correlation_context() # must not raise + + +@pytest.mark.asyncio +async def test_restore_correlation_context_works_across_asyncio_task_boundary(): + """_restore_correlation_context() must succeed even when it's called from a + different asyncio Task than the one Logging.__init__() ran in - exactly what + happens on litellm's real async success path, where async_success_handler is + dispatched via asyncio.create_task / the global logging worker rather than + awaited directly in the request's own task. + + A contextvars.Token can only be reset in the exact Context it was created in + and raises ValueError otherwise (verified separately against raw contextvars, + not just this codebase). The fix uses a plain set() of the captured pre-call + value instead, which works regardless of which Task calls it. This test + fails with a token-based implementation - the child task's reset() would + raise, get silently swallowed, and leave the child's view unrestored - and + passes with the value-based one. + """ + from litellm.litellm_core_utils.litellm_logging import Logging + + trace_id_var.set("outer-trace-cross-task") + session_id_var.set("outer-session-cross-task") + try: + # __init__ runs in THIS (outer) task's context. + log_obj = Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="acompletion", + start_time=None, + litellm_call_id="cross-task-call", + function_id="fn-cross-task", + kwargs={"litellm_session_id": "cross-task-session"}, + ) + assert trace_id_var.get() == log_obj.litellm_trace_id + assert session_id_var.get() == "cross-task-session" + + async def restore_in_new_task(): + # Simulates async_success_handler running in a task spawned after + # __init__ already ran elsewhere - a different Context object. + log_obj._restore_correlation_context() + return trace_id_var.get(), session_id_var.get() + + trace_in_child, session_in_child = await asyncio.create_task(restore_in_new_task()) + + assert trace_in_child == "outer-trace-cross-task" + assert session_in_child == "outer-session-cross-task" + finally: + trace_id_var.set("") + session_id_var.set("") diff --git a/tests/test_litellm/litellm_core_utils/test_max_streaming_duration.py b/tests/test_litellm/litellm_core_utils/test_max_streaming_duration.py index a417ad90eb7..1eb49f4859f 100644 --- a/tests/test_litellm/litellm_core_utils/test_max_streaming_duration.py +++ b/tests/test_litellm/litellm_core_utils/test_max_streaming_duration.py @@ -9,7 +9,7 @@ Covers: import os import sys import time -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -69,8 +69,12 @@ class TestCustomStreamWrapperMaxDuration: @pytest.mark.asyncio async def test_should_raise_on_async_anext_when_exceeded(self): - """__anext__ should check the limit before iterating.""" + """__anext__ should check the limit before iterating, dispatching the + same failure-callback/logging path every other stream failure goes + through (dispatch_failure_handlers is async on the real Logging class, + so the mock needs to be awaitable too).""" wrapper = _make_custom_stream_wrapper() + wrapper.logging_obj.dispatch_failure_handlers = AsyncMock() wrapper._stream_created_time = time.time() - 20 with patch("litellm.constants.LITELLM_MAX_STREAMING_DURATION_SECONDS", 10.0): with pytest.raises(litellm.Timeout): diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 5806b37539c..101935cac0a 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -14,6 +14,8 @@ import traceback from typing import Optional import litellm +from litellm import verbose_logger +from litellm._logging import session_id_var, trace_id_var from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.streaming_handler import ( AUDIO_ATTRIBUTE, @@ -3551,3 +3553,613 @@ def test_openai_custom_tool_call_stream_deltas_survive_conversion(logging_obj: L assert combined_input == "*** Begin Patch\n*** End Patch\n" finish_reasons = [chunk.choices[0].finish_reason for chunk in emitted if chunk.choices] assert "tool_calls" in finish_reasons + + +def test_sync_completion_never_stamps_correlation_context(monkeypatch): + """wrapper() (the sync entry point) does not participate in + request_correlation_in_logs at all: Logging.__init__() is called with + supports_correlation_logging=False for every sync call, so + trace_id_var/session_id_var are never touched, regardless of whether the + caller passes litellm_trace_id/litellm_session_id or the call streams. + + This is a deliberate scoping decision, not an oversight: a plain OS + thread has no per-call isolation the way an asyncio Task does, and a + thread pool's worker threads are recycled across unrelated requests, so + safely supporting this for the sync path needs its own restore mechanism + with its own tests - tracked as a separate, follow-up piece of work. + Async (acompletion/wrapper_async, the only path the proxy uses) is + unaffected - see test_async_streaming_completion_does_not_reset_context_before_iteration.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + # Reset explicitly rather than asserting a clean slate - this must hold + # regardless of what any other test left behind in these module-level + # contextvars. + trace_id_var.set("") + session_id_var.set("") + try: + litellm.completion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + mock_response="Hello there!", + litellm_trace_id="should-never-appear", + litellm_session_id="should-never-appear-either", + num_retries=0, + ) + assert trace_id_var.get() == "" + assert session_id_var.get() == "" + + response = litellm.completion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + mock_response="Hello there!", + stream=True, + litellm_trace_id="should-never-appear-stream", + litellm_session_id="should-never-appear-stream-either", + num_retries=0, + ) + for _ in response: + pass + assert trace_id_var.get() == "" + assert session_id_var.get() == "" + finally: + trace_id_var.set("") + session_id_var.set("") + + +def test_abandoned_sync_stream_cannot_contaminate_a_later_call_on_the_same_thread(monkeypatch): + """The maintainer-reported blocking bug reproduced live in this session - + request A starts a sync stream, consumes one chunk, abandons it; request + B runs next on the same forced-reuse ThreadPoolExecutor worker - is now + structurally impossible rather than merely restored-after-the-fact: since + sync calls never stamp trace_id_var/session_id_var at all + (supports_correlation_logging=False), there is nothing for request A to + leave behind for request B to inherit.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + + from concurrent.futures import ThreadPoolExecutor + + pool = ThreadPoolExecutor(max_workers=1) + try: + + def call_a_abandon_stream(): + response = litellm.completion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "call A"}], + mock_response="call A response", + stream=True, + litellm_session_id="SESSION-AAA", + litellm_trace_id="TRACE-AAA", + num_retries=0, + ) + next(response) # consume exactly one chunk, then abandon it + + def call_b_non_streaming(): + litellm.completion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "call B"}], + mock_response="call B response", + litellm_session_id="SESSION-BBB", + litellm_trace_id="TRACE-BBB", + num_retries=0, + ) + return trace_id_var.get(), session_id_var.get() + + pool.submit(call_a_abandon_stream).result() + ids_after_b = pool.submit(call_b_non_streaming).result() + + assert ids_after_b == ("", "") + finally: + pool.shutdown(wait=True) + + +@pytest.mark.asyncio +async def test_async_streaming_completion_does_not_reset_context_before_iteration(monkeypatch): + """Same as above for wrapper_async()/acompletion().""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + trace_id_var.set("outer-trace-async-stream") + session_id_var.set("outer-session-async-stream") + try: + response = await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + mock_response="Hello there!", + stream=True, + litellm_session_id="async-streaming-call-session", + num_retries=0, + ) + assert session_id_var.get() == "async-streaming-call-session" + + async for _ in response: + pass + + # Once the stream is genuinely exhausted, the *consuming* task's own + # context must be restored - async_success_handler's own dispatch (via + # asyncio.create_task) only fixes up its own detached task, not this one. + assert session_id_var.get() == "outer-session-async-stream" + assert trace_id_var.get() == "outer-trace-async-stream" + finally: + trace_id_var.set("") + session_id_var.set("") + + +def test_stream_wrapper_del_restores_correlation_context(): + """CustomStreamWrapper.__del__ is the best-effort fallback for an abandoned + stream (caller never exhausts it, so the normal terminal-handler restore + never fires). Testing this via real garbage collection is unreliable in + practice - CPython's per-chunk logging submits work to a thread pool + executor whose worker thread transiently holds its own reference to the + wrapper (a bound method argument) until that task completes, so refcount + doesn't reliably hit zero on a deterministic schedule even with polling. + Call __del__ directly instead: it's a plain method, calling it early + doesn't run actual finalization, and this exercises exactly the logic that + real garbage collection would eventually trigger. + """ + trace_id_var.set("outer-trace-abandoned") + session_id_var.set("outer-session-abandoned") + try: + log_obj = Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=None, + litellm_call_id="abandoned-stream-call", + function_id="fn-abandoned-stream", + kwargs={"litellm_session_id": "abandoned-stream-session"}, + ) + wrapper = CustomStreamWrapper( + completion_stream=iter([]), + model="gpt-3.5-turbo", + logging_obj=log_obj, + ) + wrapper.__del__() + + assert trace_id_var.get() == "outer-trace-abandoned" + assert session_id_var.get() == "outer-session-abandoned" + finally: + trace_id_var.set("") + session_id_var.set("") + + +def test_stream_wrapper_del_never_raises_with_broken_logging_obj(): + """__del__ runs during garbage collection, possibly at interpreter + shutdown - it must never raise regardless of what's wrong with logging_obj, + or Python prints an ignored "exception in __del__" warning and, worse, + could mask the real error a caller is in the middle of handling.""" + + class ExplodingLogging: + model_call_details: dict = {} + + def _restore_correlation_context(self): + raise RuntimeError("logging_obj is in a bad state") + + wrapper = CustomStreamWrapper( + completion_stream=iter([]), + model="gpt-3.5-turbo", + logging_obj=ExplodingLogging(), + ) + wrapper.__del__() # must not raise + + +def test_stream_wrapper_del_does_not_clobber_a_newer_active_call(): + """A delayed finalizer must never stomp a different, still-active call's + context. If an abandoned stream's __del__ fires late - after a new call + has already started in the same Task/thread and claimed the contextvars - + unconditionally restoring the abandoned stream's own pre-call snapshot + would corrupt the active call's subsequent log lines with stale ids.""" + trace_id_var.set("outer-trace-before-abandoned-call") + session_id_var.set("outer-session-before-abandoned-call") + try: + abandoned_log_obj = Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=None, + litellm_call_id="abandoned-stream-call", + function_id="fn-abandoned-stream", + kwargs={"litellm_session_id": "abandoned-stream-session"}, + ) + wrapper = CustomStreamWrapper( + completion_stream=iter([]), + model="gpt-3.5-turbo", + logging_obj=abandoned_log_obj, + ) + + # A new, unrelated call starts in this same Task/thread before the + # abandoned stream's __del__ ever fires, and claims the contextvars. + Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="completion", + start_time=None, + litellm_call_id="newer-active-call", + function_id="fn-newer-active-call", + kwargs={"litellm_session_id": "newer-active-session"}, + ) + assert trace_id_var.get() != abandoned_log_obj.litellm_trace_id + assert session_id_var.get() == "newer-active-session" + + # The delayed finalizer for the abandoned stream must not clobber + # the newer call's still-active ids. + wrapper.__del__() + + assert trace_id_var.get() != abandoned_log_obj.litellm_trace_id + assert session_id_var.get() == "newer-active-session" + finally: + trace_id_var.set("") + session_id_var.set("") + + +def test_stream_wrapper_del_restores_when_own_session_id_needed_sanitizing(): + """The __del__ guard must compare against the *sanitized* id actually + stored in the contextvar, not the raw litellm_session_id/litellm_trace_id + - set_session_id()/set_trace_id() strip control characters before + storing, so a caller-supplied id containing e.g. a newline would never + equal the raw attribute, and the guard would wrongly conclude some other + call has claimed the context and skip cleanup forever.""" + trace_id_var.set("outer-trace-needs-sanitizing") + session_id_var.set("outer-session-needs-sanitizing") + try: + raw_session_id = "abandoned\nsession\rwith-control-chars" + log_obj = Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=None, + litellm_call_id="abandoned-stream-needs-sanitizing", + function_id="fn-abandoned-stream-needs-sanitizing", + kwargs={"litellm_session_id": raw_session_id}, + ) + # Sanity: the contextvar holds the sanitized value, which differs + # from the raw litellm_session_id this test constructed it with. + assert session_id_var.get() != raw_session_id + assert log_obj.litellm_session_id == raw_session_id + + wrapper = CustomStreamWrapper( + completion_stream=iter([]), + model="gpt-3.5-turbo", + logging_obj=log_obj, + ) + wrapper.__del__() + + assert trace_id_var.get() == "outer-trace-needs-sanitizing" + assert session_id_var.get() == "outer-session-needs-sanitizing" + finally: + trace_id_var.set("") + session_id_var.set("") + + +def test_stream_wrapper_next_keeps_context_active_through_synthesized_finish_reason_chunk(): + """When the underlying stream ends without ever emitting an explicit + finish_reason chunk, __next__ synthesizes one via finish_reason_handler() + and returns it. That chunk is still this call's own data - the caller's + own (application-level) log statements processing it run immediately + after this return, in the same synchronous frame, so context must NOT be + restored yet or those log lines would carry the wrong ids. A caller that + keeps iterating (the common, non-early-break pattern) still gets a + correct, deterministic restore on the very next __next__() call, since + completion_stream is already exhausted and immediately re-raises + StopIteration.""" + trace_id_var.set("outer-trace-finish-reason") + session_id_var.set("outer-session-finish-reason") + try: + log_obj = Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=None, + litellm_call_id="finish-reason-call", + function_id="fn-finish-reason", + kwargs={"litellm_session_id": "finish-reason-session"}, + ) + wrapper = CustomStreamWrapper( + completion_stream=iter([]), + model="gpt-3.5-turbo", + logging_obj=log_obj, + ) + assert trace_id_var.get() == log_obj.litellm_trace_id + assert session_id_var.get() == "finish-reason-session" + + chunk = next(wrapper) + + assert chunk.choices[0].finish_reason is not None + # Still this call's own ids - not restored yet. + assert trace_id_var.get() == log_obj.litellm_trace_id + assert session_id_var.get() == "finish-reason-session" + + # A caller that keeps iterating (doesn't break early) still gets a + # deterministic restore right here, on the next real StopIteration. + with pytest.raises(StopIteration): + next(wrapper) + assert trace_id_var.get() == "outer-trace-finish-reason" + assert session_id_var.get() == "outer-session-finish-reason" + finally: + trace_id_var.set("") + session_id_var.set("") + + +def test_stream_wrapper_del_cleans_up_after_synthesized_finish_reason_chunk(): + """A caller that breaks immediately after seeing finish_reason (the + early-break pattern) never triggers the next()-driven restore above - it + relies on the best-effort __del__ guard instead, same as any other + abandoned stream. The guard must still recognize this call's own + (unrestored) ids as unclaimed and clean them up.""" + trace_id_var.set("outer-trace-finish-reason-del") + session_id_var.set("outer-session-finish-reason-del") + try: + log_obj = Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=None, + litellm_call_id="finish-reason-del-call", + function_id="fn-finish-reason-del", + kwargs={"litellm_session_id": "finish-reason-del-session"}, + ) + wrapper = CustomStreamWrapper( + completion_stream=iter([]), + model="gpt-3.5-turbo", + logging_obj=log_obj, + ) + + chunk = next(wrapper) + assert chunk.choices[0].finish_reason is not None + + wrapper.__del__() + + assert trace_id_var.get() == "outer-trace-finish-reason-del" + assert session_id_var.get() == "outer-session-finish-reason-del" + finally: + trace_id_var.set("") + session_id_var.set("") + + +@pytest.mark.asyncio +async def test_stream_wrapper_anext_keeps_context_active_through_synthesized_finish_reason_chunk(): + """Async sibling of test_stream_wrapper_next_keeps_context_active_through_synthesized_finish_reason_chunk - + _finalize_completed_stream()'s else branch must not restore before + returning the synthesized chunk either.""" + trace_id_var.set("outer-trace-anext-finish-reason") + session_id_var.set("outer-session-anext-finish-reason") + try: + log_obj = Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=None, + litellm_call_id="anext-finish-reason-call", + function_id="fn-anext-finish-reason", + kwargs={"litellm_session_id": "anext-finish-reason-session"}, + ) + + async def _empty_aiter(): + return + yield # pragma: no cover - makes this an async generator + + wrapper = CustomStreamWrapper( + completion_stream=_empty_aiter(), + model="gpt-3.5-turbo", + logging_obj=log_obj, + ) + assert trace_id_var.get() == log_obj.litellm_trace_id + assert session_id_var.get() == "anext-finish-reason-session" + + chunk = await wrapper.__anext__() + + assert chunk.choices[0].finish_reason is not None + # Still this call's own ids - not restored yet. + assert trace_id_var.get() == log_obj.litellm_trace_id + assert session_id_var.get() == "anext-finish-reason-session" + + # A caller that keeps iterating still gets a deterministic restore + # right here, on the next real StopAsyncIteration. + with pytest.raises(StopAsyncIteration): + await wrapper.__anext__() + assert trace_id_var.get() == "outer-trace-anext-finish-reason" + assert session_id_var.get() == "outer-session-anext-finish-reason" + finally: + trace_id_var.set("") + session_id_var.set("") + + +@pytest.mark.asyncio +async def test_stream_wrapper_anext_max_duration_timeout_restores_consumer_correlation_context(monkeypatch): + """_check_max_streaming_duration() raises litellm.Timeout when a client keeps + an async stream open past LITELLM_MAX_STREAMING_DURATION_SECONDS. That raise + must flow through the same except Exception -> _handle_stream_fallback_error + path as every other failure so the consumer's outer correlation context gets + restored - calling the check before entering __anext__()'s try block would + let the Timeout bypass that restoration entirely.""" + monkeypatch.setattr(litellm.constants, "LITELLM_MAX_STREAMING_DURATION_SECONDS", 1) + trace_id_var.set("outer-trace-max-duration") + session_id_var.set("outer-session-max-duration") + try: + log_obj = Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=None, + litellm_call_id="max-duration-call", + function_id="fn-max-duration", + kwargs={"litellm_session_id": "max-duration-session"}, + ) + + async def _empty_aiter(): + return + yield # pragma: no cover - makes this an async generator + + wrapper = CustomStreamWrapper( + completion_stream=_empty_aiter(), + model="gpt-3.5-turbo", + logging_obj=log_obj, + ) + assert trace_id_var.get() == log_obj.litellm_trace_id + assert session_id_var.get() == "max-duration-session" + + wrapper._stream_created_time = time.time() - 10 + + with pytest.raises(Exception): + await wrapper.__anext__() + + assert trace_id_var.get() == "outer-trace-max-duration" + assert session_id_var.get() == "outer-session-max-duration" + finally: + trace_id_var.set("") + session_id_var.set("") + + +@pytest.mark.asyncio +async def test_stream_wrapper_aclose_restores_consumer_correlation_context(): + """Explicit early termination (aclose(), e.g. on client disconnect or a + router fallback aborting an in-progress stream) must restore the caller's + correlation context too - not just __del__'s best-effort GC-timed fallback, + since aclose() is normally called deterministically by the consumer/ + framework, unlike __del__.""" + trace_id_var.set("outer-trace-aclose") + session_id_var.set("outer-session-aclose") + try: + log_obj = Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=None, + litellm_call_id="aclose-call", + function_id="fn-aclose", + kwargs={"litellm_session_id": "aclose-session"}, + ) + + async def _empty_aiter(): + return + yield # pragma: no cover - makes this an async generator + + wrapper = CustomStreamWrapper( + completion_stream=_empty_aiter(), + model="gpt-3.5-turbo", + logging_obj=log_obj, + ) + assert trace_id_var.get() == log_obj.litellm_trace_id + assert session_id_var.get() == "aclose-session" + + await wrapper.aclose() + + assert trace_id_var.get() == "outer-trace-aclose" + assert session_id_var.get() == "outer-session-aclose" + finally: + trace_id_var.set("") + session_id_var.set("") + + +@pytest.mark.asyncio +async def test_stream_wrapper_aclose_keeps_context_active_through_close_failure_diagnostic(monkeypatch): + """If closing the underlying provider stream raises, aclose()'s except + branch logs a debug diagnostic. That log line must still carry the + closing stream's own trace_id/session_id - the outer context must not be + restored until after the close attempt (and its diagnostic) completes.""" + trace_id_var.set("outer-trace-close-fail") + session_id_var.set("outer-session-close-fail") + try: + log_obj = Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=None, + litellm_call_id="close-fail-call", + function_id="fn-close-fail", + kwargs={"litellm_session_id": "close-fail-session"}, + ) + + class _RaisingAsyncCloseStream: + async def aclose(self): + raise RuntimeError("boom closing stream") + + def __aiter__(self): + return self + + async def __anext__(self): + raise StopAsyncIteration + + wrapper = CustomStreamWrapper( + completion_stream=_RaisingAsyncCloseStream(), + model="gpt-3.5-turbo", + logging_obj=log_obj, + ) + assert trace_id_var.get() == log_obj.litellm_trace_id + assert session_id_var.get() == "close-fail-session" + + captured_ids = {} + real_debug = verbose_logger.debug + + def fake_debug(msg, *args, **kwargs): + if "error closing completion_stream" in msg: + captured_ids["trace_id"] = trace_id_var.get() + captured_ids["session_id"] = session_id_var.get() + return real_debug(msg, *args, **kwargs) + + monkeypatch.setattr(verbose_logger, "debug", fake_debug) + + await wrapper.aclose() + + assert captured_ids["trace_id"] == log_obj.litellm_trace_id + assert captured_ids["session_id"] == "close-fail-session" + assert trace_id_var.get() == "outer-trace-close-fail" + assert session_id_var.get() == "outer-session-close-fail" + finally: + trace_id_var.set("") + session_id_var.set("") + + +def test_handle_stream_fallback_error_restores_context_only_after_exception_mapping(monkeypatch): + """_map_anthropic_exception/_map_aleph_alpha_exception synchronously log a + debug diagnostic (the raw status code) as part of exception_type()'s + mapping. The consumer's outer context must not be restored until that + mapping call returns, or the diagnostic log line would carry the outer + (or empty) trace_id/session_id instead of the failing stream's own.""" + trace_id_var.set("outer-trace-fallback") + session_id_var.set("outer-session-fallback") + try: + log_obj = Logging( + model="claude-3-opus", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=None, + litellm_call_id="fallback-error-call", + function_id="fn-fallback-error", + kwargs={"litellm_session_id": "fallback-error-session"}, + ) + wrapper = CustomStreamWrapper( + completion_stream=iter([]), + model="claude-3-opus", + custom_llm_provider="anthropic", + logging_obj=log_obj, + ) + + captured_ids = {} + + def fake_exception_type(**kwargs): + captured_ids["trace_id"] = trace_id_var.get() + captured_ids["session_id"] = session_id_var.get() + return ValueError("mapped boom") + + monkeypatch.setattr("litellm.litellm_core_utils.streaming_handler.exception_type", fake_exception_type) + + with pytest.raises(Exception): + wrapper._handle_stream_fallback_error(RuntimeError("boom")) + + # The mapper ran while the stream's own ids were still active. + assert captured_ids["trace_id"] == log_obj.litellm_trace_id + assert captured_ids["session_id"] == "fallback-error-session" + # Restored to the consumer's outer context once mapping/raise completes. + assert trace_id_var.get() == "outer-trace-fallback" + assert session_id_var.get() == "outer-session-fallback" + finally: + trace_id_var.set("") + session_id_var.set("") diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 22f9e6bb67a..2d37d8f8351 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -2713,6 +2713,149 @@ def test_get_chain_id_from_headers_generic_vendor_session_id(): ) +def test_trace_id_from_traceparent_valid(): + from litellm.proxy.litellm_pre_call_utils import _trace_id_from_traceparent + + assert ( + _trace_id_from_traceparent("00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01") + == "4bf92f3577b34da6a3ce929d0e0e4736" + ) + # Case-insensitive, normalized to lowercase + assert ( + _trace_id_from_traceparent("00-4BF92F3577B34DA6A3CE929D0E0E4736-00f067aa0ba902b7-01") + == "4bf92f3577b34da6a3ce929d0e0e4736" + ) + + +@pytest.mark.parametrize( + "traceparent", + [ + "not-a-traceparent", + "00-tooshort-00f067aa0ba902b7-01", + "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7", # missing flags segment + "00-4bf92f3577g34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", # non-hex char + "00-00000000000000000000000000000000-00f067aa0ba902b7-01", # all-zero trace-id, invalid per spec + "", + ], +) +def test_trace_id_from_traceparent_rejects_malformed(traceparent: str): + from litellm.proxy.litellm_pre_call_utils import _trace_id_from_traceparent + + assert _trace_id_from_traceparent(traceparent) is None + + +def test_session_id_from_baggage_valid(): + from litellm.proxy.litellm_pre_call_utils import _session_id_from_baggage + + assert _session_id_from_baggage("session.id=abc-123,user.id=42") == "abc-123" + assert _session_id_from_baggage("user.id=42, session.id=xyz-789") == "xyz-789" + + +@pytest.mark.parametrize( + "baggage", + [ + "user.id=42", + "", + "session.id=", + ], +) +def test_session_id_from_baggage_absent_or_empty(baggage: str): + from litellm.proxy.litellm_pre_call_utils import _session_id_from_baggage + + assert _session_id_from_baggage(baggage) is None + + +def test_add_litellm_metadata_from_request_headers_traceparent_sets_trace_id_only(): + """A bare traceparent header (no litellm-specific headers) sets litellm_trace_id + from its trace-id component and leaves litellm_session_id unset.""" + headers = {"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"} + data = {"metadata": {}} + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers=headers, data=data, _metadata_variable_name="metadata" + ) + assert data["litellm_trace_id"] == "4bf92f3577b34da6a3ce929d0e0e4736" + assert data["metadata"]["trace_id"] == "4bf92f3577b34da6a3ce929d0e0e4736" + assert "litellm_session_id" not in data + + +def test_add_litellm_metadata_from_request_headers_baggage_sets_session_id_only(): + """A bare baggage header (no litellm-specific headers) sets litellm_session_id + from its session.id entry and leaves litellm_trace_id unset.""" + headers = {"baggage": "session.id=baggage-session-42,user.id=7"} + data = {"metadata": {}} + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers=headers, data=data, _metadata_variable_name="metadata" + ) + assert data["litellm_session_id"] == "baggage-session-42" + assert data["metadata"]["session_id"] == "baggage-session-42" + assert "litellm_trace_id" not in data + + +def test_add_litellm_metadata_from_request_headers_baggage_session_id_not_logged_raw(caplog): + """The raw baggage session.id value must never reach the debug log line - + it isn't sanitized until set_session_id() runs much later in + Logging.__init__(), so logging it here would let a caller with control + characters or terminal escape sequences forge plaintext log output.""" + import logging + + poisoned = "poisoned\x1b[31mFAKE_RED_TEXT\x1b[0m" + headers = {"baggage": f"session.id={poisoned}"} + data = {"metadata": {}} + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers=headers, data=data, _metadata_variable_name="metadata" + ) + assert data["litellm_session_id"] == poisoned + assert not any(poisoned in record.getMessage() for record in caplog.records) + + +def test_add_litellm_metadata_from_request_headers_traceparent_and_baggage_together(): + """traceparent and baggage are resolved independently - trace_id and + session_id do not have to be the same value, unlike the chain_id path.""" + headers = { + "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", + "baggage": "session.id=baggage-session-42", + } + data = {"metadata": {}} + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers=headers, data=data, _metadata_variable_name="metadata" + ) + assert data["litellm_trace_id"] == "4bf92f3577b34da6a3ce929d0e0e4736" + assert data["litellm_session_id"] == "baggage-session-42" + + +def test_add_litellm_metadata_from_request_headers_explicit_trace_id_beats_traceparent(): + """x-litellm-trace-id must win over a traceparent header carrying a + different trace-id - explicit litellm headers are always highest priority.""" + headers = { + "x-litellm-trace-id": "explicit-trace-id-value", + "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", + } + data = {"metadata": {}} + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers=headers, data=data, _metadata_variable_name="metadata" + ) + assert data["litellm_trace_id"] == "explicit-trace-id-value" + assert data["litellm_session_id"] == "explicit-trace-id-value" + + +def test_add_litellm_metadata_from_request_headers_anthropic_metadata_beats_baggage(): + """The existing Anthropic metadata.user_id session_id path must win over a + baggage session.id fallback.""" + data = { + "metadata": { + "user_id": "user_abc123_account__session_e96634a3-fa28-4083-b354-55542e2dca01", + } + } + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers={"baggage": "session.id=baggage-session-42"}, + data=data, + _metadata_variable_name="metadata", + ) + assert data["litellm_session_id"] == "e96634a3-fa28-4083-b354-55542e2dca01" + assert "litellm_trace_id" not in data + + def test_get_internal_user_header_from_mapping_returns_expected_header(): mappings = [ {"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"}, diff --git a/tests/test_litellm/test_logging.py b/tests/test_litellm/test_logging.py index beba5794444..9ab362f6cd5 100644 --- a/tests/test_litellm/test_logging.py +++ b/tests/test_litellm/test_logging.py @@ -17,9 +17,15 @@ import sys import litellm from litellm._logging import ( ALL_LOGGERS, + CorrelationContextFilter, + CorrelationPlainFormatter, JsonFormatter, _initialize_loggers_with_handler, _turn_on_json, + session_id_var, + set_session_id, + set_trace_id, + trace_id_var, verbose_logger, verbose_proxy_logger, verbose_router_logger, @@ -393,3 +399,244 @@ def test_logging_calls_do_not_build_their_message_eagerly(): "these logging calls build their message eagerly; pass the values as %-style arguments instead:\n" + "\n".join(offenders) ) + + +class _JsonCapture(logging.Handler): + def __init__(self): + super().__init__() + self.formatter = JsonFormatter() + self.records: list[dict] = [] + self.addFilter(CorrelationContextFilter()) + + def emit(self, record): + self.records.append(json.loads(self.formatter.format(record))) + + +def _make_capture_logger(name: str) -> tuple[logging.Logger, _JsonCapture]: + lg = logging.getLogger(name) + cap = _JsonCapture() + lg.addHandler(cap) + lg.setLevel(logging.DEBUG) + return lg, cap + + +def test_trace_id_injected_into_json_record(monkeypatch): + """trace_id set via set_trace_id() appears in every JSON record in that context.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + lg, cap = _make_capture_logger("test.trace_inject") + set_trace_id("trace-abc-123") + try: + lg.info("test message") + assert len(cap.records) == 1 + assert cap.records[0]["trace_id"] == "trace-abc-123" + finally: + trace_id_var.set("") + + +def test_session_id_injected_when_set(monkeypatch): + """session_id set via set_session_id() appears in JSON record.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + lg, cap = _make_capture_logger("test.session_inject") + set_session_id("sess-xyz-456") + try: + lg.info("another message") + assert cap.records[0]["session_id"] == "sess-xyz-456" + finally: + session_id_var.set("") + + +def test_trace_id_and_session_id_cannot_be_spoofed_by_message_content(monkeypatch): + """A log message that happens to parse as JSON/dict with "trace_id"/"session_id" + keys (e.g. the proxy logging a raw request-header dict) must not override the + real correlation ids set via set_trace_id()/set_session_id().""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + lg, cap = _make_capture_logger("test.spoof_attempt") + set_trace_id("real-trace-id") + set_session_id("real-session-id") + try: + lg.info('{"trace_id": "attacker-supplied-trace", "session_id": "attacker-supplied-session"}') + assert cap.records[0]["trace_id"] == "real-trace-id" + assert cap.records[0]["session_id"] == "real-session-id" + finally: + trace_id_var.set("") + session_id_var.set("") + + +def test_trace_id_and_session_id_cannot_be_injected_with_no_active_context(monkeypatch): + """A message that happens to parse as JSON/dict with "trace_id"/"session_id" keys + must not surface those fields at all when CorrelationContextFilter hasn't stamped + this record - e.g. a log line emitted before Logging.__init__() runs for a request + (request_correlation_in_logs on, but no genuine trace/session id active yet).""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + lg, cap = _make_capture_logger("test.no_context_spoof_attempt") + trace_id_var.set("") + session_id_var.set("") + lg.info('{"trace_id": "attacker-supplied-trace", "session_id": "attacker-supplied-session"}') + assert "trace_id" not in cap.records[0] + assert "session_id" not in cap.records[0] + + +def test_trace_id_and_session_id_are_redacted_when_credential_shaped(monkeypatch): + """A caller-controlled trace_id/session_id (e.g. from x-litellm-trace-id or a W3C + baggage header) that happens to look like a real credential must not reach log + records unredacted. CorrelationContextFilter stamps trace_id/session_id onto the + record after SecretRedactionFilter has already run, so those two fields would + otherwise bypass credential redaction entirely - the fix redacts at set_trace_id()/ + set_session_id() time instead, before the value ever reaches a log record.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + lg, cap = _make_capture_logger("test.credential_shaped_correlation_id") + poisoned_trace_id = "sk-ant-api03-" + "A" * 40 + poisoned_session_id = "AKIA" + "B" * 16 + set_trace_id(poisoned_trace_id) + set_session_id(poisoned_session_id) + try: + lg.info("some benign log line") + assert cap.records[0]["trace_id"] == "REDACTED" + assert cap.records[0]["session_id"] == "REDACTED" + assert poisoned_trace_id not in json.dumps(cap.records[0]) + assert poisoned_session_id not in json.dumps(cap.records[0]) + finally: + trace_id_var.set("") + session_id_var.set("") + + +def test_session_id_absent_when_not_set(): + """session_id must NOT appear in JSON record when not set for this context.""" + lg, cap = _make_capture_logger("test.no_session") + session_id_var.set("") + lg.info("no session message") + assert "session_id" not in cap.records[0] + + +def test_trace_id_absent_when_not_set(): + """trace_id must NOT appear when not set.""" + lg, cap = _make_capture_logger("test.no_trace") + trace_id_var.set("") + lg.info("no trace message") + assert "trace_id" not in cap.records[0] + + +@pytest.mark.asyncio +async def test_contextvar_isolation_between_tasks(): + """Two concurrent async tasks each see only their own trace_id.""" + results: dict[str, str] = {} + + async def task(task_id: str, trace_id: str) -> None: + set_trace_id(trace_id) + await asyncio.sleep(0) + results[task_id] = trace_id_var.get() + + await asyncio.gather( + task("A", "trace-for-A"), + task("B", "trace-for-B"), + ) + + assert results["A"] == "trace-for-A" + assert results["B"] == "trace-for-B" + + +def test_trace_id_not_in_log_when_flag_disabled(monkeypatch): + """When request_correlation_in_logs is False (default), trace_id must not appear in JSON records even when set.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", False) + lg, cap = _make_capture_logger("test.no_trace_gated") + set_trace_id("trace-should-not-appear") + try: + lg.info("message") + assert "trace_id" not in cap.records[0] + finally: + trace_id_var.set("") + + +def test_session_id_not_in_log_when_flag_disabled(monkeypatch): + """When request_correlation_in_logs is False (default), session_id must not appear in JSON records even when set.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", False) + lg, cap = _make_capture_logger("test.no_session_gated") + set_session_id("sess-should-not-appear") + try: + lg.info("message") + assert "session_id" not in cap.records[0] + finally: + session_id_var.set("") + + +class _PlainCapture(logging.Handler): + def __init__(self): + super().__init__() + self.formatter = CorrelationPlainFormatter("%(message)s") + self.records: list[str] = [] + self.addFilter(CorrelationContextFilter()) + + def emit(self, record): + self.records.append(self.formatter.format(record)) + + +def _make_plain_capture_logger(name: str) -> tuple[logging.Logger, _PlainCapture]: + lg = logging.getLogger(name) + cap = _PlainCapture() + lg.addHandler(cap) + lg.setLevel(logging.DEBUG) + return lg, cap + + +def test_plain_formatter_appends_trace_id_and_session_id(monkeypatch): + """CorrelationPlainFormatter must append trace_id/session_id to non-JSON log lines too.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + lg, cap = _make_plain_capture_logger("test.plain_trace_session") + set_trace_id("plain-trace-1") + set_session_id("plain-session-1") + try: + lg.info("plaintext message") + assert cap.records[0] == "plaintext message [trace_id=plain-trace-1 session_id=plain-session-1]" + finally: + trace_id_var.set("") + session_id_var.set("") + + +def test_plain_formatter_appends_only_trace_id_when_session_id_absent(monkeypatch): + """Only trace_id is appended when session_id was never set.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + lg, cap = _make_plain_capture_logger("test.plain_trace_only") + set_trace_id("plain-trace-2") + session_id_var.set("") + try: + lg.info("plaintext message") + assert cap.records[0] == "plaintext message [trace_id=plain-trace-2]" + finally: + trace_id_var.set("") + + +def test_plain_formatter_unchanged_when_flag_disabled(monkeypatch): + """When request_correlation_in_logs is False, plain log lines are unmodified even if the contextvars are set.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", False) + lg, cap = _make_plain_capture_logger("test.plain_flag_off") + set_trace_id("should-not-appear") + set_session_id("should-not-appear") + try: + lg.info("plaintext message") + assert cap.records[0] == "plaintext message" + finally: + trace_id_var.set("") + session_id_var.set("") + + +def test_set_trace_id_strips_control_characters(): + """set_trace_id() must strip \\r/\\n/escape sequences so a caller-controlled + trace id can't forge fake log entries when interpolated into plain-text logs.""" + token = set_trace_id('evil\r\n{"level": "CRITICAL", "message": "forged"}') + try: + value = trace_id_var.get() + assert "\r" not in value + assert "\n" not in value + finally: + trace_id_var.reset(token) + + +def test_set_session_id_bounds_length(): + """set_session_id() must bound length so an oversized caller-supplied value + isn't repeated across every log line for the request.""" + token = set_session_id("a" * 1000) + try: + assert len(session_id_var.get()) == 256 + finally: + session_id_var.reset(token) + diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index e80960a22c9..048d7c8f3cc 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -1,4 +1,5 @@ import json +import logging import os import sys from unittest.mock import AsyncMock, MagicMock, patch @@ -11,6 +12,13 @@ sys.path.insert( ) # Adds the parent directory to the system path import litellm +from litellm._logging import ( + CorrelationContextFilter, + JsonFormatter, + session_id_var, + trace_id_var, + verbose_logger, +) from litellm.proxy.utils import is_valid_api_key from litellm.types.utils import ( CallTypes, @@ -5125,3 +5133,124 @@ def test_ai21_api_key_is_resolved_from_the_documented_env_var(monkeypatch: pytes monkeypatch.setenv("AI21_API_KEY", "sk-ai21-resolved-from-env") assert get_api_key(llm_provider="ai21", dynamic_api_key=None) == "sk-ai21-resolved-from-env" + + +class _JsonCapture(logging.Handler): + def __init__(self): + super().__init__() + self.formatter = JsonFormatter() + self.records: list[dict] = [] + self.addFilter(CorrelationContextFilter()) + + def emit(self, record): + self.records.append(json.loads(self.formatter.format(record))) + + +def _make_capture_logger(name: str) -> tuple[logging.Logger, _JsonCapture]: + lg = logging.getLogger(name) + cap = _JsonCapture() + lg.addHandler(cap) + lg.setLevel(logging.DEBUG) + return lg, cap + + +@pytest.mark.asyncio +async def test_wrapper_async_restores_originating_task_context_after_success(monkeypatch): + """A successful acompletion() dispatches async_success_handler via + asyncio.create_task + the global logging worker - a different Task than the + one running acompletion() itself (this test's own task). That handler's own + restore only fixes up the detached child task it runs in; wrapper_async's own + finally block (in litellm/utils.py) must separately restore the *originating* + task's trace_id/session_id, since nothing else does. + """ + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + trace_id_var.set("outer-trace-wrapper-test") + session_id_var.set("outer-session-wrapper-test") + try: + await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + mock_response="Hello there!", + litellm_session_id="mock-call-session", + num_retries=0, + ) + assert trace_id_var.get() == "outer-trace-wrapper-test" + assert session_id_var.get() == "outer-session-wrapper-test" + finally: + trace_id_var.set("") + session_id_var.set("") + + +def test_function_setup_failure_after_logging_construction_restores_context(monkeypatch): + """If function_setup() constructs Logging() (which already mutated + trace_id_var/session_id_var in __init__) but then raises before returning, + the caller's wrapper() never gets a logging_obj reference to restore from. + function_setup()'s own except block must restore the correlation context + itself in that case, or it leaks into every subsequent log line in this + thread/task until something unrelated happens to reset it.""" + from litellm.litellm_core_utils.litellm_logging import Logging + + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + + def _boom(self, *args, **kwargs): + raise RuntimeError("simulated failure after Logging() construction") + + monkeypatch.setattr(Logging, "update_environment_variables", _boom) + + trace_id_var.set("pre-setup-failure-trace") + session_id_var.set("pre-setup-failure-session") + try: + with pytest.raises(RuntimeError, match="simulated failure"): + litellm.completion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + mock_response="Hello there!", + litellm_session_id="doomed-call-session", + num_retries=0, + ) + assert trace_id_var.get() == "pre-setup-failure-trace" + assert session_id_var.get() == "pre-setup-failure-session" + finally: + trace_id_var.set("") + session_id_var.set("") + + +def test_function_setup_failure_log_line_shows_outer_not_doomed_ids(monkeypatch): + """The 'Error in function_setup' diagnostic log line itself must be stamped + with the outer/pre-call correlation ids, not the doomed call's own ids - + restoring context must happen *before* logging the exception, not after, + since the failed call never produces a usable logging object for anything + else to be attributed to.""" + from litellm.litellm_core_utils.litellm_logging import Logging + + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + + def _boom(self, *args, **kwargs): + raise RuntimeError("simulated failure after Logging() construction") + + monkeypatch.setattr(Logging, "update_environment_variables", _boom) + + lg, cap = _make_capture_logger("test.function_setup_failure_log_order") + # verbose_logger is a distinct, module-level logger from our throwaway one - + # temporarily attach the same capture handler so we see its own emitted record. + verbose_logger.addHandler(cap) + try: + trace_id_var.set("outer-trace") + session_id_var.set("outer-session") + with pytest.raises(RuntimeError, match="simulated failure"): + litellm.completion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + mock_response="Hello there!", + litellm_session_id="doomed-call-session", + num_retries=0, + ) + setup_failure_records = [r for r in cap.records if "Error in function_setup" in r.get("message", "")] + assert len(setup_failure_records) == 1 + record = setup_failure_records[0] + assert record.get("session_id") == "outer-session" + assert record.get("trace_id") == "outer-trace" + finally: + verbose_logger.removeHandler(cap) + trace_id_var.set("") + session_id_var.set("")