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

mcp: add EventStore (#152)

Browse source 
Introduce an EventStore interface to store events for resumable streams.

Provide an in-memory implementation.

Still to do: connect to streaming transports.

For #10
This commit is contained in:
Jonathan Amsterdam 2025-07-22 10:21:05 -04:00 • committed by GitHub
commit 1465442b25
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 474 additions and 17 deletions
Download patch file Download diff file Expand all files Collapse all files

290
mcp/event.go
View file

@ -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,285 @@ 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 {
// AppendEvent appends data for an outgoing event to given stream, which is part of the
// given session. It returns the index of the event in the stream, suitable for constructing
// an event ID to send to the client.
AppendEvent(_ context.Context, sessionID string, _ StreamID, data []byte) (int, 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]
// StreamClosed informs the store that the given stream is finished.
// A store cannot rely on this method being called for cleanup. It should institute
// additional mechanisms, such as timeouts, to reclaim storage.
StreamClosed(_ context.Context, sessionID string, streamID StreamID) 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
}
// 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
}
// lastIndex returns the index of the last data item in dl.
// It panics if there are none.
func (dl *dataList) lastIndex() int {
if len(dl.data) == 0 {
panic("empty dataList")
}
return dl.first + len(dl.data) - 1
}
// 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),
}
}
// AppendEvent implements [EventStore.AppendEvent] by recording data
// in memory.
func (s *MemoryEventStore) AppendEvent(_ context.Context, sessionID string, streamID StreamID, data []byte) (int, 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 dl.lastIndex(), nil
}
// 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)
}
if dl.first > index {
return nil, fmt.Errorf("MemoryEventStore.After: data purged at index %d, stream ID %v, session %q", index, streamID, sessionID)
}
return slices.Clone(dl.data[index-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
}
}
}
}
// StreamClosed implements [EventStore.StreamClosed].
func (s *MemoryEventStore) StreamClosed(_ context.Context, sessionID string, streamID StreamID) error {
if sessionID == "" {
panic("empty sessionID")
}
s.mu.Lock()
defer s.mu.Unlock()
sm := s.store[sessionID]
dl := sm[streamID]
s.nBytes -= dl.size
delete(sm, streamID)
if len(sm) == 0 {
delete(s.store, sessionID)
}
s.validate()
return nil
}
// 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()
}

167
mcp/event_test.go
View file

