scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-21 00:21:49 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 73

fix(mcp): retain wire aliases in guardrail inspection payloads

Browse source
This commit is contained in:
Joshua Valluru 2026-09-18 23:00:32 -07:00
parent f5ab563499
commit c91ca90477
2 changed files with 21 additions and 17 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense_mcp.py
View file

@ -34,7 +34,7 @@ def _serialize_mcp_content_item(item: object) -> dict[str, object]:
model_dump: Final = getattr(item, "model_dump", None)
if callable(model_dump):
try:
dumped: Final[dict[str, object]] = model_dump(exclude_none=True)
dumped: Final[dict[str, object]] = model_dump(exclude_none=True, by_alias=True)
return dict(dumped)
except TypeError:
dumped_fallback: Final[dict[str, object]] = model_dump()
@ -498,7 +498,7 @@ class _CiscoAIDefenseMcpMixin:
model_dump: Final = getattr(response, "model_dump", None)
if callable(model_dump):
try:
dumped = model_dump(exclude_none=True)
dumped = model_dump(exclude_none=True, by_alias=True)
except TypeError:
dumped = model_dump()
if isinstance(dumped, dict):

34
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_cisco_ai_defense_mcp.py
View file

@ -376,9 +376,10 @@ class TestCiscoAIDefenseMCPMode:
assert sent_payload["result"]["content"][0]["text"] == text_content
assert result is None
@pytest.mark.parametrize("use_wrapper", [True, False])
@pytest.mark.asyncio
async def test_mcp_response_hook_through_real_logging_wrapper(self):
from mcp.types import CallToolResult, TextContent
async def test_mcp_response_hook_through_real_logging_wrapper(self, use_wrapper):
from mcp.types import AudioContent, CallToolResult, EmbeddedResource, ImageContent, TextContent, TextResourceContents
from litellm.types.mcp import MCPPostCallResponseObject
@ -387,7 +388,14 @@ class TestCiscoAIDefenseMCPMode:
)
real_result = CallToolResult(
content=[TextContent(type="text", text="leak 9045629876")],
content=[
TextContent(type="text", text="leak 9045629876"),
ImageContent(type="image", data="aGVsbG8=", mimeType="image/png"),
AudioContent(type="audio", data="aGVsbG8=", mimeType="audio/wav"),
EmbeddedResource(type="resource", resource=TextResourceContents(
uri="memo://status", mimeType="text/plain", text="resource text"
)),
],
structuredContent={"patient": {"ssn": "123-45-6789"}},
isError=False,
)
@ -396,15 +404,6 @@ class TestCiscoAIDefenseMCPMode:
hidden_params={},
)
assert isinstance(wrapped.mcp_tool_call_response, list)
assert all(
isinstance(item, tuple) and len(item) == 2
for item in wrapped.mcp_tool_call_response
), (
"Pydantic coercion shape changed — update the normalizer to "
"match the new wire format."
)
post_mock = AsyncMock(return_value=_safe_response(url=MCP_URL))
with _patch_inspection_post(g, post_mock):
result = await g.async_post_mcp_tool_call_hook(
@ -414,7 +413,7 @@ class TestCiscoAIDefenseMCPMode:
"mcp_server_name": "vault",
"litellm_call_id": "real-wire-call",
},
response_obj=wrapped,
response_obj=wrapped if use_wrapper else real_result,
start_time=datetime.now(),
end_time=datetime.now(),
)
@ -428,8 +427,8 @@ class TestCiscoAIDefenseMCPMode:
sent_payload = post_mock.call_args.kwargs["json"]
content_items = sent_payload["result"]["content"]
assert len(content_items) == 1, (
f"expected exactly 1 content item from the real "
assert len(content_items) == 4, (
f"expected exactly 4 content items from the real "
f"CallToolResult.content list, got {len(content_items)}: "
f"{content_items!r}"
)
@ -441,6 +440,11 @@ class TestCiscoAIDefenseMCPMode:
f"``content`` field."
)
assert content_items[0].get("type") == "text"
assert content_items[1:] == [
{"type": "image", "data": "aGVsbG8=", "mimeType": "image/png"},
{"type": "audio", "data": "aGVsbG8=", "mimeType": "audio/wav"},
{"type": "resource", "resource": {"uri": "memo://status", "mimeType": "text/plain", "text": "resource text"}},
]
assert sent_payload["result"]["structuredContent"] == {
"patient": {"ssn": "123-45-6789"}
}