mcp: Implement KeepAlive for client and server

Implement KeepAlive features from the design doc.

Fixes #24.

mcp/server.go

@@ -16,6 +16,7 @@ import (
 	"path/filepath"
 	"slices"
 	"sync"
+	"time"
 
 	"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
 	"github.com/modelcontextprotocol/go-sdk/internal/util"
@@ -57,6 +58,10 @@ type ServerOptions struct {
 	RootsListChangedHandler func(context.Context, *ServerSession, *RootsListChangedParams)
 	// If non-nil, called when "notifications/progress" is received.
 	ProgressNotificationHandler func(context.Context, *ServerSession, *ProgressNotificationParams)
+	// If non-zero, defines an interval for regular "ping" requests.
+	// If the peer fails to respond to pings originating from the keepalive check,
+	// the session is automatically closed.
+	KeepAlive time.Duration
 }
 
 // NewServer creates a new MCP server. The resulting server has no features:
@@ -460,6 +465,9 @@ func (s *Server) Connect(ctx context.Context, t Transport) (*ServerSession, erro
 }
 
 func (s *Server) callInitializedHandler(ctx context.Context, ss *ServerSession, params *InitializedParams) (Result, error) {
+	if s.opts.KeepAlive > 0 {
+		ss.startKeepalive(s.opts.KeepAlive)
+	}
 	return callNotificationHandler(ctx, s.opts.InitializedHandler, ss, params)
 }
 
@@ -492,6 +500,7 @@ type ServerSession struct {
 	logLevel         LoggingLevel
 	initializeParams *InitializeParams
 	initialized      bool
+	keepaliveCancel  context.CancelFunc
 }
 
 // Ping pings the client.
@@ -678,6 +687,14 @@ func (ss *ServerSession) setLevel(_ context.Context, params *SetLevelParams) (*e
 // requests from being handled, and waiting for ongoing requests to return.
 // Close then terminates the connection.
 func (ss *ServerSession) Close() error {
+	if ss.keepaliveCancel != nil {
+		// Note: keepaliveCancel access is safe without a mutex because:
+		// 1. keepaliveCancel is only written once during startKeepalive (happens-before all Close calls)
+		// 2. context.CancelFunc is safe to call multiple times and from multiple goroutines
+		// 3. The keepalive goroutine calls Close on ping failure, but this is safe since
+		//    Close is idempotent and conn.Close() handles concurrent calls correctly
+		ss.keepaliveCancel()
+	}
 	return ss.conn.Close()
 }
 
@@ -686,6 +703,11 @@ func (ss *ServerSession) Wait() error {
 	return ss.conn.Wait()
 }
 
+// startKeepalive starts the keepalive mechanism for this server session.
+func (ss *ServerSession) startKeepalive(interval time.Duration) {
+	startKeepalive(ss, interval, &ss.keepaliveCancel)
+}
+
 // pageToken is the internal structure for the opaque pagination cursor.
 // It will be Gob-encoded and then Base64-encoded for use as a string token.
 type pageToken struct {