diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 5fe94320491..9ba9fd082f3 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -556,7 +556,7 @@ class Logging(LiteLLMLoggingBaseClass): # 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: + if supports_correlation_logging and litellm.request_correlation_in_logs: 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 @@ -2442,7 +2442,7 @@ class Logging(LiteLLMLoggingBaseClass): 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 + Uses a plain contextvar 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 @@ -2453,8 +2453,8 @@ class Logging(LiteLLMLoggingBaseClass): 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) + trace_id_var.set(self._pre_call_trace_id) + session_id_var.set(self._pre_call_session_id) def _restore_correlation_context_if_unclaimed(self) -> None: """Guarded variant for __del__-triggered cleanup only. 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 451b740dd71..70f9bae283b 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -5221,10 +5221,11 @@ def test_handle_anthropic_messages_parsed_response_logging_preserves_fast_mode_s assert getattr(result.usage, "speed", None) == "fast" -def test_logging_init_sets_trace_id(): +def test_logging_init_sets_trace_id(monkeypatch): """Logging.__init__() must call set_trace_id with self.litellm_trace_id.""" from litellm.litellm_core_utils.litellm_logging import Logging + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) trace_id_var.set("") log_obj = Logging( @@ -5240,7 +5241,7 @@ def test_logging_init_sets_trace_id(): assert trace_id_var.get() == log_obj.litellm_trace_id -def test_logging_init_skips_stamping_when_correlation_logging_unsupported(): +def test_logging_init_skips_stamping_when_correlation_logging_unsupported(monkeypatch): """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 @@ -5248,6 +5249,7 @@ def test_logging_init_skips_stamping_when_correlation_logging_unsupported(): usual - only the ambient contextvar stamping is gated.""" from litellm.litellm_core_utils.litellm_logging import Logging + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) trace_id_var.set("") session_id_var.set("") @@ -5271,10 +5273,48 @@ def test_logging_init_skips_stamping_when_correlation_logging_unsupported(): assert log_obj.litellm_session_id == "should-not-be-stamped" -def test_logging_init_sets_session_id_when_provided(): +def test_logging_init_skips_stamping_when_request_correlation_in_logs_disabled(monkeypatch): + from litellm.litellm_core_utils.litellm_logging import Logging + + monkeypatch.setattr(litellm, "request_correlation_in_logs", False) + trace_id_var.set("outer") + session_id_var.set("outer-sid") + try: + with ( + patch( # test-quality-ok: regression test verifies disabled stamping skips both setters + "litellm.litellm_core_utils.litellm_logging.set_trace_id" + ) as mock_set_trace_id, + patch( # test-quality-ok: regression test verifies disabled stamping skips both setters + "litellm.litellm_core_utils.litellm_logging.set_session_id" + ) as mock_set_session_id, + ): + 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-disabled", + function_id="fn-disabled", + kwargs={"litellm_session_id": "disabled-session"}, + supports_correlation_logging=True, + ) + + assert trace_id_var.get() == "outer" + assert session_id_var.get() == "outer-sid" + assert log_obj._own_trace_id == "outer" + mock_set_trace_id.assert_not_called() + mock_set_session_id.assert_not_called() + finally: + trace_id_var.set("") + session_id_var.set("") + + +def test_logging_init_sets_session_id_when_provided(monkeypatch): """Logging.__init__() must call set_session_id when litellm_session_id is in kwargs.""" from litellm.litellm_core_utils.litellm_logging import Logging + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) session_id_var.set("") Logging( @@ -5290,11 +5330,12 @@ def test_logging_init_sets_session_id_when_provided(): assert session_id_var.get() == "my-session-99" -def test_logging_init_resets_session_id_to_empty_when_absent(): +def test_logging_init_resets_session_id_to_empty_when_absent(monkeypatch): """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 + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) session_id_var.set("preexisting-sid") Logging( @@ -5310,7 +5351,7 @@ def test_logging_init_resets_session_id_to_empty_when_absent(): assert session_id_var.get() == "" -def test_restore_correlation_context_resets_to_pre_call_value(): +def test_restore_correlation_context_resets_to_pre_call_value(monkeypatch): """_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 @@ -5318,6 +5359,7 @@ def test_restore_correlation_context_resets_to_pre_call_value(): session_id into the outer call's subsequent log lines.""" from litellm.litellm_core_utils.litellm_logging import Logging + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) trace_id_var.set("outer-trace") session_id_var.set("outer-session") try: @@ -5343,7 +5385,7 @@ def test_restore_correlation_context_resets_to_pre_call_value(): session_id_var.set("") -def test_restore_correlation_context_safe_to_call_repeatedly(): +def test_restore_correlation_context_safe_to_call_repeatedly(monkeypatch): """Calling _restore_correlation_context() more than once must not raise. It's deliberately NOT guarded against repeat calls: wrapper()'s finally @@ -5353,6 +5395,7 @@ def test_restore_correlation_context_safe_to_call_repeatedly(): the contextvars, so repeat calls are expected, not just tolerated.""" from litellm.litellm_core_utils.litellm_logging import Logging + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) log_obj = Logging( model="gpt-3.5-turbo", messages=[{"role": "user", "content": "hi"}], @@ -5367,8 +5410,40 @@ def test_restore_correlation_context_safe_to_call_repeatedly(): log_obj._restore_correlation_context() # must not raise +def test_restore_correlation_context_does_not_resanitize(monkeypatch): + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm._logging import _sanitize_correlation_id + + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + trace_id_var.set("outer-trace") + session_id_var.set("outer-session") + try: + 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-no-resanitize", + function_id="fn-no-resanitize", + kwargs={"litellm_session_id": "inner-session"}, + ) + + with patch( # test-quality-ok: regression test verifies restore avoids sanitization + "litellm._logging._sanitize_correlation_id", wraps=_sanitize_correlation_id + ) as mock_sanitize: + log_obj._restore_correlation_context() + + mock_sanitize.assert_not_called() + assert trace_id_var.get() == "outer-trace" + assert session_id_var.get() == "outer-session" + finally: + trace_id_var.set("") + session_id_var.set("") + + @pytest.mark.asyncio -async def test_restore_correlation_context_works_across_asyncio_task_boundary(): +async def test_restore_correlation_context_works_across_asyncio_task_boundary(monkeypatch): """_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 @@ -5385,6 +5460,7 @@ async def test_restore_correlation_context_works_across_asyncio_task_boundary(): """ from litellm.litellm_core_utils.litellm_logging import Logging + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) trace_id_var.set("outer-trace-cross-task") session_id_var.set("outer-session-cross-task") try: 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 37e2031fdf4..784468c839b 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -4100,7 +4100,7 @@ async def test_async_streaming_completion_does_not_reset_context_before_iteratio session_id_var.set("") -def test_stream_wrapper_del_restores_correlation_context(): +def test_stream_wrapper_del_restores_correlation_context(monkeypatch): """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 @@ -4112,6 +4112,7 @@ def test_stream_wrapper_del_restores_correlation_context(): doesn't run actual finalization, and this exercises exactly the logic that real garbage collection would eventually trigger. """ + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) trace_id_var.set("outer-trace-abandoned") session_id_var.set("outer-session-abandoned") try: @@ -4159,12 +4160,13 @@ def test_stream_wrapper_del_never_raises_with_broken_logging_obj(): wrapper.__del__() # must not raise -def test_stream_wrapper_del_does_not_clobber_a_newer_active_call(): +def test_stream_wrapper_del_does_not_clobber_a_newer_active_call(monkeypatch): """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.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) trace_id_var.set("outer-trace-before-abandoned-call") session_id_var.set("outer-session-before-abandoned-call") try: @@ -4210,13 +4212,14 @@ def test_stream_wrapper_del_does_not_clobber_a_newer_active_call(): session_id_var.set("") -def test_stream_wrapper_del_restores_when_own_session_id_needed_sanitizing(): +def test_stream_wrapper_del_restores_when_own_session_id_needed_sanitizing(monkeypatch): """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.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) trace_id_var.set("outer-trace-needs-sanitizing") session_id_var.set("outer-session-needs-sanitizing") try: @@ -4250,7 +4253,7 @@ def test_stream_wrapper_del_restores_when_own_session_id_needed_sanitizing(): session_id_var.set("") -def test_stream_wrapper_next_keeps_context_active_through_synthesized_finish_reason_chunk(): +def test_stream_wrapper_next_keeps_context_active_through_synthesized_finish_reason_chunk(monkeypatch): """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 @@ -4261,6 +4264,7 @@ def test_stream_wrapper_next_keeps_context_active_through_synthesized_finish_rea correct, deterministic restore on the very next __next__() call, since completion_stream is already exhausted and immediately re-raises StopIteration.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) trace_id_var.set("outer-trace-finish-reason") session_id_var.set("outer-session-finish-reason") try: @@ -4300,12 +4304,13 @@ def test_stream_wrapper_next_keeps_context_active_through_synthesized_finish_rea session_id_var.set("") -def test_stream_wrapper_del_cleans_up_after_synthesized_finish_reason_chunk(): +def test_stream_wrapper_del_cleans_up_after_synthesized_finish_reason_chunk(monkeypatch): """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.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) trace_id_var.set("outer-trace-finish-reason-del") session_id_var.set("outer-session-finish-reason-del") try: @@ -4338,10 +4343,11 @@ def test_stream_wrapper_del_cleans_up_after_synthesized_finish_reason_chunk(): @pytest.mark.asyncio -async def test_stream_wrapper_anext_keeps_context_active_through_synthesized_finish_reason_chunk(): +async def test_stream_wrapper_anext_keeps_context_active_through_synthesized_finish_reason_chunk(monkeypatch): """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.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) trace_id_var.set("outer-trace-anext-finish-reason") session_id_var.set("outer-session-anext-finish-reason") try: @@ -4394,6 +4400,7 @@ async def test_stream_wrapper_anext_max_duration_timeout_restores_consumer_corre 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, "request_correlation_in_logs", True) 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") @@ -4434,12 +4441,13 @@ async def test_stream_wrapper_anext_max_duration_timeout_restores_consumer_corre @pytest.mark.asyncio -async def test_stream_wrapper_aclose_restores_consumer_correlation_context(): +async def test_stream_wrapper_aclose_restores_consumer_correlation_context(monkeypatch): """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__.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) trace_id_var.set("outer-trace-aclose") session_id_var.set("outer-session-aclose") try: @@ -4481,6 +4489,7 @@ async def test_stream_wrapper_aclose_keeps_context_active_through_close_failure_ 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.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) trace_id_var.set("outer-trace-close-fail") session_id_var.set("outer-session-close-fail") try: @@ -4541,6 +4550,7 @@ def test_handle_stream_fallback_error_restores_context_only_after_exception_mapp 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.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) trace_id_var.set("outer-trace-fallback") session_id_var.set("outer-session-fallback") try: