fix(Srv/stream): add ID field to PingRequest (#353)

cryo-zd · web-flow · commit 4b7117679253 · 2025-05-30T10:42:58.000+03:00
* fix(Srv/stream): add ID field to PingRequest

* fix: make session requestID unique and incremental
diff --git a/server/streamable_http.go b/server/streamable_http.go
@@ -115,8 +115,9 @@ func WithLogger(logger util.Logger) StreamableHTTPOption {
 //   - Batching of requests/notifications/responses in arrays.
 //   - Stream Resumability
 type StreamableHTTPServer struct {
-	server       *MCPServer
-	sessionTools *sessionToolsStore
+	server            *MCPServer
+	sessionTools      *sessionToolsStore
+	sessionRequestIDs sync.Map // sessionId --> last requestID(*atomic.Int64)
 
 	httpServer *http.Server
 	mu         sync.RWMutex
@@ -407,15 +408,16 @@ func (s *StreamableHTTPServer) handleGet(w http.ResponseWriter, r *http.Request)
 		go func() {
 			ticker := time.NewTicker(s.listenHeartbeatInterval)
 			defer ticker.Stop()
-			message := mcp.JSONRPCRequest{
-				JSONRPC: "2.0",
-				Request: mcp.Request{
-					Method: "ping",
-				},
-			}
 			for {
 				select {
 				case <-ticker.C:
+					message := mcp.JSONRPCRequest{
+						JSONRPC: "2.0",
+						ID:      mcp.NewRequestId(s.nextRequestID(sessionID)),
+						Request: mcp.Request{
+							Method: "ping",
+						},
+					}
 					select {
 					case writeChan <- message:
 					case <-done:
@@ -465,6 +467,9 @@ func (s *StreamableHTTPServer) handleDelete(w http.ResponseWriter, r *http.Reque
 	// remove the session relateddata from the sessionToolsStore
 	s.sessionTools.set(sessionID, nil)
 
+	// remove current session's requstID information
+	s.sessionRequestIDs.Delete(sessionID)
+
 	w.WriteHeader(http.StatusOK)
 }
 
@@ -496,6 +501,13 @@ func (s *StreamableHTTPServer) writeJSONRPCError(
 	}
 }
 
+// nextRequestID gets the next incrementing requestID for the current session
+func (s *StreamableHTTPServer) nextRequestID(sessionID string) int64 {
+	actual, _ := s.sessionRequestIDs.LoadOrStore(sessionID, new(atomic.Int64))
+	counter := actual.(*atomic.Int64)
+	return counter.Add(1)
+}
+
 // --- session ---
 
 type sessionToolsStore struct {
