Use backpressure to see when final event is emitted

Author: kabir

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

@@ -1,6 +1,9 @@
 package io.a2a.server.events;
 
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.Flow;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import io.a2a.spec.A2AError;
 import io.a2a.spec.A2AServerException;
@@ -57,6 +60,21 @@ public Flow.Publisher<EventQueueItem> consumeAll() {
                 .withBufferSize(256);
         return ZeroPublisher.create(conf, tube -> {
             boolean completed = false;
+
+            // Synchronization mechanism for final event write completion
+            // Used to wait for subscriber to finish writing final event before calling tube.complete()
+            final CountDownLatch finalEventWriteLatch = new CountDownLatch(1);
+            final AtomicBoolean waitingForFinalEventWrite = new AtomicBoolean(false);
+
+            // Register callback to detect when subscriber finishes writing events
+            // The subscriber (A2AServerRoutes) calls upstream.request(1) after async write completes
+            tube.whenRequested(requested -> {
+                if (waitingForFinalEventWrite.get()) {
+                    LOGGER.debug("Subscriber requested {} items after final event - async write completed", requested);
+                    finalEventWriteLatch.countDown();
+                }
+            });
+
             try {
                 while (true) {
                     // Check if cancelled by client disconnect
@@ -199,6 +217,10 @@ public Flow.Publisher<EventQueueItem> consumeAll() {
                         // and should not be exposed to API consumers
                         boolean isFinalSent = false;
                         if (!(event instanceof QueueClosedEvent)) {
+                            // Set flag BEFORE sending to ensure whenRequested callback can detect final event
+                            if (isFinalEvent) {
+                                waitingForFinalEventWrite.set(true);
+                            }
                             tube.send(item);
                             isFinalSent = isFinalEvent;
                         }
@@ -208,18 +230,24 @@ public Flow.Publisher<EventQueueItem> consumeAll() {
                             queue.close();
                             LOGGER.debug("Queue closed, breaking loop for queue {}", System.identityHashCode(queue));
 
-                            // CRITICAL: Allow tube buffer to flush before calling tube.complete()
-                            // tube.send() buffers events asynchronously. If we call tube.complete() immediately,
-                            // the stream-end signal can reach the client BEFORE the buffered final event,
-                            // causing the client to close the connection and never receive the final event.
-                            // This is especially important in replicated scenarios where events arrive via Kafka
-                            // and timing is less deterministic. A delay ensures the buffer flushes.
-                            // Increased to 150ms to account for CI environment latency and JVM scheduling delays.
+                            // CRITICAL: Wait for subscriber to process final event before calling tube.complete()
+                            // The subscriber (A2AServerRoutes) writes events asynchronously via response.write().
+                            // Only after the async write completes does it call upstream.request(1), which triggers
+                            // our whenRequested callback. By waiting for this callback, we ensure the final event
+                            // has been written to the HTTP socket before we call tube.complete(), preventing the
+                            // "Stream cancelled" error where the stream-end signal reaches the client before the
+                            // buffered final event.
                             if (isFinalSent) {
                                 try {
-                                    Thread.sleep(150);  // 150ms to allow SSE buffer flush in CI environments
+                                    boolean writeCompleted = finalEventWriteLatch.await(2, TimeUnit.SECONDS);
+                                    if (writeCompleted) {
+                                        LOGGER.debug("Final event write confirmed via backpressure signal");
+                                    } else {
+                                        LOGGER.warn("Timeout waiting for final event write confirmation - proceeding anyway");
+                                    }
                                 } catch (InterruptedException e) {
                                     Thread.currentThread().interrupt();
+                                    LOGGER.warn("Interrupted while waiting for final event write confirmation");
                                 }
                             }
                             break;