Merge branch 'main' into maeste/issue413

Author: holtskinner

src/a2a/grpc/a2a_pb2.py

@@ -55,7 +55,7 @@ class SendMessageConfiguration(_message.Message):
     push_notification: PushNotificationConfig
     history_length: int
     blocking: bool
-    def __init__(self, accepted_output_modes: _Optional[_Iterable[str]] = ..., push_notification: _Optional[_Union[PushNotificationConfig, _Mapping]] = ..., history_length: _Optional[int] = ..., blocking: bool = ...) -> None: ...
+    def __init__(self, accepted_output_modes: _Optional[_Iterable[str]] = ..., push_notification: _Optional[_Union[PushNotificationConfig, _Mapping]] = ..., history_length: _Optional[int] = ..., blocking: _Optional[bool] = ...) -> None: ...
 
 class Task(_message.Message):
     __slots__ = ("id", "context_id", "status", "artifacts", "history", "metadata")
@@ -157,7 +157,7 @@ class TaskStatusUpdateEvent(_message.Message):
     status: TaskStatus
     final: bool
     metadata: _struct_pb2.Struct
-    def __init__(self, task_id: _Optional[str] = ..., context_id: _Optional[str] = ..., status: _Optional[_Union[TaskStatus, _Mapping]] = ..., final: bool = ..., metadata: _Optional[_Union[_struct_pb2.Struct, _Mapping]] = ...) -> None: ...
+    def __init__(self, task_id: _Optional[str] = ..., context_id: _Optional[str] = ..., status: _Optional[_Union[TaskStatus, _Mapping]] = ..., final: _Optional[bool] = ..., metadata: _Optional[_Union[_struct_pb2.Struct, _Mapping]] = ...) -> None: ...
 
 class TaskArtifactUpdateEvent(_message.Message):
     __slots__ = ("task_id", "context_id", "artifact", "append", "last_chunk", "metadata")
@@ -173,7 +173,7 @@ class TaskArtifactUpdateEvent(_message.Message):
     append: bool
     last_chunk: bool
     metadata: _struct_pb2.Struct
-    def __init__(self, task_id: _Optional[str] = ..., context_id: _Optional[str] = ..., artifact: _Optional[_Union[Artifact, _Mapping]] = ..., append: bool = ..., last_chunk: bool = ..., metadata: _Optional[_Union[_struct_pb2.Struct, _Mapping]] = ...) -> None: ...
+    def __init__(self, task_id: _Optional[str] = ..., context_id: _Optional[str] = ..., artifact: _Optional[_Union[Artifact, _Mapping]] = ..., append: _Optional[bool] = ..., last_chunk: _Optional[bool] = ..., metadata: _Optional[_Union[_struct_pb2.Struct, _Mapping]] = ...) -> None: ...
 
 class PushNotificationConfig(_message.Message):
     __slots__ = ("id", "url", "token", "authentication")
@@ -204,7 +204,7 @@ class AgentInterface(_message.Message):
     def __init__(self, url: _Optional[str] = ..., transport: _Optional[str] = ...) -> None: ...
 
 class AgentCard(_message.Message):
-    __slots__ = ("protocol_version", "name", "description", "url", "preferred_transport", "additional_interfaces", "provider", "version", "documentation_url", "capabilities", "security_schemes", "security", "default_input_modes", "default_output_modes", "skills", "supports_authenticated_extended_card", "signatures")
+    __slots__ = ("protocol_version", "name", "description", "url", "preferred_transport", "additional_interfaces", "provider", "version", "documentation_url", "capabilities", "security_schemes", "security", "default_input_modes", "default_output_modes", "skills", "supports_authenticated_extended_card", "signatures", "icon_url")
     class SecuritySchemesEntry(_message.Message):
         __slots__ = ("key", "value")
         KEY_FIELD_NUMBER: _ClassVar[int]
@@ -229,6 +229,7 @@ class AgentCard(_message.Message):
     SKILLS_FIELD_NUMBER: _ClassVar[int]
     SUPPORTS_AUTHENTICATED_EXTENDED_CARD_FIELD_NUMBER: _ClassVar[int]
     SIGNATURES_FIELD_NUMBER: _ClassVar[int]
+    ICON_URL_FIELD_NUMBER: _ClassVar[int]
     protocol_version: str
     name: str
     description: str
