refactor: encapsulate websocket send queue

Author: kamikaz1k

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

@@ -935,11 +935,7 @@ private final class WebSocketConnectionState {
 
 		private final AtomicBoolean closed = new AtomicBoolean(false);
 
-		private final Object outboundLock = new Object();
-
-		private final ArrayDeque<String> outboundQueue = new ArrayDeque<>();
-
-		private boolean outboundSendInProgress = false;
+		private final SerializedWebSocketSender outboundSender = new SerializedWebSocketSender();
 
 		private volatile Session session;
 
@@ -978,92 +974,109 @@ void sendToClient(JSONRPCMessage message) {
 			try {
 				String payload = jsonMapper.writeValueAsString(message);
 				logger.debug("Sending streamable ACP WebSocket message: {}", payload);
-				enqueueOutbound(payload);
+				outboundSender.send(payload);
 			}
 			catch (Exception e) {
 				remoteConnection.signalException(e);
 				close(StatusCode.SERVER_ERROR, "failed to send ACP message");
 			}
 		}
 
-		private void enqueueOutbound(String payload) {
-			boolean shouldDrain;
-			synchronized (outboundLock) {
-				if (closed.get()) {
-					throw new AcpConnectionException("Streamable ACP WebSocket connection is closed");
+		void close() {
+			close(StatusCode.NORMAL, "server closing");
+		}
+
+		void close(int statusCode, String reason) {
+			if (!closed.compareAndSet(false, true)) {
+				return;
+			}
+			outboundSender.close();
+			webSocketConnections.remove(id, this);
+			Session currentSession = this.session;
+			if (currentSession != null && currentSession.isOpen()) {
+				currentSession.close(statusCode, reason, Callback.NOOP);
+			}
+			remoteConnection.closeGracefully().subscribe();
+		}
+
+		private final class SerializedWebSocketSender {
+
+			private final Object lock = new Object();
+
+			private final ArrayDeque<String> queue = new ArrayDeque<>();
+
+			private boolean sendInProgress = false;
+
+			void send(String payload) {
+				boolean shouldDrain;
+				synchronized (lock) {
+					if (closed.get()) {
+						throw new AcpConnectionException("Streamable ACP WebSocket connection is closed");
+					}
+					queue.addLast(payload);
+					shouldDrain = !sendInProgress;
+					if (shouldDrain) {
+						sendInProgress = true;
+					}
 				}
-				outboundQueue.addLast(payload);
-				shouldDrain = !outboundSendInProgress;
 				if (shouldDrain) {
-					outboundSendInProgress = true;
+					drain();
 				}
 			}
-			if (shouldDrain) {
-				drainOutbound();
-			}
-		}
 
-		private void drainOutbound() {
-			String payload;
-			Session currentSession;
-			synchronized (outboundLock) {
-				if (closed.get()) {
-					outboundQueue.clear();
-					outboundSendInProgress = false;
-					return;
+			/*
+			 * Jetty WebSocket sessions do not allow overlapping writes. Agent messages can
+			 * be produced by concurrent prompt handlers, so this per-connection queue sends
+			 * exactly one frame at a time and advances only after Jetty completes the
+			 * callback for the previous frame.
+			 */
+			private void drain() {
+				String payload;
+				Session currentSession;
+				synchronized (lock) {
+					if (closed.get()) {
+						clear();
+						return;
+					}
+					payload = queue.pollFirst();
+					if (payload == null) {
+						sendInProgress = false;
+						return;
+					}
+					currentSession = session;
 				}
-				payload = outboundQueue.pollFirst();
-				if (payload == null) {
-					outboundSendInProgress = false;
+
+				if (currentSession == null || !currentSession.isOpen()) {
+					fail(new AcpConnectionException("Streamable ACP WebSocket connection is closed"));
 					return;
 				}
-				currentSession = this.session;
-			}
 
-			if (currentSession == null || !currentSession.isOpen()) {
-				failOutbound(new AcpConnectionException("Streamable ACP WebSocket connection is closed"));
-				return;
+				try {
+					currentSession.sendText(payload, Callback.from(this::drain, this::fail));
+				}
+				catch (Exception e) {
+					fail(e);
+				}
 			}
 
-			try {
-				/*
-				 * Jetty WebSocket sessions do not allow overlapping writes. Agent messages can
-				 * be produced by concurrent prompt handlers, so this per-connection queue sends
-				 * exactly one frame at a time and advances only after Jetty completes the
-				 * callback for the previous frame.
-				 */
-				currentSession.sendText(payload, Callback.from(this::drainOutbound, this::failOutbound));
-			}
-			catch (Exception e) {
-				failOutbound(e);
+			private void fail(Throwable error) {
+				if (!closed.get()) {
+					remoteConnection.signalException(error);
+					WebSocketConnectionState.this.close(StatusCode.SERVER_ERROR, "failed to send ACP message");
+				}
 			}
-		}
 
-		private void failOutbound(Throwable error) {
-			if (!closed.get()) {
-				remoteConnection.signalException(error);
-				close(StatusCode.SERVER_ERROR, "failed to send ACP message");
+			void close() {
+				clear();
 			}
-		}
-
-		void close() {
-			close(StatusCode.NORMAL, "server closing");
-		}
 
-		void close(int statusCode, String reason) {
-			if (!closed.compareAndSet(false, true)) {
-				return;
-			}
-			synchronized (outboundLock) {
-				outboundQueue.clear();
-				outboundSendInProgress = false;
-			}
-			webSocketConnections.remove(id, this);
-			Session currentSession = this.session;
-			if (currentSession != null && currentSession.isOpen()) {
-				currentSession.close(statusCode, reason, Callback.NOOP);
+			private void clear() {
+				synchronized (lock) {
+					queue.clear();
+					sendInProgress = false;
+				}
 			}
-			remoteConnection.closeGracefully().subscribe();
+
 		}
 
 	}