Merge pull request #114 from kabir/async-improvements

fjuma · web-flow · commit 2e944f96f8b0 · 2025-06-19T09:35:09.000-04:00
Async improvements
diff --git a/core/src/main/java/io/a2a/server/events/EventConsumer.java b/core/src/main/java/io/a2a/server/events/EventConsumer.java
@@ -16,8 +16,6 @@
 public class EventConsumer {
     private final EventQueue queue;
     private Throwable error;
-    private final Executor executor = Executors.newCachedThreadPool();
-
 
     private static final String ERROR_MSG = "Agent did not return any response";
     private static final int NO_WAIT = -1;
diff --git a/core/src/main/java/io/a2a/server/events/EventQueue.java b/core/src/main/java/io/a2a/server/events/EventQueue.java
@@ -6,6 +6,7 @@
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import io.a2a.util.TempLoggerWrapper;
 import org.slf4j.Logger;
@@ -35,7 +36,9 @@ public static EventQueue create() {
         return new MainQueue();
     }
 
-    abstract CountDownLatch getPollingStartedLatch();
+    public abstract void awaitQueuePollerStart() throws InterruptedException ;
+
+    abstract void signalQueuePollerStarted();
 
     public void enqueueEvent(Event event) {
         if (closed) {
@@ -71,21 +74,22 @@ public Event dequeueEvent(int waitMilliSeconds) throws EventQueueClosedException
                 }
                 return event;
             } catch (InterruptedException e) {
-                log.debug("Interrupted {}", this);
+                log.debug("Interrupted dequeue (waiting) {}", this);
                 Thread.currentThread().interrupt();
                 return null;
             }
         } finally {
-            log.debug("Signalling that queue polling started {}", this);
-            getPollingStartedLatch().countDown();
+            signalQueuePollerStarted();
         }
     }
 
     public void taskDone() {
         // TODO Not sure if needed yet. BlockingQueue.poll()/.take() remove the events.
     }
 
