refactor: integrate TodoListMiddleware and update related components

- Removed the write_todos tool as it is now included by default through TodoListMiddleware in the deep agent. - Updated the system prompt and documentation to reflect the integration of TodoListMiddleware, clarifying its capabilities for managing planning and todo lists. - Enhanced the chat handling logic to extract todos directly from the deep agent's command output, ensuring seamless user experience. - Refactored UI components to align with the new data structure and improve rendering of todo items, including updates to the Plan and TodoItem components. - Cleaned up code for better maintainability and readability, following recent refactoring efforts.
2026-05-02 20:32:39 +02:00 · 2025-12-27 15:18:34 +05:30 · 2025-12-27 15:18:34 +05:30 · c28a90fc29
commit c28a90fc29
parent 8a3ab3dfac
10 changed files with 172 additions and 419 deletions
--- a/surfsense_web/components/tool-ui/index.ts
+++ b/surfsense_web/components/tool-ui/index.ts
@ -77,10 +77,4 @@ export {
 	type PlanTodo,
 	type TodoStatus,
 } from "./plan";
-export {
-	WriteTodosToolUI,
-	WriteTodosArgsSchema,
-	WriteTodosResultSchema,
-	type WriteTodosArgs,
-	type WriteTodosResult,
-} from "./write-todos";
+export { WriteTodosToolUI, WriteTodosSchema, type WriteTodosData } from "./write-todos";
--- a/surfsense_web/components/tool-ui/plan/plan.tsx
+++ b/surfsense_web/components/tool-ui/plan/plan.tsx
@ -1,22 +1,16 @@
 "use client";

