Add agent await/awaitTermination for daemon thread early-exit fix (Step 3.7.2)

- Add awaitTermination() to AcpAgentTransport interface
- Implement termination signaling in StdioAcpAgentTransport and WebSocketAcpAgentTransport
- Add awaitTermination() to AcpAsyncAgent interface and DefaultAcpAsyncAgent
- Add await() and run() methods to AcpSyncAgent
- Rename CleanShutdownTest to CleanShutdownIT and add agent await tests
- Fix StdioAcpAgentTransportTest echo handler bug (use non-echoing handler)
- Add awaitTermination() to InMemoryTransportPair for test support

Author: Mark Pollack

acp-core/src/main/java/com/agentclientprotocol/sdk/agent/AcpAsyncAgent.java

@@ -34,6 +34,21 @@ public interface AcpAsyncAgent {
 	 */
 	Mono<Void> start();
 
+	/**
+	 * Returns a Mono that completes when the agent terminates.
+	 * This is useful for blocking until the transport closes, particularly
+	 * when using daemon threads.
+	 *
+	 * <p>Example usage:
+	 * <pre>{@code
+	 * agent.start().block();
+	 * agent.awaitTermination().block(); // Block until transport closes
+	 * }</pre>
+	 *
+	 * @return A Mono that completes when the agent terminates
+	 */
+	Mono<Void> awaitTermination();
+
 	/**
 	 * Returns the capabilities negotiated with the client during initialization.
 	 *

acp-core/src/main/java/com/agentclientprotocol/sdk/agent/AcpSyncAgent.java

@@ -49,11 +49,46 @@ public AcpSyncAgent(AcpAsyncAgent asyncAgent, Duration blockTimeout) {
 
 	/**
 	 * Starts the agent, beginning to accept client connections.
+	 * This method returns immediately after setup is complete.
+	 * Use {@link #await()} or {@link #run()} to block until the transport closes.
 	 */
 	public void start() {
 		asyncAgent.start().block(blockTimeout);
 	}
 
+	/**
+	 * Blocks until the agent terminates (transport closes).
+	 * This is useful for keeping the main thread alive when using daemon threads.
+	 *
+	 * <p>Example usage:
+	 * <pre>{@code
+	 * agent.start();
+	 * agent.await(); // Blocks until stdin closes
+	 * }</pre>
+	 */
+	public void await() {
+		asyncAgent.awaitTermination().block();
+	}
+
+	/**
+	 * Starts the agent and blocks until it terminates.
+	 * This is a convenience method combining {@link #start()} and {@link #await()}.
+	 *
+	 * <p>Example usage for a standalone agent:
+	 * <pre>{@code
+	 * public static void main(String[] args) {
+	 *     AcpSyncAgent agent = AcpAgent.sync(transport)
+	 *         .promptHandler(...)
+	 *         .build();
+	 *     agent.run(); // Start and block until done
+	 * }
+	 * }</pre>
+	 */
+	public void run() {
+		start();
+		await();
+	}
+
 	/**
 	 * Sends a session update notification to the client.
 	 * @param sessionId The session ID

acp-core/src/main/java/com/agentclientprotocol/sdk/agent/DefaultAcpAsyncAgent.java

@@ -181,6 +181,11 @@ public Mono<Void> start() {
 		});
 	}
 
+	@Override
+	public Mono<Void> awaitTermination() {
+		return transport.awaitTermination();
+	}
+
 	@Override
 	public NegotiatedCapabilities getClientCapabilities() {
 		return clientCapabilities.get();

acp-core/src/main/java/com/agentclientprotocol/sdk/agent/transport/StdioAcpAgentTransport.java

@@ -69,6 +69,8 @@ public class StdioAcpAgentTransport implements AcpAgentTransport {
 
 	private final Sinks.One<Void> outboundReady = Sinks.one();
 
+	private final Sinks.One<Void> terminationSink = Sinks.one();
+
 	private Scheduler inboundScheduler;
 
 	private Scheduler outboundScheduler;
@@ -153,6 +155,11 @@ public Mono<Void> start(Function<Mono<JSONRPCMessage>, Mono<JSONRPCMessage>> han
 	private void handleIncomingMessages(Function<Mono<JSONRPCMessage>, Mono<JSONRPCMessage>> handler) {
 		this.inboundSink.asFlux()
 			.flatMap(message -> Mono.just(message).transform(handler))
+			.doOnNext(response -> {
+				if (response != null) {
+					this.outboundSink.tryEmitNext(response);
+				}
+			})
 			.doOnTerminate(() -> {
 				this.outboundSink.tryEmitComplete();
 				this.inboundScheduler.dispose();
@@ -206,6 +213,8 @@ private void startInboundProcessing() {
 			finally {
 				isClosing.set(true);
 				inboundSink.tryEmitComplete();
+				terminationSink.tryEmitValue(null);  // Signal termination for awaitTermination()
+				logger.debug("Agent transport terminated");
 			}
 		});
 	}
@@ -301,6 +310,11 @@ public void setExceptionHandler(Consumer<Throwable> handler) {
 		this.exceptionHandler = handler;
 	}
 
+	@Override
+	public Mono<Void> awaitTermination() {
+		return terminationSink.asMono();
+	}
+
 	@Override
 	public <T> T unmarshalFrom(Object data, TypeRef<T> typeRef) {
 		return jsonMapper.convertValue(data, typeRef);

acp-core/src/main/java/com/agentclientprotocol/sdk/spec/AcpAgentTransport.java

@@ -41,4 +41,19 @@ public interface AcpAgentTransport extends AcpTransport {
 	default void setExceptionHandler(Consumer<Throwable> handler) {
 	}
 
+	/**
+	 * Returns a Mono that completes when the transport terminates.
+	 * This is useful for agents that need to block until the transport is done,
+	 * particularly when using daemon threads.
+	 *
+	 * <p>Example usage:
+	 * <pre>{@code
+	 * transport.start(handler).block();
+	 * transport.awaitTermination().block(); // Block until stdin closes
+	 * }</pre>
+	 *
+	 * @return a {@link Mono} that completes when the transport terminates
+	 */
+	Mono<Void> awaitTermination();
+
 }

