Change type of Tool.InputSchema to be JsonElement (#4)

* Change type of Tool.InputSchema to be JsonElement

* Add schema validation in the InputSchema property setter.

* Simplify AIFunction tool adapter.

---------

Co-authored-by: Eirik Tsarpalis <eirik.tsarpalis@gmail.com>

samples/anthropic/tools/ToolsConsole/McpToolExtensions.cs

@@ -16,9 +16,7 @@ public static class McpToolExtensions
         List<Anthropic.SDK.Common.Tool> result = [];
         foreach (var tool in tools)
         {
-            var function = tool.InputSchema == null
+            var function = new Function(tool.Name, tool.Description, JsonSerializer.SerializeToNode(tool.InputSchema));
-                ? new Function(tool.Name, tool.Description)
-                : new Function(tool.Name, tool.Description, JsonSerializer.Serialize(tool.InputSchema));
             result.Add(function);
         }
         return result;

src/ModelContextProtocol/Client/McpClientExtensions.cs

@@ -466,8 +466,6 @@ public static class McpClientExtensions
     /// <summary>Provides an AI function that calls a tool through <see cref="IMcpClient"/>.</summary>
     private sealed class McpAIFunction(IMcpClient client, Tool tool) : AIFunction
     {
-        private JsonElement? _jsonSchema;
-
         /// <inheritdoc/>
         public override string Name => tool.Name;
 
@@ -475,15 +473,7 @@ public static class McpClientExtensions
         public override string Description => tool.Description ?? string.Empty;
 
         /// <inheritdoc/>
-        public override JsonElement JsonSchema => _jsonSchema ??=
+        public override JsonElement JsonSchema => tool.InputSchema;
-            JsonSerializer.SerializeToElement(new Dictionary<string, object>
-            {
-                ["type"] = "object",
-                ["title"] = tool.Name,
-                ["description"] = tool.Description ?? string.Empty,
-                ["properties"] = tool.InputSchema?.Properties ?? [],
-                ["required"] = tool.InputSchema?.Required ?? []
-            }, McpJsonUtilities.JsonContext.Default.DictionaryStringObject);
 
         /// <inheritdoc/>
         protected async override Task<object?> InvokeCoreAsync(

src/ModelContextProtocol/Configuration/McpServerBuilderExtensions.Tools.cs

@@ -96,7 +96,7 @@ public static partial class McpServerBuilderExtensions
             {
                 Name = function.Name,
                 Description = function.Description,
-                InputSchema = JsonSerializer.Deserialize(function.JsonSchema, McpJsonUtilities.JsonContext.Default.JsonSchema),
+                InputSchema = function.JsonSchema,
             });
 
             callbacks.Add(function.Name, async (request, cancellationToken) =>

src/ModelContextProtocol/Protocol/Types/JsonSchema.cs

@@ -1,26 +0,0 @@
-namespace ModelContextProtocol.Protocol.Types;
-
-/// <summary>
-/// Represents a JSON schema for a tool's input arguments.
-/// <see href="https://github.com/modelcontextprotocol/specification/blob/main/schema/2024-11-05/schema.json">See the schema for details</see>
-/// </summary>
-public class JsonSchema
-{
-    /// <summary>
-    /// The type of the schema, should be "object".
-    /// </summary>
-    [System.Text.Json.Serialization.JsonPropertyName("type")]
-    public string Type { get; set; } = "object";
-
-    /// <summary>
-    /// Map of property names to property definitions.
-    /// </summary>
-    [System.Text.Json.Serialization.JsonPropertyName("properties")]
-    public Dictionary<string, JsonSchemaProperty>? Properties { get; set; }
-
-    /// <summary>
-    /// List of required property names.
-    /// </summary>
-    [System.Text.Json.Serialization.JsonPropertyName("required")]
-    public List<string>? Required { get; set; }
-}

src/ModelContextProtocol/Protocol/Types/JsonSchemaProperty.cs

@@ -1,20 +0,0 @@
-namespace ModelContextProtocol.Protocol.Types;
-
-/// <summary>
-/// Represents a property in a JSON schema.
-/// <see href="https://github.com/modelcontextprotocol/specification/blob/main/schema/2024-11-05/schema.json">See the schema for details</see>
-/// </summary>
-public class JsonSchemaProperty
-{
-    /// <summary>
-    /// The type of the property. Should be a JSON Schema type and is required.
-    /// </summary>
-    [System.Text.Json.Serialization.JsonPropertyName("type")]
-    public string Type { get; set; } = string.Empty;
-
-    /// <summary>
-    /// A human-readable description of the property.
-    /// </summary>
-    [System.Text.Json.Serialization.JsonPropertyName("description")]
-    public string? Description { get; set; } = string.Empty;
-}

src/ModelContextProtocol/Protocol/Types/Tool.cs

@@ -1,4 +1,6 @@
-using System.Text.Json.Serialization;
+using ModelContextProtocol.Utils.Json;
+using System.Text.Json;
+using System.Text.Json.Serialization;
 
 namespace ModelContextProtocol.Protocol.Types;
 
@@ -23,6 +25,23 @@ public class Tool
     /// <summary>
     /// A JSON Schema object defining the expected parameters for the tool.
     /// </summary>
+    /// <remarks>
+    /// Needs to a valid JSON schema object that additionally is of type object.
+    /// </remarks>
     [JsonPropertyName("inputSchema")]
-    public JsonSchema? InputSchema { get; set; }
+    public JsonElement InputSchema 
+    { 
+        get => _inputSchema; 
+        set
+        {
+            if (!McpJsonUtilities.IsValidMcpToolSchema(value))
+            {
+                throw new ArgumentException("The specified document is not a valid MPC tool JSON schema.", nameof(InputSchema));
+            }
+
+            _inputSchema = value;
+        }
+    }
+
+    private JsonElement _inputSchema = McpJsonUtilities.DefaultMcpToolSchema;
 }

src/ModelContextProtocol/Utils/Json/McpJsonUtilities.cs

@@ -77,6 +77,31 @@ public static partial class McpJsonUtilities
     internal static JsonTypeInfo<T> GetTypeInfo<T>(this JsonSerializerOptions options) =>
         (JsonTypeInfo<T>)options.GetTypeInfo(typeof(T));
 
+    internal static JsonElement DefaultMcpToolSchema = ParseJsonElement("{\"type\":\"object\"}"u8);
+    internal static bool IsValidMcpToolSchema(JsonElement element)
+    {
+        if (element.ValueKind is not JsonValueKind.Object)
+        {
+            return false;
+        }
+
+        foreach (JsonProperty property in element.EnumerateObject())
+        {
+            if (property.NameEquals("type"))
+            {
+                if (property.Value.ValueKind is not JsonValueKind.String ||
+                    !property.Value.ValueEquals("object"))
+                {
+                    return false;
+                }
+
+                return true; // No need to check other properties
+            }
+        }
+
+        return false; // No type keyword found.
+    }
+
     // Keep in sync with CreateDefaultOptions above.
     [JsonSourceGenerationOptions(JsonSerializerDefaults.Web,
         UseStringEnumConverter = true,
@@ -96,7 +121,13 @@ public static partial class McpJsonUtilities
     [JsonSerializable(typeof(CreateMessageResult))]
     [JsonSerializable(typeof(ListRootsResult))]
     [JsonSerializable(typeof(InitializeResult))]
-    [JsonSerializable(typeof(JsonSchema))]
     [JsonSerializable(typeof(CallToolResponse))]
     internal sealed partial class JsonContext : JsonSerializerContext;
+
+    private static JsonElement ParseJsonElement(ReadOnlySpan<byte> utf8Json)
+    {
+        Utf8JsonReader reader = new(utf8Json);
+        return JsonElement.ParseValue(ref reader);
+    }
+
 }

tests/ModelContextProtocol.TestServer/Program.cs

@@ -1,10 +1,11 @@
-using System.Text;
+using Microsoft.Extensions.Logging;
 using ModelContextProtocol.Protocol.Messages;
 using ModelContextProtocol.Protocol.Transport;
 using ModelContextProtocol.Protocol.Types;
 using ModelContextProtocol.Server;
-using Microsoft.Extensions.Logging;
 using Serilog;
+using System.Text;
+using System.Text.Json;
 
 namespace ModelContextProtocol.TestServer;
 
@@ -91,28 +92,39 @@ internal static class Program
                         {
                             Name = "echo",
                             Description = "Echoes the input back to the client.",
-                            InputSchema = new JsonSchema()
+                            InputSchema = JsonSerializer.Deserialize<JsonElement>("""
-                            {
-                                Type = "object",
-                                Properties = new Dictionary<string, JsonSchemaProperty>()
                                 {
-                                    ["message"] = new JsonSchemaProperty() { Type = "string", Description = "The input to echo back." }
+                                    "type": "object",
+                                    "properties": {
+                                        "message": {
+                                            "type": "string",
+                                            "description": "The input to echo back."
+                                        }
+                                    },
+                                    "required": ["message"]
                                 }
-                            },
+                                """),
                         },
                         new Tool()
                         {
                             Name = "sampleLLM",
                             Description = "Samples from an LLM using MCP's sampling feature.",
-                            InputSchema = new JsonSchema()
+                            InputSchema = JsonSerializer.Deserialize<JsonElement>("""
-                            {
-                                Type = "object",
-                                Properties = new Dictionary<string, JsonSchemaProperty>()
                                 {
-                                    ["prompt"] = new JsonSchemaProperty() { Type = "string", Description = "The prompt to send to the LLM" },
+                                    "type": "object",
-                                    ["maxTokens"] = new JsonSchemaProperty() { Type = "number", Description = "Maximum number of tokens to generate" }
+                                    "properties": {
+                                        "prompt": {
+                                            "type": "string",
+                                            "description": "The prompt to send to the LLM"
+                                        },
+                                        "maxTokens": {
+                                            "type": "number",
+                                            "description": "Maximum number of tokens to generate"
+                                        }
+                                    },
+                                    "required": ["prompt", "maxTokens"]
                                 }
-                            },
+                                """),
                         }
                     ]
                 });

