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 36

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

Browse source 
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.
This commit is contained in:
Deepanshu 2026-07-28 22:14:59 -04:00
parent 5fe7041c8d
commit 628585ae61
3 changed files with 295 additions and 674 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -12,6 +12,7 @@ sys.path.insert(
import time
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
@ -4198,3 +4199,168 @@ 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_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("")

674
tests/test_litellm/test_logging.py
View file

@ -31,7 +31,6 @@ from litellm._logging import (
verbose_router_logger,
)
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.types.utils import StandardLoggingPayload
@ -481,64 +480,6 @@ async def test_contextvar_isolation_between_tasks():
assert results["B"] == "trace-for-B"
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_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_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)
@ -623,63 +564,6 @@ def test_plain_formatter_unchanged_when_flag_disabled(monkeypatch):
session_id_var.set("")
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
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."""
@ -701,561 +585,3 @@ def test_set_session_id_bounds_length():
finally:
session_id_var.reset(token)
@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("")
@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("")
def test_streaming_completion_does_not_reset_context_before_iteration(monkeypatch):
"""wrapper() must not restore the originating thread's correlation context
the instant a streaming completion() call returns a lazy CustomStreamWrapper -
the caller hasn't started iterating it yet, and log lines emitted while doing
so (in the same thread) should still show this call's own ids, not whatever
was ambient before the call."""
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,
)
# The call returned a lazy stream; nothing has been iterated yet, but
# this call's own ids must still be the ambient ones right now.
assert session_id_var.get() == "streaming-call-session"
for _ in response:
pass
# Once the stream is genuinely exhausted, the *consuming* thread's
# context (this test's own) must be restored - not just some detached
# executor-thread context the terminal success_handler happens to run in.
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("")
@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.
"""
from litellm.litellm_core_utils.litellm_logging import Logging
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."""
from litellm.litellm_core_utils.litellm_logging import Logging
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."""
from litellm.litellm_core_utils.litellm_logging import Logging
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."""
from litellm.litellm_core_utils.litellm_logging import Logging
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."""
from litellm.litellm_core_utils.litellm_logging import Logging
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."""
from litellm.litellm_core_utils.litellm_logging import Logging
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__."""
from litellm.litellm_core_utils.litellm_logging import Logging
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("")

129
tests/test_litellm/test_utils.py
View file

@ -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,
@ -5015,3 +5023,124 @@ async def test_builtin_string_callback_registers_when_subclass_already_active(
)
assert any(type(cb) is S3Logger for cb in litellm._async_success_callback)
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("")