Merge pull request #351 from smart-mcp-proxy/fix/session-tracking-344

fix: session tracking not updating in Recent Sessions

Author: Dumbris

internal/runtime/lifecycle.go

@@ -239,10 +239,10 @@ func (r *Runtime) backgroundToolIndexing(ctx context.Context) {
 // backgroundSessionCleanup periodically closes sessions that haven't had activity.
 // This handles the HTTP transport limitation where OnUnregisterSession is never called.
 func (r *Runtime) backgroundSessionCleanup(ctx context.Context) {
-	// Session inactivity timeout: 5 minutes
-	// This is a reasonable timeout for MCP sessions where clients typically
-	// send tool calls every few seconds during active use.
-	const sessionInactivityTimeout = 5 * time.Minute
+	// Session inactivity timeout: 30 minutes
+	// MCP clients may have gaps between tool calls (e.g., user reading results).
+	// 5 minutes was too aggressive and caused sessions to appear stale.
+	const sessionInactivityTimeout = 30 * time.Minute
 
 	// Check every minute for inactive sessions
 	ticker := time.NewTicker(1 * time.Minute)

internal/server/mcp.go

@@ -101,6 +101,9 @@ type MCPProxyServer struct {
 
 	// MCP session tracking
 	sessionStore *SessionStore
+
+	// Hooks shared across all routing mode servers
+	hooks *mcpserver.Hooks
 }
 
 // NewMCPProxyServer creates a new MCP proxy server
@@ -122,6 +125,17 @@ func NewMCPProxyServer(
 
 	// Create hooks to capture session information
 	hooks := &mcpserver.Hooks{}
+	// Update session activity on every MCP message to prevent premature cleanup.
+	// Without this, sessions that only do retrieve_tools (not tool calls) would
+	// appear inactive and get closed by the background cleanup.
+	hooks.AddBeforeAny(func(ctx context.Context, id any, method mcp.MCPMethod, message any) {
+		session := mcpserver.ClientSessionFromContext(ctx)
+		if session == nil {
+			return
+		}
+		sessionStore.UpdateActivity(session.SessionID())
+	})
+
 	hooks.AddOnRegisterSession(func(ctx context.Context, sess mcpserver.ClientSession) {
 		sessionID := sess.SessionID()
 
@@ -235,6 +249,7 @@ func NewMCPProxyServer(
 		config:          config,
 		jsPool:          jsPool,
 		sessionStore:    sessionStore,
+		hooks:           hooks,
 	}
 
 	// Register proxy tools for the default (retrieve_tools) server

internal/server/mcp_routing.go

@@ -453,28 +453,34 @@ func (p *MCPProxyServer) buildCodeExecutionTool() []mcpserver.ServerTool {
 // Each server instance has its own set of tools registered appropriate for that mode.
 // The main "server" field remains the retrieve_tools mode server (default).
 func (p *MCPProxyServer) initRoutingModeServers() {
+	// All routing mode servers share the same hooks for session tracking
+	opts := []mcpserver.ServerOption{
+		mcpserver.WithToolCapabilities(true),
+		mcpserver.WithRecovery(),
+	}
+	if p.hooks != nil {
+		opts = append(opts, mcpserver.WithHooks(p.hooks))
+	}
+
 	// Create direct mode server
 	p.directServer = mcpserver.NewMCPServer(
 		"mcpproxy-go",
 		mcpServerVersion(),
-		mcpserver.WithToolCapabilities(true),
-		mcpserver.WithRecovery(),
+		opts...,
 	)
 
 	// Create code execution mode server
 	p.codeExecServer = mcpserver.NewMCPServer(
 		"mcpproxy-go",
 		mcpServerVersion(),
-		mcpserver.WithToolCapabilities(true),
-		mcpserver.WithRecovery(),
+		opts...,
 	)
 
 	// Create call tool mode server (/mcp/call)
 	p.callToolServer = mcpserver.NewMCPServer(
 		"mcpproxy-go",
 		mcpServerVersion(),
-		mcpserver.WithToolCapabilities(true),
-		mcpserver.WithRecovery(),
+		opts...,
 	)
 
 	// Register tools for code execution mode (static tools that don't change)

internal/server/session_store.go

@@ -4,8 +4,8 @@ import (
 	"sync"
 	"time"
 
-	"go.uber.org/zap"
 	"github.com/smart-mcp-proxy/mcpproxy-go/internal/storage"
+	"go.uber.org/zap"
 )
 
 // SessionInfo holds MCP session metadata
@@ -127,6 +127,16 @@ func (s *SessionStore) UpdateSessionStats(sessionID string, tokens int) {
 	}
 }
 
+// UpdateActivity updates the last activity timestamp for a session without incrementing stats.
+func (s *SessionStore) UpdateActivity(sessionID string) {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+
+	if s.storageManager != nil {
+		_ = s.storageManager.UpdateSessionActivity(sessionID)
+	}
+}
+
 // Count returns the number of active sessions
 func (s *SessionStore) Count() int {
 	s.mu.RLock()

internal/storage/manager.go

@@ -962,8 +962,8 @@ func (m *Manager) CreateSession(session *SessionRecord) error {
 		c := bucket.Cursor()
 		for k, v := c.First(); k != nil; k, v = c.Next() {
 			keyStr := string(k)
-			// Check if key ends with the session ID (after the underscore)
-			if len(keyStr) > len(session.ID) && keyStr[len(keyStr)-len(session.ID):] == session.ID {
+			// Check if key ends with _{session_id}
+			if strings.HasSuffix(keyStr, "_"+session.ID) {
 				existingKey = k
 				if err := json.Unmarshal(v, &existingSession); err != nil {
 					m.logger.Warnw("Failed to unmarshal existing session", "error", err)
@@ -1033,8 +1033,8 @@ func (m *Manager) CloseSession(sessionID string) error {
 		for k, v := c.First(); k != nil; k, v = c.Next() {
 			// Key format: {timestamp_ns}_{session_id}
 			keyStr := string(k)
-			// Check if key ends with the session ID (after the underscore)
-			if len(keyStr) > len(sessionID) && keyStr[len(keyStr)-len(sessionID):] == sessionID {
+			// Check if key ends with _{session_id}
+			if strings.HasSuffix(keyStr, "_"+sessionID) {
 				sessionKey = k
 				if err := json.Unmarshal(v, &session); err != nil {
 					return fmt.Errorf("failed to unmarshal session: %w", err)
@@ -1115,8 +1115,8 @@ func (m *Manager) GetSessionByID(sessionID string) (*SessionRecord, error) {
 		c := bucket.Cursor()
 		for k, v := c.First(); k != nil; k, v = c.Next() {
 			keyStr := string(k)
-			// Check if key ends with the session ID (after the underscore)
-			if len(keyStr) > len(sessionID) && keyStr[len(keyStr)-len(sessionID):] == sessionID {
+			// Check if key ends with _{session_id}
+			if strings.HasSuffix(keyStr, "_"+sessionID) {
 				var s SessionRecord
 				if err := json.Unmarshal(v, &s); err != nil {
 					return fmt.Errorf("failed to unmarshal session: %w", err)
@@ -1200,8 +1200,8 @@ func (m *Manager) UpdateSessionStats(sessionID string, tokens int) error {
 		c := bucket.Cursor()
 		for k, v := c.First(); k != nil; k, v = c.Next() {
 			keyStr := string(k)
-			// Check if key ends with the session ID (after the underscore)
-			if len(keyStr) > len(sessionID) && keyStr[len(keyStr)-len(sessionID):] == sessionID {
+			// Check if key ends with _{session_id}
+			if strings.HasSuffix(keyStr, "_"+sessionID) {
 				sessionKey = k
 				if err := json.Unmarshal(v, &session); err != nil {
 					return fmt.Errorf("failed to unmarshal session: %w", err)
@@ -1228,6 +1228,38 @@ func (m *Manager) UpdateSessionStats(sessionID string, tokens int) error {
 	})
 }
 
+// UpdateSessionActivity updates LastActivity without incrementing tool call counts.
+// Call this on any MCP message to keep the session alive.
+func (m *Manager) UpdateSessionActivity(sessionID string) error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	return m.db.db.Update(func(tx *bbolt.Tx) error {
+		bucket := tx.Bucket([]byte(SessionsBucket))
+		if bucket == nil {
+			return nil
+		}
+
+		c := bucket.Cursor()
+		for k, v := c.First(); k != nil; k, v = c.Next() {
+			keyStr := string(k)
+			if strings.HasSuffix(keyStr, "_"+sessionID) {
+				var session SessionRecord
+				if err := json.Unmarshal(v, &session); err != nil {
+					return err
+				}
+				session.LastActivity = time.Now()
+				data, err := json.Marshal(session)
+				if err != nil {
+					return err
+				}
+				return bucket.Put(k, data)
+			}
+		}
+		return nil // Session not found is not an error - may not be persisted yet
+	})
+}
+
 // CloseInactiveSessions closes sessions that haven't had activity for the specified duration
 func (m *Manager) CloseInactiveSessions(inactivityTimeout time.Duration) (int, error) {
 	m.mu.Lock()