Add server tool timeouts (default & per-tool) and honor JSON-RPC $/cancelRequest (#922)

src/ModelContextProtocol.Core/McpErrorCode.cs

@@ -46,4 +46,12 @@ public enum McpErrorCode
     /// This error is used when the endpoint encounters an unexpected condition that prevents it from fulfilling the request.
     /// </remarks>
     InternalError = -32603,
+
+    /// <summary>
+� � /// Indicates that the request was cancelled by the client.
+� � /// </summary>
+� � /// <remarks>
+� � /// This error is returned when the CancellationToken passed with the request is cancelled before processing completes.
+� � /// </remarks>
+� � RequestCancelled = -32800,
 }

src/ModelContextProtocol.Core/McpSessionHandler.cs

@@ -171,6 +171,15 @@ async Task ProcessMessageAsync()
                     }
                     catch (Exception ex)
                     {
+                        // Fast-path: user-initiated cancellation → emit JSON-RPC RequestCancelled and exit.
+                        if (ex is OperationCanceledException oce
+                            && message is JsonRpcRequest cancelledReq
+                            && IsUserInitiatedCancellation(oce, cancellationToken, combinedCts))
+                        {
+                            await SendRequestCancelledErrorAsync(cancelledReq, cancellationToken).ConfigureAwait(false);
+                            return;
+                        }
+
                         // Only send responses for request errors that aren't user-initiated cancellation.
                         bool isUserCancellation =
                             ex is OperationCanceledException &&
@@ -301,8 +310,35 @@ private async Task HandleMessageAsync(JsonRpcMessage message, CancellationToken
         }
     }
 
+    /// <summary>
+    /// Handles inbound JSON-RPC notifications. Special-cases <c>$/cancelRequest</c>
+    /// to cancel the exact in-flight request, and also supports the SDK's custom
+    /// <see cref="NotificationMethods.CancelledNotification"/> for backwards compatibility.
+    /// </summary>
     private async Task HandleNotification(JsonRpcNotification notification, CancellationToken cancellationToken)
     {
+        // Handle JSON-RPC native cancellation: $/cancelRequest
+        if (notification.Method == NotificationMethods.JsonRpcCancelRequest)
+        {
+            try
+            {
+                if (TryGetJsonRpcIdFromCancelParams(notification.Params, out var reqId) &&
+                    _handlingRequests.TryGetValue(reqId, out var cts))
+                {
+                    // Request-specific CTS → cancel the in-flight handler
+                    await cts.CancelAsync().ConfigureAwait(false);
+                    LogRequestCanceled(EndpointName, reqId, reason: "jsonrpc/$/cancelRequest");
+                }
+            }
+            catch
+            {
+                // Per spec, invalid cancel messages should be ignored.
+            }
+
+            // We do not forward $/cancelRequest to user handlers.
+            return;
+        }
+
         // Special-case cancellation to cancel a pending operation. (We'll still subsequently invoke a user-specified handler if one exists.)
         if (notification.Method == NotificationMethods.CancelledNotification)
         {
@@ -567,6 +603,79 @@ private Task SendToRelatedTransportAsync(JsonRpcMessage message, CancellationTok
         }
     }
 
+    /// <summary>
+    /// Parses the <c>id</c> field from a <c>$/cancelRequest</c> notification's params.
+    /// Returns <see langword="true"/> only when the id is a valid JSON-RPC request id
+    /// (string or number).
+    /// </summary>
+    private static bool TryGetJsonRpcIdFromCancelParams(JsonNode? notificationParams, out RequestId id)
+    {
+        id = default;
+
+        if (notificationParams is not JsonObject obj)
+            return false;
+
+        if (!obj.TryGetPropertyValue("id", out var idNode) || idNode is null)
+            return false;
+
+        if (idNode.GetValueKind() == System.Text.Json.JsonValueKind.String)
+        {
+            id = new RequestId(idNode.GetValue<string>());
+            return true;
+        }
+
+        if (idNode.GetValueKind() == System.Text.Json.JsonValueKind.Number)
+        {
+            try
+            {
+                var n = idNode.GetValue<long>();
+                id = new RequestId(n);
+                return true;
+            }
+            catch
+            {
+                return false;
+            }
+        }
+
+        return false;
+    }
+
+    /// <summary>
+    /// Determines whether the caught <see cref="OperationCanceledException"/> corresponds
+    /// to a user-initiated JSON-RPC cancellation (i.e., <c>$/cancelRequest</c>).
+    /// We distinguish this from global/session shutdown by checking tokens.
+    /// </summary>
+    private static bool IsUserInitiatedCancellation(
+        OperationCanceledException _,
+        CancellationToken sessionOrLoopToken,
+        CancellationTokenSource? perRequestCts)
+    {
+        // User cancellation: per-request CTS is canceled, but the outer/session token is NOT.
+        return !sessionOrLoopToken.IsCancellationRequested
+               && perRequestCts?.IsCancellationRequested == true;
+    }
+
+    /// <summary>
+    /// Sends a standard JSON-RPC <c>RequestCancelled</c> error for the given request.
+    /// </summary>
+    private Task SendRequestCancelledErrorAsync(JsonRpcRequest request, CancellationToken ct)
+    {
+        var error = new JsonRpcError
+        {
+            Id = request.Id,
+            JsonRpc = "2.0",
+            Error = new JsonRpcErrorDetail
+            {
+                Code = (int)McpErrorCode.RequestCancelled,
+                Message = "Request was cancelled."
+            },
+            Context = new JsonRpcMessageContext { RelatedTransport = request.Context?.RelatedTransport },
+        };
+
+        return SendMessageAsync(error, ct);
+    }
+
     private static string CreateActivityName(string method) => method;
 
     private static string GetMethodName(JsonRpcMessage message) =>

src/ModelContextProtocol.Core/Protocol/NotificationMethods.cs

@@ -131,4 +131,13 @@ public static class NotificationMethods
     /// </para>
     /// </remarks>
     public const string CancelledNotification = "notifications/cancelled";
+
+    /// <summary>
+    /// JSON-RPC core cancellation method name (<c>$/cancelRequest</c>).
+    /// </summary>
+    /// <remarks>
+    /// Carries a single <c>id</c> field (string or number) identifying the in-flight
+    /// request that should be cancelled.
+    /// </remarks>
+    public const string JsonRpcCancelRequest = "$/cancelRequest";
 }

