feat: enhance chat functionality with chain-of-thought display and thinking steps management

surfsense_backend/app/services/new_streaming_service.py

@@ -450,6 +450,35 @@ class VercelStreamingService:
         """
         return self.format_data("further-questions", {"questions": questions})
 
+    def format_thinking_step(
+        self,
+        step_id: str,
+        title: str,
+        status: str = "in_progress",
+        items: list[str] | None = None,
+    ) -> str:
+        """
+        Format a thinking step for chain-of-thought display (SurfSense specific).
+
+        Args:
+            step_id: Unique identifier for the step
+            title: The step title (e.g., "Analyzing your request")
+            status: Step status - "pending", "in_progress", or "completed"
+            items: Optional list of sub-items/details for this step
+
+        Returns:
+            str: SSE formatted thinking step data part
+        """
+        return self.format_data(
+            "thinking-step",
+            {
+                "id": step_id,
+                "title": title,
+                "status": status,
+                "items": items or [],
+            },
+        )
+
     # =========================================================================
     # Error Part
     # =========================================================================

surfsense_backend/app/tasks/chat/stream_new_chat.py

@@ -154,6 +154,49 @@ async def stream_new_chat(
         # Reset text tracking for this stream
         accumulated_text = ""
 
+        # Track thinking steps for chain-of-thought display
+        thinking_step_counter = 0
+        # Map run_id -> step_id for tool calls so we can update them on completion
+        tool_step_ids: dict[str, str] = {}
+        # Track the last active step so we can mark it complete at the end
+        last_active_step_id: str | None = None
+        last_active_step_title: str = ""
+        last_active_step_items: list[str] = []
+        # Track which steps have been completed to avoid duplicate completions
+        completed_step_ids: set[str] = set()
+        # Track if we just finished a tool (text flows silently after tools)
+        just_finished_tool: bool = False
+
+        def next_thinking_step_id() -> str:
+            nonlocal thinking_step_counter
+            thinking_step_counter += 1
+            return f"thinking-{thinking_step_counter}"
+
+        def complete_current_step() -> str | None:
+            """Complete the current active step and return the completion event, if any."""
+            nonlocal last_active_step_id, last_active_step_title, last_active_step_items
+            if last_active_step_id and last_active_step_id not in completed_step_ids:
+                completed_step_ids.add(last_active_step_id)
+                return streaming_service.format_thinking_step(
+                    step_id=last_active_step_id,
+                    title=last_active_step_title,
+                    status="completed",
+                    items=last_active_step_items if last_active_step_items else None,
+                )
+            return None
+
+        # Initial thinking step - analyzing the request
+        analyze_step_id = next_thinking_step_id()
+        last_active_step_id = analyze_step_id
+        last_active_step_title = "Understanding your request"
+        last_active_step_items = [f"Processing: {user_query[:80]}{'...' if len(user_query) > 80 else ''}"]
+        yield streaming_service.format_thinking_step(
+            step_id=analyze_step_id,
+            title="Understanding your request",
+            status="in_progress",
+            items=last_active_step_items,
+        )
+
         # Stream the agent response with thread config for memory
         async for event in agent.astream_events(
             input_state, config=config, version="v2"
@@ -168,6 +211,31 @@ async def stream_new_chat(
                     if content and isinstance(content, str):
                         # Start a new text block if needed
                         if current_text_id is None:
+                            # Complete any previous step
+                            completion_event = complete_current_step()
+                            if completion_event:
+                                yield completion_event
+
+                            if just_finished_tool:
+                                # We just finished a tool - don't create a step here,
+                                # text will flow silently after tools.
+                                # Clear the active step tracking.
+                                last_active_step_id = None
+                                last_active_step_title = ""
+                                last_active_step_items = []
+                                just_finished_tool = False
+                            else:
+                                # Normal text generation (not after a tool)
+                                gen_step_id = next_thinking_step_id()
+                                last_active_step_id = gen_step_id
+                                last_active_step_title = "Generating response"
+                                last_active_step_items = []
+                                yield streaming_service.format_thinking_step(
+                                    step_id=gen_step_id,
+                                    title="Generating response",
+                                    status="in_progress",
+                                )
+
                             current_text_id = streaming_service.generate_text_id()
                             yield streaming_service.format_text_start(current_text_id)
 
@@ -188,6 +256,67 @@ async def stream_new_chat(
                     yield streaming_service.format_text_end(current_text_id)
                     current_text_id = None
 
+                # Complete any previous step EXCEPT "Synthesizing response"
+                # (we want to reuse the Synthesizing step after tools complete)
+                if last_active_step_title != "Synthesizing response":
+                    completion_event = complete_current_step()
+                    if completion_event:
+                        yield completion_event
+
+                # Reset the just_finished_tool flag since we're starting a new tool
+                just_finished_tool = False
+
+                # Create thinking step for the tool call and store it for later update
+                tool_step_id = next_thinking_step_id()
+                tool_step_ids[run_id] = tool_step_id
+                last_active_step_id = tool_step_id
+                if tool_name == "search_knowledge_base":
+                    query = (
+                        tool_input.get("query", "")
+                        if isinstance(tool_input, dict)
+                        else str(tool_input)
+                    )
+                    last_active_step_title = "Searching knowledge base"
+                    last_active_step_items = [f"Query: {query[:100]}{'...' if len(query) > 100 else ''}"]
+                    yield streaming_service.format_thinking_step(
+                        step_id=tool_step_id,
+                        title="Searching knowledge base",
+                        status="in_progress",
+                        items=last_active_step_items,
+                    )
+                elif tool_name == "generate_podcast":
+                    podcast_title = (
+                        tool_input.get("podcast_title", "SurfSense Podcast")
+                        if isinstance(tool_input, dict)
+                        else "SurfSense Podcast"
+                    )
+                    # Get content length for context
+                    content_len = len(
+                        tool_input.get("source_content", "")
+                        if isinstance(tool_input, dict)
+                        else ""
+                    )
+                    last_active_step_title = "Generating podcast"
+                    last_active_step_items = [
+                        f"Title: {podcast_title}",
+                        f"Content: {content_len:,} characters",
+                        "Preparing audio generation...",
+                    ]
+                    yield streaming_service.format_thinking_step(
+                        step_id=tool_step_id,
+                        title="Generating podcast",
+                        status="in_progress",
+                        items=last_active_step_items,
+                    )
+                else:
+                    last_active_step_title = f"Using {tool_name.replace('_', ' ')}"
+                    last_active_step_items = []
+                    yield streaming_service.format_thinking_step(
+                        step_id=tool_step_id,
+                        title=last_active_step_title,
+                        status="in_progress",
+                    )
+
                 # Stream tool info
                 tool_call_id = (
                     f"call_{run_id[:32]}"
@@ -254,6 +383,87 @@ async def stream_new_chat(
 
                 tool_call_id = f"call_{run_id[:32]}" if run_id else "call_unknown"
 
+                # Get the original tool step ID to update it (not create a new one)
+                original_step_id = tool_step_ids.get(run_id, f"thinking-unknown-{run_id[:8]}")
+
+                # Mark the tool thinking step as completed using the SAME step ID
+                # Also add to completed set so we don't try to complete it again
+                completed_step_ids.add(original_step_id)
+                if tool_name == "search_knowledge_base":
+                    # Get result count if available
+                    result_info = "Search completed"
+                    if isinstance(tool_output, dict):
+                        result_len = tool_output.get("result_length", 0)
+                        if result_len > 0:
+                            result_info = f"Found relevant information ({result_len} chars)"
+                    # Include original query in completed items
+                    completed_items = [*last_active_step_items, result_info]
+                    yield streaming_service.format_thinking_step(
+                        step_id=original_step_id,
+                        title="Searching knowledge base",
+                        status="completed",
+                        items=completed_items,
+                    )
+                elif tool_name == "generate_podcast":
+                    # Build detailed completion items based on podcast status
+                    podcast_status = (
+                        tool_output.get("status", "unknown")
+                        if isinstance(tool_output, dict)
+                        else "unknown"
+                    )
+                    podcast_title = (
+                        tool_output.get("title", "Podcast")
+                        if isinstance(tool_output, dict)
+                        else "Podcast"
+                    )
+                    
+                    if podcast_status == "processing":
+                        completed_items = [
+                            f"Title: {podcast_title}",
+                            "Audio generation started",
+                            "Processing in background...",
+                        ]
+                    elif podcast_status == "already_generating":
+                        completed_items = [
+                            f"Title: {podcast_title}",
+                            "Podcast already in progress",
+                            "Please wait for it to complete",
+                        ]
+                    elif podcast_status == "error":
+                        error_msg = (
+                            tool_output.get("error", "Unknown error")
+                            if isinstance(tool_output, dict)
+                            else "Unknown error"
+                        )
+                        completed_items = [
+                            f"Title: {podcast_title}",
+                            f"Error: {error_msg[:50]}",
+                        ]
+                    else:
+                        completed_items = last_active_step_items
+                    
+                    yield streaming_service.format_thinking_step(
+                        step_id=original_step_id,
+                        title="Generating podcast",
+                        status="completed",
+                        items=completed_items,
+                    )
+                else:
+                    yield streaming_service.format_thinking_step(
+                        step_id=original_step_id,
+                        title=f"Using {tool_name.replace('_', ' ')}",
+                        status="completed",
+                        items=last_active_step_items,
+                    )
+
+                # Mark that we just finished a tool - "Synthesizing response" will be created
+                # when text actually starts flowing (not immediately)
+                just_finished_tool = True
+                # Clear the active step since the tool is done
+                last_active_step_id = None
+                last_active_step_title = ""
+                last_active_step_items = []
+
                 # Handle different tool outputs
                 if tool_name == "generate_podcast":
                     # Stream the full podcast result so frontend can render the audio player
@@ -302,6 +512,11 @@ async def stream_new_chat(
         if current_text_id is not None:
             yield streaming_service.format_text_end(current_text_id)
 
+        # Mark the last active thinking step as completed using the same title
+        completion_event = complete_current_step()
+        if completion_event:
+            yield completion_event
+
         # Finish the step and message
         yield streaming_service.format_finish_step()
         yield streaming_service.format_finish()

surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx

@@ -11,6 +11,7 @@ import { useCallback, useEffect, useMemo, useRef, useState } from "react";
 import { toast } from "sonner";
 import { Thread } from "@/components/assistant-ui/thread";
 import { GeneratePodcastToolUI } from "@/components/tool-ui/generate-podcast";
+import type { ThinkingStep } from "@/components/tool-ui/deepagent-thinking";
 import { getBearerToken } from "@/lib/auth-utils";
 import { createAttachmentAdapter, extractAttachmentContent } from "@/lib/chat/attachment-adapter";
 import {
@@ -25,6 +26,23 @@ import {
 	type MessageRecord,
 } from "@/lib/chat/thread-persistence";
 
+/**
+ * Extract thinking steps from message content
+ */
+function extractThinkingSteps(content: unknown): ThinkingStep[] {
+	if (!Array.isArray(content)) return [];
+	
+	const thinkingPart = content.find(
+		(part: unknown) => 
+			typeof part === "object" && 
+			part !== null && 
+			"type" in part && 
+			(part as { type: string }).type === "thinking-steps"
+	) as { type: "thinking-steps"; steps: ThinkingStep[] } | undefined;
+	
+	return thinkingPart?.steps || [];
+}
+
 /**
  * Convert backend message to assistant-ui ThreadMessageLike format
  */
@@ -52,6 +70,16 @@ function convertToThreadMessage(msg: MessageRecord): ThreadMessageLike {
  */
 const TOOLS_WITH_UI = new Set(["generate_podcast"]);
 
+/**
+ * Type for thinking step data from the backend
+ */
+interface ThinkingStepData {
+	id: string;
+	title: string;
+	status: "pending" | "in_progress" | "completed";
+	items: string[];
+}
+
 export default function NewChatPage() {
 	const params = useParams();
 	const router = useRouter();
@@ -59,6 +87,10 @@ export default function NewChatPage() {
 	const [threadId, setThreadId] = useState<number | null>(null);
 	const [messages, setMessages] = useState<ThreadMessageLike[]>([]);
 	const [isRunning, setIsRunning] = useState(false);
+	// Store thinking steps per message ID
+	const [messageThinkingSteps, setMessageThinkingSteps] = useState<
+		Map<string, ThinkingStep[]>
+	>(new Map());
 	const abortControllerRef = useRef<AbortController | null>(null);
 
 	// Create the attachment adapter for file processing
@@ -95,6 +127,20 @@ export default function NewChatPage() {
 				if (response.messages && response.messages.length > 0) {
 					const loadedMessages = response.messages.map(convertToThreadMessage);
 					setMessages(loadedMessages);
+					
+					// Extract and restore thinking steps from persisted messages
+					const restoredThinkingSteps = new Map<string, ThinkingStep[]>();
+					for (const msg of response.messages) {
+						if (msg.role === "assistant") {
+							const steps = extractThinkingSteps(msg.content);
+							if (steps.length > 0) {
+								restoredThinkingSteps.set(`msg-${msg.id}`, steps);
+							}
+						}
+					}
+					if (restoredThinkingSteps.size > 0) {
+						setMessageThinkingSteps(restoredThinkingSteps);
+					}
 				}
 			} else {
 				// Create new thread
@@ -187,6 +233,7 @@ export default function NewChatPage() {
 			// Prepare assistant message
 			const assistantMsgId = `msg-assistant-${Date.now()}`;
 			let accumulatedText = "";
+			const currentThinkingSteps = new Map<string, ThinkingStepData>();
 			const toolCalls = new Map<
 				string,
 				{
@@ -197,7 +244,7 @@ export default function NewChatPage() {
 				}
 			>();
 
-			// Helper to build content
+			// Helper to build content (includes thinking steps for persistence)
 			const buildContent = (): ThreadMessageLike["content"] => {
 				const parts: Array<
 					| { type: "text"; text: string }
@@ -208,7 +255,20 @@ export default function NewChatPage() {
 							args: Record<string, unknown>;
 							result?: unknown;
 					  }
+					| {
+							type: "thinking-steps";
+							steps: ThinkingStepData[];
+					  }
 				> = [];
+				
+				// Include thinking steps for persistence
+				if (currentThinkingSteps.size > 0) {
+					parts.push({
+						type: "thinking-steps",
+						steps: Array.from(currentThinkingSteps.values()),
+					});
+				}
+				
 				if (accumulatedText) {
 					parts.push({ type: "text", text: accumulatedText });
 				}
@@ -367,6 +427,24 @@ export default function NewChatPage() {
 											break;
 										}
 
+										case "data-thinking-step": {
+											// Handle thinking step events for chain-of-thought display
+											const stepData = parsed.data as ThinkingStepData;
+											if (stepData?.id) {
+												currentThinkingSteps.set(stepData.id, stepData);
+												// Update message-specific thinking steps
+												setMessageThinkingSteps((prev) => {
+													const newMap = new Map(prev);
+													newMap.set(
+														assistantMsgId,
+														Array.from(currentThinkingSteps.values())
+													);
+													return newMap;
+												});
+											}
+											break;
+										}
+
 										case "error":
 											throw new Error(parsed.errorText || "Server error");
 									}
@@ -415,6 +493,7 @@ export default function NewChatPage() {
 			} finally {
 				setIsRunning(false);
 				abortControllerRef.current = null;
+				// Note: We no longer clear thinking steps - they persist with the message
 			}
 		},
 		[threadId, searchSpaceId, messages]
@@ -483,7 +562,7 @@ export default function NewChatPage() {
 		<AssistantRuntimeProvider runtime={runtime}>
 			<GeneratePodcastToolUI />
 			<div className="h-[calc(100vh-64px)] max-h-[calc(100vh-64px)] overflow-hidden">
-				<Thread />
+				<Thread messageThinkingSteps={messageThinkingSteps} />
 			</div>
 		</AssistantRuntimeProvider>
 	);

surfsense_web/components/assistant-ui/thread.tsx

@@ -7,10 +7,13 @@ import {
 	MessagePrimitive,
 	ThreadPrimitive,
 	useAssistantState,
+	useMessage,
 } from "@assistant-ui/react";
 import {
 	ArrowDownIcon,
 	ArrowUpIcon,
+	Brain,
+	CheckCircle2,
 	CheckIcon,
 	ChevronLeftIcon,
 	ChevronRightIcon,
@@ -19,6 +22,8 @@ import {
 	Loader2,
 	PencilIcon,
 	RefreshCwIcon,
+	Search,
+	Sparkles,
 	SquareIcon,
 } from "lucide-react";
 import Image from "next/image";
@@ -32,44 +37,134 @@ import {
 import { MarkdownText } from "@/components/assistant-ui/markdown-text";
 import { ToolFallback } from "@/components/assistant-ui/tool-fallback";
 import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
+import {
+	ChainOfThought,
+	ChainOfThoughtContent,
+	ChainOfThoughtItem,
+	ChainOfThoughtStep,
+	ChainOfThoughtTrigger,
+} from "@/components/prompt-kit/chain-of-thought";
 import { Button } from "@/components/ui/button";
 import { cn } from "@/lib/utils";
 import { currentUserAtom } from "@/atoms/user/user-query.atoms";
+import type { ThinkingStep } from "@/components/tool-ui/deepagent-thinking";
 
-export const Thread: FC = () => {
+/**
+ * Props for the Thread component
+ */
+interface ThreadProps {
+	messageThinkingSteps?: Map<string, ThinkingStep[]>;
+}
+
+// Context to pass thinking steps to AssistantMessage
+import { createContext, useContext } from "react";
+
+const ThinkingStepsContext = createContext<Map<string, ThinkingStep[]>>(new Map());
+
+/**
+ * Get icon based on step status and title
+ */
+function getStepIcon(status: "pending" | "in_progress" | "completed", title: string) {
+	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" />;
+	}
+	
+	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" />;
+}
+
+/**
+ * Chain of thought display component
+ */
+const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[] }> = ({ steps }) => {
+	if (steps.length === 0) return null;
+	
 	return (
-		<ThreadPrimitive.Root
-			className="aui-root aui-thread-root @container flex h-full flex-col bg-background"
-			style={{
-				["--thread-max-width" as string]: "44rem",
-			}}
-		>
-			<ThreadPrimitive.Viewport
-				turnAnchor="top"
-				className="aui-thread-viewport relative flex flex-1 flex-col overflow-x-auto overflow-y-scroll scroll-smooth px-4 pt-4"
+		<div className="mx-auto w-full max-w-(--thread-max-width) px-2 py-2">
+			<ChainOfThought>
+				{steps.map((step) => {
+					const icon = getStepIcon(step.status, step.title);
+					return (
+						<ChainOfThoughtStep 
+							key={step.id} 
+							defaultOpen={step.status === "in_progress"}
+						>
+							<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>
+							{step.items && step.items.length > 0 && (
+								<ChainOfThoughtContent>
+									{step.items.map((item, index) => (
+										<ChainOfThoughtItem key={`${step.id}-item-${index}`}>
+											{item}
+										</ChainOfThoughtItem>
+									))}
+								</ChainOfThoughtContent>
+							)}
+						</ChainOfThoughtStep>
+					);
+				})}
+			</ChainOfThought>
+		</div>
+	);
+};
+
+export const Thread: FC<ThreadProps> = ({ messageThinkingSteps = new Map() }) => {
+	return (
+		<ThinkingStepsContext.Provider value={messageThinkingSteps}>
+			<ThreadPrimitive.Root
+				className="aui-root aui-thread-root @container flex h-full flex-col bg-background"
+				style={{
+					["--thread-max-width" as string]: "44rem",
+				}}
 			>
-				<AssistantIf condition={({ thread }) => thread.isEmpty}>
-					<ThreadWelcome />
-				</AssistantIf>
-
-				<ThreadPrimitive.Messages
-					components={{
-						UserMessage,
-						EditComposer,
-						AssistantMessage,
-					}}
-				/>
-
-				<ThreadPrimitive.ViewportFooter className="aui-thread-viewport-footer sticky bottom-0 mx-auto mt-auto flex w-full max-w-(--thread-max-width) flex-col gap-4 overflow-visible rounded-t-3xl bg-background pb-4 md:pb-6">
-					<ThreadScrollToBottom />
-					<AssistantIf condition={({ thread }) => !thread.isEmpty}>
-						<div className="fade-in slide-in-from-bottom-4 animate-in duration-500 ease-out fill-mode-both">
-							<Composer />
-						</div>
+				<ThreadPrimitive.Viewport
+					turnAnchor="top"
+					className="aui-thread-viewport relative flex flex-1 flex-col overflow-x-auto overflow-y-scroll scroll-smooth px-4 pt-4"
+				>
+					<AssistantIf condition={({ thread }) => thread.isEmpty}>
+						<ThreadWelcome />
 					</AssistantIf>
-				</ThreadPrimitive.ViewportFooter>
-			</ThreadPrimitive.Viewport>
-		</ThreadPrimitive.Root>
+
+					<ThreadPrimitive.Messages
+						components={{
+							UserMessage,
+							EditComposer,
+							AssistantMessage,
+						}}
+					/>
+
+					<ThreadPrimitive.ViewportFooter className="aui-thread-viewport-footer sticky bottom-0 mx-auto mt-auto flex w-full max-w-(--thread-max-width) flex-col gap-4 overflow-visible rounded-t-3xl bg-background pb-4 md:pb-6">
+						<ThreadScrollToBottom />
+						<AssistantIf condition={({ thread }) => !thread.isEmpty}>
+							<div className="fade-in slide-in-from-bottom-4 animate-in duration-500 ease-out fill-mode-both">
+								<Composer />
+							</div>
+						</AssistantIf>
+					</ThreadPrimitive.ViewportFooter>
+				</ThreadPrimitive.Viewport>
+			</ThreadPrimitive.Root>
+		</ThinkingStepsContext.Provider>
 	);
 };
 
@@ -197,7 +292,7 @@ const ThreadWelcome: FC = () => {
 const Composer: FC = () => {
 	return (
 		<ComposerPrimitive.Root className="aui-composer-root relative flex w-full flex-col">
-			<ComposerPrimitive.AttachmentDropzone className="aui-composer-attachment-dropzone flex w-full flex-col rounded-2xl border-input bg-muted px-1 pt-2 outline-none transition-shadow has-[textarea:focus-visible]:border-ring has-[textarea:focus-visible]:ring-2 has-[textarea:focus-visible]:ring-ring/20 data-[dragging=true]:border-ring data-[dragging=true]:border-dashed data-[dragging=true]:bg-accent/50">
+			<ComposerPrimitive.AttachmentDropzone className="aui-composer-attachment-dropzone flex w-full flex-col rounded-2xl border-input bg-muted px-1 pt-2 outline-none transition-shadow data-[dragging=true]:border-ring data-[dragging=true]:border-dashed data-[dragging=true]:bg-accent/50">
 				<ComposerAttachments />
 				<ComposerPrimitive.Input
 					placeholder="Ask SurfSense"
@@ -297,12 +392,22 @@ const MessageError: FC = () => {
 	);
 };
 
-const AssistantMessage: FC = () => {
+const AssistantMessageInner: FC = () => {
+	const thinkingStepsMap = useContext(ThinkingStepsContext);
+	
+	// Get the current message ID to look up thinking steps
+	const messageId = useMessage((m) => m.id);
+	const thinkingSteps = thinkingStepsMap.get(messageId) || [];
+	
 	return (
-		<MessagePrimitive.Root
-			className="aui-assistant-message-root fade-in slide-in-from-bottom-1 relative mx-auto w-full max-w-(--thread-max-width) animate-in py-3 duration-150"
-			data-role="assistant"
-		>
+		<>
+			{/* Show thinking steps BEFORE the text response */}
+			{thinkingSteps.length > 0 && (
+				<div className="mb-3">
+					<ThinkingStepsDisplay steps={thinkingSteps} />
+				</div>
+			)}
+			
 			<div className="aui-assistant-message-content wrap-break-word px-2 text-foreground leading-relaxed">
 				<MessagePrimitive.Parts
 					components={{
@@ -317,6 +422,17 @@ const AssistantMessage: FC = () => {
 				<BranchPicker />
 				<AssistantActionBar />
 			</div>
+		</>
+	);
+};
+
+const AssistantMessage: FC = () => {
+	return (
+		<MessagePrimitive.Root
+			className="aui-assistant-message-root fade-in slide-in-from-bottom-1 relative mx-auto w-full max-w-(--thread-max-width) animate-in py-3 duration-150"
+			data-role="assistant"
+		>
+			<AssistantMessageInner />
 		</MessagePrimitive.Root>
 	);
 };

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

@@ -0,0 +1,148 @@
+"use client"
+
+import {
+  Collapsible,
+  CollapsibleContent,
+  CollapsibleTrigger,
+} from "@/components/ui/collapsible"
+import { cn } from "@/lib/utils"
+import { Brain, ChevronDown, Circle, Loader2, Search, Sparkles, Lightbulb, CheckCircle2 } from "lucide-react"
+import React from "react"
+
+export type ChainOfThoughtItemProps = React.ComponentProps<"div">
+
+export const ChainOfThoughtItem = ({
+  children,
+  className,
+  ...props
+}: ChainOfThoughtItemProps) => (
+  <div className={cn("text-muted-foreground text-sm", className)} {...props}>
+    {children}
+  </div>
+)
+
+export type ChainOfThoughtTriggerProps = React.ComponentProps<
+  typeof CollapsibleTrigger
+> & {
+  leftIcon?: React.ReactNode
+  swapIconOnHover?: boolean
+}
+
+export const ChainOfThoughtTrigger = ({
+  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",
+      className
+    )}
+    {...props}
+  >
+    <div className="flex items-center gap-2">
+      {leftIcon ? (
+        <span className="relative inline-flex size-4 items-center justify-center">
+          <span
+            className={cn(
+              "transition-opacity",
+              swapIconOnHover && "group-hover:opacity-0"
+            )}
+          >
+            {leftIcon}
+          </span>
+          {swapIconOnHover && (
+            <ChevronDown className="absolute size-4 opacity-0 transition-opacity group-hover:opacity-100 group-data-[state=open]:rotate-180" />
+          )}
+        </span>
+      ) : (
+        <span className="relative inline-flex size-4 items-center justify-center">
+          <Circle className="size-2 fill-current" />
+        </span>
+      )}
+      <span>{children}</span>
+    </div>
+    {!leftIcon && (
+      <ChevronDown className="size-4 transition-transform group-data-[state=open]:rotate-180" />
+    )}
+  </CollapsibleTrigger>
+)
+
+export type ChainOfThoughtContentProps = React.ComponentProps<
+  typeof CollapsibleContent
+>
+
+export const ChainOfThoughtContent = ({
+  children,
+  className,
+  ...props
+}: ChainOfThoughtContentProps) => {
+  return (
+    <CollapsibleContent
+      className={cn(
+        "text-popover-foreground data-[state=closed]:animate-collapsible-up data-[state=open]:animate-collapsible-down overflow-hidden",
+        className
+      )}
+      {...props}
+    >
+      <div className="grid grid-cols-[min-content_minmax(0,1fr)] gap-x-4">
+        <div className="bg-primary/20 ml-1.75 h-full w-px group-data-[last=true]:hidden" />
+        <div className="ml-1.75 h-full w-px bg-transparent group-data-[last=false]:hidden" />
+        <div className="mt-2 space-y-2">{children}</div>
+      </div>
+    </CollapsibleContent>
+  )
+}
+
+export type ChainOfThoughtProps = {
+  children: React.ReactNode
+  className?: string
+}
+
+export function ChainOfThought({ children, className }: ChainOfThoughtProps) {
+  const childrenArray = React.Children.toArray(children)
+
+  return (
+    <div className={cn("space-y-0", className)}>
+      {childrenArray.map((child, index) => (
+        <React.Fragment key={index}>
+          {React.isValidElement(child) &&
+            React.cloneElement(
+              child as React.ReactElement<ChainOfThoughtStepProps>,
+              {
+                isLast: index === childrenArray.length - 1,
+              }
+            )}
+        </React.Fragment>
+      ))}
+    </div>
+  )
+}
+
+export type ChainOfThoughtStepProps = {
+  children: React.ReactNode
+  className?: string
+  isLast?: boolean
+}
+
+export const ChainOfThoughtStep = ({
+  children,
+  className,
+  isLast = false,
+  ...props
+}: ChainOfThoughtStepProps & React.ComponentProps<typeof Collapsible>) => {
+  return (
+    <Collapsible
+      className={cn("group", className)}
+      data-last={isLast}
+      {...props}
+    >
+      {children}
+      <div className="flex justify-start group-data-[last=true]:hidden">
+        <div className="bg-primary/20 ml-1.75 h-4 w-px" />
+      </div>
+    </Collapsible>
+  )
+}

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

@@ -0,0 +1,203 @@
+"use client";
+
+import { makeAssistantToolUI } from "@assistant-ui/react";
+import { Brain, CheckCircle2, Loader2, Search, Sparkles } from "lucide-react";
+import { useMemo } from "react";
+import {
+  ChainOfThought,
+  ChainOfThoughtContent,
+  ChainOfThoughtItem,
+  ChainOfThoughtStep,
+  ChainOfThoughtTrigger,
+} from "@/components/prompt-kit/chain-of-thought";
+import { cn } from "@/lib/utils";
+
+/**
+ * Types for the deepagent thinking/reasoning tool
+ */
+interface ThinkingStep {
+  id: string;
+  title: string;
+  items: string[];
+  status: "pending" | "in_progress" | "completed";
+}
+
+interface DeepAgentThinkingArgs {
+  query?: string;
+  context?: string;
+}
+
+interface DeepAgentThinkingResult {
+  steps?: ThinkingStep[];
+  status?: "thinking" | "searching" | "synthesizing" | "completed";
+  summary?: string;
+}
+
+/**
+ * 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
+ */
+function ThinkingStepDisplay({ step }: { step: ThinkingStep }) {
+  const icon = useMemo(() => getStepIcon(step.status, step.title), [step.status, step.title]);
+  
+  return (
+    <ChainOfThoughtStep defaultOpen={step.status === "in_progress"}>
+      <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>
+  );
+}
+
+/**
+ * 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({ args, 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">
+        <ChainOfThought>
+          {result.steps.map((step) => (
+            <ThinkingStepDisplay key={step.id} step={step} />
+          ))}
+        </ChainOfThought>
+      </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 />
+      ) : (
+        <ChainOfThought>
+          {steps.map((step) => (
+            <ThinkingStepDisplay key={step.id} step={step} />
+          ))}
+        </ChainOfThought>
+      )}
+    </div>
+  );
+}
+
+export type { ThinkingStep, DeepAgentThinkingArgs, DeepAgentThinkingResult };
+

surfsense_web/components/tool-ui/index.ts

@@ -8,3 +8,10 @@
 
 export { Audio } from "./audio";
 export { GeneratePodcastToolUI } from "./generate-podcast";
+export {
+  DeepAgentThinkingToolUI,
+  InlineThinkingDisplay,
+  type ThinkingStep,
+  type DeepAgentThinkingArgs,
+  type DeepAgentThinkingResult,
+} from "./deepagent-thinking";

surfsense_web/components/ui/collapsible.tsx

@@ -1,21 +1,33 @@
-"use client";
+"use client"
 
-import * as CollapsiblePrimitive from "@radix-ui/react-collapsible";
+import * as CollapsiblePrimitive from "@radix-ui/react-collapsible"
 
-function Collapsible({ ...props }: React.ComponentProps<typeof CollapsiblePrimitive.Root>) {
-	return <CollapsiblePrimitive.Root data-slot="collapsible" {...props} />;
+function Collapsible({
+  ...props
+}: React.ComponentProps<typeof CollapsiblePrimitive.Root>) {
+  return <CollapsiblePrimitive.Root data-slot="collapsible" {...props} />
 }
 
 function CollapsibleTrigger({
-	...props
+  ...props
 }: React.ComponentProps<typeof CollapsiblePrimitive.CollapsibleTrigger>) {
-	return <CollapsiblePrimitive.CollapsibleTrigger data-slot="collapsible-trigger" {...props} />;
+  return (
+    <CollapsiblePrimitive.CollapsibleTrigger
+      data-slot="collapsible-trigger"
+      {...props}
+    />
+  )
 }
 
 function CollapsibleContent({
-	...props
+  ...props
 }: React.ComponentProps<typeof CollapsiblePrimitive.CollapsibleContent>) {
-	return <CollapsiblePrimitive.CollapsibleContent data-slot="collapsible-content" {...props} />;
+  return (
+    <CollapsiblePrimitive.CollapsibleContent
+      data-slot="collapsible-content"
+      {...props}
+    />
+  )
 }
 
-export { Collapsible, CollapsibleTrigger, CollapsibleContent };
+export { Collapsible, CollapsibleTrigger, CollapsibleContent }