fixes from copilot

Author: taskbot

pkg/transport/proxy/transparent/transparent_proxy.go

@@ -80,7 +80,7 @@ func NewTransparentProxy(
 		middlewares:       middlewares,
 		shutdownCh:        make(chan struct{}),
 		prometheusHandler: prometheusHandler,
-		sessionManager:    session.NewManager(30 * time.Minute),
+		sessionManager:    session.NewManager(30*time.Minute, session.NewProxySession),
 	}
 
 	// Create MCP pinger and health checker
@@ -95,12 +95,12 @@ type tracingTransport struct {
 	p    *TransparentProxy
 }
 
-func (t *tracingTransport) setServerInitialized() {
-	if !t.p.IsServerInitialized {
-		t.p.mutex.Lock()
-		t.p.IsServerInitialized = true
-		t.p.mutex.Unlock()
-		logger.Infof("Server was initialized successfully for %s", t.p.containerName)
+func (p *TransparentProxy) setServerInitialized() {
+	if !p.IsServerInitialized {
+		p.mutex.Lock()
+		p.IsServerInitialized = true
+		p.mutex.Unlock()
+		logger.Infof("Server was initialized successfully for %s", p.containerName)
 	}
 }
 
@@ -136,21 +136,16 @@ func (t *tracingTransport) RoundTrip(req *http.Request) (*http.Response, error)
 		if ct != "" {
 			logger.Infof("Detected Mcp-Session-Id header: %s", ct)
 			if _, ok := t.p.sessionManager.Get(ct); !ok {
-				fmt.Println("i get session id")
 				if err := t.p.sessionManager.AddWithID(ct); err != nil {
-					fmt.Println("i add session")
 					logger.Errorf("Failed to create session from header %s: %v", ct, err)
 				}
-				fmt.Println("i set server initialized")
 			}
-			fmt.Println("i set server initialized")
-			t.setServerInitialized()
+			t.p.setServerInitialized()
 			return resp, nil
 		}
 		// status was ok and we saw an initialize call
 		if sawInitialize && !t.p.IsServerInitialized {
-			fmt.Println("here")
-			t.setServerInitialized()
+			t.p.setServerInitialized()
 			return resp, nil
 		}
 	}
@@ -212,7 +207,7 @@ func (p *TransparentProxy) modifyForSessionID(resp *http.Response) error {
 					if sid == "" {
 						sid = string(m[2])
 					}
-					p.IsServerInitialized = true
+					p.setServerInitialized()
 					err := p.sessionManager.AddWithID(sid)
 					if err != nil {
 						logger.Errorf("Failed to create session from SSE line: %v", err)

pkg/transport/session/manager.go

@@ -20,13 +20,19 @@ type Manager struct {
 	sessions sync.Map
 	ttl      time.Duration
 	stopCh   chan struct{}
+	factory  Factory
 }
 
+// Factory defines a function type for creating new sessions.
+type Factory func(id string) *ProxySession
+
 // NewManager creates a session manager with TTL and starts cleanup worker.
-func NewManager(ttl time.Duration) *Manager {
+func NewManager(ttl time.Duration, factory Factory) *Manager {
 	m := &Manager{
-		ttl:    ttl,
-		stopCh: make(chan struct{}),
+		sessions: sync.Map{},
+		ttl:      ttl,
+		stopCh:   make(chan struct{}),
+		factory:  factory,
 	}
 	go m.cleanupRoutine()
 	return m
@@ -40,7 +46,11 @@ func (m *Manager) cleanupRoutine() {
 		case <-ticker.C:
 			cutoff := time.Now().Add(-m.ttl)
 			m.sessions.Range(func(key, val any) bool {
-				sess := val.(Session)
+				sess, ok := val.(Session)
+				if !ok {
+					// Skip invalid value
+					return true
+				}
 				if sess.UpdatedAt().Before(cutoff) {
 					m.sessions.Delete(key)
 				}
@@ -59,7 +69,8 @@ func (m *Manager) AddWithID(id string) error {
 		return fmt.Errorf("session ID cannot be empty")
 	}
 	// Use LoadOrStore: returns existing if already present
-	_, loaded := m.sessions.LoadOrStore(id, NewProxySession(id))
+	session := m.factory(id)
+	_, loaded := m.sessions.LoadOrStore(id, session)
 	if loaded {
 		return fmt.Errorf("session ID %q already exists", id)
 	}
@@ -73,7 +84,11 @@ func (m *Manager) Get(id string) (Session, bool) {
 	if !ok {
 		return nil, false
 	}
-	sess := v.(Session)
+	sess, ok := v.(Session)
+	if !ok {
+		return nil, false // Invalid session type
+	}
+
 	sess.Touch()
 	return sess, true
 }
@@ -87,3 +102,14 @@ func (m *Manager) Delete(id string) {
 func (m *Manager) Stop() {
 	close(m.stopCh)
 }
+
+func (m *Manager) cleanupExpiredOnce() {
+	cutoff := time.Now().Add(-m.ttl)
+	m.sessions.Range(func(key, val any) bool {
+		sess := val.(Session)
+		if sess.UpdatedAt().Before(cutoff) {
+			m.sessions.Delete(key)
+		}
+		return true
+	})
+}

pkg/transport/session/manager_test.go

@@ -1,123 +1,140 @@
 package session
 
 import (
+	"sync"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
-func TestAddAndGetWithStubSession(t *testing.T) {
-	t.Parallel()
+// stubFactory returns ProxySessions with fixed timestamps and records IDs.
+type stubFactory struct {
+	mu         sync.Mutex
+	createdIDs []string
+	fixedTime  time.Time
+}
 
-	orig := NewProxySession
-	NewProxySession = func(id string) *ProxySession {
-		ts := time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC)
-		return &ProxySession{id: id, created: ts, updated: ts}
+func (f *stubFactory) New(id string) *ProxySession {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	f.createdIDs = append(f.createdIDs, id)
+	return &ProxySession{
+		id:      id,
+		created: f.fixedTime,
+		updated: f.fixedTime,
 	}
-	defer func() { NewProxySession = orig }()
+}
 
-	m := NewManager(1 * time.Hour)
+func TestAddAndGetWithStubSession(t *testing.T) {
+	t.Parallel()
+	now := time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC)
+	factory := &stubFactory{fixedTime: now}
+
+	m := NewManager(time.Hour, factory.New)
 	defer m.Stop()
 
 	require.NoError(t, m.AddWithID("foo"))
 
 	sess, ok := m.Get("foo")
-	require.True(t, ok)
+	require.True(t, ok, "session foo should exist")
 	assert.Equal(t, "foo", sess.ID())
+	assert.Contains(t, factory.createdIDs, "foo")
 }
 
 func TestAddDuplicate(t *testing.T) {
 	t.Parallel()
+	factory := &stubFactory{fixedTime: time.Now()}
 
-	m := NewManager(time.Hour)
+	m := NewManager(time.Hour, factory.New)
 	defer m.Stop()
 
-	err := m.AddWithID("dup")
-	assert.NoError(t, err)
+	require.NoError(t, m.AddWithID("dup"))
 
-	err2 := m.AddWithID("dup")
-	assert.Error(t, err2)
-	assert.Contains(t, err2.Error(), "already exists")
+	err := m.AddWithID("dup")
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "already exists")
 }
 
 func TestDeleteSession(t *testing.T) {
 	t.Parallel()
+	factory := &stubFactory{fixedTime: time.Now()}
 
-	m := NewManager(time.Hour)
+	m := NewManager(time.Hour, factory.New)
 	defer m.Stop()
 
 	require.NoError(t, m.AddWithID("del"))
 	m.Delete("del")
 
 	_, ok := m.Get("del")
-	assert.False(t, ok)
+	assert.False(t, ok, "deleted session should not be found")
 }
 
 func TestGetUpdatesTimestamp(t *testing.T) {
 	t.Parallel()
+	oldTime := time.Now().Add(-1 * time.Minute)
+	factory := &stubFactory{fixedTime: oldTime}
 
-	orig := NewProxySession
-	NewProxySession = func(id string) *ProxySession {
-		ts := time.Now().Add(-1 * time.Minute)
-		return &ProxySession{id: id, created: ts, updated: ts}
-	}
-	defer func() { NewProxySession = orig }()
-
-	m := NewManager(1 * time.Hour)
+	m := NewManager(time.Hour, factory.New)
 	defer m.Stop()
 
 	require.NoError(t, m.AddWithID("touchme"))
-	s1, _ := m.Get("touchme")
+	s1, ok := m.Get("touchme")
+	require.True(t, ok)
 	t0 := s1.UpdatedAt()
 
-	time.Sleep(5 * time.Millisecond)
-	s2, _ := m.Get("touchme")
+	time.Sleep(10 * time.Millisecond)
+	s2, ok2 := m.Get("touchme")
+	require.True(t, ok2)
 	t1 := s2.UpdatedAt()
 
-	assert.True(t, t1.After(t0), "UpdatedAt should update on Get()")
+	assert.True(t, t1.After(t0), "UpdatedAt should update on repeated Get()")
 }
-
-func TestCleanupExpired(t *testing.T) {
+func TestCleanupExpired_ManualTrigger(t *testing.T) {
 	t.Parallel()
 
+	// Stub factory: all sessions start with UpdatedAt = `now`
+	now := time.Now()
+	factory := &stubFactory{fixedTime: now}
 	ttl := 50 * time.Millisecond
-	orig := NewProxySession
-	NewProxySession = func(id string) *ProxySession {
-		return &ProxySession{
-			id:      id,
-			created: time.Now(),
-			updated: time.Now(),
-		}
-	}
-	defer func() { NewProxySession = orig }()
 
-	m := NewManager(ttl)
+	m := NewManager(ttl, factory.New)
 	defer m.Stop()
 
 	require.NoError(t, m.AddWithID("old"))
-	time.Sleep(ttl * 2) // allow old to expire
 
-	require.NoError(t, m.AddWithID("new"))
-	time.Sleep(ttl) // let cleanup execute
+	// Retrieve and expire session manually
+	sess, ok := m.Get("old")
+	require.True(t, ok)
+	ps := sess.(*ProxySession)
+	ps.updated = now.Add(-ttl * 2)
 
+	// Run cleanup manually
+	m.cleanupExpiredOnce()
+
+	// Now it should be gone
 	_, okOld := m.Get("old")
+	assert.False(t, okOld, "expired session should have been cleaned")
+
+	// Add fresh session and assert it remains after cleanup
+	require.NoError(t, m.AddWithID("new"))
+	m.cleanupExpiredOnce()
 	_, okNew := m.Get("new")
-	assert.False(t, okOld, "expired session should be cleaned")
-	assert.True(t, okNew, "recent session should remain")
+	assert.True(t, okNew, "new session should still exist after cleanup")
 }
 
 func TestStopDisablesCleanup(t *testing.T) {
 	t.Parallel()
-
 	ttl := 50 * time.Millisecond
-	m := NewManager(ttl)
-	m.Stop() // stop cleanup upfront
+	factory := &stubFactory{fixedTime: time.Now()}
+
+	m := NewManager(ttl, factory.New)
+	m.Stop() // disable cleanup before any session expires
 
 	require.NoError(t, m.AddWithID("stay"))
 	time.Sleep(ttl * 2)
 
 	_, ok := m.Get("stay")
-	assert.True(t, ok, "session should persist after Stop()")
+	assert.True(t, ok, "session should still be present even after Stop() and TTL elapsed")
 }