fix: use current attachments for follow-up questions instead of first… (#1167)

Co-authored-by: Wendong-Fan <w3ndong.fan@gmail.com>
Co-authored-by: Wendong-Fan <133094783+Wendong-Fan@users.noreply.github.com>

backend/app/controller/chat_controller.py

@@ -45,6 +45,7 @@ from app.service.task import (
     ActionSkipTaskData,
     ActionStopData,
     ActionSupplementData,
+    ImprovePayload,
     delete_task_lock,
     get_or_create_task_lock,
     get_task_lock,
@@ -224,7 +225,13 @@ async def post(data: Chat, request: Request):
 
     # Put initial action in queue to start processing
     await task_lock.put_queue(
-        ActionImproveData(data=data.question, new_task_id=data.task_id)
+        ActionImproveData(
+            data=ImprovePayload(
+                question=data.question,
+                attaches=data.attaches or [],
+            ),
+            new_task_id=data.task_id,
+        )
     )
 
     chat_logger.info(
@@ -331,7 +338,13 @@ def improve(id: str, data: SupplementChat):
 
     asyncio.run(
         task_lock.put_queue(
-            ActionImproveData(data=data.question, new_task_id=data.task_id)
+            ActionImproveData(
+                data=ImprovePayload(
+                    question=data.question,
+                    attaches=data.attaches or [],
+                ),
+                new_task_id=data.task_id,
+            )
         )
     )
     chat_logger.info(

backend/app/model/chat.py

@@ -132,6 +132,7 @@ class Chat(BaseModel):
 class SupplementChat(BaseModel):
     question: str
     task_id: str | None = None
+    attaches: list[str] = []
 
 
 class HumanReply(BaseModel):

backend/app/service/chat_service.py

@@ -462,6 +462,7 @@ async def step_solve(options: Chat, request: Request, task_lock: TaskLock):
                 # tracer.start()
                 if start_event_loop is True:
                     question = options.question
+                    attaches_to_use = options.attaches
                     logger.info(
                         "[NEW-QUESTION] Initial question"
                         " from options.question: "
@@ -470,7 +471,12 @@ async def step_solve(options: Chat, request: Request, task_lock: TaskLock):
                     start_event_loop = False
                 else:
                     assert isinstance(item, ActionImproveData)
-                    question = item.data
+                    question = item.data.question
+                    attaches_to_use = (
+                        item.data.attaches
+                        if item.data.attaches
+                        else options.attaches
+                    )
                     logger.info(
                         "[NEW-QUESTION] Follow-up "
                         "question from "
@@ -508,7 +514,7 @@ async def step_solve(options: Chat, request: Request, task_lock: TaskLock):
                 # Determine task complexity: attachments
                 # mean workforce, otherwise let agent decide
                 is_complex_task: bool
-                if len(options.attaches) > 0:
+                if len(attaches_to_use) > 0:
                     is_complex_task = True
                     logger.info(
                         "[NEW-QUESTION] Has attachments"
@@ -655,10 +661,10 @@ async def step_solve(options: Chat, request: Request, task_lock: TaskLock):
                     camel_task = Task(
                         content=clean_task_content, id=options.task_id
                     )
-                    if len(options.attaches) > 0:
+                    if len(attaches_to_use) > 0:
                         camel_task.additional_info = {
                             Path(file_path).name: file_path
-                            for file_path in options.attaches
+                            for file_path in attaches_to_use
                         }
 
                     # Stream decomposition in background

backend/app/service/task.py

@@ -71,9 +71,16 @@ class Action(str, Enum):
     timeout = "timeout"  # backend -> user (task timeout error)
 
 
+class ImprovePayload(BaseModel):
+    """User input payload for an improve action."""
+
+    question: str
+    attaches: list[str] = []
+
+
 class ActionImproveData(BaseModel):
     action: Literal[Action.improve] = Action.improve
-    data: str
+    data: ImprovePayload
     new_task_id: str | None = None
 
 

backend/tests/unit/service/test_chat_service.py

@@ -39,6 +39,7 @@ from app.service.task import (
     ActionEndData,
     ActionImproveData,
     ActionInstallMcpData,
+    ImprovePayload,
     TaskLock,
 )
 
@@ -854,7 +855,10 @@ class TestChatServiceIntegration:
         mock_task_lock.get_queue = AsyncMock(
             side_effect=[
                 # First call returns improve action
-                ActionImproveData(action=Action.improve, data="Test question"),
+                ActionImproveData(
+                    action=Action.improve,
+                    data=ImprovePayload(question="Test question"),
+                ),
                 # Second call returns end action
                 ActionEndData(action=Action.end),
             ]

backend/tests/unit/service/test_task.py

@@ -34,6 +34,7 @@ from app.service.task import (
     ActionTaskStateData,
     ActionUpdateTaskData,
     Agents,
+    ImprovePayload,
     TaskLock,
     create_task_lock,
     delete_task_lock,
@@ -52,20 +53,21 @@ class TestTaskServiceModels:
 
     def test_action_improve_data_creation(self):
         """Test ActionImproveData model creation."""
-        data = ActionImproveData(data="Improve this code")
+        payload = ImprovePayload(question="Improve this code")
+        data = ActionImproveData(data=payload)
 
         assert data.action == Action.improve
-        assert data.data == "Improve this code"
+        assert data.data.question == "Improve this code"
+        assert data.data.attaches == []
         assert data.new_task_id is None
 
     def test_action_improve_data_with_new_task_id(self):
         """Test ActionImproveData model creation with new_task_id."""
-        data = ActionImproveData(
-            data="Improve this code", new_task_id="task_123"
-        )
+        payload = ImprovePayload(question="Improve this code")
+        data = ActionImproveData(data=payload, new_task_id="task_123")
 
         assert data.action == Action.improve
-        assert data.data == "Improve this code"
+        assert data.data.question == "Improve this code"
         assert data.new_task_id == "task_123"
 
     def test_action_start_data_creation(self):
@@ -568,12 +570,14 @@ class TestTaskServiceIntegration:
         task_lock.add_human_input_listen(agent_name)
 
         # Test queue operations
-        improve_data = ActionImproveData(data="Improve this")
+        improve_data = ActionImproveData(
+            data=ImprovePayload(question="Improve this")
+        )
         await task_lock.put_queue(improve_data)
 
         retrieved_data = await task_lock.get_queue()
         assert retrieved_data.action == Action.improve
-        assert retrieved_data.data == "Improve this"
+        assert retrieved_data.data.question == "Improve this"
 
         # Test human input operations
         await task_lock.put_human_input(agent_name, "User response")

src/components/ChatBox/index.tsx

@@ -23,6 +23,7 @@ import {
 import useChatStoreAdapter from '@/hooks/useChatStoreAdapter';
 import { generateUniqueId, replayActiveTask } from '@/lib';
 import { useAuthStore } from '@/store/authStore';
+import { AgentStep, ChatTaskStatus } from '@/types/constants';
 import { Square, SquareCheckBig, TriangleAlert } from 'lucide-react';
 import { useCallback, useEffect, useMemo, useRef, useState } from 'react';
 import { useTranslation } from 'react-i18next';
@@ -30,7 +31,6 @@ import { useLocation, useNavigate, useSearchParams } from 'react-router-dom';
 import { toast } from 'sonner';
 import BottomBox from './BottomBox';
 import { ProjectChatContainer } from './ProjectChatContainer';
-import { AgentStep, ChatTaskStatus } from '@/types/constants';
 
 export default function ChatBox(): JSX.Element {
   const [message, setMessage] = useState<string>('');
@@ -257,7 +257,9 @@ export default function ChatBox(): JSX.Element {
       task.status === ChatTaskStatus.RUNNING ||
       task.status === ChatTaskStatus.PAUSE ||
       // splitting phase
-      task.messages.some((m) => m.step === AgentStep.TO_SUB_TASKS && !m.isConfirm) ||
+      task.messages.some(
+        (m) => m.step === AgentStep.TO_SUB_TASKS && !m.isConfirm
+      ) ||
       // skeleton/computing phase
       (!task.messages.find((m) => m.step === AgentStep.TO_SUB_TASKS) &&
         !task.hasWaitComfirm &&
@@ -428,13 +430,17 @@ export default function ChatBox(): JSX.Element {
       (task.status === ChatTaskStatus.RUNNING && task.hasMessages) ||
       task.status === ChatTaskStatus.PAUSE ||
       // splitting phase: has to_sub_tasks not confirmed OR skeleton computing
-      task.messages.some((m) => m.step === AgentStep.TO_SUB_TASKS && !m.isConfirm) ||
+      task.messages.some(
+        (m) => m.step === AgentStep.TO_SUB_TASKS && !m.isConfirm
+      ) ||
       (!task.messages.find((m) => m.step === AgentStep.TO_SUB_TASKS) &&
         !task.hasWaitComfirm &&
         task.messages.length > 0) ||
       task.isTakeControl ||
       // explicit confirm wait while task is pending but card not confirmed yet
-      (!!task.messages.find((m) => m.step === AgentStep.TO_SUB_TASKS && !m.isConfirm) &&
+      (!!task.messages.find(
+        (m) => m.step === AgentStep.TO_SUB_TASKS && !m.isConfirm
+      ) &&
         task.status === ChatTaskStatus.PENDING);
     const isReplayChatStore = task?.type === 'replay';
     if (!requiresHumanReply && isTaskBusy && !isReplayChatStore) {
@@ -471,6 +477,7 @@ export default function ChatBox(): JSX.Element {
           agent: chatStore.tasks[_taskId].activeAsk,
           reply: tempMessageContent,
         });
+        chatStore.setAttaches(_taskId, []);
         if (chatStore.tasks[_taskId].askList.length === 0) {
           chatStore.setActiveAsk(_taskId, '');
         } else {
@@ -509,7 +516,8 @@ export default function ChatBox(): JSX.Element {
           (hasWaitComfirm && !wasTaskStopped) ||
           (isFinished && !wasTaskStopped) ||
           (hasMessages &&
-            chatStore.tasks[_taskId as string].status === ChatTaskStatus.PENDING);
+            chatStore.tasks[_taskId as string].status ===
+              ChatTaskStatus.PENDING);
 
         if (shouldContinueConversation) {
           // Check if this is the very first message and task hasn't started
@@ -528,7 +536,8 @@ export default function ChatBox(): JSX.Element {
           // Only start a new task if: pending, no messages processed yet
           // OR while or after replaying a project
           if (
-            (chatStore.tasks[_taskId as string].status === ChatTaskStatus.PENDING &&
+            (chatStore.tasks[_taskId as string].status ===
+              ChatTaskStatus.PENDING &&
               !hasSimpleResponse &&
               !hasComplexTask &&
               !isFinished) ||
@@ -549,6 +558,7 @@ export default function ChatBox(): JSX.Element {
                 tempMessageContent,
                 attachesToSend
               );
+              chatStore.setAttaches(_taskId, []);
             } catch (err: any) {
               console.error('Failed to start task:', err);
               toast.error(
@@ -564,26 +574,30 @@ export default function ChatBox(): JSX.Element {
               '[Multi-turn] Continuing conversation with improve API'
             );
 
+            const attachesForThisTurn = JSON.parse(
+              JSON.stringify(chatStore.tasks[_taskId]?.attaches || [])
+            );
+            const improveAttaches =
+              attachesForThisTurn.map(
+                (f: { filePath: string }) => f.filePath
+              ) || [];
+
             //Generate nextId in case new chatStore is created to sync with the backend beforehand
             const nextTaskId = generateUniqueId();
             chatStore.setNextTaskId(nextTaskId);
 
             // Use improve endpoint (POST /chat/{id}) - {id} is project_id
-            // This reuses the existing SSE connection and step_solve loop
             fetchPost(`/chat/${projectStore.activeProjectId}`, {
               question: tempMessageContent,
               task_id: nextTaskId,
+              attaches: improveAttaches,
             });
             chatStore.setIsPending(_taskId, true);
-            // Add the user message to show it in UI
             chatStore.addMessages(_taskId, {
               id: generateUniqueId(),
               role: 'user',
               content: tempMessageContent,
-              attaches:
-                JSON.parse(
-                  JSON.stringify(chatStore.tasks[_taskId]?.attaches)
-                ) || [],
+              attaches: attachesForThisTurn,
             });
             chatStore.setAttaches(_taskId, []);
             setMessage('');
@@ -625,6 +639,7 @@ export default function ChatBox(): JSX.Element {
               attachesToSend
             );
             chatStore.setHasWaitComfirm(_taskId as string, true);
+            chatStore.setAttaches(_taskId, []);
           } catch (err: any) {
             console.error('Failed to start task:', err);
             toast.error(
@@ -670,10 +685,13 @@ export default function ChatBox(): JSX.Element {
       if (result.success && result.files && result.files.length > 0) {
         const taskId = chatStore.activeTaskId as string;
         const files = [
-          ...chatStore.tasks[taskId].attaches.filter(
-            (f) => !result.files.find((r: File) => r.filePath === f.filePath)
+          ...(chatStore.tasks[taskId].attaches || []),
+          ...result.files.filter(
+            (r: File) =>
+              !chatStore.tasks[taskId].attaches?.some(
+                (f: File) => f.filePath === r.filePath
+              )
           ),
-          ...result.files,
         ];
         chatStore.setAttaches(taskId, files);
       }
@@ -891,7 +909,10 @@ export default function ChatBox(): JSX.Element {
     }
 
     // Check task status
-    if (task.status === ChatTaskStatus.RUNNING || task.status === ChatTaskStatus.PAUSE) {
+    if (
+      task.status === ChatTaskStatus.RUNNING ||
+      task.status === ChatTaskStatus.PAUSE
+    ) {
       return 'running';
     }
 
@@ -939,7 +960,11 @@ export default function ChatBox(): JSX.Element {
       // Note: Replay creates a new chatstore, so no conflicts
       const task = chatStore.tasks[chatStore.activeTaskId as string];
       // Only skip backend call if task is finished or hasn't started yet (no messages)
-      if (task && task.messages.length > 0 && task.status !== ChatTaskStatus.FINISHED) {
+      if (
+        task &&
+        task.messages.length > 0 &&
+        task.status !== ChatTaskStatus.FINISHED
+      ) {
         try {
           await fetchDelete(`/chat/${project_id}/remove-task/${task_id}`, {
             project_id: project_id,
@@ -1011,7 +1036,8 @@ export default function ChatBox(): JSX.Element {
               taskStatus={chatStore.tasks[chatStore.activeTaskId]?.status}
               onReplay={handleReplay}
               replayDisabled={
-                chatStore.tasks[chatStore.activeTaskId]?.status !== ChatTaskStatus.FINISHED
+                chatStore.tasks[chatStore.activeTaskId]?.status !==
+                ChatTaskStatus.FINISHED
               }
               replayLoading={isReplayLoading}
               onPauseResume={handlePauseResume}

src/store/chatStore.ts

@@ -845,17 +845,22 @@ const chatStore = (initial?: Partial<ChatStore>) =>
                   );
                 }
 
+                const attachesForNewMessage =
+                  lastMessage?.role === 'user' && lastMessage?.attaches?.length
+                    ? lastMessage.attaches
+                    : [
+                        ...(previousChatStore.tasks[currentTaskId]?.attaches ||
+                          []),
+                        ...(messageAttaches || []),
+                      ];
+
                 //Trick: by the time the question is retrieved from event,
                 //the last message from previous chatStore is at display
                 newChatStore.getState().addMessages(newTaskId, {
                   id: generateUniqueId(),
                   role: 'user',
                   content: question || (messageContent as string),
-                  //TODO: The attaches that reach here (when Improve API is called) doesn't reach the backend
-                  attaches: [
-                    ...(previousChatStore.tasks[currentTaskId]?.attaches || []),
-                    ...(messageAttaches || []),
-                  ],
+                  attaches: attachesForNewMessage,
                 });
                 console.log('[NEW CHATSTORE] Created for ', project_id);
 
@@ -1398,7 +1403,9 @@ const chatStore = (initial?: Partial<ChatStore>) =>
                   setTaskAssigning(currentTaskId, [...taskAssigning]);
                 }
               }
-              const taskIndex = taskRunning.findIndex((task) => task.id === process_task_id);
+              const taskIndex = taskRunning.findIndex(
+                (task) => task.id === process_task_id
+              );
               if (taskIndex !== -1 && taskRunning[taskIndex].agent) {
                 taskRunning[taskIndex].agent!.status = 'completed';
               }