mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
fix(redaction): redact MCP tool arguments and results
MCP tool arguments and results ride in metadata.mcp_tool_call_metadata rather than in messages or response, so every integration that logs the standard payload exported them verbatim while message logging was off. Redact them in the global and per-callback paths, and on the raw call details, which global redaction already rewrites for messages. Only assign the payload metadata when redaction actually changed it, so a payload without a metadata key does not gain one. Widen the two StandardLoggingMCPToolCall fields to dict | str so the marker is visible to type checkers.
This commit is contained in:
parent
5c1623888e
commit
40f1e3eab0
4 changed files with 151 additions and 4 deletions
|
|
@ -828,6 +828,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
|
||||
import litellm
|
||||
from litellm import Choices, Message, ModelResponse
|
||||
from litellm.litellm_core_utils.redact_messages import redacted_mcp_tool_call_metadata
|
||||
|
||||
turn_off_message_logging: Final[bool] = getattr(self, "turn_off_message_logging", False)
|
||||
excluded_fields: Final[list[str] | None] = getattr(litellm, "standard_logging_payload_excluded_fields", None)
|
||||
|
|
@ -859,6 +860,12 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
if "messages" not in (excluded_fields or []) and standard_logging_object_copy.get("messages") is not None:
|
||||
standard_logging_object_copy["messages"] = [Message(content=redacted_str).model_dump()]
|
||||
|
||||
if "metadata" not in (excluded_fields or []):
|
||||
metadata: Final = standard_logging_object_copy.get("metadata")
|
||||
redacted_metadata: Final = redacted_mcp_tool_call_metadata(metadata, redacted_str)
|
||||
if redacted_metadata is not metadata:
|
||||
standard_logging_object_copy["metadata"] = redacted_metadata
|
||||
|
||||
if "response" not in (excluded_fields or []) and standard_logging_object_copy.get("response") is not None:
|
||||
response: Final = standard_logging_object_copy["response"]
|
||||
# Check if this is a ResponsesAPIResponse (has "output" field)
|
||||
|
|
|
|||
|
|
@ -158,6 +158,41 @@ def _redact_responses_api_output_dict(output_items, redacted_str: str):
|
|||
output_item["arguments"] = redacted_str
|
||||
|
||||
|
||||
def redacted_mcp_tool_call_metadata(metadata: object, redacted_str: str) -> object:
|
||||
"""MCP tool arguments and results are user content, and they ride in
|
||||
`metadata.mcp_tool_call_metadata` rather than in `messages` / `response`,
|
||||
so every integration that logs metadata exports them unless redacted here.
|
||||
"""
|
||||
if not isinstance(metadata, dict):
|
||||
return metadata
|
||||
|
||||
mcp_tool_call: Final = metadata.get("mcp_tool_call_metadata")
|
||||
if not isinstance(mcp_tool_call, dict):
|
||||
return metadata
|
||||
|
||||
redacted_call: Final = {
|
||||
**mcp_tool_call,
|
||||
**{key: redacted_str for key in ("arguments", "result") if mcp_tool_call.get(key) is not None},
|
||||
}
|
||||
return {**metadata, "mcp_tool_call_metadata": redacted_call}
|
||||
|
||||
|
||||
def _redact_mcp_tool_call_in_place(model_call_details: dict) -> None:
|
||||
"""Redact the raw `mcp_tool_call_metadata` the MCP server stamped onto the
|
||||
call details, so callbacks reading it off `kwargs` see the same redaction
|
||||
the standard logging payload gets.
|
||||
"""
|
||||
mcp_tool_call: Final = model_call_details.get("mcp_tool_call_metadata")
|
||||
if not isinstance(mcp_tool_call, dict):
|
||||
return
|
||||
|
||||
for key in ("arguments", "result"):
|
||||
if mcp_tool_call.get(key) is not None:
|
||||
# mutable-ok: redacting the shared call details in place is the point; a copy would leave
|
||||
# callbacks that read mcp_tool_call_metadata off kwargs holding the unredacted values
|
||||
mcp_tool_call[key] = "redacted-by-litellm"
|
||||
|
||||
|
||||
def _redact_standard_logging_object(model_call_details: dict):
|
||||
"""Redact messages and response inside standard_logging_object if present."""
|
||||
standard_logging_object: Final = model_call_details.get("standard_logging_object")
|
||||
|
|
@ -169,6 +204,11 @@ def _redact_standard_logging_object(model_call_details: dict):
|
|||
if standard_logging_object.get("messages") is not None:
|
||||
standard_logging_object["messages"] = [{"role": "user", "content": redacted_str}]
|
||||
|
||||
metadata: Final = standard_logging_object.get("metadata")
|
||||
redacted_metadata: Final = redacted_mcp_tool_call_metadata(metadata, redacted_str)
|
||||
if redacted_metadata is not metadata:
|
||||
standard_logging_object["metadata"] = redacted_metadata
|
||||
|
||||
response: Final = standard_logging_object.get("response")
|
||||
if response is not None:
|
||||
if isinstance(response, dict) and "output" in response:
|
||||
|
|
@ -238,6 +278,7 @@ def perform_redaction(model_call_details: dict, result, redact_streaming_respons
|
|||
model_call_details["messages"] = [{"role": "user", "content": "redacted-by-litellm"}]
|
||||
model_call_details["prompt"] = ""
|
||||
model_call_details["input"] = ""
|
||||
_redact_mcp_tool_call_in_place(model_call_details)
|
||||
_redact_standard_logging_object(model_call_details)
|
||||
redact_vertex_ai_metadata_from_litellm_params(model_call_details)
|
||||
|
||||
|
|
|
|||
|
|
@ -2637,13 +2637,13 @@ class StandardLoggingMCPToolCall(TypedDict, total=False):
|
|||
"""
|
||||
Name of the tool to call
|
||||
"""
|
||||
arguments: dict
|
||||
arguments: dict | str
|
||||
"""
|
||||
Arguments to pass to the tool
|
||||
Arguments to pass to the tool, or the redaction marker when message logging is off
|
||||
"""
|
||||
result: dict
|
||||
result: dict | str
|
||||
"""
|
||||
Result of the tool call
|
||||
Result of the tool call, or the redaction marker when message logging is off
|
||||
"""
|
||||
|
||||
mcp_server_name: str | None
|
||||
|
|
|
|||
|
|
@ -720,3 +720,102 @@ class TestRedactStreamingResponsesForCustomLogger:
|
|||
|
||||
assert result_details is model_call_details
|
||||
assert response_obj.choices[0].message.content == "secret content"
|
||||
|
||||
|
||||
class TestRedactMcpToolCallMetadata:
|
||||
"""MCP tool arguments and results ride in metadata rather than in
|
||||
messages/response, so they need their own redaction pass."""
|
||||
|
||||
def _slp(self):
|
||||
return {
|
||||
"standard_logging_object": {
|
||||
"messages": [{"role": "user", "content": "secret prompt"}],
|
||||
"metadata": {
|
||||
"user_api_key_hash": "abc",
|
||||
"mcp_tool_call_metadata": {
|
||||
"name": "get_weather",
|
||||
"arguments": {"city": "TOPSECRETCITY"},
|
||||
"result": {"text": "sensitive"},
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
def test_global_redaction_strips_mcp_arguments_and_result(self):
|
||||
from litellm.litellm_core_utils.redact_messages import _redact_standard_logging_object
|
||||
|
||||
model_call_details = self._slp()
|
||||
_redact_standard_logging_object(model_call_details)
|
||||
|
||||
mcp_meta = model_call_details["standard_logging_object"]["metadata"]["mcp_tool_call_metadata"]
|
||||
assert mcp_meta["arguments"] == "redacted-by-litellm"
|
||||
assert mcp_meta["result"] == "redacted-by-litellm"
|
||||
assert mcp_meta["name"] == "get_weather"
|
||||
assert model_call_details["standard_logging_object"]["metadata"]["user_api_key_hash"] == "abc"
|
||||
|
||||
def test_callback_level_redaction_strips_mcp_arguments(self):
|
||||
opted_out_logger = CustomLogger(turn_off_message_logging=True)
|
||||
|
||||
redacted = opted_out_logger.redact_standard_logging_payload_from_model_call_details(self._slp())
|
||||
|
||||
mcp_meta = redacted["standard_logging_object"]["metadata"]["mcp_tool_call_metadata"]
|
||||
assert mcp_meta["arguments"] == "redacted-by-litellm"
|
||||
assert mcp_meta["result"] == "redacted-by-litellm"
|
||||
|
||||
def test_callback_level_redaction_does_not_mutate_the_original(self):
|
||||
model_call_details = self._slp()
|
||||
opted_out_logger = CustomLogger(turn_off_message_logging=True)
|
||||
|
||||
opted_out_logger.redact_standard_logging_payload_from_model_call_details(model_call_details)
|
||||
|
||||
original = model_call_details["standard_logging_object"]["metadata"]["mcp_tool_call_metadata"]
|
||||
assert original["arguments"] == {"city": "TOPSECRETCITY"}
|
||||
|
||||
def test_metadata_without_mcp_tool_call_is_untouched(self):
|
||||
from litellm.litellm_core_utils.redact_messages import redacted_mcp_tool_call_metadata
|
||||
|
||||
metadata = {"user_api_key_hash": "abc"}
|
||||
|
||||
assert redacted_mcp_tool_call_metadata(metadata, "redacted-by-litellm") is metadata
|
||||
|
||||
def test_global_redaction_also_strips_the_raw_call_details_copy(self):
|
||||
from litellm.litellm_core_utils.redact_messages import perform_redaction
|
||||
|
||||
model_call_details = {
|
||||
"messages": [{"role": "user", "content": "secret prompt"}],
|
||||
"mcp_tool_call_metadata": {
|
||||
"name": "get_weather",
|
||||
"arguments": {"city": "TOPSECRETCITY"},
|
||||
"result": {"text": "sensitive"},
|
||||
},
|
||||
"standard_logging_object": {"messages": [], "metadata": {}},
|
||||
}
|
||||
|
||||
perform_redaction(model_call_details, result=None)
|
||||
|
||||
raw = model_call_details["mcp_tool_call_metadata"]
|
||||
assert raw["arguments"] == "redacted-by-litellm"
|
||||
assert raw["result"] == "redacted-by-litellm"
|
||||
assert raw["name"] == "get_weather"
|
||||
|
||||
def test_payload_without_metadata_does_not_gain_a_metadata_key(self):
|
||||
"""The callback-logs replay endpoint seeds a payload from an unvalidated
|
||||
POST body, so `metadata` can legitimately be absent. Consumers that do
|
||||
`payload.get("metadata", {}).get(...)` break on an explicit None."""
|
||||
from litellm.litellm_core_utils.redact_messages import _redact_standard_logging_object
|
||||
|
||||
model_call_details = {"standard_logging_object": {"messages": [{"role": "user", "content": "x"}]}}
|
||||
|
||||
_redact_standard_logging_object(model_call_details)
|
||||
|
||||
assert "metadata" not in model_call_details["standard_logging_object"]
|
||||
|
||||
def test_metadata_without_mcp_tool_call_keeps_its_identity(self):
|
||||
from litellm.litellm_core_utils.redact_messages import _redact_standard_logging_object
|
||||
|
||||
metadata = {"user_api_key_hash": "abc"}
|
||||
model_call_details = {"standard_logging_object": {"messages": [], "metadata": metadata}}
|
||||
|
||||
_redact_standard_logging_object(model_call_details)
|
||||
|
||||
assert model_call_details["standard_logging_object"]["metadata"] is metadata
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue