scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 35

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

Browse source 
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.
This commit is contained in:
Deepanshu 2026-07-28 22:16:20 -04:00
parent 037fb3f11f
commit 9f3a20f4b2
3 changed files with 546 additions and 64 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

27
litellm/litellm_core_utils/streaming_handler.py
View file

@ -250,21 +250,32 @@ class CustomStreamWrapper:
restore = getattr(logging_obj, method_name, None)
if restore is not None:
restore()
except Exception:
except Exception: # noqa: BLE001 # best-effort cleanup; must never raise into the caller's actual stream handling regardless of what's wrong with logging_obj
pass
def __del__(self) -> None:
"""Best-effort correlation-context cleanup for an abandoned stream.
"""Best-effort correlation-context cleanup for an abandoned async stream.
If the caller never fully consumes the stream - stops early, drops the
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. This
is a best-effort fallback, not a guarantee: __del__ timing is
_restore_consumer_correlation_context() is called from ever run. For a
sync stream (wrapper()), this is a no-op in practice: wrapper() already
restores unconditionally before ever handing the stream back, so there
is nothing left to clean up here.
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. guarded=True additionally
ensures it never clobbers a different, still-active call's context if
this fires late.
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)

119
litellm/utils.py
View file

@ -728,29 +728,47 @@ def _remove_thought_signatures_from_messages(messages: List, thought_signature_s
return processed_messages
def _restore_correlation_context_if_supported(logging_obj: Any) -> None:
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.
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 = getattr(logging_obj, "_restore_correlation_context", None)
if restore is not None:
restore()
def _is_streaming_response_for_correlation(result: Any) -> bool:
def _is_streaming_response_for_correlation(result: object) -> bool:
"""True if `result` is a lazy stream wrapper rather than an already-complete response.
wrapper()/wrapper_async() 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. The corresponding terminal handler (async_success_handler, dispatched
once the full stream is actually assembled) is what restores it once
streaming genuinely finishes.
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 and restores
unconditionally instead: a plain OS thread has no such per-call isolation,
and a thread pool's worker threads *are* recycled across unrelated
requests, so leaving a sync stream's contextvars "open" indefinitely can
permanently misattribute every subsequent log line on that thread to an
abandoned request.
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
@ -765,6 +783,7 @@ def function_setup(
verbose_logger.warning(
"`litellm.set_verbose` is deprecated. Please set `os.environ['LITELLM_LOG'] = 'DEBUG'` for debug logs."
)
logging_obj: LiteLLMLoggingObject | None = None
try:
global callback_list, add_breadcrumb, user_logger_fn, Logging
@ -1101,9 +1120,8 @@ def function_setup(
# 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.
_logging_obj_for_correlation_cleanup = locals().get("logging_obj")
if _logging_obj_for_correlation_cleanup is not None:
_restore_correlation_context_if_supported(_logging_obj_for_correlation_cleanup)
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
@ -1318,25 +1336,6 @@ def client(original_function):
@wraps(original_function)
def wrapper(*args, **kwargs):
# Restore trace_id/session_id contextvars to their pre-call value once this
# call (in this task/thread) is fully done - see request_correlation_in_logs.
# _wrapper_body rebinds its own local `kwargs` name via function_setup(), so
# a dict-identity trick doesn't work here; it stashes logging_obj into this
# holder directly instead.
_correlation_logging_obj_holder: dict = {}
_correlation_result_holder: dict = {}
try:
result = _wrapper_body(args, kwargs, _correlation_logging_obj_holder)
_correlation_result_holder["result"] = result
return result
finally:
_correlation_logging_obj = _correlation_logging_obj_holder.get("logging_obj")
if _correlation_logging_obj is not None and not _is_streaming_response_for_correlation(
_correlation_result_holder.get("result")
):
_restore_correlation_context_if_supported(_correlation_logging_obj)
def _wrapper_body(args, kwargs, _correlation_logging_obj_holder):
# DO NOT MOVE THIS. It always needs to run first
# Check if this is an async function. If so only execute the async function
call_type = original_function.__name__
@ -1380,7 +1379,6 @@ def client(original_function):
try:
if logging_obj is None:
logging_obj, kwargs = function_setup(original_function.__name__, rules_obj, start_time, *args, **kwargs)
_correlation_logging_obj_holder["logging_obj"] = logging_obj
# 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"
@ -1636,27 +1634,19 @@ def client(original_function):
) # DO NOT MAKE THREADED - router retry fallback relies on this!
raise e
finally:
# Restore trace_id/session_id contextvars to their pre-call value once
# this call is fully done, in every case (success, retried-and-returned,
# or re-raised) - see request_correlation_in_logs. Always safe to restore
# here, even for a stream: unlike wrapper_async(), this is a plain
# synchronous function whose own thread never re-enters user code after
# this point without the caller making a brand new call, and leaving this
# ambient on a thread a pool may hand to a *different*, unrelated future
# call would misattribute that call's logs to this one.
_restore_correlation_context_if_supported(logging_obj)
@wraps(original_function)
async def wrapper_async(*args, **kwargs):
# Restore trace_id/session_id contextvars to their pre-call value once this
# call (in this task) is fully done - see request_correlation_in_logs.
# _wrapper_async_body rebinds its own local `kwargs` name via
# function_setup(), so a dict-identity trick doesn't work here; it
# stashes logging_obj into this holder directly instead.
_correlation_logging_obj_holder: dict = {}
_correlation_result_holder: dict = {}
try:
result = await _wrapper_async_body(args, kwargs, _correlation_logging_obj_holder)
_correlation_result_holder["result"] = result
return result
finally:
_correlation_logging_obj = _correlation_logging_obj_holder.get("logging_obj")
if _correlation_logging_obj is not None and not _is_streaming_response_for_correlation(
_correlation_result_holder.get("result")
):
_restore_correlation_context_if_supported(_correlation_logging_obj)
async def _wrapper_async_body(args, kwargs, _correlation_logging_obj_holder):
print_args_passed_to_litellm(original_function, args, kwargs)
start_time = datetime.datetime.now()
result = None
@ -1684,7 +1674,6 @@ def client(original_function):
# 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"
_correlation_logging_obj_holder["logging_obj"] = logging_obj
modified_kwargs = await async_pre_call_deployment_hook(kwargs, call_type)
if modified_kwargs is not None:
@ -1909,7 +1898,8 @@ 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)
return result
except Exception:
pass
elif (
@ -1922,7 +1912,8 @@ def client(original_function):
args[0] = context_window_fallback_dict[model] # type: ignore
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 {}
@ -1939,7 +1930,8 @@ 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)
return result
except Exception:
pass
@ -1951,6 +1943,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 = getattr(sys.modules[__name__], "get_coroutine_checker")
is_coroutine = get_coroutine_checker().is_async_callable(original_function)

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

@ -14,6 +14,7 @@ import traceback
from typing import Optional
import litellm
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,
@ -3355,3 +3356,466 @@ async def test_transport_read_error_before_finish_reason_raises(logging_obj: Log
if chunk.choices and chunk.choices[0].finish_reason
]
assert fabricated_finish_reasons == []
def test_sync_streaming_completion_restores_context_immediately_on_return(monkeypatch):
"""wrapper() (the sync entry point) must restore the originating thread's
correlation context the instant a streaming completion() call returns,
before the caller ever starts iterating - unlike wrapper_async(), it
cannot safely leave the contextvars "open" across the caller's own,
unboundedly long iteration.
A plain OS thread has no per-call isolation the way an asyncio Task does:
if a sync stream is abandoned mid-iteration (client disconnect, an early
break, an exception), nothing ever restores it, and a thread pool can hand
that same, now-poisoned thread to a completely unrelated future call,
which then inherits the wrong ids as its own "pre-call" baseline and
propagates them forward indefinitely. Restoring unconditionally in
wrapper()'s own finally, before the stream is ever handed to the caller,
closes that hole entirely for the sync path."""
monkeypatch.setattr(litellm, "request_correlation_in_logs", True)
trace_id_var.set("outer-trace-stream")
session_id_var.set("outer-session-stream")
try:
response = litellm.completion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "hi"}],
mock_response="Hello there!",
stream=True,
litellm_session_id="streaming-call-session",
num_retries=0,
)
# Already restored - before the caller has iterated a single chunk.
assert session_id_var.get() == "outer-session-stream"
assert trace_id_var.get() == "outer-trace-stream"
for _ in response:
pass
# Consuming the stream afterward must not disturb the already-restored
# outer context either.
assert session_id_var.get() == "outer-session-stream"
assert trace_id_var.get() == "outer-trace-stream"
finally:
trace_id_var.set("")
session_id_var.set("")
def test_abandoned_sync_stream_does_not_contaminate_a_later_call_on_the_same_thread(monkeypatch):
"""Reproduces the maintainer-reported blocking bug: request A starts a sync
stream, consumes one chunk, and abandons it (client disconnect / early
break / uncaught exception - never exhausts, never calls aclose()).
Request B is a completely separate, non-streaming call that later runs on
the *same* worker thread (a real ThreadPoolExecutor with a single worker,
forcing thread reuse, exactly like a WSGI/thread-pool-backed sync server).
Before the fix, A's contextvars stayed "open" (wrapper() skipped restoring
for streams) with no deterministic point that ever closed them, so B
inherited A's ids as its own pre-call baseline and then "restored" back to
that poison when it finished - permanently wrong-attributing every log
line on this thread to request A from then on, including B's own and any
unrelated future work. Restoring unconditionally in wrapper()'s own
finally, before a sync stream is ever handed to the caller, closes this:
B never sees A's ids in the first place."""
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,
)
# wrapper() has already restored to B's own pre-call snapshot by the
# time completion() returns - read it here, right after, to see
# exactly what B (wrongly, if the bug is present) inherited as its
# own baseline and just restored the thread back to.
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()
# The critical assertion: B must not have inherited request A's ids as
# its own pre-call baseline and restored the thread back to them - it
# must see whatever this thread started with before any of this (the
# ContextVar default), not request A's abandoned trace/session id.
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_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("")