internal/mcp: add notifications
Add notification logic for intialization and list changes. Change-Id: I1d392bb49d5995f2046c400da955a033ab2715ce Reviewed-on: https://go-review.googlesource.com/c/tools/+/673775 Reviewed-by: Sam Thanawalla <samthanawalla@google.com> LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com> Reviewed-by: Robert Findley <rfindley@google.com>
This commit is contained in:
Jonathan Amsterdam
parent
ec8fb49c6b
commit
ce5b6cd926
6 changed files with 273 additions and 78 deletions
|
|
@ -49,6 +49,10 @@ type ClientOptions struct {
|
|||
// Handler for sampling.
|
||||
// Called when a server calls CreateMessage.
|
||||
CreateMessageHandler func(context.Context, *ClientSession, *CreateMessageParams) (*CreateMessageResult, error)
|
||||
// Handlers for notifications from the server.
|
||||
ToolListChangedHandler func(context.Context, *ClientSession, *ToolListChangedParams)
|
||||
PromptListChangedHandler func(context.Context, *ClientSession, *PromptListChangedParams)
|
||||
ResourceListChangedHandler func(context.Context, *ClientSession, *ResourceListChangedParams)
|
||||
}
|
||||
|
||||
// bind implements the binder[*ClientSession] interface, so that Clients can
|
||||
|
|
@ -86,10 +90,13 @@ func (c *Client) Connect(ctx context.Context, t Transport) (cs *ClientSession, e
|
|||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
caps := &ClientCapabilities{}
|
||||
caps.Roots.ListChanged = true
|
||||
if c.opts.CreateMessageHandler != nil {
|
||||
caps.Sampling = &SamplingCapabilities{}
|
||||
}
|
||||
|
||||
params := &InitializeParams{
|
||||
ClientInfo: &implementation{Name: c.name, Version: c.version},
|
||||
Capabilities: caps,
|
||||
|
|
@ -133,11 +140,13 @@ func (c *ClientSession) Wait() error {
|
|||
// AddRoots adds the given roots to the client,
|
||||
// replacing any with the same URIs,
|
||||
// and notifies any connected servers.
|
||||
// TODO: notification
|
||||
func (c *Client) AddRoots(roots ...*Root) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.roots.add(roots...)
|
||||
// Only notify if something could change.
|
||||
if len(roots) == 0 {
|
||||
return
|
||||
}
|
||||
c.changeAndNotify(notificationRootsListChanged, &RootsListChangedParams{},
|
||||
func() bool { c.roots.add(roots...); return true })
|
||||
}
|
||||
|
||||
// RemoveRoots removes the roots with the given URIs,
|
||||
|
|
@ -145,9 +154,22 @@ func (c *Client) AddRoots(roots ...*Root) {
|
|||
// It is not an error to remove a nonexistent root.
|
||||
// TODO: notification
|
||||
func (c *Client) RemoveRoots(uris ...string) {
|
||||
c.changeAndNotify(notificationRootsListChanged, &RootsListChangedParams{},
|
||||
func() bool { return c.roots.remove(uris...) })
|
||||
}
|
||||
|
||||
// changeAndNotify is called when a feature is added or removed.
|
||||
// It calls change, which should do the work and report whether a change actually occurred.
|
||||
// If there was a change, it notifies a snapshot of the sessions.
|
||||
func (c *Client) changeAndNotify(notification string, params any, change func() bool) {
|
||||
var sessions []*ClientSession
|
||||
// Lock for the change, but not for the notification.
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.roots.remove(uris...)
|
||||
if change() {
|
||||
sessions = slices.Clone(c.sessions)
|
||||
}
|
||||
c.mu.Unlock()
|
||||
notifySessions(sessions, notification, params)
|
||||
}
|
||||
|
||||
func (c *Client) listRoots(_ context.Context, _ *ClientSession, _ *ListRootsParams) (*ListRootsResult, error) {
|
||||
|
|
@ -180,10 +202,12 @@ func (c *Client) AddMiddleware(middleware ...Middleware[ClientSession]) {
|
|||
|
||||
// clientMethodInfos maps from the RPC method name to serverMethodInfos.
|
||||
var clientMethodInfos = map[string]methodInfo[ClientSession]{
|
||||
methodPing: newMethodInfo(sessionMethod((*ClientSession).ping)),
|
||||
methodListRoots: newMethodInfo(clientMethod((*Client).listRoots)),
|
||||
methodCreateMessage: newMethodInfo(clientMethod((*Client).createMessage)),
|
||||
// TODO: notifications
|
||||
methodPing: newMethodInfo(sessionMethod((*ClientSession).ping)),
|
||||
methodListRoots: newMethodInfo(clientMethod((*Client).listRoots)),
|
||||
methodCreateMessage: newMethodInfo(clientMethod((*Client).createMessage)),
|
||||
notificationToolListChanged: newMethodInfo(clientMethod((*Client).callToolChangedHandler)),
|
||||
notificationPromptListChanged: newMethodInfo(clientMethod((*Client).callPromptChangedHandler)),
|
||||
notificationResourceListChanged: newMethodInfo(clientMethod((*Client).callResourceChangedHandler)),
|
||||
}
|
||||
|
||||
var _ session[ClientSession] = (*ClientSession)(nil)
|
||||
|
|
@ -202,6 +226,9 @@ func (cs *ClientSession) methodHandler() MethodHandler[ClientSession] {
|
|||
return cs.client.methodHandler_
|
||||
}
|
||||
|
||||
// getConn implements [session.getConn].
|
||||
func (cs *ClientSession) getConn() *jsonrpc2.Connection { return cs.conn }
|
||||
|
||||
func (c *ClientSession) ping(ct context.Context, params *PingParams) (struct{}, error) {
|
||||
return struct{}{}, nil
|
||||
}
|
||||
|
|
@ -263,6 +290,18 @@ func (c *ClientSession) ReadResource(ctx context.Context, params *ReadResourcePa
|
|||
return standardCall[ReadResourceResult](ctx, c.conn, methodReadResource, params)
|
||||
}
|
||||
|
||||
func (c *Client) callToolChangedHandler(ctx context.Context, s *ClientSession, params *ToolListChangedParams) (any, error) {
|
||||
return callNotificationHandler(ctx, c.opts.ToolListChangedHandler, s, params)
|
||||
}
|
||||
|
||||
func (c *Client) callPromptChangedHandler(ctx context.Context, s *ClientSession, params *PromptListChangedParams) (any, error) {
|
||||
return callNotificationHandler(ctx, c.opts.PromptListChangedHandler, s, params)
|
||||
}
|
||||
|
||||
func (c *Client) callResourceChangedHandler(ctx context.Context, s *ClientSession, params *ResourceListChangedParams) (any, error) {
|
||||
return callNotificationHandler(ctx, c.opts.ResourceListChangedHandler, s, params)
|
||||
}
|
||||
|
||||
func standardCall[TRes, TParams any](ctx context.Context, conn *jsonrpc2.Connection, method string, params TParams) (*TRes, error) {
|
||||
var result TRes
|
||||
if err := call(ctx, conn, method, params, &result); err != nil {
|
||||
|
|
|
|||
|
|
@ -120,8 +120,12 @@ var declarations = config{
|
|||
"Prompt": {},
|
||||
"PromptMessage": {},
|
||||
"PromptArgument": {},
|
||||
"ProgressToken": {Name: "-", Substitute: "any"}, // null|number|string
|
||||
"RequestId": {Name: "-", Substitute: "any"}, // null|number|string
|
||||
"PromptListChangedNotification": {
|
||||
Name: "-",
|
||||
Fields: config{"Params": {Name: "PromptListChangedParams"}},
|
||||
},
|
||||
"ProgressToken": {Name: "-", Substitute: "any"}, // null|number|string
|
||||
"RequestId": {Name: "-", Substitute: "any"}, // null|number|string
|
||||
"ReadResourceRequest": {
|
||||
Name: "-",
|
||||
Fields: config{"Params": {Name: "ReadResourceParams"}},
|
||||
|
|
@ -130,8 +134,16 @@ var declarations = config{
|
|||
Fields: config{"Contents": {Substitute: "*ResourceContents"}},
|
||||
},
|
||||
"Resource": {},
|
||||
"Role": {},
|
||||
"Root": {},
|
||||
"ResourceListChangedNotification": {
|
||||
Name: "-",
|
||||
Fields: config{"Params": {Name: "ResourceListChangedParams"}},
|
||||
},
|
||||
"Role": {},
|
||||
"Root": {},
|
||||
"RootsListChangedNotification": {
|
||||
Name: "-",
|
||||
Fields: config{"Params": {Name: "RootsListChangedParams"}},
|
||||
},
|
||||
|
||||
"SamplingCapabilities": {Substitute: "struct{}"},
|
||||
"SamplingMessage": {},
|
||||
|
|
@ -147,6 +159,10 @@ var declarations = config{
|
|||
Fields: config{"InputSchema": {Substitute: "*jsonschema.Schema"}},
|
||||
},
|
||||
"ToolAnnotations": {},
|
||||
"ToolListChangedNotification": {
|
||||
Name: "-",
|
||||
Fields: config{"Params": {Name: "ToolListChangedParams"}},
|
||||
},
|
||||
}
|
||||
|
||||
func main() {
|
||||
|
|
|
|||
|
|
@ -39,7 +39,26 @@ func TestEndToEnd(t *testing.T) {
|
|||
ctx := context.Background()
|
||||
ct, st := NewInMemoryTransports()
|
||||
|
||||
s := NewServer("testServer", "v1.0.0", nil)
|
||||
// Channels to check if notification callbacks happened.
|
||||
notificationChans := map[string]chan int{}
|
||||
for _, name := range []string{"initialized", "roots", "tools", "prompts", "resources"} {
|
||||
notificationChans[name] = make(chan int, 1)
|
||||
}
|
||||
waitForNotification := func(t *testing.T, name string) {
|
||||
t.Helper()
|
||||
select {
|
||||
case <-notificationChans[name]:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatalf("%s handler never called", name)
|
||||
}
|
||||
}
|
||||
|
||||
sopts := &ServerOptions{
|
||||
InitializedHandler: func(context.Context, *ServerSession, *InitializedParams) { notificationChans["initialized"] <- 0 },
|
||||
RootsListChangedHandler: func(context.Context, *ServerSession, *RootsListChangedParams) { notificationChans["roots"] <- 0 },
|
||||
}
|
||||
|
||||
s := NewServer("testServer", "v1.0.0", sopts)
|
||||
|
||||
// The 'greet' tool says hi.
|
||||
s.AddTools(NewTool("greet", "say hi", sayHi))
|
||||
|
|
@ -89,6 +108,9 @@ func TestEndToEnd(t *testing.T) {
|
|||
CreateMessageHandler: func(context.Context, *ClientSession, *CreateMessageParams) (*CreateMessageResult, error) {
|
||||
return &CreateMessageResult{Model: "aModel"}, nil
|
||||
},
|
||||
ToolListChangedHandler: func(context.Context, *ClientSession, *ToolListChangedParams) { notificationChans["tools"] <- 0 },
|
||||
PromptListChangedHandler: func(context.Context, *ClientSession, *PromptListChangedParams) { notificationChans["prompts"] <- 0 },
|
||||
ResourceListChangedHandler: func(context.Context, *ClientSession, *ResourceListChangedParams) { notificationChans["resources"] <- 0 },
|
||||
}
|
||||
c := NewClient("testClient", "v1.0.0", opts)
|
||||
rootAbs, err := filepath.Abs(filepath.FromSlash("testdata/files"))
|
||||
|
|
@ -103,6 +125,7 @@ func TestEndToEnd(t *testing.T) {
|
|||
t.Fatal(err)
|
||||
}
|
||||
|
||||
waitForNotification(t, "initialized")
|
||||
if err := cs.Ping(ctx, nil); err != nil {
|
||||
t.Fatalf("ping failed: %v", err)
|
||||
}
|
||||
|
|
@ -141,6 +164,11 @@ func TestEndToEnd(t *testing.T) {
|
|||
if _, err := cs.GetPrompt(ctx, &GetPromptParams{Name: "fail"}); err == nil || !strings.Contains(err.Error(), failure.Error()) {
|
||||
t.Errorf("fail returned unexpected error: got %v, want containing %v", err, failure)
|
||||
}
|
||||
|
||||
s.AddPrompts(&ServerPrompt{Prompt: &Prompt{Name: "T"}})
|
||||
waitForNotification(t, "prompts")
|
||||
s.RemovePrompts("T")
|
||||
waitForNotification(t, "prompts")
|
||||
})
|
||||
|
||||
t.Run("tools", func(t *testing.T) {
|
||||
|
|
@ -198,6 +226,11 @@ func TestEndToEnd(t *testing.T) {
|
|||
if diff := cmp.Diff(wantFail, gotFail); diff != "" {
|
||||
t.Errorf("tools/call 'fail' mismatch (-want +got):\n%s", diff)
|
||||
}
|
||||
|
||||
s.AddTools(&ServerTool{Tool: &Tool{Name: "T"}})
|
||||
waitForNotification(t, "tools")
|
||||
s.RemoveTools("T")
|
||||
waitForNotification(t, "tools")
|
||||
})
|
||||
|
||||
t.Run("resources", func(t *testing.T) {
|
||||
|
|
@ -254,6 +287,11 @@ func TestEndToEnd(t *testing.T) {
|
|||
}
|
||||
}
|
||||
}
|
||||
|
||||
s.AddResources(&ServerResource{Resource: &Resource{URI: "http://U"}})
|
||||
waitForNotification(t, "resources")
|
||||
s.RemoveResources("http://U")
|
||||
waitForNotification(t, "resources")
|
||||
})
|
||||
t.Run("roots", func(t *testing.T) {
|
||||
rootRes, err := ss.ListRoots(ctx, &ListRootsParams{})
|
||||
|
|
@ -265,6 +303,11 @@ func TestEndToEnd(t *testing.T) {
|
|||
if diff := cmp.Diff(wantRoots, gotRoots); diff != "" {
|
||||
t.Errorf("roots/list mismatch (-want +got):\n%s", diff)
|
||||
}
|
||||
|
||||
c.AddRoots(&Root{URI: "U"})
|
||||
waitForNotification(t, "roots")
|
||||
c.RemoveRoots("U")
|
||||
waitForNotification(t, "roots")
|
||||
})
|
||||
t.Run("sampling", func(t *testing.T) {
|
||||
// TODO: test that a client that doesn't have the handler returns CodeUnsupportedMethod.
|
||||
|
|
@ -462,6 +505,10 @@ func TestMiddleware(t *testing.T) {
|
|||
2 >initialize
|
||||
2 <initialize
|
||||
1 <initialize
|
||||
1 >notifications/initialized
|
||||
2 >notifications/initialized
|
||||
2 <notifications/initialized
|
||||
1 <notifications/initialized
|
||||
1 >tools/list
|
||||
2 >tools/list
|
||||
2 <tools/list
|
||||
|
|
|
|||
|
|
@ -323,6 +323,13 @@ type PromptArgument struct {
|
|||
Required bool `json:"required,omitempty"`
|
||||
}
|
||||
|
||||
type PromptListChangedParams struct {
|
||||
// This parameter name is reserved by MCP to allow clients and servers to attach
|
||||
// additional metadata to their notifications.
|
||||
Meta struct {
|
||||
} `json:"_meta,omitempty"`
|
||||
}
|
||||
|
||||
// Describes a message returned as part of a prompt.
|
||||
//
|
||||
// This is similar to `SamplingMessage`, but also supports the embedding of
|
||||
|
|
@ -371,6 +378,13 @@ type Resource struct {
|
|||
URI string `json:"uri"`
|
||||
}
|
||||
|
||||
type ResourceListChangedParams struct {
|
||||
// This parameter name is reserved by MCP to allow clients and servers to attach
|
||||
// additional metadata to their notifications.
|
||||
Meta struct {
|
||||
} `json:"_meta,omitempty"`
|
||||
}
|
||||
|
||||
// The sender or recipient of messages and data in a conversation.
|
||||
type Role string
|
||||
|
||||
|
|
@ -386,6 +400,13 @@ type Root struct {
|
|||
URI string `json:"uri"`
|
||||
}
|
||||
|
||||
type RootsListChangedParams struct {
|
||||
// This parameter name is reserved by MCP to allow clients and servers to attach
|
||||
// additional metadata to their notifications.
|
||||
Meta struct {
|
||||
} `json:"_meta,omitempty"`
|
||||
}
|
||||
|
||||
// Present if the client supports sampling from an LLM.
|
||||
type SamplingCapabilities struct {
|
||||
}
|
||||
|
|
@ -448,6 +469,13 @@ type ToolAnnotations struct {
|
|||
Title string `json:"title,omitempty"`
|
||||
}
|
||||
|
||||
type ToolListChangedParams struct {
|
||||
// This parameter name is reserved by MCP to allow clients and servers to attach
|
||||
// additional metadata to their notifications.
|
||||
Meta struct {
|
||||
} `json:"_meta,omitempty"`
|
||||
}
|
||||
|
||||
// Describes the name and version of an MCP implementation.
|
||||
type implementation struct {
|
||||
Name string `json:"name"`
|
||||
|
|
@ -496,28 +524,28 @@ type toolCapabilities struct {
|
|||
}
|
||||
|
||||
const (
|
||||
notificationCancelled = "notifications/cancelled"
|
||||
methodInitialize = "initialize"
|
||||
notificationProgress = "notifications/progress"
|
||||
methodSetLevel = "logging/setLevel"
|
||||
methodCreateMessage = "sampling/createMessage"
|
||||
notificationResourceListChanged = "notifications/resources/list_changed"
|
||||
notificationInitialized = "notifications/initialized"
|
||||
methodUnsubscribe = "resources/unsubscribe"
|
||||
notificationLoggingMessage = "notifications/message"
|
||||
methodSubscribe = "resources/subscribe"
|
||||
methodComplete = "completion/complete"
|
||||
methodCallTool = "tools/call"
|
||||
notificationPromptListChanged = "notifications/prompts/list_changed"
|
||||
methodReadResource = "resources/read"
|
||||
methodListResourceTemplates = "resources/templates/list"
|
||||
methodListRoots = "roots/list"
|
||||
notificationToolListChanged = "notifications/tools/list_changed"
|
||||
methodGetPrompt = "prompts/get"
|
||||
notificationResourceListChanged = "notifications/resources/list_changed"
|
||||
methodListPrompts = "prompts/list"
|
||||
methodPing = "ping"
|
||||
notificationRootsListChanged = "notifications/roots/list_changed"
|
||||
methodListTools = "tools/list"
|
||||
methodListResources = "resources/list"
|
||||
notificationPromptListChanged = "notifications/prompts/list_changed"
|
||||
notificationResourceUpdated = "notifications/resources/updated"
|
||||
notificationCancelled = "notifications/cancelled"
|
||||
methodSetLevel = "logging/setLevel"
|
||||
methodInitialize = "initialize"
|
||||
methodListRoots = "roots/list"
|
||||
notificationProgress = "notifications/progress"
|
||||
methodGetPrompt = "prompts/get"
|
||||
methodListResourceTemplates = "resources/templates/list"
|
||||
methodPing = "ping"
|
||||
methodComplete = "completion/complete"
|
||||
methodListResources = "resources/list"
|
||||
notificationInitialized = "notifications/initialized"
|
||||
methodCallTool = "tools/call"
|
||||
notificationLoggingMessage = "notifications/message"
|
||||
methodReadResource = "resources/read"
|
||||
methodListTools = "tools/list"
|
||||
notificationRootsListChanged = "notifications/roots/list_changed"
|
||||
methodSubscribe = "resources/subscribe"
|
||||
methodUnsubscribe = "resources/unsubscribe"
|
||||
)
|
||||
|
|
|
|||
130
mcp/server.go
130
mcp/server.go
|
|
@ -36,11 +36,16 @@ type Server struct {
|
|||
|
||||
// ServerOptions is used to configure behavior of the server.
|
||||
type ServerOptions struct {
|
||||
// Optional instructions for connected clients.
|
||||
Instructions string
|
||||
// If non-nil, called when "notifications/intialized" is received.
|
||||
InitializedHandler func(context.Context, *ServerSession, *InitializedParams)
|
||||
// If non-nil, called when "notifications/roots/list_changed" is received.
|
||||
RootsListChangedHandler func(context.Context, *ServerSession, *RootsListChangedParams)
|
||||
}
|
||||
|
||||
// NewServer creates a new MCP server. The resulting server has no features:
|
||||
// add features using [Server.AddTools]. (TODO: support more features).
|
||||
// add features using [Server.AddTools], [Server.AddPrompts] and [Server.AddResources].
|
||||
//
|
||||
// The server can be connected to one or more MCP clients using [Server.Start]
|
||||
// or [Server.Run].
|
||||
|
|
@ -64,43 +69,43 @@ func NewServer(name, version string, opts *ServerOptions) *Server {
|
|||
// AddPrompts adds the given prompts to the server,
|
||||
// replacing any with the same names.
|
||||
func (s *Server) AddPrompts(prompts ...*ServerPrompt) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.prompts.add(prompts...)
|
||||
// Assume there was a change, since add replaces existing prompts.
|
||||
// (It's possible a prompt was replaced with an identical one, but not worth checking.)
|
||||
// TODO(rfindley): notify connected clients
|
||||
// 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.)
|
||||
s.changeAndNotify(
|
||||
notificationPromptListChanged,
|
||||
&PromptListChangedParams{},
|
||||
func() bool { s.prompts.add(prompts...); return true })
|
||||
}
|
||||
|
||||
// RemovePrompts removes the prompts with the given names.
|
||||
// It is not an error to remove a nonexistent prompt.
|
||||
func (s *Server) RemovePrompts(names ...string) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.prompts.remove(names...) {
|
||||
// TODO: notify
|
||||
}
|
||||
s.changeAndNotify(notificationPromptListChanged, &PromptListChangedParams{},
|
||||
func() bool { return s.prompts.remove(names...) })
|
||||
}
|
||||
|
||||
// AddTools adds the given tools to the server,
|
||||
// replacing any with the same names.
|
||||
func (s *Server) AddTools(tools ...*ServerTool) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.tools.add(tools...)
|
||||
// Only notify if something could change.
|
||||
if len(tools) == 0 {
|
||||
return
|
||||
}
|
||||
// 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(rfindley): notify connected clients
|
||||
s.changeAndNotify(notificationToolListChanged, &ToolListChangedParams{},
|
||||
func() bool { s.tools.add(tools...); return true })
|
||||
}
|
||||
|
||||
// RemoveTools removes the tools with the given names.
|
||||
// It is not an error to remove a nonexistent tool.
|
||||
func (s *Server) RemoveTools(names ...string) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.tools.remove(names...) {
|
||||
// TODO: notify
|
||||
}
|
||||
s.changeAndNotify(notificationToolListChanged, &ToolListChangedParams{},
|
||||
func() bool { return s.tools.remove(names...) })
|
||||
}
|
||||
|
||||
// AddResource adds the given resource to the server and associates it with
|
||||
|
|
@ -108,27 +113,45 @@ func (s *Server) RemoveTools(names ...string) {
|
|||
// If a resource with the same URI already exists, this one replaces it.
|
||||
// AddResource panics if a resource URI is invalid or not absolute (has an empty scheme).
|
||||
func (s *Server) AddResources(resources ...*ServerResource) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
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)
|
||||
// Only notify if something could change.
|
||||
if len(resources) == 0 {
|
||||
return
|
||||
}
|
||||
// TODO: notify
|
||||
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)
|
||||
}
|
||||
return true
|
||||
})
|
||||
}
|
||||
|
||||
// RemoveResources removes the resources with the given URIs.
|
||||
// It is not an error to remove a nonexistent resource.
|
||||
func (s *Server) RemoveResources(uris ...string) {
|
||||
s.changeAndNotify(notificationResourceListChanged, &ResourceListChangedParams{},
|
||||
func() bool { return s.resources.remove(uris...) })
|
||||
}
|
||||
|
||||
// changeAndNotify is called when a feature is added or removed.
|
||||
// It calls change, which should do the work and report whether a change actually occurred.
|
||||
// If there was a change, it notifies a snapshot of the sessions.
|
||||
func (s *Server) changeAndNotify(notification string, params any, change func() bool) {
|
||||
var sessions []*ServerSession
|
||||
// Lock for the change, but not for the notification.
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.resources.remove(uris...)
|
||||
if change() {
|
||||
sessions = slices.Clone(s.sessions)
|
||||
}
|
||||
s.mu.Unlock()
|
||||
notifySessions(sessions, notification, params)
|
||||
}
|
||||
|
||||
// Sessions returns an iterator that yields the current set of server sessions.
|
||||
|
|
@ -303,6 +326,14 @@ func (s *Server) Connect(ctx context.Context, t Transport) (*ServerSession, erro
|
|||
return connect(ctx, t, s)
|
||||
}
|
||||
|
||||
func (s *Server) callInitializedHandler(ctx context.Context, ss *ServerSession, params *InitializedParams) (any, error) {
|
||||
return callNotificationHandler(ctx, s.opts.InitializedHandler, ss, params)
|
||||
}
|
||||
|
||||
func (s *Server) callRootsListChangedHandler(ctx context.Context, ss *ServerSession, params *RootsListChangedParams) (any, error) {
|
||||
return callNotificationHandler(ctx, s.opts.RootsListChangedHandler, ss, params)
|
||||
}
|
||||
|
||||
// A ServerSession is a logical connection from a single MCP client. Its
|
||||
// methods can be used to send requests or notifications to the client. Create
|
||||
// a session by calling [Server.Connect].
|
||||
|
|
@ -347,15 +378,16 @@ func (s *Server) AddMiddleware(middleware ...Middleware[ServerSession]) {
|
|||
|
||||
// serverMethodInfos maps from the RPC method name to serverMethodInfos.
|
||||
var serverMethodInfos = map[string]methodInfo[ServerSession]{
|
||||
methodInitialize: newMethodInfo(sessionMethod((*ServerSession).initialize)),
|
||||
methodPing: newMethodInfo(sessionMethod((*ServerSession).ping)),
|
||||
methodListPrompts: newMethodInfo(serverMethod((*Server).listPrompts)),
|
||||
methodGetPrompt: newMethodInfo(serverMethod((*Server).getPrompt)),
|
||||
methodListTools: newMethodInfo(serverMethod((*Server).listTools)),
|
||||
methodCallTool: newMethodInfo(serverMethod((*Server).callTool)),
|
||||
methodListResources: newMethodInfo(serverMethod((*Server).listResources)),
|
||||
methodReadResource: newMethodInfo(serverMethod((*Server).readResource)),
|
||||
// TODO: notifications
|
||||
methodInitialize: newMethodInfo(sessionMethod((*ServerSession).initialize)),
|
||||
methodPing: newMethodInfo(sessionMethod((*ServerSession).ping)),
|
||||
methodListPrompts: newMethodInfo(serverMethod((*Server).listPrompts)),
|
||||
methodGetPrompt: newMethodInfo(serverMethod((*Server).getPrompt)),
|
||||
methodListTools: newMethodInfo(serverMethod((*Server).listTools)),
|
||||
methodCallTool: newMethodInfo(serverMethod((*Server).callTool)),
|
||||
methodListResources: newMethodInfo(serverMethod((*Server).listResources)),
|
||||
methodReadResource: newMethodInfo(serverMethod((*Server).readResource)),
|
||||
notificationInitialized: newMethodInfo(serverMethod((*Server).callInitializedHandler)),
|
||||
notificationRootsListChanged: newMethodInfo(serverMethod((*Server).callRootsListChangedHandler)),
|
||||
}
|
||||
|
||||
// *ServerSession implements the session interface.
|
||||
|
|
@ -372,6 +404,9 @@ func (ss *ServerSession) methodHandler() MethodHandler[ServerSession] {
|
|||
return ss.server.methodHandler_
|
||||
}
|
||||
|
||||
// getConn implements [session.getConn].
|
||||
func (ss *ServerSession) getConn() *jsonrpc2.Connection { return ss.conn }
|
||||
|
||||
// handle invokes the method described by the given JSON RPC request.
|
||||
func (ss *ServerSession) handle(ctx context.Context, req *jsonrpc2.Request) (any, error) {
|
||||
ss.mu.Lock()
|
||||
|
|
@ -414,10 +449,13 @@ func (ss *ServerSession) initialize(ctx context.Context, params *InitializeParam
|
|||
ProtocolVersion: "2024-11-05",
|
||||
Capabilities: &serverCapabilities{
|
||||
Prompts: &promptCapabilities{
|
||||
ListChanged: false, // not yet supported
|
||||
ListChanged: true,
|
||||
},
|
||||
Tools: &toolCapabilities{
|
||||
ListChanged: false, // not yet supported
|
||||
ListChanged: true,
|
||||
},
|
||||
Resources: &resourceCapabilities{
|
||||
ListChanged: true,
|
||||
},
|
||||
},
|
||||
Instructions: ss.server.opts.Instructions,
|
||||
|
|
|
|||
|
|
@ -13,7 +13,9 @@ import (
|
|||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log"
|
||||
"slices"
|
||||
"time"
|
||||
|
||||
"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
|
||||
)
|
||||
|
|
@ -40,6 +42,7 @@ func addMiddleware[S ClientSession | ServerSession](handlerp *MethodHandler[S],
|
|||
type session[S ClientSession | ServerSession] interface {
|
||||
methodHandler() MethodHandler[S]
|
||||
methodInfos() map[string]methodInfo[S]
|
||||
getConn() *jsonrpc2.Connection
|
||||
}
|
||||
|
||||
// toSession[S] converts its argument to a session[S].
|
||||
|
|
@ -147,3 +150,27 @@ const (
|
|||
// The error code if the method exists and was called properly, but the peer does not support it.
|
||||
CodeUnsupportedMethod = -31001
|
||||
)
|
||||
|
||||
func callNotificationHandler[S ClientSession | ServerSession, P any](ctx context.Context, h func(context.Context, *S, *P), sess *S, params *P) (any, error) {
|
||||
if h != nil {
|
||||
h(ctx, sess, params)
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// notifySessions calls Notify on all the sessions.
|
||||
// Should be called on a copy of the peer sessions.
|
||||
func notifySessions[S ClientSession | ServerSession](sessions []*S, method string, params any) {
|
||||
if sessions == nil {
|
||||
return
|
||||
}
|
||||
// TODO: make this timeout configurable, or call Notify asynchronously.
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
for _, ss := range sessions {
|
||||
if err := toSession(ss).getConn().Notify(ctx, method, params); err != nil {
|
||||
// TODO(jba): surface this error better
|
||||
log.Printf("calling %s: %v", method, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue