mcp: changed streamID's from int64 to random strings

Changed StreamID's to store randomly generated strings.
Created a defaultStreamID field on the StreamableServerTransport struct
which is generated and saved on instansiation.
Updated all tests.

Fixes #259

Author: 2bitburrito

mcp/event.go

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

mcp/event_test.go

@@ -119,21 +119,21 @@ 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,
 		},
 		{
 			"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,31 +157,31 @@ 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,
 		},
 		{
 			"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 {

mcp/streamable.go

@@ -176,13 +176,14 @@ func NewStreamableServerTransport(sessionID string, opts *StreamableServerTransp
 		opts = &StreamableServerTransportOptions{}
 	}
 	t := &StreamableServerTransport{
-		sessionID:      sessionID,
-		incoming:       make(chan jsonrpc.Message, 10),
-		done:           make(chan struct{}),
-		streams:        make(map[StreamID]*stream),
-		requestStreams: make(map[jsonrpc.ID]StreamID),
-	}
-	t.streams[0] = newStream(0)
+		sessionID:       sessionID,
+		incoming:        make(chan jsonrpc.Message, 10),
+		done:            make(chan struct{}),
+		streams:         make(map[StreamID]*stream),
+		requestStreams:  make(map[jsonrpc.ID]StreamID),
+		defaultStreamID: StreamID(randText()),
+	}
+	t.streams[t.defaultStreamID] = newStream(t.defaultStreamID)
 	if opts != nil {
 		t.opts = *opts
 	}
@@ -199,12 +200,11 @@ func (t *StreamableServerTransport) SessionID() string {
 // A StreamableServerTransport implements the [Transport] interface for a
 // single session.
 type StreamableServerTransport struct {
-	nextStreamID atomic.Int64 // incrementing next stream ID
-
-	sessionID string
-	opts      StreamableServerTransportOptions
-	incoming  chan jsonrpc.Message // messages from the client to the server
-	done      chan struct{}
+	sessionID       string
+	defaultStreamID StreamID
+	opts            StreamableServerTransportOptions
+	incoming        chan jsonrpc.Message // messages from the client to the server
+	done            chan struct{}
 
 	mu sync.Mutex
 	// Sessions are closed exactly once.
@@ -242,7 +242,7 @@ type StreamableServerTransport struct {
 // 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.
+	// defaultStreamID is used for messages that don't correlate with an incoming request.
 	id StreamID
 
 	// signal is a 1-buffered channel, owned by an incoming HTTP request, that signals
@@ -283,9 +283,10 @@ 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 a random string
+// 26 characters in length
 // [ServerSession].
-type StreamID int64
+type StreamID string
 
 // Connect implements the [Transport] interface.
 //
@@ -338,8 +339,8 @@ func (t *StreamableServerTransport) ServeHTTP(w http.ResponseWriter, req *http.R
 }
 
 func (t *StreamableServerTransport) serveGET(w http.ResponseWriter, req *http.Request) (int, string) {
-	// connID 0 corresponds to the default GET request.
-	id := StreamID(0)
+	// connID = t.defaultStreamID corresponds to the default GET request.
+	id := t.defaultStreamID
 	// 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.
@@ -399,7 +400,7 @@ func (t *StreamableServerTransport) servePOST(w http.ResponseWriter, req *http.R
 	}
 
 	// Update accounting for this request.
-	stream := newStream(StreamID(t.nextStreamID.Add(1)))
+	stream := newStream(StreamID(randText()))
 	t.mu.Lock()
 	t.streams[stream.id] = stream
 	if len(requests) > 0 {
@@ -533,7 +534,7 @@ stream:
 //
 // 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
@@ -543,15 +544,15 @@ 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
+	stream := StreamID(parts[0])
+	if len(stream) == 0 {
+		return "", 0, false
 	}
-	idx, err = strconv.Atoi(parts[1])
+	idx, err := strconv.Atoi(parts[1])
 	if err != nil || idx < 0 {
-		return 0, 0, false
+		return "", 0, false
 	}
 	return StreamID(stream), idx, true
 }
@@ -592,7 +593,7 @@ func (t *StreamableServerTransport) Write(ctx context.Context, msg jsonrpc.Messa
 	// Find the logical connection corresponding to this request.
 	//
 	// For messages sent outside of a request context, this is the default
-	// connection 0.
+	// connection.
 	var forConn StreamID
 	if forRequest.IsValid() {
 		t.mu.Lock()
@@ -612,18 +613,21 @@ func (t *StreamableServerTransport) Write(ctx context.Context, msg jsonrpc.Messa
 	}
 
 	stream := t.streams[forConn]
+	if forConn == "" {
+		stream = t.streams[t.defaultStreamID]
+	}
 	if stream == nil {
-		return fmt.Errorf("no stream with ID %d", forConn)
+		return fmt.Errorf("no stream with ID %s", forConn)
 	}
-	if len(stream.requests) == 0 && forConn != 0 {
+	if len(stream.requests) == 0 && forConn != "" {
 		// No outstanding requests for this connection, which means it is logically
 		// done. This is a sequencing violation from the server, so we should report
 		// a side-channel error here. Put the message on the general queue to avoid
 		// dropping messages.
-		stream = t.streams[0]
+		stream = t.streams[t.defaultStreamID]
 	}
 
-	// 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 {

mcp/streamable_test.go

@@ -222,9 +222,10 @@ func TestServerInitiatedSSE(t *testing.T) {
 
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
-	client := NewClient(testImpl, &ClientOptions{ToolListChangedHandler: func(ctx context.Context, cc *ClientSession, params *ToolListChangedParams) {
-		notifications <- "toolListChanged"
-	},
+	client := NewClient(testImpl, &ClientOptions{
+		ToolListChangedHandler: func(ctx context.Context, cc *ClientSession, params *ToolListChangedParams) {
+			notifications <- "toolListChanged"
+		},
 	})
 	clientSession, err := client.Connect(ctx, NewStreamableClientTransport(httpServer.URL, nil))
 	if err != nil {
@@ -768,22 +769,22 @@ 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},
+		{"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)
 			}
 		})
 	}
@@ -793,9 +794,7 @@ func TestEventID(t *testing.T) {
 		"_",
 		"1_",
 		"_1",
-		"a_1",
 		"1_a",
-		"-1_1",
 		"1_-1",
 	}