Upgrade MCP SDK to 2.2.0 and tie location subscriptions to the session lifetime

CosmosDBShell.Tests/McpLocationSubscriptionTests.cs

@@ -4,13 +4,15 @@
 
 namespace CosmosShell.Tests;
 
-using System.Net;
-using System.Net.Sockets;
 using System.Text.Json;
 using Azure.Data.Cosmos.Shell.Core;
 using Azure.Data.Cosmos.Shell.Mcp;
 using Azure.Data.Cosmos.Shell.States;
+using Microsoft.AspNetCore.Hosting.Server;
+using Microsoft.AspNetCore.Hosting.Server.Features;
 using Microsoft.Azure.Cosmos;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Hosting;
 using ModelContextProtocol;
 using ModelContextProtocol.Client;
 
@@ -22,20 +24,11 @@ public async Task SubscribedClient_ReceivesInteractiveLocationChange()
     {
         using var timeout = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken);
         timeout.CancelAfter(TimeSpan.FromSeconds(10));
-        using var listener = new TcpListener(IPAddress.Loopback, 0);
-        listener.Start();
-        var port = ((IPEndPoint)listener.LocalEndpoint).Port;
-        listener.Stop();
-
-        using var host = McpServer.CreateHost(new Program.CosmosShellOptions { McpPort = port });
+        using var host = McpServer.CreateHost(new Program.CosmosShellOptions { McpPort = 0 });
         await host.StartAsync(timeout.Token);
         try
         {
-            var transport = new HttpClientTransport(new HttpClientTransportOptions
-            {
-                Endpoint = new Uri($"http://127.0.0.1:{port}/"),
-            });
-            await using var client = await McpClient.CreateAsync(transport, cancellationToken: timeout.Token);
+            await using var client = await ConnectAsync(host, timeout.Token);
             var resources = await client.ListResourcesAsync(cancellationToken: timeout.Token);
             Assert.Contains(resources, resource => resource.Uri == ResourceOperations.CurrentLocationUri);
 
@@ -61,6 +54,9 @@ public async Task SubscribedClient_ReceivesInteractiveLocationChange()
                 new CosmosClientOptions { ConnectionMode = ConnectionMode.Gateway });
             try
             {
+                // The subscription must not depend on per-request objects that the GC can reclaim.
+                GC.Collect();
+                GC.WaitForPendingFinalizers();
                 ShellInterpreter.Instance.State = new DatabaseState("McpNotificationTest", cosmosClient);
                 Assert.Equal(ResourceOperations.CurrentLocationUri, await updated.Task.WaitAsync(timeout.Token));
 
@@ -80,4 +76,54 @@ public async Task SubscribedClient_ReceivesInteractiveLocationChange()
             await host.StopAsync(TestContext.Current.CancellationToken);
         }
     }
