Skip to content

Commit 9326c87

Browse files
author
Tarek Mahmoud Sayed
committed
Harden Tasks seams in PR modelcontextprotocol#1693 review
- Reject custom RequestHandlers that collide with built-in or duplicate methods - Throw when an outgoing-request interceptor returns no result for sampling/roots - Guard the Tasks background execution against unobserved store exceptions and log failures - Add collision-guard tests for custom request handlers
1 parent 87064d3 commit 9326c87

4 files changed

Lines changed: 191 additions & 50 deletions

File tree

‎src/ModelContextProtocol.Core/Server/McpServer.Methods.cs‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -622,7 +622,18 @@ private async ValueTask<TResponse> SendRequestViaInterceptorAsync<TRequest, TRes
622622
{
623623
var paramsNode = JsonSerializer.SerializeToNode(request, requestTypeInfo);
624624
var resultNode = await interceptor(method, paramsNode, cancellationToken).ConfigureAwait(false);
625-
return resultNode is null ? default! : resultNode.Deserialize(responseTypeInfo)!;
625+
if (resultNode is null)
626+
{
627+
// A null result cannot be deserialized into a concrete TResponse (e.g. SampleAsync's
628+
// CreateMessageResult or RequestRootsAsync's ListRootsResult). Returning default! here
629+
// would hand callers a null response typed as non-null, deferring the failure to a
630+
// confusing NullReferenceException at the use site. Fail fast with a clear message.
631+
// Callers that can tolerate no result (such as ElicitAsync) do not route through this
632+
// helper and handle null themselves.
633+
throw new McpException($"The outgoing-request interceptor returned no result for the '{method}' request.");
634+
}
635+
636+
return resultNode.Deserialize(responseTypeInfo)!;
626637
}
627638

628639
private void ThrowIfElicitationUnsupported(ElicitRequestParams request)

‎src/ModelContextProtocol.Core/Server/McpServerImpl.cs‎

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -755,15 +755,32 @@ private void ConfigureCustomRequestHandlers(McpServerOptions options)
755755
{
756756
#pragma warning disable MCPEXP002
757757
if (options.RequestHandlers is not { Count: > 0 } customHandlers)
758-
#pragma warning restore MCPEXP002
759758
{
760759
return;
761760
}
762761

763762
foreach (var entry in customHandlers)
764763
{
764+
if (string.IsNullOrEmpty(entry.Method))
765+
{
766+
throw new InvalidOperationException(
767+
$"A custom request handler registered through {nameof(McpServerOptions)}.{nameof(McpServerOptions.RequestHandlers)} has a null or empty {nameof(McpServerRequestHandler.Method)}.");
768+
}
769+
770+
// Custom handlers are registered after all built-in handlers, so a method already present
771+
// belongs to a built-in method (e.g. initialize, tools/call) or an earlier custom handler.
772+
// Silently overwriting it would bypass the built-in handler's filters and protocol gating,
773+
// so reject the collision instead.
774+
if (_requestHandlers.ContainsKey(entry.Method))
775+
{
776+
throw new InvalidOperationException(
777+
$"A custom request handler registered through {nameof(McpServerOptions)}.{nameof(McpServerOptions.RequestHandlers)} " +
778+
$"uses the method '{entry.Method}', which is already handled by the server. Custom handlers cannot replace built-in methods or other custom handlers.");
779+
}
780+
765781
SetRawHandler(entry.Method, entry.Handler);
766782
}
783+
#pragma warning restore MCPEXP002
767784
}
768785

769786
private void SetRawHandler(string method, Func<JsonRpcRequest, CancellationToken, ValueTask<JsonNode?>> handler)

‎src/ModelContextProtocol.Extensions.Tasks/Server/McpTasksBuilderExtensions.cs‎

