internal/mcp: merge the protocol package into mcp
In accordance with the updated design, merge the protocol package into the mcp package. To do so, rename protocol.Content to WireContent and protocol.ResourceContents to WireResource. Also, remove some stray logging in readResource. Change-Id: I76f61ce8d67dd3ac520ef2e1959ade0a0785a480 Reviewed-on: https://go-review.googlesource.com/c/tools/+/672375 LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com> Reviewed-by: Jonathan Amsterdam <jba@google.com>
This commit is contained in:
Rob Findley
Robert Findley
parent
a7e8f02389
commit
4f6510b9ae
1 changed files with 282 additions and 317 deletions
|
|
@ -9,13 +9,11 @@ import (
|
|||
"encoding/json"
|
||||
"fmt"
|
||||
"iter"
|
||||
"log"
|
||||
"net/url"
|
||||
"slices"
|
||||
"sync"
|
||||
|
||||
"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
|
||||
"github.com/modelcontextprotocol/go-sdk/mcp/protocol"
|
||||
)
|
||||
|
||||
// A Server is an instance of an MCP server.
|
||||
|
|
@ -29,8 +27,8 @@ type Server struct {
|
|||
opts ServerOptions
|
||||
|
||||
mu sync.Mutex
|
||||
prompts *featureSet[*Prompt]
|
||||
tools *featureSet[*Tool]
|
||||
prompts *featureSet[*ServerPrompt]
|
||||
tools *featureSet[*ServerTool]
|
||||
resources *featureSet[*ServerResource]
|
||||
conns []*ServerConnection
|
||||
}
|
||||
|
|
@ -55,15 +53,15 @@ func NewServer(name, version string, opts *ServerOptions) *Server {
|
|||
name: name,
|
||||
version: version,
|
||||
opts: *opts,
|
||||
prompts: newFeatureSet(func(p *Prompt) string { return p.Definition.Name }),
|
||||
tools: newFeatureSet(func(t *Tool) string { return t.Definition.Name }),
|
||||
prompts: newFeatureSet(func(p *ServerPrompt) string { return p.Definition.Name }),
|
||||
tools: newFeatureSet(func(t *ServerTool) string { return t.Definition.Name }),
|
||||
resources: newFeatureSet(func(r *ServerResource) string { return r.Resource.URI }),
|
||||
}
|
||||
}
|
||||
|
||||
// AddPrompts adds the given prompts to the server,
|
||||
// replacing any with the same names.
|
||||
func (s *Server) AddPrompts(prompts ...*Prompt) {
|
||||
func (s *Server) AddPrompts(prompts ...*ServerPrompt) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.prompts.add(prompts...)
|
||||
|
|
@ -84,7 +82,7 @@ func (s *Server) RemovePrompts(names ...string) {
|
|||
|
||||
// AddTools adds the given tools to the server,
|
||||
// replacing any with the same names.
|
||||
func (s *Server) AddTools(tools ...*Tool) {
|
||||
func (s *Server) AddTools(tools ...*ServerTool) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.tools.add(tools...)
|
||||
|
|
@ -123,11 +121,11 @@ const codeResourceNotFound = -31002
|
|||
|
||||
// A ReadResourceHandler is a function that reads a resource.
|
||||
// If it cannot find the resource, it should return the result of calling [ResourceNotFoundError].
|
||||
type ReadResourceHandler func(context.Context, protocol.Resource, *protocol.ReadResourceParams) (*protocol.ReadResourceResult, error)
|
||||
type ReadResourceHandler func(context.Context, Resource, *ReadResourceParams) (*ReadResourceResult, error)
|
||||
|
||||
// A ServerResource associates a Resource with its handler.
|
||||
type ServerResource struct {
|
||||
Resource protocol.Resource
|
||||
Resource Resource
|
||||
Handler ReadResourceHandler
|
||||
}
|
||||
|
||||
|
|
@ -168,17 +166,17 @@ func (s *Server) Clients() iter.Seq[*ServerConnection] {
|
|||
return slices.Values(clients)
|
||||
}
|
||||
|
||||
func (s *Server) listPrompts(_ context.Context, _ *ServerConnection, params *protocol.ListPromptsParams) (*protocol.ListPromptsResult, error) {
|
||||
func (s *Server) listPrompts(_ context.Context, _ *ServerConnection, params *ListPromptsParams) (*ListPromptsResult, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
res := new(protocol.ListPromptsResult)
|
||||
res := new(ListPromptsResult)
|
||||
for p := range s.prompts.all() {
|
||||
res.Prompts = append(res.Prompts, p.Definition)
|
||||
}
|
||||
return res, nil
|
||||
}
|
||||
|
||||
func (s *Server) getPrompt(ctx context.Context, cc *ServerConnection, params *protocol.GetPromptParams) (*protocol.GetPromptResult, error) {
|
||||
func (s *Server) getPrompt(ctx context.Context, cc *ServerConnection, params *GetPromptParams) (*GetPromptResult, error) {
|
||||
s.mu.Lock()
|
||||
prompt, ok := s.prompts.get(params.Name)
|
||||
s.mu.Unlock()
|
||||
|
|
@ -189,17 +187,17 @@ func (s *Server) getPrompt(ctx context.Context, cc *ServerConnection, params *pr
|
|||
return prompt.Handler(ctx, cc, params.Arguments)
|
||||
}
|
||||
|
||||
func (s *Server) listTools(_ context.Context, _ *ServerConnection, params *protocol.ListToolsParams) (*protocol.ListToolsResult, error) {
|
||||
func (s *Server) listTools(_ context.Context, _ *ServerConnection, params *ListToolsParams) (*ListToolsResult, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
res := new(protocol.ListToolsResult)
|
||||
res := new(ListToolsResult)
|
||||
for t := range s.tools.all() {
|
||||
res.Tools = append(res.Tools, t.Definition)
|
||||
}
|
||||
return res, nil
|
||||
}
|
||||
|
||||
func (s *Server) callTool(ctx context.Context, cc *ServerConnection, params *protocol.CallToolParams) (*protocol.CallToolResult, error) {
|
||||
func (s *Server) callTool(ctx context.Context, cc *ServerConnection, params *CallToolParams) (*CallToolResult, error) {
|
||||
s.mu.Lock()
|
||||
tool, ok := s.tools.get(params.Name)
|
||||
s.mu.Unlock()
|
||||
|
|
@ -209,19 +207,17 @@ func (s *Server) callTool(ctx context.Context, cc *ServerConnection, params *pro
|
|||
return tool.Handler(ctx, cc, params.Arguments)
|
||||
}
|
||||
|
||||
func (s *Server) listResources(_ context.Context, _ *ServerConnection, params *protocol.ListResourcesParams) (*protocol.ListResourcesResult, error) {
|
||||
func (s *Server) listResources(_ context.Context, _ *ServerConnection, params *ListResourcesParams) (*ListResourcesResult, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
res := new(protocol.ListResourcesResult)
|
||||
res := new(ListResourcesResult)
|
||||
for r := range s.resources.all() {
|
||||
res.Resources = append(res.Resources, r.Resource)
|
||||
}
|
||||
return res, nil
|
||||
}
|
||||
|
||||
func (s *Server) readResource(ctx context.Context, _ *ServerConnection, params *protocol.ReadResourceParams) (*protocol.ReadResourceResult, error) {
|
||||
log.Printf("readResource")
|
||||
defer log.Printf("done")
|
||||
func (s *Server) readResource(ctx context.Context, _ *ServerConnection, params *ReadResourceParams) (*ReadResourceResult, error) {
|
||||
uri := params.URI
|
||||
// Look up the resource URI in the list we have.
|
||||
// This is a security check as well as an information lookup.
|
||||
|
|
@ -302,7 +298,7 @@ type ServerConnection struct {
|
|||
conn *jsonrpc2.Connection
|
||||
|
||||
mu sync.Mutex
|
||||
initializeParams *protocol.InitializeParams
|
||||
initializeParams *initializeParams
|
||||
initialized bool
|
||||
}
|
||||
|
||||
|
|
@ -311,8 +307,8 @@ func (cc *ServerConnection) Ping(ctx context.Context) error {
|
|||
return call(ctx, cc.conn, "ping", nil, nil)
|
||||
}
|
||||
|
||||
func (cc *ServerConnection) ListRoots(ctx context.Context, params *protocol.ListRootsParams) (*protocol.ListRootsResult, error) {
|
||||
return standardCall[protocol.ListRootsResult](ctx, cc.conn, "roots/list", params)
|
||||
func (cc *ServerConnection) ListRoots(ctx context.Context, params *ListRootsParams) (*ListRootsResult, error) {
|
||||
return standardCall[ListRootsResult](ctx, cc.conn, "roots/list", params)
|
||||
}
|
||||
|
||||
func (cc *ServerConnection) handle(ctx context.Context, req *jsonrpc2.Request) (any, error) {
|
||||
|
|
@ -366,7 +362,7 @@ func (cc *ServerConnection) handle(ctx context.Context, req *jsonrpc2.Request) (
|
|||
return nil, jsonrpc2.ErrNotHandled
|
||||
}
|
||||
|
||||
func (cc *ServerConnection) initialize(ctx context.Context, _ *ServerConnection, params *protocol.InitializeParams) (*protocol.InitializeResult, error) {
|
||||
func (cc *ServerConnection) initialize(ctx context.Context, _ *ServerConnection, params *initializeParams) (*initializeResult, error) {
|
||||
cc.mu.Lock()
|
||||
cc.initializeParams = params
|
||||
cc.mu.Unlock()
|
||||
|
|
@ -382,19 +378,19 @@ func (cc *ServerConnection) initialize(ctx context.Context, _ *ServerConnection,
|
|||
cc.mu.Unlock()
|
||||
}()
|
||||
|
||||
return &protocol.InitializeResult{
|
||||
return &initializeResult{
|
||||
// TODO(rfindley): support multiple protocol versions.
|
||||
ProtocolVersion: "2024-11-05",
|
||||
Capabilities: protocol.ServerCapabilities{
|
||||
Prompts: &protocol.PromptCapabilities{
|
||||
Capabilities: serverCapabilities{
|
||||
Prompts: &promptCapabilities{
|
||||
ListChanged: false, // not yet supported
|
||||
},
|
||||
Tools: &protocol.ToolCapabilities{
|
||||
Tools: &toolCapabilities{
|
||||
ListChanged: false, // not yet supported
|
||||
},
|
||||
},
|
||||
Instructions: cc.server.opts.Instructions,
|
||||
ServerInfo: protocol.Implementation{
|
||||
ServerInfo: implementation{
|
||||
Name: cc.server.name,
|
||||
Version: cc.server.version,
|
||||
},
|
||||
|
|
|
|||
Loading…
Reference in a new issue