diff --git a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense_mcp.py b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense_mcp.py index 15b4f713a50..67ef05fc324 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense_mcp.py +++ b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense_mcp.py @@ -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): diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_cisco_ai_defense_mcp.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_cisco_ai_defense_mcp.py index bc784923eb5..826edab694d 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_cisco_ai_defense_mcp.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_cisco_ai_defense_mcp.py @@ -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"} }