+
+    [Fact]
+    public async Task EndedSession_RemovesLocationSubscription()
+    {
+        using var timeout = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken);
+        timeout.CancelAfter(TimeSpan.FromSeconds(10));
+        using var host = McpServer.CreateHost(new Program.CosmosShellOptions { McpPort = 0 });
+        await host.StartAsync(timeout.Token);
+        try
+        {
+            var subscriptions = host.Services.GetRequiredService<LocationResourceSubscriptions>();
+            var client = await ConnectAsync(host, timeout.Token);
+            await client.SubscribeToResourceAsync(ResourceOperations.CurrentLocationUri, cancellationToken: timeout.Token);
+            Assert.Equal(1, subscriptions.SubscriberCount);
+
+            // Disposing the client ends the session with DELETE, which must release the subscription.
+            await client.DisposeAsync();
+            while (subscriptions.SubscriberCount != 0)
+            {
+                await Task.Delay(20, timeout.Token);
+            }
+        }
+        finally
+        {
+            await host.StopAsync(TestContext.Current.CancellationToken);
+        }
+    }
+
+    [Fact]
+    public void HttpTransport_UsesStatefulSessionsWithBoundedIdleTimeout()
+    {
+        var options = new ModelContextProtocol.AspNetCore.HttpServerTransportOptions();
+        McpServer.ConfigureHttpTransport(options);
+
+        Assert.Equal(ModelContextProtocol.AspNetCore.HttpServerSessionMode.Stateful, options.SessionMode);
+#pragma warning disable MCP9006, MCPEXP002
+        Assert.Equal(McpServer.SessionIdleTimeout, options.IdleTimeout);
+        Assert.NotNull(options.RunSessionHandler);
+#pragma warning restore MCP9006, MCPEXP002
+    }
+
+    private static async Task<McpClient> ConnectAsync(IHost host, CancellationToken cancellationToken)
+    {
+        var address = host.Services.GetRequiredService<IServer>().Features.Get<IServerAddressesFeature>()!.Addresses.Single();
+        var transport = new HttpClientTransport(new HttpClientTransportOptions
+        {
+            Endpoint = new Uri(address.TrimEnd('/') + "/"),
+        });
+        return await McpClient.CreateAsync(transport, cancellationToken: cancellationToken);
+    }
 }

CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/LocationResourceSubscriptions.cs

@@ -16,7 +16,11 @@ internal sealed class LocationResourceSubscriptions : BackgroundService
 {
     private readonly object sync = new();
 
-    private readonly List<WeakReference<ModelContextProtocol.Server.McpServer>> subscribers = [];
+    // Session servers keyed by session ID. Entries are added when a session starts running and
+    // removed when it ends (DELETE, idle timeout, or shutdown), so no dead session is retained.
+    private readonly Dictionary<string, ModelContextProtocol.Server.McpServer> sessions = new(StringComparer.Ordinal);
+
+    private readonly HashSet<string> subscribedSessionIds = new(StringComparer.Ordinal);
 
     // Notifications carry only the URI, so pending changes are coalesced into one.
     private readonly Channel<bool> changes = Channel.CreateBounded<bool>(
@@ -30,25 +34,81 @@ public LocationResourceSubscriptions(ILogger<LocationResourceSubscriptions> logg
         ShellInterpreter.Instance.LocationChanged += this.OnLocationChanged;
     }
 
-    public void Subscribe(ModelContextProtocol.Server.McpServer server, string uri)
+    internal int SubscriberCount
+    {
+        get
+        {
+            lock (this.sync)
+            {
+                return this.subscribedSessionIds.Count;
+            }
+        }
+    }
+
+    /// <summary>
+    /// Runs an MCP session and keeps it available for notifications until the session ends.
+    /// </summary>
+    public async Task RunSessionAsync(ModelContextProtocol.Server.McpServer server, CancellationToken cancellationToken)
+    {
+        var sessionId = server.SessionId;
+        if (sessionId is not null)
+        {
+            lock (this.sync)
+            {
+                this.sessions[sessionId] = server;
+            }
+        }
+
+        try
+        {
+            await server.RunAsync(cancellationToken);
+        }
+        finally
+        {
+            if (sessionId is not null)
+            {
+                lock (this.sync)
+                {
+                    if (this.sessions.TryGetValue(sessionId, out var current) && ReferenceEquals(current, server))
+                    {
+                        this.sessions.Remove(sessionId);
+                        this.subscribedSessionIds.Remove(sessionId);
+                    }
+                }
+            }
+        }
+    }
+
+    public void Subscribe(string? sessionId, string uri)
     {
         ValidateUri(uri);
+        if (sessionId is null)
+        {
+            throw new McpProtocolException(
+                "Resource subscriptions require a stateful MCP session.",
+                McpErrorCode.InvalidRequest);
+        }
+
         lock (this.sync)
         {
-            this.PruneSubscribers();
-            if (!this.subscribers.Any(reference => reference.TryGetTarget(out var target) && ReferenceEquals(target, server)))
+            if (this.sessions.ContainsKey(sessionId))
             {
-                this.subscribers.Add(new WeakReference<ModelContextProtocol.Server.McpServer>(server));
+                this.subscribedSessionIds.Add(sessionId);
             }
         }
     }
 
-    public void Unsubscribe(ModelContextProtocol.Server.McpServer server, string uri)
+    public void Unsubscribe(string? sessionId, string uri)
     {
         ValidateUri(uri);
+        if (sessionId is null)
+        {
+            return;
+        }
+
         lock (this.sync)
         {
-            this.subscribers.RemoveAll(reference => !reference.TryGetTarget(out var target) || ReferenceEquals(target, server));
+            this.subscribedSessionIds.Remove(sessionId);
         }
     }
 
@@ -59,9 +119,8 @@ protected override async Task ExecuteAsync(CancellationToken stoppingToken)
             ModelContextProtocol.Server.McpServer[] servers;
             lock (this.sync)
             {
-                this.PruneSubscribers();
-                servers = this.subscribers
-                    .Select(reference => reference.TryGetTarget(out var server) ? server : null)
+                servers = this.subscribedSessionIds
+                    .Select(sessionId => this.sessions.TryGetValue(sessionId, out var server) ? server : null)
                     .OfType<ModelContextProtocol.Server.McpServer>()
                     .ToArray();
             }
@@ -81,11 +140,8 @@ await server.SendNotificationAsync(
                 }
                 catch (Exception ex) when (!stoppingToken.IsCancellationRequested)
                 {
+                    // Session cleanup is tied to the session lifetime; a failed send only affects this notification.
                     this.logger.LogWarning(ex, "Could not notify an MCP client about the shell location change.");
-                    lock (this.sync)
-                    {
-                        this.subscribers.RemoveAll(reference => !reference.TryGetTarget(out var target) || ReferenceEquals(target, server));
-                    }
                 }
             }
         }
