ws: make sure MCP connection is initialized before a message is consumed

- fixes intermittent CI failures

Author: mkouba

transports/websocket/deployment/src/main/java/io/quarkiverse/mcp/server/websocket/deployment/WebSocketMcpServerProcessor.java

@@ -68,7 +68,7 @@ void generateEndpoints(McpWebSocketServersBuildTimeConfig config, BuildProducer<
                     .className(endpointClassName)
                     .superClass(WebSocketMcpMessageHandler.class)
                     .build();
-            // @WebSocket(path = "/foo/bar")
+            // @WebSocket(path = "/foo/bar", inboundProcessingMode = InboundProcessingMode.CONCURRENT)
             endpointCreator.addAnnotation(
                     AnnotationInstance.builder(WebSocket.class)
                             .add("path", e.getValue().websocket().endpointPath())

transports/websocket/runtime/src/main/java/io/quarkiverse/mcp/server/websocket/runtime/WebSocketMcpMessageHandler.java

@@ -1,6 +1,8 @@
 package io.quarkiverse.mcp.server.websocket.runtime;
 
 import java.util.List;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
 
 import jakarta.enterprise.inject.Instance;
 
@@ -30,7 +32,6 @@
 import io.quarkus.websockets.next.OnClose;
 import io.quarkus.websockets.next.OnOpen;
 import io.quarkus.websockets.next.OnTextMessage;
-import io.quarkus.websockets.next.UserData.TypedKey;
 import io.quarkus.websockets.next.WebSocketConnection;
 import io.smallrye.mutiny.Uni;
 import io.smallrye.mutiny.vertx.UniHelper;
@@ -44,8 +45,6 @@ public abstract class WebSocketMcpMessageHandler extends McpMessageHandler<WebSo
 
     private static final Logger LOG = Logger.getLogger(WebSocketMcpMessageHandler.class);
 
-    private static final String MCP_CONNECTION_ID = "mcpConnectionId";
-
     CurrentIdentityAssociation currentIdentityAssociation;
 
     protected WebSocketMcpMessageHandler(McpServersRuntimeConfig config,
@@ -70,20 +69,17 @@ protected WebSocketMcpMessageHandler(McpServersRuntimeConfig config,
 
     protected abstract String serverName();
 
+    private final ConcurrentMap<String, WebSocketMcpConnection> connections = new ConcurrentHashMap<>();
+
     @OnOpen
     void openConnection(WebSocketConnection connection) {
-        String id = ConnectionManager.connectionId();
-        WebSocketMcpConnection mcpConnection = new WebSocketMcpConnection(id, config.servers().get(serverName()), connection);
-        connectionManager.add(mcpConnection);
-        LOG.debugf("WebSocket connection initialized [%s]", id);
-        connection.userData().put(TypedKey.forString(MCP_CONNECTION_ID), id);
+        LOG.debugf("MCP WebSocket connection open [id: %s]", connection.id());
     }
 
     @SuppressWarnings("unchecked")
     @OnTextMessage
-    Uni<Void> consumeMessage(WebSocketConnection connection, String message) {
-        String connectionId = connection.userData().get(TypedKey.forString(MCP_CONNECTION_ID));
-        WebSocketMcpConnection mcpConnection = (WebSocketMcpConnection) connectionManager.get(connectionId);
+    Uni<Void> consumeMessage(WebSocketConnection connection, String message) throws InterruptedException {
+        WebSocketMcpConnection mcpConnection = connections.computeIfAbsent(connection.id(), k -> newConnection(connection));
         Object json = Json.decodeValue(message);
 
         SecuritySupport securitySupport;
@@ -106,16 +102,26 @@ public void setCurrentIdentity(CurrentIdentityAssociation currentIdentityAssocia
 
     @OnClose
     void closeConnection(WebSocketConnection connection) {
-        String id = connection.userData().get(TypedKey.forString(MCP_CONNECTION_ID));
-        connectionManager.remove(id);
-        LOG.debugf("WebSocket connection closed [%s]", id);
+        WebSocketMcpConnection mcpConnection = connections.remove(connection.id());
+        if (mcpConnection != null) {
+            connectionManager.remove(mcpConnection.id());
+            LOG.debugf("MCP WebSocket connection closed [mcpId: %s, id: %s]", mcpConnection.id(), connection.id());
+        }
     }
 
     @Override
     protected Transport transport() {
         return Transport.WEBSOCKET;
     }
 
+    private WebSocketMcpConnection newConnection(WebSocketConnection connection) {
+        String id = ConnectionManager.connectionId();
+        WebSocketMcpConnection mcpConnection = new WebSocketMcpConnection(id, config.servers().get(serverName()), connection);
+        connectionManager.add(mcpConnection);
+        LOG.debugf("MCP WebSocket connection initialized [mcpId: %s, id: %s]", mcpConnection.id(), connection.id());
+        return mcpConnection;
+    }
+
     static class WebSocketMcpRequest extends McpRequestImpl<WebSocketMcpConnection> {
 
         WebSocketMcpRequest(String serverName, Object json, WebSocketMcpConnection connection,