fix: JSONRPCHandler's republisher should happen in a separate thread. (#247)

FIxes #224

Also, add a test for JSONRPCHandler.
GrpcHandler already uses a separate thread, but add a similar test
there.

---------

Co-authored-by: chenkangqiang <[email protected]>

Author: kabir

transport/grpc/src/test/java/io/a2a/transport/grpc/handler/GrpcHandlerTest.java

@@ -6,6 +6,7 @@
 import java.util.List;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import com.google.protobuf.Empty;
 import com.google.protobuf.Struct;
@@ -744,6 +745,69 @@ public void testOnGetAuthenticatedExtendedAgentCard() throws Exception {
         // TODO - getting the authenticated extended agent card isn't support for gRPC right now
     }
 
+    @Test
+    public void testStreamingDoesNotBlockMainThread() throws Exception {
+        GrpcHandler handler = new TestGrpcHandler(AbstractA2ARequestHandlerTest.CARD, requestHandler);
+        
+        // Track if the main thread gets blocked during streaming
+        AtomicBoolean eventReceived = new AtomicBoolean(false);
+        CountDownLatch streamStarted = new CountDownLatch(1);
+        GrpcHandler.setStreamingSubscribedRunnable(streamStarted::countDown);
+        CountDownLatch eventProcessed = new CountDownLatch(1);
+        
+        agentExecutorExecute = (context, eventQueue) -> {
+            // Wait a bit to ensure the main thread continues
+            try {
+                Thread.sleep(100);
+            } catch (InterruptedException e) {
+                Thread.currentThread().interrupt();
+            }
+            eventQueue.enqueueEvent(context.getMessage());
+        };
+
+        // Start streaming with a custom StreamObserver
+        List<StreamResponse> results = new ArrayList<>();
+        List<Throwable> errors = new ArrayList<>();
+        StreamObserver<StreamResponse> streamObserver = new StreamObserver<>() {
+            @Override
+            public void onNext(StreamResponse streamResponse) {
+                results.add(streamResponse);
+                eventReceived.set(true);
+                eventProcessed.countDown();
+            }
+
+            @Override
+            public void onError(Throwable throwable) {
+                errors.add(throwable);
+                eventProcessed.countDown();
+            }
+
+            @Override
+            public void onCompleted() {
+                eventProcessed.countDown();
+            }
+        };
+
+        sendStreamingMessageRequest(handler, streamObserver);
+
+        // The main thread should not be blocked - we should be able to continue immediately
+        Assertions.assertTrue(streamStarted.await(100, TimeUnit.MILLISECONDS), 
+            "Streaming subscription should start quickly without blocking main thread");
+
+        // This proves the main thread is not blocked - we can do other work
+        // Simulate main thread doing other work
+        Thread.sleep(50);
+
+        // Wait for the actual event processing to complete
+        Assertions.assertTrue(eventProcessed.await(2, TimeUnit.SECONDS), 
+            "Event should be processed within reasonable time");
+
+        // Verify we received the event and no errors occurred
+        Assertions.assertTrue(eventReceived.get(), "Should have received streaming event");
+        Assertions.assertTrue(errors.isEmpty(), "Should not have any errors");
+        Assertions.assertEquals(1, results.size(), "Should have received exactly one event");
+    }
+
     private StreamRecorder<SendMessageResponse> sendMessageRequest(GrpcHandler handler) throws Exception {
         SendMessageRequest request = SendMessageRequest.newBuilder()
                 .setRequest(GRPC_MESSAGE)

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

@@ -6,6 +6,7 @@
 import jakarta.inject.Inject;
 
 import java.util.List;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Flow;
 
 import io.a2a.server.ExtendedAgentCard;
@@ -242,37 +243,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);
-                    }
+                CompletableFuture.runAsync(() -> {
+                    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 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())));
+                        @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();
                         }
-                        onComplete();
-                    }
 
-                    @Override
-                    public void onComplete() {
-                        tube.complete();
-                    }
+                        @Override
+                        public void onComplete() {
+                            tube.complete();
+                        }
+                    });
                 });
             });
     }