feat(logging): add opt-in session_id and trace_id correlation to JSON log records via contextvars (#34418)

* feat(logging): add opt-in session_id/trace_id correlation to JSON log records via contextvars

Adds two ContextVar instances (session_id_var, trace_id_var) to litellm/_logging.py and
two setter functions (set_session_id, set_trace_id). Logging.__init__() now calls both
setters after assigning litellm_trace_id so every JSON log record emitted within the
async request context carries trace_id and, when provided, session_id — enabling log
correlation in Loki, CloudWatch Logs Insights, and other structured-log sinks without
any changes to individual log call sites.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* fix(logging): guard session_id/trace_id injection against overwriting caller-supplied extra fields

* fix(logging): always reset session_id_var to empty string when no session_id provided

* feat: gate request correlation IDs in logs behind request_correlation_in_logs flag

* refactor: move correlation ID injection into CorrelationContextFilter

* feat(logging): extend request_correlation_in_logs to plaintext logs and StandardLoggingPayload

Plaintext log lines (json_logs off) now get the same trace_id/session_id
suffix as JSON logs via a new CorrelationPlainFormatter, so the flag has a
visible effect regardless of log format.

StandardLoggingPayload gets a new independent session_id field, populated
from litellm_session_id. trace_id's existing session_id-first fallback is
preserved when request_correlation_in_logs is off; with the flag on, an
explicit litellm_trace_id now takes priority over litellm_session_id so
the two fields carry genuinely independent values.

* fix(logging): restore correlation context after nested calls; sanitize correlation ids

Addresses two review findings on this PR.

CorrelationContextFilter's trace_id/session_id contextvars were set on every
Logging.__init__ but never reset, so 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 ids instead of its own. set_trace_id/
set_session_id now return their contextvars.Token, and Logging stores them
and resets both once its own success/failure handler actually completes,
via a new idempotent _restore_correlation_context() called from all four
terminal handlers.

set_trace_id/set_session_id also now strip control characters and bound
length before storing a caller-controlled trace_id/session_id, since these
values can originate from request input (litellm_session_id, x-litellm-
trace-id) and get interpolated into plain-text log lines - without this, a
caller could embed \r/\n or escape sequences to forge fake log entries.

* fix(logging): restore correlation context after nested calls, not before

The previous commit called _restore_correlation_context() as the first
line of each terminal handler, before that handler's own callback
dispatch loop runs. That's backwards: a nested LiteLLM call triggered
from within a callback (e.g. a guardrail's own LLM-as-judge call) would
then capture the *already-reset* value as its own pre-call baseline,
and its own reset would restore to that instead of the true outer
value - verified live to still leak.

success_handler/async_success_handler/failure_handler/async_failure_handler
are now thin wrappers: the original bodies move to
_success_handler_body/etc, called inside a try/finally that restores
context only once the full body - including any nested calls its own
callback dispatch triggers - has actually finished, mirroring proper
stack-scoped nesting semantics.

* test(logging): cover async_failure_handler's correlation-context restore

Codecov flagged the new async_failure_handler wrapper (try/finally around
_async_failure_handler_body) as uncovered - the method had no direct test
at all before this PR's refactor split it into a wrapper. Adds a test that
awaits it directly and asserts both that async_log_failure_event still
fires and that _restore_correlation_context() puts the pre-call
trace_id/session_id back.

* fix(logging): restore correlation context by value, not by contextvars.Token

veria-ai correctly flagged that contextvars.Token.reset() only works in the
exact Context it was created in, and litellm's async success path (and
streaming failure path) dispatch async_success_handler/async_failure_handler
via asyncio.create_task and the global logging worker - a different Context
than Logging.__init__ ran in. reset_trace_id/reset_session_id silently
swallowed the resulting ValueError, so the restore was a no-op for exactly
those paths. Verified independently: reproduced the raw contextvars
behavior, then confirmed litellm's async success dispatch really does go
through asyncio.create_task + GLOBAL_LOGGING_WORKER (litellm/utils.py).

Logging now captures the pre-call *value* (not a Token) and restores via a
plain set_trace_id()/set_session_id() call, which works regardless of which
Task/Context calls it. reset_trace_id/reset_session_id are removed as
dead/unreliable code. Added a regression test that spawns __init__ and the
restore in different asyncio Tasks - confirmed it fails against the prior
Token-based commit and passes here.

* fix(logging): restore correlation context in the originating task too

Greptile's re-review correctly identified a remaining gap: for a
successful acompletion(), async_success_handler is dispatched via
asyncio.create_task + the global logging worker into a *different* Task
than the one wrapper_async/Logging.__init__ ran in. The prior fix (43c164a)
only restored the handler's own (detached, throwaway) Task - it never
touched the originating request Task, which keeps this call's trace_id/
session_id set for the rest of its own execution (e.g. nested calls made
via the same Task).

wrapper()/wrapper_async() in litellm/utils.py now restore the originating
Task's correlation context in a finally block once the whole call is done,
regardless of what detached logging tasks it spawned along the way. Since
the wrapped body rebinds its own `kwargs` local via function_setup(),
sharing the dict object doesn't work here; a small mutable holder carries
the constructed Logging instance back out to the outer wrapper instead.

_restore_correlation_context() is no longer guarded against repeat calls:
with value-based (not Token-based) restoration, each distinct Task that
calls it needs its own restore to take effect in that Task's own view of
the contextvars, so multiple calls (once per Task involved in an attempt)
are required, not just tolerated.

Added a regression test using mock_response to exercise the real success
dispatch path (asyncio.create_task + GLOBAL_LOGGING_WORKER) without a live
provider call, asserting the *test's own* (originating) task context is
restored after the call - this is exactly the case Greptile flagged and
the prior commit didn't cover.

* fix(logging): restore correlation context when function_setup itself fails

Greptile's 4th finding: if function_setup() constructs Logging() (whose
__init__ already mutates trace_id_var/session_id_var) and then raises
before returning - e.g. update_environment_variables() throws - the
caller's wrapper()/wrapper_async() never receives a logging_obj reference,
so its own restore-on-finally never fires. The correlation ids leak into
every subsequent log line on that thread/task with no way to clear them.