@ -5,6 +5,9 @@
package mcp
import (
"context"
"fmt"
"slices"
"strings"
"testing"
)
@ -97,3 +100,167 @@ func TestScanEvents(t *testing.T) {
})
}
}
func TestMemoryEventStoreState(t *testing.T) {
ctx := context.Background()
appendEvent := func(s *MemoryEventStore, sess string, str StreamID, data string) {
if _, err := s.AppendEvent(ctx, sess, str, []byte(data)); err != nil {
t.Fatal(err)
}
}
for _, tt := range []struct {
name string
actions func(*MemoryEventStore)
want string // output of debugString
wantSize int // value of nBytes
}{
{
"appends",
func(s *MemoryEventStore) {
appendEvent(s, "S1", 1, "d1")
appendEvent(s, "S1", 2, "d2")
appendEvent(s, "S1", 1, "d3")
appendEvent(s, "S2", 8, "d4")
},
"S1 1 first=0 d1 d3; S1 2 first=0 d2; S2 8 first=0 d4",
8,
},
{
"stream close",
func(s *MemoryEventStore) {
appendEvent(s, "S1", 1, "d1")
appendEvent(s, "S1", 2, "d2")
appendEvent(s, "S1", 1, "d3")
appendEvent(s, "S2", 8, "d4")
s.StreamClosed(ctx, "S1", 1)
},
"S1 2 first=0 d2; S2 8 first=0 d4",
4,
},
{
"session close",
func(s *MemoryEventStore) {
appendEvent(s, "S1", 1, "d1")
appendEvent(s, "S1", 2, "d2")
appendEvent(s, "S1", 1, "d3")
appendEvent(s, "S2", 8, "d4")
s.SessionClosed(ctx, "S1")
},
"S2 8 first=0 d4",
2,
},
{
"purge",
func(s *MemoryEventStore) {
appendEvent(s, "S1", 1, "d1")
appendEvent(s, "S1", 2, "d2")
appendEvent(s, "S1", 1, "d3")
appendEvent(s, "S2", 8, "d4")
// We are using 8 bytes (d1,d2, d3, d4).
// To purge 6, we remove the first of each stream, leaving only d3.
s.SetMaxBytes(2)
},
// The other streams remain, because we may add to them.
"S1 1 first=1 d3; S1 2 first=1; S2 8 first=1",
2,
},
{
"purge append",
func(s *MemoryEventStore) {
appendEvent(s, "S1", 1, "d1")
appendEvent(s, "S1", 2, "d2")
appendEvent(s, "S1", 1, "d3")
appendEvent(s, "S2", 8, "d4")
s.SetMaxBytes(2)
// Up to here, identical to the "purge" case.
// Each of these additions will result in a purge.
appendEvent(s, "S1", 2, "d5") // remove d3
appendEvent(s, "S1", 2, "d6") // remove d5
},
"S1 1 first=2; S1 2 first=2 d6; S2 8 first=1",
2,
},
{
"purge resize append",
func(s *MemoryEventStore) {
appendEvent(s, "S1", 1, "d1")
appendEvent(s, "S1", 2, "d2")
appendEvent(s, "S1", 1, "d3")
appendEvent(s, "S2", 8, "d4")
s.SetMaxBytes(2)
// Up to here, identical to the "purge" case.
s.SetMaxBytes(6) // make room
appendEvent(s, "S1", 2, "d5")
appendEvent(s, "S1", 2, "d6")
},
// The other streams remain, because we may add to them.
"S1 1 first=1 d3; S1 2 first=1 d5 d6; S2 8 first=1",
6,
},
} {
t.Run(tt.name, func(t *testing.T) {
s := NewMemoryEventStore(nil)
tt.actions(s)
got := s.debugString()
if got != tt.want {
t.Errorf("\ngot %s\nwant %s", got, tt.want)
}
if g, w := s.nBytes, tt.wantSize; g != w {
t.Errorf("got size %d, want %d", g, w)
}
})
}
}
func TestMemoryEventStoreAfter(t *testing.T) {
ctx := context.Background()
s := NewMemoryEventStore(nil)
s.SetMaxBytes(4)
s.AppendEvent(ctx, "S1", 1, []byte("d1"))
s.AppendEvent(ctx, "S1", 1, []byte("d2"))
s.AppendEvent(ctx, "S1", 1, []byte("d3"))
s.AppendEvent(ctx, "S1", 2, []byte("d4")) // will purge d1
want := "S1 1 first=1 d2 d3; S1 2 first=0 d4"
if got := s.debugString(); got != want {
t.Fatalf("got state %q, want %q", got, want)
}
for _, tt := range []struct {
sessionID string
streamID StreamID
index int
want []string
wantErr string // if non-empty, error should contain this string
}{
{"S1", 1, 0, nil, "purge"},
{"S1", 1, 1, []string{"d2", "d3"}, ""},
{"S1", 1, 2, []string{"d3"}, ""},
{"S1", 2, 0, []string{"d4"}, ""},
{"S1", 3, 0, nil, "unknown stream ID"},
{"S2", 0, 0, nil, "unknown session ID"},
} {
t.Run(fmt.Sprintf("%s-%d-%d", tt.sessionID, tt.streamID, tt.index), func(t *testing.T) {
var got []string
for d, err := range s.After(ctx, tt.sessionID, tt.streamID, tt.index) {
if err != nil {
if tt.wantErr == "" {
t.Fatalf("unexpected error %q", err)
} else if g := err.Error(); !strings.Contains(g, tt.wantErr) {
t.Fatalf("got error %q, want it to contain %q", g, tt.wantErr)
} else {
return
}
}
got = append(got, string(d))
}
if tt.wantErr != "" {
t.Fatalf("expected error containing %q, got nil", tt.wantErr)
}
if !slices.Equal(got, tt.want) {
t.Errorf("got %v, want %v", got, tt.want)
}
})
}
}

32
mcp/streamable.go
View file

