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

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

Browse source 
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.
This commit is contained in:
Deepanshu 2026-07-23 18:00:28 -04:00
parent cc40bdb609
commit 43c164aa32
5 changed files with 94 additions and 92 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

24
litellm/_logging.py
View file

@ -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."

26
litellm/litellm_core_utils/litellm_logging.py
View file

@ -39,10 +39,10 @@ from litellm import (
from litellm._logging import (
_is_debugging_on,
_redact_string,
reset_session_id,
reset_trace_id,
session_id_var,
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)
reset_session_id(self._session_id_token)
set_trace_id(self._pre_call_trace_id)
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

73
litellm/proxy/management_endpoints/credential_migration.py
View file

@ -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(
value=value, key=key, exception_type="debug", return_original_value=False
)
decrypted = decrypt_value_helper(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(
value=value, key=key, exception_type="debug", return_original_value=False
)
decrypted = decrypt_value_helper(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(
data: dict[str, object], sensitive_keys: list[str]
) -> dict[str, object]:
def reencrypt_selective_dict(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(
where={"param_name": param_name}
)
record = await prisma_client.db.litellm_config.find_unique(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(
where={"id": "sso_config"}
)
record = await prisma_client.db.litellm_ssoconfig.find_unique(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 (
"logging" not in metadata and "callback_settings" not in metadata
):
if not isinstance(metadata, dict) or ("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(
where={"param_name": "environment_variables"}
)
record = await prisma_client.db.litellm_config.find_unique(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(
await _scan_one_table(
prisma_client, location, db_attr, json_cols, scalar_cols
)
)
reports.append(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(
prisma_client: object, user_api_key_dict: object
) -> list[LocationReport]:
async def _migrate_covered_tables(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. "
"Set LITELLM_SALT_KEY before migrating."
"Cannot migrate covered tables: no salt key / master key is set. 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(
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_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_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(
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
)
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

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

@ -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")
session_id_var.set("pre-existing-session")
# logging_obj is constructed by the fixture (before this line runs), so it
# 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 session_id_var.get() == "pre-existing-session"
assert trace_id_var.get() == logging_obj._pre_call_trace_id
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("")

50
tests/test_litellm/test_logging.py
View file

@ -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("")