@agentuity/workbench 0.0.64 → 0.0.66

package/dist/types/config.d.ts

@@ -0,0 +1,29 @@
+import type { UIMessage } from 'ai';
+import type { WorkbenchConfig } from '@agentuity/core/workbench';
+import type { AgentSchemaData, AgentSchemasResponse } from '../hooks/useAgentSchemas';
+export interface Agent {
+    id: string;
+    name: string;
+    description?: string;
+    avatar?: string;
+}
+export type ConnectionStatus = 'connected' | 'restarting' | 'disconnected';
+export interface WorkbenchContextType {
+    config: WorkbenchConfig;
+    agents: Record<string, AgentSchemaData>;
+    suggestions: string[];
+    messages: UIMessage[];
+    setMessages: (messages: UIMessage[] | ((prev: UIMessage[]) => UIMessage[])) => void;
+    selectedAgent: string;
+    setSelectedAgent: (agentId: string) => void;
+    inputMode: 'text' | 'form';
+    setInputMode: (mode: 'text' | 'form') => void;
+    isLoading: boolean;
+    submitMessage: (value: string, mode?: 'text' | 'form') => Promise<void>;
+    schemas: AgentSchemasResponse | null;
+    schemasLoading: boolean;
+    schemasError: Error | null;
+    refetchSchemas: () => void;
+    connectionStatus: ConnectionStatus;
+}
+//# sourceMappingURL=config.d.ts.map

package/dist/types/config.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"config.d.ts","sourceRoot":"","sources":["../../src/types/config.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,SAAS,EAAE,MAAM,IAAI,CAAC;AACpC,OAAO,KAAK,EAAE,eAAe,EAAE,MAAM,2BAA2B,CAAC;AACjE,OAAO,KAAK,EAAE,eAAe,EAAE,oBAAoB,EAAE,MAAM,0BAA0B,CAAC;AAEtF,MAAM,WAAW,KAAK;IACrB,EAAE,EAAE,MAAM,CAAC;IACX,IAAI,EAAE,MAAM,CAAC;IACb,WAAW,CAAC,EAAE,MAAM,CAAC;IACrB,MAAM,CAAC,EAAE,MAAM,CAAC;CAChB;AAED,MAAM,MAAM,gBAAgB,GAAG,WAAW,GAAG,YAAY,GAAG,cAAc,CAAC;AAG3E,MAAM,WAAW,oBAAoB;IACpC,MAAM,EAAE,eAAe,CAAC;IACxB,MAAM,EAAE,MAAM,CAAC,MAAM,EAAE,eAAe,CAAC,CAAC;IACxC,WAAW,EAAE,MAAM,EAAE,CAAC;IACtB,QAAQ,EAAE,SAAS,EAAE,CAAC;IACtB,WAAW,EAAE,CAAC,QAAQ,EAAE,SAAS,EAAE,GAAG,CAAC,CAAC,IAAI,EAAE,SAAS,EAAE,KAAK,SAAS,EAAE,CAAC,KAAK,IAAI,CAAC;IACpF,aAAa,EAAE,MAAM,CAAC;IACtB,gBAAgB,EAAE,CAAC,OAAO,EAAE,MAAM,KAAK,IAAI,CAAC;IAC5C,SAAS,EAAE,MAAM,GAAG,MAAM,CAAC;IAC3B,YAAY,EAAE,CAAC,IAAI,EAAE,MAAM,GAAG,MAAM,KAAK,IAAI,CAAC;IAC9C,SAAS,EAAE,OAAO,CAAC;IACnB,aAAa,EAAE,CAAC,KAAK,EAAE,MAAM,EAAE,IAAI,CAAC,EAAE,MAAM,GAAG,MAAM,KAAK,OAAO,CAAC,IAAI,CAAC,CAAC;IAExE,OAAO,EAAE,oBAAoB,GAAG,IAAI,CAAC;IACrC,cAAc,EAAE,OAAO,CAAC;IACxB,YAAY,EAAE,KAAK,GAAG,IAAI,CAAC;IAC3B,cAAc,EAAE,MAAM,IAAI,CAAC;IAE3B,gBAAgB,EAAE,gBAAgB,CAAC;CACnC"}

