Rework the streaming subscribed check so it can work from a client

Author: kabir

reference-impl/src/main/java/io/a2a/server/apps/quarkus/A2AServerRoutes.java

@@ -8,7 +8,6 @@
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.Function;
 
-import com.fasterxml.jackson.databind.JsonNode;
 import jakarta.enterprise.inject.Instance;
 import jakarta.inject.Inject;
 import jakarta.inject.Singleton;
@@ -17,8 +16,10 @@
 import com.fasterxml.jackson.core.JsonParseException;
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.core.io.JsonEOFException;
+import com.fasterxml.jackson.databind.JsonNode;
 import io.a2a.server.ExtendedAgentCard;
 import io.a2a.server.requesthandlers.JSONRPCHandler;
+import io.a2a.server.util.async.Internal;
 import io.a2a.spec.AgentCard;
 import io.a2a.spec.CancelTaskRequest;
 import io.a2a.spec.GetTaskPushNotificationConfigRequest;
@@ -44,7 +45,6 @@
 import io.a2a.spec.TaskResubscriptionRequest;
 import io.a2a.spec.UnsupportedOperationError;
 import io.a2a.util.Utils;
-import io.a2a.server.util.async.Internal;
 import io.quarkus.vertx.web.Body;
 import io.quarkus.vertx.web.ReactiveRoutes;
 import io.quarkus.vertx.web.Route;
