fix: release MCP session state so the server stops leaking heap

Makefile

@@ -58,11 +58,11 @@ tidy: ## Run go mod tidy to ensure dependencies are up to date.
 
 .PHONY: test
 test: build lint ## Run all tests with build, lint, and coverage
-	go test -tags=test -v -cover ./pkg/... ./internal/...
+	go test -tags=test -v -cover ./cmd/... ./pkg/... ./internal/...
 
 .PHONY: test-only
 test-only: ## Run tests only (without build/lint for faster iteration)
-	go test -tags=test -v -cover ./pkg/... ./internal/...
+	go test -tags=test -v -cover ./cmd/... ./pkg/... ./internal/...
 
 .PHONY: e2e
 e2e: test retag

README.md

@@ -197,6 +197,7 @@ The server runs using sse transport for MCP communication.
 | `--stdio` | `false` | Use stdio for communication instead of HTTP |
 | `--tools` | `[]` (all) | Comma-separated list of tool providers to register |
 | `--read-only` | `false` | Disable tools that perform write operations |
+| `--session-idle-ttl` | `10m` | Reclaim streamable HTTP session state after this idle duration (`0` disables the sweeper) |
 | `--kubeconfig` | `""` | Path to kubeconfig file (defaults to in-cluster config) |
 | `--version`, `-v` | `false` | Show version information and exit |
 

cmd/main.go