function_setup()'s own except block now restores the context itself in
that case, using whatever logging_obj it managed to construct before
failing (locals().get(), safe against the earlier failure modes where
logging_obj was never assigned at all).

Added a regression test that monkeypatches Logging.update_environment_variables
to raise after construction, confirmed it fails without this fix (the
leaked ids show up directly in the raised exception's own log line) and
passes with it. Broader sweep (test_utils.py, test_router.py,
test_main_module_header.py, streaming handler tests, plus all
logging-specific tests): 722 passed.

* fix(logging): don't assume every litellm_logging_obj is a real Logging instance

CI caught a real regression from the last commit: tests/test_litellm/llms/xai/test_xai_key_fallback.py
injects a minimal FakeLogging stand-in (only implementing
update_from_kwargs) as litellm_logging_obj for a narrow realtime-config
unit test, bypassing the real Logging class entirely. wrapper()/
wrapper_async()'s finally block and function_setup()'s except block both
unconditionally called _restore_correlation_context() on whatever ended up
in the holder, which doesn't exist on that stand-in.

_restore_correlation_context is new plumbing specific to this PR's
feature, not part of any pre-existing stand-in's expected interface, so
callers of it can't assume every object playing the litellm_logging_obj
role implements it. Added _restore_correlation_context_if_supported(),
a small getattr-guarded helper, and used it at all three call sites.

* fix(logging): don't restore context too early on setup failure or streaming

Two more findings from Greptile's 5th review round.

1. function_setup()'s except block restored correlation context *after*
   logging the "Error in function_setup" exception, so that diagnostic log
   line itself was stamped with the doomed call's ids instead of the outer
   ids - misleading, since the failed call never produces anything else to
   attribute those ids to. Restore now happens before the log call.

2. wrapper()/wrapper_async() restored the originating task's context as
   soon as a streaming call returned, before the caller ever starts
   iterating the CustomStreamWrapper it just got back. Any log lines
   emitted while iterating (in the same thread/task) incorrectly showed
   the pre-call ids instead of this call's own ones. The wrapper finally
   block now skips the restore when the return value is a stream wrapper,
   deferring to the terminal handler that already fires once the stream is
   actually assembled/exhausted.

Both verified with tests that fail against the prior commit and pass
against this one. Broader sweep unchanged at 829 passing.

* fix(logging): best-effort correlation cleanup on abandoned streams

Greptile's 7th finding: if a caller returns a streaming response and never
fully consumes it - stops iterating early, drops the reference, cancels
it - the terminal handler that normally restores the originating task's
trace_id/session_id never fires, since it only runs once the stream is
actually assembled/exhausted. The ids leak into every subsequent log line
in that thread/task with no bound.

There's no reliable Python hook for "this was abandoned without being
closed" - CustomStreamWrapper has no close()/__aexit__/context-manager
convention today, and the only automatic option is __del__, whose timing
is inherently unpredictable (delayed by cyclic GC, not guaranteed at
interpreter shutdown, can run on a different thread). This is a best-effort
safety net, not a guarantee, and is documented as such in the docstring.

Testing this via real garbage collection proved unreliable in practice:
per-chunk logging submits work to a thread pool executor whose worker
thread transiently holds its own bound-method reference to the wrapper
until that task completes, so refcount doesn't hit zero on a
deterministic schedule even with polling. Tests call __del__ directly
instead - a plain method, safe to invoke early - which exercises exactly
the restore logic real garbage collection would eventually trigger,
plus a case confirming a broken logging_obj can never make __del__ raise.

* fix(logging): restore consumer's context at every real stream exit point

Two more findings from this round.

Veria AI: even a *fully consumed* stream never restored the actual
consuming thread/task's correlation context. The terminal success dispatch
(dispatch_success_handlers via asyncio.create_task for async, or
success_handler via the shared executor for sync) only restores whatever
detached context it runs in - never the caller's own thread/task that's
running the for/async for loop. Same root cause as the wrapper-level fix
two rounds ago, just missed for the streaming-completion path.

Greptile: explicit aclose() (client disconnect, router fallback aborting
a partial stream) closed the underlying stream without restoring
correlation context either, since request wrappers intentionally skip
restoration for returned streams and no terminal handler runs on this
path.

Added CustomStreamWrapper._restore_consumer_correlation_context(), called
from every point control genuinely returns to the consumer: the final
raise StopIteration/StopAsyncIteration on natural exhaustion (both sync
branches, both async branches), _handle_stream_fallback_error (the shared
choke point for all three failure-raising call sites), and aclose(). __del__
now delegates to the same helper instead of duplicating it.

Verified with tests extending the existing streaming-exhaustion cases to
assert the consuming context is restored after the loop completes (fails
against the prior commit, passes now), plus a dedicated aclose() test.
Broader sweep: 832 passing.

* fix(logging): don't let a delayed __del__ finalizer clobber a newer active call

If an abandoned stream's __del__ fires late (after cyclic GC delay), a
different call may have already taken over the correlation contextvars in
the same Task/thread. Restoring unconditionally would stomp that active
call's trace_id/session_id with the abandoned stream's stale pre-call
snapshot. __del__ now only restores when the contextvars still hold the
ids this call itself set.

* fix(logging): compare sanitized ids in the __del__ ownership guard

set_trace_id()/set_session_id() sanitize (strip control chars, bound length)
before storing, so the contextvar's value can differ from the raw
litellm_trace_id/litellm_session_id. The __del__ ownership guard was
comparing against the raw values, so a caller-supplied id containing control
characters or exceeding 256 chars would never match, permanently skipping
cleanup. Capture what set_trace_id()/set_session_id() actually stored and
compare against that instead.

* fix(logging): restore consumer context on the synthesized finish_reason chunk

Both __next__ and _finalize_completed_stream() have a branch that fires when
the underlying stream ends without ever emitting an explicit finish_reason
chunk: they synthesize one via finish_reason_handler() and return it. A
consumer that stops as soon as it sees finish_reason - a common pattern -
never calls __next__()/__anext__() again, so the existing restore in the
sent_last_chunk-is-True StopIteration branch never runs for them. The
underlying stream is already exhausted at this point regardless of whether
the caller keeps iterating, so restoring here is safe.

* fix(logging): don't restore correlation context before the caller receives the final chunk

The previous fix (5147c69186) restored context immediately before returning
the synthesized finish_reason chunk from __next__/_finalize_completed_stream,
reasoning that completion_stream was already exhausted. But that chunk is
still this call's own data, and the caller's own application-level log
statements processing it run in the same synchronous frame right after the
return - restoring first made those lines carry the wrong (outer) ids,
exactly what wrapper()/wrapper_async() deliberately avoid by not restoring
while a stream is being iterated.

Revert to not restoring there. A caller that keeps iterating still gets a
correct, deterministic restore on its very next __next__()/__anext__() call
(completion_stream is exhausted, so that immediately re-raises
StopIteration/StopAsyncIteration through the already-restoring branch). A
caller that stops right after finish_reason relies on aclose() or the
best-effort __del__ guard, same as any other stream the caller doesn't fully
exhaust.

* refactor(logging): hoist a safely-hoistable function-body import to module top

CorrelationContextFilter.filter()'s `import litellm` was a function-body
import; verified it can move to module top without a circular-import failure
(litellm/__init__.py already imports from litellm._logging before setting
request_correlation_in_logs, but a bare `import litellm` only binds the
already-in-sys.modules module object - the attribute itself isn't read until
filter() actually runs, by which point litellm is fully initialized).

* test(logging): move correlation tests into their conventionally-mapped files

tests/test_litellm/ mirrors litellm/ in a parallel path. Correlation tests
for the Logging class (litellm_logging.py), function_setup/wrapper_async
(utils.py), and CustomStreamWrapper (streaming_handler.py) had all landed in
test_logging.py, which only maps to litellm/_logging.py itself. Moving each
group to its correctly-mapped file: test_litellm_logging.py (Logging class
init/restore), test_utils.py (function_setup, wrapper_async), and
test_streaming_handler.py (CustomStreamWrapper) in the next commit.
test_logging.py keeps only what actually exercises _logging.py's own
contextvars/filters/formatters/sanitization. No behavior change - same
assertions, same coverage, just relocated.

* fix(logging): restore correlation context unconditionally in wrapper()'s sync path

Blocking finding from review: a caller-visible correlation feature was
silently misattributing one request's logs to a different, unrelated one on
the sync/threaded path. wrapper()/wrapper_async() both left trace_id/session_id
"open" across a stream's entire iteration so the caller's own log lines while
consuming it would carry the right ids. That's safe for wrapper_async(): each
async call gets its own asyncio Task with its own copy of the contextvars,
and Tasks are never recycled across requests, so a leftover value can only
ever affect that one already-abandoned Task.

It is not safe for wrapper() (sync): a plain OS thread has no such per-call
isolation, and a thread pool's worker threads *are* recycled across unrelated
requests. If a sync stream was abandoned (client disconnect, early break, an
uncaught exception) without ever being exhausted or closed, nothing restored
its contextvars, and a pool could later hand that same thread to a completely
different call, which would inherit the abandoned request's ids as its own
"pre-call" baseline and then restore back to that poison when it finished -
permanently misattributing every subsequent log line on that thread,
including its own, to the abandoned request. Strengthening the __del__
finalizer already added for this can't fix it: finalizer timing is exactly
what a permanently-reused thread can't rely on.

wrapper() now restores unconditionally in its own finally, before a sync
stream is ever handed back to the caller. The trade-off: a sync stream
consumer's own application-level log statements while iterating no longer
automatically carry this call's ids (litellm's own internal per-chunk
logging is unaffected, since it's dispatched separately). That's an
acceptable cost for eliminating a silent cross-request misattribution bug.
wrapper_async() keeps the existing conditional (skip-if-streaming) behavior,
justified by the Task-isolation argument above; CustomStreamWrapper's
__del__/aclose()/next-iteration restore machinery remains meaningful and
necessary there.

This also simplifies wrapper()/wrapper_async() back toward their original
shape: both previously used a mutable-dict-holder split into a separate
_body function to smuggle logging_obj/result out to an outer finally,
working around function_setup() rebinding its own local `kwargs`. That
restructuring is no longer needed - `logging_obj` (and, for wrapper_async(),
`result`) were already function-level locals in scope for a plain
try/finally; three of wrapper_async()'s retry-return statements now assign
through `result` first so it accurately reflects what's actually returned
even on a retry path.

Regression test: test_abandoned_sync_stream_does_not_contaminate_a_later_call_on_the_same_thread
in test_streaming_handler.py reproduces the exact reported scenario with a
real single-worker ThreadPoolExecutor - confirmed it fails with the prior
(skip-restore-on-stream) wrapper() and passes with this fix.

* refactor(logging): use Mapping instead of bare dict for read-only params

_get_standard_logging_payload_trace_id/_session_id only read litellm_params
(.get() calls, no mutation) - annotate it as Mapping[str, Any] rather than a
bare mutable dict, per the repo's no-mutable-collection-in-annotation rule.

* fix(logging): scope request_correlation_in_logs to the async/proxy path only

Blocking review finding: wrapper() (the sync entry point) used the same
skip-restore-on-stream design as wrapper_async(), but 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 - an abandoned
sync stream could leave its ids stuck on a thread a pool later hands to a
completely different request, misattributing that request's logs. A fix
existed and was tested (restore unconditionally in wrapper()'s own finally),
but it doesn't benefit this feature's primary consumer - the proxy only ever
calls the async entry point - and carries sync-specific complexity this PR
doesn't need.

Scope the feature to async only instead: Logging.__init__() takes a new
supports_correlation_logging parameter (default True), threaded down from a
new function_setup(..., is_async_call: bool = True) parameter. wrapper() is
the one caller that passes is_async_call=False; every other function_setup()
call site (wrapper_async(), the router, and proxy/MCP-internal call sites)
is already async and keeps the default. With
supports_correlation_logging=False, Logging.__init__() never calls
set_trace_id()/set_session_id() at all, so a sync call has nothing to leak
in the first place. wrapper() reverts to its pre-review shape with no
correlation-specific code at all.

StandardLoggingPayload's own trace_id/session_id fields are unaffected
either way - they're a deterministic per-call read of
self.litellm_trace_id/self.litellm_session_id, not ambient contextvar state,
so they were never exposed to the cross-request bug.

Full sync/direct-SDK support (stamping + its own safe-restore mechanism) is
deferred to a follow-up PR; the fix and its regression test already exist in
this branch's history at commit 9f3a20f4b2 and can be resurrected there.

Tests: replaced the two wrapper()-level tests with ones proving the new
invariant (sync calls, streaming and non-streaming, never touch
trace_id_var/session_id_var even when the caller explicitly passes
litellm_trace_id/litellm_session_id), and added a direct unit test for the
supports_correlation_logging=False gate on Logging.__init__ itself. Verified
live: a real proxy (Postgres-backed, real OpenAI calls) shows clean
trace_id/session_id isolation across two concurrent sessions with no
cross-contamination; a standalone script confirms real sync SDK calls
against a real model never touch the correlation contextvars.

* feat(logging): fall back to W3C traceparent/baggage for trace_id/session_id

request_correlation_in_logs previously only resolved trace_id/session_id from
litellm-specific sources: x-litellm-trace-id/x-litellm-session-id headers, a
generic x-<vendor>-session-id header, or Anthropic-style metadata.user_id. If
none were present, trace_id fell back to an auto-generated UUID unrelated to
anything else, and session_id stayed empty - even when the caller already had
real distributed-tracing instrumentation sending the actual industry-standard
headers for this.

Add a fallback to the W3C Trace Context traceparent header (trace-id
component) and W3C Baggage header (session.id entry), so a request already
carrying real OpenTelemetry trace context correlates litellm's own logs with
the same trace in the caller's observability backend (Datadog, Honeycomb,
Tempo, etc.) instead of getting an unrelated generated id. Precedence is
unchanged for existing sources: explicit litellm headers and the Anthropic
metadata path both still win over this new fallback, which only fires when
neither found anything. trace_id and session_id are resolved independently
here (unlike the existing chain_id mechanism, which uses one shared value for
both), since traceparent and baggage are semantically distinct W3C concepts.

New helpers _trace_id_from_traceparent/_session_id_from_baggage in
litellm_pre_call_utils.py parse the header formats directly (no new
dependency - both are simple fixed-width/delimited strings), wired into
LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers() only
when the corresponding litellm_trace_id/litellm_session_id key isn't already
set by the existing paths.

Verified live against a real proxy: a bare traceparent header produces a log
trace_id exactly matching its trace-id component; a traceparent alongside an
explicit x-litellm-trace-id header (different value) produces a log showing
the explicit header's value, proving precedence.

* fix(logging): reserve trace_id/session_id in JsonFormatter against message-content spoofing

JsonFormatter merges keys parsed from the message body before applying extra
record attributes, and the extra-attributes loop skips a key that's already
present. A caller-controlled log message that happens to parse as JSON/dict
with a "trace_id"/"session_id" key (e.g. the proxy logging a raw request-header
dict) could therefore make the JSON record carry the attacker-supplied value
instead of the real correlation context set via CorrelationContextFilter.

trace_id/session_id are now applied from the LogRecord's own attributes after
message-content parsing, unconditionally overwriting anything the message body
claimed for those two keys.

* style(logging): fix import order (ruff I001) in _logging.py and litellm_logging.py

- _logging.py: import litellm belongs after the stdlib from-imports, grouped
  with the other litellm.* imports, not before them.
- litellm_logging.py: the refactor to Mapping introduced a second, separate
  `from collections.abc import Mapping` instead of merging it into the
  existing `from collections.abc import Callable` import.

Caught by the strict-rule budget gate (ruff-strict-budget.json caps I001 at
0 new violations); both auto-fixed with `ruff check --fix --select I001`.

* style(logging): freeze mutable-collection constructions flagged by LIT002

Five sites in this PR's diff built a mutable list/dict literal instead of a
frozen value: a plain list of optional strings in CorrelationPlainFormatter,
a `kwargs or {}` fallback, a `metadata or {}` fallback, two `[...]` candidate
orderings, and a `dict(headers)` copy feeding a dict comprehension. Each is
build-once/read-only, so this rewrites them as tuples, MappingProxyType, or a
plain conditional `.get()` instead of seeding then reading a fresh mutable
collection - no behavior change, confirmed by the existing test suite.

Caught by the type-discipline budget gate (LIT002 capped at 0 new
violations).

* fix(logging): reserve trace_id/session_id even when no correlation context is active

Live-proxy verification surfaced a gap in the earlier message-content-spoofing
fix (7f390a57fc): that fix only overwrites trace_id/session_id from the
LogRecord's own attribute, so it does nothing for a log line emitted before
CorrelationContextFilter has stamped anything on this record (e.g. the
"Request Headers" debug line, which fires before Logging.__init__() runs for
the request). On such a record, a caller-supplied header literally named
trace_id/session_id still got promoted into the JSON output via the embedded
JSON/dict-repr parser, since there was no genuine value to protect.

Fixed at the source: trace_id/session_id are now excluded unconditionally from
the message-content-parsing promotion step, not just superseded afterward.
Verified live against a real proxy - the exact adversarial request (headers
literally named trace_id/session_id) no longer leaks into any JSON log record.
Added a regression test for this no-active-context variant specifically,
confirmed it fails against the prior commit and passes now.

Also fixes an unrelated basedpyright regression from an earlier rebase's
conflict resolution: litellm/utils.py's `logging_obj` was incorrectly
re-annotated `Final` at its second assignment in function_setup() (it's first
declared `None` a few lines earlier), which basedpyright correctly rejects.

* fix(proxy): stop logging the raw W3C baggage session_id value

_session_id_from_baggage() extracts the caller-controlled session.id entry
verbatim - it isn't sanitized until set_session_id() runs later in
Logging.__init__(). The debug log line for this extraction interpolated the
raw value directly, so a caller could embed terminal control characters or
ANSI escape sequences that forge/alter plaintext log output for anyone
tailing the proxy's logs.

Verified live: a baggage header with an embedded ANSI escape reached the
terminal as a real, unescaped control sequence before this fix. Drops the
value from the log line entirely (the extraction succeeding is enough signal
on its own) rather than sanitizing-then-logging, matching veria-ai's
suggestion. Added a regression test using caplog that fails against the prior
commit and passes now.

* fix(logging): restore consumer context only after stream-failure exception mapping

_map_anthropic_exception/_map_aleph_alpha_exception synchronously log a debug
diagnostic (the raw status code) as part of exception_type()'s mapping.
_handle_stream_fallback_error restored the consumer's outer correlation
context before calling exception_type(), so that diagnostic log line carried
the outer (or empty) trace_id/session_id instead of the failing stream's own -
flagged by Greptile.

Moved the restore to run after mapping completes, matching the same
restore-after-not-before pattern already applied elsewhere in this file for
success/finish_reason handling. Added a regression test that captures the
correlation context live during a mocked exception_type() call; fails against
the prior commit, passes now.

* fix(logging): restore consumer context only after aclose()'s stream close completes

aclose() restored the consumer's outer correlation context as its first
statement, before awaiting the underlying provider stream's own aclose()/
close(). If that close attempt raises, the except branch's debug diagnostic
ran under the already-restored outer context instead of the closing stream's
own trace_id/session_id - flagged by Greptile, same restore-too-early pattern
as the stream-failure fix in f1cf9589d6.

Moved the restore to the end of aclose(), after the close attempt (and its
diagnostic logging) completes. Added a regression test with a fake stream
whose aclose() raises, capturing the correlation context live during the
diagnostic log call; fails against the prior commit, passes now.

* style(logging): satisfy new strict-lint budgets introduced upstream (Final, ANN401, S110, TRY300, kwargs typing)

Rebasing onto litellm_internal_staging pulled in 116 upstream commits that
introduced/tightened several lint gates this PR's own code now trips:

- LIT010 (every local/module-level variable must be Final): added Final
  annotations across _logging.py, litellm_logging.py, streaming_handler.py,
  litellm_pre_call_utils.py, and utils.py. Where a name is genuinely
  reassigned (logging_obj: starts None, later set to the real object) or
  branch-assigned, either restructured into a single ternary expression
  (ordered_candidates) or suppressed with `# rebind-ok: <reason>` matching
  this repo's documented escape hatch.
- LIT011 (parameter mutation): suppressed the two new `data[key] = value`
  writes in litellm_pre_call_utils.py with `# rebind-ok`, matching the
  unsuppressed precedent already used for every other `data[...]` write in
  the same function - `data` is an intentional out-param there.
- ANN001/ANN003/ANN202 (missing parameter/return type annotations): fully
  typed success_handler/_success_handler_body, their async twins, and
  failure_handler/_failure_handler_body/async variants in litellm_logging.py,
  plus function_setup in utils.py (added Rules to its existing TYPE_CHECKING
  block for the rules_obj: Rules annotation).
- ANN401 (explicit Any disallowed): suppressed with `# noqa: ANN401` on the
  handful of genuinely-heterogeneous result/*args/**kwargs parameters, since
  ordinary suppression is this repo's documented path.
- S110 (try/except/pass): added to the existing BLE001 noqa on the one
  best-effort correlation-cleanup try/except this PR added.
- TRY300 (return inside try): moved two `return result` statements into
  `else:` blocks in the retry-fallback paths this PR's own diff touched.
- reportPrivateUsage (basedpyright): renamed the two new
  StandardLoggingPayloadSetup static methods (get_standard_logging_payload_
  trace_id/session_id) to drop their leading underscore, since they're
  genuinely called from a sibling module-level function in the same file.

No behavior change - confirmed by the full existing test suite (819 passed)
plus all four lint gates (ruff format, ruff-strict, type-discipline,
basedpyright) passing clean.

* fix(lint): stop RUF100 flagging noqa suppressions the strict gate needs

CI's plain "ruff check" job uses the default ruff.toml, a narrower config
than ruff-strict.toml (used only by the strict-rule budget gate). ANN401 and
S110 aren't enabled in the default config, so RUF100 (unused-noqa) flagged
the `# noqa: ANN401`/`# noqa: ...,S110` suppressions this PR added as pointless
under that config, even though they're genuinely needed under ruff-strict.toml.

- ANN401: added to ruff.toml's existing `lint.external` list (same mechanism
  already used for C901/TID251, enforced by the strict gate but not by this
  config) - these Any usages are genuinely dynamic/forwarded, so the
  suppression itself is correct and just needed registering.
- S110: fixed the underlying code instead of registering another external
  code - the try/except/pass in
  CustomStreamWrapper._restore_consumer_correlation_context now logs at
  debug level on failure (matching the existing best-effort-cleanup pattern
  in _record_partial_usage_for_failure elsewhere in this file), which
  satisfies S110's own suggestion directly and needs no suppression at all.

Verified against both ruff.toml and ruff-strict.toml directly, plus all
three other gates (ruff format, type-discipline, basedpyright) and the full
test suite (821 passed).

* fix(lint): scope the ANN401 exemption to file level instead of a repo-wide noqa

Ruff has no per-line-scoped way to register a noqa code across configs (that
requires the default ruff.toml's lint.external list, which is repo-wide in
scope even though the noqa itself is per-line). Since ruff does support
file-level exemptions via per-file-ignores, and ANN401 only needed exempting
in exactly two files, moved the exemption there instead:

- ruff-strict.toml: added [lint.per-file-ignores] disabling ANN401 for
  litellm_logging.py and utils.py specifically, with a comment explaining
  why (heterogeneous response/forwarded-args parameters with no fitting
  concrete type - already verified by trying CostResponseTypes and hitting
  a real basedpyright mismatch).
- ruff.toml: reverted the ANN401 entry from lint.external - no longer
  needed, since there's no `# noqa: ANN401` left anywhere for RUF100 to
  second-guess.
- Removed the now-redundant `# noqa: ANN401` from the 10 affected
  parameters in both files, keeping the existing kwargs-ok reasons and
  adding a short inline comment on the `result`/`*args` lines pointing at
  the ruff-strict.toml exemption for context.

Verified against both configs directly (ANN401 clean under ruff-strict.toml
for these files, RUF100 clean under the default config), all four gates
(ruff format, ruff-strict, type-discipline, basedpyright), and the full
test suite (821 passed).

* fix(logging): redact credential-shaped trace_id/session_id before stamping log records

CorrelationContextFilter stamps trace_id/session_id onto a LogRecord after
SecretRedactionFilter has already run, so a caller-controlled value (e.g. via
x-litellm-trace-id or a W3C baggage header) that happens to look like a real
credential reached JSON and plaintext logs unredacted. Apply the same
credential redaction already used elsewhere in this module at
_sanitize_correlation_id(), the single choke point both set_trace_id() and
set_session_id() route through, so every caller-facing entry point is covered
without depending on filter ordering.

* fix(logging): restore correlation context when a stream's max-duration timeout fires

CustomStreamWrapper.__anext__() called _check_max_streaming_duration() before
entering its try block, so the litellm.Timeout it raises bypassed the except
Exception -> _handle_stream_fallback_error path entirely, leaking the timed-out
stream's own trace_id/session_id into whatever the consumer's task logs next.
Move the check inside the try so it flows through the same restoration path
every other stream failure already uses.

* test(streaming): make dispatch_failure_handlers mock awaitable for the async max-duration test

Moving _check_max_streaming_duration() inside __anext__()'s try block (prior
commit) means a max-duration Timeout now dispatches failure handlers through
the same path every other stream failure already uses, instead of bypassing
it entirely. dispatch_failure_handlers is async on the real Logging class;
the test's plain MagicMock logging_obj made asyncio.create_task() choke on a
non-coroutine return value once that path actually got exercised.

---------

Co-authored-by: Deepanshu <deepanshu.lulla@alpha-sense.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>

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

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

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
+        ordered_candidates: Final[tuple[Any, Any, Any, Any]] = (
-        # `litellm_trace_id` is an internal litellm param
+            (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:
+        metadata: Final = litellm_params.get("metadata")
-            return str(dynamic_litellm_trace_id)
+        metadata_session_id: Final = metadata.get("session_id") if metadata else None
-        # 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")
         if metadata_session_id:
             return str(metadata_session_id)
-        if metadata_trace_id:
+        return logging_obj.litellm_session_id
-            return str(metadata_trace_id)
-        return logging_obj.litellm_trace_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,
             ),

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

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

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

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
+    original_function: str,
-):  # just run once to check if user wants to send their data anywhere - PostHog/Sentry/Slack/etc.
+    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)
 

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
 

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():
+def testget_standard_logging_payload_trace_id():
-    """Test _get_standard_logging_payload_trace_id with different input scenarios"""
+    """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

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("")

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

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("")

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"},

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)
+

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("")