SurfSense/surfsense_web/components/tool-ui/deepagent-thinking.tsx

"use client";

import { makeAssistantToolUI } from "@assistant-ui/react";
import { Brain, CheckCircle2, Loader2, Search, Sparkles } from "lucide-react";
import { useEffect, useMemo, useRef, useState } from "react";
import { z } from "zod";
import {
	ChainOfThought,
	ChainOfThoughtContent,
	ChainOfThoughtItem,
	ChainOfThoughtStep,
	ChainOfThoughtTrigger,
} from "@/components/prompt-kit/chain-of-thought";
import { cn } from "@/lib/utils";

/**
 * Zod schemas for runtime validation
 */
const ThinkingStepSchema = z.object({
	id: z.string(),
	title: z.string(),
	items: z.array(z.string()).default([]),
	status: z.enum(["pending", "in_progress", "completed"]).default("pending"),
});

const DeepAgentThinkingArgsSchema = z.object({
	query: z.string().optional(),
	context: z.string().optional(),
});

const DeepAgentThinkingResultSchema = z.object({
	steps: z.array(ThinkingStepSchema).optional(),
	status: z.enum(["thinking", "searching", "synthesizing", "completed"]).optional(),
	summary: z.string().optional(),
});

/**
 * Types derived from Zod schemas
 */
type ThinkingStep = z.infer<typeof ThinkingStepSchema>;
type DeepAgentThinkingArgs = z.infer<typeof DeepAgentThinkingArgsSchema>;
type DeepAgentThinkingResult = z.infer<typeof DeepAgentThinkingResultSchema>;

/**
 * Parse and validate a single thinking step
 */
export function parseThinkingStep(data: unknown): ThinkingStep {
	const result = ThinkingStepSchema.safeParse(data);
	if (!result.success) {
		console.warn("Invalid thinking step data:", result.error.issues);
		// Return a fallback step
		return {
			id: "unknown",
			title: "Processing...",
			items: [],
			status: "pending",
		};
	}
	return result.data;
}

/**
 * Parse and validate thinking result
 */
export function parseThinkingResult(data: unknown): DeepAgentThinkingResult {
	const result = DeepAgentThinkingResultSchema.safeParse(data);
	if (!result.success) {
		console.warn("Invalid thinking result data:", result.error.issues);
		return {};
	}
	return result.data;
}

/**
 * Get icon based on step status and type
 */
function getStepIcon(status: "pending" | "in_progress" | "completed", title: string) {
	// Check for specific step types based on title keywords
	const titleLower = title.toLowerCase();

	if (status === "in_progress") {
		return <Loader2 className="size-4 animate-spin text-primary" />;
	}

	if (status === "completed") {
		return <CheckCircle2 className="size-4 text-emerald-500" />;
	}

	// Default icons based on step type
	if (titleLower.includes("search") || titleLower.includes("knowledge")) {
		return <Search className="size-4 text-muted-foreground" />;
	}

	if (titleLower.includes("analy") || titleLower.includes("understand")) {
		return <Brain className="size-4 text-muted-foreground" />;
	}

	return <Sparkles className="size-4 text-muted-foreground" />;
}

/**
 * Component to display a single thinking step with controlled open state
 */
function ThinkingStepDisplay({
	step,
	isOpen,
	onToggle,
}: {
	step: ThinkingStep;
	isOpen: boolean;
	onToggle: () => void;
}) {
	const icon = useMemo(() => getStepIcon(step.status, step.title), [step.status, step.title]);

	return (
		<ChainOfThoughtStep open={isOpen} onOpenChange={onToggle}>
			<ChainOfThoughtTrigger
				leftIcon={icon}
				swapIconOnHover={step.status !== "in_progress"}
				className={cn(
					step.status === "in_progress" && "text-foreground font-medium",
					step.status === "completed" && "text-muted-foreground"
				)}
			>
				{step.title}
			</ChainOfThoughtTrigger>
			<ChainOfThoughtContent>
				{step.items.map((item, index) => (
					<ChainOfThoughtItem key={`${step.id}-item-${index}`}>{item}</ChainOfThoughtItem>
				))}
			</ChainOfThoughtContent>
		</ChainOfThoughtStep>
	);
}

/**
 * Loading state with animated thinking indicator
 */