package/dist/types/config.js

@@ -0,0 +1,2 @@
+export {};
+//# sourceMappingURL=config.js.map

package/dist/types/config.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"config.js","sourceRoot":"","sources":["../../src/types/config.ts"],"names":[],"mappings":""}

package/dist/types.d.ts

@@ -1,4 +1,4 @@
-import type { WorkbenchConfig } from '@agentuity/core';
+import type { WorkbenchConfig } from '@agentuity/core/workbench';
 export interface WorkbenchInstance {
     config: WorkbenchConfig;
 }

package/dist/types.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"types.d.ts","sourceRoot":"","sources":["../src/types.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,eAAe,EAAE,MAAM,iBAAiB,CAAC;AAEvD,MAAM,WAAW,iBAAiB;IACjC,MAAM,EAAE,eAAe,CAAC;CACxB"}
+{"version":3,"file":"types.d.ts","sourceRoot":"","sources":["../src/types.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,eAAe,EAAE,MAAM,2BAA2B,CAAC;AAEjE,MAAM,WAAW,iBAAiB;IACjC,MAAM,EAAE,eAAe,CAAC;CACxB"}

package/dist/workbench.d.ts

@@ -1,5 +1,4 @@
-import type { WorkbenchConfig } from '@agentuity/core';
+import type { WorkbenchConfig } from '@agentuity/core/workbench';
 import type { WorkbenchInstance } from './types';
 export declare function createWorkbench(config?: WorkbenchConfig): WorkbenchInstance;
-export { Workbench } from './components';
 //# sourceMappingURL=workbench.d.ts.map

package/dist/workbench.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"workbench.d.ts","sourceRoot":"","sources":["../src/workbench.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,eAAe,EAAE,MAAM,iBAAiB,CAAC;AACvD,OAAO,KAAK,EAAE,iBAAiB,EAAE,MAAM,SAAS,CAAC;AAEjD,wBAAgB,eAAe,CAC9B,MAAM,GAAE,eAAsD,GAC5D,iBAAiB,CASnB;AAGD,OAAO,EAAE,SAAS,EAAE,MAAM,cAAc,CAAC"}
+{"version":3,"file":"workbench.d.ts","sourceRoot":"","sources":["../src/workbench.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,eAAe,EAAE,MAAM,2BAA2B,CAAC;AACjE,OAAO,KAAK,EAAE,iBAAiB,EAAE,MAAM,SAAS,CAAC;AAEjD,wBAAgB,eAAe,CAC9B,MAAM,GAAE,eAAsD,GAC5D,iBAAiB,CASnB"}

package/dist/workbench.js

@@ -7,6 +7,4 @@ export function createWorkbench(config = { route: '/workbench', headers: {} }) {
         config: finalConfig,
     };
 }
-// Export the actual component for use in HTML generation
-export { Workbench } from './components';
 //# sourceMappingURL=workbench.js.map

package/dist/workbench.js.map

@@ -1 +1 @@
-{"version":3,"file":"workbench.js","sourceRoot":"","sources":["../src/workbench.ts"],"names":[],"mappings":"AAGA,MAAM,UAAU,eAAe,CAC9B,SAA0B,EAAE,KAAK,EAAE,YAAY,EAAE,OAAO,EAAE,EAAE,EAAE;IAE9D,MAAM,WAAW,GAAoB;QACpC,KAAK,EAAE,MAAM,CAAC,KAAK;QACnB,OAAO,EAAE,MAAM,CAAC,OAAO;KACvB,CAAC;IAEF,OAAO;QACN,MAAM,EAAE,WAAW;KACnB,CAAC;AACH,CAAC;AAED,yDAAyD;AACzD,OAAO,EAAE,SAAS,EAAE,MAAM,cAAc,CAAC"}
+{"version":3,"file":"workbench.js","sourceRoot":"","sources":["../src/workbench.ts"],"names":[],"mappings":"AAGA,MAAM,UAAU,eAAe,CAC9B,SAA0B,EAAE,KAAK,EAAE,YAAY,EAAE,OAAO,EAAE,EAAE,EAAE;IAE9D,MAAM,WAAW,GAAoB;QACpC,KAAK,EAAE,MAAM,CAAC,KAAK;QACnB,OAAO,EAAE,MAAM,CAAC,OAAO;KACvB,CAAC;IAEF,OAAO;QACN,MAAM,EAAE,WAAW;KACnB,CAAC;AACH,CAAC"}

package/package.json

@@ -1,6 +1,6 @@
 {
 	"name": "@agentuity/workbench",
-	"version": "0.0.64",
+	"version": "0.0.66",
 	"license": "Apache-2.0",
 	"author": "Agentuity employees and contributors",
 	"type": "module",
@@ -11,6 +11,10 @@
 			"types": "./dist/index.d.ts",
 			"import": "./dist/index.js"
 		},
+		"./server": {
+			"types": "./dist/server.d.ts",
+			"import": "./dist/server.js"
+		},
 		"./styles": "./dist/styles.css"
 	},
 	"files": [
@@ -23,26 +27,65 @@
 		"build": "bun run build:ts && bun run build:css",
 		"build:ts": "bunx tsc --build --force",
 		"build:css": "bun x tailwindcss -i ./src/styles.css -o ./dist/styles.css",
-		"dev": "bun run build && bun run dev:watch",
-		"dev:watch": "bunx tsc --watch",
+		"dev": "bun run build && concurrently \"bun run dev:ts\" \"bun run dev:css\" --names \"ts,css\" --prefix-colors \"blue,yellow\"",
+		"dev:ts": "bunx tsc --watch",
+		"dev:css": "bun x tailwindcss -i ./src/styles.css -o ./dist/styles.css --watch",
 		"typecheck": "bunx tsc --noEmit",
 		"prepublishOnly": "bun run clean && bun run build"
 	},
 	"dependencies": {
-		"@agentuity/core": "0.0.64",
-		"@agentuity/react": "0.0.64",
+		"@agentuity/core": "0.0.66",
+		"@agentuity/react": "0.0.66",
+		"@ai-sdk/react": "^2.0.82",
+		"@hookform/resolvers": "^5.2.2",
+		"@radix-ui/react-avatar": "^1.1.10",
+		"@radix-ui/react-checkbox": "^1.3.3",
+		"@radix-ui/react-collapsible": "^1.1.12",
+		"@radix-ui/react-dialog": "^1.1.15",
+		"@radix-ui/react-dropdown-menu": "^2.1.16",
+		"@radix-ui/react-hover-card": "^1.1.15",
 		"@radix-ui/react-label": "^2.1.7",
+		"@radix-ui/react-popover": "^1.1.15",
+		"@radix-ui/react-progress": "^1.1.7",
+		"@radix-ui/react-scroll-area": "^1.2.10",
 		"@radix-ui/react-select": "^2.2.6",
+		"@radix-ui/react-separator": "^1.1.7",
 		"@radix-ui/react-slot": "^1.2.3",
+		"@radix-ui/react-switch": "^1.2.6",
+		"@radix-ui/react-tabs": "^1.1.13",
+		"@radix-ui/react-toggle": "^1.1.10",
+		"@radix-ui/react-tooltip": "^1.2.8",
+		"@radix-ui/react-use-controllable-state": "^1.2.2",
+		"@xyflow/react": "^12.9.1",
+		"ai": "^5.0.82",
+		"bun-plugin-tailwind": "^0.0.14",
 		"class-variance-authority": "^0.7.1",
 		"clsx": "^2.1.1",
-		"lucide-react": "^0.545.0",
-		"tailwind-merge": "^3.3.1"
+		"cmdk": "^1.1.1",
+		"embla-carousel-react": "^8.6.0",
+		"lucide-react": "^0.548.0",
+		"modern-monaco": "^0.3.2",
+		"motion": "^12.23.24",
+		"nanoid": "^5.1.6",
+		"non.geist": "^1.0.4",
+		"react-hook-form": "^7.65.0",
+		"react-resizable-panels": "^3.0.6",
+		"shiki": "^3.14.0",
+		"streamdown": "^1.4.0",
+		"tailwind-merge": "^3.3.1",
+		"tailwindcss-animate": "^1.0.7",
+		"tokenlens": "^1.3.1",
+		"use-stick-to-bottom": "^1.1.1",
+		"zocker": "^3.0.0",
+		"zod": "^4.1.12",
+		"zod-from-json-schema": "^0.5.2"
 	},
 	"devDependencies": {
 		"@tailwindcss/cli": "^4.1.17",
 		"@types/bun": "latest",
+		"@types/hast": "^3.0.4",
 		"bun-types": "latest",
+		"concurrently": "^9.2.1",
 		"tailwindcss": "^4.1.11",
 		"typescript": "^5.9.0"
 	},

