feat: add --session-idle-timeout-minutes CLI flag (#691)

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: tehranian

README.md

@@ -335,6 +335,9 @@ The `mcp-grafana` binary supports various command-line flags for configuration:
 - `--metrics`: Enable Prometheus metrics endpoint at `/metrics`
 - `--metrics-address`: Separate address for metrics server (e.g., `:9090`). If empty, metrics are served on the main server
 
+**Session Management:**
+- `--session-idle-timeout-minutes`: Session idle timeout in minutes. Sessions with no activity for this duration are automatically reaped - default: `30`. Set to `0` to disable session reaping. Only relevant for SSE and streamable-http transports.
+
 **Tool Configuration:**
 - `--enabled-tools`: Comma-separated list of enabled categories - default: all categories except `admin`, to enable admin tools, add `admin` to the list (e.g., `"search,datasource,...,admin"`)
 - `--max-loki-log-limit`: Maximum number of log lines returned per `query_loki_logs` call - default: `100`. Note: Set this at least 1 below Loki's server-side `max_entries_limit_per_query` to allow truncation detection (the tool requests `limit+1` internally to detect if more data exists).

cmd/mcp-grafana/main.go

@@ -131,8 +131,10 @@ func (dt *disabledTools) addTools(s *server.MCPServer) {
 	maybeAddTools(s, tools.AddRunPanelQueryTools, enabledTools, dt.runpanelquery, "runpanelquery")
 }
 
-func newServer(transport string, dt disabledTools, obs *observability.Observability) (*server.MCPServer, *mcpgrafana.ToolManager, *mcpgrafana.SessionManager) {
-	sm := mcpgrafana.NewSessionManager()
+func newServer(transport string, dt disabledTools, obs *observability.Observability, sessionIdleTimeoutMinutes int) (*server.MCPServer, *mcpgrafana.ToolManager, *mcpgrafana.SessionManager) {
+	sm := mcpgrafana.NewSessionManager(
+		mcpgrafana.WithSessionTTL(time.Duration(sessionIdleTimeoutMinutes) * time.Minute),
+	)
 
 	// Declare variable for ToolManager that will be initialized after server creation
 	var stm *mcpgrafana.ToolManager
@@ -277,7 +279,7 @@ func runMetricsServer(addr string, o *observability.Observability) {
 	}
 }
 
-func run(transport, addr, basePath, endpointPath string, logLevel slog.Level, dt disabledTools, gc mcpgrafana.GrafanaConfig, tls tlsConfig, obs observability.Config) error {
+func run(transport, addr, basePath, endpointPath string, logLevel slog.Level, dt disabledTools, gc mcpgrafana.GrafanaConfig, tls tlsConfig, obs observability.Config, sessionIdleTimeoutMinutes int) error {
 	slog.SetDefault(slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: logLevel})))
 
 	// Set up observability (metrics and tracing)
@@ -301,7 +303,7 @@ func run(transport, addr, basePath, endpointPath string, logLevel slog.Level, dt
 		defer clientCache.Close()
 	}
 
-	s, tm, sm := newServer(transport, dt, o)
+	s, tm, sm := newServer(transport, dt, o, sessionIdleTimeoutMinutes)
 	defer sm.Close()
 
 	// Create a context that will be cancelled on shutdown
@@ -416,6 +418,7 @@ func main() {
 	basePath := flag.String("base-path", "", "Base path for the sse server")
 	endpointPath := flag.String("endpoint-path", "/mcp", "Endpoint path for the streamable-http server")
 	logLevel := flag.String("log-level", "info", "Log level (debug, info, warn, error)")
+	sessionIdleTimeoutMinutes := flag.Int("session-idle-timeout-minutes", 30, "Session idle timeout in minutes. Sessions with no activity for this duration are automatically reaped. Set to 0 to disable session reaping")
 	showVersion := flag.Bool("version", false, "Print the version and exit")
 	var dt disabledTools
 	dt.addFlags()
@@ -459,7 +462,7 @@ func main() {
 		obs.NetworkTransport = mcpconv.NetworkTransportTCP
 	}
 
-	if err := run(transport, *addr, *basePath, *endpointPath, parseLevel(*logLevel), dt, grafanaConfig, tls, obs); err != nil {
+	if err := run(transport, *addr, *basePath, *endpointPath, parseLevel(*logLevel), dt, grafanaConfig, tls, obs, *sessionIdleTimeoutMinutes); err != nil {
 		panic(err)
 	}
 }

cmd/mcp-grafana/main_test.go

@@ -0,0 +1,71 @@
+package main
+
+import (
+	"context"
+	"testing"
+	"testing/synctest"
+	"time"
+
+	"github.com/mark3labs/mcp-go/mcp"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/grafana/mcp-grafana/observability"
+)
+
+// testClientSession implements server.ClientSession for unit tests.
+type testClientSession struct {
+	id string
+}
+
+func (s *testClientSession) SessionID() string                                   { return s.id }
+func (s *testClientSession) NotificationChannel() chan<- mcp.JSONRPCNotification { return nil }
+func (s *testClientSession) Initialize()                                         {}
+func (s *testClientSession) Initialized() bool                                   { return true }
+
+func newTestObservability(t *testing.T) *observability.Observability {
+	t.Helper()
+	obs, err := observability.Setup(observability.Config{})
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		_ = obs.Shutdown(context.Background())
+	})
+	return obs
+}
+
+func TestNewServer_SessionIdleTimeoutZeroDisablesReaping(t *testing.T) {
+	obs := newTestObservability(t)
+	synctest.Test(t, func(t *testing.T) {
+		_, _, sm := newServer("stdio", disabledTools{enabledTools: "search"}, obs, 0)
+		defer sm.Close()
+
+		session := &testClientSession{id: "should-persist"}
+		sm.CreateSession(context.Background(), session)
+
+		// Advance the fake clock well beyond any reasonable reaper interval.
+		// With reaper disabled (TTL=0), the session must survive.
+		time.Sleep(time.Hour)
+
+		_, exists := sm.GetSession("should-persist")
+		assert.True(t, exists, "Session should persist when idle timeout is 0 (reaper disabled)")
+	})
+}
+
+func TestNewServer_SessionIdleTimeoutCustomValue(t *testing.T) {
+	obs := newTestObservability(t)
+	synctest.Test(t, func(t *testing.T) {
+		_, _, sm := newServer("stdio", disabledTools{enabledTools: "search"}, obs, 1)
+		defer sm.Close()
+
+		session := &testClientSession{id: "custom-ttl"}
+		sm.CreateSession(context.Background(), session)
+
+		// Advance the fake clock past the 1-minute TTL.
+		// The reaper runs every TTL/2 (30s), so by 2 minutes
+		// it will have fired and reaped the idle session.
+		time.Sleep(2 * time.Minute)
+
+		_, exists := sm.GetSession("custom-ttl")
+		assert.False(t, exists, "Session should be reaped after exceeding the 1-minute idle timeout")
+	})
+}