Skip to content

Commit 542b500

Browse files
ehsavoieehsavoie
committed
fix: extract SSE writing into SseResponseWriter, switch to request(Long.MAX_VALUE)
- Replace the MultiSseSupport inner class with a shared SseResponseWriter utility. Switch from request(1) backpressure to request(Long.MAX_VALUE) with a headersSet flag for first-write detection, combining the SSE kickstart comment and the first event into a single write to avoid duplicate error callbacks. - Restore BUFFER_FLUSH_DELAY_MS in EventConsumer: Mutiny's internal demand management can fire request(1) on the underlying publisher independently of the write callback, causing onComplete to race with a pending response.write(). The sleep ensures every write callback confirms delivery before tube.complete() is called. Signed-off-by: Emmanuel Hugonnet <ehugonne@redhat.com>
1 parent dbe8f25 commit 542b500

7 files changed

Lines changed: 202 additions & 535 deletions

File tree

‎compat-0.3/reference/jsonrpc/src/main/java/org/a2aproject/sdk/compat03/server/apps/quarkus/A2AServerRoutes_v0_3.java‎

Lines changed: 3 additions & 87 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@
22

33
import static io.vertx.core.http.HttpHeaders.CONTENT_TYPE;
44
import static jakarta.ws.rs.core.MediaType.APPLICATION_JSON;
5-
import static jakarta.ws.rs.core.MediaType.SERVER_SENT_EVENTS;
65
import static org.a2aproject.sdk.compat03.transport.jsonrpc.context.JSONRPCContextKeys_v0_3.HEADERS_KEY;
76
import static org.a2aproject.sdk.compat03.transport.jsonrpc.context.JSONRPCContextKeys_v0_3.METHOD_NAME_KEY;
87

@@ -24,11 +23,6 @@
2423
import io.quarkus.security.ForbiddenException;
2524
import io.quarkus.security.UnauthorizedException;
2625
import io.smallrye.mutiny.Multi;
27-
import io.vertx.core.AsyncResult;
28-
import io.vertx.core.Handler;
29-
import io.vertx.core.MultiMap;
30-
import io.vertx.core.buffer.Buffer;
31-
import io.vertx.core.http.HttpServerResponse;
3226
import io.vertx.ext.web.Router;
3327
import io.vertx.ext.web.RoutingContext;
3428
import io.vertx.ext.web.handler.BodyHandler;
@@ -66,20 +60,19 @@
6660
import org.a2aproject.sdk.server.ServerCallContext;
6761
import org.a2aproject.sdk.server.auth.UnauthenticatedUser;
6862
import org.a2aproject.sdk.server.auth.User;
63+
import org.a2aproject.sdk.server.common.quarkus.SseResponseWriter;
6964
import org.a2aproject.sdk.server.common.quarkus.VertxSecurityHelper;
7065
import org.a2aproject.sdk.server.extensions.A2AExtensions;
7166
import org.a2aproject.sdk.spec.AgentCard;
7267
import org.jspecify.annotations.Nullable;
73-
import org.slf4j.Logger;
74-
import org.slf4j.LoggerFactory;
7568