function ThinkingLoadingState({ status }: { status?: string }) {
	const statusText = useMemo(() => {
		switch (status) {
			case "searching":
				return "Searching knowledge base...";
			case "synthesizing":
				return "Synthesizing response...";
			case "thinking":
			default:
				return "Thinking...";
		}
	}, [status]);

	return (
		<div className="my-3 flex items-center gap-2 rounded-lg border border-border/50 bg-muted/30 px-4 py-3">
			<div className="relative">
				<Brain className="size-5 text-primary" />
				<span className="absolute -right-0.5 -top-0.5 flex size-2">
					<span className="absolute inline-flex size-full animate-ping rounded-full bg-primary/60" />
					<span className="relative inline-flex size-2 rounded-full bg-primary" />
				</span>
			</div>
			<span className="text-sm text-muted-foreground">{statusText}</span>
		</div>
	);
}

/**
 * Smart chain of thought renderer with state management
 */
function SmartChainOfThought({ steps }: { steps: ThinkingStep[] }) {
	// Track which steps the user has manually toggled
	const [manualOverrides, setManualOverrides] = useState<Record<string, boolean>>({});
	// Track previous step statuses to detect changes
	const prevStatusesRef = useRef<Record<string, string>>({});

	// Check if any step is currently in progress
	const hasInProgressStep = steps.some((step) => step.status === "in_progress");

	// Find the last completed step index
	const lastCompletedIndex = steps
		.map((s, i) => (s.status === "completed" ? i : -1))
		.filter((i) => i !== -1)
		.pop();

	// Clear manual overrides when a step's status changes
	useEffect(() => {
		const currentStatuses: Record<string, string> = {};
		steps.forEach((step) => {
			currentStatuses[step.id] = step.status;
			// If status changed, clear any manual override for this step
			if (prevStatusesRef.current[step.id] && prevStatusesRef.current[step.id] !== step.status) {
				setManualOverrides((prev) => {
					const next = { ...prev };
					delete next[step.id];
					return next;
				});
			}
		});
		prevStatusesRef.current = currentStatuses;
	}, [steps]);

	const getStepOpenState = (step: ThinkingStep, index: number): boolean => {
		// If user has manually toggled, respect that
		if (manualOverrides[step.id] !== undefined) {
			return manualOverrides[step.id];
		}
		// Auto behavior: open if in progress
		if (step.status === "in_progress") {
			return true;
		}
		// Auto behavior: keep last completed step open if no in-progress step
		if (!hasInProgressStep && index === lastCompletedIndex) {
			return true;
		}
		// Default: collapsed
		return false;
	};

	const handleToggle = (stepId: string, currentOpen: boolean) => {
		setManualOverrides((prev) => ({
			...prev,
			[stepId]: !currentOpen,
		}));
	};

	return (
		<ChainOfThought>
			{steps.map((step, index) => {
				const isOpen = getStepOpenState(step, index);
				return (
					<ThinkingStepDisplay
						key={step.id}
						step={step}
						isOpen={isOpen}
						onToggle={() => handleToggle(step.id, isOpen)}
					/>
				);
			})}
		</ChainOfThought>
	);
}

/**
 * DeepAgent Thinking Tool UI Component
 *
 * This component displays the agent's chain-of-thought reasoning
 * when the deepagent is processing a query. It shows thinking steps
 * in a collapsible, hierarchical format.
 */
export const DeepAgentThinkingToolUI = makeAssistantToolUI<
	DeepAgentThinkingArgs,
	DeepAgentThinkingResult
>({
	toolName: "deepagent_thinking",
	render: function DeepAgentThinkingUI({ result, status }) {
		// Loading state - tool is still running
		if (status.type === "running" || status.type === "requires-action") {
			return <ThinkingLoadingState status={result?.status} />;
		}

		// Incomplete/cancelled state
		if (status.type === "incomplete") {
			if (status.reason === "cancelled") {
				return null; // Don't show anything if cancelled
			}
			if (status.reason === "error") {
				return null; // Don't show error for thinking - it's not critical
			}
		}

		// No result or no steps - don't render anything
		if (!result?.steps || result.steps.length === 0) {
			return null;
		}

		// Render the chain of thought
		return (
			<div className="my-3 w-full">
				<SmartChainOfThought steps={result.steps} />
			</div>
		);
	},
});

/**
 * Inline Thinking Display Component
 *
 * A simpler version that can be used inline with the message content
 * for displaying reasoning without the full tool UI infrastructure.
 */
export function InlineThinkingDisplay({
	steps,
	isStreaming = false,
	className,
}: {
	steps: ThinkingStep[];
	isStreaming?: boolean;
	className?: string;
}) {
	if (steps.length === 0 && !isStreaming) {
		return null;
	}

	return (
		<div className={cn("my-3 w-full", className)}>
			{isStreaming && steps.length === 0 ? (
				<ThinkingLoadingState />
			) : (
				<SmartChainOfThought steps={steps} />
			)}
		</div>
	);
}

export type { ThinkingStep, DeepAgentThinkingArgs, DeepAgentThinkingResult };