fix: Avoid sending a single json message in multiple chunks. (#837)

Fixes #742

Author: bartek-w

src/a2a/compat/v0_3/rest_adapter.py

@@ -1,4 +1,5 @@
 import functools
+import json
 import logging
 
 from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable
@@ -103,9 +104,9 @@ async def _handle_streaming_request(
 
         async def event_generator(
             stream: AsyncIterable[Any],
-        ) -> AsyncIterator[dict[str, dict[str, Any]]]:
+        ) -> AsyncIterator[str]:
             async for item in stream:
-                yield {'data': item}
+                yield json.dumps(item)
 
         return EventSourceResponse(
             event_generator(method(request, call_context))

src/a2a/compat/v0_3/rest_handler.py

@@ -1,9 +1,9 @@
 import logging
 
-from collections.abc import AsyncIterable, AsyncIterator
+from collections.abc import AsyncIterator
 from typing import TYPE_CHECKING, Any
 
-from google.protobuf.json_format import MessageToDict, MessageToJson, Parse
+from google.protobuf.json_format import MessageToDict, Parse
 
 
 if TYPE_CHECKING:
@@ -86,7 +86,7 @@ async def on_message_send_stream(
         self,
         request: Request,
         context: ServerCallContext,
-    ) -> AsyncIterator[str]:
+    ) -> AsyncIterator[dict[str, Any]]:
         """Handles the 'message/stream' REST method.
 
         Args:
@@ -108,7 +108,7 @@ async def on_message_send_stream(
             v03_pb_resp = proto_utils.ToProto.stream_response(
                 v03_stream_resp.result
             )
-            yield MessageToJson(v03_pb_resp)
+            yield MessageToDict(v03_pb_resp)
 
     async def on_cancel_task(
         self,
@@ -142,7 +142,7 @@ async def on_subscribe_to_task(
         self,
         request: Request,
         context: ServerCallContext,
-    ) -> AsyncIterable[str]:
+    ) -> AsyncIterator[dict[str, Any]]:
         """Handles the 'tasks/{id}:subscribe' REST method.
 
         Args:
@@ -164,7 +164,7 @@ async def on_subscribe_to_task(
             v03_pb_resp = proto_utils.ToProto.stream_response(
                 v03_stream_resp.result
             )
-            yield MessageToJson(v03_pb_resp)
+            yield MessageToDict(v03_pb_resp)
 
     async def get_push_notification(
         self,

src/a2a/server/apps/rest/rest_adapter.py

@@ -1,4 +1,5 @@
 import functools
+import json
 import logging
 
 from abc import ABC, abstractmethod
@@ -150,9 +151,9 @@ async def _handle_streaming_request(
 
         async def event_generator(
             stream: AsyncIterable[Any],
-        ) -> AsyncIterator[dict[str, dict[str, Any]]]:
+        ) -> AsyncIterator[str]:
             async for item in stream:
-                yield {'data': item}
+                yield json.dumps(item)
 
         return EventSourceResponse(
             event_generator(method(request, call_context))

src/a2a/server/request_handlers/rest_handler.py

@@ -1,11 +1,10 @@
 import logging
 
-from collections.abc import AsyncIterable, AsyncIterator
+from collections.abc import AsyncIterator
 from typing import TYPE_CHECKING, Any
 
 from google.protobuf.json_format import (
     MessageToDict,
-    MessageToJson,
     Parse,
 )
 
@@ -96,7 +95,7 @@ async def on_message_send_stream(
         self,
         request: Request,
         context: ServerCallContext,
-    ) -> AsyncIterator[str]:
+    ) -> AsyncIterator[dict[str, Any]]:
         """Handles the 'message/stream' REST method.
 
         Yields response objects as they are produced by the underlying handler's stream.
@@ -116,7 +115,7 @@ async def on_message_send_stream(
             params, context
         ):
             response = proto_utils.to_stream_response(event)
-            yield MessageToJson(response)
+            yield MessageToDict(response)
 
     async def on_cancel_task(
         self,
@@ -148,7 +147,7 @@ async def on_subscribe_to_task(
         self,
         request: Request,
         context: ServerCallContext,
-    ) -> AsyncIterable[str]:
+    ) -> AsyncIterator[dict[str, Any]]:
         """Handles the 'SubscribeToTask' REST method.
 
         Yields response objects as they are produced by the underlying handler's stream.
@@ -164,7 +163,7 @@ async def on_subscribe_to_task(
         async for event in self.request_handler.on_subscribe_to_task(
             SubscribeToTaskRequest(id=task_id), context
         ):
-            yield MessageToJson(proto_utils.to_stream_response(event))
+            yield MessageToDict(proto_utils.to_stream_response(event))
 
     async def get_push_notification(
         self,

tests/compat/v0_3/test_rest_handler.py

@@ -110,7 +110,7 @@ async def mock_stream(*args, **kwargs):
     )
 
     results = [
-        json.loads(chunk)
+        chunk
         async for chunk in rest_handler.on_message_send_stream(
             mock_request, mock_context
         )
@@ -169,7 +169,7 @@ async def mock_stream(*args, **kwargs):
     )
 
     results = [
-        json.loads(chunk)
+        chunk
         async for chunk in rest_handler.on_subscribe_to_task(
             mock_request, mock_context
         )

tests/server/apps/rest/test_rest_fastapi_app.py

@@ -1,4 +1,5 @@
 import logging
+import json
 
 from typing import Any
 from unittest.mock import MagicMock
@@ -339,6 +340,71 @@ async def mock_stream_response():
     request_handler.on_message_send_stream.assert_called_once()
 
 
+@pytest.mark.anyio
+async def test_streaming_content_verification(
+    streaming_client: AsyncClient, request_handler: MagicMock
+) -> None:
+    """Test that streaming endpoint returns correct SSE content."""
+
+    async def mock_stream_response():
+        yield Message(
+            message_id='stream_msg_1',
+            role=Role.ROLE_AGENT,
+            parts=[Part(text='First chunk')],
+        )
+        yield Message(
+            message_id='stream_msg_2',
+            role=Role.ROLE_AGENT,
+            parts=[Part(text='Second chunk')],
+        )
+
+    request_handler.on_message_send_stream.return_value = mock_stream_response()
+
+    request = a2a_pb2.SendMessageRequest(
+        message=a2a_pb2.Message(
+            message_id='test_stream_msg',
+            role=a2a_pb2.ROLE_USER,
+            parts=[a2a_pb2.Part(text='Test message')],
+        ),
+    )
+
+    response = await streaming_client.post(
+        '/message:stream',
+        json=json_format.MessageToDict(request),
+        headers={'Accept': 'text/event-stream'},
+    )
+
+    response.raise_for_status()
+
+    # Read the response content
+    lines = [line async for line in response.aiter_lines()]
+
+    # SSE format is "data: <json>\n\n"
+    # httpx.aiter_lines() will give us each line.
+    data_lines = [
+        json.loads(line[6:]) for line in lines if line.startswith('data: ')
+    ]
+
+    expected_data_lines = [
+        {
+            'message': {
+                'messageId': 'stream_msg_1',
+                'role': 'ROLE_AGENT',
+                'parts': [{'text': 'First chunk'}],
+            }
+        },
+        {
+            'message': {
+                'messageId': 'stream_msg_2',
+                'role': 'ROLE_AGENT',
+                'parts': [{'text': 'Second chunk'}],
+            }
+        },
+    ]
+
+    assert data_lines == expected_data_lines
+
+
 @pytest.mark.anyio
 async def test_streaming_endpoint_with_invalid_content_type(
     streaming_client: AsyncClient, request_handler: MagicMock