tests/ModelContextProtocol.TestSseServer/Program.cs

@@ -4,6 +4,7 @@ using ModelContextProtocol.Server;
 using Microsoft.Extensions.Logging;
 using Serilog;
 using System.Text;
+using System.Text.Json;
 
 internal class Program
 {
@@ -121,28 +122,39 @@ internal class Program
                             {
                                 Name = "echo",
                                 Description = "Echoes the input back to the client.",
-                                InputSchema = new JsonSchema()
+                                InputSchema = JsonSerializer.Deserialize<JsonElement>("""
-                                {
-                                    Type = "object",
-                                    Properties = new Dictionary<string, JsonSchemaProperty>()
                                     {
-                                        ["message"] = new JsonSchemaProperty() { Type = "string", Description = "The input to echo back." }
+                                        "type": "object",
+                                        "properties": {
+                                            "message": {
+                                                "type": "string",
+                                                "description": "The input to echo back."
+                                            }
+                                        },
+                                        "required": ["message"]
                                     }
-                                },
+                                    """),
                             },
                             new Tool()
                             {
                                 Name = "sampleLLM",
                                 Description = "Samples from an LLM using MCP's sampling feature.",
-                                InputSchema = new JsonSchema()
+                                InputSchema = JsonSerializer.Deserialize<JsonElement>("""
-                                {
-                                    Type = "object",
-                                    Properties = new Dictionary<string, JsonSchemaProperty>()
                                     {
-                                        ["prompt"] = new JsonSchemaProperty() { Type = "string", Description = "The prompt to send to the LLM" },
+                                        "type": "object",
-                                        ["maxTokens"] = new JsonSchemaProperty() { Type = "number", Description = "Maximum number of tokens to generate" }
+                                        "properties": {
+                                            "prompt": {
+                                                "type": "string",
+                                                "description": "The prompt to send to the LLM"
+                                            },
+                                            "maxTokens": {
+                                                "type": "number",
+                                                "description": "Maximum number of tokens to generate"
+                                            }
+                                        },
+                                        "required": ["prompt", "maxTokens"]
                                     }
-                                },
+                                    """),
                             }
                         ]
                     });

