WIP

Author: ishymko

src/a2a/client/transports/http_helpers.py

@@ -78,8 +78,20 @@ async def send_http_stream_request(
         async with aconnect_sse(
             httpx_client, method, url, **kwargs
         ) as event_source:
-            event_source.response.raise_for_status()
-            async for sse in event_source.aiter_sse():
-                if not sse.data:
-                    continue
-                yield sse.data
+            try:
+                event_source.response.raise_for_status()
+            except httpx.HTTPStatusError as e:
+                await event_source.response.aread()
+                raise e
+
+            try:
+                async for sse in event_source.aiter_sse():
+                    if not sse.data:
+                        continue
+                    yield sse.data
+            except SSEError as e:
+                if 'application/json' in event_source.response.headers.get('content-type', ''):
+                    content = await event_source.response.aread()
+                    yield content.decode('utf-8')
+                else:
+                    raise e

src/a2a/server/apps/jsonrpc/jsonrpc_app.py

@@ -482,7 +482,7 @@ async def _process_streaming_request(
                 request_obj, context
             )
 
-        return self._create_response(context, handler_result)
+        return await self._create_response(context, handler_result)
 
     async def _process_non_streaming_request(
         self,
@@ -562,9 +562,9 @@ async def _process_non_streaming_request(
                 )
                 return self._generate_error_response(request_id, error)
 
-        return self._create_response(context, handler_result)
+        return await self._create_response(context, handler_result)
 
-    def _create_response(
+    async def _create_response(
         self,
         context: ServerCallContext,
         handler_result: AsyncGenerator[dict[str, Any]] | dict[str, Any],
@@ -587,15 +587,35 @@ def _create_response(
         if exts := context.activated_extensions:
             headers[HTTP_EXTENSION_HEADER] = ', '.join(sorted(exts))
         if isinstance(handler_result, AsyncGenerator):
+            try:
+                # Prime to see if it fails upfront
+                first_item = await handler_result.__anext__()
+            except StopAsyncIteration:
+
+                async def empty_generator() -> AsyncGenerator[dict[str, str], None]:
+                    if False:
+                        yield {}
+
+                return EventSourceResponse(empty_generator(), headers=headers)
+            except Exception as e:
+                logger.debug('Upfront exception in streaming handler: %s', e)
+                if not isinstance(e, A2AError | JSONRPCError):
+                    e = InternalError(message=str(e))
+                request_id = context.state.get('request_id')
+                error_payload = build_error_response(request_id, e)
+                return JSONResponse(error_payload, headers=headers)
+
             # Result is a stream of dict objects
             async def event_generator(
                 stream: AsyncGenerator[dict[str, Any]],
-            ) -> AsyncGenerator[dict[str, str]]:
+                first_item: dict[str, Any],
+            ) -> AsyncGenerator[dict[str, str], None]:
+                yield {'data': json.dumps(first_item)}
                 async for item in stream:
                     yield {'data': json.dumps(item)}
 
             return EventSourceResponse(
-                event_generator(handler_result), headers=headers
+                event_generator(handler_result, first_item), headers=headers
             )
 
         # handler_result is a dict (JSON-RPC response)

src/a2a/utils/error_handlers.py

@@ -165,17 +165,26 @@ async def wrapper(*args: Any, **kwargs: Any) -> Any:
             ):
                 original_iterator = response.body_iterator
 
-                async def error_catching_iterator() -> AsyncGenerator[
-                    Any, None
-                ]:
-                    try:
-                        async for item in original_iterator:
-                            yield item
-                    except Exception as stream_error:
-                        _log_error(stream_error)
-                        raise stream_error
-
-                response.body_iterator = error_catching_iterator()
+                try:
+                    # Prime the stream to catch upfront errors
+                    first_item = await original_iterator.__anext__()
+                except StopAsyncIteration:
+                    # Stream is empty
+                    pass
+                except Exception as e:  # noqa: BLE001
+                    return _create_error_response(e)
+                else:
+
+                    async def error_catching_iterator() -> AsyncGenerator[Any, None]:
+                        yield first_item
+                        try:
+                            async for item in original_iterator:
+                                yield item
+                        except Exception as stream_error:
+                            _log_error(stream_error)
+                            raise stream_error
+
+                    response.body_iterator = error_catching_iterator()
 
         except Exception as e:  # noqa: BLE001
             return _create_error_response(e)