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>
This commit is contained in:
Stephen Toub 2025-03-20 13:16:41 -04:00 • committed by GitHub
commit 3de19d38e4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 175 additions and 128 deletions

View file

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

View file

@ -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 ??=
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);
public override JsonElement JsonSchema => tool.InputSchema;
/// <inheritdoc/>
protected async override Task<object?> InvokeCoreAsync(

View file

@ -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) =>

View file

@ -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; }
}

View file

@ -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;
}

View file

@ -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;
}

View file

@ -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);
}
}

View file

@ -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()
{
Type = "object",
Properties = new Dictionary<string, JsonSchemaProperty>()
InputSchema = JsonSerializer.Deserialize<JsonElement>("""
{
["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()
{
Type = "object",
Properties = new Dictionary<string, JsonSchemaProperty>()
InputSchema = JsonSerializer.Deserialize<JsonElement>("""
{
["prompt"] = new JsonSchemaProperty() { Type = "string", Description = "The prompt to send to the LLM" },
["maxTokens"] = new JsonSchemaProperty() { Type = "number", Description = "Maximum number of tokens to generate" }
"type": "object",
"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"]
}
},
"""),
}
]
});

View file

@ -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()
{
Type = "object",
Properties = new Dictionary<string, JsonSchemaProperty>()
InputSchema = JsonSerializer.Deserialize<JsonElement>("""
{
["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()
{
Type = "object",
Properties = new Dictionary<string, JsonSchemaProperty>()
InputSchema = JsonSerializer.Deserialize<JsonElement>("""
{
["prompt"] = new JsonSchemaProperty() { Type = "string", Description = "The prompt to send to the LLM" },
["maxTokens"] = new JsonSchemaProperty() { Type = "number", Description = "Maximum number of tokens to generate" }
"type": "object",
"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"]
}
},
"""),
}
]
});

View file

@ -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.Type);
Assert.NotNull(tool.InputSchema.Properties);
Assert.NotEmpty(tool.InputSchema.Properties);
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);
Assert.Equal("object", tool.InputSchema.GetProperty("type").GetString());
Assert.Equal(JsonValueKind.Object, tool.InputSchema.GetProperty("properties").GetProperty("message").ValueKind);
Assert.Equal("the echoes message", tool.InputSchema.GetProperty("properties").GetProperty("message").GetProperty("description").GetString());
Assert.Equal(1, tool.InputSchema.GetProperty("required").GetArrayLength());
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.Type);
Assert.NotNull(tool.InputSchema.Properties);
Assert.NotEmpty(tool.InputSchema.Properties);
Assert.Equal("object", tool.InputSchema.GetProperty("type").GetString());
Assert.Contains("number", tool.InputSchema.Properties);
Assert.Equal("integer", tool.InputSchema.Properties["number"].Type);
Assert.Contains("otherNumber", tool.InputSchema.Properties);
Assert.Equal("number", tool.InputSchema.Properties["otherNumber"].Type);
Assert.Contains("someCheck", tool.InputSchema.Properties);
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);
Assert.Contains("integer", tool.InputSchema.GetProperty("properties").GetProperty("number").GetProperty("type").GetString());
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("string", tool.InputSchema.GetProperty("properties").GetProperty("someDate").GetProperty("type").GetString());
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("object", tool.InputSchema.GetProperty("properties").GetProperty("complexObject").GetProperty("type").GetString());
}
[McpToolType]

View file

@ -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));
}
}

View file

@ -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);
}