fix: await streamable listener shutdown

Author: kamikaz1k

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

@@ -259,27 +259,36 @@ public int getPort() {
 	 * @return a mono that completes when shutdown finishes
 	 */
 	public Mono<Void> closeGracefully() {
-		return Mono.fromRunnable(() -> {
+		return Mono.defer(() -> {
 			if (!closing.compareAndSet(false, true)) {
-				return;
+				return Mono.<Void>empty();
 			}
-			connections.values().forEach(ConnectionState::close);
+			List<Mono<Void>> connectionClosures = new ArrayList<>();
+			connections.values().forEach(connection -> connectionClosures.add(connection.closeGracefully()));
 			connections.clear();
-			webSocketConnections.values().forEach(WebSocketConnectionState::close);
+			webSocketConnections.values().forEach(connection -> connectionClosures.add(connection.closeGracefully()));
 			webSocketConnections.clear();
-			Server currentServer = this.server;
-			if (currentServer != null) {
-				try {
-					currentServer.stop();
-				}
-				catch (Exception e) {
-					throw new AcpConnectionException("Failed to stop Streamable HTTP listener", e);
-				}
-			}
-			terminationSink.tryEmitValue(null);
+
+			return Mono.whenDelayError(connectionClosures)
+				.then(Mono.<Void>fromRunnable(this::stopServer))
+				.doOnSuccess(ignored -> {
+					terminationSink.tryEmitValue(null);
+				});
 		});
 	}
 
+	private void stopServer() {
+		Server currentServer = this.server;
+		if (currentServer != null) {
+			try {
+				currentServer.stop();
+			}
+			catch (Exception e) {
+				throw new AcpConnectionException("Failed to stop Streamable HTTP listener", e);
+			}
+		}
+	}
+
 	/**
 	 * Returns a mono that completes once the listener terminates.
 	 * @return termination mono
@@ -550,11 +559,15 @@ void openStream(HttpServletRequest request, HttpServletResponse response, String
 			stream.subscribe(asyncContext, response);
 		}
 
-		void close() {
+		Mono<Void> closeGracefully() {
 			connections.remove(id, this);
 			connectionStream.close();
 			sessionStreams.values().forEach(OutboundStream::close);
-			connection.closeGracefully().subscribe(v -> {
+			return connection.closeGracefully();
+		}
+
+		void close() {
+			closeGracefully().subscribe(v -> {
 			}, error -> logger.warn("Error closing Streamable HTTP ACP connection {}", id, error));
 		}
 
@@ -950,21 +963,30 @@ void sendToClient(JSONRPCMessage message) {
 			}
 		}
 
-		void close() {
-			close(StatusCode.NORMAL, "server closing");
+		Mono<Void> closeGracefully() {
+			return closeGracefully(StatusCode.NORMAL, "server closing");
 		}
 
-		void close(int statusCode, String reason) {
+		private Mono<Void> closeGracefully(int statusCode, String reason) {
 			if (!closed.compareAndSet(false, true)) {
-				return;
+				return Mono.empty();
 			}
 			outboundSender.close();
 			webSocketConnections.remove(id, this);
 			Session currentSession = this.session;
 			if (currentSession != null && currentSession.isOpen()) {
 				currentSession.close(statusCode, reason, Callback.NOOP);
 			}
-			remoteConnection.closeGracefully().subscribe(v -> {
+			return remoteConnection.closeGracefully();
+		}
+
+		void close() {
+			closeGracefully().subscribe(v -> {
+			}, error -> logger.warn("Error closing Streamable ACP WebSocket connection {}", id, error));
+		}
+
+		void close(int statusCode, String reason) {
+			closeGracefully(statusCode, reason).subscribe(v -> {
 			}, error -> logger.warn("Error closing Streamable ACP WebSocket connection {}", id, error));
 		}
 

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

@@ -32,6 +32,8 @@
 
 import com.agentclientprotocol.sdk.agent.AcpAgent;
 import com.agentclientprotocol.sdk.agent.AcpAgentFactory;
+import com.agentclientprotocol.sdk.agent.AcpAsyncAgent;
+import com.agentclientprotocol.sdk.capabilities.NegotiatedCapabilities;
 import com.agentclientprotocol.sdk.client.AcpAsyncClient;
 import com.agentclientprotocol.sdk.client.AcpClient;
 import com.agentclientprotocol.sdk.client.transport.WebSocketAcpClientTransport;
@@ -41,6 +43,7 @@
 import org.junit.jupiter.api.Test;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
+import reactor.core.publisher.Sinks;
 import reactor.core.scheduler.Schedulers;
 
 import static org.assertj.core.api.Assertions.assertThat;
@@ -325,6 +328,30 @@ void rejectsDuplicateInitializeWithoutForwardingItToTheAgent() throws Exception
 		}
 	}
 
+	@Test
+	void listenerShutdownWaitsForWebSocketAgentShutdown() throws Exception {
+		Sinks.One<Void> allowAgentShutdown = Sinks.one();
+		CountDownLatch agentShutdownStarted = new CountDownLatch(1);
+		AcpAgentFactory agentFactory = transport -> new BlockingCloseAgent(allowAgentShutdown, agentShutdownStarted);
+
+		try (FixtureServer server = FixtureServer.start(agentFactory)) {
+			MessageRecordingListener listener = new MessageRecordingListener();
+			HttpClient.newHttpClient()
+				.newWebSocketBuilder()
+				.connectTimeout(TIMEOUT)
+				.buildAsync(server.endpoint(), listener)
+				.get(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS);
+			assertThat(listener.openLatch.await(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS)).isTrue();
+
+			CompletableFuture<Void> shutdown = server.transport().closeGracefully().toFuture();
+			assertThat(agentShutdownStarted.await(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS)).isTrue();
+			assertThat(shutdown).isNotDone();
+
+			allowAgentShutdown.tryEmitEmpty();
+			shutdown.get(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS);
+		}
+	}
+
 	private static AcpAgentFactory simpleAgentFactory() {
 		AtomicInteger sessionCounter = new AtomicInteger();
 		return AcpAgentFactory.async(transport -> AcpAgent.async(transport)
@@ -462,4 +489,103 @@ public CompletionStage<?> onText(WebSocket webSocket, CharSequence data, boolean
 
 	}
 
+	private static final class BlockingCloseAgent implements AcpAsyncAgent {
+
+		private final Sinks.One<Void> allowShutdown;
+
+		private final CountDownLatch shutdownStarted;
+
+		BlockingCloseAgent(Sinks.One<Void> allowShutdown, CountDownLatch shutdownStarted) {
+			this.allowShutdown = allowShutdown;
+			this.shutdownStarted = shutdownStarted;
+		}
+
+		@Override
+		public Mono<Void> start() {
+			return Mono.empty();
+		}
+
+		@Override
+		public Mono<Void> awaitTermination() {
+			return Mono.never();
+		}
+
+		@Override
+		public NegotiatedCapabilities getClientCapabilities() {
+			return null;
+		}
+
+		@Override
+		public Mono<Void> sendSessionUpdate(String sessionId, AcpSchema.SessionUpdate update) {
+			return unsupported();
+		}
+
+		@Override
+		public Mono<AcpSchema.RequestPermissionResponse> requestPermission(AcpSchema.RequestPermissionRequest request) {
+			return unsupported();
+		}
+
+		@Override
+		public Mono<AcpSchema.ReadTextFileResponse> readTextFile(AcpSchema.ReadTextFileRequest request) {
+			return unsupported();
+		}
+
+		@Override
+		public Mono<AcpSchema.WriteTextFileResponse> writeTextFile(AcpSchema.WriteTextFileRequest request) {
+			return unsupported();
+		}
+
+		@Override
+		public Mono<AcpSchema.CreateTerminalResponse> createTerminal(AcpSchema.CreateTerminalRequest request) {
+			return unsupported();
+		}
+
+		@Override
+		public Mono<AcpSchema.TerminalOutputResponse> getTerminalOutput(AcpSchema.TerminalOutputRequest request) {
+			return unsupported();
+		}
+
+		@Override
+		public Mono<AcpSchema.ReleaseTerminalResponse> releaseTerminal(AcpSchema.ReleaseTerminalRequest request) {
+			return unsupported();
+		}
+
+		@Override
+		public Mono<AcpSchema.WaitForTerminalExitResponse> waitForTerminalExit(AcpSchema.WaitForTerminalExitRequest request) {
+			return unsupported();
+		}
+
+		@Override
+		public Mono<AcpSchema.KillTerminalCommandResponse> killTerminal(AcpSchema.KillTerminalCommandRequest request) {
+			return unsupported();
+		}
+
+		@Override
+		public Mono<AcpSchema.CreateElicitationResponse> createElicitation(AcpSchema.CreateElicitationRequest request) {
+			return unsupported();
+		}
+
+		@Override
+		public Mono<Void> completeElicitation(AcpSchema.CompleteElicitationNotification notification) {
+			return unsupported();
+		}
+
+		@Override
+		public Mono<Void> closeGracefully() {
+			return Mono.defer(() -> {
+				shutdownStarted.countDown();
+				return allowShutdown.asMono();
+			});
+		}
+
+		@Override
+		public void close() {
+		}
+
+		private static <T> Mono<T> unsupported() {
+			return Mono.error(new UnsupportedOperationException());
+		}
+
+	}
+
 }