fixes from review

Author: taskbot

pkg/transport/proxy/transparent/transparent_proxy.go

@@ -13,7 +13,6 @@ import (
 	"net/http/httputil"
 	"net/url"
 	"regexp"
-	"strconv"
 	"strings"
 	"sync"
 	"time"
@@ -56,16 +55,13 @@ type TransparentProxy struct {
 	// Optional Prometheus metrics handler
 	prometheusHandler http.Handler
 
-	// Sessions for managing client connections
-	sessions map[string]session.Session
+	// Sessions for tracking state
+	sessionManager *session.Manager
 
-	// mutex for protecting session access
-	sessionMutex sync.Mutex
+	// If mcp server has been initialized
+	IsServerInitialized bool
 }
 
-// TransparentProxySessionID is the session ID used for the transparent proxy.
-const TransparentProxySessionID = "transparent-proxy-session"
-
 // NewTransparentProxy creates a new transparent proxy with optional middlewares.
 func NewTransparentProxy(
 	host string,
@@ -83,26 +79,27 @@ func NewTransparentProxy(
 		middlewares:       middlewares,
 		shutdownCh:        make(chan struct{}),
 		prometheusHandler: prometheusHandler,
-		sessions:          make(map[string]session.Session),
+		sessionManager:    session.NewManager(30 * time.Minute),
 	}
 
 	// Create MCP pinger and health checker
 	mcpPinger := NewMCPPinger(targetURI)
 	proxy.healthChecker = healthcheck.NewHealthChecker("sse", mcpPinger)
 
-	_, err := proxy.CreateSession(TransparentProxySessionID)
-	if err != nil {
-		logger.Errorf("Failed to create session for TransparentProxy: %v", err)
-		return nil
-	}
-
 	return proxy
 }
 
+var sessionIDRegex = regexp.MustCompile(`sessionId=([\w-]+)`)
+
 func (p *TransparentProxy) handleModifyResponse(res *http.Response) error {
-	// Log headers
 	if sid := res.Header.Get("Mcp-Session-Id"); sid != "" {
-		logger.Infof("🆔 Streamable session ID from header: %s", sid)
+		logger.Infof("Detected Mcp-Session-Id header: %s", sid)
+		if _, ok := p.sessionManager.Get(sid); !ok {
+			if _, err := p.sessionManager.AddWithID(sid); err != nil {
+				logger.Errorf("Failed to create session from header %s: %v", sid, err)
+			}
+		}
+		p.IsServerInitialized = true
 	}
 
 	// Handle streaming (SSE)
@@ -114,22 +111,21 @@ func (p *TransparentProxy) handleModifyResponse(res *http.Response) error {
 		go func() {
 			defer pw.Close()
 			scanner := bufio.NewScanner(orig)
-			re := regexp.MustCompile(`sessionId=([\w-]+)`) // Capture UUID-like IDs
-
 			for scanner.Scan() {
 				line := scanner.Text()
 
-				if matches := re.FindStringSubmatch(line); len(matches) == 2 {
+				if matches := sessionIDRegex.FindStringSubmatch(line); len(matches) == 2 {
 					sessionID := matches[1]
-
-					// set session id for proxy
-					extractedSession, ok := p.GetSession(TransparentProxySessionID)
+					_, ok := p.sessionManager.Get(sessionID)
 					if !ok {
-						logger.Errorf("Failed to get session for TransparentProxy")
-						continue
+						var err error
+						_, err = p.sessionManager.AddWithID(sessionID)
+						if err != nil {
+							logger.Errorf("Failed to create session %s: %v", sessionID, err)
+							continue
+						}
 					}
-					extractedSession.SetIsInitialized(true)
-					extractedSession.SetMCPSessionID(sessionID)
+					p.IsServerInitialized = true
 				}
 				_, err := pw.Write([]byte(line + "\n"))
 				if err != nil {
@@ -140,29 +136,6 @@ func (p *TransparentProxy) handleModifyResponse(res *http.Response) error {
 		return nil
 	}
 
-	// Handle non-streaming (JSON) bodies
-	body, err := io.ReadAll(res.Body)
-	if err != nil {
-		return err
-	}
-	err = res.Body.Close()
-	if err != nil {
-		logger.Errorf("Failed to close response body: %v", err)
-	}
-
-	text := string(body)
-	logger.Infof("HTTP response body: %s", text)
-
-	// Parse sessionId if embedded in JSON
-	rejson := regexp.MustCompile(`"sessionId"\s*:\s*"([\w-]+)"`)
-	if matches := rejson.FindStringSubmatch(text); len(matches) == 2 {
-		sessionID := matches[1]
-		logger.Infof("🆔 Captured sessionId from JSON: %s", sessionID)
-	}
-
-	res.Body = io.NopCloser(bytes.NewReader(body))
-	res.ContentLength = int64(len(body))
-	res.Header.Set("Content-Length", strconv.Itoa(len(body)))
 	return nil
 }
 
@@ -176,13 +149,8 @@ func (p *TransparentProxy) handleAndDetectInitialize(w http.ResponseWriter, r *h
 			logger.Errorf("Error reading request body: %v", err)
 		} else {
 			if bytes.Contains(body, []byte(`"method":"initialize"`)) {
-				logger.Infof("🔧 Detected initialize request to %s", r.URL.Path)
-				extractedSession, ok := p.GetSession(TransparentProxySessionID)
-				if ok {
-					extractedSession.SetIsInitialized(true)
-				} else {
-					logger.Errorf("No session found to mark initialized")
-				}
+				logger.Infof("Detected initialize request to %s", r.URL.Path)
+				p.IsServerInitialized = true
 			}
 			r.Body = io.NopCloser(bytes.NewReader(body))
 			r.ContentLength = int64(len(body))
@@ -278,13 +246,7 @@ func (p *TransparentProxy) monitorHealth(parentCtx context.Context) {
 			return
 		case <-ticker.C:
 			// Perform health check only if mcp server has been initialized
-			extractedSession, isSession := p.GetSession(TransparentProxySessionID)
-			if !isSession {
-				logger.Errorf("Failed to get session for health check")
-				return
-			}
-
-			if extractedSession.IsInitialized() {
+			if p.IsServerInitialized {
 				alive := p.healthChecker.CheckHealth(parentCtx)
 				if alive.Status != healthcheck.StatusHealthy {
 					logger.Infof("Health check failed for %s; initiating proxy shutdown", p.containerName)
@@ -294,7 +256,7 @@ func (p *TransparentProxy) monitorHealth(parentCtx context.Context) {
 					return
 				}
 			} else {
-				logger.Infof("Session %s is not initialized, cannot start healthcheck", extractedSession.ID())
+				logger.Infof("MCP server not initialized yet, skipping health check for %s", p.containerName)
 			}
 		}
 	}
@@ -347,32 +309,3 @@ func (*TransparentProxy) SendMessageToDestination(_ jsonrpc2.Message) error {
 func (*TransparentProxy) ForwardResponseToClients(_ context.Context, _ jsonrpc2.Message) error {
 	return fmt.Errorf("ForwardResponseToClients not implemented for TransparentProxy")
 }
-
-// CreateSession creates a new session for the transparent proxy.
-func (p *TransparentProxy) CreateSession(id string) (session.Session, error) {
-	logger.Infof("Creating session with ID: %s", id)
-	if id == "" {
-		return nil, fmt.Errorf("session ID cannot be empty")
-	}
-	p.sessionMutex.Lock()
-	defer p.sessionMutex.Unlock()
-	if _, found := p.sessions[id]; found {
-		return nil, fmt.Errorf("session %s exists", id)
-	}
-	s := &session.ProxySession{Id: id}
-	s.Init()
-	p.sessions[id] = s
-	return s, nil
-}
-
-// GetSession retrieves a session by ID.
-func (p *TransparentProxy) GetSession(id string) (session.Session, bool) {
-	logger.Infof("Retrieving session with ID: %s", id)
-	if id == "" {
-		return nil, false
-	}
-	p.sessionMutex.Lock()
-	defer p.sessionMutex.Unlock()
-	s, ok := p.sessions[id]
-	return s, ok
-}

pkg/transport/session/manager.go

@@ -0,0 +1,106 @@
+// Package session provides a session manager with TTL cleanup.
+package session
+
+import (
+	"fmt"
+	"sync"
+	"time"
+)
+
+// Session interface
+type Session interface {
+	ID() string
+	CreatedAt() time.Time
+	UpdatedAt() time.Time
+	Touch()
+}
+
+// Manager holds sessions with TTL cleanup.
+type Manager struct {
+	sessions map[string]Session
+	mu       sync.RWMutex
+	ttl      time.Duration
+	stopCh   chan struct{}
+}
+
+// NewManager creates a session manager with TTL and starts cleanup worker.
+func NewManager(ttl time.Duration) *Manager {
+	m := &Manager{
+		sessions: make(map[string]Session),
+		ttl:      ttl,
+		stopCh:   make(chan struct{}),
+	}
+	go m.cleanupRoutine()
+	return m
+}
+
+func (m *Manager) cleanupRoutine() {
+	ticker := time.NewTicker(m.ttl / 2)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ticker.C:
+			m.CleanupExpired()
+		case <-m.stopCh:
+			return
+		}
+	}
+}
+
+// AddWithID creates (and adds) a new session with the provided ID.
+// Returns error if ID is empty or already exists.
+func (m *Manager) AddWithID(id string) (Session, error) {
+	if id == "" {
+		return nil, fmt.Errorf("session ID cannot be empty")
+	}
+
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	if _, exists := m.sessions[id]; exists {
+		return nil, fmt.Errorf("session ID %q already exists", id)
+	}
+
+	s := NewProxySession(id)
+	m.sessions[id] = s
+	return s, nil
+}
+
+// Get retrieves a session by ID. Returns (session, true) if found,
+// and also updates its UpdatedAt timestamp.
+func (m *Manager) Get(id string) (Session, bool) {
+	m.mu.RLock()
+	s, ok := m.sessions[id]
+	m.mu.RUnlock()
+
+	if !ok {
+		return nil, false
+	}
+
+	s.Touch()
+	return s, true
+}
+
+// Delete removes a session by ID.
+func (m *Manager) Delete(id string) {
+	m.mu.Lock()
+	delete(m.sessions, id)
+	m.mu.Unlock()
+}
+
+// CleanupExpired removes sessions that have not been updated within the TTL.
+func (m *Manager) CleanupExpired() {
+	cutoff := time.Now().Add(-m.ttl)
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	for id, s := range m.sessions {
+		if s.UpdatedAt().Before(cutoff) {
+			delete(m.sessions, id)
+		}
+	}
+}
+
+// Stop stops the cleanup worker.
+func (m *Manager) Stop() {
+	close(m.stopCh)
+}

pkg/transport/session/proxy_session.go

@@ -0,0 +1,28 @@
+package session
+
+import "time"
+
+// ProxySession implements the Session interface for proxy sessions.
+type ProxySession struct {
+	id      string
+	created time.Time
+	updated time.Time
+}
+
+// NewProxySession creates a new ProxySession with the given ID.
+func NewProxySession(id string) *ProxySession {
+	now := time.Now()
+	return &ProxySession{id: id, created: now, updated: now}
+}
+
+// ID returns the session ID.
+func (s *ProxySession) ID() string { return s.id }
+
+// CreatedAt returns the creation time of the session.
+func (s *ProxySession) CreatedAt() time.Time { return s.created }
+
+// UpdatedAt returns the last updated time of the session.
+func (s *ProxySession) UpdatedAt() time.Time { return s.updated }
+
+// Touch updates the session's last updated time to the current time.
+func (s *ProxySession) Touch() { s.updated = time.Now() }