diff --git a/README.md b/README.md index 3e81604..d934ef6 100644 --- a/README.md +++ b/README.md @@ -384,6 +384,11 @@ server := mcp.NewServer( | `WithGetSSEEnabled` | Allow GET for SSE connections | `true` | | `WithNotificationBufferSize` | Size of notification buffer | `10` | | `WithStatelessMode` | Run in stateless mode | `false` | +| `WithServerKeepAlive` | Enable SSE comment keepalive | `true` | +| `WithServerKeepAliveInterval` | Interval for SSE comment keepalive | `30s` | +| `WithServerPingKeepAlive` | Enable JSON-RPC ping keepalive | `false` | +| `WithServerPingInterval` | Interval for ping requests | `30s` | +| `WithServerPingTimeout` | Timeout for ping requests | `15s` | ### Client Configuration @@ -728,6 +733,89 @@ basicPromptHandler := func(ctx context.Context, req *mcp.GetPromptRequest) (*mcp server.RegisterPrompt(basicPrompt, basicPromptHandler) ``` +### Connection Keepalive + +The framework provides two keepalive mechanisms for SSE connections to prevent timeouts: + +#### 1. SSE Comment Keepalive (Default) + +Sends SSE comment lines (`: keepalive`) at regular intervals. This is lightweight and enabled by default. + +```go +server := mcp.NewServer( + "My-Server", + "1.0.0", + // Comment keepalive is enabled by default + // Optionally customize: + mcp.WithServerKeepAlive(true), + mcp.WithServerKeepAliveInterval(30*time.Second), +) +``` + +#### 2. JSON-RPC Ping Keepalive (Optional) + +Sends JSON-RPC ping requests to clients, allowing health detection. Enable this for production systems. + +```go +server := mcp.NewServer( + "My-Server", + "1.0.0", + // Enable ping keepalive (comment keepalive remains enabled by default) + mcp.WithServerPingKeepAlive(true), + mcp.WithServerPingInterval(30*time.Second), + mcp.WithServerPingTimeout(15*time.Second), +) +``` + +#### Comparison + +| Feature | SSE Comment | JSON-RPC Ping | +|---------|-------------|---------------| +| Overhead | Small (~15 bytes) | Large (~100 bytes) | +| Health Check | ❌ No | ✅ Yes | +| Client Response | Not required | Required | +| Default | ✅ Enabled | ❌ Disabled | + +#### Configuration Examples + +**Default (Comment only)**: +```go +server := mcp.NewServer("My-Server", "1.0.0") +// ✅ SSE comment keepalive enabled (30s interval) +// ❌ Ping keepalive disabled +``` + +**Both modes (Recommended for production)**: +```go +server := mcp.NewServer("My-Server", "1.0.0", + mcp.WithServerPingKeepAlive(true), +) +// ✅ SSE comment keepalive enabled (30s interval) +// ✅ Ping keepalive enabled (30s interval, 15s timeout) +``` + +**Ping only**: +```go +server := mcp.NewServer("My-Server", "1.0.0", + mcp.WithServerKeepAlive(false), + mcp.WithServerPingKeepAlive(true), +) +// ❌ SSE comment keepalive disabled +// ✅ Ping keepalive enabled (30s interval, 15s timeout) +``` + +**Custom intervals**: +```go +server := mcp.NewServer("My-Server", "1.0.0", + mcp.WithServerKeepAliveInterval(60*time.Second), + mcp.WithServerPingKeepAlive(true), + mcp.WithServerPingInterval(45*time.Second), + mcp.WithServerPingTimeout(20*time.Second), +) +``` + +For a complete example, see [`examples/ping-keepalive/`](examples/ping-keepalive/). + ## Struct-First API (Recommended) Define MCP tools using Go structs for automatic schema generation and type safety: diff --git a/server.go b/server.go index e7f6e90..bbb43fd 100644 --- a/server.go +++ b/server.go @@ -14,6 +14,7 @@ import ( "net/http" "sync" "sync/atomic" + "time" ) // Common errors @@ -71,6 +72,13 @@ type serverConfig struct { getSSEEnabled bool notificationBufferSize int + // Keepalive configuration + keepAliveEnabled bool + keepAliveInterval time.Duration + pingEnabled bool + pingInterval time.Duration + pingTimeout time.Duration + // HTTP context functions for extracting information from HTTP requests httpContextFuncs []HTTPContextFunc @@ -120,6 +128,11 @@ func NewServer(name, version string, options ...ServerOption) *Server { postSSEEnabled: true, getSSEEnabled: true, notificationBufferSize: defaultNotificationBufferSize, + keepAliveEnabled: true, // Default: SSE comment keepalive enabled + keepAliveInterval: 30 * time.Second, // Default: 30 seconds + pingEnabled: false, // Default: ping keepalive disabled + pingInterval: 30 * time.Second, // Default: 30 seconds + pingTimeout: 15 * time.Second, // Default: 15 seconds } // Create server with provided serverInfo @@ -220,6 +233,12 @@ func (s *Server) initComponents() { withTransportNotificationBufferSize(s.config.notificationBufferSize), ) + // Keepalive configuration. + httpOptions = append(httpOptions, + withKeepAliveConfig(s.config.keepAliveEnabled, s.config.keepAliveInterval), + withPingConfig(s.config.pingEnabled, s.config.pingInterval, s.config.pingTimeout), + ) + // HTTP context functions configuration. if len(s.config.httpContextFuncs) > 0 { httpOptions = append(httpOptions, withTransportHTTPContextFuncs(s.config.httpContextFuncs)) @@ -300,6 +319,53 @@ func WithStatelessMode(enabled bool) ServerOption { } } +// WithServerKeepAlive enables or disables SSE comment keepalive for all SSE connections. +// When enabled, the server will send SSE comment lines (": keepalive") at regular intervals +// to prevent connection timeouts. This is enabled by default. +func WithServerKeepAlive(enabled bool) ServerOption { + return func(s *Server) { + s.config.keepAliveEnabled = enabled + } +} + +// WithServerKeepAliveInterval sets the interval for SSE comment keepalive messages. +// This option does not automatically enable keepalive; use WithServerKeepAlive(true) to enable it. +// Default is 30 seconds. +func WithServerKeepAliveInterval(interval time.Duration) ServerOption { + return func(s *Server) { + s.config.keepAliveInterval = interval + } +} + +// WithServerPingKeepAlive enables or disables JSON-RPC ping keepalive. +// When enabled, the server will periodically send ping requests to all connected clients. +// This provides application-layer health checking in addition to (or instead of) SSE comment keepalive. +// Disabled by default. +func WithServerPingKeepAlive(enabled bool) ServerOption { + return func(s *Server) { + s.config.pingEnabled = enabled + } +} + +// WithServerPingInterval sets the interval for sending ping requests. +// This option does not automatically enable ping keepalive; use WithServerPingKeepAlive(true) to enable it. +// Default is 30 seconds. +func WithServerPingInterval(interval time.Duration) ServerOption { + return func(s *Server) { + s.config.pingInterval = interval + } +} + +// WithServerPingTimeout sets the timeout for ping requests. +// If a ping request does not receive a response within this timeout, +// it will be logged as failed but the connection will remain open. +// Default is 15 seconds. +func WithServerPingTimeout(timeout time.Duration) ServerOption { + return func(s *Server) { + s.config.pingTimeout = timeout + } +} + // WithToolListFilter sets a tool list filter that will be applied to tools/list requests. // The filter function receives the request context and all registered tools, and should // return a filtered list of tools that should be visible to the client. diff --git a/server_test.go b/server_test.go index 15fc82b..ac509b3 100644 --- a/server_test.go +++ b/server_test.go @@ -11,6 +11,7 @@ import ( "net/http" "net/http/httptest" "testing" + "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -356,3 +357,85 @@ func TestServer_UnregisterTools(t *testing.T) { tools = server.toolManager.getTools() assert.Len(t, tools, 0) } + +// TestServerDefaultKeepAliveConfig tests that server has correct default keepalive configuration +func TestServerDefaultKeepAliveConfig(t *testing.T) { + server := NewServer("test-server", "1.0.0") + + // Verify default keepalive config + assert.True(t, server.config.keepAliveEnabled) + assert.Equal(t, 30*time.Second, server.config.keepAliveInterval) + assert.False(t, server.config.pingEnabled) + assert.Equal(t, 30*time.Second, server.config.pingInterval) + assert.Equal(t, 15*time.Second, server.config.pingTimeout) +} + +// TestServerPingKeepAliveOptions tests server ping keepalive configuration options +func TestServerPingKeepAliveOptions(t *testing.T) { + server := NewServer("test-server", "1.0.0", + WithServerPingKeepAlive(true), + WithServerPingInterval(60*time.Second), + WithServerPingTimeout(30*time.Second), + ) + + // Verify ping configuration + assert.True(t, server.config.pingEnabled) + assert.Equal(t, 60*time.Second, server.config.pingInterval) + assert.Equal(t, 30*time.Second, server.config.pingTimeout) + // Verify comment keepalive is still enabled + assert.True(t, server.config.keepAliveEnabled) +} + +// TestServerCommentKeepAliveOptions tests server comment keepalive configuration options +func TestServerCommentKeepAliveOptions(t *testing.T) { + server := NewServer("test-server", "1.0.0", + WithServerKeepAlive(false), + WithServerKeepAliveInterval(45*time.Second), + ) + + // Verify comment keepalive configuration + assert.False(t, server.config.keepAliveEnabled) + assert.Equal(t, 45*time.Second, server.config.keepAliveInterval) + // Verify ping is still disabled + assert.False(t, server.config.pingEnabled) +} + +// TestServerBackwardCompatibility tests that existing server config is not affected +func TestServerBackwardCompatibility(t *testing.T) { + // Test 1: Server without any keepalive options + server1 := NewServer("test-server", "1.0.0") + assert.True(t, server1.config.keepAliveEnabled) + assert.False(t, server1.config.pingEnabled) + + // Test 2: Server with only existing options + server2 := NewServer("test-server", "1.0.0", + WithServerAddress(":8080"), + WithServerPath("/api"), + ) + assert.True(t, server2.config.keepAliveEnabled) + assert.False(t, server2.config.pingEnabled) + assert.Equal(t, ":8080", server2.config.addr) + assert.Equal(t, "/api", server2.config.path) +} + +// TestServerPingOnlyMode tests using only ping keepalive +func TestServerPingOnlyMode(t *testing.T) { + server := NewServer("test-server", "1.0.0", + WithServerKeepAlive(false), + WithServerPingKeepAlive(true), + ) + + assert.False(t, server.config.keepAliveEnabled) + assert.True(t, server.config.pingEnabled) +} + +// TestServerBothKeepAliveModes tests using both comment and ping keepalive +func TestServerBothKeepAliveModes(t *testing.T) { + server := NewServer("test-server", "1.0.0", + WithServerPingKeepAlive(true), + ) + + // Both should be enabled + assert.True(t, server.config.keepAliveEnabled) // Comment is default + assert.True(t, server.config.pingEnabled) // Explicitly enabled +} diff --git a/sse_client.go b/sse_client.go index d5b7eff..4629fab 100644 --- a/sse_client.go +++ b/sse_client.go @@ -405,6 +405,8 @@ func (t *sseClientTransport) handleIncomingRequest(data string) { switch request.Method { case MethodRootsList: t.handleRootsListRequest(&request) + case MethodPing: + t.handlePingRequest(&request) default: // Send method not found error. t.sendErrorResponse(&request, ErrCodeMethodNotFound, fmt.Sprintf("Method not found: %s", request.Method)) @@ -443,6 +445,28 @@ func (t *sseClientTransport) handleRootsListRequest(request *JSONRPCRequest) { t.sendResponseMessage(response) } +// handlePingRequest handles ping requests from the server. +// It responds with an empty result to indicate the client is alive. +func (t *sseClientTransport) handlePingRequest(request *JSONRPCRequest) { + if t.logger != nil { + t.logger.Debugf("Received ping request, ID: %v", request.ID) + } + + // Create empty response + response := &JSONRPCResponse{ + JSONRPC: JSONRPCVersion, + ID: request.ID, + Result: map[string]interface{}{}, // Empty object + } + + // Send response + t.sendResponseMessage(response) + + if t.logger != nil { + t.logger.Debugf("Responded to ping request, ID: %v", request.ID) + } +} + // sendErrorResponse sends an error response to the server. func (t *sseClientTransport) sendErrorResponse(request *JSONRPCRequest, code int, message string) { errorResp := newJSONRPCErrorResponse(request.ID, code, message, nil) diff --git a/sse_server.go b/sse_server.go index 080d5b2..66284c0 100644 --- a/sse_server.go +++ b/sse_server.go @@ -168,8 +168,15 @@ type SSEServer struct { contextFunc func(ctx context.Context, r *http.Request) context.Context // HTTP context function. sessionIDGenerator SessionIDGenerator // Custom session ID generator. sessionPubSub SessionPubSub // Optional session Pub/Sub for distributed sessions. - keepAlive bool // Whether to keep the connection alive. - keepAliveInterval time.Duration // Keep-alive interval. + keepAlive bool // Whether to keep the connection alive with SSE comments. + keepAliveInterval time.Duration // Keep-alive interval for SSE comments. + pingEnabled bool // Whether to enable JSON-RPC ping keepalive. + pingInterval time.Duration // Interval for sending ping requests. + pingTimeout time.Duration // Timeout for ping requests. + pingStarted bool // Whether ping loop has been started. + pingStartMu sync.Mutex // Mutex to ensure ping loop starts only once. + stop chan struct{} // Channel to signal server shutdown. + stopOnce sync.Once // Ensures stop channel is closed only once. logger Logger // Logger for this server. requestID atomic.Int64 // Request ID counter for generating unique request IDs. responses map[uint64]interface{} // Map for storing response channels. @@ -233,6 +240,10 @@ func NewSSEServer(name, version string, opts ...SSEOption) *SSEServer { sessionIDGenerator: &defaultSessionIDGenerator{}, // Default session ID generator keepAlive: true, keepAliveInterval: 30 * time.Second, + pingEnabled: false, // Ping keepalive disabled by default + pingInterval: 30 * time.Second, // Default ping interval + pingTimeout: 15 * time.Second, // Default ping timeout + stop: make(chan struct{}), // Stop channel for graceful shutdown logger: GetDefaultLogger(), responses: make(map[uint64]interface{}), notificationHandlers: make(map[string]ServerNotificationHandler), @@ -303,6 +314,32 @@ func WithKeepAliveInterval(interval time.Duration) SSEOption { } } +// WithPingKeepAlive enables or disables JSON-RPC ping keepalive. +// When enabled, the server will periodically send ping requests to all connected clients. +// This provides application-layer health checking in addition to (or instead of) SSE comment keepalive. +func WithPingKeepAlive(enabled bool) SSEOption { + return func(s *SSEServer) { + s.pingEnabled = enabled + } +} + +// WithPingInterval sets the interval for sending ping requests. +// This option does not automatically enable ping keepalive; use WithPingKeepAlive(true) to enable it. +func WithPingInterval(interval time.Duration) SSEOption { + return func(s *SSEServer) { + s.pingInterval = interval + } +} + +// WithPingTimeout sets the timeout for ping requests. +// If a ping request does not receive a response within this timeout, +// it will be considered failed (logged but connection remains open). +func WithPingTimeout(timeout time.Duration) SSEOption { + return func(s *SSEServer) { + s.pingTimeout = timeout + } +} + // WithSSEContextFunc sets a function to modify the context from the request. func WithSSEContextFunc(fn func(ctx context.Context, r *http.Request) context.Context) SSEOption { return func(s *SSEServer) { @@ -456,6 +493,18 @@ func (s *SSEServer) handleSSE(w http.ResponseWriter, r *http.Request) { return } + // Start ping loop if enabled (only once, on first connection) + if s.pingEnabled { + s.pingStartMu.Lock() + if !s.pingStarted { + s.pingStarted = true + s.pingStartMu.Unlock() + s.startPingLoop() + } else { + s.pingStartMu.Unlock() + } + } + // Set SSE headers and immediately flush. w.Header().Set("Content-Type", "text/event-stream") w.Header().Set("Cache-Control", "no-cache") @@ -1641,6 +1690,119 @@ func (s *SSEServer) SendRequest(ctx context.Context, sessionID string, request * } } +// PingSession sends a ping request to the specified session. +// Returns an error if the ping request fails or times out. +// Note: Ping failures are logged but do not cause the session to be disconnected. +func (s *SSEServer) PingSession(ctx context.Context, sessionID string) error { + // Create ping request + requestID := s.requestID.Add(1) + request := &JSONRPCRequest{ + JSONRPC: JSONRPCVersion, + ID: requestID, + Request: Request{ + Method: MethodPing, + }, + } + + // Send request and wait for response + result, err := s.SendRequest(ctx, sessionID, request) + if err != nil { + return fmt.Errorf("ping request failed: %w", err) + } + + // Verify response + if result == nil { + return fmt.Errorf("ping response is nil") + } + + return nil +} + +// PingAllSessions sends ping requests to all connected sessions. +// Ping failures are logged but do not cause sessions to be disconnected. +func (s *SSEServer) PingAllSessions(ctx context.Context) { + var wg sync.WaitGroup + + s.sessions.Range(func(key, value interface{}) bool { + sessionID, ok := key.(string) + if !ok { + s.logger.Warnf("Invalid session key type: %T", key) + return true + } + + wg.Add(1) + go func(sid string) { + defer wg.Done() + + // Create independent timeout context for each ping + pingCtx, cancel := context.WithTimeout(ctx, s.pingTimeout) + defer cancel() + + if err := s.PingSession(pingCtx, sid); err != nil { + // Log failure but do not disconnect session + s.logger.Warnf("Ping session %s failed: %v", sid, err) + } else { + s.logger.Debugf("Ping session %s succeeded", sid) + } + }(sessionID) + + return true + }) + + // Wait for all pings to complete (with overall timeout) + done := make(chan struct{}) + go func() { + wg.Wait() + close(done) + }() + + select { + case <-done: + // All pings completed + case <-ctx.Done(): + s.logger.Warnf("PingAllSessions context cancelled") + } +} + +// startPingLoop starts the ping keepalive loop. +// It sends ping requests to all sessions at the configured interval. +func (s *SSEServer) startPingLoop() { + if s.pingInterval <= 0 { + s.logger.Warnf("Invalid ping interval: %v, ping keepalive disabled", s.pingInterval) + return + } + + go func() { + s.logger.Infof("Ping keepalive started, interval: %v, timeout: %v", s.pingInterval, s.pingTimeout) + + ticker := time.NewTicker(s.pingInterval) + defer ticker.Stop() + + for { + select { + case <-ticker.C: + // Create context with timeout for entire ping cycle + ctx, cancel := context.WithTimeout(context.Background(), s.pingInterval) + s.PingAllSessions(ctx) + cancel() + + case <-s.stop: + s.logger.Infof("Ping keepalive stopped") + return + } + } + }() +} + +// Close gracefully closes the SSE server. +// It stops the ping keepalive loop and performs cleanup. +func (s *SSEServer) Close() error { + s.stopOnce.Do(func() { + close(s.stop) + }) + return nil +} + // formatSSEEvent formats SSE event. func formatSSEEvent(eventType string, data []byte) string { var builder strings.Builder diff --git a/sse_server_test.go b/sse_server_test.go index 00684c5..6408c3e 100644 --- a/sse_server_test.go +++ b/sse_server_test.go @@ -12,6 +12,7 @@ import ( "net/http" "strings" "testing" + "time" ) func TestSSEServer_UnregisterTools(t *testing.T) { @@ -346,14 +347,14 @@ func (g *testSessionIDGenerator) GenerateSessionID(r *http.Request) string { // mockSessionPubSub is a mock implementation of SessionPubSub for testing. type mockSessionPubSub struct { - subscriptions map[string]SessionMessageHandler - published []mockPublishedMessage - subscribeCalls int + subscriptions map[string]SessionMessageHandler + published []mockPublishedMessage + subscribeCalls int unsubscribeCalls int - publishCalls int - subscribeErr error - unsubscribeErr error - publishErr error + publishCalls int + subscribeErr error + unsubscribeErr error + publishErr error } type mockPublishedMessage struct { @@ -797,3 +798,132 @@ func TestSerializedRequestWithContextFunc(t *testing.T) { t.Errorf("Expected remote addr '10.0.0.1:54321', got '%s'", capturedRemoteAddr) } } + +// TestDefaultPingConfiguration tests that ping keepalive is disabled by default +func TestDefaultPingConfiguration(t *testing.T) { + server := NewSSEServer("test-server", "1.0.0") + + // Verify default configuration + if server.keepAlive != true { + t.Errorf("Expected keepAlive to be true by default, got %v", server.keepAlive) + } + if server.keepAliveInterval != 30*time.Second { + t.Errorf("Expected keepAliveInterval to be 30s, got %v", server.keepAliveInterval) + } + if server.pingEnabled != false { + t.Errorf("Expected pingEnabled to be false by default, got %v", server.pingEnabled) + } + if server.pingInterval != 30*time.Second { + t.Errorf("Expected pingInterval to be 30s, got %v", server.pingInterval) + } + if server.pingTimeout != 15*time.Second { + t.Errorf("Expected pingTimeout to be 15s, got %v", server.pingTimeout) + } +} + +// TestPingKeepAliveEnabled tests enabling ping keepalive +func TestPingKeepAliveEnabled(t *testing.T) { + server := NewSSEServer("test-server", "1.0.0", + WithPingKeepAlive(true), + ) + + // Verify ping is enabled + if !server.pingEnabled { + t.Error("Expected pingEnabled to be true") + } + // Verify comment keepalive is still enabled (default) + if !server.keepAlive { + t.Error("Expected keepAlive to still be true") + } +} + +// TestPingOnlyMode tests using only ping keepalive (disabling comment) +func TestPingOnlyMode(t *testing.T) { + server := NewSSEServer("test-server", "1.0.0", + WithKeepAlive(false), + WithPingKeepAlive(true), + ) + + // Verify only ping is enabled + if server.keepAlive { + t.Error("Expected keepAlive to be false") + } + if !server.pingEnabled { + t.Error("Expected pingEnabled to be true") + } +} + +// TestCustomPingConfiguration tests custom ping configuration +func TestCustomPingConfiguration(t *testing.T) { + server := NewSSEServer("test-server", "1.0.0", + WithPingKeepAlive(true), + WithPingInterval(60*time.Second), + WithPingTimeout(30*time.Second), + ) + + // Verify custom configuration + if !server.pingEnabled { + t.Error("Expected pingEnabled to be true") + } + if server.pingInterval != 60*time.Second { + t.Errorf("Expected pingInterval to be 60s, got %v", server.pingInterval) + } + if server.pingTimeout != 30*time.Second { + t.Errorf("Expected pingTimeout to be 30s, got %v", server.pingTimeout) + } +} + +// TestExistingConfigNotAffected tests that existing keepalive config is not affected by new ping options +func TestExistingConfigNotAffected(t *testing.T) { + // Test with explicit keepalive disabled + server1 := NewSSEServer("test-server", "1.0.0", + WithKeepAlive(false), + ) + if server1.keepAlive { + t.Error("Expected keepAlive to be false") + } + if server1.pingEnabled { + t.Error("Expected pingEnabled to be false (default)") + } + + // Test with keepalive interval (which auto-enables keepalive - existing behavior) + server2 := NewSSEServer("test-server", "1.0.0", + WithKeepAliveInterval(60*time.Second), + ) + if !server2.keepAlive { + t.Error("Expected keepAlive to be true (auto-enabled by WithKeepAliveInterval)") + } + if server2.keepAliveInterval != 60*time.Second { + t.Errorf("Expected keepAliveInterval to be 60s, got %v", server2.keepAliveInterval) + } + if server2.pingEnabled { + t.Error("Expected pingEnabled to be false (default)") + } +} + +// TestSSEServerBackwardCompatibility tests that existing SSE server code works without modification +func TestSSEServerBackwardCompatibility(t *testing.T) { + // Scenario 1: Server created with no options + server1 := NewSSEServer("test-server", "1.0.0") + if !server1.keepAlive { + t.Error("Expected keepAlive to be true (default)") + } + if server1.pingEnabled { + t.Error("Expected pingEnabled to be false (default)") + } + + // Scenario 2: Server with existing keepalive options + server2 := NewSSEServer("test-server", "1.0.0", + WithKeepAlive(true), + WithKeepAliveInterval(60*time.Second), + ) + if !server2.keepAlive { + t.Error("Expected keepAlive to be true") + } + if server2.keepAliveInterval != 60*time.Second { + t.Errorf("Expected keepAliveInterval to be 60s, got %v", server2.keepAliveInterval) + } + if server2.pingEnabled { + t.Error("Expected pingEnabled to be false (default, not affected by keepAlive options)") + } +} diff --git a/streamable_server.go b/streamable_server.go index 3344709..feae619 100644 --- a/streamable_server.go +++ b/streamable_server.go @@ -73,6 +73,17 @@ type httpServerHandler struct { // Response manager for server-to-client requests. responseManager *responseManager + + // Keepalive configuration + keepAliveEnabled bool // Whether SSE comment keepalive is enabled + keepAliveInterval time.Duration // Interval for SSE comment keepalive + pingEnabled bool // Whether JSON-RPC ping keepalive is enabled + pingInterval time.Duration // Interval for ping requests + pingTimeout time.Duration // Timeout for ping requests + pingStarted bool // Whether ping loop has been started + pingStartMu sync.Mutex // Mutex to ensure ping loop starts only once + stop chan struct{} // Channel to signal shutdown + stopOnce sync.Once // Ensures stop channel is closed only once } // getSSEConnection represents a GET SSE connection @@ -103,6 +114,12 @@ func newHTTPServerHandler(handler requestHandler, serverPath string, options ... getSSEConnections: make(map[string]*getSSEConnection), serverPath: serverPath, responseManager: newResponseManager(), + keepAliveEnabled: true, // Default: SSE comment keepalive enabled + keepAliveInterval: 30 * time.Second, // Default: 30 seconds + pingEnabled: false, // Default: ping keepalive disabled + pingInterval: 30 * time.Second, // Default: 30 seconds + pingTimeout: 15 * time.Second, // Default: 15 seconds + stop: make(chan struct{}), // Stop channel for graceful shutdown } // Apply options @@ -172,6 +189,23 @@ func withTransportNotificationBufferSize(size int) func(*httpServerHandler) { } } +// withKeepAliveConfig sets the SSE comment keepalive configuration +func withKeepAliveConfig(enabled bool, interval time.Duration) func(*httpServerHandler) { + return func(h *httpServerHandler) { + h.keepAliveEnabled = enabled + h.keepAliveInterval = interval + } +} + +// withPingConfig sets the JSON-RPC ping keepalive configuration +func withPingConfig(enabled bool, interval, timeout time.Duration) func(*httpServerHandler) { + return func(h *httpServerHandler) { + h.pingEnabled = enabled + h.pingInterval = interval + h.pingTimeout = timeout + } +} + // withTransportStatelessMode sets the server to stateless mode // In stateless mode, the server does not generate persistent session IDs; each request uses a temporary session func withTransportStatelessMode() func(*httpServerHandler) { @@ -589,6 +623,18 @@ func (h *httpServerHandler) handleGet(ctx context.Context, w http.ResponseWriter w.WriteHeader(http.StatusOK) flusher.Flush() + // Start ping loop if enabled (only once, on first GET SSE connection) + if h.pingEnabled { + h.pingStartMu.Lock() + if !h.pingStarted { + h.pingStarted = true + h.pingStartMu.Unlock() + h.startPingLoop() + } else { + h.pingStartMu.Unlock() + } + } + // Create context, for canceling connection connCtx, cancelConn := context.WithCancel(ctx) localCancelFunc = cancelConn // Assign to the variable captured by defer @@ -617,6 +663,11 @@ func (h *httpServerHandler) handleGet(ctx context.Context, w http.ResponseWriter // Record connection information h.logger.Infof("Established GET SSE connection, session ID: %s", session.GetID()) + // Start comment keepalive if enabled + if h.keepAliveEnabled { + go h.handleGetSSECommentKeepAlive(connCtx, conn, session.GetID()) + } + // If there's Last-Event-ID, try to resume stream if lastEventID != "" { h.handleStreamResumption(connCtx, conn, session.GetID()) @@ -632,6 +683,38 @@ func (h *httpServerHandler) handleGet(ctx context.Context, w http.ResponseWriter h.logger.Infof("GET SSE connection closed, session ID: %s", session.GetID()) } +// handleGetSSECommentKeepAlive handles SSE comment keepalive for GET SSE connections. +// It sends SSE comment lines at the configured interval to keep the connection alive. +func (h *httpServerHandler) handleGetSSECommentKeepAlive(ctx context.Context, conn *getSSEConnection, sessionID string) { + if h.keepAliveInterval <= 0 { + return + } + + ticker := time.NewTicker(h.keepAliveInterval) + defer ticker.Stop() + + for { + select { + case <-ticker.C: + conn.writeLock.Lock() + // Send SSE comment + fmt.Fprintf(conn.writer, ": keepalive\n\n") + conn.flusher.Flush() + conn.writeLock.Unlock() + + if h.logger != nil { + h.logger.Debugf("Sent keepalive comment to GET SSE session: %s", sessionID) + } + + case <-ctx.Done(): + if h.logger != nil { + h.logger.Debugf("GET SSE keepalive stopped for session: %s", sessionID) + } + return + } + } +} + // Send notification through GET SSE func (h *httpServerHandler) sendNotificationToGetSSE(sessionID string, notification *JSONRPCNotification) error { h.getSSEConnectionsLock.RLock() @@ -790,6 +873,134 @@ func (h *httpServerHandler) cleanupSession(sessionID string) { h.getSSEConnectionsLock.Unlock() } +// PingSession sends a ping request to the specified session via GET SSE connection. +// Returns an error if the session has no GET SSE connection or the ping fails. +func (h *httpServerHandler) PingSession(ctx context.Context, sessionID string) error { + // Create ping request + requestID := h.responseManager.GenerateRequestID() + request := &JSONRPCRequest{ + JSONRPC: JSONRPCVersion, + ID: requestID, + Request: Request{ + Method: MethodPing, + }, + } + + // Send request and wait for response + result, err := h.SendRequest(ctx, sessionID, request) + if err != nil { + return fmt.Errorf("ping request failed: %w", err) + } + + // Verify response + if result == nil { + return fmt.Errorf("ping response is nil") + } + + return nil +} + +// PingAllSessions sends ping requests to all sessions that have GET SSE connections. +// Ping failures are logged but do not cause sessions to be disconnected. +func (h *httpServerHandler) PingAllSessions(ctx context.Context) { + var sessionIDs []string + + // Collect all session IDs with GET SSE connections + h.getSSEConnectionsLock.RLock() + for sessionID := range h.getSSEConnections { + sessionIDs = append(sessionIDs, sessionID) + } + h.getSSEConnectionsLock.RUnlock() + + if len(sessionIDs) == 0 { + return + } + + var wg sync.WaitGroup + for _, sessionID := range sessionIDs { + wg.Add(1) + go func(sid string) { + defer wg.Done() + + // Create independent timeout context for each ping + pingCtx, cancel := context.WithTimeout(ctx, h.pingTimeout) + defer cancel() + + if err := h.PingSession(pingCtx, sid); err != nil { + // Log failure but do not disconnect session + if h.logger != nil { + h.logger.Warnf("Ping Streamable HTTP session %s failed: %v", sid, err) + } + } else { + if h.logger != nil { + h.logger.Debugf("Ping Streamable HTTP session %s succeeded", sid) + } + } + }(sessionID) + } + + // Wait for all pings to complete + done := make(chan struct{}) + go func() { + wg.Wait() + close(done) + }() + + select { + case <-done: + // All pings completed + case <-ctx.Done(): + if h.logger != nil { + h.logger.Warnf("PingAllSessions context cancelled") + } + } +} + +// startPingLoop starts the ping keepalive loop for Streamable HTTP. +// It sends ping requests to all sessions with GET SSE connections at the configured interval. +func (h *httpServerHandler) startPingLoop() { + if h.pingInterval <= 0 { + if h.logger != nil { + h.logger.Warnf("Invalid ping interval: %v, ping keepalive disabled", h.pingInterval) + } + return + } + + go func() { + if h.logger != nil { + h.logger.Infof("Streamable HTTP ping keepalive started, interval: %v, timeout: %v", h.pingInterval, h.pingTimeout) + } + + ticker := time.NewTicker(h.pingInterval) + defer ticker.Stop() + + for { + select { + case <-ticker.C: + // Create context with timeout for entire ping cycle + ctx, cancel := context.WithTimeout(context.Background(), h.pingInterval) + h.PingAllSessions(ctx) + cancel() + + case <-h.stop: + if h.logger != nil { + h.logger.Infof("Streamable HTTP ping keepalive stopped") + } + return + } + } + }() +} + +// Close gracefully closes the Streamable HTTP handler. +// It stops the ping keepalive loop and performs cleanup. +func (h *httpServerHandler) Close() error { + h.stopOnce.Do(func() { + close(h.stop) + }) + return nil +} + // isValidPath validates if the request path matches the configured server path. func (h *httpServerHandler) isValidPath(requestPath string) bool { if h.serverPath == "" {