From cc40bdb609a63a99ed6c3f1d1c17ade9cd750fac Mon Sep 17 00:00:00 2001
From: Deepanshu <deepanshu.lulla@alpha-sense.com>
Date: Thu, 23 Jul 2026 17:41:31 -0400
Subject: [PATCH] test(logging): cover async_failure_handler's
 correlation-context restore

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

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 5bffda126fe..b5e214aac51 100644
--- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py
+++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py
@@ -3018,6 +3018,47 @@ def test_failure_handler_runs_sync_callbacks_for_non_pass_through_requests(
     dummy_logger.log_failure_event.assert_called_once()
 
 
+@pytest.mark.asyncio
+async def test_async_failure_handler_runs_callbacks_and_restores_correlation_context(logging_obj):
+    """await logging_obj.async_failure_handler(...) must dispatch async failure callbacks
+    and, once its own body completes, restore trace_id/session_id contextvars via
+    _restore_correlation_context() (the fix for the nested-call context leak)."""
+    from litellm._logging import session_id_var, trace_id_var
+    from litellm.integrations.custom_logger import CustomLogger
+
+    class DummyLogger(CustomLogger):
+        pass
+
+    logging_obj.call_type = "acompletion"
+    logging_obj.stream = False
+    logging_obj.model_call_details["litellm_params"] = {}
+    logging_obj.litellm_params = {}
+
+    dummy_logger = DummyLogger()
+    dummy_logger.async_log_failure_event = AsyncMock()
+
+    trace_id_var.set("pre-existing-trace")
+    session_id_var.set("pre-existing-session")
+    try:
+        with patch.object(
+            logging_obj,
+            "get_combined_callback_list",
+            return_value=[dummy_logger],
+        ):
+            await logging_obj.async_failure_handler(
+                exception=Exception("test error"),
+                traceback_exception="",
+            )
+
+        dummy_logger.async_log_failure_event.assert_called_once()
+        assert logging_obj._correlation_context_restored is True
+        assert trace_id_var.get() == "pre-existing-trace"
+        assert session_id_var.get() == "pre-existing-session"
+    finally:
+        trace_id_var.set("")
+        session_id_var.set("")
+
+
 def test_merge_hidden_params_from_response_into_metadata_populates_metadata():
     """Streaming completion path should mirror non-stream: metadata.hidden_params from response."""
     from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj