fix: serialize streamable websocket writes

Author: kamikaz1k

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

@@ -935,6 +935,12 @@ private final class WebSocketConnectionState {
 
 		private final AtomicBoolean closed = new AtomicBoolean(false);
 
+		private final Object outboundLock = new Object();
+
+		private final ArrayDeque<String> outboundQueue = new ArrayDeque<>();
+
+		private boolean outboundSendInProgress = false;
+
 		private volatile Session session;
 
 		WebSocketConnectionState(String id) {
@@ -970,27 +976,76 @@ void acceptFromClient(JSONRPCMessage message) {
 
 		void sendToClient(JSONRPCMessage message) {
 			try {
-				Session currentSession = this.session;
-				if (closed.get() || currentSession == null || !currentSession.isOpen()) {
-					throw new AcpConnectionException("Streamable ACP WebSocket connection is closed");
-				}
 				String payload = jsonMapper.writeValueAsString(message);
 				logger.debug("Sending streamable ACP WebSocket message: {}", payload);
-				currentSession.sendText(payload, Callback.from(() -> {
-					// Jetty requires an explicit success callback; there is no
-					// follow-up work after the frame has been accepted for writing.
-				}, error -> {
-					if (!closed.get()) {
-						remoteConnection.signalException(error);
-					}
-				}));
+				enqueueOutbound(payload);
 			}
 			catch (Exception e) {
 				remoteConnection.signalException(e);
 				close(StatusCode.SERVER_ERROR, "failed to send ACP message");
 			}
 		}
 
+		private void enqueueOutbound(String payload) {
+			boolean shouldDrain;
+			synchronized (outboundLock) {
+				if (closed.get()) {
+					throw new AcpConnectionException("Streamable ACP WebSocket connection is closed");
+				}
+				outboundQueue.addLast(payload);
+				shouldDrain = !outboundSendInProgress;
+				if (shouldDrain) {
+					outboundSendInProgress = true;
+				}
+			}
+			if (shouldDrain) {
+				drainOutbound();
+			}
+		}
+
+		private void drainOutbound() {
+			String payload;
+			Session currentSession;
+			synchronized (outboundLock) {
+				if (closed.get()) {
+					outboundQueue.clear();
+					outboundSendInProgress = false;
+					return;
+				}
+				payload = outboundQueue.pollFirst();
+				if (payload == null) {
+					outboundSendInProgress = false;
+					return;
+				}
+				currentSession = this.session;
+			}
+
+			if (currentSession == null || !currentSession.isOpen()) {
+				failOutbound(new AcpConnectionException("Streamable ACP WebSocket connection is closed"));
+				return;
+			}
+
+			try {
+				/*
+				 * Jetty WebSocket sessions do not allow overlapping writes. Agent messages can
+				 * be produced by concurrent prompt handlers, so this per-connection queue sends
+				 * exactly one frame at a time and advances only after Jetty completes the
+				 * callback for the previous frame.
+				 */
+				currentSession.sendText(payload, Callback.from(this::drainOutbound, this::failOutbound));
+			}
+			catch (Exception e) {
+				failOutbound(e);
+			}
+		}
+
+		private void failOutbound(Throwable error) {
+			if (!closed.get()) {
+				remoteConnection.signalException(error);
+				close(StatusCode.SERVER_ERROR, "failed to send ACP message");
+			}
+		}
+
 		void close() {
 			close(StatusCode.NORMAL, "server closing");
 		}
@@ -999,6 +1054,10 @@ void close(int statusCode, String reason) {
 			if (!closed.compareAndSet(false, true)) {
 				return;
 			}
+			synchronized (outboundLock) {
+				outboundQueue.clear();
+				outboundSendInProgress = false;
+			}
 			webSocketConnections.remove(id, this);
 			Session currentSession = this.session;
 			if (currentSession != null && currentSession.isOpen()) {

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

@@ -37,7 +37,9 @@
 import com.agentclientprotocol.sdk.spec.AcpSchema;
 import org.eclipse.jetty.websocket.api.StatusCode;
 import org.junit.jupiter.api.Test;
+import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
+import reactor.core.scheduler.Schedulers;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
@@ -209,6 +211,59 @@ void supportsMultipleConcurrentWebSocketClients() throws Exception {
 		}
 	}
 
+	@Test
+	void serializesConcurrentAgentMessagesOnOneWebSocketConnection() throws Exception {
+		AtomicInteger sessionCounter = new AtomicInteger();
+		AtomicInteger receivedUpdates = new AtomicInteger();
+		AcpAgentFactory agentFactory = AcpAgentFactory.async(transport -> AcpAgent.async(transport)
+			.initializeHandler(request -> Mono.just(new AcpSchema.InitializeResponse(
+					AcpSchema.LATEST_PROTOCOL_VERSION, new AcpSchema.AgentCapabilities(true, null, null), List.of())))
+			.newSessionHandler(request -> Mono.just(new AcpSchema.NewSessionResponse(
+					"sess-" + sessionCounter.incrementAndGet(), null, null)))
+			.promptHandler((request, context) -> Mono.delay(Duration.ofMillis(25))
+				.thenMany(Flux.range(0, 20)
+					.flatMap(i -> context.sendMessage(request.sessionId() + ": update-" + i)
+						.subscribeOn(Schedulers.parallel()), 8))
+				.then(Mono.just(AcpSchema.PromptResponse.endTurn())))
+			.build());
+
+		try (FixtureServer server = FixtureServer.start(agentFactory)) {
+			AcpAsyncClient client = AcpClient
+				.async(new WebSocketAcpClientTransport(server.endpoint(), AcpJsonMapper.createDefault()))
+				.sessionUpdateConsumer(update -> {
+					receivedUpdates.incrementAndGet();
+					return Mono.empty();
+				})
+				.requestTimeout(TIMEOUT)
+				.build();
+			try {
+				client.initialize(new AcpSchema.InitializeRequest(
+						AcpSchema.LATEST_PROTOCOL_VERSION, new AcpSchema.ClientCapabilities()))
+					.block(TIMEOUT);
+				AcpSchema.NewSessionResponse firstSession = client
+					.newSession(new AcpSchema.NewSessionRequest("/workspace/one", List.of()))
+					.block(TIMEOUT);
+				AcpSchema.NewSessionResponse secondSession = client
+					.newSession(new AcpSchema.NewSessionRequest("/workspace/two", List.of()))
+					.block(TIMEOUT);
+
+				var prompts = Mono.zip(
+						client.prompt(new AcpSchema.PromptRequest(firstSession.sessionId(),
+								List.of(new AcpSchema.TextContent("first")))),
+						client.prompt(new AcpSchema.PromptRequest(secondSession.sessionId(),
+								List.of(new AcpSchema.TextContent("second")))))
+					.block(TIMEOUT);
+
+				assertThat(prompts.getT1().stopReason()).isEqualTo(AcpSchema.StopReason.END_TURN);
+				assertThat(prompts.getT2().stopReason()).isEqualTo(AcpSchema.StopReason.END_TURN);
+				assertThat(receivedUpdates).hasValue(40);
+			}
+			finally {
+				client.closeGracefully().block(TIMEOUT);
+			}
+		}
+	}
+
 	@Test
 	void rejectsNonInitializeFirstMessage() throws Exception {
 		try (FixtureServer server = FixtureServer.start(simpleAgentFactory())) {