From da447bb7595acbd14841c4feb54e2a7f351db7b0 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Fri, 6 Nov 2026 02:37:36 +0000 Subject: [PATCH] fix(router): write the prompt caching pin when turn_off_message_logging is on (#50276) Backport of #50276 to lts/1.108.x. tests/integration/_support/redaction_probe.py also sends the standard key the picked integration test reads (main's probe got it in #47192). Cherry-picked from 3f13c1ca34874c3bfa021e6fc14ac950c56d97b5 (main). Conflicts in litellm_logging.py and the prompt caching check test were main-only neighbours (_get_completed_success_callbacks and the tools/system affinity tests), resolved by adding only this PR's method and test. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/custom_logger.py | 5 + litellm/litellm_core_utils/litellm_logging.py | 21 +++- .../prompt_caching_deployment_check.py | 3 + tests/integration/_support/redaction_probe.py | 1 + ...ompt_cache_pin_with_message_logging_off.py | 95 +++++++++++++++++++ .../test_prompt_caching_deployment_check.py | 55 +++++++++++ 6 files changed, 179 insertions(+), 1 deletion(-) create mode 100644 tests/integration/routing/test_prompt_cache_pin_with_message_logging_off.py diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 5addc994df2..9332137e3b1 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -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. diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 1d6250fc528..95323c2ef5e 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -3751,7 +3751,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( @@ -4289,6 +4293,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 _should_run_sync_callbacks_for_async_calls(self) -> bool: """ Returns: diff --git a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py index b2f0c658161..435f9b9c5a4 100644 --- a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py +++ b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py @@ -138,6 +138,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, diff --git a/tests/integration/_support/redaction_probe.py b/tests/integration/_support/redaction_probe.py index 2a344daa73c..83d1cbd9cba 100644 --- a/tests/integration/_support/redaction_probe.py +++ b/tests/integration/_support/redaction_probe.py @@ -19,6 +19,7 @@ def _payload(kwargs: Mapping[str, object], response_obj: object) -> str: "request": request, "response": response_obj.model_dump() if isinstance(response_obj, BaseModel) else response_obj, "standard_response": standard.get("response") if isinstance(standard, Mapping) else None, + "standard": standard, }, default=str, ) diff --git a/tests/integration/routing/test_prompt_cache_pin_with_message_logging_off.py b/tests/integration/routing/test_prompt_cache_pin_with_message_logging_off.py new file mode 100644 index 00000000000..13bf67582b3 --- /dev/null +++ b/tests/integration/routing/test_prompt_cache_pin_with_message_logging_off.py @@ -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 diff --git a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index 5f018ab5dda..5b31e0e5e3b 100644 --- a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -1236,3 +1236,58 @@ async def test_router_with_prompt_caching(anthropic_messages): print("response=", response) assert response._hidden_params["model_id"] == initial_model_id + + +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