Adding support for returning collections from tools (#116)

* Adding support for returning collections from tools

This will mean we can return multiple AIContent objects from a tool, such as a mixed text/image set

Contributes to #68

* Code review feedback

* Adding tests and some more cases

* More cases in switch and tests

src/ModelContextProtocol/AIContentExtensions.cs

@@ -1,7 +1,9 @@
 using Microsoft.Extensions.AI;
 using ModelContextProtocol.Protocol.Types;
 using ModelContextProtocol.Utils;
+using ModelContextProtocol.Utils.Json;
 using System.Runtime.InteropServices;
+using System.Text.Json;
 
 namespace ModelContextProtocol;
 
@@ -101,4 +103,28 @@ public static class AIContentExtensions
             Convert.ToBase64String(dataContent.Data.ToArray());
 #endif
     }
+
+    internal static Content ToContent(this AIContent content) =>
+        content switch
+        {
+            TextContent textContent => new()
+            {
+                Text = textContent.Text,
+                Type = "text",
+            },
+            DataContent dataContent => new()
+            {
+                Data = dataContent.GetBase64Data(),
+                MimeType = dataContent.MediaType,
+                Type =
+                    dataContent.HasTopLevelMediaType("image") ? "image" :
+                    dataContent.HasTopLevelMediaType("audio") ? "audio" :
+                    "resource",
+            },
+            _ => new()
+            {
+                Text = JsonSerializer.Serialize(content, McpJsonUtilities.DefaultOptions.GetTypeInfo(typeof(object))),
+                Type = "text",
+            }
+        };
 }

src/ModelContextProtocol/Server/AIFunctionMcpServerTool.cs

@@ -195,57 +195,49 @@ internal sealed class AIFunctionMcpServerTool : McpServerTool
             };
         }
 
-        switch (result)
+        return result switch
         {
-            case null:
+            AIContent aiContent => new()
-                return new()
+            {
-                {
+                Content = [aiContent.ToContent()]
-                    Content = []
+            },
-                };
+            null => new()
-
+            {
-            case string text:
+                Content = []
-                return new()
+            },
-                {
+            string text => new()
-                    Content = [new() { Text = text, Type = "text" }]
+            {
-                };
+                Content = [new() { Text = text, Type = "text" }]
-
+            },
-            case TextContent textContent:
+            Content content => new()
-                return new()
+            {
-                {
+                Content = [content]
-                    Content = [new() { Text = textContent.Text, Type = "text" }]
+            },
-                };
+            IEnumerable<string> texts => new()
-
+            {
-            case DataContent dataContent:
+                Content = [.. texts.Select(x => new Content() { Type = "text", Text = x ?? string.Empty })]
-                return new()
+            },
-                {
+            IEnumerable<AIContent> contentItems => new()
-                    Content = [new()
+            {
-                    {
+                Content = [.. contentItems.Select(static item => item.ToContent())]
-                        Data = dataContent.GetBase64Data(),
+            },
-                        MimeType = dataContent.MediaType,
+            IEnumerable<Content> contents => new()
-                        Type = dataContent.HasTopLevelMediaType("image") ? "image" : "resource",
+            {
-                    }]
+                Content = [.. contents]
-                };
+            },
-
+            CallToolResponse callToolResponse => callToolResponse,
-            case string[] texts:
-                return new()
-                {
-                    Content = texts
-                        .Select(x => new Content() { Type = "text", Text = x ?? string.Empty })
-                        .ToList()
-                };
 
             // TODO https://github.com/modelcontextprotocol/csharp-sdk/issues/69:
             // Add specialization for annotations.
-
+            _ => new()
-            default:
+            {
-                return new()
+                Content = [new()
-                {
-                    Content = [new()
                     {
                         Text = JsonSerializer.Serialize(result, McpJsonUtilities.DefaultOptions.GetTypeInfo(typeof(object))),
                         Type = "text"
                     }]
-                };
+            },
-        }
+        };
     }
+
 }

