modelcontextprotocol/go-sdk
Watch
1
0
Fork
You've already forked go-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

internal/mcp: add notifications

Browse source 
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 2025-05-17 12:59:45 -04:00
commit ce5b6cd926
1 changed files with 273 additions and 78 deletions
Download patch file Download diff file Expand all files Collapse all files

59
mcp/client.go
View file

@ -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 {