From d5ef0d2598573578d3abf0140c58da6d4e63401d Mon Sep 17 00:00:00 2001 From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com> Date: Wed, 29 Apr 2026 19:15:46 +0530 Subject: [PATCH] feat(ui): surface pinned premium quota alerts in chat thread --- .../new-chat/[[...chat_id]]/page.tsx | 81 +++++++++++++++++-- .../atoms/chat/premium-alert.atom.ts | 33 ++++++++ .../components/assistant-ui/thread.tsx | 44 +++++++++- 3 files changed, 148 insertions(+), 10 deletions(-) create mode 100644 surfsense_web/atoms/chat/premium-alert.atom.ts 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 7773a438a..a5461e17f 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 @@ -19,6 +19,7 @@ import { currentThreadAtom, setTargetCommentIdAtom, } from "@/atoms/chat/current-thread.atom"; +import { setPremiumAlertForThreadAtom } from "@/atoms/chat/premium-alert.atom"; import { type MentionedDocumentInfo, mentionedDocumentIdsAtom, @@ -200,6 +201,19 @@ const BASE_TOOLS_WITH_UI = new Set([ // "write_todos", // Disabled for now ]); +const PINNED_PREMIUM_QUOTA_MESSAGE = "Premium token quota exceeded for this pinned model."; + +function getPinnedPremiumQuotaErrorMessage(error: unknown): string | null { + if (!(error instanceof Error)) return null; + if (!error.message.toLowerCase().includes("premium token quota exceeded")) { + return null; + } + if (!error.message.toLowerCase().includes("pinned model")) { + return null; + } + return error.message || PINNED_PREMIUM_QUOTA_MESSAGE; +} + export default function NewChatPage() { const params = useParams(); const queryClient = useQueryClient(); @@ -226,6 +240,7 @@ export default function NewChatPage() { const setMentionedDocuments = useSetAtom(mentionedDocumentsAtom); const setMessageDocumentsMap = useSetAtom(messageDocumentsMapAtom); const setCurrentThreadState = useSetAtom(currentThreadAtom); + const setPremiumAlertForThread = useSetAtom(setPremiumAlertForThreadAtom); const setTargetCommentId = useSetAtom(setTargetCommentIdAtom); const clearTargetCommentId = useSetAtom(clearTargetCommentIdAtom); const closeReportPanel = useSetAtom(closeReportPanelAtom); @@ -951,6 +966,7 @@ export default function NewChatPage() { return; } console.error("[NewChatPage] Chat error:", error); + const premiumQuotaAlertMessage = getPinnedPremiumQuotaErrorMessage(error); // Track chat error trackChatError( @@ -959,7 +975,15 @@ export default function NewChatPage() { error instanceof Error ? error.message : "Unknown error" ); - toast.error("Failed to get response. Please try again."); + if (premiumQuotaAlertMessage) { + setPremiumAlertForThread({ + threadId: currentThreadId, + message: premiumQuotaAlertMessage, + }); + toast.error(PINNED_PREMIUM_QUOTA_MESSAGE); + } else { + toast.error("Failed to get response. Please try again."); + } // Update assistant message with error setMessages((prev) => prev.map((m) => @@ -969,7 +993,9 @@ export default function NewChatPage() { content: [ { type: "text", - text: "Sorry, there was an error. Please try again.", + text: + premiumQuotaAlertMessage ?? + "Sorry, there was an error. Please try again.", }, ], } @@ -998,6 +1024,7 @@ export default function NewChatPage() { pendingUserImageUrls, setPendingUserImageUrls, toolsWithUI, + setPremiumAlertForThread, ] ); @@ -1257,13 +1284,29 @@ export default function NewChatPage() { return; } console.error("[NewChatPage] Resume error:", error); - toast.error("Failed to resume. Please try again."); + const premiumQuotaAlertMessage = getPinnedPremiumQuotaErrorMessage(error); + if (premiumQuotaAlertMessage) { + setPremiumAlertForThread({ + threadId: resumeThreadId, + message: premiumQuotaAlertMessage, + }); + toast.error(PINNED_PREMIUM_QUOTA_MESSAGE); + } else { + toast.error("Failed to resume. Please try again."); + } } finally { setIsRunning(false); abortControllerRef.current = null; } }, - [pendingInterrupt, messages, searchSpaceId, tokenUsageStore, toolsWithUI] + [ + pendingInterrupt, + messages, + searchSpaceId, + tokenUsageStore, + toolsWithUI, + setPremiumAlertForThread, + ] ); useEffect(() => { @@ -1584,18 +1627,34 @@ export default function NewChatPage() { } batcher.dispose(); console.error("[NewChatPage] Regeneration error:", error); + const premiumQuotaAlertMessage = getPinnedPremiumQuotaErrorMessage(error); trackChatError( searchSpaceId, threadId, error instanceof Error ? error.message : "Unknown error" ); - toast.error("Failed to regenerate response. Please try again."); + if (premiumQuotaAlertMessage) { + setPremiumAlertForThread({ + threadId, + message: premiumQuotaAlertMessage, + }); + toast.error(PINNED_PREMIUM_QUOTA_MESSAGE); + } else { + toast.error("Failed to regenerate response. Please try again."); + } setMessages((prev) => prev.map((m) => m.id === assistantMsgId ? { ...m, - content: [{ type: "text", text: "Sorry, there was an error. Please try again." }], + content: [ + { + type: "text", + text: + premiumQuotaAlertMessage ?? + "Sorry, there was an error. Please try again.", + }, + ], } : m ) @@ -1605,7 +1664,15 @@ export default function NewChatPage() { abortControllerRef.current = null; } }, - [threadId, searchSpaceId, messages, disabledTools, tokenUsageStore, toolsWithUI] + [ + threadId, + searchSpaceId, + messages, + disabledTools, + tokenUsageStore, + toolsWithUI, + setPremiumAlertForThread, + ] ); // Handle editing a message - truncates history and regenerates with new query diff --git a/surfsense_web/atoms/chat/premium-alert.atom.ts b/surfsense_web/atoms/chat/premium-alert.atom.ts new file mode 100644 index 000000000..c0efc174f --- /dev/null +++ b/surfsense_web/atoms/chat/premium-alert.atom.ts @@ -0,0 +1,33 @@ +import { atom } from "jotai"; + +export type PremiumAlertState = { + message: string; +}; + +export const premiumAlertByThreadAtom = atom>({}); + +export const setPremiumAlertForThreadAtom = atom( + null, + ( + get, + set, + payload: { + threadId: number; + message: string; + } + ) => { + const current = get(premiumAlertByThreadAtom); + set(premiumAlertByThreadAtom, { + ...current, + [payload.threadId]: { message: payload.message }, + }); + } +); + +export const clearPremiumAlertForThreadAtom = atom(null, (get, set, threadId: number) => { + const current = get(premiumAlertByThreadAtom); + if (!(threadId in current)) return; + const next = { ...current }; + delete next[threadId]; + set(premiumAlertByThreadAtom, next); +}); diff --git a/surfsense_web/components/assistant-ui/thread.tsx b/surfsense_web/components/assistant-ui/thread.tsx index cf99598f1..06f25f5fb 100644 --- a/surfsense_web/components/assistant-ui/thread.tsx +++ b/surfsense_web/components/assistant-ui/thread.tsx @@ -37,10 +37,13 @@ import { toggleToolAtom, } from "@/atoms/agent-tools/agent-tools.atoms"; import { chatSessionStateAtom } from "@/atoms/chat/chat-session-state.atom"; -import { - mentionedDocumentsAtom, -} from "@/atoms/chat/mentioned-documents.atom"; +import { currentThreadAtom } from "@/atoms/chat/current-thread.atom"; +import { mentionedDocumentsAtom } from "@/atoms/chat/mentioned-documents.atom"; import { pendingUserImageDataUrlsAtom } from "@/atoms/chat/pending-user-images.atom"; +import { + clearPremiumAlertForThreadAtom, + premiumAlertByThreadAtom, +} from "@/atoms/chat/premium-alert.atom"; import { connectorDialogOpenAtom } from "@/atoms/connector-dialog/connector-dialog.atoms"; import { connectorsAtom } from "@/atoms/connectors/connector-query.atoms"; import { membersAtom } from "@/atoms/members/members-query.atoms"; @@ -134,6 +137,9 @@ const ThreadContent: FC = () => { style={{ paddingBottom: "max(1rem, env(safe-area-inset-bottom))" }} > + !thread.isEmpty}> + + !thread.isEmpty}> @@ -143,6 +149,38 @@ const ThreadContent: FC = () => { ); }; +const PremiumQuotaPinnedAlert: FC = () => { + const currentThreadState = useAtomValue(currentThreadAtom); + const alertsByThread = useAtomValue(premiumAlertByThreadAtom); + const clearPremiumAlertForThread = useSetAtom(clearPremiumAlertForThreadAtom); + + const currentThreadId = currentThreadState?.id; + if (!currentThreadId) return null; + + const alert = alertsByThread[currentThreadId]; + if (!alert) return null; + + return ( +
+
+ +
+

Premium quota exhausted

+

{alert.message}

+
+ +
+
+ ); +}; + const ThreadScrollToBottom: FC = () => { return (