refactor: introduce write_todos tool for enhanced task management

- Added a new write_todos tool to facilitate the creation and management of planning lists within the chat interface.
- Updated system prompt with detailed instructions on using the write_todos tool, including usage patterns and restrictions.
- Enhanced the chat message handling to support the new tool, ensuring proper integration and user experience.
- Implemented UI components for displaying and interacting with the planning lists, including progress tracking and status indicators.

surfsense_backend/app/agents/new_chat/system_prompt.py

@@ -109,6 +109,65 @@ You have access to the following tools:
     * This makes your response more visual and engaging.
     * Prioritize showing: diagrams, charts, infographics, key illustrations, or images that help explain the content.
     * Don't show every image - just the most relevant 1-3 images that enhance understanding.
+
+6. write_todos: Create and update a planning/todo list to break down complex tasks.
+  - Use this tool when you need to plan your approach to a complex task.
+  - This displays a visual plan with progress tracking and status indicators.
+  
+  - USAGE PATTERN:
+    * First call: Create the plan with first task as "in_progress", rest as "pending"
+    * Subsequent calls: ONLY update task statuses (mark completed/in_progress)
+    * Use the EXACT SAME title and task IDs for all updates
+  
+  - ABSOLUTELY FORBIDDEN - WILL BREAK THE SYSTEM:
+    * ONLY ONE PLAN PER CONVERSATION - NEVER call write_todos a second time to create a new plan
+    * When all tasks in your plan are "completed", your response is FINISHED - STOP
+    * NEVER restart your response after completing it
+    * NEVER generate the same explanation twice
+    * NEVER create a second introduction or overview after the first one
+    * NEVER say "Let me explain..." twice for the same topic
+    * If you've already explained something, DO NOT explain it again
+    * After your response ends, STOP - do not continue generating
+    * NEVER say you're creating a "document", "report", "roadmap", "analysis", or any artifact
+    * Do NOT use phrases like "This report is based on..." or "Based on my research..."
+    * Just answer the question directly - do not roleplay producing a deliverable
+  
+  - CORRECT BEHAVIOR:
+    * Call write_todos to update statuses as you progress
+    * Each section of your response appears EXACTLY ONCE
+    * When you finish explaining all tasks, your response is COMPLETE
+    * Do NOT generate additional content after concluding
+  
+  - CONTENT QUALITY:
+    * Provide thorough, detailed explanations for each task
+    * The restriction is on DUPLICATING content, not on depth or detail
+    * Each task deserves a complete, comprehensive explanation
+    * Be as detailed as needed - just don't repeat yourself
+  
+  - When to use:
+    * Breaking down a complex multi-step task (3-5 tasks recommended)
+    * Showing the user what steps you'll take to solve their problem
+    * Creating an implementation roadmap
+  
+  - Args:
+    - todos: List of todo items, each with:
+      * id: Unique identifier (KEEP SAME IDs across updates)
+      * content: Description of the task (KEEP SAME content across updates)
+      * status: "pending", "in_progress", or "completed"
+    - title: Title for the plan (MUST BE IDENTICAL across all updates)
+    - description: Optional context description
+  
+  - Returns: A visual plan card with progress bar and status indicators
+  
+  - CORRECT PATTERN:
+    1. Create plan with task 1 as "in_progress"
+    2. Explain task 1 content in detail
+    3. Update plan: task 1 "completed", task 2 "in_progress"
+    4. Explain task 2 content (NEW content, not repeating task 1)
+    5. Continue until all tasks are "completed"
+    6. When all tasks are "completed", your response is FINISHED
+    7. STOP IMMEDIATELY - do NOT create another plan or continue generating
+    8. ONE PLAN ONLY - never call write_todos again after completing all tasks
 </tools>
 <tool_call_examples>
 - User: "Fetch all my notes and what's in them?"
@@ -166,6 +225,22 @@ You have access to the following tools:
   - Then, if the content contains useful diagrams/images like `![Neural Network Diagram](https://example.com/nn-diagram.png)`:
     - Call: `display_image(src="https://example.com/nn-diagram.png", alt="Neural Network Diagram", title="Neural Network Architecture")`
   - Then provide your explanation, referencing the displayed image
