From a0944e1638f528f3fc912d763586e312401f6316 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 28 May 2026 03:40:20 +0000
Subject: [PATCH] fix(a2a): align pydantic-ai test + docstring with
 direct-Message result shape

The non-streaming A2A response was changed so that `result` is the Message
itself (kind="message"), per spec / SendMessageResponse. Update the
PydanticAITransformation._transform_to_a2a_response test and docstring that
still described the old `result.message` envelope so internal consumers
match the producer.
---
 .../providers/pydantic_ai_agents/transformation.py   | 12 ++++++------
 .../test_pydantic_ai_agent_transformation.py         |  9 +++++----
 2 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py
index bb8a27f4eb5..bf68a01d98c 100644
--- a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py
+++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py
@@ -289,16 +289,16 @@ class PydanticAITransformation:
         Transform Pydantic AI task response to standard A2A non-streaming format.
 
         Pydantic AI returns a task with history/artifacts, but the standard A2A
-        non-streaming format expects:
+        non-streaming format expects ``result`` to be the Message directly
+        (``kind="message"``), per the A2A spec / ``SendMessageResponse``:
         {
             "jsonrpc": "2.0",
             "id": "...",
             "result": {
-                "message": {
-                    "role": "agent",
-                    "parts": [{"kind": "text", "text": "..."}],
-                    "messageId": "..."
-                }
+                "kind": "message",
+                "role": "agent",
+                "parts": [{"kind": "text", "text": "..."}],
+                "messageId": "..."
             }
         }
 
diff --git a/tests/litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py b/tests/litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py
index efbb628ee6d..717a7c902b5 100644
--- a/tests/litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py
+++ b/tests/litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py
@@ -90,9 +90,10 @@ class TestPydanticAITransformation:
             request_id="req-123",
         )
 
-        # Should return standard A2A format with message
+        # Should return standard A2A non-streaming format where `result` is the
+        # Message itself (kind="message"), per A2A spec / SendMessageResponse.
         assert result["jsonrpc"] == "2.0"
         assert result["id"] == "req-123"
-        assert "message" in result["result"]
-        assert result["result"]["message"]["role"] == "agent"
-        assert result["result"]["message"]["parts"][0]["text"] == "The answer is 4."
+        assert result["result"]["kind"] == "message"
+        assert result["result"]["role"] == "agent"
+        assert result["result"]["parts"][0]["text"] == "The answer is 4."