acp-core/src/test/java/com/agentclientprotocol/sdk/agent/transport/StdioAcpAgentTransportTest.java

@@ -105,7 +105,10 @@ void sendsMessageToOutputStream() throws Exception {
 
 		StdioAcpAgentTransport transport = new StdioAcpAgentTransport(jsonMapper, agentIn, agentOut);
 
-		transport.start(msg -> msg).subscribe();
+		// Use a non-echoing handler - the default echo handler (msg -> msg) would
+		// echo the warmup message back to output, causing the test to read the
+		// wrong message. We only want to test explicit sendMessage() output.
+		transport.start(msg -> Mono.empty()).subscribe();
 
 		// Wait for transport to be ready by sending a dummy message from "client" first
 		AcpSchema.JSONRPCRequest dummyRequest = AcpTestFixtures.createJsonRpcRequest(AcpSchema.METHOD_INITIALIZE, "0",

acp-core/src/test/java/com/agentclientprotocol/sdk/integration/CleanShutdownIT.java

@@ -0,0 +1,215 @@
+/*
+ * Copyright 2025-2025 the original author or authors.
+ */
+
+package com.agentclientprotocol.sdk.integration;
+
+import java.time.Duration;
+import java.util.List;
+import java.util.Set;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.stream.Collectors;
+
+import com.agentclientprotocol.sdk.agent.AcpAgent;
+import com.agentclientprotocol.sdk.agent.AcpSyncAgent;
+import com.agentclientprotocol.sdk.client.AcpClient;
+import com.agentclientprotocol.sdk.client.AcpSyncClient;
+import com.agentclientprotocol.sdk.spec.AcpSchema;
+import com.agentclientprotocol.sdk.test.InMemoryTransportPair;
+import org.junit.jupiter.api.Test;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * Integration tests to verify that ACP transports clean up properly on shutdown
+ * and don't leave lingering threads that prevent JVM exit.
+ *
+ * <p>
+ * These tests address GitHub issue: Transport threads prevent JVM exit when
+ * closeGracefully() isn't called properly. The fix uses daemon threads for
+ * all transport schedulers.
+ * </p>
+ *
+ * @author Mark Pollack
+ */
+class CleanShutdownIT {
+
+	/**
+	 * Verifies that no ACP-prefixed threads remain after client close.
+	 *
+	 * <p>
+	 * This test ensures that the daemon thread fix is working correctly.
+	 * Before the fix, non-daemon threads would prevent the JVM from exiting.
+	 * </p>
+	 */
+	@Test
+	void noLingeringAcpThreadsAfterClientClose() throws InterruptedException {
+		// Record threads before test
+		Set<String> acpThreadsBefore = getAcpThreadNames();
+
+		// Create and use client with in-memory transport
+		InMemoryTransportPair transportPair = InMemoryTransportPair.create();
+
+		try (AcpSyncClient client = AcpClient.sync(transportPair.clientTransport()).build()) {
+			// The transport threads are created when build() connects
+			// Give threads a moment to start
+			Thread.sleep(100);
+
+			// Verify threads were created (sanity check)
+			Set<String> threadsWhileRunning = getAcpThreadNames();
+			// Note: InMemory transport may not create threads in the same way as stdio,
+			// but we should still verify no leak occurs
+		}
+
+		// Allow time for thread cleanup
+		Thread.sleep(500);
+
+		// Verify no ACP threads remain
+		Set<String> acpThreadsAfter = getAcpThreadNames();
+		assertThat(acpThreadsAfter)
+			.describedAs("ACP threads should be cleaned up after client close")
+			.isEqualTo(acpThreadsBefore);
+	}
+
+	/**
+	 * Verifies that daemon threads are used for transport schedulers.
+	 *
+	 * <p>
+	 * This test inspects active threads to ensure any ACP-prefixed threads
+	 * are daemon threads, which allows the JVM to exit gracefully.
+	 * </p>
+	 */
+	@Test
+	void acpThreadsAreDaemonThreads() throws InterruptedException {
+		InMemoryTransportPair transportPair = InMemoryTransportPair.create();
+
+		try (AcpSyncClient client = AcpClient.sync(transportPair.clientTransport()).build()) {
+			// Give threads a moment to start
+			Thread.sleep(100);
+
+			// Check that any ACP threads are daemon threads
+			Set<Thread> acpThreads = Thread.getAllStackTraces().keySet().stream()
+				.filter(t -> t.getName().startsWith("acp-"))
+				.collect(Collectors.toSet());
+
+			for (Thread thread : acpThreads) {
+				assertThat(thread.isDaemon())
+					.describedAs("Thread '%s' should be a daemon thread", thread.getName())
+					.isTrue();
+			}
+		}
+	}
+
+	/**
+	 * Verifies that agent awaitTermination() blocks until transport closes.
+	 *
+	 * <p>
+	 * This test ensures the await() fix for daemon thread early exit is working.
+	 * Before the fix, agents using daemon threads would exit immediately after start().
+	 * With await(), agents properly block until the transport terminates.
+	 * </p>
+	 */
+	@Test
+	void agentAwaitBlocksUntilTransportCloses() throws InterruptedException {
+		InMemoryTransportPair transportPair = InMemoryTransportPair.create();
+
+		// Build agent - SyncAgentBuilder expects SyncInitializeHandler (returns response directly)
+		AcpSyncAgent agent = AcpAgent.sync(transportPair.agentTransport())
+			.requestTimeout(Duration.ofSeconds(5))
+			.initializeHandler(request ->
+				new AcpSchema.InitializeResponse(1, new AcpSchema.AgentCapabilities(), List.of()))
+			.build();
+
+		// Track whether await() returned
+		AtomicBoolean awaitReturned = new AtomicBoolean(false);
+		CountDownLatch agentStarted = new CountDownLatch(1);
+
+		// Start agent in background thread that calls await()
+		Thread agentThread = new Thread(() -> {
+			agent.start();
+			agentStarted.countDown();
+			agent.await(); // Should block until transport closes
+			awaitReturned.set(true);
+		});
+		agentThread.setDaemon(true);
+		agentThread.start();
+
+		// Wait for agent to start
+		assertThat(agentStarted.await(5, TimeUnit.SECONDS)).isTrue();
+
+		// Give a moment and verify await() has NOT returned yet
+		Thread.sleep(200);
+		assertThat(awaitReturned.get())
+			.describedAs("await() should NOT return while transport is still open")
+			.isFalse();
+
+		// Close transport - this should unblock await()
+		transportPair.closeGracefully().block(Duration.ofSeconds(5));
+
+		// Wait for agent thread to complete
+		agentThread.join(5000);
+		assertThat(awaitReturned.get())
+			.describedAs("await() should return after transport closes")
+			.isTrue();
+	}
+
+	/**
+	 * Verifies that agent run() combines start() and await().
+	 *
+	 * <p>
+	 * The run() method is a convenience for standalone agents that want to
+	 * start and block in one call, similar to HTTP server patterns.
+	 * </p>
+	 */
+	@Test
+	void agentRunBlocksUntilTransportCloses() throws InterruptedException {
+		InMemoryTransportPair transportPair = InMemoryTransportPair.create();
+
+		// Build agent - SyncAgentBuilder expects SyncInitializeHandler (returns response directly)
+		AcpSyncAgent agent = AcpAgent.sync(transportPair.agentTransport())
+			.requestTimeout(Duration.ofSeconds(5))
+			.initializeHandler(request ->
+				new AcpSchema.InitializeResponse(1, new AcpSchema.AgentCapabilities(), List.of()))
+			.build();
+
+		// Track whether run() returned
+		AtomicBoolean runReturned = new AtomicBoolean(false);
+
+		// Start agent using run() in background thread
+		Thread agentThread = new Thread(() -> {
+			agent.run(); // Should block until transport closes
+			runReturned.set(true);
+		});
+		agentThread.setDaemon(true);
+		agentThread.start();
+
+		// Give a moment for agent to start
+		Thread.sleep(200);
+		assertThat(runReturned.get())
+			.describedAs("run() should NOT return while transport is still open")
+			.isFalse();
+
+		// Close transport - this should unblock run()
+		transportPair.closeGracefully().block(Duration.ofSeconds(5));
+
+		// Wait for agent thread to complete
+		agentThread.join(5000);
+		assertThat(runReturned.get())
+			.describedAs("run() should return after transport closes")
+			.isTrue();
+	}
+
+	/**
+	 * Gets the names of all threads that start with "acp-".
+	 * @return set of ACP thread names
+	 */
+	private Set<String> getAcpThreadNames() {
+		return Thread.getAllStackTraces().keySet().stream()
+			.map(Thread::getName)
+			.filter(name -> name.startsWith("acp-"))
+			.collect(Collectors.toSet());
+	}
+
+}