Fix Gemini feedback

Author: kabir

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

@@ -471,8 +471,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
         }

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

@@ -44,7 +44,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;
 
@@ -71,8 +71,8 @@ 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

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

@@ -12,10 +12,13 @@
  * Usage in tests:
  * <pre>
  * {@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
  * }
  * }
  * </pre>

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

@@ -267,9 +267,13 @@ public void testDequeueEventWhenClosedButHasEvents() throws Exception {
         // 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
+        // Poll for event to arrive in childQueue (async MainEventBusProcessor distribution)
+        long startTime = System.currentTimeMillis();
+        long timeout = 2000;
+        while (childQueue.size() == 0 && (System.currentTimeMillis() - startTime) < timeout) {
+            Thread.sleep(50);
+        }
+        assertTrue(childQueue.size() > 0, "Event should arrive in ChildQueue within timeout");
 
         childQueue.close(); // Graceful close - events should remain
         assertTrue(childQueue.isClosed());

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

@@ -26,77 +26,43 @@ public class EventQueueUtil {
     // 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
+    /**
+     * Get an EventQueue builder pre-configured with the shared test MainEventBus and a unique taskId.
+     * <p>
+     * Note: Returns MainQueue - tests should call .tap() if they need to consume events.
+     * </p>
+     *
+     * @return builder with TEST_EVENT_BUS and unique taskId already set
+     */
     public static EventQueue.EventQueueBuilder getEventQueueBuilder() {
-        return new EventQueueBuilderWrapper(
-                EventQueue.builder(TEST_EVENT_BUS)
-                        .taskId("test-task-" + TASK_ID_COUNTER.incrementAndGet())
-        );
+        return EventQueue.builder(TEST_EVENT_BUS)
+                .taskId("test-task-" + TASK_ID_COUNTER.incrementAndGet());
     }
 
-    // Get the shared test MainEventBus instance
+    /**
+     * Get the shared test MainEventBus instance.
+     *
+     * @return the test MainEventBus
+     */
     public static MainEventBus getTestEventBus() {
         return TEST_EVENT_BUS;
     }
 
+    /**
+     * 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();
-        }
-    }
 }

transport/jsonrpc/src/test/java/io/a2a/transport/jsonrpc/handler/JSONRPCHandlerTest.java

@@ -285,7 +285,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);
-        io.a2a.server.events.MainEventBusProcessor.setCallback(new io.a2a.server.events.MainEventBusProcessorCallback() {
+        mainEventBusProcessor.setCallback(new io.a2a.server.events.MainEventBusProcessorCallback() {
             @Override
             public void onEventProcessed(String taskId, io.a2a.spec.Event event) {
                 processingLatch.countDown();
@@ -407,7 +407,7 @@ public void onComplete() {
             assertEquals(MINIMAL_TASK.getId(), receivedStatus.getTaskId());
             assertEquals(TaskState.COMPLETED, receivedStatus.getStatus().state());
         } finally {
-            io.a2a.server.events.MainEventBusProcessor.setCallback(null);
+            mainEventBusProcessor.setCallback(null);
         }
     }
 
@@ -679,7 +679,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);
-        io.a2a.server.events.MainEventBusProcessor.setCallback(new io.a2a.server.events.MainEventBusProcessorCallback() {
+        mainEventBusProcessor.setCallback(new io.a2a.server.events.MainEventBusProcessorCallback() {
             @Override
             public void onEventProcessed(String taskId, io.a2a.spec.Event event) {
                 processingLatch.countDown();
@@ -799,7 +799,7 @@ public void onComplete() {
         assertEquals(1, curr.getArtifacts().get(0).parts().size());
         assertEquals("text", ((TextPart)curr.getArtifacts().get(0).parts().get(0)).getText());
         } finally {
-            io.a2a.server.events.MainEventBusProcessor.setCallback(null);
+            mainEventBusProcessor.setCallback(null);
         }
     }
 
@@ -1267,7 +1267,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);
-        io.a2a.server.events.MainEventBusProcessor.setCallback(new io.a2a.server.events.MainEventBusProcessorCallback() {
+        mainEventBusProcessor.setCallback(new io.a2a.server.events.MainEventBusProcessorCallback() {
             @Override
             public void onEventProcessed(String taskId, io.a2a.spec.Event event) {
                 processingLatch.countDown();
@@ -1332,7 +1332,7 @@ public void onComplete() {
             Assertions.assertEquals(1, results.size());
             Assertions.assertInstanceOf(InternalError.class, results.get(0).getError());
         } finally {
-            io.a2a.server.events.MainEventBusProcessor.setCallback(null);
+            mainEventBusProcessor.setCallback(null);
         }
     }