mcp: new API for features (#88)
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
GitHub
parent
328a25d503
commit
aebd244981
14 changed files with 366 additions and 612 deletions
|
|
@ -22,12 +22,12 @@ func TestList(t *testing.T) {
|
|||
defer serverSession.Close()
|
||||
|
||||
t.Run("tools", func(t *testing.T) {
|
||||
toolA := mcp.NewServerTool("apple", "apple tool", SayHi)
|
||||
toolB := mcp.NewServerTool("banana", "banana tool", SayHi)
|
||||
toolC := mcp.NewServerTool("cherry", "cherry tool", SayHi)
|
||||
tools := []*mcp.ServerTool{toolA, toolB, toolC}
|
||||
wantTools := []*mcp.Tool{toolA.Tool, toolB.Tool, toolC.Tool}
|
||||
server.AddTools(tools...)
|
||||
var wantTools []*mcp.Tool
|
||||
for _, name := range []string{"apple", "banana", "cherry"} {
|
||||
t := &mcp.Tool{Name: name, Description: name + " tool"}
|
||||
wantTools = append(wantTools, t)
|
||||
mcp.AddTool(server, t, SayHi)
|
||||
}
|
||||
t.Run("list", func(t *testing.T) {
|
||||
res, err := clientSession.ListTools(ctx, nil)
|
||||
if err != nil {
|
||||
|
|
@ -43,12 +43,13 @@ func TestList(t *testing.T) {
|
|||
})
|
||||
|
||||
t.Run("resources", func(t *testing.T) {
|
||||
resourceA := &mcp.ServerResource{Resource: &mcp.Resource{URI: "http://apple"}}
|
||||
resourceB := &mcp.ServerResource{Resource: &mcp.Resource{URI: "http://banana"}}
|
||||
resourceC := &mcp.ServerResource{Resource: &mcp.Resource{URI: "http://cherry"}}
|
||||
wantResources := []*mcp.Resource{resourceA.Resource, resourceB.Resource, resourceC.Resource}
|
||||
resources := []*mcp.ServerResource{resourceA, resourceB, resourceC}
|
||||
server.AddResources(resources...)
|
||||
var wantResources []*mcp.Resource
|
||||
for _, name := range []string{"apple", "banana", "cherry"} {
|
||||
r := &mcp.Resource{URI: "http://" + name}
|
||||
wantResources = append(wantResources, r)
|
||||
server.AddResource(r, nil)
|
||||
}
|
||||
|
||||
t.Run("list", func(t *testing.T) {
|
||||
res, err := clientSession.ListResources(ctx, nil)
|
||||
if err != nil {
|
||||
|
|
@ -64,15 +65,12 @@ func TestList(t *testing.T) {
|
|||
})
|
||||
|
||||
t.Run("templates", func(t *testing.T) {
|
||||
resourceTmplA := &mcp.ServerResourceTemplate{ResourceTemplate: &mcp.ResourceTemplate{URITemplate: "http://apple/{x}"}}
|
||||
resourceTmplB := &mcp.ServerResourceTemplate{ResourceTemplate: &mcp.ResourceTemplate{URITemplate: "http://banana/{x}"}}
|
||||
resourceTmplC := &mcp.ServerResourceTemplate{ResourceTemplate: &mcp.ResourceTemplate{URITemplate: "http://cherry/{x}"}}
|
||||
wantResourceTemplates := []*mcp.ResourceTemplate{
|
||||
resourceTmplA.ResourceTemplate, resourceTmplB.ResourceTemplate,
|
||||
resourceTmplC.ResourceTemplate,
|
||||
var wantResourceTemplates []*mcp.ResourceTemplate
|
||||
for _, name := range []string{"apple", "banana", "cherry"} {
|
||||
rt := &mcp.ResourceTemplate{URITemplate: "http://" + name + "/{x}"}
|
||||
wantResourceTemplates = append(wantResourceTemplates, rt)
|
||||
server.AddResourceTemplate(rt, nil)
|
||||
}
|
||||
resourceTemplates := []*mcp.ServerResourceTemplate{resourceTmplA, resourceTmplB, resourceTmplC}
|
||||
server.AddResourceTemplates(resourceTemplates...)
|
||||
t.Run("list", func(t *testing.T) {
|
||||
res, err := clientSession.ListResourceTemplates(ctx, nil)
|
||||
if err != nil {
|
||||
|
|
@ -88,12 +86,12 @@ func TestList(t *testing.T) {
|
|||
})
|
||||
|
||||
t.Run("prompts", func(t *testing.T) {
|
||||
promptA := newServerPrompt("apple", "apple prompt")
|
||||
promptB := newServerPrompt("banana", "banana prompt")
|
||||
promptC := newServerPrompt("cherry", "cherry prompt")
|
||||
wantPrompts := []*mcp.Prompt{promptA.Prompt, promptB.Prompt, promptC.Prompt}
|
||||
prompts := []*mcp.ServerPrompt{promptA, promptB, promptC}
|
||||
server.AddPrompts(prompts...)
|
||||
var wantPrompts []*mcp.Prompt
|
||||
for _, name := range []string{"apple", "banana", "cherry"} {
|
||||
p := &mcp.Prompt{Name: name, Description: name + " prompt"}
|
||||
wantPrompts = append(wantPrompts, p)
|
||||
server.AddPrompt(p, testPromptHandler)
|
||||
}
|
||||
t.Run("list", func(t *testing.T) {
|
||||
res, err := clientSession.ListPrompts(ctx, 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) {
|
||||
panic("not implemented")
|
||||
}
|
||||
|
||||
func newServerPrompt(name, desc string) *mcp.ServerPrompt {
|
||||
return &mcp.ServerPrompt{
|
||||
Prompt: &mcp.Prompt{Name: name, Description: desc},
|
||||
Handler: testPromptHandler,
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -31,8 +31,7 @@ func runServer() {
|
|||
ctx := context.Background()
|
||||
|
||||
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 {
|
||||
log.Fatal(err)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -27,9 +27,9 @@ func SayHi(ctx context.Context, cc *ServerSession, params *CallToolParamsFor[Say
|
|||
}
|
||||
|
||||
func TestFeatureSetOrder(t *testing.T) {
|
||||
toolA := NewServerTool("apple", "apple tool", SayHi).Tool
|
||||
toolB := NewServerTool("banana", "banana tool", SayHi).Tool
|
||||
toolC := NewServerTool("cherry", "cherry tool", SayHi).Tool
|
||||
toolA := &Tool{Name: "apple", Description: "apple tool"}
|
||||
toolB := &Tool{Name: "banana", Description: "banana tool"}
|
||||
toolC := &Tool{Name: "cherry", Description: "cherry tool"}
|
||||
|
||||
testCases := []struct {
|
||||
tools []*Tool
|
||||
|
|
@ -52,9 +52,9 @@ func TestFeatureSetOrder(t *testing.T) {
|
|||
}
|
||||
|
||||
func TestFeatureSetAbove(t *testing.T) {
|
||||
toolA := NewServerTool("apple", "apple tool", SayHi).Tool
|
||||
toolB := NewServerTool("banana", "banana tool", SayHi).Tool
|
||||
toolC := NewServerTool("cherry", "cherry tool", SayHi).Tool
|
||||
toolA := &Tool{Name: "apple", Description: "apple tool"}
|
||||
toolB := &Tool{Name: "banana", Description: "banana tool"}
|
||||
toolC := &Tool{Name: "cherry", Description: "cherry tool"}
|
||||
|
||||
testCases := []struct {
|
||||
tools []*Tool
|
||||
|
|
|
|||
135
mcp/mcp_test.go
135
mcp/mcp_test.go
|
|
@ -21,7 +21,6 @@ import (
|
|||
"time"
|
||||
|
||||
"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/jsonschema"
|
||||
)
|
||||
|
|
@ -30,6 +29,9 @@ type hiParams struct {
|
|||
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) {
|
||||
if err := ss.Ping(ctx, nil); err != nil {
|
||||
return nil, fmt.Errorf("ping failed: %v", err)
|
||||
|
|
@ -63,9 +65,31 @@ func TestEndToEnd(t *testing.T) {
|
|||
},
|
||||
}
|
||||
s := NewServer("testServer", "v1.0.0", sopts)
|
||||
add(tools, s.AddTools, "greet", "fail")
|
||||
add(prompts, s.AddPrompts, "code_review", "fail")
|
||||
add(resources, s.AddResources, "info.txt", "fail.txt")
|
||||
AddTool(s, &Tool{
|
||||
Name: "greet",
|
||||
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.
|
||||
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)
|
||||
}
|
||||
|
||||
s.AddPrompts(&ServerPrompt{Prompt: &Prompt{Name: "T"}})
|
||||
s.AddPrompt(&Prompt{Name: "T"}, nil)
|
||||
waitForNotification(t, "prompts")
|
||||
s.RemovePrompts("T")
|
||||
waitForNotification(t, "prompts")
|
||||
})
|
||||
|
||||
t.Run("tools", func(t *testing.T) {
|
||||
res, err := cs.ListTools(ctx, nil)
|
||||
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)
|
||||
}
|
||||
|
||||
// ListTools is tested in client_list_test.go.
|
||||
gotHi, err := cs.CallTool(ctx, &CallToolParams{
|
||||
Name: "greet",
|
||||
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)
|
||||
}
|
||||
|
||||
s.AddTools(&ServerTool{Tool: &Tool{Name: "T"}, Handler: nopHandler})
|
||||
s.AddTool(&Tool{Name: "T", InputSchema: &jsonschema.Schema{}}, nopHandler)
|
||||
waitForNotification(t, "tools")
|
||||
s.RemoveTools("T")
|
||||
waitForNotification(t, "tools")
|
||||
|
|
@ -246,8 +245,7 @@ func TestEndToEnd(t *testing.T) {
|
|||
MIMEType: "text/template",
|
||||
URITemplate: "file:///{+filename}", // the '+' means that filename can contain '/'
|
||||
}
|
||||
st := &ServerResourceTemplate{ResourceTemplate: template, Handler: readHandler}
|
||||
s.AddResourceTemplates(st)
|
||||
s.AddResourceTemplate(template, readHandler)
|
||||
tres, err := cs.ListResourceTemplates(ctx, nil)
|
||||
if err != nil {
|
||||
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")
|
||||
s.RemoveResources("http://U")
|
||||
waitForNotification(t, "resources")
|
||||
|
|
@ -434,40 +432,6 @@ func TestEndToEnd(t *testing.T) {
|
|||
var (
|
||||
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{
|
||||
Name: "public",
|
||||
MIMEType: "text/plain",
|
||||
|
|
@ -484,11 +448,6 @@ var (
|
|||
URI: "embedded:info",
|
||||
}
|
||||
readHandler = fileResourceHandler("testdata/files")
|
||||
resources = map[string]*ServerResource{
|
||||
"info.txt": {resource1, readHandler},
|
||||
"fail.txt": {resource2, readHandler},
|
||||
"info": {resource3, handleEmbeddedResource},
|
||||
}
|
||||
)
|
||||
|
||||
var embeddedResources = map[string]string{
|
||||
|
|
@ -540,21 +499,21 @@ func errorCode(err error) int64 {
|
|||
return -1
|
||||
}
|
||||
|
||||
// basicConnection returns a new basic client-server connection configured with
|
||||
// the provided tools.
|
||||
// basicConnection returns a new basic client-server connection, with the server
|
||||
// configured via the provided function.
|
||||
//
|
||||
// The caller should cancel either the client connection or server connection
|
||||
// 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()
|
||||
|
||||
ctx := context.Background()
|
||||
ct, st := NewInMemoryTransports()
|
||||
|
||||
s := NewServer("testServer", "v1.0.0", nil)
|
||||
|
||||
// The 'greet' tool says hi.
|
||||
s.AddTools(tools...)
|
||||
if config != nil {
|
||||
config(s)
|
||||
}
|
||||
ss, err := s.Connect(ctx, st)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
|
|
@ -569,7 +528,9 @@ func basicConnection(t *testing.T, tools ...*ServerTool) (*ServerSession, *Clien
|
|||
}
|
||||
|
||||
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()
|
||||
|
||||
ctx := context.Background()
|
||||
|
|
@ -651,11 +612,9 @@ func TestCancellation(t *testing.T) {
|
|||
}
|
||||
return nil, nil
|
||||
}
|
||||
st := &ServerTool{
|
||||
Tool: &Tool{Name: "slow"},
|
||||
Handler: slowRequest,
|
||||
}
|
||||
_, cs := basicConnection(t, st)
|
||||
_, cs := basicConnection(t, func(s *Server) {
|
||||
s.AddTool(&Tool{Name: "slow"}, slowRequest)
|
||||
})
|
||||
defer cs.Close()
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
|
|
@ -852,7 +811,7 @@ func TestKeepAlive(t *testing.T) {
|
|||
KeepAlive: 100 * time.Millisecond,
|
||||
}
|
||||
s := NewServer("testServer", "v1.0.0", serverOpts)
|
||||
s.AddTools(NewServerTool("greet", "say hi", sayHi))
|
||||
AddTool(s, greetTool(), sayHi)
|
||||
|
||||
ss, err := s.Connect(ctx, st)
|
||||
if err != nil {
|
||||
|
|
@ -897,7 +856,7 @@ func TestKeepAliveFailure(t *testing.T) {
|
|||
|
||||
// Server without keepalive (to test one-sided keepalive)
|
||||
s := NewServer("testServer", "v1.0.0", nil)
|
||||
s.AddTools(NewServerTool("greet", "say hi", sayHi))
|
||||
AddTool(s, greetTool(), sayHi)
|
||||
ss, err := s.Connect(ctx, st)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
|
|
|
|||
|
|
@ -11,8 +11,7 @@ import (
|
|||
// A PromptHandler handles a call to prompts/get.
|
||||
type PromptHandler func(context.Context, *ServerSession, *GetPromptParams) (*GetPromptResult, error)
|
||||
|
||||
// A Prompt is a prompt definition bound to a prompt handler.
|
||||
type ServerPrompt struct {
|
||||
Prompt *Prompt
|
||||
Handler PromptHandler
|
||||
type serverPrompt struct {
|
||||
prompt *Prompt
|
||||
handler PromptHandler
|
||||
}
|
||||
|
|
|
|||
|
|
@ -20,16 +20,16 @@ import (
|
|||
"github.com/modelcontextprotocol/go-sdk/internal/util"
|
||||
)
|
||||
|
||||
// A ServerResource associates a Resource with its handler.
|
||||
type ServerResource struct {
|
||||
Resource *Resource
|
||||
Handler ResourceHandler
|
||||
// A serverResource associates a Resource with its handler.
|
||||
type serverResource struct {
|
||||
resource *Resource
|
||||
handler ResourceHandler
|
||||
}
|
||||
|
||||
// A ServerResourceTemplate associates a ResourceTemplate with its handler.
|
||||
type ServerResourceTemplate struct {
|
||||
ResourceTemplate *ResourceTemplate
|
||||
Handler ResourceHandler
|
||||
// A serverResourceTemplate associates a ResourceTemplate with its handler.
|
||||
type serverResourceTemplate struct {
|
||||
resourceTemplate *ResourceTemplate
|
||||
handler ResourceHandler
|
||||
}
|
||||
|
||||
// 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.
|
||||
// TODO: use "github.com/yosida95/uritemplate/v3"
|
||||
func (sr *ServerResourceTemplate) Matches(uri string) bool {
|
||||
re, err := uriTemplateToRegexp(sr.ResourceTemplate.URITemplate)
|
||||
func (sr *serverResourceTemplate) Matches(uri string) bool {
|
||||
re, err := uriTemplateToRegexp(sr.resourceTemplate.URITemplate)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
|
|
|
|||
170
mcp/server.go
170
mcp/server.go
|
|
@ -12,6 +12,7 @@ import (
|
|||
"encoding/json"
|
||||
"fmt"
|
||||
"iter"
|
||||
"log"
|
||||
"net/url"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
|
|
@ -20,7 +21,6 @@ import (
|
|||
|
||||
"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
|
||||
"github.com/modelcontextprotocol/go-sdk/internal/util"
|
||||
"github.com/modelcontextprotocol/go-sdk/jsonschema"
|
||||
)
|
||||
|
||||
const DefaultPageSize = 1000
|
||||
|
|
@ -36,10 +36,10 @@ type Server struct {
|
|||
opts ServerOptions
|
||||
|
||||
mu sync.Mutex
|
||||
prompts *featureSet[*ServerPrompt]
|
||||
tools *featureSet[*ServerTool]
|
||||
resources *featureSet[*ServerResource]
|
||||
resourceTemplates *featureSet[*ServerResourceTemplate]
|
||||
prompts *featureSet[*serverPrompt]
|
||||
tools *featureSet[*serverTool]
|
||||
resources *featureSet[*serverResource]
|
||||
resourceTemplates *featureSet[*serverResourceTemplate]
|
||||
sessions []*ServerSession
|
||||
sendingMethodHandler_ MethodHandler[*ServerSession]
|
||||
receivingMethodHandler_ MethodHandler[*ServerSession]
|
||||
|
|
@ -87,28 +87,23 @@ func NewServer(name, version string, opts *ServerOptions) *Server {
|
|||
name: name,
|
||||
version: version,
|
||||
opts: *opts,
|
||||
prompts: newFeatureSet(func(p *ServerPrompt) string { return p.Prompt.Name }),
|
||||
tools: newFeatureSet(func(t *ServerTool) string { return t.Tool.Name }),
|
||||
resources: newFeatureSet(func(r *ServerResource) string { return r.Resource.URI }),
|
||||
resourceTemplates: newFeatureSet(func(t *ServerResourceTemplate) string { return t.ResourceTemplate.URITemplate }),
|
||||
prompts: newFeatureSet(func(p *serverPrompt) string { return p.prompt.Name }),
|
||||
tools: newFeatureSet(func(t *serverTool) string { return t.tool.Name }),
|
||||
resources: newFeatureSet(func(r *serverResource) string { return r.resource.URI }),
|
||||
resourceTemplates: newFeatureSet(func(t *serverResourceTemplate) string { return t.resourceTemplate.URITemplate }),
|
||||
sendingMethodHandler_: defaultSendingMethodHandler[*ServerSession],
|
||||
receivingMethodHandler_: defaultReceivingMethodHandler[*ServerSession],
|
||||
}
|
||||
}
|
||||
|
||||
// AddPrompts adds the given prompts to the server,
|
||||
// replacing any with the same names.
|
||||
func (s *Server) AddPrompts(prompts ...*ServerPrompt) {
|
||||
// Only notify if something could change.
|
||||
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.)
|
||||
// AddPrompt adds a [Prompt] to the server, or replaces one with the same name.
|
||||
func (s *Server) AddPrompt(p *Prompt, h PromptHandler) {
|
||||
// Assume there was a change, since add replaces existing items.
|
||||
// (It's possible an item was replaced with an identical one, but not worth checking.)
|
||||
s.changeAndNotify(
|
||||
notificationPromptListChanged,
|
||||
&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.
|
||||
|
|
@ -118,55 +113,44 @@ func (s *Server) RemovePrompts(names ...string) {
|
|||
func() bool { return s.prompts.remove(names...) })
|
||||
}
|
||||
|
||||
// AddTools adds the given tools to the server,
|
||||
// replacing any with the same names.
|
||||
// The arguments must not be modified after this call.
|
||||
//
|
||||
// AddTools panics if errors are detected.
|
||||
func (s *Server) AddTools(tools ...*ServerTool) {
|
||||
if err := s.addToolsErr(tools...); err != nil {
|
||||
// AddTool adds a [Tool] to the server, or replaces one with the same name.
|
||||
// The tool's input schema must be non-nil.
|
||||
// The Tool argument must not be modified after this call.
|
||||
func (s *Server) AddTool(t *Tool, h ToolHandler) {
|
||||
// TODO(jba): This is a breaking behavior change. Add before v0.2.0?
|
||||
if t.InputSchema == 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)
|
||||
}
|
||||
}
|
||||
|
||||
// addToolsErr is like [AddTools], but returns an error instead of panicking.
|
||||
func (s *Server) addToolsErr(tools ...*ServerTool) error {
|
||||
// Only notify if something could change.
|
||||
if len(tools) == 0 {
|
||||
return nil
|
||||
// AddTool adds a [Tool] to the server, or replaces one with the same name.
|
||||
// If the tool's input schema is nil, it is set to the schema inferred from the In
|
||||
// type parameter, using [jsonschema.For].
|
||||
// If the tool's output schema is nil and the Out type parameter is not the empty
|
||||
// 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 {
|
||||
// st.outputResolved, err := st.Tool.OutputSchema.Resolve(&jsonschema.ResolveOptions{ValidateDefaults: true})
|
||||
// if err != nil {
|
||||
// return err
|
||||
// }
|
||||
// }
|
||||
}
|
||||
func addToolErr[In, Out any](s *Server, t *Tool, h ToolHandlerFor[In, Out]) (err error) {
|
||||
defer util.Wrapf(&err, "adding tool %q", t.Name)
|
||||
st, err := newServerTool(t, h)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// 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.)
|
||||
// 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{},
|
||||
func() bool { s.tools.add(tools...); return true })
|
||||
func() bool { s.tools.add(st); return true })
|
||||
return nil
|
||||
}
|
||||
|
||||
|
|
@ -177,26 +161,19 @@ func (s *Server) RemoveTools(names ...string) {
|
|||
func() bool { return s.tools.remove(names...) })
|
||||
}
|
||||
|
||||
// AddResources adds the given resources to the server.
|
||||
// If a resource with the same URI already exists, it is replaced.
|
||||
// AddResources panics if a resource URI is invalid or not absolute (has an empty scheme).
|
||||
func (s *Server) AddResources(resources ...*ServerResource) {
|
||||
// Only notify if something could change.
|
||||
if len(resources) == 0 {
|
||||
return
|
||||
}
|
||||
// AddResource adds a [Resource] to the server, or replaces one with the same URI.
|
||||
// AddResource panics if the resource URI is invalid or not absolute (has an empty scheme).
|
||||
func (s *Server) AddResource(r *Resource, h ResourceHandler) {
|
||||
s.changeAndNotify(notificationResourceListChanged, &ResourceListChangedParams{},
|
||||
func() bool {
|
||||
for _, r := range resources {
|
||||
u, err := url.Parse(r.Resource.URI)
|
||||
if err != nil {
|
||||
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)
|
||||
u, err := url.Parse(r.URI)
|
||||
if err != nil {
|
||||
panic(err) // url.Parse includes the URI in the error
|
||||
}
|
||||
if !u.IsAbs() {
|
||||
panic(fmt.Errorf("URI %s needs a scheme", r.URI))
|
||||
}
|
||||
s.resources.add(&serverResource{r, h})
|
||||
return true
|
||||
})
|
||||
}
|
||||
|
|
@ -208,20 +185,13 @@ func (s *Server) RemoveResources(uris ...string) {
|
|||
func() bool { return s.resources.remove(uris...) })
|
||||
}
|
||||
|
||||
// AddResourceTemplates adds the given resource templates to the server.
|
||||
// If a resource template with the same URI template already exists, it will be replaced.
|
||||
// AddResourceTemplates panics if a URI template is invalid or not absolute (has an empty scheme).
|
||||
func (s *Server) AddResourceTemplates(templates ...*ServerResourceTemplate) {
|
||||
// Only notify if something could change.
|
||||
if len(templates) == 0 {
|
||||
return
|
||||
}
|
||||
// AddResourceTemplate adds a [ResourceTemplate] to the server, or replaces on with the same URI.
|
||||
// AddResourceTemplate panics if a URI template is invalid or not absolute (has an empty scheme).
|
||||
func (s *Server) AddResourceTemplate(t *ResourceTemplate, h ResourceHandler) {
|
||||
s.changeAndNotify(notificationResourceListChanged, &ResourceListChangedParams{},
|
||||
func() bool {
|
||||
for _, t := range templates {
|
||||
// TODO: check template validity.
|
||||
s.resourceTemplates.add(t)
|
||||
}
|
||||
// TODO: check template validity.
|
||||
s.resourceTemplates.add(&serverResourceTemplate{t, h})
|
||||
return true
|
||||
})
|
||||
}
|
||||
|
|
@ -268,10 +238,10 @@ func (s *Server) listPrompts(_ context.Context, _ *ServerSession, params *ListPr
|
|||
if params == nil {
|
||||
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
|
||||
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.
|
||||
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) {
|
||||
|
|
@ -293,22 +263,22 @@ func (s *Server) listTools(_ context.Context, _ *ServerSession, params *ListTool
|
|||
if params == nil {
|
||||
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
|
||||
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) {
|
||||
s.mu.Lock()
|
||||
tool, ok := s.tools.get(params.Name)
|
||||
st, ok := s.tools.get(params.Name)
|
||||
s.mu.Unlock()
|
||||
if !ok {
|
||||
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) {
|
||||
|
|
@ -317,10 +287,10 @@ func (s *Server) listResources(_ context.Context, _ *ServerSession, params *List
|
|||
if params == nil {
|
||||
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
|
||||
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{}
|
||||
}
|
||||
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
|
||||
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()
|
||||
// Try resources first.
|
||||
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.
|
||||
for rt := range s.resourceTemplates.all() {
|
||||
if rt.Matches(uri) {
|
||||
return rt.Handler, rt.ResourceTemplate.MIMEType, true
|
||||
return rt.handler, rt.resourceTemplate.MIMEType, true
|
||||
}
|
||||
}
|
||||
return nil, "", false
|
||||
|
|
|
|||
|
|
@ -29,7 +29,7 @@ func ExampleServer() {
|
|||
clientTransport, serverTransport := mcp.NewInMemoryTransports()
|
||||
|
||||
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)
|
||||
if err != nil {
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
func TestNewServerToolValidate(t *testing.T) {
|
||||
func TestToolValidate(t *testing.T) {
|
||||
// Check that the tool returned from NewServerTool properly validates its input schema.
|
||||
|
||||
type req struct {
|
||||
|
|
@ -26,9 +26,10 @@ func TestNewServerToolValidate(t *testing.T) {
|
|||
return nil, nil
|
||||
}
|
||||
|
||||
tool := NewServerTool("test", "test", dummyHandler)
|
||||
// Need to add the tool to a server to get resolved schemas.
|
||||
// s := NewServer("", "", nil)
|
||||
st, err := newServerTool(&Tool{Name: "test", Description: "test"}, dummyHandler)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
for _, tt := range []struct {
|
||||
desc string
|
||||
|
|
@ -71,7 +72,7 @@ func TestNewServerToolValidate(t *testing.T) {
|
|||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, err = tool.rawHandler(context.Background(), nil,
|
||||
_, err = st.handler(context.Background(), nil,
|
||||
&CallToolParamsFor[json.RawMessage]{Arguments: json.RawMessage(raw)})
|
||||
if err == nil && tt.want != "" {
|
||||
t.Error("got success, wanted failure")
|
||||
|
|
|
|||
|
|
@ -28,7 +28,7 @@ func Add(ctx context.Context, cc *mcp.ServerSession, params *mcp.CallToolParamsF
|
|||
|
||||
func ExampleSSEHandler() {
|
||||
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 })
|
||||
httpServer := httptest.NewServer(handler)
|
||||
|
|
|
|||
|
|
@ -20,7 +20,7 @@ func TestSSEServer(t *testing.T) {
|
|||
t.Run(fmt.Sprintf("closeServerFirst=%t", closeServerFirst), func(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
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 })
|
||||
|
||||
|
|
|
|||
|
|
@ -32,8 +32,7 @@ func TestStreamableTransports(t *testing.T) {
|
|||
|
||||
// 1. Create a server with a simple "greet" tool.
|
||||
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
|
||||
// cookie-checking middleware.
|
||||
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
|
||||
// behavior, if any.
|
||||
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 {
|
||||
test.tool(t, ctx, ss)
|
||||
}
|
||||
return &CallToolResultFor[any]{}, nil
|
||||
})
|
||||
server.AddTools(tool)
|
||||
|
||||
// Start the streamable handler.
|
||||
handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, nil)
|
||||
|
|
|
|||
233
mcp/tool.go
233
mcp/tool.go
|
|
@ -9,7 +9,7 @@ import (
|
|||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"slices"
|
||||
"reflect"
|
||||
|
||||
"github.com/modelcontextprotocol/go-sdk/jsonschema"
|
||||
)
|
||||
|
|
@ -17,8 +17,7 @@ import (
|
|||
// A ToolHandler handles a call to tools/call.
|
||||
// [CallToolParams.Arguments] will contain a map[string]any that has been validated
|
||||
// against the input schema.
|
||||
// TODO: Perhaps this should be an alias for ToolHandlerFor[map[string]any, map[string]any]?
|
||||
type ToolHandler func(context.Context, *ServerSession, *CallToolParamsFor[map[string]any]) (*CallToolResult, error)
|
||||
type ToolHandler = ToolHandlerFor[map[string]any, any]
|
||||
|
||||
// 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)
|
||||
|
|
@ -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.
|
||||
type rawToolHandler = func(context.Context, *ServerSession, *CallToolParamsFor[json.RawMessage]) (*CallToolResult, error)
|
||||
|
||||
// A ServerTool is a tool definition that is bound to a tool handler.
|
||||
type ServerTool struct {
|
||||
Tool *Tool
|
||||
Handler ToolHandler
|
||||
// Set in NewServerTool or Server.addToolsErr.
|
||||
rawHandler rawToolHandler
|
||||
// Resolved tool schemas. Set in Server.addToolsErr.
|
||||
// A serverTool is a tool definition that is bound to a tool handler.
|
||||
type serverTool struct {
|
||||
tool *Tool
|
||||
handler rawToolHandler
|
||||
// Resolved tool schemas. Set in newServerTool.
|
||||
inputResolved, outputResolved *jsonschema.Resolved
|
||||
}
|
||||
|
||||
// NewServerTool is a helper to make a tool using reflection on the given type parameters.
|
||||
// When the tool is called, CallToolParams.Arguments will be of type In.
|
||||
//
|
||||
// If provided, variadic [ToolOption] values may be used to customize the tool.
|
||||
//
|
||||
// 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
|
||||
}
|
||||
// newServerTool creates a serverTool from a tool and a handler.
|
||||
// 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.
|
||||
func newServerTool[In, Out any](t *Tool, h ToolHandlerFor[In, Out]) (*serverTool, error) {
|
||||
st := &serverTool{tool: t}
|
||||
|
||||
func newServerToolErr[In, Out any](name, description string, handler ToolHandlerFor[In, Out], opts ...ToolOption) (*ServerTool, error) {
|
||||
// TODO: check that In is a struct.
|
||||
ischema, err := jsonschema.For[In]()
|
||||
if err != nil {
|
||||
if err := setSchema[In](&t.InputSchema, &st.inputResolved); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// TODO: uncomment when output schemas drop.
|
||||
// oschema, err := jsonschema.For[TRes]()
|
||||
// if err != nil {
|
||||
// return nil, err
|
||||
// }
|
||||
|
||||
t := &ServerTool{
|
||||
Tool: &Tool{
|
||||
Name: name,
|
||||
Description: description,
|
||||
InputSchema: ischema,
|
||||
// OutputSchema: oschema,
|
||||
},
|
||||
}
|
||||
for _, opt := range opts {
|
||||
opt.set(t)
|
||||
if reflect.TypeFor[Out]() != reflect.TypeFor[any]() {
|
||||
if err := setSchema[Out](&t.OutputSchema, &st.outputResolved); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
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
|
||||
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
|
||||
}
|
||||
}
|
||||
|
|
@ -91,45 +61,7 @@ func newServerToolErr[In, Out any](name, description string, handler ToolHandler
|
|||
Name: rparams.Name,
|
||||
Arguments: args,
|
||||
}
|
||||
res, err := handler(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)
|
||||
res, err := h(ctx, ss, params)
|
||||
// TODO(rfindley): investigate why server errors are embedded in this strange way,
|
||||
// rather than returned as jsonrpc2 server errors.
|
||||
if err != nil {
|
||||
|
|
@ -138,8 +70,32 @@ func newRawHandler(st *ServerTool) rawToolHandler {
|
|||
IsError: true,
|
||||
}, 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
|
||||
|
|
@ -169,105 +125,6 @@ func unmarshalSchema(data json.RawMessage, resolved *jsonschema.Resolved, v any)
|
|||
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.
|
||||
func schemaJSON(s *jsonschema.Schema) string {
|
||||
m, err := json.Marshal(s)
|
||||
|
|
|
|||
112
mcp/tool_test.go
112
mcp/tool_test.go
|
|
@ -16,76 +16,80 @@ import (
|
|||
)
|
||||
|
||||
// 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")
|
||||
}
|
||||
|
||||
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) {
|
||||
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 {
|
||||
tool *ServerTool
|
||||
want *jsonschema.Schema
|
||||
tool *serverTool
|
||||
wantIn, wantOut *jsonschema.Schema
|
||||
}{
|
||||
{
|
||||
NewServerTool("basic", "", testToolHandler[struct {
|
||||
Name string `json:"name"`
|
||||
}]),
|
||||
&jsonschema.Schema{
|
||||
Type: "object",
|
||||
Required: []string{"name"},
|
||||
Properties: map[string]*jsonschema.Schema{
|
||||
"name": {Type: "string"},
|
||||
},
|
||||
AdditionalProperties: &jsonschema.Schema{Not: new(jsonschema.Schema)},
|
||||
},
|
||||
srvTool(t, &Tool{Name: "basic"}, testToolHandler[Name, Size]),
|
||||
nameSchema,
|
||||
sizeSchema,
|
||||
},
|
||||
{
|
||||
NewServerTool("enum", "", testToolHandler[struct{ Name string }], Input(
|
||||
Property("Name", Enum("x", "y", "z")),
|
||||
)),
|
||||
&jsonschema.Schema{
|
||||
Type: "object",
|
||||
Required: []string{"Name"},
|
||||
Properties: map[string]*jsonschema.Schema{
|
||||
"Name": {Type: "string", Enum: []any{"x", "y", "z"}},
|
||||
},
|
||||
AdditionalProperties: &jsonschema.Schema{Not: new(jsonschema.Schema)},
|
||||
},
|
||||
srvTool(t, &Tool{
|
||||
Name: "in untouched",
|
||||
InputSchema: &jsonschema.Schema{},
|
||||
}, testToolHandler[Name, Size]),
|
||||
&jsonschema.Schema{},
|
||||
sizeSchema,
|
||||
},
|
||||
{
|
||||
NewServerTool("required", "", testToolHandler[struct {
|
||||
Name string `json:"name"`
|
||||
Language string `json:"language"`
|
||||
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)},
|
||||
},
|
||||
srvTool(t, &Tool{Name: "out untouched", OutputSchema: &jsonschema.Schema{}}, testToolHandler[Name, Size]),
|
||||
nameSchema,
|
||||
&jsonschema.Schema{},
|
||||
},
|
||||
{
|
||||
NewServerTool("set_schema", "", testToolHandler[struct {
|
||||
X int `json:"x,omitempty"`
|
||||
Y int `json:"y,omitempty"`
|
||||
}], Input(
|
||||
Schema(&jsonschema.Schema{Type: "object"})),
|
||||
),
|
||||
&jsonschema.Schema{
|
||||
Type: "object",
|
||||
},
|
||||
srvTool(t, &Tool{Name: "nil out"}, testToolHandler[Name, any]),
|
||||
nameSchema,
|
||||
nil,
|
||||
},
|
||||
}
|
||||
for _, test := range tests {
|
||||
if diff := cmp.Diff(test.want, test.tool.Tool.InputSchema, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
|
||||
t.Errorf("NewServerTool(%v) mismatch (-want +got):\n%s", test.tool.Tool.Name, diff)
|
||||
if diff := cmp.Diff(test.wantIn, test.tool.tool.InputSchema, cmpopts.IgnoreUnexported(jsonschema.Schema{})); 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)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue