From b774809086db4bcf9262d40049ff82855da5ab5a Mon Sep 17 00:00:00 2001 From: Kason Braley <59150626+KasonBraley@users.noreply.github.com> Date: Fri, 5 Sep 2025 04:58:26 -0700 Subject: [PATCH] examples: Update and fix outdated (#420) Update examples/server/rate-limiting to use the same replace directive strategy as the auth-middleware example, and fix the breakages. Also remove the unneeded type parameters in examples/client/middleware. --- examples/client/middleware/main.go | 4 +-- examples/server/rate-limiting/go.mod | 11 +++++--- examples/server/rate-limiting/go.sum | 6 +++-- examples/server/rate-limiting/main.go | 36 +++++++++++++-------------- 4 files changed, 32 insertions(+), 25 deletions(-) diff --git a/examples/client/middleware/main.go b/examples/client/middleware/main.go index 6ae87df..9b6d1bb 100644 --- a/examples/client/middleware/main.go +++ b/examples/client/middleware/main.go @@ -17,10 +17,10 @@ var nextProgressToken atomic.Int64 // from the client. func main() { c := mcp.NewClient(&mcp.Implementation{Name: "test"}, nil) - c.AddSendingMiddleware(addProgressToken[*mcp.ClientSession]) + c.AddSendingMiddleware(addProgressToken) } -func addProgressToken[S mcp.Session](h mcp.MethodHandler) mcp.MethodHandler { +func addProgressToken(h mcp.MethodHandler) mcp.MethodHandler { return func(ctx context.Context, method string, req mcp.Request) (result mcp.Result, err error) { if rp, ok := req.GetParams().(mcp.RequestParams); ok { rp.SetProgressToken(nextProgressToken.Add(1)) diff --git a/examples/server/rate-limiting/go.mod b/examples/server/rate-limiting/go.mod index 5b76b16..2b4b00d 100644 --- a/examples/server/rate-limiting/go.mod +++ b/examples/server/rate-limiting/go.mod @@ -2,9 +2,14 @@ module github.com/modelcontextprotocol/go-sdk/examples/rate-limiting go 1.23.0 -toolchain go1.24.4 - require ( - github.com/modelcontextprotocol/go-sdk v0.1.0 + github.com/modelcontextprotocol/go-sdk v0.3.0 golang.org/x/time v0.12.0 ) + +require ( + github.com/google/jsonschema-go v0.2.1-0.20250825175020-748c325cec76 // indirect + github.com/yosida95/uritemplate/v3 v3.0.2 // indirect +) + +replace github.com/modelcontextprotocol/go-sdk => ../../../ diff --git a/examples/server/rate-limiting/go.sum b/examples/server/rate-limiting/go.sum index d73f0a5..16b70ef 100644 --- a/examples/server/rate-limiting/go.sum +++ b/examples/server/rate-limiting/go.sum @@ -1,7 +1,9 @@ github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= -github.com/modelcontextprotocol/go-sdk v0.1.0 h1:ItzbFWYNt4EHcUrScX7P8JPASn1FVYb29G773Xkl+IU= -github.com/modelcontextprotocol/go-sdk v0.1.0/go.mod h1:DcXfbr7yl7e35oMpzHfKw2nUYRjhIGS2uou/6tdsTB0= +github.com/google/jsonschema-go v0.2.1-0.20250825175020-748c325cec76 h1:mBlBwtDebdDYr+zdop8N62a44g+Nbv7o2KjWyS1deR4= +github.com/google/jsonschema-go v0.2.1-0.20250825175020-748c325cec76/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE= +github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4= +github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4= golang.org/x/time v0.12.0 h1:ScB/8o8olJvc+CQPWrK3fPZNfh7qgwCrY0zJmoEQLSE= golang.org/x/time v0.12.0/go.mod h1:CDIdPxbZBQxdj6cxyCIdrNogrJKMJ7pr37NYpMcMDSg= golang.org/x/tools v0.34.0 h1:qIpSLOxeCYGg9TrcJokLBG4KFA6d795g0xkBkiESGlo= diff --git a/examples/server/rate-limiting/main.go b/examples/server/rate-limiting/main.go index c3265c4..e107183 100644 --- a/examples/server/rate-limiting/main.go +++ b/examples/server/rate-limiting/main.go @@ -18,13 +18,13 @@ import ( // GlobalRateLimiterMiddleware creates a middleware that applies a global rate limit. // Every request attempting to pass through will try to acquire a token. // If a token cannot be acquired immediately, the request will be rejected. -func GlobalRateLimiterMiddleware[S mcp.Session](limiter *rate.Limiter) mcp.Middleware[S] { - return func(next mcp.MethodHandler[S]) mcp.MethodHandler[S] { - return func(ctx context.Context, session S, method string, params mcp.Params) (mcp.Result, error) { +func GlobalRateLimiterMiddleware(limiter *rate.Limiter) mcp.Middleware { + return func(next mcp.MethodHandler) mcp.MethodHandler { + return func(ctx context.Context, method string, req mcp.Request) (mcp.Result, error) { if !limiter.Allow() { return nil, errors.New("JSON RPC overloaded") } - return next(ctx, session, method, params) + return next(ctx, method, req) } } } @@ -32,40 +32,40 @@ func GlobalRateLimiterMiddleware[S mcp.Session](limiter *rate.Limiter) mcp.Middl // PerMethodRateLimiterMiddleware creates a middleware that applies rate limiting // on a per-method basis. // Methods not specified in limiters will not be rate limited by this middleware. -func PerMethodRateLimiterMiddleware[S mcp.Session](limiters map[string]*rate.Limiter) mcp.Middleware[S] { - return func(next mcp.MethodHandler[S]) mcp.MethodHandler[S] { - return func(ctx context.Context, session S, method string, params mcp.Params) (mcp.Result, error) { +func PerMethodRateLimiterMiddleware(limiters map[string]*rate.Limiter) mcp.Middleware { + return func(next mcp.MethodHandler) mcp.MethodHandler { + return func(ctx context.Context, method string, req mcp.Request) (mcp.Result, error) { if limiter, ok := limiters[method]; ok { if !limiter.Allow() { return nil, errors.New("JSON RPC overloaded") } } - return next(ctx, session, method, params) + return next(ctx, method, req) } } } // PerSessionRateLimiterMiddleware creates a middleware that applies rate limiting // on a per-session basis for receiving requests. -func PerSessionRateLimiterMiddleware[S mcp.Session](limit rate.Limit, burst int) mcp.Middleware[S] { +func PerSessionRateLimiterMiddleware(limit rate.Limit, burst int) mcp.Middleware { // A map to store limiters, keyed by the session ID. var ( sessionLimiters = make(map[string]*rate.Limiter) mu sync.Mutex ) - return func(next mcp.MethodHandler[S]) mcp.MethodHandler[S] { - return func(ctx context.Context, session S, method string, params mcp.Params) (mcp.Result, error) { + return func(next mcp.MethodHandler) mcp.MethodHandler { + return func(ctx context.Context, method string, req mcp.Request) (mcp.Result, error) { // It's possible that session.ID() may be empty at this point in time // for some transports (e.g., stdio) or until the MCP initialize handshake // has completed. - sessionID := session.ID() + sessionID := req.GetSession().ID() if sessionID == "" { // In this situation, you could apply a single global identifier // if session ID is empty or bypass the rate limiter. // In this example, we bypass the rate limiter. log.Printf("Warning: Session ID is empty for method %q. Skipping per-session rate limiting.", method) - return next(ctx, session, method, params) // Skip limiting if ID is unavailable + return next(ctx, method, req) // Skip limiting if ID is unavailable } mu.Lock() limiter, ok := sessionLimiters[sessionID] @@ -77,19 +77,19 @@ func PerSessionRateLimiterMiddleware[S mcp.Session](limit rate.Limit, burst int) if !limiter.Allow() { return nil, errors.New("JSON RPC overloaded") } - return next(ctx, session, method, params) + return next(ctx, method, req) } } } func main() { - server := mcp.NewServer("greeter1", "v0.0.1", nil) - server.AddReceivingMiddleware(GlobalRateLimiterMiddleware[*mcp.ServerSession](rate.NewLimiter(rate.Every(time.Second/5), 10))) - server.AddReceivingMiddleware(PerMethodRateLimiterMiddleware[*mcp.ServerSession](map[string]*rate.Limiter{ + server := mcp.NewServer(&mcp.Implementation{Name: "greeter1", Version: "v0.0.1"}, nil) + server.AddReceivingMiddleware(GlobalRateLimiterMiddleware(rate.NewLimiter(rate.Every(time.Second/5), 10))) + server.AddReceivingMiddleware(PerMethodRateLimiterMiddleware(map[string]*rate.Limiter{ "callTool": rate.NewLimiter(rate.Every(time.Second), 5), // once a second with a burst up to 5 "listTools": rate.NewLimiter(rate.Every(time.Minute), 20), // once a minute with a burst up to 20 })) - server.AddReceivingMiddleware(PerSessionRateLimiterMiddleware[*mcp.ServerSession](rate.Every(time.Second/5), 10)) + server.AddReceivingMiddleware(PerSessionRateLimiterMiddleware(rate.Every(time.Second/5), 10)) // Run Server logic. log.Println("MCP Server instance created with Middleware (but not running).") log.Println("This example demonstrates configuration, not live interaction.")