feat (multi-tenancy): activity based tiered polling

Author: itsskofficial

src/server/app/app.py

@@ -733,117 +733,121 @@ async def gdrive_tool(tool_call_input: dict) -> Dict[str, Any]:
 agent_task_processor_instance: Optional[AgentTaskProcessor] = None 
 
 async def polling_scheduler_loop():
-    print(f"[POLLING_SCHEDULER] Starting polling scheduler loop (interval: {POLLING_SCHEDULER_INTERVAL_SECONDS}s)")
+    print(f"[POLLING_SCHEDULER] Starting loop (interval: {POLLING_SCHEDULER_INTERVAL_SECONDS}s)")
+    if not mongo_manager:
+        print("[POLLING_SCHEDULER_ERROR] MongoManager not initialized. Scheduler cannot run.")
+        return
+
+    await mongo_manager.reset_stale_polling_locks() 
+    
     while True:
         try:
-            print(f"[POLLING_SCHEDULER] Checking for due polling tasks at {datetime.now(timezone.utc).isoformat()}")
-            due_tasks_states = await mongo_manager.get_due_polling_tasks()
+            # print(f"[POLLING_SCHEDULER] Checking for due polling tasks at {datetime.now(timezone.utc).isoformat()}")
+            due_tasks_states = await mongo_manager.get_due_polling_tasks() 
 
             if not due_tasks_states:
                 # print(f"[POLLING_SCHEDULER] No tasks due at this time.")
                 pass
             else:
-                print(f"[POLLING_SCHEDULER] Found {len(due_tasks_states)} due tasks.")
+                print(f"[POLLING_SCHEDULER] Found {len(due_tasks_states)} due polling tasks.")
 
             for task_state in due_tasks_states:
                 user_id = task_state["user_id"]
-                engine_category = task_state["engine_category"]
+                service_name = task_state["service_name"] 
 
-                print(f"[POLLING_SCHEDULER] Attempting to process task for {user_id}/{engine_category}")
+                print(f"[POLLING_SCHEDULER] Attempting to process task for {user_id}/{service_name}")
 
-                # Try to acquire lock and get the task. If successful, proceed.
-                locked_task_state = await mongo_manager.set_polling_status_and_get(user_id, engine_category)
+                locked_task_state = await mongo_manager.set_polling_status_and_get(user_id, service_name)
 
                 if locked_task_state:
-                    print(f"[POLLING_SCHEDULER] Acquired lock for {user_id}/{engine_category}. Triggering poll.")
-                    engine_instance = active_context_engines.get(user_id, {}).get(engine_category)
+                    print(f"[POLLING_SCHEDULER] Acquired lock for {user_id}/{service_name}. Triggering poll.")
+                    
+                    engine_instance = active_context_engines.get(user_id, {}).get(service_name)
 
                     if not engine_instance:
-                        engine_class = DATA_SOURCES_CONFIG.get(engine_category, {}).get("engine_class")
-                        if engine_class:
-                            print(f"[POLLING_SCHEDULER] Creating new instance for {user_id}/{engine_category}")
+                        engine_config = DATA_SOURCES_CONFIG.get(service_name)
+                        if engine_config and engine_config.get("engine_class"):
+                            engine_class = engine_config["engine_class"]
+                            print(f"[POLLING_SCHEDULER] Creating new {engine_class.__name__} instance for {user_id}/{service_name}")
+                            # Ensure all dependencies for engine_class are passed correctly
                             engine_instance = engine_class(
                                 user_id=user_id,
-                                task_queue=task_queue, # global
-                                memory_backend=memory_backend, # global
-                                websocket_manager=manager, # global websocket_manager
-                                mongo_manager_instance=mongo_manager # global
+                                task_queue=task_queue, # Pass your global/initialized instances
+                                memory_backend=memory_backend,
+                                websocket_manager=manager, # Global websocket_manager
+                                mongo_manager_instance=mongo_manager # Global mongo_manager
                             )
                             if user_id not in active_context_engines:
                                 active_context_engines[user_id] = {}
