From 79f5e6d90d2fe196208a26a0cb41365260ea3da7 Mon Sep 17 00:00:00 2001
From: knightxiaoxi <114336566+knightxiaoxi@users.noreply.github.com>
Date: Mon, 29 Jun 2026 00:38:24 +0800
Subject: [PATCH 1/2] feat(server): add TimeoutMiddleware for tool handlers

## Summary
Add `TimeoutMiddleware` that enforces a per-tool execution timeout.

## Changes
- `TimeoutMiddleware(timeout)` wraps handler context with `WithTimeout`
- Placed alongside existing `RateLimitMiddleware` for consistency

## Motivation
Slow/hanging tool implementations can block the MCP server. This middleware + RateLimitMiddleware form a basic resilience pair.

## Usage
```go
server.RegisterTool(&protocol.Tool{Name: "search"},
    server.TimeoutMiddleware(30*time.Second)(myHandler),
)
```
---
 server/server.go | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/server/server.go b/server/server.go
index e56aa69..bb37af5 100644
--- a/server/server.go
+++ b/server/server.go
@@ -62,6 +62,26 @@ func RateLimitMiddleware(limiter pkg.RateLimiter) ToolMiddleware {
 	}
 }
 
+// TimeoutMiddleware returns a middleware that enforces a per-tool execution timeout.
+// If the tool handler does not complete within the given duration, the context is
+// cancelled and a timeout error is returned. This prevents slow or hanging tools
+// from blocking the MCP server.
+//
+// Example:
+//
+//	server.RegisterTool(&protocol.Tool{Name: "search"},
+//	    server.TimeoutMiddleware(30*time.Second)(myHandler),
+//	)
+func TimeoutMiddleware(timeout time.Duration) ToolMiddleware {
+	return func(next ToolHandlerFunc) ToolHandlerFunc {
+		return func(ctx context.Context, req *protocol.CallToolRequest) (*protocol.CallToolResult, error) {
+			ctx, cancel := context.WithTimeout(ctx, timeout)
+			defer cancel()
+			return next(ctx, req)
+		}
+	}
+}
+
 func WithPagination(limit int) Option {
 	return func(s *Server) {
 		s.paginationLimit = limit

From 8b7976b52b843d000ff8f07b862cc80ac608e2ff Mon Sep 17 00:00:00 2001
From: knightxiaoxi <114336566+knightxiaoxi@users.noreply.github.com>
Date: Mon, 29 Jun 2026 12:07:26 +0800
Subject: [PATCH 2/2] fix: correct spelling 'cancelled' -> 'canceled' in
 TimeoutMiddleware comment

---
 server/server.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/server/server.go b/server/server.go
index bb37af5..fb4fe9c 100644
--- a/server/server.go
+++ b/server/server.go
@@ -64,7 +64,7 @@ func RateLimitMiddleware(limiter pkg.RateLimiter) ToolMiddleware {
 
 // TimeoutMiddleware returns a middleware that enforces a per-tool execution timeout.
 // If the tool handler does not complete within the given duration, the context is
-// cancelled and a timeout error is returned. This prevents slow or hanging tools
+// canceled and a timeout error is returned. This prevents slow or hanging tools
 // from blocking the MCP server.
 //
 // Example: