From 6ace8d1b740695b295ac5586a86c0e7bd640a1da Mon Sep 17 00:00:00 2001 From: Kabir Khan Date: Mon, 8 Dec 2025 11:08:14 +0000 Subject: [PATCH 1/8] feat: Rework Event Queues to use central event bus --- .../core/ReplicatedQueueManager.java | 18 +- .../core/ReplicatedQueueManagerTest.java | 39 +- .../io/a2a/server/events/EventQueueUtil.java | 16 + .../src/test/resources/application.properties | 5 + .../java/io/a2a/server/events/EventQueue.java | 76 +++- .../server/events/InMemoryQueueManager.java | 32 +- .../io/a2a/server/events/MainEventBus.java | 42 ++ .../server/events/MainEventBusContext.java | 28 ++ .../server/events/MainEventBusProcessor.java | 252 +++++++++++ .../events/MainEventBusProcessorCallback.java | 65 +++ .../MainEventBusProcessorInitializer.java | 43 ++ .../io/a2a/server/events/QueueManager.java | 20 + .../DefaultRequestHandler.java | 55 +-- .../io/a2a/server/tasks/ResultAggregator.java | 44 +- .../io/a2a/server/events/EventQueueTest.java | 69 ++- .../io/a2a/server/events/EventQueueUtil.java | 8 + .../events/InMemoryQueueManagerTest.java | 19 +- .../AbstractA2ARequestHandlerTest.java | 24 +- .../DefaultRequestHandlerTest.java | 29 +- .../server/tasks/ResultAggregatorTest.java | 7 +- .../grpc/handler/GrpcHandlerTest.java | 8 +- .../jsonrpc/handler/JSONRPCHandlerTest.java | 414 +++++++++--------- 22 files changed, 973 insertions(+), 340 deletions(-) create mode 100644 extras/queue-manager-replicated/core/src/test/java/io/a2a/server/events/EventQueueUtil.java create mode 100644 server-common/src/main/java/io/a2a/server/events/MainEventBus.java create mode 100644 server-common/src/main/java/io/a2a/server/events/MainEventBusContext.java create mode 100644 server-common/src/main/java/io/a2a/server/events/MainEventBusProcessor.java create mode 100644 server-common/src/main/java/io/a2a/server/events/MainEventBusProcessorCallback.java create mode 100644 server-common/src/main/java/io/a2a/server/events/MainEventBusProcessorInitializer.java diff --git a/extras/queue-manager-replicated/core/src/main/java/io/a2a/extras/queuemanager/replicated/core/ReplicatedQueueManager.java b/extras/queue-manager-replicated/core/src/main/java/io/a2a/extras/queuemanager/replicated/core/ReplicatedQueueManager.java index 88c0c75ab..c526c81ef 100644 --- a/extras/queue-manager-replicated/core/src/main/java/io/a2a/extras/queuemanager/replicated/core/ReplicatedQueueManager.java +++ b/extras/queue-manager-replicated/core/src/main/java/io/a2a/extras/queuemanager/replicated/core/ReplicatedQueueManager.java @@ -13,6 +13,7 @@ import io.a2a.server.events.EventQueueFactory; import io.a2a.server.events.EventQueueItem; import io.a2a.server.events.InMemoryQueueManager; +import io.a2a.server.events.MainEventBus; import io.a2a.server.events.QueueManager; import io.a2a.server.tasks.TaskStateProvider; import org.slf4j.Logger; @@ -31,10 +32,12 @@ public class ReplicatedQueueManager implements QueueManager { private TaskStateProvider taskStateProvider; @Inject - public ReplicatedQueueManager(ReplicationStrategy replicationStrategy, TaskStateProvider taskStateProvider) { + public ReplicatedQueueManager(ReplicationStrategy replicationStrategy, + TaskStateProvider taskStateProvider, + MainEventBus mainEventBus) { this.replicationStrategy = replicationStrategy; this.taskStateProvider = taskStateProvider; - this.delegate = new InMemoryQueueManager(new ReplicatingEventQueueFactory(), taskStateProvider); + this.delegate = new InMemoryQueueManager(new ReplicatingEventQueueFactory(), taskStateProvider, mainEventBus); } @@ -138,12 +141,11 @@ public EventQueue.EventQueueBuilder builder(String taskId) { // which sends the QueueClosedEvent after the database transaction commits. // This ensures proper ordering and transactional guarantees. - // Return the builder with callbacks - return delegate.getEventQueueBuilder(taskId) - .taskId(taskId) - .hook(new ReplicationHook(taskId)) - .addOnCloseCallback(delegate.getCleanupCallback(taskId)) - .taskStateProvider(taskStateProvider); + // Call createBaseEventQueueBuilder() directly to avoid infinite recursion + // (getEventQueueBuilder() would delegate back to this factory, creating a loop) + // The base builder already includes: taskId, cleanup callback, taskStateProvider, mainEventBus + return delegate.createBaseEventQueueBuilder(taskId) + .hook(new ReplicationHook(taskId)); } } diff --git a/extras/queue-manager-replicated/core/src/test/java/io/a2a/extras/queuemanager/replicated/core/ReplicatedQueueManagerTest.java b/extras/queue-manager-replicated/core/src/test/java/io/a2a/extras/queuemanager/replicated/core/ReplicatedQueueManagerTest.java index 43571cd30..05b81cea0 100644 --- a/extras/queue-manager-replicated/core/src/test/java/io/a2a/extras/queuemanager/replicated/core/ReplicatedQueueManagerTest.java +++ b/extras/queue-manager-replicated/core/src/test/java/io/a2a/extras/queuemanager/replicated/core/ReplicatedQueueManagerTest.java @@ -22,7 +22,12 @@ import io.a2a.server.events.EventQueueClosedException; import io.a2a.server.events.EventQueueItem; import io.a2a.server.events.EventQueueTestHelper; +import io.a2a.server.events.EventQueueUtil; +import io.a2a.server.events.MainEventBus; +import io.a2a.server.events.MainEventBusProcessor; import io.a2a.server.events.QueueClosedEvent; +import io.a2a.server.tasks.InMemoryTaskStore; +import io.a2a.server.tasks.PushNotificationSender; import io.a2a.spec.Event; import io.a2a.spec.StreamingEventKind; import io.a2a.spec.TaskState; @@ -35,10 +40,24 @@ class ReplicatedQueueManagerTest { private ReplicatedQueueManager queueManager; private StreamingEventKind testEvent; + private MainEventBus mainEventBus; + private MainEventBusProcessor mainEventBusProcessor; + private static final PushNotificationSender NOOP_PUSHNOTIFICATION_SENDER = task -> {}; @BeforeEach void setUp() { - queueManager = new ReplicatedQueueManager(new NoOpReplicationStrategy(), new MockTaskStateProvider(true)); + // Create MainEventBus and MainEventBusProcessor for tests + InMemoryTaskStore taskStore = new InMemoryTaskStore(); + mainEventBus = new MainEventBus(); + mainEventBusProcessor = new MainEventBusProcessor(mainEventBus, taskStore, NOOP_PUSHNOTIFICATION_SENDER); + EventQueueUtil.start(mainEventBusProcessor); + + queueManager = new ReplicatedQueueManager( + new NoOpReplicationStrategy(), + new MockTaskStateProvider(true), + mainEventBus + ); + testEvent = TaskStatusUpdateEvent.builder() .taskId("test-task") .contextId("test-context") @@ -50,7 +69,7 @@ void setUp() { @Test void testReplicationStrategyTriggeredOnNormalEnqueue() throws InterruptedException { CountingReplicationStrategy strategy = new CountingReplicationStrategy(); - queueManager = new ReplicatedQueueManager(strategy, new MockTaskStateProvider(true)); + queueManager = new ReplicatedQueueManager(strategy, new MockTaskStateProvider(true), mainEventBus); String taskId = "test-task-1"; EventQueue queue = queueManager.createOrTap(taskId); @@ -65,7 +84,7 @@ void testReplicationStrategyTriggeredOnNormalEnqueue() throws InterruptedExcepti @Test void testReplicationStrategyNotTriggeredOnReplicatedEvent() throws InterruptedException { CountingReplicationStrategy strategy = new CountingReplicationStrategy(); - queueManager = new ReplicatedQueueManager(strategy, new MockTaskStateProvider(true)); + queueManager = new ReplicatedQueueManager(strategy, new MockTaskStateProvider(true), mainEventBus); String taskId = "test-task-2"; EventQueue queue = queueManager.createOrTap(taskId); @@ -79,7 +98,7 @@ void testReplicationStrategyNotTriggeredOnReplicatedEvent() throws InterruptedEx @Test void testReplicationStrategyWithCountingImplementation() throws InterruptedException { CountingReplicationStrategy countingStrategy = new CountingReplicationStrategy(); - queueManager = new ReplicatedQueueManager(countingStrategy, new MockTaskStateProvider(true)); + queueManager = new ReplicatedQueueManager(countingStrategy, new MockTaskStateProvider(true), mainEventBus); String taskId = "test-task-3"; EventQueue queue = queueManager.createOrTap(taskId); @@ -170,7 +189,7 @@ void testBasicQueueManagerFunctionality() throws InterruptedException { void testQueueToTaskIdMappingMaintained() throws InterruptedException { String taskId = "test-task-6"; CountingReplicationStrategy countingStrategy = new CountingReplicationStrategy(); - queueManager = new ReplicatedQueueManager(countingStrategy, new MockTaskStateProvider(true)); + queueManager = new ReplicatedQueueManager(countingStrategy, new MockTaskStateProvider(true), mainEventBus); EventQueue queue = queueManager.createOrTap(taskId); queue.enqueueEvent(testEvent); @@ -217,7 +236,7 @@ void testReplicatedEventJsonSerialization() throws Exception { @Test void testParallelReplicationBehavior() throws InterruptedException { CountingReplicationStrategy strategy = new CountingReplicationStrategy(); - queueManager = new ReplicatedQueueManager(strategy, new MockTaskStateProvider(true)); + queueManager = new ReplicatedQueueManager(strategy, new MockTaskStateProvider(true), mainEventBus); String taskId = "parallel-test-task"; EventQueue queue = queueManager.createOrTap(taskId); @@ -297,7 +316,7 @@ void testParallelReplicationBehavior() throws InterruptedException { void testReplicatedEventSkippedWhenTaskInactive() throws InterruptedException { // Create a task state provider that returns false (task is inactive) MockTaskStateProvider stateProvider = new MockTaskStateProvider(false); - queueManager = new ReplicatedQueueManager(new CountingReplicationStrategy(), stateProvider); + queueManager = new ReplicatedQueueManager(new CountingReplicationStrategy(), stateProvider, mainEventBus); String taskId = "inactive-task"; @@ -316,7 +335,7 @@ void testReplicatedEventSkippedWhenTaskInactive() throws InterruptedException { void testReplicatedEventProcessedWhenTaskActive() throws InterruptedException { // Create a task state provider that returns true (task is active) MockTaskStateProvider stateProvider = new MockTaskStateProvider(true); - queueManager = new ReplicatedQueueManager(new CountingReplicationStrategy(), stateProvider); + queueManager = new ReplicatedQueueManager(new CountingReplicationStrategy(), stateProvider, mainEventBus); String taskId = "active-task"; @@ -347,7 +366,7 @@ void testReplicatedEventProcessedWhenTaskActive() throws InterruptedException { void testReplicatedEventToExistingQueueWhenTaskBecomesInactive() throws InterruptedException { // Create a task state provider that returns true initially MockTaskStateProvider stateProvider = new MockTaskStateProvider(true); - queueManager = new ReplicatedQueueManager(new CountingReplicationStrategy(), stateProvider); + queueManager = new ReplicatedQueueManager(new CountingReplicationStrategy(), stateProvider, mainEventBus); String taskId = "task-becomes-inactive"; @@ -387,7 +406,7 @@ void testReplicatedEventToExistingQueueWhenTaskBecomesInactive() throws Interrup @Test void testPoisonPillSentViaTransactionAwareEvent() throws InterruptedException { CountingReplicationStrategy strategy = new CountingReplicationStrategy(); - queueManager = new ReplicatedQueueManager(strategy, new MockTaskStateProvider(true)); + queueManager = new ReplicatedQueueManager(strategy, new MockTaskStateProvider(true), mainEventBus); String taskId = "poison-pill-test"; EventQueue queue = queueManager.createOrTap(taskId); diff --git a/extras/queue-manager-replicated/core/src/test/java/io/a2a/server/events/EventQueueUtil.java b/extras/queue-manager-replicated/core/src/test/java/io/a2a/server/events/EventQueueUtil.java new file mode 100644 index 000000000..8490aa00c --- /dev/null +++ b/extras/queue-manager-replicated/core/src/test/java/io/a2a/server/events/EventQueueUtil.java @@ -0,0 +1,16 @@ +package io.a2a.server.events; + +public class EventQueueUtil { + // Since EventQueue.builder() is package protected, add a method to expose it + public static EventQueue.EventQueueBuilder getEventQueueBuilder() { + return EventQueue.builder(); + } + + public static void start(MainEventBusProcessor processor) { + processor.start(); + } + + public static void stop(MainEventBusProcessor processor) { + processor.stop(); + } +} diff --git a/reference/jsonrpc/src/test/resources/application.properties b/reference/jsonrpc/src/test/resources/application.properties index 7b9cea9cc..e612925d4 100644 --- a/reference/jsonrpc/src/test/resources/application.properties +++ b/reference/jsonrpc/src/test/resources/application.properties @@ -1 +1,6 @@ quarkus.arc.selected-alternatives=io.a2a.server.apps.common.TestHttpClient + +# Debug logging for event processing and request handling +quarkus.log.category."io.a2a.server.events".level=DEBUG +quarkus.log.category."io.a2a.server.requesthandlers".level=DEBUG +quarkus.log.category."io.a2a.server.tasks".level=DEBUG diff --git a/server-common/src/main/java/io/a2a/server/events/EventQueue.java b/server-common/src/main/java/io/a2a/server/events/EventQueue.java index a08f63084..a0bb074cb 100644 --- a/server-common/src/main/java/io/a2a/server/events/EventQueue.java +++ b/server-common/src/main/java/io/a2a/server/events/EventQueue.java @@ -95,6 +95,7 @@ public static class EventQueueBuilder { private @Nullable String taskId; private List onCloseCallbacks = new java.util.ArrayList<>(); private @Nullable TaskStateProvider taskStateProvider; + private @Nullable MainEventBus mainEventBus; /** * Sets the maximum queue size. @@ -153,14 +154,25 @@ public EventQueueBuilder taskStateProvider(TaskStateProvider taskStateProvider) return this; } + /** + * Sets the main event bus + * + * @param mainEventBus the main event bus + * @return this builder + */ + public EventQueueBuilder mainEventBus(MainEventBus mainEventBus) { + this.mainEventBus = mainEventBus; + return this; + } + /** * Builds and returns the configured EventQueue. * * @return a new MainQueue instance */ public EventQueue build() { - if (hook != null || !onCloseCallbacks.isEmpty() || taskStateProvider != null) { - return new MainQueue(queueSize, hook, taskId, onCloseCallbacks, taskStateProvider); + if (hook != null || !onCloseCallbacks.isEmpty() || taskStateProvider != null || mainEventBus != null) { + return new MainQueue(queueSize, hook, taskId, onCloseCallbacks, taskStateProvider, mainEventBus); } else { return new MainQueue(queueSize); } @@ -365,6 +377,7 @@ static class MainQueue extends EventQueue { private final @Nullable String taskId; private final List onCloseCallbacks; private final @Nullable TaskStateProvider taskStateProvider; + private final @Nullable MainEventBus mainEventBus; MainQueue() { super(); @@ -372,6 +385,7 @@ static class MainQueue extends EventQueue { this.taskId = null; this.onCloseCallbacks = List.of(); this.taskStateProvider = null; + this.mainEventBus = null; } MainQueue(int queueSize) { @@ -380,6 +394,7 @@ static class MainQueue extends EventQueue { this.taskId = null; this.onCloseCallbacks = List.of(); this.taskStateProvider = null; + this.mainEventBus = null; } MainQueue(EventEnqueueHook hook) { @@ -388,6 +403,7 @@ static class MainQueue extends EventQueue { this.taskId = null; this.onCloseCallbacks = List.of(); this.taskStateProvider = null; + this.mainEventBus = null; } MainQueue(int queueSize, EventEnqueueHook hook) { @@ -396,24 +412,40 @@ static class MainQueue extends EventQueue { this.taskId = null; this.onCloseCallbacks = List.of(); this.taskStateProvider = null; + this.mainEventBus = null; } - MainQueue(int queueSize, @Nullable EventEnqueueHook hook, @Nullable String taskId, List onCloseCallbacks, @Nullable TaskStateProvider taskStateProvider) { + MainQueue(int queueSize, + @Nullable EventEnqueueHook hook, + @Nullable String taskId, + List onCloseCallbacks, + @Nullable TaskStateProvider taskStateProvider, + @Nullable MainEventBus mainEventBus) { super(queueSize); this.enqueueHook = hook; this.taskId = taskId; this.onCloseCallbacks = List.copyOf(onCloseCallbacks); // Defensive copy this.taskStateProvider = taskStateProvider; - LOGGER.debug("Created MainQueue for task {} with {} onClose callbacks and TaskStateProvider: {}", - taskId, onCloseCallbacks.size(), taskStateProvider != null); + this.mainEventBus = mainEventBus; + LOGGER.debug("Created MainQueue for task {} with {} onClose callbacks, TaskStateProvider: {}, MainEventBus: {}", + taskId, onCloseCallbacks.size(), taskStateProvider != null, mainEventBus != null); } + public EventQueue tap() { ChildQueue child = new ChildQueue(this); children.add(child); return child; } + /** + * Returns the current number of child queues. + * Useful for debugging and logging event distribution. + */ + public int getChildCount() { + return children.size(); + } + @Override public void enqueueItem(EventQueueItem item) { // MainQueue must accept events even when closed to support: @@ -436,10 +468,15 @@ public void enqueueItem(EventQueueItem item) { queue.add(item); LOGGER.debug("Enqueued event {} {}", event instanceof Throwable ? event.toString() : event, this); - // Distribute to all ChildQueues (they will receive the event even if MainQueue is closed) - children.forEach(eq -> eq.internalEnqueueItem(item)); + // Submit to MainEventBus for centralized persistence + distribution + if (mainEventBus != null && taskId != null) { + mainEventBus.submit(taskId, this, item); + } else { + // This should not happen in properly configured systems + LOGGER.error("MainEventBus not configured for task {} - events will NOT be distributed to children!", taskId); + } - // Trigger replication hook if configured + // Trigger replication hook if configured (KEEP for inter-process replication) if (enqueueHook != null) { enqueueHook.onEnqueue(item); } @@ -493,6 +530,29 @@ void childClosing(ChildQueue child, boolean immediate) { this.doClose(immediate); } + /** + * Distribute event to all ChildQueues. + * Called by MainEventBusProcessor after TaskStore persistence. + */ + void distributeToChildren(EventQueueItem item) { + synchronized (children) { + int childCount = children.size(); + if (LOGGER.isDebugEnabled()) { + LOGGER.debug("MainQueue[{}]: Distributing event {} to {} children", + taskId, item.getEvent().getClass().getSimpleName(), childCount); + } + children.forEach(child -> { + LOGGER.debug("MainQueue[{}]: Enqueueing event {} to child queue", + taskId, item.getEvent().getClass().getSimpleName()); + child.internalEnqueueItem(item); + }); + if (LOGGER.isDebugEnabled()) { + LOGGER.debug("MainQueue[{}]: Completed distribution of {} to {} children", + taskId, item.getEvent().getClass().getSimpleName(), childCount); + } + } + } + /** * Get the count of active child queues. * Used for testing to verify reference counting mechanism. diff --git a/server-common/src/main/java/io/a2a/server/events/InMemoryQueueManager.java b/server-common/src/main/java/io/a2a/server/events/InMemoryQueueManager.java index 19179f50d..63063ba69 100644 --- a/server-common/src/main/java/io/a2a/server/events/InMemoryQueueManager.java +++ b/server-common/src/main/java/io/a2a/server/events/InMemoryQueueManager.java @@ -19,16 +19,20 @@ public class InMemoryQueueManager implements QueueManager { private final EventQueueFactory factory; private final TaskStateProvider taskStateProvider; + MainEventBus mainEventBus; + @Inject - public InMemoryQueueManager(TaskStateProvider taskStateProvider) { + public InMemoryQueueManager(TaskStateProvider taskStateProvider, MainEventBus mainEventBus) { + this.mainEventBus = mainEventBus; this.factory = new DefaultEventQueueFactory(); this.taskStateProvider = taskStateProvider; } - // For testing with custom factory - public InMemoryQueueManager(EventQueueFactory factory, TaskStateProvider taskStateProvider) { + // For testing/extensions with custom factory and MainEventBus + public InMemoryQueueManager(EventQueueFactory factory, TaskStateProvider taskStateProvider, MainEventBus mainEventBus) { this.factory = factory; this.taskStateProvider = taskStateProvider; + this.mainEventBus = mainEventBus; } @Override @@ -113,6 +117,12 @@ public void awaitQueuePollerStart(EventQueue eventQueue) throws InterruptedExcep eventQueue.awaitQueuePollerStart(); } + @Override + public EventQueue.EventQueueBuilder getEventQueueBuilder(String taskId) { + // Use the factory to ensure proper configuration (MainEventBus, callbacks, etc.) + return factory.builder(taskId); + } + @Override public int getActiveChildQueueCount(String taskId) { EventQueue queue = queues.get(taskId); @@ -127,6 +137,15 @@ public int getActiveChildQueueCount(String taskId) { return -1; } + @Override + public EventQueue.EventQueueBuilder createBaseEventQueueBuilder(String taskId) { + return EventQueue.builder() + .taskId(taskId) + .addOnCloseCallback(getCleanupCallback(taskId)) + .taskStateProvider(taskStateProvider) + .mainEventBus(mainEventBus); + } + /** * Get the cleanup callback that removes a queue from the map when it closes. * This is exposed so that subclasses (like ReplicatedQueueManager) can reuse @@ -166,11 +185,8 @@ public Runnable getCleanupCallback(String taskId) { private class DefaultEventQueueFactory implements EventQueueFactory { @Override public EventQueue.EventQueueBuilder builder(String taskId) { - // Return builder with callback that removes queue from map when closed - return EventQueue.builder() - .taskId(taskId) - .addOnCloseCallback(getCleanupCallback(taskId)) - .taskStateProvider(taskStateProvider); + // Delegate to the base builder creation method + return createBaseEventQueueBuilder(taskId); } } } diff --git a/server-common/src/main/java/io/a2a/server/events/MainEventBus.java b/server-common/src/main/java/io/a2a/server/events/MainEventBus.java new file mode 100644 index 000000000..73500254e --- /dev/null +++ b/server-common/src/main/java/io/a2a/server/events/MainEventBus.java @@ -0,0 +1,42 @@ +package io.a2a.server.events; + +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.LinkedBlockingDeque; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import jakarta.enterprise.context.ApplicationScoped; + +@ApplicationScoped +public class MainEventBus { + private static final Logger LOGGER = LoggerFactory.getLogger(MainEventBus.class); + private final BlockingQueue queue; + + public MainEventBus() { + this.queue = new LinkedBlockingDeque<>(); + } + + public void submit(String taskId, EventQueue eventQueue, EventQueueItem item) { + try { + queue.put(new MainEventBusContext(taskId, eventQueue, item)); + LOGGER.debug("Submitted event for task {} to MainEventBus (queue size: {})", + taskId, queue.size()); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException("Interrupted submitting to MainEventBus", e); + } + } + + public MainEventBusContext take() throws InterruptedException { + LOGGER.debug("MainEventBus: Waiting to take event (current queue size: {})...", queue.size()); + MainEventBusContext context = queue.take(); + LOGGER.debug("MainEventBus: Took event for task {} (remaining queue size: {})", + context.taskId(), queue.size()); + return context; + } + + public int size() { + return queue.size(); + } +} diff --git a/server-common/src/main/java/io/a2a/server/events/MainEventBusContext.java b/server-common/src/main/java/io/a2a/server/events/MainEventBusContext.java new file mode 100644 index 000000000..caa288849 --- /dev/null +++ b/server-common/src/main/java/io/a2a/server/events/MainEventBusContext.java @@ -0,0 +1,28 @@ +package io.a2a.server.events; + +import java.util.Objects; + +class MainEventBusContext { + private final String taskId; + private final EventQueue eventQueue; + private final EventQueueItem eventQueueItem; + + public MainEventBusContext(String taskId, EventQueue eventQueue, EventQueueItem eventQueueItem) { + this.taskId = Objects.requireNonNull(taskId, "taskId cannot be null"); + this.eventQueue = Objects.requireNonNull(eventQueue, "eventQueue cannot be null"); + this.eventQueueItem = Objects.requireNonNull(eventQueueItem, "eventQueueItem cannot be null"); + } + + public String taskId() { + return taskId; + } + + public EventQueue eventQueue() { + return eventQueue; + } + + public EventQueueItem eventQueueItem() { + return eventQueueItem; + } + +} diff --git a/server-common/src/main/java/io/a2a/server/events/MainEventBusProcessor.java b/server-common/src/main/java/io/a2a/server/events/MainEventBusProcessor.java new file mode 100644 index 000000000..94c3d390e --- /dev/null +++ b/server-common/src/main/java/io/a2a/server/events/MainEventBusProcessor.java @@ -0,0 +1,252 @@ +package io.a2a.server.events; + +import jakarta.annotation.Nullable; +import jakarta.annotation.PostConstruct; +import jakarta.annotation.PreDestroy; +import jakarta.enterprise.context.ApplicationScoped; +import jakarta.inject.Inject; + +import io.a2a.server.tasks.PushNotificationSender; +import io.a2a.server.tasks.TaskManager; +import io.a2a.server.tasks.TaskStore; +import io.a2a.spec.A2AServerException; +import io.a2a.spec.Event; +import io.a2a.spec.Task; +import io.a2a.spec.TaskArtifactUpdateEvent; +import io.a2a.spec.TaskStatusUpdateEvent; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Background processor for the MainEventBus. + *

+ * This processor runs in a dedicated background thread, consuming events from the MainEventBus + * and performing two critical operations in order: + *

+ *
    + *
  1. Update TaskStore with event data (persistence FIRST)
    2. + *
  2. Distribute event to ChildQueues (clients see it AFTER persistence)
    3. + *
+ *

+ * This architecture ensures clients never receive events before they're persisted, + * eliminating race conditions and enabling reliable event replay. + *

+ *

+ * Note: This bean is eagerly initialized by {@link MainEventBusProcessorInitializer} + * to ensure the background thread starts automatically when the application starts. + *

+ */ +@ApplicationScoped +public class MainEventBusProcessor implements Runnable { + private static final Logger LOGGER = LoggerFactory.getLogger(MainEventBusProcessor.class); + + /** + * Callback for testing synchronization with async event processing. + * Default is NOOP to avoid null checks in production code. + * Tests can inject their own callback via setCallback(). + */ + private static volatile MainEventBusProcessorCallback callback = MainEventBusProcessorCallback.NOOP; + + private final MainEventBus eventBus; + + private final TaskStore taskStore; + + private final PushNotificationSender pushSender; + + private volatile boolean running = true; + private @Nullable Thread processorThread; + + @Inject + public MainEventBusProcessor(MainEventBus eventBus, TaskStore taskStore, PushNotificationSender pushSender) { + this.eventBus = eventBus; + this.taskStore = taskStore; + this.pushSender = pushSender; + } + + /** + * Set a callback for testing synchronization with async event processing. + *

+ * This is primarily intended for tests that need to wait for event processing to complete. + * Pass null to reset to the default NOOP callback. + *

+ * + * @param callback the callback to invoke during event processing, or null for NOOP + */ + public static void setCallback(MainEventBusProcessorCallback callback) { + MainEventBusProcessor.callback = callback != null ? callback : MainEventBusProcessorCallback.NOOP; + } + + @PostConstruct + void start() { + processorThread = new Thread(this, "MainEventBusProcessor"); + processorThread.setDaemon(false); // Keep JVM alive + processorThread.start(); + LOGGER.info("MainEventBusProcessor started"); + } + + /** + * No-op method to force CDI proxy resolution and ensure @PostConstruct has been called. + * Called by MainEventBusProcessorInitializer during application startup. + */ + public void ensureStarted() { + // Method intentionally empty - just forces proxy resolution + } + + @PreDestroy + void stop() { + LOGGER.info("MainEventBusProcessor stopping..."); + running = false; + if (processorThread != null) { + processorThread.interrupt(); + try { + processorThread.join(5000); // Wait up to 5 seconds + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } + LOGGER.info("MainEventBusProcessor stopped"); + } + + @Override + public void run() { + LOGGER.info("MainEventBusProcessor processing loop started"); + while (running) { + try { + LOGGER.debug("MainEventBusProcessor: Waiting for event from MainEventBus..."); + MainEventBusContext context = eventBus.take(); + LOGGER.debug("MainEventBusProcessor: Retrieved event for task {} from MainEventBus", + context.taskId()); + processEvent(context); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + LOGGER.info("MainEventBusProcessor interrupted, shutting down"); + break; + } catch (Exception e) { + LOGGER.error("Error processing event from MainEventBus", e); + // Continue processing despite errors + } + } + LOGGER.info("MainEventBusProcessor processing loop ended"); + } + + private void processEvent(MainEventBusContext context) { + String taskId = context.taskId(); + Event event = context.eventQueueItem().getEvent(); + EventQueue eventQueue = context.eventQueue(); + + LOGGER.debug("MainEventBusProcessor: Processing event for task {}: {} (queue type: {})", + taskId, event.getClass().getSimpleName(), eventQueue.getClass().getSimpleName()); + + // Step 1: Update TaskStore FIRST (persistence before clients see it) + updateTaskStore(taskId, event); + + // Step 2: Send push notification AFTER persistence (ensures notification sees latest state) + sendPushNotification(taskId); + + // Step 3: Then distribute to ChildQueues (clients see it AFTER persistence + notification) + if (eventQueue instanceof EventQueue.MainQueue mainQueue) { + int childCount = mainQueue.getChildCount(); + LOGGER.debug("MainEventBusProcessor: Distributing event to {} children for task {}", childCount, taskId); + mainQueue.distributeToChildren(context.eventQueueItem()); + LOGGER.debug("MainEventBusProcessor: Distributed event {} to {} children for task {}", + event.getClass().getSimpleName(), childCount, taskId); + } else { + LOGGER.warn("MainEventBusProcessor: Expected MainQueue but got {} for task {}", + eventQueue.getClass().getSimpleName(), taskId); + } + + LOGGER.debug("MainEventBusProcessor: Completed processing event for task {}", taskId); + + // Step 4: Notify callback after all processing is complete + callback.onEventProcessed(taskId, event); + + // Step 5: If this is a final event, notify task finalization + if (isFinalEvent(event)) { + callback.onTaskFinalized(taskId); + } + } + + /** + * Updates TaskStore using TaskManager.process(). + *

+ * Creates a temporary TaskManager instance for this event and delegates to its process() method, + * which handles all event types (Task, TaskStatusUpdateEvent, TaskArtifactUpdateEvent). + * This leverages existing TaskManager logic for status updates, artifact appending, message history, etc. + *

+ */ + private void updateTaskStore(String taskId, Event event) { + try { + // Extract contextId from event (all relevant events have it) + String contextId = extractContextId(event); + + // Create temporary TaskManager instance for this event + TaskManager taskManager = new TaskManager(taskId, contextId, taskStore, null); + + // Use TaskManager.process() - handles all event types with existing logic + taskManager.process(event); + LOGGER.debug("TaskStore updated via TaskManager.process() for task {}: {}", + taskId, event.getClass().getSimpleName()); + } catch (A2AServerException e) { + LOGGER.error("Error updating TaskStore via TaskManager for task {}", taskId, e); + // Don't rethrow - we still want to distribute to ChildQueues + } catch (Exception e) { + LOGGER.error("Unexpected error updating TaskStore for task {}", taskId, e); + // Don't rethrow - we still want to distribute to ChildQueues + } + } + + /** + * Sends push notification for the task AFTER persistence. + *

+ * This is called after updateTaskStore() to ensure the notification contains + * the latest persisted state, avoiding race conditions. + *

+ */ + private void sendPushNotification(String taskId) { + try { + Task task = taskStore.get(taskId); + if (task != null) { + LOGGER.debug("Sending push notification for task {}", taskId); + pushSender.sendNotification(task); + } else { + LOGGER.debug("Skipping push notification - task {} not found in TaskStore", taskId); + } + } catch (Exception e) { + LOGGER.error("Error sending push notification for task {}", taskId, e); + // Don't rethrow - we still want to distribute to ChildQueues + } + } + + /** + * Extracts contextId from an event. + * Returns null if the event type doesn't have a contextId (e.g., Message). + */ + @Nullable + private String extractContextId(Event event) { + if (event instanceof Task task) { + return task.contextId(); + } else if (event instanceof TaskStatusUpdateEvent statusUpdate) { + return statusUpdate.contextId(); + } else if (event instanceof TaskArtifactUpdateEvent artifactUpdate) { + return artifactUpdate.contextId(); + } + // Message and other events don't have contextId + return null; + } + + /** + * Checks if an event represents a final task state. + * + * @param event the event to check + * @return true if the event represents a final state (COMPLETED, FAILED, CANCELED, REJECTED, UNKNOWN) + */ + private boolean isFinalEvent(Event event) { + if (event instanceof Task task) { + return task.status() != null && task.status().state() != null + && task.status().state().isFinal(); + } else if (event instanceof TaskStatusUpdateEvent statusUpdate) { + return statusUpdate.isFinal(); + } + return false; + } +} diff --git a/server-common/src/main/java/io/a2a/server/events/MainEventBusProcessorCallback.java b/server-common/src/main/java/io/a2a/server/events/MainEventBusProcessorCallback.java new file mode 100644 index 000000000..379acac37 --- /dev/null +++ b/server-common/src/main/java/io/a2a/server/events/MainEventBusProcessorCallback.java @@ -0,0 +1,65 @@ +package io.a2a.server.events; + +import io.a2a.spec.Event; + +/** + * Callback interface for MainEventBusProcessor events. + *

+ * This interface is primarily intended for testing, allowing tests to synchronize + * with the asynchronous MainEventBusProcessor. Production code should not rely on this. + *

+ *

+ * Usage in tests: + * 

+ * {@code
+ * @BeforeEach
+ * void setUp() {
+ *     CountDownLatch latch = new CountDownLatch(3);
+ *     MainEventBusProcessor.setCallback(new MainEventBusProcessorCallback() {
+ *         public void onEventProcessed(String taskId, Event event) {
+ *             latch.countDown();
+ *         }
+ *     });
+ * }
+ *
+ * @AfterEach
+ * void tearDown() {
+ *     MainEventBusProcessor.setCallback(null); // Reset to NOOP
+ * }
+ * }
+ *
+ *

+ */ +public interface MainEventBusProcessorCallback { + + /** + * Called after an event has been fully processed (persisted, notification sent, distributed to children). + * + * @param taskId the task ID + * @param event the event that was processed + */ + void onEventProcessed(String taskId, Event event); + + /** + * Called when a task reaches a final state (COMPLETED, FAILED, CANCELED, REJECTED). + * + * @param taskId the task ID that was finalized + */ + void onTaskFinalized(String taskId); + + /** + * No-op implementation that does nothing. + * Used as the default callback to avoid null checks. + */ + MainEventBusProcessorCallback NOOP = new MainEventBusProcessorCallback() { + @Override + public void onEventProcessed(String taskId, Event event) { + // No-op + } + + @Override + public void onTaskFinalized(String taskId) { + // No-op + } + }; +} diff --git a/server-common/src/main/java/io/a2a/server/events/MainEventBusProcessorInitializer.java b/server-common/src/main/java/io/a2a/server/events/MainEventBusProcessorInitializer.java new file mode 100644 index 000000000..ba4b300be --- /dev/null +++ b/server-common/src/main/java/io/a2a/server/events/MainEventBusProcessorInitializer.java @@ -0,0 +1,43 @@ +package io.a2a.server.events; + +import jakarta.enterprise.context.ApplicationScoped; +import jakarta.enterprise.context.Initialized; +import jakarta.enterprise.event.Observes; +import jakarta.inject.Inject; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Portable CDI initializer for MainEventBusProcessor. + *

+ * This bean observes the ApplicationScoped initialization event and injects + * MainEventBusProcessor, which triggers its eager creation and starts the background thread. + *

+ *

+ * This approach is portable across all Jakarta CDI implementations (Weld, OpenWebBeans, Quarkus, etc.) + * and ensures MainEventBusProcessor starts automatically when the application starts. + *

+ */ +@ApplicationScoped +public class MainEventBusProcessorInitializer { + private static final Logger LOGGER = LoggerFactory.getLogger(MainEventBusProcessorInitializer.class); + + @Inject + MainEventBusProcessor processor; + + /** + * Observes ApplicationScoped initialization to force eager creation of MainEventBusProcessor. + * The injection of MainEventBusProcessor in this bean triggers its creation, and calling + * ensureStarted() forces the CDI proxy to be resolved, which ensures @PostConstruct has been + * called and the background thread is running. + */ + void onStart(@Observes @Initialized(ApplicationScoped.class) Object event) { + if (processor != null) { + // Force proxy resolution to ensure @PostConstruct has been called + processor.ensureStarted(); + LOGGER.info("MainEventBusProcessor initialized and started"); + } else { + LOGGER.error("MainEventBusProcessor is null - initialization failed!"); + } + } +} diff --git a/server-common/src/main/java/io/a2a/server/events/QueueManager.java b/server-common/src/main/java/io/a2a/server/events/QueueManager.java index 01e754fcb..395fddb9a 100644 --- a/server-common/src/main/java/io/a2a/server/events/QueueManager.java +++ b/server-common/src/main/java/io/a2a/server/events/QueueManager.java @@ -180,6 +180,26 @@ default EventQueue.EventQueueBuilder getEventQueueBuilder(String taskId) { return EventQueue.builder(); } + /** + * Creates a base EventQueueBuilder with standard configuration for this QueueManager. + * This method provides the foundation for creating event queues with proper configuration + * (MainEventBus, TaskStateProvider, cleanup callbacks, etc.). + *

+ * QueueManager implementations that use custom factories can call this method directly + * to get the base builder without going through the factory (which could cause infinite + * recursion if the factory delegates back to getEventQueueBuilder()). + *

+ *

+ * Callers can then add additional configuration (hooks, callbacks) before building the queue. + *

+ * + * @param taskId the task ID for the queue + * @return a builder with base configuration specific to this QueueManager implementation + */ + default EventQueue.EventQueueBuilder createBaseEventQueueBuilder(String taskId) { + return EventQueue.builder().taskId(taskId); + } + /** * Returns the number of active ChildQueues for a task. *

diff --git a/server-common/src/main/java/io/a2a/server/requesthandlers/DefaultRequestHandler.java b/server-common/src/main/java/io/a2a/server/requesthandlers/DefaultRequestHandler.java index 61af1d43e..fa55b5142 100644 --- a/server-common/src/main/java/io/a2a/server/requesthandlers/DefaultRequestHandler.java +++ b/server-common/src/main/java/io/a2a/server/requesthandlers/DefaultRequestHandler.java @@ -56,7 +56,6 @@ import io.a2a.spec.Message; import io.a2a.spec.MessageSendParams; import io.a2a.spec.PushNotificationConfig; -import io.a2a.spec.PushNotificationNotSupportedError; import io.a2a.spec.StreamingEventKind; import io.a2a.spec.Task; import io.a2a.spec.TaskIdParams; @@ -212,7 +211,6 @@ public class DefaultRequestHandler implements RequestHandler { private final TaskStore taskStore; private final QueueManager queueManager; private final PushNotificationConfigStore pushConfigStore; - private final PushNotificationSender pushSender; private final Supplier requestContextBuilder; private final ConcurrentMap> runningAgents = new ConcurrentHashMap<>(); @@ -223,12 +221,11 @@ public class DefaultRequestHandler implements RequestHandler { @Inject public DefaultRequestHandler(AgentExecutor agentExecutor, TaskStore taskStore, QueueManager queueManager, PushNotificationConfigStore pushConfigStore, - PushNotificationSender pushSender, @Internal Executor executor) { + @Internal Executor executor) { this.agentExecutor = agentExecutor; this.taskStore = taskStore; this.queueManager = queueManager; this.pushConfigStore = pushConfigStore; - this.pushSender = pushSender; this.executor = executor; // TODO In Python this is also a constructor parameter defaulting to this SimpleRequestContextBuilder // implementation if the parameter is null. Skip that for now, since otherwise I get CDI errors, and @@ -250,9 +247,9 @@ void initConfig() { */ public static DefaultRequestHandler create(AgentExecutor agentExecutor, TaskStore taskStore, QueueManager queueManager, PushNotificationConfigStore pushConfigStore, - PushNotificationSender pushSender, Executor executor) { + Executor executor) { DefaultRequestHandler handler = - new DefaultRequestHandler(agentExecutor, taskStore, queueManager, pushConfigStore, pushSender, executor); + new DefaultRequestHandler(agentExecutor, taskStore, queueManager, pushConfigStore, executor); handler.agentCompletionTimeoutSeconds = 5; handler.consumptionCompletionTimeoutSeconds = 2; return handler; @@ -387,9 +384,6 @@ public EventKind onMessageSend(MessageSendParams params, ServerCallContext conte 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); - EventConsumer consumer = new EventConsumer(queue); // This callback must be added before we start consuming. Otherwise, @@ -483,9 +477,6 @@ public EventKind onMessageSend(MessageSendParams params, ServerCallContext conte if (kind instanceof Task taskResult && !taskId.equals(taskResult.id())) { throw new InternalError("Task ID mismatch in agent response"); } - - // Send push notification after initial return (for both blocking and non-blocking) - pushNotificationCallback.run(); } finally { // Remove agent from map immediately to prevent accumulation CompletableFuture agentFuture = runningAgents.remove(taskId); @@ -557,14 +548,6 @@ public Flow.Publisher onMessageSendStream( } } - String currentTaskId = taskId.get(); - if (pushSender != null && currentTaskId != null) { - EventKind latest = resultAggregator.getCurrentResult(); - if (latest instanceof Task latestTask) { - pushSender.sendNotification(latestTask); - } - } - return true; })); @@ -881,14 +864,23 @@ private CompletableFuture cleanupProducer(@Nullable CompletableFuture consumeAndEmit(EventConsumer consumer) { Flow.Publisher allItems = consumer.consumeAll(); - // Process items conditionally - only save non-replicated events to database + // Just stream events - no persistence needed + // TaskStore update moved to MainEventBusProcessor return processor(createTubeConfig(), allItems, (errorConsumer, item) -> { - // Only process non-replicated events to avoid duplicate database writes - if (!item.isReplicated()) { - try { - callTaskManagerProcess(item.getEvent()); - } catch (A2AServerException e) { - errorConsumer.accept(e); - return false; - } - } - // Continue processing and emit (both replicated and non-replicated) + // Continue processing and emit all events return true; }); } @@ -81,15 +72,7 @@ public EventKind consumeAll(EventConsumer consumer) throws A2AError { return false; } } - // Only process non-replicated events to avoid duplicate database writes - if (!item.isReplicated()) { - try { - callTaskManagerProcess(event); - } catch (A2AServerException e) { - error.set(e); - return false; - } - } + // TaskStore update moved to MainEventBusProcessor return true; }, error::set); @@ -146,16 +129,7 @@ public EventTypeAndInterrupt consumeAndBreakOnInterrupt(EventConsumer consumer, return false; } - // Process event through TaskManager - only for non-replicated events - if (!item.isReplicated()) { - try { - callTaskManagerProcess(event); - } catch (A2AServerException e) { - errorRef.set(e); - completionFuture.completeExceptionally(e); - return false; - } - } + // TaskStore update moved to MainEventBusProcessor // Determine interrupt behavior boolean shouldInterrupt = false; @@ -185,7 +159,6 @@ 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()) { @@ -201,8 +174,7 @@ else if (blocking) { // 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. + // This ensures all events are fully processed before cleanup. // // For non-blocking and auth-required calls, complete immediately to allow // cleanup to proceed while consumption continues in background. @@ -279,10 +251,6 @@ else if (blocking) { consumptionCompletionFuture); } - private void callTaskManagerProcess(Event event) throws A2AServerException { - taskManager.process(event); - } - private String taskIdForLogging() { Task task = taskManager.getTask(); return task != null ? task.id() : "unknown"; diff --git a/server-common/src/test/java/io/a2a/server/events/EventQueueTest.java b/server-common/src/test/java/io/a2a/server/events/EventQueueTest.java index a3dc7d916..c0e22b8d0 100644 --- a/server-common/src/test/java/io/a2a/server/events/EventQueueTest.java +++ b/server-common/src/test/java/io/a2a/server/events/EventQueueTest.java @@ -12,6 +12,8 @@ import java.util.List; +import io.a2a.server.tasks.InMemoryTaskStore; +import io.a2a.server.tasks.PushNotificationSender; import io.a2a.spec.A2AError; import io.a2a.spec.Artifact; import io.a2a.spec.Event; @@ -23,12 +25,15 @@ import io.a2a.spec.TaskStatus; import io.a2a.spec.TaskStatusUpdateEvent; import io.a2a.spec.TextPart; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; public class EventQueueTest { private EventQueue eventQueue; + private MainEventBus mainEventBus; + private MainEventBusProcessor mainEventBusProcessor; private static final String MINIMAL_TASK = """ { @@ -46,11 +51,37 @@ public class EventQueueTest { } """; + private static final PushNotificationSender NOOP_PUSHNOTIFICATION_SENDER = task -> {}; @BeforeEach public void init() { - eventQueue = EventQueue.builder().build(); + // Set up MainEventBus and processor for production-like test environment + InMemoryTaskStore taskStore = new InMemoryTaskStore(); + mainEventBus = new MainEventBus(); + mainEventBusProcessor = new MainEventBusProcessor(mainEventBus, taskStore, NOOP_PUSHNOTIFICATION_SENDER); + EventQueueUtil.start(mainEventBusProcessor); + + eventQueue = EventQueue.builder() + .taskId("test-task") + .mainEventBus(mainEventBus) + .build(); + } + + @AfterEach + public void cleanup() { + if (mainEventBusProcessor != null) { + EventQueueUtil.stop(mainEventBusProcessor); + } + } + /** + * Helper to create a queue with MainEventBus configured (for tests that need event distribution). + */ + private EventQueue createQueueWithEventBus(String taskId) { + return EventQueue.builder() + .taskId(taskId) + .mainEventBus(mainEventBus) + .build(); } @Test @@ -95,15 +126,16 @@ public void testTapOnChildQueueThrowsException() { @Test public void testEnqueueEventPropagagesToChildren() throws Exception { - EventQueue parentQueue = EventQueue.builder().build(); + EventQueue parentQueue = createQueueWithEventBus("test-propagate"); EventQueue childQueue = parentQueue.tap(); Event event = fromJson(MINIMAL_TASK, Task.class); parentQueue.enqueueEvent(event); // Event should be available in both parent and child queues - Event parentEvent = parentQueue.dequeueEventItem(-1).getEvent(); - Event childEvent = childQueue.dequeueEventItem(-1).getEvent(); + // Note: MainEventBusProcessor runs async, so we use dequeueEventItem with timeout + Event parentEvent = parentQueue.dequeueEventItem(5000).getEvent(); + Event childEvent = childQueue.dequeueEventItem(5000).getEvent(); assertSame(event, parentEvent); assertSame(event, childEvent); @@ -111,7 +143,7 @@ public void testEnqueueEventPropagagesToChildren() throws Exception { @Test public void testMultipleChildQueuesReceiveEvents() throws Exception { - EventQueue parentQueue = EventQueue.builder().build(); + EventQueue parentQueue = createQueueWithEventBus("test-multiple"); EventQueue childQueue1 = parentQueue.tap(); EventQueue childQueue2 = parentQueue.tap(); @@ -122,42 +154,43 @@ public void testMultipleChildQueuesReceiveEvents() throws Exception { parentQueue.enqueueEvent(event2); // All queues should receive both events - assertSame(event1, parentQueue.dequeueEventItem(-1).getEvent()); - assertSame(event2, parentQueue.dequeueEventItem(-1).getEvent()); + // Note: Use timeout for async processing + assertSame(event1, parentQueue.dequeueEventItem(5000).getEvent()); + assertSame(event2, parentQueue.dequeueEventItem(5000).getEvent()); - assertSame(event1, childQueue1.dequeueEventItem(-1).getEvent()); - assertSame(event2, childQueue1.dequeueEventItem(-1).getEvent()); + assertSame(event1, childQueue1.dequeueEventItem(5000).getEvent()); + assertSame(event2, childQueue1.dequeueEventItem(5000).getEvent()); - assertSame(event1, childQueue2.dequeueEventItem(-1).getEvent()); - assertSame(event2, childQueue2.dequeueEventItem(-1).getEvent()); + assertSame(event1, childQueue2.dequeueEventItem(5000).getEvent()); + assertSame(event2, childQueue2.dequeueEventItem(5000).getEvent()); } @Test public void testChildQueueDequeueIndependently() throws Exception { - EventQueue parentQueue = EventQueue.builder().build(); + EventQueue parentQueue = createQueueWithEventBus("test-independent"); EventQueue childQueue1 = parentQueue.tap(); EventQueue childQueue2 = parentQueue.tap(); Event event = fromJson(MINIMAL_TASK, Task.class); parentQueue.enqueueEvent(event); - // Dequeue from child1 first - Event child1Event = childQueue1.dequeueEventItem(-1).getEvent(); + // Dequeue from child1 first (use timeout for async processing) + Event child1Event = childQueue1.dequeueEventItem(5000).getEvent(); assertSame(event, child1Event); // child2 should still have the event available - Event child2Event = childQueue2.dequeueEventItem(-1).getEvent(); + Event child2Event = childQueue2.dequeueEventItem(5000).getEvent(); assertSame(event, child2Event); // Parent should still have the event available - Event parentEvent = parentQueue.dequeueEventItem(-1).getEvent(); + Event parentEvent = parentQueue.dequeueEventItem(5000).getEvent(); assertSame(event, parentEvent); } @Test public void testCloseImmediatePropagationToChildren() throws Exception { - EventQueue parentQueue = EventQueue.builder().build(); + EventQueue parentQueue = createQueueWithEventBus("test-close"); EventQueue childQueue = parentQueue.tap(); // Add events to both parent and child @@ -166,7 +199,7 @@ public void testCloseImmediatePropagationToChildren() throws Exception { assertFalse(childQueue.isClosed()); try { - assertNotNull(childQueue.dequeueEventItem(-1)); // Child has the event + assertNotNull(childQueue.dequeueEventItem(5000)); // Child has the event (use timeout) } catch (EventQueueClosedException e) { // This is fine if queue closed before dequeue } diff --git a/server-common/src/test/java/io/a2a/server/events/EventQueueUtil.java b/server-common/src/test/java/io/a2a/server/events/EventQueueUtil.java index 39201c1f6..8490aa00c 100644 --- a/server-common/src/test/java/io/a2a/server/events/EventQueueUtil.java +++ b/server-common/src/test/java/io/a2a/server/events/EventQueueUtil.java @@ -5,4 +5,12 @@ public class EventQueueUtil { public static EventQueue.EventQueueBuilder getEventQueueBuilder() { return EventQueue.builder(); } + + public static void start(MainEventBusProcessor processor) { + processor.start(); + } + + public static void stop(MainEventBusProcessor processor) { + processor.stop(); + } } diff --git a/server-common/src/test/java/io/a2a/server/events/InMemoryQueueManagerTest.java b/server-common/src/test/java/io/a2a/server/events/InMemoryQueueManagerTest.java index 1eca1b739..1d10be2c8 100644 --- a/server-common/src/test/java/io/a2a/server/events/InMemoryQueueManagerTest.java +++ b/server-common/src/test/java/io/a2a/server/events/InMemoryQueueManagerTest.java @@ -14,7 +14,10 @@ import java.util.concurrent.ExecutionException; import java.util.stream.IntStream; +import io.a2a.server.tasks.InMemoryTaskStore; import io.a2a.server.tasks.MockTaskStateProvider; +import io.a2a.server.tasks.PushNotificationSender; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -22,11 +25,25 @@ public class InMemoryQueueManagerTest { private InMemoryQueueManager queueManager; private MockTaskStateProvider taskStateProvider; + private InMemoryTaskStore taskStore; + private MainEventBus mainEventBus; + private MainEventBusProcessor mainEventBusProcessor; + private static final PushNotificationSender NOOP_PUSHNOTIFICATION_SENDER = task -> {}; @BeforeEach public void setUp() { taskStateProvider = new MockTaskStateProvider(); - queueManager = new InMemoryQueueManager(taskStateProvider); + taskStore = new InMemoryTaskStore(); + mainEventBus = new MainEventBus(); + mainEventBusProcessor = new MainEventBusProcessor(mainEventBus, taskStore, NOOP_PUSHNOTIFICATION_SENDER); + EventQueueUtil.start(mainEventBusProcessor); + + queueManager = new InMemoryQueueManager(taskStateProvider, mainEventBus); + } + + @AfterEach + public void tearDown() { + EventQueueUtil.stop(mainEventBusProcessor); } @Test diff --git a/server-common/src/test/java/io/a2a/server/requesthandlers/AbstractA2ARequestHandlerTest.java b/server-common/src/test/java/io/a2a/server/requesthandlers/AbstractA2ARequestHandlerTest.java index 8038bc147..406275e5d 100644 --- a/server-common/src/test/java/io/a2a/server/requesthandlers/AbstractA2ARequestHandlerTest.java +++ b/server-common/src/test/java/io/a2a/server/requesthandlers/AbstractA2ARequestHandlerTest.java @@ -26,7 +26,10 @@ import io.a2a.server.agentexecution.RequestContext; import io.a2a.server.events.EventQueue; import io.a2a.server.events.EventQueueItem; +import io.a2a.server.events.EventQueueUtil; import io.a2a.server.events.InMemoryQueueManager; +import io.a2a.server.events.MainEventBus; +import io.a2a.server.events.MainEventBusProcessor; import io.a2a.server.tasks.BasePushNotificationSender; import io.a2a.server.tasks.InMemoryPushNotificationConfigStore; import io.a2a.server.tasks.InMemoryTaskStore; @@ -66,6 +69,8 @@ public class AbstractA2ARequestHandlerTest { private static final String PREFERRED_TRANSPORT = "preferred-transport"; private static final String A2A_REQUESTHANDLER_TEST_PROPERTIES = "/a2a-requesthandler-test.properties"; + private static final PushNotificationSender NOOP_PUSHNOTIFICATION_SENDER = task -> {}; + protected AgentExecutor executor; protected TaskStore taskStore; protected RequestHandler requestHandler; @@ -73,6 +78,8 @@ public class AbstractA2ARequestHandlerTest { protected AgentExecutorMethod agentExecutorCancel; protected InMemoryQueueManager queueManager; protected TestHttpClient httpClient; + protected MainEventBus mainEventBus; + protected MainEventBusProcessor mainEventBusProcessor; protected final Executor internalExecutor = Executors.newCachedThreadPool(); @@ -96,19 +103,32 @@ public void cancel(RequestContext context, EventQueue eventQueue) throws A2AErro InMemoryTaskStore inMemoryTaskStore = new InMemoryTaskStore(); taskStore = inMemoryTaskStore; - queueManager = new InMemoryQueueManager(inMemoryTaskStore); + + // Create push notification components BEFORE MainEventBusProcessor httpClient = new TestHttpClient(); PushNotificationConfigStore pushConfigStore = new InMemoryPushNotificationConfigStore(); PushNotificationSender pushSender = new BasePushNotificationSender(pushConfigStore, httpClient); + // Create MainEventBus and MainEventBusProcessor (production code path) + mainEventBus = new MainEventBus(); + mainEventBusProcessor = new MainEventBusProcessor(mainEventBus, taskStore, pushSender); + EventQueueUtil.start(mainEventBusProcessor); + + queueManager = new InMemoryQueueManager(inMemoryTaskStore, mainEventBus); + requestHandler = DefaultRequestHandler.create( - executor, taskStore, queueManager, pushConfigStore, pushSender, internalExecutor); + executor, taskStore, queueManager, pushConfigStore, internalExecutor); } @AfterEach public void cleanup() { agentExecutorExecute = null; agentExecutorCancel = null; + + // Stop MainEventBusProcessor background thread + if (mainEventBusProcessor != null) { + EventQueueUtil.stop(mainEventBusProcessor); + } } protected static AgentCard createAgentCard(boolean streaming, boolean pushNotifications, boolean stateTransitionHistory) { diff --git a/server-common/src/test/java/io/a2a/server/requesthandlers/DefaultRequestHandlerTest.java b/server-common/src/test/java/io/a2a/server/requesthandlers/DefaultRequestHandlerTest.java index 293babe4e..662ddcab4 100644 --- a/server-common/src/test/java/io/a2a/server/requesthandlers/DefaultRequestHandlerTest.java +++ b/server-common/src/test/java/io/a2a/server/requesthandlers/DefaultRequestHandlerTest.java @@ -17,9 +17,13 @@ import io.a2a.server.agentexecution.RequestContext; import io.a2a.server.auth.UnauthenticatedUser; import io.a2a.server.events.EventQueue; +import io.a2a.server.events.EventQueueUtil; import io.a2a.server.events.InMemoryQueueManager; +import io.a2a.server.events.MainEventBus; +import io.a2a.server.events.MainEventBusProcessor; import io.a2a.server.tasks.InMemoryPushNotificationConfigStore; import io.a2a.server.tasks.InMemoryTaskStore; +import io.a2a.server.tasks.PushNotificationSender; import io.a2a.server.tasks.TaskUpdater; import io.a2a.spec.A2AError; import io.a2a.spec.ListTaskPushNotificationConfigParams; @@ -32,6 +36,7 @@ import io.a2a.spec.TaskState; import io.a2a.spec.TaskStatus; import io.a2a.spec.TextPart; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.Timeout; @@ -50,12 +55,23 @@ public class DefaultRequestHandlerTest { private InMemoryQueueManager queueManager; private TestAgentExecutor agentExecutor; private ServerCallContext serverCallContext; + private MainEventBus mainEventBus; + private MainEventBusProcessor mainEventBusProcessor; + + private static final PushNotificationSender NOOP_PUSHNOTIFICATION_SENDER = task -> {}; @BeforeEach void setUp() { taskStore = new InMemoryTaskStore(); + + // Create MainEventBus and MainEventBusProcessor (production code path) + mainEventBus = new MainEventBus(); + mainEventBusProcessor = new MainEventBusProcessor(mainEventBus, taskStore, NOOP_PUSHNOTIFICATION_SENDER); + EventQueueUtil.start(mainEventBusProcessor); + // Pass taskStore as TaskStateProvider to queueManager for task-aware queue management - queueManager = new InMemoryQueueManager(taskStore); + queueManager = new InMemoryQueueManager(taskStore, mainEventBus); + agentExecutor = new TestAgentExecutor(); requestHandler = DefaultRequestHandler.create( @@ -63,13 +79,20 @@ void setUp() { taskStore, queueManager, null, // pushConfigStore - null, // pushSender Executors.newCachedThreadPool() ); serverCallContext = new ServerCallContext(UnauthenticatedUser.INSTANCE, Map.of(), Set.of()); } + @AfterEach + void tearDown() { + // Stop MainEventBusProcessor background thread + if (mainEventBusProcessor != null) { + EventQueueUtil.stop(mainEventBusProcessor); + } + } + /** * Test that multiple blocking messages to the same task work correctly * when agent doesn't emit final events (fire-and-forget pattern). @@ -822,7 +845,6 @@ void testBlockingMessageStoresPushNotificationConfigForNewTask() throws Exceptio taskStore, queueManager, pushConfigStore, // Add push config store - null, // pushSender Executors.newCachedThreadPool() ); @@ -893,7 +915,6 @@ void testBlockingMessageStoresPushNotificationConfigForExistingTask() throws Exc taskStore, queueManager, pushConfigStore, // Add push config store - null, // pushSender Executors.newCachedThreadPool() ); diff --git a/server-common/src/test/java/io/a2a/server/tasks/ResultAggregatorTest.java b/server-common/src/test/java/io/a2a/server/tasks/ResultAggregatorTest.java index d64729077..ff801b70d 100644 --- a/server-common/src/test/java/io/a2a/server/tasks/ResultAggregatorTest.java +++ b/server-common/src/test/java/io/a2a/server/tasks/ResultAggregatorTest.java @@ -17,6 +17,7 @@ import io.a2a.server.events.EventConsumer; import io.a2a.server.events.EventQueue; import io.a2a.server.events.InMemoryQueueManager; +import io.a2a.server.events.MainEventBus; import io.a2a.spec.EventKind; import io.a2a.spec.Message; import io.a2a.spec.Task; @@ -203,8 +204,9 @@ void testConsumeAndBreakNonBlocking() throws Exception { when(mockTaskManager.getTask()).thenReturn(firstEvent); // Create an event queue using QueueManager (which has access to builder) + MainEventBus mainEventBus = new MainEventBus(); InMemoryQueueManager queueManager = - new InMemoryQueueManager(new MockTaskStateProvider()); + new InMemoryQueueManager(new MockTaskStateProvider(), mainEventBus); EventQueue queue = queueManager.getEventQueueBuilder("test-task").build(); queue.enqueueEvent(firstEvent); @@ -221,7 +223,8 @@ void testConsumeAndBreakNonBlocking() throws Exception { assertEquals(firstEvent, result.eventType()); assertTrue(result.interrupted()); - verify(mockTaskManager).process(firstEvent); + // NOTE: ResultAggregator no longer calls taskManager.process() + // That responsibility has moved to MainEventBusProcessor for centralized persistence // getTask() is called at least once for the return value (line 255) // May be called once more if debug logging executes in time (line 209) // The async consumer may or may not execute before verification, so we accept 1-2 calls diff --git a/transport/grpc/src/test/java/io/a2a/transport/grpc/handler/GrpcHandlerTest.java b/transport/grpc/src/test/java/io/a2a/transport/grpc/handler/GrpcHandlerTest.java index ad9b879f9..0229fc350 100644 --- a/transport/grpc/src/test/java/io/a2a/transport/grpc/handler/GrpcHandlerTest.java +++ b/transport/grpc/src/test/java/io/a2a/transport/grpc/handler/GrpcHandlerTest.java @@ -281,7 +281,7 @@ public void testPushNotificationsNotSupportedError() throws Exception { public void testOnGetPushNotificationNoPushNotifierConfig() throws Exception { // Create request handler without a push notifier DefaultRequestHandler requestHandler = - new DefaultRequestHandler(executor, taskStore, queueManager, null, null, internalExecutor); + new DefaultRequestHandler(executor, taskStore, queueManager,null, internalExecutor); AgentCard card = AbstractA2ARequestHandlerTest.createAgentCard(false, true, false); GrpcHandler handler = new TestGrpcHandler(card, requestHandler, internalExecutor); String NAME = "tasks/" + AbstractA2ARequestHandlerTest.MINIMAL_TASK.id() + "/pushNotificationConfigs/" + AbstractA2ARequestHandlerTest.MINIMAL_TASK.id(); @@ -293,7 +293,7 @@ public void testOnGetPushNotificationNoPushNotifierConfig() throws Exception { public void testOnSetPushNotificationNoPushNotifierConfig() throws Exception { // Create request handler without a push notifier DefaultRequestHandler requestHandler = DefaultRequestHandler.create( - executor, taskStore, queueManager, null, null, internalExecutor); + executor, taskStore, queueManager, null, internalExecutor); AgentCard card = AbstractA2ARequestHandlerTest.createAgentCard(false, true, false); GrpcHandler handler = new TestGrpcHandler(card, requestHandler, internalExecutor); String NAME = "tasks/" + AbstractA2ARequestHandlerTest.MINIMAL_TASK.id() + "/pushNotificationConfigs/" + AbstractA2ARequestHandlerTest.MINIMAL_TASK.id(); @@ -668,7 +668,7 @@ public void testListPushNotificationConfigNotSupported() throws Exception { @Test public void testListPushNotificationConfigNoPushConfigStore() { DefaultRequestHandler requestHandler = DefaultRequestHandler.create( - executor, taskStore, queueManager, null, null, internalExecutor); + executor, taskStore, queueManager, null, internalExecutor); GrpcHandler handler = new TestGrpcHandler(AbstractA2ARequestHandlerTest.CARD, requestHandler, internalExecutor); taskStore.save(AbstractA2ARequestHandlerTest.MINIMAL_TASK); agentExecutorExecute = (context, eventQueue) -> { @@ -741,7 +741,7 @@ public void testDeletePushNotificationConfigNotSupported() throws Exception { @Test public void testDeletePushNotificationConfigNoPushConfigStore() { DefaultRequestHandler requestHandler = DefaultRequestHandler.create( - executor, taskStore, queueManager, null, null, internalExecutor); + executor, taskStore, queueManager, null, internalExecutor); GrpcHandler handler = new TestGrpcHandler(AbstractA2ARequestHandlerTest.CARD, requestHandler, internalExecutor); String NAME = "tasks/" + AbstractA2ARequestHandlerTest.MINIMAL_TASK.id() + "/pushNotificationConfigs/" + AbstractA2ARequestHandlerTest.MINIMAL_TASK.id(); DeleteTaskPushNotificationConfigRequest request = DeleteTaskPushNotificationConfigRequest.newBuilder() diff --git a/transport/jsonrpc/src/test/java/io/a2a/transport/jsonrpc/handler/JSONRPCHandlerTest.java b/transport/jsonrpc/src/test/java/io/a2a/transport/jsonrpc/handler/JSONRPCHandlerTest.java index c8e69f43f..3b979666c 100644 --- a/transport/jsonrpc/src/test/java/io/a2a/transport/jsonrpc/handler/JSONRPCHandlerTest.java +++ b/transport/jsonrpc/src/test/java/io/a2a/transport/jsonrpc/handler/JSONRPCHandlerTest.java @@ -40,6 +40,8 @@ import io.a2a.server.ServerCallContext; import io.a2a.server.auth.UnauthenticatedUser; import io.a2a.server.events.EventConsumer; +import io.a2a.server.events.MainEventBusProcessor; +import io.a2a.server.events.MainEventBusProcessorCallback; import io.a2a.server.requesthandlers.AbstractA2ARequestHandlerTest; import io.a2a.server.requesthandlers.DefaultRequestHandler; import io.a2a.server.tasks.ResultAggregator; @@ -49,11 +51,10 @@ import io.a2a.spec.AgentExtension; import io.a2a.spec.AgentInterface; import io.a2a.spec.Artifact; -import io.a2a.spec.ExtendedCardNotConfiguredError; -import io.a2a.spec.ExtensionSupportRequiredError; -import io.a2a.spec.VersionNotSupportedError; import io.a2a.spec.DeleteTaskPushNotificationConfigParams; import io.a2a.spec.Event; +import io.a2a.spec.ExtendedCardNotConfiguredError; +import io.a2a.spec.ExtensionSupportRequiredError; import io.a2a.spec.GetTaskPushNotificationConfigParams; import io.a2a.spec.InternalError; import io.a2a.spec.InvalidRequestError; @@ -74,6 +75,7 @@ import io.a2a.spec.TaskStatusUpdateEvent; import io.a2a.spec.TextPart; import io.a2a.spec.UnsupportedOperationError; +import io.a2a.spec.VersionNotSupportedError; import mutiny.zero.ZeroPublisher; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Disabled; @@ -170,38 +172,9 @@ public void testOnMessageNewMessageSuccess() { SendMessageRequest request = new SendMessageRequest("1", new MessageSendParams(message, null, null)); SendMessageResponse response = handler.onMessageSend(request, callContext); assertNull(response.getError()); - // The Python implementation returns a Task here, but then again they are using hardcoded mocks and - // bypassing the whole EventQueue. - // If we were to send a Task in agentExecutorExecute EventConsumer.consumeAll() would not exit due to - // the Task not having a 'final' state - // - // See testOnMessageNewMessageSuccessMocks() for a test more similar to the Python implementation Assertions.assertSame(message, response.getResult()); } - @Test - public void testOnMessageNewMessageSuccessMocks() { - JSONRPCHandler handler = new JSONRPCHandler(CARD, requestHandler, internalExecutor); - - Message message = Message.builder(MESSAGE) - .taskId(MINIMAL_TASK.id()) - .contextId(MINIMAL_TASK.contextId()) - .build(); - - SendMessageRequest request = new SendMessageRequest("1", new MessageSendParams(message, null, null)); - SendMessageResponse response; - try (MockedConstruction mocked = Mockito.mockConstruction( - EventConsumer.class, - (mock, context) -> { - Mockito.doReturn(ZeroPublisher.fromItems(wrapEvent(MINIMAL_TASK))).when(mock).consumeAll(); - Mockito.doCallRealMethod().when(mock).createAgentRunnableDoneCallback(); - })) { - response = handler.onMessageSend(request, callContext); - } - assertNull(response.getError()); - Assertions.assertSame(MINIMAL_TASK, response.getResult()); - } - @Test public void testOnMessageNewMessageWithExistingTaskSuccess() { JSONRPCHandler handler = new JSONRPCHandler(CARD, requestHandler, internalExecutor); @@ -216,38 +189,9 @@ public void testOnMessageNewMessageWithExistingTaskSuccess() { SendMessageRequest request = new SendMessageRequest("1", new MessageSendParams(message, null, null)); SendMessageResponse response = handler.onMessageSend(request, callContext); assertNull(response.getError()); - // The Python implementation returns a Task here, but then again they are using hardcoded mocks and - // bypassing the whole EventQueue. - // If we were to send a Task in agentExecutorExecute EventConsumer.consumeAll() would not exit due to - // the Task not having a 'final' state - // - // See testOnMessageNewMessageWithExistingTaskSuccessMocks() for a test more similar to the Python implementation Assertions.assertSame(message, response.getResult()); } - @Test - public void testOnMessageNewMessageWithExistingTaskSuccessMocks() { - JSONRPCHandler handler = new JSONRPCHandler(CARD, requestHandler, internalExecutor); - taskStore.save(MINIMAL_TASK); - - Message message = Message.builder(MESSAGE) - .taskId(MINIMAL_TASK.id()) - .contextId(MINIMAL_TASK.contextId()) - .build(); - SendMessageRequest request = new SendMessageRequest("1", new MessageSendParams(message, null, null)); - SendMessageResponse response; - try (MockedConstruction mocked = Mockito.mockConstruction( - EventConsumer.class, - (mock, context) -> { - Mockito.doReturn(ZeroPublisher.fromItems(wrapEvent(MINIMAL_TASK))).when(mock).consumeAll(); - })) { - response = handler.onMessageSend(request, callContext); - } - assertNull(response.getError()); - Assertions.assertSame(MINIMAL_TASK, response.getResult()); - - } - @Test public void testOnMessageError() { // See testMessageOnErrorMocks() for a test more similar to the Python implementation, using mocks for @@ -348,9 +292,24 @@ public void onComplete() { @Test public void testOnMessageStreamNewMessageMultipleEventsSuccess() throws InterruptedException { - JSONRPCHandler handler = new JSONRPCHandler(CARD, requestHandler, internalExecutor); + // Setup callback to wait for all 3 events to be processed by MainEventBusProcessor + CountDownLatch processingLatch = new CountDownLatch(3); + MainEventBusProcessor.setCallback(new MainEventBusProcessorCallback() { + @Override + public void onEventProcessed(String taskId, Event event) { + processingLatch.countDown(); + } + + @Override + public void onTaskFinalized(String taskId) { + // Not needed for this test + } + }); + + try { + JSONRPCHandler handler = new JSONRPCHandler(CARD, requestHandler, internalExecutor); - // Create multiple events to be sent during streaming + // Create multiple events to be sent during streaming Task taskEvent = Task.builder(MINIMAL_TASK) .status(new TaskStatus(TaskState.WORKING)) .build(); @@ -425,33 +384,40 @@ public void onComplete() { } }); - // Wait for all events to be received - Assertions.assertTrue(latch.await(2, TimeUnit.SECONDS), - "Expected to receive 3 events within timeout"); - - // Assert no error occurred during streaming - Assertions.assertNull(error.get(), "No error should occur during streaming"); - - // Verify that all 3 events were received - assertEquals(3, results.size(), "Should have received exactly 3 events"); - - // Verify the first event is the task - Task receivedTask = assertInstanceOf(Task.class, results.get(0), "First event should be a Task"); - assertEquals(MINIMAL_TASK.id(), receivedTask.id()); - assertEquals(MINIMAL_TASK.contextId(), receivedTask.contextId()); - assertEquals(TaskState.WORKING, receivedTask.status().state()); - - // Verify the second event is the artifact update - TaskArtifactUpdateEvent receivedArtifact = assertInstanceOf(TaskArtifactUpdateEvent.class, results.get(1), - "Second event should be a TaskArtifactUpdateEvent"); - assertEquals(MINIMAL_TASK.id(), receivedArtifact.taskId()); - assertEquals("artifact-1", receivedArtifact.artifact().artifactId()); - - // Verify the third event is the status update - TaskStatusUpdateEvent receivedStatus = assertInstanceOf(TaskStatusUpdateEvent.class, results.get(2), - "Third event should be a TaskStatusUpdateEvent"); - assertEquals(MINIMAL_TASK.id(), receivedStatus.taskId()); - assertEquals(TaskState.COMPLETED, receivedStatus.status().state()); + // Wait for all events to be received (increased timeout for async processing) + Assertions.assertTrue(latch.await(10, TimeUnit.SECONDS), + "Expected to receive 3 events within timeout"); + + // Wait for MainEventBusProcessor to complete processing all 3 events + Assertions.assertTrue(processingLatch.await(5, TimeUnit.SECONDS), + "MainEventBusProcessor should have processed all 3 events"); + + // Assert no error occurred during streaming + Assertions.assertNull(error.get(), "No error should occur during streaming"); + + // Verify that all 3 events were received + assertEquals(3, results.size(), "Should have received exactly 3 events"); + + // Verify the first event is the task + Task receivedTask = assertInstanceOf(Task.class, results.get(0), "First event should be a Task"); + assertEquals(MINIMAL_TASK.id(), receivedTask.id()); + assertEquals(MINIMAL_TASK.contextId(), receivedTask.contextId()); + assertEquals(TaskState.WORKING, receivedTask.status().state()); + + // Verify the second event is the artifact update + TaskArtifactUpdateEvent receivedArtifact = assertInstanceOf(TaskArtifactUpdateEvent.class, results.get(1), + "Second event should be a TaskArtifactUpdateEvent"); + assertEquals(MINIMAL_TASK.id(), receivedArtifact.taskId()); + assertEquals("artifact-1", receivedArtifact.artifact().artifactId()); + + // Verify the third event is the status update + TaskStatusUpdateEvent receivedStatus = assertInstanceOf(TaskStatusUpdateEvent.class, results.get(2), + "Third event should be a TaskStatusUpdateEvent"); + assertEquals(MINIMAL_TASK.id(), receivedStatus.taskId()); + assertEquals(TaskState.COMPLETED, receivedStatus.status().state()); + } finally { + MainEventBusProcessor.setCallback(null); + } } @Test @@ -725,106 +691,130 @@ public void testGetPushNotificationConfigSuccess() { @Test public void testOnMessageStreamNewMessageSendPushNotificationSuccess() throws Exception { - JSONRPCHandler handler = new JSONRPCHandler(CARD, requestHandler, internalExecutor); - taskStore.save(MINIMAL_TASK); - - List events = List.of( - MINIMAL_TASK, - TaskArtifactUpdateEvent.builder() - .taskId(MINIMAL_TASK.id()) - .contextId(MINIMAL_TASK.contextId()) - .artifact(Artifact.builder() - .artifactId("11") - .parts(new TextPart("text")) - .build()) - .build(), - TaskStatusUpdateEvent.builder() - .taskId(MINIMAL_TASK.id()) - .contextId(MINIMAL_TASK.contextId()) - .status(new TaskStatus(TaskState.COMPLETED)) - .build()); - - agentExecutorExecute = (context, eventQueue) -> { - // Hardcode the events to send here - for (Event event : events) { - eventQueue.enqueueEvent(event); + // Setup callback to wait for all 3 events to be processed by MainEventBusProcessor + CountDownLatch processingLatch = new CountDownLatch(3); + MainEventBusProcessor.setCallback(new MainEventBusProcessorCallback() { + @Override + public void onEventProcessed(String taskId, Event event) { + processingLatch.countDown(); } - }; - - TaskPushNotificationConfig config = new TaskPushNotificationConfig( - MINIMAL_TASK.id(), - PushNotificationConfig.builder().id("c295ea44-7543-4f78-b524-7a38915ad6e4").url("http://example.com").build(), "tenant"); - - SetTaskPushNotificationConfigRequest stpnRequest = new SetTaskPushNotificationConfigRequest("1", config); - SetTaskPushNotificationConfigResponse stpnResponse = handler.setPushNotificationConfig(stpnRequest, callContext); - assertNull(stpnResponse.getError()); - - Message msg = Message.builder(MESSAGE) - .taskId(MINIMAL_TASK.id()) - .build(); - SendStreamingMessageRequest request = new SendStreamingMessageRequest("1", new MessageSendParams(msg, null, null)); - Flow.Publisher response = handler.onMessageSendStream(request, callContext); - - final List results = Collections.synchronizedList(new ArrayList<>()); - final AtomicReference subscriptionRef = new AtomicReference<>(); - final CountDownLatch latch = new CountDownLatch(6); - httpClient.latch = latch; - - Executors.newSingleThreadExecutor().execute(() -> { - response.subscribe(new Flow.Subscriber<>() { - @Override - public void onSubscribe(Flow.Subscription subscription) { - subscriptionRef.set(subscription); - subscription.request(1); - } - - @Override - public void onNext(SendStreamingMessageResponse item) { - System.out.println("-> " + item.getResult()); - results.add(item.getResult()); - System.out.println(results); - subscriptionRef.get().request(1); - latch.countDown(); - } - @Override - public void onError(Throwable throwable) { - subscriptionRef.get().cancel(); - } + @Override + public void onTaskFinalized(String taskId) { + // Not needed for this test + } + }); - @Override - public void onComplete() { - subscriptionRef.get().cancel(); + try { + JSONRPCHandler handler = new JSONRPCHandler(CARD, requestHandler, internalExecutor); + taskStore.save(MINIMAL_TASK); + + List events = List.of( + MINIMAL_TASK, + TaskArtifactUpdateEvent.builder() + .taskId(MINIMAL_TASK.id()) + .contextId(MINIMAL_TASK.contextId()) + .artifact(Artifact.builder() + .artifactId("11") + .parts(new TextPart("text")) + .build()) + .build(), + TaskStatusUpdateEvent.builder() + .taskId(MINIMAL_TASK.id()) + .contextId(MINIMAL_TASK.contextId()) + .status(new TaskStatus(TaskState.COMPLETED)) + .build()); + + + agentExecutorExecute = (context, eventQueue) -> { + // Hardcode the events to send here + for (Event event : events) { + eventQueue.enqueueEvent(event); } + }; + + TaskPushNotificationConfig config = new TaskPushNotificationConfig( + MINIMAL_TASK.id(), + PushNotificationConfig.builder().id("c295ea44-7543-4f78-b524-7a38915ad6e4").url("http://example.com").build(), "tenant"); + + SetTaskPushNotificationConfigRequest stpnRequest = new SetTaskPushNotificationConfigRequest("1", config); + SetTaskPushNotificationConfigResponse stpnResponse = handler.setPushNotificationConfig(stpnRequest, callContext); + assertNull(stpnResponse.getError()); + + Message msg = Message.builder(MESSAGE) + .taskId(MINIMAL_TASK.id()) + .build(); + SendStreamingMessageRequest request = new SendStreamingMessageRequest("1", new MessageSendParams(msg, null, null)); + Flow.Publisher response = handler.onMessageSendStream(request, callContext); + + final List results = Collections.synchronizedList(new ArrayList<>()); + final AtomicReference subscriptionRef = new AtomicReference<>(); + final CountDownLatch latch = new CountDownLatch(6); + httpClient.latch = latch; + + Executors.newSingleThreadExecutor().execute(() -> { + response.subscribe(new Flow.Subscriber<>() { + @Override + public void onSubscribe(Flow.Subscription subscription) { + subscriptionRef.set(subscription); + subscription.request(1); + } + + @Override + public void onNext(SendStreamingMessageResponse item) { + System.out.println("-> " + item.getResult()); + results.add(item.getResult()); + System.out.println(results); + subscriptionRef.get().request(1); + latch.countDown(); + } + + @Override + public void onError(Throwable throwable) { + subscriptionRef.get().cancel(); + } + + @Override + public void onComplete() { + subscriptionRef.get().cancel(); + } + }); }); - }); - Assertions.assertTrue(latch.await(5, TimeUnit.SECONDS)); - subscriptionRef.get().cancel(); - assertEquals(3, results.size()); - assertEquals(3, httpClient.tasks.size()); - - Task curr = httpClient.tasks.get(0); - assertEquals(MINIMAL_TASK.id(), curr.id()); - assertEquals(MINIMAL_TASK.contextId(), curr.contextId()); - assertEquals(MINIMAL_TASK.status().state(), curr.status().state()); - assertEquals(0, curr.artifacts() == null ? 0 : curr.artifacts().size()); - - curr = httpClient.tasks.get(1); - assertEquals(MINIMAL_TASK.id(), curr.id()); - assertEquals(MINIMAL_TASK.contextId(), curr.contextId()); - assertEquals(MINIMAL_TASK.status().state(), curr.status().state()); - assertEquals(1, curr.artifacts().size()); - assertEquals(1, curr.artifacts().get(0).parts().size()); - assertEquals("text", ((TextPart) curr.artifacts().get(0).parts().get(0)).text()); - - curr = httpClient.tasks.get(2); - assertEquals(MINIMAL_TASK.id(), curr.id()); - assertEquals(MINIMAL_TASK.contextId(), curr.contextId()); - assertEquals(TaskState.COMPLETED, curr.status().state()); - assertEquals(1, curr.artifacts().size()); - assertEquals(1, curr.artifacts().get(0).parts().size()); - assertEquals("text", ((TextPart) curr.artifacts().get(0).parts().get(0)).text()); + Assertions.assertTrue(latch.await(5, TimeUnit.SECONDS)); + + // Wait for MainEventBusProcessor to complete processing all 3 events + Assertions.assertTrue(processingLatch.await(5, TimeUnit.SECONDS), + "MainEventBusProcessor should have processed all 3 events"); + + subscriptionRef.get().cancel(); + assertEquals(3, results.size()); + assertEquals(3, httpClient.tasks.size()); + + Task curr = httpClient.tasks.get(0); + assertEquals(MINIMAL_TASK.id(), curr.id()); + assertEquals(MINIMAL_TASK.contextId(), curr.contextId()); + assertEquals(MINIMAL_TASK.status().state(), curr.status().state()); + assertEquals(0, curr.artifacts() == null ? 0 : curr.artifacts().size()); + + curr = httpClient.tasks.get(1); + assertEquals(MINIMAL_TASK.id(), curr.id()); + assertEquals(MINIMAL_TASK.contextId(), curr.contextId()); + assertEquals(MINIMAL_TASK.status().state(), curr.status().state()); + assertEquals(1, curr.artifacts().size()); + assertEquals(1, curr.artifacts().get(0).parts().size()); + assertEquals("text", ((TextPart) curr.artifacts().get(0).parts().get(0)).text()); + + curr = httpClient.tasks.get(2); + assertEquals(MINIMAL_TASK.id(), curr.id()); + assertEquals(MINIMAL_TASK.contextId(), curr.contextId()); + assertEquals(TaskState.COMPLETED, curr.status().state()); + assertEquals(1, curr.artifacts().size()); + assertEquals(1, curr.artifacts().get(0).parts().size()); + assertEquals("text", ((TextPart) curr.artifacts().get(0).parts().get(0)).text()); + } finally { + MainEventBusProcessor.setCallback(null); + } } @Test @@ -1132,7 +1122,7 @@ public void testPushNotificationsNotSupportedError() { public void testOnGetPushNotificationNoPushNotifierConfig() { // Create request handler without a push notifier DefaultRequestHandler requestHandler = DefaultRequestHandler.create( - executor, taskStore, queueManager, null, null, internalExecutor); + executor, taskStore, queueManager, null, internalExecutor); AgentCard card = createAgentCard(false, true, false); JSONRPCHandler handler = new JSONRPCHandler(card, requestHandler, internalExecutor); @@ -1151,7 +1141,7 @@ public void testOnGetPushNotificationNoPushNotifierConfig() { public void testOnSetPushNotificationNoPushNotifierConfig() { // Create request handler without a push notifier DefaultRequestHandler requestHandler = DefaultRequestHandler.create( - executor, taskStore, queueManager, null, null, internalExecutor); + executor, taskStore, queueManager, null, internalExecutor); AgentCard card = createAgentCard(false, true, false); JSONRPCHandler handler = new JSONRPCHandler(card, requestHandler, internalExecutor); @@ -1243,7 +1233,7 @@ public void testDefaultRequestHandlerWithCustomComponents() { @Test public void testOnMessageSendErrorHandling() { DefaultRequestHandler requestHandler = DefaultRequestHandler.create( - executor, taskStore, queueManager, null, null, internalExecutor); + executor, taskStore, queueManager, null, internalExecutor); AgentCard card = createAgentCard(false, true, false); JSONRPCHandler handler = new JSONRPCHandler(card, requestHandler, internalExecutor); @@ -1289,16 +1279,31 @@ public void testOnMessageSendTaskIdMismatch() { } @Test - public void testOnMessageStreamTaskIdMismatch() { - JSONRPCHandler handler = new JSONRPCHandler(CARD, requestHandler, internalExecutor); - taskStore.save(MINIMAL_TASK); + public void testOnMessageStreamTaskIdMismatch() throws InterruptedException { + // Setup callback to wait for the 1 event to be processed by MainEventBusProcessor + CountDownLatch processingLatch = new CountDownLatch(1); + MainEventBusProcessor.setCallback(new MainEventBusProcessorCallback() { + @Override + public void onEventProcessed(String taskId, Event event) { + processingLatch.countDown(); + } - agentExecutorExecute = ((context, eventQueue) -> { - eventQueue.enqueueEvent(MINIMAL_TASK); + @Override + public void onTaskFinalized(String taskId) { + // Not needed for this test + } }); - SendStreamingMessageRequest request = new SendStreamingMessageRequest("1", new MessageSendParams(MESSAGE, null, null)); - Flow.Publisher response = handler.onMessageSendStream(request, callContext); + try { + JSONRPCHandler handler = new JSONRPCHandler(CARD, requestHandler, internalExecutor); + taskStore.save(MINIMAL_TASK); + + agentExecutorExecute = ((context, eventQueue) -> { + eventQueue.enqueueEvent(MINIMAL_TASK); + }); + + SendStreamingMessageRequest request = new SendStreamingMessageRequest("1", new MessageSendParams(MESSAGE, null, null)); + Flow.Publisher response = handler.onMessageSendStream(request, callContext); CompletableFuture future = new CompletableFuture<>(); List results = new ArrayList<>(); @@ -1333,11 +1338,18 @@ public void onComplete() { } }); - future.join(); + future.join(); + + // Wait for MainEventBusProcessor to complete processing the event + Assertions.assertTrue(processingLatch.await(5, TimeUnit.SECONDS), + "MainEventBusProcessor should have processed the event"); - Assertions.assertNull(error.get()); - Assertions.assertEquals(1, results.size()); - Assertions.assertInstanceOf(InternalError.class, results.get(0).getError()); + Assertions.assertNull(error.get()); + Assertions.assertEquals(1, results.size()); + Assertions.assertInstanceOf(InternalError.class, results.get(0).getError()); + } finally { + MainEventBusProcessor.setCallback(null); + } } @Test @@ -1401,7 +1413,7 @@ public void testListPushNotificationConfigNotSupported() { @Test public void testListPushNotificationConfigNoPushConfigStore() { DefaultRequestHandler requestHandler = DefaultRequestHandler.create( - executor, taskStore, queueManager, null, null, internalExecutor); + executor, taskStore, queueManager, null, internalExecutor); JSONRPCHandler handler = new JSONRPCHandler(CARD, requestHandler, internalExecutor); taskStore.save(MINIMAL_TASK); agentExecutorExecute = (context, eventQueue) -> { @@ -1492,8 +1504,8 @@ public void testDeletePushNotificationConfigNotSupported() { @Test public void testDeletePushNotificationConfigNoPushConfigStore() { - DefaultRequestHandler requestHandler = DefaultRequestHandler.create( - executor, taskStore, queueManager, null, null, internalExecutor); + DefaultRequestHandler requestHandler = + DefaultRequestHandler.create(executor, taskStore, queueManager, null, internalExecutor); JSONRPCHandler handler = new JSONRPCHandler(CARD, requestHandler, internalExecutor); taskStore.save(MINIMAL_TASK); agentExecutorExecute = (context, eventQueue) -> { From 00739404f42e42c31db10df208deb322942d8ce9 Mon Sep 17 00:00:00 2001 From: Kabir Khan Date: Mon, 8 Dec 2025 11:09:30 +0000 Subject: [PATCH 2/8] Cleanup and ensure that we now only listen to ChildQueues. MainQueue is now only for sending to the central processor, and cannot be subscribed to --- .../core/ReplicatedQueueManagerTest.java | 8 +- .../io/a2a/server/events/EventQueueUtil.java | 5 +- .../java/io/a2a/server/events/EventQueue.java | 248 +++++++++--------- .../server/events/InMemoryQueueManager.java | 5 +- .../io/a2a/server/events/QueueManager.java | 8 +- .../DefaultRequestHandler.java | 2 +- .../a2a/server/events/EventConsumerTest.java | 30 ++- .../io/a2a/server/events/EventQueueTest.java | 138 +++++----- .../io/a2a/server/events/EventQueueUtil.java | 88 ++++++- .../events/InMemoryQueueManagerTest.java | 16 +- .../server/tasks/ResultAggregatorTest.java | 19 +- .../io/a2a/server/tasks/TaskUpdaterTest.java | 33 ++- 12 files changed, 390 insertions(+), 210 deletions(-) diff --git a/extras/queue-manager-replicated/core/src/test/java/io/a2a/extras/queuemanager/replicated/core/ReplicatedQueueManagerTest.java b/extras/queue-manager-replicated/core/src/test/java/io/a2a/extras/queuemanager/replicated/core/ReplicatedQueueManagerTest.java index 05b81cea0..ee1e4b88f 100644 --- a/extras/queue-manager-replicated/core/src/test/java/io/a2a/extras/queuemanager/replicated/core/ReplicatedQueueManagerTest.java +++ b/extras/queue-manager-replicated/core/src/test/java/io/a2a/extras/queuemanager/replicated/core/ReplicatedQueueManagerTest.java @@ -151,9 +151,11 @@ void testReplicatedEventCreatesQueueIfNeeded() throws InterruptedException { assertNotNull(queue, "Queue should be created when processing replicated event for non-existent task"); // Verify the event was enqueued by dequeuing it + // Need to tap() the MainQueue to get a ChildQueue for consumption + EventQueue childQueue = queue.tap(); Event dequeuedEvent; try { - dequeuedEvent = queue.dequeueEventItem(100).getEvent(); + dequeuedEvent = childQueue.dequeueEventItem(100).getEvent(); } catch (EventQueueClosedException e) { fail("Queue should not be closed"); return; @@ -351,9 +353,11 @@ void testReplicatedEventProcessedWhenTaskActive() throws InterruptedException { assertNotNull(queue, "Queue should be created for active task"); // Verify the event was enqueued + // Need to tap() the MainQueue to get a ChildQueue for consumption + EventQueue childQueue = queue.tap(); Event dequeuedEvent; try { - dequeuedEvent = queue.dequeueEventItem(100).getEvent(); + dequeuedEvent = childQueue.dequeueEventItem(100).getEvent(); } catch (EventQueueClosedException e) { fail("Queue should not be closed"); return; diff --git a/extras/queue-manager-replicated/core/src/test/java/io/a2a/server/events/EventQueueUtil.java b/extras/queue-manager-replicated/core/src/test/java/io/a2a/server/events/EventQueueUtil.java index 8490aa00c..e5131ad15 100644 --- a/extras/queue-manager-replicated/core/src/test/java/io/a2a/server/events/EventQueueUtil.java +++ b/extras/queue-manager-replicated/core/src/test/java/io/a2a/server/events/EventQueueUtil.java @@ -1,9 +1,12 @@ package io.a2a.server.events; public class EventQueueUtil { + // Shared MainEventBus for all tests - ensures events are properly distributed + private static final MainEventBus TEST_EVENT_BUS = new MainEventBus(); + // Since EventQueue.builder() is package protected, add a method to expose it public static EventQueue.EventQueueBuilder getEventQueueBuilder() { - return EventQueue.builder(); + return EventQueue.builder(TEST_EVENT_BUS); } public static void start(MainEventBusProcessor processor) { diff --git a/server-common/src/main/java/io/a2a/server/events/EventQueue.java b/server-common/src/main/java/io/a2a/server/events/EventQueue.java index a0bb074cb..d518ffd9e 100644 --- a/server-common/src/main/java/io/a2a/server/events/EventQueue.java +++ b/server-common/src/main/java/io/a2a/server/events/EventQueue.java @@ -1,6 +1,7 @@ package io.a2a.server.events; import java.util.List; +import java.util.Objects; import java.util.concurrent.BlockingQueue; import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.CountDownLatch; @@ -23,7 +24,7 @@ * and hierarchical queue structures via MainQueue and ChildQueue implementations. *

*

- * Use {@link #builder()} to create configured instances or extend MainQueue/ChildQueue directly. + * Use {@link #builder(MainEventBus)} to create configured instances or extend MainQueue/ChildQueue directly. *

*/ public abstract class EventQueue implements AutoCloseable { @@ -36,14 +37,17 @@ public abstract class EventQueue implements AutoCloseable { public static final int DEFAULT_QUEUE_SIZE = 1000; private final int queueSize; + /** * Internal blocking queue for storing event queue items. */ protected final BlockingQueue queue = new LinkedBlockingDeque<>(); + /** * Semaphore for backpressure control, limiting the number of pending events. */ protected final Semaphore semaphore; + private volatile boolean closed = false; /** @@ -78,8 +82,8 @@ protected EventQueue(EventQueue parent) { LOGGER.trace("Creating {}, parent: {}", this, parent); } - static EventQueueBuilder builder() { - return new EventQueueBuilder(); + static EventQueueBuilder builder(MainEventBus mainEventBus) { + return new EventQueueBuilder().mainEventBus(mainEventBus); } /** @@ -171,11 +175,14 @@ public EventQueueBuilder mainEventBus(MainEventBus mainEventBus) { * @return a new MainQueue instance */ public EventQueue build() { - if (hook != null || !onCloseCallbacks.isEmpty() || taskStateProvider != null || mainEventBus != null) { - return new MainQueue(queueSize, hook, taskId, onCloseCallbacks, taskStateProvider, mainEventBus); - } else { - return new MainQueue(queueSize); + // MainEventBus is now REQUIRED - enforce single architectural path + if (mainEventBus == null) { + throw new IllegalStateException("MainEventBus is required for EventQueue creation"); + } + if (taskId == null) { + throw new IllegalStateException("taskId is required for EventQueue creation"); } + return new MainQueue(queueSize, hook, taskId, onCloseCallbacks, taskStateProvider, mainEventBus); } } @@ -221,22 +228,7 @@ public void enqueueEvent(Event event) { * @param item the event queue item to enqueue * @throws RuntimeException if interrupted while waiting to acquire the semaphore */ - public void enqueueItem(EventQueueItem item) { - Event event = item.getEvent(); - if (closed) { - LOGGER.warn("Queue is closed. Event will not be enqueued. {} {}", this, event); - return; - } - // Call toString() since for errors we don't really want the full stacktrace - try { - semaphore.acquire(); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new RuntimeException("Unable to acquire the semaphore to enqueue the event", e); - } - queue.add(item); - LOGGER.debug("Enqueued event {} {}", event instanceof Throwable ? event.toString() : event, this); - } + public abstract void enqueueItem(EventQueueItem item); /** * Creates a child queue that shares events with this queue. @@ -256,48 +248,17 @@ public void enqueueItem(EventQueueItem item) { * This method returns the full EventQueueItem wrapper, allowing callers to check * metadata like whether the event is replicated via {@link EventQueueItem#isReplicated()}. *

+ *

+ * Note: MainQueue does not support dequeue operations - only ChildQueues can be consumed. + *

* * @param waitMilliSeconds the maximum time to wait in milliseconds * @return the EventQueueItem, or null if timeout occurs * @throws EventQueueClosedException if the queue is closed and empty + * @throws UnsupportedOperationException if called on MainQueue */ - public @Nullable EventQueueItem dequeueEventItem(int waitMilliSeconds) throws EventQueueClosedException { - if (closed && queue.isEmpty()) { - LOGGER.debug("Queue is closed, and empty. Sending termination message. {}", this); - throw new EventQueueClosedException(); - } - try { - if (waitMilliSeconds <= 0) { - EventQueueItem item = queue.poll(); - if (item != null) { - Event event = item.getEvent(); - // Call toString() since for errors we don't really want the full stacktrace - LOGGER.debug("Dequeued event item (no wait) {} {}", this, event instanceof Throwable ? event.toString() : event); - semaphore.release(); - } - return item; - } - try { - LOGGER.trace("Polling queue {} (wait={}ms)", System.identityHashCode(this), waitMilliSeconds); - EventQueueItem item = queue.poll(waitMilliSeconds, TimeUnit.MILLISECONDS); - if (item != null) { - Event event = item.getEvent(); - // Call toString() since for errors we don't really want the full stacktrace - LOGGER.debug("Dequeued event item (waiting) {} {}", this, event instanceof Throwable ? event.toString() : event); - semaphore.release(); - } else { - LOGGER.trace("Dequeue timeout (null) from queue {}", System.identityHashCode(this)); - } - return item; - } catch (InterruptedException e) { - LOGGER.debug("Interrupted dequeue (waiting) {}", this); - Thread.currentThread().interrupt(); - return null; - } - } finally { - signalQueuePollerStarted(); - } - } + @Nullable + public abstract EventQueueItem dequeueEventItem(int waitMilliSeconds) throws EventQueueClosedException; /** * Placeholder method for task completion notification. @@ -307,6 +268,17 @@ public void taskDone() { // TODO Not sure if needed yet. BlockingQueue.poll()/.take() remove the events. } + /** + * Returns the current size of the queue. + *

+ * For MainQueue: returns the size of the MainEventBus queue (events pending persistence/distribution). + * For ChildQueue: returns the size of the local consumption queue. + *

+ * + * @return the number of events currently in the queue + */ + public abstract int size(); + /** * Closes this event queue gracefully, allowing pending events to be consumed. */ @@ -360,13 +332,7 @@ protected void doClose(boolean immediate) { LOGGER.debug("Closing {} (immediate={})", this, immediate); closed = true; } - - if (immediate) { - // Immediate close: clear pending events - queue.clear(); - LOGGER.debug("Cleared queue for immediate close: {}", this); - } - // For graceful close, let the queue drain naturally through normal consumption + // Subclasses handle immediate close logic (e.g., ChildQueue clears its local queue) } static class MainQueue extends EventQueue { @@ -374,50 +340,14 @@ static class MainQueue extends EventQueue { private final CountDownLatch pollingStartedLatch = new CountDownLatch(1); private final AtomicBoolean pollingStarted = new AtomicBoolean(false); private final @Nullable EventEnqueueHook enqueueHook; - private final @Nullable String taskId; + private final String taskId; private final List onCloseCallbacks; private final @Nullable TaskStateProvider taskStateProvider; - private final @Nullable MainEventBus mainEventBus; - - MainQueue() { - super(); - this.enqueueHook = null; - this.taskId = null; - this.onCloseCallbacks = List.of(); - this.taskStateProvider = null; - this.mainEventBus = null; - } - - MainQueue(int queueSize) { - super(queueSize); - this.enqueueHook = null; - this.taskId = null; - this.onCloseCallbacks = List.of(); - this.taskStateProvider = null; - this.mainEventBus = null; - } - - MainQueue(EventEnqueueHook hook) { - super(); - this.enqueueHook = hook; - this.taskId = null; - this.onCloseCallbacks = List.of(); - this.taskStateProvider = null; - this.mainEventBus = null; - } - - MainQueue(int queueSize, EventEnqueueHook hook) { - super(queueSize); - this.enqueueHook = hook; - this.taskId = null; - this.onCloseCallbacks = List.of(); - this.taskStateProvider = null; - this.mainEventBus = null; - } + private final MainEventBus mainEventBus; MainQueue(int queueSize, @Nullable EventEnqueueHook hook, - @Nullable String taskId, + String taskId, List onCloseCallbacks, @Nullable TaskStateProvider taskStateProvider, @Nullable MainEventBus mainEventBus) { @@ -426,9 +356,9 @@ static class MainQueue extends EventQueue { this.taskId = taskId; this.onCloseCallbacks = List.copyOf(onCloseCallbacks); // Defensive copy this.taskStateProvider = taskStateProvider; - this.mainEventBus = mainEventBus; - LOGGER.debug("Created MainQueue for task {} with {} onClose callbacks, TaskStateProvider: {}, MainEventBus: {}", - taskId, onCloseCallbacks.size(), taskStateProvider != null, mainEventBus != null); + this.mainEventBus = Objects.requireNonNull(mainEventBus, "MainEventBus is required"); + LOGGER.debug("Created MainQueue for task {} with {} onClose callbacks, TaskStateProvider: {}, MainEventBus configured", + taskId, onCloseCallbacks.size(), taskStateProvider != null); } @@ -446,6 +376,17 @@ public int getChildCount() { return children.size(); } + @Override + public EventQueueItem dequeueEventItem(int waitMilliSeconds) throws EventQueueClosedException { + throw new UnsupportedOperationException("MainQueue cannot be consumed directly - use tap() to create a ChildQueue for consumption"); + } + + @Override + public int size() { + // Return size of MainEventBus queue (events pending persistence/distribution) + return mainEventBus.size(); + } + @Override public void enqueueItem(EventQueueItem item) { // MainQueue must accept events even when closed to support: @@ -464,19 +405,13 @@ public void enqueueItem(EventQueueItem item) { throw new RuntimeException("Unable to acquire the semaphore to enqueue the event", e); } - // Add to this MainQueue's internal queue - queue.add(item); LOGGER.debug("Enqueued event {} {}", event instanceof Throwable ? event.toString() : event, this); // Submit to MainEventBus for centralized persistence + distribution - if (mainEventBus != null && taskId != null) { - mainEventBus.submit(taskId, this, item); - } else { - // This should not happen in properly configured systems - LOGGER.error("MainEventBus not configured for task {} - events will NOT be distributed to children!", taskId); - } + // MainEventBus is guaranteed non-null by constructor requirement + mainEventBus.submit(taskId, this, item); - // Trigger replication hook if configured (KEEP for inter-process replication) + // Trigger replication hook if configured (for inter-process replication) if (enqueueHook != null) { enqueueHook.onEnqueue(item); } @@ -603,6 +538,7 @@ public void close(boolean immediate, boolean notifyParent) { static class ChildQueue extends EventQueue { private final MainQueue parent; + private final BlockingQueue queue = new LinkedBlockingDeque<>(); public ChildQueue(MainQueue parent) { this.parent = parent; @@ -613,8 +549,65 @@ public void enqueueEvent(Event event) { parent.enqueueEvent(event); } + @Override + public void enqueueItem(EventQueueItem item) { + // ChildQueue delegates writes to parent MainQueue + parent.enqueueItem(item); + } + private void internalEnqueueItem(EventQueueItem item) { - super.enqueueItem(item); + // Internal method called by MainEventBusProcessor to add to local queue + Event event = item.getEvent(); + if (isClosed()) { + LOGGER.warn("ChildQueue is closed. Event will not be enqueued. {} {}", this, event); + return; + } + try { + semaphore.acquire(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException("Unable to acquire the semaphore to enqueue the event", e); + } + queue.add(item); + LOGGER.debug("Enqueued event {} {}", event instanceof Throwable ? event.toString() : event, this); + } + + @Override + @Nullable + public EventQueueItem dequeueEventItem(int waitMilliSeconds) throws EventQueueClosedException { + if (isClosed() && queue.isEmpty()) { + LOGGER.debug("ChildQueue is closed, and empty. Sending termination message. {}", this); + throw new EventQueueClosedException(); + } + try { + if (waitMilliSeconds <= 0) { + EventQueueItem item = queue.poll(); + if (item != null) { + Event event = item.getEvent(); + LOGGER.debug("Dequeued event item (no wait) {} {}", this, event instanceof Throwable ? event.toString() : event); + semaphore.release(); + } + return item; + } + try { + LOGGER.trace("Polling ChildQueue {} (wait={}ms)", System.identityHashCode(this), waitMilliSeconds); + EventQueueItem item = queue.poll(waitMilliSeconds, TimeUnit.MILLISECONDS); + if (item != null) { + Event event = item.getEvent(); + LOGGER.debug("Dequeued event item (waiting) {} {}", this, event instanceof Throwable ? event.toString() : event); + semaphore.release(); + } else { + LOGGER.trace("Dequeue timeout (null) from ChildQueue {}", System.identityHashCode(this)); + } + return item; + } catch (InterruptedException e) { + LOGGER.debug("Interrupted dequeue (waiting) {}", this); + Thread.currentThread().interrupt(); + return null; + } + } finally { + signalQueuePollerStarted(); + } } @Override @@ -622,6 +615,12 @@ public EventQueue tap() { throw new IllegalStateException("Can only tap the main queue"); } + @Override + public int size() { + // Return size of local consumption queue + return queue.size(); + } + @Override public void awaitQueuePollerStart() throws InterruptedException { parent.awaitQueuePollerStart(); @@ -632,6 +631,17 @@ public void signalQueuePollerStarted() { parent.signalQueuePollerStarted(); } + @Override + protected void doClose(boolean immediate) { + super.doClose(immediate); // Sets closed flag + if (immediate) { + // Immediate close: clear pending events from local queue + queue.clear(); + LOGGER.debug("Cleared ChildQueue for immediate close: {}", this); + } + // For graceful close, let the queue drain naturally through normal consumption + } + @Override public void close() { close(false); diff --git a/server-common/src/main/java/io/a2a/server/events/InMemoryQueueManager.java b/server-common/src/main/java/io/a2a/server/events/InMemoryQueueManager.java index 63063ba69..f1a704c43 100644 --- a/server-common/src/main/java/io/a2a/server/events/InMemoryQueueManager.java +++ b/server-common/src/main/java/io/a2a/server/events/InMemoryQueueManager.java @@ -139,11 +139,10 @@ public int getActiveChildQueueCount(String taskId) { @Override public EventQueue.EventQueueBuilder createBaseEventQueueBuilder(String taskId) { - return EventQueue.builder() + return EventQueue.builder(mainEventBus) .taskId(taskId) .addOnCloseCallback(getCleanupCallback(taskId)) - .taskStateProvider(taskStateProvider) - .mainEventBus(mainEventBus); + .taskStateProvider(taskStateProvider); } /** diff --git a/server-common/src/main/java/io/a2a/server/events/QueueManager.java b/server-common/src/main/java/io/a2a/server/events/QueueManager.java index 395fddb9a..4ad30f0cb 100644 --- a/server-common/src/main/java/io/a2a/server/events/QueueManager.java +++ b/server-common/src/main/java/io/a2a/server/events/QueueManager.java @@ -177,7 +177,9 @@ public interface QueueManager { * @return a builder for creating event queues */ default EventQueue.EventQueueBuilder getEventQueueBuilder(String taskId) { - return EventQueue.builder(); + throw new UnsupportedOperationException( + "QueueManager implementations must override getEventQueueBuilder() to provide MainEventBus" + ); } /** @@ -197,7 +199,9 @@ default EventQueue.EventQueueBuilder getEventQueueBuilder(String taskId) { * @return a builder with base configuration specific to this QueueManager implementation */ default EventQueue.EventQueueBuilder createBaseEventQueueBuilder(String taskId) { - return EventQueue.builder().taskId(taskId); + throw new UnsupportedOperationException( + "QueueManager implementations must override createBaseEventQueueBuilder() to provide MainEventBus" + ); } /** diff --git a/server-common/src/main/java/io/a2a/server/requesthandlers/DefaultRequestHandler.java b/server-common/src/main/java/io/a2a/server/requesthandlers/DefaultRequestHandler.java index fa55b5142..5a6ddd7bb 100644 --- a/server-common/src/main/java/io/a2a/server/requesthandlers/DefaultRequestHandler.java +++ b/server-common/src/main/java/io/a2a/server/requesthandlers/DefaultRequestHandler.java @@ -333,7 +333,7 @@ public Task onCancelTask(TaskIdParams params, ServerCallContext context) throws EventQueue queue = queueManager.tap(task.id()); if (queue == null) { - queue = queueManager.getEventQueueBuilder(task.id()).build(); + queue = queueManager.getEventQueueBuilder(task.id()).build().tap(); } agentExecutor.cancel( requestContextBuilder.get() diff --git a/server-common/src/test/java/io/a2a/server/events/EventConsumerTest.java b/server-common/src/test/java/io/a2a/server/events/EventConsumerTest.java index 6114e8f21..511e044ee 100644 --- a/server-common/src/test/java/io/a2a/server/events/EventConsumerTest.java +++ b/server-common/src/test/java/io/a2a/server/events/EventConsumerTest.java @@ -54,7 +54,7 @@ public class EventConsumerTest { @BeforeEach public void init() { - eventQueue = EventQueue.builder().build(); + eventQueue = EventQueueUtil.getEventQueueBuilder().build().tap(); eventConsumer = new EventConsumer(eventQueue); } @@ -340,7 +340,7 @@ public void onComplete() { @Test public void testConsumeAllStopsOnQueueClosed() throws Exception { - EventQueue queue = EventQueue.builder().build(); + EventQueue queue = EventQueueUtil.getEventQueueBuilder().build().tap(); EventConsumer consumer = new EventConsumer(queue); // Close the queue immediately @@ -386,13 +386,21 @@ public void onComplete() { @Test public void testConsumeAllHandlesQueueClosedException() throws Exception { - EventQueue queue = EventQueue.builder().build(); + EventQueue queue = EventQueueUtil.getEventQueueBuilder().build().tap(); EventConsumer consumer = new EventConsumer(queue); // Add a message event (which will complete the stream) Event message = fromJson(MESSAGE_PAYLOAD, Message.class); queue.enqueueEvent(message); + // Poll for event to arrive in ChildQueue (async MainEventBusProcessor distribution) + long startTime = System.currentTimeMillis(); + long timeout = 2000; + while (queue.size() == 0 && (System.currentTimeMillis() - startTime) < timeout) { + Thread.sleep(50); + } + assertTrue(queue.size() > 0, "Event should arrive in ChildQueue within timeout"); + // Close the queue before consuming queue.close(); @@ -436,7 +444,7 @@ public void onComplete() { @Test public void testConsumeAllTerminatesOnQueueClosedEvent() throws Exception { - EventQueue queue = EventQueue.builder().build(); + EventQueue queue = EventQueueUtil.getEventQueueBuilder().build().tap(); EventConsumer consumer = new EventConsumer(queue); // Enqueue a QueueClosedEvent (poison pill) @@ -486,7 +494,19 @@ public void onComplete() { private void enqueueAndConsumeOneEvent(Event event) throws Exception { eventQueue.enqueueEvent(event); - Event result = eventConsumer.consumeOne(); + // Poll for event with 2-second timeout + long startTime = System.currentTimeMillis(); + long timeout = 2000; + Event result = null; + while (result == null && (System.currentTimeMillis() - startTime) < timeout) { + try { + result = eventConsumer.consumeOne(); + } catch (A2AServerException e) { + // Event not available yet, wait a bit and try again + Thread.sleep(50); + } + } + assertNotNull(result, "Event should arrive within timeout"); assertSame(event, result); } diff --git a/server-common/src/test/java/io/a2a/server/events/EventQueueTest.java b/server-common/src/test/java/io/a2a/server/events/EventQueueTest.java index c0e22b8d0..6494eef87 100644 --- a/server-common/src/test/java/io/a2a/server/events/EventQueueTest.java +++ b/server-common/src/test/java/io/a2a/server/events/EventQueueTest.java @@ -61,10 +61,10 @@ public void init() { mainEventBusProcessor = new MainEventBusProcessor(mainEventBus, taskStore, NOOP_PUSHNOTIFICATION_SENDER); EventQueueUtil.start(mainEventBusProcessor); - eventQueue = EventQueue.builder() + eventQueue = EventQueueUtil.getEventQueueBuilder() .taskId("test-task") .mainEventBus(mainEventBus) - .build(); + .build().tap(); } @AfterEach @@ -78,7 +78,7 @@ public void cleanup() { * Helper to create a queue with MainEventBus configured (for tests that need event distribution). */ private EventQueue createQueueWithEventBus(String taskId) { - return EventQueue.builder() + return EventQueueUtil.getEventQueueBuilder() .taskId(taskId) .mainEventBus(mainEventBus) .build(); @@ -86,29 +86,29 @@ private EventQueue createQueueWithEventBus(String taskId) { @Test public void testConstructorDefaultQueueSize() { - EventQueue queue = EventQueue.builder().build(); + EventQueue queue = EventQueueUtil.getEventQueueBuilder().build(); assertEquals(EventQueue.DEFAULT_QUEUE_SIZE, queue.getQueueSize()); } @Test public void testConstructorCustomQueueSize() { int customSize = 500; - EventQueue queue = EventQueue.builder().queueSize(customSize).build(); + EventQueue queue = EventQueueUtil.getEventQueueBuilder().queueSize(customSize).build(); assertEquals(customSize, queue.getQueueSize()); } @Test public void testConstructorInvalidQueueSize() { // Test zero queue size - assertThrows(IllegalArgumentException.class, () -> EventQueue.builder().queueSize(0).build()); + assertThrows(IllegalArgumentException.class, () -> EventQueueUtil.getEventQueueBuilder().queueSize(0).build()); // Test negative queue size - assertThrows(IllegalArgumentException.class, () -> EventQueue.builder().queueSize(-10).build()); + assertThrows(IllegalArgumentException.class, () -> EventQueueUtil.getEventQueueBuilder().queueSize(-10).build()); } @Test public void testTapCreatesChildQueue() { - EventQueue parentQueue = EventQueue.builder().build(); + EventQueue parentQueue = EventQueueUtil.getEventQueueBuilder().build(); EventQueue childQueue = parentQueue.tap(); assertNotNull(childQueue); @@ -118,7 +118,7 @@ public void testTapCreatesChildQueue() { @Test public void testTapOnChildQueueThrowsException() { - EventQueue parentQueue = EventQueue.builder().build(); + EventQueue parentQueue = EventQueueUtil.getEventQueueBuilder().build(); EventQueue childQueue = parentQueue.tap(); assertThrows(IllegalStateException.class, () -> childQueue.tap()); @@ -126,53 +126,56 @@ public void testTapOnChildQueueThrowsException() { @Test public void testEnqueueEventPropagagesToChildren() throws Exception { - EventQueue parentQueue = createQueueWithEventBus("test-propagate"); - EventQueue childQueue = parentQueue.tap(); + EventQueue mainQueue = createQueueWithEventBus("test-propagate"); + EventQueue childQueue1 = mainQueue.tap(); + EventQueue childQueue2 = mainQueue.tap(); Event event = fromJson(MINIMAL_TASK, Task.class); - parentQueue.enqueueEvent(event); + mainQueue.enqueueEvent(event); - // Event should be available in both parent and child queues + // Event should be available in all child queues // Note: MainEventBusProcessor runs async, so we use dequeueEventItem with timeout - Event parentEvent = parentQueue.dequeueEventItem(5000).getEvent(); - Event childEvent = childQueue.dequeueEventItem(5000).getEvent(); + Event child1Event = childQueue1.dequeueEventItem(5000).getEvent(); + Event child2Event = childQueue2.dequeueEventItem(5000).getEvent(); - assertSame(event, parentEvent); - assertSame(event, childEvent); + assertSame(event, child1Event); + assertSame(event, child2Event); } @Test public void testMultipleChildQueuesReceiveEvents() throws Exception { - EventQueue parentQueue = createQueueWithEventBus("test-multiple"); - EventQueue childQueue1 = parentQueue.tap(); - EventQueue childQueue2 = parentQueue.tap(); + EventQueue mainQueue = createQueueWithEventBus("test-multiple"); + EventQueue childQueue1 = mainQueue.tap(); + EventQueue childQueue2 = mainQueue.tap(); + EventQueue childQueue3 = mainQueue.tap(); Event event1 = fromJson(MINIMAL_TASK, Task.class); Event event2 = fromJson(MESSAGE_PAYLOAD, Message.class); - parentQueue.enqueueEvent(event1); - parentQueue.enqueueEvent(event2); + mainQueue.enqueueEvent(event1); + mainQueue.enqueueEvent(event2); - // All queues should receive both events + // All child queues should receive both events // Note: Use timeout for async processing - assertSame(event1, parentQueue.dequeueEventItem(5000).getEvent()); - assertSame(event2, parentQueue.dequeueEventItem(5000).getEvent()); - assertSame(event1, childQueue1.dequeueEventItem(5000).getEvent()); assertSame(event2, childQueue1.dequeueEventItem(5000).getEvent()); assertSame(event1, childQueue2.dequeueEventItem(5000).getEvent()); assertSame(event2, childQueue2.dequeueEventItem(5000).getEvent()); + + assertSame(event1, childQueue3.dequeueEventItem(5000).getEvent()); + assertSame(event2, childQueue3.dequeueEventItem(5000).getEvent()); } @Test public void testChildQueueDequeueIndependently() throws Exception { - EventQueue parentQueue = createQueueWithEventBus("test-independent"); - EventQueue childQueue1 = parentQueue.tap(); - EventQueue childQueue2 = parentQueue.tap(); + EventQueue mainQueue = createQueueWithEventBus("test-independent"); + EventQueue childQueue1 = mainQueue.tap(); + EventQueue childQueue2 = mainQueue.tap(); + EventQueue childQueue3 = mainQueue.tap(); Event event = fromJson(MINIMAL_TASK, Task.class); - parentQueue.enqueueEvent(event); + mainQueue.enqueueEvent(event); // Dequeue from child1 first (use timeout for async processing) Event child1Event = childQueue1.dequeueEventItem(5000).getEvent(); @@ -182,9 +185,9 @@ public void testChildQueueDequeueIndependently() throws Exception { Event child2Event = childQueue2.dequeueEventItem(5000).getEvent(); assertSame(event, child2Event); - // Parent should still have the event available - Event parentEvent = parentQueue.dequeueEventItem(5000).getEvent(); - assertSame(event, parentEvent); + // child3 should still have the event available + Event child3Event = childQueue3.dequeueEventItem(5000).getEvent(); + assertSame(event, child3Event); } @@ -220,27 +223,32 @@ public void testCloseImmediatePropagationToChildren() throws Exception { @Test public void testEnqueueEventWhenClosed() throws Exception { - EventQueue queue = EventQueue.builder().build(); + EventQueue mainQueue = EventQueueUtil.getEventQueueBuilder().build(); + EventQueue childQueue = mainQueue.tap(); Event event = fromJson(MINIMAL_TASK, Task.class); - queue.close(); // Close the queue first - assertTrue(queue.isClosed()); + childQueue.close(); // Close the child queue first + assertTrue(childQueue.isClosed()); // MainQueue accepts events even when closed (for replication support) // This ensures late-arriving replicated events can be enqueued to closed queues - queue.enqueueEvent(event); + mainQueue.enqueueEvent(event); - // Event should be available for dequeuing - Event dequeuedEvent = queue.dequeueEventItem(-1).getEvent(); + // Create a new child queue to receive the event (closed child won't receive it) + EventQueue newChildQueue = mainQueue.tap(); + EventQueueItem item = newChildQueue.dequeueEventItem(5000); + assertNotNull(item); + Event dequeuedEvent = item.getEvent(); assertSame(event, dequeuedEvent); - // Now queue is closed and empty, should throw exception - assertThrows(EventQueueClosedException.class, () -> queue.dequeueEventItem(-1)); + // Now new child queue is closed and empty, should throw exception + newChildQueue.close(); + assertThrows(EventQueueClosedException.class, () -> newChildQueue.dequeueEventItem(-1)); } @Test public void testDequeueEventWhenClosedAndEmpty() throws Exception { - EventQueue queue = EventQueue.builder().build(); + EventQueue queue = EventQueueUtil.getEventQueueBuilder().build().tap(); queue.close(); assertTrue(queue.isClosed()); @@ -250,19 +258,28 @@ public void testDequeueEventWhenClosedAndEmpty() throws Exception { @Test public void testDequeueEventWhenClosedButHasEvents() throws Exception { - EventQueue queue = EventQueue.builder().build(); + EventQueue mainQueue = EventQueueUtil.getEventQueueBuilder().build(); + EventQueue childQueue = mainQueue.tap(); Event event = fromJson(MINIMAL_TASK, Task.class); - queue.enqueueEvent(event); - queue.close(); // Graceful close - events should remain - assertTrue(queue.isClosed()); + // Enqueue to mainQueue + mainQueue.enqueueEvent(event); + + // Wait for event to arrive in childQueue (use peek-like behavior by dequeueing then re-checking) + // Actually, just wait a bit for async processing + Thread.sleep(100); // Give MainEventBusProcessor time to distribute event + + childQueue.close(); // Graceful close - events should remain + assertTrue(childQueue.isClosed()); - // Should still be able to dequeue existing events - Event dequeuedEvent = queue.dequeueEventItem(-1).getEvent(); + // Should still be able to dequeue existing events from closed queue + EventQueueItem item = childQueue.dequeueEventItem(5000); + assertNotNull(item); + Event dequeuedEvent = item.getEvent(); assertSame(event, dequeuedEvent); // Now queue is closed and empty, should throw exception - assertThrows(EventQueueClosedException.class, () -> queue.dequeueEventItem(-1)); + assertThrows(EventQueueClosedException.class, () -> childQueue.dequeueEventItem(-1)); } @Test @@ -277,7 +294,9 @@ public void testEnqueueAndDequeueEvent() throws Exception { public void testDequeueEventNoWait() throws Exception { Event event = fromJson(MINIMAL_TASK, Task.class); eventQueue.enqueueEvent(event); - Event dequeuedEvent = eventQueue.dequeueEventItem(-1).getEvent(); + EventQueueItem item = eventQueue.dequeueEventItem(5000); + assertNotNull(item); + Event dequeuedEvent = item.getEvent(); assertSame(event, dequeuedEvent); } @@ -380,7 +399,7 @@ public void testCloseIdempotent() throws Exception { assertTrue(eventQueue.isClosed()); // Test with immediate close as well - EventQueue eventQueue2 = EventQueue.builder().build(); + EventQueue eventQueue2 = EventQueueUtil.getEventQueueBuilder().build(); eventQueue2.close(true); assertTrue(eventQueue2.isClosed()); @@ -394,19 +413,20 @@ public void testCloseIdempotent() throws Exception { */ @Test public void testCloseChildQueues() throws Exception { - EventQueue childQueue = eventQueue.tap(); + EventQueue mainQueue = EventQueueUtil.getEventQueueBuilder().build(); + EventQueue childQueue = mainQueue.tap(); assertTrue(childQueue != null); // Graceful close - parent closes but children remain open - eventQueue.close(); - assertTrue(eventQueue.isClosed()); + mainQueue.close(); + assertTrue(mainQueue.isClosed()); assertFalse(childQueue.isClosed()); // Child NOT closed on graceful parent close // Immediate close - parent force-closes all children - EventQueue parentQueue2 = EventQueue.builder().build(); - EventQueue childQueue2 = parentQueue2.tap(); - parentQueue2.close(true); // immediate=true - assertTrue(parentQueue2.isClosed()); + EventQueue mainQueue2 = EventQueueUtil.getEventQueueBuilder().build(); + EventQueue childQueue2 = mainQueue2.tap(); + mainQueue2.close(true); // immediate=true + assertTrue(mainQueue2.isClosed()); assertTrue(childQueue2.isClosed()); // Child IS closed on immediate parent close } @@ -416,7 +436,7 @@ public void testCloseChildQueues() throws Exception { */ @Test public void testMainQueueReferenceCountingStaysOpenWithActiveChildren() throws Exception { - EventQueue mainQueue = EventQueue.builder().build(); + EventQueue mainQueue = EventQueueUtil.getEventQueueBuilder().build(); EventQueue child1 = mainQueue.tap(); EventQueue child2 = mainQueue.tap(); diff --git a/server-common/src/test/java/io/a2a/server/events/EventQueueUtil.java b/server-common/src/test/java/io/a2a/server/events/EventQueueUtil.java index 8490aa00c..beb8b3d5a 100644 --- a/server-common/src/test/java/io/a2a/server/events/EventQueueUtil.java +++ b/server-common/src/test/java/io/a2a/server/events/EventQueueUtil.java @@ -1,9 +1,44 @@ package io.a2a.server.events; +import io.a2a.server.tasks.InMemoryTaskStore; +import io.a2a.server.tasks.PushNotificationSender; +import io.a2a.server.tasks.TaskStateProvider; +import java.util.concurrent.atomic.AtomicInteger; + public class EventQueueUtil { + // Shared MainEventBus for all tests (to avoid creating one per test) + private static final MainEventBus TEST_EVENT_BUS = new MainEventBus(); + + // Shared MainEventBusProcessor for all tests (automatically processes events) + private static final MainEventBusProcessor TEST_PROCESSOR; + + static { + // Initialize and start the processor once for all tests + InMemoryTaskStore testTaskStore = new InMemoryTaskStore(); + PushNotificationSender testPushSender = taskId -> {}; // No-op for tests + TEST_PROCESSOR = new MainEventBusProcessor(TEST_EVENT_BUS, testTaskStore, testPushSender); + TEST_PROCESSOR.start(); // Start background thread + + // Register shutdown hook to stop processor + Runtime.getRuntime().addShutdownHook(new Thread(() -> TEST_PROCESSOR.stop())); + } + + // Counter for generating unique test taskIds + private static final AtomicInteger TASK_ID_COUNTER = new AtomicInteger(0); + // Since EventQueue.builder() is package protected, add a method to expose it + // Note: Now includes MainEventBus requirement and default taskId + // Returns MainQueue - tests should call .tap() if they need to consume events public static EventQueue.EventQueueBuilder getEventQueueBuilder() { - return EventQueue.builder(); + return new EventQueueBuilderWrapper( + EventQueue.builder(TEST_EVENT_BUS) + .taskId("test-task-" + TASK_ID_COUNTER.incrementAndGet()) + ); + } + + // Get the shared test MainEventBus instance + public static MainEventBus getTestEventBus() { + return TEST_EVENT_BUS; } public static void start(MainEventBusProcessor processor) { @@ -13,4 +48,55 @@ public static void start(MainEventBusProcessor processor) { public static void stop(MainEventBusProcessor processor) { processor.stop(); } + + // Wrapper that delegates to actual builder + private static class EventQueueBuilderWrapper extends EventQueue.EventQueueBuilder { + private final EventQueue.EventQueueBuilder delegate; + + EventQueueBuilderWrapper(EventQueue.EventQueueBuilder delegate) { + this.delegate = delegate; + } + + @Override + public EventQueue.EventQueueBuilder queueSize(int queueSize) { + delegate.queueSize(queueSize); + return this; + } + + @Override + public EventQueue.EventQueueBuilder hook(EventEnqueueHook hook) { + delegate.hook(hook); + return this; + } + + @Override + public EventQueue.EventQueueBuilder taskId(String taskId) { + delegate.taskId(taskId); + return this; + } + + @Override + public EventQueue.EventQueueBuilder addOnCloseCallback(Runnable onCloseCallback) { + delegate.addOnCloseCallback(onCloseCallback); + return this; + } + + @Override + public EventQueue.EventQueueBuilder taskStateProvider(TaskStateProvider taskStateProvider) { + delegate.taskStateProvider(taskStateProvider); + return this; + } + + @Override + public EventQueue.EventQueueBuilder mainEventBus(MainEventBus mainEventBus) { + delegate.mainEventBus(mainEventBus); + return this; + } + + @Override + public EventQueue build() { + // Return MainQueue directly - tests should call .tap() if they need ChildQueue + return delegate.build(); + } + } } diff --git a/server-common/src/test/java/io/a2a/server/events/InMemoryQueueManagerTest.java b/server-common/src/test/java/io/a2a/server/events/InMemoryQueueManagerTest.java index 1d10be2c8..46f97b510 100644 --- a/server-common/src/test/java/io/a2a/server/events/InMemoryQueueManagerTest.java +++ b/server-common/src/test/java/io/a2a/server/events/InMemoryQueueManagerTest.java @@ -49,7 +49,7 @@ public void tearDown() { @Test public void testAddNewQueue() { String taskId = "test_task_id"; - EventQueue queue = EventQueue.builder().build(); + EventQueue queue = EventQueueUtil.getEventQueueBuilder().build(); queueManager.add(taskId, queue); @@ -60,8 +60,8 @@ public void testAddNewQueue() { @Test public void testAddExistingQueueThrowsException() { String taskId = "test_task_id"; - EventQueue queue1 = EventQueue.builder().build(); - EventQueue queue2 = EventQueue.builder().build(); + EventQueue queue1 = EventQueueUtil.getEventQueueBuilder().build(); + EventQueue queue2 = EventQueueUtil.getEventQueueBuilder().build(); queueManager.add(taskId, queue1); @@ -73,7 +73,7 @@ public void testAddExistingQueueThrowsException() { @Test public void testGetExistingQueue() { String taskId = "test_task_id"; - EventQueue queue = EventQueue.builder().build(); + EventQueue queue = EventQueueUtil.getEventQueueBuilder().build(); queueManager.add(taskId, queue); EventQueue result = queueManager.get(taskId); @@ -90,7 +90,7 @@ public void testGetNonexistentQueue() { @Test public void testTapExistingQueue() { String taskId = "test_task_id"; - EventQueue queue = EventQueue.builder().build(); + EventQueue queue = EventQueueUtil.getEventQueueBuilder().build(); queueManager.add(taskId, queue); EventQueue tappedQueue = queueManager.tap(taskId); @@ -111,7 +111,7 @@ public void testTapNonexistentQueue() { @Test public void testCloseExistingQueue() { String taskId = "test_task_id"; - EventQueue queue = EventQueue.builder().build(); + EventQueue queue = EventQueueUtil.getEventQueueBuilder().build(); queueManager.add(taskId, queue); queueManager.close(taskId); @@ -146,7 +146,7 @@ public void testCreateOrTapNewQueue() { @Test public void testCreateOrTapExistingQueue() { String taskId = "test_task_id"; - EventQueue originalQueue = EventQueue.builder().build(); + EventQueue originalQueue = EventQueueUtil.getEventQueueBuilder().build(); queueManager.add(taskId, originalQueue); EventQueue result = queueManager.createOrTap(taskId); @@ -168,7 +168,7 @@ public void testConcurrentOperations() throws InterruptedException, ExecutionExc // Add tasks concurrently List> addFutures = taskIds.stream() .map(taskId -> CompletableFuture.supplyAsync(() -> { - EventQueue queue = EventQueue.builder().build(); + EventQueue queue = EventQueueUtil.getEventQueueBuilder().build(); queueManager.add(taskId, queue); return taskId; })) diff --git a/server-common/src/test/java/io/a2a/server/tasks/ResultAggregatorTest.java b/server-common/src/test/java/io/a2a/server/tasks/ResultAggregatorTest.java index ff801b70d..2c60cba50 100644 --- a/server-common/src/test/java/io/a2a/server/tasks/ResultAggregatorTest.java +++ b/server-common/src/test/java/io/a2a/server/tasks/ResultAggregatorTest.java @@ -16,8 +16,10 @@ import io.a2a.server.events.EventConsumer; import io.a2a.server.events.EventQueue; +import io.a2a.server.events.EventQueueUtil; import io.a2a.server.events.InMemoryQueueManager; import io.a2a.server.events.MainEventBus; +import io.a2a.server.events.MainEventBusProcessor; import io.a2a.spec.EventKind; import io.a2a.spec.Message; import io.a2a.spec.Task; @@ -205,12 +207,24 @@ void testConsumeAndBreakNonBlocking() throws Exception { // Create an event queue using QueueManager (which has access to builder) MainEventBus mainEventBus = new MainEventBus(); + InMemoryTaskStore taskStore = new InMemoryTaskStore(); + MainEventBusProcessor processor = new MainEventBusProcessor(mainEventBus, taskStore, task -> {}); + EventQueueUtil.start(processor); + InMemoryQueueManager queueManager = new InMemoryQueueManager(new MockTaskStateProvider(), mainEventBus); - EventQueue queue = queueManager.getEventQueueBuilder("test-task").build(); + EventQueue queue = queueManager.getEventQueueBuilder("test-task").build().tap(); queue.enqueueEvent(firstEvent); + // Poll for event to arrive in ChildQueue (async MainEventBusProcessor distribution) + long startTime = System.currentTimeMillis(); + long timeout = 2000; + while (queue.size() == 0 && (System.currentTimeMillis() - startTime) < timeout) { + Thread.sleep(50); + } + assertTrue(queue.size() > 0, "Event should arrive in ChildQueue within timeout"); + // Create real EventConsumer with the queue EventConsumer eventConsumer = new EventConsumer(queue); @@ -230,5 +244,8 @@ void testConsumeAndBreakNonBlocking() throws Exception { // The async consumer may or may not execute before verification, so we accept 1-2 calls verify(mockTaskManager, atLeast(1)).getTask(); verify(mockTaskManager, atMost(2)).getTask(); + + // Cleanup: stop the processor + EventQueueUtil.stop(processor); } } diff --git a/server-common/src/test/java/io/a2a/server/tasks/TaskUpdaterTest.java b/server-common/src/test/java/io/a2a/server/tasks/TaskUpdaterTest.java index 40f763569..c36ac7efd 100644 --- a/server-common/src/test/java/io/a2a/server/tasks/TaskUpdaterTest.java +++ b/server-common/src/test/java/io/a2a/server/tasks/TaskUpdaterTest.java @@ -14,6 +14,7 @@ import io.a2a.server.agentexecution.RequestContext; import io.a2a.server.events.EventQueue; +import io.a2a.server.events.EventQueueItem; import io.a2a.server.events.EventQueueUtil; import io.a2a.spec.Event; import io.a2a.spec.Message; @@ -45,7 +46,7 @@ public class TaskUpdaterTest { @BeforeEach public void init() { - eventQueue = EventQueueUtil.getEventQueueBuilder().build(); + eventQueue = EventQueueUtil.getEventQueueBuilder().build().tap(); RequestContext context = new RequestContext.Builder() .setTaskId(TEST_TASK_ID) .setContextId(TEST_TASK_CONTEXT_ID) @@ -56,7 +57,9 @@ public void init() { @Test public void testAddArtifactWithCustomIdAndName() throws Exception { taskUpdater.addArtifact(SAMPLE_PARTS, "custom-artifact-id", "Custom Artifact", null); - Event event = eventQueue.dequeueEventItem(0).getEvent(); + EventQueueItem item = eventQueue.dequeueEventItem(5000); + assertNotNull(item); + Event event = item.getEvent(); assertNotNull(event); assertInstanceOf(TaskArtifactUpdateEvent.class, event); @@ -239,7 +242,9 @@ public void testNewAgentMessageWithMetadata() throws Exception { @Test public void testAddArtifactWithAppendTrue() throws Exception { taskUpdater.addArtifact(SAMPLE_PARTS, "artifact-id", "Test Artifact", null, true, null); - Event event = eventQueue.dequeueEventItem(0).getEvent(); + EventQueueItem item = eventQueue.dequeueEventItem(5000); + assertNotNull(item); + Event event = item.getEvent(); assertNotNull(event); assertInstanceOf(TaskArtifactUpdateEvent.class, event); @@ -258,7 +263,9 @@ public void testAddArtifactWithAppendTrue() throws Exception { @Test public void testAddArtifactWithLastChunkTrue() throws Exception { taskUpdater.addArtifact(SAMPLE_PARTS, "artifact-id", "Test Artifact", null, null, true); - Event event = eventQueue.dequeueEventItem(0).getEvent(); + EventQueueItem item = eventQueue.dequeueEventItem(5000); + assertNotNull(item); + Event event = item.getEvent(); assertNotNull(event); assertInstanceOf(TaskArtifactUpdateEvent.class, event); @@ -273,7 +280,9 @@ public void testAddArtifactWithLastChunkTrue() throws Exception { @Test public void testAddArtifactWithAppendAndLastChunk() throws Exception { taskUpdater.addArtifact(SAMPLE_PARTS, "artifact-id", "Test Artifact", null, true, false); - Event event = eventQueue.dequeueEventItem(0).getEvent(); + EventQueueItem item = eventQueue.dequeueEventItem(5000); + assertNotNull(item); + Event event = item.getEvent(); assertNotNull(event); assertInstanceOf(TaskArtifactUpdateEvent.class, event); @@ -287,7 +296,9 @@ public void testAddArtifactWithAppendAndLastChunk() throws Exception { @Test public void testAddArtifactGeneratesIdWhenNull() throws Exception { taskUpdater.addArtifact(SAMPLE_PARTS, null, "Test Artifact", null); - Event event = eventQueue.dequeueEventItem(0).getEvent(); + EventQueueItem item = eventQueue.dequeueEventItem(5000); + assertNotNull(item); + Event event = item.getEvent(); assertNotNull(event); assertInstanceOf(TaskArtifactUpdateEvent.class, event); @@ -383,7 +394,9 @@ public void testConcurrentCompletionAttempts() throws Exception { thread2.join(); // Exactly one event should have been queued - Event event = eventQueue.dequeueEventItem(0).getEvent(); + EventQueueItem item = eventQueue.dequeueEventItem(5000); + assertNotNull(item); + Event event = item.getEvent(); assertNotNull(event); assertInstanceOf(TaskStatusUpdateEvent.class, event); @@ -396,7 +409,10 @@ public void testConcurrentCompletionAttempts() throws Exception { } private TaskStatusUpdateEvent checkTaskStatusUpdateEventOnQueue(boolean isFinal, TaskState state, Message statusMessage) throws Exception { - Event event = eventQueue.dequeueEventItem(0).getEvent(); + // Wait up to 5 seconds for event (async MainEventBusProcessor needs time to distribute) + EventQueueItem item = eventQueue.dequeueEventItem(5000); + assertNotNull(item); + Event event = item.getEvent(); assertNotNull(event); assertInstanceOf(TaskStatusUpdateEvent.class, event); @@ -408,6 +424,7 @@ private TaskStatusUpdateEvent checkTaskStatusUpdateEventOnQueue(boolean isFinal, assertEquals(state, tsue.status().state()); assertEquals(statusMessage, tsue.status().message()); + // Check no additional events (still use 0 timeout for this check) assertNull(eventQueue.dequeueEventItem(0)); return tsue; From e25166d644801bba7f1804ed467c3d0e82031b47 Mon Sep 17 00:00:00 2001 From: Kabir Khan Date: Mon, 8 Dec 2025 11:11:17 +0000 Subject: [PATCH 3/8] Get rid of the background processing --- .../DefaultRequestHandler.java | 114 ++++-------------- 1 file changed, 23 insertions(+), 91 deletions(-) diff --git a/server-common/src/main/java/io/a2a/server/requesthandlers/DefaultRequestHandler.java b/server-common/src/main/java/io/a2a/server/requesthandlers/DefaultRequestHandler.java index 5a6ddd7bb..b80080309 100644 --- a/server-common/src/main/java/io/a2a/server/requesthandlers/DefaultRequestHandler.java +++ b/server-common/src/main/java/io/a2a/server/requesthandlers/DefaultRequestHandler.java @@ -11,13 +11,11 @@ import java.util.Map; import java.util.Objects; import java.util.Optional; -import java.util.Set; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentHashMap; 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; @@ -122,7 +120,6 @@ *
  • {@link EventConsumer} polls and processes events on Vert.x worker thread
    • *
  • Queue closes automatically on final event (COMPLETED/FAILED/CANCELED)
    • *
  • Cleanup waits for both agent execution AND event consumption to complete
    • - *
  • Background tasks tracked via {@link #trackBackgroundTask(CompletableFuture)}
    • * * *

    Threading Model

    @@ -214,7 +211,6 @@ public class DefaultRequestHandler implements RequestHandler { private final Supplier requestContextBuilder; private final ConcurrentMap> runningAgents = new ConcurrentHashMap<>(); - private final Set> backgroundTasks = ConcurrentHashMap.newKeySet(); private final Executor executor; @@ -482,9 +478,9 @@ public EventKind onMessageSend(MessageSendParams params, ServerCallContext conte CompletableFuture agentFuture = runningAgents.remove(taskId); LOGGER.debug("Removed agent for task {} from runningAgents in finally block, size after: {}", taskId, runningAgents.size()); - // Track cleanup as background task to avoid blocking Vert.x threads + // Cleanup as background task to avoid blocking Vert.x threads // Pass the consumption future to ensure cleanup waits for background consumption to complete - trackBackgroundTask(cleanupProducer(agentFuture, etai != null ? etai.consumptionFuture() : null, taskId, queue, false)); + cleanupProducer(agentFuture, etai != null ? etai.consumptionFuture() : null, taskId, queue, false); } LOGGER.debug("Returning: {}", kind); @@ -494,8 +490,8 @@ public EventKind onMessageSend(MessageSendParams params, ServerCallContext conte @Override public Flow.Publisher onMessageSendStream( MessageSendParams params, ServerCallContext context) throws A2AError { - LOGGER.debug("onMessageSendStream START - task: {}; context: {}; runningAgents: {}; backgroundTasks: {}", - params.message().taskId(), params.message().contextId(), runningAgents.size(), backgroundTasks.size()); + LOGGER.debug("onMessageSendStream START - task: {}; context: {}; runningAgents: {}", + params.message().taskId(), params.message().contextId(), runningAgents.size()); MessageSendSetup mss = initMessageSend(params, context); @Nullable String initialTaskId = mss.requestContext.getTaskId(); @@ -503,7 +499,7 @@ public Flow.Publisher onMessageSendStream( // Use a temporary ID for queue creation if needed String queueTaskId = initialTaskId != null ? initialTaskId : "temp-" + java.util.UUID.randomUUID(); - AtomicReference<@NonNull String> taskId = new AtomicReference<>(queueTaskId); + final AtomicReference<@NonNull String> taskId = new AtomicReference<>(queueTaskId); @SuppressWarnings("NullAway") EventQueue queue = queueManager.createOrTap(taskId.get()); LOGGER.debug("Created/tapped queue for task {}: {}", taskId.get(), queue); @@ -512,12 +508,9 @@ public Flow.Publisher onMessageSendStream( EnhancedRunnable producerRunnable = registerAndExecuteAgentAsync(queueTaskId, 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()); - AtomicBoolean backgroundConsumeStarted = new AtomicBoolean(false); - try { Flow.Publisher results = resultAggregator.consumeAndEmit(consumer); @@ -556,7 +549,8 @@ public Flow.Publisher onMessageSendStream( Flow.Publisher finalPublisher = convertingProcessor(eventPublisher, event -> (StreamingEventKind) event); - // Wrap publisher to detect client disconnect and continue background consumption + // Wrap publisher to detect client disconnect and immediately close ChildQueue + // This prevents ChildQueue backpressure from blocking MainEventBusProcessor return subscriber -> { String currentTaskId = taskId.get(); LOGGER.debug("Creating subscription wrapper for task {}", currentTaskId); @@ -577,8 +571,10 @@ public void request(long n) { @Override public void cancel() { - LOGGER.debug("Client cancelled subscription for task {}, starting background consumption", taskId.get()); - startBackgroundConsumption(); + LOGGER.debug("Client cancelled subscription for task {}, closing ChildQueue immediately", taskId.get()); + // Close ChildQueue immediately to prevent backpressure + // (clears queue and releases semaphore permits) + queue.close(true); // immediate=true subscription.cancel(); } }); @@ -603,8 +599,8 @@ public void onComplete() { subscriber.onComplete(); } catch (IllegalStateException e) { // Client already disconnected and response closed - this is expected - // for streaming responses where client disconnect triggers background - // consumption. Log and ignore. + // for streaming responses where client disconnect closes ChildQueue. + // Log and ignore. if (e.getMessage() != null && e.getMessage().contains("Response has already been written")) { LOGGER.debug("Client disconnected before onComplete, response already closed for task {}", taskId.get()); } else { @@ -612,36 +608,21 @@ public void onComplete() { } } } - - private void startBackgroundConsumption() { - if (backgroundConsumeStarted.compareAndSet(false, true)) { - LOGGER.debug("Starting background consumption for task {}", taskId.get()); - // Client disconnected: continue consuming and persisting events in background - CompletableFuture bgTask = CompletableFuture.runAsync(() -> { - try { - LOGGER.debug("Background consumption thread started for task {}", taskId.get()); - 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); - } else { - LOGGER.debug("Background consumption already started for task {}", taskId.get()); - } - } }); }; } finally { - LOGGER.debug("onMessageSendStream FINALLY - task: {}; runningAgents: {}; backgroundTasks: {}", - taskId.get(), runningAgents.size(), backgroundTasks.size()); + // Needed to satisfy Nullaway + String idOfTask = taskId.get(); + if (idOfTask != null) { + LOGGER.debug("onMessageSendStream FINALLY - task: {}; runningAgents: {}", + idOfTask, runningAgents.size()); - // Remove agent from map immediately to prevent accumulation - CompletableFuture agentFuture = runningAgents.remove(taskId.get()); - LOGGER.debug("Removed agent for task {} from runningAgents in finally block, size after: {}", taskId.get(), runningAgents.size()); + // Remove agent from map immediately to prevent accumulation + CompletableFuture agentFuture = runningAgents.remove(idOfTask); + LOGGER.debug("Removed agent for task {} from runningAgents in finally block, size after: {}", taskId.get(), runningAgents.size()); - trackBackgroundTask(cleanupProducer(agentFuture, null, Objects.requireNonNull(taskId.get()), queue, true)); + cleanupProducer(agentFuture, null, idOfTask, queue, true); + } } } @@ -799,47 +780,6 @@ public void run() { return runnable; } - private void trackBackgroundTask(CompletableFuture task) { - backgroundTasks.add(task); - LOGGER.debug("Tracking background task (total: {}): {}", backgroundTasks.size(), task); - - task.whenComplete((result, throwable) -> { - try { - if (throwable != null) { - // Unwrap CompletionException to check for CancellationException - Throwable cause = throwable; - if (throwable instanceof java.util.concurrent.CompletionException && throwable.getCause() != null) { - cause = throwable.getCause(); - } - - if (cause instanceof java.util.concurrent.CancellationException) { - LOGGER.debug("Background task cancelled: {}", task); - } else { - LOGGER.error("Background task failed", throwable); - } - } - } finally { - backgroundTasks.remove(task); - LOGGER.debug("Removed background task (remaining: {}): {}", backgroundTasks.size(), task); - } - }); - } - - /** - * Wait for all background tasks to complete. - * Useful for testing to ensure cleanup completes before assertions. - * - * @return CompletableFuture that completes when all background tasks finish - */ - public CompletableFuture waitForBackgroundTasks() { - CompletableFuture[] tasks = backgroundTasks.toArray(new CompletableFuture[0]); - if (tasks.length == 0) { - return CompletableFuture.completedFuture(null); - } - LOGGER.debug("Waiting for {} background tasks to complete", tasks.length); - return CompletableFuture.allOf(tasks); - } - private CompletableFuture cleanupProducer(@Nullable CompletableFuture agentFuture, @Nullable CompletableFuture consumptionFuture, String taskId, EventQueue queue, boolean isStreaming) { LOGGER.debug("Starting cleanup for task {} (streaming={})", taskId, isStreaming); logThreadStats("CLEANUP START"); @@ -941,7 +881,6 @@ private void logThreadStats(String label) { LOGGER.debug("=== THREAD STATS: {} ===", label); LOGGER.debug("Active threads: {}", activeThreads); LOGGER.debug("Running agents: {}", runningAgents.size()); - LOGGER.debug("Background tasks: {}", backgroundTasks.size()); LOGGER.debug("Queue manager active queues: {}", queueManager.getClass().getSimpleName()); // List running agents @@ -952,13 +891,6 @@ private void logThreadStats(String label) { ); } - // List background tasks - if (!backgroundTasks.isEmpty()) { - LOGGER.debug("Background tasks:"); - backgroundTasks.forEach(task -> - LOGGER.debug(" - {}: {}", task, task.isDone() ? "DONE" : "RUNNING") - ); - } LOGGER.debug("=== END THREAD STATS ==="); } From fdc17603712a888adba6ce0a1e900b10feb21b99 Mon Sep 17 00:00:00 2001 From: Kabir Khan Date: Tue, 9 Dec 2025 20:02:48 +0000 Subject: [PATCH 4/8] Gemini feedback and fixing tests --- .../core/ReplicatedQueueManagerTest.java | 199 ++++++++++++------ .../io/a2a/server/events/EventQueueUtil.java | 8 - .../java/io/a2a/server/events/EventQueue.java | 53 ++--- .../server/events/MainEventBusContext.java | 27 +-- .../server/events/MainEventBusProcessor.java | 111 +++++++--- .../events/MainEventBusProcessorCallback.java | 7 +- .../DefaultRequestHandler.java | 14 +- .../io/a2a/server/tasks/ResultAggregator.java | 4 - .../a2a/server/events/EventConsumerTest.java | 111 +++++++--- .../io/a2a/server/events/EventQueueTest.java | 94 ++++++--- .../io/a2a/server/events/EventQueueUtil.java | 105 ++------- .../events/InMemoryQueueManagerTest.java | 16 +- .../DefaultRequestHandlerTest.java | 107 +++++++--- .../server/tasks/ResultAggregatorTest.java | 51 ++++- .../io/a2a/server/tasks/TaskUpdaterTest.java | 25 ++- .../jsonrpc/handler/JSONRPCHandlerTest.java | 13 +- 16 files changed, 578 insertions(+), 367 deletions(-) diff --git a/extras/queue-manager-replicated/core/src/test/java/io/a2a/extras/queuemanager/replicated/core/ReplicatedQueueManagerTest.java b/extras/queue-manager-replicated/core/src/test/java/io/a2a/extras/queuemanager/replicated/core/ReplicatedQueueManagerTest.java index ee1e4b88f..a339be543 100644 --- a/extras/queue-manager-replicated/core/src/test/java/io/a2a/extras/queuemanager/replicated/core/ReplicatedQueueManagerTest.java +++ b/extras/queue-manager-replicated/core/src/test/java/io/a2a/extras/queuemanager/replicated/core/ReplicatedQueueManagerTest.java @@ -33,6 +33,7 @@ import io.a2a.spec.TaskState; import io.a2a.spec.TaskStatus; import io.a2a.spec.TaskStatusUpdateEvent; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -66,6 +67,61 @@ void setUp() { .build(); } + /** + * Helper to create a test event with the specified taskId. + * This ensures taskId consistency between queue creation and event creation. + */ + private TaskStatusUpdateEvent createEventForTask(String taskId) { + return TaskStatusUpdateEvent.builder() + .taskId(taskId) + .contextId("test-context") + .status(new TaskStatus(TaskState.SUBMITTED)) + .isFinal(false) + .build(); + } + + @AfterEach + void tearDown() { + if (mainEventBusProcessor != null) { + mainEventBusProcessor.setCallback(null); // Clear any test callbacks + EventQueueUtil.stop(mainEventBusProcessor); + } + mainEventBusProcessor = null; + mainEventBus = null; + queueManager = null; + } + + /** + * Helper to wait for MainEventBusProcessor to process an event. + * Replaces polling patterns with deterministic callback-based waiting. + * + * @param action the action that triggers event processing + * @throws InterruptedException if waiting is interrupted + * @throws AssertionError if processing doesn't complete within timeout + */ + private void waitForEventProcessing(Runnable action) throws InterruptedException { + CountDownLatch processingLatch = new CountDownLatch(1); + mainEventBusProcessor.setCallback(new io.a2a.server.events.MainEventBusProcessorCallback() { + @Override + public void onEventProcessed(String taskId, io.a2a.spec.Event event) { + processingLatch.countDown(); + } + + @Override + public void onTaskFinalized(String taskId) { + // Not needed for basic event processing wait + } + }); + + try { + action.run(); + assertTrue(processingLatch.await(5, TimeUnit.SECONDS), + "MainEventBusProcessor should have processed the event within timeout"); + } finally { + mainEventBusProcessor.setCallback(null); + } + } + @Test void testReplicationStrategyTriggeredOnNormalEnqueue() throws InterruptedException { CountingReplicationStrategy strategy = new CountingReplicationStrategy(); @@ -119,48 +175,45 @@ void testReplicationStrategyWithCountingImplementation() throws InterruptedExcep @Test void testReplicatedEventDeliveredToCorrectQueue() throws InterruptedException { String taskId = "test-task-4"; + TaskStatusUpdateEvent eventForTask = createEventForTask(taskId); // Use matching taskId EventQueue queue = queueManager.createOrTap(taskId); - ReplicatedEventQueueItem replicatedEvent = new ReplicatedEventQueueItem(taskId, testEvent); - queueManager.onReplicatedEvent(replicatedEvent); + ReplicatedEventQueueItem replicatedEvent = new ReplicatedEventQueueItem(taskId, eventForTask); - Event dequeuedEvent; - try { - dequeuedEvent = queue.dequeueEventItem(100).getEvent(); - } catch (EventQueueClosedException e) { - fail("Queue should not be closed"); - return; - } - assertEquals(testEvent, dequeuedEvent); + // Use callback to wait for event processing + EventQueueItem item = dequeueEventWithRetry(queue, () -> queueManager.onReplicatedEvent(replicatedEvent)); + assertNotNull(item, "Event should be available in queue"); + Event dequeuedEvent = item.getEvent(); + assertEquals(eventForTask, dequeuedEvent); } @Test void testReplicatedEventCreatesQueueIfNeeded() throws InterruptedException { String taskId = "non-existent-task"; + TaskStatusUpdateEvent eventForTask = createEventForTask(taskId); // Use matching taskId // Verify no queue exists initially assertNull(queueManager.get(taskId)); - ReplicatedEventQueueItem replicatedEvent = new ReplicatedEventQueueItem(taskId, testEvent); - - // Process the replicated event - assertDoesNotThrow(() -> queueManager.onReplicatedEvent(replicatedEvent)); - - // Verify that a queue was created and the event was enqueued - EventQueue queue = queueManager.get(taskId); - assertNotNull(queue, "Queue should be created when processing replicated event for non-existent task"); - - // Verify the event was enqueued by dequeuing it - // Need to tap() the MainQueue to get a ChildQueue for consumption - EventQueue childQueue = queue.tap(); - Event dequeuedEvent; - try { - dequeuedEvent = childQueue.dequeueEventItem(100).getEvent(); - } catch (EventQueueClosedException e) { - fail("Queue should not be closed"); - return; - } - assertEquals(testEvent, dequeuedEvent, "The replicated event should be enqueued in the newly created queue"); + // Create a ChildQueue BEFORE processing the replicated event + // This ensures the ChildQueue exists when MainEventBusProcessor distributes the event + EventQueue childQueue = queueManager.createOrTap(taskId); + assertNotNull(childQueue, "ChildQueue should be created"); + + // Verify MainQueue was created + EventQueue mainQueue = queueManager.get(taskId); + assertNotNull(mainQueue, "MainQueue should exist after createOrTap"); + + ReplicatedEventQueueItem replicatedEvent = new ReplicatedEventQueueItem(taskId, eventForTask); + + // Process the replicated event and wait for distribution + // Use callback to wait for event processing + EventQueueItem item = dequeueEventWithRetry(childQueue, () -> { + assertDoesNotThrow(() -> queueManager.onReplicatedEvent(replicatedEvent)); + }); + assertNotNull(item, "Event should be available in queue"); + Event dequeuedEvent = item.getEvent(); + assertEquals(eventForTask, dequeuedEvent, "The replicated event should be enqueued in the newly created queue"); } @Test @@ -340,29 +393,29 @@ void testReplicatedEventProcessedWhenTaskActive() throws InterruptedException { queueManager = new ReplicatedQueueManager(new CountingReplicationStrategy(), stateProvider, mainEventBus); String taskId = "active-task"; + TaskStatusUpdateEvent eventForTask = createEventForTask(taskId); // Use matching taskId // Verify no queue exists initially assertNull(queueManager.get(taskId)); - // Process a replicated event for an active task - ReplicatedEventQueueItem replicatedEvent = new ReplicatedEventQueueItem(taskId, testEvent); - queueManager.onReplicatedEvent(replicatedEvent); + // Create a ChildQueue BEFORE processing the replicated event + // This ensures the ChildQueue exists when MainEventBusProcessor distributes the event + EventQueue childQueue = queueManager.createOrTap(taskId); + assertNotNull(childQueue, "ChildQueue should be created"); - // Queue should be created and event should be enqueued - EventQueue queue = queueManager.get(taskId); - assertNotNull(queue, "Queue should be created for active task"); + // Verify MainQueue was created + EventQueue mainQueue = queueManager.get(taskId); + assertNotNull(mainQueue, "MainQueue should exist after createOrTap"); - // Verify the event was enqueued - // Need to tap() the MainQueue to get a ChildQueue for consumption - EventQueue childQueue = queue.tap(); - Event dequeuedEvent; - try { - dequeuedEvent = childQueue.dequeueEventItem(100).getEvent(); - } catch (EventQueueClosedException e) { - fail("Queue should not be closed"); - return; - } - assertEquals(testEvent, dequeuedEvent, "Event should be enqueued for active task"); + // Process a replicated event for an active task + ReplicatedEventQueueItem replicatedEvent = new ReplicatedEventQueueItem(taskId, eventForTask); + + // Verify the event was enqueued and distributed to our ChildQueue + // Use callback to wait for event processing + EventQueueItem item = dequeueEventWithRetry(childQueue, () -> queueManager.onReplicatedEvent(replicatedEvent)); + assertNotNull(item, "Event should be available in queue"); + Event dequeuedEvent = item.getEvent(); + assertEquals(eventForTask, dequeuedEvent, "Event should be enqueued for active task"); } @@ -474,36 +527,21 @@ void testQueueClosedEventJsonSerialization() throws Exception { @Test void testReplicatedQueueClosedEventTerminatesConsumer() throws InterruptedException { String taskId = "remote-close-test"; + TaskStatusUpdateEvent eventForTask = createEventForTask(taskId); // Use matching taskId EventQueue queue = queueManager.createOrTap(taskId); - // Enqueue a normal event - queue.enqueueEvent(testEvent); - // Simulate receiving QueueClosedEvent from remote node QueueClosedEvent closedEvent = new QueueClosedEvent(taskId); ReplicatedEventQueueItem replicatedClosedEvent = new ReplicatedEventQueueItem(taskId, closedEvent); - queueManager.onReplicatedEvent(replicatedClosedEvent); - // Dequeue the normal event first - EventQueueItem item1; - try { - item1 = queue.dequeueEventItem(100); - } catch (EventQueueClosedException e) { - fail("Should not throw on first dequeue"); - return; - } - assertNotNull(item1); - assertEquals(testEvent, item1.getEvent()); + // Dequeue the normal event first (use callback to wait for async processing) + EventQueueItem item1 = dequeueEventWithRetry(queue, () -> queue.enqueueEvent(eventForTask)); + assertNotNull(item1, "First event should be available"); + assertEquals(eventForTask, item1.getEvent()); - // Next dequeue should get the QueueClosedEvent - EventQueueItem item2; - try { - item2 = queue.dequeueEventItem(100); - } catch (EventQueueClosedException e) { - fail("Should not throw on second dequeue, should return the event"); - return; - } - assertNotNull(item2); + // Next dequeue should get the QueueClosedEvent (use callback to wait for async processing) + EventQueueItem item2 = dequeueEventWithRetry(queue, () -> queueManager.onReplicatedEvent(replicatedClosedEvent)); + assertNotNull(item2, "QueueClosedEvent should be available"); assertTrue(item2.getEvent() instanceof QueueClosedEvent, "Second event should be QueueClosedEvent"); } @@ -562,4 +600,25 @@ public void setActive(boolean active) { this.active = active; } } + + /** + * Helper method to dequeue an event after waiting for MainEventBusProcessor distribution. + * Uses callback-based waiting instead of polling for deterministic synchronization. + * + * @param queue the queue to dequeue from + * @param enqueueAction the action that enqueues the event (triggers event processing) + * @return the dequeued EventQueueItem, or null if queue is closed + */ + private EventQueueItem dequeueEventWithRetry(EventQueue queue, Runnable enqueueAction) throws InterruptedException { + // Wait for event to be processed and distributed + waitForEventProcessing(enqueueAction); + + // Event is now available, dequeue directly + try { + return queue.dequeueEventItem(100); + } catch (EventQueueClosedException e) { + // Queue closed, return null + return null; + } + } } \ No newline at end of file diff --git a/extras/queue-manager-replicated/core/src/test/java/io/a2a/server/events/EventQueueUtil.java b/extras/queue-manager-replicated/core/src/test/java/io/a2a/server/events/EventQueueUtil.java index e5131ad15..a91575aaa 100644 --- a/extras/queue-manager-replicated/core/src/test/java/io/a2a/server/events/EventQueueUtil.java +++ b/extras/queue-manager-replicated/core/src/test/java/io/a2a/server/events/EventQueueUtil.java @@ -1,14 +1,6 @@ package io.a2a.server.events; public class EventQueueUtil { - // Shared MainEventBus for all tests - ensures events are properly distributed - private static final MainEventBus TEST_EVENT_BUS = new MainEventBus(); - - // Since EventQueue.builder() is package protected, add a method to expose it - public static EventQueue.EventQueueBuilder getEventQueueBuilder() { - return EventQueue.builder(TEST_EVENT_BUS); - } - public static void start(MainEventBusProcessor processor) { processor.start(); } diff --git a/server-common/src/main/java/io/a2a/server/events/EventQueue.java b/server-common/src/main/java/io/a2a/server/events/EventQueue.java index d518ffd9e..4ed694c72 100644 --- a/server-common/src/main/java/io/a2a/server/events/EventQueue.java +++ b/server-common/src/main/java/io/a2a/server/events/EventQueue.java @@ -470,22 +470,29 @@ void childClosing(ChildQueue child, boolean immediate) { * Called by MainEventBusProcessor after TaskStore persistence. */ void distributeToChildren(EventQueueItem item) { - synchronized (children) { - int childCount = children.size(); - if (LOGGER.isDebugEnabled()) { - LOGGER.debug("MainQueue[{}]: Distributing event {} to {} children", - taskId, item.getEvent().getClass().getSimpleName(), childCount); - } - children.forEach(child -> { - LOGGER.debug("MainQueue[{}]: Enqueueing event {} to child queue", - taskId, item.getEvent().getClass().getSimpleName()); - child.internalEnqueueItem(item); - }); - if (LOGGER.isDebugEnabled()) { - LOGGER.debug("MainQueue[{}]: Completed distribution of {} to {} children", - taskId, item.getEvent().getClass().getSimpleName(), childCount); - } + int childCount = children.size(); + if (LOGGER.isDebugEnabled()) { + LOGGER.debug("MainQueue[{}]: Distributing event {} to {} children", + taskId, item.getEvent().getClass().getSimpleName(), childCount); } + children.forEach(child -> { + LOGGER.debug("MainQueue[{}]: Enqueueing event {} to child queue", + taskId, item.getEvent().getClass().getSimpleName()); + child.internalEnqueueItem(item); + }); + if (LOGGER.isDebugEnabled()) { + LOGGER.debug("MainQueue[{}]: Completed distribution of {} to {} children", + taskId, item.getEvent().getClass().getSimpleName(), childCount); + } + } + + /** + * Release the semaphore after event processing is complete. + * Called by MainEventBusProcessor in finally block to ensure release even on exceptions. + * Balances the acquire() in enqueueEvent() - protects MainEventBus throughput. + */ + void releaseSemaphore() { + semaphore.release(); } /** @@ -557,17 +564,12 @@ public void enqueueItem(EventQueueItem item) { private void internalEnqueueItem(EventQueueItem item) { // Internal method called by MainEventBusProcessor to add to local queue + // Note: Semaphore is managed by parent MainQueue (acquire/release), not ChildQueue Event event = item.getEvent(); if (isClosed()) { LOGGER.warn("ChildQueue is closed. Event will not be enqueued. {} {}", this, event); return; } - try { - semaphore.acquire(); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new RuntimeException("Unable to acquire the semaphore to enqueue the event", e); - } queue.add(item); LOGGER.debug("Enqueued event {} {}", event instanceof Throwable ? event.toString() : event, this); } @@ -585,7 +587,7 @@ public EventQueueItem dequeueEventItem(int waitMilliSeconds) throws EventQueueCl if (item != null) { Event event = item.getEvent(); LOGGER.debug("Dequeued event item (no wait) {} {}", this, event instanceof Throwable ? event.toString() : event); - semaphore.release(); + // Note: Semaphore is managed by parent MainQueue, not released here } return item; } @@ -595,7 +597,7 @@ public EventQueueItem dequeueEventItem(int waitMilliSeconds) throws EventQueueCl if (item != null) { Event event = item.getEvent(); LOGGER.debug("Dequeued event item (waiting) {} {}", this, event instanceof Throwable ? event.toString() : event); - semaphore.release(); + // Note: Semaphore is managed by parent MainQueue, not released here } else { LOGGER.trace("Dequeue timeout (null) from ChildQueue {}", System.identityHashCode(this)); } @@ -636,8 +638,11 @@ protected void doClose(boolean immediate) { super.doClose(immediate); // Sets closed flag if (immediate) { // Immediate close: clear pending events from local queue + int clearedCount = queue.size(); queue.clear(); - LOGGER.debug("Cleared ChildQueue for immediate close: {}", this); + // Release semaphore permits for cleared events to prevent deadlock + semaphore.release(clearedCount); + LOGGER.debug("Cleared {} events from ChildQueue for immediate close: {}", clearedCount, this); } // For graceful close, let the queue drain naturally through normal consumption } diff --git a/server-common/src/main/java/io/a2a/server/events/MainEventBusContext.java b/server-common/src/main/java/io/a2a/server/events/MainEventBusContext.java index caa288849..f8e5e03ec 100644 --- a/server-common/src/main/java/io/a2a/server/events/MainEventBusContext.java +++ b/server-common/src/main/java/io/a2a/server/events/MainEventBusContext.java @@ -2,27 +2,10 @@ import java.util.Objects; -class MainEventBusContext { - private final String taskId; - private final EventQueue eventQueue; - private final EventQueueItem eventQueueItem; - - public MainEventBusContext(String taskId, EventQueue eventQueue, EventQueueItem eventQueueItem) { - this.taskId = Objects.requireNonNull(taskId, "taskId cannot be null"); - this.eventQueue = Objects.requireNonNull(eventQueue, "eventQueue cannot be null"); - this.eventQueueItem = Objects.requireNonNull(eventQueueItem, "eventQueueItem cannot be null"); - } - - public String taskId() { - return taskId; - } - - public EventQueue eventQueue() { - return eventQueue; - } - - public EventQueueItem eventQueueItem() { - return eventQueueItem; +record MainEventBusContext(String taskId, EventQueue eventQueue, EventQueueItem eventQueueItem) { + MainEventBusContext { + Objects.requireNonNull(taskId, "taskId cannot be null"); + Objects.requireNonNull(eventQueue, "eventQueue cannot be null"); + Objects.requireNonNull(eventQueueItem, "eventQueueItem cannot be null"); } - } diff --git a/server-common/src/main/java/io/a2a/server/events/MainEventBusProcessor.java b/server-common/src/main/java/io/a2a/server/events/MainEventBusProcessor.java index 94c3d390e..0f3033b71 100644 --- a/server-common/src/main/java/io/a2a/server/events/MainEventBusProcessor.java +++ b/server-common/src/main/java/io/a2a/server/events/MainEventBusProcessor.java @@ -11,6 +11,7 @@ import io.a2a.server.tasks.TaskStore; import io.a2a.spec.A2AServerException; import io.a2a.spec.Event; +import io.a2a.spec.InternalError; import io.a2a.spec.Task; import io.a2a.spec.TaskArtifactUpdateEvent; import io.a2a.spec.TaskStatusUpdateEvent; @@ -45,7 +46,7 @@ public class MainEventBusProcessor implements Runnable { * Default is NOOP to avoid null checks in production code. * Tests can inject their own callback via setCallback(). */ - private static volatile MainEventBusProcessorCallback callback = MainEventBusProcessorCallback.NOOP; + private volatile MainEventBusProcessorCallback callback = MainEventBusProcessorCallback.NOOP; private final MainEventBus eventBus; @@ -72,14 +73,14 @@ public MainEventBusProcessor(MainEventBus eventBus, TaskStore taskStore, PushNot * * @param callback the callback to invoke during event processing, or null for NOOP */ - public static void setCallback(MainEventBusProcessorCallback callback) { - MainEventBusProcessor.callback = callback != null ? callback : MainEventBusProcessorCallback.NOOP; + public void setCallback(MainEventBusProcessorCallback callback) { + this.callback = callback != null ? callback : MainEventBusProcessorCallback.NOOP; } @PostConstruct void start() { processorThread = new Thread(this, "MainEventBusProcessor"); - processorThread.setDaemon(false); // Keep JVM alive + processorThread.setDaemon(true); // Allow JVM to exit even if this thread is running processorThread.start(); LOGGER.info("MainEventBusProcessor started"); } @@ -99,9 +100,13 @@ void stop() { if (processorThread != null) { processorThread.interrupt(); try { + long start = System.currentTimeMillis(); processorThread.join(5000); // Wait up to 5 seconds + long elapsed = System.currentTimeMillis() - start; + LOGGER.info("MainEventBusProcessor thread stopped in {}ms", elapsed); } catch (InterruptedException e) { Thread.currentThread().interrupt(); + LOGGER.warn("Interrupted while waiting for MainEventBusProcessor thread to stop"); } } LOGGER.info("MainEventBusProcessor stopped"); @@ -137,32 +142,67 @@ private void processEvent(MainEventBusContext context) { LOGGER.debug("MainEventBusProcessor: Processing event for task {}: {} (queue type: {})", taskId, event.getClass().getSimpleName(), eventQueue.getClass().getSimpleName()); - // Step 1: Update TaskStore FIRST (persistence before clients see it) - updateTaskStore(taskId, event); - - // Step 2: Send push notification AFTER persistence (ensures notification sees latest state) - sendPushNotification(taskId); - - // Step 3: Then distribute to ChildQueues (clients see it AFTER persistence + notification) - if (eventQueue instanceof EventQueue.MainQueue mainQueue) { - int childCount = mainQueue.getChildCount(); - LOGGER.debug("MainEventBusProcessor: Distributing event to {} children for task {}", childCount, taskId); - mainQueue.distributeToChildren(context.eventQueueItem()); - LOGGER.debug("MainEventBusProcessor: Distributed event {} to {} children for task {}", - event.getClass().getSimpleName(), childCount, taskId); - } else { - LOGGER.warn("MainEventBusProcessor: Expected MainQueue but got {} for task {}", - eventQueue.getClass().getSimpleName(), taskId); - } + Event eventToDistribute = null; + try { + // Step 1: Update TaskStore FIRST (persistence before clients see it) + // If this throws, we distribute an error to ensure "persist before client visibility" + + try { + updateTaskStore(taskId, event); + eventToDistribute = event; // Success - distribute original event + } catch (InternalError e) { + // Persistence failed - create error event to distribute instead + LOGGER.error("Failed to persist event for task {}, distributing error to clients", taskId, e); + String errorMessage = "Failed to persist event: " + e.getMessage(); + eventToDistribute = e; + } catch (Exception e) { + LOGGER.error("Failed to persist event for task {}, distributing error to clients", taskId, e); + String errorMessage = "Failed to persist event: " + e.getMessage(); + eventToDistribute = new InternalError(errorMessage); + } + + // Step 2: Send push notification AFTER successful persistence + if (eventToDistribute == event) { + sendPushNotification(taskId); + } + + // Step 3: Then distribute to ChildQueues (clients see either event or error AFTER persistence attempt) + if (eventQueue instanceof EventQueue.MainQueue mainQueue) { + int childCount = mainQueue.getChildCount(); + LOGGER.debug("MainEventBusProcessor: Distributing {} to {} children for task {}", + eventToDistribute.getClass().getSimpleName(), childCount, taskId); + // Create new EventQueueItem with the event to distribute (original or error) + EventQueueItem itemToDistribute = new LocalEventQueueItem(eventToDistribute); + mainQueue.distributeToChildren(itemToDistribute); + LOGGER.debug("MainEventBusProcessor: Distributed {} to {} children for task {}", + eventToDistribute.getClass().getSimpleName(), childCount, taskId); + } else { + LOGGER.warn("MainEventBusProcessor: Expected MainQueue but got {} for task {}", + eventQueue.getClass().getSimpleName(), taskId); + } - LOGGER.debug("MainEventBusProcessor: Completed processing event for task {}", taskId); + LOGGER.debug("MainEventBusProcessor: Completed processing event for task {}", taskId); - // Step 4: Notify callback after all processing is complete - callback.onEventProcessed(taskId, event); + } finally { + try { + // Step 4: Notify callback after all processing is complete + // Call callback with the distributed event (original or error) + if (eventToDistribute != null) { + callback.onEventProcessed(taskId, eventToDistribute); - // Step 5: If this is a final event, notify task finalization - if (isFinalEvent(event)) { - callback.onTaskFinalized(taskId); + // Step 5: If this is a final event, notify task finalization + // Only for successful persistence (not for errors) + if (eventToDistribute == event && isFinalEvent(event)) { + callback.onTaskFinalized(taskId); + } + } + } finally { + // ALWAYS release semaphore, even if processing fails + // Balances the acquire() in MainQueue.enqueueEvent() + if (eventQueue instanceof EventQueue.MainQueue mainQueue) { + mainQueue.releaseSemaphore(); + } + } } } @@ -173,8 +213,15 @@ private void processEvent(MainEventBusContext context) { * which handles all event types (Task, TaskStatusUpdateEvent, TaskArtifactUpdateEvent). * This leverages existing TaskManager logic for status updates, artifact appending, message history, etc. *

    + *

    + * If persistence fails, the exception is propagated to processEvent() which distributes an + * InternalError to clients instead of the original event, ensuring "persist before visibility". + * See Gemini's comment: https://github.com/a2aproject/a2a-java/pull/515#discussion_r2604621833 + *

    + * + * @throws InternalError if persistence fails */ - private void updateTaskStore(String taskId, Event event) { + private void updateTaskStore(String taskId, Event event) throws InternalError { try { // Extract contextId from event (all relevant events have it) String contextId = extractContextId(event); @@ -186,12 +233,14 @@ private void updateTaskStore(String taskId, Event event) { taskManager.process(event); LOGGER.debug("TaskStore updated via TaskManager.process() for task {}: {}", taskId, event.getClass().getSimpleName()); - } catch (A2AServerException e) { + } catch (InternalError e) { LOGGER.error("Error updating TaskStore via TaskManager for task {}", taskId, e); - // Don't rethrow - we still want to distribute to ChildQueues + // Rethrow to prevent distributing unpersisted event to clients + throw e; } catch (Exception e) { LOGGER.error("Unexpected error updating TaskStore for task {}", taskId, e); - // Don't rethrow - we still want to distribute to ChildQueues + // Rethrow to prevent distributing unpersisted event to clients + throw new InternalError("TaskStore persistence failed: " + e.getMessage()); } } diff --git a/server-common/src/main/java/io/a2a/server/events/MainEventBusProcessorCallback.java b/server-common/src/main/java/io/a2a/server/events/MainEventBusProcessorCallback.java index 379acac37..2757e85f9 100644 --- a/server-common/src/main/java/io/a2a/server/events/MainEventBusProcessorCallback.java +++ b/server-common/src/main/java/io/a2a/server/events/MainEventBusProcessorCallback.java @@ -12,10 +12,13 @@ * Usage in tests: * 
      * {@code
+ * @Inject
+ * MainEventBusProcessor processor;
+ *
  * @BeforeEach
  * void setUp() {
  *     CountDownLatch latch = new CountDownLatch(3);
- *     MainEventBusProcessor.setCallback(new MainEventBusProcessorCallback() {
+ *     processor.setCallback(new MainEventBusProcessorCallback() {
  *         public void onEventProcessed(String taskId, Event event) {
  *             latch.countDown();
  *         }
@@ -24,7 +27,7 @@
  *
  * @AfterEach
  * void tearDown() {
- *     MainEventBusProcessor.setCallback(null); // Reset to NOOP
+ *     processor.setCallback(null); // Reset to NOOP
  * }
  * }
  *
    diff --git a/server-common/src/main/java/io/a2a/server/requesthandlers/DefaultRequestHandler.java b/server-common/src/main/java/io/a2a/server/requesthandlers/DefaultRequestHandler.java index b80080309..f908c108e 100644 --- a/server-common/src/main/java/io/a2a/server/requesthandlers/DefaultRequestHandler.java +++ b/server-common/src/main/java/io/a2a/server/requesthandlers/DefaultRequestHandler.java @@ -480,7 +480,12 @@ public EventKind onMessageSend(MessageSendParams params, ServerCallContext conte // Cleanup as background task to avoid blocking Vert.x threads // Pass the consumption future to ensure cleanup waits for background consumption to complete - cleanupProducer(agentFuture, etai != null ? etai.consumptionFuture() : null, taskId, queue, false); + cleanupProducer(agentFuture, etai != null ? etai.consumptionFuture() : null, taskId, queue, false) + .whenComplete((res, err) -> { + if (err != null) { + LOGGER.error("Error during async cleanup for task {}", taskId, err); + } + }); } LOGGER.debug("Returning: {}", kind); @@ -621,7 +626,12 @@ public void onComplete() { CompletableFuture agentFuture = runningAgents.remove(idOfTask); LOGGER.debug("Removed agent for task {} from runningAgents in finally block, size after: {}", taskId.get(), runningAgents.size()); - cleanupProducer(agentFuture, null, idOfTask, queue, true); + cleanupProducer(agentFuture, null, idOfTask, queue, true) + .whenComplete((res, err) -> { + if (err != null) { + LOGGER.error("Error during async cleanup for streaming task {}", taskId.get(), err); + } + }); } } } diff --git a/server-common/src/main/java/io/a2a/server/tasks/ResultAggregator.java b/server-common/src/main/java/io/a2a/server/tasks/ResultAggregator.java index b97f4bbb7..72f6c36b3 100644 --- a/server-common/src/main/java/io/a2a/server/tasks/ResultAggregator.java +++ b/server-common/src/main/java/io/a2a/server/tasks/ResultAggregator.java @@ -133,7 +133,6 @@ public EventTypeAndInterrupt consumeAndBreakOnInterrupt(EventConsumer consumer, // Determine interrupt behavior boolean shouldInterrupt = false; - boolean continueInBackground = false; boolean isFinalEvent = (event instanceof Task task && task.status().state().isFinal()) || (event instanceof TaskStatusUpdateEvent tsue && tsue.isFinal()); boolean isAuthRequired = (event instanceof Task task && task.status().state() == TaskState.AUTH_REQUIRED) @@ -148,19 +147,16 @@ public EventTypeAndInterrupt consumeAndBreakOnInterrupt(EventConsumer consumer, // new request is expected in order for the agent to make progress, // so the agent should exit. shouldInterrupt = true; - continueInBackground = true; } else if (!blocking) { // For non-blocking calls, interrupt as soon as a task is available. shouldInterrupt = true; - continueInBackground = true; } 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 shouldInterrupt = true; - continueInBackground = true; if (LOGGER.isDebugEnabled()) { LOGGER.debug("Blocking call for task {}: {} event, returning with background consumption", taskIdForLogging(), isFinalEvent ? "final" : "non-final"); diff --git a/server-common/src/test/java/io/a2a/server/events/EventConsumerTest.java b/server-common/src/test/java/io/a2a/server/events/EventConsumerTest.java index 511e044ee..2b4468c76 100644 --- a/server-common/src/test/java/io/a2a/server/events/EventConsumerTest.java +++ b/server-common/src/test/java/io/a2a/server/events/EventConsumerTest.java @@ -16,6 +16,8 @@ import java.util.concurrent.atomic.AtomicReference; import io.a2a.jsonrpc.common.json.JsonProcessingException; +import io.a2a.server.tasks.InMemoryTaskStore; +import io.a2a.server.tasks.PushNotificationSender; import io.a2a.spec.A2AError; import io.a2a.spec.A2AServerException; import io.a2a.spec.Artifact; @@ -27,14 +29,19 @@ import io.a2a.spec.TaskStatus; import io.a2a.spec.TaskStatusUpdateEvent; import io.a2a.spec.TextPart; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; public class EventConsumerTest { + private static final PushNotificationSender NOOP_PUSHNOTIFICATION_SENDER = task -> {}; + private static final String TASK_ID = "123"; // Must match MINIMAL_TASK id + private EventQueue eventQueue; private EventConsumer eventConsumer; - + private MainEventBus mainEventBus; + private MainEventBusProcessor mainEventBusProcessor; private static final String MINIMAL_TASK = """ { @@ -54,10 +61,58 @@ public class EventConsumerTest { @BeforeEach public void init() { - eventQueue = EventQueueUtil.getEventQueueBuilder().build().tap(); + // Set up MainEventBus and processor for production-like test environment + InMemoryTaskStore taskStore = new InMemoryTaskStore(); + mainEventBus = new MainEventBus(); + mainEventBusProcessor = new MainEventBusProcessor(mainEventBus, taskStore, NOOP_PUSHNOTIFICATION_SENDER); + EventQueueUtil.start(mainEventBusProcessor); + + eventQueue = EventQueueUtil.getEventQueueBuilder(mainEventBus) + .taskId(TASK_ID) + .mainEventBus(mainEventBus) + .build().tap(); eventConsumer = new EventConsumer(eventQueue); } + @AfterEach + public void cleanup() { + if (mainEventBusProcessor != null) { + mainEventBusProcessor.setCallback(null); // Clear any test callbacks + EventQueueUtil.stop(mainEventBusProcessor); + } + } + + /** + * Helper to wait for MainEventBusProcessor to process an event. + * Replaces polling patterns with deterministic callback-based waiting. + * + * @param action the action that triggers event processing + * @throws InterruptedException if waiting is interrupted + * @throws AssertionError if processing doesn't complete within timeout + */ + private void waitForEventProcessing(Runnable action) throws InterruptedException { + CountDownLatch processingLatch = new CountDownLatch(1); + mainEventBusProcessor.setCallback(new MainEventBusProcessorCallback() { + @Override + public void onEventProcessed(String taskId, Event event) { + processingLatch.countDown(); + } + + @Override + public void onTaskFinalized(String taskId) { + // Not needed for basic event processing wait + } + }); + + try { + action.run(); + assertTrue(processingLatch.await(5, TimeUnit.SECONDS), + "MainEventBusProcessor should have processed the event within timeout"); + } finally { + mainEventBusProcessor.setCallback(null); + } + } + @Test public void testConsumeOneTaskEvent() throws Exception { Task event = fromJson(MINIMAL_TASK, Task.class); @@ -92,7 +147,7 @@ public void testConsumeAllMultipleEvents() throws JsonProcessingException { List events = List.of( fromJson(MINIMAL_TASK, Task.class), TaskArtifactUpdateEvent.builder() - .taskId("task-123") + .taskId(TASK_ID) .contextId("session-xyz") .artifact(Artifact.builder() .artifactId("11") @@ -100,7 +155,7 @@ public void testConsumeAllMultipleEvents() throws JsonProcessingException { .build()) .build(), TaskStatusUpdateEvent.builder() - .taskId("task-123") + .taskId(TASK_ID) .contextId("session-xyz") .status(new TaskStatus(TaskState.WORKING)) .isFinal(true) @@ -153,7 +208,7 @@ public void testConsumeUntilMessage() throws Exception { List events = List.of( fromJson(MINIMAL_TASK, Task.class), TaskArtifactUpdateEvent.builder() - .taskId("task-123") + .taskId(TASK_ID) .contextId("session-xyz") .artifact(Artifact.builder() .artifactId("11") @@ -161,7 +216,7 @@ public void testConsumeUntilMessage() throws Exception { .build()) .build(), TaskStatusUpdateEvent.builder() - .taskId("task-123") + .taskId(TASK_ID) .contextId("session-xyz") .status(new TaskStatus(TaskState.WORKING)) .isFinal(true) @@ -340,7 +395,9 @@ public void onComplete() { @Test public void testConsumeAllStopsOnQueueClosed() throws Exception { - EventQueue queue = EventQueueUtil.getEventQueueBuilder().build().tap(); + EventQueue queue = EventQueueUtil.getEventQueueBuilder(mainEventBus) + .mainEventBus(mainEventBus) + .build().tap(); EventConsumer consumer = new EventConsumer(queue); // Close the queue immediately @@ -386,20 +443,16 @@ public void onComplete() { @Test public void testConsumeAllHandlesQueueClosedException() throws Exception { - EventQueue queue = EventQueueUtil.getEventQueueBuilder().build().tap(); + EventQueue queue = EventQueueUtil.getEventQueueBuilder(mainEventBus) + .mainEventBus(mainEventBus) + .build().tap(); EventConsumer consumer = new EventConsumer(queue); // Add a message event (which will complete the stream) Event message = fromJson(MESSAGE_PAYLOAD, Message.class); - queue.enqueueEvent(message); - // Poll for event to arrive in ChildQueue (async MainEventBusProcessor distribution) - long startTime = System.currentTimeMillis(); - long timeout = 2000; - while (queue.size() == 0 && (System.currentTimeMillis() - startTime) < timeout) { - Thread.sleep(50); - } - assertTrue(queue.size() > 0, "Event should arrive in ChildQueue within timeout"); + // Use callback to wait for event processing + waitForEventProcessing(() -> queue.enqueueEvent(message)); // Close the queue before consuming queue.close(); @@ -444,11 +497,13 @@ public void onComplete() { @Test public void testConsumeAllTerminatesOnQueueClosedEvent() throws Exception { - EventQueue queue = EventQueueUtil.getEventQueueBuilder().build().tap(); + EventQueue queue = EventQueueUtil.getEventQueueBuilder(mainEventBus) + .mainEventBus(mainEventBus) + .build().tap(); EventConsumer consumer = new EventConsumer(queue); // Enqueue a QueueClosedEvent (poison pill) - QueueClosedEvent queueClosedEvent = new QueueClosedEvent("task-123"); + QueueClosedEvent queueClosedEvent = new QueueClosedEvent(TASK_ID); queue.enqueueEvent(queueClosedEvent); Flow.Publisher publisher = consumer.consumeAll(); @@ -493,20 +548,12 @@ public void onComplete() { } private void enqueueAndConsumeOneEvent(Event event) throws Exception { - eventQueue.enqueueEvent(event); - // Poll for event with 2-second timeout - long startTime = System.currentTimeMillis(); - long timeout = 2000; - Event result = null; - while (result == null && (System.currentTimeMillis() - startTime) < timeout) { - try { - result = eventConsumer.consumeOne(); - } catch (A2AServerException e) { - // Event not available yet, wait a bit and try again - Thread.sleep(50); - } - } - assertNotNull(result, "Event should arrive within timeout"); + // Use callback to wait for event processing + waitForEventProcessing(() -> eventQueue.enqueueEvent(event)); + + // Event is now available, consume it directly + Event result = eventConsumer.consumeOne(); + assertNotNull(result, "Event should be available"); assertSame(event, result); } diff --git a/server-common/src/test/java/io/a2a/server/events/EventQueueTest.java b/server-common/src/test/java/io/a2a/server/events/EventQueueTest.java index 6494eef87..8e808b7cc 100644 --- a/server-common/src/test/java/io/a2a/server/events/EventQueueTest.java +++ b/server-common/src/test/java/io/a2a/server/events/EventQueueTest.java @@ -11,6 +11,8 @@ import static org.junit.jupiter.api.Assertions.assertTrue; import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; import io.a2a.server.tasks.InMemoryTaskStore; import io.a2a.server.tasks.PushNotificationSender; @@ -35,6 +37,8 @@ public class EventQueueTest { private MainEventBus mainEventBus; private MainEventBusProcessor mainEventBusProcessor; + private static final String TASK_ID = "123"; // Must match MINIMAL_TASK id + private static final String MINIMAL_TASK = """ { "id": "123", @@ -61,8 +65,8 @@ public void init() { mainEventBusProcessor = new MainEventBusProcessor(mainEventBus, taskStore, NOOP_PUSHNOTIFICATION_SENDER); EventQueueUtil.start(mainEventBusProcessor); - eventQueue = EventQueueUtil.getEventQueueBuilder() - .taskId("test-task") + eventQueue = EventQueueUtil.getEventQueueBuilder(mainEventBus) + .taskId(TASK_ID) .mainEventBus(mainEventBus) .build().tap(); } @@ -70,6 +74,7 @@ public void init() { @AfterEach public void cleanup() { if (mainEventBusProcessor != null) { + mainEventBusProcessor.setCallback(null); // Clear any test callbacks EventQueueUtil.stop(mainEventBusProcessor); } } @@ -78,37 +83,67 @@ public void cleanup() { * Helper to create a queue with MainEventBus configured (for tests that need event distribution). */ private EventQueue createQueueWithEventBus(String taskId) { - return EventQueueUtil.getEventQueueBuilder() + return EventQueueUtil.getEventQueueBuilder(mainEventBus) .taskId(taskId) - .mainEventBus(mainEventBus) .build(); } + /** + * Helper to wait for MainEventBusProcessor to process an event. + * Replaces polling patterns with deterministic callback-based waiting. + * + * @param action the action that triggers event processing + * @throws InterruptedException if waiting is interrupted + * @throws AssertionError if processing doesn't complete within timeout + */ + private void waitForEventProcessing(Runnable action) throws InterruptedException { + CountDownLatch processingLatch = new CountDownLatch(1); + mainEventBusProcessor.setCallback(new io.a2a.server.events.MainEventBusProcessorCallback() { + @Override + public void onEventProcessed(String taskId, io.a2a.spec.Event event) { + processingLatch.countDown(); + } + + @Override + public void onTaskFinalized(String taskId) { + // Not needed for basic event processing wait + } + }); + + try { + action.run(); + assertTrue(processingLatch.await(5, TimeUnit.SECONDS), + "MainEventBusProcessor should have processed the event within timeout"); + } finally { + mainEventBusProcessor.setCallback(null); + } + } + @Test public void testConstructorDefaultQueueSize() { - EventQueue queue = EventQueueUtil.getEventQueueBuilder().build(); + EventQueue queue = EventQueueUtil.getEventQueueBuilder(mainEventBus).build(); assertEquals(EventQueue.DEFAULT_QUEUE_SIZE, queue.getQueueSize()); } @Test public void testConstructorCustomQueueSize() { int customSize = 500; - EventQueue queue = EventQueueUtil.getEventQueueBuilder().queueSize(customSize).build(); + EventQueue queue = EventQueueUtil.getEventQueueBuilder(mainEventBus).queueSize(customSize).build(); assertEquals(customSize, queue.getQueueSize()); } @Test public void testConstructorInvalidQueueSize() { // Test zero queue size - assertThrows(IllegalArgumentException.class, () -> EventQueueUtil.getEventQueueBuilder().queueSize(0).build()); + assertThrows(IllegalArgumentException.class, () -> EventQueueUtil.getEventQueueBuilder(mainEventBus).queueSize(0).build()); // Test negative queue size - assertThrows(IllegalArgumentException.class, () -> EventQueueUtil.getEventQueueBuilder().queueSize(-10).build()); + assertThrows(IllegalArgumentException.class, () -> EventQueueUtil.getEventQueueBuilder(mainEventBus).queueSize(-10).build()); } @Test public void testTapCreatesChildQueue() { - EventQueue parentQueue = EventQueueUtil.getEventQueueBuilder().build(); + EventQueue parentQueue = EventQueueUtil.getEventQueueBuilder(mainEventBus).build(); EventQueue childQueue = parentQueue.tap(); assertNotNull(childQueue); @@ -118,7 +153,7 @@ public void testTapCreatesChildQueue() { @Test public void testTapOnChildQueueThrowsException() { - EventQueue parentQueue = EventQueueUtil.getEventQueueBuilder().build(); + EventQueue parentQueue = EventQueueUtil.getEventQueueBuilder(mainEventBus).build(); EventQueue childQueue = parentQueue.tap(); assertThrows(IllegalStateException.class, () -> childQueue.tap()); @@ -126,7 +161,7 @@ public void testTapOnChildQueueThrowsException() { @Test public void testEnqueueEventPropagagesToChildren() throws Exception { - EventQueue mainQueue = createQueueWithEventBus("test-propagate"); + EventQueue mainQueue = createQueueWithEventBus(TASK_ID); EventQueue childQueue1 = mainQueue.tap(); EventQueue childQueue2 = mainQueue.tap(); @@ -144,7 +179,7 @@ public void testEnqueueEventPropagagesToChildren() throws Exception { @Test public void testMultipleChildQueuesReceiveEvents() throws Exception { - EventQueue mainQueue = createQueueWithEventBus("test-multiple"); + EventQueue mainQueue = createQueueWithEventBus(TASK_ID); EventQueue childQueue1 = mainQueue.tap(); EventQueue childQueue2 = mainQueue.tap(); EventQueue childQueue3 = mainQueue.tap(); @@ -169,7 +204,7 @@ public void testMultipleChildQueuesReceiveEvents() throws Exception { @Test public void testChildQueueDequeueIndependently() throws Exception { - EventQueue mainQueue = createQueueWithEventBus("test-independent"); + EventQueue mainQueue = createQueueWithEventBus(TASK_ID); EventQueue childQueue1 = mainQueue.tap(); EventQueue childQueue2 = mainQueue.tap(); EventQueue childQueue3 = mainQueue.tap(); @@ -193,7 +228,7 @@ public void testChildQueueDequeueIndependently() throws Exception { @Test public void testCloseImmediatePropagationToChildren() throws Exception { - EventQueue parentQueue = createQueueWithEventBus("test-close"); + EventQueue parentQueue = createQueueWithEventBus(TASK_ID); EventQueue childQueue = parentQueue.tap(); // Add events to both parent and child @@ -223,7 +258,9 @@ public void testCloseImmediatePropagationToChildren() throws Exception { @Test public void testEnqueueEventWhenClosed() throws Exception { - EventQueue mainQueue = EventQueueUtil.getEventQueueBuilder().build(); + EventQueue mainQueue = EventQueueUtil.getEventQueueBuilder(mainEventBus) + .taskId(TASK_ID) + .build(); EventQueue childQueue = mainQueue.tap(); Event event = fromJson(MINIMAL_TASK, Task.class); @@ -248,7 +285,7 @@ public void testEnqueueEventWhenClosed() throws Exception { @Test public void testDequeueEventWhenClosedAndEmpty() throws Exception { - EventQueue queue = EventQueueUtil.getEventQueueBuilder().build().tap(); + EventQueue queue = EventQueueUtil.getEventQueueBuilder(mainEventBus).build().tap(); queue.close(); assertTrue(queue.isClosed()); @@ -258,17 +295,16 @@ public void testDequeueEventWhenClosedAndEmpty() throws Exception { @Test public void testDequeueEventWhenClosedButHasEvents() throws Exception { - EventQueue mainQueue = EventQueueUtil.getEventQueueBuilder().build(); + EventQueue mainQueue = EventQueueUtil.getEventQueueBuilder(mainEventBus) + .taskId(TASK_ID) + .build(); EventQueue childQueue = mainQueue.tap(); Event event = fromJson(MINIMAL_TASK, Task.class); - // Enqueue to mainQueue - mainQueue.enqueueEvent(event); - - // Wait for event to arrive in childQueue (use peek-like behavior by dequeueing then re-checking) - // Actually, just wait a bit for async processing - Thread.sleep(100); // Give MainEventBusProcessor time to distribute event + // Use callback to wait for event processing instead of polling + waitForEventProcessing(() -> mainQueue.enqueueEvent(event)); + // At this point, event has been processed and distributed to childQueue childQueue.close(); // Graceful close - events should remain assertTrue(childQueue.isClosed()); @@ -309,7 +345,7 @@ public void testDequeueEventEmptyQueueNoWait() throws Exception { @Test public void testDequeueEventWait() throws Exception { Event event = TaskStatusUpdateEvent.builder() - .taskId("task-123") + .taskId(TASK_ID) .contextId("session-xyz") .status(new TaskStatus(TaskState.WORKING)) .isFinal(true) @@ -323,7 +359,7 @@ public void testDequeueEventWait() throws Exception { @Test public void testTaskDone() throws Exception { Event event = TaskArtifactUpdateEvent.builder() - .taskId("task-123") + .taskId(TASK_ID) .contextId("session-xyz") .artifact(Artifact.builder() .artifactId("11") @@ -399,7 +435,7 @@ public void testCloseIdempotent() throws Exception { assertTrue(eventQueue.isClosed()); // Test with immediate close as well - EventQueue eventQueue2 = EventQueueUtil.getEventQueueBuilder().build(); + EventQueue eventQueue2 = EventQueueUtil.getEventQueueBuilder(mainEventBus).build(); eventQueue2.close(true); assertTrue(eventQueue2.isClosed()); @@ -413,7 +449,7 @@ public void testCloseIdempotent() throws Exception { */ @Test public void testCloseChildQueues() throws Exception { - EventQueue mainQueue = EventQueueUtil.getEventQueueBuilder().build(); + EventQueue mainQueue = EventQueueUtil.getEventQueueBuilder(mainEventBus).build(); EventQueue childQueue = mainQueue.tap(); assertTrue(childQueue != null); @@ -423,7 +459,7 @@ public void testCloseChildQueues() throws Exception { assertFalse(childQueue.isClosed()); // Child NOT closed on graceful parent close // Immediate close - parent force-closes all children - EventQueue mainQueue2 = EventQueueUtil.getEventQueueBuilder().build(); + EventQueue mainQueue2 = EventQueueUtil.getEventQueueBuilder(mainEventBus).build(); EventQueue childQueue2 = mainQueue2.tap(); mainQueue2.close(true); // immediate=true assertTrue(mainQueue2.isClosed()); @@ -436,7 +472,7 @@ public void testCloseChildQueues() throws Exception { */ @Test public void testMainQueueReferenceCountingStaysOpenWithActiveChildren() throws Exception { - EventQueue mainQueue = EventQueueUtil.getEventQueueBuilder().build(); + EventQueue mainQueue = EventQueueUtil.getEventQueueBuilder(mainEventBus).build(); EventQueue child1 = mainQueue.tap(); EventQueue child2 = mainQueue.tap(); diff --git a/server-common/src/test/java/io/a2a/server/events/EventQueueUtil.java b/server-common/src/test/java/io/a2a/server/events/EventQueueUtil.java index beb8b3d5a..6c9ed4a17 100644 --- a/server-common/src/test/java/io/a2a/server/events/EventQueueUtil.java +++ b/server-common/src/test/java/io/a2a/server/events/EventQueueUtil.java @@ -1,102 +1,39 @@ package io.a2a.server.events; -import io.a2a.server.tasks.InMemoryTaskStore; -import io.a2a.server.tasks.PushNotificationSender; -import io.a2a.server.tasks.TaskStateProvider; import java.util.concurrent.atomic.AtomicInteger; public class EventQueueUtil { - // Shared MainEventBus for all tests (to avoid creating one per test) - private static final MainEventBus TEST_EVENT_BUS = new MainEventBus(); - - // Shared MainEventBusProcessor for all tests (automatically processes events) - private static final MainEventBusProcessor TEST_PROCESSOR; - - static { - // Initialize and start the processor once for all tests - InMemoryTaskStore testTaskStore = new InMemoryTaskStore(); - PushNotificationSender testPushSender = taskId -> {}; // No-op for tests - TEST_PROCESSOR = new MainEventBusProcessor(TEST_EVENT_BUS, testTaskStore, testPushSender); - TEST_PROCESSOR.start(); // Start background thread - - // Register shutdown hook to stop processor - Runtime.getRuntime().addShutdownHook(new Thread(() -> TEST_PROCESSOR.stop())); - } - // Counter for generating unique test taskIds private static final AtomicInteger TASK_ID_COUNTER = new AtomicInteger(0); - // Since EventQueue.builder() is package protected, add a method to expose it - // Note: Now includes MainEventBus requirement and default taskId - // Returns MainQueue - tests should call .tap() if they need to consume events - public static EventQueue.EventQueueBuilder getEventQueueBuilder() { - return new EventQueueBuilderWrapper( - EventQueue.builder(TEST_EVENT_BUS) - .taskId("test-task-" + TASK_ID_COUNTER.incrementAndGet()) - ); - } - - // Get the shared test MainEventBus instance - public static MainEventBus getTestEventBus() { - return TEST_EVENT_BUS; + /** + * Get an EventQueue builder pre-configured with the shared test MainEventBus and a unique taskId. + *

    + * Note: Returns MainQueue - tests should call .tap() if they need to consume events. + *

    + * + * @return builder with TEST_EVENT_BUS and unique taskId already set + */ + public static EventQueue.EventQueueBuilder getEventQueueBuilder(MainEventBus eventBus) { + return EventQueue.builder(eventBus) + .taskId("test-task-" + TASK_ID_COUNTER.incrementAndGet()); } + /** + * Start a MainEventBusProcessor instance. + * + * @param processor the processor to start + */ public static void start(MainEventBusProcessor processor) { processor.start(); } + /** + * Stop a MainEventBusProcessor instance. + * + * @param processor the processor to stop + */ public static void stop(MainEventBusProcessor processor) { processor.stop(); } - - // Wrapper that delegates to actual builder - private static class EventQueueBuilderWrapper extends EventQueue.EventQueueBuilder { - private final EventQueue.EventQueueBuilder delegate; - - EventQueueBuilderWrapper(EventQueue.EventQueueBuilder delegate) { - this.delegate = delegate; - } - - @Override - public EventQueue.EventQueueBuilder queueSize(int queueSize) { - delegate.queueSize(queueSize); - return this; - } - - @Override - public EventQueue.EventQueueBuilder hook(EventEnqueueHook hook) { - delegate.hook(hook); - return this; - } - - @Override - public EventQueue.EventQueueBuilder taskId(String taskId) { - delegate.taskId(taskId); - return this; - } - - @Override - public EventQueue.EventQueueBuilder addOnCloseCallback(Runnable onCloseCallback) { - delegate.addOnCloseCallback(onCloseCallback); - return this; - } - - @Override - public EventQueue.EventQueueBuilder taskStateProvider(TaskStateProvider taskStateProvider) { - delegate.taskStateProvider(taskStateProvider); - return this; - } - - @Override - public EventQueue.EventQueueBuilder mainEventBus(MainEventBus mainEventBus) { - delegate.mainEventBus(mainEventBus); - return this; - } - - @Override - public EventQueue build() { - // Return MainQueue directly - tests should call .tap() if they need ChildQueue - return delegate.build(); - } - } } diff --git a/server-common/src/test/java/io/a2a/server/events/InMemoryQueueManagerTest.java b/server-common/src/test/java/io/a2a/server/events/InMemoryQueueManagerTest.java index 46f97b510..808a1107a 100644 --- a/server-common/src/test/java/io/a2a/server/events/InMemoryQueueManagerTest.java +++ b/server-common/src/test/java/io/a2a/server/events/InMemoryQueueManagerTest.java @@ -49,7 +49,7 @@ public void tearDown() { @Test public void testAddNewQueue() { String taskId = "test_task_id"; - EventQueue queue = EventQueueUtil.getEventQueueBuilder().build(); + EventQueue queue = EventQueueUtil.getEventQueueBuilder(mainEventBus).build(); queueManager.add(taskId, queue); @@ -60,8 +60,8 @@ public void testAddNewQueue() { @Test public void testAddExistingQueueThrowsException() { String taskId = "test_task_id"; - EventQueue queue1 = EventQueueUtil.getEventQueueBuilder().build(); - EventQueue queue2 = EventQueueUtil.getEventQueueBuilder().build(); + EventQueue queue1 = EventQueueUtil.getEventQueueBuilder(mainEventBus).build(); + EventQueue queue2 = EventQueueUtil.getEventQueueBuilder(mainEventBus).build(); queueManager.add(taskId, queue1); @@ -73,7 +73,7 @@ public void testAddExistingQueueThrowsException() { @Test public void testGetExistingQueue() { String taskId = "test_task_id"; - EventQueue queue = EventQueueUtil.getEventQueueBuilder().build(); + EventQueue queue = EventQueueUtil.getEventQueueBuilder(mainEventBus).build(); queueManager.add(taskId, queue); EventQueue result = queueManager.get(taskId); @@ -90,7 +90,7 @@ public void testGetNonexistentQueue() { @Test public void testTapExistingQueue() { String taskId = "test_task_id"; - EventQueue queue = EventQueueUtil.getEventQueueBuilder().build(); + EventQueue queue = EventQueueUtil.getEventQueueBuilder(mainEventBus).build(); queueManager.add(taskId, queue); EventQueue tappedQueue = queueManager.tap(taskId); @@ -111,7 +111,7 @@ public void testTapNonexistentQueue() { @Test public void testCloseExistingQueue() { String taskId = "test_task_id"; - EventQueue queue = EventQueueUtil.getEventQueueBuilder().build(); + EventQueue queue = EventQueueUtil.getEventQueueBuilder(mainEventBus).build(); queueManager.add(taskId, queue); queueManager.close(taskId); @@ -146,7 +146,7 @@ public void testCreateOrTapNewQueue() { @Test public void testCreateOrTapExistingQueue() { String taskId = "test_task_id"; - EventQueue originalQueue = EventQueueUtil.getEventQueueBuilder().build(); + EventQueue originalQueue = EventQueueUtil.getEventQueueBuilder(mainEventBus).build(); queueManager.add(taskId, originalQueue); EventQueue result = queueManager.createOrTap(taskId); @@ -168,7 +168,7 @@ public void testConcurrentOperations() throws InterruptedException, ExecutionExc // Add tasks concurrently List> addFutures = taskIds.stream() .map(taskId -> CompletableFuture.supplyAsync(() -> { - EventQueue queue = EventQueueUtil.getEventQueueBuilder().build(); + EventQueue queue = EventQueueUtil.getEventQueueBuilder(mainEventBus).build(); queueManager.add(taskId, queue); return taskId; })) diff --git a/server-common/src/test/java/io/a2a/server/requesthandlers/DefaultRequestHandlerTest.java b/server-common/src/test/java/io/a2a/server/requesthandlers/DefaultRequestHandlerTest.java index 662ddcab4..0f9aa1868 100644 --- a/server-common/src/test/java/io/a2a/server/requesthandlers/DefaultRequestHandlerTest.java +++ b/server-common/src/test/java/io/a2a/server/requesthandlers/DefaultRequestHandlerTest.java @@ -89,10 +89,42 @@ void setUp() { void tearDown() { // Stop MainEventBusProcessor background thread if (mainEventBusProcessor != null) { + mainEventBusProcessor.setCallback(null); // Clear any test callbacks EventQueueUtil.stop(mainEventBusProcessor); } } + /** + * Helper to wait for MainEventBusProcessor to finalize a task. + * Replaces polling patterns with deterministic callback-based waiting. + * + * @param action the action that triggers task finalization (e.g., enqueuing a final event) + * @throws InterruptedException if waiting is interrupted + * @throws AssertionError if finalization doesn't complete within timeout + */ + private void waitForTaskFinalization(Runnable action) throws InterruptedException { + CountDownLatch finalizationLatch = new CountDownLatch(1); + mainEventBusProcessor.setCallback(new io.a2a.server.events.MainEventBusProcessorCallback() { + @Override + public void onEventProcessed(String taskId, io.a2a.spec.Event event) { + // Not used for task finalization wait + } + + @Override + public void onTaskFinalized(String taskId) { + finalizationLatch.countDown(); + } + }); + + try { + action.run(); + assertTrue(finalizationLatch.await(5, TimeUnit.SECONDS), + "MainEventBusProcessor should have finalized the task within timeout"); + } finally { + mainEventBusProcessor.setCallback(null); + } + } + /** * Test that multiple blocking messages to the same task work correctly * when agent doesn't emit final events (fire-and-forget pattern). @@ -599,32 +631,15 @@ void testNonBlockingMessagePersistsAllEventsInBackground() throws Exception { // At this point, the non-blocking call has returned, but the agent is still running - // Allow the agent to emit the final COMPLETED event - allowCompletion.countDown(); - - // Assertion 2: Poll for the final task state to be persisted in background - // Use polling loop instead of fixed sleep for faster and more reliable test - long timeoutMs = 5000; - long startTime = System.currentTimeMillis(); - Task persistedTask = null; - boolean completedStateFound = false; - - while (System.currentTimeMillis() - startTime < timeoutMs) { - persistedTask = taskStore.get(taskId); - if (persistedTask != null && persistedTask.status().state() == TaskState.COMPLETED) { - completedStateFound = true; - break; - } - Thread.sleep(100); // Poll every 100ms - } + // Assertion 2: Wait for the final task to be processed and finalized in background + // Use callback to wait for task finalization instead of polling + waitForTaskFinalization(() -> allowCompletion.countDown()); - assertTrue(persistedTask != null, "Task should be persisted to store"); - assertTrue( - completedStateFound, - "Final task state should be COMPLETED (background consumption should have processed it), got: " + - (persistedTask != null ? persistedTask.status().state() : "null") + - " after " + (System.currentTimeMillis() - startTime) + "ms" - ); + // Verify the task was persisted with COMPLETED state + Task persistedTask = taskStore.get(taskId); + assertNotNull(persistedTask, "Task should be persisted to store"); + assertEquals(TaskState.COMPLETED, persistedTask.status().state(), + "Final task state should be COMPLETED (background consumption should have processed it)"); } /** @@ -801,16 +816,42 @@ void testBlockingCallReturnsCompleteTaskWithArtifacts() throws Exception { updater.complete(); }); - // Call blocking onMessageSend - should wait for ALL events - Object result = requestHandler.onMessageSend(params, serverCallContext); + // Use callback to ensure task finalization is complete before checking TaskStore + // This ensures MainEventBusProcessor has finished persisting the final state + CountDownLatch finalizationLatch = new CountDownLatch(1); + mainEventBusProcessor.setCallback(new io.a2a.server.events.MainEventBusProcessorCallback() { + @Override + public void onEventProcessed(String taskId, io.a2a.spec.Event event) { + // Not used + } - // The returned result should be a Task with ALL artifacts - assertTrue(result instanceof Task, "Result should be a Task"); - Task returnedTask = (Task) result; + @Override + public void onTaskFinalized(String taskId) { + finalizationLatch.countDown(); + } + }); + + Task returnedTask; + try { + // Call blocking onMessageSend - should wait for ALL events + Object result = requestHandler.onMessageSend(params, serverCallContext); + + // Wait for finalization callback to ensure TaskStore is fully updated + assertTrue(finalizationLatch.await(5, TimeUnit.SECONDS), + "Task should be finalized within timeout"); - // Verify task is completed - assertEquals(TaskState.COMPLETED, returnedTask.status().state(), - "Returned task should be COMPLETED"); + // The returned result should be a Task with ALL artifacts + assertTrue(result instanceof Task, "Result should be a Task"); + returnedTask = (Task) result; + + // Fetch final state from TaskStore (guaranteed to be persisted after callback) + returnedTask = taskStore.get(taskId); + + assertEquals(TaskState.COMPLETED, returnedTask.status().state(), + "Returned task should be COMPLETED"); + } finally { + mainEventBusProcessor.setCallback(null); + } // Verify artifacts are included in the returned task assertNotNull(returnedTask.artifacts(), diff --git a/server-common/src/test/java/io/a2a/server/tasks/ResultAggregatorTest.java b/server-common/src/test/java/io/a2a/server/tasks/ResultAggregatorTest.java index 2c60cba50..e087a8ff2 100644 --- a/server-common/src/test/java/io/a2a/server/tasks/ResultAggregatorTest.java +++ b/server-common/src/test/java/io/a2a/server/tasks/ResultAggregatorTest.java @@ -11,8 +11,10 @@ import static org.mockito.Mockito.when; import java.util.Collections; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.Executor; import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; import io.a2a.server.events.EventConsumer; import io.a2a.server.events.EventQueue; @@ -20,12 +22,14 @@ import io.a2a.server.events.InMemoryQueueManager; import io.a2a.server.events.MainEventBus; import io.a2a.server.events.MainEventBusProcessor; +import io.a2a.spec.Event; import io.a2a.spec.EventKind; import io.a2a.spec.Message; import io.a2a.spec.Task; import io.a2a.spec.TaskState; import io.a2a.spec.TaskStatus; import io.a2a.spec.TextPart; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.mockito.Mock; @@ -72,6 +76,38 @@ private Task createSampleTask(String taskId, TaskState statusState, String conte .build(); } + /** + * Helper to wait for MainEventBusProcessor to process an event. + * Replaces polling patterns with deterministic callback-based waiting. + * + * @param processor the processor to set callback on + * @param action the action that triggers event processing + * @throws InterruptedException if waiting is interrupted + * @throws AssertionError if processing doesn't complete within timeout + */ + private void waitForEventProcessing(MainEventBusProcessor processor, Runnable action) throws InterruptedException { + CountDownLatch processingLatch = new CountDownLatch(1); + processor.setCallback(new io.a2a.server.events.MainEventBusProcessorCallback() { + @Override + public void onEventProcessed(String taskId, Event event) { + processingLatch.countDown(); + } + + @Override + public void onTaskFinalized(String taskId) { + // Not needed for basic event processing wait + } + }); + + try { + action.run(); + assertTrue(processingLatch.await(5, TimeUnit.SECONDS), + "MainEventBusProcessor should have processed the event within timeout"); + } finally { + processor.setCallback(null); + } + } + // Basic functionality tests @@ -200,7 +236,8 @@ void testGetCurrentResultWithMessageTakesPrecedence() { @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"); + String taskId = "test-task"; + Task firstEvent = createSampleTask(taskId, TaskState.WORKING, "ctx1"); // After processing firstEvent, the current result will be that task when(mockTaskManager.getTask()).thenReturn(firstEvent); @@ -214,16 +251,10 @@ void testConsumeAndBreakNonBlocking() throws Exception { InMemoryQueueManager queueManager = new InMemoryQueueManager(new MockTaskStateProvider(), mainEventBus); - EventQueue queue = queueManager.getEventQueueBuilder("test-task").build().tap(); - queue.enqueueEvent(firstEvent); + EventQueue queue = queueManager.getEventQueueBuilder(taskId).build().tap(); - // Poll for event to arrive in ChildQueue (async MainEventBusProcessor distribution) - long startTime = System.currentTimeMillis(); - long timeout = 2000; - while (queue.size() == 0 && (System.currentTimeMillis() - startTime) < timeout) { - Thread.sleep(50); - } - assertTrue(queue.size() > 0, "Event should arrive in ChildQueue within timeout"); + // Use callback to wait for event processing (replaces polling) + waitForEventProcessing(processor, () -> queue.enqueueEvent(firstEvent)); // Create real EventConsumer with the queue EventConsumer eventConsumer = diff --git a/server-common/src/test/java/io/a2a/server/tasks/TaskUpdaterTest.java b/server-common/src/test/java/io/a2a/server/tasks/TaskUpdaterTest.java index c36ac7efd..fd195e0a5 100644 --- a/server-common/src/test/java/io/a2a/server/tasks/TaskUpdaterTest.java +++ b/server-common/src/test/java/io/a2a/server/tasks/TaskUpdaterTest.java @@ -16,6 +16,8 @@ import io.a2a.server.events.EventQueue; import io.a2a.server.events.EventQueueItem; import io.a2a.server.events.EventQueueUtil; +import io.a2a.server.events.MainEventBus; +import io.a2a.server.events.MainEventBusProcessor; import io.a2a.spec.Event; import io.a2a.spec.Message; import io.a2a.spec.Part; @@ -23,6 +25,7 @@ import io.a2a.spec.TaskState; import io.a2a.spec.TaskStatusUpdateEvent; import io.a2a.spec.TextPart; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -39,14 +42,27 @@ public class TaskUpdaterTest { private static final List> SAMPLE_PARTS = List.of(new TextPart("Test message")); + private static final PushNotificationSender NOOP_PUSHNOTIFICATION_SENDER = task -> {}; + EventQueue eventQueue; + private MainEventBus mainEventBus; + private MainEventBusProcessor mainEventBusProcessor; private TaskUpdater taskUpdater; @BeforeEach public void init() { - eventQueue = EventQueueUtil.getEventQueueBuilder().build().tap(); + // Set up MainEventBus and processor for production-like test environment + InMemoryTaskStore taskStore = new InMemoryTaskStore(); + mainEventBus = new MainEventBus(); + mainEventBusProcessor = new MainEventBusProcessor(mainEventBus, taskStore, NOOP_PUSHNOTIFICATION_SENDER); + EventQueueUtil.start(mainEventBusProcessor); + + eventQueue = EventQueueUtil.getEventQueueBuilder(mainEventBus) + .taskId(TEST_TASK_ID) + .mainEventBus(mainEventBus) + .build().tap(); RequestContext context = new RequestContext.Builder() .setTaskId(TEST_TASK_ID) .setContextId(TEST_TASK_CONTEXT_ID) @@ -54,6 +70,13 @@ public void init() { taskUpdater = new TaskUpdater(context, eventQueue); } + @AfterEach + public void cleanup() { + if (mainEventBusProcessor != null) { + EventQueueUtil.stop(mainEventBusProcessor); + } + } + @Test public void testAddArtifactWithCustomIdAndName() throws Exception { taskUpdater.addArtifact(SAMPLE_PARTS, "custom-artifact-id", "Custom Artifact", null); diff --git a/transport/jsonrpc/src/test/java/io/a2a/transport/jsonrpc/handler/JSONRPCHandlerTest.java b/transport/jsonrpc/src/test/java/io/a2a/transport/jsonrpc/handler/JSONRPCHandlerTest.java index 3b979666c..0e3e42b05 100644 --- a/transport/jsonrpc/src/test/java/io/a2a/transport/jsonrpc/handler/JSONRPCHandlerTest.java +++ b/transport/jsonrpc/src/test/java/io/a2a/transport/jsonrpc/handler/JSONRPCHandlerTest.java @@ -40,7 +40,6 @@ import io.a2a.server.ServerCallContext; import io.a2a.server.auth.UnauthenticatedUser; import io.a2a.server.events.EventConsumer; -import io.a2a.server.events.MainEventBusProcessor; import io.a2a.server.events.MainEventBusProcessorCallback; import io.a2a.server.requesthandlers.AbstractA2ARequestHandlerTest; import io.a2a.server.requesthandlers.DefaultRequestHandler; @@ -294,7 +293,7 @@ public void onComplete() { public void testOnMessageStreamNewMessageMultipleEventsSuccess() throws InterruptedException { // Setup callback to wait for all 3 events to be processed by MainEventBusProcessor CountDownLatch processingLatch = new CountDownLatch(3); - MainEventBusProcessor.setCallback(new MainEventBusProcessorCallback() { + mainEventBusProcessor.setCallback(new MainEventBusProcessorCallback() { @Override public void onEventProcessed(String taskId, Event event) { processingLatch.countDown(); @@ -416,7 +415,7 @@ public void onComplete() { assertEquals(MINIMAL_TASK.id(), receivedStatus.taskId()); assertEquals(TaskState.COMPLETED, receivedStatus.status().state()); } finally { - MainEventBusProcessor.setCallback(null); + mainEventBusProcessor.setCallback(null); } } @@ -693,7 +692,7 @@ public void testGetPushNotificationConfigSuccess() { public void testOnMessageStreamNewMessageSendPushNotificationSuccess() throws Exception { // Setup callback to wait for all 3 events to be processed by MainEventBusProcessor CountDownLatch processingLatch = new CountDownLatch(3); - MainEventBusProcessor.setCallback(new MainEventBusProcessorCallback() { + mainEventBusProcessor.setCallback(new MainEventBusProcessorCallback() { @Override public void onEventProcessed(String taskId, Event event) { processingLatch.countDown(); @@ -813,7 +812,7 @@ public void onComplete() { assertEquals(1, curr.artifacts().get(0).parts().size()); assertEquals("text", ((TextPart) curr.artifacts().get(0).parts().get(0)).text()); } finally { - MainEventBusProcessor.setCallback(null); + mainEventBusProcessor.setCallback(null); } } @@ -1282,7 +1281,7 @@ public void testOnMessageSendTaskIdMismatch() { public void testOnMessageStreamTaskIdMismatch() throws InterruptedException { // Setup callback to wait for the 1 event to be processed by MainEventBusProcessor CountDownLatch processingLatch = new CountDownLatch(1); - MainEventBusProcessor.setCallback(new MainEventBusProcessorCallback() { + mainEventBusProcessor.setCallback(new MainEventBusProcessorCallback() { @Override public void onEventProcessed(String taskId, Event event) { processingLatch.countDown(); @@ -1348,7 +1347,7 @@ public void onComplete() { Assertions.assertEquals(1, results.size()); Assertions.assertInstanceOf(InternalError.class, results.get(0).getError()); } finally { - MainEventBusProcessor.setCallback(null); + mainEventBusProcessor.setCallback(null); } } From 29f0b971e59cdcefabe5d2b58d76e54dc40826ca Mon Sep 17 00:00:00 2001 From: Kabir Khan Date: Wed, 10 Dec 2025 14:11:35 +0000 Subject: [PATCH 5/8] Move sepaphore into MainQueue --- .../java/io/a2a/server/events/EventQueue.java | 26 ++++++------------- 1 file changed, 8 insertions(+), 18 deletions(-) diff --git a/server-common/src/main/java/io/a2a/server/events/EventQueue.java b/server-common/src/main/java/io/a2a/server/events/EventQueue.java index 4ed694c72..73e973067 100644 --- a/server-common/src/main/java/io/a2a/server/events/EventQueue.java +++ b/server-common/src/main/java/io/a2a/server/events/EventQueue.java @@ -37,17 +37,6 @@ public abstract class EventQueue implements AutoCloseable { public static final int DEFAULT_QUEUE_SIZE = 1000; private final int queueSize; - - /** - * Internal blocking queue for storing event queue items. - */ - protected final BlockingQueue queue = new LinkedBlockingDeque<>(); - - /** - * Semaphore for backpressure control, limiting the number of pending events. - */ - protected final Semaphore semaphore; - private volatile boolean closed = false; /** @@ -68,7 +57,6 @@ protected EventQueue(int queueSize) { throw new IllegalArgumentException("Queue size must be greater than 0"); } this.queueSize = queueSize; - this.semaphore = new Semaphore(queueSize, true); LOGGER.trace("Creating {} with queue size: {}", this, queueSize); } @@ -337,6 +325,7 @@ protected void doClose(boolean immediate) { static class MainQueue extends EventQueue { private final List children = new CopyOnWriteArrayList<>(); + protected final Semaphore semaphore; private final CountDownLatch pollingStartedLatch = new CountDownLatch(1); private final AtomicBoolean pollingStarted = new AtomicBoolean(false); private final @Nullable EventEnqueueHook enqueueHook; @@ -352,6 +341,7 @@ static class MainQueue extends EventQueue { @Nullable TaskStateProvider taskStateProvider, @Nullable MainEventBus mainEventBus) { super(queueSize); + this.semaphore = new Semaphore(queueSize, true); this.enqueueHook = hook; this.taskId = taskId; this.onCloseCallbacks = List.copyOf(onCloseCallbacks); // Defensive copy @@ -570,8 +560,12 @@ private void internalEnqueueItem(EventQueueItem item) { LOGGER.warn("ChildQueue is closed. Event will not be enqueued. {} {}", this, event); return; } - queue.add(item); - LOGGER.debug("Enqueued event {} {}", event instanceof Throwable ? event.toString() : event, this); + if (!queue.offer(item)) { + LOGGER.warn("ChildQueue {} is full. Closing immediately.", this); + close(true); // immediate close + } else { + LOGGER.debug("Enqueued event {} {}", event instanceof Throwable ? event.toString() : event, this); + } } @Override @@ -587,7 +581,6 @@ public EventQueueItem dequeueEventItem(int waitMilliSeconds) throws EventQueueCl if (item != null) { Event event = item.getEvent(); LOGGER.debug("Dequeued event item (no wait) {} {}", this, event instanceof Throwable ? event.toString() : event); - // Note: Semaphore is managed by parent MainQueue, not released here } return item; } @@ -597,7 +590,6 @@ public EventQueueItem dequeueEventItem(int waitMilliSeconds) throws EventQueueCl if (item != null) { Event event = item.getEvent(); LOGGER.debug("Dequeued event item (waiting) {} {}", this, event instanceof Throwable ? event.toString() : event); - // Note: Semaphore is managed by parent MainQueue, not released here } else { LOGGER.trace("Dequeue timeout (null) from ChildQueue {}", System.identityHashCode(this)); } @@ -640,8 +632,6 @@ protected void doClose(boolean immediate) { // Immediate close: clear pending events from local queue int clearedCount = queue.size(); queue.clear(); - // Release semaphore permits for cleared events to prevent deadlock - semaphore.release(clearedCount); LOGGER.debug("Cleared {} events from ChildQueue for immediate close: {}", clearedCount, this); } // For graceful close, let the queue drain naturally through normal consumption From 8f5db109f321a574c7fb6e84fa62dec9b2dbedd6 Mon Sep 17 00:00:00 2001 From: Kabir Khan Date: Wed, 10 Dec 2025 15:56:24 +0000 Subject: [PATCH 6/8] TCK fixes --- .../io/a2a/server/events/EventConsumer.java | 3 ++ .../server/events/MainEventBusProcessor.java | 5 ++++ .../events/MainEventBusProcessorCallback.java | 2 -- .../io/a2a/server/tasks/ResultAggregator.java | 29 ++++++++++++------- 4 files changed, 26 insertions(+), 13 deletions(-) diff --git a/server-common/src/main/java/io/a2a/server/events/EventConsumer.java b/server-common/src/main/java/io/a2a/server/events/EventConsumer.java index 7f1cbd4cc..f1f2d39c9 100644 --- a/server-common/src/main/java/io/a2a/server/events/EventConsumer.java +++ b/server-common/src/main/java/io/a2a/server/events/EventConsumer.java @@ -65,7 +65,10 @@ public Flow.Publisher consumeAll() { } event = item.getEvent(); + // Defensive logging for error handling if (event instanceof Throwable thr) { + LOGGER.info("EventConsumer detected Throwable event: {} - triggering tube.fail()", + thr.getClass().getSimpleName()); tube.fail(thr); return; } diff --git a/server-common/src/main/java/io/a2a/server/events/MainEventBusProcessor.java b/server-common/src/main/java/io/a2a/server/events/MainEventBusProcessor.java index 0f3033b71..1551712ec 100644 --- a/server-common/src/main/java/io/a2a/server/events/MainEventBusProcessor.java +++ b/server-common/src/main/java/io/a2a/server/events/MainEventBusProcessor.java @@ -167,6 +167,11 @@ private void processEvent(MainEventBusContext context) { } // Step 3: Then distribute to ChildQueues (clients see either event or error AFTER persistence attempt) + if (eventToDistribute == null) { + LOGGER.error("MainEventBusProcessor: eventToDistribute is NULL for task {} - this should never happen!", taskId); + eventToDistribute = new InternalError("Internal error: event processing failed"); + } + if (eventQueue instanceof EventQueue.MainQueue mainQueue) { int childCount = mainQueue.getChildCount(); LOGGER.debug("MainEventBusProcessor: Distributing {} to {} children for task {}", diff --git a/server-common/src/main/java/io/a2a/server/events/MainEventBusProcessorCallback.java b/server-common/src/main/java/io/a2a/server/events/MainEventBusProcessorCallback.java index 2757e85f9..b0a9adbce 100644 --- a/server-common/src/main/java/io/a2a/server/events/MainEventBusProcessorCallback.java +++ b/server-common/src/main/java/io/a2a/server/events/MainEventBusProcessorCallback.java @@ -8,7 +8,6 @@ * This interface is primarily intended for testing, allowing tests to synchronize * with the asynchronous MainEventBusProcessor. Production code should not rely on this. *

    - *

    * Usage in tests: * 

      * {@code
@@ -31,7 +30,6 @@
  * }
  * }
  *
    - *

    */ public interface MainEventBusProcessorCallback { diff --git a/server-common/src/main/java/io/a2a/server/tasks/ResultAggregator.java b/server-common/src/main/java/io/a2a/server/tasks/ResultAggregator.java index 72f6c36b3..78b24cbc1 100644 --- a/server-common/src/main/java/io/a2a/server/tasks/ResultAggregator.java +++ b/server-common/src/main/java/io/a2a/server/tasks/ResultAggregator.java @@ -16,6 +16,7 @@ import io.a2a.spec.A2AError; import io.a2a.spec.Event; import io.a2a.spec.EventKind; +import io.a2a.spec.InternalError; import io.a2a.spec.Message; import io.a2a.spec.Task; import io.a2a.spec.TaskState; @@ -96,6 +97,7 @@ public EventKind consumeAll(EventConsumer consumer) throws A2AError { public EventTypeAndInterrupt consumeAndBreakOnInterrupt(EventConsumer consumer, boolean blocking) throws A2AError { Flow.Publisher allItems = consumer.consumeAll(); AtomicReference message = new AtomicReference<>(); + AtomicReference capturedTask = new AtomicReference<>(); // Capture Task events AtomicBoolean interrupted = new AtomicBoolean(false); AtomicReference errorRef = new AtomicReference<>(); CompletableFuture completionFuture = new CompletableFuture<>(); @@ -129,6 +131,11 @@ public EventTypeAndInterrupt consumeAndBreakOnInterrupt(EventConsumer consumer, return false; } + // Capture Task events (especially for new tasks where taskManager.getTask() would return null) + if (event instanceof Task t) { + capturedTask.set(t); + } + // TaskStore update moved to MainEventBusProcessor // Determine interrupt behavior @@ -229,20 +236,20 @@ else if (blocking) { Utils.rethrow(error); } - EventKind eventType; - Message msg = message.get(); - if (msg != null) { - eventType = msg; - } else { - Task task = taskManager.getTask(); - if (task == null) { - throw new io.a2a.spec.InternalError("No task or message available after consuming events"); - } - eventType = task; + // Return Message if captured, otherwise Task if captured, otherwise fetch from TaskStore + EventKind eventKind = message.get(); + if (eventKind == null) { + eventKind = capturedTask.get(); + } + if (eventKind == null) { + eventKind = taskManager.getTask(); + } + if (eventKind == null) { + throw new InternalError("Could not find a Task/Message for " + taskManager.getTaskId()); } return new EventTypeAndInterrupt( - eventType, + eventKind, interrupted.get(), consumptionCompletionFuture); } From 881bdbe553ff17c06e66aededf1ca22000eba755 Mon Sep 17 00:00:00 2001 From: Kabir Khan Date: Thu, 11 Dec 2025 09:08:30 +0000 Subject: [PATCH 7/8] TCK fixes 2 --- .../DefaultRequestHandler.java | 35 ++++++++++--------- .../io/a2a/server/tasks/ResultAggregator.java | 13 +++++++ 2 files changed, 31 insertions(+), 17 deletions(-) diff --git a/server-common/src/main/java/io/a2a/server/requesthandlers/DefaultRequestHandler.java b/server-common/src/main/java/io/a2a/server/requesthandlers/DefaultRequestHandler.java index f908c108e..934eb0cc0 100644 --- a/server-common/src/main/java/io/a2a/server/requesthandlers/DefaultRequestHandler.java +++ b/server-common/src/main/java/io/a2a/server/requesthandlers/DefaultRequestHandler.java @@ -407,7 +407,8 @@ public EventKind onMessageSend(MessageSendParams params, ServerCallContext conte // Store push notification config for newly created tasks (mirrors streaming logic) // Only for NEW tasks - existing tasks are handled by initMessageSend() if (mss.task() == null && kind instanceof Task createdTask && shouldAddPushInfo(params)) { - LOGGER.debug("Storing push notification config for new task {}", createdTask.id()); + LOGGER.debug("Storing push notification config for new task {} (original taskId from params: {})", + createdTask.id(), params.message().taskId()); pushConfigStore.setInfo(createdTask.id(), params.configuration().pushNotificationConfig()); } @@ -508,6 +509,18 @@ public Flow.Publisher onMessageSendStream( @SuppressWarnings("NullAway") EventQueue queue = queueManager.createOrTap(taskId.get()); LOGGER.debug("Created/tapped queue for task {}: {}", taskId.get(), queue); + + // Store push notification config SYNCHRONOUSLY for new tasks before agent starts + // This ensures config is available when MainEventBusProcessor sends push notifications + // For existing tasks, config was already stored in initMessageSend() + if (mss.task() == null && shouldAddPushInfo(params)) { + // Satisfy Nullaway + Objects.requireNonNull(taskId.get(), "taskId was null"); + LOGGER.debug("Storing push notification config for new streaming task {} EARLY (original taskId from params: {})", + taskId.get(), params.message().taskId()); + pushConfigStore.setInfo(taskId.get(), params.configuration().pushNotificationConfig()); + } + ResultAggregator resultAggregator = new ResultAggregator(mss.taskManager, null, executor); EnhancedRunnable producerRunnable = registerAndExecuteAgentAsync(queueTaskId, mss.requestContext, queue); @@ -536,15 +549,6 @@ public Flow.Publisher onMessageSendStream( } catch (TaskQueueExistsException e) { // TODO Log } - if (pushConfigStore != null && - params.configuration() != null && - params.configuration().pushNotificationConfig() != null) { - - pushConfigStore.setInfo( - createdTask.id(), - params.configuration().pushNotificationConfig()); - } - } return true; })); @@ -815,13 +819,10 @@ private CompletableFuture cleanupProducer(@Nullable CompletableFuture Date: Fri, 12 Dec 2025 10:09:05 +0000 Subject: [PATCH 8/8] Debug logging --- .../io/a2a/server/events/EventConsumer.java | 9 +++++ .../DefaultRequestHandler.java | 34 +++++++++++++------ .../io/a2a/server/tasks/ResultAggregator.java | 7 ++++ 3 files changed, 40 insertions(+), 10 deletions(-) diff --git a/server-common/src/main/java/io/a2a/server/events/EventConsumer.java b/server-common/src/main/java/io/a2a/server/events/EventConsumer.java index f1f2d39c9..578079185 100644 --- a/server-common/src/main/java/io/a2a/server/events/EventConsumer.java +++ b/server-common/src/main/java/io/a2a/server/events/EventConsumer.java @@ -59,11 +59,15 @@ public Flow.Publisher consumeAll() { EventQueueItem item; Event event; try { + LOGGER.debug("EventConsumer polling queue {} (error={})", System.identityHashCode(queue), error); item = queue.dequeueEventItem(QUEUE_WAIT_MILLISECONDS); if (item == null) { + LOGGER.debug("EventConsumer poll timeout (null item), continuing"); continue; } event = item.getEvent(); + LOGGER.debug("EventConsumer received event: {} (queue={})", + event.getClass().getSimpleName(), System.identityHashCode(queue)); // Defensive logging for error handling if (event instanceof Throwable thr) { @@ -125,8 +129,13 @@ public Flow.Publisher consumeAll() { public EnhancedRunnable.DoneCallback createAgentRunnableDoneCallback() { return agentRunnable -> { + LOGGER.info("EventConsumer: Agent done callback invoked (hasError={}, queue={})", + agentRunnable.getError() != null, System.identityHashCode(queue)); if (agentRunnable.getError() != null) { error = agentRunnable.getError(); + LOGGER.info("EventConsumer: Set error field from agent callback"); + } else { + LOGGER.info("EventConsumer: Agent completed successfully (no error), continuing consumption"); } }; } diff --git a/server-common/src/main/java/io/a2a/server/requesthandlers/DefaultRequestHandler.java b/server-common/src/main/java/io/a2a/server/requesthandlers/DefaultRequestHandler.java index 934eb0cc0..2b74d37a3 100644 --- a/server-common/src/main/java/io/a2a/server/requesthandlers/DefaultRequestHandler.java +++ b/server-common/src/main/java/io/a2a/server/requesthandlers/DefaultRequestHandler.java @@ -374,6 +374,17 @@ public EventKind onMessageSend(MessageSendParams params, ServerCallContext conte boolean blocking = params.configuration() != null && Boolean.TRUE.equals(params.configuration().blocking()); + // Log blocking behavior from client request + if (params.configuration() != null && params.configuration().blocking() != null) { + LOGGER.info("DefaultRequestHandler: Client requested blocking={} for task {}", + params.configuration().blocking(), taskId); + } else if (params.configuration() != null) { + LOGGER.info("DefaultRequestHandler: Client sent configuration but blocking=null, using default blocking=true for task {}", taskId); + } else { + LOGGER.info("DefaultRequestHandler: Client sent no configuration, using default blocking=true for task {}", taskId); + } + LOGGER.info("DefaultRequestHandler: Final blocking decision: {} for task {}", blocking, taskId); + boolean interruptedOrNonBlocking = false; EnhancedRunnable producerRunnable = registerAndExecuteAgentAsync(taskId, mss.requestContext, queue); @@ -395,7 +406,8 @@ public EventKind onMessageSend(MessageSendParams params, ServerCallContext conte throw new InternalError("No result"); } interruptedOrNonBlocking = etai.interrupted(); - LOGGER.debug("Was interrupted or non-blocking: {}", interruptedOrNonBlocking); + LOGGER.info("DefaultRequestHandler: interruptedOrNonBlocking={} (blocking={}, eventType={})", + interruptedOrNonBlocking, blocking, kind != null ? kind.getClass().getSimpleName() : null); // For blocking calls that were interrupted (returned on first event), // wait for agent execution and event processing BEFORE returning to client. @@ -419,16 +431,18 @@ public EventKind onMessageSend(MessageSendParams params, ServerCallContext conte // 2. Close the queue to signal consumption can complete // 3. Wait for consumption to finish processing events // 4. Fetch final task state from TaskStore + LOGGER.debug("DefaultRequestHandler: Entering blocking fire-and-forget handling for task {}", taskId); 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); + LOGGER.debug("DefaultRequestHandler: Step 1 - 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); + LOGGER.debug("DefaultRequestHandler: Step 1 - Agent still running for task {} after {}s timeout", + taskId, agentCompletionTimeoutSeconds); } } @@ -436,12 +450,12 @@ public EventKind onMessageSend(MessageSendParams params, ServerCallContext conte // 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); + LOGGER.debug("DefaultRequestHandler: Step 2 - 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); + LOGGER.debug("DefaultRequestHandler: Step 3 - Consumption completed for task {}", taskId); } } catch (InterruptedException e) { Thread.currentThread().interrupt(); @@ -464,11 +478,11 @@ public EventKind onMessageSend(MessageSendParams params, ServerCallContext conte Task updatedTask = taskStore.get(nonNullTaskId); if (updatedTask != null) { kind = updatedTask; - if (LOGGER.isDebugEnabled()) { - LOGGER.debug("Fetched final task for {} with state {} and {} artifacts", - nonNullTaskId, updatedTask.status().state(), - updatedTask.artifacts().size()); - } + LOGGER.debug("DefaultRequestHandler: Step 4 - Fetched final task for {} with state {} and {} artifacts", + taskId, updatedTask.status().state(), + updatedTask.artifacts().size()); + } else { + LOGGER.warn("DefaultRequestHandler: Step 4 - Task {} not found in TaskStore!", taskId); } } if (kind instanceof Task taskResult && !taskId.equals(taskResult.id())) { diff --git a/server-common/src/main/java/io/a2a/server/tasks/ResultAggregator.java b/server-common/src/main/java/io/a2a/server/tasks/ResultAggregator.java index 865fd6d5d..51e1783d7 100644 --- a/server-common/src/main/java/io/a2a/server/tasks/ResultAggregator.java +++ b/server-common/src/main/java/io/a2a/server/tasks/ResultAggregator.java @@ -152,6 +152,9 @@ public EventTypeAndInterrupt consumeAndBreakOnInterrupt(EventConsumer consumer, boolean isAuthRequired = (event instanceof Task task && task.status().state() == TaskState.AUTH_REQUIRED) || (event instanceof TaskStatusUpdateEvent tsue && tsue.status().state() == TaskState.AUTH_REQUIRED); + LOGGER.info("ResultAggregator: Evaluating interrupt (blocking={}, isFinal={}, isAuth={}, eventType={})", + blocking, isFinalEvent, isAuthRequired, event.getClass().getSimpleName()); + // Always interrupt on auth_required, as it needs external action. if (isAuthRequired) { // auth-required is a special state: the message should be @@ -161,16 +164,19 @@ public EventTypeAndInterrupt consumeAndBreakOnInterrupt(EventConsumer consumer, // new request is expected in order for the agent to make progress, // so the agent should exit. shouldInterrupt = true; + LOGGER.info("ResultAggregator: Setting shouldInterrupt=true (AUTH_REQUIRED)"); } else if (!blocking) { // For non-blocking calls, interrupt as soon as a task is available. shouldInterrupt = true; + LOGGER.info("ResultAggregator: Setting shouldInterrupt=true (non-blocking)"); } 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 shouldInterrupt = true; + LOGGER.info("ResultAggregator: Setting shouldInterrupt=true (blocking, isFinal={})", isFinalEvent); if (LOGGER.isDebugEnabled()) { LOGGER.debug("Blocking call for task {}: {} event, returning with background consumption", taskIdForLogging(), isFinalEvent ? "final" : "non-final"); @@ -178,6 +184,7 @@ else if (blocking) { } if (shouldInterrupt) { + LOGGER.info("ResultAggregator: Interrupting consumption (setting interrupted=true)"); // Complete the future to unblock the main thread interrupted.set(true); completionFuture.complete(null);