fix: fix flaky test testNonBlockingWithMultipleMessages

Author: cmtoan

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

@@ -1274,144 +1274,144 @@ public void testDeletePushNotificationConfigSetWithoutConfigId() throws Exceptio
     @Test
     @Timeout(value = 1, unit = TimeUnit.MINUTES)
     public void testNonBlockingWithMultipleMessages() throws Exception {
-        String multiEventTaskId = "multi-event-test";
+        String multiEventTaskId = "multi-event-test-" + java.util.UUID.randomUUID();
         try {
-            // 1. Send first non-blocking message to create task in WORKING state
-            Message message1 = Message.builder(MESSAGE)
+        // 1. Send first non-blocking message to create task in WORKING state
+        Message message1 = Message.builder(MESSAGE)
                 .taskId(multiEventTaskId)
                 .contextId("test-context")
                 .parts(new TextPart("First request"))
                 .build();
 
-            AtomicReference<String> taskIdRef = new AtomicReference<>();
-            CountDownLatch firstTaskLatch = new CountDownLatch(1);
+        AtomicReference<String> taskIdRef = new AtomicReference<>();
+        CountDownLatch firstTaskLatch = new CountDownLatch(1);
 
-            BiConsumer<ClientEvent, AgentCard> firstMessageConsumer = (event, agentCard) -> {
-                if (event instanceof TaskEvent te) {
-                    taskIdRef.set(te.getTask().id());
-                    firstTaskLatch.countDown();
-                } else if (event instanceof TaskUpdateEvent tue && tue.getUpdateEvent() instanceof TaskStatusUpdateEvent status) {
-                    taskIdRef.set(status.taskId());
-                    firstTaskLatch.countDown();
-                }
-            };
-
-            // Non-blocking message creates task in WORKING state and returns immediately
-            // Queue stays open because task is not in final state
-            getPollingClient().sendMessage(message1, List.of(firstMessageConsumer), null);
-
-            assertTrue(firstTaskLatch.await(10, TimeUnit.SECONDS));
-            String taskId = taskIdRef.get();
-            assertNotNull(taskId);
-            assertEquals(multiEventTaskId, taskId);
-
-            // 2. Resubscribe to task (queue should still be open)
-            CountDownLatch resubEventLatch = new CountDownLatch(2);  // artifact-2 + completion
-            List<io.a2a.spec.UpdateEvent> resubReceivedEvents = new CopyOnWriteArrayList<>();
-            AtomicBoolean resubUnexpectedEvent = new AtomicBoolean(false);
-            AtomicReference<Throwable> resubErrorRef = new AtomicReference<>();
+        BiConsumer<ClientEvent, AgentCard> firstMessageConsumer = (event, agentCard) -> {
+            if (event instanceof TaskEvent te) {
+                taskIdRef.set(te.getTask().id());
+                firstTaskLatch.countDown();
+            } else if (event instanceof TaskUpdateEvent tue && tue.getUpdateEvent() instanceof TaskStatusUpdateEvent status) {
+                taskIdRef.set(status.taskId());
+                firstTaskLatch.countDown();
+            }
+        };
 
-            BiConsumer<ClientEvent, AgentCard> resubConsumer = (event, agentCard) -> {
-                if (event instanceof TaskUpdateEvent tue) {
-                    resubReceivedEvents.add(tue.getUpdateEvent());
-                    resubEventLatch.countDown();
-                } else {
-                    resubUnexpectedEvent.set(true);
-                }
-            };
+        // Non-blocking message creates task in WORKING state and returns immediately
+        // Queue stays open because task is not in final state
+        getPollingClient().sendMessage(message1, List.of(firstMessageConsumer), null);
+
+        assertTrue(firstTaskLatch.await(10, TimeUnit.SECONDS));
+        String taskId = taskIdRef.get();
+        assertNotNull(taskId);
+        assertEquals(multiEventTaskId, taskId);
+
+        // 2. Resubscribe to task (queue should still be open)
+        CountDownLatch resubEventLatch = new CountDownLatch(2);  // artifact-2 + completion
+        List<io.a2a.spec.UpdateEvent> resubReceivedEvents = new CopyOnWriteArrayList<>();
+        AtomicBoolean resubUnexpectedEvent = new AtomicBoolean(false);
+        AtomicReference<Throwable> resubErrorRef = new AtomicReference<>();
+
+        BiConsumer<ClientEvent, AgentCard> resubConsumer = (event, agentCard) -> {
+            if (event instanceof TaskUpdateEvent tue) {
+                resubReceivedEvents.add(tue.getUpdateEvent());
+                resubEventLatch.countDown();
+            } else {
+                resubUnexpectedEvent.set(true);
+            }
+        };
 
-            Consumer<Throwable> resubErrorHandler = error -> {
-                if (!isStreamClosedError(error)) {
-                    resubErrorRef.set(error);
-                }
-            };
+        Consumer<Throwable> resubErrorHandler = error -> {
+            if (!isStreamClosedError(error)) {
+                resubErrorRef.set(error);
+            }
+        };
 
-            // Wait for subscription to be active
-            CountDownLatch subscriptionLatch = new CountDownLatch(1);
-            awaitStreamingSubscription()
+        // Wait for subscription to be active
+        CountDownLatch subscriptionLatch = new CountDownLatch(1);
+        awaitStreamingSubscription()
                 .whenComplete((unused, throwable) -> subscriptionLatch.countDown());
 
-            getClient().resubscribe(new TaskIdParams(taskId),
+        getClient().resubscribe(new TaskIdParams(taskId),
                 List.of(resubConsumer),
                 resubErrorHandler);
 
-            assertTrue(subscriptionLatch.await(15, TimeUnit.SECONDS));
+        assertTrue(subscriptionLatch.await(15, TimeUnit.SECONDS));
 
-            // 3. Send second streaming message to same taskId
-            Message message2 = Message.builder(MESSAGE)
+        // 3. Send second streaming message to same taskId
+        Message message2 = Message.builder(MESSAGE)
                 .taskId(multiEventTaskId) // Same taskId
                 .contextId("test-context")
                 .parts(new TextPart("Second request"))
                 .build();
 
-            CountDownLatch streamEventLatch = new CountDownLatch(2);  // artifact-2 + completion
-            List<io.a2a.spec.UpdateEvent> streamReceivedEvents = new CopyOnWriteArrayList<>();
-            AtomicBoolean streamUnexpectedEvent = new AtomicBoolean(false);
+        CountDownLatch streamEventLatch = new CountDownLatch(2);  // artifact-2 + completion
+        List<io.a2a.spec.UpdateEvent> streamReceivedEvents = new CopyOnWriteArrayList<>();
+        AtomicBoolean streamUnexpectedEvent = new AtomicBoolean(false);
 
-            BiConsumer<ClientEvent, AgentCard> streamConsumer = (event, agentCard) -> {
-                if (event instanceof TaskUpdateEvent tue) {
-                    streamReceivedEvents.add(tue.getUpdateEvent());
-                    streamEventLatch.countDown();
-                } else {
-                    streamUnexpectedEvent.set(true);
-                }
-            };
+        BiConsumer<ClientEvent, AgentCard> streamConsumer = (event, agentCard) -> {
+            if (event instanceof TaskUpdateEvent tue) {
+                streamReceivedEvents.add(tue.getUpdateEvent());
+                streamEventLatch.countDown();
+            } else {
+                streamUnexpectedEvent.set(true);
+            }
+        };
 
-            // Streaming message adds artifact-2 and completes task
-            getClient().sendMessage(message2, List.of(streamConsumer), null);
+        // Streaming message adds artifact-2 and completes task
+        getClient().sendMessage(message2, List.of(streamConsumer), null);
 
-            // 4. Verify both consumers received artifact-2 and completion
-            assertTrue(resubEventLatch.await(10, TimeUnit.SECONDS));
-            assertTrue(streamEventLatch.await(10, TimeUnit.SECONDS));
+        // 4. Verify both consumers received artifact-2 and completion
+        assertTrue(resubEventLatch.await(10, TimeUnit.SECONDS));
+        assertTrue(streamEventLatch.await(10, TimeUnit.SECONDS));
 
-            assertFalse(resubUnexpectedEvent.get());
-            assertFalse(streamUnexpectedEvent.get());
-            assertNull(resubErrorRef.get());
+        assertFalse(resubUnexpectedEvent.get());
+        assertFalse(streamUnexpectedEvent.get());
+        assertNull(resubErrorRef.get());
 
-            // Both should have received 2 events: artifact-2 and completion
-            assertEquals(2, resubReceivedEvents.size());
-            assertEquals(2, streamReceivedEvents.size());
+        // Both should have received 2 events: artifact-2 and completion
+        assertEquals(2, resubReceivedEvents.size());
+        assertEquals(2, streamReceivedEvents.size());
 
-            // Verify resubscription events
-            long resubArtifactCount = resubReceivedEvents.stream()
+        // Verify resubscription events
+        long resubArtifactCount = resubReceivedEvents.stream()
                 .filter(e -> e instanceof TaskArtifactUpdateEvent)
                 .count();
-            assertEquals(1, resubArtifactCount);
+        assertEquals(1, resubArtifactCount);
 
-            long resubCompletionCount = resubReceivedEvents.stream()
+        long resubCompletionCount = resubReceivedEvents.stream()
                 .filter(e -> e instanceof TaskStatusUpdateEvent)
                 .filter(e -> ((TaskStatusUpdateEvent) e).isFinal())
                 .count();
-            assertEquals(1, resubCompletionCount);
+        assertEquals(1, resubCompletionCount);
 
-            // Verify streaming events
-            long streamArtifactCount = streamReceivedEvents.stream()
+        // Verify streaming events
+        long streamArtifactCount = streamReceivedEvents.stream()
                 .filter(e -> e instanceof TaskArtifactUpdateEvent)
                 .count();
-            assertEquals(1, streamArtifactCount);
+        assertEquals(1, streamArtifactCount);
 
-            long streamCompletionCount = streamReceivedEvents.stream()
+        long streamCompletionCount = streamReceivedEvents.stream()
                 .filter(e -> e instanceof TaskStatusUpdateEvent)
                 .filter(e -> ((TaskStatusUpdateEvent) e).isFinal())
                 .count();
-            assertEquals(1, streamCompletionCount);
+        assertEquals(1, streamCompletionCount);
 
-            // Verify artifact-2 details from resubscription
-            TaskArtifactUpdateEvent resubArtifact = (TaskArtifactUpdateEvent) resubReceivedEvents.stream()
+        // Verify artifact-2 details from resubscription
+        TaskArtifactUpdateEvent resubArtifact = (TaskArtifactUpdateEvent) resubReceivedEvents.stream()
                 .filter(e -> e instanceof TaskArtifactUpdateEvent)
                 .findFirst()
                 .orElseThrow();
-            assertEquals("artifact-2", resubArtifact.artifact().artifactId());
-            assertEquals("Second message artifact",
+        assertEquals("artifact-2", resubArtifact.artifact().artifactId());
+        assertEquals("Second message artifact",
                 ((TextPart) resubArtifact.artifact().parts().get(0)).text());
 
-            // Verify artifact-2 details from streaming
-            TaskArtifactUpdateEvent streamArtifact = (TaskArtifactUpdateEvent) streamReceivedEvents.stream()
+        // Verify artifact-2 details from streaming
+        TaskArtifactUpdateEvent streamArtifact = (TaskArtifactUpdateEvent) streamReceivedEvents.stream()
                 .filter(e -> e instanceof TaskArtifactUpdateEvent)
                 .findFirst()
                 .orElseThrow();
-            assertEquals("artifact-2", streamArtifact.artifact().artifactId());
-            assertEquals("Second message artifact",
+        assertEquals("artifact-2", streamArtifact.artifact().artifactId());
+        assertEquals("Second message artifact",
                 ((TextPart) streamArtifact.artifact().parts().get(0)).text());
         } finally {
             deleteTaskInTaskStore(multiEventTaskId);

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

@@ -27,7 +27,7 @@ public void execute(RequestContext context, EventQueue eventQueue) throws A2AErr
                 String taskId = context.getTaskId();
 
                 // Special handling for multi-event test
-                if ("multi-event-test".equals(taskId)) {
+                if (taskId != null && taskId.startsWith("multi-event-test")) {
                     // First call: context.getTask() == null (new task)
                     if (context.getTask() == null) {
                         updater.startWork();