src/ModelContextProtocol.Core/Server/IMcpToolWithTimeout.cs

@@ -0,0 +1,17 @@
+namespace ModelContextProtocol.Server;
+
+/// <summary>
+/// Optional contract for tools that expose a per-tool execution timeout.
+/// </summary>
+/// <remarks>
+/// When specified, this value overrides the server-level
+/// <see cref="McpServerOptions.DefaultToolTimeout"/> for this tool only.
+/// </remarks>
+public interface IMcpToolWithTimeout
+{
+    /// <summary>
+    /// Gets the per-tool timeout. When <see langword="null"/>, the server's
+    /// default applies (if any).
+    /// </summary>
+    TimeSpan? Timeout { get; }
+}

src/ModelContextProtocol.Core/Server/McpServerImpl.cs

@@ -2,7 +2,9 @@
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Logging.Abstractions;
 using ModelContextProtocol.Protocol;
+using System.Collections.Concurrent;
 using System.Runtime.CompilerServices;
+using System.Text.Json.Nodes;
 using System.Text.Json.Serialization.Metadata;
 
 namespace ModelContextProtocol.Server;
@@ -41,7 +43,7 @@ internal sealed partial class McpServerImpl : McpServer
     /// rather than a nullable to be able to manipulate it atomically.
     /// </remarks>
     private StrongBox<LoggingLevel>? _loggingLevel;
-
+    
     /// <summary>
     /// Creates a new instance of <see cref="McpServerImpl"/>.
     /// </summary>
@@ -508,6 +510,12 @@ await originalListPromptsHandler(request, cancellationToken).ConfigureAwait(fals
             McpJsonUtilities.JsonContext.Default.GetPromptResult);
     }
 
+    /// <summary>
+    /// Wires up tools capability: listing, invocation, DI-provided collections,
+    /// and the filter pipeline. Invocation enforces per-tool timeouts (when
+    /// <see cref="IMcpToolWithTimeout.Timeout"/> is set) or falls back to
+    /// <see cref="McpServerOptions.DefaultToolTimeout"/> when present.
+    /// </summary>
     private void ConfigureTools(McpServerOptions options)
     {
         var listToolsHandler = options.Handlers.ListToolsHandler;
@@ -578,12 +586,58 @@ await originalListToolsHandler(request, cancellationToken).ConfigureAwait(false)
                     request.MatchedPrimitive = tool;
                 }
 
+                TimeSpan? effectiveTimeout = null;
+
                 try
                 {
+                    // Determine effective timeout: per-tool overrides server default
+                    effectiveTimeout = (request.MatchedPrimitive as IMcpToolWithTimeout)?.Timeout
+                                                               ?? ServerOptions.DefaultToolTimeout;
+
+                    if (effectiveTimeout is { } ts)
+                    {
+                        // Create and link cancellation tokens to enforce the timeout.
+                        // The 'using' statements ensure these disposable resources are cleaned up 
+                        // when the try block exits, regardless of success or exception.
+                        using var timeoutCts = new CancellationTokenSource(ts);
+                        using var linked = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, timeoutCts.Token);
+
+                        // Execute the next handler in the pipeline with the linked token.
+                        return await handler(request, linked.Token);
+                    }
+
+                    // If no timeout is configured, use the original request cancellation token.
                     return await handler(request, cancellationToken);
                 }
