refactor: improve write_todos tool and UI components

- Refactored the write_todos tool to enhance argument and result schemas using Zod for better validation and type safety.
- Updated the WriteTodosToolUI to streamline the rendering logic and improve loading states, ensuring a smoother user experience.
- Enhanced the Plan and TodoItem components to better handle streaming states and display progress, providing clearer feedback during task management.
- Cleaned up code formatting and structure for improved readability and maintainability.
This commit is contained in:
Anish Sarkar 2025-12-26 17:49:56 +05:30
parent 2c86287264
commit ebc04f590e
18 changed files with 833 additions and 751 deletions

View file

@ -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<typeof DisplayImageArgsSchema>;
type DisplayImageResult = z.infer<typeof DisplayImageResultSchema>;
/**
* Error state component shown when image display fails
@ -142,4 +157,9 @@ export const DisplayImageToolUI = makeAssistantToolUI<DisplayImageArgs, DisplayI
},
});
export type { DisplayImageArgs, DisplayImageResult };
export {
DisplayImageArgsSchema,
DisplayImageResultSchema,
type DisplayImageArgs,
type DisplayImageResult,
};

View file

@ -24,6 +24,8 @@ export {
type ThinkingStep,
} from "./deepagent-thinking";
export {
DisplayImageArgsSchema,
DisplayImageResultSchema,
type DisplayImageArgs,
type DisplayImageResult,
DisplayImageToolUI,
@ -39,9 +41,13 @@ export {
type SerializableImage,
} from "./image";
export {
LinkPreviewArgsSchema,
LinkPreviewResultSchema,
type LinkPreviewArgs,
type LinkPreviewResult,
LinkPreviewToolUI,
MultiLinkPreviewArgsSchema,
MultiLinkPreviewResultSchema,
type MultiLinkPreviewArgs,
type MultiLinkPreviewResult,
MultiLinkPreviewToolUI,
@ -56,6 +62,8 @@ export {
type SerializableMediaCard,
} from "./media-card";
export {
ScrapeWebpageArgsSchema,
ScrapeWebpageResultSchema,
type ScrapeWebpageArgs,
type ScrapeWebpageResult,
ScrapeWebpageToolUI,
@ -71,6 +79,8 @@ export {
} from "./plan";
export {
WriteTodosToolUI,
WriteTodosArgsSchema,
WriteTodosResultSchema,
type WriteTodosArgs,
type WriteTodosResult,
} from "./write-todos";

View file

@ -2,6 +2,7 @@
import { makeAssistantToolUI } from "@assistant-ui/react";
import { AlertCircleIcon, ExternalLinkIcon, LinkIcon } from "lucide-react";
import { z } from "zod";
import {
MediaCard,
MediaCardErrorBoundary,
@ -10,25 +11,39 @@ import {
type SerializableMediaCard,
} from "@/components/tool-ui/media-card";
/**
* Type definitions for the link_preview tool
*/
interface LinkPreviewArgs {
url: string;
title?: string;
}
// ============================================================================
// Zod Schemas
// ============================================================================
interface LinkPreviewResult {
id: string;
assetId: string;
kind: "link";
href: string;
title: string;
description?: string;
thumb?: string;
domain?: string;
error?: string;
}
/**
* Schema for link_preview tool arguments
*/
const LinkPreviewArgsSchema = z.object({
url: z.string(),
title: z.string().nullish(),
});
/**
* Schema for link_preview tool result
*/
const LinkPreviewResultSchema = z.object({
id: z.string(),
assetId: z.string(),
kind: z.literal("link"),
href: z.string(),
title: z.string(),
description: z.string().nullish(),
thumb: z.string().nullish(),
domain: z.string().nullish(),
error: z.string().nullish(),
});
// ============================================================================
// Types
// ============================================================================
type LinkPreviewArgs = z.infer<typeof LinkPreviewArgsSchema>;
type LinkPreviewResult = z.infer<typeof LinkPreviewResultSchema>;
/**
* Error state component shown when link preview fails
@ -150,20 +165,35 @@ export const LinkPreviewToolUI = makeAssistantToolUI<LinkPreviewArgs, LinkPrevie
},
});
/**
* Multiple Link Previews Tool UI Component
*
* This component handles cases where multiple links need to be previewed.
* It renders a grid of link preview cards.
*/
interface MultiLinkPreviewArgs {
urls: string[];
}
// ============================================================================
// Multi Link Preview Schemas
// ============================================================================
interface MultiLinkPreviewResult {
previews: LinkPreviewResult[];
errors?: { url: string; error: string }[];
}
/**
* Schema for multi_link_preview tool arguments
*/
const MultiLinkPreviewArgsSchema = z.object({
urls: z.array(z.string()),
});
/**
* Schema for error items in multi_link_preview result
*/
const MultiLinkPreviewErrorSchema = z.object({
url: z.string(),
error: z.string(),
});
/**
* Schema for multi_link_preview tool result
*/
const MultiLinkPreviewResultSchema = z.object({
previews: z.array(LinkPreviewResultSchema),
errors: z.array(MultiLinkPreviewErrorSchema).nullish(),
});
type MultiLinkPreviewArgs = z.infer<typeof MultiLinkPreviewArgsSchema>;
type MultiLinkPreviewResult = z.infer<typeof MultiLinkPreviewResultSchema>;
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,
};

