diff --git a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx index f569940d7..b303a45f5 100644 --- a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx +++ b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx @@ -34,10 +34,10 @@ import { closeEditorPanelAtom } from "@/atoms/editor/editor-panel.atom"; import { membersAtom } from "@/atoms/members/members-query.atoms"; import { currentUserAtom } from "@/atoms/user/user-query.atoms"; import { Thread } from "@/components/assistant-ui/thread"; +import { ThinkingStepsDataUI } from "@/components/assistant-ui/thinking-steps"; import { MobileEditorPanel } from "@/components/editor-panel/editor-panel"; import { MobileHitlEditPanel } from "@/components/hitl-edit-panel/hitl-edit-panel"; import { MobileReportPanel } from "@/components/report-panel/report-panel"; -import type { ThinkingStep } from "@/components/tool-ui/deepagent-thinking"; import { Skeleton } from "@/components/ui/skeleton"; import { useChatSessionStateSync } from "@/hooks/use-chat-session-state"; import { useMessagesElectric } from "@/hooks/use-messages-electric"; @@ -57,6 +57,7 @@ import { type ContentPartsState, readSSEStream, type ThinkingStepData, + updateThinkingSteps, updateToolCall, } from "@/lib/chat/streaming-state"; import { @@ -93,23 +94,6 @@ function markInterruptsCompleted(contentParts: Array<{ type: string; result?: un } } -/** - * 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 || []; -} - /** * Zod schema for mentioned document info (for type-safe parsing) */ @@ -183,11 +167,6 @@ export default function NewChatPage() { const [currentThread, setCurrentThread] = useState(null); const [messages, setMessages] = useState([]); const [isRunning, setIsRunning] = useState(false); - // Store thinking steps per message ID - kept separate from content to avoid - // "unsupported part type" errors from assistant-ui - const [messageThinkingSteps, setMessageThinkingSteps] = useState>( - new Map() - ); const abortControllerRef = useRef(null); const [pendingInterrupt, setPendingInterrupt] = useState<{ threadId: number; @@ -295,7 +274,6 @@ export default function NewChatPage() { setMessages([]); setThreadId(null); setCurrentThread(null); - setMessageThinkingSteps(new Map()); setMentionedDocuments([]); setSidebarDocuments([]); setMessageDocumentsMap({}); @@ -320,18 +298,8 @@ export default function NewChatPage() { const loadedMessages = messagesResponse.messages.map(convertToThreadMessage); setMessages(loadedMessages); - // Extract and restore thinking steps from persisted messages - const restoredThinkingSteps = new Map(); - // Extract and restore mentioned documents from persisted messages const restoredDocsMap: Record = {}; - for (const msg of messagesResponse.messages) { - if (msg.role === "assistant") { - const steps = extractThinkingSteps(msg.content); - if (steps.length > 0) { - restoredThinkingSteps.set(`msg-${msg.id}`, steps); - } - } if (msg.role === "user") { const docs = extractMentionedDocuments(msg.content); if (docs.length > 0) { @@ -339,9 +307,6 @@ export default function NewChatPage() { } } } - if (restoredThinkingSteps.size > 0) { - setMessageThinkingSteps(restoredThinkingSteps); - } if (Object.keys(restoredDocsMap).length > 0) { setMessageDocumentsMap(restoredDocsMap); } @@ -745,18 +710,17 @@ export default function NewChatPage() { } 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 thinking steps state for rendering - // The ThinkingStepsScrollHandler in Thread component - // will handle auto-scrolling when this state changes - setMessageThinkingSteps((prev) => { - const newMap = new Map(prev); - newMap.set(assistantMsgId, Array.from(currentThinkingSteps.values())); - return newMap; - }); + updateThinkingSteps(contentPartsState, currentThinkingSteps); + setMessages((prev) => + prev.map((m) => + m.id === assistantMsgId + ? { ...m, content: buildContentForUI(contentPartsState, TOOLS_WITH_UI) } + : m + ) + ); } break; } @@ -821,13 +785,8 @@ export default function NewChatPage() { } } - // Persist assistant message (with thinking steps for restoration on refresh) // Skip persistence for interrupted messages -- handleResume will persist the final version - const finalContent = buildContentForPersistence( - contentPartsState, - TOOLS_WITH_UI, - currentThinkingSteps - ); + const finalContent = buildContentForPersistence(contentPartsState, TOOLS_WITH_UI); if (contentParts.length > 0 && !wasInterrupted) { try { const savedMessage = await appendMessage(currentThreadId, { @@ -847,18 +806,6 @@ export default function NewChatPage() { ? { ...prev, assistantMsgId: newMsgId } : prev ); - - // Also update thinking steps map with new ID - setMessageThinkingSteps((prev) => { - const steps = prev.get(assistantMsgId); - if (steps) { - const newMap = new Map(prev); - newMap.delete(assistantMsgId); - newMap.set(newMsgId, steps); - return newMap; - } - return prev; - }); } catch (err) { console.error("Failed to persist assistant message:", err); } @@ -875,11 +822,7 @@ export default function NewChatPage() { (part.type === "tool-call" && TOOLS_WITH_UI.has(part.toolName)) ); if (hasContent && currentThreadId) { - const partialContent = buildContentForPersistence( - contentPartsState, - TOOLS_WITH_UI, - currentThinkingSteps - ); + const partialContent = buildContentForPersistence(contentPartsState, TOOLS_WITH_UI); try { const savedMessage = await appendMessage(currentThreadId, { role: "assistant", @@ -926,7 +869,6 @@ export default function NewChatPage() { } finally { setIsRunning(false); abortControllerRef.current = null; - // Note: We no longer clear thinking steps - they persist with the message } }, [ @@ -969,9 +911,7 @@ export default function NewChatPage() { const controller = new AbortController(); abortControllerRef.current = controller; - const currentThinkingSteps = new Map( - (messageThinkingSteps.get(assistantMsgId) ?? []).map((s) => [s.id, s]) - ); + const currentThinkingSteps = new Map(); const contentPartsState: ContentPartsState = { contentParts: [], @@ -998,6 +938,15 @@ export default function NewChatPage() { result: p.result as unknown, }); contentPartsState.currentTextPartIndex = -1; + } else if (p.type === "data-thinking-steps") { + const stepsData = p.data as { steps: ThinkingStepData[] } | undefined; + contentParts.push({ + type: "data-thinking-steps", + data: { steps: stepsData?.steps ?? [] }, + }); + for (const step of stepsData?.steps ?? []) { + currentThinkingSteps.set(step.id, step); + } } } } @@ -1115,11 +1064,14 @@ export default function NewChatPage() { const stepData = parsed.data as ThinkingStepData; if (stepData?.id) { currentThinkingSteps.set(stepData.id, stepData); - setMessageThinkingSteps((prev) => { - const newMap = new Map(prev); - newMap.set(assistantMsgId, Array.from(currentThinkingSteps.values())); - return newMap; - }); + updateThinkingSteps(contentPartsState, currentThinkingSteps); + setMessages((prev) => + prev.map((m) => + m.id === assistantMsgId + ? { ...m, content: buildContentForUI(contentPartsState, TOOLS_WITH_UI) } + : m + ) + ); } break; } @@ -1173,11 +1125,7 @@ export default function NewChatPage() { } } - const finalContent = buildContentForPersistence( - contentPartsState, - TOOLS_WITH_UI, - currentThinkingSteps - ); + const finalContent = buildContentForPersistence(contentPartsState, TOOLS_WITH_UI); if (contentParts.length > 0) { try { const savedMessage = await appendMessage(resumeThreadId, { @@ -1188,16 +1136,6 @@ export default function NewChatPage() { setMessages((prev) => prev.map((m) => (m.id === assistantMsgId ? { ...m, id: newMsgId } : m)) ); - setMessageThinkingSteps((prev) => { - const steps = prev.get(assistantMsgId); - if (steps) { - const newMap = new Map(prev); - newMap.delete(assistantMsgId); - newMap.set(newMsgId, steps); - return newMap; - } - return prev; - }); } catch (err) { console.error("Failed to persist resumed assistant message:", err); } @@ -1213,7 +1151,7 @@ export default function NewChatPage() { abortControllerRef.current = null; } }, - [pendingInterrupt, messages, searchSpaceId, messageThinkingSteps] + [pendingInterrupt, messages, searchSpaceId] ); useEffect(() => { @@ -1332,20 +1270,6 @@ export default function NewChatPage() { return prev; }); - // Clear thinking steps for the removed messages - setMessageThinkingSteps((prev) => { - const newMap = new Map(prev); - // Remove thinking steps for the last two messages - const lastTwoIds = messages - .slice(-2) - .map((m) => m.id) - .filter((id): id is string => !!id); - for (const id of lastTwoIds) { - newMap.delete(id); - } - return newMap; - }); - // Start streaming setIsRunning(true); const controller = new AbortController(); @@ -1476,11 +1400,14 @@ export default function NewChatPage() { const stepData = parsed.data as ThinkingStepData; if (stepData?.id) { currentThinkingSteps.set(stepData.id, stepData); - setMessageThinkingSteps((prev) => { - const newMap = new Map(prev); - newMap.set(assistantMsgId, Array.from(currentThinkingSteps.values())); - return newMap; - }); + updateThinkingSteps(contentPartsState, currentThinkingSteps); + setMessages((prev) => + prev.map((m) => + m.id === assistantMsgId + ? { ...m, content: buildContentForUI(contentPartsState, TOOLS_WITH_UI) } + : m + ) + ); } break; } @@ -1491,11 +1418,7 @@ export default function NewChatPage() { } // Persist messages after streaming completes - const finalContent = buildContentForPersistence( - contentPartsState, - TOOLS_WITH_UI, - currentThinkingSteps - ); + const finalContent = buildContentForPersistence(contentPartsState, TOOLS_WITH_UI); if (contentParts.length > 0) { try { // Persist user message (for both edit and reload modes, since backend deleted it) @@ -1526,18 +1449,6 @@ export default function NewChatPage() { prev.map((m) => (m.id === assistantMsgId ? { ...m, id: newMsgId } : m)) ); - setMessageThinkingSteps((prev) => { - const steps = prev.get(assistantMsgId); - if (steps) { - const newMap = new Map(prev); - newMap.delete(assistantMsgId); - newMap.set(newMsgId, steps); - return newMap; - } - return prev; - }); - - // Track successful response trackChatResponseReceived(searchSpaceId, threadId); } catch (err) { console.error("Failed to persist regenerated message:", err); @@ -1570,7 +1481,7 @@ export default function NewChatPage() { abortControllerRef.current = null; } }, - [threadId, searchSpaceId, messages, setMessageThinkingSteps, disabledTools] + [threadId, searchSpaceId, messages, disabledTools] ); // Handle editing a message - truncates history and regenerates with new query @@ -1675,9 +1586,10 @@ export default function NewChatPage() { return ( +
- +
diff --git a/surfsense_web/components/assistant-ui/assistant-message.tsx b/surfsense_web/components/assistant-ui/assistant-message.tsx index 72fed42ae..223c4fc37 100644 --- a/surfsense_web/components/assistant-ui/assistant-message.tsx +++ b/surfsense_web/components/assistant-ui/assistant-message.tsx @@ -8,20 +8,15 @@ import { import { useAtomValue } from "jotai"; import { CheckIcon, CopyIcon, DownloadIcon, MessageSquare, RefreshCwIcon } from "lucide-react"; import type { FC } from "react"; -import { useContext, useEffect, useMemo, useRef, useState } from "react"; +import { useEffect, useMemo, useRef, useState } from "react"; import { commentsEnabledAtom, targetCommentIdAtom } from "@/atoms/chat/current-thread.atom"; import { activeSearchSpaceIdAtom } from "@/atoms/search-spaces/search-space-query.atoms"; import { MarkdownText } from "@/components/assistant-ui/markdown-text"; -import { - ThinkingStepsContext, - ThinkingStepsDisplay, -} from "@/components/assistant-ui/thinking-steps"; import { ToolFallback } from "@/components/assistant-ui/tool-fallback"; import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button"; import { CommentPanelContainer } from "@/components/chat-comments/comment-panel-container/comment-panel-container"; import { CommentSheet } from "@/components/chat-comments/comment-sheet/comment-sheet"; import { CreateConfluencePageToolUI, DeleteConfluencePageToolUI, UpdateConfluencePageToolUI } from "@/components/tool-ui/confluence"; -import { DeepAgentThinkingToolUI } from "@/components/tool-ui/deepagent-thinking"; import { DisplayImageToolUI } from "@/components/tool-ui/display-image"; import { GeneratePodcastToolUI } from "@/components/tool-ui/generate-podcast"; import { GenerateReportToolUI } from "@/components/tool-ui/generate-report"; @@ -50,44 +45,15 @@ export const MessageError: FC = () => { ); }; -/** - * Custom component to render thinking steps from Context - */ -const ThinkingStepsPart: FC = () => { - const thinkingStepsMap = useContext(ThinkingStepsContext); - - // Get the current message ID to look up thinking steps - const messageId = useAuiState(({ message }) => message?.id); - const thinkingSteps = thinkingStepsMap.get(messageId) || []; - - // Check if this specific message is currently streaming - // A message is streaming if: thread is running AND this is the last assistant message - const isThreadRunning = useAuiState(({ thread }) => thread.isRunning); - const isLastMessage = useAuiState(({ message }) => message?.isLast ?? false); - const isMessageStreaming = isThreadRunning && isLastMessage; - - if (thinkingSteps.length === 0) return null; - - return ( -
- -
- ); -}; - const AssistantMessageInner: FC = () => { return ( <> - {/* Render thinking steps from message content - this ensures proper scroll tracking */} - -
>(new Map()); - /** * Chain of thought display component - single collapsible dropdown design */ @@ -18,7 +16,6 @@ export const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?: }) => { const [isOpen, setIsOpen] = useState(true); - // Derive effective status for each step const getEffectiveStatus = useCallback( (step: ThinkingStep): "pending" | "in_progress" | "completed" => { if (step.status === "in_progress" && !isThreadRunning) { @@ -36,7 +33,6 @@ export const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?: steps.every((s) => getEffectiveStatus(s) === "completed"); const isProcessing = isThreadRunning && !allCompleted; - // Auto-collapse when all tasks are completed useEffect(() => { if (allCompleted) { setIsOpen(false); @@ -61,7 +57,6 @@ export const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?: return (
- {/* Main collapsible header */} - {/* Collapsible content with CSS grid animation */}
- {/* Dot and line column */}
- {/* Vertical connection line - extends to next dot */} {!isLast && (
)} - {/* Step dot - on top of line */}
{effectiveStatus === "in_progress" ? ( @@ -117,9 +106,7 @@ export const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?:
- {/* Step content */}
- {/* Step title */}
- {/* Step items (sub-content) */} {step.items && step.items.length > 0 && (
{step.items.map((item, idx) => ( @@ -153,3 +139,28 @@ export const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?: ); }; +/** + * assistant-ui data UI component that renders thinking steps from message content. + * Registered globally via makeAssistantDataUI — renders inside MessagePrimitive.Parts + * at the position of the data part in the content array. + */ +function ThinkingStepsDataRenderer({ data }: { name: string; data: unknown }) { + const isThreadRunning = useAuiState(({ thread }) => thread.isRunning); + const isLastMessage = useAuiState(({ message }) => message?.isLast ?? false); + const isMessageStreaming = isThreadRunning && isLastMessage; + + const steps = (data as { steps: ThinkingStep[] } | null)?.steps ?? []; + if (steps.length === 0) return null; + + return ( +
+ +
+ ); +} + +export const ThinkingStepsDataUI = makeAssistantDataUI({ + name: "thinking-steps", + render: ThinkingStepsDataRenderer, +}); + diff --git a/surfsense_web/components/assistant-ui/thread.tsx b/surfsense_web/components/assistant-ui/thread.tsx index e8c765d8c..02e57ba20 100644 --- a/surfsense_web/components/assistant-ui/thread.tsx +++ b/surfsense_web/components/assistant-ui/thread.tsx @@ -1,9 +1,6 @@ import { - ActionBarPrimitive, AuiIf, - BranchPickerPrimitive, ComposerPrimitive, - ErrorPrimitive, MessagePrimitive, ThreadPrimitive, useAui, @@ -14,14 +11,8 @@ import { AlertCircle, ArrowDownIcon, ArrowUpIcon, - CheckIcon, - ChevronLeftIcon, - ChevronRightIcon, - CopyIcon, - DownloadIcon, Globe, Plus, - RefreshCwIcon, Settings2, SquareIcon, Unplug, @@ -32,7 +23,7 @@ import { import { AnimatePresence, motion } from "motion/react"; import Image from "next/image"; import { useParams } from "next/navigation"; -import { type FC, useCallback, useContext, useEffect, useMemo, useRef, useState } from "react"; +import { type FC, useCallback, useEffect, useMemo, useRef, useState } from "react"; import { createPortal } from "react-dom"; import { agentToolsAtom, @@ -63,12 +54,6 @@ import { InlineMentionEditor, type InlineMentionEditorRef, } from "@/components/assistant-ui/inline-mention-editor"; -import { MarkdownText } from "@/components/assistant-ui/markdown-text"; -import { - ThinkingStepsContext, - ThinkingStepsDisplay, -} from "@/components/assistant-ui/thinking-steps"; -import { ToolFallback } from "@/components/assistant-ui/tool-fallback"; import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button"; import { UserMessage } from "@/components/assistant-ui/user-message"; import { SLIDEOUT_PANEL_OPENED_EVENT } from "@/components/layout/ui/sidebar/SidebarSlideOutPanel"; @@ -76,7 +61,6 @@ import { DocumentMentionPicker, type DocumentMentionPickerRef, } from "@/components/new-chat/document-mention-picker"; -import type { ThinkingStep } from "@/components/tool-ui/deepagent-thinking"; import { Avatar, AvatarFallback, AvatarGroup } from "@/components/ui/avatar"; import { Button } from "@/components/ui/button"; import { Drawer, DrawerContent, DrawerHandle, DrawerTitle } from "@/components/ui/drawer"; @@ -111,16 +95,8 @@ const CYCLING_PLACEHOLDERS = [ "Check if this week's Slack messages reference any GitHub issues", ]; -interface ThreadProps { - messageThinkingSteps?: Map; -} - -export const Thread: FC = ({ messageThinkingSteps = new Map() }) => { - return ( - - - - ); +export const Thread: FC = () => { + return ; }; const ThreadContent: FC = () => { @@ -1132,97 +1108,6 @@ const TOOL_GROUPS: ToolGroup[] = [ }, ]; -const MessageError: FC = () => { - return ( - - - - - - ); -}; - -/** - * Custom component to render thinking steps from Context - */ -const ThinkingStepsPart: FC = () => { - const thinkingStepsMap = useContext(ThinkingStepsContext); - - // Get the current message ID to look up thinking steps - const messageId = useAuiState(({ message }) => message?.id); - const thinkingSteps = thinkingStepsMap.get(messageId) || []; - - // Check if this specific message is currently streaming - // A message is streaming if: thread is running AND this is the last assistant message - const isThreadRunning = useAuiState(({ thread }) => thread.isRunning); - const isLastMessage = useAuiState(({ message }) => message?.isLast ?? false); - const isMessageStreaming = isThreadRunning && isLastMessage; - - if (thinkingSteps.length === 0) return null; - - return ( -
- -
- ); -}; - -const AssistantMessageInner: FC = () => { - return ( - <> - {/* Render thinking steps from message content - this ensures proper scroll tracking */} - - -
- - -
- -
- - -
- - ); -}; - -const AssistantActionBar: FC = () => { - return ( - - - - message.isCopied}> - - - !message.isCopied}> - - - - - - - - - - - - - - - - ); -}; - const EditComposer: FC = () => { return ( @@ -1245,30 +1130,3 @@ const EditComposer: FC = () => { ); }; - -const BranchPicker: FC = ({ className, ...rest }) => { - return ( - - - - - - - - / - - - - - - - - ); -}; diff --git a/surfsense_web/components/public-chat/public-chat-view.tsx b/surfsense_web/components/public-chat/public-chat-view.tsx index 52830f601..f8dd6db5a 100644 --- a/surfsense_web/components/public-chat/public-chat-view.tsx +++ b/surfsense_web/components/public-chat/public-chat-view.tsx @@ -1,6 +1,7 @@ "use client"; import { AssistantRuntimeProvider } from "@assistant-ui/react"; +import { ThinkingStepsDataUI } from "@/components/assistant-ui/thinking-steps"; import { Navbar } from "@/components/homepage/navbar"; import { ReportPanel } from "@/components/report-panel/report-panel"; import { Spinner } from "@/components/ui/spinner"; @@ -39,6 +40,7 @@ export function PublicChatView({ shareToken }: PublicChatViewProps) {
+
} /> diff --git a/surfsense_web/lib/chat/message-utils.ts b/surfsense_web/lib/chat/message-utils.ts index 81538731b..7c0da03c4 100644 --- a/surfsense_web/lib/chat/message-utils.ts +++ b/surfsense_web/lib/chat/message-utils.ts @@ -2,8 +2,8 @@ import type { ThreadMessageLike } from "@assistant-ui/react"; import type { MessageRecord } from "./thread-persistence"; /** - * Convert backend message to assistant-ui ThreadMessageLike format - * Filters out 'thinking-steps' part as it's handled separately via messageThinkingSteps + * Convert backend message to assistant-ui ThreadMessageLike format. + * Migrates legacy `thinking-steps` parts to `data-thinking-steps` (assistant-ui data parts). */ export function convertToThreadMessage(msg: MessageRecord): ThreadMessageLike { let content: ThreadMessageLike["content"]; @@ -11,26 +11,34 @@ export function convertToThreadMessage(msg: MessageRecord): ThreadMessageLike { if (typeof msg.content === "string") { content = [{ type: "text", text: msg.content }]; } else if (Array.isArray(msg.content)) { - // Filter out custom metadata parts - they're handled separately - const filteredContent = msg.content.filter((part: unknown) => { - if (typeof part !== "object" || part === null || !("type" in part)) return true; - const partType = (part as { type: string }).type; - // Filter out metadata parts not directly renderable by assistant-ui - return ( - partType !== "thinking-steps" && - partType !== "mentioned-documents" && - partType !== "attachments" - ); - }); + const convertedContent = msg.content + .filter((part: unknown) => { + if (typeof part !== "object" || part === null || !("type" in part)) return true; + const partType = (part as { type: string }).type; + return partType !== "mentioned-documents" && partType !== "attachments"; + }) + .map((part: unknown) => { + if ( + typeof part === "object" && + part !== null && + "type" in part && + (part as { type: string }).type === "thinking-steps" + ) { + return { + type: "data-thinking-steps", + data: { steps: (part as { steps: unknown[] }).steps ?? [] }, + }; + } + return part; + }); content = - filteredContent.length > 0 - ? (filteredContent as ThreadMessageLike["content"]) + convertedContent.length > 0 + ? (convertedContent as ThreadMessageLike["content"]) : [{ type: "text", text: "" }]; } else { content = [{ type: "text", text: String(msg.content) }]; } - // Build metadata.custom for author display in shared chats const metadata = msg.author_id ? { custom: { diff --git a/surfsense_web/lib/chat/streaming-state.ts b/surfsense_web/lib/chat/streaming-state.ts index 4364fd515..3f1c498b6 100644 --- a/surfsense_web/lib/chat/streaming-state.ts +++ b/surfsense_web/lib/chat/streaming-state.ts @@ -15,6 +15,10 @@ export type ContentPart = toolName: string; args: Record; result?: unknown; + } + | { + type: "data-thinking-steps"; + data: { steps: ThinkingStepData[] }; }; export interface ContentPartsState { @@ -23,6 +27,32 @@ export interface ContentPartsState { toolCallIndices: Map; } +export function updateThinkingSteps( + state: ContentPartsState, + steps: Map +): void { + const stepsArray = Array.from(steps.values()); + const existingIdx = state.contentParts.findIndex((p) => p.type === "data-thinking-steps"); + + if (existingIdx >= 0) { + state.contentParts[existingIdx] = { + type: "data-thinking-steps", + data: { steps: stepsArray }, + }; + } else { + state.contentParts.unshift({ + type: "data-thinking-steps", + data: { steps: stepsArray }, + }); + if (state.currentTextPartIndex >= 0) { + state.currentTextPartIndex += 1; + } + for (const [id, idx] of state.toolCallIndices) { + state.toolCallIndices.set(id, idx + 1); + } + } +} + export function appendText(state: ContentPartsState, delta: string): void { if ( state.currentTextPartIndex >= 0 && @@ -75,6 +105,7 @@ export function buildContentForUI( const filtered = state.contentParts.filter((part) => { if (part.type === "text") return part.text.length > 0; if (part.type === "tool-call") return toolsWithUI.has(part.toolName); + if (part.type === "data-thinking-steps") return true; return false; }); return filtered.length > 0 @@ -84,23 +115,17 @@ export function buildContentForUI( export function buildContentForPersistence( state: ContentPartsState, - toolsWithUI: Set, - currentThinkingSteps: Map + toolsWithUI: Set ): unknown[] { const parts: unknown[] = []; - if (currentThinkingSteps.size > 0) { - parts.push({ - type: "thinking-steps", - steps: Array.from(currentThinkingSteps.values()), - }); - } - for (const part of state.contentParts) { if (part.type === "text" && part.text.length > 0) { parts.push(part); } else if (part.type === "tool-call" && toolsWithUI.has(part.toolName)) { parts.push(part); + } else if (part.type === "data-thinking-steps") { + parts.push(part); } }