-                catch (Exception e) when (e is not OperationCanceledException and not McpProtocolException)
+                catch (Exception e) when (e is not McpProtocolException)
                 {
+                    // Handle Cancellation Exceptions
+                    if (e is OperationCanceledException)
+                    {
+                        // Distinguishing between a server-side timeout and a client-side cancellation 
+                        // is necessary here (both use the linked token).
+
+                        if (effectiveTimeout.HasValue) // Was a server-side timeout configured?
+                        {
+                            // If a timeout was configured and cancellation occurred, report it as a timeout error.
+                            return new()
+                            {
+                                IsError = true,
+
+                                // Machine readable timeout indication in 'Meta' property.
+                                // Required structural data for robust test assertions (checking 'Meta.IsTimeout' instead of parsing the 'Content' string).
+                                Meta = new JsonObject { ["IsTimeout"] = true },
+
+                                Content = [new TextContentBlock { Text = $"Tool '{request.Params?.Name}' timed out after {effectiveTimeout.Value.TotalMilliseconds}ms." }],
+                            };
+                        }
+
+                        // If no server timeout was set, the cancellation must originate from the client's CancellationToken 
+                        // passed into RunAsync (or the network layer). We re-throw the exception to allow the 
+                        // JSON-RPC handler to process it into a standard protocol-level cancellation error (JsonRpcError).
+                        throw;
+                    }
+
                     ToolCallError(request.Params?.Name ?? string.Empty, e);
 
                     string errorMessage = e is McpException ?

src/ModelContextProtocol.Core/Server/McpServerOptions.cs

@@ -93,12 +93,12 @@ public sealed class McpServerOptions
     /// <summary>
     /// Gets or sets the container of handlers used by the server for processing protocol messages.
     /// </summary>
-    public McpServerHandlers Handlers 
-    { 
+    public McpServerHandlers Handlers
+    {
         get => field ??= new();
         set
-        { 
-            Throw.IfNull(value); 
+        {
+            Throw.IfNull(value);
             field = value;
         }
     }
@@ -166,4 +166,15 @@ public McpServerHandlers Handlers
     /// </para>
     /// </remarks>
     public int MaxSamplingOutputTokens { get; set; } = 1000;
+
+    /// <summary>
+    /// Gets or sets the default timeout applied to tool invocations.
+    /// </summary>
+    /// <remarks>
+    /// When set, the server enforces this timeout for all tools that do not define
+    /// their own timeout. Tools implementing <see cref="IMcpToolWithTimeout"/> can
+    /// override this value on a per-tool basis. When <see langword="null"/>, no
+    /// server-enforced timeout is applied.
+    /// </remarks>
+    public TimeSpan? DefaultToolTimeout { get; set; }
 }

src/ModelContextProtocol.Core/Server/McpServerToolAttribute.cs

@@ -177,10 +177,10 @@ public McpServerToolAttribute()
     /// The default is <see langword="true"/>.
     /// </para>
     /// </remarks>
-    public bool Destructive 
+    public bool Destructive
     {
-        get => _destructive ?? DestructiveDefault; 
-        set => _destructive = value; 
+        get => _destructive ?? DestructiveDefault;
+        set => _destructive = value;
     }
 
     /// <summary>
@@ -195,10 +195,10 @@ public bool Destructive
     /// The default is <see langword="false"/>.
     /// </para>
     /// </remarks>
-    public bool Idempotent  
+    public bool Idempotent
     {
         get => _idempotent ?? IdempotentDefault;
-        set => _idempotent = value; 
+        set => _idempotent = value;
     }
 
     /// <summary>
@@ -215,8 +215,8 @@ public bool Idempotent
     /// </remarks>
     public bool OpenWorld
     {
-        get => _openWorld ?? OpenWorldDefault; 
-        set => _openWorld = value; 
+        get => _openWorld ?? OpenWorldDefault;
+        set => _openWorld = value;
     }
 
     /// <summary>
@@ -235,10 +235,10 @@ public bool OpenWorld
     /// The default is <see langword="false"/>.
     /// </para>
     /// </remarks>
-    public bool ReadOnly  
+    public bool ReadOnly
     {
-        get => _readOnly ?? ReadOnlyDefault; 
-        set => _readOnly = value; 
+        get => _readOnly ?? ReadOnlyDefault;
+        set => _readOnly = value;
     }
 
     /// <summary>
@@ -269,4 +269,10 @@ public bool ReadOnly
     /// </para>
     /// </remarks>
     public string? IconSource { get; set; }
+
+    /// <summary>
+    /// Optional timeout for this tool in seconds.
+    /// If null, the global default (if any) applies.
+    /// </summary>
+    public int? TimeoutSeconds { get; set; }
 }