Auto-populate completion handlers from AllowedValuesAttribute on prompt/resource parameters

src/ModelContextProtocol.Core/Server/McpServerImpl.cs

@@ -254,12 +254,61 @@ private void ConfigureCompletion(McpServerOptions options)
         var completeHandler = options.Handlers.CompleteHandler;
         var completionsCapability = options.Capabilities?.Completions;
 
-        if (completeHandler is null && completionsCapability is null)
+        // Build completion value lookups from prompt/resource collections' [AllowedValues]-attributed parameters.
+        Dictionary<string, Dictionary<string, string[]>>? promptCompletions = BuildAllowedValueCompletions(options.PromptCollection);
+        Dictionary<string, Dictionary<string, string[]>>? resourceCompletions = BuildAllowedValueCompletions(options.ResourceCollection);
+        bool hasCollectionCompletions = promptCompletions is not null || resourceCompletions is not null;
+
+        if (completeHandler is null && completionsCapability is null && !hasCollectionCompletions)
         {
             return;
         }
 
         completeHandler ??= (static async (_, __) => new CompleteResult());
+
+        // Augment the completion handler with allowed values from prompt/resource collections.
+        if (hasCollectionCompletions)
+        {
+            var originalCompleteHandler = completeHandler;
+            completeHandler = async (request, cancellationToken) =>
+            {
+                CompleteResult result = await originalCompleteHandler(request, cancellationToken).ConfigureAwait(false);
+
+                string[]? allowedValues = null;
+                switch (request.Params?.Ref)
+                {
+                    case PromptReference pr when promptCompletions is not null:
+                        if (promptCompletions.TryGetValue(pr.Name, out var promptParams))
+                        {
+                            promptParams.TryGetValue(request.Params.Argument.Name, out allowedValues);
+                        }
+                        break;
+
+                    case ResourceTemplateReference rtr when resourceCompletions is not null:
+                        if (rtr.Uri is not null && resourceCompletions.TryGetValue(rtr.Uri, out var resourceParams))
+                        {
+                            resourceParams.TryGetValue(request.Params.Argument.Name, out allowedValues);
+                        }
+                        break;
+                }
+
+                if (allowedValues is not null)
+                {
+                    string partialValue = request.Params!.Argument.Value;
+                    var filtered = Array.FindAll(allowedValues, v => v.StartsWith(partialValue, StringComparison.OrdinalIgnoreCase));
+
+                    foreach (var v in filtered)
+                    {
+                        result.Completion.Values.Add(v);
+                    }
+
+                    result.Completion.Total = result.Completion.Values.Count;
+                }
+
+                return result;
+            };
+        }
+
         completeHandler = BuildFilterPipeline(completeHandler, options.Filters.Request.CompleteFilters);
 
         ServerCapabilities.Completions = new();
@@ -271,6 +320,75 @@ private void ConfigureCompletion(McpServerOptions options)
             McpJsonUtilities.JsonContext.Default.CompleteResult);
     }
 