7669
@Singleton
7770
public class A2AServerRoutes_v0_3 {
7871

7972
@Inject
8073
JSONRPCHandler_v0_3 jsonRpcHandler;
8174

82-
// Hook so testing can wait until the MultiSseSupport is subscribed.
75+
// Hook so testing can wait until the SSE subscriber is attached.
8376
// Without this we get intermittent failures
8477
private static volatile @Nullable Runnable streamingMultiSseSupportSubscribedRunnable;
8578

@@ -206,7 +199,7 @@ public void invokeJSONRPCHandler(String body, RoutingContext rc) {
206199
AtomicLong eventIdCounter = new AtomicLong(0);
207200
Multi<String> sseEvents = streamingResponse
208201
.map(response -> formatSseEvent(response, eventIdCounter.getAndIncrement()));
209-
MultiSseSupport.writeSseStrings(sseEvents, rc, context);
202+
SseResponseWriter.writeSseStrings(sseEvents, rc, context, streamingMultiSseSupportSubscribedRunnable);
210203
} else {
211204
rc.response()
212205
.setStatusCode(200)
@@ -335,81 +328,4 @@ public String getUsername() {
335328
}
336329
}
337330

338-
private static class MultiSseSupport {
339-
private static final Logger logger = LoggerFactory.getLogger(MultiSseSupport.class);
340-
341-
private MultiSseSupport() {
342-
// Avoid direct instantiation.
343-
}
344-
345-
public static void writeSseStrings(Multi<String> sseStrings, RoutingContext rc, ServerCallContext context) {
346-
HttpServerResponse response = rc.response();
347-
348-
sseStrings.subscribe().withSubscriber(new Flow.Subscriber<String>() {
349-
Flow.@Nullable Subscription upstream;
350-
351-
@Override
352-
public void onSubscribe(Flow.Subscription subscription) {
353-
this.upstream = subscription;
354-
this.upstream.request(1);
355-
356-
response.closeHandler(v -> {
357-
logger.info("SSE connection closed by client, calling EventConsumer.cancel() to stop polling loop");
358-
context.invokeEventConsumerCancelCallback();
359-
subscription.cancel();
360-
});
361-
362-
// Notify tests that we are subscribed
363-
Runnable runnable = streamingMultiSseSupportSubscribedRunnable;
364-
if (runnable != null) {
365-
runnable.run();
366-
}
367-
}
368-
369-
@Override
370-
public void onNext(String sseEvent) {
371-
if (response.bytesWritten() == 0) {
372-
MultiMap headers = response.headers();
373-
if (headers.get(CONTENT_TYPE) == null) {
374-
headers.set(CONTENT_TYPE, SERVER_SENT_EVENTS);
375-
}
376-
headers.set("Cache-Control", "no-cache");
377-
headers.set("X-Accel-Buffering", "no");
378-
response.setChunked(true);
379-
response.setWriteQueueMaxSize(1);
380-
response.write(": SSE stream started\n\n");
381-
}
382-
383-
response.write(Buffer.buffer(sseEvent), new Handler<AsyncResult<Void>>() {
384-
@Override
385-
public void handle(AsyncResult<Void> ar) {
386-
if (ar.failed()) {
387-
java.util.Objects.requireNonNull(upstream).cancel();
388-
rc.fail(ar.cause());
389-
} else {
390-
java.util.Objects.requireNonNull(upstream).request(1);
391-
}
392-
}
393-
});
394-
}
395-
396-
@Override
397-
public void onError(Throwable throwable) {
398-
java.util.Objects.requireNonNull(upstream).cancel();
399-
rc.fail(throwable);
400-
}
401-
402-
@Override
403-
public void onComplete() {
404-
if (response.bytesWritten() == 0) {
405-
MultiMap headers = response.headers();
406-
if (headers.get(CONTENT_TYPE) == null) {
407-
headers.set(CONTENT_TYPE, SERVER_SENT_EVENTS);
408-
}
409-
}
410-
response.end();
411-
}
412-
});
413-
}
414-
}
415331
}

‎compat-0.3/reference/rest/src/main/java/org/a2aproject/sdk/compat03/server/rest/quarkus/A2AServerRoutes_v0_3.java‎

Lines changed: 4 additions & 88 deletions
Original file line numberDiff line numberDiff line change
@@ -2,15 +2,13 @@
22

33
import static io.vertx.core.http.HttpHeaders.CONTENT_TYPE;
44
import static jakarta.ws.rs.core.MediaType.APPLICATION_JSON;
5-
import static jakarta.ws.rs.core.MediaType.SERVER_SENT_EVENTS;
65
import static org.a2aproject.sdk.compat03.transport.rest.context.RestContextKeys_v0_3.HEADERS_KEY;
76
import static org.a2aproject.sdk.compat03.transport.rest.context.RestContextKeys_v0_3.METHOD_NAME_KEY;
87

98
import java.util.HashMap;
109
import java.util.List;
1110
import java.util.Map;
1211
import java.util.Set;
13-
import java.util.concurrent.Flow;
1412
import java.util.concurrent.atomic.AtomicLong;
1513
import java.util.function.Consumer;
1614

@@ -24,11 +22,7 @@
2422
import io.quarkus.security.ForbiddenException;
2523
import io.quarkus.security.UnauthorizedException;
2624
import io.smallrye.mutiny.Multi;
27-
import io.vertx.core.AsyncResult;
2825
import io.vertx.core.Handler;
29-
import io.vertx.core.MultiMap;
30-
import io.vertx.core.buffer.Buffer;
31-
import io.vertx.core.http.HttpServerResponse;
3226
import io.vertx.ext.web.Router;
3327
import io.vertx.ext.web.RoutingContext;
3428
import io.vertx.ext.web.handler.BodyHandler;
@@ -53,20 +47,19 @@
5347
import org.a2aproject.sdk.server.ServerCallContext;
5448
import org.a2aproject.sdk.server.auth.UnauthenticatedUser;
5549
import org.a2aproject.sdk.server.auth.User;
50+
import org.a2aproject.sdk.server.common.quarkus.SseResponseWriter;
5651
import org.a2aproject.sdk.server.common.quarkus.VertxSecurityHelper;
5752
import org.a2aproject.sdk.server.extensions.A2AExtensions;
5853
import org.a2aproject.sdk.spec.AgentCard;
5954
import org.jspecify.annotations.Nullable;
60-
import org.slf4j.Logger;
61-
import org.slf4j.LoggerFactory;
6255