-                            active_context_engines[user_id][engine_category] = engine_instance
-                            # No need to call engine_instance.initialize_polling_state() here,
-                            # run_poll_cycle will handle it if state is missing or needs reset.
+                            active_context_engines[user_id][service_name] = engine_instance
                         else:
-                            print(f"[POLLING_SCHEDULER_ERROR] No engine class configured for category: {engine_category}")
-                            # Release lock if instance can't be created
-                            await mongo_manager.update_polling_state(user_id, engine_category, {"is_currently_polling": False})
+                            print(f"[POLLING_SCHEDULER_ERROR] No engine class configured for service: {service_name}")
+                            await mongo_manager.update_polling_state(user_id, service_name, {"is_currently_polling": False})
                             continue
 
-                    # Run the poll cycle in a new task so the scheduler doesn't block
-                    asyncio.create_task(engine_instance.run_poll_cycle())
+                    asyncio.create_task(engine_instance.run_poll_cycle()) 
                 else:
-                    print(f"[POLLING_SCHEDULER] Could not acquire lock for {user_id}/{engine_category} (already processing or not due).")
+                    print(f"[POLLING_SCHEDULER] Could not acquire lock for {user_id}/{service_name} (already processing or no longer due).")
 
         except Exception as e:
             print(f"[POLLING_SCHEDULER_ERROR] Error in scheduler loop: {e}")
-            traceback.print_exc()
+            traceback.print_exc() 
 
         await asyncio.sleep(POLLING_SCHEDULER_INTERVAL_SECONDS)
 
-
 async def start_user_context_engines(user_id: str):
     """Ensures polling state exists for all enabled services for a user."""
-    if user_id not in active_context_engines: # This dict might be less critical now for *running* engines
+    print(f"[CONTEXT_ENGINE_MGR] Starting/Ensuring context engines and polling states for user {user_id}.")
+    if user_id not in active_context_engines: 
         active_context_engines[user_id] = {}
 
+    if not mongo_manager:
+        print(f"[CONTEXT_ENGINE_MGR_ERROR] MongoManager not initialized. Cannot start engines for user {user_id}.")
+        return
+
     user_profile = await mongo_manager.get_user_profile(user_id)
     user_settings = user_profile.get("userData", {}) if user_profile else {}
 
     for service_name, config in DATA_SOURCES_CONFIG.items():
-        # Determine if the service is enabled for the user
-        # This could come from user_profile's userData, or a default.
-        # Example: is_service_enabled = user_settings.get(f"{service_name}_polling_enabled", config["enabled_by_default"])
-        # For now, let's assume we check a specific field or default
-        
-        is_service_enabled_in_db = False
-        polling_state_doc = await mongo_manager.get_polling_state(user_id, service_name)
-        if polling_state_doc:
-            is_service_enabled_in_db = polling_state_doc.get("is_enabled", False) # Default to False if key missing
-        else: # No state yet, use default from config
-            is_service_enabled_in_db = config.get("enabled_by_default", True)
+        # Default to enabled_by_default if the specific setting isn't in user_profile
+        # This user-specific setting in user_profile would be set by /utils/set_data_source_enabled
+        is_service_explicitly_enabled_in_profile = user_settings.get(f"{service_name}_polling_enabled")
 
+        is_service_considered_enabled = False
+        if is_service_explicitly_enabled_in_profile is not None:
+            is_service_considered_enabled = is_service_explicitly_enabled_in_profile
+        else: # Not set in profile, use default from config
+            is_service_considered_enabled = config.get("enabled_by_default", True)
 
-        if is_service_enabled_in_db:
-            # Check if an engine instance exists (less critical now, but can be kept for potential direct calls)
+
+        if is_service_considered_enabled:
             engine_instance = active_context_engines.get(user_id, {}).get(service_name)
-            if not engine_instance:
+            if not engine_instance: # Engine instance not in memory, create one for initialization
                 engine_class = config["engine_class"]
                 print(f"[CONTEXT_ENGINE_MGR] Creating transient {engine_class.__name__} instance for {user_id}/{service_name} for state init.")
                 engine_instance = engine_class(
                     user_id=user_id,
-                    task_queue=task_queue, # Pass your global/initialized instances
+                    task_queue=task_queue, 
                     memory_backend=memory_backend,
-                    websocket_manager=manager,
-                    mongo_manager_instance=mongo_manager
+                    websocket_manager=manager, # global websocket_manager
+                    mongo_manager_instance=mongo_manager # global mongo_manager
                 )
-                # active_context_engines[user_id][service_name] = engine_instance # Optional to store
+                # We might not need to store it in active_context_engines if scheduler recreates on demand
+                # However, initialize_polling_state needs an instance.
 
-            # Crucially, ensure the polling state is initialized in the DB
-            # This will set it up for the central scheduler if it's new or enable it.
             await engine_instance.initialize_polling_state() 
             print(f"[CONTEXT_ENGINE_MGR] Polling state ensured for {service_name} for user {user_id}.")
         else:
             print(f"[CONTEXT_ENGINE_MGR] Service {service_name} is disabled for user {user_id}. Skipping engine start/state init.")
+            # If disabling, also ensure the polling state in DB is marked as is_enabled=False
+            await mongo_manager.update_polling_state(user_id, service_name, {"is_enabled": False})
+
 
 @app.on_event("startup")
 async def startup_event():
-    # ... (your existing startup code: STT, TTS, DB init, agent_task_processor, memory_backend) ...
-    print(f"[FASTAPI_LIFECYCLE] App startup.")
-    global stt_model, tts_model, agent_task_processor_instance, polling_scheduler_task_handle
-
-    agent_task_processor_instance = agent_task_processor
-
+    print(f"[FASTAPI_LIFECYCLE] App startup at {datetime.now(timezone.utc).isoformat()}")
+    global agent_task_processor_instance, polling_scheduler_task_handle, mongo_manager, task_queue, memory_backend, manager
+    # ... (STT, TTS init as before) ...
     print("[LIFECYCLE] Loading STT...")
     try:
         stt_model = FasterWhisperSTT(model_size="base", device="cpu", compute_type="int8")
@@ -870,11 +874,29 @@ async def startup_event():
     except Exception as e:
         print(f"[ERROR] TTS model failed to load. Voice features will be unavailable. Details: {e}")
         tts_model = None
+
+    mongo_manager = MongoManager() 
+    await mongo_manager.initialize_db() 
 
-    await task_queue.initialize_db()
-    await mongo_manager.initialize_db() # This now includes polling state indexes
-    await memory_backend.memory_queue.initialize_db()
+    # Assuming task_queue and memory_backend are initialized globally or here using mongo_manager
+    # task_queue = TaskQueue() # If it uses its own DB connection
+    # memory_backend = MemoryBackend(mongo_manager=mongo_manager) # Pass the initialized manager
 
+    # Re-initialize these if they depend on mongo_manager being set
+    if 'task_queue' not in globals() or task_queue is None:
+        task_queue = TaskQueue() # Initialize if not already
+        await task_queue.initialize_db()
+
+    if 'memory_backend' not in globals() or memory_backend is None:
+        memory_backend = MemoryBackend(mongo_manager=mongo_manager)
+        await memory_backend.memory_queue.initialize_db()
+
+    if 'manager' not in globals() or manager is None: # Assuming 'manager' is your WebSocketManager
+        manager = WebSocketManager()
+
+
+    agent_task_processor_instance = agent_task_processor 
+
     if agent_task_processor_instance:
         asyncio.create_task(agent_task_processor_instance.process_queue())
         asyncio.create_task(agent_task_processor_instance.cleanup_tasks_periodically())
@@ -883,11 +905,10 @@ async def startup_event():
 
     asyncio.create_task(memory_backend.process_memory_operations())
 
-    # Start the central polling scheduler
     polling_scheduler_task_handle = asyncio.create_task(polling_scheduler_loop())
     print(f"[FASTAPI_LIFECYCLE] Central polling scheduler started.")
 
-    print(f"[FASTAPI_LIFECYCLE] App startup complete.")
+    print(f"[FASTAPI_LIFECYCLE] App startup complete at {datetime.now(timezone.utc).isoformat()}")
 
 
 @app.on_event("shutdown")