scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

Merge pull request #41054 from BerriAI/litellm_gate_correlation_contextvar_stamp

Browse source 
perf(logging): skip correlation contextvar stamping when request_correlation_in_logs is off
This commit is contained in:
Yassin Kortam 2026-09-14 11:14:44 -07:00 • committed by GitHub
commit 08a78a3982
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 104 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
litellm/litellm_core_utils/litellm_logging.py
View file

@ -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.

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

@ -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:

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

@ -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: