Skip to content

Commit c257762

Browse files
chopmob-cloudkabir
chopmob-cloud
and
kabir
authored
fix(client): deliver exactly one terminal SSE callback (a2aproject#1170) (a2aproject#1173)
* fix(client): deliver exactly one terminal SSE callback (a2aproject#1170) The JSON-RPC SSE listeners guarded only onComplete() with a volatile completed flag, so onError, a parse failure, and post-cancellation signals could each still reach the error/completion consumer, delivering more than one terminal callback (including a null completion after an error or cancellation). Introduce one shared atomic transition. AbstractSSEEventListener now owns an AtomicBoolean and a signalTerminal(Throwable) helper that lets the first caller win via compareAndSet and delivers exactly one outcome (a non-null failure or a null normal completion); every later signal is dropped. onError, onComplete and the parse-error path all route through it. The 0.3 compatibility JSON-RPC listener, which does not share the base class, mirrors the same AtomicBoolean pattern. REST is left as is (both native and 0.3 use a no-op completion callback, an API decision). Tests cover complete-then-error, error-then-complete, repeated completion, a 32-thread concurrent race, and a final event followed by completion, asserting exactly one terminal callback in each. Signal normal completion before auto-closing streams for final events so cancellation cannot report a contradictory terminal outcome. Apply this to the shared listener and the v0.3 JSON-RPC listener. Verify that an onError callback caused by closing a stream after a final event does not replace the normal completion outcome in native and v0.3 JSON-RPC listeners. This fixes a2aproject#1170 --------- Signed-off-by: AlgoVoi <chopmob@gmail.com> Co-authored-by: Kabir Khan <kkhan@redhat.com>
1 parent a5e76a0 commit c257762

5 files changed

Lines changed: 285 additions & 46 deletions

File tree

‎client/transport/jsonrpc/src/main/java/org/a2aproject/sdk/client/transport/jsonrpc/sse/SSEEventListener.java‎

Lines changed: 2 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,6 @@
2121
public class SSEEventListener extends AbstractSSEEventListener {
2222

2323
private static final Logger LOGGER = Logger.getLogger(SSEEventListener.class.getName());
24-
private volatile boolean completed = false;
25-
2624
public SSEEventListener(Consumer<StreamingEventKind> eventHandler,
2725
@Nullable Consumer<Throwable> errorHandler) {
2826
super(eventHandler, errorHandler);
@@ -34,21 +32,8 @@ public void onMessage(ServerSentEvent event, @Nullable Future<Void> completableF
3432
}
3533

3634
public void onComplete() {
37-
// Idempotent: only signal completion once, even if called multiple times
38-
if (completed) {
39-
LOGGER.fine("SSEEventListener.onComplete() called again - ignoring (already completed)");
40-
return;
41-
}
42-
completed = true;
43-
44-
// Signal normal stream completion (null error means successful completion)
4535
LOGGER.fine("SSEEventListener.onComplete() called - signaling successful stream completion");
46-
if (getErrorHandler() != null) {
47-
LOGGER.fine("Calling errorHandler.accept(null) to signal successful completion");
48-
getErrorHandler().accept(null);
49-
} else {
50-
LOGGER.warning("errorHandler is null, cannot signal completion");
51-
}
36+
signalTerminal(null);
5237
}
5338

5439
/**
@@ -65,9 +50,7 @@ private void parseAndHandleMessage(String message, @Nullable Future<Void> future
6550
// Delegate to base class for common event handling and auto-close logic
6651
handleEvent(event, future);
6752
} catch (A2AError error) {
68-
if (getErrorHandler() != null) {
69-
getErrorHandler().accept(error);
70-
}
53+
signalTerminal(error);
7154
} catch (JsonProcessingException e) {
7255
throw new RuntimeException(e);
7356
}

‎client/transport/jsonrpc/src/test/java/org/a2aproject/sdk/client/transport/jsonrpc/sse/SSEEventListenerTest.java‎

Lines changed: 129 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,18 @@
44
import static org.junit.jupiter.api.Assertions.assertFalse;
55
import static org.junit.jupiter.api.Assertions.assertInstanceOf;
66
import static org.junit.jupiter.api.Assertions.assertNotNull;
7+
import static org.junit.jupiter.api.Assertions.assertNull;
8+
import static org.junit.jupiter.api.Assertions.assertSame;
79
import static org.junit.jupiter.api.Assertions.assertTrue;
810

11+
import java.util.concurrent.CountDownLatch;
12+
import java.util.concurrent.ExecutorService;
13+
import java.util.concurrent.Executors;
914
import java.util.concurrent.ExecutionException;
1015
import java.util.concurrent.Future;
1116
import java.util.concurrent.TimeUnit;
1217
import java.util.concurrent.TimeoutException;
18+
import java.util.concurrent.atomic.AtomicInteger;
1319
import java.util.concurrent.atomic.AtomicBoolean;
1420
import java.util.concurrent.atomic.AtomicReference;
1521

@@ -248,6 +254,128 @@ public void testOnEventWithFinalTaskStatusUpdateEventEventCancels() throws Excep
248254
}
249255

250256

257+
258+
@Test
259+
public void testOnCompleteThenOnErrorDeliversSingleTerminalCallback() {
260+
AtomicInteger terminalCount = new AtomicInteger(0);
261+
AtomicReference<Throwable> lastArg = new AtomicReference<>();
262+
SSEEventListener listener = new SSEEventListener(
263+
event -> {},
264+
error -> { terminalCount.incrementAndGet(); lastArg.set(error); });
265+
266+
listener.onComplete();
267+
listener.onError(new RuntimeException("late error"), new CancelCapturingFuture());
268+
269+
assertEquals(1, terminalCount.get());
270+
assertNull(lastArg.get());
271+
}
272+
273+
@Test
274+
public void testOnErrorThenOnCompleteDeliversSingleTerminalCallback() {
275+
AtomicInteger terminalCount = new AtomicInteger(0);
276+
AtomicReference<Throwable> lastArg = new AtomicReference<>();
277+
SSEEventListener listener = new SSEEventListener(
278+
event -> {},
279+
error -> { terminalCount.incrementAndGet(); lastArg.set(error); });
280+
281+
RuntimeException boom = new RuntimeException("first error");
282+
listener.onError(boom, new CancelCapturingFuture());
283+
listener.onComplete();
284+
285+
assertEquals(1, terminalCount.get());
286+
assertSame(boom, lastArg.get());
287+
}
288+
289+
@Test
290+
public void testRepeatedOnCompleteDeliversSingleTerminalCallback() {
291+
AtomicInteger terminalCount = new AtomicInteger(0);
292+
AtomicReference<Throwable> lastArg = new AtomicReference<>();
293+
SSEEventListener listener = new SSEEventListener(
294+
event -> {},
295+
error -> { terminalCount.incrementAndGet(); lastArg.set(error); });
296+
297+
listener.onComplete();
298+
listener.onComplete();
299+
listener.onComplete();
300+
301+
assertEquals(1, terminalCount.get());
302+
}
303+
304+
@Test
305+
public void testConcurrentTerminalSignalsDeliverExactlyOneCallback() throws Exception {
306+
AtomicInteger terminalCount = new AtomicInteger(0);
307+
AtomicReference<Throwable> lastArg = new AtomicReference<>();
308+
SSEEventListener listener = new SSEEventListener(
309+
event -> {},
310+
error -> { terminalCount.incrementAndGet(); lastArg.set(error); });
311+
312+
int threads = 32;
313+
ExecutorService pool = Executors.newFixedThreadPool(threads);
314+
CountDownLatch start = new CountDownLatch(1);
315+
CountDownLatch done = new CountDownLatch(threads);
316+
for (int i = 0; i < threads; i++) {
317+
final boolean complete = (i % 2 == 0);
318+
pool.submit(() -> {
319+
try {
320+
start.await();
321+
if (complete) {
322+
listener.onComplete();
323+
} else {
324+
listener.onError(new RuntimeException("race"), new CancelCapturingFuture());
325+
}
326+
} catch (InterruptedException e) {
327+
Thread.currentThread().interrupt();
328+
} finally {
329+
done.countDown();
330+
}
331+
});
332+
}
333+
start.countDown();
334+
assertTrue(done.await(10, TimeUnit.SECONDS));
335+
pool.shutdownNow();
336+
337+
assertEquals(1, terminalCount.get());
338+
}
339+
340+
@Test
341+
public void testFinalEventThenOnCompleteDeliversSingleTerminalCallback() throws Exception {
342+
AtomicInteger terminalCount = new AtomicInteger(0);
343+
SSEEventListener listener = new SSEEventListener(
344+
event -> {},
345+
error -> terminalCount.incrementAndGet());
346+
347+
String eventData = JsonStreamingMessages.STREAMING_STATUS_UPDATE_EVENT_FINAL.substring(
348+
JsonStreamingMessages.STREAMING_STATUS_UPDATE_EVENT_FINAL.indexOf("{"));
349+
CancelCapturingFuture future = new CancelCapturingFuture();
350+
listener.onMessage(new ServerSentEvent(eventData), future);
351+
listener.onComplete();
352+
353+
assertTrue(future.cancelHandlerCalled);
354+
assertEquals(1, terminalCount.get());
355+
}
356+
357+
@Test
358+
public void testFinalEventThenOnErrorDeliversNormalCompletionOnly() {
359+
AtomicInteger terminalCount = new AtomicInteger(0);
360+
AtomicReference<Throwable> terminalError = new AtomicReference<>();
361+
SSEEventListener listener = new SSEEventListener(
362+
event -> {},
363+
error -> {
364+
terminalCount.incrementAndGet();
365+
terminalError.set(error);
366+
});
367+
368+
String eventData = JsonStreamingMessages.STREAMING_STATUS_UPDATE_EVENT_FINAL.substring(
369+
JsonStreamingMessages.STREAMING_STATUS_UPDATE_EVENT_FINAL.indexOf("{"));
370+
CancelCapturingFuture future = new CancelCapturingFuture();
371+
listener.onMessage(new ServerSentEvent(eventData), future);
372+
listener.onError(new RuntimeException("cancelled after final event"), future);
373+
374+
assertTrue(future.cancelHandlerCalled);
375+
assertEquals(1, terminalCount.get());
376+
assertNull(terminalError.get());
377+
}
378+
251379
private static class CancelCapturingFuture implements Future<Void> {
252380
private boolean cancelHandlerCalled;
253381

@@ -280,4 +408,4 @@ public Void get(long timeout, TimeUnit unit) throws InterruptedException, Execut
280408
return null;
281409
}
282410
}
283-
}
411+
}

‎client/transport/spi/src/main/java/org/a2aproject/sdk/client/transport/spi/sse/AbstractSSEEventListener.java‎

Lines changed: 25 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package org.a2aproject.sdk.client.transport.spi.sse;
22

33
import java.util.concurrent.Future;
4+
import java.util.concurrent.atomic.AtomicBoolean;
45
import java.util.function.Consumer;
56
import java.util.logging.Logger;
67

@@ -25,6 +26,7 @@ public abstract class AbstractSSEEventListener {
2526

2627
private final Consumer<StreamingEventKind> eventHandler;
2728
private final @Nullable Consumer<Throwable> errorHandler;
29+
private final AtomicBoolean terminalSignaled = new AtomicBoolean(false);
2830

2931
/**
3032
* Creates a new SSE event listener with the specified handlers.
@@ -73,14 +75,33 @@ protected Consumer<StreamingEventKind> getEventHandler() {
7375
* @param future Optional future for closing the SSE connection
7476
*/
7577
public void onError(Throwable throwable, @Nullable Future<Void> future) {
76-
if (errorHandler != null) {
77-
errorHandler.accept(throwable);
78-
}
78+
signalTerminal(throwable);
7979
if (future != null) {
8080
future.cancel(true); // close SSE channel
8181
}
8282
}
8383

84+
/**
85+
* Delivers exactly one terminal callback for the stream. The first caller to win
86+
* the atomic transition delivers its outcome to the error/completion consumer (a
87+
* non-null {@code error} is a failure, {@code null} is normal completion); every
88+
* later completion, error or post-cancellation signal is dropped, so a single
89+
* streaming request yields exactly one terminal callback.
90+
*
91+
* @param error the failure to report, or {@code null} to signal normal completion
92+
*/
93+
protected void signalTerminal(@Nullable Throwable error) {
94+
if (!terminalSignaled.compareAndSet(false, true)) {
95+
LOGGER.fine("Terminal callback already delivered, ignoring subsequent signal");
96+
return;
97+
}
98+
if (errorHandler != null) {
99+
errorHandler.accept(error);
100+
} else if (error != null) {
101+
LOGGER.warning("errorHandler is null, cannot report terminal error");
102+
}
103+
}
104+
84105
/**
85106
* Processes a parsed streaming event and handles auto-close logic for final events.
86107
* This method encapsulates the common logic for handling events and determining
@@ -97,6 +118,7 @@ protected void handleEvent(StreamingEventKind event, @Nullable Future<Void> futu
97118
// This covers late subscriptions to completed tasks and ensures no connection leaks
98119
if (shouldAutoClose(event) && future != null) {
99120
LOGGER.fine("Auto-closing SSE connection for final event: " + event.getClass().getSimpleName());
121+
signalTerminal(null);
100122
future.cancel(true); // close SSE channel
101123
}
102124
}

‎compat-0.3/client/transport/jsonrpc/src/main/java/org/a2aproject/sdk/compat03/client/transport/jsonrpc/sse/SSEEventListener_v0_3.java‎

Lines changed: 17 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -10,14 +10,15 @@
1010
import org.a2aproject.sdk.compat03.spec.TaskStatusUpdateEvent_v0_3;
1111

1212
import java.util.concurrent.Future;
13+
import java.util.concurrent.atomic.AtomicBoolean;
1314
import java.util.function.Consumer;
1415
import java.util.logging.Logger;
1516

1617
public class SSEEventListener_v0_3 {
1718
private static final Logger LOGGER = Logger.getLogger(SSEEventListener_v0_3.class.getName());
1819
private final Consumer<StreamingEventKind_v0_3> eventHandler;
1920
private final Consumer<Throwable> errorHandler;
20-
private volatile boolean completed = false;
21+
private final AtomicBoolean terminalSignaled = new AtomicBoolean(false);
2122

2223
public SSEEventListener_v0_3(Consumer<StreamingEventKind_v0_3> eventHandler,
2324
Consumer<Throwable> errorHandler) {
@@ -34,50 +35,44 @@ public void onMessage(String message, Future<Void> completableFuture) {
3435
LOGGER.warning("Failed to process JSON message: " + message);
3536
} catch (IllegalArgumentException e) {
3637
LOGGER.warning("Invalid message format: " + message);
37-
if (errorHandler != null) {
38-
errorHandler.accept(e);
39-
}
38+
signalTerminal(e);
4039
completableFuture.cancel(true); // close SSE channel
4140
}
4241
}
4342

4443
public void onError(Throwable throwable, Future<Void> future) {
45-
if (errorHandler != null) {
46-
errorHandler.accept(throwable);
47-
}
44+
signalTerminal(throwable);
4845
future.cancel(true); // close SSE channel
4946
}
5047

51-
public void onComplete() {
52-
// Idempotent: only signal completion once, even if called multiple times
53-
if (completed) {
54-
LOGGER.fine("SSEEventListener.onComplete() called again - ignoring (already completed)");
48+
private void signalTerminal(Throwable error) {
49+
if (!terminalSignaled.compareAndSet(false, true)) {
50+
LOGGER.fine("Terminal callback already delivered, ignoring subsequent signal");
5551
return;
5652
}
57-
completed = true;
58-
59-
// Signal normal stream completion (null error means successful completion)
60-
LOGGER.fine("SSEEventListener.onComplete() called - signaling successful stream completion");
6153
if (errorHandler != null) {
62-
LOGGER.fine("Calling errorHandler.accept(null) to signal successful completion");
63-
errorHandler.accept(null);
64-
} else {
65-
LOGGER.warning("errorHandler is null, cannot signal completion");
54+
errorHandler.accept(error);
55+
} else if (error != null) {
56+
LOGGER.warning("errorHandler is null, cannot report terminal error");
6657
}
6758
}
6859

60+
public void onComplete() {
61+
LOGGER.fine("SSEEventListener.onComplete() called - signaling successful stream completion");
62+
signalTerminal(null);
63+
}
64+
6965
private void handleMessage(JsonObject jsonObject, Future<Void> future) throws JsonProcessingException_v0_3 {
7066
if (jsonObject.has("error")) {
7167
JSONRPCError_v0_3 error = JsonUtil_v0_3.fromJson(jsonObject.get("error").toString(), JSONRPCError_v0_3.class);
72-
if (errorHandler != null) {
73-
errorHandler.accept(error);
74-
}
68+
signalTerminal(error);
7569
} else if (jsonObject.has("result")) {
7670
// result can be a Task, Message, TaskStatusUpdateEvent, or TaskArtifactUpdateEvent
7771
String resultJson = jsonObject.get("result").toString();
7872
StreamingEventKind_v0_3 event = JsonUtil_v0_3.fromJson(resultJson, StreamingEventKind_v0_3.class);
7973
eventHandler.accept(event);
8074
if (event instanceof TaskStatusUpdateEvent_v0_3 tsue && tsue.isFinal()) {
75+
signalTerminal(null);
8176
future.cancel(true); // close SSE channel
8277
}
8378
} else {

0 commit comments

Comments
 (0)