@@ -246,7 +247,8 @@ class AgentCard(_message.Message):
     skills: _containers.RepeatedCompositeFieldContainer[AgentSkill]
     supports_authenticated_extended_card: bool
     signatures: _containers.RepeatedCompositeFieldContainer[AgentCardSignature]
-    def __init__(self, protocol_version: _Optional[str] = ..., name: _Optional[str] = ..., description: _Optional[str] = ..., url: _Optional[str] = ..., preferred_transport: _Optional[str] = ..., additional_interfaces: _Optional[_Iterable[_Union[AgentInterface, _Mapping]]] = ..., provider: _Optional[_Union[AgentProvider, _Mapping]] = ..., version: _Optional[str] = ..., documentation_url: _Optional[str] = ..., capabilities: _Optional[_Union[AgentCapabilities, _Mapping]] = ..., security_schemes: _Optional[_Mapping[str, SecurityScheme]] = ..., security: _Optional[_Iterable[_Union[Security, _Mapping]]] = ..., default_input_modes: _Optional[_Iterable[str]] = ..., default_output_modes: _Optional[_Iterable[str]] = ..., skills: _Optional[_Iterable[_Union[AgentSkill, _Mapping]]] = ..., supports_authenticated_extended_card: bool = ..., signatures: _Optional[_Iterable[_Union[AgentCardSignature, _Mapping]]] = ...) -> None: ...
+    icon_url: str
+    def __init__(self, protocol_version: _Optional[str] = ..., name: _Optional[str] = ..., description: _Optional[str] = ..., url: _Optional[str] = ..., preferred_transport: _Optional[str] = ..., additional_interfaces: _Optional[_Iterable[_Union[AgentInterface, _Mapping]]] = ..., provider: _Optional[_Union[AgentProvider, _Mapping]] = ..., version: _Optional[str] = ..., documentation_url: _Optional[str] = ..., capabilities: _Optional[_Union[AgentCapabilities, _Mapping]] = ..., security_schemes: _Optional[_Mapping[str, SecurityScheme]] = ..., security: _Optional[_Iterable[_Union[Security, _Mapping]]] = ..., default_input_modes: _Optional[_Iterable[str]] = ..., default_output_modes: _Optional[_Iterable[str]] = ..., skills: _Optional[_Iterable[_Union[AgentSkill, _Mapping]]] = ..., supports_authenticated_extended_card: _Optional[bool] = ..., signatures: _Optional[_Iterable[_Union[AgentCardSignature, _Mapping]]] = ..., icon_url: _Optional[str] = ...) -> None: ...
 
 class AgentProvider(_message.Message):
     __slots__ = ("url", "organization")
@@ -264,7 +266,7 @@ class AgentCapabilities(_message.Message):
     streaming: bool
     push_notifications: bool
     extensions: _containers.RepeatedCompositeFieldContainer[AgentExtension]
-    def __init__(self, streaming: bool = ..., push_notifications: bool = ..., extensions: _Optional[_Iterable[_Union[AgentExtension, _Mapping]]] = ...) -> None: ...
+    def __init__(self, streaming: _Optional[bool] = ..., push_notifications: _Optional[bool] = ..., extensions: _Optional[_Iterable[_Union[AgentExtension, _Mapping]]] = ...) -> None: ...
 
 class AgentExtension(_message.Message):
     __slots__ = ("uri", "description", "required", "params")
@@ -276,7 +278,7 @@ class AgentExtension(_message.Message):
     description: str
     required: bool
     params: _struct_pb2.Struct
-    def __init__(self, uri: _Optional[str] = ..., description: _Optional[str] = ..., required: bool = ..., params: _Optional[_Union[_struct_pb2.Struct, _Mapping]] = ...) -> None: ...
+    def __init__(self, uri: _Optional[str] = ..., description: _Optional[str] = ..., required: _Optional[bool] = ..., params: _Optional[_Union[_struct_pb2.Struct, _Mapping]] = ...) -> None: ...
 
 class AgentSkill(_message.Message):
     __slots__ = ("id", "name", "description", "tags", "examples", "input_modes", "output_modes", "security")

src/a2a/grpc/a2a_pb2.pyi