Lines changed: 82 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
11
using Microsoft.Extensions.DependencyInjection;
2+
using Microsoft.Extensions.Logging;
3+
using Microsoft.Extensions.Logging.Abstractions;
24
using Microsoft.Extensions.Options;
35
using ModelContextProtocol;
46
using ModelContextProtocol.Protocol;
@@ -30,13 +32,18 @@ public static IMcpServerBuilder WithTasks(this IMcpServerBuilder builder, IMcpTa
3032
if (store is null) throw new ArgumentNullException(nameof(store));
3133
#endif
3234

33-
builder.Services.AddSingleton<IPostConfigureOptions<McpServerOptions>>(new McpTasksPostConfigureOptions(store));
35+
// Resolve ILoggerFactory from the provider (rather than requiring the caller to pass one) so the
36+
// background task body has somewhere to report failures. It is optional: if no logging is
37+
// registered, the options fall back to NullLoggerFactory.
38+
builder.Services.AddSingleton<IPostConfigureOptions<McpServerOptions>>(
39+
sp => new McpTasksPostConfigureOptions(store, sp.GetService<ILoggerFactory>()));
3440
return builder;
3541
}
3642

37-
private sealed class McpTasksPostConfigureOptions(IMcpTaskStore store) : IPostConfigureOptions<McpServerOptions>
43+
private sealed class McpTasksPostConfigureOptions(IMcpTaskStore store, ILoggerFactory? loggerFactory) : IPostConfigureOptions<McpServerOptions>
3844
{
3945
private readonly IMcpTaskStore _store = store;
46+
private readonly ILogger _logger = (loggerFactory ?? NullLoggerFactory.Instance).CreateLogger<McpTasksPostConfigureOptions>();
4047
private readonly ConcurrentDictionary<string, CancellationTokenSource> _cancellationSources = new(StringComparer.Ordinal);
4148

4249
public void PostConfigure(string? name, McpServerOptions options)
@@ -74,71 +81,98 @@ public void PostConfigure(string? name, McpServerOptions options)
7481

7582
_ = Task.Run(async () =>
7683
{
77-
using (McpTasksServerExtensions.CreateMcpTaskScope(request.Server, taskId, _store))
84+
try
7885
{
79-
try
86+
using (McpTasksServerExtensions.CreateMcpTaskScope(request.Server, taskId, _store))
8087
{
81-
var augmented = await next(request, taskCancellationToken).ConfigureAwait(false);
88+
try
89+
{
90+
var augmented = await next(request, taskCancellationToken).ConfigureAwait(false);
8291

83-
if (augmented.IsAlternate)
92+
if (augmented.IsAlternate)
93+
{
94+
var error = new JsonRpcErrorDetail
95+
{
96+
Code = (int)McpErrorCode.InternalError,
97+
Message = $"{nameof(IMcpTaskStore)} is configured and the {nameof(McpServerHandlers.CallToolWithAlternateHandler)} returned IsAlternate = true. Use only one mechanism.",
98+
};
99+
var errorJson = JsonSerializer.SerializeToElement(error, McpJsonUtilities.DefaultOptions.GetTypeInfo<JsonRpcErrorDetail>());
100+
await _store.SetFailedAsync(taskId, errorJson).ConfigureAwait(false);
101+
return;
102+
}
103+
104+
var resultJson = JsonSerializer.SerializeToElement(augmented.Result!, McpJsonUtilities.DefaultOptions.GetTypeInfo<CallToolResult>());
105+
await _store.SetCompletedAsync(taskId, resultJson).ConfigureAwait(false);
106+
}
107+
catch (OperationCanceledException) when (taskCancellationToken.IsCancellationRequested)
108+
{
109+
await _store.SetCancelledAsync(taskId, CancellationToken.None).ConfigureAwait(false);
110+
}
111+
catch (InputRequiredException)
84112
{
85113
var error = new JsonRpcErrorDetail
86114
{
87-
Code = (int)McpErrorCode.InternalError,
88-
Message = $"{nameof(IMcpTaskStore)} is configured and the {nameof(McpServerHandlers.CallToolWithAlternateHandler)} returned IsAlternate = true. Use only one mechanism.",
115+
Code = (int)McpErrorCode.InvalidRequest,
116+
Message = "MRTR and tasks cannot be composed via [McpServerTool] yet.",
89117
};
90118
var errorJson = JsonSerializer.SerializeToElement(error, McpJsonUtilities.DefaultOptions.GetTypeInfo<JsonRpcErrorDetail>());
91119
await _store.SetFailedAsync(taskId, errorJson).ConfigureAwait(false);
92-
return;
93120
}
94-
95-
var resultJson = JsonSerializer.SerializeToElement(augmented.Result!, McpJsonUtilities.DefaultOptions.GetTypeInfo<CallToolResult>());
96-
await _store.SetCompletedAsync(taskId, resultJson).ConfigureAwait(false);
97-
}
98-
catch (OperationCanceledException) when (taskCancellationToken.IsCancellationRequested)
99-
{
100-
await _store.SetCancelledAsync(taskId, CancellationToken.None).ConfigureAwait(false);
101-
}
102-
catch (InputRequiredException)
103-
{
104-
var error = new JsonRpcErrorDetail
121+
catch (McpProtocolException mcpEx)
105122
{
106-
Code = (int)McpErrorCode.InvalidRequest,
107-
Message = "MRTR and tasks cannot be composed via [McpServerTool] yet.",
108-
};
109-
var errorJson = JsonSerializer.SerializeToElement(error, McpJsonUtilities.DefaultOptions.GetTypeInfo<JsonRpcErrorDetail>());
110-
await _store.SetFailedAsync(taskId, errorJson).ConfigureAwait(false);
123+
// SEP-2663 §186: protocol exceptions store as failed with JSON-RPC error shape.
124+
var error = new JsonRpcErrorDetail { Code = (int)mcpEx.ErrorCode, Message = mcpEx.Message };
125+
var errorJson = JsonSerializer.SerializeToElement(error, McpJsonUtilities.DefaultOptions.GetTypeInfo<JsonRpcErrorDetail>());
126+
await _store.SetFailedAsync(taskId, errorJson).ConfigureAwait(false);
127+
}
128+
catch (Exception ex)
129+
{
130+
// Non-protocol exceptions are wrapped as CallToolResult { IsError = true },
131+
// matching Core's BuildInitialAlternateToolFilter behavior.
132+
var errorResult = new CallToolResult
133+
{
134+
IsError = true,
135+
Content = [new TextContentBlock
136+
{
137+
Text = ex is McpException
138+
? $"An error occurred invoking '{request.Params?.Name}': {ex.Message}"
139+
: $"An error occurred invoking '{request.Params?.Name}'.",
140+
}],
141+
};
142+
var resultJson = JsonSerializer.SerializeToElement(errorResult, McpJsonUtilities.DefaultOptions.GetTypeInfo<CallToolResult>());
143+
await _store.SetCompletedAsync(taskId, resultJson).ConfigureAwait(false);
144+
}
145+
finally
146+
{
147+
if (_cancellationSources.TryRemove(taskId, out var registeredCts))
148+
{
149+
registeredCts.Dispose();
150+
}
151+
}
111152
}
112-
catch (McpProtocolException mcpEx)
153+
}
154+
catch (Exception outer)
155+
{
156+
// The inner handlers above record every expected outcome. Reaching here means a
157+
// store operation inside one of those handlers (or the task scope) threw, most
158+
// likely from a custom IMcpTaskStore. Record the failure best-effort and never let
159+
// it surface as an unobserved task exception.
160+
_logger.LogError(outer, "Background execution of task '{TaskId}' terminated unexpectedly while recording its result.", taskId);
161+
162+
try
113163
{
114-
// SEP-2663 §186: protocol exceptions store as failed with JSON-RPC error shape.
115-
var error = new JsonRpcErrorDetail { Code = (int)mcpEx.ErrorCode, Message = mcpEx.Message };
164+
var error = new JsonRpcErrorDetail { Code = (int)McpErrorCode.InternalError, Message = outer.Message };
116165
var errorJson = JsonSerializer.SerializeToElement(error, McpJsonUtilities.DefaultOptions.GetTypeInfo<JsonRpcErrorDetail>());
117166
await _store.SetFailedAsync(taskId, errorJson).ConfigureAwait(false);
118167
}
119-
catch (Exception ex)
168+
catch (Exception storeEx)
120169
{
121-
// Non-protocol exceptions are wrapped as CallToolResult { IsError = true },
122-
// matching Core's BuildInitialAlternateToolFilter behavior.
123-
var errorResult = new CallToolResult
124-
{
125-
IsError = true,
126-
Content = [new TextContentBlock
127-
{
128-
Text = ex is McpException
129-
? $"An error occurred invoking '{request.Params?.Name}': {ex.Message}"
130-
: $"An error occurred invoking '{request.Params?.Name}'.",
131-
}],
132-
};
133-
var resultJson = JsonSerializer.SerializeToElement(errorResult, McpJsonUtilities.DefaultOptions.GetTypeInfo<CallToolResult>());
134-
await _store.SetCompletedAsync(taskId, resultJson).ConfigureAwait(false);
170+
_logger.LogError(storeEx, "Failed to record the failure of background task '{TaskId}'.", taskId);
135171
}
136-
finally
172+
173+
if (_cancellationSources.TryRemove(taskId, out var leftoverCts))
137174
{
138-
if (_cancellationSources.TryRemove(taskId, out var registeredCts))
139-
{
140-
registeredCts.Dispose();
141-
}
175+
leftoverCts.Dispose();
142176
}
143177
}
144178
}, CancellationToken.None);

