mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-06-06 20:15:17 +02:00
feat(chat): add cached thread prefetching for faster navigation
This commit is contained in:
Anish Sarkar
parent
8b704b2fef
commit
168c0d2f89
9 changed files with 356 additions and 139 deletions
|
|
@ -51,6 +51,7 @@ import {
|
|||
TokenUsageProvider,
|
||||
} from "@/components/assistant-ui/token-usage-context";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { Skeleton } from "@/components/ui/skeleton";
|
||||
import {
|
||||
type HitlDecision,
|
||||
PendingInterruptProvider,
|
||||
|
|
@ -65,6 +66,7 @@ import {
|
|||
} from "@/hooks/use-agent-actions-query";
|
||||
import { useChatSessionStateSync } from "@/hooks/use-chat-session-state";
|
||||
import { useMessagesSync } from "@/hooks/use-messages-sync";
|
||||
import { useThreadDetail, useThreadMessages } from "@/hooks/use-thread-queries";
|
||||
import { getAgentFilesystemSelection } from "@/lib/agent-filesystem";
|
||||
import { documentsApiService } from "@/lib/apis/documents-api.service";
|
||||
import { getBearerToken } from "@/lib/auth-utils";
|
||||
|
|
@ -101,8 +103,6 @@ import {
|
|||
appendMessage,
|
||||
createThread,
|
||||
getRegenerateUrl,
|
||||
getThreadFull,
|
||||
getThreadMessages,
|
||||
type ThreadListItem,
|
||||
type ThreadListResponse,
|
||||
type ThreadRecord,
|
||||
|
|
@ -120,7 +120,7 @@ import {
|
|||
trackChatMessageSent,
|
||||
trackChatResponseReceived,
|
||||
} from "@/lib/posthog/events";
|
||||
import Loading from "../loading";
|
||||
import { cacheKeys } from "@/lib/query-client/cache-keys";
|
||||
|
||||
const MobileEditorPanel = dynamic(
|
||||
() =>
|
||||
|
|
@ -288,11 +288,78 @@ function computeFallbackTurnCancellingRetryDelay(attempt: number): number {
|
|||
return Math.min(raw, TURN_CANCELLING_MAX_DELAY_MS);
|
||||
}
|
||||
|
||||
function parseUrlChatId(id: string | string[] | undefined): number {
|
||||
let parsed = 0;
|
||||
if (Array.isArray(id) && id.length > 0) {
|
||||
parsed = Number.parseInt(id[0], 10);
|
||||
} else if (typeof id === "string") {
|
||||
parsed = Number.parseInt(id, 10);
|
||||
}
|
||||
return Number.isNaN(parsed) ? 0 : parsed;
|
||||
}
|
||||
|
||||
function ThreadMessagesSkeleton() {
|
||||
return (
|
||||
<div
|
||||
className="aui-root aui-thread-root @container flex h-full min-h-0 flex-col bg-panel"
|
||||
style={{
|
||||
["--thread-max-width" as string]: "42rem",
|
||||
}}
|
||||
>
|
||||
<div
|
||||
className="aui-thread-viewport relative flex flex-1 min-h-0 flex-col overflow-y-auto px-4 scroll-smooth"
|
||||
style={{ scrollbarGutter: "stable" }}
|
||||
>
|
||||
<div
|
||||
aria-hidden
|
||||
className="aui-chat-viewport-top-fade pointer-events-none sticky top-0 z-10 -mx-4 h-2 shrink-0 bg-gradient-to-b from-panel from-20% to-transparent"
|
||||
/>
|
||||
<div className="mx-auto w-full max-w-(--thread-max-width) flex flex-1 flex-col gap-6 py-8">
|
||||
<div className="flex justify-end">
|
||||
<Skeleton className="h-12 w-[65%] max-w-56 rounded-2xl" />
|
||||
</div>
|
||||
|
||||
<div className="flex flex-col gap-2">
|
||||
<Skeleton className="h-4 w-full" />
|
||||
<Skeleton className="h-4 w-[85%]" />
|
||||
<Skeleton className="h-18 w-[40%]" />
|
||||
</div>
|
||||
|
||||
<div className="flex gap-2 justify-end">
|
||||
<Skeleton className="h-12 w-[78%] max-w-72 rounded-2xl" />
|
||||
</div>
|
||||
|
||||
<div className="flex flex-col gap-2">
|
||||
<Skeleton className="h-10 w-[30%]" />
|
||||
<Skeleton className="h-4 w-[90%]" />
|
||||
<Skeleton className="h-6 w-[60%]" />
|
||||
</div>
|
||||
|
||||
<div className="flex gap-2 justify-end">
|
||||
<Skeleton className="h-12 w-[85%] max-w-96 rounded-2xl" />
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div
|
||||
className="aui-chat-composer-footer sticky bottom-0 z-20 -mx-4 mt-auto flex flex-col items-stretch bg-gradient-to-t from-panel from-60% to-transparent px-4 pt-6"
|
||||
style={{ paddingBottom: "max(0.5rem, env(safe-area-inset-bottom))" }}
|
||||
>
|
||||
<div className="aui-chat-composer-area relative mx-auto flex w-full max-w-(--thread-max-width) flex-col gap-3 overflow-visible">
|
||||
<Skeleton className="h-28 w-full rounded-3xl" />
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
export default function NewChatPage() {
|
||||
const params = useParams();
|
||||
const queryClient = useQueryClient();
|
||||
const [isInitializing, setIsInitializing] = useState(true);
|
||||
const [threadId, setThreadId] = useState<number | null>(null);
|
||||
const urlChatId = useMemo(() => parseUrlChatId(params.chat_id), [params.chat_id]);
|
||||
const [threadId, setThreadId] = useState<number | null>(() => (urlChatId > 0 ? urlChatId : null));
|
||||
const activeThreadId = urlChatId > 0 ? urlChatId : threadId;
|
||||
const handledLoadErrorThreadRef = useRef<number | null>(null);
|
||||
const [currentThread, setCurrentThread] = useState<ThreadRecord | null>(null);
|
||||
const [messages, setMessages] = useState<ThreadMessageLike[]>([]);
|
||||
const [isRunning, setIsRunning] = useState(false);
|
||||
|
|
@ -404,9 +471,11 @@ export default function NewChatPage() {
|
|||
const { data: currentUser } = useAtomValue(currentUserAtom);
|
||||
const { data: agentFlags } = useAtomValue(agentFlagsAtom);
|
||||
const localFilesystemEnabled = agentFlags?.enable_desktop_local_filesystem === true;
|
||||
const threadDetailQuery = useThreadDetail(activeThreadId);
|
||||
const threadMessagesQuery = useThreadMessages(activeThreadId);
|
||||
|
||||
// Live collaboration: sync session state and messages via Zero
|
||||
useChatSessionStateSync(threadId);
|
||||
useChatSessionStateSync(activeThreadId);
|
||||
const { data: membersData } = useAtomValue(membersAtom);
|
||||
|
||||
const handleSyncedMessagesUpdate = useCallback(
|
||||
|
|
@ -467,7 +536,7 @@ export default function NewChatPage() {
|
|||
[isRunning, membersData]
|
||||
);
|
||||
|
||||
useMessagesSync(threadId, handleSyncedMessagesUpdate);
|
||||
useMessagesSync(activeThreadId, handleSyncedMessagesUpdate);
|
||||
|
||||
// Extract search_space_id from URL params
|
||||
const searchSpaceId = useMemo(() => {
|
||||
|
|
@ -481,19 +550,7 @@ export default function NewChatPage() {
|
|||
// per-turn Revert button all read). Hydrates from
|
||||
// ``GET /threads/{id}/actions`` and is updated incrementally by the
|
||||
// SSE handlers + revert-batch results below — no atom side-channel.
|
||||
const { items: agentActionItems } = useAgentActionsQuery(threadId);
|
||||
|
||||
// Extract chat_id from URL params
|
||||
const urlChatId = useMemo(() => {
|
||||
const id = params.chat_id;
|
||||
let parsed = 0;
|
||||
if (Array.isArray(id) && id.length > 0) {
|
||||
parsed = Number.parseInt(id[0], 10);
|
||||
} else if (typeof id === "string") {
|
||||
parsed = Number.parseInt(id, 10);
|
||||
}
|
||||
return Number.isNaN(parsed) ? 0 : parsed;
|
||||
}, [params.chat_id]);
|
||||
const { items: agentActionItems } = useAgentActionsQuery(activeThreadId);
|
||||
|
||||
const handleChatFailure = useCallback(
|
||||
async ({
|
||||
|
|
@ -632,14 +689,19 @@ export default function NewChatPage() {
|
|||
});
|
||||
}, []);
|
||||
|
||||
// Initialize thread and load messages
|
||||
// For new chats (no urlChatId), we use lazy creation - thread is created on first message
|
||||
const initializeThread = useCallback(async () => {
|
||||
setIsInitializing(true);
|
||||
const hydratedMessagesRef = useRef<{
|
||||
threadId: number | null;
|
||||
data: typeof threadMessagesQuery.data;
|
||||
}>({ threadId: null, data: undefined });
|
||||
|
||||
// Reset all state when switching between chats/search spaces to prevent stale data
|
||||
// Reset thread-local runtime state on route/search-space changes. Data fetching
|
||||
// is handled by React Query below so the chat shell can render immediately.
|
||||
useEffect(() => {
|
||||
const nextThreadId = urlChatId > 0 ? urlChatId : null;
|
||||
handledLoadErrorThreadRef.current = null;
|
||||
hydratedMessagesRef.current = { threadId: null, data: undefined };
|
||||
setThreadId(nextThreadId);
|
||||
setMessages([]);
|
||||
setThreadId(null);
|
||||
setCurrentThread(null);
|
||||
setMentionedDocuments([]);
|
||||
tokenUsageStore.clear();
|
||||
|
|
@ -649,82 +711,96 @@ export default function NewChatPage() {
|
|||
closeEditorPanel();
|
||||
// Note: agent-action data is keyed by threadId in react-query so
|
||||
// switching threads naturally swaps caches; no explicit reset.
|
||||
|
||||
try {
|
||||
if (urlChatId > 0) {
|
||||
// Thread exists - load thread data and messages
|
||||
setThreadId(urlChatId);
|
||||
|
||||
// Load thread data (for visibility info) and messages in parallel
|
||||
const [threadData, messagesResponse] = await Promise.all([
|
||||
getThreadFull(urlChatId),
|
||||
getThreadMessages(urlChatId),
|
||||
]);
|
||||
|
||||
setCurrentThread(threadData);
|
||||
|
||||
if (messagesResponse.messages && messagesResponse.messages.length > 0) {
|
||||
const loadedMessages = reconcileInterruptedAssistantMessages(
|
||||
messagesResponse.messages
|
||||
).map(convertToThreadMessage);
|
||||
setMessages(loadedMessages);
|
||||
|
||||
for (const msg of messagesResponse.messages) {
|
||||
if (msg.token_usage) {
|
||||
tokenUsageStore.set(`msg-${msg.id}`, msg.token_usage as TokenUsageData);
|
||||
}
|
||||
}
|
||||
|
||||
const restoredDocsMap: Record<string, MentionedDocumentInfo[]> = {};
|
||||
for (const msg of messagesResponse.messages) {
|
||||
if (msg.role === "user") {
|
||||
const docs = extractMentionedDocuments(msg.content);
|
||||
if (docs.length > 0) {
|
||||
restoredDocsMap[`msg-${msg.id}`] = docs;
|
||||
}
|
||||
}
|
||||
}
|
||||
if (Object.keys(restoredDocsMap).length > 0) {
|
||||
setMessageDocumentsMap(restoredDocsMap);
|
||||
}
|
||||
}
|
||||
}
|
||||
// For new chats (urlChatId === 0), don't create thread yet
|
||||
// Thread will be created lazily when user sends first message
|
||||
// This improves UX (instant load) and avoids orphan threads
|
||||
} catch (error) {
|
||||
console.error("[NewChatPage] Failed to initialize thread:", error);
|
||||
if (urlChatId > 0 && error instanceof NotFoundError) {
|
||||
removeChatTab(urlChatId);
|
||||
if (typeof window !== "undefined") {
|
||||
window.history.replaceState(null, "", `/dashboard/${searchSpaceId}/new-chat`);
|
||||
}
|
||||
toast.error("This chat was deleted.");
|
||||
return;
|
||||
}
|
||||
// Keep threadId as null - don't use Date.now() as it creates an invalid ID
|
||||
// that will cause 404 errors on subsequent API calls
|
||||
setThreadId(null);
|
||||
setCurrentThread(null);
|
||||
toast.error("Failed to load chat. Please try again.");
|
||||
} finally {
|
||||
setIsInitializing(false);
|
||||
}
|
||||
}, [
|
||||
urlChatId,
|
||||
setMessageDocumentsMap,
|
||||
setMentionedDocuments,
|
||||
setMessageDocumentsMap,
|
||||
tokenUsageStore,
|
||||
closeReportPanel,
|
||||
closeEditorPanel,
|
||||
removeChatTab,
|
||||
searchSpaceId,
|
||||
]);
|
||||
|
||||
useEffect(() => {
|
||||
if (!activeThreadId) {
|
||||
setCurrentThread(null);
|
||||
return;
|
||||
}
|
||||
if (threadDetailQuery.data?.id === activeThreadId) {
|
||||
setCurrentThread(threadDetailQuery.data);
|
||||
}
|
||||
}, [activeThreadId, threadDetailQuery.data]);
|
||||
|
||||
useEffect(() => {
|
||||
const messagesResponse = threadMessagesQuery.data;
|
||||
if (!activeThreadId || !messagesResponse) return;
|
||||
|
||||
if (
|
||||
hydratedMessagesRef.current.threadId === activeThreadId &&
|
||||
hydratedMessagesRef.current.data === messagesResponse
|
||||
) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (isRunning) {
|
||||
return;
|
||||
}
|
||||
|
||||
const loadedMessages = reconcileInterruptedAssistantMessages(messagesResponse.messages).map(
|
||||
convertToThreadMessage
|
||||
);
|
||||
setMessages(loadedMessages);
|
||||
|
||||
tokenUsageStore.clear();
|
||||
const restoredDocsMap: Record<string, MentionedDocumentInfo[]> = {};
|
||||
for (const msg of messagesResponse.messages) {
|
||||
if (msg.token_usage) {
|
||||
tokenUsageStore.set(`msg-${msg.id}`, msg.token_usage as TokenUsageData);
|
||||
}
|
||||
if (msg.role === "user") {
|
||||
const docs = extractMentionedDocuments(msg.content);
|
||||
if (docs.length > 0) {
|
||||
restoredDocsMap[`msg-${msg.id}`] = docs;
|
||||
}
|
||||
}
|
||||
}
|
||||
setMessageDocumentsMap(restoredDocsMap);
|
||||
hydratedMessagesRef.current = { threadId: activeThreadId, data: messagesResponse };
|
||||
}, [
|
||||
activeThreadId,
|
||||
isRunning,
|
||||
setMessageDocumentsMap,
|
||||
threadMessagesQuery.data,
|
||||
tokenUsageStore,
|
||||
]);
|
||||
|
||||
// Initialize on mount, and re-init when switching search spaces (even if urlChatId is the same)
|
||||
useEffect(() => {
|
||||
initializeThread();
|
||||
}, [initializeThread]);
|
||||
const loadError = threadDetailQuery.error ?? threadMessagesQuery.error;
|
||||
if (!activeThreadId || !loadError) return;
|
||||
if (handledLoadErrorThreadRef.current === activeThreadId) return;
|
||||
|
||||
handledLoadErrorThreadRef.current = activeThreadId;
|
||||
console.error("[NewChatPage] Failed to load thread:", loadError);
|
||||
|
||||
if (loadError instanceof NotFoundError) {
|
||||
removeChatTab(activeThreadId);
|
||||
if (typeof window !== "undefined") {
|
||||
window.history.replaceState(null, "", `/dashboard/${searchSpaceId}/new-chat`);
|
||||
}
|
||||
setThreadId(null);
|
||||
setCurrentThread(null);
|
||||
setMessages([]);
|
||||
toast.error("This chat was deleted.");
|
||||
return;
|
||||
}
|
||||
|
||||
toast.error("Failed to load chat. Please try again.");
|
||||
}, [
|
||||
activeThreadId,
|
||||
removeChatTab,
|
||||
searchSpaceId,
|
||||
threadDetailQuery.error,
|
||||
threadMessagesQuery.error,
|
||||
]);
|
||||
|
||||
// Prefetch document titles for @ mention picker
|
||||
// Runs when user lands on page so data is ready when they type @
|
||||
|
|
@ -752,7 +828,7 @@ export default function NewChatPage() {
|
|||
const readAndApplyCommentId = () => {
|
||||
const params = new URLSearchParams(window.location.search);
|
||||
const raw = params.get("commentId");
|
||||
if (raw && !isInitializing) {
|
||||
if (raw && activeThreadId) {
|
||||
const commentId = Number.parseInt(raw, 10);
|
||||
if (!Number.isNaN(commentId)) {
|
||||
setTargetCommentId(commentId);
|
||||
|
|
@ -770,11 +846,14 @@ export default function NewChatPage() {
|
|||
window.removeEventListener("popstate", readAndApplyCommentId);
|
||||
clearTargetCommentId();
|
||||
};
|
||||
}, [isInitializing, setTargetCommentId, clearTargetCommentId]);
|
||||
}, [activeThreadId, setTargetCommentId, clearTargetCommentId]);
|
||||
|
||||
// Sync current thread state to atom
|
||||
useEffect(() => {
|
||||
if (!currentThread) {
|
||||
if (activeThreadId) {
|
||||
return;
|
||||
}
|
||||
setCurrentThreadMetadata({
|
||||
id: null,
|
||||
visibility: null,
|
||||
|
|
@ -794,6 +873,7 @@ export default function NewChatPage() {
|
|||
hasComments: currentThread.has_comments ?? false,
|
||||
});
|
||||
}, [
|
||||
activeThreadId,
|
||||
currentThread,
|
||||
currentThreadState.id,
|
||||
currentThreadState.visibility,
|
||||
|
|
@ -882,6 +962,8 @@ export default function NewChatPage() {
|
|||
setThreadId(currentThreadId);
|
||||
// Set currentThread so share button in header appears immediately
|
||||
setCurrentThread(newThread);
|
||||
queryClient.setQueryData(cacheKeys.threads.detail(newThread.id), newThread);
|
||||
queryClient.setQueryData(cacheKeys.threads.messages(newThread.id), { messages: [] });
|
||||
|
||||
// Track chat creation
|
||||
trackChatCreated(searchSpaceId, currentThreadId);
|
||||
|
|
@ -1389,6 +1471,14 @@ export default function NewChatPage() {
|
|||
} finally {
|
||||
setIsRunning(false);
|
||||
abortControllerRef.current = null;
|
||||
if (currentThreadId) {
|
||||
void queryClient.invalidateQueries({
|
||||
queryKey: cacheKeys.threads.messages(currentThreadId),
|
||||
});
|
||||
void queryClient.invalidateQueries({
|
||||
queryKey: cacheKeys.threads.detail(currentThreadId),
|
||||
});
|
||||
}
|
||||
}
|
||||
},
|
||||
[
|
||||
|
|
@ -1737,6 +1827,12 @@ export default function NewChatPage() {
|
|||
} finally {
|
||||
setIsRunning(false);
|
||||
abortControllerRef.current = null;
|
||||
void queryClient.invalidateQueries({
|
||||
queryKey: cacheKeys.threads.messages(resumeThreadId),
|
||||
});
|
||||
void queryClient.invalidateQueries({
|
||||
queryKey: cacheKeys.threads.detail(resumeThreadId),
|
||||
});
|
||||
}
|
||||
},
|
||||
[
|
||||
|
|
@ -2230,6 +2326,12 @@ export default function NewChatPage() {
|
|||
} finally {
|
||||
setIsRunning(false);
|
||||
abortControllerRef.current = null;
|
||||
void queryClient.invalidateQueries({
|
||||
queryKey: cacheKeys.threads.messages(threadId),
|
||||
});
|
||||
void queryClient.invalidateQueries({
|
||||
queryKey: cacheKeys.threads.detail(threadId),
|
||||
});
|
||||
}
|
||||
},
|
||||
[
|
||||
|
|
@ -2416,22 +2518,25 @@ export default function NewChatPage() {
|
|||
onCancel: cancelRun,
|
||||
});
|
||||
|
||||
// Show loading state only when loading an existing thread
|
||||
if (isInitializing) {
|
||||
return <Loading />;
|
||||
}
|
||||
const threadLoadError = activeThreadId
|
||||
? (threadDetailQuery.error ?? threadMessagesQuery.error)
|
||||
: null;
|
||||
const shouldShowThreadLoadError =
|
||||
!!threadLoadError && !!activeThreadId && !currentThread && messages.length === 0;
|
||||
const isThreadMessagesLoading =
|
||||
!!activeThreadId &&
|
||||
threadMessagesQuery.isPending &&
|
||||
messages.length === 0 &&
|
||||
!threadMessagesQuery.error;
|
||||
|
||||
// Show error state only if we tried to load an existing thread but failed
|
||||
// For new chats (urlChatId === 0), threadId being null is expected (lazy creation)
|
||||
if (!threadId && urlChatId > 0) {
|
||||
if (shouldShowThreadLoadError) {
|
||||
return (
|
||||
<div className="flex h-full flex-col items-center justify-center gap-4">
|
||||
<div className="text-destructive">Failed to load chat</div>
|
||||
<Button
|
||||
type="button"
|
||||
onClick={() => {
|
||||
setIsInitializing(true);
|
||||
initializeThread();
|
||||
void Promise.all([threadDetailQuery.refetch(), threadMessagesQuery.refetch()]);
|
||||
}}
|
||||
>
|
||||
Try Again
|
||||
|
|
@ -2450,8 +2555,13 @@ export default function NewChatPage() {
|
|||
onSubmit={handleApprovalSubmit}
|
||||
>
|
||||
<div key={searchSpaceId} className="flex h-full overflow-hidden">
|
||||
<div className="flex-1 flex flex-col min-w-0 overflow-hidden">
|
||||
<div className="relative flex-1 flex flex-col min-w-0 overflow-hidden">
|
||||
<Thread />
|
||||
{isThreadMessagesLoading ? (
|
||||
<div className="absolute inset-0 z-10 bg-panel">
|
||||
<ThreadMessagesSkeleton />
|
||||
</div>
|
||||
) : null}
|
||||
</div>
|
||||
<MobileReportPanel />
|
||||
<MobileEditorPanel />
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue