Bug-fix : Async Tool Schema Generation

dev/io.openliberty.mcp.internal/bnd.bnd

@@ -50,11 +50,11 @@ Bundle-Localization: \
 
 
 -testpath: \
-    org.hamcrest:hamcrest-all;version=1.3, \
-    ../build.sharedResources/lib/junit/old/junit.jar;version=file, \
-    io.openliberty.jakarta.jsonb.3.0;version=latest,\
-    io.openliberty.org.eclipse.yasson.3.0;version=latest,\
-    io.openliberty.org.eclipse.parsson.1.1;version=latest,\
-    io.openliberty.mcp;version=latest,\
-    org.json:json;version=20080701,\
-    org.skyscreamer:jsonassert
+	org.hamcrest:hamcrest-all;version=1.3,\
+	../build.sharedResources/lib/junit/old/junit.jar;version=file,\
+	io.openliberty.jakarta.jsonb.3.0;version=latest,\
+	io.openliberty.org.eclipse.yasson.3.0;version=latest,\
+	io.openliberty.org.eclipse.parsson.1.1;version=latest,\
+	io.openliberty.mcp;version=latest,\
+	org.json:json;version=20080701,\
+	org.skyscreamer:jsonassert

dev/io.openliberty.mcp.internal/src/io/openliberty/mcp/internal/ToolMetadata.java

@@ -100,6 +100,9 @@ public static ToolMetadata createFrom(Tool annotation, Bean<?> bean, AnnotatedMe
         String description = annotation.description().isEmpty() ? null : annotation.description();
 
         Type returnType = method.getJavaMember().getGenericReturnType();
+        Type outputType = resolveOutputType(returnType);
+        Type unwrappedOutputType = resolveOutputType(outputType);
+
         Class<?> returnTypeClass = method.getJavaMember().getReturnType();
 
         WrapBusinessError wrapAnnotation = method.getAnnotation(WrapBusinessError.class);
@@ -108,20 +111,42 @@ public static ToolMetadata createFrom(Tool annotation, Bean<?> bean, AnnotatedMe
         SchemaRegistry sr = SchemaRegistry.get();
 
         JsonObject inputSchema = sr.getToolInputSchema(method);
+        boolean isString = unwrappedOutputType.equals(String.class);
+        boolean isContent = unwrappedOutputType.equals(Content.class);
 
         boolean hasContentListReturn = (returnType instanceof ParameterizedType pt && ((Class<?>) pt.getRawType()).isAssignableFrom(List.class)
                                         && (pt.getActualTypeArguments()[0] instanceof Class<?>) && ((Class<?>) pt.getActualTypeArguments()[0]).isAssignableFrom(Content.class));
         boolean hasOutputSchema = (!returnTypeClass.isAssignableFrom(ToolResponse.class) && !hasContentListReturn && !returnTypeClass.isAssignableFrom(Content.class)
                                    && !returnTypeClass.isAssignableFrom(String.class) && annotation.structuredContent());
 
+        if (annotation.structuredContent()) {
+            // Skip schema for basic string/content types
+            if (isString || isContent || hasContentListReturn) {
+                hasOutputSchema = false;
+            } else if (returnTypeClass.equals(ToolResponse.class)) {
+                // Only generate schema if @Schema explicitly provided
+                hasOutputSchema = method.isAnnotationPresent(Schema.class) &&
+                                  !method.getAnnotation(Schema.class).value().equals(Schema.UNSET);
+            } else if (returnType instanceof ParameterizedType pt &&
+                       CompletionStage.class.isAssignableFrom((Class<?>) pt.getRawType())) {
+
+                Type inner = pt.getActualTypeArguments()[0];
+                unwrappedOutputType = resolveOutputType(inner);
+                hasOutputSchema = true;
+            } else {
+                // if it's not a String/Content, we want schema
+                hasOutputSchema = true;
+            }
+        }
+
         if (!hasOutputSchema && returnTypeClass.isAssignableFrom(ToolResponse.class) && annotation.structuredContent()
             && method.isAnnotationPresent(Schema.class)
             && method.getAnnotation(Schema.class).value() != Schema.UNSET) {
 
             hasOutputSchema = true;
 
         }
-        JsonObject outputSchema = hasOutputSchema ? sr.getToolOutputSchema(method) : null;
+        JsonObject outputSchema = hasOutputSchema ? sr.getToolOutputSchema(method, unwrappedOutputType) : null;
 
         outputSchema = (outputSchema == null || outputSchema.isEmpty()) ? null : outputSchema;
 
@@ -246,4 +271,22 @@ public String getToolQualifiedName() {
     public static String getToolQualifiedName(Bean<?> bean, AnnotatedMethod<?> method) {
         return bean.getBeanClass() + "." + method.getJavaMember().getName();
     }
+
+    /**
+     * Resolves the type inside CompletionStage if applicable.
+     *
+     * @param returnType the return type to resolve
+     * @return the inner type or the original type if not wrapped
+     */
+    private static Type resolveOutputType(Type returnType) {
+        if (returnType instanceof ParameterizedType pt) {
+            Type raw = pt.getRawType();
+            if (raw instanceof Class<?> rawClass &&
+                CompletionStage.class.isAssignableFrom(rawClass)) {
+                return pt.getActualTypeArguments()[0];
+            }
+        }
+        return returnType;
+    }
+
 }

dev/io.openliberty.mcp.internal/src/io/openliberty/mcp/internal/schemas/SchemaGenerator.java

@@ -132,23 +132,22 @@ public static JsonObject generateToolInputSchema(AnnotatedMethod<?> toolMethod,
     /**
      * Generate the output schema for a tool
      *
-     * @param tool the tool to generate the schema for
+     * @param toolMethod the tool method to generate a schema for
+     * @param toolOutputType the output type of the tool, same as {@code toolMethod} for sync tools, unwrapped for async tools
      * @param blueprintRegistry the blueprint registry to use
      * @return the schema as a json object
      */
-    public static JsonObject generateToolOutputSchema(AnnotatedMethod<?> toolMethod, SchemaCreationBlueprintRegistry blueprintRegistry) {
-
-        Type returnType = toolMethod.getJavaMember().getGenericReturnType();
+    public static JsonObject generateToolOutputSchema(AnnotatedMethod<?> toolMethod, Type toolOutputType, SchemaCreationBlueprintRegistry blueprintRegistry) {
 
         Method method = toolMethod.getJavaMember();
         SchemaAnnotation schemaAnnotation = SchemaAnnotation.read(method);
 
         SchemaGenerationContext ctx = new SchemaGenerationContext(blueprintRegistry, OUTPUT);
         if (!method.getReturnType().isAssignableFrom(ToolResponse.class)) {
-            calculateClassFrequency(returnType, SchemaDirection.OUTPUT, ctx);
+            calculateClassFrequency(toolOutputType, SchemaDirection.OUTPUT, ctx);
         }
 
-        JsonObjectBuilder outputSchema = generateSubSchema(returnType, ctx, schemaAnnotation);
+        JsonObjectBuilder outputSchema = generateSubSchema(toolOutputType, ctx, schemaAnnotation);
         addDefs(outputSchema, ctx);
 
         return outputSchema.build();

dev/io.openliberty.mcp.internal/src/io/openliberty/mcp/internal/schemas/SchemaRegistry.java

@@ -9,6 +9,7 @@
  *******************************************************************************/
 package io.openliberty.mcp.internal.schemas;
 
+import java.lang.reflect.Type;
 import java.util.HashMap;
 
 import io.openliberty.mcp.internal.McpCdiExtension;
@@ -78,9 +79,9 @@ public JsonObject getToolInputSchema(AnnotatedMethod<?> toolMethod) {
      * @param toolMetadata the tool to get the schema for
      * @return the json schema
      */
-    public JsonObject getToolOutputSchema(AnnotatedMethod<?> toolMethod) {
+    public JsonObject getToolOutputSchema(AnnotatedMethod<?> toolMethod, Type toolOutputType) {
         ToolKey key = new ToolKey(toolMethod, SchemaDirection.OUTPUT);
-        return schemaCache.computeIfAbsent(key, k -> SchemaGenerator.generateToolOutputSchema(toolMethod, blueprintRegistry));
+        return schemaCache.computeIfAbsent(key, k -> SchemaGenerator.generateToolOutputSchema(toolMethod, toolOutputType, blueprintRegistry));
     }
 
     /**

dev/io.openliberty.mcp.internal/test/io/openliberty/mcp/internal/test/schema/SchemaTest.java

@@ -9,6 +9,7 @@
  *******************************************************************************/
 package io.openliberty.mcp.internal.test.schema;
 
+import java.lang.reflect.Type;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -330,8 +331,9 @@ public void testToolInputSchema() throws NoSuchMethodException, SecurityExceptio
     @Test
     public void testToolOutputSchema() throws NoSuchMethodException, SecurityException {
         MockAnnotatedMethod<Object> toolMethod = TestUtils.findMethod(SchemaTest.class, "updateWidget");
+        Type returnType = toolMethod.getJavaMember().getGenericReturnType();
 
-        String toolInputSchema = registry.getToolOutputSchema(toolMethod).toString();
+        String toolInputSchema = registry.getToolOutputSchema(toolMethod, returnType).toString();
         String expectedSchema = """
                         {
                             "type": "object",
@@ -414,7 +416,9 @@ public void testToolInputRecursive() {
     @Test
     public void testToolOutputRecursive() {
         MockAnnotatedMethod<Object> toolMethod = TestUtils.findMethod(SchemaTest.class, "combineWidgets");
-        String toolInputSchema = registry.getToolOutputSchema(toolMethod).toString();
+        Type returnType = toolMethod.getJavaMember().getGenericReturnType();
+
+        String toolInputSchema = registry.getToolOutputSchema(toolMethod, returnType).toString();
         String expectedSchema = """
                         {
                             "$defs": {
@@ -999,7 +1003,9 @@ public void testPersonAddtoListToolInputSchema() throws NoSuchMethodException, S
     @Test
     public void testPersonAddtoListToolOutputSchema() throws NoSuchMethodException, SecurityException {
         MockAnnotatedMethod<Object> toolMethod = TestUtils.findMethod(SchemaTest.class, "addPersonToList");
-        String response = registry.getToolOutputSchema(toolMethod).toString();
+        Type returnType = toolMethod.getJavaMember().getGenericReturnType();
+
+        String response = registry.getToolOutputSchema(toolMethod, returnType).toString();
         String expectedResponseString = """
                             {
                             "$defs": {
@@ -2258,9 +2264,10 @@ public int[] primitiveArrayTest(@ToolArg(name = "name", description = "name") St
     @Test
     public void testPrimitiveArray() {
         MockAnnotatedMethod<Object> toolMethod = TestUtils.findMethod(SchemaTest.class, "primitiveArrayTest");
-        String response = registry.getToolOutputSchema(toolMethod).toString();
+        Type returnType = toolMethod.getJavaMember().getGenericReturnType();
+        String response = registry.getToolOutputSchema(toolMethod, returnType).toString();
         String expectedResponseString = """
-                                    {"type":"array","items":{"type":"integer"}}
+                            {"type":"array","items":{"type":"integer"}}
                         """;
         JSONAssert.assertEquals(expectedResponseString, response, true);
     }

dev/io.openliberty.mcp.internal_fat/fat/src/io/openliberty/mcp/internal/fat/tool/AsyncToolsTest.java

@@ -10,16 +10,22 @@
 package io.openliberty.mcp.internal.fat.tool;
 
 import static com.ibm.websphere.simplicity.ShrinkHelper.DeployOptions.SERVER_ONLY;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 
+import java.util.List;
+
 import org.jboss.shrinkwrap.api.ShrinkWrap;
 import org.jboss.shrinkwrap.api.spec.WebArchive;
+import org.json.JSONArray;
+import org.json.JSONObject;
 import org.junit.AfterClass;
 import org.junit.BeforeClass;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.skyscreamer.jsonassert.JSONAssert;
+import org.skyscreamer.jsonassert.JSONParser;
 
 import com.ibm.websphere.simplicity.ShrinkHelper;
 
@@ -38,7 +44,6 @@
 public class AsyncToolsTest extends FATServletClient {
 
     private static final String EXPECTED_ERROR = "Method call caused runtime exception. This is expected if the input was 'throw error'";
-
     @Server("mcp-server-async")
     public static LibertyServer server;
 
@@ -301,4 +306,124 @@ public void testCompletionStageThatNeverCompletes() throws Exception {
                         """;
         client.callMCP(request);
     }
+
+    @Test
+    public void testAsyncObjectToolHasExpectedOutputSchema() throws Exception {
+        String response = client.callMCP("""
+                            {
+                              "jsonrpc": "2.0",
+                              "id": 1,
+                              "method": "tools/list"
+                            }
+                        """);
+
+        JSONObject root = (JSONObject) JSONParser.parseJSON(response);
+        JSONArray tools = root.getJSONObject("result").getJSONArray("tools");
+
+        JSONObject asyncObjectTool = null;
+
+        for (int i = 0; i < tools.length(); i++) {
+            JSONObject tool = tools.getJSONObject(i);
+            if ("asyncObjectTool".equals(tool.getString("name"))) {
+                asyncObjectTool = tool;
+                break;
+            }
+        }
+
+        assertNotNull("Tool 'asyncObjectTool' should be present in tool list", asyncObjectTool);
+
+        String actualToolJson = asyncObjectTool.toString();
+
+        String expectedToolJson = """
+                                                    {
+                          "name": "asyncObjectTool",
+                          "title": "Async asyncObjectTool",
+                          "description": "A tool to return an object of cities asynchronously",
+                          "inputSchema": {
+                            "type": "object",
+                            "properties": {
+                              "name": {
+                                "type": "string",
+                                "description": "name of your city"
+                              }
+                            },
+                            "required": ["name"]
+                          },
+                          "outputSchema": {
+                            "type": "object",
+                            "properties": {
+                              "country": {"type": "string"},
+                              "isCapital": {"type": "boolean"},
+                              "name": {"type": "string"},
+                              "population": {"type": "integer"}
+                            },
+                            "required": ["name", "country", "population", "isCapital"]
+                          }
+                        }
+                         """;
+
+        JSONAssert.assertEquals(expectedToolJson, actualToolJson, true);
+    }
+
+    @Test
+    public void testAsyncToolsWithoutStructuredContentHaveNoOutputSchema() throws Exception {
+        String response = client.callMCP("""
+                            {
+                              "jsonrpc": "2.0",
+                              "id": 1,
+                              "method": "tools/list"
+                            }
+                        """);
+
+        JSONObject root = (JSONObject) JSONParser.parseJSON(response);
+        JSONArray tools = root.getJSONObject("result").getJSONArray("tools");
+
+        List<String> toolsThatShouldNotHaveOutputSchema = List.of(
+                                                                  "asyncStringTool",
+                                                                  "asyncContentTool",
+                                                                  "asyncListContentTool",
+                                                                  "asyncObjectToolWithoutStructuredContent");
+
+        for (int i = 0; i < tools.length(); i++) {
+            JSONObject tool = tools.getJSONObject(i);
+            String name = tool.getString("name");
+
+            if (toolsThatShouldNotHaveOutputSchema.contains(name)) {
+                assertFalse("Tool '" + name + "' should NOT have an output schema", tool.has("outputSchema"));
+            }
+        }
+    }
+
+    @Test
+    public void testAsyncToolsWithBasicOutputTypesShouldNotHaveOutputSchema() throws Exception {
+        String response = client.callMCP("""
+                            {
+                              "jsonrpc": "2.0",
+                              "id": 1,
+                              "method": "tools/list"
+                            }
+                        """);
+
+        JSONObject root = (JSONObject) JSONParser.parseJSON(response);
+        JSONArray tools = root.getJSONObject("result").getJSONArray("tools");
+
+        // Tools that SHOULD NOT have an output schema
+        List<String> toolsWithoutSchema = List.of(
+                                                  "asyncEcho",
+                                                  "asyncDelayedEcho",
+                                                  "asyncCancellationTool",
+                                                  "asyncToolThatNeverCompletes");
+
+        for (int i = 0; i < tools.length(); i++) {
+            JSONObject tool = tools.getJSONObject(i);
+            String name = tool.getString("name");
+
+            if (toolsWithoutSchema.contains(name)) {
+                assertFalse(
+                            "Tool '" + name + "' should NOT have an output schema",
+                            tool.has("outputSchema"));
+            }
+        }
+    }
+
 }

dev/io.openliberty.mcp.internal_fat/fat/src/io/openliberty/mcp/internal/fat/tool/asyncToolApp/AsyncTools.java

@@ -48,6 +48,11 @@ public CompletionStage<String> asyncEcho(@ToolArg(name = "input", description =
         return CompletableFuture.completedStage(input + ": (async)");
     }
 
+    @Tool(name = "asyncObjectToolWithoutStructuredContent", title = "Async Object Tool No Schema", description = "Returns a city object but no structuredContent")
+    public CompletionStage<City> asyncObjectToolWithoutStructuredContent() {
+        return executor.supplyAsync(() -> new City("Leeds", "England", 7000, false));
+    }
+
     @Tool(name = "asyncDelayedEcho", title = "Async Echo", description = "Echoes input asynchronously")
     public CompletionStage<String> asyncDelayedEcho(@ToolArg(name = "input", description = "input to echo") String input) {
         return executor.supplyAsync(() -> {

dev/io.openliberty.mcp.internal_fat/fat/src/io/openliberty/mcp/internal/fat/tool/basicToolApp/BasicTools.java

@@ -440,9 +440,10 @@ public boolean checkPerson(@ToolArg(name = "person", description = "Person objec
     }
 
     @Tool(name = "addPersonToList", title = "adds person to people list", description = "adds person to people list", structuredContent = true)
-    public @Schema(description = "Returns list of person object") List<Person> addPersonToList(@ToolArg(name = "employeeList",
-                                                                                                        description = "List of people") List<Person> employeeList,
-                                                                                               @ToolArg(name = "person", description = "Person object") Optional<Person> person) {
+    @Schema(description = "Returns list of person object")
+    public List<Person> addPersonToList(
+                                        @ToolArg(name = "employeeList", description = "List of people") List<Person> employeeList,
+                                        @ToolArg(name = "person", description = "Person object") Optional<Person> person) {
         employeeList.add(person.get());
         return employeeList;
     }
@@ -461,6 +462,7 @@ public boolean checkPerson(@ToolArg(name = "person", description = "Person objec
         _meta.put(MetaKey.from("api.ibmtest.org/location"), "Hursley");
         _meta.put(MetaKey.from("api.libertytest.org/person"), personInstance);
         return new ToolResponse(false, List.of(new TextContent(jsonb.toJson(employeeList))), employeeList, _meta);
+
     }
 
     @Tool(name = "addPersonToListToolResponseWithMetaRequest", title = "adds person to people list", description = "adds person to people list", structuredContent = true)