refactor: streamline ThinkingStepsDisplay and ChainOfThought components

- Removed unnecessary calculations for in-progress and last completed step states in ThinkingStepsDisplay.
- Simplified the getStepOpenState function to focus on current step status.
- Introduced constants for animation timing and file extension categories in ChainOfThought for better maintainability.
- Enhanced icon retrieval logic based on file types and step statuses, improving clarity and performance.
- Updated component props to use TypeScript interfaces for better type safety and readability.

surfsense_web/components/assistant-ui/thread.tsx

@@ -147,15 +147,6 @@ const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?: boolea
 		return step.status;
 	};
 
-	// Check if any step is effectively in progress
-	const hasInProgressStep = steps.some((step) => getEffectiveStatus(step) === "in_progress");
-
-	// Find the last completed step index (using effective status)
-	const lastCompletedIndex = steps
-		.map((s, i) => (getEffectiveStatus(s) === "completed" ? i : -1))
-		.filter((i) => i !== -1)
-		.pop();
-
 	// Clear manual overrides when a step's status changes
 	useEffect(() => {
 		const currentStatuses: Record<string, string> = {};
@@ -175,7 +166,7 @@ const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?: boolea
 
 	if (steps.length === 0) return null;
 
-	const getStepOpenState = (step: ThinkingStep, index: number): boolean => {
+	const getStepOpenState = (step: ThinkingStep): boolean => {
 		const effectiveStatus = getEffectiveStatus(step);
 		// If user has manually toggled, respect that
 		if (manualOverrides[step.id] !== undefined) {
@@ -185,11 +176,7 @@ const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?: boolea
 		if (effectiveStatus === "in_progress") {
 			return true;
 		}
-		// Auto behavior: keep last completed step open if no in-progress step
+		// Default: collapsed (all steps collapse when processing is done)
-		if (!hasInProgressStep && index === lastCompletedIndex) {
-			return true;
-		}
-		// Default: collapsed
 		return false;
 	};
 
@@ -203,10 +190,10 @@ const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?: boolea
 	return (
 		<div className="mx-auto w-full max-w-(--thread-max-width) px-2 py-2">
 			<ChainOfThought>
-				{steps.map((step, index) => {
+				{steps.map((step) => {
 					const effectiveStatus = getEffectiveStatus(step);
 					const icon = getStepIcon(effectiveStatus, step.title);
-					const isOpen = getStepOpenState(step, index);
+					const isOpen = getStepOpenState(step);
 					return (
 						<ChainOfThoughtStep
 							key={step.id}

surfsense_web/components/prompt-kit/chain-of-thought.tsx

@@ -15,11 +15,43 @@ import React, { useEffect, useState } from "react";
 import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible";
 import { cn } from "@/lib/utils";
 
+// ============================================================================
+// Constants
+// ============================================================================
+
+/** Animation timing constants (in milliseconds) */
+const ANIMATION = {
+	/** Delay between each step appearing */
+	STAGGER_DELAY_MS: 50,
+	/** Additional delay for connection line animation */
+	CONNECTION_LINE_DELAY_MS: 150,
+} as const;
+
+/** File extension categories for icon mapping */
+const FILE_EXTENSIONS = {
+	DOCUMENT: ["pdf", "doc", "docx"] as const,
+	SPREADSHEET: ["xls", "xlsx", "csv"] as const,
+	IMAGE: ["png", "jpg", "jpeg", "gif", "webp", "svg"] as const,
+	AUDIO: ["mp3", "wav", "m4a", "ogg", "webm"] as const,
+	VIDEO: ["mp4", "mov", "avi", "mkv"] as const,
+	CODE: ["js", "ts", "tsx", "jsx", "py", "html", "css", "json", "md"] as const,
+} as const;
+
+/** Type for file extension categories */
+type FileExtensionCategory = keyof typeof FILE_EXTENSIONS;
+
+/** Icon size class for file icons */
+const FILE_ICON_SIZE_CLASS = "size-3.5" as const;
+
+// ============================================================================
+// Hooks
+// ============================================================================
+
 /**
  * Custom hook for entrance animation
  * Returns true after mount to trigger CSS transitions
  */
-function useEntranceAnimation(delay = 0) {
+function useEntranceAnimation(delay = 0): boolean {
 	const [isVisible, setIsVisible] = useState(false);
 
 	useEffect(() => {
@@ -30,44 +62,60 @@ function useEntranceAnimation(delay = 0) {
 	return isVisible;
 }
 
+// ============================================================================
+// File Icon Utilities
+// ============================================================================
+
+/**
+ * Check if an extension belongs to a specific category
+ */
+function isExtensionInCategory(
+	ext: string,
+	category: FileExtensionCategory
+): boolean {
+	return (FILE_EXTENSIONS[category] as readonly string[]).includes(ext);
+}
+
 /**
  * Get file icon based on file extension (all icons are muted/gray)
  */
 function getFileIcon(name: string): React.ReactNode {
-	const ext = name.split(".").pop()?.toLowerCase() || "";
+	const ext = name.split(".").pop()?.toLowerCase() ?? "";
 
-	// PDF / Word documents
+	if (isExtensionInCategory(ext, "DOCUMENT")) {
-	if (ext === "pdf" || ["doc", "docx"].includes(ext)) {
+		return <FileText className={FILE_ICON_SIZE_CLASS} />;
-		return <FileText className="size-3.5" />;
 	}
-	// Spreadsheets
+	if (isExtensionInCategory(ext, "SPREADSHEET")) {
-	if (["xls", "xlsx", "csv"].includes(ext)) {
+		return <FileSpreadsheet className={FILE_ICON_SIZE_CLASS} />;
-		return <FileSpreadsheet className="size-3.5" />;
 	}
-	// Images
+	if (isExtensionInCategory(ext, "IMAGE")) {
-	if (["png", "jpg", "jpeg", "gif", "webp", "svg"].includes(ext)) {
+		return <FileImage className={FILE_ICON_SIZE_CLASS} />;
-		return <FileImage className="size-3.5" />;
 	}
-	// Audio
+	if (isExtensionInCategory(ext, "AUDIO")) {
-	if (["mp3", "wav", "m4a", "ogg", "webm"].includes(ext)) {
+		return <FileAudio className={FILE_ICON_SIZE_CLASS} />;
-		return <FileAudio className="size-3.5" />;
 	}
-	// Video
+	if (isExtensionInCategory(ext, "VIDEO")) {
-	if (["mp4", "mov", "avi", "mkv"].includes(ext)) {
+		return <FileVideo className={FILE_ICON_SIZE_CLASS} />;
-		return <FileVideo className="size-3.5" />;
 	}
-	// Code files
+	if (isExtensionInCategory(ext, "CODE")) {
-	if (["js", "ts", "tsx", "jsx", "py", "html", "css", "json", "md"].includes(ext)) {
+		return <FileCode className={FILE_ICON_SIZE_CLASS} />;
-		return <FileCode className="size-3.5" />;
 	}
-	// Default
+	return <File className={FILE_ICON_SIZE_CLASS} />;
-	return <File className="size-3.5" />;
+}
+
+// ============================================================================
+// Attachment Components
+// ============================================================================
+
+interface AttachmentTileProps {
+	/** File name to display */
+	name: string;
 }
 
 /**
  * Compact attachment tile component - matches the chat UI style
  */
-const AttachmentTile: React.FC<{ name: string }> = ({ name }) => {
+const AttachmentTile: React.FC<AttachmentTileProps> = ({ name }) => {
 	const icon = getFileIcon(name);
 
 	return (
@@ -120,26 +168,46 @@ function parseAndRenderWithBadges(text: string): React.ReactNode {
 	return parts;
 }
 
-export type ChainOfThoughtItemProps = React.ComponentProps<"div">;
+// ============================================================================
+// Chain of Thought Components
+// ============================================================================
 
-export const ChainOfThoughtItem = ({ children, className, ...props }: ChainOfThoughtItemProps) => (
+export interface ChainOfThoughtItemProps
-	<div className={cn("text-muted-foreground text-sm flex flex-wrap items-center gap-1", className)} {...props}>
+	extends React.HTMLAttributes<HTMLDivElement> {
+	children: React.ReactNode;
+}
+
+export const ChainOfThoughtItem: React.FC<ChainOfThoughtItemProps> = ({
+	children,
+	className,
+	...props
+}) => (
+	<div
+		className={cn(
+			"text-muted-foreground text-sm flex flex-wrap items-center gap-1",
+			className
+		)}
+		{...props}
+	>
 		{typeof children === "string" ? parseAndRenderWithBadges(children) : children}
 	</div>
 );
 
-export type ChainOfThoughtTriggerProps = React.ComponentProps<typeof CollapsibleTrigger> & {
+export interface ChainOfThoughtTriggerProps
+	extends React.ComponentProps<typeof CollapsibleTrigger> {
+	/** Optional icon to display on the left side */
 	leftIcon?: React.ReactNode;
+	/** Whether to swap the icon with chevron on hover */
 	swapIconOnHover?: boolean;
-};
+}
 
-export const ChainOfThoughtTrigger = ({
+export const ChainOfThoughtTrigger: React.FC<ChainOfThoughtTriggerProps> = ({
 	children,
 	className,
 	leftIcon,
 	swapIconOnHover = true,
 	...props
-}: ChainOfThoughtTriggerProps) => (
+}) => (
 	<CollapsibleTrigger
 		className={cn(
 			"group text-muted-foreground hover:text-foreground flex cursor-pointer items-center justify-start gap-1 text-left text-sm transition-colors",
@@ -170,13 +238,14 @@ export const ChainOfThoughtTrigger = ({
 	</CollapsibleTrigger>
 );
 
-export type ChainOfThoughtContentProps = React.ComponentProps<typeof CollapsibleContent>;
+export interface ChainOfThoughtContentProps
+	extends React.ComponentProps<typeof CollapsibleContent> {}
 
-export const ChainOfThoughtContent = ({
+export const ChainOfThoughtContent: React.FC<ChainOfThoughtContentProps> = ({
 	children,
 	className,
 	...props
-}: ChainOfThoughtContentProps) => {
+}) => {
 	return (
 		<CollapsibleContent
 			className={cn(
@@ -201,7 +270,10 @@ export const ChainOfThoughtContent = ({
 							<div
 								key={key}
 								className="animate-in fade-in slide-in-from-left-2 duration-200"
-								style={{ animationDelay: `${index * 50}ms`, animationFillMode: "backwards" }}
+								style={{
+									animationDelay: `${index * ANIMATION.STAGGER_DELAY_MS}ms`,
+									animationFillMode: "backwards",
+								}}
 							>
 								{child}
 							</div>
@@ -213,12 +285,15 @@ export const ChainOfThoughtContent = ({
 	);
 };
 
-export type ChainOfThoughtProps = {
+export interface ChainOfThoughtProps {
 	children: React.ReactNode;
 	className?: string;
-};
+}
 
-export function ChainOfThought({ children, className }: ChainOfThoughtProps) {
+export const ChainOfThought: React.FC<ChainOfThoughtProps> = ({
+	children,
+	className,
+}) => {
 	const childrenArray = React.Children.toArray(children);
 
 	return (
@@ -238,25 +313,31 @@ export function ChainOfThought({ children, className }: ChainOfThoughtProps) {
 			})}
 		</div>
 	);
-}
-
-export type ChainOfThoughtStepProps = {
-	children: React.ReactNode;
-	className?: string;
-	isLast?: boolean;
-	/** Index of the step for staggered animation */
-	stepIndex?: number;
 };
 
-export const ChainOfThoughtStep = ({
+export interface ChainOfThoughtStepProps
+	extends Omit<React.ComponentProps<typeof Collapsible>, "children"> {
+	children: React.ReactNode;
+	className?: string;
+	/** Whether this is the last step (hides connection line) */
+	isLast?: boolean;
+	/** Index of the step for staggered animation timing */
+	stepIndex?: number;
+}
+
+export const ChainOfThoughtStep: React.FC<ChainOfThoughtStepProps> = ({
 	children,
 	className,
 	isLast = false,
 	stepIndex = 0,
 	...props
-}: ChainOfThoughtStepProps & React.ComponentProps<typeof Collapsible>) => {
+}) => {
 	// Staggered entrance animation based on step index
-	const isVisible = useEntranceAnimation(stepIndex * 50);
+	const isVisible = useEntranceAnimation(stepIndex * ANIMATION.STAGGER_DELAY_MS);
+
+	// Calculate connection line delay: step delay + additional offset
+	const connectionLineDelay =
+		stepIndex * ANIMATION.STAGGER_DELAY_MS + ANIMATION.CONNECTION_LINE_DELAY_MS;
 
 	return (
 		<Collapsible
@@ -278,7 +359,7 @@ export const ChainOfThoughtStep = ({
 						// Animate line height from 0 to full
 						isVisible ? "h-4 scale-y-100" : "h-0 scale-y-0"
 					)}
-					style={{ transitionDelay: `${stepIndex * 50 + 150}ms` }}
+					style={{ transitionDelay: `${connectionLineDelay}ms` }}
 				/>
 			</div>
 		</Collapsible>

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

@@ -2,7 +2,8 @@
 
 import { makeAssistantToolUI } from "@assistant-ui/react";
 import { Brain, CheckCircle2, Loader2, Search, Sparkles } from "lucide-react";
-import { useEffect, useMemo, useRef, useState } from "react";
+import type { FC, ReactNode } from "react";
+import { useCallback, useEffect, useMemo, useRef, useState } from "react";
 import { z } from "zod";
 import {
 	ChainOfThought,
@@ -13,14 +14,59 @@ import {
 } from "@/components/prompt-kit/chain-of-thought";
 import { cn } from "@/lib/utils";
 
-/**
+// ============================================================================
- * Zod schemas for runtime validation
+// Constants
- */
+// ============================================================================
+
+/** Step status values */
+const STEP_STATUS = {
+	PENDING: "pending",
+	IN_PROGRESS: "in_progress",
+	COMPLETED: "completed",
+} as const;
+
+/** Agent thinking status values */
+const THINKING_STATUS = {
+	THINKING: "thinking",
+	SEARCHING: "searching",
+	SYNTHESIZING: "synthesizing",
+	COMPLETED: "completed",
+} as const;
+
+/** Keywords for icon detection */
+const STEP_KEYWORDS = {
+	SEARCH: ["search", "knowledge"] as const,
+	ANALYSIS: ["analy", "understand"] as const,
+} as const;
+
+/** Icon size class */
+const ICON_SIZE_CLASS = "size-4" as const;
+
+/** Status text mapping */
+const STATUS_TEXT_MAP: Record<string, string> = {
+	[THINKING_STATUS.SEARCHING]: "Searching knowledge base...",
+	[THINKING_STATUS.SYNTHESIZING]: "Synthesizing response...",
+	[THINKING_STATUS.THINKING]: "Thinking...",
+} as const;
+
+// ============================================================================
+// Type Definitions
+// ============================================================================
+
+type StepStatus = (typeof STEP_STATUS)[keyof typeof STEP_STATUS];
+type ThinkingStatus = (typeof THINKING_STATUS)[keyof typeof THINKING_STATUS];
+
+// ============================================================================
+// Zod Schemas
+// ============================================================================
+
 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"),
+	status: z
+		.enum([STEP_STATUS.PENDING, STEP_STATUS.IN_PROGRESS, STEP_STATUS.COMPLETED])
+		.default(STEP_STATUS.PENDING),
 });
 
 const DeepAgentThinkingArgsSchema = z.object({
@@ -30,17 +76,34 @@ const DeepAgentThinkingArgsSchema = z.object({
 
 const DeepAgentThinkingResultSchema = z.object({
 	steps: z.array(ThinkingStepSchema).optional(),
-	status: z.enum(["thinking", "searching", "synthesizing", "completed"]).optional(),
+	status: z
+		.enum([
+			THINKING_STATUS.THINKING,
+			THINKING_STATUS.SEARCHING,
+			THINKING_STATUS.SYNTHESIZING,
+			THINKING_STATUS.COMPLETED,
+		])
+		.optional(),
 	summary: z.string().optional(),
 });
 
-/**
+/** Types derived from Zod schemas */
- * Types derived from Zod schemas
- */
 type ThinkingStep = z.infer<typeof ThinkingStepSchema>;
 type DeepAgentThinkingArgs = z.infer<typeof DeepAgentThinkingArgsSchema>;
 type DeepAgentThinkingResult = z.infer<typeof DeepAgentThinkingResultSchema>;
 
+// ============================================================================
+// Parser Functions
+// ============================================================================
+
+/** Default fallback step when parsing fails */
+const DEFAULT_FALLBACK_STEP: ThinkingStep = {
+	id: "unknown",
+	title: "Processing...",
+	items: [],
+	status: STEP_STATUS.PENDING,
+} as const;
+
 /**
  * Parse and validate a single thinking step
  */
@@ -48,13 +111,7 @@ 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 DEFAULT_FALLBACK_STEP;
-		return {
-			id: "unknown",
-			title: "Processing...",
-			items: [],
-			status: "pending",
-		};
 	}
 	return result.data;
 }
@@ -71,55 +128,79 @@ export function parseThinkingResult(data: unknown): DeepAgentThinkingResult {
 	return result.data;
 }
 
+// ============================================================================
+// Icon Utilities
+// ============================================================================
+
 /**
- * Get icon based on step status and type
+ * Check if title contains any of the keywords
  */
-function getStepIcon(status: "pending" | "in_progress" | "completed", title: string) {
+function titleContainsKeywords(
-	// Check for specific step types based on title keywords
+	title: string,
+	keywords: readonly string[]
+): boolean {
 	const titleLower = title.toLowerCase();
+	return keywords.some((keyword) => titleLower.includes(keyword));
+}
 
-	if (status === "in_progress") {
+/**
-		return <Loader2 className="size-4 animate-spin text-primary" />;
+ * Get icon based on step status and title
+ */
+function getStepIcon(status: StepStatus, title: string): ReactNode {
+	if (status === STEP_STATUS.IN_PROGRESS) {
+		return <Loader2 className={cn(ICON_SIZE_CLASS, "animate-spin text-primary")} />;
 	}
 
-	if (status === "completed") {
+	if (status === STEP_STATUS.COMPLETED) {
-		return <CheckCircle2 className="size-4 text-emerald-500" />;
+		return <CheckCircle2 className={cn(ICON_SIZE_CLASS, "text-emerald-500")} />;
 	}
 
-	// Default icons based on step type
+	// Default icons based on step type keywords
-	if (titleLower.includes("search") || titleLower.includes("knowledge")) {
+	if (titleContainsKeywords(title, STEP_KEYWORDS.SEARCH)) {
-		return <Search className="size-4 text-muted-foreground" />;
+		return <Search className={cn(ICON_SIZE_CLASS, "text-muted-foreground")} />;
 	}
 
-	if (titleLower.includes("analy") || titleLower.includes("understand")) {
+	if (titleContainsKeywords(title, STEP_KEYWORDS.ANALYSIS)) {
-		return <Brain className="size-4 text-muted-foreground" />;
+		return <Brain className={cn(ICON_SIZE_CLASS, "text-muted-foreground")} />;
 	}
 
-	return <Sparkles className="size-4 text-muted-foreground" />;
+	return <Sparkles className={cn(ICON_SIZE_CLASS, "text-muted-foreground")} />;
+}
+
+// ============================================================================
+// Sub-Components
+// ============================================================================
+
+interface ThinkingStepDisplayProps {
+	step: ThinkingStep;
+	isOpen: boolean;
+	onToggle: () => void;
 }
 
 /**
  * Component to display a single thinking step with controlled open state
  */
-function ThinkingStepDisplay({
+const ThinkingStepDisplay: FC<ThinkingStepDisplayProps> = ({
 	step,
 	isOpen,
 	onToggle,
-}: {
+}) => {
-	step: ThinkingStep;
+	const icon = useMemo(
-	isOpen: boolean;
+		() => getStepIcon(step.status, step.title),
-	onToggle: () => void;
+		[step.status, step.title]
-}) {
+	);
-	const icon = useMemo(() => getStepIcon(step.status, step.title), [step.status, step.title]);
+
+	const isInProgress = step.status === STEP_STATUS.IN_PROGRESS;
+	const isCompleted = step.status === STEP_STATUS.COMPLETED;
 
 	return (
 		<ChainOfThoughtStep open={isOpen} onOpenChange={onToggle}>
 			<ChainOfThoughtTrigger
 				leftIcon={icon}
-				swapIconOnHover={step.status !== "in_progress"}
+				swapIconOnHover={!isInProgress}
 				className={cn(
-					step.status === "in_progress" && "text-foreground font-medium",
+					isInProgress && "text-foreground font-medium",
-					step.status === "completed" && "text-muted-foreground"
+					isCompleted && "text-muted-foreground"
 				)}
 			>
 				{step.title}
@@ -131,22 +212,21 @@ function ThinkingStepDisplay({
 			</ChainOfThoughtContent>
 		</ChainOfThoughtStep>
 	);
+};
+
+interface ThinkingLoadingStateProps {
+	status?: ThinkingStatus | string;
 }
 
 /**
  * Loading state with animated thinking indicator
  */
-function ThinkingLoadingState({ status }: { status?: string }) {
+const ThinkingLoadingState: FC<ThinkingLoadingStateProps> = ({ status }) => {
 	const statusText = useMemo(() => {
-		switch (status) {
+		if (status && status in STATUS_TEXT_MAP) {
-			case "searching":
+			return STATUS_TEXT_MAP[status];
-				return "Searching knowledge base...";
-			case "synthesizing":
-				return "Synthesizing response...";
-			case "thinking":
-			default:
-				return "Thinking...";
 		}
+		return STATUS_TEXT_MAP[THINKING_STATUS.THINKING];
 	}, [status]);
 
 	return (
@@ -161,33 +241,35 @@ function ThinkingLoadingState({ status }: { status?: string }) {
 			<span className="text-sm text-muted-foreground">{statusText}</span>
 		</div>
 	);
+};
+
+interface SmartChainOfThoughtProps {
+	steps: ThinkingStep[];
 }
 
+/** Type for tracking step override states */
+type StepOverrides = Record<string, boolean>;
+
+/** Type for tracking step status history */
+type StepStatusHistory = Record<string, StepStatus>;
+
 /**
  * Smart chain of thought renderer with state management
  */
-function SmartChainOfThought({ steps }: { steps: ThinkingStep[] }) {
+const SmartChainOfThought: FC<SmartChainOfThoughtProps> = ({ steps }) => {
 	// Track which steps the user has manually toggled
-	const [manualOverrides, setManualOverrides] = useState<Record<string, boolean>>({});
+	const [manualOverrides, setManualOverrides] = useState<StepOverrides>({});
 	// Track previous step statuses to detect changes
-	const prevStatusesRef = useRef<Record<string, string>>({});
+	const prevStatusesRef = useRef<StepStatusHistory>({});
-
-	// 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> = {};
+		const currentStatuses: StepStatusHistory = {};
 		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) {
+			const prevStatus = prevStatusesRef.current[step.id];
+			if (prevStatus && prevStatus !== step.status) {
 				setManualOverrides((prev) => {
 					const next = { ...prev };
 					delete next[step.id];
@@ -198,34 +280,33 @@ function SmartChainOfThought({ steps }: { steps: ThinkingStep[] }) {
 		prevStatusesRef.current = currentStatuses;
 	}, [steps]);
 
-	const getStepOpenState = (step: ThinkingStep, index: number): boolean => {
+	const getStepOpenState = useCallback(
-		// If user has manually toggled, respect that
+		(step: ThinkingStep): boolean => {
-		if (manualOverrides[step.id] !== undefined) {
+			// If user has manually toggled, respect that
-			return manualOverrides[step.id];
+			if (manualOverrides[step.id] !== undefined) {
-		}
+				return manualOverrides[step.id];
-		// Auto behavior: open if in progress
+			}
-		if (step.status === "in_progress") {
+			// Auto behavior: open if in progress
-			return true;
+			if (step.status === STEP_STATUS.IN_PROGRESS) {
-		}
+				return true;
-		// Auto behavior: keep last completed step open if no in-progress step
+			}
-		if (!hasInProgressStep && index === lastCompletedIndex) {
+			// Default: collapsed (all steps collapse when processing is done)
-			return true;
+			return false;
-		}
+		},
-		// Default: collapsed
+		[manualOverrides]
-		return false;
+	);
-	};
 
-	const handleToggle = (stepId: string, currentOpen: boolean) => {
+	const handleToggle = useCallback((stepId: string, currentOpen: boolean) => {
 		setManualOverrides((prev) => ({
 			...prev,
 			[stepId]: !currentOpen,
 		}));
-	};
+	}, []);
 
 	return (
 		<ChainOfThought>
-			{steps.map((step, index) => {
+			{steps.map((step) => {
-				const isOpen = getStepOpenState(step, index);
+				const isOpen = getStepOpenState(step);
 				return (
 					<ThinkingStepDisplay
 						key={step.id}
@@ -237,7 +318,7 @@ function SmartChainOfThought({ steps }: { steps: ThinkingStep[] }) {
 			})}
 		</ChainOfThought>
 	);
-}
+};
 
 /**
  * DeepAgent Thinking Tool UI Component
@@ -281,21 +362,30 @@ export const DeepAgentThinkingToolUI = makeAssistantToolUI<
 	},
 });
 
+// ============================================================================
+// Public Components
+// ============================================================================
+
+export interface InlineThinkingDisplayProps {
+	/** The thinking steps to display */
+	steps: ThinkingStep[];
+	/** Whether content is currently streaming */
+	isStreaming?: boolean;
+	/** Additional CSS class names */
+	className?: string;
+}
+
 /**
  * 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({
+export const InlineThinkingDisplay: FC<InlineThinkingDisplayProps> = ({
 	steps,
 	isStreaming = false,
 	className,
-}: {
+}) => {
-	steps: ThinkingStep[];
-	isStreaming?: boolean;
-	className?: string;
-}) {
 	if (steps.length === 0 && !isStreaming) {
 		return null;
 	}
@@ -309,6 +399,18 @@ export function InlineThinkingDisplay({
 			)}
 		</div>
 	);
-}
+};
 
-export type { ThinkingStep, DeepAgentThinkingArgs, DeepAgentThinkingResult };
+// ============================================================================
+// Exports
+// ============================================================================
+
+export type {
+	ThinkingStep,
+	DeepAgentThinkingArgs,
+	DeepAgentThinkingResult,
+	StepStatus,
+	ThinkingStatus,
+};
+
+export { STEP_STATUS, THINKING_STATUS };