tests/ModelContextProtocol.Tests/Configuration/McpServerBuilderExtensionsToolsTests.cs

@@ -48,16 +48,10 @@ public class McpServerBuilderExtensionsToolsTests
         var tool = result.Tools[0];
         Assert.Equal("Echo", tool.Name);
         Assert.Equal("Echoes the input back to the client.", tool.Description);
-        Assert.NotNull(tool.InputSchema);
+        Assert.Equal("object", tool.InputSchema.GetProperty("type").GetString());
-        Assert.Equal("object", tool.InputSchema.Type);
+        Assert.Equal(JsonValueKind.Object, tool.InputSchema.GetProperty("properties").GetProperty("message").ValueKind);
-        Assert.NotNull(tool.InputSchema.Properties);
+        Assert.Equal("the echoes message", tool.InputSchema.GetProperty("properties").GetProperty("message").GetProperty("description").GetString());
-        Assert.NotEmpty(tool.InputSchema.Properties);
+        Assert.Equal(1, tool.InputSchema.GetProperty("required").GetArrayLength());
-        Assert.Contains("message", tool.InputSchema.Properties);
-        Assert.Equal("string", tool.InputSchema.Properties["message"].Type);
-        Assert.Equal("the echoes message", tool.InputSchema.Properties["message"].Description);
-        Assert.NotNull(tool.InputSchema.Required);
-        Assert.NotEmpty(tool.InputSchema.Required);
-        Assert.Contains("message", tool.InputSchema.Required);
 
         tool = result.Tools[1];
         Assert.Equal("double_echo", tool.Name);
@@ -288,31 +282,15 @@ public class McpServerBuilderExtensionsToolsTests
         var tool = result.Tools.First(t => t.Name == "TestTool");
         Assert.Equal("TestTool", tool.Name);
         Assert.Empty(tool.Description!);
-        Assert.NotNull(tool.InputSchema);
+        Assert.Equal("object", tool.InputSchema.GetProperty("type").GetString());
-        Assert.Equal("object", tool.InputSchema.Type);
-        Assert.NotNull(tool.InputSchema.Properties);
-        Assert.NotEmpty(tool.InputSchema.Properties);
 
-        Assert.Contains("number", tool.InputSchema.Properties);
+        Assert.Contains("integer", tool.InputSchema.GetProperty("properties").GetProperty("number").GetProperty("type").GetString());
-        Assert.Equal("integer", tool.InputSchema.Properties["number"].Type);
+        Assert.Contains("number", tool.InputSchema.GetProperty("properties").GetProperty("otherNumber").GetProperty("type").GetString());
-
+        Assert.Contains("boolean", tool.InputSchema.GetProperty("properties").GetProperty("someCheck").GetProperty("type").GetString());
-        Assert.Contains("otherNumber", tool.InputSchema.Properties);
+        Assert.Contains("string", tool.InputSchema.GetProperty("properties").GetProperty("someDate").GetProperty("type").GetString());
-        Assert.Equal("number", tool.InputSchema.Properties["otherNumber"].Type);
+        Assert.Contains("string", tool.InputSchema.GetProperty("properties").GetProperty("someOtherDate").GetProperty("type").GetString());
-
+        Assert.Contains("array", tool.InputSchema.GetProperty("properties").GetProperty("data").GetProperty("type").GetString());
-        Assert.Contains("someCheck", tool.InputSchema.Properties);
+        Assert.Contains("object", tool.InputSchema.GetProperty("properties").GetProperty("complexObject").GetProperty("type").GetString());
-        Assert.Equal("boolean", tool.InputSchema.Properties["someCheck"].Type);
-
-        Assert.Contains("someDate", tool.InputSchema.Properties);
-        Assert.Equal("string", tool.InputSchema.Properties["someDate"].Type);
-
-        Assert.Contains("someOtherDate", tool.InputSchema.Properties);
-        Assert.Equal("string", tool.InputSchema.Properties["someOtherDate"].Type);
-
-        Assert.Contains("data", tool.InputSchema.Properties);
-        Assert.Equal("array", tool.InputSchema.Properties["data"].Type);
-
-        Assert.Contains("complexObject", tool.InputSchema.Properties);
-        Assert.Equal("object", tool.InputSchema.Properties["complexObject"].Type);
     }
 
     [McpToolType]