@@ -125,7 +125,7 @@ async def consume_all(self) -> AsyncGenerator[Event]:
                 # other part is waiting for an event or a closed queue.
                 if is_final_event:
                     logger.debug('Stopping event consumption in consume_all.')
-                    await self.queue.close()
+                    await self.queue.close(True)
                     yield event
                     break
                 yield event
@@ -135,7 +135,7 @@ async def consume_all(self) -> AsyncGenerator[Event]:
             except asyncio.TimeoutError:  # pyright: ignore [reportUnusedExcept]
                 # This class was made an alias of build-in TimeoutError after 3.11
                 continue
-            except QueueClosed:
+            except (QueueClosed, asyncio.QueueEmpty):
                 # Confirm that the queue is closed, e.g. we aren't on
                 # python 3.12 and get a queue empty error on an open queue
                 if self.queue.is_closed():

src/a2a/server/events/event_consumer.py

@@ -90,7 +90,12 @@ async def dequeue_event(self, no_wait: bool = False) -> Event:
             asyncio.QueueShutDown: If the queue has been closed and is empty.
         """
         async with self._lock:
-            if self._is_closed and self.queue.empty():
+            if (
+                sys.version_info < (3, 13)
+                and self._is_closed
+                and self.queue.empty()
+            ):
+                # On 3.13+, skip early raise; await self.queue.get() will raise QueueShutDown after shutdown()
                 logger.warning('Queue is closed. Event will not be dequeued.')
                 raise asyncio.QueueEmpty('Queue is closed.')
 
@@ -127,25 +132,38 @@ def tap(self) -> 'EventQueue':
         self._children.append(queue)
         return queue
 
-    async def close(self) -> None:
-        """Closes the queue for future push events.
+    async def close(self, immediate: bool = False) -> None:
+        """Closes the queue for future push events and also closes all child queues.
+
+        Once closed, no new events can be enqueued. For Python 3.13+, this will trigger
+        `asyncio.QueueShutDown` when the queue is empty and a consumer tries to dequeue.
+        For lower versions, the queue will be marked as closed and optionally cleared.
+
+        Args:
+            immediate (bool):
+                - True: Immediately closes the queue and clears all unprocessed events without waiting for them to be consumed. This is suitable for scenarios where you need to forcefully interrupt and quickly release resources.
+                - False (default): Gracefully closes the queue, waiting for all queued events to be processed (i.e., the queue is drained) before closing. This is suitable when you want to ensure all events are handled.
 
-        Once closed, `dequeue_event` will eventually raise `asyncio.QueueShutDown`
-        when the queue is empty. Also closes all child queues.
         """
         logger.debug('Closing EventQueue.')
         async with self._lock:
             # If already closed, just return.
-            if self._is_closed:
+            if self._is_closed and not immediate:
                 return
-            self._is_closed = True
+            if not self._is_closed:
+                self._is_closed = True
         # If using python 3.13 or higher, use the shutdown method
         if sys.version_info >= (3, 13):
-            self.queue.shutdown()
+            self.queue.shutdown(immediate)
             for child in self._children:
-                await child.close()
+                await child.close(immediate)
         # Otherwise, join the queue
         else:
