diff --git a/surfsense_backend/app/agents/new_chat/system_prompt.py b/surfsense_backend/app/agents/new_chat/system_prompt.py index 3c5ab55a0..d85a89db7 100644 --- a/surfsense_backend/app/agents/new_chat/system_prompt.py +++ b/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 - 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 """ diff --git a/surfsense_backend/app/agents/new_chat/tools/registry.py b/surfsense_backend/app/agents/new_chat/tools/registry.py index 3b0c2ddac..5c746d726 100644 --- a/surfsense_backend/app/agents/new_chat/tools/registry.py +++ b/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 # ========================================================================= diff --git a/surfsense_backend/app/agents/new_chat/tools/write_todos.py b/surfsense_backend/app/agents/new_chat/tools/write_todos.py new file mode 100644 index 000000000..95b5cb155 --- /dev/null +++ b/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 + diff --git a/surfsense_backend/app/tasks/chat/stream_new_chat.py b/surfsense_backend/app/tasks/chat/stream_new_chat.py index 8b326e1d1..5bb33e399 100644 --- a/surfsense_backend/app/tasks/chat/stream_new_chat.py +++ b/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( diff --git a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx index 20b71dd63..df092a630 100644 --- a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx +++ b/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() { +
latest plan state + * Using title as key since it stays constant across updates + */ +export const planStatesAtom = atom>(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; +} + diff --git a/surfsense_web/components/tool-ui/index.ts b/surfsense_web/components/tool-ui/index.ts index c863d8722..93474acdf 100644 --- a/surfsense_web/components/tool-ui/index.ts +++ b/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"; diff --git a/surfsense_web/components/tool-ui/plan/index.tsx b/surfsense_web/components/tool-ui/plan/index.tsx new file mode 100644 index 000000000..989daede9 --- /dev/null +++ b/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 { + 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 ( + + +
+ Failed to render plan +
+ + + ); + } + + return this.props.children; + } +} + diff --git a/surfsense_web/components/tool-ui/plan/plan.tsx b/surfsense_web/components/tool-ui/plan/plan.tsx new file mode 100644 index 000000000..a520ea416 --- /dev/null +++ b/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 = ({ status, className }) => { + const baseClass = cn("size-4 shrink-0", className); + + switch (status) { + case "completed": + return ; + case "in_progress": + return ( + + ); + case "cancelled": + return ; + case "pending": + default: + return ; + } +}; + +// ============================================================================ +// Todo Item Component +// ============================================================================ + +interface TodoItemProps { + todo: PlanTodo; +} + +const TodoItem: FC = ({ todo }) => { + const isStrikethrough = todo.status === "completed" || todo.status === "cancelled"; + const isShimmer = todo.status === "in_progress"; + + if (todo.description) { + return ( + + +
+ + + {todo.label} + +
+ + +

{todo.description}

+ + + ); + } + + return ( +
+ + + {todo.label} + +
+ ); +}; + +// ============================================================================ +// 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 = ({ + 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 ( + + {items.map((todo) => ( + + ))} + + ); + } + + return ( +
+ {items.map((todo) => ( + + ))} +
+ ); + }; + + return ( + + +
+
+ {title} + {description && ( + {description} + )} +
+ {isAllComplete && ( +
+ +
+ )} +
+ + {showProgress && ( +
+
+ + {progress.completed} of {progress.total} complete + + {Math.round(progress.percentage)}% +
+ +
+ )} + + + + + + {hasHiddenTodos && ( + + + + + + + + + )} + + {actionArray.length > 0 && ( +
+ {actionArray.map((action) => ( + + ))} +
+ )} + + + ); +}; + diff --git a/surfsense_web/components/tool-ui/plan/schema.ts b/surfsense_web/components/tool-ui/plan/schema.ts new file mode 100644 index 000000000..e72233d03 --- /dev/null +++ b/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; + +/** + * 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; + +/** + * 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; + +/** + * 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; + + 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; +} + diff --git a/surfsense_web/components/tool-ui/shared/action-buttons.tsx b/surfsense_web/components/tool-ui/shared/action-buttons.tsx new file mode 100644 index 000000000..61141647f --- /dev/null +++ b/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 = ({ 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 ( +
+ {actionArray.map((action) => ( + + ))} +
+ ); +}; + diff --git a/surfsense_web/components/tool-ui/shared/index.ts b/surfsense_web/components/tool-ui/shared/index.ts new file mode 100644 index 000000000..fae3af451 --- /dev/null +++ b/surfsense_web/components/tool-ui/shared/index.ts @@ -0,0 +1,3 @@ +export * from "./schema"; +export * from "./action-buttons"; + diff --git a/surfsense_web/components/tool-ui/shared/schema.ts b/surfsense_web/components/tool-ui/shared/schema.ts new file mode 100644 index 000000000..71c2422b9 --- /dev/null +++ b/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; + +/** + * Actions configuration schema + */ +export const ActionsConfigSchema = z.object({ + confirm: ActionSchema.optional(), + cancel: ActionSchema.optional(), +}); + +export type ActionsConfig = z.infer; + diff --git a/surfsense_web/components/tool-ui/write-todos.tsx b/surfsense_web/components/tool-ui/write-todos.tsx new file mode 100644 index 000000000..a8a7eaf43 --- /dev/null +++ b/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 ( +
+
+ + Creating plan... +
+
+ ); +} + +/** + * 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({ + 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 ( +
+ + + +
+ ); + } + return ; + } + + // 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 ( +
+ + + +
+ ); + } + 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 ; + } + + const plan = parseSerializablePlan(planToRender); + return ( +
+ + + +
+ ); + }, +}); + +export type { WriteTodosArgs, WriteTodosResult }; +