Repository: a2aproject/a2a-python

Issue:
[Feat]: Non blocking with intermediary responses (via push notification)
### Is your feature request related to a problem? Please describe.

I want to avoid long-lived persistent connections (SSE) but also be able to receive multiple Agent Responses.

Imagine I sended a request to an Agent to perform a heavy task with multiple steps. I want to receive intermediary messages where the agent can notify steps beeing finished (between some tool calls for example).

Today, the push-notification only works for "final" states like task completed or input required. And streaming responses require persistent connections via SSE.

### Describe the solution you'd like

Use the already defined `blocking` attribute of `MessageSendConfiguration` to indicate whether the request processing should last for the entire event stream or stop immediately at the task `submitted` status.

Then, push notifications should also be able to send artifacts generated during processing.

### Describe alternatives you've considered

_No response_

### Additional context
Some points to help us decide the best solution.

The protocol specification metion "partial responses" as a good use case for streaming [spec reference](https://a2aproject.github.io/A2A/latest/topics/streaming-and-async/#1-streaming-with-server-sent-events-sse)
> # When to Use Streaming:
> - Interactive, conversational exchanges where immediate feedback or partial responses are beneficial.

Also in protocol specification is mentioned that push notification payloads are not strictly task objects, so I belive that TaskStatusUpdateEvent and TaskArtifactUpdateEvent could also be sent in streaming request push notification ? [spec reference](https://a2aproject.github.io/A2A/latest/topics/streaming-and-async/#2-push-notifications-for-disconnected-scenarios)

> Notification Payload: The A2A protocol itself does not strictly define the HTTP body payload of the push notification sent by the server to the client's webhook. However, the notification SHOULD contain sufficient information for the client to identify the Task ID and understand the general nature of the update (e.g., the new TaskState). Servers might send a minimal payload (just Task ID and new state) or a more comprehensive one (e.g., a summary or even the full [Task](https://a2aproject.github.io/A2A/latest/specification/#61-task-object) object).

### Code of Conduct

- [x] I agree to follow this project's Code of Conduct


Discussion:
Ok, if I understand you correctly: you are saying that when a client invokes the 'execute' function of the executor, if the client request is `send_message` and not `send_message_stream`, all of the partial notifications don't go out to anyone who may be subscribed to push notifications about the task?

So if we look at my hacky example script https://github.com/mozilla-ai/any-agent/blob/414-a2a-push-not/tests/integration/a2a/test_push_notification.py#L54

In that test, 

* I send a non-streaming message and a task is created. 
* I make another call to subscribe to push notifications about the task.
* I make another non-streaming message with the same task id, and verify that I get `at least one` push notification

What you are saying is that in this scenario, even if there are multiple task updates caused by me making that second non-streaming message, I would only get one push notification about the final result of the call? 

I haven't checked this behavior so I don't know if this is true or false. Thanks!


Exactly !

Using the actual sdk, I am unable to define an implementation capable of carrying out the flow below
```
[Request - Response]
Client --Message--> Remote Agent --Task Submitted--> Client

[Push Notifications]
Remote Agent --Artifact--> Client (This beeing a complete artifact response, not a chunk)
Remote Agent --Artifact--> Client (This beeing a complete artifact response, not a chunk)
Remote Agent --Artifact--> Client (This beeing a complete artifact response, not a chunk)
Remote Agent --Task Completed--> Client
```
I’ve been debugging the SDK and exploring possible implementations, but I’ve run into a couple of roadblocks.

# Option 1: Using message/send
I don’t think this will work because of the consume_and_break_on_interrupt method in (see [source](https://github.com/a2aproject/a2a-python/blob/main/src/a2a/server/request_handlers/default_request_handler.py#L281)).
```python
await result_aggregator.consume_and_break_on_interrupt(consumer)
```
That method loops over the generator returned by consumer.consume_all() (see [source](https://github.com/a2aproject/a2a-python/blob/main/src/a2a/server/tasks/result_aggregator.py#L113)), and consume_all() only stops when it encounters an interrupt or terminal state:

```python
is_final_event = (
    (isinstance(event, TaskStatusUpdateEvent) and event.final)
    or isinstance(event, Message)
    or (
        isinstance(event, Task)
        and event.status.state
        in (
            TaskState.completed,
            TaskState.canceled,
            TaskState.failed,
            TaskState.rejected,
            TaskState.unknown,
            TaskState.input_required,
        )
    )
)
```
Because of this, you can’t send intermediate events—only “final” ones.

# Option 2: Using message/stream (with SSE)
This approach relies on a persistent SSE connection, which I’m trying to avoid. It seems there’s no non-blocking request option. 

Also, result_aggregator.current_result always returns a Task, so only Task objects ever get pushed:

```python
latest_task = await result_aggregator.current_result
if isinstance(latest_task, Task):
    await self._push_notifier.send_notification(latest_task)
```
That means the client will receive a Task both when you call updater.add_artifact() and again when you call updater.complete(), making it hard to tell when actual “artifact” events should be forwarded to the end user.

Am I missing something here?
I if understand correctly, I did come across this issue too. `DefaultRequestHandler` has a gap in push notification support for non-streaming `message/send` requests. When a new task is created (as opposed to updating an existing task), the push notification configuration is not stored, preventing webhook delivery for task completion events.

This was my own hacky approach, but I am unsure if I am foot gunning myself here as I had to drop the work with a plan to return later and really dig in. I apologies if you're having a different issue and I am clouding things up:

```python
import logging
from typing import Union

from a2a.server.request_handlers import DefaultRequestHandler
from a2a.server.context import ServerCallContext
from a2a.types import (
    Message,
    Task,
    MessageSendParams,
)

logger = logging.getLogger(__name__)


class CustomRequestHandler(DefaultRequestHandler):
    """
    Enhanced request handler that fixes push notification storage for non-streaming endpoints.

    The default a2a-sdk DefaultRequestHandler only stores push notification configs
    for existing tasks in non-streaming mode, not for new tasks. This fixes that, but
    we should eventually upstream this change to the a2a-sdk, or find out why it
    has not been done already.
    """

    async def on_message_send(
        self,
        params: MessageSendParams,
        context: ServerCallContext | None = None,
    ) -> Union[Message, Task]:
        """
        Override to ensure push notification configs are saved for new tasks.
        """
        # Call parent implementation
        result = await super().on_message_send(params, context)

        # If we got a task back and have push notification config, save it
        if (
            isinstance(result, Task)
            and self._push_notifier
            and params.configuration
            and params.configuration.pushNotificationConfig
        ):
            logger.info(f"Saving push notification config for new task {result.id}")
            await self._push_notifier.set_info(
                result.id,
                params.configuration.pushNotificationConfig,
            )

            # Also send the initial notification until upsteam fixes this
            await self._push_notifier.send_notification(result)

        return result
```
Ahh, @njbrake has already got this fixed! noice!
> Exactly !
> 
> Using the actual sdk, I am unable to define an implementation capable of carrying out the flow below
> 
> ```
> [Request - Response]
> Client --Message--> Remote Agent --Task Submitted--> Client
> 
> [Push Notifications]
> Remote Agent --Artifact--> Client (This beeing a complete artifact response, not a chunk)
> Remote Agent --Artifact--> Client (This beeing a complete artifact response, not a chunk)
> Remote Agent --Artifact--> Client (This beeing a complete artifact response, not a chunk)
> Remote Agent --Task Completed--> Client
> ```

Been trying to implement something like the above. I can see the other day they fixed the issue of push notifications simply not being sent with message/send but it still seems like this flow is not possible. I can't seem to find any way to send any messages before agent execution is finished (agentexecutor execute function) when I send non-streaming requests.

Is this functionality possible? Or is it even intended? I'm imagining an agent that takes multiple hours to run and we want it to just respond with an acknowledgement when the client sends a message, then can send push notifications the rest of the way - that seems like what the push notifications should have been for, but it seems nothing gets sent to the push notification server until the execute function is finished.

My hacky solution is to just manually send the push notifications inside the execution function - otherwise they only get sent at the end.. basically just run this function every time I update the task status: `await self.push_sender._dispatch_notification(task, push_config)`


I have encountered the same problem: only got the first push notification from server but miss the following notifications.

I think this issue should be fixed at a2a-sdk. After doing some modifications, the problem is sovled(tested well in my scenario). Hope these changes can be merged.



Patch:
diff --git a/src/a2a/server/request_handlers/default_request_handler.py b/src/a2a/server/request_handlers/default_request_handler.py
index 25fe2525..afe77712 100644
--- a/src/a2a/server/request_handlers/default_request_handler.py
+++ b/src/a2a/server/request_handlers/default_request_handler.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
diff --git a/src/a2a/server/tasks/result_aggregator.py b/src/a2a/server/tasks/result_aggregator.py
index 147c3202..fb1ab62e 100644
--- a/src/a2a/server/tasks/result_aggregator.py
+++ b/src/a2a/server/tasks/result_aggregator.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()


Test patch:
diff --git a/tests/server/request_handlers/test_default_request_handler.py b/tests/server/request_handlers/test_default_request_handler.py
index 6cb21662..e8906554 100644
--- a/tests/server/request_handlers/test_default_request_handler.py
+++ b/tests/server/request_handlers/test_default_request_handler.py
@@ -405,6 +405,134 @@ async def get_current_result():
     mock_agent_executor.execute.assert_awaited_once()
 
 
+@pytest.mark.asyncio
+async def test_on_message_send_with_push_notification_in_non_blocking_request():
+    """Test that push notification callback is called during background event processing for non-blocking requests."""
+    mock_task_store = AsyncMock(spec=TaskStore)
+    mock_push_notification_store = AsyncMock(spec=PushNotificationConfigStore)
+    mock_agent_executor = AsyncMock(spec=AgentExecutor)
+    mock_request_context_builder = AsyncMock(spec=RequestContextBuilder)
+    mock_push_sender = AsyncMock()
+
+    task_id = 'non_blocking_task_1'
+    context_id = 'non_blocking_ctx_1'
+
+    # Create a task that will be returned after the first event
+    initial_task = create_sample_task(
+        task_id=task_id, context_id=context_id, status_state=TaskState.working
+    )
+
+    # Create a final task that will be available during background processing
+    final_task = create_sample_task(
+        task_id=task_id, context_id=context_id, status_state=TaskState.completed
+    )
+
+    mock_task_store.get.return_value = None
+
+    # Mock request context
+    mock_request_context = MagicMock(spec=RequestContext)
+    mock_request_context.task_id = task_id
+    mock_request_context.context_id = context_id
+    mock_request_context_builder.build.return_value = mock_request_context
+
+    request_handler = DefaultRequestHandler(
+        agent_executor=mock_agent_executor,
+        task_store=mock_task_store,
+        push_config_store=mock_push_notification_store,
+        request_context_builder=mock_request_context_builder,
+        push_sender=mock_push_sender,
+    )
+
+    # Configure push notification
+    push_config = PushNotificationConfig(url='http://callback.com/push')
+    message_config = MessageSendConfiguration(
+        push_notification_config=push_config,
+        accepted_output_modes=['text/plain'],
+        blocking=False,  # Non-blocking request
+    )
+    params = MessageSendParams(
+        message=Message(
+            role=Role.user,
+            message_id='msg_non_blocking',
+            parts=[],
+            task_id=task_id,
+            context_id=context_id,
+        ),
+        configuration=message_config,
+    )
+
+    # Mock ResultAggregator with custom behavior
+    mock_result_aggregator_instance = AsyncMock(spec=ResultAggregator)
+
+    # First call returns the initial task and indicates interruption (non-blocking)
+    mock_result_aggregator_instance.consume_and_break_on_interrupt.return_value = (
+        initial_task,
+        True,  # interrupted = True for non-blocking
+    )
+
+    # Mock the current_result property to return the final task
+    async def get_current_result():
+        return final_task
+
+    type(mock_result_aggregator_instance).current_result = PropertyMock(
+        return_value=get_current_result()
+    )
+
+    # Track if the event_callback was passed to consume_and_break_on_interrupt
+    event_callback_passed = False
+    event_callback_received = None
+
+    async def mock_consume_and_break_on_interrupt(
+        consumer, blocking=True, event_callback=None
+    ):
+        nonlocal event_callback_passed, event_callback_received
+        event_callback_passed = event_callback is not None
+        event_callback_received = event_callback
+        return initial_task, True  # interrupted = True for non-blocking
+
+    mock_result_aggregator_instance.consume_and_break_on_interrupt = (
+        mock_consume_and_break_on_interrupt
+    )
+
+    with (
+        patch(
+            'a2a.server.request_handlers.default_request_handler.ResultAggregator',
+            return_value=mock_result_aggregator_instance,
+        ),
+        patch(
+            'a2a.server.request_handlers.default_request_handler.TaskManager.get_task',
+            return_value=initial_task,
+        ),
+        patch(
+            'a2a.server.request_handlers.default_request_handler.TaskManager.update_with_message',
+            return_value=initial_task,
+        ),
+    ):
+        # Execute the non-blocking request
+        result = await request_handler.on_message_send(
+            params, create_server_call_context()
+        )
+
+    # Verify the result is the initial task (non-blocking behavior)
+    assert result == initial_task
+
+    # Verify that the event_callback was passed to consume_and_break_on_interrupt
+    assert event_callback_passed, (
+        'event_callback should have been passed to consume_and_break_on_interrupt'
+    )
+    assert event_callback_received is not None, (
+        'event_callback should not be None'
+    )
+
+    # Verify that the push notification was sent with the final task
+    mock_push_sender.send_notification.assert_called_with(final_task)
+
+    # Verify that the push notification config was stored
+    mock_push_notification_store.set_info.assert_awaited_once_with(
+        task_id, push_config
+    )
+
+
 @pytest.mark.asyncio
 async def test_on_message_send_with_push_notification_no_existing_Task():
     """Test on_message_send for new task sets push notification info if provided."""