diff --git a/surfsense_backend/app/services/new_streaming_service.py b/surfsense_backend/app/services/new_streaming_service.py index 52a215997..3e24c1376 100644 --- a/surfsense_backend/app/services/new_streaming_service.py +++ b/surfsense_backend/app/services/new_streaming_service.py @@ -565,20 +565,24 @@ class VercelStreamingService: # Error Part # ========================================================================= - def format_error(self, error_text: str) -> str: + def format_error(self, error_text: str, error_code: str | None = None) -> str: """ Format an error message. Args: error_text: The error message text + error_code: Optional machine-readable error code for frontend branching Returns: str: SSE formatted error part Example output: - data: {"type":"error","errorText":"Something went wrong"} + data: {"type":"error","errorText":"Something went wrong","errorCode":"SOME_CODE"} """ - return self._format_sse({"type": "error", "errorText": error_text}) + payload: dict[str, str] = {"type": "error", "errorText": error_text} + if error_code: + payload["errorCode"] = error_code + return self._format_sse(payload) # ========================================================================= # Tool Parts diff --git a/surfsense_backend/app/tasks/chat/stream_new_chat.py b/surfsense_backend/app/tasks/chat/stream_new_chat.py index edc5aa763..060dd23c6 100644 --- a/surfsense_backend/app/tasks/chat/stream_new_chat.py +++ b/surfsense_backend/app/tasks/chat/stream_new_chat.py @@ -1581,7 +1581,8 @@ async def stream_new_chat( ) else: yield streaming_service.format_error( - "Premium tokens exhausted. Buy more tokens to continue with this model, or switch to a free model." + "Buy more tokens to continue with this model, or switch to a free model.", + error_code="PREMIUM_QUOTA_EXHAUSTED", ) yield streaming_service.format_done() return @@ -2348,7 +2349,8 @@ async def stream_resume_chat( ) else: yield streaming_service.format_error( - "Premium tokens exhausted. Buy more tokens to continue with this model, or switch to a free model." + "Buy more tokens to continue with this model, or switch to a free model.", + error_code="PREMIUM_QUOTA_EXHAUSTED", ) yield streaming_service.format_done() return 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 ed0611ee9..f775e1f06 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 @@ -206,10 +206,15 @@ const PREMIUM_QUOTA_ASSISTANT_MESSAGE = function getPinnedPremiumQuotaErrorMessage(error: unknown): string | null { if (!(error instanceof Error)) return null; + const withCode = error as Error & { errorCode?: string }; + if (withCode.errorCode === "PREMIUM_QUOTA_EXHAUSTED") { + return error.message; + } const normalized = error.message.toLowerCase(); if ( !normalized.includes("premium tokens exhausted") && !normalized.includes("premium token quota exceeded") + && !normalized.includes("buy more tokens") ) { return null; } @@ -233,6 +238,50 @@ export default function NewChatPage() { } | null>(null); const toolsWithUI = useMemo(() => new Set([...BASE_TOOLS_WITH_UI]), []); + const persistAssistantErrorMessage = useCallback( + async ({ + threadId, + assistantMsgId, + text, + }: { + threadId: number | null; + assistantMsgId: string; + text: string; + }) => { + setMessages((prev) => + prev.map((m) => + m.id === assistantMsgId + ? { + ...m, + content: [{ type: "text", text }], + } + : m + ) + ); + + if (!threadId) return; + + // Persist only temporary assistant placeholders to avoid duplicate rows + // when the message already has a database-backed ID. + if (!assistantMsgId.startsWith("msg-assistant-")) return; + + try { + const savedMessage = await appendMessage(threadId, { + role: "assistant", + content: [{ type: "text", text }], + }); + const newMsgId = `msg-${savedMessage.id}`; + tokenUsageStore.rename(assistantMsgId, newMsgId); + setMessages((prev) => + prev.map((m) => (m.id === assistantMsgId ? { ...m, id: newMsgId } : m)) + ); + } catch (persistErr) { + console.error("Failed to persist assistant error message:", persistErr); + } + }, + [tokenUsageStore] + ); + // Get disabled tools from the tool toggle UI const disabledTools = useAtomValue(disabledToolsAtom); @@ -903,7 +952,9 @@ export default function NewChatPage() { break; case "error": - throw new Error(parsed.errorText || "Server error"); + throw Object.assign(new Error(parsed.errorText || "Server error"), { + errorCode: parsed.errorCode, + }); } } @@ -985,26 +1036,14 @@ export default function NewChatPage() { } else { toast.error("Failed to get response. Please try again."); } - // Update assistant message with error - setMessages((prev) => - prev.map((m) => - m.id === assistantMsgId - ? { - ...m, - content: [ - { - type: "text", - text: - (premiumQuotaAlertMessage - ? PREMIUM_QUOTA_ASSISTANT_MESSAGE - : undefined) ?? - "Sorry, there was an error. Please try again.", - }, - ], - } - : m - ) - ); + await persistAssistantErrorMessage({ + threadId: currentThreadId, + assistantMsgId, + text: + (premiumQuotaAlertMessage + ? PREMIUM_QUOTA_ASSISTANT_MESSAGE + : undefined) ?? "Sorry, there was an error. Please try again.", + }); } finally { setIsRunning(false); abortControllerRef.current = null; @@ -1028,6 +1067,7 @@ export default function NewChatPage() { setPendingUserImageUrls, toolsWithUI, setPremiumAlertForThread, + persistAssistantErrorMessage, ] ); @@ -1258,7 +1298,9 @@ export default function NewChatPage() { break; case "error": - throw new Error(parsed.errorText || "Server error"); + throw Object.assign(new Error(parsed.errorText || "Server error"), { + errorCode: parsed.errorCode, + }); } } @@ -1293,19 +1335,17 @@ export default function NewChatPage() { threadId: resumeThreadId, message: premiumQuotaAlertMessage, }); - setMessages((prev) => - prev.map((m) => - m.id === assistantMsgId - ? { - ...m, - content: [{ type: "text", text: PREMIUM_QUOTA_ASSISTANT_MESSAGE }], - } - : m - ) - ); } else { toast.error("Failed to resume. Please try again."); } + await persistAssistantErrorMessage({ + threadId: resumeThreadId, + assistantMsgId, + text: + (premiumQuotaAlertMessage + ? PREMIUM_QUOTA_ASSISTANT_MESSAGE + : undefined) ?? "Sorry, there was an error. Please try again.", + }); } finally { setIsRunning(false); abortControllerRef.current = null; @@ -1318,6 +1358,7 @@ export default function NewChatPage() { tokenUsageStore, toolsWithUI, setPremiumAlertForThread, + persistAssistantErrorMessage, ] ); @@ -1589,7 +1630,9 @@ export default function NewChatPage() { break; case "error": - throw new Error(parsed.errorText || "Server error"); + throw Object.assign(new Error(parsed.errorText || "Server error"), { + errorCode: parsed.errorCode, + }); } } @@ -1653,25 +1696,14 @@ export default function NewChatPage() { } else { toast.error("Failed to regenerate response. Please try again."); } - setMessages((prev) => - prev.map((m) => - m.id === assistantMsgId - ? { - ...m, - content: [ - { - type: "text", - text: - (premiumQuotaAlertMessage - ? PREMIUM_QUOTA_ASSISTANT_MESSAGE - : undefined) ?? - "Sorry, there was an error. Please try again.", - }, - ], - } - : m - ) - ); + await persistAssistantErrorMessage({ + threadId, + assistantMsgId, + text: + (premiumQuotaAlertMessage + ? PREMIUM_QUOTA_ASSISTANT_MESSAGE + : undefined) ?? "Sorry, there was an error. Please try again.", + }); } finally { setIsRunning(false); abortControllerRef.current = null; @@ -1685,6 +1717,7 @@ export default function NewChatPage() { tokenUsageStore, toolsWithUI, setPremiumAlertForThread, + persistAssistantErrorMessage, ] ); diff --git a/surfsense_web/lib/chat/streaming-state.ts b/surfsense_web/lib/chat/streaming-state.ts index ff8fdfbd4..9f2ac87a5 100644 --- a/surfsense_web/lib/chat/streaming-state.ts +++ b/surfsense_web/lib/chat/streaming-state.ts @@ -256,7 +256,7 @@ export type SSEEvent = }>; }; } - | { type: "error"; errorText: string }; + | { type: "error"; errorText: string; errorCode?: string }; /** * Async generator that reads an SSE stream and yields parsed JSON objects.