feat: integrate token usage context and enhance message info display in chat UI

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<ThreadRecord | null>(null);
 	const [messages, setMessages] = useState<ThreadMessageLike[]>([]);
 	const [isRunning, setIsRunning] = useState(false);
+	const [tokenUsageStore] = useState(() => createTokenUsageStore());
 	const abortControllerRef = useRef<AbortController | null>(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<string, MentionedDocumentInfo[]> = {};
 					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<string, unknown> ?? {}), 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<string, unknown> ?? {}), 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<string, unknown> ?? {}), 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 (
+		<TokenUsageProvider store={tokenUsageStore}>
 		<AssistantRuntimeProvider runtime={runtime}>
 			<ThinkingStepsDataUI />
 			<div key={searchSpaceId} className="flex h-full overflow-hidden">
@@ -1627,5 +1614,6 @@ export default function NewChatPage() {
 				<MobileHitlEditPanel />
 			</div>
 		</AssistantRuntimeProvider>
+		</TokenUsageProvider>
 	);
 }

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<string, unknown> | undefined;
-		return custom?.usage as Record<string, unknown> | 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<string, { prompt_tokens: number; completion_tokens: number; total_tokens: number }>
-		| 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 (
 		<DropdownMenu>
@@ -401,24 +405,31 @@ const TokenUsageDropdown: FC = () => {
 				</Button>
 			</DropdownMenuTrigger>
 			<DropdownMenuContent align="start" className="min-w-[180px]">
-				<DropdownMenuLabel className="text-xs text-muted-foreground font-normal">
-					Token Usage
-				</DropdownMenuLabel>
-				{models.length > 0 ? (
-					models.map(([model, counts]) => (
-						<DropdownMenuItem key={model} className="flex-col items-start gap-0.5 cursor-default" onSelect={(e) => e.preventDefault()}>
-							<span className="text-xs font-medium">{model}</span>
-							<span className="text-xs text-muted-foreground">
-								{counts.total_tokens.toLocaleString()} tokens
-							</span>
-						</DropdownMenuItem>
-					))
-				) : (
-					<DropdownMenuItem className="flex-col items-start gap-0.5 cursor-default" onSelect={(e) => e.preventDefault()}>
-						<span className="text-xs text-muted-foreground">
-							{totalTokens.toLocaleString()} tokens
-						</span>
-					</DropdownMenuItem>
+				{createdAt && (
+					<DropdownMenuLabel className="text-xs text-muted-foreground font-normal select-none">
+						{formatMessageDate(createdAt)}
+					</DropdownMenuLabel>
+				)}
+				{hasUsage && (
+					<>
+						<DropdownMenuSeparator />
+						{models.length > 0 ? (
+							models.map(([model, counts]) => (
+								<DropdownMenuItem key={model} className="flex-col items-start gap-0.5 cursor-default" onSelect={(e) => e.preventDefault()}>
+									<span className="text-xs font-medium">{model}</span>
+									<span className="text-xs text-muted-foreground">
+										{counts.total_tokens.toLocaleString()} tokens
+									</span>
+								</DropdownMenuItem>
+							))
+						) : (
+							<DropdownMenuItem className="flex-col items-start gap-0.5 cursor-default" onSelect={(e) => e.preventDefault()}>
+								<span className="text-xs text-muted-foreground">
+									{usage.total_tokens.toLocaleString()} tokens
+								</span>
+							</DropdownMenuItem>
+						)}
+					</>
 				)}
 			</DropdownMenuContent>
 		</DropdownMenu>
@@ -683,7 +694,7 @@ const AssistantActionBar: FC = () => {
 					<ClipboardPaste />
 				</TooltipIconButton>
 			)}
-			<TokenUsageDropdown />
+			<MessageInfoDropdown />
 		</ActionBarPrimitive.Root>
 	);
 };

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<string, { prompt_tokens: number; completion_tokens: number; total_tokens: number }>;
+	model_breakdown?: Record<string, { prompt_tokens: number; completion_tokens: number; total_tokens: number }>;
+}
+
+type Listener = () => void;
+
+class TokenUsageStore {
+	private data = new Map<string, TokenUsageData>();
+	private listeners = new Set<Listener>();
+
+	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<TokenUsageStore | null>(null);
+
+export const TokenUsageProvider: FC<{ store: TokenUsageStore; children: ReactNode }> = ({ store, children }) => (
+	<TokenUsageContext.Provider value={store}>{children}</TokenUsageContext.Provider>
+);
+
+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();
+}