Merge branch 'main' into dev

Author: liushuangls

README.md

@@ -12,6 +12,10 @@
 
 [![Tutorial](http://img.youtube.com/vi/qoaeYMrXJH0/0.jpg)](http://www.youtube.com/watch?v=qoaeYMrXJH0 "Tutorial")
 
+<br>
+
+Discuss the SDK on [Discord](https://discord.gg/RqSS2NQVsY)
+
 </div>
 
 ```go
@@ -122,6 +126,7 @@ func main() {
         "1.0.0",
         server.WithResourceCapabilities(true, true),
         server.WithLogging(),
+        server.WithRecovery(),
     )
 
     // Add a calculator tool
@@ -522,6 +527,12 @@ initialization.
 Add the `Hooks` to the server at the time of creation using the
 `server.WithHooks` option.
 
+### Tool Handler Middleware
+
+Add middleware to tool call handlers using the `server.WithToolHandlerMiddleware` option. Middlewares can be registered on server creation and are applied on every tool call.
+
+A recovery middleware option is available to recover from panics in a tool call and can be added to the server with the `server.WithRecovery` option.
+
 ## Contributing
 
 <details>

client/client.go

@@ -3,6 +3,8 @@ package client
 
 import (
 	"context"
+	"encoding/json"
+	"fmt"
 
 	"github.com/mark3labs/mcp-go/mcp"
 )
@@ -18,12 +20,25 @@ type MCPClient interface {
 	// Ping checks if the server is alive
 	Ping(ctx context.Context) error
 
+	// ListResourcesByPage manually list resources by page.
+	ListResourcesByPage(
+		ctx context.Context,
+		request mcp.ListResourcesRequest,
+	) (*mcp.ListResourcesResult, error)
+
 	// ListResources requests a list of available resources from the server
 	ListResources(
 		ctx context.Context,
 		request mcp.ListResourcesRequest,
 	) (*mcp.ListResourcesResult, error)
 
+	// ListResourceTemplatesByPage manually list resource templates by page.
+	ListResourceTemplatesByPage(
+		ctx context.Context,
+		request mcp.ListResourceTemplatesRequest,
+	) (*mcp.ListResourceTemplatesResult,
+		error)
+
 	// ListResourceTemplates requests a list of available resource templates from the server
 	ListResourceTemplates(
 		ctx context.Context,
@@ -43,6 +58,12 @@ type MCPClient interface {
 	// Unsubscribe cancels notifications for a specific resource
 	Unsubscribe(ctx context.Context, request mcp.UnsubscribeRequest) error
 
+	// ListPromptsByPage manually list prompts by page.
+	ListPromptsByPage(
+		ctx context.Context,
+		request mcp.ListPromptsRequest,
+	) (*mcp.ListPromptsResult, error)
+
 	// ListPrompts requests a list of available prompts from the server
 	ListPrompts(
 		ctx context.Context,
@@ -55,6 +76,12 @@ type MCPClient interface {
 		request mcp.GetPromptRequest,
 	) (*mcp.GetPromptResult, error)
 
+	// ListToolsByPage manually list tools by page.
+	ListToolsByPage(
+		ctx context.Context,
+		request mcp.ListToolsRequest,
+	) (*mcp.ListToolsResult, error)
+
 	// ListTools requests a list of available tools from the server
 	ListTools(
 		ctx context.Context,
@@ -82,3 +109,26 @@ type MCPClient interface {
 	// OnNotification registers a handler for notifications
 	OnNotification(handler func(notification mcp.JSONRPCNotification))
 }
+
+type mcpClient interface {
+	MCPClient
+
+	sendRequest(ctx context.Context, method string, params interface{}) (*json.RawMessage, error)
+}
+
+func listByPage[T any](
+	ctx context.Context,
+	client mcpClient,
+	request mcp.PaginatedRequest,
+	method string,
+) (*T, error) {
+	response, err := client.sendRequest(ctx, method, request.Params)
+	if err != nil {
+		return nil, err
+	}
+	var result T
+	if err := json.Unmarshal(*response, &result); err != nil {
+		return nil, fmt.Errorf("failed to unmarshal response: %w", err)
+	}
+	return &result, nil
+}

client/sse.go

@@ -23,20 +23,19 @@ import (
 // while sending requests over regular HTTP POST calls. The client handles
 // automatic reconnection and message routing between requests and responses.
 type SSEMCPClient struct {
-	baseURL        *url.URL
-	endpoint       *url.URL
-	httpClient     *http.Client
-	requestID      atomic.Int64
-	responses      map[int64]chan RPCResponse
-	mu             sync.RWMutex
-	done           chan struct{}
-	initialized    bool
-	notifications  []func(mcp.JSONRPCNotification)
-	notifyMu       sync.RWMutex
-	endpointChan   chan struct{}
-	capabilities   mcp.ServerCapabilities
-	headers        map[string]string
-	sseReadTimeout time.Duration
+	baseURL       *url.URL
+	endpoint      *url.URL
+	httpClient    *http.Client
+	requestID     atomic.Int64
+	responses     map[int64]chan RPCResponse
+	mu            sync.RWMutex
+	done          chan struct{}
+	initialized   bool
+	notifications []func(mcp.JSONRPCNotification)
+	notifyMu      sync.RWMutex
+	endpointChan  chan struct{}
+	capabilities  mcp.ServerCapabilities
+	headers       map[string]string
 }
 
 type ClientOption func(*SSEMCPClient)
@@ -68,13 +67,12 @@ func NewSSEMCPClient(baseURL string, options ...ClientOption) (*SSEMCPClient, er
 	}
 
 	smc := &SSEMCPClient{
-		baseURL:        parsedURL,
-		httpClient:     &http.Client{},
-		responses:      make(map[int64]chan RPCResponse),
-		done:           make(chan struct{}),
-		endpointChan:   make(chan struct{}),
-		sseReadTimeout: 30 * time.Second,
-		headers:        make(map[string]string),
+		baseURL:      parsedURL,
+		httpClient:   &http.Client{},
+		responses:    make(map[int64]chan RPCResponse),
+		done:         make(chan struct{}),
+		endpointChan: make(chan struct{}),
+		headers:      make(map[string]string),
 	}
 
 	for _, opt := range options {
@@ -99,6 +97,9 @@ func (c *SSEMCPClient) Start(ctx context.Context) error {
 	req.Header.Set("Accept", "text/event-stream")
 	req.Header.Set("Cache-Control", "no-cache")
 	req.Header.Set("Connection", "keep-alive")
+	for k, v := range c.headers {
+		req.Header.Set(k, v)
+	}
 
 	resp, err := c.httpClient.Do(req)
 	if err != nil {
@@ -134,12 +135,9 @@ func (c *SSEMCPClient) readSSE(reader io.ReadCloser) {
 	br := bufio.NewReader(reader)
 	var event, data string
 
-	ctx, cancel := context.WithTimeout(context.Background(), c.sseReadTimeout)
-	defer cancel()
-
 	for {
 		select {
-		case <-ctx.Done():
+		case <-c.done:
 			return
 		default:
 			line, err := br.ReadString('\n')
@@ -405,7 +403,7 @@ func (c *SSEMCPClient) Initialize(
 			err,
 		)
 	}
-	resp.Body.Close()
+	defer resp.Body.Close()
 
 	c.initialized = true
 	return &result, nil
@@ -416,42 +414,77 @@ func (c *SSEMCPClient) Ping(ctx context.Context) error {
 	return err
 }
 
-func (c *SSEMCPClient) ListResources(
+// ListResourcesByPage manually list resources by page.
+func (c *SSEMCPClient) ListResourcesByPage(
 	ctx context.Context,
 	request mcp.ListResourcesRequest,
 ) (*mcp.ListResourcesResult, error) {
-	response, err := c.sendRequest(ctx, "resources/list", request.Params)
+	result, err := listByPage[mcp.ListResourcesResult](ctx, c, request.PaginatedRequest, "resources/list")
 	if err != nil {
 		return nil, err
 	}
+	return result, nil
+}
 
-	var result mcp.ListResourcesResult
-	if err := json.Unmarshal(*response, &result); err != nil {
-		return nil, fmt.Errorf("failed to unmarshal response: %w", err)
+func (c *SSEMCPClient) ListResources(
+	ctx context.Context,
+	request mcp.ListResourcesRequest,
+) (*mcp.ListResourcesResult, error) {
+	result, err := c.ListResourcesByPage(ctx, request)
+	if err != nil {
+		return nil, err
 	}
+	for result.NextCursor != "" {
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		default:
+			request.Params.Cursor = result.NextCursor
+			newPageRes, err := c.ListResourcesByPage(ctx, request)
+			if err != nil {
+				return nil, err
+			}
+			result.Resources = append(result.Resources, newPageRes.Resources...)
+			result.NextCursor = newPageRes.NextCursor
+		}
+	}
+	return result, nil
+}
 
-	return &result, nil
+func (c *SSEMCPClient) ListResourceTemplatesByPage(
+	ctx context.Context,
+	request mcp.ListResourceTemplatesRequest,
+) (*mcp.ListResourceTemplatesResult, error) {
+	result, err := listByPage[mcp.ListResourceTemplatesResult](ctx, c, request.PaginatedRequest, "resources/templates/list")
+	if err != nil {
+		return nil, err
+	}
+	return result, nil
 }
 
 func (c *SSEMCPClient) ListResourceTemplates(
 	ctx context.Context,
 	request mcp.ListResourceTemplatesRequest,
 ) (*mcp.ListResourceTemplatesResult, error) {
-	response, err := c.sendRequest(
-		ctx,
-		"resources/templates/list",
-		request.Params,
-	)
+	result, err := c.ListResourceTemplatesByPage(ctx, request)
 	if err != nil {
 		return nil, err
 	}
-
-	var result mcp.ListResourceTemplatesResult
-	if err := json.Unmarshal(*response, &result); err != nil {
-		return nil, fmt.Errorf("failed to unmarshal response: %w", err)
+	for result.NextCursor != "" {
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		default:
+			request.Params.Cursor = result.NextCursor
+			newPageRes, err := c.ListResourceTemplatesByPage(ctx, request)
+			if err != nil {
+				return nil, err
+			}
+			result.ResourceTemplates = append(result.ResourceTemplates, newPageRes.ResourceTemplates...)
+			result.NextCursor = newPageRes.NextCursor
+		}
 	}
-
-	return &result, nil
+	return result, nil
 }
 
 func (c *SSEMCPClient) ReadResource(
@@ -482,21 +515,40 @@ func (c *SSEMCPClient) Unsubscribe(
 	return err
 }
 
-func (c *SSEMCPClient) ListPrompts(
+func (c *SSEMCPClient) ListPromptsByPage(
 	ctx context.Context,
 	request mcp.ListPromptsRequest,
 ) (*mcp.ListPromptsResult, error) {
-	response, err := c.sendRequest(ctx, "prompts/list", request.Params)
+	result, err := listByPage[mcp.ListPromptsResult](ctx, c, request.PaginatedRequest, "prompts/list")
 	if err != nil {
 		return nil, err
 	}
+	return result, nil
+}
 
-	var result mcp.ListPromptsResult
-	if err := json.Unmarshal(*response, &result); err != nil {
-		return nil, fmt.Errorf("failed to unmarshal response: %w", err)
+func (c *SSEMCPClient) ListPrompts(
+	ctx context.Context,
+	request mcp.ListPromptsRequest,
+) (*mcp.ListPromptsResult, error) {
+	result, err := c.ListPromptsByPage(ctx, request)
+	if err != nil {
+		return nil, err
 	}
-
-	return &result, nil
+	for result.NextCursor != "" {
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		default:
+			request.Params.Cursor = result.NextCursor
+			newPageRes, err := c.ListPromptsByPage(ctx, request)
+			if err != nil {
+				return nil, err
+			}
+			result.Prompts = append(result.Prompts, newPageRes.Prompts...)
+			result.NextCursor = newPageRes.NextCursor
+		}
+	}
+	return result, nil
 }
 
 func (c *SSEMCPClient) GetPrompt(
@@ -511,21 +563,40 @@ func (c *SSEMCPClient) GetPrompt(
 	return mcp.ParseGetPromptResult(response)
 }
 
-func (c *SSEMCPClient) ListTools(
+func (c *SSEMCPClient) ListToolsByPage(
 	ctx context.Context,
 	request mcp.ListToolsRequest,
 ) (*mcp.ListToolsResult, error) {
-	response, err := c.sendRequest(ctx, "tools/list", request.Params)
+	result, err := listByPage[mcp.ListToolsResult](ctx, c, request.PaginatedRequest, "tools/list")
 	if err != nil {
 		return nil, err
 	}
+	return result, nil
+}
 
-	var result mcp.ListToolsResult
-	if err := json.Unmarshal(*response, &result); err != nil {
-		return nil, fmt.Errorf("failed to unmarshal response: %w", err)
+func (c *SSEMCPClient) ListTools(
+	ctx context.Context,
+	request mcp.ListToolsRequest,
+) (*mcp.ListToolsResult, error) {
+	result, err := c.ListToolsByPage(ctx, request)
+	if err != nil {
+		return nil, err
 	}
-
-	return &result, nil
+	for result.NextCursor != "" {
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		default:
+			request.Params.Cursor = result.NextCursor
+			newPageRes, err := c.ListToolsByPage(ctx, request)
+			if err != nil {
+				return nil, err
+			}
+			result.Tools = append(result.Tools, newPageRes.Tools...)
+			result.NextCursor = newPageRes.NextCursor
+		}
+	}
+	return result, nil
 }
 
 func (c *SSEMCPClient) CallTool(