diff --git a/server-common/src/main/java/org/a2aproject/sdk/server/events/EventQueue.java b/server-common/src/main/java/org/a2aproject/sdk/server/events/EventQueue.java index c018bee5a..c35d35cab 100644 --- a/server-common/src/main/java/org/a2aproject/sdk/server/events/EventQueue.java +++ b/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; diff --git a/server-common/src/test/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandlerTest.java b/server-common/src/test/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandlerTest.java index 66340fd0f..4401f89f0 100644 --- a/server-common/src/test/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandlerTest.java +++ b/server-common/src/test/java/org/a2aproject/sdk/server/requesthandlers/DefaultRequestHandlerTest.java @@ -24,6 +24,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 +1751,66 @@ 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 { + CountDownLatch agentReturned = new CountDownLatch(1); + AtomicBoolean agentReturnedInTime = new AtomicBoolean(); + holdEventsAfterWorkingUntilAgentReturned(agentReturned, agentReturnedInTime); + agentExecutorExecute = (context, emitter) -> { + emitter.startWork(); + emitter.complete(); + agentReturned.countDown(); + }; + + EventKind result = requestHandler.onMessageSend(blockingParams("msg-blocking-final-event"), NULL_CONTEXT); + + assertTrue(agentReturnedInTime.get(), "the agent did not return while its next event was held"); + Task task = assertInstanceOf(Task.class, result); + assertEquals(TaskState.TASK_STATE_COMPLETED, task.status().state()); + } + + private void holdEventsAfterWorkingUntilAgentReturned(CountDownLatch agentReturned, AtomicBoolean agentReturnedInTime) { + mainEventBusProcessor.setCallback(new MainEventBusProcessorCallback() { + @Override + public void onEventProcessed(String taskId, Event event) { + if (event instanceof TaskStatusUpdateEvent statusUpdate + && statusUpdate.status().state() == TaskState.TASK_STATE_WORKING) { + try { + // Hold the agent's next event in the bus until the handler has closed the queue + agentReturnedInTime.set(agentReturned.await(5, TimeUnit.SECONDS)); + Thread.sleep(200); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } + } + + @Override + public void onTaskFinalized(String taskId) { + } + }); + } + + 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(