Remove queues when closed

Author: kabir

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

@@ -49,6 +49,8 @@ static EventQueueBuilder builder() {
     public static class EventQueueBuilder {
         private int queueSize = DEFAULT_QUEUE_SIZE;
         private EventEnqueueHook hook;
+        private String taskId;
+        private Runnable onCloseCallback;
 
         public EventQueueBuilder queueSize(int queueSize) {
             this.queueSize = queueSize;
@@ -60,9 +62,19 @@ public EventQueueBuilder hook(EventEnqueueHook hook) {
             return this;
         }
 
+        public EventQueueBuilder taskId(String taskId) {
+            this.taskId = taskId;
+            return this;
+        }
+
+        public EventQueueBuilder onClose(Runnable onCloseCallback) {
+            this.onCloseCallback = onCloseCallback;
+            return this;
+        }
+
         public EventQueue build() {
-            if (hook != null) {
-                return new MainQueue(queueSize, hook);
+            if (hook != null || onCloseCallback != null) {
+                return new MainQueue(queueSize, hook, taskId, onCloseCallback);
             } else {
                 return new MainQueue(queueSize);
             }
@@ -154,11 +166,11 @@ public boolean isClosed() {
         return closed;
     }
 
-    public void doClose() {
+    protected void doClose() {
         doClose(false);
     }
 
-    public void doClose(boolean immediate) {
+    protected void doClose(boolean immediate) {
         synchronized (this) {
             if (closed) {
                 return;
@@ -180,25 +192,43 @@ static class MainQueue extends EventQueue {
         private final CountDownLatch pollingStartedLatch = new CountDownLatch(1);
         private final AtomicBoolean pollingStarted = new AtomicBoolean(false);
         private final EventEnqueueHook enqueueHook;
+        private final String taskId;
+        private final Runnable onCloseCallback;
 
         MainQueue() {
             super();
             this.enqueueHook = null;
+            this.taskId = null;
+            this.onCloseCallback = null;
         }
 
         MainQueue(int queueSize) {
             super(queueSize);
             this.enqueueHook = null;
+            this.taskId = null;
+            this.onCloseCallback = null;
         }
 
         MainQueue(EventEnqueueHook hook) {
             super();
             this.enqueueHook = hook;
+            this.taskId = null;
+            this.onCloseCallback = null;
         }
 
         MainQueue(int queueSize, EventEnqueueHook hook) {
             super(queueSize);
             this.enqueueHook = hook;
+            this.taskId = null;
+            this.onCloseCallback = null;
+        }
+
+        MainQueue(int queueSize, EventEnqueueHook hook, String taskId, Runnable onCloseCallback) {
+            super(queueSize);
+            this.enqueueHook = hook;
+            this.taskId = taskId;
+            this.onCloseCallback = onCloseCallback;
+            LOGGER.debug("Created MainQueue for task {} with onClose callback: {}", taskId, onCloseCallback != null);
         }
 
         EventQueue tap() {
@@ -254,6 +284,20 @@ public int getActiveChildCount() {
             return children.size();
         }
 
+        @Override
+        protected void doClose(boolean immediate) {
+            super.doClose(immediate);
+            // Invoke callback after closing to notify QueueManager
+            if (onCloseCallback != null) {
+                LOGGER.debug("Invoking onClose callback for task {}", taskId);
+                try {
+                    onCloseCallback.run();
+                } catch (Exception e) {
+                    LOGGER.error("Error in onClose callback for task {}", taskId, e);
+                }
+            }
+        }
+
         @Override
         public void close() {
             close(false);

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

@@ -54,6 +54,7 @@ public void close(String taskId) {
 
     @Override
     public EventQueue createOrTap(String taskId) {
+        LOGGER.info("createOrTap called for task {}, current map size: {}", taskId, queues.size());
         EventQueue existing = queues.get(taskId);
 
         // Lazy cleanup: remove closed queues from map
@@ -76,8 +77,8 @@ public EventQueue createOrTap(String taskId) {
         EventQueue result = main.tap();  // Always return ChildQueue
 
         if (existing == null) {
-            LOGGER.debug("Created new MainQueue {} for task {}, returning ChildQueue {}",
-                System.identityHashCode(main), taskId, System.identityHashCode(result));
+            LOGGER.info("Created new MainQueue {} for task {}, returning ChildQueue {} (map size: {})",
+                System.identityHashCode(main), taskId, System.identityHashCode(result), queues.size());
         } else {
             LOGGER.debug("Tapped existing MainQueue {} -> ChildQueue {} for task {}",
                 System.identityHashCode(main), System.identityHashCode(result), taskId);
@@ -104,11 +105,22 @@ public int getActiveChildQueueCount(String taskId) {
         return -1;
     }
 
-    private static class DefaultEventQueueFactory implements EventQueueFactory {
+    private class DefaultEventQueueFactory implements EventQueueFactory {
         @Override
         public EventQueue.EventQueueBuilder builder(String taskId) {
-            // Default implementation doesn't need task-specific configuration
-            return EventQueue.builder();
+            // Return builder with callback that removes queue from map when closed
+            return EventQueue.builder()
+                    .taskId(taskId)
+                    .onClose(() -> {
+                        LOGGER.debug("Queue close callback invoked for task {}, removing from map", taskId);
+                        EventQueue removed = queues.remove(taskId);
+                        if (removed != null) {
+                            LOGGER.info("Removed closed queue for task {} from QueueManager (map size: {})",
+                                    taskId, queues.size());
+                        } else {
+                            LOGGER.debug("Queue for task {} was already removed from map", taskId);
+                        }
+                    });
         }
     }
 }