From 53729cf747d8dd725ffdd9acf9939dc3eb5bac4a Mon Sep 17 00:00:00 2001 From: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> Date: Mon, 12 Oct 2026 12:35:11 -0700 Subject: [PATCH] fix(logging): keep replayed instructions and router usage marker under redaction The shared request body keeps system and instructions again, so spend logs store them and previous_response_id follow-ups replay the real instructions. Per-callback copies still redact them. Provider-specific fields in metadata.hidden_params are cleared in place, so the metadata dict keeps its identity and the router's counted-usage marker stays visible to deployment_callback_on_success (no double tpm/rpm count). --- litellm/litellm_core_utils/redact_messages.py | 46 +++++-- .../test_custom_logger_redaction.py | 116 +++++++++++++++++- .../test_redact_messages.py | 27 +++- 3 files changed, 174 insertions(+), 15 deletions(-) diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index 12a4384efa7..aa8373ab300 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -281,7 +281,7 @@ def _redact_model_response_dict_choices(choices, redacted_str: str): def _hidden_params_without_provider_fields(hidden_params: object) -> Mapping[object, object] | None: """Copy of ``hidden_params`` with ``provider_specific_fields`` cleared, or None when it carries none.""" - if not (isinstance(hidden_params, Mapping) and "provider_specific_fields" in hidden_params): + if not (isinstance(hidden_params, Mapping) and hidden_params.get("provider_specific_fields") is not None): return None return {**_REQUEST_BODY.validate_python(hidden_params), "provider_specific_fields": None} @@ -302,12 +302,29 @@ def _redacted_metadata(metadata: Mapping[str, object], *, redact_content: bool) return {**metadata, **redacted_entries} if redacted_entries else metadata +def _clear_hidden_provider_fields(params: Mapping[str, object]) -> None: + """ + Clear ``provider_specific_fields`` from the shared metadata ``hidden_params`` in place, so the + metadata dict keeps its identity: the router writes its counted-usage marker onto it later. + """ + for metadata in (params.get("metadata"), params.get("litellm_metadata")): + if not isinstance(metadata, dict): + continue + hidden_params = _hidden_params_without_provider_fields( + _REQUEST_BODY.validate_python(metadata).get("hidden_params") + ) + if hidden_params is not None: + metadata["hidden_params"] = hidden_params + + def _without_hidden_provider_fields(metadata: Mapping[object, object]) -> Mapping[object, object]: hidden_params: Final = _hidden_params_without_provider_fields(metadata.get("hidden_params")) return metadata if hidden_params is None else {**metadata, "hidden_params": hidden_params} -def _redacted_proxy_request(request: Mapping[str, object], *, redact_content: bool) -> Mapping[str, object]: +def _redacted_proxy_request( + request: Mapping[str, object], *, redact_content: bool, redact_instructions: bool +) -> Mapping[str, object]: raw_body: Final = request.get("body") if not (redact_content and isinstance(raw_body, Mapping)): return without_classifier_audit(request) @@ -324,16 +341,23 @@ def _redacted_proxy_request(request: Mapping[str, object], *, redact_content: bo **({"messages": [{"role": "user", "content": REDACTED_BY_LITELLM}]} if "messages" in body else {}), **({"prompt": ""} if "prompt" in body else {}), **({"input": ""} if "input" in body else {}), - **{key: REDACTED_BY_LITELLM for key in ("system", "instructions") if body.get(key) is not None}, + **{ + key: REDACTED_BY_LITELLM + for key in ("system", "instructions") + if redact_instructions and body.get(key) is not None + }, }, } -def redacted_litellm_params(params: Mapping[str, object], *, redact_content: bool = True) -> dict[str, object]: +def redacted_litellm_params( + params: Mapping[str, object], *, redact_content: bool = True, redact_instructions: bool = True +) -> dict[str, object]: """ Copy of litellm_params without classifier audit fields or baseline route request parameters. redact_content also redacts the nested proxy request body and the provider-specific fields - copied into ``hidden_params``. + copied into ``hidden_params``; redact_instructions also redacts the body's ``system`` and + ``instructions``. """ request: Final = params.get("proxy_server_request") return { @@ -344,7 +368,11 @@ def redacted_litellm_params(params: Mapping[str, object], *, redact_content: boo if key in ("metadata", "litellm_metadata") and isinstance(value, Mapping) }, **( - {"proxy_server_request": _redacted_proxy_request(request, redact_content=redact_content)} + { + "proxy_server_request": _redacted_proxy_request( + request, redact_content=redact_content, redact_instructions=redact_instructions + ) + } if isinstance(request, Mapping) else {} ), @@ -358,12 +386,16 @@ def perform_redaction(model_call_details: dict, result, redact_streaming_respons redact_streaming_responses=False is the per-callback path: it leaves the shared streaming response entries, proxy request body and hidden provider fields in place, and each opted-out callback gets its own redacted copy via redact_streaming_responses_for_custom_logger instead. + The shared request body keeps ``system`` and ``instructions``: spend logs store it and + ``previous_response_id`` follow-ups replay the stored instructions. """ # Redact model_call_details params: Final = model_call_details.get("litellm_params") if isinstance(params, Mapping): + if redact_streaming_responses: + _clear_hidden_provider_fields(params) model_call_details["litellm_params"] = redacted_litellm_params( - params, redact_content=redact_streaming_responses + params, redact_content=redact_streaming_responses, redact_instructions=False ) model_call_details["messages"] = [{"role": "user", "content": REDACTED_BY_LITELLM}] model_call_details["prompt"] = "" diff --git a/tests/integration/observability/test_custom_logger_redaction.py b/tests/integration/observability/test_custom_logger_redaction.py index f2e717df236..e6d2ee8b919 100644 --- a/tests/integration/observability/test_custom_logger_redaction.py +++ b/tests/integration/observability/test_custom_logger_redaction.py @@ -1,5 +1,7 @@ import asyncio import json +import os +import uuid from collections.abc import Iterator from concurrent.futures import ThreadPoolExecutor from pathlib import Path @@ -7,17 +9,24 @@ from typing import Final import pytest import yaml -from integration._support.client import Gateway, eventually, gateway_from_environment +from integration._support.client import Gateway, eventually, gateway_from_environment, held +from integration._support.database import read_rows from integration._support.process import owned_proxy from integration._support.redaction_probe import RedactionProbe, SyncRedactionProbe from integration._support.wire import Reply, Request, Wire, wire_server +from redis import Redis import litellm +from litellm.constants import REDACTED_BY_LITELLM from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER INPUT: Final = "private-input-redaction-canary" REASONING: Final = "private-reasoning-redaction-canary" OUTPUT: Final = "private-output-redaction-canary" +INSTRUCTIONS: Final = "Reply in French. redaction-instructions-canary" +BRIDGED: Final = "bridged-redaction-canary" +USAGE_DEPLOYMENT: Final = f"redaction-usage-{uuid.uuid4().hex}" +REDACT: Final = {"x-litellm-enable-message-redaction": "true"} pytestmark = pytest.mark.timeout(240) @@ -75,7 +84,8 @@ def _provider(request: Request) -> Reply: ).encode() ) if request.target == "/v1/messages": - assert body["messages"] == [{"role": "user", "content": INPUT}], body + if BRIDGED not in request.body.decode(): + assert body["messages"] == [{"role": "user", "content": INPUT}], body return Reply( body=json.dumps( { @@ -164,6 +174,19 @@ def redaction_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[ ("redaction-responses", "openai/gpt-5.6-sol"), ("redaction-messages", "anthropic/claude-sonnet-5-5"), ) + ] + + [ + { + "model_name": "redaction-usage", + "litellm_params": { + "model": "anthropic/claude-sonnet-5-5", + "api_base": provider.url, + "api_key": "synthetic-provider-key", + "rpm": 100000, + "tpm": 100000000, + }, + "model_info": {"id": USAGE_DEPLOYMENT}, + } ], "litellm_settings": { "callbacks": ["integration._support.redaction_probe.probe"], @@ -179,7 +202,9 @@ def redaction_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[ "master_key": "os.environ/LITELLM_MASTER_KEY", "database_url": "os.environ/DATABASE_URL", "disable_model_info_refresh": True, + "store_prompts_in_spend_logs": True, }, + "router_settings": {"redis_host": "os.environ/REDIS_HOST", "redis_port": "os.environ/REDIS_PORT"}, } ) ) @@ -331,6 +356,93 @@ def test_concurrent_redacted_and_unredacted_callbacks_remain_isolated( assert len(tuple(request for request in provider.drain() if request.method == "POST")) == 4 +def _drained(wire: Wire, count: int) -> tuple[Request, ...]: + batches: Final[list[Request]] = [] # mutable-ok: drain consumes the queue, so polls retain previous events + + def observed() -> tuple[Request, ...]: + batches.extend(wire.drain()) + return tuple(batches) + + return eventually(observed, lambda values: len(values) >= count, seconds=30) + + +def test_redacted_responses_follow_up_replays_original_instructions( + redaction_proxy: tuple[Gateway, Wire, Wire], +) -> None: + proxy, provider, sink = redaction_proxy + end_user: Final = f"redaction-replay-{uuid.uuid4().hex}" + first: Final = proxy.request( + "POST", + "/v1/responses", + { + "model": "redaction-messages", + "instructions": INSTRUCTIONS, + "input": f"{BRIDGED} turn one", + "user": end_user, + "cache": {"no-cache": True}, + }, + headers=REDACT, + ) + assert first.status_code == 200, first.text + eventually( + lambda: read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE end_user = %s', (end_user,)), + lambda rows: len(rows) == 1, + seconds=70, + ) + assert len(_drained(provider, 1)) == 1 + follow_up: Final = proxy.request( + "POST", + "/v1/responses", + { + "model": "redaction-messages", + "previous_response_id": first.json()["id"], + "input": f"{BRIDGED} turn two", + "user": end_user, + "cache": {"no-cache": True}, + }, + headers=REDACT, + ) + assert follow_up.status_code == 200, follow_up.text + sent: Final = tuple(json.loads(request.body) for request in _drained(provider, 1)) + assert len(_drained(sink, 2)) == 2 + assert len(sent) == 1, sent + assert len(sent[0]["messages"]) == 3, sent[0] + system: Final = json.dumps(sent[0].get("system")) + assert INSTRUCTIONS in system and REDACTED_BY_LITELLM not in system, sent[0] + + +def test_redacted_router_usage_is_counted_once(redaction_proxy: tuple[Gateway, Wire, Wire]) -> None: + proxy, provider, sink = redaction_proxy + requests: Final = 48 + + def send(index: int): + return proxy.request( + "POST", + "/v1/chat/completions", + { + "model": "redaction-usage", + "messages": [{"role": "user", "content": f"{BRIDGED} usage {index}"}], + "cache": {"no-cache": True}, + }, + headers=REDACT, + ) + + with ThreadPoolExecutor(max_workers=16) as pool: + responses: Final = tuple(pool.map(send, range(requests))) + for response in responses: + assert response.status_code == 200, response.text + assert response.json()["usage"]["total_tokens"] == 15, response.json() + + def counted_tokens() -> int: + with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as cache: + keys: Final = tuple(cache.scan_iter(f"global_router:{USAGE_DEPLOYMENT}:*:tpm:*")) + return sum(int(float(cache.get(key) or 0)) for key in keys) + + _drained(sink, requests) + held(counted_tokens, lambda total: total == 15 * requests, holding=3, seconds=30) + assert len(_drained(provider, requests)) == requests + + @pytest.mark.asyncio @pytest.mark.parametrize("asynchronous", [False, True]) @pytest.mark.parametrize("redact", [False, True]) diff --git a/tests/unit/litellm_core_utils/test_redact_messages.py b/tests/unit/litellm_core_utils/test_redact_messages.py index 23243ae694f..8e9a9bc514d 100644 --- a/tests/unit/litellm_core_utils/test_redact_messages.py +++ b/tests/unit/litellm_core_utils/test_redact_messages.py @@ -23,7 +23,7 @@ import pytest_asyncio from pydantic import JsonValue import litellm -from litellm.constants import LOGGING_WORKER_MAX_TIME_PER_COROUTINE +from litellm.constants import LOGGING_WORKER_MAX_TIME_PER_COROUTINE, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.litellm_core_utils.redact_messages import ( @@ -34,6 +34,7 @@ from litellm.litellm_core_utils.redact_messages import ( should_redact_message_logging, ) from litellm.responses.main import mock_responses_api_response +from litellm.router_utils.router_callbacks.track_deployment_metrics import get_counted_usage_tokens from litellm.types.router import BaselineRouteStamp from litellm.types.utils import ( ModelResponse, @@ -107,8 +108,8 @@ def _request_with_hidden_reasoning(secret: str) -> tuple[dict, dict]: metadata: Final = {"session_id": "session", "hidden_params": hidden_params} body: Final = { "model": "redaction-model", - "system": secret, - "instructions": secret, + "system": "kept-instructions", + "instructions": "kept-instructions", "messages": [{"role": "user", "content": secret}], "metadata": metadata, "litellm_metadata": metadata, @@ -132,14 +133,26 @@ def test_redaction_removes_reasoning_copied_into_hidden_params_and_request_body( redacted: Final = details["litellm_params"] assert secret not in str(redacted) body: Final = redacted["proxy_server_request"]["body"] - assert body["system"] == "redacted-by-litellm" - assert body["instructions"] == "redacted-by-litellm" + assert body["system"] == body["instructions"] == "kept-instructions" assert body["model"] == "redaction-model" for metadata in (redacted["metadata"], redacted["litellm_metadata"], body["metadata"], body["litellm_metadata"]): assert metadata["session_id"] == "session" assert metadata["hidden_params"]["provider_specific_fields"] is None assert metadata["hidden_params"]["additional_headers"] == hidden_params["additional_headers"] - assert params == snapshot + assert redacted["metadata"] is params["metadata"] + assert hidden_params == snapshot["metadata"]["hidden_params"] + + +@pytest.mark.parametrize("channel", ["metadata", "litellm_metadata"]) +def test_router_usage_marker_written_after_redaction_stays_visible(channel: str) -> None: + params, _ = _request_with_hidden_reasoning("private-hidden-reasoning") + params[channel] = copy.deepcopy(params[channel]) + details: Final = {"litellm_params": params} + + perform_redaction(details, None) + params[channel][ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY] = 7 + + assert get_counted_usage_tokens(details["litellm_params"]) == 7 def test_per_callback_redaction_keeps_shared_request_body_for_other_readers() -> None: @@ -155,6 +168,8 @@ def test_per_callback_redaction_keeps_shared_request_body_for_other_readers() -> assert details["litellm_params"] == snapshot assert details["messages"] == [{"role": "user", "content": "redacted-by-litellm"}] assert secret not in str(callback_details["litellm_params"]) + callback_body: Final = callback_details["litellm_params"]["proxy_server_request"]["body"] + assert callback_body["system"] == callback_body["instructions"] == "redacted-by-litellm" assert redact_streaming_responses_for_custom_logger(details, CustomLogger()) is details