+            if immediate:
+                await self.clear_events(True)
+                for child in self._children:
+                    await child.close(immediate)
+                return
             tasks = [asyncio.create_task(self.queue.join())]
             tasks.extend(
                 asyncio.create_task(child.close()) for child in self._children
@@ -155,3 +173,53 @@ async def close(self) -> None:
     def is_closed(self) -> bool:
         """Checks if the queue is closed."""
         return self._is_closed
+
+    async def clear_events(self, clear_child_queues: bool = True) -> None:
+        """Clears all events from the current queue and optionally all child queues.
+
+        This method removes all pending events from the queue without processing them.
+        Child queues can be optionally cleared based on the clear_child_queues parameter.
+
+        Args:
+            clear_child_queues: If True (default), clear all child queues as well.
+                              If False, only clear the current queue, leaving child queues untouched.
+        """
+        logger.debug('Clearing all events from EventQueue and child queues.')
+
+        # Clear all events from the queue, even if closed
+        cleared_count = 0
+        async with self._lock:
+            try:
+                while True:
+                    event = self.queue.get_nowait()
+                    logger.debug(
+                        f'Discarding unprocessed event of type: {type(event)}, content: {event}'
+                    )
+                    self.queue.task_done()
+                    cleared_count += 1
+            except asyncio.QueueEmpty:
+                pass
+            except Exception as e:
+                # Handle Python 3.13+ QueueShutDown
+                if (
+                    sys.version_info >= (3, 13)
+                    and type(e).__name__ == 'QueueShutDown'
+                ):
+                    pass
+                else:
+                    raise
+
+            if cleared_count > 0:
+                logger.debug(
+                    f'Cleared {cleared_count} unprocessed events from EventQueue.'
+                )
+
+        # Clear all child queues (lock released before awaiting child tasks)
+        if clear_child_queues and self._children:
+            child_tasks = [
+                asyncio.create_task(child.clear_events())
+                for child in self._children
+            ]
+
+            if child_tasks:
+                await asyncio.gather(*child_tasks, return_exceptions=True)

src/a2a/server/events/event_queue.py

@@ -286,11 +286,19 @@ async def on_message_send(
 
         interrupted_or_non_blocking = False
         try:
+            # Create async callback for push notifications
+            async def push_notification_callback() -> None:
+                await self._send_push_notification_if_needed(
+                    task_id, result_aggregator
+                )
+
             (
                 result,
                 interrupted_or_non_blocking,
             ) = await result_aggregator.consume_and_break_on_interrupt(
-                consumer, blocking=blocking
+                consumer,
+                blocking=blocking,
+                event_callback=push_notification_callback,
             )
             if not result:
                 raise ServerError(error=InternalError())  # noqa: TRY301

src/a2a/server/request_handlers/default_request_handler.py

@@ -1,7 +1,7 @@
 import asyncio
 import logging
 
-from collections.abc import AsyncGenerator, AsyncIterator
+from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable
 
 from a2a.server.events import Event, EventConsumer
 from a2a.server.tasks.task_manager import TaskManager
@@ -24,7 +24,10 @@ class ResultAggregator:
        Task object and emit that Task object.
     """
 
-    def __init__(self, task_manager: TaskManager):
+    def __init__(
+        self,
+        task_manager: TaskManager,
+    ) -> None:
         """Initializes the ResultAggregator.
 
         Args:
@@ -92,7 +95,10 @@ async def consume_all(
         return await self.task_manager.get_task()
 
     async def consume_and_break_on_interrupt(
-        self, consumer: EventConsumer, blocking: bool = True
+        self,
+        consumer: EventConsumer,
+        blocking: bool = True,
+        event_callback: Callable[[], Awaitable[None]] | None = None,
     ) -> tuple[Task | Message | None, bool]:
         """Processes the event stream until completion or an interruptable state is encountered.
 
@@ -105,6 +111,9 @@ async def consume_and_break_on_interrupt(
             consumer: The `EventConsumer` to read events from.
             blocking: If `False`, the method returns as soon as a task/message
                       is available. If `True`, it waits for a terminal state.
+            event_callback: Optional async callback function to be called after each event
+                           is processed in the background continuation.
+                           Mainly used for push notifications currently.
 
         Returns:
             A tuple containing:
@@ -150,13 +159,17 @@ async def consume_and_break_on_interrupt(
             if should_interrupt:
                 # Continue consuming the rest of the events in the background.
                 # TODO: We should track all outstanding tasks to ensure they eventually complete.
-                asyncio.create_task(self._continue_consuming(event_stream))  # noqa: RUF006
+                asyncio.create_task(  # noqa: RUF006
+                    self._continue_consuming(event_stream, event_callback)
+                )
                 interrupted = True
                 break
         return await self.task_manager.get_task(), interrupted
 
     async def _continue_consuming(
-        self, event_stream: AsyncIterator[Event]
+        self,
+        event_stream: AsyncIterator[Event],
+        event_callback: Callable[[], Awaitable[None]] | None = None,
     ) -> None:
         """Continues processing an event stream in a background task.
 
@@ -165,6 +178,9 @@ async def _continue_consuming(
 
         Args:
             event_stream: The remaining `AsyncIterator` of events from the consumer.
+            event_callback: Optional async callback function to be called after each event is processed.
         """
         async for event in event_stream:
             await self.task_manager.process(event)
+            if event_callback:
+                await event_callback()