package/src/components/App.tsx

@@ -0,0 +1,29 @@
+import React from 'react';
+import { WorkbenchProvider } from './internal/WorkbenchProvider';
+import { Header } from './internal/Header';
+import { Chat } from './internal/Chat';
+import { ThemeProvider } from './ui/theme-provider';
+import { ResizableProvider } from './ui/resizable-provider';
+import { decodeWorkbenchConfig } from '@agentuity/core/workbench';
+
+export interface AppProps {
+	configBase64: string;
+}
+
+export function App({ configBase64 }: AppProps) {
+	const decodedConfig = decodeWorkbenchConfig(configBase64);
+	return (
+		<WorkbenchProvider config={decodedConfig}>
+			<ThemeProvider>
+				<ResizableProvider>
+					<div className="flex flex-col h-screen">
+						<Header />
+						<Chat />
+					</div>
+				</ResizableProvider>
+			</ThemeProvider>
+		</WorkbenchProvider>
+	);
+}
+
+export default App;

package/src/components/ConnectionStatus.tsx

@@ -0,0 +1,67 @@
+import React from 'react';
+import type { ConnectionStatus as ConnectionStatusType } from '../types/config';
+
+interface ConnectionStatusProps {
+	status: ConnectionStatusType;
+}
+
+export function ConnectionStatus({ status }: ConnectionStatusProps) {
+	if (status === 'connected') {
+		return null; // Don't show anything when connected normally
+	}
+
+	const getStatusConfig = (status: ConnectionStatusType) => {
+		switch (status) {
+			case 'restarting':
+				return {
+					text: 'Server restarting...',
+					bgColor: 'bg-amber-100',
+					textColor: 'text-amber-800',
+					borderColor: 'border-amber-200',
+					icon: '🔄',
+				};
+			case 'disconnected':
+				return {
+					text: 'Disconnected from server',
+					bgColor: 'bg-red-100',
+					textColor: 'text-red-800',
+					borderColor: 'border-red-200',
+					icon: '⚠️',
+				};
+			case 'connected':
+				return {
+					text: 'Connected',
+					bgColor: 'bg-green-100',
+					textColor: 'text-green-800',
+					borderColor: 'border-green-200',
+					icon: '✅',
+				};
+			default:
+				return {
+					text: 'Unknown status',
+					bgColor: 'bg-gray-100',
+					textColor: 'text-gray-800',
+					borderColor: 'border-gray-200',
+					icon: '❓',
+				};
+		}
+	};
+
+	const config = getStatusConfig(status);
+
+	return (
+		<div
+			className={`
+				fixed top-0 left-0 right-0 z-50
+				px-4 py-2
+				${config.bgColor} ${config.textColor} ${config.borderColor}
+				border-b
+				text-center text-sm font-medium
+				transition-all duration-200 ease-in-out
+			`}
+		>
+			<span className="mr-2">{config.icon}</span>
+			{config.text}
+		</div>
+	);
+}

