diff --git a/src/ModelContextProtocol.Core/Server/McpServerPrimitiveCollection.cs b/src/ModelContextProtocol.Core/Server/McpServerPrimitiveCollection.cs index aa281989b..02174f249 100644 --- a/src/ModelContextProtocol.Core/Server/McpServerPrimitiveCollection.cs +++ b/src/ModelContextProtocol.Core/Server/McpServerPrimitiveCollection.cs @@ -6,11 +6,24 @@ namespace ModelContextProtocol.Server; /// Provides a thread-safe collection of instances, indexed by their names. /// The type of primitive stored in the collection. +/// +/// Enumerating the collection, including via , , and , +/// yields primitives in the order they were added. A primitive that is removed and then added again moves to the end. +/// public class McpServerPrimitiveCollection : ICollection, IReadOnlyCollection where T : IMcpServerPrimitive { /// Concurrent dictionary of primitives, indexed by their names. - private readonly ConcurrentDictionary _primitives; + private readonly ConcurrentDictionary _primitives; + + /// The entries in , in the order they were added. + private readonly List _orderedEntries = []; + + /// Lock serializing additions and removals across and . + private readonly object _mutationLock = new(); + + /// Cached copy of used for enumeration, or if it needs rebuilding. + private volatile Entry[]? _snapshot; /// Lock protecting and . private readonly object _deferralLock = new(); @@ -132,14 +145,20 @@ public T this[string name] get { Throw.IfNull(name); - return _primitives[name]; + return _primitives[name].Primitive; } } /// Clears all primitives from the collection. public virtual void Clear() { - _primitives.Clear(); + lock (_mutationLock) + { + _primitives.Clear(); + _orderedEntries.Clear(); + _snapshot = null; + } + RaiseChanged(); } @@ -163,13 +182,21 @@ public virtual bool TryAdd(T primitive) { Throw.IfNull(primitive); - bool added = _primitives.TryAdd(primitive.Id, primitive); - if (added) + string key = primitive.Id; + lock (_mutationLock) { - RaiseChanged(); + Entry entry = new(key, primitive); + if (!_primitives.TryAdd(key, entry)) + { + return false; + } + + _orderedEntries.Add(entry); + _snapshot = null; } - return added; + RaiseChanged(); + return true; } /// Removes the specified primitive from the collection. @@ -182,13 +209,28 @@ public virtual bool Remove(T primitive) { Throw.IfNull(primitive); - bool removed = ((ICollection>)_primitives).Remove(new(primitive.Id, primitive)); - if (removed) + string key = primitive.Id; + lock (_mutationLock) { - RaiseChanged(); + if (!_primitives.TryGetValue(key, out Entry? entry) || + !EqualityComparer.Default.Equals(entry.Primitive, primitive)) + { + return false; + } + + // 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))) + { + return false; + } + + _orderedEntries.Remove(entry); + _snapshot = null; } - return removed; + RaiseChanged(); + return true; } /// Attempts to get the primitive with the specified name from the collection. @@ -201,7 +243,15 @@ public virtual bool Remove(T primitive) public virtual bool TryGetPrimitive(string name, [NotNullWhen(true)] out T? primitive) { Throw.IfNull(name); - return _primitives.TryGetValue(name, out primitive); + + if (_primitives.TryGetValue(name, out Entry? entry)) + { + primitive = entry.Primitive; + return true; + } + + primitive = default; + return false; } /// Checks if a specific primitive is present in the collection of primitives. @@ -211,31 +261,71 @@ public virtual bool TryGetPrimitive(string name, [NotNullWhen(true)] out T? prim public virtual bool Contains(T primitive) { Throw.IfNull(primitive); - return ((ICollection>)_primitives).Contains(new(primitive.Id, primitive)); + return + _primitives.TryGetValue(primitive.Id, out Entry? entry) && + EqualityComparer.Default.Equals(entry.Primitive, primitive); } - /// Gets the names of all of the primitives in the collection. - public virtual ICollection PrimitiveNames => _primitives.Keys; + /// 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); /// Creates an array containing all of the primitives in the collection. - /// An array containing all of the primitives in the collection. - public virtual T[] ToArray() => _primitives.Values.ToArray(); + /// An array containing all of the primitives in the collection, in the order they were added. + public virtual T[] ToArray() => Array.ConvertAll(GetSnapshot(), static e => e.Primitive); /// public virtual void CopyTo(T[] array, int arrayIndex) { Throw.IfNull(array); + Throw.IfNegative(arrayIndex); - _primitives.Values.CopyTo(array, arrayIndex); + Entry[] entries = GetSnapshot(); + if (array.Length - arrayIndex < entries.Length) + { + throw new ArgumentException("The destination array is not long enough to copy all the primitives in the collection.", nameof(array)); + } + + for (int i = 0; i < entries.Length; i++) + { + array[arrayIndex + i] = entries[i].Primitive; + } } /// public virtual IEnumerator GetEnumerator() { - foreach (var entry in _primitives) + foreach (var entry in GetSnapshot()) + { + yield return entry.Primitive; + } + } + + /// Gets the collection's entries in insertion order, rebuilding the cached snapshot if a mutation invalidated it. + /// The returned array is shared between callers and must not be modified or exposed. + private Entry[] GetSnapshot() + { + Entry[]? snapshot = _snapshot; + if (snapshot is null) { - yield return entry.Value; + lock (_mutationLock) + { + snapshot = _snapshot ??= _orderedEntries.ToArray(); + } } + + return snapshot; + } + + /// 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) + { + public string Key { get; } = key; + + public T Primitive { get; } = primitive; } /// diff --git a/tests/ModelContextProtocol.Tests/Configuration/McpServerPrimitiveListOrderTests.cs b/tests/ModelContextProtocol.Tests/Configuration/McpServerPrimitiveListOrderTests.cs new file mode 100644 index 000000000..27bd619e0 --- /dev/null +++ b/tests/ModelContextProtocol.Tests/Configuration/McpServerPrimitiveListOrderTests.cs @@ -0,0 +1,117 @@ +using Microsoft.Extensions.DependencyInjection; +using ModelContextProtocol.Client; +using ModelContextProtocol.Server; + +namespace ModelContextProtocol.Tests.Configuration; + +public class McpServerPrimitiveListOrderTests : ClientServerTestBase +{ + public McpServerPrimitiveListOrderTests(ITestOutputHelper testOutputHelper) + : base(testOutputHelper) + { + } + + protected override void ConfigureServices(ServiceCollection services, IMcpServerBuilder mcpServerBuilder) + { + mcpServerBuilder + .WithTools() + .WithPrompts() + .WithResources(); + } + + [Fact] + public async Task ListTools_ReturnsToolsInDeclarationOrder() + { + await using McpClient client = await CreateMcpClientForServer(); + + var tools = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken); + + Assert.Equal( + ["zulu", "alpha", "mike", "echo", "yankee", "bravo", "kilo", "xray", "charlie", "lima"], + tools.Select(t => t.Name)); + } + + [Fact] + public async Task ListPrompts_ReturnsPromptsInDeclarationOrder() + { + await using McpClient client = await CreateMcpClientForServer(); + + var prompts = await client.ListPromptsAsync(cancellationToken: TestContext.Current.CancellationToken); + + Assert.Equal( + ["zulu", "alpha", "mike", "echo", "yankee", "bravo", "kilo", "xray", "charlie", "lima"], + prompts.Select(p => p.Name)); + } + + [Fact] + public async Task ListResources_ReturnsResourcesInDeclarationOrder() + { + await using McpClient client = await CreateMcpClientForServer(); + + var resources = await client.ListResourcesAsync(cancellationToken: TestContext.Current.CancellationToken); + + Assert.Equal( + ["test://zulu", "test://alpha", "test://mike", "test://echo", "test://yankee", "test://bravo", "test://kilo", "test://xray"], + resources.Select(r => r.Uri)); + } + + [Fact] + public async Task ListResourceTemplates_ReturnsTemplatesInDeclarationOrder() + { + await using McpClient client = await CreateMcpClientForServer(); + + var templates = await client.ListResourceTemplatesAsync(cancellationToken: TestContext.Current.CancellationToken); + + Assert.Equal( + ["test://zulu/{id}", "test://alpha/{id}", "test://mike/{id}", "test://echo/{id}", "test://yankee/{id}"], + templates.Select(t => t.UriTemplate)); + } + + [McpServerToolType] + public sealed class OrderedTools + { + [McpServerTool(Name = "zulu")] public static string Zulu() => "z"; + [McpServerTool(Name = "alpha")] public static string Alpha() => "a"; + [McpServerTool(Name = "mike")] public static string Mike() => "m"; + [McpServerTool(Name = "echo")] public static string Echo() => "e"; + [McpServerTool(Name = "yankee")] public static string Yankee() => "y"; + [McpServerTool(Name = "bravo")] public static string Bravo() => "b"; + [McpServerTool(Name = "kilo")] public static string Kilo() => "k"; + [McpServerTool(Name = "xray")] public static string Xray() => "x"; + [McpServerTool(Name = "charlie")] public static string Charlie() => "c"; + [McpServerTool(Name = "lima")] public static string Lima() => "l"; + } + + [McpServerPromptType] + public sealed class OrderedPrompts + { + [McpServerPrompt(Name = "zulu")] public static string Zulu() => "z"; + [McpServerPrompt(Name = "alpha")] public static string Alpha() => "a"; + [McpServerPrompt(Name = "mike")] public static string Mike() => "m"; + [McpServerPrompt(Name = "echo")] public static string Echo() => "e"; + [McpServerPrompt(Name = "yankee")] public static string Yankee() => "y"; + [McpServerPrompt(Name = "bravo")] public static string Bravo() => "b"; + [McpServerPrompt(Name = "kilo")] public static string Kilo() => "k"; + [McpServerPrompt(Name = "xray")] public static string Xray() => "x"; + [McpServerPrompt(Name = "charlie")] public static string Charlie() => "c"; + [McpServerPrompt(Name = "lima")] public static string Lima() => "l"; + } + + [McpServerResourceType] + public sealed class OrderedResources + { + [McpServerResource(UriTemplate = "test://zulu")] public static string Zulu() => "z"; + [McpServerResource(UriTemplate = "test://zulu/{id}")] public static string ZuluItem(string id) => id; + [McpServerResource(UriTemplate = "test://alpha")] public static string Alpha() => "a"; + [McpServerResource(UriTemplate = "test://alpha/{id}")] public static string AlphaItem(string id) => id; + [McpServerResource(UriTemplate = "test://mike")] public static string Mike() => "m"; + [McpServerResource(UriTemplate = "test://mike/{id}")] public static string MikeItem(string id) => id; + [McpServerResource(UriTemplate = "test://echo")] public static string Echo() => "e"; + [McpServerResource(UriTemplate = "test://echo/{id}")] public static string EchoItem(string id) => id; + [McpServerResource(UriTemplate = "test://yankee")] public static string Yankee() => "y"; + [McpServerResource(UriTemplate = "test://yankee/{id}")] public static string YankeeItem(string id) => id; + [McpServerResource(UriTemplate = "test://bravo")] public static string Bravo() => "b"; + [McpServerResource(UriTemplate = "test://kilo")] public static string Kilo() => "k"; + [McpServerResource(UriTemplate = "test://xray")] public static string Xray() => "x"; + } +} diff --git a/tests/ModelContextProtocol.Tests/Server/McpServerPrimitiveCollectionTests.cs b/tests/ModelContextProtocol.Tests/Server/McpServerPrimitiveCollectionTests.cs index 7acac660e..ea13c93df 100644 --- a/tests/ModelContextProtocol.Tests/Server/McpServerPrimitiveCollectionTests.cs +++ b/tests/ModelContextProtocol.Tests/Server/McpServerPrimitiveCollectionTests.cs @@ -1,5 +1,7 @@ using Microsoft.Extensions.AI; using ModelContextProtocol.Server; +using ModelContextProtocol.Tests.Utils; +using System.Runtime.CompilerServices; namespace ModelContextProtocol.Tests.Server; @@ -592,4 +594,300 @@ public void DeferChangedEvents_PromptCollection_MultipleMutations_FiresOneChange Assert.Equal(1, changeCount); Assert.Equal(3, collection.Count); } + + // ------------------------------------------------------------------------- + // Enumeration order + // ------------------------------------------------------------------------- + + private static readonly string[] s_unorderedNames = + [ + "zulu", "alpha", "mike", "echo", "yankee", "bravo", "kilo", "xray", "charlie", "lima", + "whiskey", "delta", "november", "victor", "foxtrot", "oscar", "uniform", "golf", "papa", "tango", + ]; + + private static McpServerResource CreateResource(string uri) => + McpServerResource.Create(() => uri, new() { UriTemplate = uri, Name = uri }); + + [Fact] + public void Enumeration_YieldsPrimitivesInInsertionOrder() + { + var collection = new McpServerPrimitiveCollection(); + foreach (var name in s_unorderedNames) + { + collection.Add(CreateTool(name)); + } + + Assert.Equal(s_unorderedNames, collection.Select(t => t.ProtocolTool.Name)); + Assert.Equal(s_unorderedNames, collection.ToArray().Select(t => t.ProtocolTool.Name)); + Assert.Equal(s_unorderedNames, collection.PrimitiveNames); + + var copy = new McpServerTool[collection.Count + 1]; + collection.CopyTo(copy, 1); + Assert.Null(copy[0]); + Assert.Equal(s_unorderedNames, copy.Skip(1).Select(t => t.ProtocolTool.Name)); + } + + [Fact] + public void Enumeration_PromptCollection_YieldsPromptsInInsertionOrder() + { + var collection = new McpServerPrimitiveCollection(); + foreach (var name in s_unorderedNames) + { + collection.Add(CreatePrompt(name)); + } + + Assert.Equal(s_unorderedNames, collection.Select(p => p.ProtocolPrompt.Name)); + } + + [Fact] + public void Enumeration_ResourceCollection_YieldsResourcesInInsertionOrder() + { + string[] uris = [.. s_unorderedNames.Select(name => $"test://{name}")]; + + var collection = new McpServerResourceCollection(); + foreach (var uri in uris) + { + collection.Add(CreateResource(uri)); + } + + Assert.Equal(uris, collection.Select(r => r.ProtocolResourceTemplate.UriTemplate)); + Assert.Equal(uris, collection.PrimitiveNames); + } + + [Fact] + public void Enumeration_RemovedAndReaddedPrimitive_MovesToEnd() + { + var collection = new McpServerPrimitiveCollection(); + var first = CreateTool("first"); + collection.Add(first); + collection.Add(CreateTool("second")); + collection.Add(CreateTool("third")); + + Assert.True(collection.Remove(first)); + collection.Add(CreateTool("first")); + + Assert.Equal(["second", "third", "first"], collection.Select(t => t.ProtocolTool.Name)); + } + + [Fact] + public void Enumeration_AfterClear_UsesNewInsertionOrder() + { + var collection = new McpServerPrimitiveCollection(); + collection.Add(CreateTool("alpha")); + collection.Add(CreateTool("bravo")); + + collection.Clear(); + collection.Add(CreateTool("bravo")); + collection.Add(CreateTool("alpha")); + + Assert.Equal(["bravo", "alpha"], collection.Select(t => t.ProtocolTool.Name)); + } + + [Fact] + public void Remove_DifferentInstanceWithSameName_ReturnsFalse() + { + var collection = new McpServerPrimitiveCollection(); + var stored = CreateTool("tool1"); + collection.Add(stored); + + var other = CreateTool("tool1"); + + Assert.DoesNotContain(other, collection); + Assert.False(collection.Remove(other)); + Assert.Contains(stored, collection); + Assert.Same(stored, collection["tool1"]); + } + + [Fact] + public void PrimitiveNames_PrimitiveRenamedAfterAdd_ReturnsNameItWasAddedUnder() + { + var collection = new McpServerPrimitiveCollection(); + var tool = CreateTool("original"); + collection.Add(tool); + + tool.ProtocolTool.Name = "renamed"; + + Assert.Equal(["original"], collection.PrimitiveNames); + Assert.True(collection.TryGetPrimitive("original", out var found)); + Assert.Same(tool, found); + } + + [Fact] + public void ToArrayAndPrimitiveNames_ReturnCopiesTheCallerOwns() + { + var collection = new McpServerPrimitiveCollection(); + collection.Add(CreateTool("alpha")); + collection.Add(CreateTool("bravo")); + + var array = collection.ToArray(); + array[0] = array[1]; + var names = collection.PrimitiveNames; + ((string[])names)[0] = "changed"; + + Assert.Equal(["alpha", "bravo"], collection.Select(t => t.ProtocolTool.Name)); + Assert.Equal(["alpha", "bravo"], collection.PrimitiveNames); + Assert.NotSame(collection.ToArray(), collection.ToArray()); + } + + [Fact] + public void CopyTo_DestinationTooSmall_Throws() + { + var collection = new McpServerPrimitiveCollection(); + collection.Add(CreateTool("alpha")); + collection.Add(CreateTool("bravo")); + + Assert.Throws(() => collection.CopyTo(new McpServerTool[2], 1)); + Assert.Throws(() => collection.CopyTo(new McpServerTool[2], -1)); + } + + [Fact] + public void Remove_EqualsReplacesEntryUnderSameKey_LeavesReplacementInPlace() + { + var collection = new McpServerPrimitiveCollection(); + var original = new ReentrantEqualityPrimitive("tool1"); + var replacement = new ReentrantEqualityPrimitive("tool1"); + collection.Add(original); + + original.OnEquals = () => + { + collection.Clear(); + collection.Add(replacement); + }; + + Assert.False(collection.Remove(original)); + Assert.False(collection.IsEmpty); + Assert.Same(replacement, collection["tool1"]); + Assert.Same(replacement, Assert.Single(collection.ToArray())); + } + + [Fact] + public void Enumeration_AlreadyVisiblePrimitive_KeepsItsPosition() + { + // The comparer reenters the collection on the same thread before "first" is inserted, which deterministically + // forces the interleaving where a later add becomes visible while an earlier add is still in progress. + var comparer = new InsertionInterleavingComparer(); + var collection = new McpServerPrimitiveCollection(comparer); + string[] whileFirstPending = []; + + comparer.BeforeFirstInsertion = () => + { + collection.Add(CreateTool("second")); + whileFirstPending = [.. collection.Select(t => t.ProtocolTool.Name)]; + }; + + collection.Add(CreateTool("first")); + string[] afterward = [.. collection.Select(t => t.ProtocolTool.Name)]; + + Assert.Equal(["second"], whileFirstPending); + Assert.Equal(["second", "first"], afterward); + } + + [Fact] + public async Task Enumeration_AddStalledMidway_LaterAddDoesNotOvertakeIt() + { + // Best-effort threaded coverage: the second add is given a window to complete while the first is stalled, but + // isn't guaranteed to use it. Enumeration_AlreadyVisiblePrimitive_KeepsItsPosition covers the interleaving + // deterministically. + CancellationToken cancellationToken = TestContext.Current.CancellationToken; + using var comparer = new StallingComparer("first", cancellationToken); + var collection = new McpServerPrimitiveCollection(comparer); + + // Enumerate once so listing while an add is stalled reads the cached snapshot rather than waiting for the add. + Assert.Empty(collection); + + Task addFirst = Task.Run(() => collection.Add(CreateTool("first")), cancellationToken); + Task addSecond = Task.CompletedTask; + string[] whileStalled; + try + { + Assert.True(comparer.Stalled.Wait(TestConstants.DefaultTimeout, cancellationToken)); + + addSecond = Task.Run(() => collection.Add(CreateTool("second")), cancellationToken); + await Task.WhenAny(addSecond, Task.Delay(TimeSpan.FromMilliseconds(200), cancellationToken)); + whileStalled = [.. collection.Select(t => t.ProtocolTool.Name)]; + } + finally + { + comparer.Release.Set(); + await Task.WhenAll(addFirst, addSecond); + } + + string[] afterward = [.. collection.Select(t => t.ProtocolTool.Name)]; + + // A primitive that was already listed must not be reordered by an add that started before it. + Assert.Equal(whileStalled, afterward.Take(whileStalled.Length)); + } + + /// A primitive whose runs once before comparing by reference. + private sealed class ReentrantEqualityPrimitive(string id) : IMcpServerPrimitive + { + public string Id => id; + + public IReadOnlyList Metadata => []; + + public Action? OnEquals { get; set; } + + public override bool Equals(object? obj) + { + Action? onEquals = OnEquals; + OnEquals = null; + onEquals?.Invoke(); + return ReferenceEquals(this, obj); + } + + public override int GetHashCode() => RuntimeHelpers.GetHashCode(this); + } + + /// A comparer that runs once, the first time it hashes "first". + private sealed class InsertionInterleavingComparer : IEqualityComparer + { + private bool _callbackInvoked; + + public Action? BeforeFirstInsertion { get; set; } + + public bool Equals(string? x, string? y) => StringComparer.Ordinal.Equals(x, y); + + public int GetHashCode(string value) + { + if (value == "first" && !_callbackInvoked) + { + _callbackInvoked = true; + BeforeFirstInsertion?.Invoke(); + } + + return StringComparer.Ordinal.GetHashCode(value); + } + } + + /// A comparer that blocks the first time it hashes stallKey, until is set. + private sealed class StallingComparer(string stallKey, CancellationToken cancellationToken) : IEqualityComparer, IDisposable + { + private int _hasStalled; + + public ManualResetEventSlim Stalled { get; } = new(); + + public ManualResetEventSlim Release { get; } = new(); + + public bool Equals(string? x, string? y) => StringComparer.Ordinal.Equals(x, y); + + public int GetHashCode(string obj) + { + if (obj == stallKey && Interlocked.Exchange(ref _hasStalled, 1) == 0) + { + Stalled.Set(); + if (!Release.Wait(TestConstants.DefaultTimeout, cancellationToken)) + { + throw new TimeoutException($"The stalled add of '{stallKey}' was never released."); + } + } + + return StringComparer.Ordinal.GetHashCode(obj); + } + + public void Dispose() + { + Stalled.Dispose(); + Release.Dispose(); + } + } }