From 628585ae613aeb7b30049fb3927ec264f885cacc Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Tue, 28 Jul 2026 22:14:59 -0400 Subject: [PATCH] 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. --- .../test_litellm_logging.py | 166 +++++ tests/test_litellm/test_logging.py | 674 ------------------ tests/test_litellm/test_utils.py | 129 ++++ 3 files changed, 295 insertions(+), 674 deletions(-) diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 4a32e80acf2..1fd5b2d5571 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -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("") diff --git a/tests/test_litellm/test_logging.py b/tests/test_litellm/test_logging.py index 04df571fe43..85719db5cc4 100644 --- a/tests/test_litellm/test_logging.py +++ b/tests/test_litellm/test_logging.py @@ -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("") diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index b23d3333ea7..db5b51215aa 100644 --- a/tests/test_litellm/test_utils.py +++ b/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, @@ -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("")