fix(server): deliver awaited final events to gracefully closed child queues

server-common/src/main/java/org/a2aproject/sdk/server/events/EventQueue.java

@@ -769,14 +769,14 @@ void distributeToChildren(EventQueueItem item) {
                         taskId, item.getEvent().getClass().getSimpleName(), childCount);
             }
             children.forEach(child -> {
-                // Skip children that have already been closed. CopyOnWriteArrayList gives
-                // snapshot semantics, so a child removed by closeIfNotAwaitingFinalEvent()
-                // may still appear in this iteration. Checking isClosed() here narrows the
-                // TOCTOU window to the time between the check and internalEnqueueItem();
-                // a narrow residual race remains for graceful close but is harmless for
-                // non-final events (they'd land in the draining queue) and impossible for
-                // final events (submittedFinalEvents prevents close while one is in flight).
-                if (child.isClosed()) {
+                // Skip children that have already been closed, unless a child is still awaiting
+                // a final event. CopyOnWriteArrayList gives snapshot semantics, so a child removed
+                // by closeIfNotAwaitingFinalEvent() may still appear in this iteration; for a
+                // non-final event that is harmless (it lands in the draining queue). A child that
+                // was closed gracefully while awaiting a final event must still receive it: a
+                // blocking call closes the caller's queue before the event bus may have
+                // distributed the final event.
+                if (child.isClosed() && !child.isAwaitingFinalEvent()) {
                     LOGGER.debug("MainQueue[{}]: Skipping closed child queue for event {}",
                             taskId, item.getEvent().getClass().getSimpleName());
                     return;

server-common/src/test/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandlerTest.java

@@ -10,6 +10,7 @@
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
+import java.time.Duration;
 import java.time.Instant;
 import java.util.Arrays;
 import java.util.List;
@@ -24,6 +25,7 @@
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 
 import org.a2aproject.sdk.server.ServerCallContext;
@@ -1750,6 +1752,104 @@ public void cancel(RequestContext context, AgentEmitter emitter) {
         assertEquals(TaskState.TASK_STATE_FAILED, storedTask.status().state());
     }
 
+    /**
+     * A blocking call returns to the handler on the agent's first event, and the handler closes the
+     * caller's queue as soon as the agent returns. A final event the event bus distributes after that
+     * must still reach the queue, which was told to expect it; otherwise the consumer waits out its
+     * grace period for it and the call fails or returns late.
+     */
+    @Test
+    void blockingSendStillReturnsCompletedTaskWhenFinalEventArrivesLate() throws Exception {
+        AtomicReference<EventQueue> callerQueue = new AtomicReference<>();
+        useQueueManager(new InMemoryQueueManager((TaskStateProvider) taskStore, mainEventBus) {
+            @Override
+            public EventQueue createOrTap(String taskId) {
+                EventQueue queue = super.createOrTap(taskId);
+                callerQueue.compareAndSet(null, queue);
+                return queue;
+            }
+        });
+        AtomicBoolean callerQueueClosedInTime = new AtomicBoolean();
+        holdEventsAfterWorkingUntilClosed(callerQueue, callerQueueClosedInTime);
+        agentExecutorExecute = (context, emitter) -> {
+            emitter.startWork();
+            emitter.complete();
+        };
+
+        EventKind result = requestHandler.onMessageSend(blockingParams("msg-blocking-final-event"), NULL_CONTEXT);
+
+        assertTrue(callerQueueClosedInTime.get(), "the caller's queue was not closed while the final event was held");
+        Task task = assertInstanceOf(Task.class, result);
+        assertEquals(TaskState.TASK_STATE_COMPLETED, task.status().state());
+    }
+
+    private void useQueueManager(InMemoryQueueManager manager) {
+        EventQueueUtil.stop(mainEventBusProcessor);
+        queueManager = manager;
+        mainEventBusProcessor = new MainEventBusProcessor(mainEventBus, taskStore,
+                PushNotificationSender.NOOP, queueManager);
+        EventQueueUtil.start(mainEventBusProcessor);
+        requestHandler = DefaultRequestHandler.builder()
+                .agentExecutor(executor)
+                .taskStore(taskStore)
+                .queueManager(queueManager)
+                .pushConfigStore(pushConfigStore)
+                .pushNotificationsEnabled(true)
+                .mainEventBusProcessor(mainEventBusProcessor)
+                .executor(internalExecutor)
+                .eventConsumerExecutor(internalExecutor)
+                .authorizationRequired(false)
+                .build();
+    }
+
+    private void holdEventsAfterWorkingUntilClosed(AtomicReference<EventQueue> callerQueue, AtomicBoolean closedInTime) {
+        mainEventBusProcessor.setCallback(new MainEventBusProcessorCallback() {
+            @Override
+            public void onEventProcessed(String taskId, Event event) {
+                if (event instanceof TaskStatusUpdateEvent statusUpdate
+                        && statusUpdate.status().state() == TaskState.TASK_STATE_WORKING) {
+                    // Hold the agent's next event in the bus until the handler has closed the caller's queue
+                    closedInTime.set(awaitClosed(callerQueue, Duration.ofSeconds(5)));
+                }
+            }
+
+            @Override
+            public void onTaskFinalized(String taskId) {
+            }
+        });
+    }
+
+    private static boolean awaitClosed(AtomicReference<EventQueue> queue, Duration timeout) {
+        long deadline = System.nanoTime() + timeout.toNanos();
+        while (System.nanoTime() < deadline) {
+            EventQueue current = queue.get();
+            if (current != null && current.isClosed()) {
+                return true;
+            }
+            try {
+                Thread.sleep(10);
+            } catch (InterruptedException e) {
+                Thread.currentThread().interrupt();
+                return false;
+            }
+        }
+        return false;
+    }
+
+    private static MessageSendParams blockingParams(String messageId) {
+        return MessageSendParams.builder()
+            .message(Message.builder()
+                .messageId(messageId)
+                .role(Message.Role.ROLE_USER)
+                .parts(new TextPart("hello"))
+                .build())
+            .configuration(MessageSendConfiguration.builder()
+                .returnImmediately(false)
+                .acceptedOutputModes(List.of())
+                .build())
+            .build();
+    }
+
     @Test
     void interruptedTerminalEnqueueStillReportsAgentFailure() throws Exception {
         assertInterruptedTerminalEnqueueReportsError(