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
1 changed files with 366 additions and 612 deletions
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
|
||||
|
|
|
|||
Loading…
Reference in a new issue