diff --git a/surfsense_backend/app/agents/new_chat/tools/write_todos.py b/surfsense_backend/app/agents/new_chat/tools/write_todos.py index 95b5cb155..f747d891f 100644 --- a/surfsense_backend/app/agents/new_chat/tools/write_todos.py +++ b/surfsense_backend/app/agents/new_chat/tools/write_todos.py @@ -5,7 +5,7 @@ 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 typing import Any from langchain_core.tools import tool @@ -91,4 +91,3 @@ def create_write_todos_tool(): } return write_todos - diff --git a/surfsense_backend/app/connectors/webcrawler_connector.py b/surfsense_backend/app/connectors/webcrawler_connector.py index 411f99a51..7ffc66644 100644 --- a/surfsense_backend/app/connectors/webcrawler_connector.py +++ b/surfsense_backend/app/connectors/webcrawler_connector.py @@ -76,12 +76,17 @@ class WebCrawlerConnector: return result, None except Exception as firecrawl_error: # Firecrawl failed, fallback to Chromium - logger.warning(f"[webcrawler] Firecrawl failed, falling back to Chromium+Trafilatura for: {url}") + logger.warning( + f"[webcrawler] Firecrawl failed, falling back to Chromium+Trafilatura for: {url}" + ) try: result = await self._crawl_with_chromium(url) return result, None except Exception as chromium_error: - return None, f"Both Firecrawl and Chromium failed. Firecrawl error: {firecrawl_error!s}, Chromium error: {chromium_error!s}" + return ( + None, + f"Both Firecrawl and Chromium failed. Firecrawl error: {firecrawl_error!s}, Chromium error: {chromium_error!s}", + ) else: # No Firecrawl API key, use Chromium directly logger.info(f"[webcrawler] Using Chromium+Trafilatura for: {url}") 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 1a6d173f3..35a096497 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 @@ -149,7 +149,11 @@ function convertToThreadMessage(msg: MessageRecord): ThreadMessageLike { if (typeof part !== "object" || part === null || !("type" in part)) return true; const partType = (part as { type: string }).type; // Filter out thinking-steps, mentioned-documents, and attachments - return partType !== "thinking-steps" && partType !== "mentioned-documents" && partType !== "attachments"; + return ( + partType !== "thinking-steps" && + partType !== "mentioned-documents" && + partType !== "attachments" + ); }); content = filteredContent.length > 0 @@ -319,7 +323,13 @@ export default function NewChatPage() { } finally { setIsInitializing(false); } - }, [urlChatId, setMessageDocumentsMap, setMentionedDocumentIds, setMentionedDocuments, hydratePlanState]); + }, [ + urlChatId, + setMessageDocumentsMap, + setMentionedDocumentIds, + setMentionedDocuments, + hydratePlanState, + ]); // Initialize on mount useEffect(() => { @@ -786,9 +796,7 @@ export default function NewChatPage() { appendMessage(currentThreadId, { role: "assistant", content: partialContent, - }).catch((err) => - console.error("Failed to persist partial assistant message:", err) - ); + }).catch((err) => console.error("Failed to persist partial assistant message:", err)); } return; } diff --git a/surfsense_web/atoms/chat/plan-state.atom.ts b/surfsense_web/atoms/chat/plan-state.atom.ts index c8a0815ce..22c33ff90 100644 --- a/surfsense_web/atoms/chat/plan-state.atom.ts +++ b/surfsense_web/atoms/chat/plan-state.atom.ts @@ -10,20 +10,20 @@ import { atom } from "jotai"; export interface PlanTodo { - id: string; - label: string; - status: "pending" | "in_progress" | "completed" | "cancelled"; - description?: string; + 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; + id: string; + title: string; + description?: string; + todos: PlanTodo[]; + lastUpdated: number; + /** The toolCallId of the first component that rendered this plan */ + ownerToolCallId: string; } /** @@ -38,14 +38,14 @@ let firstPlanOwner: { toolCallId: string; title: string } | null = null; * 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; + if (!firstPlanOwner) { + // First plan in this conversation - claim ownership + firstPlanOwner = { toolCallId, title }; + return true; + } + + // Check if we're the owner + return firstPlanOwner.toolCallId === toolCallId; } /** @@ -53,35 +53,35 @@ export function registerPlanOwner(title: string, toolCallId: string): boolean { * Returns the first plan's title if one exists, otherwise the provided title */ export function getCanonicalPlanTitle(title: string): string { - return firstPlanOwner?.title || title; + return firstPlanOwner?.title || title; } /** * Check if a plan already exists in this conversation */ export function hasPlan(): boolean { - return firstPlanOwner !== null; + return firstPlanOwner !== null; } /** * Get the first plan's info */ export function getFirstPlanInfo(): { toolCallId: string; title: string } | null { - return firstPlanOwner; + return firstPlanOwner; } /** * Check if a toolCallId is the owner of the plan SYNCHRONOUSLY */ export function isPlanOwner(toolCallId: string): boolean { - return !firstPlanOwner || firstPlanOwner.toolCallId === toolCallId; + return !firstPlanOwner || firstPlanOwner.toolCallId === toolCallId; } /** * Clear ownership registry (call when starting a new chat) */ export function clearPlanOwnerRegistry(): void { - firstPlanOwner = null; + firstPlanOwner = null; } /** @@ -94,56 +94,53 @@ 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; + 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); - } -); +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); + 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()); + clearPlanOwnerRegistry(); + set(planStatesAtom, new Map()); }); /** @@ -151,84 +148,80 @@ export const clearPlanStatesAtom = atom(null, (get, set) => { * 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; - }>; - }; + 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); - } - } -); +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; -} + 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/assistant-ui/thread.tsx b/surfsense_web/components/assistant-ui/thread.tsx index 04b2f28bd..23dc0bbad 100644 --- a/surfsense_web/components/assistant-ui/thread.tsx +++ b/surfsense_web/components/assistant-ui/thread.tsx @@ -910,7 +910,7 @@ const AssistantMessageInner: FC = () => { -
+
diff --git a/surfsense_web/components/prompt-kit/loader.tsx b/surfsense_web/components/prompt-kit/loader.tsx index a51876e64..435a6a136 100644 --- a/surfsense_web/components/prompt-kit/loader.tsx +++ b/surfsense_web/components/prompt-kit/loader.tsx @@ -3,16 +3,16 @@ import { cn } from "@/lib/utils"; export interface LoaderProps { - variant?: "text-shimmer"; - size?: "sm" | "md" | "lg"; - text?: string; - className?: string; + variant?: "text-shimmer"; + size?: "sm" | "md" | "lg"; + text?: string; + className?: string; } const textSizes = { - sm: "text-xs", - md: "text-sm", - lg: "text-base", + sm: "text-xs", + md: "text-sm", + lg: "text-base", } as const; /** @@ -20,55 +20,47 @@ const textSizes = { * Used for in-progress states in write_todos and chain-of-thought */ export function TextShimmerLoader({ - text = "Thinking", - className, - size = "md", + text = "Thinking", + className, + size = "md", }: { - text?: string; - className?: string; - size?: "sm" | "md" | "lg"; + text?: string; + className?: string; + size?: "sm" | "md" | "lg"; }) { - return ( - <> - - - {text} - - - ); + + + {text} + + + ); } /** * Loader component - currently only supports text-shimmer variant * Can be extended with more variants if needed in the future */ -export function Loader({ - variant = "text-shimmer", - size = "md", - text, - className, -}: LoaderProps) { - switch (variant) { - case "text-shimmer": - default: - return ( - - ); - } +export function Loader({ variant = "text-shimmer", size = "md", text, className }: LoaderProps) { + switch (variant) { + case "text-shimmer": + default: + return ; + } } - diff --git a/surfsense_web/components/sidebar/AppSidebarProvider.tsx b/surfsense_web/components/sidebar/AppSidebarProvider.tsx index 0ee0bb230..f5146c427 100644 --- a/surfsense_web/components/sidebar/AppSidebarProvider.tsx +++ b/surfsense_web/components/sidebar/AppSidebarProvider.tsx @@ -3,7 +3,7 @@ import { useQuery, useQueryClient } from "@tanstack/react-query"; import { useAtomValue, useSetAtom } from "jotai"; import { Trash2 } from "lucide-react"; -import { useRouter } from "next/navigation"; +import { useParams, useRouter } from "next/navigation"; import { useTranslations } from "next-intl"; import { useCallback, useMemo, useState } from "react"; import { hasUnsavedEditorChangesAtom, pendingEditorNavigationAtom } from "@/atoms/editor/ui.atoms"; @@ -50,7 +50,13 @@ export function AppSidebarProvider({ const t = useTranslations("dashboard"); const tCommon = useTranslations("common"); const router = useRouter(); + const params = useParams(); const queryClient = useQueryClient(); + + // Get current chat ID from URL params + const currentChatId = params?.chat_id + ? Number(Array.isArray(params.chat_id) ? params.chat_id[0] : params.chat_id) + : null; const [isDeletingThread, setIsDeletingThread] = useState(false); // Editor state for handling unsaved changes @@ -61,7 +67,6 @@ export function AppSidebarProvider({ const { data: threadsData, error: threadError, - isLoading: isLoadingThreads, refetch: refetchThreads, } = useQuery({ queryKey: ["threads", searchSpaceId], @@ -73,7 +78,6 @@ export function AppSidebarProvider({ data: searchSpace, isLoading: isLoadingSearchSpace, error: searchSpaceError, - refetch: fetchSearchSpace, } = useQuery({ queryKey: cacheKeys.searchSpaces.detail(searchSpaceId), queryFn: () => searchSpacesApiService.getSearchSpace({ id: Number(searchSpaceId) }), @@ -83,12 +87,7 @@ export function AppSidebarProvider({ const { data: user } = useAtomValue(currentUserAtom); // Fetch notes - const { - data: notesData, - error: notesError, - isLoading: isLoadingNotes, - refetch: refetchNotes, - } = useQuery({ + const { data: notesData, refetch: refetchNotes } = useQuery({ queryKey: ["notes", searchSpaceId], queryFn: () => notesApiService.getNotes({ @@ -108,11 +107,6 @@ export function AppSidebarProvider({ } | null>(null); const [isDeletingNote, setIsDeletingNote] = useState(false); - // Retry function - const retryFetch = useCallback(() => { - fetchSearchSpace(); - }, [fetchSearchSpace]); - // Transform threads to the format expected by AppSidebar const recentChats = useMemo(() => { if (!threadsData?.threads) return []; @@ -149,8 +143,10 @@ export function AppSidebarProvider({ await deleteThread(threadToDelete.id); // Invalidate threads query to refresh the list queryClient.invalidateQueries({ queryKey: ["threads", searchSpaceId] }); - // Navigate to new-chat after successful deletion - router.push(`/dashboard/${searchSpaceId}/new-chat`); + // Only navigate to new-chat if the deleted chat is currently open + if (currentChatId === threadToDelete.id) { + router.push(`/dashboard/${searchSpaceId}/new-chat`); + } } catch (error) { console.error("Error deleting thread:", error); } finally { @@ -158,7 +154,7 @@ export function AppSidebarProvider({ setShowDeleteDialog(false); setThreadToDelete(null); } - }, [threadToDelete, queryClient, searchSpaceId, router]); + }, [threadToDelete, queryClient, searchSpaceId, router, currentChatId]); // Handle delete note with confirmation const handleDeleteNote = useCallback(async () => { diff --git a/surfsense_web/components/tool-ui/display-image.tsx b/surfsense_web/components/tool-ui/display-image.tsx index 333cd496c..660e95bca 100644 --- a/surfsense_web/components/tool-ui/display-image.tsx +++ b/surfsense_web/components/tool-ui/display-image.tsx @@ -2,6 +2,7 @@ import { makeAssistantToolUI } from "@assistant-ui/react"; import { AlertCircleIcon, ImageIcon } from "lucide-react"; +import { z } from "zod"; import { Image, ImageErrorBoundary, @@ -9,27 +10,41 @@ import { parseSerializableImage, } from "@/components/tool-ui/image"; -/** - * Type definitions for the display_image tool - */ -interface DisplayImageArgs { - src: string; - alt?: string; - title?: string; - description?: string; -} +// ============================================================================ +// Zod Schemas +// ============================================================================ -interface DisplayImageResult { - id: string; - assetId: string; - src: string; - alt?: string; // Made optional - parseSerializableImage provides fallback - title?: string; - description?: string; - domain?: string; - ratio?: string; - error?: string; -} +/** + * Schema for display_image tool arguments + */ +const DisplayImageArgsSchema = z.object({ + src: z.string(), + alt: z.string().nullish(), + title: z.string().nullish(), + description: z.string().nullish(), +}); + +/** + * Schema for display_image tool result + */ +const DisplayImageResultSchema = z.object({ + id: z.string(), + assetId: z.string(), + src: z.string(), + alt: z.string().nullish(), + title: z.string().nullish(), + description: z.string().nullish(), + domain: z.string().nullish(), + ratio: z.string().nullish(), + error: z.string().nullish(), +}); + +// ============================================================================ +// Types +// ============================================================================ + +type DisplayImageArgs = z.infer; +type DisplayImageResult = z.infer; /** * Error state component shown when image display fails @@ -142,4 +157,9 @@ export const DisplayImageToolUI = makeAssistantToolUI; +type LinkPreviewResult = z.infer; /** * Error state component shown when link preview fails @@ -150,20 +165,35 @@ export const LinkPreviewToolUI = makeAssistantToolUI; +type MultiLinkPreviewResult = z.infer; export const MultiLinkPreviewToolUI = makeAssistantToolUI< MultiLinkPreviewArgs, @@ -217,4 +247,13 @@ export const MultiLinkPreviewToolUI = makeAssistantToolUI< }, }); -export type { LinkPreviewArgs, LinkPreviewResult, MultiLinkPreviewArgs, MultiLinkPreviewResult }; +export { + LinkPreviewArgsSchema, + LinkPreviewResultSchema, + MultiLinkPreviewArgsSchema, + MultiLinkPreviewResultSchema, + type LinkPreviewArgs, + type LinkPreviewResult, + type MultiLinkPreviewArgs, + type MultiLinkPreviewResult, +}; diff --git a/surfsense_web/components/tool-ui/plan/index.tsx b/surfsense_web/components/tool-ui/plan/index.tsx index 989daede9..8467b0af9 100644 --- a/surfsense_web/components/tool-ui/plan/index.tsx +++ b/surfsense_web/components/tool-ui/plan/index.tsx @@ -11,43 +11,42 @@ export * from "./schema"; // ============================================================================ interface PlanErrorBoundaryProps { - children: ReactNode; - fallback?: ReactNode; + children: ReactNode; + fallback?: ReactNode; } interface PlanErrorBoundaryState { - hasError: boolean; - error?: Error; + hasError: boolean; + error?: Error; } export class PlanErrorBoundary extends Component { - constructor(props: PlanErrorBoundaryProps) { - super(props); - this.state = { hasError: false }; - } + constructor(props: PlanErrorBoundaryProps) { + super(props); + this.state = { hasError: false }; + } - static getDerivedStateFromError(error: Error): PlanErrorBoundaryState { - return { hasError: true, error }; - } + static getDerivedStateFromError(error: Error): PlanErrorBoundaryState { + return { hasError: true, error }; + } - render() { - if (this.state.hasError) { - if (this.props.fallback) { - return this.props.fallback; - } + render() { + if (this.state.hasError) { + if (this.props.fallback) { + return this.props.fallback; + } - return ( - - -
- Failed to render plan -
- - - ); - } + return ( + + +
+ Failed to render plan +
+ + + ); + } - return this.props.children; - } + return this.props.children; + } } - diff --git a/surfsense_web/components/tool-ui/plan/plan.tsx b/surfsense_web/components/tool-ui/plan/plan.tsx index 90ac5c614..169749356 100644 --- a/surfsense_web/components/tool-ui/plan/plan.tsx +++ b/surfsense_web/components/tool-ui/plan/plan.tsx @@ -1,19 +1,13 @@ "use client"; -import { - CheckCircle2, - Circle, - CircleDashed, - PartyPopper, - XCircle, -} from "lucide-react"; +import { CheckCircle2, Circle, CircleDashed, PartyPopper, XCircle } from "lucide-react"; import type { FC } from "react"; import { useMemo, useState } from "react"; import { - Accordion, - AccordionContent, - AccordionItem, - AccordionTrigger, + Accordion, + AccordionContent, + AccordionItem, + AccordionTrigger, } from "@/components/ui/accordion"; import { Button } from "@/components/ui/button"; import { Card, CardContent, CardDescription, CardHeader, CardTitle } from "@/components/ui/card"; @@ -29,37 +23,33 @@ import type { PlanTodo, TodoStatus } from "./schema"; // ============================================================================ interface StatusIconProps { - status: TodoStatus; - className?: string; - /** When false, in_progress items show as static (no spinner) */ - isStreaming?: boolean; + status: TodoStatus; + className?: string; + /** When false, in_progress items show as static (no spinner) */ + isStreaming?: boolean; } const StatusIcon: FC = ({ status, className, isStreaming = true }) => { - const baseClass = cn("size-4 shrink-0", className); + const baseClass = cn("size-4 shrink-0", className); - switch (status) { - case "completed": - return ; - case "in_progress": - // Only animate the spinner if we're actively streaming - // When streaming is stopped, show as a static dashed circle - return ( - - ); - case "cancelled": - return ; - case "pending": - default: - return ; - } + switch (status) { + case "completed": + return ; + case "in_progress": + // Only animate the spinner if we're actively streaming + // When streaming is stopped, show as a static dashed circle + return ( + + ); + case "cancelled": + return ; + case "pending": + default: + return ; + } }; // ============================================================================ @@ -67,55 +57,50 @@ const StatusIcon: FC = ({ status, className, isStreaming = true // ============================================================================ interface TodoItemProps { - todo: PlanTodo; - /** When false, in_progress items show as static (no spinner/pulse) */ - isStreaming?: boolean; + todo: PlanTodo; + /** When false, in_progress items show as static (no spinner/pulse) */ + isStreaming?: boolean; } const TodoItem: FC = ({ todo, isStreaming = true }) => { - const isStrikethrough = todo.status === "completed" || todo.status === "cancelled"; - // Only show shimmer animation if streaming and in progress - const isShimmer = todo.status === "in_progress" && isStreaming; + const isStrikethrough = todo.status === "completed" || todo.status === "cancelled"; + // 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 = () => { - if (isShimmer) { - return ; - } - return ( - - {todo.label} - - ); - }; + // Render the label with optional shimmer effect + const renderLabel = () => { + if (isShimmer) { + return ; + } + return ( + + {todo.label} + + ); + }; - if (todo.description) { - return ( - - -
- - {renderLabel()} -
- - -

