Track active prompts per ACP session

Author: kamikaz1k

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

@@ -6,11 +6,11 @@
 
 import java.time.Duration;
 import java.util.Map;
+import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.Executors;
 import java.util.concurrent.atomic.AtomicLong;
-import java.util.concurrent.atomic.AtomicReference;
 
 import reactor.core.scheduler.Scheduler;
 import reactor.core.scheduler.Schedulers;
@@ -78,10 +78,17 @@ public class AcpAgentSession implements AcpSession {
 	private final AtomicLong requestCounter = new AtomicLong(0);
 
 	/**
-	 * Active prompt tracking for single-turn enforcement.
-	 * Only ONE prompt can be active at a time per ACP session.
+	 * Active prompt tracking for single-turn enforcement, keyed by logical ACP
+	 * sessionId.
+	 *
+	 * <p>
+	 * Kotlin SDK precedent: its Agent.SessionWrapper owns a single active prompt guard
+	 * per logical session wrapper. This Java session can multiplex multiple logical ACP
+	 * sessionIds over one transport connection, so the same single-turn rule needs to
+	 * be applied per sessionId instead of once for the whole connection.
+	 * </p>
 	 */
-	private final AtomicReference<ActivePrompt> activePrompt = new AtomicReference<>(null);
+	private final ConcurrentHashMap<String, ActivePrompt> activePrompts = new ConcurrentHashMap<>();
 
 	/**
 	 * Represents an active prompt session for single-turn enforcement.
@@ -235,12 +242,12 @@ private Mono<AcpSchema.JSONRPCResponse> handleIncomingRequest(AcpSchema.JSONRPCR
 				String sessionId = extractSessionId(request.params());
 				ActivePrompt newPrompt = new ActivePrompt(sessionId, request.id());
 
-				// Try to set as active prompt - fails if another prompt is active
-				if (!activePrompt.compareAndSet(null, newPrompt)) {
-					ActivePrompt current = activePrompt.get();
-					logger.warn("Rejected concurrent prompt request. Active prompt: sessionId={}, requestId={}",
-							current != null ? current.sessionId() : "unknown",
-							current != null ? current.requestId() : "unknown");
+				// Try to set as active prompt - fails if this logical session already has
+				// a prompt active.
+				ActivePrompt current = activePrompts.putIfAbsent(sessionId, newPrompt);
+				if (current != null) {
+					logger.warn("Rejected concurrent prompt request for sessionId={}. Active requestId={}", sessionId,
+							current.requestId());
 					return Mono.just(new AcpSchema.JSONRPCResponse(AcpSchema.JSONRPC_VERSION, request.id(), null,
 							new AcpSchema.JSONRPCError(-32000, "There is already an active prompt execution", null)));
 				}
@@ -249,8 +256,8 @@ private Mono<AcpSchema.JSONRPCResponse> handleIncomingRequest(AcpSchema.JSONRPCR
 				return handler.handle(request.params())
 					.map(result -> new AcpSchema.JSONRPCResponse(AcpSchema.JSONRPC_VERSION, request.id(), result, null))
 					.doFinally(signal -> {
-						activePrompt.compareAndSet(newPrompt, null);
-						logger.debug("Prompt completed with signal: {}", signal);
+						activePrompts.remove(sessionId, newPrompt);
+						logger.debug("Prompt completed for sessionId={} with signal: {}", sessionId, signal);
 					});
 			}
 
@@ -262,8 +269,13 @@ private Mono<AcpSchema.JSONRPCResponse> handleIncomingRequest(AcpSchema.JSONRPCR
 	/**
 	 * Extracts the sessionId from request parameters.
 	 */
-	@SuppressWarnings("unchecked")
 	private String extractSessionId(Object params) {
+		if (params instanceof AcpSchema.PromptRequest promptRequest) {
+			return promptRequest.sessionId() != null ? promptRequest.sessionId() : "unknown";
+		}
+		if (params instanceof AcpSchema.CancelNotification cancelNotification) {
+			return cancelNotification.sessionId() != null ? cancelNotification.sessionId() : "unknown";
+		}
 		if (params instanceof Map<?, ?> map) {
 			Object sessionId = map.get("sessionId");
 			return sessionId != null ? sessionId.toString() : "unknown";
@@ -289,9 +301,8 @@ private Mono<Void> handleIncomingNotification(AcpSchema.JSONRPCNotification noti
 			// Handle cancel notification specially
 			if (AcpSchema.METHOD_SESSION_CANCEL.equals(notification.method())) {
 				String sessionId = extractSessionId(notification.params());
-				ActivePrompt current = activePrompt.get();
-				if (current != null && sessionId.equals(current.sessionId())) {
-					activePrompt.compareAndSet(current, null);
+				ActivePrompt current = activePrompts.remove(sessionId);
+				if (current != null) {
 					logger.debug("Cancelled active prompt for session: {}", sessionId);
 				}
 			}
@@ -372,16 +383,39 @@ public Mono<Void> sendNotification(String method, Object params) {
 	 * @return true if a prompt is currently active
 	 */
 	public boolean hasActivePrompt() {
-		return activePrompt.get() != null;
+		return !activePrompts.isEmpty();
+	}
+
+	/**
+	 * Checks if there is an active prompt being processed for the specified logical
+	 * ACP session.
+	 * @param sessionId the logical ACP session ID
+	 * @return true if a prompt is currently active for the session
+	 */
+	public boolean hasActivePrompt(String sessionId) {
+		Assert.hasText(sessionId, "The sessionId can not be empty");
+		return activePrompts.containsKey(sessionId);
 	}
 
 	/**
-	 * Gets the session ID of the active prompt, if any.
-	 * @return the session ID or null if no prompt is active
+	 * Gets one active prompt session ID, if any.
+	 *
+	 * <p>
+	 * This is a legacy aggregate view. When multiple logical ACP sessions are active on
+	 * the same transport connection, the returned session ID is arbitrary.
+	 * </p>
+	 * @return one active session ID or null if no prompt is active
 	 */
 	public String getActivePromptSessionId() {
-		ActivePrompt current = activePrompt.get();
-		return current != null ? current.sessionId() : null;
+		return activePrompts.keySet().stream().findFirst().orElse(null);
+	}
+
+	/**
+	 * Gets the logical ACP session IDs that currently have active prompts.
+	 * @return an immutable snapshot of active prompt session IDs
+	 */
+	public Set<String> getActivePromptSessionIds() {
+		return Set.copyOf(activePrompts.keySet());
 	}
 
 	/**
@@ -391,7 +425,7 @@ public String getActivePromptSessionId() {
 	@Override
 	public Mono<Void> closeGracefully() {
 		return Mono.fromRunnable(() -> {
-			activePrompt.set(null);
+			activePrompts.clear();
 			dismissPendingResponses();
 			timeoutScheduler.dispose();
 		}).then(this.transport.closeGracefully());
@@ -402,7 +436,7 @@ public Mono<Void> closeGracefully() {
 	 */
 	@Override
 	public void close() {
-		activePrompt.set(null);
+		activePrompts.clear();
 		dismissPendingResponses();
 		timeoutScheduler.dispose();
 		transport.close();