package/src/components/Inline.tsx

@@ -0,0 +1,16 @@
+import React from 'react';
+import { Chat } from './internal/Chat';
+
+export interface InlineProps {
+	className?: string;
+}
+
+/**
+ * Inline component - just the chat area
+ * Must be used within WorkbenchProvider
+ */
+export function Inline({ className }: InlineProps) {
+	return <Chat className={className} />;
+}
+
+export default Inline;

package/src/components/ai-elements/actions.tsx

@@ -0,0 +1,60 @@
+'use client';
+
+import { Button } from '../ui/button';
+import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from '../ui/tooltip';
+import { cn } from '../../lib/utils';
+import type { ComponentProps } from 'react';
+
+export type ActionsProps = ComponentProps<'div'>;
+
+export const Actions = ({ className, children, ...props }: ActionsProps) => (
+	<div className={cn('flex items-center gap-1', className)} {...props}>
+		{children}
+	</div>
+);
+
+export type ActionProps = ComponentProps<typeof Button> & {
+	tooltip?: string;
+	label?: string;
+};
+
+export const Action = ({
+	tooltip,
+	children,
+	label,
+	className,
+	variant = 'ghost',
+	size = 'sm',
+	...props
+}: ActionProps) => {
+	const button = (
+		<Button
+			className={cn(
+				'relative size-9 p-1.5 text-muted-foreground hover:text-foreground',
+				className
+			)}
+			size={size}
+			type="button"
+			variant={variant}
+			{...props}
+		>
+			{children}
+			<span className="sr-only">{label || tooltip}</span>
+		</Button>
+	);
+
+	if (tooltip) {
+		return (
+			<TooltipProvider>
+				<Tooltip>
+					<TooltipTrigger asChild>{button}</TooltipTrigger>
+					<TooltipContent>
+						<p>{tooltip}</p>
+					</TooltipContent>
+				</Tooltip>
+			</TooltipProvider>
+		);
+	}
+
+	return button;
+};

