feat(chat): refine shared todo card

src/components/chat/hooks/useAgentChatSession.ts

@@ -139,19 +139,21 @@ const completeRunningProgress = (progress: ChatProgress[] | undefined) =>
     };
   });
 
-const cancelRunningTodos = (todos: AgentTodoUpdate[] | undefined) =>
-  todos?.map((todoUpdate) => ({
-    ...todoUpdate,
-    todos: todoUpdate.todos.map((todo) =>
-      todo.status === "pending" || todo.status === "in_progress"
-        ? {
-            ...todo,
-            status: "cancelled" as const,
-            updatedAt: Date.now(),
-          }
-        : todo,
-    ),
-  }));
+const cancelRunningTodos = (todoUpdate: AgentTodoUpdate | undefined) =>
+  todoUpdate
+    ? {
+        ...todoUpdate,
+        todos: todoUpdate.todos.map((todo) =>
+          todo.status === "pending" || todo.status === "in_progress"
+            ? {
+                ...todo,
+                status: "cancelled" as const,
+                updatedAt: Date.now(),
+              }
+            : todo,
+        ),
+      }
+    : undefined;
 
 const upsertPermission = (
   permissions: AgentPermissionRequest[] | undefined,
@@ -354,17 +356,64 @@ const applyQuestionResponse = (
       : question,
   );
 
-const upsertTodoUpdate = (
-  todos: AgentTodoUpdate[] | undefined,
+const createTodoUpdateFromEvent = (
   event: StreamEvent & { type: "todo_update" },
-) => [
-  {
-    sessionId: event.sessionId,
-    messageId: event.messageId,
-    todos: event.todos,
-    createdAt: event.createdAt,
-  },
-];
+): AgentTodoUpdate => ({
+  sessionId: event.sessionId,
+  messageId: event.messageId,
+  todos: event.todos,
+  createdAt: event.createdAt,
+});
+
+const normalizeSessionTodos = (
+  messages: Message[],
+  nextTodoUpdate?: AgentTodoUpdate,
+  targetAssistantMessageId?: string,
+) => {
+  let latestTodoUpdate = nextTodoUpdate;
+  if (!latestTodoUpdate) {
+    for (const message of messages) {
+      if (message.todos) {
+        latestTodoUpdate = message.todos;
+      }
+    }
+  }
+
+  if (!latestTodoUpdate) {
+    return messages;
+  }
+
+  const targetMessageId =
+    targetAssistantMessageId ??
+    [...messages].reverse().find((message) => message.role === "assistant")?.id;
+  if (!targetMessageId) {
+    return messages;
+  }
+
+  let changed = false;
+  const nextMessages = messages.map((message) => {
+    if (message.id === targetMessageId) {
+      if (message.todos === latestTodoUpdate) {
+        return message;
+      }
+      changed = true;
+      return {
+        ...message,
+        todos: latestTodoUpdate,
+      };
+    }
+    if (!message.todos) {
+      return message;
+    }
+    changed = true;
+    return {
+      ...message,
+      todos: undefined,
+    };
+  });
+
+  return changed ? nextMessages : messages;
+};
 
 const rejectOpenPermissionsAfterAbort = (
   permissions: AgentPermissionRequest[] | undefined,
@@ -425,7 +474,7 @@ const finalizeAssistantMessageAfterAbort = (message: Message): Message => {
     Boolean(rejectedPermissions?.length) ||
     Boolean(rejectedQuestions?.length) ||
     Boolean(completedProgress?.length) ||
-    Boolean(cancelledTodos?.length);
+    Boolean(cancelledTodos);
 
   if (!hasVisibleOutput) {
     return message;
@@ -544,7 +593,9 @@ export const useAgentChatSession = ({
         hydrationNonceRef.current += 1;
         titleUpdateNonceRef.current += 1;
 
-        setMessages(dedupeQuestionsAcrossMessages(loadedState.messages));
+        setMessages(
+          normalizeSessionTodos(dedupeQuestionsAcrossMessages(loadedState.messages)),
+        );
         setSessionTitle(loadedState.title);
         setIsSessionTitleManuallyEdited(loadedState.isTitleManuallyEdited ?? false);
         setSessionId(loadedState.sessionId);
@@ -654,8 +705,8 @@ export const useAgentChatSession = ({
       }
 
       if (event.type === "state") {
-        const nextMessages = dedupeQuestionsAcrossMessages(
-          cloneMessages(event.messages as Message[]),
+        const nextMessages = normalizeSessionTodos(
+          dedupeQuestionsAcrossMessages(cloneMessages(event.messages as Message[])),
         );
         messagesRef.current = nextMessages;
         setMessages(nextMessages);
@@ -774,13 +825,10 @@ export const useAgentChatSession = ({
         );
       } else if (event.type === "todo_update") {
         setMessages((prev) =>
-          prev.map((message) =>
-            message.id === assistantMessageId
-              ? {
-                  ...message,
-                  todos: upsertTodoUpdate(message.todos, event),
-                }
-              : message,
+          normalizeSessionTodos(
+            prev,
+            createTodoUpdateFromEvent(event),
+            assistantMessageId,
           ),
         );
       } else if (event.type === "done") {
@@ -914,7 +962,7 @@ export const useAgentChatSession = ({
                     message.content.trim().length === 0 &&
                     !(message.artifacts?.length) &&
                     !(message.progress?.length) &&
-                    !(message.todos?.length)
+                    !message.todos
                   ),
               ),
           );