Add remote WebSocket agent transport

Author: kamikaz1k

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

@@ -0,0 +1,244 @@
+/*
+ * Copyright 2025-2026 the original author or authors.
+ */
+
+package com.agentclientprotocol.sdk.agent.transport;
+
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Consumer;
+import java.util.function.Function;
+
+import com.agentclientprotocol.sdk.agent.AcpAgentFactory;
+import com.agentclientprotocol.sdk.agent.AcpAsyncAgent;
+import com.agentclientprotocol.sdk.error.AcpConnectionException;
+import com.agentclientprotocol.sdk.json.AcpJsonMapper;
+import com.agentclientprotocol.sdk.json.TypeRef;
+import com.agentclientprotocol.sdk.spec.AcpAgentTransport;
+import com.agentclientprotocol.sdk.spec.AcpSchema;
+import com.agentclientprotocol.sdk.spec.AcpSchema.JSONRPCMessage;
+import com.agentclientprotocol.sdk.util.Assert;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import reactor.core.publisher.Mono;
+import reactor.core.publisher.Sinks;
+
+/**
+ * Shared per-connection core for listener-backed remote ACP agent transports.
+ *
+ * <p>
+ * Remote transports such as Streamable HTTP and WebSocket have different wire-level
+ * framing, but they both need the same agent-side shape once a remote ACP connection
+ * exists: one connection-bound {@link AcpAgentTransport}, one fresh agent runtime from
+ * {@link AcpAgentFactory}, inbound JSON-RPC delivery to the agent, and outbound JSON-RPC
+ * delivery back to the wire adapter.
+ * </p>
+ *
+ * <p>
+ * This class intentionally does not know about HTTP headers, SSE streams, WebSocket
+ * sessions, or route maps. Those remain transport-adapter concerns.
+ * </p>
+ *
+ * @author Kaiser Dandangi
+ */
+public final class RemoteAcpConnection {
+
+	private static final Logger logger = LoggerFactory.getLogger(RemoteAcpConnection.class);
+
+	private final String id;
+
+	private final AcpJsonMapper jsonMapper;
+
+	private final ConnectionTransport transport;
+
+	private final AtomicBoolean started = new AtomicBoolean(false);
+
+	private final AtomicBoolean closing = new AtomicBoolean(false);
+
+	private volatile AcpAsyncAgent agent;
+
+	/**
+	 * Creates a new remote ACP connection core.
+	 * @param id stable transport connection id
+	 * @param jsonMapper JSON mapper used by the connection transport
+	 * @param outboundConsumer callback that receives agent-originated outbound messages
+	 */
+	public RemoteAcpConnection(String id, AcpJsonMapper jsonMapper, Consumer<JSONRPCMessage> outboundConsumer) {
+		Assert.hasText(id, "The id can not be empty");
+		Assert.notNull(jsonMapper, "The jsonMapper can not be null");
+		Assert.notNull(outboundConsumer, "The outboundConsumer can not be null");
+		this.id = id;
+		this.jsonMapper = jsonMapper;
+		this.transport = new ConnectionTransport(outboundConsumer);
+	}
+
+	/**
+	 * Returns the transport-level connection id.
+	 * @return connection id
+	 */
+	public String id() {
+		return id;
+	}
+
+	/**
+	 * Starts a fresh agent runtime for this connection.
+	 * @param agentFactory factory used to create the connection-bound agent runtime
+	 * @return mono that completes when the agent runtime is started
+	 */
+	public Mono<Void> start(AcpAgentFactory agentFactory) {
+		Assert.notNull(agentFactory, "The agentFactory can not be null");
+		if (!started.compareAndSet(false, true)) {
+			return Mono.error(new IllegalStateException("Already started"));
+		}
+		return Mono.defer(() -> {
+			this.agent = agentFactory.create(transport);
+			return this.agent.start();
+		}).doOnError(this::signalException);
+	}
+
+	/**
+	 * Accepts one client-originated JSON-RPC message for delivery to the connection's
+	 * agent runtime.
+	 * @param message inbound message
+	 */
+	public void acceptInbound(JSONRPCMessage message) {
+		transport.acceptInbound(message);
+	}
+
+	/**
+	 * Reports a transport adapter exception to the agent transport exception handler.
+	 * @param error exception to report
+	 */
+	public void signalException(Throwable error) {
+		transport.signalException(error);
+	}
+
+	/**
+	 * Closes the connection and its agent runtime gracefully.
+	 * @return mono that completes when close work has been requested
+	 */
+	public Mono<Void> closeGracefully() {
+		return Mono.defer(() -> {
+			if (!closing.compareAndSet(false, true)) {
+				return Mono.empty();
+			}
+			AcpAsyncAgent currentAgent = this.agent;
+			if (currentAgent != null) {
+				return currentAgent.closeGracefully()
+					.onErrorResume(error -> {
+						signalException(error);
+						return Mono.empty();
+					})
+					.then(transport.closeGracefully());
+			}
+			return transport.closeGracefully();
+		});
+	}
+
+	/**
+	 * Closes the connection and its agent runtime immediately.
+	 */
+	public void close() {
+		if (!closing.compareAndSet(false, true)) {
+			return;
+		}
+		AcpAsyncAgent currentAgent = this.agent;
+		if (currentAgent != null) {
+			currentAgent.close();
+		}
+		transport.close();
+	}
+
+	private final class ConnectionTransport implements AcpAgentTransport {
+
+		private final Consumer<JSONRPCMessage> outboundConsumer;
+
+		private final Sinks.Many<JSONRPCMessage> inboundSink = Sinks.many().unicast().onBackpressureBuffer();
+
+		private final Sinks.One<Void> terminationSink = Sinks.one();
+
+		private final AtomicBoolean transportStarted = new AtomicBoolean(false);
+
+		private final AtomicBoolean transportClosing = new AtomicBoolean(false);
+
+		private volatile Consumer<Throwable> exceptionHandler = t -> logger.error("Remote ACP transport error", t);
+
+		ConnectionTransport(Consumer<JSONRPCMessage> outboundConsumer) {
+			this.outboundConsumer = outboundConsumer;
+		}
+
+		@Override
+		public Mono<Void> start(Function<Mono<JSONRPCMessage>, Mono<JSONRPCMessage>> handler) {
+			Assert.notNull(handler, "The handler can not be null");
+			if (!transportStarted.compareAndSet(false, true)) {
+				return Mono.error(new IllegalStateException("Already started"));
+			}
+			inboundSink.asFlux()
+				.flatMap(message -> Mono.just(message).transform(handler))
+				.doOnNext(response -> {
+					if (response != null) {
+						outboundConsumer.accept(response);
+					}
+				})
+				.doOnError(this::signalException)
+				.doFinally(signal -> terminationSink.tryEmitValue(null))
+				.subscribe();
+			return Mono.empty();
+		}
+
+		void acceptInbound(JSONRPCMessage message) {
+			Assert.notNull(message, "The message can not be null");
+			if (transportClosing.get()) {
+				throw new AcpConnectionException("Remote ACP connection is closing");
+			}
+			Sinks.EmitResult result = inboundSink.tryEmitNext(message);
+			if (result.isFailure()) {
+				throw new AcpConnectionException("Failed to enqueue inbound message: " + result);
+			}
+		}
+
+		void signalException(Throwable error) {
+			exceptionHandler.accept(error);
+		}
+
+		@Override
+		public Mono<Void> sendMessage(JSONRPCMessage message) {
+			return Mono.fromRunnable(() -> {
+				if (transportClosing.get()) {
+					throw new AcpConnectionException("Remote ACP connection is closing");
+				}
+				outboundConsumer.accept(message);
+			});
+		}
+
+		@Override
+		public <T> T unmarshalFrom(Object data, TypeRef<T> typeRef) {
+			return jsonMapper.convertValue(data, typeRef);
+		}
+
+		@Override
+		public Mono<Void> closeGracefully() {
+			return Mono.fromRunnable(this::close);
+		}
+
+		@Override
+		public void close() {
+			if (transportClosing.compareAndSet(false, true)) {
+				inboundSink.tryEmitComplete();
+				terminationSink.tryEmitValue(null);
+			}
+		}
+
+		@Override
+		public void setExceptionHandler(Consumer<Throwable> handler) {
+			Assert.notNull(handler, "The handler can not be null");
+			this.exceptionHandler = handler;
+		}
+
+		@Override
+		public Mono<Void> awaitTermination() {
+			return terminationSink.asMono();
+		}
+
+	}
+
+}