tests/ModelContextProtocol.Tests/Server/McpServerToolReturnTests.cs

@@ -0,0 +1,194 @@
+using Microsoft.Extensions.AI;
+using ModelContextProtocol.Protocol.Types;
+using ModelContextProtocol.Server;
+using Moq;
+
+namespace ModelContextProtocol.Tests.Server;
+public  class McpServerToolReturnTests
+{
+    [Fact]
+    public async Task CanReturnCollectionOfAIContent()
+    {
+        Mock<IMcpServer> mockServer = new();
+        McpServerTool tool = McpServerTool.Create((IMcpServer server) =>
+        {
+            Assert.Same(mockServer.Object, server);
+            return new List<AIContent>() {
+                new TextContent("text"),
+                new DataContent("data:image/png;base64,1234"),
+                new DataContent("data:audio/wav;base64,1234")
+            };
+        });
+
+        var result = await tool.InvokeAsync(
+            new RequestContext<CallToolRequestParams>(mockServer.Object, null),
+            TestContext.Current.CancellationToken);
+
+        Assert.Equal(3, result.Content.Count);
+
+        Assert.Equal("text", result.Content[0].Text);
+        Assert.Equal("text", result.Content[0].Type);
+
+        Assert.Equal("1234", result.Content[1].Data);
+        Assert.Equal("image/png", result.Content[1].MimeType);
+        Assert.Equal("image", result.Content[1].Type);
+
+        Assert.Equal("1234", result.Content[2].Data);
+        Assert.Equal("audio/wav", result.Content[2].MimeType);
+        Assert.Equal("audio", result.Content[2].Type);
+    }
+
+    [Theory]
+    [InlineData("text", "text")]
+    [InlineData("data:image/png;base64,1234", "image")]
+    [InlineData("data:audio/wav;base64,1234", "audio")]
+    public async Task CanReturnSingleAIContent(string data, string type)
+    {
+        Mock<IMcpServer> mockServer = new();
+        McpServerTool tool = McpServerTool.Create((IMcpServer server) =>
+        {
+            Assert.Same(mockServer.Object, server);
+            return type switch 
+            {
+                "text" => (AIContent)new TextContent(data),
+                "image" => new DataContent(data),
+                "audio" => new DataContent(data),
+                _ => throw new ArgumentException("Invalid type")
+            };
+        });
+
+        var result = await tool.InvokeAsync(
+            new RequestContext<CallToolRequestParams>(mockServer.Object, null),
+            TestContext.Current.CancellationToken);
+
+        Assert.Single(result.Content);
+        Assert.Equal(type, result.Content[0].Type);
+        
+        if (type != "text")
+        {
+            Assert.NotNull(result.Content[0].MimeType);
+            Assert.Equal(data.Split(',').Last(), result.Content[0].Data);
+        }
+        else
+        {
+            Assert.Null(result.Content[0].MimeType);
+            Assert.Equal(data, result.Content[0].Text);
+        }
+    }
+
+    [Fact]
+    public async Task CanReturnNullAIContent()
+    {
+        Mock<IMcpServer> mockServer = new();
+        McpServerTool tool = McpServerTool.Create((IMcpServer server) =>
+        {
+            Assert.Same(mockServer.Object, server);
+            return (string?)null;
+        });
+        var result = await tool.InvokeAsync(
+            new RequestContext<CallToolRequestParams>(mockServer.Object, null),
+            TestContext.Current.CancellationToken);
+        Assert.Empty(result.Content);
+    }
+
+    [Fact]
+    public async Task CanReturnString()
+    {
+        Mock<IMcpServer> mockServer = new();
+        McpServerTool tool = McpServerTool.Create((IMcpServer server) =>
+        {
+            Assert.Same(mockServer.Object, server);
+            return "42";
+        });
+        var result = await tool.InvokeAsync(
+            new RequestContext<CallToolRequestParams>(mockServer.Object, null),
+            TestContext.Current.CancellationToken);
+        Assert.Single(result.Content);
+        Assert.Equal("42", result.Content[0].Text);
+        Assert.Equal("text", result.Content[0].Type);
+    }
+
+    [Fact]
+    public async Task CanReturnCollectionOfStrings()
+    {
+        Mock<IMcpServer> mockServer = new();
+        McpServerTool tool = McpServerTool.Create((IMcpServer server) =>
+        {
+            Assert.Same(mockServer.Object, server);
+            return new List<string>() { "42", "43" };
+        });
+        var result = await tool.InvokeAsync(
+            new RequestContext<CallToolRequestParams>(mockServer.Object, null),
+            TestContext.Current.CancellationToken);
+        Assert.Equal(2, result.Content.Count);
+        Assert.Equal("42", result.Content[0].Text);
+        Assert.Equal("text", result.Content[0].Type);
+        Assert.Equal("43", result.Content[1].Text);
+        Assert.Equal("text", result.Content[1].Type);
+    }
+
+    [Fact]
+    public async Task CanReturnMcpContent()
+    {
+        Mock<IMcpServer> mockServer = new();
+        McpServerTool tool = McpServerTool.Create((IMcpServer server) =>
+        {
+            Assert.Same(mockServer.Object, server);
+            return new Content { Text = "42", Type = "text" };
+        });
+        var result = await tool.InvokeAsync(
+            new RequestContext<CallToolRequestParams>(mockServer.Object, null),
+            TestContext.Current.CancellationToken);
+        Assert.Single(result.Content);
+        Assert.Equal("42", result.Content[0].Text);
+        Assert.Equal("text", result.Content[0].Type);
+    }
+
+    [Fact]
+    public async Task CanReturnCollectionOfMcpContent()
+    {
+        Mock<IMcpServer> mockServer = new();
+        McpServerTool tool = McpServerTool.Create((IMcpServer server) =>
+        {
+            Assert.Same(mockServer.Object, server);
+            return new List<Content>() { new() { Text = "42", Type = "text" }, new() { Data = "1234", Type = "image", MimeType = "image/png" } };
+        });
+        var result = await tool.InvokeAsync(
+            new RequestContext<CallToolRequestParams>(mockServer.Object, null),
+            TestContext.Current.CancellationToken);
+        Assert.Equal(2, result.Content.Count);
+        Assert.Equal("42", result.Content[0].Text);
+        Assert.Equal("text", result.Content[0].Type);
+        Assert.Equal("1234", result.Content[1].Data);
+        Assert.Equal("image", result.Content[1].Type);
+        Assert.Equal("image/png", result.Content[1].MimeType);
+        Assert.Null(result.Content[1].Text);
+    }
+
+    [Fact]
+    public async Task CanReturnCallToolResponse()
+    {
+        CallToolResponse response = new()
+        {
+            Content = [new() { Text = "text", Type = "text" }, new() { Data = "1234", Type = "image" }]
+        };
+
+        Mock<IMcpServer> mockServer = new();
+        McpServerTool tool = McpServerTool.Create((IMcpServer server) =>
+        {
+            Assert.Same(mockServer.Object, server);
+            return response;
+        });
+        var result = await tool.InvokeAsync(
+            new RequestContext<CallToolRequestParams>(mockServer.Object, null),
+            TestContext.Current.CancellationToken);
+
+        Assert.Same(response, result);
+
+        Assert.Equal(2, result.Content.Count);
+        Assert.Equal("text", result.Content[0].Text);
+        Assert.Equal("text", result.Content[0].Type);
+        Assert.Equal("1234", result.Content[1].Data);
+        Assert.Equal("image", result.Content[1].Type);
+    }
+}

tests/ModelContextProtocol.Tests/Server/McpServerToolTests.cs

@@ -1,4 +1,5 @@
-using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.AI;
+using Microsoft.Extensions.DependencyInjection;
 using ModelContextProtocol.Protocol.Types;
 using ModelContextProtocol.Server;
 using Moq;