apunkt/SurfSense
1
0
Fork 0
mirror of https://github.com/MODSetter/SurfSense.git synced 2026-05-17 18:35:19 +02:00
Code Issues Projects Releases Packages Wiki Activity

refactor: migrate thinking steps handling to new data structure and streamline related components

Browse source
This commit is contained in:
Anish Sarkar 2026-03-24 02:23:05 +05:30
parent b8f3f41326
commit e587b588c9
7 changed files with 135 additions and 353 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -34,10 +34,10 @@ import { closeEditorPanelAtom } from "@/atoms/editor/editor-panel.atom";
import { membersAtom } from "@/atoms/members/members-query.atoms"; import { membersAtom } from "@/atoms/members/members-query.atoms";
import { currentUserAtom } from "@/atoms/user/user-query.atoms"; import { currentUserAtom } from "@/atoms/user/user-query.atoms";
import { Thread } from "@/components/assistant-ui/thread"; import { Thread } from "@/components/assistant-ui/thread";
import { ThinkingStepsDataUI } from "@/components/assistant-ui/thinking-steps";
import { MobileEditorPanel } from "@/components/editor-panel/editor-panel"; import { MobileEditorPanel } from "@/components/editor-panel/editor-panel";
import { MobileHitlEditPanel } from "@/components/hitl-edit-panel/hitl-edit-panel"; import { MobileHitlEditPanel } from "@/components/hitl-edit-panel/hitl-edit-panel";
import { MobileReportPanel } from "@/components/report-panel/report-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 { Skeleton } from "@/components/ui/skeleton";
import { useChatSessionStateSync } from "@/hooks/use-chat-session-state"; import { useChatSessionStateSync } from "@/hooks/use-chat-session-state";
import { useMessagesElectric } from "@/hooks/use-messages-electric"; import { useMessagesElectric } from "@/hooks/use-messages-electric";
@ -57,6 +57,7 @@ import {
type ContentPartsState, type ContentPartsState,
readSSEStream, readSSEStream,
type ThinkingStepData, type ThinkingStepData,
updateThinkingSteps,
updateToolCall, updateToolCall,
} from "@/lib/chat/streaming-state"; } from "@/lib/chat/streaming-state";
import { 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) * Zod schema for mentioned document info (for type-safe parsing)
*/ */
@ -183,11 +167,6 @@ export default function NewChatPage() {
const [currentThread, setCurrentThread] = useState<ThreadRecord | null>(null); const [currentThread, setCurrentThread] = useState<ThreadRecord | null>(null);
const [messages, setMessages] = useState<ThreadMessageLike[]>([]); const [messages, setMessages] = useState<ThreadMessageLike[]>([]);
const [isRunning, setIsRunning] = useState(false); 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<Map<string, ThinkingStep[]>>(
new Map()
);
const abortControllerRef = useRef<AbortController | null>(null); const abortControllerRef = useRef<AbortController | null>(null);
const [pendingInterrupt, setPendingInterrupt] = useState<{ const [pendingInterrupt, setPendingInterrupt] = useState<{
threadId: number; threadId: number;
@ -295,7 +274,6 @@ export default function NewChatPage() {
setMessages([]); setMessages([]);
setThreadId(null); setThreadId(null);
setCurrentThread(null); setCurrentThread(null);
setMessageThinkingSteps(new Map());
setMentionedDocuments([]); setMentionedDocuments([]);
setSidebarDocuments([]); setSidebarDocuments([]);
setMessageDocumentsMap({}); setMessageDocumentsMap({});
@ -320,18 +298,8 @@ export default function NewChatPage() {
const loadedMessages = messagesResponse.messages.map(convertToThreadMessage); const loadedMessages = messagesResponse.messages.map(convertToThreadMessage);
setMessages(loadedMessages); setMessages(loadedMessages);
// Extract and restore thinking steps from persisted messages
const restoredThinkingSteps = new Map<string, ThinkingStep[]>();
// Extract and restore mentioned documents from persisted messages
const restoredDocsMap: Record<string, MentionedDocumentInfo[]> = {}; const restoredDocsMap: Record<string, MentionedDocumentInfo[]> = {};
for (const msg of messagesResponse.messages) { 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") { if (msg.role === "user") {
const docs = extractMentionedDocuments(msg.content); const docs = extractMentionedDocuments(msg.content);
if (docs.length > 0) { if (docs.length > 0) {
@ -339,9 +307,6 @@ export default function NewChatPage() {
} }
} }
} }
if (restoredThinkingSteps.size > 0) {
setMessageThinkingSteps(restoredThinkingSteps);
}
if (Object.keys(restoredDocsMap).length > 0) { if (Object.keys(restoredDocsMap).length > 0) {
setMessageDocumentsMap(restoredDocsMap); setMessageDocumentsMap(restoredDocsMap);
} }
@ -745,18 +710,17 @@ export default function NewChatPage() {
} }
case "data-thinking-step": { case "data-thinking-step": {
// Handle thinking step events for chain-of-thought display
const stepData = parsed.data as ThinkingStepData; const stepData = parsed.data as ThinkingStepData;
if (stepData?.id) { if (stepData?.id) {
currentThinkingSteps.set(stepData.id, stepData); currentThinkingSteps.set(stepData.id, stepData);
// Update thinking steps state for rendering updateThinkingSteps(contentPartsState, currentThinkingSteps);
// The ThinkingStepsScrollHandler in Thread component setMessages((prev) =>
// will handle auto-scrolling when this state changes prev.map((m) =>
setMessageThinkingSteps((prev) => { m.id === assistantMsgId
const newMap = new Map(prev); ? { ...m, content: buildContentForUI(contentPartsState, TOOLS_WITH_UI) }
newMap.set(assistantMsgId, Array.from(currentThinkingSteps.values())); : m
return newMap; )
}); );
} }
break; 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 // Skip persistence for interrupted messages -- handleResume will persist the final version
const finalContent = buildContentForPersistence( const finalContent = buildContentForPersistence(contentPartsState, TOOLS_WITH_UI);
contentPartsState,
TOOLS_WITH_UI,
currentThinkingSteps
);
if (contentParts.length > 0 && !wasInterrupted) { if (contentParts.length > 0 && !wasInterrupted) {
try { try {
const savedMessage = await appendMessage(currentThreadId, { const savedMessage = await appendMessage(currentThreadId, {
@ -847,18 +806,6 @@ export default function NewChatPage() {
? { ...prev, assistantMsgId: newMsgId } ? { ...prev, assistantMsgId: newMsgId }
: prev : 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) { } catch (err) {
console.error("Failed to persist assistant message:", 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)) (part.type === "tool-call" && TOOLS_WITH_UI.has(part.toolName))
); );
if (hasContent && currentThreadId) { if (hasContent && currentThreadId) {
const partialContent = buildContentForPersistence( const partialContent = buildContentForPersistence(contentPartsState, TOOLS_WITH_UI);
contentPartsState,
TOOLS_WITH_UI,
currentThinkingSteps
);
try { try {
const savedMessage = await appendMessage(currentThreadId, { const savedMessage = await appendMessage(currentThreadId, {
role: "assistant", role: "assistant",
@ -926,7 +869,6 @@ export default function NewChatPage() {
} finally { } finally {
setIsRunning(false); setIsRunning(false);
abortControllerRef.current = null; 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(); const controller = new AbortController();
abortControllerRef.current = controller; abortControllerRef.current = controller;
const currentThinkingSteps = new Map<string, ThinkingStepData>( const currentThinkingSteps = new Map<string, ThinkingStepData>();
(messageThinkingSteps.get(assistantMsgId) ?? []).map((s) => [s.id, s])
);
const contentPartsState: ContentPartsState = { const contentPartsState: ContentPartsState = {
contentParts: [], contentParts: [],
@ -998,6 +938,15 @@ export default function NewChatPage() {
result: p.result as unknown, result: p.result as unknown,
}); });
contentPartsState.currentTextPartIndex = -1; 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; const stepData = parsed.data as ThinkingStepData;
if (stepData?.id) { if (stepData?.id) {
currentThinkingSteps.set(stepData.id, stepData); currentThinkingSteps.set(stepData.id, stepData);
setMessageThinkingSteps((prev) => { updateThinkingSteps(contentPartsState, currentThinkingSteps);
const newMap = new Map(prev); setMessages((prev) =>
newMap.set(assistantMsgId, Array.from(currentThinkingSteps.values())); prev.map((m) =>
return newMap; m.id === assistantMsgId
}); ? { ...m, content: buildContentForUI(contentPartsState, TOOLS_WITH_UI) }
: m
)
);
} }
break; break;
} }
@ -1173,11 +1125,7 @@ export default function NewChatPage() {
} }
} }
const finalContent = buildContentForPersistence( const finalContent = buildContentForPersistence(contentPartsState, TOOLS_WITH_UI);
contentPartsState,
TOOLS_WITH_UI,
currentThinkingSteps
);
if (contentParts.length > 0) { if (contentParts.length > 0) {
try { try {
const savedMessage = await appendMessage(resumeThreadId, { const savedMessage = await appendMessage(resumeThreadId, {
@ -1188,16 +1136,6 @@ export default function NewChatPage() {
setMessages((prev) => setMessages((prev) =>
prev.map((m) => (m.id === assistantMsgId ? { ...m, id: newMsgId } : m)) 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) { } catch (err) {
console.error("Failed to persist resumed assistant message:", err); console.error("Failed to persist resumed assistant message:", err);
} }
@ -1213,7 +1151,7 @@ export default function NewChatPage() {
abortControllerRef.current = null; abortControllerRef.current = null;
} }
}, },
[pendingInterrupt, messages, searchSpaceId, messageThinkingSteps] [pendingInterrupt, messages, searchSpaceId]
); );
useEffect(() => { useEffect(() => {
@ -1332,20 +1270,6 @@ export default function NewChatPage() {
return prev; 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 // Start streaming
setIsRunning(true); setIsRunning(true);
const controller = new AbortController(); const controller = new AbortController();
@ -1476,11 +1400,14 @@ export default function NewChatPage() {
const stepData = parsed.data as ThinkingStepData; const stepData = parsed.data as ThinkingStepData;
if (stepData?.id) { if (stepData?.id) {
currentThinkingSteps.set(stepData.id, stepData); currentThinkingSteps.set(stepData.id, stepData);
setMessageThinkingSteps((prev) => { updateThinkingSteps(contentPartsState, currentThinkingSteps);
const newMap = new Map(prev); setMessages((prev) =>
newMap.set(assistantMsgId, Array.from(currentThinkingSteps.values())); prev.map((m) =>
return newMap; m.id === assistantMsgId
}); ? { ...m, content: buildContentForUI(contentPartsState, TOOLS_WITH_UI) }
: m
)
);
} }
break; break;
} }
@ -1491,11 +1418,7 @@ export default function NewChatPage() {
} }
// Persist messages after streaming completes // Persist messages after streaming completes
const finalContent = buildContentForPersistence( const finalContent = buildContentForPersistence(contentPartsState, TOOLS_WITH_UI);
contentPartsState,
TOOLS_WITH_UI,
currentThinkingSteps
);
if (contentParts.length > 0) { if (contentParts.length > 0) {
try { try {
// Persist user message (for both edit and reload modes, since backend deleted it) // 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)) 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); trackChatResponseReceived(searchSpaceId, threadId);
} catch (err) { } catch (err) {
console.error("Failed to persist regenerated message:", err); console.error("Failed to persist regenerated message:", err);
@ -1570,7 +1481,7 @@ export default function NewChatPage() {
abortControllerRef.current = null; abortControllerRef.current = null;
} }
}, },
[threadId, searchSpaceId, messages, setMessageThinkingSteps, disabledTools] [threadId, searchSpaceId, messages, disabledTools]
); );
// Handle editing a message - truncates history and regenerates with new query // Handle editing a message - truncates history and regenerates with new query
@ -1675,9 +1586,10 @@ export default function NewChatPage() {
return ( return (
<AssistantRuntimeProvider runtime={runtime}> <AssistantRuntimeProvider runtime={runtime}>
<ThinkingStepsDataUI />
<div key={searchSpaceId} className="flex h-[calc(100dvh-64px)] overflow-hidden"> <div key={searchSpaceId} className="flex h-[calc(100dvh-64px)] overflow-hidden">
<div className="flex-1 flex flex-col min-w-0 overflow-hidden"> <div className="flex-1 flex flex-col min-w-0 overflow-hidden">
<Thread messageThinkingSteps={messageThinkingSteps} /> <Thread />
</div> </div>
<MobileReportPanel /> <MobileReportPanel />
<MobileEditorPanel /> <MobileEditorPanel />

36
surfsense_web/components/assistant-ui/assistant-message.tsx
View file

@ -8,20 +8,15 @@ import {
import { useAtomValue } from "jotai"; import { useAtomValue } from "jotai";
import { CheckIcon, CopyIcon, DownloadIcon, MessageSquare, RefreshCwIcon } from "lucide-react"; import { CheckIcon, CopyIcon, DownloadIcon, MessageSquare, RefreshCwIcon } from "lucide-react";
import type { FC } from "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 { commentsEnabledAtom, targetCommentIdAtom } from "@/atoms/chat/current-thread.atom";
import { activeSearchSpaceIdAtom } from "@/atoms/search-spaces/search-space-query.atoms"; import { activeSearchSpaceIdAtom } from "@/atoms/search-spaces/search-space-query.atoms";
import { MarkdownText } from "@/components/assistant-ui/markdown-text"; 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 { ToolFallback } from "@/components/assistant-ui/tool-fallback";
import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button"; import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
import { CommentPanelContainer } from "@/components/chat-comments/comment-panel-container/comment-panel-container"; import { CommentPanelContainer } from "@/components/chat-comments/comment-panel-container/comment-panel-container";
import { CommentSheet } from "@/components/chat-comments/comment-sheet/comment-sheet"; import { CommentSheet } from "@/components/chat-comments/comment-sheet/comment-sheet";
import { CreateConfluencePageToolUI, DeleteConfluencePageToolUI, UpdateConfluencePageToolUI } from "@/components/tool-ui/confluence"; 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 { DisplayImageToolUI } from "@/components/tool-ui/display-image";
import { GeneratePodcastToolUI } from "@/components/tool-ui/generate-podcast"; import { GeneratePodcastToolUI } from "@/components/tool-ui/generate-podcast";
import { GenerateReportToolUI } from "@/components/tool-ui/generate-report"; 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 (
<div className="mb-3">
<ThinkingStepsDisplay steps={thinkingSteps} isThreadRunning={isMessageStreaming} />
</div>
);
};
const AssistantMessageInner: FC = () => { const AssistantMessageInner: FC = () => {
return ( return (
<> <>
{/* Render thinking steps from message content - this ensures proper scroll tracking */}
<ThinkingStepsPart />
<div className="aui-assistant-message-content wrap-break-word px-2 text-foreground leading-relaxed"> <div className="aui-assistant-message-content wrap-break-word px-2 text-foreground leading-relaxed">
<MessagePrimitive.Parts <MessagePrimitive.Parts
components={{ components={{
Text: MarkdownText, Text: MarkdownText,
tools: { tools: {
by_name: { by_name: {
deepagent_thinking: DeepAgentThinkingToolUI,
generate_report: GenerateReportToolUI, generate_report: GenerateReportToolUI,
generate_podcast: GeneratePodcastToolUI, generate_podcast: GeneratePodcastToolUI,
generate_video_presentation: GenerateVideoPresentationToolUI, generate_video_presentation: GenerateVideoPresentationToolUI,

43
surfsense_web/components/assistant-ui/thinking-steps.tsx
View file

@ -1,14 +1,12 @@
import { makeAssistantDataUI, useAuiState } from "@assistant-ui/react";
import { ChevronRightIcon } from "lucide-react"; import { ChevronRightIcon } from "lucide-react";
import type { FC } from "react"; import type { FC } from "react";
import { createContext, useCallback, useEffect, useState } from "react"; import { useCallback, useEffect, useState } from "react";
import { ChainOfThoughtItem } from "@/components/prompt-kit/chain-of-thought"; import { ChainOfThoughtItem } from "@/components/prompt-kit/chain-of-thought";
import { TextShimmerLoader } from "@/components/prompt-kit/loader"; import { TextShimmerLoader } from "@/components/prompt-kit/loader";
import type { ThinkingStep } from "@/components/tool-ui/deepagent-thinking"; import type { ThinkingStep } from "@/components/tool-ui/deepagent-thinking";
import { cn } from "@/lib/utils"; import { cn } from "@/lib/utils";
// Context to pass thinking steps to AssistantMessage
export const ThinkingStepsContext = createContext<Map<string, ThinkingStep[]>>(new Map());
/** /**
* Chain of thought display component - single collapsible dropdown design * 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); const [isOpen, setIsOpen] = useState(true);
// Derive effective status for each step
const getEffectiveStatus = useCallback( const getEffectiveStatus = useCallback(
(step: ThinkingStep): "pending" | "in_progress" | "completed" => { (step: ThinkingStep): "pending" | "in_progress" | "completed" => {
if (step.status === "in_progress" && !isThreadRunning) { if (step.status === "in_progress" && !isThreadRunning) {
@ -36,7 +33,6 @@ export const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?:
steps.every((s) => getEffectiveStatus(s) === "completed"); steps.every((s) => getEffectiveStatus(s) === "completed");
const isProcessing = isThreadRunning && !allCompleted; const isProcessing = isThreadRunning && !allCompleted;
// Auto-collapse when all tasks are completed
useEffect(() => { useEffect(() => {
if (allCompleted) { if (allCompleted) {
setIsOpen(false); setIsOpen(false);
@ -61,7 +57,6 @@ export const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?:
return ( return (
<div className="mx-auto w-full max-w-(--thread-max-width) px-2 py-2"> <div className="mx-auto w-full max-w-(--thread-max-width) px-2 py-2">
<div className="rounded-lg"> <div className="rounded-lg">
{/* Main collapsible header */}
<button <button
type="button" type="button"
onClick={() => setIsOpen(!isOpen)} onClick={() => setIsOpen(!isOpen)}
@ -70,20 +65,17 @@ export const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?:
"text-muted-foreground hover:text-foreground" "text-muted-foreground hover:text-foreground"
)} )}
> >
{/* Header text with shimmer if processing (streaming) */}
{isProcessing ? ( {isProcessing ? (
<TextShimmerLoader text={getHeaderText()} size="sm" /> <TextShimmerLoader text={getHeaderText()} size="sm" />
) : ( ) : (
<span>{getHeaderText()}</span> <span>{getHeaderText()}</span>
)} )}
{/* Chevron */}
<ChevronRightIcon <ChevronRightIcon
className={cn("size-4 transition-transform duration-200", isOpen && "rotate-90")} className={cn("size-4 transition-transform duration-200", isOpen && "rotate-90")}
/> />
</button> </button>
{/* Collapsible content with CSS grid animation */}
<div <div
className={cn( className={cn(
"grid transition-[grid-template-rows] duration-300 ease-out", "grid transition-[grid-template-rows] duration-300 ease-out",
@ -98,13 +90,10 @@ export const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?:
return ( return (
<div key={step.id} className="relative flex gap-3"> <div key={step.id} className="relative flex gap-3">
{/* Dot and line column */}
<div className="relative flex flex-col items-center w-2"> <div className="relative flex flex-col items-center w-2">
{/* Vertical connection line - extends to next dot */}
{!isLast && ( {!isLast && (
<div className="absolute left-1/2 top-[15px] -bottom-[7px] w-px -translate-x-1/2 bg-muted-foreground/30" /> <div className="absolute left-1/2 top-[15px] -bottom-[7px] w-px -translate-x-1/2 bg-muted-foreground/30" />
)} )}
{/* Step dot - on top of line */}
<div className="relative z-10 mt-[7px] flex shrink-0 items-center justify-center"> <div className="relative z-10 mt-[7px] flex shrink-0 items-center justify-center">
{effectiveStatus === "in_progress" ? ( {effectiveStatus === "in_progress" ? (
<span className="relative flex size-2"> <span className="relative flex size-2">
@ -117,9 +106,7 @@ export const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?:
</div> </div>
</div> </div>
{/* Step content */}
<div className="flex-1 min-w-0 pb-4"> <div className="flex-1 min-w-0 pb-4">
{/* Step title */}
<div <div
className={cn( className={cn(
"text-sm leading-5", "text-sm leading-5",
@ -131,7 +118,6 @@ export const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?:
{step.title} {step.title}
</div> </div>
{/* Step items (sub-content) */}
{step.items && step.items.length > 0 && ( {step.items && step.items.length > 0 && (
<div className="mt-1 space-y-0.5"> <div className="mt-1 space-y-0.5">
{step.items.map((item, idx) => ( {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 (
<div className="mb-3 -mx-2 leading-normal">
<ThinkingStepsDisplay steps={steps} isThreadRunning={isMessageStreaming} />
</div>
);
}
export const ThinkingStepsDataUI = makeAssistantDataUI({
name: "thinking-steps",
render: ThinkingStepsDataRenderer,
});

148
surfsense_web/components/assistant-ui/thread.tsx
View file

@ -1,9 +1,6 @@
import { import {
ActionBarPrimitive,
AuiIf, AuiIf,
BranchPickerPrimitive,
ComposerPrimitive, ComposerPrimitive,
ErrorPrimitive,
MessagePrimitive, MessagePrimitive,
ThreadPrimitive, ThreadPrimitive,
useAui, useAui,
@ -14,14 +11,8 @@ import {
AlertCircle, AlertCircle,
ArrowDownIcon, ArrowDownIcon,
ArrowUpIcon, ArrowUpIcon,
CheckIcon,
ChevronLeftIcon,
ChevronRightIcon,
CopyIcon,
DownloadIcon,
Globe, Globe,
Plus, Plus,
RefreshCwIcon,
Settings2, Settings2,
SquareIcon, SquareIcon,
Unplug, Unplug,
@ -32,7 +23,7 @@ import {
import { AnimatePresence, motion } from "motion/react"; import { AnimatePresence, motion } from "motion/react";
import Image from "next/image"; import Image from "next/image";
import { useParams } from "next/navigation"; 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 { createPortal } from "react-dom";
import { import {
agentToolsAtom, agentToolsAtom,
@ -63,12 +54,6 @@ import {
InlineMentionEditor, InlineMentionEditor,
type InlineMentionEditorRef, type InlineMentionEditorRef,
} from "@/components/assistant-ui/inline-mention-editor"; } 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 { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
import { UserMessage } from "@/components/assistant-ui/user-message"; import { UserMessage } from "@/components/assistant-ui/user-message";
import { SLIDEOUT_PANEL_OPENED_EVENT } from "@/components/layout/ui/sidebar/SidebarSlideOutPanel"; import { SLIDEOUT_PANEL_OPENED_EVENT } from "@/components/layout/ui/sidebar/SidebarSlideOutPanel";
@ -76,7 +61,6 @@ import {
DocumentMentionPicker, DocumentMentionPicker,
type DocumentMentionPickerRef, type DocumentMentionPickerRef,
} from "@/components/new-chat/document-mention-picker"; } 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 { Avatar, AvatarFallback, AvatarGroup } from "@/components/ui/avatar";
import { Button } from "@/components/ui/button"; import { Button } from "@/components/ui/button";
import { Drawer, DrawerContent, DrawerHandle, DrawerTitle } from "@/components/ui/drawer"; 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", "Check if this week's Slack messages reference any GitHub issues",
]; ];
interface ThreadProps { export const Thread: FC = () => {
messageThinkingSteps?: Map<string, ThinkingStep[]>; return <ThreadContent />;
}
export const Thread: FC<ThreadProps> = ({ messageThinkingSteps = new Map() }) => {
return (
<ThinkingStepsContext.Provider value={messageThinkingSteps}>
<ThreadContent />
</ThinkingStepsContext.Provider>
);
}; };
const ThreadContent: FC = () => { const ThreadContent: FC = () => {
@ -1132,97 +1108,6 @@ const TOOL_GROUPS: ToolGroup[] = [
}, },
]; ];
const MessageError: FC = () => {
return (
<MessagePrimitive.Error>
<ErrorPrimitive.Root className="aui-message-error-root mt-2 rounded-md border border-destructive bg-destructive/10 p-3 text-destructive text-sm dark:bg-destructive/5 dark:text-red-200">
<ErrorPrimitive.Message className="aui-message-error-message line-clamp-2" />
</ErrorPrimitive.Root>
</MessagePrimitive.Error>
);
};
/**
* 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 (
<div className="mb-3">
<ThinkingStepsDisplay steps={thinkingSteps} isThreadRunning={isMessageStreaming} />
</div>
);
};
const AssistantMessageInner: FC = () => {
return (
<>
{/* Render thinking steps from message content - this ensures proper scroll tracking */}
<ThinkingStepsPart />
<div className="aui-assistant-message-content wrap-break-word px-2 text-foreground leading-relaxed">
<MessagePrimitive.Parts
components={{
Text: MarkdownText,
tools: { Fallback: ToolFallback },
}}
/>
<MessageError />
</div>
<div className="aui-assistant-message-footer mt-1 mb-5 ml-2 flex">
<BranchPicker />
<AssistantActionBar />
</div>
</>
);
};
const AssistantActionBar: FC = () => {
return (
<ActionBarPrimitive.Root
hideWhenRunning
autohide="not-last"
autohideFloat="single-branch"
className="aui-assistant-action-bar-root -ml-1 col-start-3 row-start-2 flex gap-1 text-muted-foreground data-floating:absolute data-floating:rounded-md data-floating:border data-floating:bg-background data-floating:p-1 data-floating:shadow-sm"
>
<ActionBarPrimitive.Copy asChild>
<TooltipIconButton tooltip="Copy">
<AuiIf condition={({ message }) => message.isCopied}>
<CheckIcon />
</AuiIf>
<AuiIf condition={({ message }) => !message.isCopied}>
<CopyIcon />
</AuiIf>
</TooltipIconButton>
</ActionBarPrimitive.Copy>
<ActionBarPrimitive.ExportMarkdown asChild>
<TooltipIconButton tooltip="Export as Markdown">
<DownloadIcon />
</TooltipIconButton>
</ActionBarPrimitive.ExportMarkdown>
<ActionBarPrimitive.Reload asChild>
<TooltipIconButton tooltip="Refresh">
<RefreshCwIcon />
</TooltipIconButton>
</ActionBarPrimitive.Reload>
</ActionBarPrimitive.Root>
);
};
const EditComposer: FC = () => { const EditComposer: FC = () => {
return ( return (
<MessagePrimitive.Root className="aui-edit-composer-wrapper mx-auto flex w-full max-w-(--thread-max-width) flex-col px-2 py-3"> <MessagePrimitive.Root className="aui-edit-composer-wrapper mx-auto flex w-full max-w-(--thread-max-width) flex-col px-2 py-3">
@ -1245,30 +1130,3 @@ const EditComposer: FC = () => {
</MessagePrimitive.Root> </MessagePrimitive.Root>
); );
}; };
const BranchPicker: FC<BranchPickerPrimitive.Root.Props> = ({ className, ...rest }) => {
return (
<BranchPickerPrimitive.Root
hideWhenSingleBranch
className={cn(
"aui-branch-picker-root -ml-2 mr-2 inline-flex items-center text-muted-foreground text-xs",
className
)}
{...rest}
>
<BranchPickerPrimitive.Previous asChild>
<TooltipIconButton tooltip="Previous">
<ChevronLeftIcon />
</TooltipIconButton>
</BranchPickerPrimitive.Previous>
<span className="aui-branch-picker-state font-medium">
<BranchPickerPrimitive.Number /> / <BranchPickerPrimitive.Count />
</span>
<BranchPickerPrimitive.Next asChild>
<TooltipIconButton tooltip="Next">
<ChevronRightIcon />
</TooltipIconButton>
</BranchPickerPrimitive.Next>
</BranchPickerPrimitive.Root>
);
};

2
surfsense_web/components/public-chat/public-chat-view.tsx
View file

@ -1,6 +1,7 @@
"use client"; "use client";
import { AssistantRuntimeProvider } from "@assistant-ui/react"; import { AssistantRuntimeProvider } from "@assistant-ui/react";
import { ThinkingStepsDataUI } from "@/components/assistant-ui/thinking-steps";
import { Navbar } from "@/components/homepage/navbar"; import { Navbar } from "@/components/homepage/navbar";
import { ReportPanel } from "@/components/report-panel/report-panel"; import { ReportPanel } from "@/components/report-panel/report-panel";
import { Spinner } from "@/components/ui/spinner"; import { Spinner } from "@/components/ui/spinner";
@ -39,6 +40,7 @@ export function PublicChatView({ shareToken }: PublicChatViewProps) {
<main className="min-h-screen bg-main-panel text-foreground overflow-x-hidden"> <main className="min-h-screen bg-main-panel text-foreground overflow-x-hidden">
<Navbar scrolledBgClassName={navbarScrolledBg} /> <Navbar scrolledBgClassName={navbarScrolledBg} />
<AssistantRuntimeProvider runtime={runtime}> <AssistantRuntimeProvider runtime={runtime}>
<ThinkingStepsDataUI />
<div className="flex h-screen pt-16 overflow-hidden"> <div className="flex h-screen pt-16 overflow-hidden">
<div className="flex-1 flex flex-col min-w-0 overflow-hidden"> <div className="flex-1 flex flex-col min-w-0 overflow-hidden">
<PublicThread footer={<PublicChatFooter shareToken={shareToken} />} /> <PublicThread footer={<PublicChatFooter shareToken={shareToken} />} />

40
surfsense_web/lib/chat/message-utils.ts
View file

@ -2,8 +2,8 @@ import type { ThreadMessageLike } from "@assistant-ui/react";
import type { MessageRecord } from "./thread-persistence"; import type { MessageRecord } from "./thread-persistence";
/** /**
* Convert backend message to assistant-ui ThreadMessageLike format * Convert backend message to assistant-ui ThreadMessageLike format.
* Filters out 'thinking-steps' part as it's handled separately via messageThinkingSteps * Migrates legacy `thinking-steps` parts to `data-thinking-steps` (assistant-ui data parts).
*/ */
export function convertToThreadMessage(msg: MessageRecord): ThreadMessageLike { export function convertToThreadMessage(msg: MessageRecord): ThreadMessageLike {
let content: ThreadMessageLike["content"]; let content: ThreadMessageLike["content"];
@ -11,26 +11,34 @@ export function convertToThreadMessage(msg: MessageRecord): ThreadMessageLike {
if (typeof msg.content === "string") { if (typeof msg.content === "string") {
content = [{ type: "text", text: msg.content }]; content = [{ type: "text", text: msg.content }];
} else if (Array.isArray(msg.content)) { } else if (Array.isArray(msg.content)) {
// Filter out custom metadata parts - they're handled separately const convertedContent = msg.content
const filteredContent = msg.content.filter((part: unknown) => { .filter((part: unknown) => {
if (typeof part !== "object" || part === null || !("type" in part)) return true; if (typeof part !== "object" || part === null || !("type" in part)) return true;
const partType = (part as { type: string }).type; const partType = (part as { type: string }).type;
// Filter out metadata parts not directly renderable by assistant-ui return partType !== "mentioned-documents" && partType !== "attachments";
return ( })
partType !== "thinking-steps" && .map((part: unknown) => {
partType !== "mentioned-documents" && if (
partType !== "attachments" 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 = content =
filteredContent.length > 0 convertedContent.length > 0
? (filteredContent as ThreadMessageLike["content"]) ? (convertedContent as ThreadMessageLike["content"])
: [{ type: "text", text: "" }]; : [{ type: "text", text: "" }];
} else { } else {
content = [{ type: "text", text: String(msg.content) }]; content = [{ type: "text", text: String(msg.content) }];
} }
// Build metadata.custom for author display in shared chats
const metadata = msg.author_id const metadata = msg.author_id
? { ? {
custom: { custom: {

43
surfsense_web/lib/chat/streaming-state.ts
View file

@ -15,6 +15,10 @@ export type ContentPart =
toolName: string; toolName: string;
args: Record<string, unknown>; args: Record<string, unknown>;
result?: unknown; result?: unknown;
}
| {
type: "data-thinking-steps";
data: { steps: ThinkingStepData[] };
}; };
export interface ContentPartsState { export interface ContentPartsState {
@ -23,6 +27,32 @@ export interface ContentPartsState {
toolCallIndices: Map<string, number>; toolCallIndices: Map<string, number>;
} }
export function updateThinkingSteps(
state: ContentPartsState,
steps: Map<string, ThinkingStepData>
): 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 { export function appendText(state: ContentPartsState, delta: string): void {
if ( if (
state.currentTextPartIndex >= 0 && state.currentTextPartIndex >= 0 &&
@ -75,6 +105,7 @@ export function buildContentForUI(
const filtered = state.contentParts.filter((part) => { const filtered = state.contentParts.filter((part) => {
if (part.type === "text") return part.text.length > 0; if (part.type === "text") return part.text.length > 0;
if (part.type === "tool-call") return toolsWithUI.has(part.toolName); if (part.type === "tool-call") return toolsWithUI.has(part.toolName);
if (part.type === "data-thinking-steps") return true;
return false; return false;
}); });
return filtered.length > 0 return filtered.length > 0
@ -84,23 +115,17 @@ export function buildContentForUI(
export function buildContentForPersistence( export function buildContentForPersistence(
state: ContentPartsState, state: ContentPartsState,
toolsWithUI: Set<string>, toolsWithUI: Set<string>
currentThinkingSteps: Map<string, ThinkingStepData>
): unknown[] { ): unknown[] {
const parts: unknown[] = []; const parts: unknown[] = [];
if (currentThinkingSteps.size > 0) {
parts.push({
type: "thinking-steps",
steps: Array.from(currentThinkingSteps.values()),
});
}
for (const part of state.contentParts) { for (const part of state.contentParts) {
if (part.type === "text" && part.text.length > 0) { if (part.type === "text" && part.text.length > 0) {
parts.push(part); parts.push(part);
} else if (part.type === "tool-call" && toolsWithUI.has(part.toolName)) { } else if (part.type === "tool-call" && toolsWithUI.has(part.toolName)) {
parts.push(part); parts.push(part);
} else if (part.type === "data-thinking-steps") {
parts.push(part);
} }
} }