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

mcp: new API for features (#88)

Browse source 
BREAKING CHANGE

Change the API for adding tools, prompts, resources and resource
templates to a server.

The ServerXXX types are gone along with the plural Add methods
(AddTools, AddPrompts etc.)
Instead there are single Add methods (AddTool, AddPrompt).
In addition, instead of NewServerTool, there is AddTool[In, Out].

Fixes #73.
This commit is contained in:
Jonathan Amsterdam 2025-07-07 14:54:52 -04:00 • committed by GitHub
commit aebd244981
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
20 changed files with 366 additions and 612 deletions
Download patch file Download diff file Expand all files Collapse all files

24
README.md
View file

@ -1,6 +1,21 @@
<!-- Autogenerated by weave; DO NOT EDIT --> <!-- Autogenerated by weave; DO NOT EDIT -->
# MCP Go SDK # MCP Go SDK
***BREAKING CHANGES***
The latest version contains breaking changes:
- Server.AddTools is replaced by Server.AddTool.
- NewServerTool is replaced by AddTool. AddTool takes a Tool rather than a name and description, so you can
set any field on the Tool that you want before associating it with a handler.
- Tool options have been removed. If you don't want AddTool to infer a JSON Schema for you, you can construct one
as a struct literal, or using any other code that suits you.
- AddPrompts, AddResources and AddResourceTemplates are similarly replaced by singular methods which pair the
feature with a handler. The ServerXXX types have been removed.
[![PkgGoDev](https://pkg.go.dev/badge/github.com/modelcontextprotocol/go-sdk)](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 This repository contains an unreleased implementation of the official Go
@ -99,7 +114,7 @@ import (
) )
type HiParams struct { type HiParams struct {
Name string `json:"name"` Name string `json:"name", mcp:"the name of the person to greet"`
} }
func SayHi(ctx context.Context, cc *mcp.ServerSession, params *mcp.CallToolParamsFor[HiParams]) (*mcp.CallToolResultFor[any], error) { func SayHi(ctx context.Context, cc *mcp.ServerSession, params *mcp.CallToolParamsFor[HiParams]) (*mcp.CallToolResultFor[any], error) {
@ -111,11 +126,8 @@ func SayHi(ctx context.Context, cc *mcp.ServerSession, params *mcp.CallToolParam
func main() { func main() {
// Create a server with a single tool. // Create a server with a single tool.
server := mcp.NewServer("greeter", "v1.0.0", nil) server := mcp.NewServer("greeter", "v1.0.0", nil)
server.AddTools(
mcp.NewServerTool("greet", "say hi", SayHi, mcp.Input( mcp.AddTool(server, &mcp.Tool{Name: "greet", Description: "say hi"}, SayHi)
mcp.Property("name", mcp.Description("the name of the person to greet")),
)),
)
// Run the server over stdin/stdout, until the client disconnects // Run the server over stdin/stdout, until the client disconnects
if err := server.Run(context.Background(), mcp.NewStdioTransport()); err != nil { if err := server.Run(context.Background(), mcp.NewStdioTransport()); err != nil {
log.Fatal(err) log.Fatal(err)

145
design/design.md
View file

@ -372,12 +372,11 @@ A server that can handle that client call would look like this:
```go ```go
// Create a server with a single tool. // Create a server with a single tool.
server := mcp.NewServer("greeter", "v1.0.0", nil) server := mcp.NewServer("greeter", "v1.0.0", nil)
server.AddTools(mcp.NewServerTool("greet", "say hi", SayHi)) mcp.AddTool(server, &mcp.Tool{Name: "greet", Description: "say hi"}, SayHi)
// Run the server over stdin/stdout, until the client disconnects. // Run the server over stdin/stdout, until the client disconnects.
transport := mcp.NewStdioTransport() if err := server.Run(context.Background(), mcp.NewStdioTransport()); err != nil {
session, err := server.Connect(ctx, transport) log.Fatal(err)
... }
return session.Wait()
``` ```
For convenience, we provide `Server.Run` to handle the common case of running a session until the client disconnects: For convenience, we provide `Server.Run` to handle the common case of running a session until the client disconnects:
@ -603,14 +602,14 @@ type ClientOptions struct {
### Tools ### Tools
A `Tool` is a logical MCP tool, generated from the MCP spec, and a `ServerTool` is a tool bound to a tool handler. A `Tool` is a logical MCP tool, generated from the MCP spec.
A tool handler accepts `CallToolParams` and returns a `CallToolResult`. However, since we want to bind tools to Go input types, it is convenient in associated APIs to make `CallToolParams` generic, with a type parameter `TArgs` for the tool argument type. This allows tool APIs to manage the marshalling and unmarshalling of tool inputs for their caller. The bound `ServerTool` type expects a `json.RawMessage` for its tool arguments, but the `NewServerTool` constructor described below provides a mechanism to bind a typed handler. A tool handler accepts `CallToolParams` and returns a `CallToolResult`. However, since we want to bind tools to Go input types, it is convenient in associated APIs to have a generic version of `CallToolParams`, with a type parameter `In` for the tool argument type, as well as a generic version of for `CallToolResult`. This allows tool APIs to manage the marshalling and unmarshalling of tool inputs for their caller.
```go ```go
type CallToolParams[TArgs any] struct { type CallToolParamsFor[In any] struct {
Meta Meta `json:"_meta,omitempty"` Meta Meta `json:"_meta,omitempty"`
Arguments TArgs `json:"arguments,omitempty"` Arguments In `json:"arguments,omitempty"`
Name string `json:"name"` Name string `json:"name"`
} }
@ -621,23 +620,31 @@ type Tool struct {
Name string `json:"name"` Name string `json:"name"`
} }
type ToolHandler[TArgs] func(context.Context, *ServerSession, *CallToolParams[TArgs]) (*CallToolResult, error) type ToolHandlerFor[In, Out any] func(context.Context, *ServerSession, *CallToolParamsFor[In]) (*CallToolResultFor[Out], error)
type ToolHandler = ToolHandlerFor[map[string]any, any]
type ServerTool struct {
Tool Tool
Handler ToolHandler[json.RawMessage]
}
``` ```
Add tools to a server with `AddTools`: Add tools to a server with the `AddTool` method or function. The function is generic and infers schemas from the handler
arguments:
```go ```go
server.AddTools( func (s *Server) AddTool(t *Tool, h ToolHandler)
mcp.NewServerTool("add", "add numbers", addHandler), func AddTool[In, Out any](s *Server, t *Tool, h ToolHandlerFor[In, Out])
mcp.NewServerTool("subtract, subtract numbers", subHandler))
``` ```
Remove them by name with `RemoveTools`: ```go
mcp.AddTool(server, &mcp.Tool{Name: "add", Description: "add numbers"}, addHandler)
mcp.AddTool(server, &mcp.Tool{Name: "subtract", Description: "subtract numbers"}, subHandler)
```
The `AddTool` method requires an input schema, and optionally an output one. It will not modify them.
The handler should accept a `CallToolParams` and return a `CallToolResult`.
```go
t := &Tool{Name: ..., Description: ..., InputSchema: &jsonschema.Schema{...}}
server.AddTool(t, myHandler)
```
Tools can be removed by name with `RemoveTools`:
```go ```go
server.RemoveTools("add", "subtract") server.RemoveTools("add", "subtract")
@ -650,53 +657,30 @@ A tool's input schema, expressed as a [JSON Schema](https://json-schema.org), pr
Both of these have their advantages and disadvantages. Reflection is nice, because it allows you to bind directly to a Go API, and means that the JSON schema of your API is compatible with your Go types by construction. It also means that concerns like parsing and validation can be handled automatically. However, it can become cumbersome to express the full breadth of JSON schema using Go types or struct tags, and sometimes you want to express things that aren’t naturally modeled by Go types, like unions. Explicit schemas are simple and readable, and give the caller full control over their tool definition, but involve significant boilerplate. Both of these have their advantages and disadvantages. Reflection is nice, because it allows you to bind directly to a Go API, and means that the JSON schema of your API is compatible with your Go types by construction. It also means that concerns like parsing and validation can be handled automatically. However, it can become cumbersome to express the full breadth of JSON schema using Go types or struct tags, and sometimes you want to express things that aren’t naturally modeled by Go types, like unions. Explicit schemas are simple and readable, and give the caller full control over their tool definition, but involve significant boilerplate.
We have found that a hybrid model works well, where the _initial_ schema is derived using reflection, but any customization on top of that schema is applied using variadic options. We achieve this using a `NewServerTool` helper, which generates the schema from the input type, and wraps the handler to provide parsing and validation. The schema (and potentially other features) can be customized using ToolOptions. We provide both ways. The `jsonschema.For[T]` function will infer a schema, and it is called by the `AddTool` generic function.
Users can also call it themselves, or build a schema directly as a struct literal. They can still use the `AddTool` function to
create a typed handler, since `AddTool` doesn't touch schemas that are already present.
If the tool's `InputSchema` is nil, it is inferred from the `In` type parameter. If the `OutputSchema` is nil, it is inferred from the `Out` type parameter (unless `Out` is `any`).
For example, given this handler:
```go ```go
// NewServerTool creates a Tool using reflection on the given handler. type AddParams struct {
func NewServerTool[TArgs any](name, description string, handler ToolHandler[TArgs], opts …ToolOption) *ServerTool X int `json:"x"`
Y int `json:"y"`
}
type ToolOption interface { /* ... */ } func addHandler(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[AddParams]) (*mcp.CallToolResultFor[int], error) {
``` return &mcp.CallToolResultFor[int]{StructuredContent: params.Arguments.X + params.Arguments.Y}, nil
`NewServerTool` determines the input schema for a Tool from the `TArgs` type. Each struct field that would be marshaled by `encoding/json.Marshal` becomes a property of the schema. The property is required unless the field's `json` tag specifies "omitempty" or "omitzero" (new in Go 1.24). For example, given this struct:
```go
struct {
Name string `json:"name"`
Count int `json:"count,omitempty"`
Choices []string
Password []byte `json:"-"`
} }
``` ```
"name" and "Choices" are required, while "count" is optional. You can add it to a server like this:
As of this writing, the only `ToolOption` is `Input`, which allows customizing the input schema of the tool using schema options. These schema options are recursive, in the sense that they may also be applied to properties.
```go ```go
func Input(...SchemaOption) ToolOption mcp.AddTool(server, &mcp.Tool{Name: "add", Description: "add numbers"}, addHandler)
type Property(name string, opts ...SchemaOption) SchemaOption
type Description(desc string) SchemaOption
// etc.
``` ```
The input schema will be inferred from `AddParams`, and the output schema from `int`.
For example:
```go
NewServerTool(name, description, handler,
Input(Property("count", Description("size of the inventory"))))
```
The most recent JSON Schema spec defines over 40 keywords. Providing them all as options would bloat the API despite the fact that most would be very rarely used. For less common keywords, use the `Schema` option to set the schema explicitly:
```go
NewServerTool(name, description, handler,
Input(Property("Choices", Schema(&jsonschema.Schema{UniqueItems: true}))))
```
Schemas are validated on the server before the tool handler is called.
Since all the fields of the Tool struct are exported, a Tool can also be created directly with assignment or a struct literal. Since all the fields of the Tool struct are exported, a Tool can also be created directly with assignment or a struct literal.
@ -718,15 +702,7 @@ For registering tools, we provide only `AddTools`; mcp-go's `SetTools`, `AddTool
### Prompts ### Prompts
Use `NewServerPrompt` to create a prompt. As with tools, prompt argument schemas can be inferred from a struct, or obtained from options. Use `AddPrompt` to add a prompt to the server, and `RemovePrompts`
```go
func NewServerPrompt[TReq any](name, description string,
handler func(context.Context, *ServerSession, TReq) (*GetPromptResult, error),
opts ...PromptOption) *ServerPrompt
```
Use `AddPrompts` to add prompts to the server, and `RemovePrompts`
to remove them by name. to remove them by name.
```go ```go
@ -734,11 +710,12 @@ type codeReviewArgs struct {
Code string `json:"code"` Code string `json:"code"`
} }
func codeReviewHandler(context.Context, *ServerSession, codeReviewArgs) {...} func codeReviewHandler(context.Context, *ServerSession, *mcp.GetPromptParams) (*mcp.GetPromptResult, error) {...}
server.AddPrompts( server.AddPrompt(
NewServerPrompt("code_review", "review code", codeReviewHandler, &mcp.Prompt{Name: "code_review", Description: "review code"},
Argument("code", Description("the code to review")))) codeReviewHandler,
)
server.RemovePrompts("code_review") server.RemovePrompts("code_review")
``` ```
@ -757,25 +734,11 @@ type ResourceHandler func(context.Context, *ServerSession, *ReadResourceParams)
The arguments include the `ServerSession` so the handler can observe the client's roots. The handler should return the resource contents in a `ReadResourceResult`, calling either `NewTextResourceContents` or `NewBlobResourceContents`. If the handler omits the URI or MIME type, the server will populate them from the resource. The arguments include the `ServerSession` so the handler can observe the client's roots. The handler should return the resource contents in a `ReadResourceResult`, calling either `NewTextResourceContents` or `NewBlobResourceContents`. If the handler omits the URI or MIME type, the server will populate them from the resource.
The `ServerResource` and `ServerResourceTemplate` types hold the association between the resource and its handler: To add a resource or resource template to a server, users call the `AddResource` and `AddResourceTemplate` methods. We also provide methods to remove them.
```go ```go
type ServerResource struct { func (*Server) AddResource(*Resource, ResourceHandler)
Resource Resource func (*Server) AddResourceTemplate(*ResourceTemplate, ResourceHandler)
Handler ResourceHandler
}
type ServerResourceTemplate struct {
Template ResourceTemplate
Handler ResourceHandler
}
```
To add a resource or resource template to a server, users call the `AddResources` and `AddResourceTemplates` methods with one or more `ServerResource`s or `ServerResourceTemplate`s. We also provide methods to remove them.
```go
func (*Server) AddResources(...*ServerResource)
func (*Server) AddResourceTemplates(...*ServerResourceTemplate)
func (s *Server) RemoveResources(uris ...string) func (s *Server) RemoveResources(uris ...string)
func (s *Server) RemoveResourceTemplates(uriTemplates ...string) func (s *Server) RemoveResourceTemplates(uriTemplates ...string)
@ -796,9 +759,7 @@ Here is an example:
```go ```go
// Safely read "/public/puppies.txt". // Safely read "/public/puppies.txt".
s.AddResources(&mcp.ServerResource{ s.AddResource(&mcp.Resource{URI: "file:///puppies.txt"}, s.FileReadResourceHandler("/public"))
Resource: mcp.Resource{URI: "file:///puppies.txt"},
Handler: s.FileReadResourceHandler("/public")})
``` ```
Server sessions also support the spec methods `ListResources` and `ListResourceTemplates`, and the corresponding iterator methods `Resources` and `ResourceTemplates`. Server sessions also support the spec methods `ListResources` and `ListResourceTemplates`, and the corresponding iterator methods `Resources` and `ResourceTemplates`.

24
examples/hello/main.go
View file

@ -19,7 +19,7 @@ import (
var httpAddr = flag.String("http", "", "if set, use streamable HTTP at this address, instead of stdin/stdout") var httpAddr = flag.String("http", "", "if set, use streamable HTTP at this address, instead of stdin/stdout")
type HiArgs struct { type HiArgs struct {
Name string `json:"name"` Name string `json:"name" mcp:"the name to say hi to"`
} }
func SayHi(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[HiArgs]) (*mcp.CallToolResultFor[struct{}], error) { func SayHi(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[HiArgs]) (*mcp.CallToolResultFor[struct{}], error) {
@ -43,21 +43,13 @@ func main() {
flag.Parse() flag.Parse()
server := mcp.NewServer("greeter", "v0.0.1", nil) server := mcp.NewServer("greeter", "v0.0.1", nil)
server.AddTools(mcp.NewServerTool("greet", "say hi", SayHi, mcp.Input( mcp.AddTool(server, &mcp.Tool{Name: "greet", Description: "say hi"}, SayHi)
mcp.Property("name", mcp.Description("the name to say hi to")), server.AddPrompt(&mcp.Prompt{Name: "greet"}, PromptHi)
))) server.AddResource(&mcp.Resource{
server.AddPrompts(&mcp.ServerPrompt{ Name: "info",
Prompt: &mcp.Prompt{Name: "greet"}, MIMEType: "text/plain",
Handler: PromptHi, URI: "embedded:info",
}) }, handleEmbeddedResource)
server.AddResources(&mcp.ServerResource{
Resource: &mcp.Resource{
Name: "info",
MIMEType: "text/plain",
URI: "embedded:info",
},
Handler: handleEmbeddedResource,
})
if *httpAddr != "" { if *httpAddr != "" {
handler := mcp.NewStreamableHTTPHandler(func(*http.Request) *mcp.Server { handler := mcp.NewStreamableHTTPHandler(func(*http.Request) *mcp.Server {

6
examples/sse/main.go
View file

@ -35,12 +35,12 @@ func main() {
} }
server1 := mcp.NewServer("greeter1", "v0.0.1", nil) server1 := mcp.NewServer("greeter1", "v0.0.1", nil)
server1.AddTools(mcp.NewServerTool("greet1", "say hi", SayHi)) mcp.AddTool(server1, &mcp.Tool{Name: "greet1", Description: "say hi"}, SayHi)
server2 := mcp.NewServer("greeter2", "v0.0.1", nil) server2 := mcp.NewServer("greeter2", "v0.0.1", nil)
server2.AddTools(mcp.NewServerTool("greet2", "say hello", SayHi)) mcp.AddTool(server2, &mcp.Tool{Name: "greet2", Description: "say hello"}, SayHi)
log.Printf("MCP servers serving at %s\n", *httpAddr) log.Printf("MCP servers serving at %s", *httpAddr)
handler := mcp.NewSSEHandler(func(request *http.Request) *mcp.Server { handler := mcp.NewSSEHandler(func(request *http.Request) *mcp.Server {
url := request.URL.Path url := request.URL.Path
log.Printf("Handling request for URL %s\n", url) log.Printf("Handling request for URL %s\n", url)

15
internal/readme/README.src.md
View file

@ -1,5 +1,20 @@
# MCP Go SDK # MCP Go SDK
***BREAKING CHANGES***
The latest version contains breaking changes:
- Server.AddTools is replaced by Server.AddTool.
- NewServerTool is replaced by AddTool. AddTool takes a Tool rather than a name and description, so you can
set any field on the Tool that you want before associating it with a handler.
- Tool options have been removed. If you don't want AddTool to infer a JSON Schema for you, you can construct one
as a struct literal, or using any other code that suits you.
- AddPrompts, AddResources and AddResourceTemplates are similarly replaced by singular methods which pair the
feature with a handler. The ServerXXX types have been removed.
[![PkgGoDev](https://pkg.go.dev/badge/github.com/modelcontextprotocol/go-sdk)](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 This repository contains an unreleased implementation of the official Go

9
internal/readme/server/server.go
View file

@ -13,7 +13,7 @@ import (
) )
type HiParams struct { type HiParams struct {
Name string `json:"name"` Name string `json:"name", mcp:"the name of the person to greet"`
} }
func SayHi(ctx context.Context, cc *mcp.ServerSession, params *mcp.CallToolParamsFor[HiParams]) (*mcp.CallToolResultFor[any], error) { func SayHi(ctx context.Context, cc *mcp.ServerSession, params *mcp.CallToolParamsFor[HiParams]) (*mcp.CallToolResultFor[any], error) {
@ -25,11 +25,8 @@ func SayHi(ctx context.Context, cc *mcp.ServerSession, params *mcp.CallToolParam
func main() { func main() {
// Create a server with a single tool. // Create a server with a single tool.
server := mcp.NewServer("greeter", "v1.0.0", nil) server := mcp.NewServer("greeter", "v1.0.0", nil)
server.AddTools(
mcp.NewServerTool("greet", "say hi", SayHi, mcp.Input( mcp.AddTool(server, &mcp.Tool{Name: "greet", Description: "say hi"}, SayHi)
mcp.Property("name", mcp.Description("the name of the person to greet")),
)),
)
// Run the server over stdin/stdout, until the client disconnects // Run the server over stdin/stdout, until the client disconnects
if err := server.Run(context.Background(), mcp.NewStdioTransport()); err != nil { if err := server.Run(context.Background(), mcp.NewStdioTransport()); err != nil {
log.Fatal(err) log.Fatal(err)

58
mcp/client_list_test.go
View file

@ -22,12 +22,12 @@ func TestList(t *testing.T) {
defer serverSession.Close() defer serverSession.Close()
t.Run("tools", func(t *testing.T) { t.Run("tools", func(t *testing.T) {
toolA := mcp.NewServerTool("apple", "apple tool", SayHi) var wantTools []*mcp.Tool
toolB := mcp.NewServerTool("banana", "banana tool", SayHi) for _, name := range []string{"apple", "banana", "cherry"} {
toolC := mcp.NewServerTool("cherry", "cherry tool", SayHi) t := &mcp.Tool{Name: name, Description: name + " tool"}
tools := []*mcp.ServerTool{toolA, toolB, toolC} wantTools = append(wantTools, t)
wantTools := []*mcp.Tool{toolA.Tool, toolB.Tool, toolC.Tool} mcp.AddTool(server, t, SayHi)
server.AddTools(tools...) }
t.Run("list", func(t *testing.T) { t.Run("list", func(t *testing.T) {
res, err := clientSession.ListTools(ctx, nil) res, err := clientSession.ListTools(ctx, nil)
if err != nil { if err != nil {
@ -43,12 +43,13 @@ func TestList(t *testing.T) {
}) })
t.Run("resources", func(t *testing.T) { t.Run("resources", func(t *testing.T) {
resourceA := &mcp.ServerResource{Resource: &mcp.Resource{URI: "http://apple"}} var wantResources []*mcp.Resource
resourceB := &mcp.ServerResource{Resource: &mcp.Resource{URI: "http://banana"}} for _, name := range []string{"apple", "banana", "cherry"} {
resourceC := &mcp.ServerResource{Resource: &mcp.Resource{URI: "http://cherry"}} r := &mcp.Resource{URI: "http://" + name}
wantResources := []*mcp.Resource{resourceA.Resource, resourceB.Resource, resourceC.Resource} wantResources = append(wantResources, r)
resources := []*mcp.ServerResource{resourceA, resourceB, resourceC} server.AddResource(r, nil)
server.AddResources(resources...) }
t.Run("list", func(t *testing.T) { t.Run("list", func(t *testing.T) {
res, err := clientSession.ListResources(ctx, nil) res, err := clientSession.ListResources(ctx, nil)
if err != nil { if err != nil {
@ -64,15 +65,12 @@ func TestList(t *testing.T) {
}) })
t.Run("templates", func(t *testing.T) { t.Run("templates", func(t *testing.T) {
resourceTmplA := &mcp.ServerResourceTemplate{ResourceTemplate: &mcp.ResourceTemplate{URITemplate: "http://apple/{x}"}} var wantResourceTemplates []*mcp.ResourceTemplate
resourceTmplB := &mcp.ServerResourceTemplate{ResourceTemplate: &mcp.ResourceTemplate{URITemplate: "http://banana/{x}"}} for _, name := range []string{"apple", "banana", "cherry"} {
resourceTmplC := &mcp.ServerResourceTemplate{ResourceTemplate: &mcp.ResourceTemplate{URITemplate: "http://cherry/{x}"}} rt := &mcp.ResourceTemplate{URITemplate: "http://" + name + "/{x}"}
wantResourceTemplates := []*mcp.ResourceTemplate{ wantResourceTemplates = append(wantResourceTemplates, rt)
resourceTmplA.ResourceTemplate, resourceTmplB.ResourceTemplate, server.AddResourceTemplate(rt, nil)
resourceTmplC.ResourceTemplate,
} }
resourceTemplates := []*mcp.ServerResourceTemplate{resourceTmplA, resourceTmplB, resourceTmplC}
server.AddResourceTemplates(resourceTemplates...)
t.Run("list", func(t *testing.T) { t.Run("list", func(t *testing.T) {
res, err := clientSession.ListResourceTemplates(ctx, nil) res, err := clientSession.ListResourceTemplates(ctx, nil)
if err != nil { if err != nil {
@ -88,12 +86,12 @@ func TestList(t *testing.T) {
}) })
t.Run("prompts", func(t *testing.T) { t.Run("prompts", func(t *testing.T) {
promptA := newServerPrompt("apple", "apple prompt") var wantPrompts []*mcp.Prompt
promptB := newServerPrompt("banana", "banana prompt") for _, name := range []string{"apple", "banana", "cherry"} {
promptC := newServerPrompt("cherry", "cherry prompt") p := &mcp.Prompt{Name: name, Description: name + " prompt"}
wantPrompts := []*mcp.Prompt{promptA.Prompt, promptB.Prompt, promptC.Prompt} wantPrompts = append(wantPrompts, p)
prompts := []*mcp.ServerPrompt{promptA, promptB, promptC} server.AddPrompt(p, testPromptHandler)
server.AddPrompts(prompts...) }
t.Run("list", func(t *testing.T) { t.Run("list", func(t *testing.T) {
res, err := clientSession.ListPrompts(ctx, nil) res, err := clientSession.ListPrompts(ctx, nil)
if err != nil { if err != nil {
@ -123,14 +121,6 @@ func testIterator[T any](ctx context.Context, t *testing.T, seq iter.Seq2[*T, er
} }
} }
// testPromptHandler is used for type inference newServerPrompt.
func testPromptHandler(context.Context, *mcp.ServerSession, *mcp.GetPromptParams) (*mcp.GetPromptResult, error) { func testPromptHandler(context.Context, *mcp.ServerSession, *mcp.GetPromptParams) (*mcp.GetPromptResult, error) {
panic("not implemented") panic("not implemented")
} }
func newServerPrompt(name, desc string) *mcp.ServerPrompt {
return &mcp.ServerPrompt{
Prompt: &mcp.Prompt{Name: name, Description: desc},
Handler: testPromptHandler,
}
}

3
mcp/cmd_test.go
View file

@ -31,8 +31,7 @@ func runServer() {
ctx := context.Background() ctx := context.Background()
server := mcp.NewServer("greeter", "v0.0.1", nil) server := mcp.NewServer("greeter", "v0.0.1", nil)
server.AddTools(mcp.NewServerTool("greet", "say hi", SayHi)) mcp.AddTool(server, &mcp.Tool{Name: "greet", Description: "say hi"}, SayHi)
if err := server.Run(ctx, mcp.NewStdioTransport()); err != nil { if err := server.Run(ctx, mcp.NewStdioTransport()); err != nil {
log.Fatal(err) log.Fatal(err)
} }

12
mcp/features_test.go
View file

@ -27,9 +27,9 @@ func SayHi(ctx context.Context, cc *ServerSession, params *CallToolParamsFor[Say
} }
func TestFeatureSetOrder(t *testing.T) { func TestFeatureSetOrder(t *testing.T) {
toolA := NewServerTool("apple", "apple tool", SayHi).Tool toolA := &Tool{Name: "apple", Description: "apple tool"}
toolB := NewServerTool("banana", "banana tool", SayHi).Tool toolB := &Tool{Name: "banana", Description: "banana tool"}
toolC := NewServerTool("cherry", "cherry tool", SayHi).Tool toolC := &Tool{Name: "cherry", Description: "cherry tool"}
testCases := []struct { testCases := []struct {
tools []*Tool tools []*Tool
@ -52,9 +52,9 @@ func TestFeatureSetOrder(t *testing.T) {
} }
func TestFeatureSetAbove(t *testing.T) { func TestFeatureSetAbove(t *testing.T) {
toolA := NewServerTool("apple", "apple tool", SayHi).Tool toolA := &Tool{Name: "apple", Description: "apple tool"}
toolB := NewServerTool("banana", "banana tool", SayHi).Tool toolB := &Tool{Name: "banana", Description: "banana tool"}
toolC := NewServerTool("cherry", "cherry tool", SayHi).Tool toolC := &Tool{Name: "cherry", Description: "cherry tool"}
testCases := []struct { testCases := []struct {
tools []*Tool tools []*Tool

135
mcp/mcp_test.go
View file

@ -21,7 +21,6 @@ import (
"time" "time"
"github.com/google/go-cmp/cmp" "github.com/google/go-cmp/cmp"
"github.com/google/go-cmp/cmp/cmpopts"
"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2" "github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
"github.com/modelcontextprotocol/go-sdk/jsonschema" "github.com/modelcontextprotocol/go-sdk/jsonschema"
) )
@ -30,6 +29,9 @@ type hiParams struct {
Name string Name string
} }
// TODO(jba): after schemas are stateless (WIP), this can be a variable.
func greetTool() *Tool { return &Tool{Name: "greet", Description: "say hi"} }
func sayHi(ctx context.Context, ss *ServerSession, params *CallToolParamsFor[hiParams]) (*CallToolResultFor[any], error) { func sayHi(ctx context.Context, ss *ServerSession, params *CallToolParamsFor[hiParams]) (*CallToolResultFor[any], error) {
if err := ss.Ping(ctx, nil); err != nil { if err := ss.Ping(ctx, nil); err != nil {
return nil, fmt.Errorf("ping failed: %v", err) return nil, fmt.Errorf("ping failed: %v", err)
@ -63,9 +65,31 @@ func TestEndToEnd(t *testing.T) {
}, },
} }
s := NewServer("testServer", "v1.0.0", sopts) s := NewServer("testServer", "v1.0.0", sopts)
add(tools, s.AddTools, "greet", "fail") AddTool(s, &Tool{
add(prompts, s.AddPrompts, "code_review", "fail") Name: "greet",
add(resources, s.AddResources, "info.txt", "fail.txt") Description: "say hi",
}, sayHi)
s.AddTool(&Tool{Name: "fail", InputSchema: &jsonschema.Schema{}},
func(context.Context, *ServerSession, *CallToolParamsFor[map[string]any]) (*CallToolResult, error) {
return nil, errTestFailure
})
s.AddPrompt(&Prompt{
Name: "code_review",
Description: "do a code review",
Arguments: []*PromptArgument{{Name: "Code", Required: true}},
}, func(_ context.Context, _ *ServerSession, params *GetPromptParams) (*GetPromptResult, error) {
return &GetPromptResult{
Description: "Code review prompt",
Messages: []*PromptMessage{
{Role: "user", Content: &TextContent{Text: "Please review the following code: " + params.Arguments["Code"]}},
},
}, nil
})
s.AddPrompt(&Prompt{Name: "fail"}, func(_ context.Context, _ *ServerSession, _ *GetPromptParams) (*GetPromptResult, error) {
return nil, errTestFailure
})
s.AddResource(resource1, readHandler)
s.AddResource(resource2, readHandler)
// Connect the server. // Connect the server.
ss, err := s.Connect(ctx, st) ss, err := s.Connect(ctx, st)
@ -154,39 +178,14 @@ func TestEndToEnd(t *testing.T) {
t.Errorf("fail returned unexpected error: got %v, want containing %v", err, errTestFailure) t.Errorf("fail returned unexpected error: got %v, want containing %v", err, errTestFailure)
} }
s.AddPrompts(&ServerPrompt{Prompt: &Prompt{Name: "T"}}) s.AddPrompt(&Prompt{Name: "T"}, nil)
waitForNotification(t, "prompts") waitForNotification(t, "prompts")
s.RemovePrompts("T") s.RemovePrompts("T")
waitForNotification(t, "prompts") waitForNotification(t, "prompts")
}) })
t.Run("tools", func(t *testing.T) { t.Run("tools", func(t *testing.T) {
res, err := cs.ListTools(ctx, nil) // ListTools is tested in client_list_test.go.
if err != nil {
t.Errorf("tools/list failed: %v", err)
}
wantTools := []*Tool{
{
Name: "fail",
InputSchema: nil,
},
{
Name: "greet",
Description: "say hi",
InputSchema: &jsonschema.Schema{
Type: "object",
Required: []string{"Name"},
Properties: map[string]*jsonschema.Schema{
"Name": {Type: "string"},
},
AdditionalProperties: falseSchema(),
},
},
}
if diff := cmp.Diff(wantTools, res.Tools, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
t.Fatalf("tools/list mismatch (-want +got):\n%s", diff)
}
gotHi, err := cs.CallTool(ctx, &CallToolParams{ gotHi, err := cs.CallTool(ctx, &CallToolParams{
Name: "greet", Name: "greet",
Arguments: map[string]any{"name": "user"}, Arguments: map[string]any{"name": "user"},
@ -222,7 +221,7 @@ func TestEndToEnd(t *testing.T) {
t.Errorf("tools/call 'fail' mismatch (-want +got):\n%s", diff) t.Errorf("tools/call 'fail' mismatch (-want +got):\n%s", diff)
} }
s.AddTools(&ServerTool{Tool: &Tool{Name: "T"}, Handler: nopHandler}) s.AddTool(&Tool{Name: "T", InputSchema: &jsonschema.Schema{}}, nopHandler)
waitForNotification(t, "tools") waitForNotification(t, "tools")
s.RemoveTools("T") s.RemoveTools("T")
waitForNotification(t, "tools") waitForNotification(t, "tools")
@ -246,8 +245,7 @@ func TestEndToEnd(t *testing.T) {
MIMEType: "text/template", MIMEType: "text/template",
URITemplate: "file:///{+filename}", // the '+' means that filename can contain '/' URITemplate: "file:///{+filename}", // the '+' means that filename can contain '/'
} }
st := &ServerResourceTemplate{ResourceTemplate: template, Handler: readHandler} s.AddResourceTemplate(template, readHandler)
s.AddResourceTemplates(st)
tres, err := cs.ListResourceTemplates(ctx, nil) tres, err := cs.ListResourceTemplates(ctx, nil)
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
@ -292,7 +290,7 @@ func TestEndToEnd(t *testing.T) {
} }
} }
s.AddResources(&ServerResource{Resource: &Resource{URI: "http://U"}}) s.AddResource(&Resource{URI: "http://U"}, nil)
waitForNotification(t, "resources") waitForNotification(t, "resources")
s.RemoveResources("http://U") s.RemoveResources("http://U")
waitForNotification(t, "resources") waitForNotification(t, "resources")
@ -434,40 +432,6 @@ func TestEndToEnd(t *testing.T) {
var ( var (
errTestFailure = errors.New("mcp failure") errTestFailure = errors.New("mcp failure")
tools = map[string]*ServerTool{
"greet": NewServerTool("greet", "say hi", sayHi),
"fail": {
Tool: &Tool{Name: "fail"},
Handler: func(context.Context, *ServerSession, *CallToolParamsFor[map[string]any]) (*CallToolResult, error) {
return nil, errTestFailure
},
},
}
prompts = map[string]*ServerPrompt{
"code_review": {
Prompt: &Prompt{
Name: "code_review",
Description: "do a code review",
Arguments: []*PromptArgument{{Name: "Code", Required: true}},
},
Handler: func(_ context.Context, _ *ServerSession, params *GetPromptParams) (*GetPromptResult, error) {
return &GetPromptResult{
Description: "Code review prompt",
Messages: []*PromptMessage{
{Role: "user", Content: &TextContent{Text: "Please review the following code: " + params.Arguments["Code"]}},
},
}, nil
},
},
"fail": {
Prompt: &Prompt{Name: "fail"},
Handler: func(_ context.Context, _ *ServerSession, _ *GetPromptParams) (*GetPromptResult, error) {
return nil, errTestFailure
},
},
}
resource1 = &Resource{ resource1 = &Resource{
Name: "public", Name: "public",
MIMEType: "text/plain", MIMEType: "text/plain",
@ -484,11 +448,6 @@ var (
URI: "embedded:info", URI: "embedded:info",
} }
readHandler = fileResourceHandler("testdata/files") readHandler = fileResourceHandler("testdata/files")
resources = map[string]*ServerResource{
"info.txt": {resource1, readHandler},
"fail.txt": {resource2, readHandler},
"info": {resource3, handleEmbeddedResource},
}
) )
var embeddedResources = map[string]string{ var embeddedResources = map[string]string{
@ -540,21 +499,21 @@ func errorCode(err error) int64 {
return -1 return -1
} }
// basicConnection returns a new basic client-server connection configured with // basicConnection returns a new basic client-server connection, with the server
// the provided tools. // configured via the provided function.
// //
// The caller should cancel either the client connection or server connection // The caller should cancel either the client connection or server connection
// when the connections are no longer needed. // when the connections are no longer needed.
func basicConnection(t *testing.T, tools ...*ServerTool) (*ServerSession, *ClientSession) { func basicConnection(t *testing.T, config func(*Server)) (*ServerSession, *ClientSession) {
t.Helper() t.Helper()
ctx := context.Background() ctx := context.Background()
ct, st := NewInMemoryTransports() ct, st := NewInMemoryTransports()
s := NewServer("testServer", "v1.0.0", nil) s := NewServer("testServer", "v1.0.0", nil)
if config != nil {
// The 'greet' tool says hi. config(s)
s.AddTools(tools...) }
ss, err := s.Connect(ctx, st) ss, err := s.Connect(ctx, st)
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
@ -569,7 +528,9 @@ func basicConnection(t *testing.T, tools ...*ServerTool) (*ServerSession, *Clien
} }
func TestServerClosing(t *testing.T) { func TestServerClosing(t *testing.T) {
cc, cs := basicConnection(t, NewServerTool("greet", "say hi", sayHi)) cc, cs := basicConnection(t, func(s *Server) {
AddTool(s, greetTool(), sayHi)
})
defer cs.Close() defer cs.Close()
ctx := context.Background() ctx := context.Background()
@ -651,11 +612,9 @@ func TestCancellation(t *testing.T) {
} }
return nil, nil return nil, nil
} }
st := &ServerTool{ _, cs := basicConnection(t, func(s *Server) {
Tool: &Tool{Name: "slow"}, s.AddTool(&Tool{Name: "slow"}, slowRequest)
Handler: slowRequest, })
}
_, cs := basicConnection(t, st)
defer cs.Close() defer cs.Close()
ctx, cancel := context.WithCancel(context.Background()) ctx, cancel := context.WithCancel(context.Background())
@ -852,7 +811,7 @@ func TestKeepAlive(t *testing.T) {
KeepAlive: 100 * time.Millisecond, KeepAlive: 100 * time.Millisecond,
} }
s := NewServer("testServer", "v1.0.0", serverOpts) s := NewServer("testServer", "v1.0.0", serverOpts)
s.AddTools(NewServerTool("greet", "say hi", sayHi)) AddTool(s, greetTool(), sayHi)
ss, err := s.Connect(ctx, st) ss, err := s.Connect(ctx, st)
if err != nil { if err != nil {
@ -897,7 +856,7 @@ func TestKeepAliveFailure(t *testing.T) {
// Server without keepalive (to test one-sided keepalive) // Server without keepalive (to test one-sided keepalive)
s := NewServer("testServer", "v1.0.0", nil) s := NewServer("testServer", "v1.0.0", nil)
s.AddTools(NewServerTool("greet", "say hi", sayHi)) AddTool(s, greetTool(), sayHi)
ss, err := s.Connect(ctx, st) ss, err := s.Connect(ctx, st)
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)

7
mcp/prompt.go
View file

@ -11,8 +11,7 @@ import (
// A PromptHandler handles a call to prompts/get. // A PromptHandler handles a call to prompts/get.
type PromptHandler func(context.Context, *ServerSession, *GetPromptParams) (*GetPromptResult, error) type PromptHandler func(context.Context, *ServerSession, *GetPromptParams) (*GetPromptResult, error)
// A Prompt is a prompt definition bound to a prompt handler. type serverPrompt struct {
type ServerPrompt struct { prompt *Prompt
Prompt *Prompt handler PromptHandler
Handler PromptHandler
} }

20
mcp/resource.go
View file

@ -20,16 +20,16 @@ import (
"github.com/modelcontextprotocol/go-sdk/internal/util" "github.com/modelcontextprotocol/go-sdk/internal/util"
) )
// A ServerResource associates a Resource with its handler. // A serverResource associates a Resource with its handler.
type ServerResource struct { type serverResource struct {
Resource *Resource resource *Resource
Handler ResourceHandler handler ResourceHandler
} }
// A ServerResourceTemplate associates a ResourceTemplate with its handler. // A serverResourceTemplate associates a ResourceTemplate with its handler.
type ServerResourceTemplate struct { type serverResourceTemplate struct {
ResourceTemplate *ResourceTemplate resourceTemplate *ResourceTemplate
Handler ResourceHandler handler ResourceHandler
} }
// A ResourceHandler is a function that reads a resource. // A ResourceHandler is a function that reads a resource.
@ -156,8 +156,8 @@ func fileRoot(root *Root) (_ string, err error) {
// Matches reports whether the receiver's uri template matches the uri. // Matches reports whether the receiver's uri template matches the uri.
// TODO: use "github.com/yosida95/uritemplate/v3" // TODO: use "github.com/yosida95/uritemplate/v3"
func (sr *ServerResourceTemplate) Matches(uri string) bool { func (sr *serverResourceTemplate) Matches(uri string) bool {
re, err := uriTemplateToRegexp(sr.ResourceTemplate.URITemplate) re, err := uriTemplateToRegexp(sr.resourceTemplate.URITemplate)
if err != nil { if err != nil {
return false return false
} }

170
mcp/server.go
View file

@ -12,6 +12,7 @@ import (
"encoding/json" "encoding/json"
"fmt" "fmt"
"iter" "iter"
"log"
"net/url" "net/url"
"path/filepath" "path/filepath"
"slices" "slices"
@ -20,7 +21,6 @@ import (
"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2" "github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
"github.com/modelcontextprotocol/go-sdk/internal/util" "github.com/modelcontextprotocol/go-sdk/internal/util"
"github.com/modelcontextprotocol/go-sdk/jsonschema"
) )
const DefaultPageSize = 1000 const DefaultPageSize = 1000
@ -36,10 +36,10 @@ type Server struct {
opts ServerOptions opts ServerOptions
mu sync.Mutex mu sync.Mutex
prompts *featureSet[*ServerPrompt] prompts *featureSet[*serverPrompt]
tools *featureSet[*ServerTool] tools *featureSet[*serverTool]
resources *featureSet[*ServerResource] resources *featureSet[*serverResource]
resourceTemplates *featureSet[*ServerResourceTemplate] resourceTemplates *featureSet[*serverResourceTemplate]
sessions []*ServerSession sessions []*ServerSession
sendingMethodHandler_ MethodHandler[*ServerSession] sendingMethodHandler_ MethodHandler[*ServerSession]
receivingMethodHandler_ MethodHandler[*ServerSession] receivingMethodHandler_ MethodHandler[*ServerSession]
@ -87,28 +87,23 @@ func NewServer(name, version string, opts *ServerOptions) *Server {
name: name, name: name,
version: version, version: version,
opts: *opts, opts: *opts,
prompts: newFeatureSet(func(p *ServerPrompt) string { return p.Prompt.Name }), prompts: newFeatureSet(func(p *serverPrompt) string { return p.prompt.Name }),
tools: newFeatureSet(func(t *ServerTool) string { return t.Tool.Name }), tools: newFeatureSet(func(t *serverTool) string { return t.tool.Name }),
resources: newFeatureSet(func(r *ServerResource) string { return r.Resource.URI }), resources: newFeatureSet(func(r *serverResource) string { return r.resource.URI }),
resourceTemplates: newFeatureSet(func(t *ServerResourceTemplate) string { return t.ResourceTemplate.URITemplate }), resourceTemplates: newFeatureSet(func(t *serverResourceTemplate) string { return t.resourceTemplate.URITemplate }),
sendingMethodHandler_: defaultSendingMethodHandler[*ServerSession], sendingMethodHandler_: defaultSendingMethodHandler[*ServerSession],
receivingMethodHandler_: defaultReceivingMethodHandler[*ServerSession], receivingMethodHandler_: defaultReceivingMethodHandler[*ServerSession],
} }
} }
// AddPrompts adds the given prompts to the server, // AddPrompt adds a [Prompt] to the server, or replaces one with the same name.
// replacing any with the same names. func (s *Server) AddPrompt(p *Prompt, h PromptHandler) {
func (s *Server) AddPrompts(prompts ...*ServerPrompt) { // Assume there was a change, since add replaces existing items.
// Only notify if something could change. // (It's possible an item was replaced with an identical one, but not worth checking.)
if len(prompts) == 0 {
return
}
// Assume there was a change, since add replaces existing roots.
// (It's possible a root was replaced with an identical one, but not worth checking.)
s.changeAndNotify( s.changeAndNotify(
notificationPromptListChanged, notificationPromptListChanged,
&PromptListChangedParams{}, &PromptListChangedParams{},
func() bool { s.prompts.add(prompts...); return true }) func() bool { s.prompts.add(&serverPrompt{p, h}); return true })
} }
// RemovePrompts removes the prompts with the given names. // RemovePrompts removes the prompts with the given names.
@ -118,55 +113,44 @@ func (s *Server) RemovePrompts(names ...string) {
func() bool { return s.prompts.remove(names...) }) func() bool { return s.prompts.remove(names...) })
} }
// AddTools adds the given tools to the server, // AddTool adds a [Tool] to the server, or replaces one with the same name.
// replacing any with the same names. // The tool's input schema must be non-nil.
// The arguments must not be modified after this call. // The Tool argument must not be modified after this call.
// func (s *Server) AddTool(t *Tool, h ToolHandler) {
// AddTools panics if errors are detected. // TODO(jba): This is a breaking behavior change. Add before v0.2.0?
func (s *Server) AddTools(tools ...*ServerTool) { if t.InputSchema == nil {
if err := s.addToolsErr(tools...); err != nil { log.Printf("mcp: tool %q has a nil input schema. This will panic in a future release.", t.Name)
// panic(fmt.Sprintf("adding tool %q: nil input schema", t.Name))
}
if err := addToolErr(s, t, h); err != nil {
panic(err) panic(err)
} }
} }
// addToolsErr is like [AddTools], but returns an error instead of panicking. // AddTool adds a [Tool] to the server, or replaces one with the same name.
func (s *Server) addToolsErr(tools ...*ServerTool) error { // If the tool's input schema is nil, it is set to the schema inferred from the In
// Only notify if something could change. // type parameter, using [jsonschema.For].
if len(tools) == 0 { // If the tool's output schema is nil and the Out type parameter is not the empty
return nil // interface, then the output schema is set to the schema inferred from Out.
// The Tool argument must not be modified after this call.
func AddTool[In, Out any](s *Server, t *Tool, h ToolHandlerFor[In, Out]) {
if err := addToolErr(s, t, h); err != nil {
panic(err)
} }
// Wrap the user's Handlers with rawHandlers that take a json.RawMessage. }
for _, st := range tools {
if st.rawHandler == nil {
// This ServerTool was not created with NewServerTool.
if st.Handler == nil {
return fmt.Errorf("AddTools: tool %q has no handler", st.Tool.Name)
}
st.rawHandler = newRawHandler(st)
// Resolve the schemas, with no base URI. We don't expect tool schemas to
// refer outside of themselves.
if st.Tool.InputSchema != nil {
r, err := st.Tool.InputSchema.Resolve(&jsonschema.ResolveOptions{ValidateDefaults: true})
if err != nil {
return err
}
st.inputResolved = r
}
// if st.Tool.OutputSchema != nil { func addToolErr[In, Out any](s *Server, t *Tool, h ToolHandlerFor[In, Out]) (err error) {
// st.outputResolved, err := st.Tool.OutputSchema.Resolve(&jsonschema.ResolveOptions{ValidateDefaults: true}) defer util.Wrapf(&err, "adding tool %q", t.Name)
// if err != nil { st, err := newServerTool(t, h)
// return err if err != nil {
// } return err
// }
}
} }
// Assume there was a change, since add replaces existing tools. // Assume there was a change, since add replaces existing tools.
// (It's possible a tool was replaced with an identical one, but not worth checking.) // (It's possible a tool was replaced with an identical one, but not worth checking.)
// TODO: surface notify error here? // TODO: Batch these changes by size and time? The typescript SDK doesn't.
// TODO: Surface notify error here? best not, in case we need to batch.
s.changeAndNotify(notificationToolListChanged, &ToolListChangedParams{}, s.changeAndNotify(notificationToolListChanged, &ToolListChangedParams{},
func() bool { s.tools.add(tools...); return true }) func() bool { s.tools.add(st); return true })
return nil return nil
} }
@ -177,26 +161,19 @@ func (s *Server) RemoveTools(names ...string) {
func() bool { return s.tools.remove(names...) }) func() bool { return s.tools.remove(names...) })
} }
// AddResources adds the given resources to the server. // AddResource adds a [Resource] to the server, or replaces one with the same URI.
// If a resource with the same URI already exists, it is replaced. // AddResource panics if the resource URI is invalid or not absolute (has an empty scheme).
// AddResources panics if a resource URI is invalid or not absolute (has an empty scheme). func (s *Server) AddResource(r *Resource, h ResourceHandler) {
func (s *Server) AddResources(resources ...*ServerResource) {
// Only notify if something could change.
if len(resources) == 0 {
return
}
s.changeAndNotify(notificationResourceListChanged, &ResourceListChangedParams{}, s.changeAndNotify(notificationResourceListChanged, &ResourceListChangedParams{},
func() bool { func() bool {
for _, r := range resources { u, err := url.Parse(r.URI)
u, err := url.Parse(r.Resource.URI) if err != nil {
if err != nil { panic(err) // url.Parse includes the URI in the error
panic(err) // url.Parse includes the URI in the error
}
if !u.IsAbs() {
panic(fmt.Errorf("URI %s needs a scheme", r.Resource.URI))
}
s.resources.add(r)
} }
if !u.IsAbs() {
panic(fmt.Errorf("URI %s needs a scheme", r.URI))
}
s.resources.add(&serverResource{r, h})
return true return true
}) })
} }
@ -208,20 +185,13 @@ func (s *Server) RemoveResources(uris ...string) {
func() bool { return s.resources.remove(uris...) }) func() bool { return s.resources.remove(uris...) })
} }
// AddResourceTemplates adds the given resource templates to the server. // AddResourceTemplate adds a [ResourceTemplate] to the server, or replaces on with the same URI.
// If a resource template with the same URI template already exists, it will be replaced. // AddResourceTemplate panics if a URI template is invalid or not absolute (has an empty scheme).
// AddResourceTemplates panics if a URI template is invalid or not absolute (has an empty scheme). func (s *Server) AddResourceTemplate(t *ResourceTemplate, h ResourceHandler) {
func (s *Server) AddResourceTemplates(templates ...*ServerResourceTemplate) {
// Only notify if something could change.
if len(templates) == 0 {
return
}
s.changeAndNotify(notificationResourceListChanged, &ResourceListChangedParams{}, s.changeAndNotify(notificationResourceListChanged, &ResourceListChangedParams{},
func() bool { func() bool {
for _, t := range templates { // TODO: check template validity.
// TODO: check template validity. s.resourceTemplates.add(&serverResourceTemplate{t, h})
s.resourceTemplates.add(t)
}
return true return true
}) })
} }
@ -268,10 +238,10 @@ func (s *Server) listPrompts(_ context.Context, _ *ServerSession, params *ListPr
if params == nil { if params == nil {
params = &ListPromptsParams{} params = &ListPromptsParams{}
} }
return paginateList(s.prompts, s.opts.PageSize, params, &ListPromptsResult{}, func(res *ListPromptsResult, prompts []*ServerPrompt) { return paginateList(s.prompts, s.opts.PageSize, params, &ListPromptsResult{}, func(res *ListPromptsResult, prompts []*serverPrompt) {
res.Prompts = []*Prompt{} // avoid JSON null res.Prompts = []*Prompt{} // avoid JSON null
for _, p := range prompts { for _, p := range prompts {
res.Prompts = append(res.Prompts, p.Prompt) res.Prompts = append(res.Prompts, p.prompt)
} }
}) })
} }
@ -284,7 +254,7 @@ func (s *Server) getPrompt(ctx context.Context, cc *ServerSession, params *GetPr
// TODO: surface the error code over the wire, instead of flattening it into the string. // TODO: surface the error code over the wire, instead of flattening it into the string.
return nil, fmt.Errorf("%s: unknown prompt %q", jsonrpc2.ErrInvalidParams, params.Name) return nil, fmt.Errorf("%s: unknown prompt %q", jsonrpc2.ErrInvalidParams, params.Name)
} }
return prompt.Handler(ctx, cc, params) return prompt.handler(ctx, cc, params)
} }
func (s *Server) listTools(_ context.Context, _ *ServerSession, params *ListToolsParams) (*ListToolsResult, error) { func (s *Server) listTools(_ context.Context, _ *ServerSession, params *ListToolsParams) (*ListToolsResult, error) {
@ -293,22 +263,22 @@ func (s *Server) listTools(_ context.Context, _ *ServerSession, params *ListTool
if params == nil { if params == nil {
params = &ListToolsParams{} params = &ListToolsParams{}
} }
return paginateList(s.tools, s.opts.PageSize, params, &ListToolsResult{}, func(res *ListToolsResult, tools []*ServerTool) { return paginateList(s.tools, s.opts.PageSize, params, &ListToolsResult{}, func(res *ListToolsResult, tools []*serverTool) {
res.Tools = []*Tool{} // avoid JSON null res.Tools = []*Tool{} // avoid JSON null
for _, t := range tools { for _, t := range tools {
res.Tools = append(res.Tools, t.Tool) res.Tools = append(res.Tools, t.tool)
} }
}) })
} }
func (s *Server) callTool(ctx context.Context, cc *ServerSession, params *CallToolParamsFor[json.RawMessage]) (*CallToolResult, error) { func (s *Server) callTool(ctx context.Context, cc *ServerSession, params *CallToolParamsFor[json.RawMessage]) (*CallToolResult, error) {
s.mu.Lock() s.mu.Lock()
tool, ok := s.tools.get(params.Name) st, ok := s.tools.get(params.Name)
s.mu.Unlock() s.mu.Unlock()
if !ok { if !ok {
return nil, fmt.Errorf("%s: unknown tool %q", jsonrpc2.ErrInvalidParams, params.Name) return nil, fmt.Errorf("%s: unknown tool %q", jsonrpc2.ErrInvalidParams, params.Name)
} }
return tool.rawHandler(ctx, cc, params) return st.handler(ctx, cc, params)
} }
func (s *Server) listResources(_ context.Context, _ *ServerSession, params *ListResourcesParams) (*ListResourcesResult, error) { func (s *Server) listResources(_ context.Context, _ *ServerSession, params *ListResourcesParams) (*ListResourcesResult, error) {
@ -317,10 +287,10 @@ func (s *Server) listResources(_ context.Context, _ *ServerSession, params *List
if params == nil { if params == nil {
params = &ListResourcesParams{} params = &ListResourcesParams{}
} }
return paginateList(s.resources, s.opts.PageSize, params, &ListResourcesResult{}, func(res *ListResourcesResult, resources []*ServerResource) { return paginateList(s.resources, s.opts.PageSize, params, &ListResourcesResult{}, func(res *ListResourcesResult, resources []*serverResource) {
res.Resources = []*Resource{} // avoid JSON null res.Resources = []*Resource{} // avoid JSON null
for _, r := range resources { for _, r := range resources {
res.Resources = append(res.Resources, r.Resource) res.Resources = append(res.Resources, r.resource)
} }
}) })
} }
@ -332,10 +302,10 @@ func (s *Server) listResourceTemplates(_ context.Context, _ *ServerSession, para
params = &ListResourceTemplatesParams{} params = &ListResourceTemplatesParams{}
} }
return paginateList(s.resourceTemplates, s.opts.PageSize, params, &ListResourceTemplatesResult{}, return paginateList(s.resourceTemplates, s.opts.PageSize, params, &ListResourceTemplatesResult{},
func(res *ListResourceTemplatesResult, rts []*ServerResourceTemplate) { func(res *ListResourceTemplatesResult, rts []*serverResourceTemplate) {
res.ResourceTemplates = []*ResourceTemplate{} // avoid JSON null res.ResourceTemplates = []*ResourceTemplate{} // avoid JSON null
for _, rt := range rts { for _, rt := range rts {
res.ResourceTemplates = append(res.ResourceTemplates, rt.ResourceTemplate) res.ResourceTemplates = append(res.ResourceTemplates, rt.resourceTemplate)
} }
}) })
} }
@ -376,12 +346,12 @@ func (s *Server) lookupResourceHandler(uri string) (ResourceHandler, string, boo
defer s.mu.Unlock() defer s.mu.Unlock()
// Try resources first. // Try resources first.
if r, ok := s.resources.get(uri); ok { if r, ok := s.resources.get(uri); ok {
return r.Handler, r.Resource.MIMEType, true return r.handler, r.resource.MIMEType, true
} }
// Look for matching template. // Look for matching template.
for rt := range s.resourceTemplates.all() { for rt := range s.resourceTemplates.all() {
if rt.Matches(uri) { if rt.Matches(uri) {
return rt.Handler, rt.ResourceTemplate.MIMEType, true return rt.handler, rt.resourceTemplate.MIMEType, true
} }
} }
return nil, "", false return nil, "", false

2
mcp/server_example_test.go
View file

@ -29,7 +29,7 @@ func ExampleServer() {
clientTransport, serverTransport := mcp.NewInMemoryTransports() clientTransport, serverTransport := mcp.NewInMemoryTransports()
server := mcp.NewServer("greeter", "v0.0.1", nil) server := mcp.NewServer("greeter", "v0.0.1", nil)
server.AddTools(mcp.NewServerTool("greet", "say hi", SayHi)) mcp.AddTool(server, &mcp.Tool{Name: "greet", Description: "say hi"}, SayHi)
serverSession, err := server.Connect(ctx, serverTransport) serverSession, err := server.Connect(ctx, serverTransport)
if err != nil { if err != nil {

11
mcp/shared_test.go
View file

@ -12,7 +12,7 @@ import (
) )
// TODO(jba): this shouldn't be in this file, but tool_test.go doesn't have access to unexported symbols. // TODO(jba): this shouldn't be in this file, but tool_test.go doesn't have access to unexported symbols.
func TestNewServerToolValidate(t *testing.T) { func TestToolValidate(t *testing.T) {
// Check that the tool returned from NewServerTool properly validates its input schema. // Check that the tool returned from NewServerTool properly validates its input schema.
type req struct { type req struct {
@ -26,9 +26,10 @@ func TestNewServerToolValidate(t *testing.T) {
return nil, nil return nil, nil
} }
tool := NewServerTool("test", "test", dummyHandler) st, err := newServerTool(&Tool{Name: "test", Description: "test"}, dummyHandler)
// Need to add the tool to a server to get resolved schemas. if err != nil {
// s := NewServer("", "", nil) t.Fatal(err)
}
for _, tt := range []struct { for _, tt := range []struct {
desc string desc string
@ -71,7 +72,7 @@ func TestNewServerToolValidate(t *testing.T) {
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
_, err = tool.rawHandler(context.Background(), nil, _, err = st.handler(context.Background(), nil,
&CallToolParamsFor[json.RawMessage]{Arguments: json.RawMessage(raw)}) &CallToolParamsFor[json.RawMessage]{Arguments: json.RawMessage(raw)})
if err == nil && tt.want != "" { if err == nil && tt.want != "" {
t.Error("got success, wanted failure") t.Error("got success, wanted failure")

2
mcp/sse_example_test.go
View file

@ -28,7 +28,7 @@ func Add(ctx context.Context, cc *mcp.ServerSession, params *mcp.CallToolParamsF
func ExampleSSEHandler() { func ExampleSSEHandler() {
server := mcp.NewServer("adder", "v0.0.1", nil) server := mcp.NewServer("adder", "v0.0.1", nil)
server.AddTools(mcp.NewServerTool("add", "add two numbers", Add)) mcp.AddTool(server, &mcp.Tool{Name: "add", Description: "add two numbers"}, Add)
handler := mcp.NewSSEHandler(func(*http.Request) *mcp.Server { return server }) handler := mcp.NewSSEHandler(func(*http.Request) *mcp.Server { return server })
httpServer := httptest.NewServer(handler) httpServer := httptest.NewServer(handler)

2
mcp/sse_test.go
View file

@ -20,7 +20,7 @@ func TestSSEServer(t *testing.T) {
t.Run(fmt.Sprintf("closeServerFirst=%t", closeServerFirst), func(t *testing.T) { t.Run(fmt.Sprintf("closeServerFirst=%t", closeServerFirst), func(t *testing.T) {
ctx := context.Background() ctx := context.Background()
server := NewServer("testServer", "v1.0.0", nil) server := NewServer("testServer", "v1.0.0", nil)
server.AddTools(NewServerTool("greet", "say hi", sayHi)) AddTool(server, &Tool{Name: "greet"}, sayHi)
sseHandler := NewSSEHandler(func(*http.Request) *Server { return server }) sseHandler := NewSSEHandler(func(*http.Request) *Server { return server })

6
mcp/streamable_test.go
View file

@ -32,8 +32,7 @@ func TestStreamableTransports(t *testing.T) {
// 1. Create a server with a simple "greet" tool. // 1. Create a server with a simple "greet" tool.
server := NewServer("testServer", "v1.0.0", nil) server := NewServer("testServer", "v1.0.0", nil)
server.AddTools(NewServerTool("greet", "say hi", sayHi)) AddTool(server, &Tool{Name: "greet", Description: "say hi"}, sayHi)
// 2. Start an httptest.Server with the StreamableHTTPHandler, wrapped in a // 2. Start an httptest.Server with the StreamableHTTPHandler, wrapped in a
// cookie-checking middleware. // cookie-checking middleware.
handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, nil) handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, nil)
@ -323,13 +322,12 @@ func TestStreamableServerTransport(t *testing.T) {
// Create a server containing a single tool, which runs the test tool // Create a server containing a single tool, which runs the test tool
// behavior, if any. // behavior, if any.
server := NewServer("testServer", "v1.0.0", nil) server := NewServer("testServer", "v1.0.0", nil)
tool := NewServerTool("tool", "test tool", func(ctx context.Context, ss *ServerSession, params *CallToolParamsFor[any]) (*CallToolResultFor[any], error) { AddTool(server, &Tool{Name: "tool"}, func(ctx context.Context, ss *ServerSession, params *CallToolParamsFor[any]) (*CallToolResultFor[any], error) {
if test.tool != nil { if test.tool != nil {
test.tool(t, ctx, ss) test.tool(t, ctx, ss)
} }
return &CallToolResultFor[any]{}, nil return &CallToolResultFor[any]{}, nil
}) })
server.AddTools(tool)
// Start the streamable handler. // Start the streamable handler.
handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, nil) handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, nil)

233
mcp/tool.go
View file

@ -9,7 +9,7 @@ import (
"context" "context"
"encoding/json" "encoding/json"
"fmt" "fmt"
"slices" "reflect"
"github.com/modelcontextprotocol/go-sdk/jsonschema" "github.com/modelcontextprotocol/go-sdk/jsonschema"
) )
@ -17,8 +17,7 @@ import (
// A ToolHandler handles a call to tools/call. // A ToolHandler handles a call to tools/call.
// [CallToolParams.Arguments] will contain a map[string]any that has been validated // [CallToolParams.Arguments] will contain a map[string]any that has been validated
// against the input schema. // against the input schema.
// TODO: Perhaps this should be an alias for ToolHandlerFor[map[string]any, map[string]any]? type ToolHandler = ToolHandlerFor[map[string]any, any]
type ToolHandler func(context.Context, *ServerSession, *CallToolParamsFor[map[string]any]) (*CallToolResult, error)
// A ToolHandlerFor handles a call to tools/call with typed arguments and results. // A ToolHandlerFor handles a call to tools/call with typed arguments and results.
type ToolHandlerFor[In, Out any] func(context.Context, *ServerSession, *CallToolParamsFor[In]) (*CallToolResultFor[Out], error) type ToolHandlerFor[In, Out any] func(context.Context, *ServerSession, *CallToolParamsFor[In]) (*CallToolResultFor[Out], error)
@ -26,62 +25,33 @@ type ToolHandlerFor[In, Out any] func(context.Context, *ServerSession, *CallTool
// A rawToolHandler is like a ToolHandler, but takes the arguments as as json.RawMessage. // A rawToolHandler is like a ToolHandler, but takes the arguments as as json.RawMessage.
type rawToolHandler = func(context.Context, *ServerSession, *CallToolParamsFor[json.RawMessage]) (*CallToolResult, error) type rawToolHandler = func(context.Context, *ServerSession, *CallToolParamsFor[json.RawMessage]) (*CallToolResult, error)
// A ServerTool is a tool definition that is bound to a tool handler. // A serverTool is a tool definition that is bound to a tool handler.
type ServerTool struct { type serverTool struct {
Tool *Tool tool *Tool
Handler ToolHandler handler rawToolHandler
// Set in NewServerTool or Server.addToolsErr. // Resolved tool schemas. Set in newServerTool.
rawHandler rawToolHandler
// Resolved tool schemas. Set in Server.addToolsErr.
inputResolved, outputResolved *jsonschema.Resolved inputResolved, outputResolved *jsonschema.Resolved
} }
// NewServerTool is a helper to make a tool using reflection on the given type parameters. // newServerTool creates a serverTool from a tool and a handler.
// When the tool is called, CallToolParams.Arguments will be of type In. // If the tool doesn't have an input schema, it is inferred from In.
// // If the tool doesn't have an output schema and Out != any, it is inferred from Out.
// If provided, variadic [ToolOption] values may be used to customize the tool. func newServerTool[In, Out any](t *Tool, h ToolHandlerFor[In, Out]) (*serverTool, error) {
// st := &serverTool{tool: t}
// 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 {
panic(fmt.Errorf("NewServerTool(%q): %w", name, err))
}
return st
}
func newServerToolErr[In, Out any](name, description string, handler ToolHandlerFor[In, Out], opts ...ToolOption) (*ServerTool, error) { if err := setSchema[In](&t.InputSchema, &st.inputResolved); err != nil {
// TODO: check that In is a struct.
ischema, err := jsonschema.For[In]()
if err != nil {
return nil, err return nil, err
} }
// TODO: uncomment when output schemas drop. if reflect.TypeFor[Out]() != reflect.TypeFor[any]() {
// oschema, err := jsonschema.For[TRes]() if err := setSchema[Out](&t.OutputSchema, &st.outputResolved); err != nil {
// if err != nil { return nil, err
// return nil, err }
// }
t := &ServerTool{
Tool: &Tool{
Name: name,
Description: description,
InputSchema: ischema,
// OutputSchema: oschema,
},
}
for _, opt := range opts {
opt.set(t)
} }
t.rawHandler = func(ctx context.Context, ss *ServerSession, rparams *CallToolParamsFor[json.RawMessage]) (*CallToolResult, error) { st.handler = func(ctx context.Context, ss *ServerSession, rparams *CallToolParamsFor[json.RawMessage]) (*CallToolResult, error) {
var args In var args In
if rparams.Arguments != nil { if rparams.Arguments != nil {
if err := unmarshalSchema(rparams.Arguments, t.inputResolved, &args); err != nil { if err := unmarshalSchema(rparams.Arguments, st.inputResolved, &args); err != nil {
return nil, err return nil, err
} }
} }
@ -91,45 +61,7 @@ func newServerToolErr[In, Out any](name, description string, handler ToolHandler
Name: rparams.Name, Name: rparams.Name,
Arguments: args, Arguments: args,
} }
res, err := handler(ctx, ss, params) res, err := h(ctx, ss, params)
if err != nil {
return nil, err
}
var ctr CallToolResult
if res != nil {
// TODO(jba): future-proof this copy.
ctr.Meta = res.Meta
ctr.Content = res.Content
ctr.IsError = res.IsError
}
return &ctr, nil
}
return t, nil
}
// newRawHandler creates a rawToolHandler for tools not created through NewServerTool.
// It unmarshals the arguments into a map[string]any and validates them against the
// schema, then calls the ServerTool's handler.
func newRawHandler(st *ServerTool) rawToolHandler {
if st.Handler == nil {
panic("st.Handler is nil")
}
return func(ctx context.Context, ss *ServerSession, rparams *CallToolParamsFor[json.RawMessage]) (*CallToolResult, error) {
// Unmarshal the args into what should be a map.
var args map[string]any
if rparams.Arguments != nil {
if err := unmarshalSchema(rparams.Arguments, st.inputResolved, &args); err != nil {
return nil, err
}
}
// TODO: generate copy
params := &CallToolParamsFor[map[string]any]{
Meta: rparams.Meta,
Name: rparams.Name,
Arguments: args,
}
res, err := st.Handler(ctx, ss, params)
// TODO(rfindley): investigate why server errors are embedded in this strange way, // TODO(rfindley): investigate why server errors are embedded in this strange way,
// rather than returned as jsonrpc2 server errors. // rather than returned as jsonrpc2 server errors.
if err != nil { if err != nil {
@ -138,8 +70,32 @@ func newRawHandler(st *ServerTool) rawToolHandler {
IsError: true, IsError: true,
}, nil }, nil
} }
return res, nil var ctr CallToolResult
// TODO(jba): What if res == nil? Is that valid?
// TODO(jba): if t.OutputSchema != nil, check that StructuredContent is present and validates.
if res != nil {
// TODO(jba): future-proof this copy.
ctr.Meta = res.Meta
ctr.Content = res.Content
ctr.IsError = res.IsError
ctr.StructuredContent = res.StructuredContent
}
return &ctr, nil
} }
return st, nil
}
func setSchema[T any](sfield **jsonschema.Schema, rfield **jsonschema.Resolved) error {
var err error
if *sfield == nil {
*sfield, err = jsonschema.For[T]()
}
if err != nil {
return err
}
*rfield, err = (*sfield).Resolve(&jsonschema.ResolveOptions{ValidateDefaults: true})
return err
} }
// unmarshalSchema unmarshals data into v and validates the result according to // unmarshalSchema unmarshals data into v and validates the result according to
@ -169,105 +125,6 @@ func unmarshalSchema(data json.RawMessage, resolved *jsonschema.Resolved, v any)
return nil return nil
} }
// A ToolOption configures the behavior of a Tool.
type ToolOption interface {
set(*ServerTool)
}
type toolSetter func(*ServerTool)
func (s toolSetter) set(t *ServerTool) { s(t) }
// Input applies the provided [SchemaOption] configuration to the tool's input
// schema.
func Input(opts ...SchemaOption) ToolOption {
return toolSetter(func(t *ServerTool) {
for _, opt := range opts {
opt.set(t.Tool.InputSchema)
}
})
}
// A SchemaOption configures a jsonschema.Schema.
type SchemaOption interface {
set(s *jsonschema.Schema)
}
type schemaSetter func(*jsonschema.Schema)
func (s schemaSetter) set(schema *jsonschema.Schema) { s(schema) }
// Property configures the schema for the property of the given name.
// If there is no such property in the schema, it is created.
func Property(name string, opts ...SchemaOption) SchemaOption {
return schemaSetter(func(schema *jsonschema.Schema) {
propSchema, ok := schema.Properties[name]
if !ok {
propSchema = new(jsonschema.Schema)
schema.Properties[name] = propSchema
}
// Apply the options, with special handling for Required, as it needs to be
// set on the parent schema.
for _, opt := range opts {
if req, ok := opt.(required); ok {
if req {
if !slices.Contains(schema.Required, name) {
schema.Required = append(schema.Required, name)
}
} else {
schema.Required = slices.DeleteFunc(schema.Required, func(s string) bool {
return s == name
})
}
} else {
opt.set(propSchema)
}
}
})
}
// Required sets whether the associated property is required. It is only valid
// when used in a [Property] option: using Required outside of Property panics.
func Required(v bool) SchemaOption {
return required(v)
}
// required must be a distinguished type as it needs special handling to mutate
// the parent schema, and to mutate prompt arguments.
type required bool
func (required) set(s *jsonschema.Schema) {
panic("use of required outside of Property")
}
// Enum sets the provided values as the "enum" value of the schema.
func Enum(values ...any) SchemaOption {
return schemaSetter(func(s *jsonschema.Schema) {
s.Enum = values
})
}
// Description sets the provided schema description.
func Description(desc string) SchemaOption {
return description(desc)
}
// description must be a distinguished type so that it can be handled by prompt
// options.
type description string
func (d description) set(s *jsonschema.Schema) {
s.Description = string(d)
}
// Schema overrides the inferred schema with a shallow copy of the given
// schema.
func Schema(schema *jsonschema.Schema) SchemaOption {
return schemaSetter(func(s *jsonschema.Schema) {
*s = *schema
})
}
// schemaJSON returns the JSON value for s as a string, or a string indicating an error. // schemaJSON returns the JSON value for s as a string, or a string indicating an error.
func schemaJSON(s *jsonschema.Schema) string { func schemaJSON(s *jsonschema.Schema) string {
m, err := json.Marshal(s) m, err := json.Marshal(s)

112
mcp/tool_test.go
View file

@ -16,76 +16,80 @@ import (
) )
// testToolHandler is used for type inference in TestNewServerTool. // testToolHandler is used for type inference in TestNewServerTool.
func testToolHandler[T any](context.Context, *ServerSession, *CallToolParamsFor[T]) (*CallToolResultFor[any], error) { func testToolHandler[In, Out any](context.Context, *ServerSession, *CallToolParamsFor[In]) (*CallToolResultFor[Out], error) {
panic("not implemented") panic("not implemented")
} }
func srvTool[In, Out any](t *testing.T, tool *Tool, handler ToolHandlerFor[In, Out]) *serverTool {
t.Helper()
st, err := newServerTool(tool, handler)
if err != nil {
t.Fatal(err)
}
return st
}
func TestNewServerTool(t *testing.T) { func TestNewServerTool(t *testing.T) {
type (
Name struct {
Name string `json:"name"`
}
Size struct {
Size int `json:"size"`
}
)
nameSchema := &jsonschema.Schema{
Type: "object",
Required: []string{"name"},
Properties: map[string]*jsonschema.Schema{
"name": {Type: "string"},
},
AdditionalProperties: &jsonschema.Schema{Not: new(jsonschema.Schema)},
}
sizeSchema := &jsonschema.Schema{
Type: "object",
Required: []string{"size"},
Properties: map[string]*jsonschema.Schema{
"size": {Type: "integer"},
},
AdditionalProperties: &jsonschema.Schema{Not: new(jsonschema.Schema)},
}
tests := []struct { tests := []struct {
tool *ServerTool tool *serverTool
want *jsonschema.Schema wantIn, wantOut *jsonschema.Schema
}{ }{
{ {
NewServerTool("basic", "", testToolHandler[struct { srvTool(t, &Tool{Name: "basic"}, testToolHandler[Name, Size]),
Name string `json:"name"` nameSchema,
}]), sizeSchema,
&jsonschema.Schema{
Type: "object",
Required: []string{"name"},
Properties: map[string]*jsonschema.Schema{
"name": {Type: "string"},
},
AdditionalProperties: &jsonschema.Schema{Not: new(jsonschema.Schema)},
},
}, },
{ {
NewServerTool("enum", "", testToolHandler[struct{ Name string }], Input( srvTool(t, &Tool{
Property("Name", Enum("x", "y", "z")), Name: "in untouched",
)), InputSchema: &jsonschema.Schema{},
&jsonschema.Schema{ }, testToolHandler[Name, Size]),
Type: "object", &jsonschema.Schema{},
Required: []string{"Name"}, sizeSchema,
Properties: map[string]*jsonschema.Schema{
"Name": {Type: "string", Enum: []any{"x", "y", "z"}},
},
AdditionalProperties: &jsonschema.Schema{Not: new(jsonschema.Schema)},
},
}, },
{ {
NewServerTool("required", "", testToolHandler[struct { srvTool(t, &Tool{Name: "out untouched", OutputSchema: &jsonschema.Schema{}}, testToolHandler[Name, Size]),
Name string `json:"name"` nameSchema,
Language string `json:"language"` &jsonschema.Schema{},
X int `json:"x,omitempty"`
Y int `json:"y,omitempty"`
}], Input(
Property("x", Required(true)))),
&jsonschema.Schema{
Type: "object",
Required: []string{"name", "language", "x"},
Properties: map[string]*jsonschema.Schema{
"language": {Type: "string"},
"name": {Type: "string"},
"x": {Type: "integer"},
"y": {Type: "integer"},
},
AdditionalProperties: &jsonschema.Schema{Not: new(jsonschema.Schema)},
},
}, },
{ {
NewServerTool("set_schema", "", testToolHandler[struct { srvTool(t, &Tool{Name: "nil out"}, testToolHandler[Name, any]),
X int `json:"x,omitempty"` nameSchema,
Y int `json:"y,omitempty"` nil,
}], Input(
Schema(&jsonschema.Schema{Type: "object"})),
),
&jsonschema.Schema{
Type: "object",
},
}, },
} }
for _, test := range tests { for _, test := range tests {
if diff := cmp.Diff(test.want, test.tool.Tool.InputSchema, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" { if diff := cmp.Diff(test.wantIn, test.tool.tool.InputSchema, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
t.Errorf("NewServerTool(%v) mismatch (-want +got):\n%s", test.tool.Tool.Name, diff) t.Errorf("newServerTool(%q) input schema mismatch (-want +got):\n%s", test.tool.tool.Name, diff)
}
if diff := cmp.Diff(test.wantOut, test.tool.tool.OutputSchema, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
t.Errorf("newServerTool(%q) output schema mismatch (-want +got):\n%s", test.tool.tool.Name, diff)
} }
} }
} }