Skip to content

Commit d34f698

Browse files
WMK666666WMK666666
committed
fix(compat-0.3): signal normal stream completion from the REST transport
The 0.3 REST transport never reported a normally finished stream back to the caller: both completion callbacks handed to postAsyncSSE were empty lambdas and RestSSEEventListener_v0_3 had no onComplete() at all, so a client could not tell a finished stream apart from one that is still running. Do what the native REST transport already does since #1236, and what the 0.3 JSON-RPC transport has always done: turn the connection-end callback into a single terminal errorHandler(null) call. signalTerminal() keeps that to exactly one terminal outcome, so a completion that arrives after a parsing error is dropped instead of producing a second callback. The listener contract is covered directly, including error-then-completion, repeated completion and concurrent terminal signals, all of which fail when the atomic guard is removed. The errorConsumer Javadoc in ClientTransport_v0_3 now spells out the null-on-success convention, matching the native interface. Fixes #1194
1 parent bc9344a commit d34f698

5 files changed

Lines changed: 298 additions & 9 deletions

File tree

‎compat-0.3/client/transport/rest/src/main/java/org/a2aproject/sdk/compat03/client/transport/rest/RestTransport_v0_3.java‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -119,7 +119,8 @@ public void sendMessageStreaming(MessageSendParams_v0_3 messageSendParams, Consu
119119
event -> sseEventListener.onMessage(event.data(), ref.get()),
120120
throwable -> sseEventListener.onError(throwable, ref.get()),
121121
() -> {
122-
// We don't need to do anything special on completion
122+
// Signal normal stream completion to error handler (null error means success)
123+
sseEventListener.onComplete();
123124
}));
124125
} catch (IOException e) {
125126
throw new A2AClientException_v0_3("Failed to send streaming message request: " + e, e);
@@ -310,7 +311,8 @@ public void resubscribe(TaskIdParams_v0_3 request, Consumer<StreamingEventKind_v
310311
event -> sseEventListener.onMessage(event.data(), ref.get()),
311312
throwable -> sseEventListener.onError(throwable, ref.get()),
312313
() -> {
313-
// We don't need to do anything special on completion
314+
// Signal normal stream completion to error handler (null error means success)
315+
sseEventListener.onComplete();
314316
}));
315317
} catch (IOException e) {
316318
throw new A2AClientException_v0_3("Failed to send streaming message request: " + e, e);

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

Lines changed: 33 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package org.a2aproject.sdk.compat03.client.transport.rest.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

@@ -17,6 +18,7 @@ public class RestSSEEventListener_v0_3 {
1718
private static final Logger LOGGER = Logger.getLogger(RestSSEEventListener_v0_3.class.getName());
1819
private final Consumer<StreamingEventKind_v0_3> eventHandler;
1920
private final Consumer<Throwable> errorHandler;
21+
private final AtomicBoolean terminalSignaled = new AtomicBoolean(false);
2022

2123
public RestSSEEventListener_v0_3(Consumer<StreamingEventKind_v0_3> eventHandler,
2224
Consumer<Throwable> errorHandler) {
@@ -31,19 +33,45 @@ public void onMessage(String message, @Nullable Future<Void> completableFuture)
3133
JsonFormat.parser().merge(message, builder);
3234
handleMessage(builder.build());
3335
} catch (InvalidProtocolBufferException e) {
34-
errorHandler.accept(RestErrorMapper_v0_3.mapRestError(message, 500));
36+
signalTerminal(RestErrorMapper_v0_3.mapRestError(message, 500));
3537
}
3638
}
3739

3840
public void onError(Throwable throwable, @Nullable Future<Void> future) {
39-
if (errorHandler != null) {
40-
errorHandler.accept(throwable);
41-
}
41+
signalTerminal(throwable);
4242
if (future != null) {
4343
future.cancel(true); // close SSE channel
4444
}
4545
}
4646

47+
/**
48+
* Signals that the server ended the stream normally. The terminal callback is delivered with
49+
* {@code null}, the same way the native REST listener and the 0.3 JSON-RPC listener report a
50+
* successful completion, so a caller can tell "stream finished" apart from "stream failed".
51+
*/
52+
public void onComplete() {
53+
LOGGER.fine("RestSSEEventListener_v0_3.onComplete() called - signaling successful stream completion");
54+
signalTerminal(null);
55+
}
56+
57+
/**
58+
* Delivers exactly one terminal callback for the stream. The first caller wins, so a completion
59+
* that arrives after an error (or the other way round) does not deliver a second outcome.
60+
*
61+
* @param error the failure to report, or {@code null} to signal normal completion
62+
*/
63+
private void signalTerminal(@Nullable Throwable error) {
64+
if (!terminalSignaled.compareAndSet(false, true)) {
65+
LOGGER.fine("Terminal callback already delivered, ignoring subsequent signal");
66+
return;
67+
}
68+
if (errorHandler != null) {
69+
errorHandler.accept(error);
70+
} else if (error != null) {
71+
LOGGER.warning("errorHandler is null, cannot report terminal error");
72+
}
73+
}
74+
4775
private void handleMessage(StreamResponse response) {
4876
StreamingEventKind_v0_3 event;
4977
switch (response.getPayloadCase()) {
@@ -57,7 +85,7 @@ private void handleMessage(StreamResponse response) {
5785
event = ProtoUtils_v0_3.FromProto.taskArtifactUpdateEvent(response.getArtifactUpdate());
5886
default -> {
5987
LOGGER.warning("Invalid stream response " + response.getPayloadCase());
60-
errorHandler.accept(new IllegalStateException("Invalid stream response from server: " + response.getPayloadCase()));
88+
signalTerminal(new IllegalStateException("Invalid stream response from server: " + response.getPayloadCase()));
6189
return;
6290
}
6391
}

‎compat-0.3/client/transport/rest/src/test/java/org/a2aproject/sdk/compat03/client/transport/rest/RestTransport_v0_3_Test.java‎

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
import java.util.Map;
3030
import java.util.concurrent.CountDownLatch;
3131
import java.util.concurrent.TimeUnit;
32+
import java.util.concurrent.atomic.AtomicInteger;
3233
import java.util.concurrent.atomic.AtomicReference;
3334
import java.util.function.Consumer;
3435
import java.util.logging.Logger;
@@ -70,6 +71,7 @@
7071
public class RestTransport_v0_3_Test {
7172

7273
private static final Logger LOGGER = Logger.getLogger(RestTransport_v0_3_Test.class.getName());
74+
private static final Throwable NOT_INVOKED = new Throwable("terminal callback was not delivered");
7375
private ClientAndServer server;
7476
private static final AgentCard_v0_3 CARD = new AgentCard_v0_3.Builder()
7577
.name("Hello World Agent")
@@ -479,4 +481,96 @@ public void testGetTaskPreservesDataAndMetadataNulls() throws Exception {
479481
assertEquals(Map.of(), part.data().get("nested"));
480482
assertEquals(Arrays.asList(null, nullField, Map.of(), List.of()), part.data().get("items"));
481483
}
484+
485+
/**
486+
* A stream that ends normally has to be reported through the error consumer with a null
487+
* throwable, the same way the native REST transport and the 0.3 JSON-RPC transport already do
488+
* it. Without that signal a caller cannot tell a finished stream from a still running one.
489+
*/
490+
@Test
491+
public void testSendMessageStreamingSignalsNormalCompletion() throws Exception {
492+
this.server.when(
493+
request()
494+
.withMethod("POST")
495+
.withPath("/v1/message:stream")
496+
.withBody(JsonBody.json(SEND_MESSAGE_STREAMING_TEST_REQUEST, MatchType.ONLY_MATCHING_FIELDS))
497+
)
498+
.respond(
499+
response()
500+
.withStatusCode(200)
501+
.withHeader("Content-Type", "text/event-stream")
502+
.withBody(SEND_MESSAGE_STREAMING_TEST_RESPONSE)
503+
);
504+
505+
RestTransport_v0_3 client = new RestTransport_v0_3(CARD);
506+
Message_v0_3 message = new Message_v0_3.Builder()
507+
.role(Message_v0_3.Role.USER)
508+
.parts(Collections.singletonList(new TextPart_v0_3("tell me some jokes")))
509+
.contextId("context-1234")
510+
.messageId("message-1234")
511+
.build();
512+
MessageSendConfiguration_v0_3 configuration = new MessageSendConfiguration_v0_3.Builder()
513+
.acceptedOutputModes(List.of("text"))
514+
.blocking(false)
515+
.build();
516+
MessageSendParams_v0_3 params = new MessageSendParams_v0_3.Builder()
517+
.message(message)
518+
.configuration(configuration)
519+
.build();
520+
521+
CountDownLatch eventLatch = new CountDownLatch(1);
522+
CountDownLatch terminalLatch = new CountDownLatch(1);
523+
AtomicInteger terminalCalls = new AtomicInteger();
524+
AtomicReference<Throwable> terminalError = new AtomicReference<>(NOT_INVOKED);
525+
Consumer<StreamingEventKind_v0_3> eventHandler = event -> eventLatch.countDown();
526+
Consumer<Throwable> errorHandler = error -> {
527+
terminalError.set(error);
528+
terminalCalls.incrementAndGet();
529+
terminalLatch.countDown();
530+
};
531+
client.sendMessageStreaming(params, eventHandler, errorHandler, null);
532+
533+
assertTrue(eventLatch.await(10, TimeUnit.SECONDS), "Should receive streaming event");
534+
assertTrue(terminalLatch.await(10, TimeUnit.SECONDS), "Error consumer should be called when the stream ends");
535+
assertNull(terminalError.get(), "Normal completion has to be reported as a null error");
536+
assertEquals(1, terminalCalls.get(), "A single stream has to deliver exactly one terminal callback");
537+
}
538+
539+
/**
540+
* Same contract as above for a resubscription stream.
541+
*/
542+
@Test
543+
public void testResubscribeSignalsNormalCompletion() throws Exception {
544+
this.server.when(
545+
request()
546+
.withMethod("POST")
547+
.withPath("/v1/tasks/task-1234:subscribe")
548+
)
549+
.respond(
550+
response()
551+
.withStatusCode(200)
552+
.withHeader("Content-Type", "text/event-stream")
553+
.withBody(TASK_RESUBSCRIPTION_REQUEST_TEST_RESPONSE)
554+
);
555+
556+
RestTransport_v0_3 client = new RestTransport_v0_3(CARD);
557+
TaskIdParams_v0_3 taskIdParams = new TaskIdParams_v0_3("task-1234");
558+
559+
CountDownLatch eventLatch = new CountDownLatch(1);
560+
CountDownLatch terminalLatch = new CountDownLatch(1);
561+
AtomicInteger terminalCalls = new AtomicInteger();
562+
AtomicReference<Throwable> terminalError = new AtomicReference<>(NOT_INVOKED);
563+
Consumer<StreamingEventKind_v0_3> eventHandler = event -> eventLatch.countDown();
564+
Consumer<Throwable> errorHandler = error -> {
565+
terminalError.set(error);
566+
terminalCalls.incrementAndGet();
567+
terminalLatch.countDown();
568+
};
569+
client.resubscribe(taskIdParams, eventHandler, errorHandler, null);
570+
571+
assertTrue(eventLatch.await(10, TimeUnit.SECONDS), "Should receive streaming event");
572+
assertTrue(terminalLatch.await(10, TimeUnit.SECONDS), "Error consumer should be called when the stream ends");
573+
assertNull(terminalError.get(), "Normal completion has to be reported as a null error");
574+
assertEquals(1, terminalCalls.get(), "A single stream has to deliver exactly one terminal callback");
575+
}
482576
}

‎compat-0.3/client/transport/rest/src/test/java/org/a2aproject/sdk/compat03/client/transport/rest/sse/RestSSEEventListener_v0_3_Test.java‎

Lines changed: 163 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,163 @@
1+
package org.a2aproject.sdk.compat03.client.transport.rest.sse;
2+
3+
import static org.junit.jupiter.api.Assertions.assertEquals;
4+
import static org.junit.jupiter.api.Assertions.assertNotNull;
5+
import static org.junit.jupiter.api.Assertions.assertNull;
6+
import static org.junit.jupiter.api.Assertions.assertSame;
7+
import static org.junit.jupiter.api.Assertions.assertTrue;
8+
9+
import java.util.concurrent.CountDownLatch;
10+
import java.util.concurrent.ExecutorService;
11+
import java.util.concurrent.Executors;
12+
import java.util.concurrent.Future;
13+
import java.util.concurrent.TimeUnit;
14+
import java.util.concurrent.atomic.AtomicInteger;
15+
import java.util.concurrent.atomic.AtomicReference;
16+
17+
import org.a2aproject.sdk.compat03.spec.StreamingEventKind_v0_3;
18+
import org.junit.jupiter.api.Test;
19+
20+
/**
21+
* Tests for the terminal outcome contract of {@link RestSSEEventListener_v0_3}: a single stream has
22+
* to deliver exactly one terminal callback, no matter how the connection ends.
23+
*/
24+
public class RestSSEEventListener_v0_3_Test {
25+
26+
private static final String INVALID_PAYLOAD = "{ this is not a protobuf StreamResponse";
27+
28+
@Test
29+
public void testOnErrorThenOnCompleteDeliversSingleTerminalCallback() {
30+
AtomicInteger terminalCount = new AtomicInteger(0);
31+
AtomicReference<Throwable> lastArg = new AtomicReference<>();
32+
RestSSEEventListener_v0_3 listener = new RestSSEEventListener_v0_3(
33+
event -> {},
34+
error -> {
35+
terminalCount.incrementAndGet();
36+
lastArg.set(error);
37+
});
38+
39+
RuntimeException boom = new RuntimeException("first error");
40+
listener.onError(boom, null);
41+
listener.onComplete();
42+
43+
assertEquals(1, terminalCount.get());
44+
assertSame(boom, lastArg.get());
45+
}
46+
47+
@Test
48+
public void testOnCompleteThenOnErrorDeliversSingleTerminalCallback() {
49+
AtomicInteger terminalCount = new AtomicInteger(0);
50+
AtomicReference<Throwable> lastArg = new AtomicReference<>();
51+
RestSSEEventListener_v0_3 listener = new RestSSEEventListener_v0_3(
52+
event -> {},
53+
error -> {
54+
terminalCount.incrementAndGet();
55+
lastArg.set(error);
56+
});
57+
58+
listener.onComplete();
59+
listener.onError(new RuntimeException("late error"), null);
60+
61+
assertEquals(1, terminalCount.get());
62+
assertNull(lastArg.get());
63+
}
64+
65+
@Test
66+
public void testRepeatedOnCompleteDeliversSingleTerminalCallback() {
67+
AtomicInteger terminalCount = new AtomicInteger(0);
68+
AtomicReference<Throwable> lastArg = new AtomicReference<>();
69+
RestSSEEventListener_v0_3 listener = new RestSSEEventListener_v0_3(
70+
event -> {},
71+
error -> {
72+
terminalCount.incrementAndGet();
73+
lastArg.set(error);
74+
});
75+
76+
listener.onComplete();
77+
listener.onComplete();
78+
listener.onComplete();
79+
80+
assertEquals(1, terminalCount.get());
81+
assertNull(lastArg.get());
82+
}
83+
84+
@Test
85+
public void testConcurrentTerminalSignalsDeliverExactlyOneCallback() throws Exception {
86+
AtomicInteger terminalCount = new AtomicInteger(0);
87+
RestSSEEventListener_v0_3 listener = new RestSSEEventListener_v0_3(
88+
event -> {},
89+
error -> terminalCount.incrementAndGet());
90+
91+
int threads = 32;
92+
ExecutorService pool = Executors.newFixedThreadPool(threads);
93+
CountDownLatch start = new CountDownLatch(1);
94+
CountDownLatch done = new CountDownLatch(threads);
95+
for (int i = 0; i < threads; i++) {
96+
final boolean complete = (i % 2 == 0);
97+
pool.submit(() -> {
98+
try {
99+
start.await();
100+
if (complete) {
101+
listener.onComplete();
102+
} else {
103+
listener.onError(new RuntimeException("race"), null);
104+
}
105+
} catch (InterruptedException e) {
106+
Thread.currentThread().interrupt();
107+
} finally {
108+
done.countDown();
109+
}
110+
});
111+
}
112+
start.countDown();
113+
assertTrue(done.await(10, TimeUnit.SECONDS));
114+
pool.shutdownNow();
115+
116+
assertEquals(1, terminalCount.get());
117+
}
118+
119+
@Test
120+
public void testParsingFailureThenOnCompleteDeliversSingleTerminalCallback() {
121+
AtomicInteger terminalCount = new AtomicInteger(0);
122+
AtomicReference<Throwable> lastArg = new AtomicReference<>();
123+
RestSSEEventListener_v0_3 listener = new RestSSEEventListener_v0_3(
124+
event -> {},
125+
error -> {
126+
terminalCount.incrementAndGet();
127+
lastArg.set(error);
128+
});
129+
130+
listener.onMessage(INVALID_PAYLOAD, null);
131+
listener.onComplete();
132+
133+
assertEquals(1, terminalCount.get());
134+
assertNotNull(lastArg.get(), "The parsing failure has to stay the single terminal outcome");
135+
}
136+
137+
@Test
138+
public void testOnCompleteDeliversNullToErrorHandler() {
139+
AtomicInteger terminalCount = new AtomicInteger(0);
140+
AtomicReference<Throwable> lastArg = new AtomicReference<>(new Throwable("not invoked"));
141+
AtomicReference<StreamingEventKind_v0_3> receivedEvent = new AtomicReference<>();
142+
RestSSEEventListener_v0_3 listener = new RestSSEEventListener_v0_3(
143+
receivedEvent::set,
144+
error -> {
145+
terminalCount.incrementAndGet();
146+
lastArg.set(error);
147+
});
148+
149+
listener.onComplete();
150+
151+
assertEquals(1, terminalCount.get());
152+
assertNull(lastArg.get(), "A clean completion has to be reported as a null error");
153+
assertNull(receivedEvent.get(), "Completion must not fabricate an event");
154+
}
155+
156+
@Test
157+
public void testOnCompleteWithoutErrorHandlerDoesNotThrow() {
158+
RestSSEEventListener_v0_3 listener = new RestSSEEventListener_v0_3(event -> {}, null);
159+
160+
listener.onComplete();
161+
listener.onComplete();
162+
}
163+
}

‎compat-0.3/client/transport/spi/src/main/java/org/a2aproject/sdk/compat03/client/transport/spi/ClientTransport_v0_3.java‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,8 @@ EventKind_v0_3 sendMessage(MessageSendParams_v0_3 request, @Nullable ClientCallC
3939
*
4040
* @param request the message send parameters
4141
* @param eventConsumer consumer that will receive streaming events as they arrive
42-
* @param errorConsumer consumer that will be called if an error occurs during streaming
42+
* @param errorConsumer consumer called when the stream ends; called with {@code null} to signal
43+
* clean completion, or with a non-null {@link Throwable} if an error occurs
4344
* @param context optional client call context for the request (may be {@code null})
4445
* @throws A2AClientException_v0_3 if setting up the streaming connection fails
4546
*/
@@ -117,7 +118,8 @@ void deleteTaskPushNotificationConfigurations(
117118
*
118119
* @param request the task ID parameters specifying which task to resubscribe to
119120
* @param eventConsumer consumer that will receive streaming events as they arrive
120-
* @param errorConsumer consumer that will be called if an error occurs during streaming
121+
* @param errorConsumer consumer called when the stream ends; called with {@code null} to signal
122+
* clean completion, or with a non-null {@link Throwable} if an error occurs
121123
* @param context optional client call context for the request (may be {@code null})
122124
* @throws A2AClientException_v0_3 if resubscribing to the task fails for any reason
123125
*/

0 commit comments

Comments
 (0)