‎tests/ModelContextProtocol.Tests/Server/CustomRequestHandlerCollisionTests.cs‎

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
1+
using ModelContextProtocol.Protocol;
2+
using ModelContextProtocol.Server;
3+
using ModelContextProtocol.Tests.Utils;
4+
using System.Text.Json.Nodes;
5+
6+
namespace ModelContextProtocol.Tests.Server;
7+
8+
/// <summary>
9+
/// Verifies that custom request handlers registered through
10+
/// <see cref="McpServerOptions.RequestHandlers"/> cannot silently replace a built-in method
11+
/// or another custom handler.
12+
/// </summary>
13+
public class CustomRequestHandlerCollisionTests(ITestOutputHelper testOutputHelper) : LoggedTest(testOutputHelper)
14+
{
15+
#pragma warning disable MCPEXP002
16+
private static McpServerRequestHandler CreateHandler(string method) => new()
17+
{
18+
Method = method,
19+
Handler = (request, cancellationToken) => new ValueTask<JsonNode?>((JsonNode?)null),
20+
};
21+
22+
[Fact]
23+
public async Task CustomHandler_CollidingWithBuiltInMethod_Throws()
24+
{
25+
await using var transport = new StreamServerTransport(Stream.Null, Stream.Null);
26+
var options = new McpServerOptions
27+
{
28+
Capabilities = new() { Tools = new() },
29+
RequestHandlers = [CreateHandler("tools/call")],
30+
};
31+
32+
var ex = Assert.Throws<InvalidOperationException>(
33+
() => McpServer.Create(transport, options, LoggerFactory));
34+
35+
Assert.Contains("tools/call", ex.Message);
36+
}
37+
38+
[Fact]
39+
public async Task CustomHandler_CollidingWithInitialize_Throws()
40+
{
41+
await using var transport = new StreamServerTransport(Stream.Null, Stream.Null);
42+
var options = new McpServerOptions
43+
{
44+
RequestHandlers = [CreateHandler("initialize")],
45+
};
46+
47+
Assert.Throws<InvalidOperationException>(
48+
() => McpServer.Create(transport, options, LoggerFactory));
49+
}
50+
51+
[Fact]
52+
public async Task CustomHandler_DuplicateCustomMethod_Throws()
53+
{
54+
await using var transport = new StreamServerTransport(Stream.Null, Stream.Null);
55+
var options = new McpServerOptions
56+
{
57+
RequestHandlers = [CreateHandler("custom/method"), CreateHandler("custom/method")],
58+
};
59+
60+
var ex = Assert.Throws<InvalidOperationException>(
61+
() => McpServer.Create(transport, options, LoggerFactory));
62+
63+
Assert.Contains("custom/method", ex.Message);
64+
}
65+
66+
[Fact]
67+
public async Task CustomHandler_UniqueMethod_Succeeds()
68+
{
69+
await using var transport = new StreamServerTransport(Stream.Null, Stream.Null);
70+
var options = new McpServerOptions
71+
{
72+
RequestHandlers = [CreateHandler("custom/method")],
73+
};
74+
75+
await using var server = McpServer.Create(transport, options, LoggerFactory);
76+
Assert.NotNull(server);
77+
}
78+
#pragma warning restore MCPEXP002
79+
}

0 commit comments

Comments
 (0)