feat: no login experience and prem tokens

surfsense_web/components/free-chat/free-chat-page.tsx

@@ -0,0 +1,395 @@
+"use client";
+
+import {
+	type AppendMessage,
+	AssistantRuntimeProvider,
+	type ThreadMessageLike,
+	useExternalStoreRuntime,
+} from "@assistant-ui/react";
+import { Turnstile, type TurnstileInstance } from "@marsidev/react-turnstile";
+import { ShieldCheck } from "lucide-react";
+import { useCallback, useEffect, useRef, useState } from "react";
+import { ThinkingStepsDataUI } from "@/components/assistant-ui/thinking-steps";
+import {
+	createTokenUsageStore,
+	type TokenUsageData,
+	TokenUsageProvider,
+} from "@/components/assistant-ui/token-usage-context";
+import { useAnonymousMode } from "@/contexts/anonymous-mode";
+import {
+	addToolCall,
+	appendText,
+	buildContentForUI,
+	type ContentPartsState,
+	FrameBatchedUpdater,
+	readSSEStream,
+	type ThinkingStepData,
+	updateThinkingSteps,
+	updateToolCall,
+} from "@/lib/chat/streaming-state";
+import { BACKEND_URL } from "@/lib/env-config";
+import { FreeModelSelector } from "./free-model-selector";
+import { FreeThread } from "./free-thread";
+
+const TOOLS_WITH_UI = new Set(["web_search", "document_qna"]);
+const TURNSTILE_SITE_KEY = process.env.NEXT_PUBLIC_TURNSTILE_SITE_KEY ?? "";
+
+/** Try to parse a CAPTCHA_REQUIRED or CAPTCHA_INVALID code from a non-ok response. */
+function parseCaptchaError(status: number, body: string): string | null {
+	if (status !== 403) return null;
+	try {
+		const json = JSON.parse(body);
+		const code = json?.detail?.code ?? json?.error?.code;
+		if (code === "CAPTCHA_REQUIRED" || code === "CAPTCHA_INVALID") return code;
+	} catch {
+		/* not JSON */
+	}
+	return null;
+}
+
+export function FreeChatPage() {
+	const anonMode = useAnonymousMode();
+	const modelSlug = anonMode.isAnonymous ? anonMode.modelSlug : "";
+	const resetKey = anonMode.isAnonymous ? anonMode.resetKey : 0;
+
+	const [messages, setMessages] = useState<ThreadMessageLike[]>([]);
+	const [isRunning, setIsRunning] = useState(false);
+	const [tokenUsageStore] = useState(() => createTokenUsageStore());
+	const abortControllerRef = useRef<AbortController | null>(null);
+
+	// Turnstile CAPTCHA state
+	const [captchaRequired, setCaptchaRequired] = useState(false);
+	const turnstileRef = useRef<TurnstileInstance | null>(null);
+	const turnstileTokenRef = useRef<string | null>(null);
+	const pendingRetryRef = useRef<{
+		messageHistory: { role: string; content: string }[];
+		userMsgId: string;
+	} | null>(null);
+
+	useEffect(() => {
+		setMessages([]);
+		tokenUsageStore.clear();
+		if (abortControllerRef.current) {
+			abortControllerRef.current.abort();
+			abortControllerRef.current = null;
+		}
+		setIsRunning(false);
+		setCaptchaRequired(false);
+		turnstileTokenRef.current = null;
+		pendingRetryRef.current = null;
+	}, [resetKey, modelSlug, tokenUsageStore]);
+
+	const cancelRun = useCallback(() => {
+		if (abortControllerRef.current) {
+			abortControllerRef.current.abort();
+			abortControllerRef.current = null;
+		}
+		setIsRunning(false);
+	}, []);
+
+	/**
+	 * Core streaming logic shared by initial sends and CAPTCHA retries.
+	 * Returns "captcha" if the server demands a CAPTCHA, otherwise void.
+	 */
+	const doStream = useCallback(
+		async (
+			messageHistory: { role: string; content: string }[],
+			assistantMsgId: string,
+			signal: AbortSignal,
+			turnstileToken: string | null
+		): Promise<"captcha" | void> => {
+			const reqBody: Record<string, unknown> = {
+				model_slug: modelSlug,
+				messages: messageHistory,
+			};
+			if (turnstileToken) reqBody.turnstile_token = turnstileToken;
+
+			const response = await fetch(`${BACKEND_URL}/api/v1/public/anon-chat/stream`, {
+				method: "POST",
+				headers: { "Content-Type": "application/json" },
+				credentials: "include",
+				body: JSON.stringify(reqBody),
+				signal,
+			});
+
+			if (!response.ok) {
+				const body = await response.text().catch(() => "");
+				const captchaCode = parseCaptchaError(response.status, body);
+				if (captchaCode) return "captcha";
+				throw new Error(body || `Server error: ${response.status}`);
+			}
+
+			const currentThinkingSteps = new Map<string, ThinkingStepData>();
+			const batcher = new FrameBatchedUpdater();
+			const contentPartsState: ContentPartsState = {
+				contentParts: [],
+				currentTextPartIndex: -1,
+				toolCallIndices: new Map(),
+			};
+			const { toolCallIndices } = contentPartsState;
+
+			const flushMessages = () => {
+				setMessages((prev) =>
+					prev.map((m) =>
+						m.id === assistantMsgId
+							? { ...m, content: buildContentForUI(contentPartsState, TOOLS_WITH_UI) }
+							: m
+					)
+				);
+			};
+			const scheduleFlush = () => batcher.schedule(flushMessages);
+
+			try {
+				for await (const parsed of readSSEStream(response)) {
+					switch (parsed.type) {
+						case "text-delta":
+							appendText(contentPartsState, parsed.delta);
+							scheduleFlush();
+							break;
+
+						case "tool-input-start":
+							addToolCall(contentPartsState, TOOLS_WITH_UI, parsed.toolCallId, parsed.toolName, {});
+							batcher.flush();
+							break;
+
+						case "tool-input-available":
+							if (toolCallIndices.has(parsed.toolCallId)) {
+								updateToolCall(contentPartsState, parsed.toolCallId, { args: parsed.input || {} });
+							} else {
+								addToolCall(
+									contentPartsState,
+									TOOLS_WITH_UI,
+									parsed.toolCallId,
+									parsed.toolName,
+									parsed.input || {}
+								);
+							}
+							batcher.flush();
+							break;
+
+						case "tool-output-available":
+							updateToolCall(contentPartsState, parsed.toolCallId, { result: parsed.output });
+							batcher.flush();
+							break;
+
+						case "data-thinking-step": {
+							const stepData = parsed.data as ThinkingStepData;
+							if (stepData?.id) {
+								currentThinkingSteps.set(stepData.id, stepData);
+								if (updateThinkingSteps(contentPartsState, currentThinkingSteps)) scheduleFlush();
+							}
+							break;
+						}
+
+						case "data-token-usage":
+							tokenUsageStore.set(assistantMsgId, parsed.data as TokenUsageData);
+							break;
+
+						case "error":
+							throw new Error(parsed.errorText || "Server error");
+					}
+				}
+				batcher.flush();
+			} catch (err) {
+				batcher.dispose();
+				throw err;
+			}
+		},
+		[modelSlug, tokenUsageStore]
+	);
+
+	const onNew = useCallback(
+		async (message: AppendMessage) => {
+			let userQuery = "";
+			for (const part of message.content) {
+				if (part.type === "text") userQuery += part.text;
+			}
+			if (!userQuery.trim()) return;
+
+			const userMsgId = `msg-user-${Date.now()}`;
+			setMessages((prev) => [
+				...prev,
+				{
+					id: userMsgId,
+					role: "user" as const,
+					content: [{ type: "text" as const, text: userQuery }],
+					createdAt: new Date(),
+				},
+			]);
+
+			setIsRunning(true);
+			const controller = new AbortController();
+			abortControllerRef.current = controller;
+
+			const assistantMsgId = `msg-assistant-${Date.now()}`;
+			setMessages((prev) => [
+				...prev,
+				{
+					id: assistantMsgId,
+					role: "assistant" as const,
+					content: [{ type: "text" as const, text: "" }],
+					createdAt: new Date(),
+				},
+			]);
+
+			const messageHistory = messages
+				.filter((m) => m.role === "user" || m.role === "assistant")
+				.map((m) => {
+					let text = "";
+					for (const part of m.content) {
+						if (typeof part === "object" && part.type === "text" && "text" in part) {
+							text += (part as { type: "text"; text: string }).text;
+						}
+					}
+					return { role: m.role as string, content: text };
+				})
+				.filter((m) => m.content.length > 0);
+			messageHistory.push({ role: "user", content: userQuery.trim() });
+
+			try {
+				const result = await doStream(
+					messageHistory,
+					assistantMsgId,
+					controller.signal,
+					turnstileTokenRef.current
+				);
+
+				// Consume the token after use regardless of outcome
+				turnstileTokenRef.current = null;
+
+				if (result === "captcha" && TURNSTILE_SITE_KEY) {
+					// Remove the empty assistant placeholder; keep the user message
+					setMessages((prev) => prev.filter((m) => m.id !== assistantMsgId));
+					pendingRetryRef.current = { messageHistory, userMsgId };
+					setCaptchaRequired(true);
+					setIsRunning(false);
+					abortControllerRef.current = null;
+					return;
+				}
+			} catch (error) {
+				if (error instanceof Error && error.name === "AbortError") return;
+				console.error("[FreeChatPage] Chat error:", error);
+				const errorText = error instanceof Error ? error.message : "An unexpected error occurred";
+				setMessages((prev) =>
+					prev.map((m) =>
+						m.id === assistantMsgId
+							? { ...m, content: [{ type: "text" as const, text: `Error: ${errorText}` }] }
+							: m
+					)
+				);
+			} finally {
+				setIsRunning(false);
+				abortControllerRef.current = null;
+			}
+		},
+		[messages, doStream]
+	);
+
+	/** Called when Turnstile resolves successfully. Stores the token and auto-retries. */
+	const handleTurnstileSuccess = useCallback(
+		async (token: string) => {
+			turnstileTokenRef.current = token;
+			setCaptchaRequired(false);
+
+			const pending = pendingRetryRef.current;
+			if (!pending) return;
+			pendingRetryRef.current = null;
+
+			setIsRunning(true);
+			const controller = new AbortController();
+			abortControllerRef.current = controller;
+
+			const assistantMsgId = `msg-assistant-${Date.now()}`;
+			setMessages((prev) => [
+				...prev,
+				{
+					id: assistantMsgId,
+					role: "assistant" as const,
+					content: [{ type: "text" as const, text: "" }],
+					createdAt: new Date(),
+				},
+			]);
+
+			try {
+				const result = await doStream(
+					pending.messageHistory,
+					assistantMsgId,
+					controller.signal,
+					token
+				);
+				turnstileTokenRef.current = null;
+
+				if (result === "captcha") {
+					setMessages((prev) => prev.filter((m) => m.id !== assistantMsgId));
+					pendingRetryRef.current = pending;
+					setCaptchaRequired(true);
+					turnstileRef.current?.reset();
+				}
+			} catch (error) {
+				if (error instanceof Error && error.name === "AbortError") return;
+				console.error("[FreeChatPage] Retry error:", error);
+				const errorText = error instanceof Error ? error.message : "An unexpected error occurred";
+				setMessages((prev) =>
+					prev.map((m) =>
+						m.id === assistantMsgId
+							? { ...m, content: [{ type: "text" as const, text: `Error: ${errorText}` }] }
+							: m
+					)
+				);
+			} finally {
+				setIsRunning(false);
+				abortControllerRef.current = null;
+			}
+		},
+		[doStream]
+	);
+
+	const convertMessage = useCallback(
+		(message: ThreadMessageLike): ThreadMessageLike => message,
+		[]
+	);
+
+	const runtime = useExternalStoreRuntime({
+		messages,
+		isRunning,
+		onNew,
+		convertMessage,
+		onCancel: cancelRun,
+	});
+
+	return (
+		<TokenUsageProvider store={tokenUsageStore}>
+			<AssistantRuntimeProvider runtime={runtime}>
+				<ThinkingStepsDataUI />
+				<div className="flex h-full flex-col overflow-hidden">
+					<div className="flex h-14 shrink-0 items-center justify-between border-b border-border/40 px-4">
+						<FreeModelSelector />
+					</div>
+
+					{captchaRequired && TURNSTILE_SITE_KEY && (
+						<div className="flex flex-col items-center gap-3 border-b border-border/40 bg-muted/30 py-4">
+							<div className="flex items-center gap-2 text-sm text-muted-foreground">
+								<ShieldCheck className="h-4 w-4" />
+								<span>Quick verification to continue chatting</span>
+							</div>
+							<Turnstile
+								ref={turnstileRef}
+								siteKey={TURNSTILE_SITE_KEY}
+								onSuccess={handleTurnstileSuccess}
+								onError={() => turnstileRef.current?.reset()}
+								onExpire={() => turnstileRef.current?.reset()}
+								options={{ theme: "auto", size: "normal" }}
+							/>
+						</div>
+					)}
+
+					<div className="flex flex-1 min-h-0 overflow-hidden">
+						<div className="flex-1 flex flex-col min-w-0">
+							<FreeThread />
+						</div>
+					</div>
+				</div>
+			</AssistantRuntimeProvider>
+		</TokenUsageProvider>
+	);
+}