@@ -69,6 +69,7 @@ public class A2AServerRoutes {
     Instance<AgentCard> extendedAgentCard;
 
     // Hook so testing can wait until the MultiSseSupport is subscribed.
+    // Without this we get intermittent failures
     private static volatile Runnable streamingMultiSseSupportSubscribedRunnable;
 
     @Inject

reference-impl/src/test/java/io/a2a/server/apps/quarkus/A2ATestRoutes.java

@@ -4,6 +4,9 @@
 import static jakarta.ws.rs.core.MediaType.APPLICATION_JSON;
 import static jakarta.ws.rs.core.MediaType.TEXT_PLAIN;
 
+import java.util.concurrent.atomic.AtomicInteger;
+
+import jakarta.annotation.PostConstruct;
 import jakarta.inject.Inject;
 import jakarta.inject.Singleton;
 
@@ -25,6 +28,17 @@ public class A2ATestRoutes {
     @Inject
     TestUtilsBean testUtilsBean;
 
+    @Inject
+    A2AServerRoutes a2AServerRoutes;
+
+    AtomicInteger streamingSubscribedCount = new AtomicInteger(0);
+
+    @PostConstruct
+    public void init() {
+        A2AServerRoutes.setStreamingMultiSseSupportSubscribedRunnable(() -> streamingSubscribedCount.incrementAndGet());
+    }
+
+
     @Route(path = "/test/task", methods = {Route.HttpMethod.POST}, consumes = {APPLICATION_JSON}, type = Route.HandlerType.BLOCKING)
     public void saveTask(@Body String body, RoutingContext rc) {
         try {
@@ -33,7 +47,7 @@ public void saveTask(@Body String body, RoutingContext rc) {
             rc.response()
                 .setStatusCode(200)
                 .end();
-    } catch (Throwable t) {
+        } catch (Throwable t) {
             errorResponse(t, rc);
         }
     }
@@ -118,6 +132,13 @@ public void enqueueTaskArtifactUpdateEvent(@Param String taskId, @Body String bo
         }
     }
 
+    @Route(path = "test/streamingSubscribedCount", methods = {Route.HttpMethod.GET}, produces = {TEXT_PLAIN})
+    public void getStreamingSubscribedCount(RoutingContext rc) {
+        rc.response()
+                .setStatusCode(200)
+                .end(String.valueOf(streamingSubscribedCount.get()));
+    }
+
     private void errorResponse(Throwable t, RoutingContext rc) {
         rc.response()
                 .setStatusCode(200)

reference-impl/src/test/java/io/a2a/server/apps/quarkus/QuarkusA2AServerTest.java

@@ -1,10 +1,6 @@
 package io.a2a.server.apps.quarkus;
 
-import jakarta.inject.Inject;
-
 import io.a2a.server.apps.common.AbstractA2AServerTest;
-import io.a2a.server.events.InMemoryQueueManager;
-import io.a2a.server.tasks.TaskStore;
 import io.quarkus.test.junit.QuarkusTest;
 
 @QuarkusTest
@@ -13,9 +9,4 @@ public class QuarkusA2AServerTest extends AbstractA2AServerTest {
     public QuarkusA2AServerTest() {
         super(8081);
     }
-
-    @Override
-    protected void setStreamingSubscribedRunnable(Runnable runnable) {
-        A2AServerRoutes.setStreamingMultiSseSupportSubscribedRunnable(runnable);
-    }
 }

tests/server-common/src/test/java/io/a2a/server/apps/common/AbstractA2AServerTest.java

@@ -9,6 +9,7 @@
 import static org.wildfly.common.Assert.assertTrue;
 
 import java.io.EOFException;
+import java.io.IOException;
 import java.net.URI;
 import java.net.http.HttpClient;
 import java.net.http.HttpRequest;
@@ -20,13 +21,13 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Stream;
 
 import jakarta.ws.rs.core.MediaType;
 
 import com.fasterxml.jackson.core.JsonProcessingException;
-import io.a2a.server.events.InMemoryQueueManager;
 import io.a2a.spec.AgentCard;
 import io.a2a.spec.Artifact;
 import io.a2a.spec.CancelTaskRequest;
@@ -641,7 +642,8 @@ public void testResubscribeExistingTaskSuccess() throws Exception {
             TaskResubscriptionRequest taskResubscriptionRequest = new TaskResubscriptionRequest("1", new TaskIdParams(MINIMAL_TASK.getId()));
 
             // Count down the latch when the MultiSseSupport on the server has started subscribing
-            setStreamingSubscribedRunnable(taskResubscriptionRequestSent::countDown);
+            awaitStreamingSubscription()
+                    .whenComplete((unused, throwable) -> taskResubscriptionRequestSent.countDown());
 
             CompletableFuture<HttpResponse<Stream<String>>> responseFuture = initialiseStreamingRequest(taskResubscriptionRequest, null);
 
@@ -650,7 +652,6 @@ public void testResubscribeExistingTaskSuccess() throws Exception {
             responseFuture.thenAccept(response -> {
 
                 if (response.statusCode() != 200) {
-                    //errorRef.set(new IllegalStateException("Status code was " + response.statusCode()));
                     throw new IllegalStateException("Status code was " + response.statusCode());
                 }
                 try {
@@ -729,7 +730,7 @@ public void testResubscribeExistingTaskSuccess() throws Exception {
             assertEquals(TaskState.COMPLETED, taskStatusUpdateEvent.getStatus().state());
             assertNotNull(taskStatusUpdateEvent.getStatus().timestamp());
         } finally {
-            setStreamingSubscribedRunnable(null);
+            //setStreamingSubscribedRunnable(null);
             deleteTaskInTaskStore(MINIMAL_TASK.getId());
             executorService.shutdown();
             if (!executorService.awaitTermination(10, TimeUnit.SECONDS)) {
@@ -956,7 +957,7 @@ protected void ensureQueueForTask(String taskId) throws Exception {
     }
 
     protected void enqueueEventOnServer(Event event) throws Exception {
-        String path = null;
+        String path;
         if (event instanceof TaskArtifactUpdateEvent e) {
             path = "test/queue/enqueueTaskArtifactUpdateEvent/" + e.getTaskId();
         } else if (event instanceof TaskStatusUpdateEvent e) {
@@ -979,7 +980,42 @@ protected void enqueueEventOnServer(Event event) throws Exception {
         }
     }
 
-    protected abstract void setStreamingSubscribedRunnable(Runnable runnable);
+    private CompletableFuture<Void> awaitStreamingSubscription() {
+        int cnt = getStreamingSubscribedCount();
+        AtomicInteger initialCount = new AtomicInteger(cnt);
+
+        return CompletableFuture.runAsync(() -> {
+            try {
+                while (true) {
+                    int count = getStreamingSubscribedCount();
+                    if (count > initialCount.get()) {
+                        break;
+                    }
+                    Thread.sleep(500);
+                }
+            } catch (InterruptedException e) {
+                Thread.currentThread().interrupt();
+            }
+        });
+    }
+
+    private int getStreamingSubscribedCount() {
+        HttpClient client = HttpClient.newBuilder()
+                .version(HttpClient.Version.HTTP_2)
+                .build();
+        HttpRequest request = HttpRequest.newBuilder()
+                .uri(URI.create("http://localhost:" + serverPort + "/test/streamingSubscribedCount"))
+                .GET()
+                .build();
+        try {
+            HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString(StandardCharsets.UTF_8));
+            String body = response.body().trim();
+            System.out.println(body);
+            return Integer.valueOf(body);
+        } catch (IOException | InterruptedException e) {
+            throw new RuntimeException(e);
+        }
+    }
 
     private static class BreakException extends RuntimeException {