feat (multi-tenancy): kafka integration for dynamic polling

Author: itsskofficial

src/server/app/app.py

@@ -148,6 +148,88 @@
 
 oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
 
+# Define DATA_SOURCES_CONFIG to map service_name to engine class
+DATA_SOURCES_CONFIG: Dict[str, Dict[str, Any]] = {
+    "gmail": {
+        "engine_class": GmailContextEngine,
+        "enabled_by_default": True, # Or fetch from user settings
+        # Add other service-specific configs if needed
+    },
+    # "gcalendar": {
+    #     "engine_class": GCalendarContextEngine,
+    #     "enabled_by_default": True,
+    # },
+    # "internet_search": {
+    #     "engine_class": InternetSearchContextEngine,
+    #     "enabled_by_default": False, # Example: off by default
+    # }
+}
+
+POLLING_SCHEDULER_INTERVAL_SECONDS = int(os.getenv("POLLING_SCHEDULER_INTERVAL_SECONDS", 30))
+async def polling_scheduler_loop():
+    print(f"[POLLING_SCHEDULER] Starting loop (interval: {POLLING_SCHEDULER_INTERVAL_SECONDS}s)")
+    await mongo_manager.reset_stale_polling_locks() # Reset any locks from crashed previous runs
+    
+    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() # Fetches tasks that are enabled, due, and not locked
+            
+            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 polling tasks.")
+
+            for task_state in due_tasks_states:
+                user_id = task_state["user_id"]
+                service_name = task_state["service_name"] # Changed from engine_category
+                
+                print(f"[POLLING_SCHEDULER] Attempting to process task for {user_id}/{service_name}")
+
+                # Atomically try to acquire the lock for this specific task
+                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}/{service_name}. Triggering poll.")
+                    
+                    engine_instance = active_context_engines.get(user_id, {}).get(service_name)
+                    
+                    if not engine_instance:
+                        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}")
+                            engine_instance = engine_class(
+                                user_id=user_id,
+                                task_queue=task_queue, 
+                                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][service_name] = engine_instance
+                        else:
+                            print(f"[POLLING_SCHEDULER_ERROR] No engine class configured for service: {service_name}")
+                            # Release lock as we can't process it
+                            await mongo_manager.update_polling_state(user_id, service_name, {"is_currently_polling": False})
+                            continue
+                    
+                    # Run the poll cycle in a new asyncio task so the scheduler doesn't block
+                    # The run_poll_cycle itself will handle releasing the lock via calculate_and_schedule_next_poll
+                    asyncio.create_task(engine_instance.run_poll_cycle()) 
+                else:
+                    # This means another scheduler instance/worker picked it up, or it's no longer 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() # Log full error
+        
+        await asyncio.sleep(POLLING_SCHEDULER_INTERVAL_SECONDS)
+        
+
 class Auth:
     async def _validate_token_and_get_payload(self, token: str) -> dict:
         # ... (Auth class definition remains the same) ...
@@ -711,43 +793,48 @@ async def polling_scheduler_loop():
 
 
 async def start_user_context_engines(user_id: str):