{todo.description}

- - - ); - } + if (todo.description) { + return ( + + +
+ + {renderLabel()} +
+ + +

{todo.description}

+ + + ); + } - return ( -
- - {renderLabel()} -
- ); + return ( +
+ + {renderLabel()} +
+ ); }; // ============================================================================ @@ -123,159 +108,158 @@ const TodoItem: FC = ({ todo, isStreaming = true }) => { // ============================================================================ export interface PlanProps { - id: string; - title: string; - description?: string; - todos: PlanTodo[]; - maxVisibleTodos?: number; - showProgress?: boolean; - /** When false, in_progress items show as static (no spinner/pulse animations) */ - isStreaming?: boolean; - responseActions?: Action[] | ActionsConfig; - className?: string; - onResponseAction?: (actionId: string) => void; - onBeforeResponseAction?: (actionId: string) => boolean; + id: string; + title: string; + description?: string; + todos: PlanTodo[]; + maxVisibleTodos?: number; + showProgress?: boolean; + /** When false, in_progress items show as static (no spinner/pulse animations) */ + isStreaming?: boolean; + responseActions?: Action[] | ActionsConfig; + className?: string; + onResponseAction?: (actionId: string) => void; + onBeforeResponseAction?: (actionId: string) => boolean; } export const Plan: FC = ({ - id, - title, - description, - todos, - maxVisibleTodos = 4, - showProgress = true, - isStreaming = true, - responseActions, - className, - onResponseAction, - onBeforeResponseAction, + id, + title, + description, + todos, + maxVisibleTodos = 4, + showProgress = true, + isStreaming = true, + responseActions, + className, + onResponseAction, + onBeforeResponseAction, }) => { - const [isExpanded, setIsExpanded] = useState(false); + 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]); + // 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; + 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; + // 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); + // 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); - }; + // 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]); + // 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) => ( - - ))} - - ); - } + const TodoList: FC<{ items: PlanTodo[] }> = ({ items }) => { + if (hasDescriptions) { + return ( + + {items.map((todo) => ( + + ))} + + ); + } - return ( -
- {items.map((todo) => ( - - ))} -
- ); - }; + return ( +
+ {items.map((todo) => ( + + ))} +
+ ); + }; - return ( - - -
-
- {title} - {description && ( - {description} - )} -
- {isAllComplete && ( -
- -
- )} -
+ return ( + + +
+
+ {title} + {description && ( + {description} + )} +
+ {isAllComplete && ( +
+ +
+ )} +
- {showProgress && ( -
-
- - {progress.completed} of {progress.total} complete - - {Math.round(progress.percentage)}% -
- -
- )} - + {showProgress && ( +
+
+ + {progress.completed} of {progress.total} complete + + {Math.round(progress.percentage)}% +
+ +
+ )} + - - + + - {hasHiddenTodos && ( - - - - - - - - - )} + {hasHiddenTodos && ( + + + + + + + + + )} - {actionArray.length > 0 && ( -
- {actionArray.map((action) => ( - - ))} -
- )} - - - ); + {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 index e72233d03..fed49128a 100644 --- a/surfsense_web/components/tool-ui/plan/schema.ts +++ b/surfsense_web/components/tool-ui/plan/schema.ts @@ -1,5 +1,4 @@ import { z } from "zod"; -import { ActionSchema } from "../shared/schema"; /** * Todo item status @@ -11,10 +10,10 @@ 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(), + id: z.string(), + label: z.string(), + status: TodoStatusSchema, + description: z.string().optional(), }); export type PlanTodo = z.infer; @@ -23,12 +22,12 @@ 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(), + 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; @@ -37,31 +36,31 @@ 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; -} + 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/scrape-webpage.tsx b/surfsense_web/components/tool-ui/scrape-webpage.tsx index 025328235..29e7094db 100644 --- a/surfsense_web/components/tool-ui/scrape-webpage.tsx +++ b/surfsense_web/components/tool-ui/scrape-webpage.tsx @@ -2,6 +2,7 @@ import { makeAssistantToolUI } from "@assistant-ui/react"; import { AlertCircleIcon, FileTextIcon } from "lucide-react"; +import { z } from "zod"; import { Article, ArticleErrorBoundary, @@ -9,30 +10,44 @@ import { parseSerializableArticle, } from "@/components/tool-ui/article"; -/** - * Type definitions for the scrape_webpage tool - */ -interface ScrapeWebpageArgs { - url: string; - max_length?: number; -} +// ============================================================================ +// Zod Schemas +// ============================================================================ -interface ScrapeWebpageResult { - id: string; - assetId: string; - kind: "article"; - href: string; - title: string; - description?: string; - content?: string; - domain?: string; - author?: string; - date?: string; - word_count?: number; - was_truncated?: boolean; - crawler_type?: string; - error?: string; -} +/** + * Schema for scrape_webpage tool arguments + */ +const ScrapeWebpageArgsSchema = z.object({ + url: z.string(), + max_length: z.number().nullish(), +}); + +/** + * Schema for scrape_webpage tool result + */ +const ScrapeWebpageResultSchema = z.object({ + id: z.string(), + assetId: z.string(), + kind: z.literal("article"), + href: z.string(), + title: z.string(), + description: z.string().nullish(), + content: z.string().nullish(), + domain: z.string().nullish(), + author: z.string().nullish(), + date: z.string().nullish(), + word_count: z.number().nullish(), + was_truncated: z.boolean().nullish(), + crawler_type: z.string().nullish(), + error: z.string().nullish(), +}); + +// ============================================================================ +// Types +// ============================================================================ + +type ScrapeWebpageArgs = z.infer; +type ScrapeWebpageResult = z.infer; /** * Error state component shown when webpage scraping fails @@ -154,4 +169,9 @@ export const ScrapeWebpageToolUI = makeAssistantToolUI void; - disabled?: boolean; + actions?: Action[] | ActionsConfig; + onAction?: (actionId: string) => void; + disabled?: boolean; } export const ActionButtons: FC = ({ actions, onAction, disabled }) => { - if (!actions) return null; + 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[]; + // 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; + if (actionArray.length === 0) return null; - return ( -
- {actionArray.map((action) => ( - - ))} -
- ); + return ( +
+ {actionArray.map((action) => ( + + ))} +
+ ); }; - diff --git a/surfsense_web/components/tool-ui/shared/index.ts b/surfsense_web/components/tool-ui/shared/index.ts index fae3af451..1d1f1275e 100644 --- a/surfsense_web/components/tool-ui/shared/index.ts +++ b/surfsense_web/components/tool-ui/shared/index.ts @@ -1,3 +1,2 @@ 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 index 71c2422b9..8076a8e45 100644 --- a/surfsense_web/components/tool-ui/shared/schema.ts +++ b/surfsense_web/components/tool-ui/shared/schema.ts @@ -4,10 +4,10 @@ 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(), + 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; @@ -16,9 +16,8 @@ export type Action = z.infer; * Actions configuration schema */ export const ActionsConfigSchema = z.object({ - confirm: ActionSchema.optional(), - cancel: ActionSchema.optional(), + 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 index 6d8e0446b..4d2a44dcb 100644 --- a/surfsense_web/components/tool-ui/write-todos.tsx +++ b/surfsense_web/components/tool-ui/write-todos.tsx @@ -4,54 +4,76 @@ import { makeAssistantToolUI, useAssistantState } from "@assistant-ui/react"; import { useAtomValue, useSetAtom } from "jotai"; import { Loader2 } from "lucide-react"; import { useEffect, useMemo } from "react"; +import { z } from "zod"; import { - getCanonicalPlanTitle, - planStatesAtom, - registerPlanOwner, - updatePlanStateAtom, + getCanonicalPlanTitle, + planStatesAtom, + registerPlanOwner, + updatePlanStateAtom, } from "@/atoms/chat/plan-state.atom"; -import { Plan, PlanErrorBoundary, parseSerializablePlan } from "./plan"; +import { Plan, PlanErrorBoundary, parseSerializablePlan, TodoStatusSchema } from "./plan"; + +// ============================================================================ +// Zod Schemas +// ============================================================================ /** - * Tool arguments for write_todos + * Schema for a single todo item in the args */ -interface WriteTodosArgs { - title?: string; - description?: string; - todos?: Array<{ - id: string; - content: string; - status: "pending" | "in_progress" | "completed" | "cancelled"; - }>; -} +const WriteTodosArgsTodoSchema = z.object({ + id: z.string(), + content: z.string(), + status: TodoStatusSchema, +}); /** - * Tool result for write_todos + * Schema for write_todos tool arguments */ -interface WriteTodosResult { - id: string; - title: string; - description?: string; - todos: Array<{ - id: string; - label: string; - status: "pending" | "in_progress" | "completed" | "cancelled"; - description?: string; - }>; -} +const WriteTodosArgsSchema = z.object({ + title: z.string().nullish(), + description: z.string().nullish(), + todos: z.array(WriteTodosArgsTodoSchema).nullish(), +}); + +/** + * Schema for a single todo item in the result + */ +const WriteTodosResultTodoSchema = z.object({ + id: z.string(), + label: z.string(), + status: TodoStatusSchema, + description: z.string().nullish(), +}); + +/** + * Schema for write_todos tool result + */ +const WriteTodosResultSchema = z.object({ + id: z.string(), + title: z.string(), + description: z.string().nullish(), + todos: z.array(WriteTodosResultTodoSchema), +}); + +// ============================================================================ +// Types +// ============================================================================ + +type WriteTodosArgs = z.infer; +type WriteTodosResult = z.infer; /** * Loading state component */ function WriteTodosLoading() { - return ( -
-
- - Creating plan... -
-
- ); + return ( +
+
+ + Creating plan... +
+
+ ); } /** @@ -59,20 +81,20 @@ function WriteTodosLoading() { * 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; - } + 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", - })), - }; + 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", + })), + }; } /** @@ -87,116 +109,115 @@ function transformArgsToResult(args: WriteTodosArgs): WriteTodosResult | null { * 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); - - // Check if the THREAD is running (not just this tool) - // This hook subscribes to state changes, so it re-renders when thread stops - const isThreadRunning = useAssistantState(({ thread }) => thread.isRunning); + 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"; + // Check if the THREAD is running (not just this tool) + // This hook subscribes to state changes, so it re-renders when thread stops + const isThreadRunning = useAssistantState(({ thread }) => thread.isRunning); - // 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]); + // Get the plan data (from result or args) + const planData = result || transformArgsToResult(args); + const rawTitle = planData?.title || args.title || "Planning Approach"; - // 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]); + // 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]); - // 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 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 the current plan state (may be updated by other components) - const currentPlanState = planStates.get(planTitle); + // 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]); - // If we're NOT the owner, render nothing (the owner will render) - if (!isOwner) { - return null; - } + // 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]); - // Loading state - tool is still running (no data yet) - if (status.type === "running" || status.type === "requires-action") { - // Try to show partial results from args while streaming - const partialResult = transformArgsToResult(args); - if (partialResult) { - const plan = parseSerializablePlan(partialResult); - return ( -
- - - -
- ); - } - return ; - } + // Get the current plan state (may be updated by other components) + const currentPlanState = planStates.get(planTitle); - // 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); - return ( -
- - - -
- ); - } - return null; - } + // If we're NOT the owner, render nothing (the owner will render) + if (!isOwner) { + 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; - if (!planToRender) { - return ; - } - - const plan = parseSerializablePlan(planToRender); - return ( -
- - - -
- ); - }, + // Loading state - tool is still running (no data yet) + if (status.type === "running" || status.type === "requires-action") { + // Try to show partial results from args while streaming + const partialResult = transformArgsToResult(args); + if (partialResult) { + const plan = parseSerializablePlan(partialResult); + return ( +
+ + + +
+ ); + } + return ; + } + + // 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); + return ( +
+ + + +
+ ); + } + 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; + if (!planToRender) { + return ; + } + + const plan = parseSerializablePlan(planToRender); + return ( +
+ + + +
+ ); + }, }); -export type { WriteTodosArgs, WriteTodosResult }; - +export { WriteTodosArgsSchema, WriteTodosResultSchema, type WriteTodosArgs, type WriteTodosResult };