fix: avoid blocking streamable SSE writes

Author: kamikaz1k

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

@@ -5,7 +5,6 @@
 package com.agentclientprotocol.sdk.agent.transport;
 
 import java.io.IOException;
-import java.io.PrintWriter;
 import java.nio.channels.ClosedChannelException;
 import java.nio.charset.StandardCharsets;
 import java.time.Duration;
@@ -30,7 +29,11 @@
 import com.agentclientprotocol.sdk.spec.AcpSchema.JSONRPCMessage;
 import com.agentclientprotocol.sdk.util.Assert;
 import jakarta.servlet.AsyncContext;
+import jakarta.servlet.AsyncEvent;
+import jakarta.servlet.AsyncListener;
 import jakarta.servlet.ServletException;
+import jakarta.servlet.ServletOutputStream;
+import jakarta.servlet.WriteListener;
 import jakarta.servlet.http.HttpServlet;
 import jakarta.servlet.http.HttpServletRequest;
 import jakarta.servlet.http.HttpServletResponse;
@@ -92,6 +95,8 @@ public class StreamableHttpAcpAgentTransport {
 
 	private static final int MAX_REPLAY_EVENTS = 1024;
 
+	private static final int MAX_PENDING_SSE_EVENTS = 1024;
+
 	private static final Duration INITIALIZE_TIMEOUT = Duration.ofSeconds(30);
 
 	private enum ScopeKind {
@@ -827,68 +832,118 @@ synchronized void subscribe(AsyncContext asyncContext, HttpServletResponse respo
 			}
 			SseSubscriber subscriber = new SseSubscriber(this, asyncContext, response);
 			subscribers.add(subscriber);
+			subscriber.start();
 			if (replayOpen) {
 				for (String payload : new ArrayList<>(replay)) {
 					subscriber.send(payload);
 				}
 				replay.clear();
 				replayOpen = false;
 			}
-			subscriber.flush();
+			subscriber.drain();
 		}
 
 		void remove(SseSubscriber subscriber) {
 			subscribers.remove(subscriber);
 		}
 
-		void close() {
+		synchronized void close() {
 			if (closed.compareAndSet(false, true)) {
 				subscribers.forEach(SseSubscriber::close);
 				subscribers.clear();
-				synchronized (this) {
-					replay.clear();
-				}
+				replay.clear();
 			}
 		}
 
 	}
 
-	private final class SseSubscriber {
+	private final class SseSubscriber implements AsyncListener, WriteListener {
 
 		private final OutboundStream parent;
 
 		private final AsyncContext asyncContext;
 
-		private final PrintWriter writer;
+		private final ServletOutputStream output;
+
+		private final ArrayDeque<byte[]> pendingEvents = new ArrayDeque<>();
 
 		private final AtomicBoolean closed = new AtomicBoolean(false);
 
 		SseSubscriber(OutboundStream parent, AsyncContext asyncContext, HttpServletResponse response) throws IOException {
 			this.parent = parent;
 			this.asyncContext = asyncContext;
-			this.writer = response.getWriter();
+			this.output = response.getOutputStream();
+		}
+
+		void start() {
+			asyncContext.addListener(this);
+			output.setWriteListener(this);
 		}
 
 		synchronized void send(String payload) {
 			if (closed.get()) {
 				return;
 			}
-			writer.write("data: ");
-			writer.write(payload);
-			writer.write("\n\n");
-			writer.flush();
-			if (writer.checkError()) {
+			if (pendingEvents.size() == MAX_PENDING_SSE_EVENTS) {
+				logger.warn("Closing backpressured SSE subscriber after {} pending events", MAX_PENDING_SSE_EVENTS);
 				close();
+				return;
 			}
+			pendingEvents.addLast(("data: " + payload + "\n\n").getBytes(StandardCharsets.UTF_8));
+			drain();
 		}
 
-		synchronized void flush() {
-			writer.flush();
+		synchronized void drain() {
+			try {
+				while (!closed.get() && output.isReady()) {
+					byte[] event = pendingEvents.pollFirst();
+					if (event == null) {
+						return;
+					}
+					output.write(event);
+				}
+			}
+			catch (IOException e) {
+				close();
+			}
+		}
+
+		@Override
+		public void onWritePossible() {
+			drain();
+		}
+
+		@Override
+		public void onError(Throwable error) {
+			close();
+		}
+
+		@Override
+		public void onComplete(AsyncEvent event) {
+			close();
+		}
+
+		@Override
+		public void onTimeout(AsyncEvent event) {
+			close();
+		}
+
+		@Override
+		public void onError(AsyncEvent event) {
+			close();
+		}
+
+		@Override
+		public void onStartAsync(AsyncEvent event) {
+			event.getAsyncContext().addListener(this);
 		}
 
 		void close() {
 			if (closed.compareAndSet(false, true)) {
 				parent.remove(this);
+				synchronized (this) {
+					pendingEvents.clear();
+				}
 				try {
 					asyncContext.complete();
 				}

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

@@ -383,6 +383,48 @@ void deleteClosesSseAndRemovesConnection() throws Exception {
 		}
 	}
 
+	@Test
+	void concurrentSseOpenAndDeleteDoesNotLeaveStreamOpen() throws Exception {
+		try (FixtureServer server = FixtureServer.start()) {
+			HttpClient rawClient = HttpClient.newHttpClient();
+			String connectionId = initializeRaw(rawClient, server.endpoint());
+			ExecutorService executor = Executors.newFixedThreadPool(2);
+			CountDownLatch start = new CountDownLatch(1);
+			try {
+				Future<HttpResponse<InputStream>> openStream = executor.submit(() -> {
+					start.await();
+					return rawClient.send(HttpRequest.newBuilder(server.endpoint())
+						.header("Accept", "text/event-stream")
+						.header("Acp-Connection-Id", connectionId)
+						.GET()
+						.build(), HttpResponse.BodyHandlers.ofInputStream());
+				});
+				Future<HttpResponse<Void>> delete = executor.submit(() -> {
+					start.await();
+					return rawClient.send(HttpRequest.newBuilder(server.endpoint())
+						.header("Acp-Connection-Id", connectionId)
+						.DELETE()
+						.build(), HttpResponse.BodyHandlers.discarding());
+				});
+
+				start.countDown();
+				HttpResponse<InputStream> streamResponse = openStream.get(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS);
+				assertThat(delete.get(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS).statusCode()).isEqualTo(202);
+				assertThat(streamResponse.statusCode()).isIn(200, 404);
+				if (streamResponse.statusCode() == 200) {
+					try (InputStream body = streamResponse.body()) {
+						Future<Integer> endOfStream = executor.submit((java.util.concurrent.Callable<Integer>) body::read);
+						assertThat(endOfStream.get(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS)).isEqualTo(-1);
+					}
+				}
+				assertEventuallyPostStatus(rawClient, server.endpoint(), connectionId, 404);
+			}
+			finally {
+				executor.shutdownNow();
+			}
+		}
+	}
+
 	@Test
 	void replayOverflowClosesConnectionInsteadOfDroppingMessages() throws Exception {
 		try (FixtureServer server = FixtureServer.start()) {