diff --git a/backend/app/controller/chat_controller.py b/backend/app/controller/chat_controller.py index 69b4ce9bb..87e42a70a 100644 --- a/backend/app/controller/chat_controller.py +++ b/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( diff --git a/backend/app/model/chat.py b/backend/app/model/chat.py index e7e27a462..dac358dfd 100644 --- a/backend/app/model/chat.py +++ b/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): diff --git a/backend/app/service/chat_service.py b/backend/app/service/chat_service.py index 56300f8bf..3fc7de6ad 100644 --- a/backend/app/service/chat_service.py +++ b/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 diff --git a/backend/app/service/task.py b/backend/app/service/task.py index ee349cc66..604fbc717 100644 --- a/backend/app/service/task.py +++ b/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 diff --git a/backend/tests/unit/service/test_chat_service.py b/backend/tests/unit/service/test_chat_service.py index 9c4057432..36d753aa6 100644 --- a/backend/tests/unit/service/test_chat_service.py +++ b/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), ] diff --git a/backend/tests/unit/service/test_task.py b/backend/tests/unit/service/test_task.py index 219523380..e27dcbb5f 100644 --- a/backend/tests/unit/service/test_task.py +++ b/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") diff --git a/src/components/ChatBox/index.tsx b/src/components/ChatBox/index.tsx index 2c172496c..495b2fcd9 100644 --- a/src/components/ChatBox/index.tsx +++ b/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(''); @@ -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} diff --git a/src/store/chatStore.ts b/src/store/chatStore.ts index c0598e2ee..cd5acd287 100644 --- a/src/store/chatStore.ts +++ b/src/store/chatStore.ts @@ -845,17 +845,22 @@ const chatStore = (initial?: Partial) => ); } + 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) => 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'; }