refactor: enhance chat and tool UI with improved streaming and error handling

- Updated NewChatPage to persist partial responses when requests are cancelled by the user.
- Enhanced WriteTodosToolUI to check if the thread is running, improving the display of loading states and error handling.
- Modified Plan and TodoItem components to conditionally animate in-progress items based on streaming status, providing a clearer user experience during task management.

surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx

@@ -755,7 +755,21 @@ export default function NewChatPage() {
 				}
 			} catch (error) {
 				if (error instanceof Error && error.name === "AbortError") {
-					// Request was cancelled
+					// Request was cancelled by user - persist partial response if any content was received
+					const hasContent = contentParts.some(
+						(part) =>
+							(part.type === "text" && part.text.length > 0) ||
+							(part.type === "tool-call" && TOOLS_WITH_UI.has(part.toolName))
+					);
+					if (hasContent && currentThreadId) {
+						const partialContent = buildContentForPersistence();
+						appendMessage(currentThreadId, {
+							role: "assistant",
+							content: partialContent,
+						}).catch((err) =>
+							console.error("Failed to persist partial assistant message:", err)
+						);
+					}
 					return;
 				}
 				console.error("[NewChatPage] Chat error:", error);

surfsense_web/components/tool-ui/plan/plan.tsx

@@ -30,19 +30,27 @@ import type { PlanTodo, TodoStatus } from "./schema";
 interface StatusIconProps {
   status: TodoStatus;
   className?: string;
+  /** When false, in_progress items show as static (no spinner) */
+  isStreaming?: boolean;
 }
 
-const StatusIcon: FC<StatusIconProps> = ({ status, className }) => {
+const StatusIcon: FC<StatusIconProps> = ({ status, className, isStreaming = true }) => {
   const baseClass = cn("size-4 shrink-0", className);
 
   switch (status) {
     case "completed":
       return <CheckCircle2 className={cn(baseClass, "text-emerald-500")} />;
     case "in_progress":
+      // Only animate the spinner if we're actively streaming
+      // When streaming is stopped, show as a static dashed circle
       return (
         <CircleDashed
-          className={cn(baseClass, "text-primary animate-spin")}
+          className={cn(
-          style={{ animationDuration: "3s" }}
+            baseClass,
+            "text-primary",
+            isStreaming && "animate-spin"
+          )}
+          style={isStreaming ? { animationDuration: "3s" } : undefined}
         />
       );
     case "cancelled":
@@ -59,18 +67,21 @@ const StatusIcon: FC<StatusIconProps> = ({ status, className }) => {
 
 interface TodoItemProps {
   todo: PlanTodo;
+  /** When false, in_progress items show as static (no spinner/pulse) */
+  isStreaming?: boolean;
 }
 
-const TodoItem: FC<TodoItemProps> = ({ todo }) => {
+const TodoItem: FC<TodoItemProps> = ({ todo, isStreaming = true }) => {
   const isStrikethrough = todo.status === "completed" || todo.status === "cancelled";
-  const isShimmer = todo.status === "in_progress";
+  // Only show pulse animation if streaming and in progress
+  const isShimmer = todo.status === "in_progress" && isStreaming;
 
   if (todo.description) {
     return (
       <AccordionItem value={todo.id} className="border-0">
         <AccordionTrigger className="py-2 hover:no-underline">
           <div className="flex items-center gap-2">
-            <StatusIcon status={todo.status} />
+            <StatusIcon status={todo.status} isStreaming={isStreaming} />
             <span
               className={cn(
                 "text-sm text-left",
@@ -91,7 +102,7 @@ const TodoItem: FC<TodoItemProps> = ({ todo }) => {
 
   return (
     <div className="flex items-center gap-2 py-2">
-      <StatusIcon status={todo.status} />
+      <StatusIcon status={todo.status} isStreaming={isStreaming} />
       <span
         className={cn(
           "text-sm",
@@ -116,6 +127,8 @@ export interface PlanProps {
   todos: PlanTodo[];
   maxVisibleTodos?: number;
   showProgress?: boolean;
+  /** When false, in_progress items show as static (no spinner/pulse animations) */
+  isStreaming?: boolean;
   responseActions?: Action[] | ActionsConfig;
   className?: string;
   onResponseAction?: (actionId: string) => void;
@@ -129,6 +142,7 @@ export const Plan: FC<PlanProps> = ({
   todos,
   maxVisibleTodos = 4,
   showProgress = true,
+  isStreaming = true,
   responseActions,
   className,
   onResponseAction,
@@ -176,7 +190,7 @@ export const Plan: FC<PlanProps> = ({
       return (
         <Accordion type="single" collapsible className="w-full">
           {items.map((todo) => (
-            <TodoItem key={todo.id} todo={todo} />
+            <TodoItem key={todo.id} todo={todo} isStreaming={isStreaming} />
           ))}
         </Accordion>
       );
@@ -185,7 +199,7 @@ export const Plan: FC<PlanProps> = ({
     return (
       <div className="space-y-0">
         {items.map((todo) => (
-          <TodoItem key={todo.id} todo={todo} />
+          <TodoItem key={todo.id} todo={todo} isStreaming={isStreaming} />
         ))}
       </div>
     );

surfsense_web/components/tool-ui/write-todos.tsx

@@ -1,6 +1,6 @@
 "use client";
 
-import { makeAssistantToolUI } from "@assistant-ui/react";
+import { makeAssistantToolUI, useAssistantState } from "@assistant-ui/react";
 import { useAtomValue, useSetAtom } from "jotai";
 import { Loader2 } from "lucide-react";
 import { useEffect, useMemo } from "react";
@@ -92,6 +92,10 @@ export const WriteTodosToolUI = makeAssistantToolUI<WriteTodosArgs, WriteTodosRe
     const updatePlanState = useSetAtom(updatePlanStateAtom);
     const planStates = useAtomValue(planStatesAtom);
     
+    // Check if the THREAD is running (not just this tool)
+    // This hook subscribes to state changes, so it re-renders when thread stops
+    const isThreadRunning = useAssistantState(({ thread }) => thread.isRunning);
+
     // Get the plan data (from result or args)
     const planData = result || transformArgsToResult(args);
     const rawTitle = planData?.title || args.title || "Planning Approach";
@@ -140,7 +144,7 @@ export const WriteTodosToolUI = makeAssistantToolUI<WriteTodosArgs, WriteTodosRe
       return null;
     }
 
-    // Loading state - tool is still running
+    // Loading state - tool is still running (no data yet)
     if (status.type === "running" || status.type === "requires-action") {
       // Try to show partial results from args while streaming
       const partialResult = transformArgsToResult(args);
@@ -149,7 +153,7 @@ export const WriteTodosToolUI = makeAssistantToolUI<WriteTodosArgs, WriteTodosRe
         return (
           <div className="my-4">
             <PlanErrorBoundary>
-              <Plan {...plan} showProgress={true} />
+              <Plan {...plan} showProgress={true} isStreaming={isThreadRunning} />
             </PlanErrorBoundary>
           </div>
         );
@@ -159,17 +163,15 @@ export const WriteTodosToolUI = makeAssistantToolUI<WriteTodosArgs, WriteTodosRe
 
     // Incomplete/cancelled state
     if (status.type === "incomplete") {
-      if (status.reason === "cancelled") {
+      // For cancelled or errors, try to show what we have from args or shared state
-        return null;
+      // Use isThreadRunning to determine if we should still animate
-      }
-      // For errors, try to show what we have from args or shared state
       const fallbackResult = currentPlanState || transformArgsToResult(args);
       if (fallbackResult) {
         const plan = parseSerializablePlan(fallbackResult);
         return (
           <div className="my-4">
             <PlanErrorBoundary>
-              <Plan {...plan} showProgress={true} />
+              <Plan {...plan} showProgress={true} isStreaming={isThreadRunning} />
             </PlanErrorBoundary>
           </div>
         );
@@ -178,7 +180,8 @@ export const WriteTodosToolUI = makeAssistantToolUI<WriteTodosArgs, WriteTodosRe
     }
 
     // Success - render the plan using the LATEST shared state
-    // This way, even if our result is old, we show the latest data
+    // Use isThreadRunning to determine if we should animate in_progress items
+    // (LLM may still be working on tasks even though this tool call completed)
     const planToRender = currentPlanState || result;
     if (!planToRender) {
       return <WriteTodosLoading />;
@@ -188,7 +191,7 @@ export const WriteTodosToolUI = makeAssistantToolUI<WriteTodosArgs, WriteTodosRe
     return (
       <div className="my-4">
         <PlanErrorBoundary>
-          <Plan {...plan} showProgress={true} />
+          <Plan {...plan} showProgress={true} isStreaming={isThreadRunning} />
         </PlanErrorBoundary>
       </div>
     );