77using System . Diagnostics . CodeAnalysis ;
88using System . Reflection ;
99using System . Text . Json ;
10+ using System . Text . Json . Nodes ;
1011
1112namespace ModelContextProtocol . Server ;
1213
1314/// <summary>Provides an <see cref="McpServerTool"/> that's implemented via an <see cref="AIFunction"/>.</summary>
1415internal sealed partial class AIFunctionMcpServerTool : McpServerTool
1516{
1617 private readonly ILogger _logger ;
18+ private readonly bool _structuredOutputRequiresWrapping ;
1719
1820 /// <summary>
1921 /// Creates an <see cref="McpServerTool"/> instance for a method, specified via a <see cref="Delegate"/> instance.
@@ -176,7 +178,8 @@ private static AIFunctionFactoryOptions CreateAIFunctionFactoryOptions(
176178 {
177179 Name = options ? . Name ?? function . Name ,
178180 Description = options ? . Description ?? function . Description ,
179- InputSchema = function . JsonSchema ,
181+ InputSchema = function . JsonSchema ,
182+ OutputSchema = CreateOutputSchema ( function , options , out bool structuredOutputRequiresWrapping ) ,
180183 } ;
181184
182185 if ( options is not null )
@@ -198,7 +201,7 @@ options.OpenWorld is not null ||
198201 }
199202 }
200203
201- return new AIFunctionMcpServerTool ( function , tool , options ? . Services ) ;
204+ return new AIFunctionMcpServerTool ( function , tool , options ? . Services , structuredOutputRequiresWrapping ) ;
202205 }
203206
204207 private static McpServerToolCreateOptions DeriveOptions ( MethodInfo method , McpServerToolCreateOptions ? options )
@@ -243,11 +246,12 @@ private static McpServerToolCreateOptions DeriveOptions(MethodInfo method, McpSe
243246 internal AIFunction AIFunction { get ; }
244247
245248 /// <summary>Initializes a new instance of the <see cref="McpServerTool"/> class.</summary>
246- private AIFunctionMcpServerTool ( AIFunction function , Tool tool , IServiceProvider ? serviceProvider )
249+ private AIFunctionMcpServerTool ( AIFunction function , Tool tool , IServiceProvider ? serviceProvider , bool structuredOutputRequiresWrapping )
247250 {
248251 AIFunction = function ;
249252 ProtocolTool = tool ;
250253 _logger = serviceProvider ? . GetService < ILoggerFactory > ( ) ? . CreateLogger < McpServerTool > ( ) ?? ( ILogger ) NullLogger . Instance ;
254+ _structuredOutputRequiresWrapping = structuredOutputRequiresWrapping ;
251255 }
252256
253257 /// <inheritdoc />
@@ -295,39 +299,46 @@ public override async ValueTask<CallToolResponse> InvokeAsync(
295299 } ;
296300 }
297301
302+ JsonObject ? structuredContent = CreateStructuredResponse ( result ) ;
298303 return result switch
299304 {
300305 AIContent aiContent => new ( )
301306 {
302307 Content = [ aiContent . ToContent ( ) ] ,
308+ StructuredContent = structuredContent ,
303309 IsError = aiContent is ErrorContent
304310 } ,
305311
306312 null => new ( )
307313 {
308- Content = [ ]
314+ Content = [ ] ,
315+ StructuredContent = structuredContent ,
309316 } ,
310317
311318 string text => new ( )
312319 {
313- Content = [ new ( ) { Text = text , Type = "text" } ]
320+ Content = [ new ( ) { Text = text , Type = "text" } ] ,
321+ StructuredContent = structuredContent ,
314322 } ,
315323
316324 Content content => new ( )
317325 {
318- Content = [ content ]
326+ Content = [ content ] ,
327+ StructuredContent = structuredContent ,
319328 } ,
320329
321330 IEnumerable < string > texts => new ( )
322331 {
323- Content = [ .. texts . Select ( x => new Content ( ) { Type = "text" , Text = x ?? string . Empty } ) ]
332+ Content = [ .. texts . Select ( x => new Content ( ) { Type = "text" , Text = x ?? string . Empty } ) ] ,
333+ StructuredContent = structuredContent ,
324334 } ,
325335
326- IEnumerable < AIContent > contentItems => ConvertAIContentEnumerableToCallToolResponse ( contentItems ) ,
336+ IEnumerable < AIContent > contentItems => ConvertAIContentEnumerableToCallToolResponse ( contentItems , structuredContent ) ,
327337
328338 IEnumerable < Content > contents => new ( )
329339 {
330- Content = [ .. contents ]
340+ Content = [ .. contents ] ,
341+ StructuredContent = structuredContent ,
331342 } ,
332343
333344 CallToolResponse callToolResponse => callToolResponse ,
@@ -338,12 +349,111 @@ public override async ValueTask<CallToolResponse> InvokeAsync(
338349 {
339350 Text = JsonSerializer . Serialize ( result , AIFunction . JsonSerializerOptions . GetTypeInfo ( typeof ( object ) ) ) ,
340351 Type = "text"
341- } ]
352+ } ] ,
353+ StructuredContent = structuredContent ,
342354 } ,
343355 } ;
344356 }
345357
346- private static CallToolResponse ConvertAIContentEnumerableToCallToolResponse ( IEnumerable < AIContent > contentItems )
358+ private static JsonElement ? CreateOutputSchema ( AIFunction function , McpServerToolCreateOptions ? toolCreateOptions , out bool structuredOutputRequiresWrapping )
359+ {
360+ // TODO replace with https://github.com/dotnet/extensions/pull/6447 once merged.
361+
362+ structuredOutputRequiresWrapping = false ;
363+
364+ if ( toolCreateOptions ? . UseStructuredContent is not true )
365+ {
366+ return null ;
367+ }
368+
369+ if ( function . UnderlyingMethod ? . ReturnType is not Type returnType )
370+ {
371+ return null ;
372+ }
373+
374+ if ( returnType == typeof ( void ) || returnType == typeof ( Task ) || returnType == typeof ( ValueTask ) )
375+ {
376+ // Do not report an output schema for void or Task methods.
377+ return null ;
378+ }
379+
380+ if ( returnType . IsGenericType && returnType . GetGenericTypeDefinition ( ) is Type genericTypeDef &&
381+ ( genericTypeDef == typeof ( Task < > ) || genericTypeDef == typeof ( ValueTask < > ) ) )
382+ {
383+ // Extract the real type from Task<T> or ValueTask<T> if applicable.
384+ returnType = returnType . GetGenericArguments ( ) [ 0 ] ;
385+ }
386+
387+ JsonElement outputSchema = AIJsonUtilities . CreateJsonSchema ( returnType , serializerOptions : function . JsonSerializerOptions , inferenceOptions : toolCreateOptions ? . SchemaCreateOptions ) ;
388+
389+ if ( outputSchema . ValueKind is not JsonValueKind . Object ||
390+ ! outputSchema . TryGetProperty ( "type" , out JsonElement typeProperty ) ||
391+ typeProperty . ValueKind is not JsonValueKind . String ||
392+ typeProperty . GetString ( ) is not "object" )
393+ {
394+ // If the output schema is not an object, need to modify to be a valid MCP output schema.
395+ JsonNode ? schemaNode = JsonSerializer . SerializeToNode ( outputSchema , McpJsonUtilities . JsonContext . Default . JsonElement ) ;
396+
397+ if ( schemaNode is JsonObject objSchema &&
398+ objSchema . TryGetPropertyValue ( "type" , out JsonNode ? typeNode ) &&
399+ typeNode is JsonArray { Count : 2 } typeArray && typeArray . Any ( type => ( string ? ) type is "object" ) && typeArray . Any ( type => ( string ? ) type is "null" ) )
400+ {
401+ // For schemas that are of type ["object", "null"], replace with just "object" to be conformant.
402+ objSchema [ "type" ] = "object" ;
403+ }
404+ else
405+ {
406+ // For anything else, wrap the schema in an envelope with a "result" property.
407+ schemaNode = new JsonObject
408+ {
409+ [ "type" ] = "object" ,
410+ [ "properties" ] = new JsonObject
411+ {
412+ [ "result" ] = schemaNode
413+ } ,
414+ [ "required" ] = new JsonArray { ( JsonNode ) "result" }
415+ } ;
416+
417+ structuredOutputRequiresWrapping = true ;
418+ }
419+
420+ outputSchema = JsonSerializer . Deserialize ( schemaNode , McpJsonUtilities . JsonContext . Default . JsonElement ) ;
421+ }
422+
423+ return outputSchema ;
424+ }
425+
426+ private JsonObject ? CreateStructuredResponse ( object ? aiFunctionResult )
427+ {
428+ if ( ProtocolTool . OutputSchema is null )
429+ {
430+ return null ;
431+ }
432+
433+ JsonNode ? nodeResult = aiFunctionResult switch
434+ {
435+ JsonNode node => node ,
436+ JsonElement jsonElement => JsonSerializer . SerializeToNode ( jsonElement , McpJsonUtilities . JsonContext . Default . JsonElement ) ,
437+ _ => JsonSerializer . SerializeToNode ( aiFunctionResult , AIFunction . JsonSerializerOptions . GetTypeInfo ( typeof ( object ) ) ) ,
438+ } ;
439+
440+ if ( _structuredOutputRequiresWrapping )
441+ {
442+ return new JsonObject
443+ {
444+ [ "result" ] = nodeResult
445+ } ;
446+ }
447+
448+ if ( nodeResult is JsonObject jsonObject )
449+ {
450+ return jsonObject ;
451+ }
452+
453+ throw new InvalidOperationException ( "The result of the AIFunction does not match its declared output schema." ) ;
454+ }
455+
456+ private static CallToolResponse ConvertAIContentEnumerableToCallToolResponse ( IEnumerable < AIContent > contentItems , JsonObject ? structuredContent )
347457 {
348458 List < Content > contentList = [ ] ;
349459 bool allErrorContent = true ;
@@ -363,6 +473,7 @@ private static CallToolResponse ConvertAIContentEnumerableToCallToolResponse(IEn
363473 return new ( )
364474 {
365475 Content = contentList ,
476+ StructuredContent = structuredContent ,
366477 IsError = allErrorContent && hasAny
367478 } ;
368479 }
0 commit comments