modelcontextprotocol/go-sdk
Watch
1
0
Fork
You've already forked go-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

examples: Update and fix outdated (#420)

Browse source 
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.
This commit is contained in:
Kason Braley 2025-09-05 04:58:26 -07:00 • committed by GitHub
commit b774809086
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 32 additions and 25 deletions
Download patch file Download diff file Expand all files Collapse all files

4
examples/client/middleware/main.go
View file

@ -17,10 +17,10 @@ var nextProgressToken atomic.Int64
// from the client. // from the client.
func main() { func main() {
c := mcp.NewClient(&mcp.Implementation{Name: "test"}, nil) 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) { return func(ctx context.Context, method string, req mcp.Request) (result mcp.Result, err error) {
if rp, ok := req.GetParams().(mcp.RequestParams); ok { if rp, ok := req.GetParams().(mcp.RequestParams); ok {
rp.SetProgressToken(nextProgressToken.Add(1)) rp.SetProgressToken(nextProgressToken.Add(1))

11
examples/server/rate-limiting/go.mod
View file

@ -2,9 +2,14 @@ module github.com/modelcontextprotocol/go-sdk/examples/rate-limiting
go 1.23.0 go 1.23.0
toolchain go1.24.4
require ( require (
github.com/modelcontextprotocol/go-sdk v0.1.0 github.com/modelcontextprotocol/go-sdk v0.3.0
golang.org/x/time v0.12.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 => ../../../

6
examples/server/rate-limiting/go.sum
View file

@ -1,7 +1,9 @@
github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= 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/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
github.com/modelcontextprotocol/go-sdk v0.1.0 h1:ItzbFWYNt4EHcUrScX7P8JPASn1FVYb29G773Xkl+IU= github.com/google/jsonschema-go v0.2.1-0.20250825175020-748c325cec76 h1:mBlBwtDebdDYr+zdop8N62a44g+Nbv7o2KjWyS1deR4=
github.com/modelcontextprotocol/go-sdk v0.1.0/go.mod h1:DcXfbr7yl7e35oMpzHfKw2nUYRjhIGS2uou/6tdsTB0= 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 h1:ScB/8o8olJvc+CQPWrK3fPZNfh7qgwCrY0zJmoEQLSE=
golang.org/x/time v0.12.0/go.mod h1:CDIdPxbZBQxdj6cxyCIdrNogrJKMJ7pr37NYpMcMDSg= golang.org/x/time v0.12.0/go.mod h1:CDIdPxbZBQxdj6cxyCIdrNogrJKMJ7pr37NYpMcMDSg=
golang.org/x/tools v0.34.0 h1:qIpSLOxeCYGg9TrcJokLBG4KFA6d795g0xkBkiESGlo= golang.org/x/tools v0.34.0 h1:qIpSLOxeCYGg9TrcJokLBG4KFA6d795g0xkBkiESGlo=

36
examples/server/rate-limiting/main.go
View file

@ -18,13 +18,13 @@ import (
// GlobalRateLimiterMiddleware creates a middleware that applies a global rate limit. // GlobalRateLimiterMiddleware creates a middleware that applies a global rate limit.
// Every request attempting to pass through will try to acquire a token. // Every request attempting to pass through will try to acquire a token.
// If a token cannot be acquired immediately, the request will be rejected. // If a token cannot be acquired immediately, the request will be rejected.
func GlobalRateLimiterMiddleware[S mcp.Session](limiter *rate.Limiter) mcp.Middleware[S] { func GlobalRateLimiterMiddleware(limiter *rate.Limiter) mcp.Middleware {
return func(next mcp.MethodHandler[S]) mcp.MethodHandler[S] { return func(next mcp.MethodHandler) mcp.MethodHandler {
return func(ctx context.Context, session S, method string, params mcp.Params) (mcp.Result, error) { return func(ctx context.Context, method string, req mcp.Request) (mcp.Result, error) {
if !limiter.Allow() { if !limiter.Allow() {
return nil, errors.New("JSON RPC overloaded") 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 // PerMethodRateLimiterMiddleware creates a middleware that applies rate limiting
// on a per-method basis. // on a per-method basis.
// Methods not specified in limiters will not be rate limited by this middleware. // 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] { func PerMethodRateLimiterMiddleware(limiters map[string]*rate.Limiter) mcp.Middleware {
return func(next mcp.MethodHandler[S]) mcp.MethodHandler[S] { return func(next mcp.MethodHandler) mcp.MethodHandler {
return func(ctx context.Context, session S, method string, params mcp.Params) (mcp.Result, error) { return func(ctx context.Context, method string, req mcp.Request) (mcp.Result, error) {
if limiter, ok := limiters[method]; ok { if limiter, ok := limiters[method]; ok {
if !limiter.Allow() { if !limiter.Allow() {
return nil, errors.New("JSON RPC overloaded") 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 // PerSessionRateLimiterMiddleware creates a middleware that applies rate limiting
// on a per-session basis for receiving requests. // 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. // A map to store limiters, keyed by the session ID.
var ( var (
sessionLimiters = make(map[string]*rate.Limiter) sessionLimiters = make(map[string]*rate.Limiter)
mu sync.Mutex mu sync.Mutex
) )
return func(next mcp.MethodHandler[S]) mcp.MethodHandler[S] { return func(next mcp.MethodHandler) mcp.MethodHandler {
return func(ctx context.Context, session S, method string, params mcp.Params) (mcp.Result, error) { 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 // 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 // for some transports (e.g., stdio) or until the MCP initialize handshake
// has completed. // has completed.
sessionID := session.ID() sessionID := req.GetSession().ID()
if sessionID == "" { if sessionID == "" {
// In this situation, you could apply a single global identifier // In this situation, you could apply a single global identifier
// if session ID is empty or bypass the rate limiter. // if session ID is empty or bypass the rate limiter.
// In this example, we 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) 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() mu.Lock()
limiter, ok := sessionLimiters[sessionID] limiter, ok := sessionLimiters[sessionID]
@ -77,19 +77,19 @@ func PerSessionRateLimiterMiddleware[S mcp.Session](limit rate.Limit, burst int)
if !limiter.Allow() { if !limiter.Allow() {
return nil, errors.New("JSON RPC overloaded") return nil, errors.New("JSON RPC overloaded")
} }
return next(ctx, session, method, params) return next(ctx, method, req)
} }
} }
} }
func main() { func main() {
server := mcp.NewServer("greeter1", "v0.0.1", nil) server := mcp.NewServer(&mcp.Implementation{Name: "greeter1", Version: "v0.0.1"}, nil)
server.AddReceivingMiddleware(GlobalRateLimiterMiddleware[*mcp.ServerSession](rate.NewLimiter(rate.Every(time.Second/5), 10))) server.AddReceivingMiddleware(GlobalRateLimiterMiddleware(rate.NewLimiter(rate.Every(time.Second/5), 10)))
server.AddReceivingMiddleware(PerMethodRateLimiterMiddleware[*mcp.ServerSession](map[string]*rate.Limiter{ server.AddReceivingMiddleware(PerMethodRateLimiterMiddleware(map[string]*rate.Limiter{
"callTool": rate.NewLimiter(rate.Every(time.Second), 5), // once a second with a burst up to 5 "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 "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. // Run Server logic.
log.Println("MCP Server instance created with Middleware (but not running).") log.Println("MCP Server instance created with Middleware (but not running).")
log.Println("This example demonstrates configuration, not live interaction.") log.Println("This example demonstrates configuration, not live interaction.")