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

Author: kabir

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;

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) {

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

@@ -21,7 +21,6 @@ public abstract class EventQueue implements AutoCloseable {
     public static final int DEFAULT_QUEUE_SIZE = 1000;
 
     private final int queueSize;
-    protected final BlockingQueue<EventQueueItem> queue = new LinkedBlockingDeque<>();
     protected final Semaphore semaphore;
     private volatile boolean closed = false;
 
@@ -43,8 +42,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);
     }
 
     public static class EventQueueBuilder {
@@ -88,11 +87,11 @@ public EventQueueBuilder mainEventBus(MainEventBus mainEventBus) {
         }
 
         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");
             }
+            return new MainQueue(queueSize, hook, taskId, onCloseCallbacks, taskStateProvider, mainEventBus);
         }
     }
 
@@ -108,22 +107,7 @@ public void enqueueEvent(Event event) {
         enqueueItem(new LocalEventQueueItem(event));
     }
 
-    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);
 
     public abstract EventQueue tap();
 
@@ -133,53 +117,32 @@ 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()}.
      * </p>
+     * <p>
+     * Note: MainQueue does not support dequeue operations - only ChildQueues can be consumed.
+     * </p>
      *
      * @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 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();
-        }
-    }
+    public abstract EventQueueItem dequeueEventItem(int waitMilliSeconds) throws EventQueueClosedException;
 
     public void taskDone() {
         // TODO Not sure if needed yet. BlockingQueue.poll()/.take() remove the events.
     }
 
+    /**
+     * Returns the current size of the queue.
+     * <p>
+     * For MainQueue: returns the size of the MainEventBus queue (events pending persistence/distribution).
+     * For ChildQueue: returns the size of the local consumption queue.
+     * </p>
+     *
+     * @return the number of events currently in the queue
+     */
+    public abstract int size();
+
     public abstract void close();
 
     public abstract void close(boolean immediate);
@@ -211,13 +174,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 {
@@ -230,62 +187,15 @@ static class MainQueue extends EventQueue {
         private final TaskStateProvider taskStateProvider;
         private final 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;
-        }
-
-        MainQueue(int queueSize, EventEnqueueHook hook, String taskId, List<Runnable> onCloseCallbacks, TaskStateProvider taskStateProvider) {
-            super(queueSize);
-            this.enqueueHook = hook;
-            this.taskId = taskId;
-            this.onCloseCallbacks = List.copyOf(onCloseCallbacks);  // Defensive copy
-            this.taskStateProvider = taskStateProvider;
-            this.mainEventBus = null;
-            LOGGER.debug("Created MainQueue for task {} with {} onClose callbacks and TaskStateProvider: {}",
-                    taskId, onCloseCallbacks.size(), taskStateProvider != null);
-        }
-
         MainQueue(int queueSize, EventEnqueueHook hook, String taskId, List<Runnable> onCloseCallbacks, TaskStateProvider taskStateProvider, MainEventBus mainEventBus) {
             super(queueSize);
             this.enqueueHook = hook;
             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 = java.util.Objects.requireNonNull(mainEventBus, "MainEventBus is required");
+            LOGGER.debug("Created MainQueue for task {} with {} onClose callbacks, TaskStateProvider: {}, MainEventBus configured",
+                    taskId, onCloseCallbacks.size(), taskStateProvider != null);
         }
 
         public EventQueue tap() {
@@ -302,6 +212,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:
@@ -320,19 +241,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);
             }
@@ -459,6 +374,7 @@ public void close(boolean immediate, boolean notifyParent) {
 
     static class ChildQueue extends EventQueue {
         private final MainQueue parent;
+        private final BlockingQueue<EventQueueItem> queue = new LinkedBlockingDeque<>();
 
         public ChildQueue(MainQueue parent) {
             this.parent = parent;
@@ -469,15 +385,77 @@ 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
+        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
         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();
@@ -488,6 +466,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);