WIP

Author: ishymko

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

@@ -482,7 +482,7 @@ async def _process_streaming_request(
                 request_obj, context
             )
 
-        return await self._create_response(context, handler_result)
+        return 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 await self._create_response(context, handler_result)
+        return self._create_response(context, handler_result)
 
-    async def _create_response(
+    def _create_response(
         self,
         context: ServerCallContext,
         handler_result: AsyncGenerator[dict[str, Any]] | dict[str, Any],
@@ -587,35 +587,15 @@ async 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]],
-                first_item: dict[str, Any],
-            ) -> AsyncGenerator[dict[str, str], None]:
-                yield {'data': json.dumps(first_item)}
+            ) -> AsyncGenerator[dict[str, str]]:
                 async for item in stream:
                     yield {'data': json.dumps(item)}
 
             return EventSourceResponse(
-                event_generator(handler_result, first_item), headers=headers
+                event_generator(handler_result), headers=headers
             )
 
         # handler_result is a dict (JSON-RPC response)

src/a2a/utils/error_handlers.py

@@ -165,26 +165,17 @@ async def wrapper(*args: Any, **kwargs: Any) -> Any:
             ):
                 original_iterator = response.body_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()
+                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()
 
         except Exception as e:  # noqa: BLE001
             return _create_error_response(e)

src/a2a/utils/helpers.py

@@ -304,13 +304,12 @@ def validate_async_generator(
 
     def decorator(function):
         @functools.wraps(function)
-        async def wrapper(self, *args, **kwargs):
+        def wrapper(self, *args, **kwargs):
             if not expression(self):
                 final_message = error_message or str(expression)
                 logger.error('Unsupported Operation: %s', final_message)
                 raise UnsupportedOperationError(message=final_message)
-            async for i in function(self, *args, **kwargs):
-                yield i
+            return function(self, *args, **kwargs)
 
         return wrapper