-    public void close() {
+    public abstract void close();
+
+    public void doClose() {
         synchronized (this) {
             if (closed) {
                 return;
@@ -103,7 +107,8 @@ public void close() {
 
     static class MainQueue extends EventQueue {
         private final List<ChildQueue> children = new CopyOnWriteArrayList<>();
-        private CountDownLatch pollingStartedLatch = new CountDownLatch(1);
+        private final CountDownLatch pollingStartedLatch = new CountDownLatch(1);
+        private final AtomicBoolean pollingStarted = new AtomicBoolean(false);
 
         EventQueue tap() {
             ChildQueue child = new ChildQueue(this);
@@ -116,16 +121,27 @@ public void enqueueEvent(Event event) {
             children.forEach(eq -> eq.internalEnqueueEvent(event));
         }
 
-        CountDownLatch getPollingStartedLatch() {
-            return pollingStartedLatch;
+        @Override
+        public void awaitQueuePollerStart() throws InterruptedException {
+            log.debug("Waiting for queue poller to start on {}", this);
+            pollingStartedLatch.await(10, TimeUnit.SECONDS);
+            log.debug("Queue poller started on {}", this);
         }
 
-
+        @Override
+        void signalQueuePollerStarted() {
+            if (pollingStarted.get()) {
+                return;
+            }
+            log.debug("Signalling that queue polling started {}", this);
+            pollingStartedLatch.countDown();
+            pollingStarted.set(true);
+          }
 
         @Override
         public void close() {
-            super.close();
-            children.forEach(EventQueue::close);
+            doClose();
+            children.forEach(EventQueue::doClose);
         }
     }
 
@@ -151,8 +167,18 @@ EventQueue tap() {
         }
 
         @Override
-        CountDownLatch getPollingStartedLatch() {
-            return parent.getPollingStartedLatch();
+        public void awaitQueuePollerStart() throws InterruptedException {
+            parent.awaitQueuePollerStart();
+        }
+
+        @Override
+        void signalQueuePollerStarted() {
+            parent.signalQueuePollerStarted();
+        }
+
+        @Override
+        public void close() {
+            parent.close();
         }
     }
 }
diff --git a/core/src/main/java/io/a2a/server/events/InMemoryQueueManager.java b/core/src/main/java/io/a2a/server/events/InMemoryQueueManager.java
@@ -7,10 +7,12 @@
 
 import jakarta.enterprise.context.ApplicationScoped;
 
+import io.a2a.util.TempLoggerWrapper;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
 @ApplicationScoped
 public class InMemoryQueueManager implements QueueManager {
-
-
     private final Map<String, EventQueue> queues = Collections.synchronizedMap(new HashMap<>());
 
     @Override
@@ -63,7 +65,7 @@ public EventQueue createOrTap(String taskId) {
     }
 
     @Override
-    public void signalPollingStarted(EventQueue eventQueue) throws InterruptedException {
-        eventQueue.getPollingStartedLatch().await(10, TimeUnit.SECONDS);
+    public void awaitQueuePollerStart(EventQueue eventQueue) throws InterruptedException {
+        eventQueue.awaitQueuePollerStart();
     }
 }
diff --git a/core/src/main/java/io/a2a/server/events/QueueManager.java b/core/src/main/java/io/a2a/server/events/QueueManager.java
@@ -11,5 +11,5 @@ public interface QueueManager {
 
     EventQueue createOrTap(String taskId);
 
-    void signalPollingStarted(EventQueue eventQueue) throws InterruptedException;
+    void awaitQueuePollerStart(EventQueue eventQueue) throws InterruptedException;
 }
diff --git a/core/src/main/java/io/a2a/server/requesthandlers/DefaultRequestHandler.java b/core/src/main/java/io/a2a/server/requesthandlers/DefaultRequestHandler.java
@@ -355,7 +355,7 @@ private EnhancedRunnable registerAndExecuteAgentAsync(String taskId, RequestCont
             public void run() {
                 agentExecutor.execute(requestContext, queue);
                 try {
-                    queueManager.signalPollingStarted(queue);
+                    queueManager.awaitQueuePollerStart(queue);
                 } catch (InterruptedException e) {
                     Thread.currentThread().interrupt();
                 }
diff --git a/core/src/main/java/io/a2a/server/requesthandlers/JSONRPCHandler.java b/core/src/main/java/io/a2a/server/requesthandlers/JSONRPCHandler.java
@@ -1,6 +1,5 @@
 package io.a2a.server.requesthandlers;
 
-import static io.a2a.util.AsyncUtils.convertingProcessor;
 import static io.a2a.util.AsyncUtils.createTubeConfig;
 
 import java.util.concurrent.Flow;
diff --git a/quarkus-sdk/src/main/java/io/a2a/server/apps/quarkus/A2AServerRoutes.java b/quarkus-sdk/src/main/java/io/a2a/server/apps/quarkus/A2AServerRoutes.java
@@ -3,6 +3,8 @@
 import static io.vertx.core.http.HttpHeaders.CONTENT_TYPE;
 import static jakarta.ws.rs.core.MediaType.APPLICATION_JSON;
 
+import java.util.concurrent.Executor;
+import java.util.concurrent.Executors;
 import java.util.concurrent.Flow;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.Function;
@@ -66,6 +68,11 @@ public class A2AServerRoutes {
     @ExtendedAgentCard
     Instance<AgentCard> extendedAgentCard;
 
+    // Hook so testing can wait until the MultiSseSupport is subscribes.
+    private static volatile Runnable streamingMultiSseSupportSubscribedRunnable;
+
+    private final Executor executor = Executors.newCachedThreadPool();
+
     @Route(path = "/", methods = {Route.HttpMethod.POST}, consumes = {APPLICATION_JSON}, type = Route.HandlerType.BLOCKING)
     public void invokeJSONRPCHandler(@Body String body, RoutingContext rc) {
         boolean streaming = false;
@@ -93,7 +100,12 @@ public void invokeJSONRPCHandler(@Body String body, RoutingContext rc) {
                         .putHeader(CONTENT_TYPE, APPLICATION_JSON)
                         .end(Json.encodeToBuffer(error));
             } else if (streaming) {
-                MultiSseSupport.subscribeObject(streamingResponse.map(i -> (Object)i), rc);
+                final Multi<? extends JSONRPCResponse<?>> finalStreamingResponse = streamingResponse;
+                executor.execute(() -> {
+                        MultiSseSupport.subscribeObject(
+                                finalStreamingResponse.map(i -> (Object)i), rc);
+                });
+
             } else {
                 rc.response()
                         .setStatusCode(200)
@@ -212,6 +224,10 @@ private static boolean isNonStreamingRequest(String requestBody) {
                 requestBody.contains(A2A.GET_TASK_PUSH_NOTIFICATION_CONFIG_METHOD);
     }
 
+    static void setStreamingMultiSseSupportSubscribedRunnable(Runnable runnable) {
+        streamingMultiSseSupportSubscribedRunnable = runnable;
+    }
+
     // Port of import io.quarkus.vertx.web.runtime.MultiSseSupport, which is considered internal API
     private static class MultiSseSupport {
 
@@ -246,6 +262,12 @@ public static void write(Multi<Buffer> multi, RoutingContext rc) {
                 public void onSubscribe(Flow.Subscription subscription) {
                     this.upstream = subscription;
                     this.upstream.request(1);
+
+                    // Notify tests that we are subscribed
+                    Runnable runnable = streamingMultiSseSupportSubscribedRunnable;
+                    if (runnable != null) {
+                        runnable.run();
+                    }
                 }
 
                 @Override
diff --git a/quarkus-sdk/src/test/java/io/a2a/server/apps/quarkus/A2AServerResourceTest.java b/quarkus-sdk/src/test/java/io/a2a/server/apps/quarkus/A2AServerResourceTest.java

Original file line number	Diff line number	Diff line change
`@@ -11,5 +11,5 @@ public interface QueueManager {`
`11`	`11`
`12`	`12`	`EventQueue createOrTap(String taskId);`
`13`	`13`
`14`		`- void signalPollingStarted(EventQueue eventQueue) throws InterruptedException;`
	`14`	`+ void awaitQueuePollerStart(EventQueue eventQueue) throws InterruptedException;`
`15`	`15`	`}`
Original file line number	Diff line number	Diff line change
`@@ -355,7 +355,7 @@ private EnhancedRunnable registerAndExecuteAgentAsync(String taskId, RequestCont`
`355`	`355`	`public void run() {`
`356`	`356`	`agentExecutor.execute(requestContext, queue);`
`357`	`357`	`try {`
`358`		`- queueManager.signalPollingStarted(queue);`
	`358`	`+ queueManager.awaitQueuePollerStart(queue);`
`359`	`359`	`} catch (InterruptedException e) {`
`360`	`360`	`Thread.currentThread().interrupt();`
`361`	`361`	`}`