package/src/components/ai-elements/artifact.tsx

@@ -0,0 +1,118 @@
+'use client';
+
+import { Button } from '../ui/button';
+import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from '../ui/tooltip';
+import { cn } from '../../lib/utils';
+import { type LucideIcon, XIcon } from 'lucide-react';
+import type { ComponentProps, HTMLAttributes } from 'react';
+
+export type ArtifactProps = HTMLAttributes<HTMLDivElement>;
+
+export const Artifact = ({ className, ...props }: ArtifactProps) => (
+	<div
+		className={cn(
+			'flex flex-col overflow-hidden rounded-lg border bg-background shadow-sm',
+			className
+		)}
+		{...props}
+	/>
+);
+
+export type ArtifactHeaderProps = HTMLAttributes<HTMLDivElement>;
+
+export const ArtifactHeader = ({ className, ...props }: ArtifactHeaderProps) => (
+	<div
+		className={cn('flex items-center justify-between border-b bg-muted/50 px-4 py-3', className)}
+		{...props}
+	/>
+);
+
+export type ArtifactCloseProps = ComponentProps<typeof Button>;
+
+export const ArtifactClose = ({
+	className,
+	children,
+	size = 'sm',
+	variant = 'ghost',
+	...props
+}: ArtifactCloseProps) => (
+	<Button
+		className={cn('size-8 p-0 text-muted-foreground hover:text-foreground', className)}
+		size={size}
+		type="button"
+		variant={variant}
+		{...props}
+	>
+		{children ?? <XIcon className="size-4" />}
+		<span className="sr-only">Close</span>
+	</Button>
+);
+
+export type ArtifactTitleProps = HTMLAttributes<HTMLParagraphElement>;
+
+export const ArtifactTitle = ({ className, ...props }: ArtifactTitleProps) => (
+	<p className={cn('font-medium text-foreground text-sm', className)} {...props} />
+);
+
+export type ArtifactDescriptionProps = HTMLAttributes<HTMLParagraphElement>;
+
+export const ArtifactDescription = ({ className, ...props }: ArtifactDescriptionProps) => (
+	<p className={cn('text-muted-foreground text-sm', className)} {...props} />
+);
+
+export type ArtifactActionsProps = HTMLAttributes<HTMLDivElement>;
+
+export const ArtifactActions = ({ className, ...props }: ArtifactActionsProps) => (
+	<div className={cn('flex items-center gap-1', className)} {...props} />
+);
+
+export type ArtifactActionProps = ComponentProps<typeof Button> & {
+	tooltip?: string;
+	label?: string;
+	icon?: LucideIcon;
+};
+
+export const ArtifactAction = ({
+	tooltip,
+	label,
+	icon: Icon,
+	children,
+	className,
+	size = 'sm',
+	variant = 'ghost',
+	...props
+}: ArtifactActionProps) => {
+	const button = (
+		<Button
+			className={cn('size-8 p-0 text-muted-foreground hover:text-foreground', className)}
+			size={size}
+			type="button"
+			variant={variant}
+			{...props}
+		>
+			{Icon ? <Icon className="size-4" /> : children}
+			<span className="sr-only">{label || tooltip}</span>
+		</Button>
+	);
+
+	if (tooltip) {
+		return (
+			<TooltipProvider>
+				<Tooltip>
+					<TooltipTrigger asChild>{button}</TooltipTrigger>
+					<TooltipContent>
+						<p>{tooltip}</p>
+					</TooltipContent>
+				</Tooltip>
+			</TooltipProvider>
+		);
+	}
+
+	return button;
+};
+
+export type ArtifactContentProps = HTMLAttributes<HTMLDivElement>;
+
+export const ArtifactContent = ({ className, ...props }: ArtifactContentProps) => (
+	<div className={cn('flex-1 overflow-auto p-4', className)} {...props} />
+);

