Merge branch 'modelcontextprotocol:main' into feat-devcontainer-config

Author: cr2007

design/design.md

@@ -471,7 +471,7 @@ server.AddReceivingMiddleware(withLogging)
 
 #### Rate Limiting
 
-Rate limiting can be configured using middleware. Please see [examples/rate-limiting](<https://github.com/modelcontextprotocol/go-sdk/tree/main/examples/rate-limiting>] for an example on how to implement this.
+Rate limiting can be configured using middleware. Please see [examples/rate-limiting](<https://github.com/modelcontextprotocol/go-sdk/tree/main/examples/rate-limiting>) for an example on how to implement this.
 
 ### Errors
 
@@ -609,7 +609,7 @@ A tool handler accepts `CallToolParams` and returns a `CallToolResult`. However,
 ```go
 type CallToolParamsFor[In any] struct {
 	Meta      Meta   `json:"_meta,omitempty"`
-	Arguments In `json:"arguments,omitempty"`
+	Arguments In     `json:"arguments,omitempty"`
 	Name      string `json:"name"`
 }
 
@@ -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.

mcp/client.go

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

mcp/event.go

@@ -10,14 +10,22 @@ package mcp
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"errors"
 	"fmt"
 	"io"
 	"iter"
+	"maps"
 	"net/http"
+	"slices"
 	"strings"
+	"sync"
 )
 
+// If true, MemoryEventStore will do frequent validation to check invariants, slowing it down.
+// Remove when we're confident in the code.
+const validateMemoryEventStore = true
+
 // An Event is a server-sent event.
 // See https://developer.mozilla.org/en-US/docs/Web/API/Server-sent_events/Using_server-sent_events#fields.
 type Event struct {
@@ -136,3 +144,259 @@ func scanEvents(r io.Reader) iter.Seq2[Event, error] {
 		}
 	}
 }
+
+// An EventStore tracks data for SSE streams.
+// A single EventStore suffices for all sessions, since session IDs are
+// globally unique. So one EventStore can be created per process, for
+// all Servers in the process.
+// Such a store is able to bound resource usage for the entire process.
+//
+// All of an EventStore's methods must be safe for use by multiple goroutines.
+type EventStore interface {
+	// Append appends data for an outgoing event to given stream, which is part of the
+	// given session.
+	Append(_ context.Context, sessionID string, _ StreamID, data []byte) error
+
+	// After returns an iterator over the data for the given session and stream, beginning
+	// just after the given index.
+	// Once the iterator yields a non-nil error, it will stop.
+	// After's iterator must return an error immediately if any data after index was
+	// dropped; it must not return partial results.
+	After(_ context.Context, sessionID string, _ StreamID, index int) iter.Seq2[[]byte, error]
+
+	// SessionClosed informs the store that the given session is finished, along
+	// with all of its streams.
+	// A store cannot rely on this method being called for cleanup. It should institute
+	// additional mechanisms, such as timeouts, to reclaim storage.
+	//
+	SessionClosed(_ context.Context, sessionID string) error
+
+	// There is no StreamClosed method. A server doesn't know when a stream is finished, because
+	// the client can always send a GET with a Last-Event-ID referring to the stream.
+}
+
+// A dataList is a list of []byte.
+// The zero dataList is ready to use.
+type dataList struct {
+	size  int // total size of data bytes
+	first int // the stream index of the first element in data
+	data  [][]byte
+}
+
+func (dl *dataList) appendData(d []byte) {
+	// If we allowed empty data, we would consume memory without incrementing the size.
+	// We could of course account for that, but we keep it simple and assume there is no
+	// empty data.
+	if len(d) == 0 {
+		panic("empty data item")
+	}
+	dl.data = append(dl.data, d)
+	dl.size += len(d)
+}
+
+// removeFirst removes the first data item in dl, returning the size of the item.
+// It panics if dl is empty.
+func (dl *dataList) removeFirst() int {
+	if len(dl.data) == 0 {
+		panic("empty dataList")
+	}
+	r := len(dl.data[0])
+	dl.size -= r
+	dl.data[0] = nil // help GC
+	dl.data = dl.data[1:]
+	dl.first++
+	return r
+}
+
+// A MemoryEventStore is an [EventStore] backed by memory.
+type MemoryEventStore struct {
+	mu       sync.Mutex
+	maxBytes int                               // max total size of all data
+	nBytes   int                               // current total size of all data
+	store    map[string]map[StreamID]*dataList // session ID -> stream ID -> *dataList
+}
+
+// MemoryEventStoreOptions are options for a [MemoryEventStore].
+type MemoryEventStoreOptions struct{}
+
+// MaxBytes returns the maximum number of bytes that the store will retain before
+// purging data.
+func (s *MemoryEventStore) MaxBytes() int {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return s.maxBytes
+}
+
+// SetMaxBytes sets the maximum number of bytes the store will retain before purging
+// data. The argument must not be negative. If it is zero, a suitable default will be used.
+// SetMaxBytes can be called at any time. The size of the store will be adjusted
+// immediately.
+func (s *MemoryEventStore) SetMaxBytes(n int) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	switch {
+	case n < 0:
+		panic("negative argument")
+	case n == 0:
+		s.maxBytes = defaultMaxBytes
+	default:
+		s.maxBytes = n
+	}
+	s.purge()
+}
+
+const defaultMaxBytes = 10 << 20 // 10 MiB
+
+// NewMemoryEventStore creates a [MemoryEventStore] with the default value
+// for MaxBytes.
+func NewMemoryEventStore(opts *MemoryEventStoreOptions) *MemoryEventStore {
+	return &MemoryEventStore{
+		maxBytes: defaultMaxBytes,
+		store:    make(map[string]map[StreamID]*dataList),
+	}
+}
+
+// Append implements [EventStore.Append] by recording data in memory.
+func (s *MemoryEventStore) Append(_ context.Context, sessionID string, streamID StreamID, data []byte) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	streamMap, ok := s.store[sessionID]
+	if !ok {
+		streamMap = make(map[StreamID]*dataList)
+		s.store[sessionID] = streamMap
+	}
+	dl, ok := streamMap[streamID]
+	if !ok {
+		dl = &dataList{}
+		streamMap[streamID] = dl
+	}
+	// Purge before adding, so at least the current data item will be present.
+	// (That could result in nBytes > maxBytes, but we'll live with that.)
+	s.purge()
+	dl.appendData(data)
+	s.nBytes += len(data)
+	return nil
+}
+
+// ErrEventsPurged is the error that [EventStore.After] should return if the event just after the
+// index is no longer available.
+var ErrEventsPurged = errors.New("data purged")
+
+// After implements [EventStore.After].
+func (s *MemoryEventStore) After(_ context.Context, sessionID string, streamID StreamID, index int) iter.Seq2[[]byte, error] {
+	// Return the data items to yield.
+	// We must copy, because dataList.removeFirst nils out slice elements.
+	copyData := func() ([][]byte, error) {
+		s.mu.Lock()
+		defer s.mu.Unlock()
+		streamMap, ok := s.store[sessionID]
+		if !ok {
+			return nil, fmt.Errorf("MemoryEventStore.After: unknown session ID %q", sessionID)
+		}
+		dl, ok := streamMap[streamID]
+		if !ok {
+			return nil, fmt.Errorf("MemoryEventStore.After: unknown stream ID %v in session %q", streamID, sessionID)
+		}
+		start := index + 1
+		if dl.first > start {
+			return nil, fmt.Errorf("MemoryEventStore.After: index %d, stream ID %v, session %q: %w",
+				index, streamID, sessionID, ErrEventsPurged)
+		}
+		return slices.Clone(dl.data[start-dl.first:]), nil
+	}
+
+	return func(yield func([]byte, error) bool) {
+		ds, err := copyData()
+		if err != nil {
+			yield(nil, err)
+			return
+		}
+		for _, d := range ds {
+			if !yield(d, nil) {
+				return
+			}
+		}
+	}
+}
+
+// SessionClosed implements [EventStore.SessionClosed].
+func (s *MemoryEventStore) SessionClosed(_ context.Context, sessionID string) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	for _, dl := range s.store[sessionID] {
+		s.nBytes -= dl.size
+	}
+	delete(s.store, sessionID)
+	s.validate()
+	return nil
+}
+
+// purge removes data until no more than s.maxBytes bytes are in use.
+// It must be called with s.mu held.
+func (s *MemoryEventStore) purge() {
+	// Remove the first element of every dataList until below the max.
+	for s.nBytes > s.maxBytes {
+		changed := false
+		for _, sm := range s.store {
+			for _, dl := range sm {
+				if dl.size > 0 {
+					r := dl.removeFirst()
+					if r > 0 {
+						changed = true
+						s.nBytes -= r
+					}
+				}
+			}
+		}
+		if !changed {
+			panic("no progress during purge")
+		}
+	}
+	s.validate()
+}
+
+// validate checks that the store's data structures are valid.
+// It must be called with s.mu held.
+func (s *MemoryEventStore) validate() {
+	if !validateMemoryEventStore {
+		return
+	}
+	// Check that we're accounting for the size correctly.
+	n := 0
+	for _, sm := range s.store {
+		for _, dl := range sm {
+			for _, d := range dl.data {
+				n += len(d)
+			}
+		}
+	}
+	if n != s.nBytes {
+		panic("sizes don't add up")
+	}
+}
+
+// debugString returns a string containing the state of s.
+// Used in tests.
+func (s *MemoryEventStore) debugString() string {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	var b strings.Builder
+	for i, sess := range slices.Sorted(maps.Keys(s.store)) {
+		if i > 0 {
+			fmt.Fprintf(&b, "; ")
+		}
+		sm := s.store[sess]
+		for i, sid := range slices.Sorted(maps.Keys(sm)) {
+			if i > 0 {
+				fmt.Fprintf(&b, "; ")
+			}
+			dl := sm[sid]
+			fmt.Fprintf(&b, "%s %d first=%d", sess, sid, dl.first)
+			for _, d := range dl.data {
+				fmt.Fprintf(&b, " %s", d)
+			}
+		}
+	}
+	return b.String()
+}