fix(router): write the prompt caching pin when turn_off_message_logging is on (#50276)

* fix(router): write the prompt caching pin when turn_off_message_logging is on

The success event PromptCachingDeploymentCheck reads had its messages redacted in place by
turn_off_message_logging before any callback ran, so the pin was keyed off redacted-by-litellm,
never written, and every follow-up bounced across the group. Callbacks that route on the request
and never log it now opt out of that redaction and get the request messages back, while every
other callback and the spend log keep the redacted payload.

* fix(logging): type the request-messages payload helper as dict[str, object]

The bare dict return grew the LIT001 count past the merge base and left the returned fields unchecked

* test(integration): type the pin rig's Redis client as Redis[bytes]

The untyped generic grew the basedpyright reportUnknownMemberType count past the merge base

* test(integration): quote the Redis[bytes] annotation so the rig collects at runtime

redis.Redis is not generic at runtime, so the unquoted subscript raised TypeError on import

* fix(logging): read the request messages once as object in the exempt payload helper

Two partially unknown reads of self.messages grew the basedpyright reportUnknownMemberType count past the merge base

---------

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Kerry Lu 2026-11-05 15:59:57 -08:00 • committed by GitHub
parent 05e4190918
commit 3f13c1ca34
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 178 additions and 1 deletions

View file

@ -898,6 +898,11 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
def redacts_messages_itself(self) -> bool:
return False
def exempt_from_message_redaction(self) -> bool:
"""Router pre-call checks that key routing off the request and never log it return True, so
`turn_off_message_logging` leaves the success event they receive unredacted."""
return False
def redact_standard_logging_payload_from_model_call_details(self, model_call_details: dict) -> dict:
"""
Redacts or excludes fields from StandardLoggingPayload before callbacks receive it.

View file

@ -4405,7 +4405,11 @@ class Logging(LiteLLMLoggingBaseClass):
if isinstance(callback, CustomLogger): # custom logger class
from litellm.integrations.shadow_eval_logger import ShadowEvalLogger
model_call_details: dict = self.model_call_details
model_call_details: dict = (
self._model_call_details_with_request_messages()
if callback.exempt_from_message_redaction()
else self.model_call_details
)
##################################
# call redaction hook for custom logger
model_call_details = callback.redact_standard_logging_payload_from_model_call_details(
@ -4969,6 +4973,21 @@ class Logging(LiteLLMLoggingBaseClass):
) -> None:
executor.submit(self.success_handler, result, start_time, end_time, cache_hit)
def _model_call_details_with_request_messages(self) -> dict[str, object]: # mutable-ok: loggers take a dict
"""The success-event payload with the request messages put back where `turn_off_message_logging`
redacted them in place, for a callback that routes on the request and never logs it."""
request_messages: Final = cast(object, self.messages) # cast-ok: the request messages are only put back as-is
standard_logging_object: Final = self.model_call_details.get("standard_logging_object")
return {
**self.model_call_details,
"messages": request_messages,
**(
{"standard_logging_object": {**standard_logging_object, "messages": request_messages}}
if isinstance(standard_logging_object, dict)
else {}
),
}
def _get_completed_success_callbacks(
self, *, asynchronous: bool
) -> tuple[str | Callable[..., object] | CustomLogger, ...]:

View file

@ -163,6 +163,9 @@ class PromptCachingDeploymentCheck(CustomLogger):
return healthy_deployments
def exempt_from_message_redaction(self) -> bool:
return True
async def async_log_success_event(
self,
kwargs: _PromptCachingSuccessEventKwargs,

View file

@ -0,0 +1,95 @@
import collections
from collections.abc import Iterator
from dataclasses import dataclass
from typing import Final
import pytest
from integration._support.client import eventually, gateway_from_environment, object_value
from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
from integration._support.redis_process import owned_redis
from integration._support.wire import Reply, Wire, wire_server
from integration.routing._prompt_cache_pin_support import (
DEFAULT_MARK,
anthropic_deployment,
ask_turn,
chat_body,
claude_peer,
new_marker,
one_new_pin,
pin_config,
post_chat,
read_pins,
system_turn,
)
from pydantic import JsonValue, TypeAdapter
from redis import Redis
from litellm.constants import REDACTED_BY_LITELLM
pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 300)
_GROUP: Final = "pin-logging-off"
_LEFT: Final = "pin-logging-off-left"
_RIGHT: Final = "pin-logging-off-right"
_FOLLOW_UPS: Final = 20
_REDACTED_MESSAGES: Final = [{"role": "user", "content": REDACTED_BY_LITELLM}]
_LOGGED_EVENT: Final = TypeAdapter(dict[str, JsonValue])
@dataclass(frozen=True, slots=True)
class _Rig:
owned: OwnedProxy
redis: "Redis[bytes]"
logger_sink: Wire
@pytest.fixture(scope="module")
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]:
directory: Final = tmp_path_factory.mktemp("prompt-cache-pin-message-logging-off")
with (
gateway_from_environment() as gateway,
owned_redis(directory) as cache,
wire_server(claude_peer) as left,
wire_server(claude_peer) as right,
wire_server(lambda request: Reply()) as logger_sink,
):
config: Final = pin_config(
directory,
[anthropic_deployment(_GROUP, _LEFT, left.url), anthropic_deployment(_GROUP, _RIGHT, right.url)],
litellm_settings={
"turn_off_message_logging": True,
"callbacks": ["integration._support.redaction_probe.probe"],
"cache": False,
"enable_redis_auth_cache": False,
},
)
environment: Final = {
"REDIS_HOST": cache.host,
"REDIS_PORT": str(cache.port),
"REDACTION_PROBE_URL": logger_sink.url,
}
with (
owned_proxy_process(gateway, directory, environment, config=config, workers=1) as owned,
Redis(host=cache.host, port=cache.port, socket_timeout=2) as redis,
):
yield _Rig(owned, redis, logger_sink)
def test_follow_ups_stay_on_the_first_deployment_while_the_logger_sees_redacted_messages(rig: _Rig) -> None:
before: Final = frozenset(read_pins(rig.redis))
body: Final = chat_body(_GROUP, [system_turn(new_marker(), DEFAULT_MARK), ask_turn(new_marker())])
_, first_served = post_chat(rig.owned.gateway, body)
_, pin = one_new_pin(rig.redis, before)
assert pin.model_id == first_served, (pin, first_served)
served: Final = tuple(post_chat(rig.owned.gateway, body)[1] for _ in range(_FOLLOW_UPS))
assert served == (first_served,) * _FOLLOW_UPS, collections.Counter(served)
eventually(rig.logger_sink.received.qsize, lambda count: count >= 1 + _FOLLOW_UPS, seconds=30)
logged: Final = tuple(
object_value(_LOGGED_EVENT.validate_json(request.body)["standard"])["messages"]
for request in rig.logger_sink.drain()
)
assert len(logged) == 1 + _FOLLOW_UPS, len(logged)
assert all(messages == _REDACTED_MESSAGES for messages in logged), logged

View file

@ -1480,3 +1480,58 @@ async def test_a_pin_written_without_tools_is_read_when_tools_is_an_empty_list()
)
assert filtered == [deployments[0]]
class _LoggedMessagesProbe(CustomLogger):
def __init__(self) -> None:
self.logged_messages: list[object] = []
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
self.logged_messages.append((kwargs["standard_logging_object"]["messages"], kwargs["messages"]))
@pytest.mark.asyncio
async def test_pin_is_written_with_message_logging_off_while_loggers_still_see_redacted_messages(
monkeypatch, local_model_cost_map
):
"""
`turn_off_message_logging` redacted the shared success-event payload before the check read it, so the
pin was keyed off `redacted-by-litellm`, never written, and a cacheable conversation bounced across the
group. The pin must land whatever the flag says, while a real logger alongside still sees only the
redacted messages.
"""
monkeypatch.setattr(litellm, "turn_off_message_logging", True)
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", False)
probe = _LoggedMessagesProbe()
monkeypatch.setattr(litellm, "callbacks", [probe])
router = litellm.Router(
model_list=[
{
"model_name": MODEL_GROUP_ALIAS,
"litellm_params": {"model": "anthropic/claude-opus-4-6", "api_key": "sk-fake"},
"model_info": {"id": model_id},
}
for model_id in ("dep-1", "dep-2")
],
optional_pre_call_checks=["prompt_caching"],
)
messages = _messages(word_count=5000)
first = await router.acompletion(model=MODEL_GROUP_ALIAS, messages=messages, mock_response="ok")
served_by = first._hidden_params["model_id"]
affinity_key = PromptCachingCache.get_prompt_caching_cache_key(messages, None)
pin = await _eventually(lambda: router.prompt_caching_pin_cache.get_cache(key=affinity_key))
assert pin == {"model_id": served_by}, pin
subsequent = [
(await router.acompletion(model=MODEL_GROUP_ALIAS, messages=messages, mock_response="ok"))._hidden_params[
"model_id"
]
for _ in range(4)
]
assert subsequent == [served_by] * 4
await _eventually(lambda: len(probe.logged_messages) == 5)
redacted = [{"role": "user", "content": "redacted-by-litellm"}]
assert probe.logged_messages == [(redacted, redacted)] * 5, probe.logged_messages