fix (tasks): new history management for task-scoped chats

Author: itsskofficial

src/client/components/tasks/TaskDetailsContent.js

@@ -144,6 +144,7 @@ const TaskChatSection = ({ task, onSendChatMessage }) => {
 					<ChatBubble
 						key={index}
 						role={msg.role}
+						turn_steps={msg.turn_steps || []}
 						content={msg.content}
 						message={msg}
 					/>

src/server/main/chat/utils.py

@@ -126,6 +126,15 @@ def parse_assistant_response(messages: List[Dict[str, Any]]) -> Dict[str, Any]:
         role = msg.get("role", "")
         content = msg.get("content", "")
 
+        # ✅ Thoughts from <think> tags
+        if role == "assistant" and content and "<think>" in content:
+            think_matches = re.findall(r"<think>([\s\S]*?)</think>", content, re.DOTALL)
+            for match in think_matches:
+                turn_steps.append({
+                    "type": "thought",
+                    "content": match.strip()
+                })
+
         # ✅ Thoughts from <think> tags
         if role == "assistant" and content and "<think>" in content:
             think_matches = re.findall(r"<think>([\s\S]*?)</think>", content, re.DOTALL)

src/server/workers/planner/db.py

@@ -98,11 +98,8 @@ async def update_task_field(self, task_id: str, fields: dict):
         )
         logger.info(f"Updated fields for task {task_id}: {list(fields.keys())}")
 
-    async def update_task_with_plan(self, task_id: str, plan_data: dict, is_change_request: bool = False):
+    async def update_task_with_plan(self, task_id: str, plan_data: dict, is_change_request: bool = False, chat_history: Optional[List[Dict]] = None):
         """Updates a task with a generated plan and sets it to pending approval."""
-        SENSITIVE_PLAN_FIELDS = ["name", "description", "plan"]
-        encrypt_doc(plan_data, SENSITIVE_PLAN_FIELDS)
-
         plan_steps = plan_data.get("plan", [])
 
         update_doc = {
@@ -111,12 +108,19 @@ async def update_task_with_plan(self, task_id: str, plan_data: dict, is_change_r
             "updated_at": datetime.datetime.now(datetime.timezone.utc)
         }
 
+        if chat_history is not None:
+            update_doc["chat_history"] = chat_history
+
         # Only set the main description for the very first run.
         if not is_change_request:
             name = plan_data.get("name", "Proactively generated plan")
             update_doc["name"] = name
             update_doc["description"] = plan_data.get("description", "")
 
+        # Encrypt all sensitive fields at once before updating
+        SENSITIVE_TASK_FIELDS = ["name", "description", "plan", "chat_history"]
+        encrypt_doc(update_doc, SENSITIVE_TASK_FIELDS)
+
         result = await self.tasks_collection.update_one(
             {"task_id": task_id},
             {"$set": update_doc}

src/server/workers/planner/prompts.py

@@ -58,11 +58,12 @@
   ]
 }}
 
-Final Instructions:
+Final Instructions & Output Format:
 - Create a concise `name` for the task.
 - Create a concise `description` summarizing the overall goal.
 - Break down the goal into logical steps, choosing the most appropriate tool for each.
 - If an action item is not actionable with the given tools (e.g., "Think about the marketing report"), do not create a plan for it.
 - Do not include any text outside of the JSON object. Your response must begin with `{{` and end with `}}`.
 - ALWAYS RETURN THE JSON OBJECT.
