Skip to content

Commit bbf0c02

Browse files
authored
fix(server): deliver awaited final events to gracefully closed child queues (#1227)
* fix(server): deliver awaited final events to gracefully closed child queues A blocking message/send closes the caller's child queue as soon as the agent returns, while the final event can still be in the MainEventBus. Since dabca6c, distributeToChildren skipped every closed child, so the final event was lost and the EventConsumer waited out its 3-second grace period for it. Deliver to a closed child that is still awaiting a final event; a gracefully closed child still drains its queue. Fixes #1226 * fix(server): correct the child-queue comment and assert the held event on the main thread * test(server): wait for the caller's queue to close instead of sleeping --------- Co-authored-by: Jeroen Bellen <1204135+jeroenbellen@users.noreply.github.com>
1 parent 9bc9134 commit bbf0c02

2 files changed

Lines changed: 108 additions & 8 deletions

File tree

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

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

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

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import static org.mockito.Mockito.mock;
1111
import static org.mockito.Mockito.when;
1212

13+
import java.time.Duration;
1314
import java.time.Instant;
1415
import java.util.Arrays;
1516
import java.util.List;
@@ -24,6 +25,7 @@
2425
import java.util.concurrent.Future;
2526
import java.util.concurrent.TimeUnit;
2627
import java.util.concurrent.TimeoutException;
28+
import java.util.concurrent.atomic.AtomicBoolean;
2729
import java.util.concurrent.atomic.AtomicReference;
2830

2931
import org.a2aproject.sdk.server.ServerCallContext;
@@ -1750,6 +1752,104 @@ public void cancel(RequestContext context, AgentEmitter emitter) {
17501752
assertEquals(TaskState.TASK_STATE_FAILED, storedTask.status().state());
17511753
}
17521754

1755+
/**
1756+
* A blocking call returns to the handler on the agent's first event, and the handler closes the
1757+
* caller's queue as soon as the agent returns. A final event the event bus distributes after that
1758+
* must still reach the queue, which was told to expect it; otherwise the consumer waits out its
1759+
* grace period for it and the call fails or returns late.
1760+
*/
1761+
@Test
1762+
void blockingSendStillReturnsCompletedTaskWhenFinalEventArrivesLate() throws Exception {
1763+
AtomicReference<EventQueue> callerQueue = new AtomicReference<>();
1764+
useQueueManager(new InMemoryQueueManager((TaskStateProvider) taskStore, mainEventBus) {
1765+
@Override
1766+
public EventQueue createOrTap(String taskId) {
1767+
EventQueue queue = super.createOrTap(taskId);
1768+
callerQueue.compareAndSet(null, queue);
1769+
return queue;
1770+
}
1771+
});
1772+
AtomicBoolean callerQueueClosedInTime = new AtomicBoolean();
1773+
holdEventsAfterWorkingUntilClosed(callerQueue, callerQueueClosedInTime);
1774+
agentExecutorExecute = (context, emitter) -> {
1775+
emitter.startWork();
1776+
emitter.complete();
1777+
};
1778+
1779+
EventKind result = requestHandler.onMessageSend(blockingParams("msg-blocking-final-event"), NULL_CONTEXT);
1780+
1781+
assertTrue(callerQueueClosedInTime.get(), "the caller's queue was not closed while the final event was held");
1782+
Task task = assertInstanceOf(Task.class, result);
1783+
assertEquals(TaskState.TASK_STATE_COMPLETED, task.status().state());
1784+
}
1785+
1786+
private void useQueueManager(InMemoryQueueManager manager) {
1787+
EventQueueUtil.stop(mainEventBusProcessor);
1788+
queueManager = manager;
1789+
mainEventBusProcessor = new MainEventBusProcessor(mainEventBus, taskStore,
1790+
PushNotificationSender.NOOP, queueManager);
1791+
EventQueueUtil.start(mainEventBusProcessor);
1792+
requestHandler = DefaultRequestHandler.builder()
1793+
.agentExecutor(executor)
1794+
.taskStore(taskStore)
1795+
.queueManager(queueManager)
1796+
.pushConfigStore(pushConfigStore)
1797+
.pushNotificationsEnabled(true)
1798+
.mainEventBusProcessor(mainEventBusProcessor)
1799+
.executor(internalExecutor)
1800+
.eventConsumerExecutor(internalExecutor)
1801+
.authorizationRequired(false)
1802+
.build();
1803+
}
1804+
1805+
private void holdEventsAfterWorkingUntilClosed(AtomicReference<EventQueue> callerQueue, AtomicBoolean closedInTime) {
1806+
mainEventBusProcessor.setCallback(new MainEventBusProcessorCallback() {
1807+
@Override
1808+
public void onEventProcessed(String taskId, Event event) {
1809+
if (event instanceof TaskStatusUpdateEvent statusUpdate
1810+
&& statusUpdate.status().state() == TaskState.TASK_STATE_WORKING) {
1811+
// Hold the agent's next event in the bus until the handler has closed the caller's queue
1812+
closedInTime.set(awaitClosed(callerQueue, Duration.ofSeconds(5)));
1813+
}
1814+
}
1815+
1816+
@Override
1817+
public void onTaskFinalized(String taskId) {
1818+
}
1819+
});
1820+
}
1821+
1822+
private static boolean awaitClosed(AtomicReference<EventQueue> queue, Duration timeout) {
1823+
long deadline = System.nanoTime() + timeout.toNanos();
1824+
while (System.nanoTime() < deadline) {
1825+
EventQueue current = queue.get();
1826+
if (current != null && current.isClosed()) {
1827+
return true;
1828+
}
1829+
try {
1830+
Thread.sleep(10);
1831+
} catch (InterruptedException e) {
1832+
Thread.currentThread().interrupt();
1833+
return false;
1834+
}
1835+
}
1836+
return false;
1837+
}
1838+
1839+
private static MessageSendParams blockingParams(String messageId) {
1840+
return MessageSendParams.builder()
1841+
.message(Message.builder()
1842+
.messageId(messageId)
1843+
.role(Message.Role.ROLE_USER)
1844+
.parts(new TextPart("hello"))
1845+
.build())
1846+
.configuration(MessageSendConfiguration.builder()
1847+
.returnImmediately(false)
1848+
.acceptedOutputModes(List.of())
1849+
.build())
1850+
.build();
1851+
}
1852+
17531853
@Test
17541854
void interruptedTerminalEnqueueStillReportsAgentFailure() throws Exception {
17551855
assertInterruptedTerminalEnqueueReportsError(

0 commit comments

Comments
 (0)