Fix server-common tests

Author: kabir

server-common/src/main/java/io/a2a/server/events/EventQueue.java

@@ -316,13 +316,12 @@ public void enqueueItem(EventQueueItem item) {
             queue.add(item);
             LOGGER.debug("Enqueued event {} {}", event instanceof Throwable ? event.toString() : event, this);
 
-            // Submit to MainEventBus instead of immediate distribution to children
+            // Submit to MainEventBus for centralized persistence + distribution
             if (mainEventBus != null && taskId != null) {
                 mainEventBus.submit(taskId, this, item);
             } else {
-                // Fallback: immediate distribution (for tests/non-bus mode)
-                LOGGER.warn("MainEventBus not configured for task {}, using immediate distribution", taskId);
-                children.forEach(eq -> eq.internalEnqueueItem(item));
+                // 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 (KEEP for inter-process replication)

server-common/src/main/java/io/a2a/server/events/InMemoryQueueManager.java

@@ -18,11 +18,11 @@ public class InMemoryQueueManager implements QueueManager {
     private final EventQueueFactory factory;
     private final TaskStateProvider taskStateProvider;
 
-    @Inject
     MainEventBus mainEventBus;
 
     @Inject
-    public InMemoryQueueManager(TaskStateProvider taskStateProvider) {
+    public InMemoryQueueManager(TaskStateProvider taskStateProvider, MainEventBus mainEventBus) {
+        this.mainEventBus = mainEventBus;
         this.factory = new DefaultEventQueueFactory();
         this.taskStateProvider = taskStateProvider;
     }

server-common/src/main/java/io/a2a/server/events/MainEventBusProcessor.java

@@ -33,15 +33,20 @@
 public class MainEventBusProcessor implements Runnable {
     private static final Logger LOGGER = LoggerFactory.getLogger(MainEventBusProcessor.class);
 
-    @Inject
-    MainEventBus eventBus;
 
-    @Inject
-    TaskStore taskStore;
+    private final MainEventBus eventBus;
+
+    private final TaskStore taskStore;
 
     private volatile boolean running = true;
     private Thread processorThread;
 
+    @Inject
+    public MainEventBusProcessor(MainEventBus eventBus, TaskStore taskStore) {
+        this.eventBus = eventBus;
+        this.taskStore = taskStore;
+    }
+
     @PostConstruct
     void start() {
         processorThread = new Thread(this, "MainEventBusProcessor");
@@ -91,10 +96,10 @@ private void processEvent(MainEventBusContext context) {
 
         LOGGER.debug("Processing event for task {}: {}", taskId, event.getClass().getSimpleName());
 
-        // Step 1: Update TaskStore FIRST
+        // Step 1: Update TaskStore FIRST (persistence before clients see it)
         updateTaskStore(taskId, event);
 
-        // Step 2: Then distribute to ChildQueues
+        // Step 2: Then distribute to ChildQueues (clients see it AFTER persistence)
         if (eventQueue instanceof EventQueue.MainQueue mainQueue) {
             mainQueue.distributeToChildren(context.eventQueueItem());
             LOGGER.debug("Distributed event to children for task {}", taskId);

server-common/src/test/java/io/a2a/server/events/EventQueueTest.java

@@ -20,15 +20,19 @@
 import io.a2a.spec.TaskNotFoundError;
 import io.a2a.spec.TaskState;
 import io.a2a.spec.TaskStatus;
+import io.a2a.server.tasks.InMemoryTaskStore;
 import io.a2a.spec.TaskStatusUpdateEvent;
 import io.a2a.spec.TextPart;
 import io.a2a.util.Utils;
+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 = """
             {
@@ -51,8 +55,33 @@ public class EventQueueTest {
 
     @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);
+        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
@@ -97,23 +126,24 @@ 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 = Utils.unmarshalFrom(MINIMAL_TASK, Task.TYPE_REFERENCE);
         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);
     }
 
     @Test
     public void testMultipleChildQueuesReceiveEvents() throws Exception {
-        EventQueue parentQueue = EventQueue.builder().build();
+        EventQueue parentQueue = createQueueWithEventBus("test-multiple");
         EventQueue childQueue1 = parentQueue.tap();
         EventQueue childQueue2 = parentQueue.tap();
 
@@ -124,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 = Utils.unmarshalFrom(MINIMAL_TASK, Task.TYPE_REFERENCE);
         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
@@ -168,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
         }

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();
+    }
 }

server-common/src/test/java/io/a2a/server/events/InMemoryQueueManagerTest.java

@@ -14,19 +14,34 @@
 import java.util.concurrent.ExecutionException;
 import java.util.stream.IntStream;
 
+import io.a2a.server.tasks.InMemoryTaskStore;
 import io.a2a.server.tasks.MockTaskStateProvider;
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
 public class InMemoryQueueManagerTest {
 
     private InMemoryQueueManager queueManager;
     private MockTaskStateProvider taskStateProvider;
+    private InMemoryTaskStore taskStore;
+    private MainEventBus mainEventBus;
+    private MainEventBusProcessor mainEventBusProcessor;
 
     @BeforeEach
     public void setUp() {
         taskStateProvider = new MockTaskStateProvider();
-        queueManager = new InMemoryQueueManager(taskStateProvider);
+        taskStore = new InMemoryTaskStore();
+        mainEventBus = new MainEventBus();
+        mainEventBusProcessor = new MainEventBusProcessor(mainEventBus, taskStore);
+        EventQueueUtil.start(mainEventBusProcessor);
+
+        queueManager = new InMemoryQueueManager(taskStateProvider, mainEventBus);
+    }
+
+    @AfterEach
+    public void tearDown() {
+        EventQueueUtil.stop(mainEventBusProcessor);
     }
 
     @Test

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

@@ -21,7 +21,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;
@@ -71,6 +74,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();
 
@@ -94,7 +99,14 @@ public void cancel(RequestContext context, EventQueue eventQueue) throws JSONRPC
 
         InMemoryTaskStore inMemoryTaskStore = new InMemoryTaskStore();
         taskStore = inMemoryTaskStore;
-        queueManager = new InMemoryQueueManager(inMemoryTaskStore);
+
+        // Create MainEventBus and MainEventBusProcessor (production code path)
+        mainEventBus = new MainEventBus();
+        mainEventBusProcessor = new MainEventBusProcessor(mainEventBus, taskStore);
+        EventQueueUtil.start(mainEventBusProcessor);
+
+        queueManager = new InMemoryQueueManager(inMemoryTaskStore, mainEventBus);
+
         httpClient = new TestHttpClient();
         PushNotificationConfigStore pushConfigStore = new InMemoryPushNotificationConfigStore();
         PushNotificationSender pushSender = new BasePushNotificationSender(pushConfigStore, httpClient);
@@ -107,6 +119,11 @@ public void cancel(RequestContext context, EventQueue eventQueue) throws JSONRPC
     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) {

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

@@ -19,7 +19,10 @@
 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.TaskUpdater;
@@ -32,6 +35,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 +54,21 @@ public class DefaultRequestHandlerTest {
     private InMemoryQueueManager queueManager;
     private TestAgentExecutor agentExecutor;
     private ServerCallContext serverCallContext;
+    private MainEventBus mainEventBus;
+    private MainEventBusProcessor mainEventBusProcessor;
 
     @BeforeEach
     void setUp() {
         taskStore = new InMemoryTaskStore();
+
+        // Create MainEventBus and MainEventBusProcessor (production code path)
+        mainEventBus = new MainEventBus();
+        mainEventBusProcessor = new MainEventBusProcessor(mainEventBus, taskStore);
+        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(
@@ -70,6 +83,14 @@ void setUp() {
         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).