@@ -46,6 +46,8 @@ var (
 	showVersion bool
 	readOnly    bool
 
+	sessionIdleTTL time.Duration
+
 	// These variables should be set during build time using -ldflags
 	Name      = "kagent-tools-server"
 	Version   = version.Version
@@ -66,6 +68,7 @@ func init() {
 	rootCmd.Flags().StringSliceVar(&tools, "tools", []string{}, "List of tools to register. If empty, all tools are registered.")
 	rootCmd.Flags().BoolVarP(&showVersion, "version", "v", false, "Show version information and exit")
 	rootCmd.Flags().BoolVar(&readOnly, "read-only", false, "Run in read-only mode (disable tools that perform write operations)")
+	rootCmd.Flags().DurationVar(&sessionIdleTTL, "session-idle-ttl", 10*time.Minute, "Reclaim streamable HTTP session state after this idle duration (0 disables the sweeper)")
 	kubeconfig = rootCmd.Flags().String("kubeconfig", "", "kubeconfig file path (optional, defaults to in-cluster config)")
 
 	// if found .env file, load it
@@ -160,7 +163,8 @@ func run(cmd *cobra.Command, args []string) {
 
 	// HTTP server reference (only used when not in stdio mode)
 	var httpServer *http.Server
-	var metricsServer *http.Server // Separate server for metrics if metricsPort is different from main port
+	var metricsServer *http.Server                    // Separate server for metrics if metricsPort is different from main port
+	var streamableServer *server.StreamableHTTPServer // Streamable HTTP transport, shut down to stop its session sweeper
 
 	// Start server based on chosen mode
 	wg.Add(1)
@@ -170,9 +174,7 @@ func run(cmd *cobra.Command, args []string) {
 			runStdioServer(ctx, mcp)
 		}()
 	} else {
-		sseServer := server.NewStreamableHTTPServer(mcp,
-			server.WithHeartbeatInterval(30*time.Second),
-		)
+		streamableServer = newStreamableHTTPServer(mcp, sessionIdleTTL)
 
 		// Create a mux to handle different routes
 		mux := http.NewServeMux()
@@ -198,9 +200,7 @@ func run(cmd *cobra.Command, args []string) {
 				Handler: metricsMux,
 			}
 
-			wg.Add(1)
-			go func() {
-				defer wg.Done()
+			wg.Go(func() {
 				logger.Get().Info("Starting Prometheus metrics endpoint on /metrics", "port", strconv.Itoa(metricsPort))
 				if err := metricsServer.ListenAndServe(); err != nil {
 					if !errors.Is(err, http.ErrServerClosed) {
@@ -209,7 +209,7 @@ func run(cmd *cobra.Command, args []string) {
 						logger.Get().Info("Metrics server closed gracefully.")
 					}
 				}
-			}()
+			})
 		} else {
 			logger.Get().Info("Starting Prometheus metrics endpoint on /metrics", "port", strconv.Itoa(port))
 			mux.Handle("/metrics", promhttp.HandlerFor(registry, promhttp.HandlerOpts{}))
@@ -222,7 +222,7 @@ func run(cmd *cobra.Command, args []string) {
 
 		// Handle all other routes with the MCP server wrapped in telemetry middleware
 		mux.Handle("/", telemetry.HTTPMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			sseServer.ServeHTTP(w, r)
+			streamableServer.ServeHTTP(w, r)
 		})))
 
 		httpServer = &http.Server{
@@ -232,7 +232,7 @@ func run(cmd *cobra.Command, args []string) {
 
 		go func() {
 			defer wg.Done()
-			logger.Get().Info("Running KAgent Tools Server", "port", fmt.Sprintf(":%d", port), "tools", strings.Join(tools, ","))
+			logger.Get().Info("Running KAgent Tools Server", "port", fmt.Sprintf(":%d", port), "tools", strings.Join(tools, ","), "session_idle_ttl", sessionIdleTTL.String())
 			if err := httpServer.ListenAndServe(); err != nil {
 				if !errors.Is(err, http.ErrServerClosed) {
 					logger.Get().Error("Failed to start HTTP server", "error", err)
@@ -268,6 +268,18 @@ func run(cmd *cobra.Command, args []string) {
 			}
 		}
 
+		// Stop the session sweeper and close sessions still registered with the
+		// MCP server, so nothing outlives the transport it belonged to
+		if !stdio && streamableServer != nil {
+			shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 5*time.Second)
+			defer shutdownCancel()
+
+			if err := streamableServer.Shutdown(shutdownCtx); err != nil {
+				logger.Get().Error("Failed to shutdown MCP streamable HTTP server gracefully", "error", err)
+				rootSpan.RecordError(err)
+			}
+		}
+
 		// Gracefully shutdown metrics server if running separately
 		if !stdio && metricsServer != nil {
 			shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 5*time.Second)
@@ -287,6 +299,19 @@ func run(cmd *cobra.Command, args []string) {
 	logger.Get().Info("Server shutdown complete")
 }
 
+// newStreamableHTTPServer builds the streamable HTTP transport.
+//
+// A session is registered on `initialize` and released only when the client
+// sends DELETE, so a client that goes away without one leaves its session
+// behind forever. idleTTL bounds that by reclaiming sessions that have seen no
+// traffic for that long; zero or less disables the sweeper.
+func newStreamableHTTPServer(mcpServer *server.MCPServer, idleTTL time.Duration) *server.StreamableHTTPServer {
+	return server.NewStreamableHTTPServer(mcpServer,
+		server.WithHeartbeatInterval(30*time.Second),
+		server.WithSessionIdleTTL(idleTTL),
+	)
+}
+
 // writeResponse writes data to an HTTP response writer with proper error handling
 func writeResponse(w http.ResponseWriter, data []byte) error {
 	_, err := w.Write(data)

cmd/streamable_http_test.go

@@ -0,0 +1,151 @@
+package main
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/mark3labs/mcp-go/server"
+)
+
+const initializeRequest = `{"jsonrpc":"2.0","id":1,"method":"initialize",` +
+	`"params":{"protocolVersion":"2025-03-26","capabilities":{},` +
+	`"clientInfo":{"name":"test","version":"1"}}}`
+
+// sessionRecorder counts the sessions the MCP server registers and releases.
+// A session that is registered but never unregistered is retained state, which
+// is what leaks when a client never sends DELETE.
+type sessionRecorder struct {
+	mu           sync.Mutex
+	registered   int
+	unregistered int
+}
+
+func (r *sessionRecorder) hooks() *server.Hooks {
+	hooks := &server.Hooks{}
+	hooks.AddOnRegisterSession(func(_ context.Context, _ server.ClientSession) {
+		r.mu.Lock()
+		defer r.mu.Unlock()
+		r.registered++
+	})
+	hooks.AddOnUnregisterSession(func(_ context.Context, _ server.ClientSession) {
+		r.mu.Lock()
+		defer r.mu.Unlock()
+		r.unregistered++
+	})
+	return hooks
+}
+
+func (r *sessionRecorder) counts() (int, int) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	return r.registered, r.unregistered
+}
+
+// newRecordedServer starts the streamable HTTP transport wired exactly as the
+// server wires it, with session hooks attached.
+func newRecordedServer(t *testing.T, idleTTL time.Duration) (*httptest.Server, *sessionRecorder) {
+	t.Helper()
+
+	recorder := &sessionRecorder{}
+	mcpServer := server.NewMCPServer("test-server", "test", server.WithHooks(recorder.hooks()))
+	streamableServer := newStreamableHTTPServer(mcpServer, idleTTL)
+	t.Cleanup(func() {
+		shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+		_ = streamableServer.Shutdown(shutdownCtx)
+	})
+
+	httpServer := httptest.NewServer(streamableServer)
+	t.Cleanup(httpServer.Close)
+	return httpServer, recorder
+}
+
+// initializeSession sends an initialize request and returns the session ID the
+// server assigned to it.
+func initializeSession(t *testing.T, httpServer *httptest.Server) string {
+	t.Helper()
+
+	req, err := http.NewRequest(http.MethodPost, httpServer.URL, strings.NewReader(initializeRequest))
+	if err != nil {
+		t.Fatalf("build initialize request: %v", err)
+	}
+	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("Accept", "application/json, text/event-stream")
+
+	resp, err := httpServer.Client().Do(req)
+	if err != nil {
+		t.Fatalf("send initialize request: %v", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		t.Fatalf("initialize returned status %d, want %d", resp.StatusCode, http.StatusOK)
+	}
+	sessionID := resp.Header.Get("Mcp-Session-Id")
+	if sessionID == "" {
+		t.Fatal("initialize response carried no Mcp-Session-Id header")
+	}
+	return sessionID
+}
+
+// waitForUnregistered polls until the expected number of sessions has been
+// released, so the sweeper's own tick interval does not make the test flaky.
+func waitForUnregistered(t *testing.T, recorder *sessionRecorder, want int, timeout time.Duration) {
+	t.Helper()
+
+	deadline := time.Now().Add(timeout)
+	for {
+		_, unregistered := recorder.counts()
+		if unregistered >= want {
+			return
+		}
+		if time.Now().After(deadline) {
+			t.Fatalf("released %d sessions after %s, want %d", unregistered, timeout, want)
+		}
+		time.Sleep(20 * time.Millisecond)
+	}
+}
+
+// A client that ends its session explicitly must have it released.
+func TestStreamableHTTPServerReleasesSessionOnDelete(t *testing.T) {
+	httpServer, recorder := newRecordedServer(t, 0)
+	sessionID := initializeSession(t, httpServer)
+
+	if registered, _ := recorder.counts(); registered != 1 {
+		t.Fatalf("registered %d sessions, want 1", registered)
+	}
+
+	req, err := http.NewRequest(http.MethodDelete, httpServer.URL, nil)
+	if err != nil {
+		t.Fatalf("build delete request: %v", err)
+	}
+	req.Header.Set("Mcp-Session-Id", sessionID)
+
+	resp, err := httpServer.Client().Do(req)
+	if err != nil {
+		t.Fatalf("send delete request: %v", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		t.Fatalf("delete returned status %d, want %d", resp.StatusCode, http.StatusOK)
+	}
+	waitForUnregistered(t, recorder, 1, 2*time.Second)
+}
+
+// A client that goes away without a DELETE must not retain its session: the
+// idle sweeper is what bounds memory for POST-only clients.
+func TestStreamableHTTPServerSweepsIdleSession(t *testing.T) {
+	httpServer, recorder := newRecordedServer(t, 100*time.Millisecond)
+	initializeSession(t, httpServer)
+
+	if registered, _ := recorder.counts(); registered != 1 {
+		t.Fatalf("registered %d sessions, want 1", registered)
+	}
+	waitForUnregistered(t, recorder, 1, 10*time.Second)
+}

go.mod

@@ -6,7 +6,7 @@ require (
 	github.com/joho/godotenv v1.5.1
 	github.com/kubescape/k8s-interface v0.0.203
 	github.com/kubescape/storage v0.0.239
-	github.com/mark3labs/mcp-go v0.43.2
+	github.com/mark3labs/mcp-go v1.1.0
 	github.com/onsi/ginkgo/v2 v2.27.2
 	github.com/onsi/gomega v1.38.2
 	github.com/prometheus/client_golang v1.23.2
@@ -38,13 +38,11 @@ require (
 	github.com/armosec/gojay v1.2.17 // indirect
 	github.com/armosec/utils-go v0.0.58 // indirect
 	github.com/armosec/utils-k8s-go v0.0.35 // indirect
-	github.com/bahlo/generic-list-go v0.2.0 // indirect
 	github.com/becheran/wildmatch-go v1.0.0 // indirect
 	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/blang/semver/v4 v4.0.0 // indirect
 	github.com/bmatcuk/doublestar/v4 v4.9.1 // indirect
 	github.com/briandowns/spinner v1.23.2 // indirect
-	github.com/buger/jsonparser v1.1.1 // indirect
 	github.com/cenkalti/backoff/v4 v4.3.0 // indirect
 	github.com/cenkalti/backoff/v5 v5.0.3 // indirect
 	github.com/cespare/xxhash/v2 v2.3.0 // indirect
@@ -99,21 +97,20 @@ require (
 	github.com/google/gnostic-models v0.7.1 // indirect
 	github.com/google/go-cmp v0.7.0 // indirect
 	github.com/google/go-containerregistry v0.20.6 // indirect
+	github.com/google/jsonschema-go v0.4.2 // indirect
 	github.com/google/licensecheck v0.3.1 // indirect
 	github.com/google/pprof v0.0.0-20251114195745-4902fdda35c8 // indirect
 	github.com/google/uuid v1.6.0 // indirect
 	github.com/grpc-ecosystem/grpc-gateway/v2 v2.27.8 // indirect
 	github.com/hashicorp/errwrap v1.1.0 // indirect
 	github.com/hashicorp/go-multierror v1.1.1 // indirect
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
-	github.com/invopop/jsonschema v0.13.0 // indirect
 	github.com/jinzhu/copier v0.4.0 // indirect
 	github.com/json-iterator/go v1.1.12 // indirect
 	github.com/klauspost/compress v1.18.4 // indirect
 	github.com/kubescape/go-logger v0.0.26 // indirect
 	github.com/kylelemons/godebug v1.1.0 // indirect
 	github.com/mackerelio/go-osstat v0.2.6 // indirect
-	github.com/mailru/easyjson v0.9.1 // indirect
 	github.com/mattn/go-colorable v0.1.14 // indirect
 	github.com/mattn/go-isatty v0.0.20 // indirect
 	github.com/mitchellh/go-homedir v1.1.0 // indirect
@@ -135,6 +132,7 @@ require (
 	github.com/prometheus/common v0.67.5 // indirect
 	github.com/prometheus/procfs v0.19.2 // indirect
 	github.com/sagikazarmark/locafero v0.12.0 // indirect
+	github.com/santhosh-tekuri/jsonschema/v6 v6.0.2 // indirect
 	github.com/sasha-s/go-deadlock v0.3.6 // indirect
 	github.com/scylladb/go-set v1.0.3-0.20200225121959-cc7b2070d91e // indirect
 	github.com/seccomp/libseccomp-golang v0.10.0 // indirect
@@ -155,7 +153,6 @@ require (
 	github.com/vishvananda/netns v0.0.5 // indirect
 	github.com/wagoodman/go-partybus v0.0.0-20230516145632-8ccac152c651 // indirect
 	github.com/wagoodman/go-progress v0.0.0-20230925121702-07e42b3cdba0 // indirect
-	github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect
 	github.com/x448/float16 v0.8.4 // indirect
 	github.com/yl2chen/cidranger v1.0.2 // indirect
 	github.com/yosida95/uritemplate/v3 v3.0.2 // indirect