Send the initialization notification

quarkiverse · Feb 5, 2025 · 62987b4 · 62987b4
1 parent 1542582
commit 62987b4
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 16 deletions.
diff --git a/docs/modules/ROOT/pages/includes/attributes.adoc b/docs/modules/ROOT/pages/includes/attributes.adoc
@@ -1,3 +1,3 @@
 :project-version: 0.24.0.CR1
-:langchain4j-version: 1.0.0-alpha1
-:examples-dir: ./../examples/
+:langchain4j-version: 1.0.0-alpha2-SNAPSHOT
+:examples-dir: ./../examples/
diff --git a/mcp/deployment/src/test/java/io/quarkiverse/langchain4j/mcp/test/MockHttpMcpServer.java b/mcp/deployment/src/test/java/io/quarkiverse/langchain4j/mcp/test/MockHttpMcpServer.java
@@ -85,6 +85,7 @@ public class MockHttpMcpServer {
     private volatile SseEventSink sink;
     private volatile Sse sse;
     private final ObjectMapper objectMapper = new ObjectMapper();
+    private volatile boolean initializationNotificationReceived = false;
 
     @Inject
     ScheduledExecutorService scheduledExecutorService;
@@ -111,12 +112,21 @@ public Response post(JsonNode message) {
         if (method.equals("notifications/cancelled")) {
             return Response.ok().build();
         }
+        if (method.equals("notifications/initialized")) {
+            if (initializationNotificationReceived) {
+                return Response.serverError().entity("Duplicate 'notifications/initialized' message").build();
+            }
+            initializationNotificationReceived = true;
+            return Response.ok().build();
+        }
         String operationId = message.get("id").asText();
         if (method.equals("initialize")) {
             initialize(operationId);
         } else if (method.equals("tools/list")) {
+            ensureInitialized();
             listTools(operationId);
         } else if (method.equals("tools/call")) {
+            ensureInitialized();
             if (message.get("params").get("name").asText().equals("add")) {
                 executeAddOperation(message, operationId);
             } else if (message.get("params").get("name").asText().equals("longRunningOperation")) {
@@ -128,6 +138,13 @@ public Response post(JsonNode message) {
         return Response.accepted().build();
     }
 
+    // throw an exception if we haven't received the 'notifications/initialized' message yet
+    private void ensureInitialized() {
+        if (!initializationNotificationReceived) {
+            throw new IllegalStateException("The client has not sent the 'notifications/initialized' message yet");
+        }
+    }
+
     private void listTools(String operationId) {
         String response = TOOLS_LIST_RESPONSE.formatted(operationId);
         sink.send(sse.newEventBuilder()

diff --git a/...me/src/main/java/io/quarkiverse/langchain4j/mcp/runtime/http/QuarkusHttpMcpTransport.java b/...me/src/main/java/io/quarkiverse/langchain4j/mcp/runtime/http/QuarkusHttpMcpTransport.java
@@ -16,6 +16,7 @@
 import com.fasterxml.jackson.databind.JsonNode;
 
 import dev.langchain4j.mcp.client.protocol.CancellationNotification;
+import dev.langchain4j.mcp.client.protocol.InitializationNotification;
 import dev.langchain4j.mcp.client.protocol.McpCallToolRequest;
 import dev.langchain4j.mcp.client.protocol.McpClientMessage;
 import dev.langchain4j.mcp.client.protocol.McpInitializeRequest;
@@ -24,6 +25,7 @@
 import dev.langchain4j.mcp.client.transport.McpTransport;
 import io.quarkiverse.langchain4j.QuarkusJsonCodecFactory;
 import io.quarkus.rest.client.reactive.QuarkusRestClientBuilder;
+import io.smallrye.mutiny.Uni;
 
 public class QuarkusHttpMcpTransport implements McpTransport {
 
@@ -79,37 +81,49 @@ public void start(McpOperationHandler messageHandler) {
 
     @Override
     public CompletableFuture<JsonNode> initialize(McpInitializeRequest request) {
-        return execute(request, request.getId());
+        return execute(request, request.getId()).onItem()
+                .transformToUni(
+                        response -> execute(new InitializationNotification(), null).onItem().transform(ignored -> response))
+                .subscribeAsCompletionStage();
     }
 
     @Override
     public CompletableFuture<JsonNode> listTools(McpListToolsRequest operation) {
-        return execute(operation, operation.getId());
+        return execute(operation, operation.getId()).subscribeAsCompletionStage();
     }
 
     @Override
     public void cancelOperation(long operationId) {
         CancellationNotification cancellationNotification = new CancellationNotification(operationId, "Timeout");
-        execute(cancellationNotification, null);
+        execute(cancellationNotification, null).subscribe().with(ignored -> {
+        });
     }
 
     @Override
     public CompletableFuture<JsonNode> executeTool(McpCallToolRequest operation) {
-        return execute(operation, operation.getId());
+        return execute(operation, operation.getId()).subscribeAsCompletionStage();
     }
 
-    private CompletableFuture<JsonNode> execute(McpClientMessage request, Long id) {
+    private Uni<JsonNode> execute(McpClientMessage request, Long id) {
         CompletableFuture<JsonNode> future = new CompletableFuture<>();
+        Uni<JsonNode> uni = Uni.createFrom().completionStage(future);
         if (id != null) {
             operationHandler.startOperation(id, future);
         }
-        postEndpoint.post(request).onItem().invoke(response -> {
-            int statusCode = response.getStatus();
-            if (!isExpectedStatusCode(statusCode)) {
-                throw new RuntimeException("Unexpected status code: " + statusCode);
-            }
-        }).subscribeAsCompletionStage();
-        return future;
+        postEndpoint.post(request)
+                .onFailure().invoke(future::completeExceptionally)
+                .onItem().invoke(response -> {
+                    int statusCode = response.getStatus();
+                    if (!isExpectedStatusCode(statusCode)) {
+                        future.completeExceptionally(new RuntimeException("Unexpected status code: " + statusCode));
+                    }
+                    // For messages with null ID, we don't wait for a response in the SSE channel,
+                    // so if the server accepted the request, we consider the operation done
+                    if (id == null) {
+                        future.complete(null);
+                    }
+                }).subscribeAsCompletionStage();
+        return uni;
     }
 
     private boolean isExpectedStatusCode(int statusCode) {

diff --git a/pom.xml b/pom.xml
@@ -34,8 +34,8 @@
   <properties>
     <quarkus.version>3.15.2</quarkus.version>
     <quarkus.extension-processor.version>3.18.0</quarkus.extension-processor.version>
-    <langchain4j.version>1.0.0-alpha1</langchain4j.version>
-    <langchain4j-embeddings.version>1.0.0-alpha1</langchain4j-embeddings.version>
+    <langchain4j.version>1.0.0-beta1</langchain4j.version>
+    <langchain4j-embeddings.version>1.0.0-beta1</langchain4j-embeddings.version>
     <quarkus-antora.version>1.0.2</quarkus-antora.version>
     <quarkus-poi.version>2.0.4</quarkus-poi.version> <!-- we need to use this version because langchain4j uses POI 5.2.3 instead of 5.2.5 and the substitution needed is different in the two versions -->
     <assertj.version>3.27.3</assertj.version>