diff --git a/CHANGELOG.md b/CHANGELOG.md index ae49951..733f592 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ ### Fixes +- Restore legacy MCP HTTP+SSE endpoints alongside Streamable HTTP on the same `--mcp` server using SDK 2.2. No extra option is required. Bound outstanding legacy requests to 16 across all sessions and reject excess requests without queuing. - Parse `exec` options and shell words like direct commands. Bind built-in options normally and pass option-shaped words to functions and script files as positional text. - `for` and `do` loops now reject misspelled `in` and `while` keywords before any statements in the input execute. They previously accepted any identifier in those positions. - `filter` now rejects array index literals larger than `2147483647`, including optional access, instead of silently selecting element zero. Representable indexes beyond the array length still return `null`. diff --git a/CosmosDBShell.Tests/LegacyHttpAdmissionTests.cs b/CosmosDBShell.Tests/LegacyHttpAdmissionTests.cs new file mode 100644 index 0000000..f200d50 --- /dev/null +++ b/CosmosDBShell.Tests/LegacyHttpAdmissionTests.cs @@ -0,0 +1,142 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace CosmosShell.Tests; + +using System.Text.Json; +using Azure.Data.Cosmos.Shell.Mcp; +using Microsoft.AspNetCore.Http; +using ModelContextProtocol.Protocol; +using ModelContextProtocol.Server; +using NSubstitute; + +public class LegacyHttpAdmissionTests : IDisposable +{ + private readonly List requestBodies = []; + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task SlowSseResponse_RejectsExcessHttpRequestsBeforeDispatch(bool errorResponse) + { + var admission = new LegacyHttpAdmission(); + var dispatched = 0; + RequestDelegate dispatch = context => + { + dispatched++; + context.Response.StatusCode = StatusCodes.Status202Accepted; + return Task.CompletedTask; + }; + for (var i = 0; i < LegacyRequestLimiter.MaxOutstandingRequests; i++) + { + await admission.InvokeAsync(CreateContext(CreateRequest(i)), dispatch); + } + + var server = Substitute.For(); + server.SessionId.Returns("legacy-test"); + var releaseStream = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var outgoing = admission.TrackOutgoing( + async (_, cancellationToken) => await releaseStream.Task.WaitAsync(cancellationToken)); + JsonRpcMessage response = errorResponse + ? new JsonRpcError { Id = new RequestId(0), Error = new JsonRpcErrorDetail { Code = -32601, Message = "Unknown method." } } + : new JsonRpcResponse { Id = new RequestId(0), Result = null }; + var pendingResponse = outgoing(new MessageContext(server, response), TestContext.Current.CancellationToken); + try + { + for (var i = 0; i < 32; i++) + { + var excess = CreateContext(CreateRequest(100 + i)); + await admission.InvokeAsync(excess, dispatch); + Assert.Equal(StatusCodes.Status429TooManyRequests, excess.Response.StatusCode); + Assert.Equal("1", excess.Response.Headers.RetryAfter.ToString()); + } + + Assert.Equal(LegacyRequestLimiter.MaxOutstandingRequests, dispatched); + Assert.Equal(LegacyRequestLimiter.MaxOutstandingRequests, admission.OutstandingRequests); + Assert.False(pendingResponse.IsCompleted); + + var cancellation = CreateContext(new JsonRpcNotification { Method = NotificationMethods.CancelledNotification }); + await admission.InvokeAsync(cancellation, dispatch); + var reply = CreateContext(new JsonRpcResponse { Id = new RequestId(200), Result = null }); + await admission.InvokeAsync(reply, dispatch); + Assert.Equal(StatusCodes.Status202Accepted, cancellation.Response.StatusCode); + Assert.Equal(StatusCodes.Status202Accepted, reply.Response.StatusCode); + Assert.Equal(LegacyRequestLimiter.MaxOutstandingRequests, admission.OutstandingRequests); + } + finally + { + releaseStream.TrySetResult(true); + await pendingResponse.WaitAsync(TestContext.Current.CancellationToken); + } + + Assert.Equal(LegacyRequestLimiter.MaxOutstandingRequests - 1, admission.OutstandingRequests); + var retry = CreateContext(CreateRequest(300)); + await admission.InvokeAsync(retry, dispatch); + Assert.Equal(StatusCodes.Status202Accepted, retry.Response.StatusCode); + Assert.Equal(LegacyRequestLimiter.MaxOutstandingRequests, admission.OutstandingRequests); + admission.EndSession("legacy-test"); + Assert.Equal(0, admission.OutstandingRequests); + } + + [Theory] + [InlineData(StatusCodes.Status400BadRequest)] + [InlineData(StatusCodes.Status500InternalServerError)] + public async Task RejectedHttpRequest_ReleasesAdmission(int statusCode) + { + var admission = new LegacyHttpAdmission(); + var context = CreateContext(CreateRequest(1)); + await admission.InvokeAsync(context, current => + { + current.Response.StatusCode = statusCode; + return Task.CompletedTask; + }); + Assert.Equal(0, admission.OutstandingRequests); + } + + [Fact] + public async Task CancelledHandler_ReleasesAdmission() + { + var admission = new LegacyHttpAdmission(); + await admission.InvokeAsync(CreateContext(CreateRequest(1)), context => + { + context.Response.StatusCode = StatusCodes.Status202Accepted; + return Task.CompletedTask; + }); + var server = Substitute.For(); + server.SessionId.Returns("legacy-test"); + using var cancellation = new CancellationTokenSource(); + await cancellation.CancelAsync(); + var handler = admission.TrackIncoming((_, token) => Task.FromCanceled(token)); + await Assert.ThrowsAnyAsync( + () => handler(new MessageContext(server, CreateRequest(1)), cancellation.Token)); + Assert.Equal(0, admission.OutstandingRequests); + } + + public void Dispose() + { + foreach (var body in this.requestBodies) + { + body.Dispose(); + } + } + + private static JsonRpcRequest CreateRequest(int id) + { + return new JsonRpcRequest { Id = new RequestId(id), Method = RequestMethods.ToolsCall }; + } + + private DefaultHttpContext CreateContext(JsonRpcMessage message) + { + var body = new MemoryStream(JsonSerializer.SerializeToUtf8Bytes(message)); + this.requestBodies.Add(body); + var context = new DefaultHttpContext(); + context.Request.Path = "/message"; + context.Request.Method = HttpMethods.Post; + context.Request.QueryString = new QueryString("?sessionId=legacy-test"); + context.Request.ContentType = "application/json"; + context.Request.Body = body; + context.Response.Body = Stream.Null; + return context; + } +} diff --git a/CosmosDBShell.Tests/LegacyRequestLimiterTests.cs b/CosmosDBShell.Tests/LegacyRequestLimiterTests.cs new file mode 100644 index 0000000..42fc9b9 --- /dev/null +++ b/CosmosDBShell.Tests/LegacyRequestLimiterTests.cs @@ -0,0 +1,67 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace CosmosShell.Tests; + +using Azure.Data.Cosmos.Shell.Mcp; +using ModelContextProtocol; +using ModelContextProtocol.Protocol; +using ModelContextProtocol.Server; +using NSubstitute; + +public class LegacyRequestLimiterTests +{ + [Fact] + public async Task SaturatedLimiter_AllowsRepliesAndNotifications() + { + using var limiter = new LegacyRequestLimiter(); + var server = Substitute.For(); + var releaseRequests = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var controlMessages = 0; + var handler = limiter.Limit(async (context, cancellationToken) => + { + if (context.JsonRpcMessage is JsonRpcRequest) + { + await releaseRequests.Task.WaitAsync(cancellationToken); + } + else + { + controlMessages++; + } + }); + var requests = new List(); + try + { + for (var i = 0; i < LegacyRequestLimiter.MaxOutstandingRequests; i++) + { + requests.Add(handler( + new MessageContext(server, new JsonRpcRequest { Id = new RequestId(i), Method = RequestMethods.ToolsCall }), + TestContext.Current.CancellationToken)); + } + + Assert.Equal(LegacyRequestLimiter.MaxOutstandingRequests, limiter.OutstandingRequests); + var error = await Assert.ThrowsAsync( + async () => await handler( + new MessageContext(server, new JsonRpcRequest { Id = new RequestId(100), Method = RequestMethods.ToolsCall }), + TestContext.Current.CancellationToken)); + Assert.Equal(LegacyRequestLimiter.OverloadErrorCode, error.ErrorCode); + + await handler( + new MessageContext(server, new JsonRpcResponse { Id = new RequestId(200), Result = null }), + TestContext.Current.CancellationToken); + await handler( + new MessageContext(server, new JsonRpcNotification { Method = NotificationMethods.CancelledNotification }), + TestContext.Current.CancellationToken); + Assert.Equal(2, controlMessages); + Assert.Equal(LegacyRequestLimiter.MaxOutstandingRequests, limiter.OutstandingRequests); + } + finally + { + releaseRequests.TrySetResult(true); + await Task.WhenAll(requests); + } + + Assert.Equal(0, limiter.OutstandingRequests); + } +} diff --git a/CosmosDBShell.Tests/McpLocationSubscriptionTests.cs b/CosmosDBShell.Tests/McpLocationSubscriptionTests.cs index 2df82ed..ea9a18a 100644 --- a/CosmosDBShell.Tests/McpLocationSubscriptionTests.cs +++ b/CosmosDBShell.Tests/McpLocationSubscriptionTests.cs @@ -20,8 +20,10 @@ namespace CosmosShell.Tests; [Collection(CosmosShell.Tests.Shell.ThemeStateTestCollection.Name)] public class McpLocationSubscriptionTests { - [Fact] - public async Task SubscribedClient_ReceivesInteractiveLocationChange() + [Theory] + [InlineData(HttpTransportMode.StreamableHttp, "2025-11-25")] + [InlineData(HttpTransportMode.Sse, "2024-11-05")] + public async Task SubscribedClient_ReceivesInteractiveLocationChange(HttpTransportMode transportMode, string protocolVersion) { using var timeout = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken); timeout.CancelAfter(TimeSpan.FromSeconds(10)); @@ -29,7 +31,7 @@ public async Task SubscribedClient_ReceivesInteractiveLocationChange() await host.StartAsync(timeout.Token); try { - await using var client = await ConnectAsync(host, timeout.Token); + await using var client = await ConnectAsync(host, timeout.Token, protocolVersion, transportMode); Assert.NotNull(client.ServerCapabilities.Resources); Assert.True(client.ServerCapabilities.Resources.Subscribe); Assert.False(client.ServerCapabilities.Resources.ListChanged); @@ -89,8 +91,10 @@ public async Task SubscribedClient_ReceivesInteractiveLocationChange() } } - [Fact] - public async Task EndedSession_RemovesLocationSubscription() + [Theory] + [InlineData(HttpTransportMode.StreamableHttp, "2025-11-25")] + [InlineData(HttpTransportMode.Sse, "2024-11-05")] + public async Task EndedSession_RemovesLocationSubscription(HttpTransportMode transportMode, string protocolVersion) { using var timeout = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken); timeout.CancelAfter(TimeSpan.FromSeconds(10)); @@ -99,11 +103,11 @@ public async Task EndedSession_RemovesLocationSubscription() try { var subscriptions = host.Services.GetRequiredService(); - var client = await ConnectAsync(host, timeout.Token); + await using var client = await ConnectAsync(host, timeout.Token, protocolVersion, transportMode); 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. + // Closing either transport must release its session's subscription. await client.DisposeAsync(); while (subscriptions.SubscriberCount != 0) { @@ -329,6 +333,153 @@ public async Task HostStop_EndsOpenListenStreamWithoutWaitingForShutdownTimeout( Assert.Equal(0, subscriptions.ListenerCount); } + [Fact] + public async Task LegacyAndStreamableHttpClients_ShareServerToolsAndResources() + { + 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 + { + await using var modern = await ConnectAsync(host, timeout.Token); + await using var legacy = await ConnectAsync(host, timeout.Token, "2024-11-05", HttpTransportMode.Sse); + var modernTools = await modern.ListToolsAsync(cancellationToken: timeout.Token); + var legacyTools = await legacy.ListToolsAsync(cancellationToken: timeout.Token); + Assert.NotEmpty(modernTools); + Assert.Equal(modernTools.Select(tool => tool.Name), legacyTools.Select(tool => tool.Name)); + + var modernLocation = await modern.ReadResourceAsync(ResourceOperations.CurrentLocationUri, cancellationToken: timeout.Token); + var legacyLocation = await legacy.ReadResourceAsync(ResourceOperations.CurrentLocationUri, cancellationToken: timeout.Token); + Assert.Equal( + Assert.IsType(Assert.Single(modernLocation.Contents)).Text, + Assert.IsType(Assert.Single(legacyLocation.Contents)).Text); + + foreach (var client in new[] { modern, legacy }) + { + var result = await client.CallToolAsync("version", cancellationToken: timeout.Token); + Assert.False(result.IsError == true); + Assert.NotEmpty(result.Content); + } + } + finally + { + await host.StopAsync(TestContext.Current.CancellationToken); + } + } + + [Theory] + [InlineData("/sse", "GET")] + [InlineData("/message", "POST")] + public async Task LegacyEndpoints_RejectNonLoopbackOrigins(string path, string method) + { + 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 address = host.Services.GetRequiredService().Features.Get()!.Addresses.Single(); + using var http = new HttpClient(); + using var request = new HttpRequestMessage(new HttpMethod(method), address.TrimEnd('/') + path); + request.Headers.Add("Origin", "https://untrusted.example"); + using var response = await http.SendAsync(request, timeout.Token); + Assert.Equal(System.Net.HttpStatusCode.Forbidden, response.StatusCode); + } + finally + { + await host.StopAsync(TestContext.Current.CancellationToken); + } + } + + [Fact] + public async Task LegacyRequests_RejectExcessAcrossSessionsAndReleaseCancelledRequests() + { + using var timeout = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken); + timeout.CancelAfter(TimeSpan.FromSeconds(15)); + using var host = McpServer.CreateHost(new Program.CosmosShellOptions { McpPort = 0 }); + await host.StartAsync(timeout.Token); + var releaseGate = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + Task? heldGate = null; + var calls = new List>(); + try + { + await using var first = await ConnectAsync(host, timeout.Token, "2024-11-05", HttpTransportMode.Sse); + await using var second = await ConnectAsync(host, timeout.Token, "2024-11-05", HttpTransportMode.Sse); + await using var modern = await ConnectAsync(host, timeout.Token); + var limiter = host.Services.GetRequiredService(); + var admission = host.Services.GetRequiredService(); + var enteredGate = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + heldGate = ShellInterpreter.Instance.RunSerializedAsync( + async () => + { + enteredGate.SetResult(true); + return await releaseGate.Task; + }, + timeout.Token); + await enteredGate.Task.WaitAsync(timeout.Token); + + var cancelledRequestId = new RequestId("cancelled-limit-test"); + var cancelledCall = first.SendRequestAsync( + new JsonRpcRequest + { + Id = cancelledRequestId, + Method = RequestMethods.ToolsCall, + Params = JsonSerializer.SerializeToNode(new CallToolRequestParams { Name = "version" }), + }, + timeout.Token); + for (var i = 1; i < LegacyRequestLimiter.MaxOutstandingRequests; i++) + { + var client = i % 2 == 0 ? first : second; + calls.Add(client.CallToolAsync("version", cancellationToken: timeout.Token).AsTask()); + } + + await WaitForOutstandingRequestsAsync(limiter, LegacyRequestLimiter.MaxOutstandingRequests, timeout.Token); + var error = await Assert.ThrowsAsync( + async () => await second.CallToolAsync("version", cancellationToken: timeout.Token)); + Assert.Equal(System.Net.HttpStatusCode.TooManyRequests, error.StatusCode); + Assert.Contains("Retry", error.Message); + Assert.Equal(LegacyRequestLimiter.MaxOutstandingRequests, limiter.OutstandingRequests); + Assert.Equal(LegacyRequestLimiter.MaxOutstandingRequests, admission.OutstandingRequests); + Assert.NotEmpty(await modern.ListToolsAsync(cancellationToken: timeout.Token)); + + await first.SendMessageAsync( + new JsonRpcNotification + { + Method = NotificationMethods.CancelledNotification, + Params = JsonSerializer.SerializeToNode(new CancelledNotificationParams { RequestId = cancelledRequestId }), + }, + timeout.Token); + await Assert.ThrowsAnyAsync(async () => await cancelledCall.WaitAsync(timeout.Token)); + await WaitForOutstandingRequestsAsync(limiter, LegacyRequestLimiter.MaxOutstandingRequests - 1, timeout.Token); + calls.Add(second.CallToolAsync("version", cancellationToken: timeout.Token).AsTask()); + await WaitForOutstandingRequestsAsync(limiter, LegacyRequestLimiter.MaxOutstandingRequests, timeout.Token); + + releaseGate.SetResult(true); + foreach (var call in calls) + { + Assert.False((await call.WaitAsync(timeout.Token)).IsError == true); + } + + await WaitForOutstandingRequestsAsync(limiter, 0, timeout.Token); + var resourceError = await Assert.ThrowsAsync( + async () => await first.ReadResourceAsync("cosmos://docs/missing", cancellationToken: timeout.Token)); + Assert.Equal(McpErrorCode.ResourceNotFound, resourceError.ErrorCode); + await WaitForOutstandingRequestsAsync(limiter, 0, timeout.Token); + Assert.False((await first.CallToolAsync("version", cancellationToken: timeout.Token)).IsError == true); + } + finally + { + releaseGate.TrySetResult(true); + if (heldGate is not null) + { + await heldGate; + } + + await host.StopAsync(TestContext.Current.CancellationToken); + } + } + [Fact] public void HttpTransport_UsesSessionsOnlyForInitializeClientsWithBoundedIdleTimeout() { @@ -336,19 +487,36 @@ public void HttpTransport_UsesSessionsOnlyForInitializeClientsWithBoundedIdleTim McpServer.ConfigureHttpTransport(options); Assert.Equal(ModelContextProtocol.AspNetCore.HttpServerSessionMode.StatefulForInitializeClients, options.SessionMode); + Assert.NotNull(options.ConfigureSessionOptions); +#pragma warning disable MCP9004 // Verify the intentionally enabled legacy transport. + Assert.True(options.EnableLegacySse); +#pragma warning restore MCP9004 #pragma warning disable MCP9006, MCPEXP002 Assert.Equal(McpServer.SessionIdleTimeout, options.IdleTimeout); Assert.NotNull(options.RunSessionHandler); #pragma warning restore MCP9006, MCPEXP002 } + private static async Task WaitForOutstandingRequestsAsync(LegacyRequestLimiter limiter, int expected, CancellationToken cancellationToken) + { + while (limiter.OutstandingRequests != expected) + { + await Task.Delay(20, cancellationToken); + } + } + // resources/subscribe and session lifetimes exist only for clients that use the initialize handshake. - private static async Task ConnectAsync(IHost host, CancellationToken cancellationToken, string? protocolVersion = "2025-11-25") + private static async Task ConnectAsync( + IHost host, + CancellationToken cancellationToken, + string? protocolVersion = "2025-11-25", + HttpTransportMode transportMode = HttpTransportMode.StreamableHttp) { var address = host.Services.GetRequiredService().Features.Get()!.Addresses.Single(); var transport = new HttpClientTransport(new HttpClientTransportOptions { - Endpoint = new Uri(address.TrimEnd('/') + "/"), + Endpoint = new Uri(address.TrimEnd('/') + (transportMode == HttpTransportMode.Sse ? "/sse" : "/")), + TransportMode = transportMode, }); return await McpClient.CreateAsync(transport, new McpClientOptions { ProtocolVersion = protocolVersion }, cancellationToken: cancellationToken); } diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/LegacyHttpAdmission.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/LegacyHttpAdmission.cs new file mode 100644 index 0000000..aaeee4b --- /dev/null +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/LegacyHttpAdmission.cs @@ -0,0 +1,160 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Azure.Data.Cosmos.Shell.Mcp; + +using System.Text.Json; +using Microsoft.AspNetCore.Http; +using ModelContextProtocol; +using ModelContextProtocol.Protocol; +using ModelContextProtocol.Server; + +internal sealed class LegacyHttpAdmission +{ + private readonly object sync = new(); + private readonly Dictionary<(string SessionId, RequestId RequestId), Admission> outstanding = []; + + internal int OutstandingRequests + { + get + { + lock (this.sync) + { + return this.outstanding.Count; + } + } + } + + internal async Task InvokeAsync(HttpContext context, RequestDelegate next) + { + if (context.Request.Path != "/message" || !HttpMethods.IsPost(context.Request.Method)) + { + await next(context); + return; + } + + context.Request.EnableBuffering(); + JsonRpcMessage? message; + try + { + message = await context.Request.ReadFromJsonAsync( + McpJsonUtilities.DefaultOptions, context.RequestAborted); + } + catch (JsonException) + { + context.Response.StatusCode = StatusCodes.Status400BadRequest; + await context.Response.WriteAsync("Invalid JSON-RPC message.", context.RequestAborted); + return; + } + + context.Request.Body.Position = 0; + if (message is not JsonRpcRequest request) + { + await next(context); + return; + } + + var key = (context.Request.Query["sessionId"].ToString(), request.Id); + Admission? admission = null; + lock (this.sync) + { + if (this.outstanding.Count < LegacyRequestLimiter.MaxOutstandingRequests && !this.outstanding.ContainsKey(key)) + { + admission = new Admission(this, key); + this.outstanding.Add(key, admission); + } + } + + if (admission is null) + { + context.Response.StatusCode = StatusCodes.Status429TooManyRequests; + context.Response.Headers.RetryAfter = "1"; + await context.Response.WriteAsync( + "Legacy SSE is busy or this request ID is already outstanding. Retry after an outstanding request completes.", + context.RequestAborted); + return; + } + + try + { + await next(context); + } + finally + { + // Accepted requests retain admission until the SSE response is written or the session ends. + if (context.Response.StatusCode != StatusCodes.Status202Accepted) + { + admission.Dispose(); + } + } + } + + internal McpMessageHandler TrackIncoming(McpMessageHandler next) + { + return async (context, cancellationToken) => + { + var admission = context.JsonRpcMessage is JsonRpcRequest request + ? this.Find(context.Server.SessionId, request.Id) + : null; + try + { + await next(context, cancellationToken); + } + catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) + { + admission?.Dispose(); + throw; + } + }; + } + + internal McpMessageHandler TrackOutgoing(McpMessageHandler next) + { + return async (context, cancellationToken) => + { + using var admission = context.JsonRpcMessage switch + { + JsonRpcResponse response => this.Find(context.Server.SessionId, response.Id), + JsonRpcError error => this.Find(context.Server.SessionId, error.Id), + _ => null, + }; + await next(context, cancellationToken); + }; + } + + internal void EndSession(string? sessionId) + { + lock (this.sync) + { + foreach (var key in this.outstanding.Keys.Where(key => key.SessionId == sessionId).ToArray()) + { + this.outstanding.Remove(key); + } + } + } + + private Admission? Find(string? sessionId, RequestId requestId) + { + lock (this.sync) + { + return sessionId is not null && this.outstanding.TryGetValue((sessionId, requestId), out var admission) + ? admission + : null; + } + } + + private sealed class Admission(LegacyHttpAdmission owner, (string SessionId, RequestId RequestId) key) : IDisposable + { + public void Dispose() + { + lock (owner.sync) + { + if (owner.outstanding.TryGetValue(key, out var current) && ReferenceEquals(current, this)) + { + owner.outstanding.Remove(key); + } + } + } + } +} diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/LegacyRequestLimiter.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/LegacyRequestLimiter.cs new file mode 100644 index 0000000..ac19984 --- /dev/null +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/LegacyRequestLimiter.cs @@ -0,0 +1,51 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Azure.Data.Cosmos.Shell.Mcp; + +using ModelContextProtocol; +using ModelContextProtocol.Protocol; +using ModelContextProtocol.Server; + +internal sealed class LegacyRequestLimiter : IDisposable +{ + internal const int MaxOutstandingRequests = 16; + internal const McpErrorCode OverloadErrorCode = (McpErrorCode)(-32000); + + private readonly SemaphoreSlim permits = new(MaxOutstandingRequests, MaxOutstandingRequests); + + internal int OutstandingRequests => MaxOutstandingRequests - this.permits.CurrentCount; + + internal McpMessageHandler Limit(McpMessageHandler next) + { + return async (context, cancellationToken) => + { + // Replies and notifications must remain available for elicitation and cancellation. + if (context.JsonRpcMessage is not JsonRpcRequest) + { + await next(context, cancellationToken); + return; + } + + if (!await this.permits.WaitAsync(0, cancellationToken)) + { + throw new McpProtocolException( + $"Legacy SSE is busy: at most {MaxOutstandingRequests} requests may be outstanding. Retry after an outstanding request completes.", + OverloadErrorCode); + } + + try + { + // The SDK message pipeline awaits the handler, unlike the legacy HTTP POST. + await next(context, cancellationToken); + } + finally + { + this.permits.Release(); + } + }; + } + + public void Dispose() => this.permits.Dispose(); +} diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/McpServer.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/McpServer.cs index 5a8a8e6..0e45344 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/McpServer.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/McpServer.cs @@ -46,6 +46,8 @@ public static IHost CreateHost(CosmosShellOptions serverArguments) }); var application = builder.Build(); application.UseOriginValidation(); + application.Use((context, next) => + context.RequestServices.GetRequiredService().InvokeAsync(context, next)); application.MapMcp(); return application; } @@ -53,6 +55,8 @@ public static IHost CreateHost(CosmosShellOptions serverArguments) private static void ConfigureMcpServer(IServiceCollection services) { services.AddSingleton(); + services.AddSingleton(); + services.AddSingleton(); services.AddSingleton(); services.AddHostedService(services => services.GetRequiredService()); services.AddOptions() @@ -98,6 +102,22 @@ internal static void ConfigureHttpTransport(HttpServerTransportOptions options) // 2026-07-28 clients are served statelessly (confirmation via MRTR, updates via subscriptions/listen). // Clients that use the initialize handshake still get a session for elicitation and resources/subscribe. options.SessionMode = HttpServerSessionMode.StatefulForInitializeClients; +#pragma warning disable MCP9004 // Legacy transport compatibility is intentional for trusted loopback clients. + options.EnableLegacySse = true; +#pragma warning restore MCP9004 + options.ConfigureSessionOptions = (httpContext, serverOptions, cancellationToken) => + { + if (httpContext.Request.Path == "/sse") + { + var limiter = httpContext.RequestServices.GetRequiredService(); + var admission = httpContext.RequestServices.GetRequiredService(); + serverOptions.Filters.Message.IncomingFilters.Add(admission.TrackIncoming); + serverOptions.Filters.Message.IncomingFilters.Add(limiter.Limit); + serverOptions.Filters.Message.OutgoingFilters.Add(admission.TrackOutgoing); + } + + return Task.CompletedTask; + }; // 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. @@ -106,8 +126,20 @@ internal static void ConfigureHttpTransport(HttpServerTransportOptions options) #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); + options.RunSessionHandler = async (httpContext, server, cancellationToken) => + { + try + { + await httpContext.RequestServices.GetRequiredService().RunSessionAsync(server, cancellationToken); + } + finally + { + if (httpContext.Request.Path == "/sse") + { + httpContext.RequestServices.GetRequiredService().EndSession(server.SessionId); + } + } + }; #pragma warning restore MCPEXP002 } diff --git a/CosmosDBShell/lang/en.ftl b/CosmosDBShell/lang/en.ftl index d371835..714afad 100644 --- a/CosmosDBShell/lang/en.ftl +++ b/CosmosDBShell/lang/en.ftl @@ -1185,7 +1185,7 @@ help-StartupDatabase = Navigate to this database after connecting at startup. help-StartupContainer = Navigate to this container after connecting at startup. Requires --database. help-EnableMcpServer = Enable MCP server for programmatic control of the shell help-EnableLspServer = Enable Language Server Protocol (LSP) server for editor integration -help-McpPort = Enable MCP HTTP server. Optionally specify a port with --mcp ; default is 6128. +help-McpPort = Enable MCP HTTP server with Streamable HTTP and legacy SSE. Optionally specify a port with --mcp ; default is 6128. help-Verbose = Print full exception details instead of only the message. help-Theme = Color theme profile to apply at startup. Falls back to the COSMOSDB_SHELL_THEME environment variable. help-Diagnostics = Write timestamped diagnostic logs to a file. Optionally specify a path with --diagnostics ; defaults to a timestamped file in the shell configuration directory. diff --git a/README.md b/README.md index 48444bd..a040c97 100644 --- a/README.md +++ b/README.md @@ -41,6 +41,10 @@ MCP command execution is serialized with the shell, and destructive confirmation MCP clients supporting resource subscriptions can watch `cosmos://shell/current-location` for interactive navigation and connection changes; the resource includes the current account endpoint separately from the location. The resource list is static, so the server does not advertise `resources.listChanged`. See [MCP location updates](docs/mcp.md#shell-location-updates). +`--mcp` supports Streamable HTTP at `http://127.0.0.1:6128/` and legacy SSE at `http://127.0.0.1:6128/sse` on the same server, without an extra option. Both share the interactive shell's state. Use trusted local clients; prefer Streamable HTTP when available. See [MCP transports](docs/mcp.md#start-mcp-server). + +Legacy SSE allows at most 16 outstanding MCP requests across all sessions; excess requests receive HTTP 429 before reaching the SDK, without being queued. Replies and cancellation notifications remain available. + ## Quick Start **Requirements:** .NET SDK 10.0+. diff --git a/docs/mcp.md b/docs/mcp.md index a8960c8..c94ecda 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -11,6 +11,29 @@ dotnet run --project CosmosDBShell -- --mcp 5050 Bare `--mcp` starts the HTTP server on the default port `6128`. +The same server supports both transports without an extra option: + +- **Streamable HTTP:** `http://127.0.0.1:6128/` (unchanged). +- **Legacy HTTP+SSE:** `http://127.0.0.1:6128/sse`. Clients send requests to + the `/message` URL announced by the SSE endpoint event, including its session + query parameter. + +Both transports use the MCP SDK's protocol handling and share the interactive +shell's connection, navigation, tools, and resources. For a custom port, replace +`6128` with that port. Prefer Streamable HTTP when the client supports it. +The server remains bound to loopback with Origin/Host validation; use only +trusted local clients. Legacy SSE acknowledges request POSTs immediately and +does not provide HTTP-level backpressure on outstanding handlers. The shell +therefore admits at most **16 outstanding legacy MCP requests across all SSE +sessions**, including requests waiting for the shell execution gate or user +confirmation. Excess requests receive **HTTP 429** with `Retry-After: 1` before +reaching the SDK, without being queued; retry after an outstanding request +completes. Admission remains held until the SSE response has been written, +including error responses, or the request is cancelled or its session ends. +Accepted HTTP POSTs return `202 Accepted`, which is not command success. +Replies and notifications, including cancellation, do not consume permits. +This limit does not apply to Streamable HTTP. + ## VS Code Setup > **Requires VS Code 1.103+** diff --git a/docs/navigation.md b/docs/navigation.md index f56dddd..1669f5f 100644 --- a/docs/navigation.md +++ b/docs/navigation.md @@ -279,7 +279,7 @@ Start the shell with options to customize behavior: | `--connect-resource-group ` | Azure resource group name for ARM database and container operations at startup | | `--database ` | Navigate to this database after connecting at startup | | `--container ` | Navigate to this container after connecting at startup. Requires `--database` | -| `--mcp [port]` | Enable MCP (Model Context Protocol) server on the given port, or `6128` by default | +| `--mcp [port]` | Enable MCP server with Streamable HTTP and legacy SSE on the given port, or `6128` by default | | `--diagnostics [path]` | Write timestamped diagnostic logs (commands, timing, errors, connection events) to a file, or to a timestamped file in the config directory by default | | `--otel [endpoint]` | Enable distributed tracing so requests carry a sampled W3C `traceparent`. Optionally export spans to an OTLP `endpoint`; falls back to the `OTEL_EXPORTER_OTLP_ENDPOINT` environment variable | | `--color-system ` | Color scheme: 0=off, 1=standard, 2=truecolor (alias: `--cs`) | @@ -352,6 +352,8 @@ cosmosdbshell --mcp # Start with MCP server enabled on a custom port cosmosdbshell --mcp 5050 +# The same server also accepts legacy clients at http://127.0.0.1:5050/sse + # Capture a diagnostic log to the default location in the config directory cosmosdbshell --diagnostics diff --git a/l10n/CosmosDBShell.json b/l10n/CosmosDBShell.json index 38a5fc5..49ac614 100644 --- a/l10n/CosmosDBShell.json +++ b/l10n/CosmosDBShell.json @@ -1305,7 +1305,7 @@ "help-HelpCommandMoreText": "Display more information on a specific command.", "help-HelpCommandScreenText": "Display this help screen.", "help-HelpOptionDescription": "Show this help text and exit.", - "help-McpPort": "Enable MCP HTTP server. Optionally specify a port with --mcp \u003Cport\u003E; default is 6128.", + "help-McpPort": "Enable MCP HTTP server with Streamable HTTP and legacy SSE. Optionally specify a port with --mcp \u003Cport\u003E; default is 6128.", "help-NotesHeadingText": "NOTES:", "help-OptionGroupWord": "Group", "help-OptionsHeadingText": "OPTIONS:",