Route SendRequestAsync logic through outgoing message filters (#1465)

halter73 · web-flow · commit afc40d1e2eee · 2026-03-27T09:09:45.000-07:00
diff --git a/src/ModelContextProtocol.Core/McpSessionHandler.cs b/src/ModelContextProtocol.Core/McpSessionHandler.cs
@@ -576,15 +576,6 @@ public async Task<JsonRpcResponse> SendRequestAsync(JsonRpcRequest request, Canc
                 AddTags(ref tags, activity, request, method, target);
             }
 
-            if (_logger.IsEnabled(LogLevel.Trace))
-            {
-                LogSendingRequestSensitive(EndpointName, request.Method, JsonSerializer.Serialize(request, McpJsonUtilities.JsonContext.Default.JsonRpcMessage));
-            }
-            else
-            {
-                LogSendingRequest(EndpointName, request.Method);
-            }
-
             await SendToRelatedTransportAsync(request, cancellationToken).ConfigureAwait(false);
 
             // Now that the request has been sent, register for cancellation. If we registered before,
@@ -671,29 +662,17 @@ public async Task SendMessageAsync(JsonRpcMessage message, CancellationToken can
                 AddTags(ref tags, activity, message, method, target);
             }
 
-            await _outgoingMessageFilter(async (msg, ct) =>
-            {
-                if (_logger.IsEnabled(LogLevel.Trace))
-                {
-                    LogSendingMessageSensitive(EndpointName, JsonSerializer.Serialize(msg, McpJsonUtilities.JsonContext.Default.JsonRpcMessage));
-                }
-                else
-                {
-                    LogSendingMessage(EndpointName);
-                }
-
-                await SendToRelatedTransportAsync(msg, ct).ConfigureAwait(false);
+            await SendToRelatedTransportAsync(message, cancellationToken).ConfigureAwait(false);
 
-                // If the sent notification was a cancellation notification, cancel the pending request's await, as either the
-                // server won't be sending a response, or per the specification, the response should be ignored. There are inherent
-                // race conditions here, so it's possible and allowed for the operation to complete before we get to this point.
-                if (msg is JsonRpcNotification { Method: NotificationMethods.CancelledNotification } notification &&
-                    GetCancelledNotificationParams(notification.Params) is CancelledNotificationParams cn &&
-                    _pendingRequests.TryRemove(cn.RequestId, out var tcs))
-                {
-                    tcs.TrySetCanceled(default);
-                }
-            })(message, cancellationToken).ConfigureAwait(false);
+            // If the sent notification was a cancellation notification, cancel the pending request's await, as either the
+            // server won't be sending a response, or per the specification, the response should be ignored. There are inherent
+            // race conditions here, so it's possible and allowed for the operation to complete before we get to this point.
+            if (message is JsonRpcNotification { Method: NotificationMethods.CancelledNotification } notification &&
+                GetCancelledNotificationParams(notification.Params) is CancelledNotificationParams cn &&
+                _pendingRequests.TryRemove(cn.RequestId, out var tcs))
+            {
+                tcs.TrySetCanceled(default);
+            }
         }
         catch (Exception ex) when (addTags)
         {
@@ -710,7 +689,33 @@ await _outgoingMessageFilter(async (msg, ct) =>
     // Streamable HTTP transport where the specification states that the server SHOULD include JSON-RPC responses in
     // the HTTP response body for the POST request containing the corresponding JSON-RPC request.
     private Task SendToRelatedTransportAsync(JsonRpcMessage message, CancellationToken cancellationToken)
-        => (message.Context?.RelatedTransport ?? _transport).SendMessageAsync(message, cancellationToken);
+        => _outgoingMessageFilter((msg, ct) =>
+        {
+            if (msg is JsonRpcRequest request)
+            {
+                if (_logger.IsEnabled(LogLevel.Trace))
+                {
+                    LogSendingRequestSensitive(EndpointName, request.Method, JsonSerializer.Serialize(msg, McpJsonUtilities.JsonContext.Default.JsonRpcMessage));
+                }
+                else
+                {
+                    LogSendingRequest(EndpointName, request.Method);
+                }
+            }
+            else
+            {
+                if (_logger.IsEnabled(LogLevel.Trace))
+                {
+                    LogSendingMessageSensitive(EndpointName, JsonSerializer.Serialize(msg, McpJsonUtilities.JsonContext.Default.JsonRpcMessage));
+                }
+                else
+                {
+                    LogSendingMessage(EndpointName);
+                }
+            }
+
+            return (msg.Context?.RelatedTransport ?? _transport).SendMessageAsync(msg, ct);
+        })(message, cancellationToken);
 
     private static CancelledNotificationParams? GetCancelledNotificationParams(JsonNode? notificationParams)
     {
diff --git a/tests/ModelContextProtocol.AspNetCore.Tests/MapMcpTests.cs b/tests/ModelContextProtocol.AspNetCore.Tests/MapMcpTests.cs
@@ -9,6 +9,7 @@
 using System.Diagnostics;
 using System.Net;
 using System.Security.Claims;
+using System.Text.Json.Nodes;
 
 namespace ModelContextProtocol.AspNetCore.Tests;
 
@@ -290,6 +291,200 @@ public async Task LongRunningToolCall_DoesNotTimeout_WhenNoEventStreamStore()
 
     }
 
+    [Fact]
+    public async Task IncomingFilter_SeesClientRequests()
+    {
+        var observedMethods = new List<string>();
+
+        Builder.Services.AddMcpServer()
+            .WithHttpTransport(ConfigureStateless)
+            .WithMessageFilters(filters => filters.AddIncomingFilter((next) => async (context, cancellationToken) =>
+            {
+                if (context.JsonRpcMessage is JsonRpcRequest request)
+                {
+                    observedMethods.Add(request.Method);
+                }
+
+                await next(context, cancellationToken);
+            }))
+            .WithTools<EchoHttpContextUserTools>();
+
+        Builder.Services.AddHttpContextAccessor();
+
+        await using var app = Builder.Build();
+        app.MapMcp();
+        await app.StartAsync(TestContext.Current.CancellationToken);
+
+        await using var client = await ConnectAsync();
+
+        await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
+        await client.CallToolAsync("echo_with_user_name",
+            new Dictionary<string, object?> { ["message"] = "hi" },
+            cancellationToken: TestContext.Current.CancellationToken);
+
+        Assert.Contains(RequestMethods.Initialize, observedMethods);
+        Assert.Contains(RequestMethods.ToolsList, observedMethods);
+        Assert.Contains(RequestMethods.ToolsCall, observedMethods);
+    }
+
+    [Fact]
+    public async Task OutgoingFilter_SeesResponsesAndRequests()
+    {
+        Assert.SkipWhen(Stateless, "Server-originated requests are not supported in stateless mode.");
+
+        var observedMessageTypes = new List<string>();
+
+        Builder.Services.AddMcpServer()
+            .WithHttpTransport(ConfigureStateless)
+            .WithMessageFilters(filters => filters.AddOutgoingFilter((next) => async (context, cancellationToken) =>
+            {
+                var typeName = context.JsonRpcMessage switch
+                {
+                    JsonRpcRequest request => $"request:{request.Method}",
+                    JsonRpcResponse r when r.Result is JsonObject obj && obj.ContainsKey("protocolVersion") => "initialize-response",
+                    JsonRpcResponse r when r.Result is JsonObject obj && obj.ContainsKey("tools") => "tools-list-response",
+                    JsonRpcResponse r when r.Result is JsonObject obj && obj.ContainsKey("content") => "tool-call-response",
+                    _ => null,
+                };
+
+                if (typeName is not null)
+                {
+                    observedMessageTypes.Add(typeName);
+                }
+
+                await next(context, cancellationToken);
+            }))
+            .WithTools<ClaimsPrincipalTools>()
+            .WithTools<SamplingRegressionTools>();
+
+        await using var app = Builder.Build();
+        app.MapMcp();
+        await app.StartAsync(TestContext.Current.CancellationToken);
+
+        var clientOptions = new McpClientOptions
+        {
+            Capabilities = new() { Sampling = new() },
+            Handlers = new()
+            {
+                SamplingHandler = (_, _, _) => new(new CreateMessageResult
+                {
+                    Content = [new TextContentBlock { Text = "sampled response" }],
+                    Model = "test-model",
+                }),
+            },
+        };
+
+        await using var client = await ConnectAsync(clientOptions: clientOptions);
+
+        await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
+        await client.CallToolAsync("echo_claims_principal",
+            new Dictionary<string, object?> { ["message"] = "hi" },
+            cancellationToken: TestContext.Current.CancellationToken);
+        await client.CallToolAsync("sampling-tool",
+            new Dictionary<string, object?> { ["prompt"] = "Hello" },
+            cancellationToken: TestContext.Current.CancellationToken);
+
+        Assert.Contains("initialize-response", observedMessageTypes);
+        Assert.Contains("tools-list-response", observedMessageTypes);
+        Assert.Contains("tool-call-response", observedMessageTypes);
+        Assert.Contains($"request:{RequestMethods.SamplingCreateMessage}", observedMessageTypes);
+    }
+
+    [Fact]
+    public async Task OutgoingFilter_MultipleFilters_ExecuteInOrder()
+    {
+        var executionOrder = new List<string>();
+
+        Builder.Services.AddMcpServer()
+            .WithHttpTransport(ConfigureStateless)
+            .WithMessageFilters(filters =>
+            {
+                filters.AddOutgoingFilter((next) => async (context, cancellationToken) =>
+                {
+                    if (context.JsonRpcMessage is JsonRpcResponse r && r.Result is JsonObject obj && obj.ContainsKey("tools"))
+                    {
+                        executionOrder.Add("filter1-before");
+                    }
+
+                    await next(context, cancellationToken);
+
+                    if (context.JsonRpcMessage is JsonRpcResponse r2 && r2.Result is JsonObject obj2 && obj2.ContainsKey("tools"))
+                    {
+                        executionOrder.Add("filter1-after");
+                    }
+                });
+
+                filters.AddOutgoingFilter((next) => async (context, cancellationToken) =>
+                {
+                    if (context.JsonRpcMessage is JsonRpcResponse r && r.Result is JsonObject obj && obj.ContainsKey("tools"))
+                    {
+                        executionOrder.Add("filter2-before");
+                    }
+
+                    await next(context, cancellationToken);
+
+                    if (context.JsonRpcMessage is JsonRpcResponse r2 && r2.Result is JsonObject obj2 && obj2.ContainsKey("tools"))
+                    {
+                        executionOrder.Add("filter2-after");
+                    }
+                });
+            })
+            .WithTools<ClaimsPrincipalTools>();
+
+        await using var app = Builder.Build();
+        app.MapMcp();
+        await app.StartAsync(TestContext.Current.CancellationToken);
+
+        await using var client = await ConnectAsync();
+
+        await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
+
+        Assert.Equal(["filter1-before", "filter2-before", "filter2-after", "filter1-after"], executionOrder);
+    }
+
+    [Fact]
+    public async Task OutgoingFilter_CanSendAdditionalMessages()
+    {
+        Builder.Services.AddMcpServer()
+            .WithHttpTransport(ConfigureStateless)
+            .WithMessageFilters(filters => filters.AddOutgoingFilter((next) => async (context, cancellationToken) =>
+            {
+                if (context.JsonRpcMessage is JsonRpcResponse response &&
+                    response.Result is JsonObject result && result.ContainsKey("tools"))
+                {
+                    var extraNotification = new JsonRpcNotification
+                    {
+                        Method = "test/extra",
+                        Params = new JsonObject { ["message"] = "injected" },
+                        Context = new JsonRpcMessageContext { RelatedTransport = context.JsonRpcMessage.Context?.RelatedTransport },
+                    };
+
+                    await next(new MessageContext(context.Server, extraNotification), cancellationToken);
+                }
+
+                await next(context, cancellationToken);
+            }))
+            .WithTools<ClaimsPrincipalTools>();
+
+        await using var app = Builder.Build();
+        app.MapMcp();
+        await app.StartAsync(TestContext.Current.CancellationToken);
+
+        await using var client = await ConnectAsync();
+
+        var extraReceived = new TaskCompletionSource<string?>(TaskCreationOptions.RunContinuationsAsynchronously);
+        await using var registration = client.RegisterNotificationHandler("test/extra", (notification, _) =>
+        {
+            extraReceived.TrySetResult(notification.Params?["message"]?.GetValue<string>());
+            return default;
+        });
+
+        await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
+
+        var extraMessage = await extraReceived.Task.WaitAsync(TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken);
+        Assert.Equal("injected", extraMessage);
+    }
+
     private ClaimsPrincipal CreateUser(string name)
         => new(new ClaimsIdentity(
             [new Claim("name", name), new Claim(ClaimTypes.NameIdentifier, name)],
diff --git a/tests/ModelContextProtocol.Tests/Configuration/McpServerBuilderExtensionsMessageFilterTests.cs b/tests/ModelContextProtocol.Tests/Configuration/McpServerBuilderExtensionsMessageFilterTests.cs