+
+- User: "Help me implement a user authentication system"
+  - Step 1: Create plan with task 1 in_progress:
+    `write_todos(title="Auth Plan", todos=[{"id": "1", "content": "Design database schema", "status": "in_progress"}, {"id": "2", "content": "Set up password hashing", "status": "pending"}, {"id": "3", "content": "Create endpoints", "status": "pending"}])`
+  - Step 2: Provide DETAILED explanation of database schema design
+  - Step 3: Update plan (task 1 done, task 2 in_progress):
+    `write_todos(title="Auth Plan", todos=[{"id": "1", "content": "Design database schema", "status": "completed"}, {"id": "2", "content": "Set up password hashing", "status": "in_progress"}, {"id": "3", "content": "Create endpoints", "status": "pending"}])`
+  - Step 4: Provide DETAILED explanation of password hashing (NEW content only)
+  - Step 5: Update plan, explain endpoints in detail
+  - Step 6: Mark all complete, END response - DO NOT restart or regenerate
+  - FORBIDDEN: Do not go back and explain schema again after step 2
+
+- User: "How should I approach refactoring this large codebase?"
+  - Create plan, explain each step with thorough detail, update statuses as you go
+  - Each explanation is comprehensive but appears ONLY ONCE
+  - When finished with all tasks, STOP - do not continue generating
 </tool_call_examples>
 """
 

surfsense_backend/app/agents/new_chat/tools/registry.py

@@ -48,6 +48,7 @@ from .knowledge_base import create_search_knowledge_base_tool
 from .link_preview import create_link_preview_tool
 from .podcast import create_generate_podcast_tool
 from .scrape_webpage import create_scrape_webpage_tool
+from .write_todos import create_write_todos_tool
 
 # =============================================================================
 # Tool Definition
@@ -125,6 +126,13 @@ BUILTIN_TOOLS: list[ToolDefinition] = [
         ),
         requires=[],  # firecrawl_api_key is optional
     ),
+    # Planning/Todo tool - creates visual todo lists
+    ToolDefinition(
+        name="write_todos",
+        description="Create a planning/todo list to break down complex tasks",
+        factory=lambda deps: create_write_todos_tool(),
+        requires=[],
+    ),
     # =========================================================================
     # ADD YOUR CUSTOM TOOLS BELOW
     # =========================================================================

surfsense_backend/app/agents/new_chat/tools/write_todos.py

@@ -0,0 +1,94 @@
+"""
+Write todos tool for the SurfSense agent.
+
+This module provides a tool for creating and displaying a planning/todo list
+in the chat UI. It helps the agent break down complex tasks into steps.
+"""
+
+from typing import Any, Literal
+
+from langchain_core.tools import tool
+
+
+def create_write_todos_tool():
+    """
+    Factory function to create the write_todos tool.
+
+    Returns:
+        A configured tool function for writing todos/plans.
+    """
+
+    @tool
+    async def write_todos(
+        todos: list[dict[str, Any]],
+        title: str = "Planning Approach",
+        description: str | None = None,
+    ) -> dict[str, Any]:
+        """
+        Create a planning/todo list to break down a complex task.
+
+        Use this tool when you need to plan your approach to a complex task
+        or show the user a step-by-step breakdown of what you'll do.
+
+        This displays a visual plan with:
+        - Progress tracking (X of Y complete)
+        - Status indicators (pending, in progress, completed, cancelled)
+        - Expandable details for each step
+
+        Args:
+            todos: List of todo items. Each item should have:
+                - id: Unique identifier for the todo
+                - content: Description of the task
+                - status: One of "pending", "in_progress", "completed", "cancelled"
+            title: Title for the plan (default: "Planning Approach")
+            description: Optional description providing context
+
+        Returns:
+            A dictionary containing the plan data for the UI to render.
+
+        Example:
+            write_todos(
+                title="Implementation Plan",
+                description="Steps to add the new feature",
+                todos=[
+                    {"id": "1", "content": "Analyze requirements", "status": "completed"},
+                    {"id": "2", "content": "Design solution", "status": "in_progress"},
+                    {"id": "3", "content": "Write code", "status": "pending"},
+                    {"id": "4", "content": "Add tests", "status": "pending"},
+                ]
+            )
+        """
+        # Generate a unique plan ID
+        import uuid
+
+        plan_id = f"plan-{uuid.uuid4().hex[:8]}"
+
+        # Transform todos to the expected format for the UI
+        formatted_todos = []
+        for i, todo in enumerate(todos):
+            todo_id = todo.get("id", f"todo-{i}")
+            content = todo.get("content", "")
+            status = todo.get("status", "pending")
+
+            # Validate status
+            valid_statuses = ["pending", "in_progress", "completed", "cancelled"]
+            if status not in valid_statuses:
+                status = "pending"
+
+            formatted_todos.append(
+                {
+                    "id": todo_id,
+                    "label": content,
+                    "status": status,
+                }
+            )
+
+        return {
+            "id": plan_id,
+            "title": title,
+            "description": description,
+            "todos": formatted_todos,
+        }
+
+    return write_todos
+

surfsense_backend/app/tasks/chat/stream_new_chat.py

@@ -718,6 +718,25 @@ async def stream_new_chat(
                         status="completed",
                         items=completed_items,
                     )
+                elif tool_name == "write_todos":
+                    # Build completion items for planning
+                    if isinstance(tool_output, dict):
+                        plan_title = tool_output.get("title", "Plan")
+                        todos = tool_output.get("todos", [])
+                        todo_count = len(todos) if isinstance(todos, list) else 0
+                        completed_items = [
+                            *last_active_step_items,
+                            f"Plan: {plan_title[:50]}{'...' if len(plan_title) > 50 else ''}",
+                            f"Tasks: {todo_count} steps defined",
+                        ]
+                    else:
+                        completed_items = [*last_active_step_items, "Plan created"]
+                    yield streaming_service.format_thinking_step(
+                        step_id=original_step_id,
+                        title="Creating plan",
+                        status="completed",
+                        items=completed_items,
+                    )
                 else:
                     yield streaming_service.format_thinking_step(
                         step_id=original_step_id,
@@ -854,6 +873,28 @@ async def stream_new_chat(
                     yield streaming_service.format_terminal_info(
                         "Knowledge base search completed", "success"
                     )
+                elif tool_name == "write_todos":
+                    # Stream the full write_todos result so frontend can render the Plan component
+                    yield streaming_service.format_tool_output_available(
+                        tool_call_id,
+                        tool_output
+                        if isinstance(tool_output, dict)
+                        else {"result": tool_output},
+                    )
+                    # Send terminal message with plan info
+                    if isinstance(tool_output, dict):
+                        title = tool_output.get("title", "Plan")
+                        todos = tool_output.get("todos", [])
+                        todo_count = len(todos) if isinstance(todos, list) else 0
+                        yield streaming_service.format_terminal_info(
+                            f"Plan created: {title} ({todo_count} tasks)",
+                            "success",
+                        )
+                    else:
+                        yield streaming_service.format_terminal_info(
+                            "Plan created",
+                            "success",
+                        )
                 else:
                     # Default handling for other tools
                     yield streaming_service.format_tool_output_available(

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

@@ -18,6 +18,11 @@ import {
 	mentionedDocumentsAtom,
 	messageDocumentsMapAtom,
 } from "@/atoms/chat/mentioned-documents.atom";
+import {
+	clearPlanOwnerRegistry,
+	extractWriteTodosFromContent,
+	hydratePlanStateAtom,
+} from "@/atoms/chat/plan-state.atom";
 import { Thread } from "@/components/assistant-ui/thread";
 import { ChatHeader } from "@/components/new-chat/chat-header";
 import type { ThinkingStep } from "@/components/tool-ui/deepagent-thinking";
@@ -25,6 +30,7 @@ import { DisplayImageToolUI } from "@/components/tool-ui/display-image";
 import { GeneratePodcastToolUI } from "@/components/tool-ui/generate-podcast";
 import { LinkPreviewToolUI } from "@/components/tool-ui/link-preview";
 import { ScrapeWebpageToolUI } from "@/components/tool-ui/scrape-webpage";
+import { WriteTodosToolUI } from "@/components/tool-ui/write-todos";
 import { getBearerToken } from "@/lib/auth-utils";
 import { createAttachmentAdapter, extractAttachmentContent } from "@/lib/chat/attachment-adapter";
 import {
@@ -93,6 +99,7 @@ const PersistedAttachmentSchema = z.object({
 	id: z.string(),
 	name: z.string(),
 	type: z.string(),
+	contentType: z.string().optional(),
 	imageDataUrl: z.string().optional(),
 	extractedContent: z.string().optional(),
 });
@@ -155,6 +162,7 @@ function convertToThreadMessage(msg: MessageRecord): ThreadMessageLike {
 				id: att.id,
 				name: att.name,
 				type: att.type as "document" | "image" | "file",
+				contentType: att.contentType || "application/octet-stream",
 				status: { type: "complete" as const },
 				content: [],
 				// Custom fields for our ChatAttachment interface
@@ -181,6 +189,7 @@ const TOOLS_WITH_UI = new Set([
 	"link_preview",
 	"display_image",
 	"scrape_webpage",
+	"write_todos",
 ]);
 
 /**
@@ -213,6 +222,7 @@ export default function NewChatPage() {
 	const setMentionedDocumentIds = useSetAtom(mentionedDocumentIdsAtom);
 	const setMentionedDocuments = useSetAtom(mentionedDocumentsAtom);
 	const setMessageDocumentsMap = useSetAtom(messageDocumentsMapAtom);
+	const hydratePlanState = useSetAtom(hydratePlanStateAtom);
 
 	// Create the attachment adapter for file processing
 	const attachmentAdapter = useMemo(() => createAttachmentAdapter(), []);
@@ -248,6 +258,7 @@ export default function NewChatPage() {
 		setMentionedDocumentIds([]);
 		setMentionedDocuments([]);
 		setMessageDocumentsMap({});
+		clearPlanOwnerRegistry(); // Reset plan ownership for new chat
 
 		try {
 			if (urlChatId > 0) {
@@ -269,6 +280,11 @@ export default function NewChatPage() {
 							if (steps.length > 0) {
 								restoredThinkingSteps.set(`msg-${msg.id}`, steps);
 							}
+							// Hydrate write_todos plan state from persisted tool calls
+							const writeTodosCalls = extractWriteTodosFromContent(msg.content);
+							for (const todoData of writeTodosCalls) {
+								hydratePlanState(todoData);
+							}
 						}
 						if (msg.role === "user") {
 							const docs = extractMentionedDocuments(msg.content);
@@ -297,7 +313,7 @@ export default function NewChatPage() {
 		} finally {
 			setIsInitializing(false);
 		}
-	}, [urlChatId, setMessageDocumentsMap, setMentionedDocumentIds, setMentionedDocuments]);
+	}, [urlChatId, setMessageDocumentsMap, setMentionedDocumentIds, setMentionedDocuments, hydratePlanState]);
 
 	// Initialize on mount
 	useEffect(() => {
@@ -425,6 +441,7 @@ export default function NewChatPage() {
 						id: att.id,
 						name: att.name,
 						type: att.type,
+						contentType: (att as { contentType?: string }).contentType,
 						// Include imageDataUrl for images so they can be displayed after reload
 						imageDataUrl: (att as { imageDataUrl?: string }).imageDataUrl,
 						// Include extractedContent for context (already extracted, no re-processing needed)
@@ -844,6 +861,7 @@ export default function NewChatPage() {
 			<LinkPreviewToolUI />
 			<DisplayImageToolUI />
 			<ScrapeWebpageToolUI />
+			<WriteTodosToolUI />
 			<div className="flex flex-col h-[calc(100vh-64px)] overflow-hidden">
 				<Thread
 					messageThinkingSteps={messageThinkingSteps}

surfsense_web/atoms/chat/plan-state.atom.ts

@@ -0,0 +1,234 @@
+/**
+ * Plan State Atom
+ *
+ * Tracks the latest state of each plan by title.
+ * When write_todos is called multiple times with the same title,
+ * only the FIRST component renders (stays fixed in position),
+ * subsequent calls just update the shared state.
+ */
+
+import { atom } from "jotai";
+
+export interface PlanTodo {
+  id: string;
+  label: string;
+  status: "pending" | "in_progress" | "completed" | "cancelled";
+  description?: string;
+}
+
+export interface PlanState {
+  id: string;
+  title: string;
+  description?: string;
+  todos: PlanTodo[];
+  lastUpdated: number;
+  /** The toolCallId of the first component that rendered this plan */
+  ownerToolCallId: string;
+}
+
+/**
+ * SYNCHRONOUS ownership registry - prevents race conditions
+ * Only ONE plan allowed per conversation - first plan wins
+ */
+let firstPlanOwner: { toolCallId: string; title: string } | null = null;
+
+/**
+ * Register as owner of a plan SYNCHRONOUSLY
+ * ONE PLAN PER CONVERSATION: Only the first write_todos call renders.
+ * All subsequent calls update the state but don't render their own card.
+ */
+export function registerPlanOwner(title: string, toolCallId: string): boolean {
+  if (!firstPlanOwner) {
+    // First plan in this conversation - claim ownership
+    firstPlanOwner = { toolCallId, title };
+    return true;
+  }
+  
+  // Check if we're the owner
+  return firstPlanOwner.toolCallId === toolCallId;
+}
+
+/**
+ * Get the canonical title for a plan
+ * Returns the first plan's title if one exists, otherwise the provided title
+ */
+export function getCanonicalPlanTitle(title: string): string {
+  return firstPlanOwner?.title || title;
+}
+
+/**
+ * Check if a plan already exists in this conversation
+ */
+export function hasPlan(): boolean {
+  return firstPlanOwner !== null;
+}
+
+/**
+ * Get the first plan's info
+ */
+export function getFirstPlanInfo(): { toolCallId: string; title: string } | null {
+  return firstPlanOwner;
+}
+
+/**
+ * Check if a toolCallId is the owner of the plan SYNCHRONOUSLY
+ */
+export function isPlanOwner(toolCallId: string): boolean {
+  return !firstPlanOwner || firstPlanOwner.toolCallId === toolCallId;
+}
+
+/**
+ * Clear ownership registry (call when starting a new chat)
+ */
+export function clearPlanOwnerRegistry(): void {
+  firstPlanOwner = null;
+}
+
+/**
+ * Map of plan title -> latest plan state
+ * Using title as key since it stays constant across updates
+ */
+export const planStatesAtom = atom<Map<string, PlanState>>(new Map());
+
+/**
+ * Input type for updating plan state
+ */
+export interface UpdatePlanInput {
+  id: string;
+  title: string;
+  description?: string;
+  todos: PlanTodo[];
+  toolCallId: string;
+}
+
+/**
+ * Helper atom to update a plan state
+ */
+export const updatePlanStateAtom = atom(
+  null,
+  (get, set, plan: UpdatePlanInput) => {
+    const states = new Map(get(planStatesAtom));
+    
+    // Register ownership synchronously if not already done
+    registerPlanOwner(plan.title, plan.toolCallId);
+    
+    // Get the actual owner from the first plan
+    const ownerToolCallId = firstPlanOwner?.toolCallId || plan.toolCallId;
+    
+    // Always use the canonical (first) title for the plan key
+    const canonicalTitle = getCanonicalPlanTitle(plan.title);
+    
+    states.set(canonicalTitle, {
+      id: plan.id,
+      title: canonicalTitle,
+      description: plan.description,
+      todos: plan.todos,
+      lastUpdated: Date.now(),
+      ownerToolCallId,
+    });
+    set(planStatesAtom, states);
+  }
+);
+
+/**
+ * Helper atom to get the latest plan state by title
+ */
+export const getPlanStateAtom = atom((get) => {
+  const states = get(planStatesAtom);
+  return (title: string) => states.get(title);
+});
+
+/**
+ * Helper atom to clear all plan states (useful when starting a new chat)
+ */
+export const clearPlanStatesAtom = atom(null, (get, set) => {
+  clearPlanOwnerRegistry();
+  set(planStatesAtom, new Map());
+});
+
+/**
+ * Hydrate plan state from persisted message content
+ * Call this when loading messages from the database to restore plan state
+ */
+export interface HydratePlanInput {
+  toolCallId: string;
+  result: {
+    id?: string;
+    title?: string;
+    description?: string;
+    todos?: Array<{
+      id: string;
+      label: string;
+      status: "pending" | "in_progress" | "completed" | "cancelled";
+      description?: string;
+    }>;
+  };
+}
+
+export const hydratePlanStateAtom = atom(
+  null,
+  (get, set, plan: HydratePlanInput) => {
+    if (!plan.result?.todos || plan.result.todos.length === 0) return;
+    
+    const states = new Map(get(planStatesAtom));
+    const title = plan.result.title || "Planning Approach";
+    
+    // Register this as the owner if no plan exists yet
+    registerPlanOwner(title, plan.toolCallId);
+    
+    // Get the canonical title
+    const canonicalTitle = getCanonicalPlanTitle(title);
+    const ownerToolCallId = firstPlanOwner?.toolCallId || plan.toolCallId;
+    
+    // Only set if this is newer or doesn't exist
+    const existing = states.get(canonicalTitle);
+    if (!existing) {
+      states.set(canonicalTitle, {
+        id: plan.result.id || `plan-${Date.now()}`,
+        title: canonicalTitle,
+        description: plan.result.description,
+        todos: plan.result.todos,
+        lastUpdated: Date.now(),
+        ownerToolCallId,
+      });
+      set(planStatesAtom, states);
+    }
+  }
+);
+
+/**
+ * Extract write_todos tool call data from message content
+ * Returns an array of { toolCallId, result } for each write_todos call found
+ */
+export function extractWriteTodosFromContent(content: unknown): HydratePlanInput[] {
+  if (!Array.isArray(content)) return [];
+  
+  const results: HydratePlanInput[] = [];
+  
+  for (const part of content) {
+    if (
+      typeof part === "object" &&
+      part !== null &&
+      "type" in part &&
+      (part as { type: string }).type === "tool-call" &&
+      "toolName" in part &&
+      (part as { toolName: string }).toolName === "write_todos" &&
+      "toolCallId" in part &&
+      "result" in part
+    ) {
+      const toolCall = part as {
+        toolCallId: string;
+        result: HydratePlanInput["result"];
+      };
+      if (toolCall.result) {
+        results.push({
+          toolCallId: toolCall.toolCallId,
+          result: toolCall.result,
+        });
+      }
+    }
+  }
+  
+  return results;
+}
+

surfsense_web/components/tool-ui/index.ts

@@ -60,3 +60,17 @@ export {
 	type ScrapeWebpageResult,
 	ScrapeWebpageToolUI,
 } from "./scrape-webpage";
+export {
+	Plan,
+	PlanErrorBoundary,
+	type PlanProps,
+	parseSerializablePlan,
+	type SerializablePlan,
+	type PlanTodo,
+	type TodoStatus,
+} from "./plan";
+export {
+	WriteTodosToolUI,
+	type WriteTodosArgs,
+	type WriteTodosResult,
+} from "./write-todos";

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

@@ -0,0 +1,53 @@
+"use client";
+
+import { Component, type ReactNode } from "react";
+import { Card, CardContent } from "@/components/ui/card";
+
+export * from "./plan";
+export * from "./schema";
+
+// ============================================================================
+// Error Boundary
+// ============================================================================
+
+interface PlanErrorBoundaryProps {
+  children: ReactNode;
+  fallback?: ReactNode;
+}
+
+interface PlanErrorBoundaryState {
+  hasError: boolean;
+  error?: Error;
+}
+
+export class PlanErrorBoundary extends Component<PlanErrorBoundaryProps, PlanErrorBoundaryState> {
+  constructor(props: PlanErrorBoundaryProps) {
+    super(props);
+    this.state = { hasError: false };
+  }
+
+  static getDerivedStateFromError(error: Error): PlanErrorBoundaryState {
+    return { hasError: true, error };
+  }
+
+  render() {
+    if (this.state.hasError) {
+      if (this.props.fallback) {
+        return this.props.fallback;
+      }
+
+      return (
+        <Card className="w-full max-w-xl border-destructive/50">
+          <CardContent className="pt-6">
+            <div className="flex items-center gap-2 text-destructive">
+              <span className="text-sm">Failed to render plan</span>
+            </div>
+          </CardContent>
+        </Card>
+      );
+    }
+
+    return this.props.children;
+  }
+}
+

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

@@ -0,0 +1,265 @@
+"use client";
+
+import {
+  CheckCircle2,
+  Circle,
+  CircleDashed,
+  PartyPopper,
+  XCircle,
+} from "lucide-react";
+import type { FC, ReactNode } from "react";
+import { useMemo, useState } from "react";
+import {
+  Accordion,
+  AccordionContent,
+  AccordionItem,
+  AccordionTrigger,
+} from "@/components/ui/accordion";
+import { Button } from "@/components/ui/button";
+import { Card, CardContent, CardDescription, CardHeader, CardTitle } from "@/components/ui/card";
+import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible";
+import { Progress } from "@/components/ui/progress";
+import { cn } from "@/lib/utils";
+import type { Action, ActionsConfig } from "../shared/schema";
+import type { PlanTodo, TodoStatus } from "./schema";
+
+// ============================================================================
+// Status Icon Component
+// ============================================================================
+
+interface StatusIconProps {
+  status: TodoStatus;
+  className?: string;
+}
+
+const StatusIcon: FC<StatusIconProps> = ({ status, className }) => {
+  const baseClass = cn("size-4 shrink-0", className);
+
+  switch (status) {
+    case "completed":
+      return <CheckCircle2 className={cn(baseClass, "text-emerald-500")} />;
+    case "in_progress":
+      return (
+        <CircleDashed
+          className={cn(baseClass, "text-primary animate-spin")}
+          style={{ animationDuration: "3s" }}
+        />
+      );
+    case "cancelled":
+      return <XCircle className={cn(baseClass, "text-destructive")} />;
+    case "pending":
+    default:
+      return <Circle className={cn(baseClass, "text-muted-foreground")} />;
+  }
+};
+
+// ============================================================================
+// Todo Item Component
+// ============================================================================
+
+interface TodoItemProps {
+  todo: PlanTodo;
+}
+
+const TodoItem: FC<TodoItemProps> = ({ todo }) => {
+  const isStrikethrough = todo.status === "completed" || todo.status === "cancelled";
+  const isShimmer = todo.status === "in_progress";
+
+  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} />
+            <span
+              className={cn(
+                "text-sm text-left",
+                isStrikethrough && "line-through text-muted-foreground",
+                isShimmer && "animate-pulse"
+              )}
+            >
+              {todo.label}
+            </span>
+          </div>
+        </AccordionTrigger>
+        <AccordionContent className="pb-2 pl-6">
+          <p className="text-sm text-muted-foreground">{todo.description}</p>
+        </AccordionContent>
+      </AccordionItem>
+    );
+  }
+
+  return (
+    <div className="flex items-center gap-2 py-2">
+      <StatusIcon status={todo.status} />
+      <span
+        className={cn(
+          "text-sm",
+          isStrikethrough && "line-through text-muted-foreground",
+          isShimmer && "animate-pulse"
+        )}
+      >
+        {todo.label}
+      </span>
+    </div>
+  );
+};
+
+// ============================================================================
+// Plan Component
+// ============================================================================
+
+export interface PlanProps {
+  id: string;
+  title: string;
+  description?: string;
+  todos: PlanTodo[];
+  maxVisibleTodos?: number;
+  showProgress?: boolean;
+  responseActions?: Action[] | ActionsConfig;
+  className?: string;
+  onResponseAction?: (actionId: string) => void;
+  onBeforeResponseAction?: (actionId: string) => boolean;
+}
+
+export const Plan: FC<PlanProps> = ({
+  id,
+  title,
+  description,
+  todos,
+  maxVisibleTodos = 4,
+  showProgress = true,
+  responseActions,
+  className,
+  onResponseAction,
+  onBeforeResponseAction,
+}) => {
+  const [isExpanded, setIsExpanded] = useState(false);
+
+  // Calculate progress
+  const progress = useMemo(() => {
+    const completed = todos.filter((t) => t.status === "completed").length;
+    const total = todos.filter((t) => t.status !== "cancelled").length;
+    return { completed, total, percentage: total > 0 ? (completed / total) * 100 : 0 };
+  }, [todos]);
+
+  const isAllComplete = progress.completed === progress.total && progress.total > 0;
+
+  // Split todos for collapsible display
+  const visibleTodos = todos.slice(0, maxVisibleTodos);
+  const hiddenTodos = todos.slice(maxVisibleTodos);
+  const hasHiddenTodos = hiddenTodos.length > 0;
+
+  // Check if any todo has a description (for accordion mode)
+  const hasDescriptions = todos.some((t) => t.description);
+
+  // Handle action click
+  const handleAction = (actionId: string) => {
+    if (onBeforeResponseAction && !onBeforeResponseAction(actionId)) {
+      return;
+    }
+    onResponseAction?.(actionId);
+  };
+
+  // Normalize actions to array
+  const actionArray: Action[] = useMemo(() => {
+    if (!responseActions) return [];
+    if (Array.isArray(responseActions)) return responseActions;
+    return [
+      responseActions.confirm && { ...responseActions.confirm, id: "confirm" },
+      responseActions.cancel && { ...responseActions.cancel, id: "cancel" },
+    ].filter(Boolean) as Action[];
+  }, [responseActions]);
+
+  const TodoList: FC<{ items: PlanTodo[] }> = ({ items }) => {
+    if (hasDescriptions) {
+      return (
+        <Accordion type="single" collapsible className="w-full">
+          {items.map((todo) => (
+            <TodoItem key={todo.id} todo={todo} />
+          ))}
+        </Accordion>
+      );
+    }
+
+    return (
+      <div className="space-y-0">
+        {items.map((todo) => (
+          <TodoItem key={todo.id} todo={todo} />
+        ))}
+      </div>
+    );
+  };
+
+  return (
+    <Card id={id} className={cn("w-full max-w-xl", className)}>
+      <CardHeader className="pb-3">
+        <div className="flex items-start justify-between gap-2">
+          <div className="flex-1 min-w-0">
+            <CardTitle className="text-base font-semibold">{title}</CardTitle>
+            {description && (
+              <CardDescription className="mt-1 text-sm">{description}</CardDescription>
+            )}
+          </div>
+          {isAllComplete && (
+            <div className="flex items-center gap-1 text-emerald-500">
+              <PartyPopper className="size-5" />
+            </div>
+          )}
+        </div>
+
+        {showProgress && (
+          <div className="mt-3 space-y-1.5">
+            <div className="flex items-center justify-between text-xs text-muted-foreground">
+              <span>
+                {progress.completed} of {progress.total} complete
+              </span>
+              <span>{Math.round(progress.percentage)}%</span>
+            </div>
+            <Progress value={progress.percentage} className="h-1.5" />
+          </div>
+        )}
+      </CardHeader>
+
+      <CardContent className="pt-0">
+        <TodoList items={visibleTodos} />
+
+        {hasHiddenTodos && (
+          <Collapsible open={isExpanded} onOpenChange={setIsExpanded}>
+            <CollapsibleTrigger asChild>
+              <Button
+                variant="ghost"
+                size="sm"
+                className="w-full mt-2 text-xs text-muted-foreground hover:text-foreground"
+              >
+                {isExpanded
+                  ? "Show less"
+                  : `Show ${hiddenTodos.length} more ${hiddenTodos.length === 1 ? "task" : "tasks"}`}
+              </Button>
+            </CollapsibleTrigger>
+            <CollapsibleContent>
+              <TodoList items={hiddenTodos} />
+            </CollapsibleContent>
+          </Collapsible>
+        )}
+
+        {actionArray.length > 0 && (
+          <div className="flex flex-wrap gap-2 pt-4 mt-2 border-t">
+            {actionArray.map((action) => (
+              <Button
+                key={action.id}
+                variant={action.variant || "default"}
+                size="sm"
+                disabled={action.disabled}
+                onClick={() => handleAction(action.id)}
+              >
+                {action.label}
+              </Button>
+            ))}
+          </div>
+        )}
+      </CardContent>
+    </Card>
+  );
+};
+

surfsense_web/components/tool-ui/plan/schema.ts

@@ -0,0 +1,67 @@
+import { z } from "zod";
+import { ActionSchema } from "../shared/schema";
+
+/**
+ * Todo item status
+ */
+export const TodoStatusSchema = z.enum(["pending", "in_progress", "completed", "cancelled"]);
+export type TodoStatus = z.infer<typeof TodoStatusSchema>;
+
+/**
+ * Single todo item in a plan
+ */
+export const PlanTodoSchema = z.object({
+  id: z.string(),
+  label: z.string(),
+  status: TodoStatusSchema,
+  description: z.string().optional(),
+});
+
+export type PlanTodo = z.infer<typeof PlanTodoSchema>;
+
+/**
+ * Serializable plan schema for tool results
+ */
+export const SerializablePlanSchema = z.object({
+  id: z.string(),
+  title: z.string(),
+  description: z.string().optional(),
+  todos: z.array(PlanTodoSchema).min(1),
+  maxVisibleTodos: z.number().optional(),
+  showProgress: z.boolean().optional(),
+});
+
+export type SerializablePlan = z.infer<typeof SerializablePlanSchema>;
+
+/**
+ * Parse and validate a serializable plan from tool result
+ */
+export function parseSerializablePlan(data: unknown): SerializablePlan {
+  const result = SerializablePlanSchema.safeParse(data);
+  
+  if (!result.success) {
+    console.warn("Invalid plan data:", result.error.issues);
+    
+    // Try to extract basic info for fallback
+    const obj = (data && typeof data === "object" ? data : {}) as Record<string, unknown>;
+    
+    return {
+      id: typeof obj.id === "string" ? obj.id : "unknown",
+      title: typeof obj.title === "string" ? obj.title : "Plan",
+      description: typeof obj.description === "string" ? obj.description : undefined,
+      todos: Array.isArray(obj.todos)
+        ? obj.todos.map((t, i) => ({
+            id: typeof (t as any)?.id === "string" ? (t as any).id : `todo-${i}`,
+            label: typeof (t as any)?.label === "string" ? (t as any).label : "Task",
+            status: TodoStatusSchema.safeParse((t as any)?.status).success
+              ? (t as any).status
+              : "pending",
+            description: typeof (t as any)?.description === "string" ? (t as any).description : undefined,
+          }))
+        : [{ id: "1", label: "No tasks", status: "pending" as const }],
+    };
+  }
+  
+  return result.data;
+}
+

surfsense_web/components/tool-ui/shared/action-buttons.tsx

@@ -0,0 +1,42 @@
+"use client";
+
+import type { FC } from "react";
+import { Button } from "@/components/ui/button";
+import type { Action, ActionsConfig } from "./schema";
+
+interface ActionButtonsProps {
+  actions?: Action[] | ActionsConfig;
+  onAction?: (actionId: string) => void;
+  disabled?: boolean;
+}
+
+export const ActionButtons: FC<ActionButtonsProps> = ({ actions, onAction, disabled }) => {
+  if (!actions) return null;
+
+  // Normalize actions to array format
+  const actionArray: Action[] = Array.isArray(actions)
+    ? actions
+    : [
+        actions.confirm && { ...actions.confirm, id: "confirm" },
+        actions.cancel && { ...actions.cancel, id: "cancel" },
+      ].filter(Boolean) as Action[];
+
+  if (actionArray.length === 0) return null;
+
+  return (
+    <div className="flex flex-wrap gap-2 pt-3">
+      {actionArray.map((action) => (
+        <Button
+          key={action.id}
+          variant={action.variant || "default"}
+          size="sm"
+          disabled={disabled || action.disabled}
+          onClick={() => onAction?.(action.id)}
+        >
+          {action.label}
+        </Button>
+      ))}
+    </div>
+  );
+};
+

surfsense_web/components/tool-ui/shared/index.ts

@@ -0,0 +1,3 @@
+export * from "./schema";
+export * from "./action-buttons";
+

surfsense_web/components/tool-ui/shared/schema.ts

@@ -0,0 +1,24 @@
+import { z } from "zod";
+
+/**
+ * Shared action schema for tool UI components
+ */
+export const ActionSchema = z.object({
+  id: z.string(),
+  label: z.string(),
+  variant: z.enum(["default", "secondary", "destructive", "outline", "ghost", "link"]).optional(),
+  disabled: z.boolean().optional(),
+});
+
+export type Action = z.infer<typeof ActionSchema>;
+
+/**
+ * Actions configuration schema
+ */
+export const ActionsConfigSchema = z.object({
+  confirm: ActionSchema.optional(),
+  cancel: ActionSchema.optional(),
+});
+
+export type ActionsConfig = z.infer<typeof ActionsConfigSchema>;
+

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

@@ -0,0 +1,199 @@
+"use client";
+
+import { makeAssistantToolUI } from "@assistant-ui/react";
+import { useAtomValue, useSetAtom } from "jotai";
+import { Loader2 } from "lucide-react";
+import { useEffect, useMemo } from "react";
+import {
+  getCanonicalPlanTitle,
+  planStatesAtom,
+  registerPlanOwner,
+  updatePlanStateAtom,
+} from "@/atoms/chat/plan-state.atom";
+import { Plan, PlanErrorBoundary, parseSerializablePlan } from "./plan";
+
+/**
+ * Tool arguments for write_todos
+ */
+interface WriteTodosArgs {
+  title?: string;
+  description?: string;
+  todos?: Array<{
+    id: string;
+    content: string;
+    status: "pending" | "in_progress" | "completed" | "cancelled";
+  }>;
+}
+
+/**
+ * Tool result for write_todos
+ */
+interface WriteTodosResult {
+  id: string;
+  title: string;
+  description?: string;
+  todos: Array<{
+    id: string;
+    label: string;
+    status: "pending" | "in_progress" | "completed" | "cancelled";
+    description?: string;
+  }>;
+}
+
+/**
+ * Loading state component
+ */
+function WriteTodosLoading() {
+  return (
+    <div className="my-4 w-full max-w-xl rounded-2xl border bg-card/60 px-5 py-4 shadow-sm">
+      <div className="flex items-center gap-3">
+        <Loader2 className="size-5 animate-spin text-primary" />
+        <span className="text-sm text-muted-foreground">Creating plan...</span>
+      </div>
+    </div>
+  );
+}
+
+/**
+ * Transform tool args to result format
+ * This handles the case where the LLM is streaming the tool call
+ */
+function transformArgsToResult(args: WriteTodosArgs): WriteTodosResult | null {
+  if (!args.todos || !Array.isArray(args.todos) || args.todos.length === 0) {
+    return null;
+  }
+
+  return {
+    id: `plan-${Date.now()}`,
+    title: args.title || "Planning Approach",
+    description: args.description,
+    todos: args.todos.map((todo, index) => ({
+      id: todo.id || `todo-${index}`,
+      label: todo.content || "Task",
+      status: todo.status || "pending",
+    })),
+  };
+}
+
+/**
+ * WriteTodos Tool UI Component
+ *
+ * Displays the agent's planning/todo list with a beautiful UI.
+ * Shows progress, status indicators, and expandable details.
+ *
+ * FIXED POSITION: When the same plan (by title) is updated multiple times,
+ * only the FIRST component renders. Subsequent updates just update the
+ * shared state, and the first component reads from it. This prevents
+ * layout shift when plans are updated.
+ */
+export const WriteTodosToolUI = makeAssistantToolUI<WriteTodosArgs, WriteTodosResult>({
+  toolName: "write_todos",
+  render: function WriteTodosUI({ args, result, status, toolCallId }) {
+    const updatePlanState = useSetAtom(updatePlanStateAtom);
+    const planStates = useAtomValue(planStatesAtom);
+
+    // Get the plan data (from result or args)
+    const planData = result || transformArgsToResult(args);
+    const rawTitle = planData?.title || args.title || "Planning Approach";
+
+    // SYNCHRONOUS ownership check - happens immediately, no race conditions
+    // ONE PLAN PER CONVERSATION: Only first write_todos call becomes owner
+    const isOwner = useMemo(() => {
+      return registerPlanOwner(rawTitle, toolCallId);
+    }, [rawTitle, toolCallId]);
+    
+    // Get canonical title - always use the FIRST plan's title
+    // This ensures all updates go to the same plan state
+    const planTitle = useMemo(() => getCanonicalPlanTitle(rawTitle), [rawTitle]);
+
+    // Register/update the plan state - ALWAYS use canonical title
+    useEffect(() => {
+      if (planData) {
+        updatePlanState({
+          id: planData.id,
+          title: planTitle, // Use canonical title, not raw title
+          description: planData.description,
+          todos: planData.todos,
+          toolCallId,
+        });
+      }
+    }, [planData, planTitle, updatePlanState, toolCallId]);
+
+    // Update when result changes (for streaming updates)
+    useEffect(() => {
+      if (result) {
+        updatePlanState({
+          id: result.id,
+          title: planTitle, // Use canonical title, not raw title
+          description: result.description,
+          todos: result.todos,
+          toolCallId,
+        });
+      }
+    }, [result, planTitle, updatePlanState, toolCallId]);
+
+    // Get the current plan state (may be updated by other components)
+    const currentPlanState = planStates.get(planTitle);
+
+    // If we're NOT the owner, render nothing (the owner will render)
+    if (!isOwner) {
+      return null;
+    }
+
+    // Loading state - tool is still running
+    if (status.type === "running" || status.type === "requires-action") {
+      // Try to show partial results from args while streaming
+      const partialResult = transformArgsToResult(args);
+      if (partialResult) {
+        const plan = parseSerializablePlan(partialResult);
+        return (
+          <div className="my-4">
+            <PlanErrorBoundary>
+              <Plan {...plan} showProgress={true} />
+            </PlanErrorBoundary>
+          </div>
+        );
+      }
+      return <WriteTodosLoading />;
+    }
+
+    // Incomplete/cancelled state
+    if (status.type === "incomplete") {
+      if (status.reason === "cancelled") {
+        return null;
+      }
+      // 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} />
+            </PlanErrorBoundary>
+          </div>
+        );
+      }
+      return null;
+    }
+
+    // Success - render the plan using the LATEST shared state
+    // This way, even if our result is old, we show the latest data
+    const planToRender = currentPlanState || result;
+    if (!planToRender) {
+      return <WriteTodosLoading />;
+    }
+    
+    const plan = parseSerializablePlan(planToRender);
+    return (
+      <div className="my-4">
+        <PlanErrorBoundary>
+          <Plan {...plan} showProgress={true} />
+        </PlanErrorBoundary>
+      </div>
+    );
+  },
+});
+
+export type { WriteTodosArgs, WriteTodosResult };
+