tests/ModelContextProtocol.Tests/Protocol/ProtocolTypeTests.cs

@@ -0,0 +1,52 @@
+using ModelContextProtocol.Protocol.Types;
+using System.Text.Json;
+
+namespace ModelContextProtocol.Tests.Protocol;
+
+public static class ProtocolTypeTests
+{
+    [Fact]
+    public static void ToolInputSchema_HasValidDefaultSchema()
+    {
+        var tool = new Tool();
+        JsonElement jsonElement = tool.InputSchema;
+
+        Assert.Equal(JsonValueKind.Object, jsonElement.ValueKind);
+        Assert.Single(jsonElement.EnumerateObject());
+        Assert.True(jsonElement.TryGetProperty("type", out JsonElement typeElement));
+        Assert.Equal(JsonValueKind.String, typeElement.ValueKind);
+        Assert.Equal("object", typeElement.GetString());
+    }
+
+    [Theory]
+    [InlineData("null")]
+    [InlineData("false")]
+    [InlineData("true")]
+    [InlineData("3.5e3")]
+    [InlineData("[]")]
+    [InlineData("{}")]
+    [InlineData("""{"properties":{}}""")]
+    [InlineData("""{"type":"number"}""")]
+    [InlineData("""{"type":"array"}""")]
+    [InlineData("""{"type":["object"]}""")]
+    public static void ToolInputSchema_RejectsInvalidSchemaDocuments(string invalidSchema)
+    {
+        using var document = JsonDocument.Parse(invalidSchema);
+        var tool = new Tool();
+
+        Assert.Throws<ArgumentException>(() => tool.InputSchema = document.RootElement);
+    }
+
+    [Theory]
+    [InlineData("""{"type":"object"}""")]
+    [InlineData("""{"type":"object", "properties": {}, "required" : [] }""")]
+    [InlineData("""{"type":"object", "title": "MyAwesomeTool", "description": "It's awesome!", "properties": {}, "required" : ["NotAParam"] }""")]
+    public static void ToolInputSchema_AcceptsValidSchemaDocuments(string validSchema)
+    {
+        using var document = JsonDocument.Parse(validSchema);
+        var tool = new Tool();
+
+        tool.InputSchema = document.RootElement;
+        Assert.True(JsonElement.DeepEquals(document.RootElement, tool.InputSchema));
+    }
+}

tests/ModelContextProtocol.Tests/Utils/TestServerTransport.cs

@@ -1,6 +1,7 @@
 using System.Threading.Channels;
 using ModelContextProtocol.Protocol.Messages;
 using ModelContextProtocol.Protocol.Transport;
+using ModelContextProtocol.Protocol.Types;
 
 namespace ModelContextProtocol.Tests.Utils;
 
@@ -71,7 +72,7 @@ public class TestServerTransport : IServerTransport
         await WriteMessageAsync(new JsonRpcResponse
         {
             Id = request.Id,
-            Result = new Protocol.Types.CreateMessageResult { Content = new(), Model = "model", Role = "role" }
+            Result = new CreateMessageResult { Content = new(), Model = "model", Role = "role" }
         }, cancellationToken);
     }