From 914edad6d95c7c07a1de2b9fd2dd95f899860f8f Mon Sep 17 00:00:00 2001 From: "Kyle Rubenok (from Dev Box)" Date: Fri, 6 Nov 2026 21:38:14 -0800 Subject: [PATCH] Avoid rebuilding registered primitive collections per stateless request Stateless Streamable HTTP creates McpServerOptions through the options factory on every request. Each time, McpServerOptionsSetup re-added every DI-registered tool, prompt and resource to new collections (locking and allocating per primitive, and parsing a Uri to hash each resource key), the Apps post-configure step re-processed every tool, and every server re-derived allowed-value completions from the JSON schemas. - Add copy constructors to McpServerPrimitiveCollection and McpServerResourceCollection. Copies share an immutable view of the source until either side is mutated, so copying is O(1) and per-session mutation and Changed semantics are unchanged. - McpServerOptionsSetup builds template collections once and gives each options instance a copy. It rebuilds a template if a primitive's name changes and keeps adding to a pre-assigned collection as before. - Cache each key's hash code inside the collection, add a lock-free, allocation-free TryAdd path for existing keys, and compare identical strings ordinally in UriTemplateComparer before parsing Uris. - Cache allowed-value completions per AIFunction-based prompt/resource. - Apply [McpAppUi] metadata once per tool instance, which also stops concurrent requests from mutating a shared tool's Meta object. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../Server/AIFunctionMcpServerPrompt.cs | 7 + .../Server/AIFunctionMcpServerResource.cs | 7 + .../Server/AllowedValueCompletions.cs | 68 +++++ .../Server/McpServerImpl.cs | 40 +-- .../Server/McpServerPrimitiveCollection.cs | 221 ++++++++++++-- .../Server/McpServerResourceCollection.cs | 25 +- src/ModelContextProtocol.Core/UriTemplate.cs | 8 +- .../Server/McpAppsBuilderExtensions.cs | 22 +- .../McpServerOptionsSetup.cs | 135 +++++++-- .../StatelessServerTests.cs | 45 +++ .../McpServerOptionsSetupTests.cs | 124 ++++++++ .../Server/McpAppsTests.cs | 48 +++ .../McpServerPrimitiveCollectionTests.cs | 276 ++++++++++++++++++ .../Server/McpServerTests.cs | 46 +++ 14 files changed, 985 insertions(+), 87 deletions(-) create mode 100644 src/ModelContextProtocol.Core/Server/AllowedValueCompletions.cs diff --git a/src/ModelContextProtocol.Core/Server/AIFunctionMcpServerPrompt.cs b/src/ModelContextProtocol.Core/Server/AIFunctionMcpServerPrompt.cs index c755cd9e5..d9073d4b5 100644 --- a/src/ModelContextProtocol.Core/Server/AIFunctionMcpServerPrompt.cs +++ b/src/ModelContextProtocol.Core/Server/AIFunctionMcpServerPrompt.cs @@ -177,6 +177,13 @@ private static McpServerPromptCreateOptions DeriveOptions(MethodInfo method, Mcp /// Gets the wrapped by this prompt. internal AIFunction AIFunction { get; } + /// Backing cache for . + private object? _allowedValueCompletions; + + /// Gets the values that each restricted parameter of the prompt allows, or if none are restricted. + internal Dictionary? AllowedValueCompletions => + Server.AllowedValueCompletions.GetOrCreate(ref _allowedValueCompletions, AIFunction.JsonSchema); + /// Initializes a new instance of the class. private AIFunctionMcpServerPrompt(AIFunction function, Prompt prompt, IReadOnlyList metadata) { diff --git a/src/ModelContextProtocol.Core/Server/AIFunctionMcpServerResource.cs b/src/ModelContextProtocol.Core/Server/AIFunctionMcpServerResource.cs index 7f0f29140..75e2edfac 100644 --- a/src/ModelContextProtocol.Core/Server/AIFunctionMcpServerResource.cs +++ b/src/ModelContextProtocol.Core/Server/AIFunctionMcpServerResource.cs @@ -289,6 +289,13 @@ private static string DeriveUriTemplate(string name, AIFunction function) /// Gets the wrapped by this resource. internal AIFunction AIFunction { get; } + /// Backing cache for . + private object? _allowedValueCompletions; + + /// Gets the values that each restricted parameter of the resource allows, or if none are restricted. + internal Dictionary? AllowedValueCompletions => + Server.AllowedValueCompletions.GetOrCreate(ref _allowedValueCompletions, AIFunction.JsonSchema); + /// Initializes a new instance of the class. private AIFunctionMcpServerResource(AIFunction function, ResourceTemplate resourceTemplate, IReadOnlyList metadata) { diff --git a/src/ModelContextProtocol.Core/Server/AllowedValueCompletions.cs b/src/ModelContextProtocol.Core/Server/AllowedValueCompletions.cs new file mode 100644 index 000000000..85aa41752 --- /dev/null +++ b/src/ModelContextProtocol.Core/Server/AllowedValueCompletions.cs @@ -0,0 +1,68 @@ +using System.Text.Json; + +namespace ModelContextProtocol.Server; + +/// +/// Extracts the values that parameters of an AIFunction-based prompt or resource are restricted to, from the enum values +/// in the function's JSON schema, for use as completions. +/// +internal static class AllowedValueCompletions +{ + /// Marks a cache whose function was found to have no restricted parameters. + private static readonly object s_none = new(); + + /// + /// Gets the parameter name → allowed values lookup for , computing and caching it in + /// on first use. The returned dictionary is shared and must not be modified. + /// + /// + /// The schema of an AIFunction is immutable, so the result can be computed once per primitive rather than once per + /// session. Concurrent first calls may each compute it; the results are equivalent. + /// + public static Dictionary? GetOrCreate(ref object? cache, JsonElement schema) + { + object? cached = Volatile.Read(ref cache); + if (cached is null) + { + cached = (object?)Create(schema) ?? s_none; + Volatile.Write(ref cache, cached); + } + + return cached as Dictionary; + } + + private static Dictionary? Create(JsonElement schema) + { + if (!schema.TryGetProperty("properties", out JsonElement properties) || + properties.ValueKind is not JsonValueKind.Object) + { + return null; + } + + Dictionary? paramValues = null; + foreach (var param in properties.EnumerateObject()) + { + if (param.Value.TryGetProperty("enum", out JsonElement enumValues) && + enumValues.ValueKind is JsonValueKind.Array) + { + List? values = null; + foreach (var item in enumValues.EnumerateArray()) + { + if (item.ValueKind is JsonValueKind.String && item.GetString() is { } str) + { + values ??= []; + values.Add(str); + } + } + + if (values is not null) + { + paramValues ??= new(StringComparer.Ordinal); + paramValues[param.Name] = [.. values]; + } + } + } + + return paramValues; + } +} diff --git a/src/ModelContextProtocol.Core/Server/McpServerImpl.cs b/src/ModelContextProtocol.Core/Server/McpServerImpl.cs index db11a84de..67973c840 100644 --- a/src/ModelContextProtocol.Core/Server/McpServerImpl.cs +++ b/src/ModelContextProtocol.Core/Server/McpServerImpl.cs @@ -1119,16 +1119,16 @@ private void ConfigureCompletion(McpServerOptions options) Dictionary>? result = null; foreach (var primitive in primitives) { - JsonElement schema; + Dictionary? paramValues; string id; if (primitive is AIFunctionMcpServerPrompt aiPrompt) { - schema = aiPrompt.AIFunction.JsonSchema; + paramValues = aiPrompt.AllowedValueCompletions; id = aiPrompt.ProtocolPrompt.Name; } else if (primitive is AIFunctionMcpServerResource aiResource && aiResource.IsTemplated) { - schema = aiResource.AIFunction.JsonSchema; + paramValues = aiResource.AllowedValueCompletions; id = aiResource.ProtocolResourceTemplate.UriTemplate; } else @@ -1136,38 +1136,10 @@ private void ConfigureCompletion(McpServerOptions options) continue; } - if (schema.TryGetProperty("properties", out JsonElement properties) && - properties.ValueKind is JsonValueKind.Object) + if (paramValues is not null) { - Dictionary? paramValues = null; - foreach (var param in properties.EnumerateObject()) - { - if (param.Value.TryGetProperty("enum", out JsonElement enumValues) && - enumValues.ValueKind is JsonValueKind.Array) - { - List? values = null; - foreach (var item in enumValues.EnumerateArray()) - { - if (item.ValueKind is JsonValueKind.String && 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; - } + result ??= new(StringComparer.Ordinal); + result[id] = paramValues; } } diff --git a/src/ModelContextProtocol.Core/Server/McpServerPrimitiveCollection.cs b/src/ModelContextProtocol.Core/Server/McpServerPrimitiveCollection.cs index 02174f249..f855591f1 100644 --- a/src/ModelContextProtocol.Core/Server/McpServerPrimitiveCollection.cs +++ b/src/ModelContextProtocol.Core/Server/McpServerPrimitiveCollection.cs @@ -1,5 +1,6 @@ using System.Collections; using System.Collections.Concurrent; +using System.Diagnostics; using System.Diagnostics.CodeAnalysis; namespace ModelContextProtocol.Server; @@ -13,11 +14,26 @@ namespace ModelContextProtocol.Server; public class McpServerPrimitiveCollection : ICollection, IReadOnlyCollection where T : IMcpServerPrimitive { - /// Concurrent dictionary of primitives, indexed by their names. - private readonly ConcurrentDictionary _primitives; + /// Compares keys using the key comparer supplied at construction, with each key's hash code cached. + private readonly PrimitiveKeyComparer _keyComparer; + + /// + /// Concurrent dictionary of primitives, indexed by their names, or if this collection was + /// copied from another and hasn't been mutated since, in which case holds its contents. + /// + /// Only assigned under , before is cleared. + private ConcurrentDictionary? _primitives; /// The entries in , in the order they were added. - private readonly List _orderedEntries = []; + /// Only accessed under . exactly when is. + private List? _orderedEntries; + + /// + /// An immutable view of the collection's current contents, shared with any collections copied from this one, or + /// if the collection has been mutated since the view was created. While set, reads are served + /// from it, and the first mutation detaches from it. + /// + private volatile FrozenEntries? _frozen; /// Lock serializing additions and removals across and . private readonly object _mutationLock = new(); @@ -39,7 +55,38 @@ public class McpServerPrimitiveCollection : ICollection, IReadOnlyCollecti /// public McpServerPrimitiveCollection(IEqualityComparer? keyComparer = null) { - _primitives = new(keyComparer ?? EqualityComparer.Default); + _keyComparer = new(keyComparer ?? EqualityComparer.Default); + _primitives = new(_keyComparer); + _orderedEntries = []; + } + + /// + /// Initializes a new instance of the class that contains the same + /// primitives as , in the same order and under the same names, and uses the same key comparer. + /// + /// The collection whose primitives are copied. + /// is . + /// + /// + /// The new collection is independent of : adding or removing primitives in either one + /// doesn't affect the other, and handlers aren't copied. + /// + /// + /// Copying is inexpensive. Both collections share an immutable view of the primitives until either is mutated, so + /// repeatedly copying an unchanging collection, such as once per session, doesn't copy the primitives each time. + /// + /// + /// The copy is taken from the primitives stored by this base type. Members overridden by a type derived from + /// aren't consulted. + /// + /// + public McpServerPrimitiveCollection(McpServerPrimitiveCollection collection) + { + Throw.IfNull(collection); + + FrozenEntries frozen = collection.Freeze(); + _keyComparer = frozen.KeyComparer; + _frozen = frozen; } /// Occurs when the collection is changed. @@ -50,10 +97,10 @@ public McpServerPrimitiveCollection(IEqualityComparer? keyComparer = nul public event EventHandler? Changed; /// Gets the number of primitives in the collection. - public int Count => _primitives.Count; + public int Count => _frozen is { } frozen ? frozen.Entries.Length : _primitives!.Count; /// Gets a value that indicates whether there are any primitives in the collection. - public bool IsEmpty => _primitives.IsEmpty; + public bool IsEmpty => _frozen is { } frozen ? frozen.Entries.Length == 0 : _primitives!.IsEmpty; /// /// Begins a deferred-change scope. notifications are suppressed @@ -145,7 +192,9 @@ public T this[string name] get { Throw.IfNull(name); - return _primitives[name].Primitive; + return TryGetEntry(_keyComparer.CreateKey(name), out Entry? entry) ? + entry.Primitive : + throw new KeyNotFoundException($"The given key '{name}' was not present in the dictionary."); } } @@ -154,8 +203,19 @@ public virtual void Clear() { lock (_mutationLock) { - _primitives.Clear(); - _orderedEntries.Clear(); + if (_primitives is null) + { + // The contents are only held by the shared frozen view, so start over rather than copying it first. + _primitives = new(_keyComparer); + _orderedEntries = []; + } + else + { + _primitives.Clear(); + _orderedEntries!.Clear(); + } + + _frozen = null; _snapshot = null; } @@ -182,16 +242,25 @@ public virtual bool TryAdd(T primitive) { Throw.IfNull(primitive); - string key = primitive.Id; + PrimitiveKey key = _keyComparer.CreateKey(primitive.Id); + + // Most attempts to add an existing primitive come from re-adding the same set of primitives, so check + // without locking or allocating first. The check is repeated under the lock below. + if (TryGetEntry(key, out _)) + { + return false; + } + lock (_mutationLock) { + ConcurrentDictionary primitives = DetachFromFrozen(); Entry entry = new(key, primitive); - if (!_primitives.TryAdd(key, entry)) + if (!primitives.TryAdd(key, entry)) { return false; } - _orderedEntries.Add(entry); + _orderedEntries!.Add(entry); _snapshot = null; } @@ -209,10 +278,10 @@ public virtual bool Remove(T primitive) { Throw.IfNull(primitive); - string key = primitive.Id; + PrimitiveKey key = _keyComparer.CreateKey(primitive.Id); lock (_mutationLock) { - if (!_primitives.TryGetValue(key, out Entry? entry) || + if (!TryGetEntry(key, out Entry? entry) || !EqualityComparer.Default.Equals(entry.Primitive, primitive)) { return false; @@ -220,12 +289,13 @@ public virtual bool Remove(T primitive) // Remove only the exact entry compared above. A primitive's Equals can reenter the collection while the // lock is held and replace the entry under the same key; the replacement must then stay in place. - if (!((ICollection>)_primitives).Remove(new(key, entry))) + ConcurrentDictionary primitives = DetachFromFrozen(); + if (!((ICollection>)primitives).Remove(new(key, entry))) { return false; } - _orderedEntries.Remove(entry); + _orderedEntries!.Remove(entry); _snapshot = null; } @@ -244,7 +314,7 @@ public virtual bool TryGetPrimitive(string name, [NotNullWhen(true)] out T? prim { Throw.IfNull(name); - if (_primitives.TryGetValue(name, out Entry? entry)) + if (TryGetEntry(_keyComparer.CreateKey(name), out Entry? entry)) { primitive = entry.Primitive; return true; @@ -262,12 +332,12 @@ public virtual bool Contains(T primitive) { Throw.IfNull(primitive); return - _primitives.TryGetValue(primitive.Id, out Entry? entry) && + TryGetEntry(_keyComparer.CreateKey(primitive.Id), out Entry? entry) && EqualityComparer.Default.Equals(entry.Primitive, primitive); } /// Gets the names of all of the primitives in the collection, in the order the primitives were added. - public virtual ICollection PrimitiveNames => Array.ConvertAll(GetSnapshot(), static e => e.Key); + public virtual ICollection PrimitiveNames => Array.ConvertAll(GetSnapshot(), static e => e.Key.Name); /// Creates an array containing all of the primitives in the collection. /// An array containing all of the primitives in the collection, in the order they were added. @@ -304,30 +374,137 @@ public virtual IEnumerator GetEnumerator() /// The returned array is shared between callers and must not be modified or exposed. private Entry[] GetSnapshot() { + if (_frozen is { } frozen) + { + return frozen.Entries; + } + Entry[]? snapshot = _snapshot; if (snapshot is null) { lock (_mutationLock) { - snapshot = _snapshot ??= _orderedEntries.ToArray(); + snapshot = _frozen?.Entries ?? (_snapshot ??= _orderedEntries!.ToArray()); } } return snapshot; } + /// Looks up the entry stored under without locking. + private bool TryGetEntry(PrimitiveKey key, [NotNullWhen(true)] out Entry? entry) + { + // _primitives is assigned before _frozen is cleared, so reading _frozen first always observes one or the other. + if (_frozen is { } frozen) + { + return frozen.Lookup.TryGetValue(key, out entry); + } + + return _primitives!.TryGetValue(key, out entry); + } + + /// + /// Returns an immutable view of the collection's current contents that copies can share, creating it if the + /// collection has been mutated since the view was last created. + /// + private FrozenEntries Freeze() + { + if (_frozen is { } frozen) + { + return frozen; + } + + lock (_mutationLock) + { + if (_frozen is { } existing) + { + return existing; + } + + Entry[] entries = _orderedEntries!.ToArray(); + Dictionary lookup = new(entries.Length, _keyComparer); + foreach (Entry entry in entries) + { + lookup.Add(entry.Key, entry); + } + + // The mutable storage is kept, so a later mutation of this collection only needs to stop using the view. + return _frozen = new FrozenEntries(lookup, entries, _keyComparer); + } + } + + /// + /// Stops serving reads from the shared frozen view, if any, so the collection can be mutated, creating mutable + /// storage from the view if this collection doesn't have its own yet. Must be called under . + /// + private ConcurrentDictionary DetachFromFrozen() + { + Debug.Assert(Monitor.IsEntered(_mutationLock), "Must be called under the mutation lock."); + + if (_frozen is { } frozen) + { + if (_primitives is null) + { + _orderedEntries = [.. frozen.Entries]; + _primitives = new(frozen.Lookup, _keyComparer); + } + + _frozen = null; + } + + return _primitives!; + } + /// A primitive together with the key it was added under. /// /// The key is captured when the primitive is added, because may be backed by /// mutable protocol properties. /// - private sealed class Entry(string key, T primitive) + private sealed class Entry(PrimitiveKey key, T primitive) { - public string Key { get; } = key; + public PrimitiveKey Key { get; } = key; public T Primitive { get; } = primitive; } + /// An immutable view of a collection's entries, shared by the collection and any copies made from it. + private sealed class FrozenEntries(Dictionary lookup, Entry[] entries, PrimitiveKeyComparer keyComparer) + { + /// Gets the entries, indexed by key. Only read after construction, so safe for concurrent readers. + public Dictionary Lookup { get; } = lookup; + + /// Gets the entries in the order they were added. + public Entry[] Entries { get; } = entries; + + public PrimitiveKeyComparer KeyComparer { get; } = keyComparer; + } + + /// A primitive name together with its hash code under the collection's key comparer. + /// + /// Caching the hash code means stored names are hashed only once, which matters for comparers whose hashing is + /// expensive, such as the URI-aware comparer used for resources. + /// + private readonly struct PrimitiveKey(string name, int hashCode) + { + public string Name { get; } = name; + + public int HashCode { get; } = hashCode; + + public override string ToString() => Name; + } + + /// Compares values using their cached hash codes and the collection's key comparer. + private sealed class PrimitiveKeyComparer(IEqualityComparer nameComparer) : IEqualityComparer + { + public PrimitiveKey CreateKey(string name) => new(name, nameComparer.GetHashCode(name)); + + public bool Equals(PrimitiveKey x, PrimitiveKey y) => + x.HashCode == y.HashCode && + (ReferenceEquals(x.Name, y.Name) || nameComparer.Equals(x.Name, y.Name)); + + public int GetHashCode(PrimitiveKey key) => key.HashCode; + } + /// IEnumerator IEnumerable.GetEnumerator() => GetEnumerator(); diff --git a/src/ModelContextProtocol.Core/Server/McpServerResourceCollection.cs b/src/ModelContextProtocol.Core/Server/McpServerResourceCollection.cs index fb5f6b4e2..9404f558e 100644 --- a/src/ModelContextProtocol.Core/Server/McpServerResourceCollection.cs +++ b/src/ModelContextProtocol.Core/Server/McpServerResourceCollection.cs @@ -1,5 +1,26 @@ namespace ModelContextProtocol.Server; /// Provides a thread-safe collection of instances, indexed by their URI templates. -public sealed class McpServerResourceCollection() - : McpServerPrimitiveCollection(UriTemplate.UriTemplateComparer.Instance); \ No newline at end of file +public sealed class McpServerResourceCollection : McpServerPrimitiveCollection +{ + /// Initializes a new instance of the class. + public McpServerResourceCollection() + : base(UriTemplate.UriTemplateComparer.Instance) + { + } + + /// + /// Initializes a new instance of the class that contains the same + /// resources as , in the same order. + /// + /// The collection whose resources are copied. + /// is . + /// + /// The new collection is independent of , and copying is inexpensive. For details, see + /// . + /// + public McpServerResourceCollection(McpServerResourceCollection collection) + : base(collection) + { + } +} \ No newline at end of file diff --git a/src/ModelContextProtocol.Core/UriTemplate.cs b/src/ModelContextProtocol.Core/UriTemplate.cs index fc577d5aa..62cf9e6ce 100644 --- a/src/ModelContextProtocol.Core/UriTemplate.cs +++ b/src/ModelContextProtocol.Core/UriTemplate.cs @@ -513,13 +513,19 @@ internal sealed class UriTemplateComparer : IEqualityComparer public bool Equals(string? uriTemplate1, string? uriTemplate2) { + // Strings that are identical are equal under either rule, so avoid parsing them as URIs. + if (string.Equals(uriTemplate1, uriTemplate2, StringComparison.Ordinal)) + { + return true; + } + if (TryParseAsNonTemplatedUri(uriTemplate1, out Uri? uri1) && TryParseAsNonTemplatedUri(uriTemplate2, out Uri? uri2)) { return uri1 == uri2; } - return string.Equals(uriTemplate1, uriTemplate2, StringComparison.Ordinal); + return false; } public int GetHashCode([DisallowNull] string uriTemplate) diff --git a/src/ModelContextProtocol.Extensions.Apps/Server/McpAppsBuilderExtensions.cs b/src/ModelContextProtocol.Extensions.Apps/Server/McpAppsBuilderExtensions.cs index a68d8fe50..cf46a0b2e 100644 --- a/src/ModelContextProtocol.Extensions.Apps/Server/McpAppsBuilderExtensions.cs +++ b/src/ModelContextProtocol.Extensions.Apps/Server/McpAppsBuilderExtensions.cs @@ -3,6 +3,7 @@ using ModelContextProtocol.Protocol; using ModelContextProtocol.Server; using System.Diagnostics.CodeAnalysis; +using System.Runtime.CompilerServices; namespace ModelContextProtocol.Extensions.Apps; @@ -49,6 +50,15 @@ public static IMcpServerBuilder WithMcpApps(this IMcpServerBuilder builder) private sealed class McpAppsPostConfigureOptions : IPostConfigureOptions { + /// The tools whose has already been applied. + /// + /// Options can be created many times, for example once per request by a stateless HTTP server, and registered + /// tools are typically shared between them, or even between service providers. Applying the attribute once per + /// tool avoids repeating the work and avoids concurrently mutating a shared tool's metadata. Tools are held + /// weakly so that tools created for a single session aren't kept alive. + /// + private static readonly ConditionalWeakTable s_processedTools = new(); + public void PostConfigure(string? name, McpServerOptions options) { // Advertise server-side MCP Apps support in capabilities. @@ -63,7 +73,17 @@ public void PostConfigure(string? name, McpServerOptions options) { foreach (var tool in tools) { - McpApps.ApplyAppUiAttributes(tool); + if (!s_processedTools.TryGetValue(tool, out _)) + { + lock (s_processedTools) + { + if (!s_processedTools.TryGetValue(tool, out _)) + { + McpApps.ApplyAppUiAttributes(tool); + s_processedTools.Add(tool, s_processedTools); + } + } + } } } } diff --git a/src/ModelContextProtocol/McpServerOptionsSetup.cs b/src/ModelContextProtocol/McpServerOptionsSetup.cs index c46854460..44d50a1c5 100644 --- a/src/ModelContextProtocol/McpServerOptionsSetup.cs +++ b/src/ModelContextProtocol/McpServerOptionsSetup.cs @@ -9,11 +9,20 @@ namespace ModelContextProtocol; /// The individually registered tools. /// The individually registered prompts. /// The individually registered resources. +/// +/// Options can be created many times from the same instance, for example once per request by a stateless HTTP server. +/// The registered primitives don't change for the lifetime of this instance, so they're collected once into template +/// collections, and each options instance gets an inexpensive, independently mutable copy of the relevant template. +/// internal sealed class McpServerOptionsSetup( IEnumerable serverTools, IEnumerable serverPrompts, IEnumerable serverResources) : IConfigureOptions { + private Template>? _toolTemplate; + private Template>? _promptTemplate; + private Template? _resourceTemplate; + /// /// Configures the given McpServerOptions instance by setting server information /// and collecting registered server primitives. @@ -23,49 +32,121 @@ public void Configure(McpServerOptions options) { Throw.IfNull(options); - // Collect all of the provided tools into a tools collection. If the options already has - // a collection, add to it, otherwise create a new one. We want to maintain the identity - // of an existing collection in case someone has provided their own derived type, wants - // change notifications, etc. - McpServerPrimitiveCollection toolCollection = options.ToolCollection ?? []; - foreach (var tool in serverTools) + // If the options already has a collection, add to it rather than replacing it. We want to maintain the identity + // of an existing collection in case someone has provided their own derived type, wants change notifications, etc. + // Otherwise, give the options its own copy of the registered primitives, if there are any. + if (options.ToolCollection is { } toolCollection) + { + AddAll(toolCollection, serverTools); + } + else if (GetTemplate(ref _toolTemplate, serverTools, static () => []) is { IsEmpty: false } tools) + { + options.ToolCollection = new(tools); + } + + if (options.PromptCollection is { } promptCollection) { - toolCollection.TryAdd(tool); + AddAll(promptCollection, serverPrompts); + } + else if (GetTemplate(ref _promptTemplate, serverPrompts, static () => []) is { IsEmpty: false } prompts) + { + options.PromptCollection = new(prompts); } - if (!toolCollection.IsEmpty) + if (options.ResourceCollection is { } resourceCollection) + { + AddAll(resourceCollection, serverResources); + } + else if (GetTemplate(ref _resourceTemplate, serverResources, static () => []) is { IsEmpty: false } resources) { - options.ToolCollection = toolCollection; + options.ResourceCollection = new(resources); } + } - // Collect all of the provided prompts into a prompts collection. If the options already has - // a collection, add to it, otherwise create a new one. We want to maintain the identity - // of an existing collection in case someone has provided their own derived type, wants - // change notifications, etc. - McpServerPrimitiveCollection promptCollection = options.PromptCollection ?? []; - foreach (var prompt in serverPrompts) + private static void AddAll(McpServerPrimitiveCollection collection, IEnumerable primitives) + where T : IMcpServerPrimitive + { + foreach (var primitive in primitives) { - promptCollection.TryAdd(prompt); + collection.TryAdd(primitive); } + } - if (!promptCollection.IsEmpty) + /// Gets the template collection for , building or rebuilding it if needed. + private static TCollection GetTemplate( + ref Template? field, IEnumerable primitives, Func createCollection) + where T : class, IMcpServerPrimitive + where TCollection : McpServerPrimitiveCollection + { + // Concurrent callers may each build a template; they're equivalent, so whichever is stored last wins. + Template? template = Volatile.Read(ref field); + if (template is null || !template.IsCurrent(primitives)) { - options.PromptCollection = promptCollection; + template = Template.Create(primitives, createCollection()); + Volatile.Write(ref field, template); } - // Collect all of the provided resources into a resources collection. If the options already has - // a collection, add to it, otherwise create a new one. We want to maintain the identity - // of an existing collection in case someone has provided their own derived type, wants - // change notifications, etc. - McpServerResourceCollection resourceCollection = options.ResourceCollection ?? []; - foreach (var resource in serverResources) + return template.Collection; + } + + /// A collection of the registered primitives, together with what it was built from. + private sealed class Template(T[] primitives, string[] ids, TCollection collection) + where T : class, IMcpServerPrimitive + where TCollection : McpServerPrimitiveCollection + { + /// Gets the collection of registered primitives. It's never mutated, only copied. + public TCollection Collection { get; } = collection; + + public static Template Create(IEnumerable source, TCollection collection) { - resourceCollection.TryAdd(resource); + T[] primitives = source.ToArray(); + string[] ids = Array.ConvertAll(primitives, static p => p.Id); + AddAll(collection, primitives); + return new(primitives, ids, collection); } - if (!resourceCollection.IsEmpty) + /// + /// Determines whether the template still reflects : the same primitive instances in the + /// same order, each with the name it had when the template was built. The name is checked because it's the + /// primitive's key, and it can be backed by a mutable protocol property. + /// + public bool IsCurrent(IEnumerable source) { - options.ResourceCollection = resourceCollection; + if (source is T[] array) + { + if (array.Length != primitives.Length) + { + return false; + } + + for (int i = 0; i < array.Length; i++) + { + if (!IsCurrent(i, array[i])) + { + return false; + } + } + + return true; + } + + int count = 0; + foreach (T primitive in source) + { + if (count >= primitives.Length || !IsCurrent(count, primitive)) + { + return false; + } + + count++; + } + + return count == primitives.Length; } + + private bool IsCurrent(int index, T primitive) => + ReferenceEquals(primitive, primitives[index]) && + primitive.Id is var id && + (ReferenceEquals(id, ids[index]) || string.Equals(id, ids[index], StringComparison.Ordinal)); } } diff --git a/tests/ModelContextProtocol.AspNetCore.Tests/StatelessServerTests.cs b/tests/ModelContextProtocol.AspNetCore.Tests/StatelessServerTests.cs index c310e5475..891af3408 100644 --- a/tests/ModelContextProtocol.AspNetCore.Tests/StatelessServerTests.cs +++ b/tests/ModelContextProtocol.AspNetCore.Tests/StatelessServerTests.cs @@ -335,6 +335,51 @@ public async Task ConfigureSessionOptions_RunsPerRequest_InStatelessMode() Assert.Equal("configured-beta", Assert.IsType(content2).Text); } + [Fact] + public async Task ConfigureSessionOptions_MutatingRegisteredTools_DoesNotAffectOtherRequests() + { + Builder.Services.AddMcpServer() + .WithHttpTransport(options => + { + options.Stateless = true; + options.ConfigureSessionOptions = (httpContext, mcpServerOptions, cancellationToken) => + { + if (httpContext.Request.Headers["X-Tool-Mode"] == "swap" && + mcpServerOptions.ToolCollection is { } tools && + tools.TryGetPrimitive("sharedTool", out var sharedTool)) + { + tools.Remove(sharedTool); + tools.Add(McpServerTool.Create(() => "session", new() { Name = "sessionTool" })); + } + + return Task.CompletedTask; + }; + }) + .WithTools([McpServerTool.Create(() => "shared", new() { Name = "sharedTool" })]); + + _app = Builder.Build(); + _app.MapMcp(); + await _app.StartAsync(TestContext.Current.CancellationToken); + + HttpClient.DefaultRequestHeaders.Accept.Add(new("application/json")); + HttpClient.DefaultRequestHeaders.Accept.Add(new("text/event-stream")); + + HttpClient.DefaultRequestHeaders.Add("X-Tool-Mode", "swap"); + await using (var swappedClient = await ConnectMcpClientAsync()) + { + var swappedTools = await swappedClient.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken); + Assert.Equal(["sessionTool"], swappedTools.Select(t => t.Name)); + } + + HttpClient.DefaultRequestHeaders.Remove("X-Tool-Mode"); + await using var client = await ConnectMcpClientAsync(); + var tools = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken); + Assert.Equal(["sharedTool"], tools.Select(t => t.Name)); + + var result = await client.CallToolAsync("sharedTool", cancellationToken: TestContext.Current.CancellationToken); + Assert.Equal("shared", Assert.IsType(Assert.Single(result.Content)).Text); + } + [Fact] public async Task StatelessMode_DoesNotAdvertise_ListChangedCapabilities() { diff --git a/tests/ModelContextProtocol.Tests/Configuration/McpServerOptionsSetupTests.cs b/tests/ModelContextProtocol.Tests/Configuration/McpServerOptionsSetupTests.cs index dc2eaf805..492db72e8 100644 --- a/tests/ModelContextProtocol.Tests/Configuration/McpServerOptionsSetupTests.cs +++ b/tests/ModelContextProtocol.Tests/Configuration/McpServerOptionsSetupTests.cs @@ -284,4 +284,128 @@ public void Configure_WithCompleteHandler_CreatesCompletionsCapability() Assert.NotNull(options.Capabilities?.Completions); } #endregion + + #region Registered Primitive Tests + [Fact] + public void Configure_RepeatedCreate_GivesEachOptionsItsOwnCollections() + { + var tool = McpServerTool.Create(() => "result", new() { Name = "tool" }); + var prompt = McpServerPrompt.Create(() => "prompt", new() { Name = "prompt" }); + var resource = McpServerResource.Create(() => "content", new() { UriTemplate = "test://resource", Name = "resource" }); + + var services = new ServiceCollection(); + services.AddMcpServer().WithTools([tool]).WithPrompts([prompt]).WithResources([resource]); + using var serviceProvider = services.BuildServiceProvider(); + var factory = serviceProvider.GetRequiredService>(); + + var options1 = factory.Create(Options.DefaultName); + var options2 = factory.Create(Options.DefaultName); + + Assert.NotSame(options1.ToolCollection, options2.ToolCollection); + Assert.NotSame(options1.PromptCollection, options2.PromptCollection); + Assert.NotSame(options1.ResourceCollection, options2.ResourceCollection); + Assert.Same(tool, Assert.Single(options1.ToolCollection!)); + Assert.Same(tool, Assert.Single(options2.ToolCollection!)); + Assert.Same(prompt, Assert.Single(options2.PromptCollection!)); + Assert.Same(resource, Assert.Single(options2.ResourceCollection!)); + } + + [Fact] + public void Configure_MutatingOneOptionsCollections_DoesNotAffectOtherOptions() + { + var tool = McpServerTool.Create(() => "result", new() { Name = "tool" }); + var resource = McpServerResource.Create(() => "content", new() { UriTemplate = "test://resource", Name = "resource" }); + + var services = new ServiceCollection(); + services.AddMcpServer().WithTools([tool]).WithResources([resource]); + using var serviceProvider = services.BuildServiceProvider(); + var factory = serviceProvider.GetRequiredService>(); + + var options1 = factory.Create(Options.DefaultName); + Assert.True(options1.ToolCollection!.Remove(tool)); + options1.ToolCollection.Add(McpServerTool.Create(() => "session", new() { Name = "sessionTool" })); + options1.ResourceCollection!.Clear(); + + var options2 = factory.Create(Options.DefaultName); + + Assert.Equal(["sessionTool"], options1.ToolCollection.PrimitiveNames); + Assert.Empty(options1.ResourceCollection); + Assert.Same(tool, Assert.Single(options2.ToolCollection!)); + Assert.Same(resource, Assert.Single(options2.ResourceCollection!)); + } + + [Fact] + public void Configure_PreassignedCollection_IsKeptAndRegisteredPrimitivesAdded() + { + var registered = McpServerTool.Create(() => "result", new() { Name = "registered" }); + var preassigned = McpServerTool.Create(() => "result", new() { Name = "preassigned" }); + McpServerPrimitiveCollection collection = [preassigned]; + + var services = new ServiceCollection(); + services.Configure(options => options.ToolCollection = collection); + services.AddMcpServer().WithTools([registered]); + using var serviceProvider = services.BuildServiceProvider(); + var factory = serviceProvider.GetRequiredService>(); + + for (int i = 0; i < 2; i++) + { + var options = factory.Create(Options.DefaultName); + Assert.Same(collection, options.ToolCollection); + Assert.Equal(["preassigned", "registered"], collection.PrimitiveNames); + } + } + + [Fact] + public void Configure_DuplicateRegisteredNames_FirstRegistrationWins() + { + var first = McpServerTool.Create(() => "first", new() { Name = "tool" }); + var second = McpServerTool.Create(() => "second", new() { Name = "tool" }); + + var services = new ServiceCollection(); + services.AddMcpServer().WithTools([first, second]); + using var serviceProvider = services.BuildServiceProvider(); + var factory = serviceProvider.GetRequiredService>(); + + for (int i = 0; i < 2; i++) + { + Assert.Same(first, Assert.Single(factory.Create(Options.DefaultName).ToolCollection!)); + } + } + + [Fact] + public void Configure_RegisteredPrimitiveRenamed_LaterOptionsUseNewName() + { + var tool = McpServerTool.Create(() => "result", new() { Name = "before" }); + + var services = new ServiceCollection(); + services.AddMcpServer().WithTools([tool]); + using var serviceProvider = services.BuildServiceProvider(); + var factory = serviceProvider.GetRequiredService>(); + + var options1 = factory.Create(Options.DefaultName); + tool.ProtocolTool.Name = "after"; + var options2 = factory.Create(Options.DefaultName); + + Assert.Equal(["before"], options1.ToolCollection!.PrimitiveNames); + Assert.Equal(["after"], options2.ToolCollection!.PrimitiveNames); + Assert.Same(tool, options2.ToolCollection["after"]); + } + + [Fact] + public void Configure_NoRegisteredPrimitives_LeavesCollectionsNull() + { + var services = new ServiceCollection(); + services.AddMcpServer(); + using var serviceProvider = services.BuildServiceProvider(); + var factory = serviceProvider.GetRequiredService>(); + + for (int i = 0; i < 2; i++) + { + var options = factory.Create(Options.DefaultName); + Assert.Null(options.ToolCollection); + Assert.Null(options.PromptCollection); + Assert.Null(options.ResourceCollection); + } + } + #endregion } \ No newline at end of file diff --git a/tests/ModelContextProtocol.Tests/Server/McpAppsTests.cs b/tests/ModelContextProtocol.Tests/Server/McpAppsTests.cs index c0d3194f1..68ae2f105 100644 --- a/tests/ModelContextProtocol.Tests/Server/McpAppsTests.cs +++ b/tests/ModelContextProtocol.Tests/Server/McpAppsTests.cs @@ -505,6 +505,54 @@ public void WithMcpApps_AppliesAppUiAttributes_ViaOptions() Assert.NotNull(weatherTool); } + [Fact] + public async Task WithMcpApps_RepeatedConcurrentOptionsCreation_AppliesAppUiOncePerTool() + { + var sc = new ServiceCollection(); + sc.AddMcpServer() + .WithTools([typeof(TestToolsWithAppUi)]) + .WithMcpApps(); + + using var sp = sc.BuildServiceProvider(); + var factory = sp.GetRequiredService>(); + + var creations = Enumerable.Range(0, 8).Select(_ => Task.Run(() => + { + for (int i = 0; i < 25; i++) + { + var options = factory.Create(Options.DefaultName); + Assert.True(options.Capabilities!.Extensions!.ContainsKey(McpApps.ExtensionId)); + Assert.Equal(2, options.ToolCollection!.Count(t => t.ProtocolTool.Meta?["ui"] is not null)); + } + }, TestContext.Current.CancellationToken)); + + await Task.WhenAll(creations); + + var tools = factory.Create(Options.DefaultName).ToolCollection!; + var weatherTool = tools.Single(t => t.ProtocolTool.Meta?["ui"]?["resourceUri"]?.GetValue() == "ui://weather/view.html"); + Assert.Single(weatherTool.ProtocolTool.Meta!, kvp => kvp.Key == "ui"); + } + + [Fact] + public void WithMcpApps_ToolCreatedPerOptions_IsProcessed() + { + var method = typeof(TestToolsWithAppUi).GetMethod(nameof(TestToolsWithAppUi.WeatherTool))!; + + var sc = new ServiceCollection(); + sc.AddMcpServer(options => options.ToolCollection = [McpServerTool.Create(method)]) + .WithMcpApps(); + + using var sp = sc.BuildServiceProvider(); + var factory = sp.GetRequiredService>(); + + var tool1 = Assert.Single(factory.Create(Options.DefaultName).ToolCollection!); + var tool2 = Assert.Single(factory.Create(Options.DefaultName).ToolCollection!); + + Assert.NotSame(tool1, tool2); + Assert.NotNull(tool1.ProtocolTool.Meta?["ui"]); + Assert.NotNull(tool2.ProtocolTool.Meta?["ui"]); + } + [Fact] public void WithMcpApps_EmptyToolCollection_DoesNotThrow() { diff --git a/tests/ModelContextProtocol.Tests/Server/McpServerPrimitiveCollectionTests.cs b/tests/ModelContextProtocol.Tests/Server/McpServerPrimitiveCollectionTests.cs index ea13c93df..0e8b514c3 100644 --- a/tests/ModelContextProtocol.Tests/Server/McpServerPrimitiveCollectionTests.cs +++ b/tests/ModelContextProtocol.Tests/Server/McpServerPrimitiveCollectionTests.cs @@ -818,6 +818,282 @@ public async Task Enumeration_AddStalledMidway_LaterAddDoesNotOvertakeIt() Assert.Equal(whileStalled, afterward.Take(whileStalled.Length)); } + // ------------------------------------------------------------------------- + // Copying + // ------------------------------------------------------------------------- + + [Fact] + public void Copy_NullCollection_Throws() + { + Assert.Throws("collection", () => new McpServerPrimitiveCollection((McpServerPrimitiveCollection)null!)); + Assert.Throws("collection", () => new McpServerResourceCollection(null!)); + } + + [Fact] + public void Copy_ContainsSamePrimitivesInSameOrder() + { + var source = new McpServerPrimitiveCollection(); + McpServerTool[] tools = [CreateTool("c"), CreateTool("a"), CreateTool("b")]; + foreach (var tool in tools) + { + source.Add(tool); + } + + var copy = new McpServerPrimitiveCollection(source); + + Assert.Equal(3, copy.Count); + Assert.False(copy.IsEmpty); + Assert.Equal(tools, copy); + Assert.Equal(tools, copy.ToArray()); + Assert.Equal(["c", "a", "b"], copy.PrimitiveNames); + Assert.Same(tools[1], copy["a"]); + Assert.True(copy.TryGetPrimitive("b", out var b)); + Assert.Same(tools[2], b); + Assert.Contains(tools[0], copy); + Assert.False(copy.TryGetPrimitive("missing", out _)); + Assert.Throws(() => copy["missing"]); + + var array = new McpServerTool[4]; + copy.CopyTo(array, 1); + Assert.Null(array[0]); + Assert.Equal(tools, array.Skip(1)); + } + + [Fact] + public void Copy_OfEmptyCollection_IsEmpty() + { + var copy = new McpServerPrimitiveCollection(new McpServerPrimitiveCollection()); + + Assert.True(copy.IsEmpty); + Assert.Empty(copy); + Assert.True(copy.TryAdd(CreateTool("tool1"))); + Assert.Single(copy); + } + + [Fact] + public void Copy_MutatingCopy_DoesNotAffectSourceOrOtherCopies() + { + var source = new McpServerPrimitiveCollection(); + var shared = CreateTool("shared"); + source.Add(shared); + source.Add(CreateTool("other")); + + int sourceChanges = 0; + source.Changed += (_, _) => sourceChanges++; + + var copy1 = new McpServerPrimitiveCollection(source); + var copy2 = new McpServerPrimitiveCollection(source); + int copy1Changes = 0; + copy1.Changed += (_, _) => copy1Changes++; + + Assert.True(copy1.Remove(shared)); + Assert.True(copy1.TryAdd(CreateTool("added"))); + + Assert.Equal(["other", "added"], copy1.PrimitiveNames); + Assert.Equal(2, copy1Changes); + Assert.Equal(["shared", "other"], source.PrimitiveNames); + Assert.Equal(["shared", "other"], copy2.PrimitiveNames); + Assert.Same(shared, copy2["shared"]); + Assert.Equal(0, sourceChanges); + } + + [Fact] + public void Copy_MutatingSource_DoesNotAffectExistingCopies_AndLaterCopiesSeeTheMutation() + { + var source = new McpServerPrimitiveCollection(); + source.Add(CreateTool("tool1")); + + var before = new McpServerPrimitiveCollection(source); + source.Add(CreateTool("tool2")); + var after = new McpServerPrimitiveCollection(source); + source.Clear(); + + Assert.Equal(["tool1"], before.PrimitiveNames); + Assert.Equal(["tool1", "tool2"], after.PrimitiveNames); + Assert.Empty(source); + } + + [Fact] + public void Copy_ClearCopy_FiresChangedAndLeavesSourceIntact() + { + var source = new McpServerPrimitiveCollection(); + source.Add(CreateTool("tool1")); + var copy = new McpServerPrimitiveCollection(source); + int changeCount = 0; + copy.Changed += (_, _) => changeCount++; + + copy.Clear(); + + Assert.Equal(1, changeCount); + Assert.Empty(copy); + Assert.Single(source); + Assert.True(copy.TryAdd(CreateTool("tool1"))); + Assert.Single(copy); + } + + [Fact] + public void Copy_DuplicateTryAddAndFailedRemove_DoNotFireChanged() + { + var source = new McpServerPrimitiveCollection(); + source.Add(CreateTool("tool1")); + var copy = new McpServerPrimitiveCollection(source); + int changeCount = 0; + copy.Changed += (_, _) => changeCount++; + + Assert.False(copy.TryAdd(CreateTool("tool1"))); + Assert.False(copy.Remove(CreateTool("tool1"))); + Assert.False(copy.Remove(CreateTool("missing"))); + + Assert.Equal(0, changeCount); + Assert.Single(copy); + } + + [Fact] + public void Copy_DeferChangedEvents_CoalescesMutationsOfCopy() + { + var source = new McpServerPrimitiveCollection(); + source.Add(CreateTool("tool1")); + var copy = new McpServerPrimitiveCollection(source); + int changeCount = 0; + copy.Changed += (_, _) => changeCount++; + + using (copy.DeferChangedEvents()) + { + copy.Add(CreateTool("tool2")); + copy.Remove(copy["tool1"]); + Assert.Equal(0, changeCount); + } + + Assert.Equal(1, changeCount); + Assert.Equal(["tool2"], copy.PrimitiveNames); + } + + [Fact] + public void Copy_OfCopy_IsIndependent() + { + var source = new McpServerPrimitiveCollection(); + source.Add(CreateTool("tool1")); + var copy = new McpServerPrimitiveCollection(source); + var copyOfCopy = new McpServerPrimitiveCollection(copy); + + copy.Add(CreateTool("tool2")); + copyOfCopy.Add(CreateTool("tool3")); + + Assert.Equal(["tool1"], source.PrimitiveNames); + Assert.Equal(["tool1", "tool2"], copy.PrimitiveNames); + Assert.Equal(["tool1", "tool3"], copyOfCopy.PrimitiveNames); + } + + [Fact] + public void Copy_PreservesNamesPrimitivesWereAddedUnder() + { + var source = new McpServerPrimitiveCollection(); + var tool = CreateTool("original"); + source.Add(tool); + tool.ProtocolTool.Name = "renamed"; + + var copy = new McpServerPrimitiveCollection(source); + + Assert.Equal(["original"], copy.PrimitiveNames); + Assert.Same(tool, copy["original"]); + } + + [Fact] + public void Copy_UsesSourceKeyComparer() + { + var source = new McpServerPrimitiveCollection(StringComparer.OrdinalIgnoreCase); + var tool = CreateTool("Tool"); + source.Add(tool); + + var copy = new McpServerPrimitiveCollection(source); + + Assert.Same(tool, copy["TOOL"]); + Assert.False(copy.TryAdd(CreateTool("tool"))); + Assert.True(copy.Remove(tool)); + Assert.True(copy.TryAdd(CreateTool("tool"))); + } + + [Fact] + public void Copy_ResourceCollection_PreservesUriEquality() + { + var source = new McpServerResourceCollection(); + var resource = McpServerResource.Create(() => "content", new() { UriTemplate = "https://example.com/resource", Name = "resource" }); + var template = McpServerResource.Create((string id) => id, new() { UriTemplate = "test://items/{id}", Name = "template" }); + source.Add(resource); + source.Add(template); + + var copy = new McpServerResourceCollection(source); + + Assert.True(copy.TryGetPrimitive("HTTPS://EXAMPLE.COM/resource", out var found)); + Assert.Same(resource, found); + Assert.Same(template, copy["test://items/{id}"]); + Assert.False(copy.TryGetPrimitive("test://ITEMS/{id}", out _)); + Assert.False(copy.TryAdd(McpServerResource.Create(() => "other", new() { UriTemplate = "https://EXAMPLE.com/resource", Name = "dup" }))); + + Assert.True(copy.Remove(resource)); + Assert.False(copy.TryGetPrimitive("https://example.com/resource", out _)); + Assert.Same(resource, source["https://example.com/resource"]); + } + + [Fact] + public async Task Copy_ConcurrentMutationOfCopies_IsIsolated() + { + var source = new McpServerPrimitiveCollection(); + for (int i = 0; i < 20; i++) + { + source.Add(CreateTool($"tool{i}")); + } + + var tasks = Enumerable.Range(0, 8).Select(worker => Task.Run(() => + { + for (int iteration = 0; iteration < 50; iteration++) + { + var copy = new McpServerPrimitiveCollection(source); + Assert.True(copy.Remove(copy[$"tool{iteration % 20}"])); + Assert.True(copy.TryAdd(CreateTool($"worker{worker}"))); + Assert.Equal(20, copy.Count); + Assert.Equal(20, copy.ToArray().Length); + Assert.False(copy.TryGetPrimitive($"tool{iteration % 20}", out _)); + } + }, TestContext.Current.CancellationToken)); + + await Task.WhenAll(tasks); + + Assert.Equal(Enumerable.Range(0, 20).Select(i => $"tool{i}"), source.PrimitiveNames); + } + + [Fact] + public async Task Copy_ConcurrentReadsDuringFirstMutation_SeeConsistentContents() + { + var source = new McpServerPrimitiveCollection(); + for (int i = 0; i < 50; i++) + { + source.Add(CreateTool($"tool{i}")); + } + + for (int iteration = 0; iteration < 20; iteration++) + { + var copy = new McpServerPrimitiveCollection(source); + using var start = new ManualResetEventSlim(); + var readers = Enumerable.Range(0, 4).Select(reader => Task.Run(() => + { + start.Wait(TestContext.Current.CancellationToken); + for (int i = 0; i < 200; i++) + { + Assert.True(copy.TryGetPrimitive("tool0", out _)); + int count = copy.ToArray().Length; + Assert.InRange(count, 50, 51); + } + }, TestContext.Current.CancellationToken)).ToArray(); + + start.Set(); + copy.Add(CreateTool("added")); + await Task.WhenAll(readers); + + Assert.Equal(51, copy.Count); + } + } + /// A primitive whose runs once before comparing by reference. private sealed class ReentrantEqualityPrimitive(string id) : IMcpServerPrimitive { diff --git a/tests/ModelContextProtocol.Tests/Server/McpServerTests.cs b/tests/ModelContextProtocol.Tests/Server/McpServerTests.cs index bc7221a27..8bb26cb65 100644 --- a/tests/ModelContextProtocol.Tests/Server/McpServerTests.cs +++ b/tests/ModelContextProtocol.Tests/Server/McpServerTests.cs @@ -520,6 +520,52 @@ await transport.SendMessageAsync(new JsonRpcRequest await runTask; } + [Fact] + public async Task Completion_AutoPopulated_FromPromptAllowedValues_SharedAcrossServers() + { + McpServerPrimitiveCollection prompts = [McpServerPrompt.Create( + ( + [System.ComponentModel.DataAnnotations.AllowedValues("dog", "cat", "cow")] string animal + ) => animal, + new McpServerPromptCreateOptions { Name = "test-prompt" })]; + + foreach (var (partial, expected) in new[] { ("c", new[] { "cat", "cow" }), ("d", new[] { "dog" }) }) + { + await using var transport = new TestServerTransport(); + var options = CreateOptions(); + options.PromptCollection = new(prompts); + + await using var server = McpServer.Create(transport, options, LoggerFactory); + var runTask = server.RunAsync(TestContext.Current.CancellationToken); + + var receivedMessage = new TaskCompletionSource(); + 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 = partial } + }, McpJsonUtilities.DefaultOptions) + }, TestContext.Current.CancellationToken); + + var response = await receivedMessage.Task.WaitAsync(TestConstants.DefaultTimeout, TestContext.Current.CancellationToken); + var result = JsonSerializer.Deserialize(response.Result, McpJsonUtilities.DefaultOptions); + Assert.NotNull(result?.Completion); + Assert.Equal(expected, result.Completion.Values); + + await transport.DisposeAsync(); + await runTask; + } + } + [Fact] public async Task Completion_AutoPopulated_FromPromptAllowedValues_NoMatch() {