6356
@Singleton
6457
public class A2AServerRoutes_v0_3 {
6558

6659
@Inject
6760
RestHandler_v0_3 jsonRestHandler;
6861

69-
// Hook so testing can wait until the MultiSseSupport is subscribed.
62+
// Hook so testing can wait until the SSE subscriber is attached.
7063
// Without this we get intermittent failures
7164
private static volatile @Nullable Runnable streamingMultiSseSupportSubscribedRunnable;
7265

@@ -202,7 +195,7 @@ public void sendMessageStreaming(String body, RoutingContext rc) {
202195
AtomicLong eventIdCounter = new AtomicLong(0);
203196
Multi<String> sseEvents = Multi.createFrom().publisher(streamingResponse.getPublisher())
204197
.map(json -> formatSseEvent(json, eventIdCounter.getAndIncrement()));
205-
MultiSseSupport.writeSseStrings(sseEvents, rc, context);
198+
SseResponseWriter.writeSseStrings(sseEvents, rc, context, streamingMultiSseSupportSubscribedRunnable);
206199
}
207200
}
208201
}
@@ -302,7 +295,7 @@ public void resubscribeTask(RoutingContext rc) {
302295
AtomicLong eventIdCounter = new AtomicLong(0);
303296
Multi<String> sseEvents = Multi.createFrom().publisher(streamingResponse.getPublisher())
304297
.map(json -> formatSseEvent(json, eventIdCounter.getAndIncrement()));
305-
MultiSseSupport.writeSseStrings(sseEvents, rc, context);
298+
SseResponseWriter.writeSseStrings(sseEvents, rc, context, streamingMultiSseSupportSubscribedRunnable);
306299
}
307300
}
308301
}
@@ -463,81 +456,4 @@ private static boolean hasNonDefaultV10AgentCard() {
463456
return false;
464457
}
465458

466-
private static class MultiSseSupport {
467-
private static final Logger logger = LoggerFactory.getLogger(MultiSseSupport.class);
468-
469-
private MultiSseSupport() {
470-
// Avoid direct instantiation.
471-
}
472-
473-
public static void writeSseStrings(Multi<String> sseStrings, RoutingContext rc, ServerCallContext context) {
474-
HttpServerResponse response = rc.response();
475-
476-
sseStrings.subscribe().withSubscriber(new Flow.Subscriber<String>() {
477-
Flow.@Nullable Subscription upstream;
478-
479-
@Override
480-
public void onSubscribe(Flow.Subscription subscription) {
481-
this.upstream = subscription;
482-
this.upstream.request(1);
483-
484-
response.closeHandler(v -> {
485-
logger.debug("REST SSE connection closed by client, calling EventConsumer.cancel() to stop polling loop");
486-
context.invokeEventConsumerCancelCallback();
487-
subscription.cancel();
488-
});
489-
490-
// Notify tests that we are subscribed
491-
Runnable runnable = streamingMultiSseSupportSubscribedRunnable;
492-
if (runnable != null) {
493-
runnable.run();
494-
}
495-
}
496-
497-
@Override
498-
public void onNext(String sseEvent) {
499-
if (response.bytesWritten() == 0) {
500-
MultiMap headers = response.headers();
501-
if (headers.get(CONTENT_TYPE) == null) {
502-
headers.set(CONTENT_TYPE, SERVER_SENT_EVENTS);
503-
}
504-
headers.set("Cache-Control", "no-cache");
505-
headers.set("X-Accel-Buffering", "no");
506-
response.setChunked(true);
507-
response.setWriteQueueMaxSize(1);
508-
response.write(": SSE stream started\n\n");
509-
}
510-
511-
response.write(Buffer.buffer(sseEvent), new Handler<AsyncResult<Void>>() {
512-
@Override
513-
public void handle(AsyncResult<Void> ar) {
514-
if (ar.failed()) {
515-
java.util.Objects.requireNonNull(upstream).cancel();
516-
rc.fail(ar.cause());
517-
} else {
518-
java.util.Objects.requireNonNull(upstream).request(1);
519-
}
520-
}
521-
});
522-
}
523-
524-
@Override
525-
public void onError(Throwable throwable) {
526-
java.util.Objects.requireNonNull(upstream).cancel();
527-
rc.fail(throwable);
528-
}
529-
530-
@Override
531-
public void onComplete() {
532-
if (response.bytesWritten() == 0) {
533-
MultiMap headers = response.headers();
534-
if (headers.get(CONTENT_TYPE) == null) {
535-
headers.set(CONTENT_TYPE, SERVER_SENT_EVENTS);
536-
}
537-
}
538-
response.end();
539-
}
540-
});
541-
}
542-
}
543459
}

0 commit comments

Comments
 (0)