Fix $ref pointer resolution after output schema wrapping (#1435)
This commit is contained in:
Weinong Wang
GitHub
parent
480979b1b5
commit
20ddd2fb10
2 changed files with 255 additions and 0 deletions
|
|
@ -543,6 +543,11 @@ internal sealed partial class AIFunctionMcpServerTool : McpServerTool
|
|||
["required"] = new JsonArray { (JsonNode)"result" }
|
||||
};
|
||||
|
||||
// After wrapping, any internal $ref pointers that used absolute JSON Pointer
|
||||
// paths (e.g., "#/items/..." or "#") are now invalid because the original schema
|
||||
// has moved under "#/properties/result". Rewrite them to account for the new location.
|
||||
RewriteRefPointers(schemaNode["properties"]!["result"]);
|
||||
|
||||
structuredOutputRequiresWrapping = true;
|
||||
}
|
||||
|
||||
|
|
@ -552,6 +557,52 @@ internal sealed partial class AIFunctionMcpServerTool : McpServerTool
|
|||
return outputSchema;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Recursively rewrites all <c>$ref</c> JSON Pointer values in the given node
|
||||
/// to account for the schema having been wrapped under <c>properties.result</c>.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// <c>System.Text.Json</c>'s <see cref="System.Text.Json.Schema.JsonSchemaExporter"/> uses absolute
|
||||
/// JSON Pointer paths (e.g., <c>#/items/properties/foo</c>) to deduplicate types that appear at
|
||||
/// multiple locations in the schema. When the original schema is moved under
|
||||
/// <c>#/properties/result</c> by the wrapping logic above, these pointers become unresolvable.
|
||||
/// This method prepends <c>/properties/result</c> to every <c>$ref</c> that starts with <c>#/</c>,
|
||||
/// and rewrites bare <c>#</c> (root self-references from recursive types) to <c>#/properties/result</c>,
|
||||
/// so the pointers remain valid after wrapping.
|
||||
/// </remarks>
|
||||
private static void RewriteRefPointers(JsonNode? node)
|
||||
{
|
||||
if (node is JsonObject obj)
|
||||
{
|
||||
if (obj.TryGetPropertyValue("$ref", out JsonNode? refNode) &&
|
||||
refNode?.GetValue<string>() is string refValue)
|
||||
{
|
||||
if (refValue == "#")
|
||||
{
|
||||
obj["$ref"] = "#/properties/result";
|
||||
}
|
||||
else if (refValue.StartsWith("#/", StringComparison.Ordinal))
|
||||
{
|
||||
obj["$ref"] = "#/properties/result" + refValue.Substring(1);
|
||||
}
|
||||
}
|
||||
|
||||
// Safe to iterate without snapshot: the $ref assignment above completes before
|
||||
// this enumerator is created, and recursive calls only mutate descendant objects.
|
||||
foreach (var property in obj)
|
||||
{
|
||||
RewriteRefPointers(property.Value);
|
||||
}
|
||||
}
|
||||
else if (node is JsonArray arr)
|
||||
{
|
||||
foreach (var item in arr)
|
||||
{
|
||||
RewriteRefPointers(item);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private JsonElement? CreateStructuredResponse(object? aiFunctionResult)
|
||||
{
|
||||
if (ProtocolTool.OutputSchema is null)
|
||||
|
|
|
|||
|
|
@ -794,6 +794,186 @@ public partial class McpServerToolTests
|
|||
Assert.True(JsonElement.DeepEquals(expectedSchema, tool.ProtocolTool.InputSchema));
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task StructuredOutput_WithDuplicateTypeRefs_RewritesRefPointers()
|
||||
{
|
||||
// When a non-object return type contains the same type at multiple locations,
|
||||
// System.Text.Json's schema exporter emits $ref pointers for deduplication.
|
||||
// After wrapping the schema under properties.result, those $ref pointers must
|
||||
// be rewritten to remain valid. This test verifies that fix.
|
||||
var data = new List<ContactInfo>
|
||||
{
|
||||
new()
|
||||
{
|
||||
WorkPhones = [new() { Number = "555-0100", Type = "work" }],
|
||||
HomePhones = [new() { Number = "555-0200", Type = "home" }],
|
||||
}
|
||||
};
|
||||
|
||||
JsonSerializerOptions options = new() { TypeInfoResolver = new DefaultJsonTypeInfoResolver() };
|
||||
McpServerTool tool = McpServerTool.Create(() => data, new() { Name = "tool", UseStructuredContent = true, SerializerOptions = options });
|
||||
var mockServer = new Mock<McpServer>();
|
||||
var result = await tool.InvokeAsync(
|
||||
new RequestContext<CallToolRequestParams>(mockServer.Object, CreateTestJsonRpcRequest(), new() { Name = "tool" }),
|
||||
TestContext.Current.CancellationToken);
|
||||
|
||||
Assert.NotNull(tool.ProtocolTool.OutputSchema);
|
||||
Assert.Equal("object", tool.ProtocolTool.OutputSchema.Value.GetProperty("type").GetString());
|
||||
Assert.NotNull(result.StructuredContent);
|
||||
|
||||
// Verify $ref pointers in the schema point to valid locations after wrapping.
|
||||
// Without the fix, $ref values like "#/items/..." would be unresolvable because
|
||||
// the original schema was moved under "#/properties/result".
|
||||
AssertMatchesJsonSchema(tool.ProtocolTool.OutputSchema.Value, result.StructuredContent);
|
||||
|
||||
// Also verify that any $ref in the schema starts with #/properties/result
|
||||
// (confirming the rewrite happened).
|
||||
string schemaJson = tool.ProtocolTool.OutputSchema.Value.GetRawText();
|
||||
var schemaNode = JsonNode.Parse(schemaJson)!;
|
||||
int refCount = AssertAllRefsStartWith(schemaNode, "#/properties/result");
|
||||
Assert.True(refCount > 0, "Expected at least one $ref in the schema to validate the rewrite, but none were found.");
|
||||
int resolvableCount = AssertAllRefsResolvable(schemaNode, schemaNode);
|
||||
Assert.True(resolvableCount > 0, "Expected at least one resolvable $ref in the schema, but none were found.");
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task StructuredOutput_WithRecursiveTypeRefs_RewritesRefPointers()
|
||||
{
|
||||
// When a non-object return type contains a recursive type, System.Text.Json's
|
||||
// schema exporter emits $ref pointers (including potentially bare "#") for the
|
||||
// recursive reference. After wrapping, these must be rewritten. For List<TreeNode>,
|
||||
// Children's items emit "$ref": "#/items" which must become "#/properties/result/items".
|
||||
var data = new List<TreeNode>
|
||||
{
|
||||
new()
|
||||
{
|
||||
Name = "root",
|
||||
Children = [new() { Name = "child" }],
|
||||
}
|
||||
};
|
||||
|
||||
JsonSerializerOptions options = new() { TypeInfoResolver = new DefaultJsonTypeInfoResolver() };
|
||||
McpServerTool tool = McpServerTool.Create(() => data, new() { Name = "tool", UseStructuredContent = true, SerializerOptions = options });
|
||||
var mockServer = new Mock<McpServer>();
|
||||
var result = await tool.InvokeAsync(
|
||||
new RequestContext<CallToolRequestParams>(mockServer.Object, CreateTestJsonRpcRequest(), new() { Name = "tool" }),
|
||||
TestContext.Current.CancellationToken);
|
||||
|
||||
Assert.NotNull(tool.ProtocolTool.OutputSchema);
|
||||
Assert.Equal("object", tool.ProtocolTool.OutputSchema.Value.GetProperty("type").GetString());
|
||||
Assert.NotNull(result.StructuredContent);
|
||||
|
||||
AssertMatchesJsonSchema(tool.ProtocolTool.OutputSchema.Value, result.StructuredContent);
|
||||
|
||||
string schemaJson = tool.ProtocolTool.OutputSchema.Value.GetRawText();
|
||||
var schemaNode = JsonNode.Parse(schemaJson)!;
|
||||
int refCount = AssertAllRefsStartWith(schemaNode, "#/properties/result");
|
||||
Assert.True(refCount > 0, "Expected at least one $ref in the schema to validate the rewrite, but none were found.");
|
||||
int resolvableCount = AssertAllRefsResolvable(schemaNode, schemaNode);
|
||||
Assert.True(resolvableCount > 0, "Expected at least one resolvable $ref in the schema, but none were found.");
|
||||
}
|
||||
|
||||
private static int AssertAllRefsStartWith(JsonNode? node, string expectedPrefix)
|
||||
{
|
||||
int count = 0;
|
||||
if (node is JsonObject obj)
|
||||
{
|
||||
if (obj.TryGetPropertyValue("$ref", out JsonNode? refNode) &&
|
||||
refNode?.GetValue<string>() is string refValue)
|
||||
{
|
||||
Assert.StartsWith(expectedPrefix, refValue);
|
||||
count++;
|
||||
}
|
||||
|
||||
foreach (var property in obj)
|
||||
{
|
||||
count += AssertAllRefsStartWith(property.Value, expectedPrefix);
|
||||
}
|
||||
}
|
||||
else if (node is JsonArray arr)
|
||||
{
|
||||
foreach (var item in arr)
|
||||
{
|
||||
count += AssertAllRefsStartWith(item, expectedPrefix);
|
||||
}
|
||||
}
|
||||
|
||||
return count;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Walks the JSON tree and verifies that every <c>$ref</c> pointer resolves to a valid node.
|
||||
/// </summary>
|
||||
private static int AssertAllRefsResolvable(JsonNode root, JsonNode? node)
|
||||
{
|
||||
int count = 0;
|
||||
if (node is JsonObject obj)
|
||||
{
|
||||
if (obj.TryGetPropertyValue("$ref", out JsonNode? refNode) &&
|
||||
refNode?.GetValue<string>() is string refValue &&
|
||||
refValue.StartsWith("#", StringComparison.Ordinal))
|
||||
{
|
||||
var resolved = ResolveJsonPointer(root, refValue);
|
||||
Assert.True(resolved is not null, $"$ref \"{refValue}\" does not resolve to a valid node in the schema.");
|
||||
count++;
|
||||
}
|
||||
|
||||
foreach (var property in obj)
|
||||
{
|
||||
count += AssertAllRefsResolvable(root, property.Value);
|
||||
}
|
||||
}
|
||||
else if (node is JsonArray arr)
|
||||
{
|
||||
foreach (var item in arr)
|
||||
{
|
||||
count += AssertAllRefsResolvable(root, item);
|
||||
}
|
||||
}
|
||||
|
||||
return count;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Resolves a JSON Pointer (e.g., <c>#/properties/result/items</c>) against a root node.
|
||||
/// Returns <c>null</c> if the pointer cannot be resolved.
|
||||
/// </summary>
|
||||
private static JsonNode? ResolveJsonPointer(JsonNode root, string pointer)
|
||||
{
|
||||
if (pointer == "#")
|
||||
{
|
||||
return root;
|
||||
}
|
||||
|
||||
if (!pointer.StartsWith("#/", StringComparison.Ordinal))
|
||||
{
|
||||
return null;
|
||||
}
|
||||
|
||||
JsonNode? current = root;
|
||||
string[] segments = pointer.Substring(2).Split('/');
|
||||
foreach (string segment in segments)
|
||||
{
|
||||
if (current is JsonObject obj)
|
||||
{
|
||||
if (!obj.TryGetPropertyValue(segment, out current))
|
||||
{
|
||||
return null;
|
||||
}
|
||||
}
|
||||
else if (current is JsonArray arr && int.TryParse(segment, out int index) && index >= 0 && index < arr.Count)
|
||||
{
|
||||
current = arr[index];
|
||||
}
|
||||
else
|
||||
{
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
return current;
|
||||
}
|
||||
|
||||
public static IEnumerable<object[]> StructuredOutput_ReturnsExpectedSchema_Inputs()
|
||||
{
|
||||
yield return new object[] { "string" };
|
||||
|
|
@ -926,6 +1106,30 @@ public partial class McpServerToolTests
|
|||
return options;
|
||||
}
|
||||
|
||||
// Types used by StructuredOutput_WithDuplicateTypeRefs_RewritesRefPointers.
|
||||
// ContactInfo has two properties of the same type (PhoneNumber) which causes
|
||||
// System.Text.Json's schema exporter to emit $ref pointers for deduplication.
|
||||
private sealed class PhoneNumber
|
||||
{
|
||||
public string? Number { get; set; }
|
||||
public string? Type { get; set; }
|
||||
}
|
||||
|
||||
private sealed class ContactInfo
|
||||
{
|
||||
public List<PhoneNumber>? WorkPhones { get; set; }
|
||||
public List<PhoneNumber>? HomePhones { get; set; }
|
||||
}
|
||||
|
||||
// Recursive type used by StructuredOutput_WithRecursiveTypeRefs_RewritesRefPointers.
|
||||
// When List<TreeNode> is the return type, Children's items emit "$ref": "#/items"
|
||||
// pointing back to the first TreeNode definition, which must be rewritten after wrapping.
|
||||
private sealed class TreeNode
|
||||
{
|
||||
public string? Name { get; set; }
|
||||
public List<TreeNode>? Children { get; set; }
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public void SupportsIconsInCreateOptions()
|
||||
{
|
||||
|
|
|
|||
Loading…
Reference in a new issue