+    /// <summary>
+    /// Builds a lookup of primitive name/URI → (parameter name → allowed values) from the enum values
+    /// in the JSON schemas of AIFunction-based prompts or resources.
+    /// </summary>
+    private static Dictionary<string, Dictionary<string, string[]>>? BuildAllowedValueCompletions<T>(
+        McpServerPrimitiveCollection<T>? primitives) where T : class, IMcpServerPrimitive
+    {
+        if (primitives is null)
+        {
+            return null;
+        }
+
+        Dictionary<string, Dictionary<string, string[]>>? result = null;
+        foreach (var primitive in primitives)
+        {
+            JsonElement schema;
+            string id;
+            if (primitive is AIFunctionMcpServerPrompt aiPrompt)
+            {
+                schema = aiPrompt.AIFunction.JsonSchema;
+                id = aiPrompt.ProtocolPrompt.Name;
+            }
+            else if (primitive is AIFunctionMcpServerResource aiResource && aiResource.IsTemplated)
+            {
+                schema = aiResource.AIFunction.JsonSchema;
+                id = aiResource.ProtocolResourceTemplate.UriTemplate;
+            }
+            else
+            {
+                continue;
+            }
+
+            if (schema.TryGetProperty("properties", out JsonElement properties))
+            {
+                Dictionary<string, string[]>? paramValues = null;
+                foreach (var param in properties.EnumerateObject())
+                {
+                    if (param.Value.TryGetProperty("enum", out JsonElement enumValues) &&
+                        enumValues.ValueKind == JsonValueKind.Array)
+                    {
+                        List<string>? values = null;
+                        foreach (var item in enumValues.EnumerateArray())
+                        {
+                            if (item.GetString() is { } str)
+                            {
+                                values ??= [];
+                                values.Add(str);
+                            }
+                        }
+
+                        if (values is not null)
+                        {
+                            paramValues ??= new(StringComparer.Ordinal);
+                            paramValues[param.Name] = [.. values];
+                        }
+                    }
+                }
+
+                if (paramValues is not null)
+                {
+                    result ??= new(StringComparer.Ordinal);
+                    result[id] = paramValues;
+                }
+            }
+        }
+
+        return result;
+    }
+
     private void ConfigureExperimental(McpServerOptions options)
     {
         ServerCapabilities.Experimental = options.Capabilities?.Experimental;

tests/ModelContextProtocol.Tests/Server/McpServerTests.cs

@@ -322,6 +322,242 @@
             });
     }
 
