Merge pull request #38 from EternisAI/joeldrotleff/fix-client-dropped-connection

fix deep research disconnect issue

Author: joeldrotleff

cmd/server/main.go

@@ -129,8 +129,9 @@ func main() {
 	mcpService := mcp.NewService()
 	searchService := search.NewService(logger.WithComponent("search"))
 
-	// Initialize deep research storage
+	// Initialize deep research storage and session manager
 	deeprStorage := deepr.NewDBStorage(logger.WithComponent("deepr-storage"), db.DB)
+	deeprSessionManager := deepr.NewSessionManager(logger.WithComponent("deepr-session"))
 
 	// Initialize handlers
 	oauthHandler := oauth.NewHandler(oauthService, logger.WithComponent("oauth"))
@@ -194,6 +195,7 @@ func main() {
 		mcpHandler:             mcpHandler,
 		searchHandler:          searchHandler,
 		deeprStorage:           deeprStorage,
+		deeprSessionManager:    deeprSessionManager,
 	})
 
 	// Initialize GraphQL server for Telegram
@@ -295,6 +297,7 @@ type restServerInput struct {
 	mcpHandler             *mcp.Handler
 	searchHandler          *search.Handler
 	deeprStorage           deepr.MessageStorage
+	deeprSessionManager    *deepr.SessionManager
 }
 
 func setupRESTServer(input restServerInput) *gin.Engine {
@@ -369,7 +372,7 @@ func setupRESTServer(input restServerInput) *gin.Engine {
 		api.POST("/exa/search", input.searchHandler.PostExaSearchHandler) // POST /api/v1/exa/search (Exa AI)
 
 		// Deep Research WebSocket endpoint (protected)
-		api.GET("/deepresearch/ws", deepr.DeepResearchHandler(input.logger, input.requestTrackingService, input.firebaseClient, input.deeprStorage)) // WebSocket proxy for deep research
+		api.GET("/deepresearch/ws", deepr.DeepResearchHandler(input.logger, input.requestTrackingService, input.firebaseClient, input.deeprStorage, input.deeprSessionManager)) // WebSocket proxy for deep research
 	}
 
 	// Protected proxy routes

internal/deepr/handlers.go

@@ -18,7 +18,7 @@ var upgrader = websocket.Upgrader{
 }
 
 // DeepResearchHandler handles WebSocket connections for deep research streaming
-func DeepResearchHandler(logger *logger.Logger, trackingService *request_tracking.Service, firebaseClient *auth.FirebaseClient, storage MessageStorage) gin.HandlerFunc {
+func DeepResearchHandler(logger *logger.Logger, trackingService *request_tracking.Service, firebaseClient *auth.FirebaseClient, storage MessageStorage, sessionManager *SessionManager) gin.HandlerFunc {
 	return func(c *gin.Context) {
 		log := logger.WithContext(c.Request.Context()).WithComponent("deepr")
 
@@ -74,8 +74,8 @@ func DeepResearchHandler(logger *logger.Logger, trackingService *request_trackin
 			slog.String("chat_id", chatID),
 			slog.String("remote_addr", c.Request.RemoteAddr))
 
-		// Create service instance with database storage
-		service := NewService(logger, trackingService, firebaseClient, storage)
+		// Create service instance with shared session manager
+		service := NewService(logger, trackingService, firebaseClient, storage, sessionManager)
 
 		// Handle the WebSocket connection
 		service.HandleConnection(c.Request.Context(), conn, userID, chatID)

internal/deepr/service.go

@@ -205,13 +205,13 @@ func (s *Service) validateFreemiumAccess(ctx context.Context, clientConn *websoc
 }
 
 // NewService creates a new deep research service with database storage
-func NewService(logger *logger.Logger, trackingService *request_tracking.Service, firebaseClient *auth.FirebaseClient, storage MessageStorage) *Service {
+func NewService(logger *logger.Logger, trackingService *request_tracking.Service, firebaseClient *auth.FirebaseClient, storage MessageStorage, sessionManager *SessionManager) *Service {
 	return &Service{
 		logger:          logger,
 		trackingService: trackingService,
 		firebaseClient:  firebaseClient,
 		storage:         storage,
-		sessionManager:  NewSessionManager(logger),
+		sessionManager:  sessionManager,
 	}
 }
 
@@ -229,6 +229,12 @@ func (s *Service) HandleConnection(ctx context.Context, clientConn *websocket.Co
 	// Check if this is a reconnection
 	isReconnection := s.sessionManager.HasActiveBackend(userID, chatID)
 
+	log.Info("reconnection check performed",
+		slog.String("user_id", userID),
+		slog.String("chat_id", chatID),
+		slog.String("client_id", clientID),
+		slog.Bool("has_active_backend", isReconnection))
+
 	if isReconnection {
 		log.Info("reconnection to existing session detected",
 			slog.String("user_id", userID),
@@ -666,9 +672,10 @@ func (s *Service) handleNewConnection(ctx context.Context, clientConn *websocket
 		}()
 	}
 
-	// Create session context derived from incoming context but independent of any single client
-	// This ensures cleanup on server shutdown while allowing session to outlive individual clients
-	sessionCtx, cancel := context.WithCancel(ctx)
+	// Create session context independent of any single client's request context
+	// This allows the backend connection to outlive individual client disconnections
+	// while still allowing cleanup when the session completes
+	sessionCtx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
 	// Create and register session

internal/deepr/session_manager.go

@@ -58,6 +58,19 @@ func (sm *SessionManager) CreateSession(userID, chatID string, backendConn *webs
 
 	key := sm.getSessionKey(userID, chatID)
 
+	// Check if session already exists
+	if existingSession, exists := sm.sessions[key]; exists {
+		sm.logger.WithComponent("deepr-session").Warn("OVERWRITING existing session",
+			slog.String("user_id", userID),
+			slog.String("chat_id", chatID),
+			slog.String("session_key", key),
+			slog.Int("existing_client_count", len(existingSession.clientConns)))
+		// Cancel the existing session's context
+		if existingSession.CancelFunc != nil {
+			existingSession.CancelFunc()
+		}
+	}
+
 	session := &ActiveSession{
 		UserID:      userID,
 		ChatID:      chatID,
@@ -85,6 +98,11 @@ func (sm *SessionManager) RemoveSession(userID, chatID string) {
 
 	key := sm.getSessionKey(userID, chatID)
 
+	sm.logger.WithComponent("deepr-session").Info("RemoveSession called",
+		slog.String("user_id", userID),
+		slog.String("chat_id", chatID),
+		slog.String("session_key", key))
+
 	if session, exists := sm.sessions[key]; exists {
 		// Close all client connections
 		session.mu.Lock()
@@ -238,7 +256,23 @@ func (sm *SessionManager) HasActiveBackend(userID, chatID string) bool {
 	defer sm.mu.RUnlock()
 
 	key := sm.getSessionKey(userID, chatID)
-	_, exists := sm.sessions[key]
+	session, exists := sm.sessions[key]
+
+	var clientCount int
+	if exists {
+		session.mu.RLock()
+		clientCount = len(session.clientConns)
+		session.mu.RUnlock()
+	}
+
+	sm.logger.WithComponent("deepr-session").Debug("HasActiveBackend called",
+		slog.String("user_id", userID),
+		slog.String("chat_id", chatID),
+		slog.String("session_key", key),
+		slog.Bool("session_exists", exists),
+		slog.Int("client_count", clientCount),
+		slog.Int("total_sessions", len(sm.sessions)))
+
 	return exists
 }