From 1465442b250629a7108d05f168c7f06faad707bc Mon Sep 17 00:00:00 2001 From: Jonathan Amsterdam Date: Tue, 22 Jul 2025 10:21:05 -0400 Subject: [PATCH] mcp: add EventStore (#152) Introduce an EventStore interface to store events for resumable streams. Provide an in-memory implementation. Still to do: connect to streaming transports. For #10 --- mcp/event.go | 290 +++++++++++++++++++++++++++++++++++++++++ mcp/event_test.go | 167 ++++++++++++++++++++++++ mcp/streamable.go | 32 ++--- mcp/streamable_test.go | 2 +- 4 files changed, 474 insertions(+), 17 deletions(-) diff --git a/mcp/event.go b/mcp/event.go index 4562b81..fbbe194 100644 --- a/mcp/event.go +++ b/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,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() +} diff --git a/mcp/event_test.go b/mcp/event_test.go index 2e2e3d2..9b01555 100644 --- a/mcp/event_test.go +++ b/mcp/event_test.go @@ -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) + } + }) + } +} diff --git a/mcp/streamable.go b/mcp/streamable.go index ef740d3..728da85 100644 --- a/mcp/streamable.go +++ b/mcp/streamable.go @@ -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] diff --git a/mcp/streamable_test.go b/mcp/streamable_test.go index 05bf59e..87664da 100644 --- a/mcp/streamable_test.go +++ b/mcp/streamable_test.go @@ -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},