Fix resource leak in session storage cleanup

Implements fix by calling Close() on sessions that implement
io.Closer before deleting them from storage during cleanup.
This prevents connection leaks when vMCP sessions with backend clients
expire.

Closes: #3871

Author: taskbot

pkg/transport/session/storage_local.go

@@ -6,6 +6,8 @@ package session
 import (
 	"context"
 	"fmt"
+	"io"
+	"log/slog"
 	"sync"
 	"time"
 )
@@ -66,9 +68,12 @@ func (s *LocalStorage) Delete(_ context.Context, id string) error {
 
 // DeleteExpired removes all sessions that haven't been updated since the given time.
 func (s *LocalStorage) DeleteExpired(ctx context.Context, before time.Time) error {
-	var toDelete []string
+	var toDelete []struct {
+		id      string
+		session Session
+	}
 
-	// First pass: collect IDs of expired sessions
+	// First pass: collect expired sessions
 	s.sessions.Range(func(key, val any) bool {
 		// Check for context cancellation
 		select {
@@ -80,16 +85,44 @@ func (s *LocalStorage) DeleteExpired(ctx context.Context, before time.Time) erro
 		if session, ok := val.(Session); ok {
 			if session.UpdatedAt().Before(before) {
 				if id, ok := key.(string); ok {
-					toDelete = append(toDelete, id)
+					toDelete = append(toDelete, struct {
+						id      string
+						session Session
+					}{id, session})
 				}
 			}
 		}
 		return true
 	})
 
-	// Second pass: delete expired sessions
-	for _, id := range toDelete {
-		s.sessions.Delete(id)
+	// Second pass: close and delete expired sessions
+	for _, item := range toDelete {
+		// Check for context cancellation before processing each session
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		default:
+		}
+
+		// Re-check expiration and use CompareAndDelete to handle race conditions:
+		// - Session may have been touched via Manager.Get().Touch() and is no longer expired
+		// - Session may have been replaced via Store/UpsertSession with a new object
+		// Only proceed if the stored value is still the same session object and still expired
+		if item.session.UpdatedAt().Before(before) {
+			// CompareAndDelete ensures we only delete if the value hasn't been replaced
+			if deleted := s.sessions.CompareAndDelete(item.id, item.session); deleted {
+				// Successfully deleted - now close if implements io.Closer
+				if closer, ok := item.session.(io.Closer); ok {
+					if err := closer.Close(); err != nil {
+						slog.Warn("failed to close session during cleanup",
+							"session_id", item.id,
+							"error", err)
+					}
+				}
+			}
+			// If CompareAndDelete returned false, the session was already replaced/deleted - skip it
+		}
+		// If re-check shows session is no longer expired (was touched), skip it
 	}
 
 	return nil

pkg/transport/session/storage_test.go

@@ -5,6 +5,7 @@ package session
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"testing"
 	"time"
@@ -13,6 +14,24 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+// mockClosableSession is a test session that implements io.Closer
+type mockClosableSession struct {
+	*ProxySession
+	closeCalled bool
+	closeError  error
+}
+
+func newMockClosableSession(id string) *mockClosableSession {
+	return &mockClosableSession{
+		ProxySession: NewProxySession(id),
+	}
+}
+
+func (m *mockClosableSession) Close() error {
+	m.closeCalled = true
+	return m.closeError
+}
+
 // TestLocalStorage tests the LocalStorage implementation
 func TestLocalStorage(t *testing.T) {
 	t.Parallel()
@@ -302,6 +321,289 @@ func TestLocalStorage(t *testing.T) {
 		// Should not error, just stop early
 		assert.NoError(t, err)
 	})
+
+	t.Run("DeleteExpired calls Close on io.Closer sessions", func(t *testing.T) {
+		t.Parallel()
+		storage := NewLocalStorage()
+		defer storage.Close()
+
+		ctx := context.Background()
+
+		// Create a closable session (implements io.Closer)
+		closableSession := newMockClosableSession("closable-session")
+		closableSession.updated = time.Now().Add(-2 * time.Hour)
+
+		// Create a regular session (does not implement io.Closer)
+		regularSession := NewProxySession("regular-session")
+		regularSession.updated = time.Now().Add(-2 * time.Hour)
+
+		// Store both sessions
+		err := storage.Store(ctx, closableSession)
+		require.NoError(t, err)
+		err = storage.Store(ctx, regularSession)
+		require.NoError(t, err)
+
+		// Delete sessions older than 1 hour
+		cutoff := time.Now().Add(-1 * time.Hour)
+		err = storage.DeleteExpired(ctx, cutoff)
+		require.NoError(t, err)
+
+		// Both sessions should be deleted
+		_, err = storage.Load(ctx, "closable-session")
+		assert.Equal(t, ErrSessionNotFound, err)
+		_, err = storage.Load(ctx, "regular-session")
+		assert.Equal(t, ErrSessionNotFound, err)
+
+		// Close() is called synchronously, so it should already be done
+		assert.True(t, closableSession.closeCalled,
+			"Close() should have been called on closable session")
+	})
+
+	t.Run("DeleteExpired continues deletion even if Close fails", func(t *testing.T) {
+		t.Parallel()
+		storage := NewLocalStorage()
+		defer storage.Close()
+
+		ctx := context.Background()
+
+		// Create a closable session that returns an error on Close()
+		failingSession := newMockClosableSession("failing-session")
+		failingSession.closeError = errors.New("close failed")
+		failingSession.updated = time.Now().Add(-2 * time.Hour)
+
+		// Store the session
+		err := storage.Store(ctx, failingSession)
+		require.NoError(t, err)
+
+		// Delete expired sessions - should not fail even if Close() returns an error
+		cutoff := time.Now().Add(-1 * time.Hour)
+		err = storage.DeleteExpired(ctx, cutoff)
+		require.NoError(t, err)
+
+		// Session should be deleted from storage even though Close() failed
+		_, err = storage.Load(ctx, "failing-session")
+		assert.Equal(t, ErrSessionNotFound, err)
+
+		// Close() is called synchronously, so it should already be done
+		assert.True(t, failingSession.closeCalled,
+			"Close() should have been called even though it returned an error")
+
+		// Note: We don't verify log output to maintain t.Parallel() compatibility.
+		// The important behavior is that deletion continues even when Close() fails.
+	})
+
+	t.Run("DeleteExpired handles non-io.Closer sessions without error", func(t *testing.T) {
+		t.Parallel()
+		storage := NewLocalStorage()
+		defer storage.Close()
+
+		ctx := context.Background()
+
+		// Create multiple regular sessions (do not implement io.Closer)
+		for i := 0; i < 5; i++ {
+			session := NewProxySession(fmt.Sprintf("session-%d", i))
+			session.updated = time.Now().Add(-2 * time.Hour)
+			err := storage.Store(ctx, session)
+			require.NoError(t, err)
+		}
+
+		// Delete expired sessions
+		cutoff := time.Now().Add(-1 * time.Hour)
+		err := storage.DeleteExpired(ctx, cutoff)
+		require.NoError(t, err)
+
+		// All sessions should be deleted
+		for i := 0; i < 5; i++ {
+			_, err := storage.Load(ctx, fmt.Sprintf("session-%d", i))
+			assert.Equal(t, ErrSessionNotFound, err)
+		}
+	})
+
+	t.Run("DeleteExpired with mixed session types", func(t *testing.T) {
+		t.Parallel()
+		storage := NewLocalStorage()
+		defer storage.Close()
+
+		ctx := context.Background()
+
+		// Create a mix of closable and regular expired sessions
+		closable1 := newMockClosableSession("closable-1")
+		closable1.updated = time.Now().Add(-2 * time.Hour)
+		closable2 := newMockClosableSession("closable-2")
+		closable2.updated = time.Now().Add(-2 * time.Hour)
+
+		regular1 := NewProxySession("regular-1")
+		regular1.updated = time.Now().Add(-2 * time.Hour)
+		regular2 := NewProxySession("regular-2")
+		regular2.updated = time.Now().Add(-2 * time.Hour)
+
+		// Store all sessions
+		err := storage.Store(ctx, closable1)
+		require.NoError(t, err)
+		err = storage.Store(ctx, closable2)
+		require.NoError(t, err)
+		err = storage.Store(ctx, regular1)
+		require.NoError(t, err)
+		err = storage.Store(ctx, regular2)
+		require.NoError(t, err)
+
+		// Delete expired sessions
+		cutoff := time.Now().Add(-1 * time.Hour)
+		err = storage.DeleteExpired(ctx, cutoff)
+		require.NoError(t, err)
+
+		// All sessions should be deleted
+		_, err = storage.Load(ctx, "closable-1")
+		assert.Equal(t, ErrSessionNotFound, err)
+		_, err = storage.Load(ctx, "closable-2")
+		assert.Equal(t, ErrSessionNotFound, err)
+		_, err = storage.Load(ctx, "regular-1")
+		assert.Equal(t, ErrSessionNotFound, err)
+		_, err = storage.Load(ctx, "regular-2")
+		assert.Equal(t, ErrSessionNotFound, err)
+
+		// Close() is called synchronously, so it should already be done
+		assert.True(t, closable1.closeCalled,
+			"Close() should have been called on closable-1")
+		assert.True(t, closable2.closeCalled,
+			"Close() should have been called on closable-2")
+	})
+
+	t.Run("DeleteExpired respects context cancellation during deletion", func(t *testing.T) {
+		t.Parallel()
+		storage := NewLocalStorage()
+		defer storage.Close()
+
+		ctx := context.Background()
+
+		// Create many expired sessions to increase chance of context check
+		for i := 0; i < 10000; i++ {
+			session := NewProxySession(fmt.Sprintf("session-%d", i))
+			session.updated = time.Now().Add(-2 * time.Hour)
+			err := storage.Store(ctx, session)
+			require.NoError(t, err)
+		}
+
+		// Create a context with a very short timeout
+		timeoutCtx, cancel := context.WithTimeout(context.Background(), 1*time.Nanosecond)
+		defer cancel()
+
+		// Wait a bit to ensure context times out
+		time.Sleep(10 * time.Millisecond)
+
+		// DeleteExpired should respect context timeout
+		cutoff := time.Now().Add(-1 * time.Hour)
+		err := storage.DeleteExpired(timeoutCtx, cutoff)
+
+		// With 10000 sessions, the context check should trigger during cleanup
+		// If it completes too quickly, that's also acceptable behavior
+		if err != nil {
+			assert.Equal(t, context.DeadlineExceeded, err)
+			// Some sessions deleted, but not all due to timeout
+			remaining := storage.Count()
+			assert.Greater(t, remaining, 0, "Some sessions should remain due to context timeout")
+		}
+	})
+
+	t.Run("DeleteExpired handles concurrent Touch() race condition", func(t *testing.T) {
+		t.Parallel()
+		storage := NewLocalStorage()
+		defer storage.Close()
+
+		ctx := context.Background()
+
+		// Create an expired session
+		session := NewProxySession("race-session")
+		session.updated = time.Now().Add(-2 * time.Hour)
+		err := storage.Store(ctx, session)
+		require.NoError(t, err)
+
+		// Create many other expired sessions to slow down the deletion loop
+		for i := 0; i < 1000; i++ {
+			dummySession := NewProxySession(fmt.Sprintf("dummy-%d", i))
+			dummySession.updated = time.Now().Add(-2 * time.Hour)
+			err := storage.Store(ctx, dummySession)
+			require.NoError(t, err)
+		}
+
+		// Start DeleteExpired in a goroutine
+		done := make(chan error, 1)
+		go func() {
+			cutoff := time.Now().Add(-1 * time.Hour)
+			done <- storage.DeleteExpired(ctx, cutoff)
+		}()
+
+		// Concurrently touch the session to make it non-expired
+		// This simulates Manager.Get().Touch() being called during cleanup
+		session.Touch()
+
+		// Wait for DeleteExpired to complete
+		err = <-done
+		require.NoError(t, err)
+
+		// The session should NOT be deleted because it was touched
+		// (CompareAndDelete would fail due to updated timestamp or re-check would skip it)
+		loaded, err := storage.Load(ctx, "race-session")
+		if err == nil {
+			// Session still exists - this is correct behavior
+			assert.NotNil(t, loaded)
+			assert.True(t, loaded.UpdatedAt().After(time.Now().Add(-1*time.Hour)),
+				"Session should have recent timestamp after Touch()")
+		}
+		// Note: Due to timing, the session might still be deleted if Touch() happened
+		// after the re-check but before CompareAndDelete. This is acceptable as the
+		// important thing is we don't close a session that's been replaced.
+	})
+
+	t.Run("DeleteExpired handles concurrent Store() replacement race condition", func(t *testing.T) {
+		t.Parallel()
+		storage := NewLocalStorage()
+		defer storage.Close()
+
+		ctx := context.Background()
+
+		// Create an expired closable session
+		oldSession := newMockClosableSession("replace-session")
+		oldSession.updated = time.Now().Add(-2 * time.Hour)
+		err := storage.Store(ctx, oldSession)
+		require.NoError(t, err)
+
+		// Create many other expired sessions to slow down the deletion loop
+		for i := 0; i < 1000; i++ {
+			dummySession := NewProxySession(fmt.Sprintf("dummy-%d", i))
+			dummySession.updated = time.Now().Add(-2 * time.Hour)
+			err := storage.Store(ctx, dummySession)
+			require.NoError(t, err)
+		}
+
+		// Start DeleteExpired in a goroutine
+		done := make(chan error, 1)
+		go func() {
+			cutoff := time.Now().Add(-1 * time.Hour)
+			done <- storage.DeleteExpired(ctx, cutoff)
+		}()
+
+		// Concurrently replace the session with a new one (same ID, different object)
+		// This simulates UpsertSession being called during cleanup
+		newSession := newMockClosableSession("replace-session")
+		err = storage.Store(ctx, newSession)
+		require.NoError(t, err)
+
+		// Wait for DeleteExpired to complete
+		err = <-done
+		require.NoError(t, err)
+
+		// The new session should still exist (CompareAndDelete prevents deleting it)
+		loaded, err := storage.Load(ctx, "replace-session")
+		require.NoError(t, err)
+		assert.NotNil(t, loaded)
+
+		// The old session's Close() may or may not have been called depending on timing
+		// The important thing is the new session is not closed
+		// Since Close() is now synchronous, we can check immediately
+		assert.False(t, newSession.closeCalled,
+			"New session should not be closed (CompareAndDelete should prevent this)")
+	})
 }
 
 // TestManagerWithStorage tests the Manager with the Storage interface