@ -163,10 +163,10 @@ func NewStreamableServerTransport(sessionID string) *StreamableServerTransport {
id: sessionID,
incoming: make(chan jsonrpc.Message, 10),
done: make(chan struct{}),
outgoingMessages: make(map[streamID][]*streamableMsg),
signals: make(map[streamID]chan struct{}),
requestStreams: make(map[jsonrpc.ID]streamID),
streamRequests: make(map[streamID]map[jsonrpc.ID]struct{}),
outgoingMessages: make(map[StreamID][]*streamableMsg),
signals: make(map[StreamID]chan struct{}),
requestStreams: make(map[jsonrpc.ID]StreamID),
streamRequests: make(map[StreamID]map[jsonrpc.ID]struct{}),
}
}
@ -212,7 +212,7 @@ type StreamableServerTransport struct {
//
// TODO(rfindley): garbage collect this data. For now, we save all outgoingMessages
// messages for the lifespan of the transport.
outgoingMessages map[streamID][]*streamableMsg
outgoingMessages map[StreamID][]*streamableMsg
// signals maps a logical stream ID to a 1-buffered channel, owned by an
// incoming HTTP request, that signals that there are messages available to
@ -223,14 +223,14 @@ type StreamableServerTransport struct {
//
// Lifecycle: signals persists for the duration of an HTTP POST or GET
// request for the given streamID.
signals map[streamID]chan struct{}
signals map[StreamID]chan struct{}
// requestStreams maps incoming requests to their logical stream ID.
//
// Lifecycle: requestStreams persists for the duration of the session.
//
// TODO(rfindley): clean up once requests are handled.
requestStreams map[jsonrpc.ID]streamID
requestStreams map[jsonrpc.ID]StreamID
// streamRequests tracks the set of unanswered incoming RPCs for each logical
// stream.
@ -241,10 +241,10 @@ type StreamableServerTransport struct {
// Lifecycle: streamRequests values persist as until the requests have been
// replied to by the server. Notably, NOT until they are sent to an HTTP
// response, as delivery is not guaranteed.
streamRequests map[streamID]map[jsonrpc.ID]struct{}
streamRequests map[StreamID]map[jsonrpc.ID]struct{}
}
type streamID int64
type StreamID int64
// a streamableMsg is an SSE event with an index into its logical stream.
type streamableMsg struct {
@ -298,7 +298,7 @@ func (t *StreamableServerTransport) ServeHTTP(w http.ResponseWriter, req *http.R
func (t *StreamableServerTransport) serveGET(w http.ResponseWriter, req *http.Request) {
// connID 0 corresponds to the default GET request.
id, nextIdx := streamID(0), 0
id, nextIdx := StreamID(0), 0
if len(req.Header.Values("Last-Event-ID")) > 0 {
eid := req.Header.Get("Last-Event-ID")
var ok bool
@ -352,7 +352,7 @@ func (t *StreamableServerTransport) servePOST(w http.ResponseWriter, req *http.R
}
// Update accounting for this request.
id := streamID(t.nextStreamID.Add(1))
id := StreamID(t.nextStreamID.Add(1))
signal := make(chan struct{}, 1)
t.mu.Lock()
if len(requests) > 0 {
@ -376,7 +376,7 @@ func (t *StreamableServerTransport) servePOST(w http.ResponseWriter, req *http.R
t.streamResponse(w, req, id, 0, signal)
}
func (t *StreamableServerTransport) streamResponse(w http.ResponseWriter, req *http.Request, id streamID, nextIndex int, signal chan struct{}) {
func (t *StreamableServerTransport) streamResponse(w http.ResponseWriter, req *http.Request, id StreamID, nextIndex int, signal chan struct{}) {
defer func() {
t.mu.Lock()
delete(t.signals, id)
@ -463,7 +463,7 @@ stream:
// streamID and message index idx.
//
// See also [parseEventID].
func formatEventID(sid streamID, idx int) string {
func formatEventID(sid StreamID, idx int) string {
return fmt.Sprintf("%d_%d", sid, idx)
}
@ -471,7 +471,7 @@ func formatEventID(sid streamID, idx int) string {
// index.
//
// See also [formatEventID].
func parseEventID(eventID string) (sid streamID, idx int, ok bool) {
func parseEventID(eventID string) (sid StreamID, idx int, ok bool) {
parts := strings.Split(eventID, "_")
if len(parts) != 2 {
return 0, 0, false
@ -484,7 +484,7 @@ func parseEventID(eventID string) (sid streamID, idx int, ok bool) {
if err != nil || idx < 0 {
return 0, 0, false
}
return streamID(stream), idx, true
return StreamID(stream), idx, true
}
// Read implements the [Connection] interface.
@ -523,7 +523,7 @@ func (t *StreamableServerTransport) Write(ctx context.Context, msg jsonrpc.Messa
//
// For messages sent outside of a request context, this is the default
// connection 0.
var forConn streamID
var forConn StreamID
if forRequest.IsValid() {
t.mu.Lock()
forConn = t.requestStreams[forRequest]

2
mcp/streamable_test.go
View file

@ -555,7 +555,7 @@ func mustMarshal(t *testing.T, v any) json.RawMessage {
func TestEventID(t *testing.T) {
tests := []struct {
sid streamID
sid StreamID
idx int
}{
{0, 0},