Preserve insertion order when listing server primitives

src/ModelContextProtocol.Core/Server/McpServerPrimitiveCollection.cs

@@ -6,11 +6,24 @@ namespace ModelContextProtocol.Server;
 
 /// <summary>Provides a thread-safe collection of <typeparamref name="T"/> instances, indexed by their names.</summary>
 /// <typeparam name="T">The type of primitive stored in the collection.</typeparam>
+/// <remarks>
+/// Enumerating the collection, including via <see cref="ToArray"/>, <see cref="CopyTo"/>, and <see cref="PrimitiveNames"/>,
+/// yields primitives in the order they were added. A primitive that is removed and then added again moves to the end.
+/// </remarks>
 public class McpServerPrimitiveCollection<T> : ICollection<T>, IReadOnlyCollection<T>
     where T : IMcpServerPrimitive
 {
     /// <summary>Concurrent dictionary of primitives, indexed by their names.</summary>
-    private readonly ConcurrentDictionary<string, T> _primitives;
+    private readonly ConcurrentDictionary<string, Entry> _primitives;
+
+    /// <summary>The entries in <see cref="_primitives"/>, in the order they were added.</summary>
+    private readonly List<Entry> _orderedEntries = [];
+
+    /// <summary>Lock serializing additions and removals across <see cref="_primitives"/> and <see cref="_orderedEntries"/>.</summary>
+    private readonly object _mutationLock = new();
+
+    /// <summary>Cached copy of <see cref="_orderedEntries"/> used for enumeration, or <see langword="null"/> if it needs rebuilding.</summary>
+    private volatile Entry[]? _snapshot;
 
     /// <summary>Lock protecting <see cref="_activeDeferralScopes"/> and <see cref="_hasDeferredChangeEvents"/>.</summary>
     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;
         }
     }
 
     /// <summary>Clears all primitives from the collection.</summary>
     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;
     }
 
     /// <summary>Removes the specified primitive from the collection.</summary>
@@ -182,13 +209,28 @@ public virtual bool Remove(T primitive)
     {
         Throw.IfNull(primitive);
 
-        bool removed = ((ICollection<KeyValuePair<string, T>>)_primitives).Remove(new(primitive.Id, primitive));
-        if (removed)
+        string key = primitive.Id;
+        lock (_mutationLock)
         {
-            RaiseChanged();
+            if (!_primitives.TryGetValue(key, out Entry? entry) ||
+                !EqualityComparer<T>.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<KeyValuePair<string, Entry>>)_primitives).Remove(new(key, entry)))
+            {
+                return false;
+            }
+
+            _orderedEntries.Remove(entry);
+            _snapshot = null;
         }
 
-        return removed;
+        RaiseChanged();
+        return true;
     }
 
     /// <summary>Attempts to get the primitive with the specified name from the collection.</summary>
@@ -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;
     }
 
     /// <summary>Checks if a specific primitive is present in the collection of primitives.</summary>
@@ -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<KeyValuePair<string, T>>)_primitives).Contains(new(primitive.Id, primitive));
+        return
+            _primitives.TryGetValue(primitive.Id, out Entry? entry) &&
+            EqualityComparer<T>.Default.Equals(entry.Primitive, primitive);
     }
 
-    /// <summary>Gets the names of all of the primitives in the collection.</summary>
-    public virtual ICollection<string> PrimitiveNames => _primitives.Keys;
+    /// <summary>Gets the names of all of the primitives in the collection, in the order the primitives were added.</summary>
+    public virtual ICollection<string> PrimitiveNames => Array.ConvertAll(GetSnapshot(), static e => e.Key);
 
     /// <summary>Creates an array containing all of the primitives in the collection.</summary>
-    /// <returns>An array containing all of the primitives in the collection.</returns>
-    public virtual T[] ToArray() => _primitives.Values.ToArray();
+    /// <returns>An array containing all of the primitives in the collection, in the order they were added.</returns>
+    public virtual T[] ToArray() => Array.ConvertAll(GetSnapshot(), static e => e.Primitive);
 
     /// <inheritdoc/>
     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;
+        }
     }
 
     /// <inheritdoc/>
     public virtual IEnumerator<T> GetEnumerator()
     {
-        foreach (var entry in _primitives)
+        foreach (var entry in GetSnapshot())
+        {
+            yield return entry.Primitive;
+        }
+    }
+
+    /// <summary>Gets the collection's entries in insertion order, rebuilding the cached snapshot if a mutation invalidated it.</summary>
+    /// <remarks>The returned array is shared between callers and must not be modified or exposed.</remarks>
+    private Entry[] GetSnapshot()
+    {
+        Entry[]? snapshot = _snapshot;
+        if (snapshot is null)
         {
-            yield return entry.Value;
+            lock (_mutationLock)
+            {
+                snapshot = _snapshot ??= _orderedEntries.ToArray();
+            }
         }
+
+        return snapshot;
+    }
+
+    /// <summary>A primitive together with the key it was added under.</summary>
+    /// <remarks>
+    /// The key is captured when the primitive is added, because <see cref="IMcpServerPrimitive.Id"/> may be backed by
+    /// mutable protocol properties.
+    /// </remarks>
+    private sealed class Entry(string key, T primitive)
+    {
+        public string Key { get; } = key;
+
+        public T Primitive { get; } = primitive;
     }
 
     /// <inheritdoc/>

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<OrderedTools>()
+            .WithPrompts<OrderedPrompts>()
+            .WithResources<OrderedResources>();
+    }
+
+    [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";
+    }
+}