Merge

Author: ishymko

src/a2a/server/routes/rest_routes.py

@@ -14,6 +14,7 @@
 from a2a.server.routes import CallContextBuilder, DefaultCallContextBuilder
 from a2a.types.a2a_pb2 import AgentCard
 from a2a.utils.error_handlers import (
+    build_rest_error_payload,
     rest_error_handler,
     rest_stream_error_handler,
 )
@@ -25,6 +26,7 @@
 
 
 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
@@ -33,6 +35,7 @@
     _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
@@ -47,6 +50,7 @@
         Route = Any
         Mount = Any
         BaseRoute = Any
+        ServerSentEvent = Any
 
         _package_starlette_installed = False
 
@@ -168,10 +172,17 @@ async def _handle_streaming_request(
         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())
 

src/a2a/utils/error_handlers.py

@@ -173,7 +173,7 @@ async def error_catching_iterator() -> AsyncGenerator[
                     try:
                         async for item in original_iterator:
                             yield item
-                    except Exception as stream_error:  # noqa: BLE001
+                    except Exception as stream_error:
                         _log_error(stream_error)
                         raise stream_error
 

tests/server/routes/test_rest_routes.py

@@ -1,13 +1,26 @@
+import json
+
 from unittest.mock import AsyncMock
 
+import httpx
 import pytest
+
 from starlette.applications import Starlette
+from starlette.routing import BaseRoute
 from starlette.testclient import TestClient
-from starlette.routing import BaseRoute, Route
 
 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 AgentCard, Task, ListTasksResponse
+from a2a.types.a2a_pb2 import (
+    AgentCapabilities,
+    AgentCard,
+    ListTasksResponse,
+    Message,
+    Part,
+    Role,
+    Task,
+)
+from a2a.utils.errors import InternalError
 
 
 @pytest.fixture
@@ -103,3 +116,125 @@ def test_rest_list_tasks(agent_card, mock_handler):
     response = client.get('/tasks', headers={'A2A-Version': '1.0'})
     assert response.status_code == 200
     assert mock_handler.on_list_tasks.called
+
+
+@pytest.fixture
+def streaming_agent_card():
+    return AgentCard(
+        capabilities=AgentCapabilities(streaming=True),
+    )
+
+
+@pytest.fixture
+def streaming_app(streaming_agent_card, mock_handler):
+    routes = create_rest_routes(
+        agent_card=streaming_agent_card, request_handler=mock_handler
+    )
+    return Starlette(routes=routes)
+
+
+@pytest.fixture
+def streaming_client(streaming_app):
+    return httpx.AsyncClient(
+        transport=httpx.ASGITransport(app=streaming_app),
+        base_url='http://test',
+        headers={'A2A-Version': '1.0'},
+    )
+
+
+@pytest.mark.asyncio
+async def test_streaming_mid_stream_error_emits_sse_error_event(
+    streaming_client, mock_handler
+):
+    """Test that mid-stream errors are sent as SSE error events."""
+
+    async def mock_stream_then_error(*args, **kwargs):
+        yield Message(
+            message_id='stream_msg_1',
+            role=Role.ROLE_AGENT,
+            parts=[Part(text='First chunk')],
+        )
+        raise InternalError(message='Something went wrong mid-stream')
+
+    mock_handler.on_message_send_stream.side_effect = mock_stream_then_error
+
+    response = await streaming_client.post(
+        '/message:stream',
+        headers={'Accept': 'text/event-stream'},
+        json={},
+    )
+
+    response.raise_for_status()
+    assert 'text/event-stream' in response.headers.get('content-type', '')
+
+    lines = [line.strip() for line in response.text.strip().splitlines()]
+
+    # Should have a normal data event followed by an error event
+    data_lines = [
+        json.loads(line[6:]) for line in lines if line.startswith('data: ')
+    ]
+    assert len(data_lines) >= 1
+    assert 'message' in data_lines[0]
+    assert data_lines[0]['message']['messageId'] == 'stream_msg_1'
+
+    # Should contain an SSE error event
+    error_event_lines = [line for line in lines if line == 'event: error']
+    assert len(error_event_lines) == 1
+
+    # Find the error data after the error event
+    error_data = None
+    for i, line in enumerate(lines):
+        if line == 'event: error':
+            for j in range(i + 1, len(lines)):
+                if lines[j].startswith('data: '):
+                    error_data = json.loads(lines[j][6:])
+                    break
+            break
+
+    assert error_data is not None
+    assert error_data['error']['code'] == 500
+    assert error_data['error']['status'] == 'INTERNAL'
+    assert 'Something went wrong mid-stream' in error_data['error']['message']
+
+
+@pytest.mark.asyncio
+async def test_streaming_mid_stream_unknown_error_emits_sse_error_event(
+    streaming_client, mock_handler
+):
+    """Test that non-A2AError mid-stream errors also produce SSE error events."""
+
+    async def mock_stream_then_error(*args, **kwargs):
+        yield Message(
+            message_id='stream_msg_1',
+            role=Role.ROLE_AGENT,
+            parts=[Part(text='First chunk')],
+        )
+        raise RuntimeError('Unexpected failure')
+
+    mock_handler.on_message_send_stream.side_effect = mock_stream_then_error
+
+    response = await streaming_client.post(
+        '/message:stream',
+        headers={'Accept': 'text/event-stream'},
+        json={},
+    )
+
+    response.raise_for_status()
+
+    lines = [line.strip() for line in response.text.strip().splitlines()]
+
+    error_event_lines = [line for line in lines if line == 'event: error']
+    assert len(error_event_lines) == 1
+
+    error_data = None
+    for i, line in enumerate(lines):
+        if line == 'event: error':
+            for j in range(i + 1, len(lines)):
+                if lines[j].startswith('data: '):
+                    error_data = json.loads(lines[j][6:])
+                    break
+            break
+
+    assert error_data is not None
+    assert error_data['error']['code'] == 500
+    assert error_data['error']['status'] == 'INTERNAL'