-    """Starts context engines for a given user if not already running and initializes polling state."""
-    if user_id not in active_context_engines:
+    """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
         active_context_engines[user_id] = {}
 
-    user_profile = await load_user_profile(user_id) # load_user_profile is your existing helper
-    user_settings = user_profile.get("userData", {})
-
-    for source_name, config in DATA_SOURCES_CONFIG.items():
-        is_enabled = user_settings.get(f"{source_name}Enabled", config["enabled_by_default"])
+    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
 
-        if is_enabled:
-            if source_name not in active_context_engines[user_id]:
-                print(f"[CONTEXT_ENGINE_MGR] Starting {source_name} engine for user {user_id}...")
+        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)
+
+
+        if is_service_enabled_in_db:
+            # Check if an engine instance exists (less critical now, but can be kept for potential direct calls)
+            engine_instance = active_context_engines.get(user_id, {}).get(service_name)
+            if not engine_instance:
                 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,
+                    task_queue=task_queue, # Pass your global/initialized instances
                     memory_backend=memory_backend,
                     websocket_manager=manager,
                     mongo_manager_instance=mongo_manager
                 )
-                active_context_engines[user_id][source_name] = engine_instance
-                # Initialize polling state (will set next_poll_at to now if new)
-                await engine_instance.initialize_polling_state() 
-                print(f"[CONTEXT_ENGINE_MGR] {source_name} engine started and polling state initialized for user {user_id}.")
-            else:
-                # Engine already active, ensure its polling state is initialized (e.g. if server restarted)
-                await active_context_engines[user_id][source_name].initialize_polling_state()
-                print(f"[CONTEXT_ENGINE_MGR] {source_name} engine for user {user_id} already active. Ensured polling state.")
-
-        elif not is_enabled and source_name in active_context_engines[user_id]:
-            print(f"[CONTEXT_ENGINE_MGR] {source_name} engine for user {user_id} is disabled. Stopping (if implemented) and removing.")
-            # Add engine stop logic if available:
-            # if hasattr(active_context_engines[user_id][source_name], 'stop'):
-            #     await active_context_engines[user_id][source_name].stop()
-            del active_context_engines[user_id][source_name]
-
+                # active_context_engines[user_id][service_name] = engine_instance # Optional to store
+            
+            # 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.")
 
 @app.on_event("startup")
 async def startup_event():

src/server/common/routes.py

@@ -474,23 +474,73 @@ async def websocket_endpoint(websocket: WebSocket):
         print(f"[WS /ws] WebSocket connection cleaned up for user: {authenticated_user_id or 'unknown'}")
 
 @router.post("/users/activity/heartbeat", status_code=status.HTTP_200_OK, summary="User Activity Heartbeat")
-async def user_activity_heartbeat(user_id: str = Depends(PermissionChecker(required_permissions=["write:profile"]))): # Assuming basic profile write permission
+async def user_activity_heartbeat(user_id: str = Depends(PermissionChecker(required_permissions=["write:profile"]))):
     print(f"[ENDPOINT /users/activity/heartbeat] Received heartbeat for user {user_id}")
     try:
         success = await mongo_manager.update_user_last_active(user_id)
         if success:
             return JSONResponse(content={"message": "User activity timestamp updated."})
         else:
-            # This might happen if the user profile doesn't exist yet, or DB error.
-            # update_user_last_active has upsert=True for userData, so it should create it.
-            # If it still fails, it's likely a DB connection issue or more serious.
-            print(f"[ERROR] Failed to update last_active_timestamp for user {user_id}")
+            print(f"[ERROR] Failed to update last_active_timestamp for user {user_id} via MongoManager.")
+            # This could be due to user_id not existing in user_profiles or a DB issue.
+            # update_user_last_active has upsert=True, so it should create if not exist.
             raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Failed to update user activity.")
     except Exception as e:
         print(f"[ERROR] /users/activity/heartbeat for {user_id}: {e}")
         traceback.print_exc()
         raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Error processing heartbeat.")
 
+@router.post("/set_data_source_enabled", status_code=status.HTTP_200_OK, summary="Enable/Disable Data Source")
+async def set_data_source_enabled_endpoint(request: SetDataSourceEnabledRequest, user_id: str = Depends(PermissionChecker(required_permissions=["write:config"]))):
+    # ... (existing code from your utils/routes.py) ...
+    # ADD an import from server.app.app import DATA_SOURCES_CONFIG, task_queue, memory_backend, manager as websocket_manager, mongo_manager
+    # Inside the try block, after successfully setting in user profile:
+    try:
+        # ... (your existing logic to update profile in mongo_manager.user_profiles_collection) ...
+        # profile = await mongo_manager.get_user_profile(user_id)
+        # if "userData" not in profile: profile["userData"] = {}
+        # profile["userData"][f"{request.source}Enabled"] = request.enabled
+        # success = await mongo_manager.update_user_profile(user_id, profile)
+        # if not success:
+        #     raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Failed to save data source status in user profile.")
+
+        # NEW: Update the polling_state_store directly for is_enabled
+        # and trigger re-initialization if being enabled.
+        polling_state_update = {"is_enabled": request.enabled}
+        if request.enabled:
+            # If enabling, also set next_scheduled_poll_time to now to trigger immediate consideration by scheduler
+            polling_state_update["next_scheduled_poll_time"] = datetime.datetime.now(datetime.timezone.utc)
+            polling_state_update["is_currently_polling"] = False # Ensure it's not locked
+            polling_state_update["error_backoff_until_timestamp"] = None # Clear any error backoff
+            polling_state_update["consecutive_empty_polls_count"] = 0 # Reset empty polls
+            # Optionally reset interval to default initial, or let it continue from where it was
+            # polling_state_update["current_polling_interval_seconds"] = DEFAULT_INITIAL_INTERVAL_SECONDS
+
+
+        success_polling_state = await mongo_manager.update_polling_state(user_id, request.source, polling_state_update)
+        
+        if not success_polling_state:
+            print(f"[WARN] Failed to update is_enabled in polling_state_store for {user_id}/{request.source}. May need manual sync if enabling.")
+        
+        # If it was just enabled and state didn't exist, initialize it fully
+        if request.enabled:
+            from server.app.app import DATA_SOURCES_CONFIG, task_queue, memory_backend, manager, mongo_manager as global_mongo_manager # Add imports
+            engine_config = DATA_SOURCES_CONFIG.get(request.source)
+            if engine_config and engine_config.get("engine_class"):
+                engine_class = engine_config["engine_class"]
+                # Create a temporary instance just to initialize state
+                temp_engine = engine_class(user_id, task_queue, memory_backend, manager, global_mongo_manager)
+                await temp_engine.initialize_polling_state() # This will create if not exists and set to poll soon
+                print(f"[set_data_source_enabled] Initialized polling state for {user_id}/{request.source} after enabling.")
+
+
+        return JSONResponse(content={"status": "success", "message": f"Data source '{request.source}' status set to {request.enabled}."})
+    except Exception as e:
+        # ... (your existing error handling) ...
+        print(f"[ERROR] /utils/set_data_source_enabled {user_id}: {e}")
+        traceback.print_exc()
+        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Failed to set data source status.")
+
 
 @router.post("/users/force-sync/{engine_category}", status_code=status.HTTP_200_OK, summary="Force Sync for a Service")
 async def force_sync_service(engine_category: str, user_id: str = Depends(PermissionChecker(required_permissions=["write:config"]))): # Assuming config write for this