Skip to content

Commit 003e6b3

Browse files
kamikaz1kkamikaz1k
committed
test: cover streamable HTTP edge cases
1 parent 0b2a5a8 commit 003e6b3

4 files changed

Lines changed: 386 additions & 18 deletions

File tree

‎acp-core/src/test/java/com/agentclientprotocol/sdk/client/transport/StreamableHttpAcpClientTransportIntegrationTest.java‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
import java.util.concurrent.TimeUnit;
2727
import java.util.concurrent.atomic.AtomicBoolean;
2828
import java.util.concurrent.atomic.AtomicInteger;
29+
import java.util.concurrent.atomic.AtomicReference;
2930

3031
import com.agentclientprotocol.sdk.AcpTestFixtures;
3132
import com.agentclientprotocol.sdk.client.AcpAsyncClient;
@@ -111,6 +112,7 @@ void permissionRequestRoundTripsOnSessionStream() throws Exception {
111112
assertThat(prompt.stopReason()).isEqualTo(AcpSchema.StopReason.END_TURN);
112113
assertThat(permissionRequests).hasValue(1);
113114
assertThat(fixture.permissionResponseReceived()).isTrue();
115+
assertThat(fixture.permissionResponseSessionHeader()).isEqualTo(session.sessionId());
114116

115117
client.closeGracefully().block(TIMEOUT);
116118
}
@@ -221,6 +223,8 @@ private static final class FixtureServer implements AutoCloseable {
221223

222224
private final AtomicBoolean permissionResponseReceived = new AtomicBoolean(false);
223225

226+
private final AtomicReference<String> permissionResponseSessionHeader = new AtomicReference<>();
227+
224228
private final AtomicBoolean sessionLoadStreamWasOpenBeforePost = new AtomicBoolean(false);
225229

226230
private final CompletableFuture<AcpSchema.JSONRPCResponse> permissionResponse = new CompletableFuture<>();
@@ -260,6 +264,10 @@ boolean permissionResponseReceived() {
260264
return permissionResponseReceived.get();
261265
}
262266

267+
String permissionResponseSessionHeader() {
268+
return permissionResponseSessionHeader.get();
269+
}
270+
263271
boolean sessionLoadStreamWasOpenBeforePost() {
264272
return sessionLoadStreamWasOpenBeforePost.get();
265273
}
@@ -316,6 +324,7 @@ private void handlePost(HttpExchange exchange) throws IOException {
316324
private void handleAcceptedMessage(AcpSchema.JSONRPCMessage message, String sessionHeader) {
317325
if (message instanceof AcpSchema.JSONRPCResponse response) {
318326
permissionResponseReceived.set(true);
327+
permissionResponseSessionHeader.set(sessionHeader);
319328
permissionResponse.complete(response);
320329
return;
321330
}

‎acp-core/src/test/java/com/agentclientprotocol/sdk/client/transport/StreamableHttpAcpClientTransportTest.java‎

Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,11 @@
1818
import java.util.Map;
1919
import java.util.concurrent.BlockingQueue;
2020
import java.util.concurrent.CompletableFuture;
21+
import java.util.concurrent.ConcurrentHashMap;
2122
import java.util.concurrent.CountDownLatch;
23+
import java.util.concurrent.ExecutorService;
24+
import java.util.concurrent.Executors;
25+
import java.util.concurrent.Future;
2226
import java.util.concurrent.LinkedBlockingQueue;
2327
import java.util.concurrent.TimeUnit;
2428
import java.util.concurrent.atomic.AtomicInteger;
@@ -289,11 +293,155 @@ void sessionNewErrorResponseIsDeliveredWithoutOpeningSessionStream() throws Exce
289293
}
290294
}
291295

296+
@Test
297+
void concurrentSessionSseEventsAreSerializedIntoInboundSink() throws Exception {
298+
HttpClient httpClient = mock(HttpClient.class);
299+
Map<String, PipedOutputStream> sessionWriters = new ConcurrentHashMap<>();
300+
BlockingQueue<AcpSchema.JSONRPCMessage> inboundMessages = new LinkedBlockingQueue<>();
301+
302+
when(httpClient.sendAsync(any(), any())).thenAnswer(invocation -> {
303+
HttpRequest request = invocation.getArgument(0);
304+
if ("POST".equals(request.method())
305+
&& request.headers().firstValue("Acp-Connection-Id").isEmpty()) {
306+
String initializeResponse = jsonMapper.writeValueAsString(AcpTestFixtures
307+
.createJsonRpcResponse("init-1", AcpTestFixtures.createInitializeResponse()));
308+
return CompletableFuture.completedFuture(response(200,
309+
Map.of("Content-Type", "application/json", "Acp-Connection-Id", "conn-1"),
310+
initializeResponse));
311+
}
312+
if ("GET".equals(request.method())
313+
&& request.headers().firstValue("Acp-Session-Id").isEmpty()) {
314+
return CompletableFuture.completedFuture(
315+
response(200, Map.of("Content-Type", "text/event-stream"), emptyBody()));
316+
}
317+
if ("GET".equals(request.method())) {
318+
String sessionId = request.headers().firstValue("Acp-Session-Id").orElseThrow();
319+
try {
320+
PipedInputStream sessionBody = new PipedInputStream(16 * 1024);
321+
PipedOutputStream writer = new PipedOutputStream(sessionBody);
322+
sessionWriters.put(sessionId, writer);
323+
return CompletableFuture.completedFuture(
324+
response(200, Map.of("Content-Type", "text/event-stream"), sessionBody));
325+
}
326+
catch (Exception e) {
327+
CompletableFuture<HttpResponse<InputStream>> failed = new CompletableFuture<>();
328+
failed.completeExceptionally(e);
329+
return failed;
330+
}
331+
}
332+
return CompletableFuture.completedFuture(response(202, Map.of(), null));
333+
});
334+
335+
StreamableHttpAcpClientTransport transport = new StreamableHttpAcpClientTransport(
336+
URI.create("https://localhost:8443/acp"), jsonMapper, httpClient);
337+
ExecutorService executor = Executors.newFixedThreadPool(2);
338+
try {
339+
transport.connect(message -> message.doOnNext(inboundMessages::add).then(Mono.empty())).block();
340+
transport.sendMessage(AcpTestFixtures.createJsonRpcRequest(AcpSchema.METHOD_INITIALIZE, "init-1",
341+
AcpTestFixtures.createInitializeRequest()))
342+
.block();
343+
awaitResponse(inboundMessages, "init-1");
344+
345+
transport.sendMessage(AcpTestFixtures.createJsonRpcRequest(AcpSchema.METHOD_SESSION_LOAD, "load-1",
346+
new AcpSchema.LoadSessionRequest("sess-1", "/workspace/one", List.of())))
347+
.block();
348+
transport.sendMessage(AcpTestFixtures.createJsonRpcRequest(AcpSchema.METHOD_SESSION_LOAD, "load-2",
349+
new AcpSchema.LoadSessionRequest("sess-2", "/workspace/two", List.of())))
350+
.block();
351+
352+
Future<?> first = executor.submit(() -> {
353+
writeSessionUpdates(sessionWriters.get("sess-1"), "sess-1", 50);
354+
return null;
355+
});
356+
Future<?> second = executor.submit(() -> {
357+
writeSessionUpdates(sessionWriters.get("sess-2"), "sess-2", 50);
358+
return null;
359+
});
360+
first.get(1, TimeUnit.SECONDS);
361+
second.get(1, TimeUnit.SECONDS);
362+
363+
awaitNotifications(inboundMessages, 100);
364+
}
365+
finally {
366+
sessionWriters.values().forEach(writer -> {
367+
try {
368+
writer.close();
369+
}
370+
catch (Exception ignored) {
371+
}
372+
});
373+
executor.shutdownNow();
374+
transport.close();
375+
}
376+
}
377+
378+
@Test
379+
void malformedSseEventDoesNotStopConnectionReader() throws Exception {
380+
HttpClient httpClient = mock(HttpClient.class);
381+
PipedInputStream connectionStreamBody = new PipedInputStream();
382+
PipedOutputStream connectionStreamWriter = new PipedOutputStream(connectionStreamBody);
383+
BlockingQueue<AcpSchema.JSONRPCMessage> inboundMessages = new LinkedBlockingQueue<>();
384+
385+
when(httpClient.sendAsync(any(), any())).thenAnswer(invocation -> {
386+
HttpRequest request = invocation.getArgument(0);
387+
if ("POST".equals(request.method())
388+
&& request.headers().firstValue("Acp-Connection-Id").isEmpty()) {
389+
String initializeResponse = jsonMapper.writeValueAsString(AcpTestFixtures
390+
.createJsonRpcResponse("init-1", AcpTestFixtures.createInitializeResponse()));
391+
return CompletableFuture.completedFuture(response(200,
392+
Map.of("Content-Type", "application/json", "Acp-Connection-Id", "conn-1"),
393+
initializeResponse));
394+
}
395+
if ("GET".equals(request.method())
396+
&& request.headers().firstValue("Acp-Session-Id").isEmpty()) {
397+
return CompletableFuture.completedFuture(
398+
response(200, Map.of("Content-Type", "text/event-stream"), connectionStreamBody));
399+
}
400+
return CompletableFuture.completedFuture(response(202, Map.of(), null));
401+
});
402+
403+
StreamableHttpAcpClientTransport transport = new StreamableHttpAcpClientTransport(
404+
URI.create("https://localhost:8443/acp"), jsonMapper, httpClient);
405+
try {
406+
transport.connect(message -> message.doOnNext(inboundMessages::add).then(Mono.empty())).block();
407+
transport.sendMessage(AcpTestFixtures.createJsonRpcRequest(AcpSchema.METHOD_INITIALIZE, "init-1",
408+
AcpTestFixtures.createInitializeRequest()))
409+
.block();
410+
awaitResponse(inboundMessages, "init-1");
411+
412+
writeRawSse(connectionStreamWriter, "{ nope");
413+
transport.sendMessage(new AcpSchema.JSONRPCRequest(AcpSchema.JSONRPC_VERSION, "ping-1",
414+
"extension/ping", Map.of()))
415+
.block();
416+
writeSse(connectionStreamWriter, AcpTestFixtures.createJsonRpcResponse("ping-1", Map.of()));
417+
418+
assertThat(awaitResponse(inboundMessages, "ping-1")).isNotNull();
419+
}
420+
finally {
421+
connectionStreamWriter.close();
422+
transport.close();
423+
}
424+
}
425+
426+
private void writeSessionUpdates(PipedOutputStream writer, String sessionId, int count) throws Exception {
427+
for (int i = 0; i < count; i++) {
428+
writeSse(writer, new AcpSchema.JSONRPCNotification(AcpSchema.METHOD_SESSION_UPDATE,
429+
new AcpSchema.SessionNotification(sessionId,
430+
new AcpSchema.AgentMessageChunk("agent_message_chunk",
431+
new AcpSchema.TextContent(sessionId + "-" + i)))));
432+
}
433+
}
434+
292435
private void writeSse(PipedOutputStream writer, AcpSchema.JSONRPCMessage message) throws Exception {
293436
writer.write(("data: " + jsonMapper.writeValueAsString(message) + "\n\n").getBytes(StandardCharsets.UTF_8));
294437
writer.flush();
295438
}
296439

440+
private void writeRawSse(PipedOutputStream writer, String data) throws Exception {
441+
writer.write(("data: " + data + "\n\n").getBytes(StandardCharsets.UTF_8));
442+
writer.flush();
443+
}
444+
297445
private AcpSchema.JSONRPCResponse awaitResponse(BlockingQueue<AcpSchema.JSONRPCMessage> messages, Object id)
298446
throws Exception {
299447
long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(2);
@@ -306,6 +454,22 @@ private AcpSchema.JSONRPCResponse awaitResponse(BlockingQueue<AcpSchema.JSONRPCM
306454
throw new AssertionError("Timed out waiting for response " + id);
307455
}
308456

457+
private void awaitNotifications(BlockingQueue<AcpSchema.JSONRPCMessage> messages, int expected) throws Exception {
458+
long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(2);
459+
int count = 0;
460+
while (System.nanoTime() < deadline) {
461+
AcpSchema.JSONRPCMessage message = messages.poll(50, TimeUnit.MILLISECONDS);
462+
if (message instanceof AcpSchema.JSONRPCNotification notification
463+
&& AcpSchema.METHOD_SESSION_UPDATE.equals(notification.method())) {
464+
count++;
465+
if (count == expected) {
466+
return;
467+
}
468+
}
469+
}
470+
throw new AssertionError("Timed out waiting for " + expected + " notifications; received " + count);
471+
}
472+
309473
private InputStream emptyBody() {
310474
return new ByteArrayInputStream(new byte[0]);
311475
}

‎acp-streamable-http-jetty/src/main/java/com/agentclientprotocol/sdk/agent/transport/StreamableHttpAcpAgentTransport.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -475,7 +475,8 @@ private boolean isInitialize(JSONRPCMessage message) {
475475
private boolean hasContentType(HttpServletRequest request, String expected) {
476476
return Optional.ofNullable(request.getContentType())
477477
.map(String::toLowerCase)
478-
.filter(contentType -> contentType.contains(expected))
478+
.map(contentType -> contentType.split(";", 2)[0].trim())
479+
.filter(contentType -> contentType.equals(expected))
479480
.isPresent();
480481
}
481482

0 commit comments

Comments
 (0)