fix: Support non-blocking sendMessage (#327)

This is a backport of the following PRs from the a2a-python repo: 349,
294, 449, 440 and 472

These include:
- Non blocking sendMessage should invoke push notification
- Not possible to cancel Task in final state
- Handle concurrent task completion during cancellation
- Persist task state after client disconnect

Also add DefaultRequestHandlerTest, and backport the parts that make
sense.

Author: kabir

.serena/cache/java/document_symbols_cache_v23-06-25.pkl

@@ -14,6 +14,7 @@
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.Executor;
 import java.util.concurrent.Flow;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Supplier;
 
@@ -48,6 +49,8 @@
 import io.a2a.spec.StreamingEventKind;
 import io.a2a.spec.Task;
 import io.a2a.spec.TaskIdParams;
+import io.a2a.spec.TaskNotCancelableError;
+import io.a2a.spec.TaskState;
 import io.a2a.spec.TaskNotFoundError;
 import io.a2a.spec.TaskPushNotificationConfig;
 import io.a2a.spec.TaskQueryParams;
@@ -122,6 +125,13 @@ public Task onCancelTask(TaskIdParams params, ServerCallContext context) throws
         if (task == null) {
             throw new TaskNotFoundError();
         }
+
+        // Check if task is in a non-cancelable state (completed, canceled, failed, rejected)
+        if (task.getStatus().state().isFinal()) {
+            throw new TaskNotCancelableError(
+                    "Task cannot be canceled - current state: " + task.getStatus().state().asString());
+        }
+
         TaskManager taskManager = new TaskManager(
                 task.getId(),
                 task.getContextId(),
@@ -148,11 +158,17 @@ public Task onCancelTask(TaskIdParams params, ServerCallContext context) throws
 
         EventConsumer consumer = new EventConsumer(queue);
         EventKind type = resultAggregator.consumeAll(consumer);
-        if (type instanceof Task tempTask) {
-            return tempTask;
+        if (!(type instanceof Task tempTask)) {
+            throw new InternalError("Agent did not return valid response for cancel");
         }
 
-        throw new InternalError("Agent did not return a valid response");
+        // Verify task was actually canceled (not completed concurrently)
+        if (tempTask.getStatus().state() != TaskState.CANCELED) {
+            throw new TaskNotCancelableError(
+                    "Task cannot be canceled - current state: " + tempTask.getStatus().state().asString());
+        }
+
+        return tempTask;
     }
 
     @Override
@@ -166,32 +182,42 @@ public EventKind onMessageSend(MessageSendParams params, ServerCallContext conte
         EventQueue queue = queueManager.createOrTap(taskId);
         ResultAggregator resultAggregator = new ResultAggregator(mss.taskManager, null);
 
-        boolean interrupted = false;
+        boolean blocking = true; // Default to blocking behavior
+        if (params.configuration() != null && Boolean.FALSE.equals(params.configuration().blocking())) {
+            blocking = false;
+        }
+
+        boolean interruptedOrNonBlocking = false;
 
         EnhancedRunnable producerRunnable = registerAndExecuteAgentAsync(taskId, mss.requestContext, queue);
         ResultAggregator.EventTypeAndInterrupt etai = null;
         try {
+            // Create callback for push notifications during background event processing
+            Runnable pushNotificationCallback = () -> sendPushNotification(taskId, resultAggregator);
+
             EventConsumer consumer = new EventConsumer(queue);
 
             // This callback must be added before we start consuming. Otherwise,
             // any errors thrown by the producerRunnable are not picked up by the consumer
             producerRunnable.addDoneCallback(consumer.createAgentRunnableDoneCallback());
-            etai = resultAggregator.consumeAndBreakOnInterrupt(consumer);
+            etai = resultAggregator.consumeAndBreakOnInterrupt(consumer, blocking, pushNotificationCallback);
 
             if (etai == null) {
                 LOGGER.debug("No result, throwing InternalError");
                 throw new InternalError("No result");
             }
-            interrupted = etai.interrupted();
-            LOGGER.debug("Was interrupted: {}", interrupted);
+            interruptedOrNonBlocking = etai.interrupted();
+            LOGGER.debug("Was interrupted or non-blocking: {}", interruptedOrNonBlocking);
 
             EventKind kind = etai.eventType();
             if (kind instanceof Task taskResult && !taskId.equals(taskResult.getId())) {
                 throw new InternalError("Task ID mismatch in agent response");
             }
 
+            // Send push notification after initial return (for both blocking and non-blocking)
+            pushNotificationCallback.run();
         } finally {
-            if (interrupted) {
+            if (interruptedOrNonBlocking) {
                 CompletableFuture<Void> cleanupTask = CompletableFuture.runAsync(() -> cleanupProducer(taskId), executor);
                 trackBackgroundTask(cleanupTask);
             } else {
@@ -214,13 +240,15 @@ public Flow.Publisher<StreamingEventKind> onMessageSendStream(
         ResultAggregator resultAggregator = new ResultAggregator(mss.taskManager, null);
 
         EnhancedRunnable producerRunnable = registerAndExecuteAgentAsync(taskId.get(), mss.requestContext, queue);
+
+        // Move consumer creation and callback registration outside try block
+        // so consumer is available for background consumption on client disconnect
         EventConsumer consumer = new EventConsumer(queue);
+        producerRunnable.addDoneCallback(consumer.createAgentRunnableDoneCallback());
 
-        try {
+        AtomicBoolean backgroundConsumeStarted = new AtomicBoolean(false);
 
-            // This callback must be added before we start consuming. Otherwise,
-            // any errors thrown by the producerRunnable are not picked up by the consumer
-            producerRunnable.addDoneCallback(consumer.createAgentRunnableDoneCallback());
+        try {
             Flow.Publisher<Event> results = resultAggregator.consumeAndEmit(consumer);
 
             Flow.Publisher<Event> eventPublisher =
@@ -258,7 +286,61 @@ public Flow.Publisher<StreamingEventKind> onMessageSendStream(
                 return true;
             }));
 
-            return convertingProcessor(eventPublisher, event -> (StreamingEventKind) event);
+            Flow.Publisher<StreamingEventKind> finalPublisher = convertingProcessor(eventPublisher, event -> (StreamingEventKind) event);
+
+            // Wrap publisher to detect client disconnect and continue background consumption
+            return subscriber -> finalPublisher.subscribe(new Flow.Subscriber<StreamingEventKind>() {
+                private Flow.Subscription subscription;
+
+                @Override
+                public void onSubscribe(Flow.Subscription subscription) {
+                    this.subscription = subscription;
+                    // Wrap subscription to detect cancellation
+                    subscriber.onSubscribe(new Flow.Subscription() {
+                        @Override
+                        public void request(long n) {
+                            subscription.request(n);
+                        }
+
+                        @Override
+                        public void cancel() {
+                            LOGGER.debug("Client cancelled subscription for task {}, starting background consumption", taskId.get());
+                            startBackgroundConsumption();
+                            subscription.cancel();
+                        }
+                    });
+                }
+
+                @Override
+                public void onNext(StreamingEventKind item) {
+                    subscriber.onNext(item);
+                }
+
+                @Override
+                public void onError(Throwable throwable) {
+                    subscriber.onError(throwable);
+                }
+
+                @Override
+                public void onComplete() {
+                    subscriber.onComplete();
+                }
+
+                private void startBackgroundConsumption() {
+                    if (backgroundConsumeStarted.compareAndSet(false, true)) {
+                        // Client disconnected: continue consuming and persisting events in background
+                        CompletableFuture<Void> bgTask = CompletableFuture.runAsync(() -> {
+                            try {
+                                resultAggregator.consumeAll(consumer);
+                                LOGGER.debug("Background consumption completed for task {}", taskId.get());
+                            } catch (Exception e) {
+                                LOGGER.error("Error during background consumption for task {}", taskId.get(), e);
+                            }
+                        }, executor);
+                        trackBackgroundTask(bgTask);
+                    }
+                }
+            });
         } finally {
             CompletableFuture<Void> cleanupTask = CompletableFuture.runAsync(() -> cleanupProducer(taskId.get()), executor);
             trackBackgroundTask(cleanupTask);
@@ -454,5 +536,14 @@ private MessageSendSetup initMessageSend(MessageSendParams params, ServerCallCon
         return new MessageSendSetup(taskManager, task, requestContext);
     }
 
+    private void sendPushNotification(String taskId, ResultAggregator resultAggregator) {
+        if (pushSender != null && taskId != null) {
+            EventKind latest = resultAggregator.getCurrentResult();
+            if (latest instanceof Task latestTask) {
+                pushSender.sendNotification(latestTask);
+            }
+        }
+    }
+
     private record MessageSendSetup(TaskManager taskManager, Task task, RequestContext requestContext) {}
 }

server-common/src/main/java/io/a2a/server/requesthandlers/DefaultRequestHandler.java

@@ -45,7 +45,7 @@ public Flow.Publisher<Event> consumeAndEmit(EventConsumer consumer) {
         }));
     }
 
-    public EventKind consumeAll(EventConsumer consumer) {
+    public EventKind consumeAll(EventConsumer consumer) throws JSONRPCError {
         AtomicReference<EventKind> returnedEvent = new AtomicReference<>();
         Flow.Publisher<Event> all = consumer.consumeAll();
         AtomicReference<Throwable> error = new AtomicReference<>();
@@ -65,13 +65,22 @@ public EventKind consumeAll(EventConsumer consumer) {
                 },
                 error::set);
 
+        Throwable err = error.get();
+        if (err != null) {
+            Utils.rethrow(err);
+        }
+
         if (returnedEvent.get() != null) {
             return returnedEvent.get();
         }
         return taskManager.getTask();
     }
 
-    public EventTypeAndInterrupt consumeAndBreakOnInterrupt(EventConsumer consumer) throws JSONRPCError {
+    public EventTypeAndInterrupt consumeAndBreakOnInterrupt(EventConsumer consumer, boolean blocking) throws JSONRPCError {
+        return consumeAndBreakOnInterrupt(consumer, blocking, null);
+    }
+
+    public EventTypeAndInterrupt consumeAndBreakOnInterrupt(EventConsumer consumer, boolean blocking, Runnable eventCallback) throws JSONRPCError {
         Flow.Publisher<Event> all = consumer.consumeAll();
         AtomicReference<Message> message = new AtomicReference<>();
         AtomicBoolean interrupted = new AtomicBoolean(false);
@@ -92,16 +101,28 @@ public EventTypeAndInterrupt consumeAndBreakOnInterrupt(EventConsumer consumer)
 
                     callTaskManagerProcess(event);
 
-                    if ((event instanceof Task task && task.getStatus().state() == TaskState.AUTH_REQUIRED)
-                            || (event instanceof TaskStatusUpdateEvent tsue && tsue.getStatus().state() == TaskState.AUTH_REQUIRED)) {
+                    boolean shouldInterrupt = false;
+                    boolean isAuthRequired = (event instanceof Task task && task.getStatus().state() == TaskState.AUTH_REQUIRED)
+                            || (event instanceof TaskStatusUpdateEvent tsue && tsue.getStatus().state() == TaskState.AUTH_REQUIRED);
+
+                    // Always interrupt on auth_required, as it needs external action.
+                    if (isAuthRequired) {
                         // auth-required is a special state: the message should be
                         // escalated back to the caller, but the agent is expected to
                         // continue producing events once the authorization is received
                         // out-of-band. This is in contrast to input-required, where a
                         // new request is expected in order for the agent to make progress,
                         // so the agent should exit.
+                        shouldInterrupt = true;
+                    }
+                    // For non-blocking calls, interrupt as soon as a task is available.
+                    else if (!blocking) {
+                        shouldInterrupt = true;
+                    }
 
-                        CompletableFuture.runAsync(() -> continueConsuming(all));
+                    if (shouldInterrupt) {
+                        // Continue consuming the rest of the events in the background.
+                        CompletableFuture.runAsync(() -> continueConsuming(all, eventCallback));
                         interrupted.set(true);
                         return false;
                     }
@@ -118,11 +139,14 @@ public EventTypeAndInterrupt consumeAndBreakOnInterrupt(EventConsumer consumer)
                 message.get() != null ? message.get() : taskManager.getTask(), interrupted.get());
     }
 
-    private void continueConsuming(Flow.Publisher<Event> all) {
+    private void continueConsuming(Flow.Publisher<Event> all, Runnable eventCallback) {
         consumer(createTubeConfig(),
                 all,
                 event -> {
                     callTaskManagerProcess(event);
+                    if (eventCallback != null) {
+                        eventCallback.run();
+                    }
                     return true;
                 },
                 t -> {});

server-common/src/main/java/io/a2a/server/tasks/ResultAggregator.java

@@ -28,10 +28,13 @@
 import org.junit.jupiter.api.Timeout;
 
 /**
- * Tests for DefaultRequestHandler background cleanup and task tracking functionality,
- * backported from Python PR #440 and #472.
+ * Comprehensive tests for DefaultRequestHandler, backported from Python's
+ * test_default_request_handler.py. These tests cover core functionality that
+ * is transport-agnostic and should work across JSON-RPC, gRPC, and REST.
+ *
+ * Background cleanup and task tracking tests are from Python PR #440 and #472.
  */
-public class DefaultRequestHandlerBackgroundTest {
+public class DefaultRequestHandlerTest {
 
     private DefaultRequestHandler requestHandler;
     private InMemoryTaskStore taskStore;

server-common/src/test/java/io/a2a/server/requesthandlers/DefaultRequestHandlerBackgroundTest.java renamed to server-common/src/test/java/io/a2a/server/requesthandlers/DefaultRequestHandlerTest.java

@@ -2,6 +2,7 @@
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.Mockito.reset;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
@@ -10,6 +11,9 @@
 
 import java.util.Collections;
 
+import io.a2a.server.events.EventConsumer;
+import io.a2a.server.events.EventQueue;
+import io.a2a.server.events.InMemoryQueueManager;
 import io.a2a.spec.EventKind;
 import io.a2a.spec.Message;
 import io.a2a.spec.Task;
@@ -183,4 +187,35 @@ void testGetCurrentResultWithMessageTakesPrecedence() {
         // Task manager should not be called when message is present
         verifyNoInteractions(mockTaskManager);
     }
+
+    @Test
+    void testConsumeAndBreakNonBlocking() throws Exception {
+        // Test that with blocking=false, the method returns after the first event
+        Task firstEvent = createSampleTask("non_blocking_task", TaskState.WORKING, "ctx1");
+
+        // After processing firstEvent, the current result will be that task
+        when(mockTaskManager.getTask()).thenReturn(firstEvent);
+
+        // Create an event queue using QueueManager (which has access to builder)
+        InMemoryQueueManager queueManager =
+            new InMemoryQueueManager();
+
+        EventQueue queue = queueManager.getEventQueueBuilder("test-task").build();
+        queue.enqueueEvent(firstEvent);
+
+        // Create real EventConsumer with the queue
+        EventConsumer eventConsumer =
+            new EventConsumer(queue);
+
+        // Close queue after first event to simulate stream ending after processing
+        queue.close();
+
+        ResultAggregator.EventTypeAndInterrupt result =
+            aggregator.consumeAndBreakOnInterrupt(eventConsumer, false);
+
+        assertEquals(firstEvent, result.eventType());
+        assertTrue(result.interrupted());
+        verify(mockTaskManager).process(firstEvent);
+        verify(mockTaskManager).getTask();
+    }
 }

server-common/src/test/java/io/a2a/server/tasks/ResultAggregatorTest.java

@@ -11,7 +11,7 @@
 @JsonInclude(JsonInclude.Include.NON_ABSENT)
 @JsonIgnoreProperties(ignoreUnknown = true)
 public record MessageSendConfiguration(List<String> acceptedOutputModes, Integer historyLength,
-                                       PushNotificationConfig pushNotificationConfig, boolean blocking) {
+                                       PushNotificationConfig pushNotificationConfig, Boolean blocking) {
 
     public MessageSendConfiguration {
         if (historyLength != null && historyLength < 0) {
@@ -23,7 +23,7 @@ public static class Builder {
         List<String> acceptedOutputModes;
         Integer historyLength;
         PushNotificationConfig pushNotificationConfig;
-        boolean blocking;
+        Boolean blocking = true;
 
         public Builder acceptedOutputModes(List<String> acceptedOutputModes) {
             this.acceptedOutputModes = acceptedOutputModes;
@@ -40,7 +40,7 @@ public Builder historyLength(Integer historyLength) {
             return this;
         }
 
-        public Builder blocking(boolean blocking) {
+        public Builder blocking(Boolean blocking) {
             this.blocking = blocking;
             return this;
         }

spec/src/main/java/io/a2a/spec/MessageSendConfiguration.java

@@ -31,4 +31,8 @@ public TaskNotCancelableError(
                 data);
     }
 
+    public TaskNotCancelableError(@JsonProperty("message") String message) {
+        this(null, message, null);
+    }
+
 }