Merge branch '1.0-dev' into fix-run-itk-locally

Author: ishymko

src/a2a/client/transports/http_helpers.py

@@ -12,6 +12,10 @@
 from a2a.client.errors import A2AClientError, A2AClientTimeoutError
 
 
+def _default_sse_error_handler(sse_data: str) -> NoReturn:
+    raise A2AClientError(f'SSE stream error event received: {sse_data}')
+
+
 @contextmanager
 def handle_http_exceptions(
     status_error_handler: Callable[[httpx.HTTPStatusError], NoReturn]
@@ -71,9 +75,22 @@ async def send_http_stream_request(
     url: str,
     status_error_handler: Callable[[httpx.HTTPStatusError], NoReturn]
     | None = None,
+    sse_error_handler: Callable[[str], NoReturn] = _default_sse_error_handler,
     **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 _SSEEventSource(
             httpx_client, method, url, **kwargs
@@ -97,6 +114,8 @@ 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
@@ -350,6 +350,7 @@ async def _send_stream_request(
             'POST',
             self.url,
             None,
+            self._handle_sse_error,
             json=rpc_request_payload,
             **http_kwargs,
         ):
@@ -360,3 +361,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/server/agent_execution/active_task.py

@@ -374,30 +374,33 @@ async def _run_consumer(self) -> None:  # noqa: PLR0915, PLR0912
                                 await self._task_manager.process(event)
 
                                 # Check for AUTH_REQUIRED or INPUT_REQUIRED or TERMINAL states
-                                res = await self._task_manager.get_task()
+                                new_task = await self._task_manager.get_task()
+                                if new_task is None:
+                                    raise RuntimeError(
+                                        f'Task {self.task_id} not found'
+                                    )
                                 is_interrupted = (
-                                    res
-                                    and res.status.state
+                                    new_task.status.state
                                     in INTERRUPTED_TASK_STATES
                                 )
                                 is_terminal = (
-                                    res
-                                    and res.status.state in TERMINAL_TASK_STATES
+                                    new_task.status.state
+                                    in TERMINAL_TASK_STATES
                                 )
 
                                 # If we hit a breakpoint or terminal state, lock in the result.
-                                if (is_interrupted or is_terminal) and res:
+                                if is_interrupted or is_terminal:
                                     logger.debug(
                                         'Consumer[%s]: Setting first result as Task (state=%s)',
                                         self._task_id,
-                                        res.status.state,
+                                        new_task.status.state,
                                     )
 
                                 if is_terminal:
                                     logger.debug(
                                         'Consumer[%s]: Reached terminal state %s',
                                         self._task_id,
-                                        res.status.state if res else 'unknown',
+                                        new_task.status.state,
                                     )
                                     if not self._is_finished.is_set():
                                         async with self._lock:
@@ -413,7 +416,7 @@ async def _run_consumer(self) -> None:  # noqa: PLR0915, PLR0912
                                     logger.debug(
                                         'Consumer[%s]: Interrupted with state %s',
                                         self._task_id,
-                                        res.status.state if res else 'unknown',
+                                        new_task.status.state,
                                     )
 
                                 if (

src/a2a/server/events/event_consumer.py

@@ -5,7 +5,7 @@
 
 from pydantic import ValidationError
 
-from a2a.server.events.event_queue import Event, EventQueue, QueueShutDown
+from a2a.server.events.event_queue import Event, EventQueueLegacy, QueueShutDown
 from a2a.types.a2a_pb2 import (
     Message,
     Task,
@@ -22,7 +22,7 @@
 class EventConsumer:
     """Consumer to read events from the agent event queue."""
 
-    def __init__(self, queue: EventQueue):
+    def __init__(self, queue: EventQueueLegacy):
         """Initializes the EventConsumer.
 
         Args:

src/a2a/server/events/event_queue.py

@@ -92,73 +92,6 @@ async def enqueue_event(self, event: Event) -> None:
         Only main queue can enqueue events. Child queues can only dequeue events.
         """
 
-    @abstractmethod
-    async def dequeue_event(self) -> Event:
-        """Pulls an event from the queue."""
-
-    @abstractmethod
-    def task_done(self) -> None:
-        """Signals that a work on dequeued event is complete."""
-
-    @abstractmethod
-    async def tap(
-        self, max_queue_size: int = DEFAULT_MAX_QUEUE_SIZE
-    ) -> 'EventQueue':
-        """Creates a child queue that receives future events.
-
-        Note: The tapped queue may receive some old events if the incoming event
-        queue is lagging behind and hasn't dispatched them yet.
-        """
-
-    @abstractmethod
-    async def close(self, immediate: bool = False) -> None:
-        """Closes the queue.
-
-        For parent queue: it closes the main queue and all its child queues.
-        For child queue: it closes only child queue.
-
-        It is safe to call it multiple times.
-        If immediate is True, the queue will be closed without waiting for all events to be processed.
-        If immediate is False, the queue will be closed after all events are processed (and confirmed with task_done() calls).
-
-        WARNING: Closing the parent queue with immediate=False is a deadlock risk if there are unconsumed events
-        in any of the child sinks and the consumer has crashed without draining its queue.
-        It is highly recommended to wrap graceful shutdowns with a timeout, e.g.,
-        `asyncio.wait_for(queue.close(immediate=False), timeout=...)`.
-        """
-
-    @abstractmethod
-    def is_closed(self) -> bool:
-        """[DEPRECATED] Checks if the queue is closed.
-
-        NOTE: Relying on this for enqueue logic introduces race conditions.
-        It is maintained primarily for backwards compatibility, workarounds for
-        Python 3.10/3.12 async queues in consumers, and for the test suite.
-        """
-
-    @abstractmethod
-    async def __aenter__(self) -> Self:
-        """Enters the async context manager, returning the queue itself.
-
-        WARNING: See `__aexit__` for important deadlock risks associated with
-        exiting this context manager if unconsumed events remain.
-        """
-
-    @abstractmethod
-    async def __aexit__(
-        self,
-        exc_type: type[BaseException] | None,
-        exc_val: BaseException | None,
-        exc_tb: TracebackType | None,
-    ) -> None:
-        """Exits the async context manager, ensuring close() is called.
-
-        WARNING: The context manager calls `close(immediate=False)` by default.
-        If a consumer exits the `async with` block early (e.g., due to an exception
-        or an explicit `break`) while unconsumed events remain in the queue,
-        `__aexit__` will deadlock waiting for `task_done()` to be called on those events.
-        """
-
 
 @trace_class(kind=SpanKind.SERVER)
 class EventQueueLegacy(EventQueue):
@@ -180,7 +113,7 @@ def __init__(self, max_queue_size: int = DEFAULT_MAX_QUEUE_SIZE) -> None:
         self._queue: AsyncQueue[Event] = _create_async_queue(
             maxsize=max_queue_size
         )
-        self._children: list[EventQueue] = []
+        self._children: list[EventQueueLegacy] = []
         self._is_closed = False
         self._lock = asyncio.Lock()
         logger.debug('EventQueue initialized.')