fix: initialize streamable connections asynchronously

Author: kamikaz1k

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

@@ -57,6 +57,7 @@
 import org.slf4j.LoggerFactory;
 import reactor.core.publisher.Mono;
 import reactor.core.publisher.Sinks;
+import reactor.core.scheduler.Schedulers;
 
 /**
  * Listener-backed ACP Streamable HTTP transport for agents.
@@ -97,6 +98,9 @@ public class StreamableHttpAcpAgentTransport {
 
 	private static final int MAX_PENDING_SSE_EVENTS = 1024;
 
+	// Commits the SSE response when there are no replay events, without emitting an ACP message.
+	private static final byte[] SSE_OPEN_COMMENT = ": connected\n\n".getBytes(StandardCharsets.UTF_8);
+
 	private static final Duration INITIALIZE_TIMEOUT = Duration.ofSeconds(30);
 
 	private enum ScopeKind {
@@ -308,9 +312,7 @@ int activeConnectionCount() {
 
 	private ConnectionState createConnection() {
 		String connectionId = UUID.randomUUID().toString();
-		ConnectionState connection = new ConnectionState(connectionId);
-		connection.start();
-		return connection;
+		return new ConnectionState(connectionId);
 	}
 
 	private WebSocketConnectionState createWebSocketConnection() {
@@ -430,9 +432,47 @@ private void handleInitialize(HttpServletRequest request, HttpServletResponse re
 			}
 
 			ConnectionState connection = createConnection();
+			AsyncContext asyncContext = request.startAsync();
+			asyncContext.setTimeout(INITIALIZE_TIMEOUT.toMillis());
+			AtomicBoolean completed = new AtomicBoolean(false);
+			asyncContext.addListener(new AsyncListener() {
+
+				@Override
+				public void onComplete(AsyncEvent event) {
+				}
+
+				@Override
+				public void onTimeout(AsyncEvent event) {
+					completeInitializeFailure(asyncContext, response, connection, completed);
+				}
+
+				@Override
+				public void onError(AsyncEvent event) {
+					completeInitializeFailure(asyncContext, response, connection, completed);
+				}
+
+				@Override
+				public void onStartAsync(AsyncEvent event) {
+					event.getAsyncContext().addListener(this);
+				}
+
+			});
+
+			connection.start()
+				.then(Mono.defer(() -> connection.initialize(initializeRequest)))
+				.timeout(INITIALIZE_TIMEOUT)
+				.subscribeOn(Schedulers.boundedElastic())
+				.subscribe(initializeResponse -> completeInitializeSuccess(asyncContext, response, connection,
+						completed, initializeResponse),
+					error -> completeInitializeFailure(asyncContext, response, connection, completed));
+		}
+
+		private void completeInitializeSuccess(AsyncContext asyncContext, HttpServletResponse response,
+				ConnectionState connection, AtomicBoolean completed, JSONRPCMessage initializeResponse) {
+			if (!completed.compareAndSet(false, true)) {
+				return;
+			}
 			try {
-				JSONRPCMessage initializeResponse = connection.initialize(initializeRequest)
-					.block(INITIALIZE_TIMEOUT);
 				if (!(initializeResponse instanceof AcpSchema.JSONRPCResponse)) {
 					throw new AcpConnectionException("initialize did not produce a JSON-RPC response");
 				}
@@ -443,9 +483,44 @@ private void handleInitialize(HttpServletRequest request, HttpServletResponse re
 				response.getWriter().write(jsonMapper.writeValueAsString(initializeResponse));
 			}
 			catch (Exception e) {
+				connections.remove(connection.id(), connection);
 				connection.close();
+				try {
+					writeText(response, HttpServletResponse.SC_INTERNAL_SERVER_ERROR, "initialize failed");
+				}
+				catch (IOException writeError) {
+					logger.warn("Failed to write Streamable HTTP initialize failure", writeError);
+				}
+			}
+			finally {
+				completeAsyncContext(asyncContext);
+			}
+		}
+
+		private void completeInitializeFailure(AsyncContext asyncContext, HttpServletResponse response,
+				ConnectionState connection, AtomicBoolean completed) {
+			if (!completed.compareAndSet(false, true)) {
+				return;
+			}
+			connections.remove(connection.id(), connection);
+			connection.close();
+			try {
 				writeText(response, HttpServletResponse.SC_INTERNAL_SERVER_ERROR, "initialize failed");
 			}
+			catch (IOException writeError) {
+				logger.warn("Failed to write Streamable HTTP initialize failure", writeError);
+			}
+			finally {
+				completeAsyncContext(asyncContext);
+			}
+		}
+
+		private void completeAsyncContext(AsyncContext asyncContext) {
+			try {
+				asyncContext.complete();
+			}
+			catch (IllegalStateException ignored) {
+			}
 		}
 
 	}
@@ -513,8 +588,8 @@ String id() {
 			return id;
 		}
 
-		void start() {
-			this.connection.start(agentFactory).block(INITIALIZE_TIMEOUT);
+		Mono<Void> start() {
+			return this.connection.start(agentFactory);
 		}
 
 		Mono<JSONRPCMessage> initialize(AcpSchema.JSONRPCRequest request) {
@@ -869,14 +944,17 @@ private final class SseSubscriber implements AsyncListener, WriteListener {
 
 		private final AtomicBoolean closed = new AtomicBoolean(false);
 
+		private boolean flushPending;
+
 		SseSubscriber(OutboundStream parent, AsyncContext asyncContext, HttpServletResponse response) throws IOException {
 			this.parent = parent;
 			this.asyncContext = asyncContext;
 			this.output = response.getOutputStream();
 		}
 
-		void start() {
+		synchronized void start() {
 			asyncContext.addListener(this);
+			pendingEvents.addLast(SSE_OPEN_COMMENT);
 			output.setWriteListener(this);
 		}
 
@@ -895,19 +973,29 @@ synchronized void send(String payload) {
 
 		synchronized void drain() {
 			try {
+				flushIfReady();
 				while (!closed.get() && output.isReady()) {
 					byte[] event = pendingEvents.pollFirst();
 					if (event == null) {
-						return;
+						break;
 					}
 					output.write(event);
+					flushPending = true;
 				}
+				flushIfReady();
 			}
-			catch (IOException e) {
+			catch (IOException | IllegalStateException e) {
 				close();
 			}
 		}
 
+		private void flushIfReady() throws IOException {
+			if (flushPending && output.isReady()) {
+				output.flush();
+				flushPending = false;
+			}
+		}
+
 		@Override
 		public void onWritePossible() {
 			drain();

acp-streamable-http-jetty/src/test/java/com/agentclientprotocol/sdk/agent/transport/StreamableHttpAcpAgentTransportIntegrationTest.java

@@ -47,6 +47,56 @@ class StreamableHttpAcpAgentTransportIntegrationTest {
 
 	private static final AcpJsonMapper JSON_MAPPER = AcpJsonMapper.createDefault();
 
+	@Test
+	void delayedInitializePublishesConnectionOnlyAfterSuccess() throws Exception {
+		CountDownLatch initializeEntered = new CountDownLatch(1);
+		AcpAgentFactory agentFactory = AcpAgentFactory.async(transport -> AcpAgent.async(transport)
+			.initializeHandler(request -> Mono.defer(() -> {
+				initializeEntered.countDown();
+				return Mono.delay(Duration.ofMillis(250)).thenReturn(new AcpSchema.InitializeResponse(
+						AcpSchema.LATEST_PROTOCOL_VERSION, new AcpSchema.AgentCapabilities(true, null, null), null));
+			}))
+			.build());
+		try (FixtureServer server = FixtureServer.start(agentFactory)) {
+			HttpClient rawClient = HttpClient.newHttpClient();
+			var initialize = rawClient.sendAsync(HttpRequest.newBuilder(server.endpoint())
+				.header("Content-Type", "application/json")
+				.header("Accept", "application/json")
+				.POST(HttpRequest.BodyPublishers.ofString("""
+						{"jsonrpc":"2.0","id":"init-delayed","method":"initialize","params":{"protocolVersion":1,"clientCapabilities":{}}}
+						"""))
+				.build(), HttpResponse.BodyHandlers.ofString(StandardCharsets.UTF_8));
+
+			assertThat(initializeEntered.await(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS)).isTrue();
+			assertThat(server.transport.activeConnectionCount()).isZero();
+
+			HttpResponse<String> response = initialize.get(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS);
+			assertThat(response.statusCode()).isEqualTo(200);
+			assertThat(response.headers().firstValue("Acp-Connection-Id")).isPresent();
+			assertThat(server.transport.activeConnectionCount()).isEqualTo(1);
+		}
+	}
+
+	@Test
+	void failedInitializeDoesNotPublishConnection() throws Exception {
+		AcpAgentFactory agentFactory = transport -> {
+			throw new IllegalStateException("agent creation failed");
+		};
+		try (FixtureServer server = FixtureServer.start(agentFactory)) {
+			HttpClient rawClient = HttpClient.newHttpClient();
+			HttpResponse<String> response = rawClient.send(HttpRequest.newBuilder(server.endpoint())
+				.header("Content-Type", "application/json")
+				.header("Accept", "application/json")
+				.POST(HttpRequest.BodyPublishers.ofString("""
+						{"jsonrpc":"2.0","id":"init-failed","method":"initialize","params":{"protocolVersion":1,"clientCapabilities":{}}}
+						"""))
+				.build(), HttpResponse.BodyHandlers.ofString(StandardCharsets.UTF_8));
+
+			assertThat(response.statusCode()).isEqualTo(500);
+			assertThat(server.transport.activeConnectionCount()).isZero();
+		}
+	}
+
 	@Test
 	void javaClientCanTalkToRunningJavaServer() throws Exception {
 		try (FixtureServer server = FixtureServer.start()) {
@@ -527,6 +577,10 @@ static FixtureServer start() throws Exception {
 						.thenReturn(AcpSchema.PromptResponse.endTurn());
 				})
 				.build());
+			return start(agentFactory);
+		}
+
+		static FixtureServer start(AcpAgentFactory agentFactory) throws Exception {
 			StreamableHttpAcpAgentTransport transport = new StreamableHttpAcpAgentTransport(
 					freePort(), AcpJsonMapper.createDefault(), agentFactory);
 			transport.start().block(TIMEOUT);