feat(chat): enhance error formatting to include optional error codes for better frontend handling

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

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

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
+				await persistAssistantErrorMessage({
-				setMessages((prev) =>
+					threadId: currentThreadId,
-					prev.map((m) =>
+					assistantMsgId,
-						m.id === assistantMsgId
-							? {
-									...m,
-									content: [
-										{
-											type: "text",
 					text:
 						(premiumQuotaAlertMessage
 							? PREMIUM_QUOTA_ASSISTANT_MESSAGE
-													: undefined) ??
+							: undefined) ?? "Sorry, there was an error. Please try again.",
-												"Sorry, there was an error. Please try again.",
+				});
-										},
-									],
-								}
-							: m
-					)
-				);
 			} 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) =>
+				await persistAssistantErrorMessage({
-					prev.map((m) =>
+					threadId,
-						m.id === assistantMsgId
+					assistantMsgId,
-							? {
-									...m,
-									content: [
-										{
-											type: "text",
 					text:
 						(premiumQuotaAlertMessage
 							? PREMIUM_QUOTA_ASSISTANT_MESSAGE
-													: undefined) ??
+							: undefined) ?? "Sorry, there was an error. Please try again.",
-												"Sorry, there was an error. Please try again.",
+				});
-										},
-									],
-								}
-							: m
-					)
-				);
 			} finally {
 				setIsRunning(false);
 				abortControllerRef.current = null;
@@ -1685,6 +1717,7 @@ export default function NewChatPage() {
 			tokenUsageStore,
 			toolsWithUI,
 			setPremiumAlertForThread,
+			persistAssistantErrorMessage,
 		]
 	);
 

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.