Merge

Author: ishymko

src/a2a/server/routes/rest_dispatcher.py

@@ -20,6 +20,7 @@
 )
 from a2a.utils import constants, proto_utils
 from a2a.utils.error_handlers import (
+    build_rest_error_payload,
     rest_error_handler,
     rest_stream_error_handler,
 )
@@ -32,20 +33,23 @@
 
 
 if TYPE_CHECKING:
+    from sse_starlette.event import ServerSentEvent
     from sse_starlette.sse import EventSourceResponse
     from starlette.requests import Request
     from starlette.responses import JSONResponse, Response
 
     _package_starlette_installed = True
 else:
     try:
+        from sse_starlette.event import ServerSentEvent
         from sse_starlette.sse import EventSourceResponse
         from starlette.requests import Request
         from starlette.responses import JSONResponse, Response
 
         _package_starlette_installed = True
     except ImportError:
         EventSourceResponse = Any
+        ServerSentEvent = Any
         Request = Any
         JSONResponse = Any
         Response = Any
@@ -135,10 +139,17 @@ async def _handle_streaming(
         except StopAsyncIteration:
             return EventSourceResponse(iter([]))
 
-        async def event_generator() -> AsyncIterator[str]:
+        async def event_generator() -> AsyncIterator[str | ServerSentEvent]:
             yield json.dumps(first_item)
-            async for item in stream:
-                yield json.dumps(item)
+            try:
+                async for item in stream:
+                    yield json.dumps(item)
+            except Exception as e:
+                logger.exception('Error during REST SSE stream')
+                yield ServerSentEvent(
+                    data=json.dumps(build_rest_error_payload(e)),
+                    event='error',
+                )
 
         return EventSourceResponse(event_generator())
 

tests/server/routes/test_rest_routes.py

@@ -12,7 +12,6 @@
 from a2a.server.request_handlers.request_handler import RequestHandler
 from a2a.server.routes.rest_routes import create_rest_routes
 from a2a.types.a2a_pb2 import (
-    AgentCapabilities,
     AgentCard,
     ListTasksResponse,
     Message,
@@ -108,16 +107,9 @@ def test_rest_list_tasks(agent_card, mock_handler):
 
 
 @pytest.fixture
-def streaming_agent_card():
-    return AgentCard(
-        capabilities=AgentCapabilities(streaming=True),
-    )
-
-
-@pytest.fixture
-def streaming_app(streaming_agent_card, mock_handler):
+def streaming_app(mock_handler):
     routes = create_rest_routes(
-        agent_card=streaming_agent_card, request_handler=mock_handler
+        request_handler=mock_handler,
     )
     return Starlette(routes=routes)