+- CRITICAL: Your final response MUST be the JSON object described above, wrapped in `<answer>` tags. For example: `<answer>{{"name": "...", ...}}</answer>`.
 """

src/server/workers/tasks.py

@@ -14,6 +14,7 @@
 from main.analytics import capture_event
 from json_extractor import JsonExtractor
 from workers.utils.api_client import notify_user, push_task_list_update
+from main.chat.utils import parse_assistant_response
 from main.plans import PLAN_LIMITS
 from main.config import INTEGRATIONS_CONFIG
 from main.tasks.prompts import TASK_CREATION_PROMPT
@@ -464,6 +465,7 @@ def generate_plan_from_context(task_id: str, user_id: str):
 
 async def async_generate_plan(task_id: str, user_id: str):
     """Async logic for plan generation."""
+    # MODIFICATION: This function now also saves the full agent turn to chat_history
     db_manager = PlannerMongoManager()
     try:
         task = await db_manager.get_task(task_id)
@@ -489,29 +491,16 @@ async def async_generate_plan(task_id: str, user_id: str):
         # If the document is missing it for some reason, log a warning and proceed.
         if not task.get("user_id"):
             logger.warning(f"Task {task_id} document is missing user_id. Proceeding with passed user_id '{user_id}'.")
-            # Attempt to heal the document
             await db_manager.update_task_field(task_id, {"user_id": user_id})
 
         original_context = task.get("original_context", {})
 
-        # For re-planning, add previous results and chat history to the context
-        if task.get("chat_history"):
-            original_context["chat_history"] = task.get("chat_history")
-            original_context["previous_plan"] = task.get("plan")
-            original_context["previous_result"] = task.get("result")
-
         user_profile = await db_manager.user_profiles_collection.find_one(
             {"user_id": user_id},
             {"userData.personalInfo": 1} # Projection to get only necessary data
         )
         if not user_profile:
             logger.error(f"User profile not found for user_id '{user_id}' associated with task {task_id}. Cannot generate plan.")
-            await db_manager.update_task_status(task_id, "error", {"error": f"User profile not found for user_id '{user_id}'."})
-            return
-
-        if not user_profile:
-            logger.error(f"User profile not found for user_id '{user_id}' associated with task {task_id}. Cannot generate plan.")
-            await db_manager.update_task_status(task_id, "error", {"error": f"User profile not found for user_id '{user_id}'."})
             return
 
         personal_info = user_profile.get("userData", {}).get("personalInfo", {})
@@ -527,36 +516,82 @@ async def async_generate_plan(task_id: str, user_id: str):
             user_timezone = ZoneInfo(user_timezone_str)
         except ZoneInfoNotFoundError:
             logger.warning(f"Invalid timezone '{user_timezone_str}' for user {user_id}. Defaulting to UTC.")
+            user_timezone_str = "UTC"
             user_timezone = ZoneInfo("UTC")
 
         current_user_time = datetime.datetime.now(user_timezone).strftime('%Y-%m-%d %H:%M:%S %Z')
 
-        action_items = task.get("action_items", [])
-        if not action_items:
-            # This is likely a manually created task. Use its description as the action item.
-            logger.info(f"Task {task_id}: No 'action_items' field found. Using main description as the action.")
-            action_items = [task.get("description", "")]
-
         available_tools = get_all_mcp_descriptions()
-
         agent_config = get_planner_agent(available_tools, current_user_time, user_name, user_location)
 
-        user_prompt_content = "Please create a plan for the following action items:\n- " + "\n- ".join(action_items)
-        messages = [{'role': 'user', 'content': user_prompt_content}]
+        messages = []
+        if is_change_request:
+            logger.info(f"Task {task_id} has chat history. Constructing messages from history for re-planning.")
+            
+            previous_plan_str = json.dumps(task.get("plan", []), indent=2)
+            # Use default=str to handle non-serializable types like datetime
+            previous_result_str = json.dumps(task.get("result", "No previous result."), indent=2, default=str)
+            
+            context_message = (
+                "You are re-planning a task based on user feedback. Here is the context of the previous run:\n\n"
+                f"**Previous Plan:**\n```json\n{previous_plan_str}\n```\n\n"
+                f"**Previous Result:**\n```json\n{previous_result_str}\n```\n\n"
+                "Now, review the following conversation and generate a new plan based on the user's latest request."
+            )
+            messages.append({"role": "system", "content": context_message})
+
+            for msg in task["chat_history"]:
+                # This part is tricky. The planner doesn't need the full turn_steps of previous turns.
+                # It just needs the user/assistant content.
+                role = msg.get("role")
+                if role not in ["user", "assistant"]:
+                    continue
+                messages.append({
+                    "role": role,
+                    "content": msg.get("content")
+                })
+        else:
+            action_items = task.get("action_items", []) or [task.get("description", "")]
+            user_prompt_content = "Please create a plan for the following action items:\n- " + "\n- ".join(action_items)
+            messages = [{'role': 'user', 'content': user_prompt_content}]
 
         final_response_str = ""
+        final_history = None
         for chunk in run_main_agent(system_message=agent_config["system_message"], function_list=agent_config["function_list"], messages=messages):
+            final_history = chunk # Keep track of the latest state
             if isinstance(chunk, list) and chunk and chunk[-1].get("role") == "assistant":
                 final_response_str = chunk[-1].get("content", "")
 
         if not final_response_str:
             raise Exception("Planner agent returned no response.")
 
-        plan_data = JsonExtractor.extract_valid_json(clean_llm_output(final_response_str))
+        if not final_history:
+            raise Exception("Planner agent returned no history.")
+
+        # --- NEW: Parse the full agent turn and save it to chat_history ---
+        assistant_turn_start_index = next((i for i, msg in reversed(list(enumerate(final_history))) if msg.get("role") in ["user", "system"]), -1) + 1
+        assistant_messages = final_history[assistant_turn_start_index:]
+        parsed_turn = parse_assistant_response(assistant_messages)
+        plan_json_str = parsed_turn.get("final_content", "")
+        turn_steps = parsed_turn.get("turn_steps", [])
+
+        plan_data = JsonExtractor.extract_valid_json(plan_json_str)
         if not plan_data or "plan" not in plan_data:
-            raise Exception(f"Planner agent returned invalid JSON: {final_response_str}")
+            raise Exception(f"Planner agent returned invalid JSON inside <answer> tag: {plan_json_str}")
+
+        new_assistant_message = {
+            "message_id": str(uuid.uuid4()),
+            "role": "assistant",
+            "content": "I have generated a new plan based on your request. Please review it for approval.",
+            "turn_steps": turn_steps,
+            "timestamp": datetime.datetime.now(datetime.timezone.utc)
+        }
+
+        chat_history = task.get("chat_history", [])
+        if not isinstance(chat_history, list): chat_history = []
+        chat_history.append(new_assistant_message)
 
-        await db_manager.update_task_with_plan(task_id, plan_data, is_change_request)
+        await db_manager.update_task_with_plan(task_id, plan_data, is_change_request, chat_history)
         capture_event(user_id, "proactive_task_generated", {
             "task_id": task_id,
             "source": task.get("original_context", {}).get("source", "unknown"),