diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md new file mode 100644 index 00000000..be13052c --- /dev/null +++ b/.github/pull_request_template.md @@ -0,0 +1,12 @@ +### PR Tips + +Typically, PRs should consist of a single commit, and so should generally follow +the [rules for Go commit messages](https://go.dev/wiki/CommitMessage), with the following +changes and additions: + +- Markdown is allowed. + +- For a pervasive change, use "all" in the title instead of a package name. + +- The PR description should provide context (why this change?) and describe the changes + at a high level. Changes that are obvious from the diffs don't need to be mentioned. diff --git a/.github/workflows/readme-check.yml b/.github/workflows/readme-check.yml new file mode 100644 index 00000000..b3398ead --- /dev/null +++ b/.github/workflows/readme-check.yml @@ -0,0 +1,38 @@ +name: README Check +on: + workflow_dispatch: + pull_request: + paths: + - 'internal/readme/**' + - 'README.md' + +permissions: + contents: read + +jobs: + readme-check: + runs-on: ubuntu-latest + steps: + - name: Set up Go + uses: actions/setup-go@v5 + - name: Check out code + uses: actions/checkout@v4 + - name: Check README is up-to-date + run: | + cd internal/readme + make + if [ -n "$(git status --porcelain)" ]; then + echo "ERROR: README.md is not up-to-date!" + echo "" + echo "The README.md file differs from what would be generated by running 'make' in internal/readme/." + echo "Please update internal/readme/README.src.md instead of README.md directly," + echo "then run 'make' in the internal/readme/ directory to regenerate README.md." + echo "" + echo "Changes:" + git status --porcelain + echo "" + echo "Diff:" + git diff + exit 1 + fi + echo "README.md is up-to-date" diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index b603512d..f2f8cb8e 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -8,8 +8,25 @@ on: permissions: contents: read - + jobs: + lint: + runs-on: ubuntu-latest + steps: + - name: Set up Go + uses: actions/setup-go@v5 + - name: Check out code + uses: actions/checkout@v4 + - name: Check formatting + run: | + unformatted=$(gofmt -l .) + if [ -n "$unformatted" ]; then + echo "The following files are not properly formatted:" + echo "$unformatted" + exit 1 + fi + echo "All Go files are properly formatted" + test: runs-on: ubuntu-latest strategy: diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 0c876da3..40739735 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -68,6 +68,15 @@ This process is similar to the [Go proposal process](https://github.com/golang/proposal), but is necessarily lighter weight to accommodate the greater rate of change expected for the SDK. +### Design discussion + +For open ended design discussion (anything that doesn't fall into the issue +categories above), use [GitHub +Discussions](https://github.com/modelcontextprotocol/go-sdk/discussions). +Ideally, each discussion should be focused on one aspect of the design. For +example: Tool Binding and Session APIs would be two separate discussions. +When discussions reach a consensus, they should be promoted into proposals. + ## Contributing code The project uses GitHub pull requests (PRs) to review changes. @@ -96,6 +105,19 @@ copyright header following the format below: // license that can be found in the LICENSE file. ``` +### Updating the README + +The top-level `README.md` file is generated from `internal/readme/README.src.md` +and should not be edited directly. To update the README: + +1. Make your changes to `internal/readme/README.src.md` +2. Run `make` in the `internal/readme/` directory to regenerate `README.md` +3. Commit both files together + +The CI system will automatically check that the README is up-to-date by running +`make` and verifying no changes result. If you see a CI failure about the +README being out of sync, follow the steps above to regenerate it. + ## Code of conduct This project follows the [Go Community Code of Conduct](https://go.dev/conduct). diff --git a/README.md b/README.md index 504eccc6..d4900674 100644 --- a/README.md +++ b/README.md @@ -1,10 +1,7 @@ # MCP Go SDK - - -[![PkgGoDev](https://pkg.go.dev/badge/golang.org/x/tools)](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk) +[![PkgGoDev](https://pkg.go.dev/badge/github.com/modelcontextprotocol/go-sdk)](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk) This repository contains an unreleased implementation of the official Go software development kit (SDK) for the Model Context Protocol (MCP). @@ -15,6 +12,18 @@ proposals, but don't use it in real projects. See the issue tracker for known issues and missing features. We aim to release a stable version of the SDK in August, 2025. +## Design + +The design doc for this SDK is at [design.md](./design/design.md), which was +initially reviewed at +[modelcontextprotocol/discussions/364](https://github.com/orgs/modelcontextprotocol/discussions/364). + +Further design discussion should occur in +[issues](https://github.com/modelcontextprotocol/go-sdk/issues) (for concrete +proposals) or +[discussions](https://github.com/modelcontextprotocol/go-sdk/discussions) for +open-ended discussion. See CONTRIBUTING.md for details. + ## Package documentation The SDK consists of two importable packages: @@ -71,7 +80,7 @@ func main() { log.Fatal("tool failed") } for _, c := range res.Content { - log.Print(c.Text) + log.Print(c.(*mcp.TextContent).Text) } } ``` @@ -95,14 +104,18 @@ type HiParams struct { func SayHi(ctx context.Context, cc *mcp.ServerSession, params *mcp.CallToolParamsFor[HiParams]) (*mcp.CallToolResultFor[any], error) { return &mcp.CallToolResultFor[any]{ - Content: []*mcp.ContentBlock{mcp.NewTextContent("Hi " + params.Name)}, + Content: []mcp.Content{&mcp.TextContent{Text: "Hi " + params.Arguments.Name}}, }, nil } func main() { // Create a server with a single tool. server := mcp.NewServer("greeter", "v1.0.0", nil) - server.AddTools(mcp.NewServerTool("greet", "say hi", SayHi)) + server.AddTools( + mcp.NewServerTool("greet", "say hi", SayHi, mcp.Input( + mcp.Property("name", mcp.Description("the name of the person to greet")), + )), + ) // Run the server over stdin/stdout, until the client disconnects if err := server.Run(context.Background(), mcp.NewStdioTransport()); err != nil { log.Fatal(err) @@ -112,15 +125,6 @@ func main() { The `examples/` directory contains more example clients and servers. -## Design - -The design doc for this SDK is at [design.md](./design/design.md), which was -initially reviewed at -[modelcontextprotocol/discussions/364](https://github.com/orgs/modelcontextprotocol/discussions/364). - -Further design discussion should occur in GitHub issues. See CONTRIBUTING.md -for details. - ## Acknowledgements Several existing Go MCP SDKs inspired the development and design of this diff --git a/design/design.md b/design/design.md index 5c77b7de..b52e9c10 100644 --- a/design/design.md +++ b/design/design.md @@ -470,6 +470,10 @@ server.AddReceivingMiddleware(withLogging) **Differences from mcp-go**: Version 0.26.0 of mcp-go defines 24 server hooks. Each hook consists of a field in the `Hooks` struct, a `Hooks.Add` method, and a type for the hook function. These are rarely used. The most common is `OnError`, which occurs fewer than ten times in open-source code. +#### Rate Limiting + +Rate limiting can be configured using middleware. Please see [examples/rate-limiting](] for an example on how to implement this. + ### Errors With the exception of tool handler errors, protocol errors are handled transparently as Go errors: errors in server-side feature handlers are propagated as errors from calls from the `ClientSession`, and vice-versa. @@ -850,7 +854,30 @@ type ClientOptions struct { ### Completion -Clients call the spec method `Complete` to request completions. Servers automatically handle these requests based on their collections of prompts and resources. +Clients call the spec method `Complete` to request completions. If a server installs a `CompletionHandler`, it will be called when the client sends a completion request. + +```go +// A CompletionHandler handles a call to completion/complete. +type CompletionHandler func(context.Context, *ServerSession, *CompleteParams) (*CompleteResult, error) + +type ServerOptions struct { + ... + // If non-nil, called when a client sends a completion request. + CompletionHandler CompletionHandler +} +``` + +#### Security Considerations + +Implementors of the CompletionHandler MUST adhere to the following security guidelines: + +- **Validate all completion inputs**: The CompleteRequest received by the handler may contain arbitrary data from the client. Before processing, thoroughly validate all fields. + +- **Implement appropriate rate limiting**: Completion requests can be high-frequency, especially in interactive user experiences. Without rate limiting, a malicious client could potentially overload the server, leading to denial-of-service (DoS) attacks. Consider applying rate limits per client session, IP address, or API key, depending on your deployment model. This can be implemented within the CompletionHandler itself or via middleware (see [Middleware](#middleware)) that wraps the handler. + +- **Control access to sensitive suggestions**: Completion suggestions should only expose information that the requesting client is authorized to access. If your completion logic draws from sensitive data sources (e.g., internal file paths, user data, restricted API endpoints), ensure that the CompletionHandler performs proper authorization checks before generating or returning suggestions. This might involve checking the ServerSession context for authentication details or consulting an external authorization system. + +- **Prevent completion-based information disclosure**: Be mindful that even seemingly innocuous completion suggestions can inadvertently reveal sensitive information. For example, suggesting internal system paths or confidential identifiers could be an attack vector. Ensure that the generated CompleteResult contains only appropriate and non-sensitive information for the given client and context. Avoid revealing internal data structures or error messages that could aid an attacker. **Differences from mcp-go**: the client API is similar. mcp-go has not yet defined its server-side behavior. @@ -929,7 +956,7 @@ In addition to the `List` methods, the SDK provides an iterator method for each # Governance and Community -While the sections above propose an initial implementation of the Go SDK, MCP is evolving rapidly. SDKs need to keep pace, by implementing changes to the spec, fixing bugs, and accomodating new and emerging use-cases. This section proposes how the SDK project can be managed so that it can change safely and transparently. +While the sections above propose an initial implementation of the Go SDK, MCP is evolving rapidly. SDKs need to keep pace, by implementing changes to the spec, fixing bugs, and accommodating new and emerging use-cases. This section proposes how the SDK project can be managed so that it can change safely and transparently. Initially, the Go SDK repository will be administered by the Go team and Anthropic, and they will be the Approvers (the set of people able to merge PRs to the SDK). The policies here are also intended to satisfy necessary constraints of the Go team's participation in the project. @@ -957,7 +984,7 @@ A proposal is an issue that proposes a new API for the SDK, or a change to the s Proposals that are straightforward and uncontroversial may be approved based on GitHub discussion. However, proposals that are deemed to be sufficiently unclear or complicated will be deferred to a regular steering meeting (see below). -This process is similar to the [Go proposal process](https://github.com/golang/proposal), but is necessarily lighter weight to accomodate the greater rate of change expected for the SDK. +This process is similar to the [Go proposal process](https://github.com/golang/proposal), but is necessarily lighter weight to accommodate the greater rate of change expected for the SDK. ### Steering meetings diff --git a/examples/hello/main.go b/examples/hello/main.go index 1abb0e2d..9af34cc3 100644 --- a/examples/hello/main.go +++ b/examples/hello/main.go @@ -25,7 +25,7 @@ type HiArgs struct { func SayHi(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[HiArgs]) (*mcp.CallToolResultFor[struct{}], error) { return &mcp.CallToolResultFor[struct{}]{ Content: []mcp.Content{ - &mcp.TextContent{Text: "Hi " + params.Name}, + &mcp.TextContent{Text: "Hi " + params.Arguments.Name}, }, }, nil } diff --git a/examples/rate-limiting/go.mod b/examples/rate-limiting/go.mod new file mode 100644 index 00000000..5ec49ddc --- /dev/null +++ b/examples/rate-limiting/go.mod @@ -0,0 +1,8 @@ +module github.com/modelcontextprotocol/go-sdk/examples/rate-limiting + +go 1.25 + +require ( + github.com/modelcontextprotocol/go-sdk v0.0.0-20250625185707-09181c2c2e89 + golang.org/x/time v0.12.0 +) diff --git a/examples/rate-limiting/go.sum b/examples/rate-limiting/go.sum new file mode 100644 index 00000000..c7027682 --- /dev/null +++ b/examples/rate-limiting/go.sum @@ -0,0 +1,8 @@ +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.0.0-20250625185707-09181c2c2e89 h1:kUGBYP25FTv3ZRBhLT4iQvtx4FDl7hPkWe3isYrMxyo= +github.com/modelcontextprotocol/go-sdk v0.0.0-20250625185707-09181c2c2e89/go.mod h1:DcXfbr7yl7e35oMpzHfKw2nUYRjhIGS2uou/6tdsTB0= +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= +golang.org/x/tools v0.34.0/go.mod h1:pAP9OwEaY1CAW3HOmg3hLZC5Z0CCmzjAF2UQMSqNARg= diff --git a/examples/rate-limiting/main.go b/examples/rate-limiting/main.go new file mode 100644 index 00000000..7e91b79f --- /dev/null +++ b/examples/rate-limiting/main.go @@ -0,0 +1,54 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +package main + +import ( + "context" + "errors" + "time" + + "github.com/modelcontextprotocol/go-sdk/mcp" + "golang.org/x/time/rate" +) + +// 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) { + if !limiter.Allow() { + return nil, errors.New("JSON RPC overloaded") + } + return next(ctx, session, method, params) + } + } +} + +// 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) { + if limiter, ok := limiters[method]; ok { + if !limiter.Allow() { + return nil, errors.New("JSON RPC overloaded") + } + } + return next(ctx, session, method, params) + } + } +} + +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{ + "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 + })) + // Run Server logic. +} diff --git a/examples/sse/main.go b/examples/sse/main.go index 5e4c851e..97ea1bd0 100644 --- a/examples/sse/main.go +++ b/examples/sse/main.go @@ -22,7 +22,7 @@ type SayHiParams struct { func SayHi(ctx context.Context, cc *mcp.ServerSession, params *mcp.CallToolParamsFor[SayHiParams]) (*mcp.CallToolResultFor[any], error) { return &mcp.CallToolResultFor[any]{ Content: []mcp.Content{ - &mcp.TextContent{Text: "Hi " + params.Name}, + &mcp.TextContent{Text: "Hi " + params.Arguments.Name}, }, }, nil } diff --git a/internal/jsonrpc2/wire.go b/internal/jsonrpc2/wire.go index 58ee41f0..309b8002 100644 --- a/internal/jsonrpc2/wire.go +++ b/internal/jsonrpc2/wire.go @@ -34,7 +34,7 @@ var ( // ErrUnknown should be used for all non coded errors. ErrUnknown = NewError(-32001, "JSON RPC unknown error") // ErrServerClosing is returned for calls that arrive while the server is closing. - ErrServerClosing = NewError(-32002, "JSON RPC server is closing") + ErrServerClosing = NewError(-32004, "JSON RPC server is closing") // ErrClientClosing is a dummy error returned for calls initiated while the client is closing. ErrClientClosing = NewError(-32003, "JSON RPC client is closing") ) diff --git a/internal/readme/README.src.md b/internal/readme/README.src.md index 9938366e..629629a4 100644 --- a/internal/readme/README.src.md +++ b/internal/readme/README.src.md @@ -1,9 +1,6 @@ # MCP Go SDK - - -[![PkgGoDev](https://pkg.go.dev/badge/golang.org/x/tools)](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk) +[![PkgGoDev](https://pkg.go.dev/badge/github.com/modelcontextprotocol/go-sdk)](https://pkg.go.dev/github.com/modelcontextprotocol/go-sdk) This repository contains an unreleased implementation of the official Go software development kit (SDK) for the Model Context Protocol (MCP). @@ -14,6 +11,18 @@ proposals, but don't use it in real projects. See the issue tracker for known issues and missing features. We aim to release a stable version of the SDK in August, 2025. +## Design + +The design doc for this SDK is at [design.md](./design/design.md), which was +initially reviewed at +[modelcontextprotocol/discussions/364](https://github.com/orgs/modelcontextprotocol/discussions/364). + +Further design discussion should occur in +[issues](https://github.com/modelcontextprotocol/go-sdk/issues) (for concrete +proposals) or +[discussions](https://github.com/modelcontextprotocol/go-sdk/discussions) for +open-ended discussion. See CONTRIBUTING.md for details. + ## Package documentation The SDK consists of two importable packages: @@ -41,15 +50,6 @@ with its client over stdin/stdout: The `examples/` directory contains more example clients and servers. -## Design - -The design doc for this SDK is at [design.md](./design/design.md), which was -initially reviewed at -[modelcontextprotocol/discussions/364](https://github.com/orgs/modelcontextprotocol/discussions/364). - -Further design discussion should occur in GitHub issues. See CONTRIBUTING.md -for details. - ## Acknowledgements Several existing Go MCP SDKs inspired the development and design of this diff --git a/internal/readme/server/server.go b/internal/readme/server/server.go index b709cd75..534e0798 100644 --- a/internal/readme/server/server.go +++ b/internal/readme/server/server.go @@ -18,14 +18,18 @@ type HiParams struct { func SayHi(ctx context.Context, cc *mcp.ServerSession, params *mcp.CallToolParamsFor[HiParams]) (*mcp.CallToolResultFor[any], error) { return &mcp.CallToolResultFor[any]{ - Content: []mcp.Content{&mcp.TextContent{Text: "Hi " + params.Name}}, + Content: []mcp.Content{&mcp.TextContent{Text: "Hi " + params.Arguments.Name}}, }, nil } func main() { // Create a server with a single tool. server := mcp.NewServer("greeter", "v1.0.0", nil) - server.AddTools(mcp.NewServerTool("greet", "say hi", SayHi)) + server.AddTools( + mcp.NewServerTool("greet", "say hi", SayHi, mcp.Input( + mcp.Property("name", mcp.Description("the name of the person to greet")), + )), + ) // Run the server over stdin/stdout, until the client disconnects if err := server.Run(context.Background(), mcp.NewStdioTransport()); err != nil { log.Fatal(err) diff --git a/mcp/client.go b/mcp/client.go index 17875c38..856df567 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -9,6 +9,7 @@ import ( "iter" "slices" "sync" + "time" "github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2" ) @@ -56,6 +57,10 @@ type ClientOptions struct { ResourceListChangedHandler func(context.Context, *ClientSession, *ResourceListChangedParams) LoggingMessageHandler func(context.Context, *ClientSession, *LoggingMessageParams) ProgressNotificationHandler func(context.Context, *ClientSession, *ProgressNotificationParams) + // If non-zero, defines an interval for regular "ping" requests. + // If the peer fails to respond to pings originating from the keepalive check, + // the session is automatically closed. + KeepAlive time.Duration } // bind implements the binder[*ClientSession] interface, so that Clients can @@ -118,6 +123,11 @@ func (c *Client) Connect(ctx context.Context, t Transport) (cs *ClientSession, e _ = cs.Close() return nil, err } + + if c.opts.KeepAlive > 0 { + cs.startKeepalive(c.opts.KeepAlive) + } + return cs, nil } @@ -131,12 +141,21 @@ type ClientSession struct { conn *jsonrpc2.Connection client *Client initializeResult *InitializeResult + keepaliveCancel context.CancelFunc } // Close performs a graceful close of the connection, preventing new requests // from being handled, and waiting for ongoing requests to return. Close then // terminates the connection. func (cs *ClientSession) Close() error { + // Note: keepaliveCancel access is safe without a mutex because: + // 1. keepaliveCancel is only written once during startKeepalive (happens-before all Close calls) + // 2. context.CancelFunc is safe to call multiple times and from multiple goroutines + // 3. The keepalive goroutine calls Close on ping failure, but this is safe since + // Close is idempotent and conn.Close() handles concurrent calls correctly + if cs.keepaliveCancel != nil { + cs.keepaliveCancel() + } return cs.conn.Close() } @@ -146,6 +165,11 @@ func (cs *ClientSession) Wait() error { return cs.conn.Wait() } +// startKeepalive starts the keepalive mechanism for this client session. +func (cs *ClientSession) startKeepalive(interval time.Duration) { + startKeepalive(cs, interval, &cs.keepaliveCancel) +} + // AddRoots adds the given roots to the client, // replacing any with the same URIs, // and notifies any connected servers. @@ -167,7 +191,7 @@ func (c *Client) RemoveRoots(uris ...string) { func() bool { return c.roots.remove(uris...) }) } -// changeAndNotifyClient is called when a feature is added or removed. +// changeAndNotify is called when a feature is added or removed. // It calls change, which should do the work and report whether a change actually occurred. // If there was a change, it notifies a snapshot of the sessions. func (c *Client) changeAndNotify(notification string, params Params, change func() bool) { @@ -323,7 +347,7 @@ func (cs *ClientSession) ListResourceTemplates(ctx context.Context, params *List return handleSend[*ListResourceTemplatesResult](ctx, cs, methodListResourceTemplates, orZero[Params](params)) } -// ReadResource ask the server to read a resource and return its contents. +// ReadResource asks the server to read a resource and return its contents. func (cs *ClientSession) ReadResource(ctx context.Context, params *ReadResourceParams) (*ReadResourceResult, error) { return handleSend[*ReadResourceResult](ctx, cs, methodReadResource, orZero[Params](params)) } diff --git a/mcp/client_list_test.go b/mcp/client_list_test.go index 6153ca54..7e6da95a 100644 --- a/mcp/client_list_test.go +++ b/mcp/client_list_test.go @@ -11,8 +11,8 @@ import ( "github.com/google/go-cmp/cmp" "github.com/google/go-cmp/cmp/cmpopts" - "github.com/modelcontextprotocol/go-sdk/mcp" "github.com/modelcontextprotocol/go-sdk/jsonschema" + "github.com/modelcontextprotocol/go-sdk/mcp" ) func TestList(t *testing.T) { diff --git a/mcp/logging.go b/mcp/logging.go index fb8e5719..4880e179 100644 --- a/mcp/logging.go +++ b/mcp/logging.go @@ -137,7 +137,7 @@ func (h *LoggingHandler) WithGroup(name string) slog.Handler { } // Handle implements [slog.Handler.Handle] by writing the Record to a JSONHandler, -// then calling [ServerSession.LoggingMesssage] with the result. +// then calling [ServerSession.LoggingMessage] with the result. func (h *LoggingHandler) Handle(ctx context.Context, r slog.Record) error { err := h.handle(ctx, r) // TODO(jba): find a way to surface the error. @@ -185,5 +185,5 @@ func (h *LoggingHandler) handle(ctx context.Context, r slog.Record) error { // documentation says not to. // In this case logging is a service to clients, not a means for debugging the // server, so we want to cancel the log message. - return h.ss.LoggingMessage(ctx, params) + return h.ss.Log(ctx, params) } diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index 71b99a4a..5f42b1b9 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -370,7 +370,7 @@ func TestEndToEnd(t *testing.T) { mustLog := func(level LoggingLevel, data any) { t.Helper() - if err := ss.LoggingMessage(ctx, &LoggingMessageParams{ + if err := ss.Log(ctx, &LoggingMessageParams{ Logger: "test", Level: level, Data: data, @@ -838,3 +838,101 @@ func falseSchema() *jsonschema.Schema { return &jsonschema.Schema{Not: &jsonsche func nopHandler(context.Context, *ServerSession, *CallToolParamsFor[map[string]any]) (*CallToolResult, error) { return nil, nil } + +func TestKeepAlive(t *testing.T) { + // TODO: try to use the new synctest package for this test once we upgrade to Go 1.24+. + // synctest would allow us to control time and avoid the time.Sleep calls, making the test + // faster and more deterministic. + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + ct, st := NewInMemoryTransports() + + serverOpts := &ServerOptions{ + KeepAlive: 100 * time.Millisecond, + } + s := NewServer("testServer", "v1.0.0", serverOpts) + s.AddTools(NewServerTool("greet", "say hi", sayHi)) + + ss, err := s.Connect(ctx, st) + if err != nil { + t.Fatal(err) + } + defer ss.Close() + + clientOpts := &ClientOptions{ + KeepAlive: 100 * time.Millisecond, + } + c := NewClient("testClient", "v1.0.0", clientOpts) + cs, err := c.Connect(ctx, ct) + if err != nil { + t.Fatal(err) + } + defer cs.Close() + + // Wait for a few keepalive cycles to ensure pings are working + time.Sleep(300 * time.Millisecond) + + // Test that the connection is still alive by making a call + result, err := cs.CallTool(ctx, &CallToolParams{ + Name: "greet", + Arguments: map[string]any{"Name": "user"}, + }) + if err != nil { + t.Fatalf("call failed after keepalive: %v", err) + } + if len(result.Content) == 0 { + t.Fatal("expected content in result") + } + if textContent, ok := result.Content[0].(*TextContent); !ok || textContent.Text != "hi user" { + t.Fatalf("unexpected result: %v", result.Content[0]) + } +} + +func TestKeepAliveFailure(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + ct, st := NewInMemoryTransports() + + // Server without keepalive (to test one-sided keepalive) + s := NewServer("testServer", "v1.0.0", nil) + s.AddTools(NewServerTool("greet", "say hi", sayHi)) + ss, err := s.Connect(ctx, st) + if err != nil { + t.Fatal(err) + } + + // Client with short keepalive + clientOpts := &ClientOptions{ + KeepAlive: 50 * time.Millisecond, + } + c := NewClient("testClient", "v1.0.0", clientOpts) + cs, err := c.Connect(ctx, ct) + if err != nil { + t.Fatal(err) + } + defer cs.Close() + + // Let the connection establish properly first + time.Sleep(30 * time.Millisecond) + + // simulate ping failure + ss.Close() + + // Wait for keepalive to detect the failure and close the client + // check periodically instead of just waiting + deadline := time.Now().Add(1 * time.Second) + for time.Now().Before(deadline) { + _, err = cs.CallTool(ctx, &CallToolParams{ + Name: "greet", + Arguments: map[string]any{"Name": "user"}, + }) + if errors.Is(err, ErrConnectionClosed) { + return // Test passed + } + time.Sleep(25 * time.Millisecond) + } + + t.Errorf("expected connection to be closed by keepalive, but it wasn't. Last error: %v", err) +} diff --git a/mcp/protocol.go b/mcp/protocol.go index 439d07a0..babfd55a 100644 --- a/mcp/protocol.go +++ b/mcp/protocol.go @@ -22,7 +22,7 @@ type Annotations struct { // Describes who the intended customer of this object or data is. // // It can include multiple entries to indicate content useful for multiple - // audiences (e.g., `["user", "assistant"]`). + // audiences (e.g., []Role{"user", "assistant"}). Audience []Role `json:"audience,omitempty"` // The moment the resource was last modified, as an ISO 8601 formatted string. // @@ -70,12 +70,12 @@ type CallToolResult struct { // // If not set, this is assumed to be false (the call was successful). // - // Any errors that originate from the tool SHOULD be reported inside the result - // object, with `isError` set to true, _not_ as an MCP protocol-level error + // Any errors that originate from the tool should be reported inside the result + // object, with isError set to true, not as an MCP protocol-level error // response. Otherwise, the LLM would not be able to see that an error occurred // and self-correct. // - // However, any errors in _finding_ the tool, an error indicating that the + // However, any errors in finding the tool, an error indicating that the // server does not support tool calls, or any other exceptional conditions, // should be reported as an MCP error response. IsError bool `json:"isError,omitempty"` @@ -115,8 +115,8 @@ type CallToolResultFor[Out any] struct { // // If not set, this is assumed to be false (the call was successful). // - // Any errors that originate from the tool SHOULD be reported inside the result - // object, with `isError` set to true, not as an MCP protocol-level error + // Any errors that originate from the tool should be reported inside the result + // object, with isError set to true, not as an MCP protocol-level error // response. Otherwise, the LLM would not be able to see that an error occurred // and self-correct. // @@ -129,16 +129,35 @@ type CallToolResultFor[Out any] struct { StructuredContent Out `json:"structuredContent,omitempty"` } +// UnmarshalJSON handles the unmarshalling of content into the Content +// interface. +func (x *CallToolResultFor[Out]) UnmarshalJSON(data []byte) error { + type res CallToolResultFor[Out] // avoid recursion + var wire struct { + res + Content []*wireContent `json:"content"` + } + if err := json.Unmarshal(data, &wire); err != nil { + return err + } + var err error + if wire.res.Content, err = contentsFromWire(wire.Content, nil); err != nil { + return err + } + *x = CallToolResultFor[Out](wire.res) + return nil +} + type CancelledParams struct { // This property is reserved by the protocol to allow clients and servers to // attach additional metadata to their responses. Meta `json:"_meta,omitempty"` - // An optional string describing the reason for the cancellation. This MAY be + // An optional string describing the reason for the cancellation. This may be // logged or presented to the user. Reason string `json:"reason,omitempty"` // The ID of the request to cancel. // - // This MUST correspond to the ID of a request previously issued in the same + // This must correspond to the ID of a request previously issued in the same // direction. RequestID any `json:"requestId"` } @@ -168,21 +187,21 @@ type CreateMessageParams struct { // attach additional metadata to their responses. Meta `json:"_meta,omitempty"` // A request to include context from one or more MCP servers (including the - // caller), to be attached to the prompt. The client MAY ignore this request. + // caller), to be attached to the prompt. The client may ignore this request. IncludeContext string `json:"includeContext,omitempty"` // The maximum number of tokens to sample, as requested by the server. The - // client MAY choose to sample fewer tokens than requested. + // client may choose to sample fewer tokens than requested. MaxTokens int64 `json:"maxTokens"` Messages []*SamplingMessage `json:"messages"` // Optional metadata to pass through to the LLM provider. The format of this // metadata is provider-specific. Metadata struct{} `json:"metadata,omitempty"` - // The server's preferences for which model to select. The client MAY ignore + // The server's preferences for which model to select. The client may ignore // these preferences. ModelPreferences *ModelPreferences `json:"modelPreferences,omitempty"` StopSequences []string `json:"stopSequences,omitempty"` // An optional system prompt the server wants to use for sampling. The client - // MAY modify or omit this prompt. + // may modify or omit this prompt. SystemPrompt string `json:"systemPrompt,omitempty"` Temperature float64 `json:"temperature,omitempty"` } @@ -236,7 +255,7 @@ type InitializeParams struct { Capabilities *ClientCapabilities `json:"capabilities"` ClientInfo *implementation `json:"clientInfo"` // The latest version of the Model Context Protocol that the client supports. - // The client MAY decide to support older versions as well. + // The client may decide to support older versions as well. ProtocolVersion string `json:"protocolVersion"` } @@ -254,11 +273,11 @@ type InitializeResult struct { // // This can be used by clients to improve the LLM's understanding of available // tools, resources, etc. It can be thought of like a "hint" to the model. For - // example, this information MAY be added to the system prompt. + // example, this information may be added to the system prompt. Instructions string `json:"instructions,omitempty"` // The version of the Model Context Protocol that the server wants to use. This // may not match the version that the client requested. If the client cannot - // support this version, it MUST disconnect. + // support this version, it must disconnect. ProtocolVersion string `json:"protocolVersion"` ServerInfo *implementation `json:"serverInfo"` } @@ -424,12 +443,12 @@ func (x *LoggingMessageParams) SetProgressToken(t any) { setProgressToken(x, t) type ModelHint struct { // A hint for a model name. // - // The client SHOULD treat this as a substring of a model name; for example: - + // The client should treat this as a substring of a model name; for example: - // `claude-3-5-sonnet` should match `claude-3-5-sonnet-20241022` - `sonnet` // should match `claude-3-5-sonnet-20241022`, `claude-3-sonnet-20240229`, etc. - // `claude` should match any Claude model // - // The client MAY also map the string to a different provider's model name or a + // The client may also map the string to a different provider's model name or a // different model family, as long as it fills a similar niche; for example: - // `gemini-1.5-flash` could match `claude-3-haiku-20240307` Name string `json:"name,omitempty"` @@ -444,7 +463,7 @@ type ModelHint struct { // on. This interface allows servers to express their priorities across multiple // dimensions to help clients make an appropriate selection for their use case. // -// These preferences are always advisory. The client MAY ignore them. It is also +// These preferences are always advisory. The client may ignore them. It is also // up to the client to decide how to interpret these preferences and how to // balance them against other considerations. type ModelPreferences struct { @@ -453,10 +472,10 @@ type ModelPreferences struct { CostPriority float64 `json:"costPriority,omitempty"` // Optional hints to use for model selection. // - // If multiple hints are specified, the client MUST evaluate them in order (such + // If multiple hints are specified, the client must evaluate them in order (such // that the first match is taken). // - // The client SHOULD prioritize these hints over the numeric priorities, but MAY + // The client should prioritize these hints over the numeric priorities, but may // still use the priorities to select from ambiguous matches. Hints []*ModelHint `json:"hints,omitempty"` // How much to prioritize intelligence and capabilities when selecting a model. @@ -536,7 +555,7 @@ func (x *PromptListChangedParams) SetProgressToken(t any) { setProgressToken(x, // Describes a message returned as part of a prompt. // -// This is similar to `SamplingMessage`, but also supports the embedding of +// This is similar to SamplingMessage, but also supports the embedding of // resources from the MCP server. type PromptMessage struct { Content Content `json:"content"` @@ -609,7 +628,7 @@ type Resource struct { // easily understood, even by those unfamiliar with domain-specific terminology. // // If not provided, the name should be used for display (except for Tool, where - // `annotations.title` should be given precedence over using `name`, if + // Annotations.Title should be given precedence over using name, if // present). Title string `json:"title,omitempty"` // The URI of this resource. @@ -647,7 +666,7 @@ type ResourceTemplate struct { // easily understood, even by those unfamiliar with domain-specific terminology. // // If not provided, the name should be used for display (except for Tool, where - // `annotations.title` should be given precedence over using `name`, if + // Annotations.Title should be given precedence over using name, if // present). Title string `json:"title,omitempty"` // A URI template (according to RFC 6570) that can be used to construct resource @@ -757,9 +776,9 @@ type Tool struct { // Additional properties describing a Tool to clients. // -// NOTE: all properties in ToolAnnotations are **hints**. They are not +// NOTE: all properties in ToolAnnotations are hints. They are not // guaranteed to provide a faithful description of tool behavior (including -// descriptive properties like `title`). +// descriptive properties like title). // // Clients should never make tool use decisions based on ToolAnnotations // received from untrusted servers. @@ -767,14 +786,14 @@ type ToolAnnotations struct { // If true, the tool may perform destructive updates to its environment. If // false, the tool performs only additive updates. // - // (This property is meaningful only when `readOnlyHint == false`) + // (This property is meaningful only when ReadOnlyHint == false.) // // Default: true - DestructiveHint bool `json:"destructiveHint,omitempty"` + DestructiveHint *bool `json:"destructiveHint,omitempty"` // If true, calling the tool repeatedly with the same arguments will have no // additional effect on the its environment. // - // (This property is meaningful only when `readOnlyHint == false`) + // (This property is meaningful only when ReadOnlyHint == false.) // // Default: false IdempotentHint bool `json:"idempotentHint,omitempty"` @@ -783,7 +802,7 @@ type ToolAnnotations struct { // a web search tool is open, whereas that of a memory tool is not. // // Default: true - OpenWorldHint bool `json:"openWorldHint,omitempty"` + OpenWorldHint *bool `json:"openWorldHint,omitempty"` // If true, the tool does not modify its environment. // // Default: false diff --git a/mcp/protocol_test.go b/mcp/protocol_test.go index d68fb738..823c409c 100644 --- a/mcp/protocol_test.go +++ b/mcp/protocol_test.go @@ -8,6 +8,8 @@ import ( "encoding/json" "maps" "testing" + + "github.com/google/go-cmp/cmp" ) func TestParamsMeta(t *testing.T) { @@ -67,3 +69,47 @@ func TestParamsMeta(t *testing.T) { p.SetProgressToken(int32(1)) p.SetProgressToken(int64(1)) } + +func TestContentUnmarshal(t *testing.T) { + // Verify that types with a Content field round-trip properly. + roundtrip := func(in, out any) { + t.Helper() + data, err := json.Marshal(in) + if err != nil { + t.Fatal(err) + } + if err := json.Unmarshal(data, out); err != nil { + t.Fatal(err) + } + if diff := cmp.Diff(in, out); diff != "" { + t.Errorf("mismatch (-want, +got):\n%s", diff) + } + } + + content := []Content{&TextContent{Text: "t"}} + + ctr := &CallToolResult{ + Meta: Meta{"m": true}, + Content: content, + IsError: true, + StructuredContent: map[string]any{"s": "x"}, + } + var got CallToolResult + roundtrip(ctr, &got) + + ctrf := &CallToolResultFor[int]{ + Meta: Meta{"m": true}, + Content: content, + IsError: true, + StructuredContent: 3, + } + var gotf CallToolResultFor[int] + roundtrip(ctrf, &gotf) + + pm := &PromptMessage{ + Content: content[0], + Role: "", + } + var gotpm PromptMessage + roundtrip(pm, &gotpm) +} diff --git a/mcp/server.go b/mcp/server.go index c4aeeed9..44ec7aa7 100644 --- a/mcp/server.go +++ b/mcp/server.go @@ -16,6 +16,7 @@ import ( "path/filepath" "slices" "sync" + "time" "github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2" "github.com/modelcontextprotocol/go-sdk/internal/util" @@ -57,6 +58,10 @@ type ServerOptions struct { RootsListChangedHandler func(context.Context, *ServerSession, *RootsListChangedParams) // If non-nil, called when "notifications/progress" is received. ProgressNotificationHandler func(context.Context, *ServerSession, *ProgressNotificationParams) + // If non-zero, defines an interval for regular "ping" requests. + // If the peer fails to respond to pings originating from the keepalive check, + // the session is automatically closed. + KeepAlive time.Duration } // NewServer creates a new MCP server. The resulting server has no features: @@ -413,7 +418,7 @@ func fileResourceHandler(dir string) ResourceHandler { } // TODO(jba): figure out mime type. Omit for now: Server.readResource will fill it in. return &ReadResourceResult{Contents: []*ResourceContents{ - &ResourceContents{URI: params.URI, Blob: data}, + {URI: params.URI, Blob: data}, }}, nil } } @@ -460,6 +465,9 @@ func (s *Server) Connect(ctx context.Context, t Transport) (*ServerSession, erro } func (s *Server) callInitializedHandler(ctx context.Context, ss *ServerSession, params *InitializedParams) (Result, error) { + if s.opts.KeepAlive > 0 { + ss.startKeepalive(s.opts.KeepAlive) + } return callNotificationHandler(ctx, s.opts.InitializedHandler, ss, params) } @@ -492,6 +500,7 @@ type ServerSession struct { logLevel LoggingLevel initializeParams *InitializeParams initialized bool + keepaliveCancel context.CancelFunc } // Ping pings the client. @@ -510,12 +519,10 @@ func (ss *ServerSession) CreateMessage(ctx context.Context, params *CreateMessag return handleSend[*CreateMessageResult](ctx, ss, methodCreateMessage, orZero[Params](params)) } -// LoggingMessage sends a logging message to the client. +// Log sends a log message to the client. // The message is not sent if the client has not called SetLevel, or if its level // is below that of the last SetLevel. -// -// TODO(jba): rename to Log or LogMessage. A logging message is the thing that is sent to logging. -func (ss *ServerSession) LoggingMessage(ctx context.Context, params *LoggingMessageParams) error { +func (ss *ServerSession) Log(ctx context.Context, params *LoggingMessageParams) error { ss.mu.Lock() logLevel := ss.logLevel ss.mu.Unlock() @@ -680,6 +687,14 @@ func (ss *ServerSession) setLevel(_ context.Context, params *SetLevelParams) (*e // requests from being handled, and waiting for ongoing requests to return. // Close then terminates the connection. func (ss *ServerSession) Close() error { + if ss.keepaliveCancel != nil { + // Note: keepaliveCancel access is safe without a mutex because: + // 1. keepaliveCancel is only written once during startKeepalive (happens-before all Close calls) + // 2. context.CancelFunc is safe to call multiple times and from multiple goroutines + // 3. The keepalive goroutine calls Close on ping failure, but this is safe since + // Close is idempotent and conn.Close() handles concurrent calls correctly + ss.keepaliveCancel() + } return ss.conn.Close() } @@ -688,6 +703,11 @@ func (ss *ServerSession) Wait() error { return ss.conn.Wait() } +// startKeepalive starts the keepalive mechanism for this server session. +func (ss *ServerSession) startKeepalive(interval time.Duration) { + startKeepalive(ss, interval, &ss.keepaliveCancel) +} + // pageToken is the internal structure for the opaque pagination cursor. // It will be Gob-encoded and then Base64-encoded for use as a string token. type pageToken struct { diff --git a/mcp/shared.go b/mcp/shared.go index 07ea5dff..a2d51470 100644 --- a/mcp/shared.go +++ b/mcp/shared.go @@ -207,13 +207,7 @@ func sessionMethod[S Session, P Params, R Result](f func(S, context.Context, P) // Error codes const ( - // The error code to return when a resource isn't found. - // See https://modelcontextprotocol.io/specification/2025-03-26/server/resources#error-handling - // However, the code they chose is in the wrong space - // (see https://github.com/modelcontextprotocol/modelcontextprotocol/issues/509). - // so we pick a different one, arbitrarily for now (until they fix it). - // The immediate problem is that jsonprc2 defines -32002 as "server closing". - CodeResourceNotFound = -31002 + CodeResourceNotFound = -32002 // The error code if the method exists and was called properly, but the peer does not support it. CodeUnsupportedMethod = -31001 ) @@ -242,9 +236,13 @@ func notifySessions[S Session](sessions []S, method string, params Params) { } } +// Meta is additional metadata for requests, responses and other types. type Meta map[string]any -func (m Meta) GetMeta() map[string]any { return m } +// GetMeta returns metadata from a value. +func (m Meta) GetMeta() map[string]any { return m } + +// SetMeta sets the metadata on a value. func (m *Meta) SetMeta(x map[string]any) { *m = x } const progressTokenKey = "progressToken" @@ -269,7 +267,9 @@ func setProgressToken(p Params, pt any) { // Params is a parameter (input) type for an MCP call or notification. type Params interface { + // GetMeta returns metadata from a value. GetMeta() map[string]any + // SetMeta sets the metadata on a value. SetMeta(map[string]any) } @@ -288,7 +288,9 @@ type RequestParams interface { // Result is a result of an MCP call. type Result interface { + // GetMeta returns metadata from a value. GetMeta() map[string]any + // SetMeta sets the metadata on a value. SetMeta(map[string]any) } @@ -308,3 +310,41 @@ type listResult[T any] interface { // Returns a pointer to the param's NextCursor field. nextCursorPtr() *string } + +// keepaliveSession represents a session that supports keepalive functionality. +type keepaliveSession interface { + Ping(ctx context.Context, params *PingParams) error + Close() error +} + +// startKeepalive starts the keepalive mechanism for a session. +// It assigns the cancel function to the provided cancelPtr and starts a goroutine +// that sends ping messages at the specified interval. +func startKeepalive(session keepaliveSession, interval time.Duration, cancelPtr *context.CancelFunc) { + ctx, cancel := context.WithCancel(context.Background()) + // Assign cancel function before starting goroutine to avoid race condition. + // We cannot return it because the caller may need to cancel during the + // window between goroutine scheduling and function return. + *cancelPtr = cancel + + go func() { + ticker := time.NewTicker(interval) + defer ticker.Stop() + + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + pingCtx, pingCancel := context.WithTimeout(context.Background(), interval/2) + err := session.Ping(pingCtx, nil) + pingCancel() + if err != nil { + // Ping failed, close the session + _ = session.Close() + return + } + } + } + }() +} diff --git a/mcp/streamable.go b/mcp/streamable.go index 63211c25..a1952f7e 100644 --- a/mcp/streamable.go +++ b/mcp/streamable.go @@ -254,7 +254,7 @@ func (s *StreamableServerTransport) Connect(context.Context) (Connection, error) // // Currently, this is implemented in [ServerSession.handle]. This is not ideal, // because it means that a user of the MCP package couldn't implement the -// streamable transport, as they'd lack this priviledged access. +// streamable transport, as they'd lack this privileged access. // // If we ever wanted to expose this mechanism, we have a few options: // 1. Make ServerSession an interface, and provide an implementation of diff --git a/mcp/tool.go b/mcp/tool.go index ee4967b7..5bf9440b 100644 --- a/mcp/tool.go +++ b/mcp/tool.go @@ -44,6 +44,8 @@ type ServerTool struct { // The input schema for the tool is extracted from the request type for the // handler, and used to unmmarshal and validate requests to the handler. This // schema may be customized using the [Input] option. +// +// TODO(jba): check that structured content is set in response. func NewServerTool[In, Out any](name, description string, handler ToolHandlerFor[In, Out], opts ...ToolOption) *ServerTool { st, err := newServerToolErr[In, Out](name, description, handler, opts...) if err != nil {