fix: JSONRPCHandler stream block

Author: chenkangqiang

transport/jsonrpc/src/main/java/io/a2a/jsonrpc/handler/JSONRPCHandler.java

@@ -5,6 +5,7 @@
 import jakarta.inject.Inject;
 
 import java.util.List;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Flow;
 
 import io.a2a.server.PublicAgentCard;
@@ -212,37 +213,39 @@ private Flow.Publisher<SendStreamingMessageResponse> convertToSendStreamingMessa
             // We can't use the normal convertingProcessor since that propagates any errors as an error handled
             // via Subscriber.onError() rather than as part of the SendStreamingResponse payload
             return ZeroPublisher.create(createTubeConfig(), tube -> {
-                publisher.subscribe(new Flow.Subscriber<StreamingEventKind>() {
-                    Flow.Subscription subscription;
-                    @Override
-                    public void onSubscribe(Flow.Subscription subscription) {
-                        this.subscription = subscription;
-                        subscription.request(1);
-                    }
-
-                    @Override
-                    public void onNext(StreamingEventKind item) {
-                        tube.send(new SendStreamingMessageResponse(requestId, item));
-                        subscription.request(1);
-                    }
-
-                    @Override
-                    public void onError(Throwable throwable) {
-                        if (throwable instanceof JSONRPCError jsonrpcError) {
-                            tube.send(new SendStreamingMessageResponse(requestId, jsonrpcError));
-                        } else {
-                            tube.send(
-                                    new SendStreamingMessageResponse(
-                                            requestId, new
-                                            InternalError(throwable.getMessage())));
+                CompletableFuture.runAsync(() -> {
+                    publisher.subscribe(new Flow.Subscriber<StreamingEventKind>() {
+                        Flow.Subscription subscription;
+                        @Override
+                        public void onSubscribe(Flow.Subscription subscription) {
+                            this.subscription = subscription;
+                            subscription.request(1);
                         }
-                        onComplete();
-                    }
 
-                    @Override
-                    public void onComplete() {
-                        tube.complete();
-                    }
+                        @Override
+                        public void onNext(StreamingEventKind item) {
+                            tube.send(new SendStreamingMessageResponse(requestId, item));
+                            subscription.request(1);
+                        }
+
+                        @Override
+                        public void onError(Throwable throwable) {
+                            if (throwable instanceof JSONRPCError jsonrpcError) {
+                                tube.send(new SendStreamingMessageResponse(requestId, jsonrpcError));
+                            } else {
+                                tube.send(
+                                        new SendStreamingMessageResponse(
+                                                requestId, new
+                                                InternalError(throwable.getMessage())));
+                            }
+                            onComplete();
+                        }
+
+                        @Override
+                        public void onComplete() {
+                            tube.complete();
+                        }
+                    });
                 });
             });
     }

transport/jsonrpc/src/test/java/io/a2a/jsonrpc/handler/JSONRPCHandlerTest.java

@@ -4,10 +4,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
-import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.Executors;
-import java.util.concurrent.Flow;
-import java.util.concurrent.TimeUnit;
+import java.util.concurrent.*;
 import java.util.concurrent.atomic.AtomicReference;
 
 import io.a2a.server.ServerCallContext;
@@ -269,7 +266,7 @@ public void testOnMessageErrorMocks() {
     }
 
     @Test
-    public void testOnMessageStreamNewMessageSuccess() {
+    public void testOnMessageStreamNewMessageSuccess() throws InterruptedException {
         JSONRPCHandler handler = new JSONRPCHandler(CARD, requestHandler);
         agentExecutorExecute = (context, eventQueue) -> {
             eventQueue.enqueueEvent(context.getTask() != null ? context.getTask() : context.getMessage());
@@ -314,6 +311,8 @@ public void onComplete() {
             }
         });
 
+        latch.await();
+
         // The Python implementation has several events emitted since it uses mocks. Also, in the
         // implementation, a Message is considered a 'final' Event in EventConsumer.consumeAll()
         // so there would be no more Events.
@@ -359,6 +358,7 @@ public void testOnMessageStreamNewMessageSuccessMocks() {
             response = handler.onMessageSendStream(request, callContext);
         }
 
+        CompletableFuture<Void> future = new CompletableFuture<>();
         List<Event> results = new ArrayList<>();
 
         response.subscribe(new Flow.Subscriber<SendStreamingMessageResponse>() {
@@ -378,15 +378,17 @@ public void onNext(SendStreamingMessageResponse item) {
 
             @Override
             public void onError(Throwable throwable) {
-
+                future.completeExceptionally(throwable);
             }
 
             @Override
             public void onComplete() {
-
+                future.complete(null);
             }
         });
 
+        future.join();
+
         Assertions.assertEquals(events, results);
     }
 
@@ -504,6 +506,7 @@ public void testOnMessageStreamNewMessageExistingTaskSuccessMocks() {
             response = handler.onMessageSendStream(request, callContext);
         }
 
+        CompletableFuture<Void> future = new CompletableFuture<>();
         List<Event> results = new ArrayList<>();
 
         // Unlike testOnMessageStreamNewMessageExistingTaskSuccess() the ZeroPublisher.fromIterable()
@@ -526,15 +529,17 @@ public void onNext(SendStreamingMessageResponse item) {
 
             @Override
             public void onError(Throwable throwable) {
-
+                future.completeExceptionally(throwable);
             }
 
             @Override
             public void onComplete() {
-
+                future.complete(null);
             }
         });
 
+        future.join();
+
         Assertions.assertEquals(events, results);
     }
 
@@ -710,7 +715,7 @@ public void testOnResubscribeExistingTaskSuccess() {
                         callContext);
         Assertions.assertNull(smr.getError());
 
-
+        CompletableFuture<Void> future = new CompletableFuture<>();
         List<StreamingEventKind> results = new ArrayList<>();
 
         response.subscribe(new Flow.Subscriber<>() {
@@ -731,14 +736,18 @@ public void onNext(SendStreamingMessageResponse item) {
             @Override
             public void onError(Throwable throwable) {
                 subscription.cancel();
+                future.completeExceptionally(throwable);
             }
 
             @Override
             public void onComplete() {
                 subscription.cancel();
+                future.complete(null);
             }
         });
 
+        future.join();
+
         // The Python implementation has several events emitted since it uses mocks.
         //
         // See testOnMessageStreamNewMessageExistingTaskSuccessMocks() for a test more similar to the Python implementation
@@ -776,6 +785,7 @@ public void testOnResubscribeExistingTaskSuccessMocks() throws Exception {
             response = handler.onResubscribeToTask(request, callContext);
         }
 
+        CompletableFuture<Void> future = new CompletableFuture<>();
         List<StreamingEventKind> results = new ArrayList<>();
 
         // Unlike testOnResubscribeExistingTaskSuccess() the ZeroPublisher.fromIterable()
@@ -799,14 +809,18 @@ public void onNext(SendStreamingMessageResponse item) {
             @Override
             public void onError(Throwable throwable) {
                 subscription.cancel();
+                future.completeExceptionally(throwable);
             }
 
             @Override
             public void onComplete() {
                 subscription.cancel();
+                future.complete(null);
             }
         });
 
+        future.join();
+
         // The Python implementation has several events emitted since it uses mocks.
         //
         // See testOnMessageStreamNewMessageExistingTaskSuccessMocks() for a test more similar to the Python implementation
@@ -1140,6 +1154,7 @@ public void testOnMessageStreamTaskIdMismatch() {
         SendStreamingMessageRequest request = new SendStreamingMessageRequest("1", new MessageSendParams(MESSAGE, null, null));
         Flow.Publisher<SendStreamingMessageResponse> response = handler.onMessageSendStream(request, callContext);
 
+        CompletableFuture<Void> future = new CompletableFuture<>();
         List<SendStreamingMessageResponse> results = new ArrayList<>();
         AtomicReference<Throwable> error = new AtomicReference<>();
 
@@ -1161,14 +1176,18 @@ public void onNext(SendStreamingMessageResponse item) {
             public void onError(Throwable throwable) {
                 error.set(throwable);
                 subscription.cancel();
+                future.completeExceptionally(throwable);
             }
 
             @Override
             public void onComplete() {
                 subscription.cancel();
+                future.complete(null);
             }
         });
 
+        future.join();
+
         Assertions.assertNull(error.get());
         Assertions.assertEquals(1, results.size());
         Assertions.assertInstanceOf(InternalError.class, results.get(0).getError());