package/src/components/ai-elements/branch.tsx

@@ -0,0 +1,187 @@
+'use client';
+
+import { Button } from '../ui/button';
+import { cn } from '../../lib/utils';
+import type { UIMessage } from 'ai';
+import { ChevronLeftIcon, ChevronRightIcon } from 'lucide-react';
+import type { ComponentProps, HTMLAttributes, ReactElement } from 'react';
+import { createContext, useContext, useEffect, useState } from 'react';
+
+type BranchContextType = {
+	currentBranch: number;
+	totalBranches: number;
+	goToPrevious: () => void;
+	goToNext: () => void;
+	branches: ReactElement[];
+	setBranches: (branches: ReactElement[]) => void;
+};
+
+const BranchContext = createContext<BranchContextType | null>(null);
+
+const useBranch = () => {
+	const context = useContext(BranchContext);
+
+	if (!context) {
+		throw new Error('Branch components must be used within Branch');
+	}
+
+	return context;
+};
+
+export type BranchProps = HTMLAttributes<HTMLDivElement> & {
+	defaultBranch?: number;
+	onBranchChange?: (branchIndex: number) => void;
+};
+
+export const Branch = ({ defaultBranch = 0, onBranchChange, className, ...props }: BranchProps) => {
+	const [currentBranch, setCurrentBranch] = useState(defaultBranch);
+	const [branches, setBranches] = useState<ReactElement[]>([]);
+
+	const handleBranchChange = (newBranch: number) => {
+		setCurrentBranch(newBranch);
+		onBranchChange?.(newBranch);
+	};
+
+	const goToPrevious = () => {
+		const newBranch = currentBranch > 0 ? currentBranch - 1 : branches.length - 1;
+		handleBranchChange(newBranch);
+	};
+
+	const goToNext = () => {
+		const newBranch = currentBranch < branches.length - 1 ? currentBranch + 1 : 0;
+		handleBranchChange(newBranch);
+	};
+
+	const contextValue: BranchContextType = {
+		currentBranch,
+		totalBranches: branches.length,
+		goToPrevious,
+		goToNext,
+		branches,
+		setBranches,
+	};
+
+	return (
+		<BranchContext.Provider value={contextValue}>
+			<div className={cn('grid w-full gap-2 [&>div]:pb-0', className)} {...props} />
+		</BranchContext.Provider>
+	);
+};
+
+export type BranchMessagesProps = HTMLAttributes<HTMLDivElement>;
+
+export const BranchMessages = ({ children, ...props }: BranchMessagesProps) => {
+	const { currentBranch, setBranches, branches } = useBranch();
+	const childrenArray = Array.isArray(children) ? children : [children];
+
+	// Use useEffect to update branches when they change
+	useEffect(() => {
+		if (branches.length !== childrenArray.length) {
+			setBranches(childrenArray);
+		}
+	}, [childrenArray, branches, setBranches]);
+
+	return childrenArray.map((branch, index) => (
+		<div
+			className={cn(
+				'grid gap-2 overflow-hidden [&>div]:pb-0',
+				index === currentBranch ? 'block' : 'hidden'
+			)}
+			key={branch.key}
+			{...props}
+		>
+			{branch}
+		</div>
+	));
+};
+
+export type BranchSelectorProps = HTMLAttributes<HTMLDivElement> & {
+	from: UIMessage['role'];
+};
+
+export const BranchSelector = ({ className, from, ...props }: BranchSelectorProps) => {
+	const { totalBranches } = useBranch();
+
+	// Don't render if there's only one branch
+	if (totalBranches <= 1) {
+		return null;
+	}
+
+	return (
+		<div
+			className={cn(
+				'flex items-center gap-2 self-end px-10',
+				from === 'assistant' ? 'justify-start' : 'justify-end',
+				className
+			)}
+			{...props}
+		/>
+	);
+};
+
+export type BranchPreviousProps = ComponentProps<typeof Button>;
+
+export const BranchPrevious = ({ className, children, ...props }: BranchPreviousProps) => {
+	const { goToPrevious, totalBranches } = useBranch();
+
+	return (
+		<Button
+			aria-label="Previous branch"
+			className={cn(
+				'size-7 shrink-0 rounded-full text-muted-foreground transition-colors',
+				'hover:bg-accent hover:text-foreground',
+				'disabled:pointer-events-none disabled:opacity-50',
+				className
+			)}
+			disabled={totalBranches <= 1}
+			onClick={goToPrevious}
+			size="icon"
+			type="button"
+			variant="ghost"
+			{...props}
+		>
+			{children ?? <ChevronLeftIcon size={14} />}
+		</Button>
+	);
+};
+
+export type BranchNextProps = ComponentProps<typeof Button>;
+
+export const BranchNext = ({ className, children, ...props }: BranchNextProps) => {
+	const { goToNext, totalBranches } = useBranch();
+
+	return (
+		<Button
+			aria-label="Next branch"
+			className={cn(
+				'size-7 shrink-0 rounded-full text-muted-foreground transition-colors',
+				'hover:bg-accent hover:text-foreground',
+				'disabled:pointer-events-none disabled:opacity-50',
+				className
+			)}
+			disabled={totalBranches <= 1}
+			onClick={goToNext}
+			size="icon"
+			type="button"
+			variant="ghost"
+			{...props}
+		>
+			{children ?? <ChevronRightIcon size={14} />}
+		</Button>
+	);
+};
+
+export type BranchPageProps = HTMLAttributes<HTMLSpanElement>;
+
+export const BranchPage = ({ className, ...props }: BranchPageProps) => {
+	const { currentBranch, totalBranches } = useBranch();
+
+	return (
+		<span
+			className={cn('font-medium text-muted-foreground text-xs tabular-nums', className)}
+			{...props}
+		>
+			{currentBranch + 1} of {totalBranches}
+		</span>
+	);
+};

package/src/components/ai-elements/canvas.tsx

@@ -0,0 +1,24 @@
+import { Background, ReactFlow, type ReactFlowProps } from '@xyflow/react';
+import type { ReactNode } from 'react';
+import '@xyflow/react/dist/style.css';
+import { Controls } from './controls';
+
+type CanvasProps = ReactFlowProps & {
+	children?: ReactNode;
+};
+
+export const Canvas = ({ children, ...props }: CanvasProps) => (
+	<ReactFlow
+		deleteKeyCode={['Backspace', 'Delete']}
+		fitView
+		panOnDrag={false}
+		panOnScroll
+		selectionOnDrag={true}
+		zoomOnDoubleClick={false}
+		{...props}
+	>
+		<Background bgColor="var(--sidebar)" />
+		<Controls />
+		{children}
+	</ReactFlow>
+);