perf(logging): skip correlation contextvar stamping when request_correlation_in_logs is off

litellm/litellm_core_utils/litellm_logging.py

@@ -556,7 +556,7 @@ def __init__(
         # 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 @@ def _restore_correlation_context(self) -> None:
         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 @@ def _restore_correlation_context(self) -> None:
         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.

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,14 +5241,15 @@ 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
     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
 
+    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,14 +5351,15 @@ 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
     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
 
+    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:

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 _restore_correlation_context(self):
     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 _empty_aiter():
 
 
 @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: