fix: encryption issues

Author: itsskofficial

src/server/mcp_hub/orchestrator/tools.py

@@ -208,25 +208,31 @@ async def ask_user_clarification(ctx: Context, question: str, urgency: str = "no
     task_id = auth.get_task_id_from_context(ctx)
     db = state_manager.MongoManager()
     try:
+        # --- FIX: Perform a read-modify-write to handle encryption correctly ---
+        task = await db.get_task(task_id, user_id)
+        if not task:
+            raise ToolError(f"Task {task_id} not found.")
+
+        # `get_task` already decrypts, so `clarification_requests` is a list
+        current_requests = task.get("clarification_requests", [])
+        if not isinstance(current_requests, list):
+            logger.warning(f"Task {task_id} had a non-list 'clarification_requests' field. Overwriting.")
+            current_requests = []
+
         request_id = str(uuid.uuid4())
-        clarification_request = {
-            "request_id": request_id,
-            "question": question,
+        current_requests.append({
+            "request_id": request_id, "question": question,
             "asked_at": datetime.datetime.now(datetime.timezone.utc),
-            "response": None,
-            "responded_at": None,
-            "status": "pending"
-        }
-        await db.tasks_collection.update_one(
-            {"task_id": task_id, "user_id": user_id},
-            {
-                "$push": {"clarification_requests": clarification_request},
-                "$set": {
-                    "orchestrator_state.current_state": "SUSPENDED",
-                    "status": "clarification_pending"
-                }
-            }
-        )
+            "response": None, "responded_at": None, "status": "pending"
+        })
+
+        orchestrator_state = task.get("orchestrator_state", {})
+        if not isinstance(orchestrator_state, dict): orchestrator_state = {}
+        orchestrator_state["current_state"] = "SUSPENDED"
+
+        update_payload = {"clarification_requests": current_requests, "orchestrator_state": orchestrator_state, "status": "clarification_pending"}
+        await db.update_task(task_id, user_id, update_payload)
+
         await state_manager.add_execution_log(task_id, user_id, "clarification_requested", {"question": question}, reasoning)
 
         await notify_user(

src/server/workers/executor/tasks.py

@@ -20,6 +20,7 @@
 from celery import chord, group
 from main.llm import run_agent as run_main_agent, LLMProviderDownError
 from workers.utils.crypto import aes_decrypt, encrypt_doc, decrypt_doc
+from main.db import MongoManager
 
 # Load environment variables for the worker from its own config
 from workers.executor.config import (MONGO_URI, MONGO_DB_NAME,
@@ -267,7 +268,7 @@ async def async_execute_task_plan(task_id: str, user_id: str, run_id: str):
     await update_task_run_status(db, task_id, run_id, "processing", user_id, block_id=block_id)
     await add_progress_update(db, task_id, run_id, user_id, {"type": "info", "content": "Executor has picked up the task and is starting execution."}, block_id=block_id)
 
-    user_profile = await db.user_profiles.find_one({"user_id": user_id})
+    user_profile = await MongoManager().get_user_profile(user_id)
     personal_info = user_profile.get("userData", {}).get("personalInfo", {}) if user_profile else {}
     user_name = personal_info.get("name", "User")
     user_location_raw = personal_info.get("location", "Not specified")
@@ -757,7 +758,7 @@ async def add_sub_task_run_update(update_type: str, content: Any):
         await update_sub_task_safely(sub_task_id, {"status": "processing", "runs": [run_doc]})
 
         # 2. Get user context and integrations
-        user_profile = await db.user_profiles.find_one({"user_id": user_id})
+        user_profile = await MongoManager().get_user_profile(user_id)
         user_integrations = user_profile.get("userData", {}).get("integrations", {}) if user_profile else {} # noqa
 
         # 3. Configure tools for the sub-agent

src/server/workers/retention_tasks.py

@@ -1,6 +1,7 @@
 import logging
 import random
 import random
+from workers.utils.crypto import decrypt_doc
 from datetime import datetime, timedelta, timezone
 
 from workers.celery_app import celery_app
@@ -84,6 +85,9 @@ async def async_run_retention_campaigns():
 
         async for user_profile in users_cursor:
             user_id = user_profile.get("user_id")
+            # Decrypt the sensitive userData field before accessing it
+            decrypt_doc(user_profile, ["userData"])
+
             user_data = user_profile.get("userData", {})
             if not user_id:
                 continue

src/server/workers/tasks.py

@@ -520,10 +520,7 @@ async def async_generate_plan(task_id: str, user_id: str):
 
         original_context = task.get("original_context", {})
 
-        user_profile = await db_manager.user_profiles_collection.find_one(
-            {"user_id": user_id},
-            {"userData.personalInfo": 1, "userData.integrations": 1} # Projection to get only necessary data
-        )
+        user_profile = await db_manager.get_user_profile(user_id)
         if not user_profile:
             logger.error(f"User profile not found for user_id '{user_id}' associated with task {task_id}. Cannot generate plan.")
             return