feat: add ping for sse server

server/sse.go

@@ -10,6 +10,7 @@ import (
 	"strings"
 	"sync"
 	"sync/atomic"
+	"time"
 
 	"github.com/google/uuid"
 	"github.com/mark3labs/mcp-go/mcp"
@@ -52,15 +53,18 @@ var _ ClientSession = (*sseSession)(nil)
 // SSEServer implements a Server-Sent Events (SSE) based MCP server.
 // It provides real-time communication capabilities over HTTP using the SSE protocol.
 type SSEServer struct {
-	server                       *MCPServer
-	baseURL                      string
-	basePath                     string
-	messageEndpoint              string
-	useFullURLForMessageEndpoint bool
-	sseEndpoint                  string
-	sessions                     sync.Map
-	srv                          *http.Server
-	contextFunc                  SSEContextFunc
+	server          *MCPServer
+	baseURL         string
+	basePath        string
+  useFullURLForMessageEndpoint bool
+	messageEndpoint string
+	sseEndpoint     string
+	sessions        sync.Map
+	srv             *http.Server
+	contextFunc     SSEContextFunc
+
+	keepAlive         bool
+	keepAliveInterval time.Duration
 }
 
 // SSEOption defines a function type for configuring SSEServer
@@ -130,6 +134,19 @@ func WithHTTPServer(srv *http.Server) SSEOption {
 	}
 }
 
+func WithKeepAliveInterval(keepAliveInterval time.Duration) SSEOption {
+	return func(s *SSEServer) {
+		s.keepAlive = true
+		s.keepAliveInterval = keepAliveInterval
+	}
+}
+
+func WithKeepAlive(keepAlive bool) SSEOption {
+	return func(s *SSEServer) {
+		s.keepAlive = keepAlive
+	}
+}
+
 // WithContextFunc sets a function that will be called to customise the context
 // to the server using the incoming request.
 func WithSSEContextFunc(fn SSEContextFunc) SSEOption {
@@ -141,10 +158,12 @@ func WithSSEContextFunc(fn SSEContextFunc) SSEOption {
 // NewSSEServer creates a new SSE server instance with the given MCP server and options.
 func NewSSEServer(server *MCPServer, opts ...SSEOption) *SSEServer {
 	s := &SSEServer{
-		server:                       server,
-		sseEndpoint:                  "/sse",
-		messageEndpoint:              "/message",
-		useFullURLForMessageEndpoint: true,
+		server:            server,
+		sseEndpoint:       "/sse",
+		messageEndpoint:   "/message",
+    useFullURLForMessageEndpoint: true,
+		keepAlive:         false,
+		keepAliveInterval: 10 * time.Second,
 	}
 
 	// Apply all options
@@ -255,6 +274,27 @@ func (s *SSEServer) handleSSE(w http.ResponseWriter, r *http.Request) {
 		}
 	}()
 
+
+	// Start keep alive : ping
+	if s.keepAlive {
+		go func() {
+			ticker := time.NewTicker(s.keepAliveInterval)
+			defer ticker.Stop()
+			for {
+				select {
+				case <-ticker.C:
+					session.eventQueue <- fmt.Sprintf(":ping - %s\n\n", time.Now().Format(time.RFC3339))
+				case <-session.done:
+					return
+				case <-r.Context().Done():
+					return
+				}
+			}
+		}()
+	}
+
+	messageEndpoint := fmt.Sprintf("%s?sessionId=%s", s.CompleteMessageEndpoint(), sessionID)
+
 	// Send the initial endpoint event
 	fmt.Fprintf(w, "event: endpoint\ndata: %s\r\n\r\n", s.GetMessageEndpointForClient(sessionID))
 	flusher.Flush()