fix(router): leave callback-started calls unmarked and stop client flags from turning mirroring off

Logging handlers now run outside the mirrored-call context, so an LLM call a success or failure callback starts while logging a mirrored call is not flagged as mirrored. The guard against mirroring a mirror reads the same context variable instead of request metadata, so a client-sent is_silent_experiment no longer skips the mirrors. The unused module-level get_standard_logging_metadata stops copying resolved keys from raw metadata.
This commit is contained in:
mateo-berri 2026-10-28 18:10:59 -07:00
parent 054f220e53
commit e4bf3f356b
6 changed files with 160 additions and 47 deletions

View file

@ -139,6 +139,16 @@ def in_mirrored_call() -> bool:
return _mirrored_call.get()
@contextmanager
def outside_mirrored_call() -> Generator[None]:
"""Logging a mirrored call: the calls its callbacks start are not mirrors themselves."""
token: Final = _mirrored_call.set(False)
try:
yield
finally:
_restore(_mirrored_call, token)
@contextmanager
def pinned_billing_time(moment: datetime) -> Generator[None]:
"""Price every rate lookup inside this block at ``moment`` rather than at each one's own clock read."""

View file

@ -23,7 +23,7 @@ from httpx import Response, Timeout
from pydantic import BaseModel, JsonValue, TypeAdapter
import litellm
from litellm._internal_context import in_mirrored_call, post_response_phase
from litellm._internal_context import in_mirrored_call, outside_mirrored_call, post_response_phase
from litellm._logging import (
is_debugging_on,
redact_string,
@ -3246,7 +3246,7 @@ class Logging(LiteLLMLoggingBaseClass):
"""Restores trace_id/session_id contextvars once this attempt's own success
logging (including any nested calls its callbacks trigger) is fully done."""
try:
with post_response_phase():
with post_response_phase(), outside_mirrored_call():
return self._success_handler_body(
result=result,
start_time=start_time,
@ -3746,7 +3746,7 @@ class Logging(LiteLLMLoggingBaseClass):
"""Restores trace_id/session_id contextvars once this attempt's own success
logging (including any nested calls its callbacks trigger) is fully done."""
try:
with post_response_phase():
with post_response_phase(), outside_mirrored_call():
return await self._async_success_handler_body(
result=result, start_time=start_time, end_time=end_time, cache_hit=cache_hit, **kwargs
)
@ -4263,13 +4263,14 @@ class Logging(LiteLLMLoggingBaseClass):
``for_async_request`` marks the sync run of a failure whose async run also logs it, so a
CustomLogger that gets ``async_log_failure_event`` from that async run skips its sync hook."""
try:
return self._failure_handler_body(
exception=exception,
traceback_exception=traceback_exception,
start_time=start_time,
end_time=end_time,
for_async_request=for_async_request,
)
with outside_mirrored_call():
return self._failure_handler_body(
exception=exception,
traceback_exception=traceback_exception,
start_time=start_time,
end_time=end_time,
for_async_request=for_async_request,
)
finally:
self._restore_correlation_context()
@ -4462,12 +4463,13 @@ class Logging(LiteLLMLoggingBaseClass):
"""Restores trace_id/session_id contextvars once this attempt's own failure
logging (including any nested calls its callbacks trigger) is fully done."""
try:
return await self._async_failure_handler_body(
exception=exception,
traceback_exception=traceback_exception,
start_time=start_time,
end_time=end_time,
)
with outside_mirrored_call():
return await self._async_failure_handler_body(
exception=exception,
traceback_exception=traceback_exception,
start_time=start_time,
end_time=end_time,
)
finally:
self._restore_correlation_context()
@ -7714,9 +7716,8 @@ def get_standard_logging_metadata(
)
if isinstance(metadata, dict):
# Update the clean_metadata with values from input metadata that match StandardLoggingMetadata fields
for key in StandardLoggingMetadata.__annotations__:
if key in metadata:
clean_metadata[key] = metadata[key]
for key in metadata.keys() & _STANDARD_LOGGING_METADATA_KEYS:
clean_metadata[key] = metadata[key]
if metadata.get("user_api_key") is not None:
if is_valid_sha256_hash(str(metadata.get("user_api_key"))):

View file

@ -52,7 +52,7 @@ from typing_extensions import assert_never, overload
import litellm
import litellm.litellm_core_utils.exception_mapping_utils
from litellm import get_secret_str
from litellm._internal_context import mirrored_call, service_target, with_service_target
from litellm._internal_context import in_mirrored_call, mirrored_call, service_target, with_service_target
from litellm._logging import verbose_router_logger
from litellm._uuid import uuid
from litellm.caching.caching import (
@ -702,7 +702,7 @@ def _every_deployment_refused(
def _silent_experiment_targets(silent_model: object) -> tuple[str, ...]:
if silent_model is None:
if silent_model is None or in_mirrored_call():
return ()
try:
targets: Final = _SILENT_MODEL_ADAPTER.validate_python(silent_model)
@ -3325,10 +3325,6 @@ class Router:
Run a silent experiment in the background (thread).
"""
try:
# Prevent infinite recursion if silent model also has a silent model
if kwargs.get("metadata", {}).get("is_silent_experiment", False):
return
messages = copy.deepcopy(messages)
verbose_router_logger.info("Starting silent experiment for model %s", silent_model)
@ -4513,10 +4509,6 @@ class Router:
Run a silent experiment in the background.
"""
try:
# Prevent infinite recursion if silent model also has a silent model
if kwargs.get("metadata", {}).get("is_silent_experiment", False):
return
messages = copy.deepcopy(messages)
verbose_router_logger.info("Starting silent experiment for model %s", silent_model)

View file

@ -82,9 +82,9 @@ class _ShadowRecorder(CustomLogger):
async def async_log_success_event(
self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object
) -> None:
metadata: Final = _OBJECT.validate_python(_OBJECT.validate_python(kwargs["litellm_params"])["metadata"] or {})
if metadata.get("is_silent_experiment") is True:
self.shadow_groups.put(str(metadata.get("model_group")))
payload: Final = _OBJECT.validate_python(kwargs["standard_logging_object"])
if _OBJECT.validate_python(payload["metadata"]).get("is_silent_experiment") is True:
self.shadow_groups.put(str(payload["model_group"]))
def drained(self) -> tuple[str, ...]:
return tuple(self.shadow_groups.get_nowait() for _ in range(self.shadow_groups.qsize()))

View file

@ -11975,6 +11975,16 @@ def test_get_standard_logging_metadata(metadata_key, metadata_value):
assert standard_logging_metadata[metadata_key] == metadata_value
@pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown")
def test_module_level_get_standard_logging_metadata_ignores_a_caller_sent_silent_experiment_flag():
from litellm.litellm_core_utils.litellm_logging import get_standard_logging_metadata
result = get_standard_logging_metadata({"is_silent_experiment": True, "user_api_key_alias": "alias"})
assert "is_silent_experiment" not in result
assert result["user_api_key_alias"] == "alias"
@pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown")
def test_get_standard_logging_metadata_user_api_key_hash():
valid_hash = "a" * 64 # 64 character string

View file

@ -2,7 +2,7 @@ import asyncio
import datetime
import json
import time
from collections.abc import Callable, Mapping
from collections.abc import Callable, Generator, Mapping
from types import SimpleNamespace
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
@ -42,15 +42,43 @@ async def _settle_shared_logging_worker() -> None:
await GLOBAL_LOGGING_WORKER.stop()
@pytest.fixture
def recording_logger():
class _JudgingLogger(_RecordingLogger):
async def _judge(self, kwargs: Mapping[str, object]) -> None:
metadata: Final = kwargs["litellm_params"]["metadata"]
if "judge_of" in metadata:
return
await litellm.acompletion(
model="openai/gpt-5.4-nano",
api_key="fake-key",
messages=[{"role": "user", "content": "score it"}],
mock_response="ok",
metadata={"judge_of": metadata["model_group"]},
)
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
await super().async_log_success_event(kwargs, response_obj, start_time, end_time)
await self._judge(kwargs)
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
await self._judge(kwargs)
def judge_flags(self) -> dict[str, object]:
return {
call["litellm_params"]["metadata"]["judge_of"]: call["standard_logging_object"]["metadata"].get(
"is_silent_experiment"
)
for call in self.success_kwargs
if "judge_of" in call["litellm_params"]["metadata"]
}
def _installed(logger: _RecordingLogger) -> Generator[_RecordingLogger]:
settle_loop: Final = asyncio.new_event_loop()
try:
settle_loop.run_until_complete(_settle_shared_logging_worker())
finally:
settle_loop.close()
original_callbacks: Final = litellm.callbacks
logger: Final = _RecordingLogger()
litellm.callbacks = [logger]
try:
yield logger
@ -58,6 +86,16 @@ def recording_logger():
litellm.callbacks = original_callbacks
@pytest.fixture
def recording_logger():
yield from _installed(_RecordingLogger())
@pytest.fixture
def judging_logger():
yield from _installed(_JudgingLogger())
async def _wait_for_shadow_successes(logger: _RecordingLogger, expected: int, timeout: float = 5.0) -> None:
deadline: Final = time.monotonic() + timeout
while len(logger.shadow_successes()) < expected and time.monotonic() < deadline:
@ -70,6 +108,21 @@ async def _wait_for_all_successes(logger: _RecordingLogger, expected: int, timeo
await asyncio.sleep(0.05)
async def _wait_for_judges(logger: _JudgingLogger, judged: set[str], timeout: float = 5.0) -> None:
deadline: Final = time.monotonic() + timeout
while not judged <= logger.judge_flags().keys() and time.monotonic() < deadline:
await asyncio.sleep(0.05)
def _silent_flag_by_group(logger: _RecordingLogger) -> dict[str, object]:
return {
call["standard_logging_object"]["model_group"]: call["standard_logging_object"]["metadata"].get(
"is_silent_experiment"
)
for call in logger.success_kwargs
}
def _wait_for_shadow_successes_sync(logger: _RecordingLogger, expected: int, timeout: float = 5.0) -> None:
deadline: Final = time.monotonic() + timeout
while len(logger.shadow_successes()) < expected and time.monotonic() < deadline:
@ -806,15 +859,62 @@ async def test_aembedding_with_silent_model_sends_provider_body_without_it(
request_body: Final = json.loads(route.calls.last.request.read())
assert request_body == {"model": "embed-model", "input": ["black dresses"], "input_type": "query"}
assert response.data[0]["embedding"] == [0.1, 0.2]
@pytest.mark.parametrize("run_silent_experiment", SILENT_EXPERIMENT_RUNNERS)
def test_silent_experiment_does_not_launch_from_a_shadow_request(run_silent_experiment):
router = Router(model_list=_streaming_model_list(["shadow-a"]))
with patch.object(router, "acompletion", new_callable=AsyncMock, return_value=None) as acompletion:
run_silent_experiment(
router,
silent_model="shadow-a",
messages=[{"role": "user", "content": "hi"}],
metadata={"is_silent_experiment": True},
)
@pytest.mark.asyncio
async def test_a_caller_sent_silent_experiment_flag_still_mirrors_the_request_async(
recording_logger: _RecordingLogger,
) -> None:
router: Final = Router(model_list=_streaming_model_list(["shadow-a"]))
await router.acompletion(
model="primary-model",
messages=[{"role": "user", "content": "hi"}],
mock_response="pong",
metadata={"is_silent_experiment": True},
)
await _wait_for_all_successes(recording_logger, expected=2)
acompletion.assert_not_awaited()
assert _silent_flag_by_group(recording_logger) == {"primary-model": None, "shadow-a": True}
def test_a_caller_sent_silent_experiment_flag_still_mirrors_the_request_sync(
recording_logger: _RecordingLogger,
) -> None:
router: Final = Router(model_list=_streaming_model_list(["shadow-a"]))
router.completion(
model="primary-model",
messages=[{"role": "user", "content": "hi"}],
mock_response="pong",
metadata={"is_silent_experiment": True},
)
deadline: Final = time.monotonic() + 5.0
while "shadow-a" not in _silent_flag_by_group(recording_logger) and time.monotonic() < deadline:
time.sleep(0.05)
assert _silent_flag_by_group(recording_logger) == {"shadow-a": True}
@pytest.mark.asyncio
@pytest.mark.parametrize("shadow_outcome", ["pong", "litellm.InternalServerError"])
async def test_calls_a_callback_starts_while_logging_a_mirrored_call_are_not_marked_as_mirrored(
judging_logger: _JudgingLogger, shadow_outcome: str
) -> None:
router: Final = Router(
model_list=[
{
"model_name": "primary-model",
"litellm_params": {
"model": "openai/gpt-5.4-mini",
"api_key": "fake-key",
"silent_model": "shadow-a",
"mock_response": "pong",
},
},
{
"model_name": "shadow-a",
"litellm_params": {"model": "openai/gpt-5.4-nano", "api_key": "fake-key", "mock_response": shadow_outcome},
},
]
)
await router.acompletion(model="primary-model", messages=[{"role": "user", "content": "hi"}])
await _wait_for_judges(judging_logger, {"primary-model", "shadow-a"})
assert judging_logger.judge_flags() == {"primary-model": None, "shadow-a": None}