View file

@ -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<PlanErrorBoundaryProps, PlanErrorBoundaryState> {
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 (
<Card className="w-full max-w-xl border-destructive/50">
<CardContent className="pt-6">
<div className="flex items-center gap-2 text-destructive">
<span className="text-sm">Failed to render plan</span>
</div>
</CardContent>
</Card>
);
}
return (
<Card className="w-full max-w-xl border-destructive/50">
<CardContent className="pt-6">
<div className="flex items-center gap-2 text-destructive">
<span className="text-sm">Failed to render plan</span>
</div>
</CardContent>
</Card>
);
}
return this.props.children;
}
return this.props.children;
}
}

View file

@ -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<StatusIconProps> = ({ 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 <CheckCircle2 className={cn(baseClass, "text-emerald-500")} />;
case "in_progress":
// Only animate the spinner if we're actively streaming
// When streaming is stopped, show as a static dashed circle
return (
<CircleDashed
className={cn(
baseClass,
"text-primary",
isStreaming && "animate-spin"
)}
style={isStreaming ? { animationDuration: "3s" } : undefined}
/>
);
case "cancelled":
return <XCircle className={cn(baseClass, "text-destructive")} />;
case "pending":
default:
return <Circle className={cn(baseClass, "text-muted-foreground")} />;
}
switch (status) {
case "completed":
return <CheckCircle2 className={cn(baseClass, "text-emerald-500")} />;
case "in_progress":
// Only animate the spinner if we're actively streaming
// When streaming is stopped, show as a static dashed circle
return (
<CircleDashed
className={cn(baseClass, "text-primary", isStreaming && "animate-spin")}
style={isStreaming ? { animationDuration: "3s" } : undefined}
/>
);
case "cancelled":
return <XCircle className={cn(baseClass, "text-destructive")} />;
case "pending":
default:
return <Circle className={cn(baseClass, "text-muted-foreground")} />;
}
};
// ============================================================================
@ -67,55 +57,50 @@ const StatusIcon: FC<StatusIconProps> = ({ 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<TodoItemProps> = ({ 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 <TextShimmerLoader text={todo.label} size="md" />;
}
return (
<span
className={cn(
"text-sm",
isStrikethrough && "line-through text-muted-foreground"
)}
>
{todo.label}
</span>
);
};
// Render the label with optional shimmer effect
const renderLabel = () => {
if (isShimmer) {
return <TextShimmerLoader text={todo.label} size="md" />;
}
return (
<span className={cn("text-sm", isStrikethrough && "line-through text-muted-foreground")}>
{todo.label}
</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>
);
}
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()}
</div>
);
return (
<div className="flex items-center gap-2 py-2">
<StatusIcon status={todo.status} isStreaming={isStreaming} />
{renderLabel()}
</div>
);
};
// ============================================================================
@ -123,159 +108,158 @@ const TodoItem: FC<TodoItemProps> = ({ 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<PlanProps> = ({
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 (
<Accordion type="single" collapsible className="w-full">
{items.map((todo) => (
<TodoItem key={todo.id} todo={todo} isStreaming={isStreaming} />
))}
</Accordion>
);
}
const TodoList: FC<{ items: PlanTodo[] }> = ({ items }) => {
if (hasDescriptions) {
return (
<Accordion type="single" collapsible className="w-full">
{items.map((todo) => (
<TodoItem key={todo.id} todo={todo} isStreaming={isStreaming} />
))}
</Accordion>
);
}
return (
<div className="space-y-0">
{items.map((todo) => (
<TodoItem key={todo.id} todo={todo} isStreaming={isStreaming} />
))}
</div>
);
};
return (
<div className="space-y-0">
{items.map((todo) => (
<TodoItem key={todo.id} todo={todo} isStreaming={isStreaming} />
))}
</div>
);
};
return (
<Card id={id} className={cn("w-full max-w-xl", className)}>
<CardHeader className="pb-3">
<div className="flex items-start justify-between gap-2">
<div className="flex-1 min-w-0">
<CardTitle className="text-base font-semibold">{title}</CardTitle>
{description && (
<CardDescription className="mt-1 text-sm">{description}</CardDescription>
)}
</div>
{isAllComplete && (
<div className="flex items-center gap-1 text-emerald-500">
<PartyPopper className="size-5" />
</div>
)}
</div>
return (
<Card id={id} className={cn("w-full max-w-xl", className)}>
<CardHeader className="pb-3">
<div className="flex items-start justify-between gap-2">
<div className="flex-1 min-w-0">
<CardTitle className="text-base font-semibold">{title}</CardTitle>
{description && (
<CardDescription className="mt-1 text-sm">{description}</CardDescription>
)}
</div>
{isAllComplete && (
<div className="flex items-center gap-1 text-emerald-500">
<PartyPopper className="size-5" />
</div>
)}
</div>
{showProgress && (
<div className="mt-3 space-y-1.5">
<div className="flex items-center justify-between text-xs text-muted-foreground">
<span>
{progress.completed} of {progress.total} complete
</span>
<span>{Math.round(progress.percentage)}%</span>
</div>
<Progress value={progress.percentage} className="h-1.5" />
</div>
)}
</CardHeader>
{showProgress && (
<div className="mt-3 space-y-1.5">
<div className="flex items-center justify-between text-xs text-muted-foreground">
<span>
{progress.completed} of {progress.total} complete
</span>
<span>{Math.round(progress.percentage)}%</span>
</div>
<Progress value={progress.percentage} className="h-1.5" />
</div>
)}
</CardHeader>
<CardContent className="pt-0">
<TodoList items={visibleTodos} />
<CardContent className="pt-0">
<TodoList items={visibleTodos} />
{hasHiddenTodos && (
<Collapsible open={isExpanded} onOpenChange={setIsExpanded}>
<CollapsibleTrigger asChild>
<Button
variant="ghost"
size="sm"
className="w-full mt-2 text-xs text-muted-foreground hover:text-foreground"
>
{isExpanded
? "Show less"
: `Show ${hiddenTodos.length} more ${hiddenTodos.length === 1 ? "task" : "tasks"}`}
</Button>
</CollapsibleTrigger>
<CollapsibleContent>
<TodoList items={hiddenTodos} />
</CollapsibleContent>
</Collapsible>
)}
{hasHiddenTodos && (
<Collapsible open={isExpanded} onOpenChange={setIsExpanded}>
<CollapsibleTrigger asChild>
<Button
variant="ghost"
size="sm"
className="w-full mt-2 text-xs text-muted-foreground hover:text-foreground"
>
{isExpanded
? "Show less"
: `Show ${hiddenTodos.length} more ${hiddenTodos.length === 1 ? "task" : "tasks"}`}
</Button>
</CollapsibleTrigger>
<CollapsibleContent>
<TodoList items={hiddenTodos} />
</CollapsibleContent>
</Collapsible>
)}
{actionArray.length > 0 && (
<div className="flex flex-wrap gap-2 pt-4 mt-2 border-t">
{actionArray.map((action) => (
<Button
key={action.id}
variant={action.variant || "default"}
size="sm"
disabled={action.disabled}
onClick={() => handleAction(action.id)}
>
{action.label}
</Button>
))}
</div>
)}
</CardContent>
</Card>
);
{actionArray.length > 0 && (
<div className="flex flex-wrap gap-2 pt-4 mt-2 border-t">
{actionArray.map((action) => (
<Button
key={action.id}
variant={action.variant || "default"}
size="sm"
disabled={action.disabled}
onClick={() => handleAction(action.id)}
>
{action.label}
</Button>
))}
</div>
)}
</CardContent>
</Card>
);
};

View file

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

View file

@ -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<typeof ScrapeWebpageArgsSchema>;
type ScrapeWebpageResult = z.infer<typeof ScrapeWebpageResultSchema>;
/**
* Error state component shown when webpage scraping fails
@ -154,4 +169,9 @@ export const ScrapeWebpageToolUI = makeAssistantToolUI<ScrapeWebpageArgs, Scrape
},
});
export type { ScrapeWebpageArgs, ScrapeWebpageResult };
export {
ScrapeWebpageArgsSchema,
ScrapeWebpageResultSchema,
type ScrapeWebpageArgs,
type ScrapeWebpageResult,
};

View file

@ -5,38 +5,37 @@ import { Button } from "@/components/ui/button";
import type { Action, ActionsConfig } from "./schema";
interface ActionButtonsProps {
actions?: Action[] | ActionsConfig;
onAction?: (actionId: string) => void;
disabled?: boolean;
actions?: Action[] | ActionsConfig;
onAction?: (actionId: string) => void;
disabled?: boolean;
}
export const ActionButtons: FC<ActionButtonsProps> = ({ 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 (
<div className="flex flex-wrap gap-2 pt-3">
{actionArray.map((action) => (
<Button
key={action.id}
variant={action.variant || "default"}
size="sm"
disabled={disabled || action.disabled}
onClick={() => onAction?.(action.id)}
>
{action.label}
</Button>
))}
</div>
);
return (
<div className="flex flex-wrap gap-2 pt-3">
{actionArray.map((action) => (
<Button
key={action.id}
variant={action.variant || "default"}
size="sm"
disabled={disabled || action.disabled}
onClick={() => onAction?.(action.id)}
>
{action.label}
</Button>
))}
</div>
);
};

View file

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

View file

@ -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<typeof ActionSchema>;
@ -16,9 +16,8 @@ export type Action = z.infer<typeof ActionSchema>;
* 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<typeof ActionsConfigSchema>;

View file

@ -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<typeof WriteTodosArgsSchema>;
type WriteTodosResult = z.infer<typeof WriteTodosResultSchema>;
/**
* Loading state component
*/
function WriteTodosLoading() {
return (
<div className="my-4 w-full max-w-xl rounded-2xl border bg-card/60 px-5 py-4 shadow-sm">
<div className="flex items-center gap-3">
<Loader2 className="size-5 animate-spin text-primary" />
<span className="text-sm text-muted-foreground">Creating plan...</span>
</div>
</div>
);
return (
<div className="my-4 w-full max-w-xl rounded-2xl border bg-card/60 px-5 py-4 shadow-sm">
<div className="flex items-center gap-3">
<Loader2 className="size-5 animate-spin text-primary" />
<span className="text-sm text-muted-foreground">Creating plan...</span>
</div>
</div>
);
}
/**
@ -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<WriteTodosArgs, WriteTodosResult>({
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 (
<div className="my-4">
<PlanErrorBoundary>
<Plan {...plan} showProgress={true} isStreaming={isThreadRunning} />
</PlanErrorBoundary>
</div>
);
}
return <WriteTodosLoading />;
}
// 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 (
<div className="my-4">
<PlanErrorBoundary>
<Plan {...plan} showProgress={true} isStreaming={isThreadRunning} />
</PlanErrorBoundary>
</div>
);
}
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 <WriteTodosLoading />;
}
const plan = parseSerializablePlan(planToRender);
return (
<div className="my-4">
<PlanErrorBoundary>
<Plan {...plan} showProgress={true} isStreaming={isThreadRunning} />
</PlanErrorBoundary>
</div>
);
},
// 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 (
<div className="my-4">
<PlanErrorBoundary>
<Plan {...plan} showProgress={true} isStreaming={isThreadRunning} />
</PlanErrorBoundary>
</div>
);
}
return <WriteTodosLoading />;
}
// 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 (
<div className="my-4">
<PlanErrorBoundary>
<Plan {...plan} showProgress={true} isStreaming={isThreadRunning} />
</PlanErrorBoundary>
</div>
);
}
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 <WriteTodosLoading />;
}
const plan = parseSerializablePlan(planToRender);
return (
<div className="my-4">
<PlanErrorBoundary>
<Plan {...plan} showProgress={true} isStreaming={isThreadRunning} />
</PlanErrorBoundary>
</div>
);
},
});
export type { WriteTodosArgs, WriteTodosResult };
export { WriteTodosArgsSchema, WriteTodosResultSchema, type WriteTodosArgs, type WriteTodosResult };