fix: Wait for agent completion and ensure all events processed in blo… (#431)

…cking calls

Fixes race condition where DefaultRequestHandler.onMessageSend() returns
before all task events are fully processed and persisted to TaskStore,
resulting in incomplete Task objects being returned to clients (missing
artifacts, incorrect state).

Root Cause:
- Blocking calls interrupted immediately after first event and returned
to client before background event consumption completed
- Agent execution and event processing happened asynchronously in
background
- No synchronization to ensure all events were consumed and persisted
before returning Task to client

Solution (4-step process):
1. Wait for agent to finish enqueueing events (configurable timeout)
2. Close the queue to signal consumption can complete (breaks
dependency)
3. Wait for consumption to finish processing events (configurable
timeout)
4. Fetch final task state from TaskStore (has all artifacts and correct
state)

This ensures blocking calls return complete Task objects with all
artifacts and correct state, including support for fire-and-forget tasks
that never emit final state events.

Added unit tests:
- testBlockingFireAndForgetReturnsNonFinalTask: Validates
fire-and-forget pattern
- testBlockingCallReturnsCompleteTaskWithArtifacts: Ensures all
artifacts included

Fixes #428 🦕

Author: kabir

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

@@ -3,6 +3,7 @@
 import static io.a2a.server.util.async.AsyncUtils.convertingProcessor;
 import static io.a2a.server.util.async.AsyncUtils.createTubeConfig;
 import static io.a2a.server.util.async.AsyncUtils.processor;
+import static java.util.concurrent.TimeUnit.*;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -51,11 +52,12 @@
 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;
+import io.a2a.spec.TaskState;
 import io.a2a.spec.UnsupportedOperationError;
+import org.eclipse.microprofile.config.inject.ConfigProperty;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -64,6 +66,26 @@ public class DefaultRequestHandler implements RequestHandler {
 
     private static final Logger LOGGER = LoggerFactory.getLogger(DefaultRequestHandler.class);
 
+    /**
+     * Timeout in seconds to wait for agent execution to complete in blocking calls.
+     * This allows slow agents (LLM-based, data processing, external APIs) sufficient time.
+     * Configurable via: a2a.blocking.agent.timeout.seconds
+     * Default: 30 seconds
+     */
+    @Inject
+    @ConfigProperty(name = "a2a.blocking.agent.timeout.seconds", defaultValue = "30")
+    int agentCompletionTimeoutSeconds;
+
+    /**
+     * Timeout in seconds to wait for event consumption to complete in blocking calls.
+     * This ensures all events are processed and persisted before returning to client.
+     * Configurable via: a2a.blocking.consumption.timeout.seconds
+     * Default: 5 seconds
+     */
+    @Inject
+    @ConfigProperty(name = "a2a.blocking.consumption.timeout.seconds", defaultValue = "5")
+    int consumptionCompletionTimeoutSeconds;
+
     private final AgentExecutor agentExecutor;
     private final TaskStore taskStore;
     private final QueueManager queueManager;
@@ -93,6 +115,19 @@ public DefaultRequestHandler(AgentExecutor agentExecutor, TaskStore taskStore,
         this.requestContextBuilder = () -> new SimpleRequestContextBuilder(taskStore, false);
     }
 
+    /**
+     * For testing
+     */
+    public static DefaultRequestHandler create(AgentExecutor agentExecutor, TaskStore taskStore,
+                         QueueManager queueManager, PushNotificationConfigStore pushConfigStore,
+                         PushNotificationSender pushSender, Executor executor) {
+        DefaultRequestHandler handler =
+                new DefaultRequestHandler(agentExecutor, taskStore, queueManager, pushConfigStore, pushSender, executor);
+        handler.agentCompletionTimeoutSeconds = 5;
+        handler.consumptionCompletionTimeoutSeconds = 2;
+        return handler;
+    }
+
     @Override
     public Task onGetTask(TaskQueryParams params, ServerCallContext context) throws JSONRPCError {
         LOGGER.debug("onGetTask {}", params.id());
@@ -192,6 +227,7 @@ public EventKind onMessageSend(MessageSendParams params, ServerCallContext conte
 
         EnhancedRunnable producerRunnable = registerAndExecuteAgentAsync(taskId, mss.requestContext, queue);
         ResultAggregator.EventTypeAndInterrupt etai = null;
+        EventKind kind = null;  // Declare outside try block so it's in scope for return
         try {
             // Create callback for push notifications during background event processing
             Runnable pushNotificationCallback = () -> sendPushNotification(taskId, resultAggregator);
@@ -201,7 +237,10 @@ public EventKind onMessageSend(MessageSendParams params, ServerCallContext conte
             // 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, blocking, pushNotificationCallback);
+
+            // Get agent future before consuming (for blocking calls to wait for agent completion)
+            CompletableFuture<Void> agentFuture = runningAgents.get(taskId);
+            etai = resultAggregator.consumeAndBreakOnInterrupt(consumer, blocking);
 
             if (etai == null) {
                 LOGGER.debug("No result, throwing InternalError");
@@ -210,7 +249,69 @@ public EventKind onMessageSend(MessageSendParams params, ServerCallContext conte
             interruptedOrNonBlocking = etai.interrupted();
             LOGGER.debug("Was interrupted or non-blocking: {}", interruptedOrNonBlocking);
 
-            EventKind kind = etai.eventType();
+            // For blocking calls that were interrupted (returned on first event),
+            // wait for agent execution and event processing BEFORE returning to client.
+            // This ensures the returned Task has all artifacts and current state.
+            // We do this HERE (not in ResultAggregator) to avoid blocking Vert.x worker threads
+            // during the consumption loop itself.
+            kind = etai.eventType();
+            if (blocking && interruptedOrNonBlocking) {
+                // For blocking calls: ensure all events are processed before returning
+                // Order of operations is critical to avoid circular dependency:
+                // 1. Wait for agent to finish enqueueing events
+                // 2. Close the queue to signal consumption can complete
+                // 3. Wait for consumption to finish processing events
+                // 4. Fetch final task state from TaskStore
+
+                try {
+                    // Step 1: Wait for agent to finish (with configurable timeout)
+                    if (agentFuture != null) {
+                        try {
+                            agentFuture.get(agentCompletionTimeoutSeconds, SECONDS);
+                            LOGGER.debug("Agent completed for task {}", taskId);
+                        } catch (java.util.concurrent.TimeoutException e) {
+                            // Agent still running after timeout - that's fine, events already being processed
+                            LOGGER.debug("Agent still running for task {} after {}s", taskId, agentCompletionTimeoutSeconds);
+                        }
+                    }
+
+                    // Step 2: Close the queue to signal consumption can complete
+                    // For fire-and-forget tasks, there's no final event, so we need to close the queue
+                    // This allows EventConsumer.consumeAll() to exit
+                    queue.close(false, false);  // graceful close, don't notify parent yet
+                    LOGGER.debug("Closed queue for task {} to allow consumption completion", taskId);
+
+                    // Step 3: Wait for consumption to complete (now that queue is closed)
+                    if (etai.consumptionFuture() != null) {
+                        etai.consumptionFuture().get(consumptionCompletionTimeoutSeconds, SECONDS);
+                        LOGGER.debug("Consumption completed for task {}", taskId);
+                    }
+                } catch (InterruptedException e) {
+                    Thread.currentThread().interrupt();
+                    String msg = String.format("Error waiting for task %s completion", taskId);
+                    LOGGER.warn(msg, e);
+                    throw new InternalError(msg);
+                } catch (java.util.concurrent.ExecutionException e) {
+                    String msg = String.format("Error during task %s execution", taskId);
+                    LOGGER.warn(msg, e.getCause());
+                    throw new InternalError(msg);
+                } catch (java.util.concurrent.TimeoutException e) {
+                    String msg = String.format("Timeout waiting for consumption to complete for task %s", taskId);
+                    LOGGER.warn(msg, taskId);
+                    throw new InternalError(msg);
+                }
+
+                // Step 4: Fetch the final task state from TaskStore (all events have been processed)
+                Task updatedTask = taskStore.get(taskId);
+                if (updatedTask != null) {
+                    kind = updatedTask;
+                    if (LOGGER.isDebugEnabled()) {
+                        LOGGER.debug("Fetched final task for {} with state {} and {} artifacts",
+                            taskId, updatedTask.getStatus().state(),
+                            updatedTask.getArtifacts().size());
+                    }
+                }
+            }
             if (kind instanceof Task taskResult && !taskId.equals(taskResult.getId())) {
                 throw new InternalError("Task ID mismatch in agent response");
             }
@@ -227,8 +328,8 @@ public EventKind onMessageSend(MessageSendParams params, ServerCallContext conte
             trackBackgroundTask(cleanupProducer(agentFuture, etai != null ? etai.consumptionFuture() : null, taskId, queue, false));
         }
 
-        LOGGER.debug("Returning: {}", etai.eventType());
-        return etai.eventType();
+        LOGGER.debug("Returning: {}", kind);
+        return kind;
     }
 
     @Override

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

@@ -11,20 +11,20 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
 import io.a2a.server.events.EventConsumer;
 import io.a2a.server.events.EventQueueItem;
 import io.a2a.spec.A2AServerException;
 import io.a2a.spec.Event;
 import io.a2a.spec.EventKind;
+import io.a2a.spec.InternalError;
 import io.a2a.spec.JSONRPCError;
 import io.a2a.spec.Message;
 import io.a2a.spec.Task;
 import io.a2a.spec.TaskState;
 import io.a2a.spec.TaskStatusUpdateEvent;
 import io.a2a.util.Utils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 public class ResultAggregator {
     private static final Logger LOGGER = LoggerFactory.getLogger(ResultAggregator.class);
@@ -106,10 +106,6 @@ public EventKind consumeAll(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<EventQueueItem> allItems = consumer.consumeAll();
         AtomicReference<Message> message = new AtomicReference<>();
         AtomicBoolean interrupted = new AtomicBoolean(false);
@@ -180,11 +176,11 @@ else if (!blocking) {
                         shouldInterrupt = true;
                         continueInBackground = true;
                     }
-                    else {
-                        // For ALL blocking calls (both final and non-final events), use background consumption
-                        // This ensures all events are processed and persisted to TaskStore in background
-                        // Queue lifecycle is now managed by DefaultRequestHandler.cleanupProducer()
-                        // which waits for BOTH agent and consumption futures before closing queues
+                    else if (blocking) {
+                        // For blocking calls: Interrupt to free Vert.x thread, but continue in background
+                        // Python's async consumption doesn't block threads, but Java's does
+                        // So we interrupt to return quickly, then rely on background consumption
+                        // DefaultRequestHandler will fetch the final state from TaskStore
                         shouldInterrupt = true;
                         continueInBackground = true;
                         if (LOGGER.isDebugEnabled()) {
@@ -198,10 +194,17 @@ else if (!blocking) {
                         interrupted.set(true);
                         completionFuture.complete(null);
 
-                        // Signal that cleanup can proceed while consumption continues in background.
-                        // This prevents infinite hangs for fire-and-forget agents that never emit final events.
-                        // Processing continues (return true below) and all events are still persisted to TaskStore.
-                        consumptionCompletionFuture.complete(null);
+                        // For blocking calls, DON'T complete consumptionCompletionFuture here.
+                        // Let it complete naturally when subscription finishes (onComplete callback below).
+                        // This ensures all events are processed and persisted to TaskStore before
+                        // DefaultRequestHandler.cleanupProducer() proceeds with cleanup.
+                        //
+                        // For non-blocking and auth-required calls, complete immediately to allow
+                        // cleanup to proceed while consumption continues in background.
+                        if (!blocking) {
+                            consumptionCompletionFuture.complete(null);
+                        }
+                        // else: blocking calls wait for actual consumption completion in onComplete
 
                         // Continue consuming in background - keep requesting events
                         // Note: continueInBackground is always true when shouldInterrupt is true
@@ -244,8 +247,8 @@ else if (!blocking) {
             }
         }
 
-        // Background consumption continues automatically via the subscription
-        // returning true in the consumer function keeps the subscription alive
+        // Note: For blocking calls that were interrupted, the wait logic has been moved
+        // to DefaultRequestHandler.onMessageSend() to avoid blocking Vert.x worker threads.
         // Queue lifecycle is managed by DefaultRequestHandler.cleanupProducer()
 
         Throwable error = errorRef.get();

server-common/src/test/java/io/a2a/server/requesthandlers/AbstractA2ARequestHandlerTest.java

@@ -98,7 +98,8 @@ public void cancel(RequestContext context, EventQueue eventQueue) throws JSONRPC
         PushNotificationConfigStore pushConfigStore = new InMemoryPushNotificationConfigStore();
         PushNotificationSender pushSender = new BasePushNotificationSender(pushConfigStore, httpClient);
 
-        requestHandler = new DefaultRequestHandler(executor, taskStore, queueManager, pushConfigStore, pushSender, internalExecutor);
+        requestHandler = DefaultRequestHandler.create(
+                executor, taskStore, queueManager, pushConfigStore, pushSender, internalExecutor);
     }
 
     @AfterEach