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_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 };
+