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

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

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

litellm/_logging.py

@@ -40,30 +40,6 @@ def set_trace_id(trace_id: str) -> "contextvars.Token[str]":
     return trace_id_var.set(_sanitize_correlation_id(trace_id))
 
 
-def reset_session_id(token: "contextvars.Token[str]") -> None:
-    """Restore session_id_var to its pre-call value.
-
-    Best-effort: swallows errors since this is observability plumbing, not
-    call-correctness - a failed reset must never break the actual LLM call.
-    """
-    try:
-        session_id_var.reset(token)
-    except (ValueError, RuntimeError):
-        pass
-
-
-def reset_trace_id(token: "contextvars.Token[str]") -> None:
-    """Restore trace_id_var to its pre-call value.
-
-    Best-effort: swallows errors since this is observability plumbing, not
-    call-correctness - a failed reset must never break the actual LLM call.
-    """
-    try:
-        trace_id_var.reset(token)
-    except (ValueError, RuntimeError):
-        pass
-
-
 if set_verbose is True:
     logging.warning(
         "`litellm.set_verbose` is deprecated. Please set `os.environ['LITELLM_LOG'] = 'DEBUG'` for debug logs."

litellm/litellm_core_utils/litellm_logging.py

@@ -39,10 +39,10 @@ from litellm import (
 from litellm._logging import (
     _is_debugging_on,
     _redact_string,
-    reset_session_id,
+    session_id_var,
-    reset_trace_id,
     set_session_id,
     set_trace_id,
+    trace_id_var,
     verbose_logger,
 )
 from litellm.exceptions import (
@@ -356,10 +356,16 @@ class Logging(LiteLLMLoggingBaseClass):
         self.litellm_call_id = litellm_call_id
         self.litellm_trace_id: str = litellm_trace_id if litellm_trace_id else str(uuid.uuid4())
 
-        self._trace_id_token = set_trace_id(self.litellm_trace_id)
+        # Capture the pre-call *value* (not a contextvars.Token) so restoration works
+        # even if this attempt's own logging ends up dispatched onto a different
+        # asyncio Task/context (e.g. via asyncio.create_task or the logging worker) -
+        # a Token can only be reset in the exact Context where it was created.
+        self._pre_call_trace_id: str = trace_id_var.get()
+        self._pre_call_session_id: str = session_id_var.get()
+        set_trace_id(self.litellm_trace_id)
         _sid = (kwargs or {}).get("litellm_session_id")
         self.litellm_session_id: str = str(_sid) if _sid else ""
-        self._session_id_token = set_session_id(self.litellm_session_id)
+        set_session_id(self.litellm_session_id)
         self._correlation_context_restored = False
 
         self.function_id = function_id
@@ -1978,21 +1984,27 @@ class Logging(LiteLLMLoggingBaseClass):
         return
 
     def _restore_correlation_context(self) -> None:
-        """Reset trace_id/session_id contextvars to their pre-call value.
+        """Restore trace_id/session_id contextvars to their pre-call value.
 
         Without this, a nested LiteLLM call sharing the same asyncio Task as an
         outer request (e.g. a guardrail's own LLM-as-judge call, an MCP sampling
         call) would leave the outer request's subsequent log lines stamped with
         the nested call's trace_id/session_id instead of its own.
 
+        Uses a plain set() of the captured pre-call value rather than
+        contextvars.Token-based reset(), since this handler can end up running
+        in a different asyncio Task/context than __init__ did (e.g. dispatched
+        via asyncio.create_task or the logging worker) - reset() only works in
+        the exact Context a Token was created in and raises otherwise.
+
         Idempotent - safe to call from every terminal handler regardless of
         which one ends up firing for this attempt.
         """
         if self._correlation_context_restored:
             return
         self._correlation_context_restored = True
-        reset_trace_id(self._trace_id_token)
+        set_trace_id(self._pre_call_trace_id)
-        reset_session_id(self._session_id_token)
+        set_session_id(self._pre_call_session_id)
 
     def success_handler(self, result=None, start_time=None, end_time=None, cache_hit=None, **kwargs):
         """Restores trace_id/session_id contextvars once this attempt's own success

litellm/proxy/management_endpoints/credential_migration.py

@@ -130,9 +130,7 @@ def classify_value(value: object, key: str = "scan") -> ValueClass:
         return "plaintext"
     if value.startswith(_V2_GCM_PREFIX):
         return "migrated"
-    decrypted = decrypt_value_helper(
+    decrypted = decrypt_value_helper(value=value, key=key, exception_type="debug", return_original_value=False)
-        value=value, key=key, exception_type="debug", return_original_value=False
-    )
     if decrypted is None:
         # Did not decrypt under nacl and has no v2 marker: legacy plaintext.
         return "plaintext"
@@ -151,9 +149,7 @@ def reencrypt_value(value: object, key: str = "migrate") -> object:
         return value
     if value.startswith(_V2_GCM_PREFIX):
         return value  # idempotent: already migrated
-    decrypted = decrypt_value_helper(
+    decrypted = decrypt_value_helper(value=value, key=key, exception_type="debug", return_original_value=False)
-        value=value, key=key, exception_type="debug", return_original_value=False
-    )
     if decrypted is None:
         # Either legacy plaintext (no ciphertext to migrate) or corrupt. Either
         # way, do not overwrite — preserve the value as stored.
@@ -161,9 +157,7 @@ def reencrypt_value(value: object, key: str = "migrate") -> object:
     return encrypt_value_helper(decrypted)
 
 
-def reencrypt_selective_dict(
+def reencrypt_selective_dict(data: dict[str, object], sensitive_keys: list[str]) -> dict[str, object]:
-    data: dict[str, object], sensitive_keys: list[str]
-) -> dict[str, object]:
     """Return a copy of ``data`` with only ``sensitive_keys`` re-encrypted.
 
     Non-sensitive fields (e.g. ``base_url``, ``connection_id``) are left as-is.
@@ -212,9 +206,7 @@ async def _migrate_config_settings_row(
     dict with selected sensitive fields (vantage_settings / cloudzero_settings).
     """
     report = LocationReport(location=param_name)
-    record = await prisma_client.db.litellm_config.find_unique(
+    record = await prisma_client.db.litellm_config.find_unique(where={"param_name": param_name})
-        where={"param_name": param_name}
-    )
     if record is None or record.param_value is None:
         return report
 
@@ -266,9 +258,7 @@ async def _migrate_sso_config(prisma_client: object, dry_run: bool) -> LocationR
     every present string field.
     """
     report = LocationReport(location="sso_config")
-    record = await prisma_client.db.litellm_ssoconfig.find_unique(
+    record = await prisma_client.db.litellm_ssoconfig.find_unique(where={"id": "sso_config"})
-        where={"id": "sso_config"}
-    )
     if record is None or record.sso_settings is None:
         return report
 
@@ -344,9 +334,7 @@ async def _migrate_callback_vars_table(
     rows = await table.find_many()
     for row in rows or []:
         metadata = getattr(row, "metadata", None)
-        if not isinstance(metadata, dict) or (
+        if not isinstance(metadata, dict) or ("logging" not in metadata and "callback_settings" not in metadata):
-            "logging" not in metadata and "callback_settings" not in metadata
-        ):
             continue
 
         # Classify every callback-var value directly (strip the litellm_enc::
@@ -534,9 +522,7 @@ async def _scan_config_env_vars(prisma_client: object) -> LocationReport:
     """Scan the ``environment_variables`` config row (``param_value`` dict)."""
     report = LocationReport(location="config_environment_variables")
     try:
-        record = await prisma_client.db.litellm_config.find_unique(
+        record = await prisma_client.db.litellm_config.find_unique(where={"param_name": "environment_variables"})
-            where={"param_name": "environment_variables"}
-        )
     except Exception as e:  # pragma: no cover - defensive
         verbose_proxy_logger.debug("scan: config env vars unavailable: %s", str(e))
         return report
@@ -557,11 +543,7 @@ async def _scan_covered_tables(prisma_client: object) -> list[LocationReport]:
     """Read-only classification of every rotation-covered table. No writes."""
     reports: list[LocationReport] = []
     for location, db_attr, json_cols, scalar_cols in _COVERED_TABLE_SPECS:
-        reports.append(
+        reports.append(await _scan_one_table(prisma_client, location, db_attr, json_cols, scalar_cols))
-            await _scan_one_table(
-                prisma_client, location, db_attr, json_cols, scalar_cols
-            )
-        )
     reports.append(await _scan_config_env_vars(prisma_client))
     return reports
 
@@ -575,9 +557,7 @@ _VANTAGE_SENSITIVE = ["api_key", "integration_token"]
 _CLOUDZERO_SENSITIVE = ["api_key"]
 
 
-async def _migrate_covered_tables(
+async def _migrate_covered_tables(prisma_client: object, user_api_key_dict: object) -> list[LocationReport]:
-    prisma_client: object, user_api_key_dict: object
-) -> list[LocationReport]:
     """Re-encrypt the tables already covered by ``_rotate_master_key`` (model
     table, credentials, MCP credential/env tables, config environment_variables)
     by running that orchestrator in *same-key* mode. With the AES gate on, the
@@ -597,8 +577,7 @@ async def _migrate_covered_tables(
     current_key = _get_salt_key()
     if current_key is None:
         raise RuntimeError(
-            "Cannot migrate covered tables: no salt key / master key is set. "
+            "Cannot migrate covered tables: no salt key / master key is set. Set LITELLM_SALT_KEY before migrating."
-            "Set LITELLM_SALT_KEY before migrating."
         )
     await _rotate_master_key(
         prisma_client=cast("PrismaClient", prisma_client),
@@ -648,19 +627,9 @@ async def migrate_encryption(
 
     # Net-new walkers (items 3, 4, 11, 12, 13).
     report.add(await _migrate_callback_vars_table(prisma_client, "team", dry_run))
-    report.add(
+    report.add(await _migrate_callback_vars_table(prisma_client, "verification_token", dry_run))
-        await _migrate_callback_vars_table(prisma_client, "verification_token", dry_run)
+    report.add(await _migrate_config_settings_row(prisma_client, "vantage_settings", _VANTAGE_SENSITIVE, dry_run))
-    )
+    report.add(await _migrate_config_settings_row(prisma_client, "cloudzero_settings", _CLOUDZERO_SENSITIVE, dry_run))
-    report.add(
-        await _migrate_config_settings_row(
-            prisma_client, "vantage_settings", _VANTAGE_SENSITIVE, dry_run
-        )
-    )
-    report.add(
-        await _migrate_config_settings_row(
-            prisma_client, "cloudzero_settings", _CLOUDZERO_SENSITIVE, dry_run
-        )
-    )
     report.add(await _migrate_sso_config(prisma_client, dry_run))
 
     return report
@@ -683,20 +652,10 @@ async def check_encryption(prisma_client: object) -> MigrationReport:
 
     # Net-new walker locations, in dry-run (read-only) mode.
     report.add(await _migrate_callback_vars_table(prisma_client, "team", dry_run=True))
+    report.add(await _migrate_callback_vars_table(prisma_client, "verification_token", dry_run=True))
+    report.add(await _migrate_config_settings_row(prisma_client, "vantage_settings", _VANTAGE_SENSITIVE, dry_run=True))
     report.add(
-        await _migrate_callback_vars_table(
+        await _migrate_config_settings_row(prisma_client, "cloudzero_settings", _CLOUDZERO_SENSITIVE, dry_run=True)
-            prisma_client, "verification_token", dry_run=True
-        )
-    )
-    report.add(
-        await _migrate_config_settings_row(
-            prisma_client, "vantage_settings", _VANTAGE_SENSITIVE, dry_run=True
-        )
-    )
-    report.add(
-        await _migrate_config_settings_row(
-            prisma_client, "cloudzero_settings", _CLOUDZERO_SENSITIVE, dry_run=True
-        )
     )
     report.add(await _migrate_sso_config(prisma_client, dry_run=True))
     return report

tests/test_litellm/litellm_core_utils/test_litellm_logging.py

@@ -3037,8 +3037,12 @@ async def test_async_failure_handler_runs_callbacks_and_restores_correlation_con
     dummy_logger = DummyLogger()
     dummy_logger.async_log_failure_event = AsyncMock()
 
-    trace_id_var.set("pre-existing-trace")
+    # logging_obj is constructed by the fixture (before this line runs), so it
-    session_id_var.set("pre-existing-session")
+    # already captured whatever was ambient at that point as its own pre-call
+    # value - assert restoration lands back on THAT captured value, not a
+    # value set here (which would be too late to affect __init__'s snapshot).
+    trace_id_var.set("mutated-during-call")
+    session_id_var.set("mutated-during-call")
     try:
         with patch.object(
             logging_obj,
@@ -3052,8 +3056,9 @@ async def test_async_failure_handler_runs_callbacks_and_restores_correlation_con
 
         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 trace_id_var.get() == logging_obj._pre_call_trace_id
-        assert session_id_var.get() == "pre-existing-session"
+        assert session_id_var.get() == logging_obj._pre_call_session_id
+        assert trace_id_var.get() != "mutated-during-call"
     finally:
         trace_id_var.set("")
         session_id_var.set("")

tests/test_litellm/test_logging.py

@@ -630,3 +630,53 @@ def test_set_session_id_bounds_length():
         assert len(session_id_var.get()) == 256
     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("")