fix: handle SSE errors occurred after stream started

Currently it'd close the connection.

Author: ishymko

src/a2a/client/transports/http_helpers.py

@@ -69,11 +69,23 @@ async def send_http_stream_request(
     httpx_client: httpx.AsyncClient,
     method: str,
     url: str,
-    status_error_handler: Callable[[httpx.HTTPStatusError], NoReturn]
-    | None = None,
+    status_error_handler: Callable[[httpx.HTTPStatusError], NoReturn],
+    sse_error_handler: Callable[[str], NoReturn],
     **kwargs: Any,
 ) -> AsyncGenerator[str]:
-    """Sends a streaming HTTP request, yielding SSE data strings and handling exceptions."""
+    """Sends a streaming HTTP request, yielding SSE data strings and handling exceptions.
+
+    Args:
+        httpx_client: The async HTTP client.
+        method: The HTTP method (e.g. 'POST', 'GET').
+        url: The URL to send the request to.
+        status_error_handler: Handler for HTTP status errors. Should raise an
+            appropriate domain-specific exception.
+        sse_error_handler: Handler for SSE error events. Called with the
+            raw SSE data string when an ``event: error`` SSE event is received.
+            Should raise an appropriate domain-specific exception.
+        **kwargs: Additional keyword arguments forwarded to ``aconnect_sse``.
+    """
     with handle_http_exceptions(status_error_handler):
         async with aconnect_sse(
             httpx_client, method, url, **kwargs
@@ -97,4 +109,6 @@ async def send_http_stream_request(
             async for sse in event_source.aiter_sse():
                 if not sse.data:
                     continue
+                if sse.event == 'error':
+                    sse_error_handler(sse.data)
                 yield sse.data

src/a2a/client/transports/jsonrpc.py

@@ -1,7 +1,7 @@
 import logging
 
 from collections.abc import AsyncGenerator
-from typing import Any
+from typing import Any, NoReturn
 from uuid import uuid4
 
 import httpx
@@ -349,6 +349,7 @@ async def _send_stream_request(
             'POST',
             self.url,
             None,
+            self._handle_sse_error,
             json=rpc_request_payload,
             **http_kwargs,
         ):
@@ -359,3 +360,10 @@ async def _send_stream_request(
                 json_rpc_response.result, StreamResponse()
             )
             yield response
+
+    def _handle_sse_error(self, sse_data: str) -> NoReturn:
+        """Handles SSE error events by parsing JSON-RPC error payload and raising the appropriate domain error."""
+        json_rpc_response = JSONRPC20Response.from_json(sse_data)
+        if json_rpc_response.error:
+            raise self._create_jsonrpc_error(json_rpc_response.error)
+        raise A2AClientError(f'SSE stream error: {sse_data}')

src/a2a/client/transports/rest.py

@@ -41,6 +41,47 @@
 logger = logging.getLogger(__name__)
 
 
+def _parse_rest_error(
+    error_payload: dict[str, Any],
+    fallback_message: str,
+) -> Exception | None:
+    """Parses a REST error payload and returns the appropriate A2AError.
+
+    Args:
+        error_payload: The parsed JSON error payload.
+        fallback_message: Message to use if the payload has no ``message``.
+
+    Returns:
+        The mapped A2AError if a known reason was found, otherwise ``None``.
+    """
+    error_data = error_payload.get('error', {})
+    message = error_data.get('message', fallback_message)
+    details = error_data.get('details', [])
+    if not isinstance(details, list):
+        return None
+
+    # The `details` array can contain multiple different error objects.
+    # We extract the first `ErrorInfo` object because it contains the
+    # specific `reason` code needed to map this back to a Python A2AError.
+    for d in details:
+        if (
+            isinstance(d, dict)
+            and d.get('@type') == 'type.googleapis.com/google.rpc.ErrorInfo'
+        ):
+            reason = d.get('reason')
+            metadata = d.get('metadata') or {}
+            if isinstance(reason, str):
+                exception_cls = A2A_REASON_TO_ERROR.get(reason)
+                if exception_cls:
+                    exc = exception_cls(message)
+                    if metadata:
+                        exc.data = metadata
+                    return exc
+            break
+
+    return None
+
+
 @trace_class(kind=SpanKind.CLIENT)
 class RestTransport(ClientTransport):
     """A REST transport for the A2A client."""
@@ -294,39 +335,12 @@ def _handle_http_error(self, e: httpx.HTTPStatusError) -> NoReturn:
         """Handles HTTP status errors and raises the appropriate A2AError."""
         try:
             error_payload = e.response.json()
-            error_data = error_payload.get('error', {})
-
-            message = error_data.get('message', str(e))
-            details = error_data.get('details', [])
-            if not isinstance(details, list):
-                details = []
-
-            # The `details` array can contain multiple different error objects.
-            # We extract the first `ErrorInfo` object because it contains the
-            # specific `reason` code needed to map this back to a Python A2AError.
-            error_info = {}
-            for d in details:
-                if (
-                    isinstance(d, dict)
-                    and d.get('@type')
-                    == 'type.googleapis.com/google.rpc.ErrorInfo'
-                ):
-                    error_info = d
-                    break
-            reason = error_info.get('reason')
-            metadata = error_info.get('metadata') or {}
-
-            if isinstance(reason, str):
-                exception_cls = A2A_REASON_TO_ERROR.get(reason)
-                if exception_cls:
-                    exc = exception_cls(message)
-                    if metadata:
-                        exc.data = metadata
-                    raise exc from e
+            mapped = _parse_rest_error(error_payload, str(e))
+            if mapped:
+                raise mapped from e
         except (json.JSONDecodeError, ValueError):
             pass
 
-        # Fallback mappings for status codes if 'type' is missing or unknown
         status_code = e.response.status_code
         if status_code == httpx.codes.NOT_FOUND:
             raise MethodNotFoundError(
@@ -335,6 +349,14 @@ def _handle_http_error(self, e: httpx.HTTPStatusError) -> NoReturn:
 
         raise A2AClientError(f'HTTP Error {status_code}: {e}') from e
 
+    def _handle_sse_error(self, sse_data: str) -> NoReturn:
+        """Handles SSE error events by parsing the REST error payload and raising the appropriate A2AError."""
+        error_payload = json.loads(sse_data)
+        mapped = _parse_rest_error(error_payload, sse_data)
+        if mapped:
+            raise mapped
+        raise A2AClientError(sse_data)
+
     async def _send_stream_request(
         self,
         method: str,
@@ -352,6 +374,7 @@ async def _send_stream_request(
             method,
             f'{self.url}{path}',
             self._handle_http_error,
+            self._handle_sse_error,
             json=json,
             **http_kwargs,
         ):

src/a2a/compat/v0_3/jsonrpc_adapter.py

@@ -306,9 +306,10 @@ async def event_generator(
                     )
                 )
                 yield {
+                    'event': 'error',
                     'data': err_resp.model_dump_json(
                         by_alias=True, exclude_none=True
-                    )
+                    ),
                 }
 
         return EventSourceResponse(event_generator(stream_gen))

src/a2a/compat/v0_3/jsonrpc_transport.py

@@ -415,6 +415,13 @@ def _handle_http_error(self, e: httpx.HTTPStatusError) -> NoReturn:
         """Handles HTTP errors for standard requests."""
         raise A2AClientError(f'HTTP Error: {e.response.status_code}') from e
 
+    def _handle_sse_error(self, sse_data: str) -> NoReturn:
+        """Handles SSE error events by parsing JSON-RPC error payload and raising the appropriate domain error."""
+        data = json.loads(sse_data)
+        if 'error' in data:
+            raise self._create_jsonrpc_error(data['error'])
+        raise A2AClientError(f'SSE stream error: {sse_data}')
+
     async def _send_stream_request(
         self,
         json_data: dict[str, Any],
@@ -430,6 +437,7 @@ async def _send_stream_request(
             'POST',
             self.url,
             self._handle_http_error,
+            self._handle_sse_error,
             json=json_data,
             **http_kwargs,
         ):

src/a2a/compat/v0_3/rest_adapter.py

@@ -7,6 +7,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
@@ -17,6 +18,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
@@ -27,6 +29,7 @@
         Request = Any
         JSONResponse = Any
         Response = Any
+        ServerSentEvent = Any
 
         _package_starlette_installed = False
 
@@ -37,6 +40,7 @@
 from a2a.server.context import ServerCallContext
 from a2a.server.routes import CallContextBuilder, DefaultCallContextBuilder
 from a2a.utils.error_handlers import (
+    build_rest_error_payload,
     rest_error_handler,
     rest_stream_error_handler,
 )
@@ -101,9 +105,16 @@ async def _handle_streaming_request(
 
         async def event_generator(
             stream: AsyncIterable[Any],
-        ) -> AsyncIterator[str]:
-            async for item in stream:
-                yield json.dumps(item)
+        ) -> AsyncIterator[str | ServerSentEvent]:
+            try:
+                async for item in stream:
+                    yield json.dumps(item)
+            except Exception as e:
+                logger.exception('Error during v0.3 REST SSE stream')
+                yield ServerSentEvent(
+                    data=json.dumps(build_rest_error_payload(e)),
+                    event='error',
+                )
 
         return EventSourceResponse(
             event_generator(method(request, call_context))

src/a2a/compat/v0_3/rest_transport.py

@@ -44,7 +44,11 @@
     TaskPushNotificationConfig,
 )
 from a2a.utils.constants import PROTOCOL_VERSION_0_3, VERSION_HEADER
-from a2a.utils.errors import JSON_RPC_ERROR_CODE_MAP, MethodNotFoundError
+from a2a.utils.errors import (
+    A2A_REASON_TO_ERROR,
+    JSON_RPC_ERROR_CODE_MAP,
+    MethodNotFoundError,
+)
 from a2a.utils.telemetry import SpanKind, trace_class
 
 
@@ -369,6 +373,30 @@ def _handle_http_error(self, e: httpx.HTTPStatusError) -> NoReturn:
 
         raise A2AClientError(f'HTTP Error {status_code}: {e}') from e
 
+    def _handle_sse_error(self, sse_data: str) -> NoReturn:
+        """Handles SSE error events by parsing the REST error payload and raising the appropriate A2AError."""
+        error_payload = json.loads(sse_data)
+        error_data = error_payload.get('error', {})
+
+        message = error_data.get('message', sse_data)
+        details = error_data.get('details', [])
+        if not isinstance(details, list):
+            details = []
+
+        for d in details:
+            if (
+                isinstance(d, dict)
+                and d.get('@type') == 'type.googleapis.com/google.rpc.ErrorInfo'
+            ):
+                reason = d.get('reason')
+                if isinstance(reason, str):
+                    exception_cls = A2A_REASON_TO_ERROR.get(reason)
+                    if exception_cls:
+                        raise exception_cls(message)
+                break
+
+        raise A2AClientError(message)
+
     async def _send_stream_request(
         self,
         method: str,
@@ -386,6 +414,7 @@ async def _send_stream_request(
             method,
             f'{self.url}{path}',
             self._handle_http_error,
+            self._handle_sse_error,
             json=json,
             **http_kwargs,
         ):

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

@@ -12,6 +12,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
@@ -20,6 +21,7 @@
 
 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
@@ -30,6 +32,7 @@
         Request = Any
         JSONResponse = Any
         Response = Any
+        ServerSentEvent = Any
 
         _package_starlette_installed = False
 
@@ -42,6 +45,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,
 )
@@ -163,10 +167,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/server/routes/jsonrpc_dispatcher.py

@@ -559,8 +559,30 @@ def _create_response(
             async def event_generator(
                 stream: AsyncGenerator[dict[str, Any]],
             ) -> AsyncGenerator[dict[str, str]]:
-                async for item in stream:
-                    yield {'data': json.dumps(item)}
+                try:
+                    async for item in stream:
+                        event: dict[str, str] = {
+                            'data': json.dumps(item),
+                        }
+                        if 'error' in item:
+                            event['event'] = 'error'
+                        yield event
+                except Exception as e:
+                    logger.exception(
+                        'Unhandled error during JSON-RPC SSE stream'
+                    )
+                    rpc_error: A2AError | JSONRPCError = (
+                        e
+                        if isinstance(e, A2AError | JSONRPCError)
+                        else InternalError(message=str(e))
+                    )
+                    error_response = build_error_response(
+                        context.state.get('request_id'), rpc_error
+                    )
+                    yield {
+                        'event': 'error',
+                        'data': json.dumps(error_response),
+                    }
 
             return EventSourceResponse(
                 event_generator(handler_result), headers=headers