+    [Fact]
+    public async Task Completion_AutoPopulated_FromPromptAllowedValues()
+    {
+        await using var transport = new TestServerTransport();
+        var options = CreateOptions();
+        options.PromptCollection = [McpServerPrompt.Create(
+            (
+                [System.ComponentModel.DataAnnotations.AllowedValues("dog", "cat", "fish")] string animal
+            ) => animal,
+            new McpServerPromptCreateOptions { Name = "test-prompt" })];
+
+        await using var server = McpServer.Create(transport, options, LoggerFactory);
+        var runTask = server.RunAsync(TestContext.Current.CancellationToken);
+
+        var receivedMessage = new TaskCompletionSource<JsonRpcResponse>();
+        transport.OnMessageSent = (message) =>
+        {
+            if (message is JsonRpcResponse response && response.Id.ToString() == "55")
+                receivedMessage.SetResult(response);
+        };
+
+        await transport.SendMessageAsync(new JsonRpcRequest
+        {
+            Method = RequestMethods.CompletionComplete,
+            Id = new RequestId(55),
+            Params = JsonSerializer.SerializeToNode(new CompleteRequestParams
+            {
+                Ref = new PromptReference { Name = "test-prompt" },
+                Argument = new Argument { Name = "animal", Value = "c" }
+            }, McpJsonUtilities.DefaultOptions)
+        }, TestContext.Current.CancellationToken);
+
+        var response = await receivedMessage.Task.WaitAsync(TestConstants.DefaultTimeout, TestContext.Current.CancellationToken);
+        Assert.NotNull(response);
+        var result = JsonSerializer.Deserialize<CompleteResult>(response.Result, McpJsonUtilities.DefaultOptions);
+        Assert.NotNull(result?.Completion);
+        Assert.Equal(["cat"], result.Completion.Values);
+        Assert.Equal(1, result.Completion.Total);
+
+        await transport.DisposeAsync();
+        await runTask;
+    }
+
+    [Fact]
+    public async Task Completion_AutoPopulated_FromPromptAllowedValues_NoMatch()
+    {
+        await using var transport = new TestServerTransport();
+        var options = CreateOptions();
+        options.PromptCollection = [McpServerPrompt.Create(
+            (
+                [System.ComponentModel.DataAnnotations.AllowedValues("dog", "cat")] string animal
+            ) => animal,
+            new McpServerPromptCreateOptions { Name = "test-prompt" })];
+
+        await using var server = McpServer.Create(transport, options, LoggerFactory);
+        var runTask = server.RunAsync(TestContext.Current.CancellationToken);
+
+        var receivedMessage = new TaskCompletionSource<JsonRpcResponse>();
+        transport.OnMessageSent = (message) =>
+        {
+            if (message is JsonRpcResponse response && response.Id.ToString() == "55")
+                receivedMessage.SetResult(response);
+        };
+
+        await transport.SendMessageAsync(new JsonRpcRequest
+        {
+            Method = RequestMethods.CompletionComplete,
+            Id = new RequestId(55),
+            Params = JsonSerializer.SerializeToNode(new CompleteRequestParams
+            {
+                Ref = new PromptReference { Name = "test-prompt" },
+                Argument = new Argument { Name = "animal", Value = "z" }
+            }, McpJsonUtilities.DefaultOptions)
+        }, TestContext.Current.CancellationToken);
+
+        var response = await receivedMessage.Task.WaitAsync(TestConstants.DefaultTimeout, TestContext.Current.CancellationToken);
+        Assert.NotNull(response);
+        var result = JsonSerializer.Deserialize<CompleteResult>(response.Result, McpJsonUtilities.DefaultOptions);
+        Assert.NotNull(result?.Completion);
+        Assert.Empty(result.Completion.Values);
+
+        await transport.DisposeAsync();
+        await runTask;
+    }
+
+    [Fact]
+    public async Task Completion_AutoPopulated_FromResourceAllowedValues()
+    {
+        await using var transport = new TestServerTransport();
+        var options = CreateOptions();
+        options.ResourceCollection =
+        [
+            McpServerResource.Create(
+                (
+                    [System.ComponentModel.DataAnnotations.AllowedValues("us-east-1", "us-west-2", "eu-west-1")] string region
+                ) => $"Resource for {region}",
+                new McpServerResourceCreateOptions
+                {
+                    UriTemplate = "resource://regions/{region}",
+                    Name = "regions"
+                })
+        ];
+
+        await using var server = McpServer.Create(transport, options, LoggerFactory);
+        var runTask = server.RunAsync(TestContext.Current.CancellationToken);
+
+        var receivedMessage = new TaskCompletionSource<JsonRpcResponse>();
+        transport.OnMessageSent = (message) =>
+        {
+            if (message is JsonRpcResponse response && response.Id.ToString() == "55")
+                receivedMessage.SetResult(response);
+        };
+
+        await transport.SendMessageAsync(new JsonRpcRequest
+        {
+            Method = RequestMethods.CompletionComplete,
+            Id = new RequestId(55),
+            Params = JsonSerializer.SerializeToNode(new CompleteRequestParams
+            {
+                Ref = new ResourceTemplateReference { Uri = "resource://regions/{region}" },
+                Argument = new Argument { Name = "region", Value = "us" }
+            }, McpJsonUtilities.DefaultOptions)
+        }, TestContext.Current.CancellationToken);
+
+        var response = await receivedMessage.Task.WaitAsync(TestConstants.DefaultTimeout, TestContext.Current.CancellationToken);
+        Assert.NotNull(response);
+        var result = JsonSerializer.Deserialize<CompleteResult>(response.Result, McpJsonUtilities.DefaultOptions);
+        Assert.NotNull(result?.Completion);
+        Assert.Equal(["us-east-1", "us-west-2"], result.Completion.Values);
+        Assert.Equal(2, result.Completion.Total);
+
+        await transport.DisposeAsync();
+        await runTask;
+    }
+
+    [Fact]
+    public async Task Completion_AutoPopulated_CombinedWithCustomHandler()
+    {
+        await using var transport = new TestServerTransport();
+        var options = CreateOptions();
+        options.PromptCollection = [McpServerPrompt.Create(
+            (
+                [System.ComponentModel.DataAnnotations.AllowedValues("dog", "cat")] string animal
+            ) => animal,
+            new McpServerPromptCreateOptions { Name = "test-prompt" })];
+
+        // Add a custom handler that provides additional completions
+        options.Handlers.CompleteHandler = async (request, ct) =>
+            new CompleteResult
+            {
+                Completion = new()
+                {
+                    Values = ["custom-value"],
+                    Total = 1,
+                    HasMore = false
+                }
+            };
+
+        await using var server = McpServer.Create(transport, options, LoggerFactory);
+        var runTask = server.RunAsync(TestContext.Current.CancellationToken);
+
+        var receivedMessage = new TaskCompletionSource<JsonRpcResponse>();
+        transport.OnMessageSent = (message) =>
+        {
+            if (message is JsonRpcResponse response && response.Id.ToString() == "55")
+                receivedMessage.SetResult(response);
+        };
+
+        await transport.SendMessageAsync(new JsonRpcRequest
+        {
+            Method = RequestMethods.CompletionComplete,
+            Id = new RequestId(55),
+            Params = JsonSerializer.SerializeToNode(new CompleteRequestParams
+            {
+                Ref = new PromptReference { Name = "test-prompt" },
+                Argument = new Argument { Name = "animal", Value = "" }
+            }, McpJsonUtilities.DefaultOptions)
+        }, TestContext.Current.CancellationToken);
+
+        var response = await receivedMessage.Task.WaitAsync(TestConstants.DefaultTimeout, TestContext.Current.CancellationToken);
+        Assert.NotNull(response);
+        var result = JsonSerializer.Deserialize<CompleteResult>(response.Result, McpJsonUtilities.DefaultOptions);
+        Assert.NotNull(result?.Completion);
+        // Custom handler values + auto-populated values should be combined
+        Assert.Equal(["custom-value", "dog", "cat"], result.Completion.Values);
+        Assert.Equal(3, result.Completion.Total);
+
+        await transport.DisposeAsync();
+        await runTask;
+    }
+
+    [Fact]
+    public async Task Completion_AutoPopulated_EnablesCompletionsCapabilityAutomatically()
+    {
+        // When prompts with AllowedValues are registered but no explicit Completions capability is set,
+        // the server should still handle completion requests (i.e., the capability is auto-enabled).
+        // This is verified by the fact that sending a completion request succeeds rather than failing.
+        await using var transport = new TestServerTransport();
+        var options = CreateOptions();
+        options.PromptCollection = [McpServerPrompt.Create(
+            (
+                [System.ComponentModel.DataAnnotations.AllowedValues("a", "b")] string param
+            ) => param,
+            new McpServerPromptCreateOptions { Name = "test-prompt" })];
+
+        await using var server = McpServer.Create(transport, options, LoggerFactory);
+        var runTask = server.RunAsync(TestContext.Current.CancellationToken);
+
+        var receivedMessage = new TaskCompletionSource<JsonRpcResponse>();
+        transport.OnMessageSent = (message) =>
+        {
+            if (message is JsonRpcResponse response && response.Id.ToString() == "55")
+                receivedMessage.SetResult(response);
+        };
+
+        await transport.SendMessageAsync(new JsonRpcRequest
+        {
+            Method = RequestMethods.CompletionComplete,
+            Id = new RequestId(55),
+            Params = JsonSerializer.SerializeToNode(new CompleteRequestParams
+            {
+                Ref = new PromptReference { Name = "test-prompt" },
+                Argument = new Argument { Name = "param", Value = "" }
+            }, McpJsonUtilities.DefaultOptions)
+        }, TestContext.Current.CancellationToken);
+
+        var response = await receivedMessage.Task.WaitAsync(TestConstants.DefaultTimeout, TestContext.Current.CancellationToken);
+        Assert.NotNull(response);
+        var result = JsonSerializer.Deserialize<CompleteResult>(response.Result, McpJsonUtilities.DefaultOptions);
+        Assert.NotNull(result?.Completion);
+        Assert.Equal(["a", "b"], result.Completion.Values);
+
+        await transport.DisposeAsync();
+        await runTask;
+    }
+
     [Fact]
     public async Task Can_Handle_ResourceTemplates_List_Requests()
     {