From 09709dc05363c6f5a687f53ae7e75c70a6bd76fb Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 1 Oct 2026 07:42:50 +0000 Subject: [PATCH 1/2] Upgrade ModelContextProtocol SDK to 2.2.0 and keep stateful HTTP sessions Co-authored-by: mkrueger <341098+mkrueger@users.noreply.github.com> --- CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/McpServer.cs | 6 +++++- Directory.Packages.props | 4 ++-- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/McpServer.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/McpServer.cs index 1193cea1..5018048e 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/McpServer.cs +++ b/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; @@ -80,7 +81,10 @@ private static void ConfigureMcpServer(IServiceCollection services) var mcpServerBuilder = services.AddMcpServer(); mcpServerBuilder.WithResources(); - 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(options => options.SessionMode = HttpServerSessionMode.Stateful); } private static string LoadServerInstructions() diff --git a/Directory.Packages.props b/Directory.Packages.props index c555da62..b379a907 100644 --- a/Directory.Packages.props +++ b/Directory.Packages.props @@ -24,8 +24,8 @@ - - + + From 17de1b370a5997ef2e8d1771c730ed21f7e00c1d Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 1 Oct 2026 09:09:40 +0000 Subject: [PATCH 2/2] Tie MCP location subscriptions to session lifetime and fix test port race Co-authored-by: mkrueger <341098+mkrueger@users.noreply.github.com> --- .../McpLocationSubscriptionTests.cs | 72 ++++++++++++--- .../LocationResourceSubscriptions.cs | 89 +++++++++++++++---- .../Azure.Data.Cosmos.Shell.Mcp/McpServer.cs | 20 ++++- .../ToolOperations.cs | 4 +- docs/mcp.md | 2 +- 5 files changed, 151 insertions(+), 36 deletions(-) diff --git a/CosmosDBShell.Tests/McpLocationSubscriptionTests.cs b/CosmosDBShell.Tests/McpLocationSubscriptionTests.cs index 2115838c..cfcb055e 100644 --- a/CosmosDBShell.Tests/McpLocationSubscriptionTests.cs +++ b/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(); + 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 ConnectAsync(IHost host, CancellationToken cancellationToken) + { + var address = host.Services.GetRequiredService().Features.Get()!.Addresses.Single(); + var transport = new HttpClientTransport(new HttpClientTransportOptions + { + Endpoint = new Uri(address.TrimEnd('/') + "/"), + }); + return await McpClient.CreateAsync(transport, cancellationToken: cancellationToken); + } } diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/LocationResourceSubscriptions.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/LocationResourceSubscriptions.cs index 5759d729..b892c0d6 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/LocationResourceSubscriptions.cs +++ b/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> 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 sessions = new(StringComparer.Ordinal); + + private readonly HashSet subscribedSessionIds = new(StringComparer.Ordinal); // Notifications carry only the URI, so pending changes are coalesced into one. private readonly Channel changes = Channel.CreateBounded( @@ -30,25 +34,81 @@ public LocationResourceSubscriptions(ILogger 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; + } + } + } + + /// + /// Runs an MCP session and keeps it available for notifications until the session ends. + /// + 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(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() .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 _)); - } } diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/McpServer.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/McpServer.cs index 38c7933c..18a9282b 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/McpServer.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/McpServer.cs @@ -27,6 +27,8 @@ namespace Azure.Data.Cosmos.Shell.Mcp; /// internal class McpServer { + internal static readonly TimeSpan SessionIdleTimeout = TimeSpan.FromMinutes(10); + public static IHost CreateHost(CosmosShellOptions serverArguments) { var builder = WebApplication.CreateBuilder([]); @@ -88,7 +90,23 @@ private static void ConfigureMcpServer(IServiceCollection services) // 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(options => options.SessionMode = HttpServerSessionMode.Stateful); + 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().RunSessionAsync(server, cancellationToken); +#pragma warning restore MCPEXP002 } private static string LoadServerInstructions() diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/ToolOperations.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/ToolOperations.cs index 612779e3..2fca20dd 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/ToolOperations.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/ToolOperations.cs @@ -54,13 +54,13 @@ public ToolOperations(ILogger logger, LocationResourceSubscripti private ValueTask SubscribeToResourcesAsync(RequestContext 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 UnsubscribeFromResourcesAsync(RequestContext 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()); } diff --git a/docs/mcp.md b/docs/mcp.md index 80f18c8f..26766803 100644 --- a/docs/mcp.md +++ b/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