acp-core/src/test/java/com/agentclientprotocol/sdk/spec/AcpAgentSessionTest.java

@@ -16,6 +16,7 @@
 import com.agentclientprotocol.sdk.test.InMemoryTransportPair;
 import org.junit.jupiter.api.Test;
 import reactor.core.publisher.Mono;
+import reactor.core.publisher.Sinks;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.junit.jupiter.api.Assertions.assertThrows;
@@ -228,13 +229,17 @@ void singleTurnEnforcementAllowsConcurrentPromptsForDifferentSessions() throws E
 		try {
 			CountDownLatch handlersStarted = new CountDownLatch(2);
 			AtomicInteger handlerInvocations = new AtomicInteger();
+			Sinks.One<Void> session1Release = Sinks.one();
+			Sinks.One<Void> session2Release = Sinks.one();
 
 			Map<String, AcpAgentSession.RequestHandler<?>> requestHandlers = Map.of(AcpSchema.METHOD_SESSION_PROMPT,
 					params -> Mono.defer(() -> {
 						handlerInvocations.incrementAndGet();
 						handlersStarted.countDown();
-						return Mono.delay(PROMPT_RESPONSE_DELAY)
-							.map(ignored -> new AcpSchema.PromptResponse(AcpSchema.StopReason.END_TURN));
+						String sessionId = sessionId(params);
+						Sinks.One<Void> release = SESSION_1.equals(sessionId) ? session1Release : session2Release;
+						return release.asMono()
+							.thenReturn(new AcpSchema.PromptResponse(AcpSchema.StopReason.END_TURN));
 					}));
 
 			AcpAgentSession session = new AcpAgentSession(TIMEOUT, transportPair.agentTransport(), requestHandlers,
@@ -262,6 +267,13 @@ void singleTurnEnforcementAllowsConcurrentPromptsForDifferentSessions() throws E
 			assertThat(session.hasActivePrompt(SESSION_2)).isTrue();
 			assertThat(session.getActivePromptSessionIds()).containsExactlyInAnyOrder(SESSION_1, SESSION_2);
 
+			// Release the responses one at a time. The in-memory test transport uses a
+			// unicast sink, so simultaneous emissions from concurrent prompt handlers can
+			// fail with FAIL_NON_SERIALIZED and obscure the behavior under test.
+			session1Release.tryEmitValue(null);
+			awaitResponse(responses, "1");
+			session2Release.tryEmitValue(null);
+
 			assertThat(responseLatch.await(5, TimeUnit.SECONDS)).isTrue();
 
 			assertThat(responseById(responses, "1").error()).isNull();
@@ -383,6 +395,23 @@ private static AcpSchema.JSONRPCResponse responseById(List<AcpSchema.JSONRPCResp
 		return responses.stream().filter(response -> id.equals(response.id())).findFirst().orElseThrow();
 	}
 
+	private static void awaitResponse(List<AcpSchema.JSONRPCResponse> responses, Object id) throws InterruptedException {
+		long deadline = System.nanoTime() + TIMEOUT.toNanos();
+		while (System.nanoTime() < deadline) {
+			if (responses.stream().anyMatch(response -> id.equals(response.id()))) {
+				return;
+			}
+			Thread.sleep(10);
+		}
+	}
+
+	private static String sessionId(Object params) {
+		if (params instanceof AcpSchema.PromptRequest promptRequest) {
+			return promptRequest.sessionId();
+		}
+		throw new IllegalArgumentException("Expected PromptRequest params but received " + params);
+	}
+
 	private static void allowAgentTransportSubscription() throws InterruptedException {
 		// AcpAgentSession subscribes to the in-memory transport in its constructor.
 		// subscribe() is asynchronous, so give the unicast sink subscriber a short