mcp: add resource subscriptions (#138)

This CL adds the ability for clients to subscribe and receive updates
for resources as described in
https://modelcontextprotocol.io/specification/2025-06-18/server/resources#subscriptions

Fixes: #23
This commit is contained in:
Sam Thanawalla 2025-07-23 15:23:50 -04:00 • committed by GitHub
commit bcbb31f072
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 214 additions and 6 deletions

View file

@ -748,13 +748,26 @@ Server sessions also support the spec methods `ListResources` and `ListResourceT
#### Subscriptions
ClientSessions can manage change notifications on particular resources:
##### Client-Side Usage
Use the Subscribe and Unsubscribe methods on a ClientSession to start or stop receiving updates for a specific resource.
```go
func (*ClientSession) Subscribe(context.Context, *SubscribeParams) error
func (*ClientSession) Unsubscribe(context.Context, *UnsubscribeParams) error
```
To process incoming update notifications, you must provide a ResourceUpdatedHandler in your ClientOptions. The SDK calls this function automatically whenever the server sends a notification for a resource you're subscribed to.
```go
type ClientOptions struct {
...
ResourceUpdatedHandler func(context.Context, *ClientSession, *ResourceUpdatedNotificationParams)
}
```
##### Server-Side Implementation
The server does not implement resource subscriptions. It passes along subscription requests to the user, and supplies a method to notify clients of changes. It tracks which sessions have subscribed to which resources so the user doesn't have to.
If a server author wants to support resource subscriptions, they must provide handlers to be called when clients subscribe and unsubscribe. It is an error to provide only one of these handlers.
@ -772,7 +785,7 @@ type ServerOptions struct {
User code should call `ResourceUpdated` when a subscribed resource changes.
```go
func (*Server) ResourceUpdated(context.Context, *ResourceUpdatedNotification) error
func (*Server) ResourceUpdated(context.Context, *ResourceUpdatedNotificationParams) error
```
The server routes these notifications to the server sessions that subscribed to the resource.

View file

@ -60,6 +60,7 @@ type ClientOptions struct {
ToolListChangedHandler func(context.Context, *ClientSession, *ToolListChangedParams)
PromptListChangedHandler func(context.Context, *ClientSession, *PromptListChangedParams)
ResourceListChangedHandler func(context.Context, *ClientSession, *ResourceListChangedParams)
ResourceUpdatedHandler func(context.Context, *ClientSession, *ResourceUpdatedNotificationParams)
LoggingMessageHandler func(context.Context, *ClientSession, *LoggingMessageParams)
ProgressNotificationHandler func(context.Context, *ClientSession, *ProgressNotificationParams)
// If non-zero, defines an interval for regular "ping" requests.
@ -293,6 +294,7 @@ var clientMethodInfos = map[string]methodInfo{
notificationToolListChanged: newMethodInfo(clientMethod((*Client).callToolChangedHandler)),
notificationPromptListChanged: newMethodInfo(clientMethod((*Client).callPromptChangedHandler)),
notificationResourceListChanged: newMethodInfo(clientMethod((*Client).callResourceChangedHandler)),
notificationResourceUpdated: newMethodInfo(clientMethod((*Client).callResourceUpdatedHandler)),
notificationLoggingMessage: newMethodInfo(clientMethod((*Client).callLoggingHandler)),
notificationProgress: newMethodInfo(sessionMethod((*ClientSession).callProgressNotificationHandler)),
}
@ -386,6 +388,20 @@ func (cs *ClientSession) Complete(ctx context.Context, params *CompleteParams) (
return handleSend[*CompleteResult](ctx, cs, methodComplete, orZero[Params](params))
}
// Subscribe sends a "resources/subscribe" request to the server, asking for
// notifications when the specified resource changes.
func (cs *ClientSession) Subscribe(ctx context.Context, params *SubscribeParams) error {
_, err := handleSend[*emptyResult](ctx, cs, methodSubscribe, orZero[Params](params))
return err
}
// Unsubscribe sends a "resources/unsubscribe" request to the server, cancelling
// a previous subscription.
func (cs *ClientSession) Unsubscribe(ctx context.Context, params *UnsubscribeParams) error {
_, err := handleSend[*emptyResult](ctx, cs, methodUnsubscribe, orZero[Params](params))
return err
}
func (c *Client) callToolChangedHandler(ctx context.Context, s *ClientSession, params *ToolListChangedParams) (Result, error) {
return callNotificationHandler(ctx, c.opts.ToolListChangedHandler, s, params)
}
@ -398,6 +414,10 @@ func (c *Client) callResourceChangedHandler(ctx context.Context, s *ClientSessio
return callNotificationHandler(ctx, c.opts.ResourceListChangedHandler, s, params)
}
func (c *Client) callResourceUpdatedHandler(ctx context.Context, s *ClientSession, params *ResourceUpdatedNotificationParams) (Result, error) {
return callNotificationHandler(ctx, c.opts.ResourceUpdatedHandler, s, params)
}
func (c *Client) callLoggingHandler(ctx context.Context, cs *ClientSession, params *LoggingMessageParams) (Result, error) {
if h := c.opts.LoggingMessageHandler; h != nil {
h(ctx, cs, params)

View file

@ -60,7 +60,7 @@ func TestEndToEnd(t *testing.T) {
// Channels to check if notification callbacks happened.
notificationChans := map[string]chan int{}
for _, name := range []string{"initialized", "roots", "tools", "prompts", "resources", "progress_server", "progress_client"} {
for _, name := range []string{"initialized", "roots", "tools", "prompts", "resources", "progress_server", "progress_client", "resource_updated", "subscribe", "unsubscribe"} {
notificationChans[name] = make(chan int, 1)
}
waitForNotification := func(t *testing.T, name string) {
@ -78,6 +78,14 @@ func TestEndToEnd(t *testing.T) {
ProgressNotificationHandler: func(context.Context, *ServerSession, *ProgressNotificationParams) {
notificationChans["progress_server"] <- 0
},
SubscribeHandler: func(context.Context, *SubscribeParams) error {
notificationChans["subscribe"] <- 0
return nil
},
UnsubscribeHandler: func(context.Context, *UnsubscribeParams) error {
notificationChans["unsubscribe"] <- 0
return nil
},
}
s := NewServer(testImpl, sopts)
AddTool(s, &Tool{
@ -128,6 +136,9 @@ func TestEndToEnd(t *testing.T) {
ProgressNotificationHandler: func(context.Context, *ClientSession, *ProgressNotificationParams) {
notificationChans["progress_client"] <- 0
},
ResourceUpdatedHandler: func(context.Context, *ClientSession, *ResourceUpdatedNotificationParams) {
notificationChans["resource_updated"] <- 0
},
}
c := NewClient(testImpl, opts)
rootAbs, err := filepath.Abs(filepath.FromSlash("testdata/files"))
@ -421,6 +432,37 @@ func TestEndToEnd(t *testing.T) {
waitForNotification(t, "progress_server")
})
t.Run("resource_subscriptions", func(t *testing.T) {
err := cs.Subscribe(ctx, &SubscribeParams{
URI: "test",
})
if err != nil {
t.Fatal(err)
}
waitForNotification(t, "subscribe")
s.ResourceUpdated(ctx, &ResourceUpdatedNotificationParams{
URI: "test",
})
waitForNotification(t, "resource_updated")
err = cs.Unsubscribe(ctx, &UnsubscribeParams{
URI: "test",
})
if err != nil {
t.Fatal(err)
}
waitForNotification(t, "unsubscribe")
// Verify the client does not receive the update after unsubscribing.
s.ResourceUpdated(ctx, &ResourceUpdatedNotificationParams{
URI: "test",
})
select {
case <-notificationChans["resource_updated"]:
t.Fatalf("resource updated after unsubscription")
case <-time.After(time.Second):
}
})
// Disconnect.
cs.Close()
clientWG.Wait()

View file

@ -859,6 +859,38 @@ type ToolListChangedParams struct {
func (x *ToolListChangedParams) GetProgressToken() any { return getProgressToken(x) }
func (x *ToolListChangedParams) SetProgressToken(t any) { setProgressToken(x, t) }
// Sent from the client to request resources/updated notifications from the
// server whenever a particular resource changes.
type SubscribeParams struct {
// This property is reserved by the protocol to allow clients and servers to
// attach additional metadata to their responses.
Meta `json:"_meta,omitempty"`
// The URI of the resource to subscribe to.
URI string `json:"uri"`
}
// Sent from the client to request cancellation of resources/updated
// notifications from the server. This should follow a previous
// resources/subscribe request.
type UnsubscribeParams struct {
// This property is reserved by the protocol to allow clients and servers to
// attach additional metadata to their responses.
Meta `json:"_meta,omitempty"`
// The URI of the resource to unsubscribe from.
URI string `json:"uri"`
}
// A notification from the server to the client, informing it that a resource
// has changed and may need to be read again. This should only be sent if the
// client previously sent a resources/subscribe request.
type ResourceUpdatedNotificationParams struct {
// This property is reserved by the protocol to allow clients and servers to
// attach additional metadata to their responses.
Meta `json:"_meta,omitempty"`
// The URI of the resource that has been updated. This might be a sub-resource of the one that the client actually subscribed to.
URI string `json:"uri"`
}
// TODO(jba): add CompleteRequest and related types.
// TODO(jba): add ElicitRequest and related types.

View file

@ -13,6 +13,7 @@ import (
"fmt"
"iter"
"log"
"maps"
"net/url"
"path/filepath"
"slices"
@ -43,6 +44,7 @@ type Server struct {
sessions []*ServerSession
sendingMethodHandler_ MethodHandler[*ServerSession]
receivingMethodHandler_ MethodHandler[*ServerSession]
resourceSubscriptions map[string]map[*ServerSession]bool // uri -> session -> bool
}
// ServerOptions is used to configure behavior of the server.
@ -64,6 +66,10 @@ type ServerOptions struct {
// If the peer fails to respond to pings originating from the keepalive check,
// the session is automatically closed.
KeepAlive time.Duration
// Function called when a client session subscribes to a resource.
SubscribeHandler func(context.Context, *SubscribeParams) error
// Function called when a client session unsubscribes from a resource.
UnsubscribeHandler func(context.Context, *UnsubscribeParams) error
}
// NewServer creates a new MCP server. The resulting server has no features:
@ -89,7 +95,12 @@ func NewServer(impl *Implementation, opts *ServerOptions) *Server {
if opts.PageSize == 0 {
opts.PageSize = DefaultPageSize
}
if opts.SubscribeHandler != nil && opts.UnsubscribeHandler == nil {
panic("SubscribeHandler requires UnsubscribeHandler")
}
if opts.UnsubscribeHandler != nil && opts.SubscribeHandler == nil {
panic("UnsubscribeHandler requires SubscribeHandler")
}
return &Server{
impl: impl,
opts: *opts,
@ -99,6 +110,7 @@ func NewServer(impl *Implementation, opts *ServerOptions) *Server {
resourceTemplates: newFeatureSet(func(t *serverResourceTemplate) string { return t.resourceTemplate.URITemplate }),
sendingMethodHandler_: defaultSendingMethodHandler[*ServerSession],
receivingMethodHandler_: defaultReceivingMethodHandler[*ServerSession],
resourceSubscriptions: make(map[string]map[*ServerSession]bool),
}
}
@ -225,6 +237,9 @@ func (s *Server) capabilities() *serverCapabilities {
}
if s.resources.len() > 0 || s.resourceTemplates.len() > 0 {
caps.Resources = &resourceCapabilities{ListChanged: true}
if s.opts.SubscribeHandler != nil {
caps.Resources.Subscribe = true
}
}
return caps
}
@ -428,6 +443,57 @@ func fileResourceHandler(dir string) ResourceHandler {
}
}
// ResourceUpdated sends a notification to all clients that have subscribed to the
// resource specified in params. This method is the primary way for a
// server author to signal that a resource has changed.
func (s *Server) ResourceUpdated(ctx context.Context, params *ResourceUpdatedNotificationParams) error {
s.mu.Lock()
subscribedSessions := s.resourceSubscriptions[params.URI]
sessions := slices.Collect(maps.Keys(subscribedSessions))
s.mu.Unlock()
notifySessions(sessions, notificationResourceUpdated, params)
return nil
}
func (s *Server) subscribe(ctx context.Context, ss *ServerSession, params *SubscribeParams) (*emptyResult, error) {
if s.opts.SubscribeHandler == nil {
return nil, fmt.Errorf("%w: server does not support resource subscriptions", jsonrpc2.ErrMethodNotFound)
}
if err := s.opts.SubscribeHandler(ctx, params); err != nil {
return nil, err
}
s.mu.Lock()
defer s.mu.Unlock()
if s.resourceSubscriptions[params.URI] == nil {
s.resourceSubscriptions[params.URI] = make(map[*ServerSession]bool)
}
s.resourceSubscriptions[params.URI][ss] = true
return &emptyResult{}, nil
}
func (s *Server) unsubscribe(ctx context.Context, ss *ServerSession, params *UnsubscribeParams) (*emptyResult, error) {
if s.opts.UnsubscribeHandler == nil {
return nil, jsonrpc2.ErrMethodNotFound
}
if err := s.opts.UnsubscribeHandler(ctx, params); err != nil {
return nil, err
}
s.mu.Lock()
defer s.mu.Unlock()
if subscribedSessions, ok := s.resourceSubscriptions[params.URI]; ok {
delete(subscribedSessions, ss)
if len(subscribedSessions) == 0 {
delete(s.resourceSubscriptions, params.URI)
}
}
return &emptyResult{}, nil
}
// Run runs the server over the given transport, which must be persistent.
//
// Run blocks until the client terminates the connection or the provided
@ -475,6 +541,10 @@ func (s *Server) disconnect(cc *ServerSession) {
s.sessions = slices.DeleteFunc(s.sessions, func(cc2 *ServerSession) bool {
return cc2 == cc
})
for _, subscribedSessions := range s.resourceSubscriptions {
delete(subscribedSessions, cc)
}
}
// Connect connects the MCP server over the given transport and starts handling
@ -616,6 +686,8 @@ var serverMethodInfos = map[string]methodInfo{
methodListResourceTemplates: newMethodInfo(serverMethod((*Server).listResourceTemplates)),
methodReadResource: newMethodInfo(serverMethod((*Server).readResource)),
methodSetLevel: newMethodInfo(sessionMethod((*ServerSession).setLevel)),
methodSubscribe: newMethodInfo(serverMethod((*Server).subscribe)),
methodUnsubscribe: newMethodInfo(serverMethod((*Server).unsubscribe)),
notificationInitialized: newMethodInfo(serverMethod((*Server).callInitializedHandler)),
notificationRootsListChanged: newMethodInfo(serverMethod((*Server).callRootsListChangedHandler)),
notificationProgress: newMethodInfo(sessionMethod((*ServerSession).callProgressNotificationHandler)),

View file

@ -5,6 +5,7 @@
package mcp
import (
"context"
"log"
"slices"
"testing"
@ -232,6 +233,7 @@ func TestServerCapabilities(t *testing.T) {
testCases := []struct {
name string
configureServer func(s *Server)
serverOpts ServerOptions
wantCapabilities *serverCapabilities
}{
{
@ -275,6 +277,25 @@ func TestServerCapabilities(t *testing.T) {
Resources: &resourceCapabilities{ListChanged: true},
},
},
{
name: "With resource subscriptions",
configureServer: func(s *Server) {
s.AddResourceTemplate(&ResourceTemplate{URITemplate: "file:///rt"}, nil)
},
serverOpts: ServerOptions{
SubscribeHandler: func(ctx context.Context, sp *SubscribeParams) error {
return nil
},
UnsubscribeHandler: func(ctx context.Context, up *UnsubscribeParams) error {
return nil
},
},
wantCapabilities: &serverCapabilities{
Completions: &completionCapabilities{},
Logging: &loggingCapabilities{},
Resources: &resourceCapabilities{ListChanged: true, Subscribe: true},
},
},
{
name: "With tools",
configureServer: func(s *Server) {
@ -294,11 +315,19 @@ func TestServerCapabilities(t *testing.T) {
s.AddResourceTemplate(&ResourceTemplate{URITemplate: "file:///rt"}, nil)
s.AddTool(&Tool{Name: "t"}, nil)
},
serverOpts: ServerOptions{
SubscribeHandler: func(ctx context.Context, sp *SubscribeParams) error {
return nil
},
UnsubscribeHandler: func(ctx context.Context, up *UnsubscribeParams) error {
return nil
},
},
wantCapabilities: &serverCapabilities{
Completions: &completionCapabilities{},
Logging: &loggingCapabilities{},
Prompts: &promptCapabilities{ListChanged: true},
Resources: &resourceCapabilities{ListChanged: true},
Resources: &resourceCapabilities{ListChanged: true, Subscribe: true},
Tools: &toolCapabilities{ListChanged: true},
},
},
@ -306,7 +335,7 @@ func TestServerCapabilities(t *testing.T) {
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
server := NewServer(testImpl, nil)
server := NewServer(testImpl, &tc.serverOpts)
tc.configureServer(server)
gotCapabilities := server.capabilities()
if diff := cmp.Diff(tc.wantCapabilities, gotCapabilities); diff != "" {