Address feedback.

Author: eiriktsarpalis

src/ModelContextProtocol.Core/McpJsonUtilities.cs

@@ -1,5 +1,6 @@
 using Microsoft.Extensions.AI;
 using ModelContextProtocol.Protocol;
+using ModelContextProtocol.Server;
 using System.Diagnostics.CodeAnalysis;
 using System.Text.Json;
 using System.Text.Json.Serialization;
@@ -75,6 +76,30 @@ internal static bool IsValidMcpToolSchema(JsonElement element)
         return false; // No type keyword found.
     }
 
+    internal static JsonElement? GetReturnSchema(this AIFunction function, AIJsonSchemaCreateOptions? schemaCreateOptions)
+    {
+        // TODO replace with https://github.com/dotnet/extensions/pull/6447 once merged.
+        if (function.UnderlyingMethod?.ReturnType is not Type returnType)
+        {
+            return null;
+        }
+
+        if (returnType == typeof(void) || returnType == typeof(Task) || returnType == typeof(ValueTask))
+        {
+            // Do not report an output schema for void or Task methods.
+            return null;
+        }
+
+        if (returnType.IsGenericType && returnType.GetGenericTypeDefinition() is Type genericTypeDef &&
+            (genericTypeDef == typeof(Task<>) || genericTypeDef == typeof(ValueTask<>)))
+        {
+            // Extract the real type from Task<T> or ValueTask<T> if applicable.
+            returnType = returnType.GetGenericArguments()[0];
+        }
+
+        return AIJsonUtilities.CreateJsonSchema(returnType, serializerOptions: function.JsonSerializerOptions, inferenceOptions: schemaCreateOptions);
+    }
+
     // Keep in sync with CreateDefaultOptions above.
     [JsonSourceGenerationOptions(JsonSerializerDefaults.Web,
         DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull,

src/ModelContextProtocol.Core/Protocol/CallToolResponse.cs

@@ -32,7 +32,7 @@ public class CallToolResponse
     /// Gets or sets an optional JSON object representing the structured result of the tool call.
     /// </summary>
     [JsonPropertyName("structuredContent")]
-    public JsonObject? StructuredContent { get; set; }
+    public JsonNode? StructuredContent { get; set; }
 
     /// <summary>
     /// Gets or sets an indication of whether the tool call was unsuccessful.

src/ModelContextProtocol.Core/Server/AIFunctionMcpServerTool.cs

@@ -299,7 +299,7 @@ public override async ValueTask<CallToolResponse> InvokeAsync(
             };
         }
 
-        JsonObject? structuredContent = CreateStructuredResponse(result);
+        JsonNode? structuredContent = CreateStructuredResponse(result);
         return result switch
         {
             AIContent aiContent => new()
@@ -357,35 +357,18 @@ public override async ValueTask<CallToolResponse> InvokeAsync(
 
     private static JsonElement? CreateOutputSchema(AIFunction function, McpServerToolCreateOptions? toolCreateOptions, out bool structuredOutputRequiresWrapping)
     {
-        // TODO replace with https://github.com/dotnet/extensions/pull/6447 once merged.
-
         structuredOutputRequiresWrapping = false;
 
         if (toolCreateOptions?.UseStructuredContent is not true)
         {
             return null;
         }
 
-        if (function.UnderlyingMethod?.ReturnType is not Type returnType)
-        {
-            return null;
-        }
-
-        if (returnType == typeof(void) || returnType == typeof(Task) || returnType == typeof(ValueTask))
+        if (function.GetReturnSchema(toolCreateOptions?.SchemaCreateOptions) is not JsonElement outputSchema)
         {
-            // Do not report an output schema for void or Task methods.
             return null;
         }
 
-        if (returnType.IsGenericType && returnType.GetGenericTypeDefinition() is Type genericTypeDef &&
-            (genericTypeDef == typeof(Task<>) || genericTypeDef == typeof(ValueTask<>)))
-        {
-            // Extract the real type from Task<T> or ValueTask<T> if applicable.
-            returnType = returnType.GetGenericArguments()[0];
-        }
-
-        JsonElement outputSchema = AIJsonUtilities.CreateJsonSchema(returnType, serializerOptions: function.JsonSerializerOptions, inferenceOptions: toolCreateOptions?.SchemaCreateOptions);
-
         if (outputSchema.ValueKind is not JsonValueKind.Object ||
             !outputSchema.TryGetProperty("type", out JsonElement typeProperty) ||
             typeProperty.ValueKind is not JsonValueKind.String ||
@@ -423,10 +406,11 @@ typeProperty.ValueKind is not JsonValueKind.String ||
         return outputSchema;
     }
 
-    private JsonObject? CreateStructuredResponse(object? aiFunctionResult)
+    private JsonNode? CreateStructuredResponse(object? aiFunctionResult)
     {
         if (ProtocolTool.OutputSchema is null)
         {
+            // Only provide structured responses if the tool has an output schema defined.
             return null;
         }
 
@@ -445,15 +429,10 @@ typeProperty.ValueKind is not JsonValueKind.String ||
             };
         }
 
-        if (nodeResult is JsonObject jsonObject)
-        {
-            return jsonObject;
-        }
-
-        throw new InvalidOperationException("The result of the AIFunction does not match its declared output schema.");
+        return nodeResult;
     }
 
-    private static CallToolResponse ConvertAIContentEnumerableToCallToolResponse(IEnumerable<AIContent> contentItems, JsonObject? structuredContent)
+    private static CallToolResponse ConvertAIContentEnumerableToCallToolResponse(IEnumerable<AIContent> contentItems, JsonNode? structuredContent)
     {
         List<Content> contentList = [];
         bool allErrorContent = true;