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 34bf0c09e..ff953eaf9 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 @@ -42,6 +42,7 @@ import { useChatSessionStateSync } from "@/hooks/use-chat-session-state"; import { useMessagesSync } from "@/hooks/use-messages-sync"; import { documentsApiService } from "@/lib/apis/documents-api.service"; import { getBearerToken } from "@/lib/auth-utils"; +import { createTokenUsageStore, TokenUsageProvider, type TokenUsageData } from "@/components/assistant-ui/token-usage-context"; import { convertToThreadMessage } from "@/lib/chat/message-utils"; import { isPodcastGenerating, @@ -195,6 +196,7 @@ export default function NewChatPage() { const [currentThread, setCurrentThread] = useState(null); const [messages, setMessages] = useState([]); const [isRunning, setIsRunning] = useState(false); + const [tokenUsageStore] = useState(() => createTokenUsageStore()); const abortControllerRef = useRef(null); const [pendingInterrupt, setPendingInterrupt] = useState<{ threadId: number; @@ -307,6 +309,7 @@ export default function NewChatPage() { setThreadId(null); setCurrentThread(null); setMentionedDocuments([]); + tokenUsageStore.clear(); setSidebarDocuments([]); setMessageDocumentsMap({}); clearPlanOwnerRegistry(); @@ -330,6 +333,12 @@ export default function NewChatPage() { const loadedMessages = 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 = {}; for (const msg of messagesResponse.messages) { if (msg.role === "user") { @@ -374,6 +383,7 @@ export default function NewChatPage() { closeEditorPanel, removeChatTab, searchSpaceId, + tokenUsageStore, ]); // Initialize on mount, and re-init when switching search spaces (even if urlChatId is the same) @@ -824,6 +834,7 @@ export default function NewChatPage() { case "data-token-usage": tokenUsageData = parsed.data; + tokenUsageStore.set(assistantMsgId, parsed.data as TokenUsageData); break; case "error": @@ -833,16 +844,6 @@ export default function NewChatPage() { batcher.flush(); - if (tokenUsageData) { - setMessages((prev) => - prev.map((m) => - m.id === assistantMsgId - ? { ...m, metadata: { ...m.metadata, custom: { ...(m.metadata?.custom as Record ?? {}), usage: tokenUsageData } } } - : m - ) - ); - } - // Skip persistence for interrupted messages -- handleResume will persist the final version const finalContent = buildContentForPersistence(contentPartsState, TOOLS_WITH_UI); if (contentParts.length > 0 && !wasInterrupted) { @@ -855,8 +856,9 @@ export default function NewChatPage() { // Update message ID from temporary to database ID so comments work immediately const newMsgId = `msg-${savedMessage.id}`; + tokenUsageStore.rename(assistantMsgId, newMsgId); setMessages((prev) => - prev.map((m) => (m.id === assistantMsgId ? { ...m, id: newMsgId } : m)) + prev.map((m) => (m.id === assistantMsgId ? { ...m, id: newMsgId } : m)), ); // Update pending interrupt with the new persisted message ID @@ -946,6 +948,7 @@ export default function NewChatPage() { currentUser, disabledTools, updateChatTabTitle, + tokenUsageStore, ] ); @@ -1168,6 +1171,7 @@ export default function NewChatPage() { case "data-token-usage": tokenUsageData = parsed.data; + tokenUsageStore.set(assistantMsgId, parsed.data as TokenUsageData); break; case "error": @@ -1177,16 +1181,6 @@ export default function NewChatPage() { batcher.flush(); - if (tokenUsageData) { - setMessages((prev) => - prev.map((m) => - m.id === assistantMsgId - ? { ...m, metadata: { ...m.metadata, custom: { ...(m.metadata?.custom as Record ?? {}), usage: tokenUsageData } } } - : m - ) - ); - } - const finalContent = buildContentForPersistence(contentPartsState, TOOLS_WITH_UI); if (contentParts.length > 0) { try { @@ -1196,8 +1190,9 @@ export default function NewChatPage() { token_usage: tokenUsageData ?? undefined, }); const newMsgId = `msg-${savedMessage.id}`; + tokenUsageStore.rename(assistantMsgId, newMsgId); setMessages((prev) => - prev.map((m) => (m.id === assistantMsgId ? { ...m, id: newMsgId } : m)) + prev.map((m) => (m.id === assistantMsgId ? { ...m, id: newMsgId } : m)), ); } catch (err) { console.error("Failed to persist resumed assistant message:", err); @@ -1215,7 +1210,7 @@ export default function NewChatPage() { abortControllerRef.current = null; } }, - [pendingInterrupt, messages, searchSpaceId] + [pendingInterrupt, messages, searchSpaceId, tokenUsageStore] ); useEffect(() => { @@ -1463,6 +1458,7 @@ export default function NewChatPage() { case "data-token-usage": tokenUsageData = parsed.data; + tokenUsageStore.set(assistantMsgId, parsed.data as TokenUsageData); break; case "error": @@ -1472,16 +1468,6 @@ export default function NewChatPage() { batcher.flush(); - if (tokenUsageData) { - setMessages((prev) => - prev.map((m) => - m.id === assistantMsgId - ? { ...m, metadata: { ...m.metadata, custom: { ...(m.metadata?.custom as Record ?? {}), usage: tokenUsageData } } } - : m - ) - ); - } - // Persist messages after streaming completes const finalContent = buildContentForPersistence(contentPartsState, TOOLS_WITH_UI); if (contentParts.length > 0) { @@ -1509,10 +1495,10 @@ export default function NewChatPage() { token_usage: tokenUsageData ?? undefined, }); - // Update assistant message ID to database ID const newMsgId = `msg-${savedMessage.id}`; + tokenUsageStore.rename(assistantMsgId, newMsgId); setMessages((prev) => - prev.map((m) => (m.id === assistantMsgId ? { ...m, id: newMsgId } : m)) + prev.map((m) => (m.id === assistantMsgId ? { ...m, id: newMsgId } : m)), ); trackChatResponseReceived(searchSpaceId, threadId); @@ -1547,7 +1533,7 @@ export default function NewChatPage() { abortControllerRef.current = null; } }, - [threadId, searchSpaceId, messages, disabledTools] + [threadId, searchSpaceId, messages, disabledTools, tokenUsageStore] ); // Handle editing a message - truncates history and regenerates with new query @@ -1616,6 +1602,7 @@ export default function NewChatPage() { } return ( +
@@ -1627,5 +1614,6 @@ export default function NewChatPage() {
+ ); } diff --git a/surfsense_web/components/assistant-ui/assistant-message.tsx b/surfsense_web/components/assistant-ui/assistant-message.tsx index 25a579947..dff52c3f5 100644 --- a/surfsense_web/components/assistant-ui/assistant-message.tsx +++ b/surfsense_web/components/assistant-ui/assistant-message.tsx @@ -45,12 +45,14 @@ import { DropdownMenuContent, DropdownMenuItem, DropdownMenuLabel, + DropdownMenuSeparator, DropdownMenuTrigger, } from "@/components/ui/dropdown-menu"; import { Button } from "@/components/ui/button"; import { useComments } from "@/hooks/use-comments"; import { useMediaQuery } from "@/hooks/use-media-query"; import { useElectronAPI } from "@/hooks/use-platform"; +import { useTokenUsage } from "@/components/assistant-ui/token-usage-context"; import { cn } from "@/lib/utils"; // Captured once at module load — survives client-side navigations that strip the query param. @@ -375,22 +377,24 @@ export const MessageError: FC = () => { ); }; -const TokenUsageDropdown: FC = () => { - const usage = useAuiState(({ message }) => { - const custom = message?.metadata?.custom as Record | undefined; - return custom?.usage as Record | undefined; +function formatMessageDate(date: Date): string { + return date.toLocaleDateString(undefined, { + month: "short", + day: "numeric", + hour: "numeric", + minute: "2-digit", + hour12: true, }); +} - if (!usage) return null; - - const totalTokens = (usage.total_tokens as number) ?? 0; - if (totalTokens === 0) return null; - - const modelBreakdown = (usage.usage ?? usage.model_breakdown) as - | Record - | undefined; +const MessageInfoDropdown: FC = () => { + const messageId = useAuiState(({ message }) => message?.id); + const createdAt = useAuiState(({ message }) => message?.createdAt); + const usage = useTokenUsage(messageId); + const modelBreakdown = usage ? (usage.usage ?? usage.model_breakdown) : undefined; const models = modelBreakdown ? Object.entries(modelBreakdown) : []; + const hasUsage = usage && usage.total_tokens > 0; return ( @@ -401,24 +405,31 @@ const TokenUsageDropdown: FC = () => { - - Token Usage - - {models.length > 0 ? ( - models.map(([model, counts]) => ( - e.preventDefault()}> - {model} - - {counts.total_tokens.toLocaleString()} tokens - - - )) - ) : ( - e.preventDefault()}> - - {totalTokens.toLocaleString()} tokens - - + {createdAt && ( + + {formatMessageDate(createdAt)} + + )} + {hasUsage && ( + <> + + {models.length > 0 ? ( + models.map(([model, counts]) => ( + e.preventDefault()}> + {model} + + {counts.total_tokens.toLocaleString()} tokens + + + )) + ) : ( + e.preventDefault()}> + + {usage.total_tokens.toLocaleString()} tokens + + + )} + )} @@ -683,7 +694,7 @@ const AssistantActionBar: FC = () => { )} - + ); }; diff --git a/surfsense_web/components/assistant-ui/token-usage-context.tsx b/surfsense_web/components/assistant-ui/token-usage-context.tsx new file mode 100644 index 000000000..8b82f33ff --- /dev/null +++ b/surfsense_web/components/assistant-ui/token-usage-context.tsx @@ -0,0 +1,79 @@ +"use client"; + +import { createContext, useContext, useCallback, useSyncExternalStore, type FC, type ReactNode } from "react"; + +export interface TokenUsageData { + prompt_tokens: number; + completion_tokens: number; + total_tokens: number; + usage?: Record; + model_breakdown?: Record; +} + +type Listener = () => void; + +class TokenUsageStore { + private data = new Map(); + private listeners = new Set(); + + get(messageId: string): TokenUsageData | undefined { + return this.data.get(messageId); + } + + set(messageId: string, usage: TokenUsageData): void { + this.data.set(messageId, usage); + this.notify(); + } + + rename(oldId: string, newId: string): void { + const usage = this.data.get(oldId); + if (usage) { + this.data.delete(oldId); + this.data.set(newId, usage); + this.notify(); + } + } + + clear(): void { + this.data.clear(); + this.notify(); + } + + subscribe = (listener: Listener): (() => void) => { + this.listeners.add(listener); + return () => this.listeners.delete(listener); + }; + + private notify(): void { + for (const l of this.listeners) l(); + } +} + +const TokenUsageContext = createContext(null); + +export const TokenUsageProvider: FC<{ store: TokenUsageStore; children: ReactNode }> = ({ store, children }) => ( + {children} +); + +export function useTokenUsageStore(): TokenUsageStore { + const store = useContext(TokenUsageContext); + if (!store) throw new Error("useTokenUsageStore must be used within TokenUsageProvider"); + return store; +} + +export function useTokenUsage(messageId: string | undefined): TokenUsageData | undefined { + const store = useContext(TokenUsageContext); + const getSnapshot = useCallback( + () => (store && messageId ? store.get(messageId) : undefined), + [store, messageId], + ); + const subscribe = useCallback( + (onStoreChange: () => void) => (store ? store.subscribe(onStoreChange) : () => {}), + [store], + ); + return useSyncExternalStore(subscribe, getSnapshot, getSnapshot); +} + +export function createTokenUsageStore(): TokenUsageStore { + return new TokenUsageStore(); +}