@@ -111,9 +167,4 @@ private void OnLocationChanged()
     {
         this.changes.Writer.TryWrite(true);
     }
-
-    private void PruneSubscribers()
-    {
-        this.subscribers.RemoveAll(reference => !reference.TryGetTarget(out _));
-    }
 }

CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/McpServer.cs

@@ -14,6 +14,7 @@ namespace Azure.Data.Cosmos.Shell.Mcp;
 using Microsoft.Extensions.Hosting;
 using Microsoft.Extensions.Logging;
 
+using ModelContextProtocol.AspNetCore;
 using ModelContextProtocol.Protocol;
 using ModelContextProtocol.Server;
 
@@ -26,6 +27,8 @@ namespace Azure.Data.Cosmos.Shell.Mcp;
 /// </summary>
 internal class McpServer
 {
+    internal static readonly TimeSpan SessionIdleTimeout = TimeSpan.FromMinutes(10);
+
     public static IHost CreateHost(CosmosShellOptions serverArguments)
     {
         var builder = WebApplication.CreateBuilder([]);
@@ -84,7 +87,26 @@ private static void ConfigureMcpServer(IServiceCollection services)
 
         var mcpServerBuilder = services.AddMcpServer();
         mcpServerBuilder.WithResources<ResourceOperations>();
-        mcpServerBuilder.WithHttpTransport();
+
+        // Destructive-command confirmation uses elicitation, a server-to-client request that
+        // requires a session. SDK 2.x defaults to stateless, so keep sessions enabled.
+        mcpServerBuilder.WithHttpTransport(ConfigureHttpTransport);
+    }
+
+    internal static void ConfigureHttpTransport(HttpServerTransportOptions options)
+    {
+        options.SessionMode = HttpServerSessionMode.Stateful;
+
+        // Sessions with an open GET stream never go idle. Once a client disconnects without DELETE,
+        // the session is disposed after this timeout, which also ends its location subscription.
+#pragma warning disable MCP9006 // Stateful Streamable HTTP options are required for session-bound features.
+        options.IdleTimeout = SessionIdleTimeout;
+#pragma warning restore MCP9006
+
+#pragma warning disable MCPEXP002 // RunSessionHandler is the only hook that observes the session lifetime.
+        options.RunSessionHandler = (httpContext, server, cancellationToken) =>
+            httpContext.RequestServices.GetRequiredService<LocationResourceSubscriptions>().RunSessionAsync(server, cancellationToken);
+#pragma warning restore MCPEXP002
     }
 
     private static string LoadServerInstructions()

CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/ToolOperations.cs

@@ -54,13 +54,13 @@ public ToolOperations(ILogger<ToolOperations> logger, LocationResourceSubscripti
 
     private ValueTask<EmptyResult> SubscribeToResourcesAsync(RequestContext<SubscribeRequestParams> context, CancellationToken cancellationToken)
     {
-        this.locationSubscriptions.Subscribe(context.Server, context.Params?.Uri ?? string.Empty);
+        this.locationSubscriptions.Subscribe(context.Server.SessionId, context.Params?.Uri ?? string.Empty);
         return ValueTask.FromResult(new EmptyResult());
     }
 
     private ValueTask<EmptyResult> UnsubscribeFromResourcesAsync(RequestContext<UnsubscribeRequestParams> context, CancellationToken cancellationToken)
     {
-        this.locationSubscriptions.Unsubscribe(context.Server, context.Params?.Uri ?? string.Empty);
+        this.locationSubscriptions.Unsubscribe(context.Server.SessionId, context.Params?.Uri ?? string.Empty);
         return ValueTask.FromResult(new EmptyResult());
     }
 

Directory.Packages.props

@@ -24,8 +24,8 @@
     <PackageVersion Include="Microsoft.Extensions.Hosting" Version="11.0.0-preview.2.26159.112" />
     <PackageVersion Include="OpenTelemetry" Version="1.15.3" />
     <PackageVersion Include="OpenTelemetry.Exporter.OpenTelemetryProtocol" Version="1.15.3" />
-    <PackageVersion Include="ModelContextProtocol" Version="1.1.0" />
-    <PackageVersion Include="ModelContextProtocol.AspNetCore" Version="1.1.0" />
+    <PackageVersion Include="ModelContextProtocol" Version="2.2.0" />
+    <PackageVersion Include="ModelContextProtocol.AspNetCore" Version="2.2.0" />
     <PackageVersion Include="OmniSharp.Extensions.LanguageServer" Version="0.19.9" />
     <PackageVersion Include="RadLine" Version="0.9.0" />
     <PackageVersion Include="SharpFuzz" Version="2.2.0" />

docs/mcp.md

@@ -99,7 +99,7 @@ Clients can read the `cosmos://shell/current-location` MCP resource. Its JSON co
 
 Only `cosmos://shell/current-location` supports subscriptions; subscribing to any other URI, including the documentation resources, returns an invalid-params error.
 
-This server uses the subscription protocol supported by its MCP SDK; clients must support subscriptions and server-to-client notifications over the HTTP connection. A notification does not guarantee that a client refreshes the model's context. Every tool response also includes `currentLocation`, and explicit `database` / `container` arguments remain the reliable way to target independent operations.
+This server uses the subscription protocol supported by its MCP SDK; clients must support subscriptions and server-to-client notifications over the HTTP connection. Notifications are delivered on the session's GET stream. A subscription lasts as long as its MCP session: it ends when the client deletes the session or when the session has had no open request (including the GET stream) for 10 minutes. After that the server returns 404 for the session and the client must start a new session and subscribe again. A notification does not guarantee that a client refreshes the model's context. Every tool response also includes `currentLocation`, and explicit `database` / `container` arguments remain the reliable way to target independent operations.
 
 ### Data Exposure