-import { CheckCircle2, Circle, CircleDashed, PartyPopper, XCircle } from "lucide-react";
+import { CheckCircle2, Circle, CircleDashed, ListTodo, PartyPopper, XCircle } from "lucide-react";
 import type { FC } 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 { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
 import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible";
 import { TextShimmerLoader } from "@/components/prompt-kit/loader";
 import { Progress } from "@/components/ui/progress";
 import { cn } from "@/lib/utils";
 import type { Action, ActionsConfig } from "../shared/schema";
-import type { PlanTodo, TodoStatus } from "./schema";
+import type { TodoStatus } from "./schema";

 // ============================================================================
 // Status Icon Component
@ -57,7 +51,7 @@ const StatusIcon: FC<StatusIconProps> = ({ status, className, isStreaming = true
 // ============================================================================

 interface TodoItemProps {
-	todo: PlanTodo;
+	todo: { id: string; content: string; status: TodoStatus };
 	/** When false, in_progress items show as static (no spinner/pulse) */
 	isStreaming?: boolean;
 }
@ -67,38 +61,22 @@ const TodoItem: FC<TodoItemProps> = ({ todo, isStreaming = true }) => {
 	// Only show shimmer animation if streaming and in progress
 	const isShimmer = todo.status === "in_progress" && isStreaming;

-	// Render the label with optional shimmer effect
-	const renderLabel = () => {
+	// Render the content with optional shimmer effect
+	const renderContent = () => {
 		if (isShimmer) {
-			return <TextShimmerLoader text={todo.label} size="md" />;
+			return <TextShimmerLoader text={todo.content} size="md" />;
 		}
 		return (
-			<span className={cn("text-sm", isStrikethrough && "line-through text-muted-foreground")}>
-				{todo.label}
+			<span className={cn("text-sm text-muted-foreground", isStrikethrough && "line-through")}>
+				{todo.content}
 			</span>
 		);
 	};

-	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} isStreaming={isStreaming} />
-						{renderLabel()}
-					</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} isStreaming={isStreaming} />
-			{renderLabel()}
+			{renderContent()}
 		</div>
 	);
 };
@ -110,8 +88,7 @@ const TodoItem: FC<TodoItemProps> = ({ todo, isStreaming = true }) => {
 export interface PlanProps {
 	id: string;
 	title: string;
-	description?: string;
-	todos: PlanTodo[];
+	todos: Array<{ id: string; content: string; status: TodoStatus }>;
 	maxVisibleTodos?: number;
 	showProgress?: boolean;
 	/** When false, in_progress items show as static (no spinner/pulse animations) */
@ -125,7 +102,6 @@ export interface PlanProps {
 export const Plan: FC<PlanProps> = ({
 	id,
 	title,
-	description,
 	todos,
 	maxVisibleTodos = 4,
 	showProgress = true,
@ -151,9 +127,6 @@ export const Plan: FC<PlanProps> = ({
 	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)) {
@ -172,22 +145,7 @@ export const Plan: FC<PlanProps> = ({
 		].filter(Boolean) as Action[];
 	}, [responseActions]);

-	// Get default expanded items (in_progress items with descriptions)
-	const defaultExpandedIds = useMemo(() => {
-		return todos.filter((t) => t.description && t.status === "in_progress").map((t) => t.id);
-	}, [todos]);
-
-	const TodoList: FC<{ items: PlanTodo[] }> = ({ items }) => {
-		if (hasDescriptions) {
-			return (
-				<Accordion type="multiple" defaultValue={defaultExpandedIds} className="w-full">
-					{items.map((todo) => (
-						<TodoItem key={todo.id} todo={todo} isStreaming={isStreaming} />
-					))}
-				</Accordion>
-			);
-		}
-
+	const TodoList: FC<{ items: typeof todos }> = ({ items }) => {
 		return (
 			<div className="space-y-0">
 				{items.map((todo) => (
@ -201,11 +159,9 @@ export const Plan: FC<PlanProps> = ({
 		<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 className="flex-1 min-w-0 flex items-center gap-2">
+						<ListTodo className="size-5 text-muted-foreground shrink-0" />
+						<CardTitle className="text-base font-semibold text-muted-foreground">{title}</CardTitle>
 					</div>
 					{isAllComplete && (
 						<div className="flex items-center gap-1 text-emerald-500">
@ -216,13 +172,13 @@ export const Plan: FC<PlanProps> = ({

 				{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 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 bg-muted [&>div]:bg-muted-foreground" />
 					</div>
 				)}
 			</CardHeader>
--- a/surfsense_web/components/tool-ui/plan/schema.ts
+++ b/surfsense_web/components/tool-ui/plan/schema.ts
@ -8,23 +8,25 @@ export type TodoStatus = z.infer<typeof TodoStatusSchema>;

 /**
 * Single todo item in a plan
+ * Matches deepagents TodoListMiddleware output: { content, status }
+ * id is auto-generated if not provided
 */
 export const PlanTodoSchema = z.object({
-	id: z.string(),
-	label: z.string(),
+	id: z.string().optional(),
+	content: z.string(),
 	status: TodoStatusSchema,
-	description: z.string().optional(),
 });

 export type PlanTodo = z.infer<typeof PlanTodoSchema>;

 /**
 * Serializable plan schema for tool results
+ * Matches deepagents TodoListMiddleware output format
+ * id/title are auto-generated if not provided
 */
 export const SerializablePlanSchema = z.object({
-	id: z.string(),
-	title: z.string(),
-	description: z.string().optional(),
+	id: z.string().optional(),
+	title: z.string().optional(),
 	todos: z.array(PlanTodoSchema).min(1),
 	maxVisibleTodos: z.number().optional(),
 	showProgress: z.boolean().optional(),
@ -33,9 +35,21 @@ export const SerializablePlanSchema = z.object({
 export type SerializablePlan = z.infer<typeof SerializablePlanSchema>;

 /**
- * Parse and validate a serializable plan from tool result
+ * Normalized plan with required fields (after auto-generation)
 */
-export function parseSerializablePlan(data: unknown): SerializablePlan {
+export interface NormalizedPlan {
+	id: string;
+	title: string;
+	todos: Array<{ id: string; content: string; status: TodoStatus }>;
+	maxVisibleTodos?: number;
+	showProgress?: boolean;
+}
+
+/**
+ * Parse and normalize a plan from tool result
+ * Auto-generates id/title if not provided (for deepagents compatibility)
+ */
+export function parseSerializablePlan(data: unknown): NormalizedPlan {
 	const result = SerializablePlanSchema.safeParse(data);

 	if (!result.success) {
@ -45,22 +59,33 @@ export function parseSerializablePlan(data: unknown): SerializablePlan {
 		const obj = (data && typeof data === "object" ? data : {}) as Record<string, unknown>;

 		return {
-			id: typeof obj.id === "string" ? obj.id : "unknown",
+			id: typeof obj.id === "string" ? obj.id : `plan-${Date.now()}`,
 			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 }],
+				? obj.todos.map((t: unknown, i: number) => {
+						const todo = t as Record<string, unknown>;
+						return {
+							id: typeof todo?.id === "string" ? todo.id : `todo-${i}`,
+							content: typeof todo?.content === "string" ? todo.content : "Task",
+							status: TodoStatusSchema.safeParse(todo?.status).success
+								? (todo.status as TodoStatus)
+								: ("pending" as const),
+						};
+					})
+				: [{ id: "1", content: "No tasks", status: "pending" as const }],
 		};
 	}

-	return result.data;
+	// Normalize: add id/title if missing
+	return {
+		id: result.data.id || `plan-${Date.now()}`,
+		title: result.data.title || "Plan",
+		todos: result.data.todos.map((t, i) => ({
+			id: t.id || `todo-${i}`,
+			content: t.content,
+			status: t.status,
+		})),
+		maxVisibleTodos: result.data.maxVisibleTodos,
+		showProgress: result.data.showProgress,
+	};
 }
--- a/surfsense_web/components/tool-ui/write-todos.tsx
+++ b/surfsense_web/components/tool-ui/write-todos.tsx
@ -14,73 +14,30 @@ import {
 import { Plan, PlanErrorBoundary, parseSerializablePlan, TodoStatusSchema } from "./plan";

 // ============================================================================
-// Zod Schemas
+// Zod Schemas - Matching deepagents TodoListMiddleware output
 // ============================================================================

 /**
- * Schema for a single todo item in the args
- * Note: Using nullish() with transform to convert null → undefined for Plan compatibility
+ * Schema for a single todo item (matches deepagents output)
 */
-const WriteTodosArgsTodoSchema = z.object({
-	id: z.string(),
+const TodoItemSchema = z.object({
 	content: z.string(),
 	status: TodoStatusSchema,
-	description: z
-		.string()
-		.nullish()
-		.transform((v) => v ?? undefined),
 });

 /**
- * Schema for write_todos tool arguments
- * Note: Using nullish() with transform to convert null → undefined for Plan compatibility
+ * Schema for write_todos tool args/result (matches deepagents output)
+ * deepagents provides: { todos: [{ content, status }] }
 */
-const WriteTodosArgsSchema = z.object({
-	title: z
-		.string()
-		.nullish()
-		.transform((v) => v ?? undefined),
-	description: z
-		.string()
-		.nullish()
-		.transform((v) => v ?? undefined),
-	todos: z.array(WriteTodosArgsTodoSchema).nullish(),
-});
-
-/**
- * Schema for a single todo item in the result
- * Note: Using nullish() with transform to convert null → undefined for Plan compatibility
- */
-const WriteTodosResultTodoSchema = z.object({
-	id: z.string(),
-	label: z.string(),
-	status: TodoStatusSchema,
-	description: z
-		.string()
-		.nullish()
-		.transform((v) => v ?? undefined),
-});
-
-/**
- * Schema for write_todos tool result
- * Note: Using nullish() with transform to convert null → undefined for Plan compatibility
- */
-const WriteTodosResultSchema = z.object({
-	id: z.string(),
-	title: z.string(),
-	description: z
-		.string()
-		.nullish()
-		.transform((v) => v ?? undefined),
-	todos: z.array(WriteTodosResultTodoSchema),
+const WriteTodosSchema = z.object({
+	todos: z.array(TodoItemSchema).nullish(),
 });

 // ============================================================================
 // Types
 // ============================================================================

-type WriteTodosArgs = z.infer<typeof WriteTodosArgsSchema>;
-type WriteTodosResult = z.infer<typeof WriteTodosResultSchema>;
+type WriteTodosData = z.infer<typeof WriteTodosSchema>;

 /**
 * Loading state component
@ -96,103 +53,65 @@ function WriteTodosLoading() {
 	);
 }

-/**
- * 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",
-			description: todo.description,
-		})),
-	};
-}
-
 /**
 * WriteTodos Tool UI Component
 *
 * Displays the agent's planning/todo list with a beautiful UI.
- * Shows progress, status indicators, and expandable details.
+ * Uses deepagents TodoListMiddleware output directly: { todos: [{ content, status }] }
 *
- * FIXED POSITION: When the same plan (by title) is updated multiple times,
+ * FIXED POSITION: When multiple write_todos calls happen in a conversation,
 * 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.
+ * shared state, and the first component reads from it.
 */
-export const WriteTodosToolUI = makeAssistantToolUI<WriteTodosArgs, WriteTodosResult>({
+export const WriteTodosToolUI = makeAssistantToolUI<WriteTodosData, WriteTodosData>({
 	toolName: "write_todos",
 	render: function WriteTodosUI({ args, result, status, toolCallId }) {
 		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
+		// Check if the THREAD is running
 		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";
+		// Use result if available, otherwise args (for streaming)
+		const data = result || args;
+		const hasTodos = data?.todos && data.todos.length > 0;

-		// SYNCHRONOUS ownership check - happens immediately, no race conditions
-		// ONE PLAN PER CONVERSATION: Only first write_todos call becomes owner
+		// Fixed title for all plans in conversation
+		const planTitle = "Plan";
+
+		// SYNCHRONOUS ownership check
 		const isOwner = useMemo(() => {
-			return registerPlanOwner(rawTitle, toolCallId);
-		}, [rawTitle, toolCallId]);
+			return registerPlanOwner(planTitle, toolCallId);
+		}, [planTitle, 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]);
+		// Get canonical title
+		const canonicalTitle = useMemo(() => getCanonicalPlanTitle(planTitle), [planTitle]);

-		// Register/update the plan state - ALWAYS use canonical title
+		// Register/update the plan state
 		useEffect(() => {
-			if (planData) {
+			if (hasTodos) {
+				const normalizedPlan = parseSerializablePlan({ todos: data.todos });
 				updatePlanState({
-					id: planData.id,
-					title: planTitle, // Use canonical title, not raw title
-					description: planData.description,
-					todos: planData.todos,
+					id: normalizedPlan.id,
+					title: canonicalTitle,
+					todos: normalizedPlan.todos,
 					toolCallId,
 				});
 			}
-		}, [planData, planTitle, updatePlanState, toolCallId]);
+		}, [data, hasTodos, canonicalTitle, 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
+		const currentPlanState = planStates.get(canonicalTitle);

-		// 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 we're NOT the owner, render nothing
 		if (!isOwner) {
 			return null;
 		}

-		// Loading state - tool is still running (no data yet)
+		// Loading state
 		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);
+			if (hasTodos) {
+				const plan = parseSerializablePlan({ todos: data.todos });
 				return (
 					<div className="my-4">
 						<PlanErrorBoundary>
@ -206,11 +125,8 @@ export const WriteTodosToolUI = makeAssistantToolUI<WriteTodosArgs, WriteTodosRe

 		// Incomplete/cancelled state
 		if (status.type === "incomplete") {
-			// For cancelled or errors, try to show what we have from args or shared state
-			// Use isThreadRunning to determine if we should still animate
-			const fallbackResult = currentPlanState || transformArgsToResult(args);
-			if (fallbackResult) {
-				const plan = parseSerializablePlan(fallbackResult);
+			if (currentPlanState || hasTodos) {
+				const plan = currentPlanState || parseSerializablePlan({ todos: data?.todos || [] });
 				return (
 					<div className="my-4">
 						<PlanErrorBoundary>
@ -222,23 +138,20 @@ export const WriteTodosToolUI = makeAssistantToolUI<WriteTodosArgs, WriteTodosRe
 			return null;
 		}

-		// Success - render the plan using the LATEST shared state
-		// 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;
+		// Success - render the plan
+		const planToRender = currentPlanState || (hasTodos ? parseSerializablePlan({ todos: data.todos }) : null);
 		if (!planToRender) {
 			return <WriteTodosLoading />;
 		}

-		const plan = parseSerializablePlan(planToRender);
 		return (
 			<div className="my-4">
 				<PlanErrorBoundary>
-					<Plan {...plan} showProgress={true} isStreaming={isThreadRunning} />
+					<Plan {...planToRender} showProgress={true} isStreaming={isThreadRunning} />
 				</PlanErrorBoundary>
 			</div>
 		);
 	},
 });

-export { WriteTodosArgsSchema, WriteTodosResultSchema, type WriteTodosArgs, type WriteTodosResult };
+export { WriteTodosSchema, type WriteTodosData };