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: changed streamID's from int64 to random strings (#266)

Browse source 
- Changed StreamID's to store randomly generated strings.
- Updated all tests.
- Resolved conflicts
This commit is contained in:
Hugh Palmer 2025-08-08 11:51:53 +02:00 • committed by Robert Findley
commit 0a8fe40c51
4 changed files with 61 additions and 66 deletions
Download patch file Download diff file Expand all files Collapse all files

2
mcp/event.go
View file

@ -392,7 +392,7 @@ func (s *MemoryEventStore) debugString() string {
fmt.Fprintf(&b, "; ")
}
dl := sm[sid]
fmt.Fprintf(&b, "%s %d first=%d", sess, sid, dl.first)
fmt.Fprintf(&b, "%s %s first=%d", sess, sid, dl.first)
for _, d := range dl.data {
fmt.Fprintf(&b, " %s", d)
}

70
mcp/event_test.go
View file

@ -119,10 +119,10 @@ func TestMemoryEventStoreState(t *testing.T) {
{
"appends",
func(s *MemoryEventStore) {
appendEvent(s, "S1", 1, "d1")
appendEvent(s, "S1", 2, "d2")
appendEvent(s, "S1", 1, "d3")
appendEvent(s, "S2", 8, "d4")
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,
@ -130,10 +130,10 @@ func TestMemoryEventStoreState(t *testing.T) {
{
"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")
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",
@ -142,10 +142,10 @@ func TestMemoryEventStoreState(t *testing.T) {
{
"purge",
func(s *MemoryEventStore) {
appendEvent(s, "S1", 1, "d1")
appendEvent(s, "S1", 2, "d2")
appendEvent(s, "S1", 1, "d3")
appendEvent(s, "S2", 8, "d4")
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)
@ -157,15 +157,15 @@ func TestMemoryEventStoreState(t *testing.T) {
{
"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")
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
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,
@ -173,15 +173,15 @@ func TestMemoryEventStoreState(t *testing.T) {
{
"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")
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")
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",
@ -206,10 +206,10 @@ func TestMemoryEventStoreAfter(t *testing.T) {
ctx := context.Background()
s := NewMemoryEventStore(nil)
s.SetMaxBytes(4)
s.Append(ctx, "S1", 1, []byte("d1"))
s.Append(ctx, "S1", 1, []byte("d2"))
s.Append(ctx, "S1", 1, []byte("d3"))
s.Append(ctx, "S1", 2, []byte("d4")) // will purge d1
s.Append(ctx, "S1", "1", []byte("d1"))
s.Append(ctx, "S1", "1", []byte("d2"))
s.Append(ctx, "S1", "1", []byte("d3"))
s.Append(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)
@ -222,14 +222,14 @@ func TestMemoryEventStoreAfter(t *testing.T) {
want []string
wantErr string // if non-empty, error should contain this string
}{
{"S1", 1, 0, []string{"d2", "d3"}, ""},
{"S1", 1, 1, []string{"d3"}, ""},
{"S1", 1, 2, nil, ""},
{"S1", 2, 0, nil, ""},
{"S1", 3, 0, nil, "unknown stream ID"},
{"S2", 0, 0, nil, "unknown session ID"},
{"S1", "1", 0, []string{"d2", "d3"}, ""},
{"S1", "1", 1, []string{"d3"}, ""},
{"S1", "1", 2, nil, ""},
{"S1", "2", 0, nil, ""},
{"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) {
t.Run(fmt.Sprintf("%s-%s-%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 {

37
mcp/streamable.go
View file

@ -276,7 +276,7 @@ func (t *StreamableServerTransport) Connect(context.Context) (Connection, error)
//
// It is always text/event-stream, since it must carry arbitrarily many
// messages.
t.connection.streams[0] = newStream(0, false)
t.connection.streams[""] = newStream("", false)
if t.connection.eventStore == nil {
t.connection.eventStore = NewMemoryEventStore(nil)
}
@ -334,7 +334,7 @@ func (c *streamableServerConn) SessionID() string {
// at any time.
type stream struct {
// id is the logical ID for the stream, unique within a session.
// ID 0 is used for messages that don't correlate with an incoming request.
// an empty string is used for messages that don't correlate with an incoming request.
id StreamID
// jsonResponse records whether this stream should respond with application/json
@ -382,9 +382,9 @@ func signalChanPtr() *chan struct{} {
return &c
}
// A StreamID identifies a stream of SSE events. It is unique within the stream's
// A StreamID identifies a stream of SSE events. It is globally unique.
// [ServerSession].
type StreamID int64
type StreamID string
// We track the incoming request ID inside the handler context using
// idContextValue, so that notifications and server->client calls that occur in
@ -434,7 +434,7 @@ func (t *StreamableServerTransport) ServeHTTP(w http.ResponseWriter, req *http.R
// It returns an HTTP status code and error message.
func (c *streamableServerConn) serveGET(w http.ResponseWriter, req *http.Request) {
// connID 0 corresponds to the default GET request.
id := StreamID(0)
id := StreamID("")
// By default, we haven't seen a last index. Since indices start at 0, we represent
// that by -1. This is incremented just before each event is written, in streamResponse
// around L407.
@ -462,7 +462,7 @@ func (c *streamableServerConn) serveGET(w http.ResponseWriter, req *http.Request
return
}
defer stream.signal.Store(nil)
persistent := id == 0 // Only the special stream 0 is a hanging get.
persistent := id == "" // Only the special stream "" is a hanging get.
c.respondSSE(stream, w, req, lastIdx, persistent)
}
@ -520,7 +520,7 @@ func (c *streamableServerConn) servePOST(w http.ResponseWriter, req *http.Reques
// notifications or server->client requests made in the course of handling.
// Update accounting for this incoming payload.
if len(requests) > 0 {
stream = newStream(StreamID(c.lastStreamID.Add(1)), c.jsonResponse)
stream = newStream(StreamID(randText()), c.jsonResponse)
c.mu.Lock()
c.streams[stream.id] = stream
stream.requests = requests
@ -719,7 +719,7 @@ func (c *streamableServerConn) messages(ctx context.Context, stream *stream, per
//
// See also [parseEventID].
func formatEventID(sid StreamID, idx int) string {
return fmt.Sprintf("%d_%d", sid, idx)
return fmt.Sprintf("%s_%d", sid, idx)
}
// parseEventID parses a Last-Event-ID value into a logical stream id and
@ -729,15 +729,12 @@ func formatEventID(sid StreamID, idx int) string {
func parseEventID(eventID string) (sid StreamID, idx int, ok bool) {
parts := strings.Split(eventID, "_")
if len(parts) != 2 {
return 0, 0, false
return "", 0, false
}
stream, err := strconv.ParseInt(parts[0], 10, 64)
if err != nil || stream < 0 {
return 0, 0, false
}
idx, err = strconv.Atoi(parts[1])
stream := StreamID(parts[0])
idx, err := strconv.Atoi(parts[1])
if err != nil || idx < 0 {
return 0, 0, false
return "", 0, false
}
return StreamID(stream), idx, true
}
@ -778,7 +775,7 @@ func (c *streamableServerConn) Write(ctx context.Context, msg jsonrpc.Message) e
// Find the logical connection corresponding to this request.
//
// For messages sent outside of a request context, this is the default
// connection 0.
// connection "".
var forStream StreamID
if forRequest.IsValid() {
c.mu.Lock()
@ -799,7 +796,7 @@ func (c *streamableServerConn) Write(ctx context.Context, msg jsonrpc.Message) e
stream := c.streams[forStream]
if stream == nil {
return fmt.Errorf("no stream with ID %d", forStream)
return fmt.Errorf("no stream with ID %s", forStream)
}
// Special case a few conditions where we fall back on stream 0 (the hanging GET):
@ -809,11 +806,11 @@ func (c *streamableServerConn) Write(ctx context.Context, msg jsonrpc.Message) e
//
// TODO(rfindley): either of these, particularly the first, might be
// considered a bug in the server. Report it through a side-channel?
if len(stream.requests) == 0 && forStream != 0 || stream.jsonResponse && !isResponse {
stream = c.streams[0]
if len(stream.requests) == 0 && forStream != "" || stream.jsonResponse && !isResponse {
stream = c.streams[""]
}
// TODO: if there is nothing to send these messages to (as would happen, for example, if forConn == 0
// TODO: if there is nothing to send these messages to (as would happen, for example, if forConn == ""
// and the client never did a GET), then memory will grow without bound. Consider a mitigation.
stream.outgoing = append(stream.outgoing, data)
if isResponse {

18
mcp/streamable_test.go
View file

@ -888,22 +888,23 @@ func TestEventID(t *testing.T) {
sid StreamID
idx int
}{
{0, 0},
{0, 1},
{1, 0},
{1, 1},
{1234, 5678},
{"0", 0},
{"0", 1},
{"1", 0},
{"1", 1},
{"", 1},
{"1234", 5678},
}
for _, test := range tests {
t.Run(fmt.Sprintf("%d_%d", test.sid, test.idx), func(t *testing.T) {
t.Run(fmt.Sprintf("%s_%d", test.sid, test.idx), func(t *testing.T) {
eventID := formatEventID(test.sid, test.idx)
gotSID, gotIdx, ok := parseEventID(eventID)
if !ok {
t.Fatalf("parseEventID(%q) failed, want ok", eventID)
}
if gotSID != test.sid || gotIdx != test.idx {
t.Errorf("parseEventID(%q) = %d, %d, want %d, %d", eventID, gotSID, gotIdx, test.sid, test.idx)
t.Errorf("parseEventID(%q) = %s, %d, want %s, %d", eventID, gotSID, gotIdx, test.sid, test.idx)
}
})
}
@ -912,10 +913,7 @@ func TestEventID(t *testing.T) {
"",
"_",
"1_",
"_1",
"a_1",
"1_a",
"-1_1",
"1_-1",
}