mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-06-18 21:15:16 +02:00
531 lines
16 KiB
TypeScript
531 lines
16 KiB
TypeScript
"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 { StepSeparatorDataUI } from "@/components/assistant-ui/step-separator";
|
||
import {
|
||
createTokenUsageStore,
|
||
type TokenUsageData,
|
||
TokenUsageProvider,
|
||
} from "@/components/assistant-ui/token-usage-context";
|
||
import { Alert, AlertDescription, AlertTitle } from "@/components/ui/alert";
|
||
import { useAnonymousMode } from "@/contexts/anonymous-mode";
|
||
import { TimelineDataUI } from "@/features/chat-messages/timeline";
|
||
import {
|
||
addStepSeparator,
|
||
addToolCall,
|
||
appendReasoning,
|
||
appendText,
|
||
appendToolInputDelta,
|
||
buildContentForUI,
|
||
type ContentPartsState,
|
||
endReasoning,
|
||
FrameBatchedUpdater,
|
||
readSSEStream,
|
||
type ThinkingStepData,
|
||
updateThinkingSteps,
|
||
updateToolCall,
|
||
} from "@/lib/chat/streaming-state";
|
||
import { buildBackendUrl } from "@/lib/env-config";
|
||
import { trackAnonymousChatMessageSent } from "@/lib/posthog/events";
|
||
import { FreeThread } from "./free-thread";
|
||
import { RemoveAdsBanner } from "./remove-ads-banner";
|
||
|
||
// Render all tool calls via ToolFallback; backend keeps persisted
|
||
// payloads bounded by summarising / truncating outputs.
|
||
const TOOLS_WITH_UI = "all" as const;
|
||
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;
|
||
}
|
||
|
||
function normalizeFreeChatErrorMessage(error: unknown): string {
|
||
if (!(error instanceof Error)) return "An unexpected error occurred";
|
||
const code = (error as Error & { errorCode?: string }).errorCode;
|
||
if (code === "THREAD_BUSY") {
|
||
return "A previous response is still stopping. Please try again in a moment.";
|
||
}
|
||
if (code === "MODEL_AUTH_FAILED") {
|
||
return "This model’s API key is invalid or expired. Switch models, or update the API key.";
|
||
}
|
||
if (code === "MODEL_NOT_FOUND") {
|
||
return "This model is unavailable or no longer exists. Please switch models.";
|
||
}
|
||
if (code === "MODEL_CONTEXT_LIMIT") {
|
||
return "This request is too large for the selected model. Reduce the input or switch models.";
|
||
}
|
||
if (code === "MODEL_PROVIDER_UNAVAILABLE") {
|
||
return "The selected model provider is temporarily unavailable. Please try again or switch models.";
|
||
}
|
||
if (code === "RATE_LIMITED") {
|
||
return "This model is temporarily rate-limited. Please try again in a few seconds or switch models.";
|
||
}
|
||
return error.message || "An unexpected error occurred";
|
||
}
|
||
|
||
function toFreeChatHttpError(status: number, body: string): Error & { errorCode?: string } {
|
||
let errorCode: string | undefined;
|
||
let message = body || `Server error: ${status}`;
|
||
try {
|
||
const parsed = JSON.parse(body) as Record<string, unknown>;
|
||
const detail =
|
||
typeof parsed.detail === "object" && parsed.detail !== null
|
||
? (parsed.detail as Record<string, unknown>)
|
||
: null;
|
||
errorCode =
|
||
(typeof detail?.error_code === "string" ? detail.error_code : undefined) ??
|
||
(typeof detail?.errorCode === "string" ? detail.errorCode : undefined) ??
|
||
(typeof parsed.error_code === "string" ? parsed.error_code : undefined) ??
|
||
(typeof parsed.errorCode === "string" ? parsed.errorCode : undefined);
|
||
message =
|
||
(typeof detail?.message === "string" ? detail.message : undefined) ??
|
||
(typeof parsed.message === "string" ? parsed.message : undefined) ??
|
||
(typeof parsed.detail === "string" ? parsed.detail : undefined) ??
|
||
message;
|
||
} catch {
|
||
// non-json response
|
||
}
|
||
|
||
if (!errorCode) {
|
||
if (status === 409) errorCode = "THREAD_BUSY";
|
||
else if (status === 429) errorCode = "RATE_LIMITED";
|
||
else if (status === 401 || status === 403) errorCode = "AUTH_EXPIRED";
|
||
else errorCode = "SERVER_ERROR";
|
||
}
|
||
|
||
return Object.assign(new Error(message), { errorCode });
|
||
}
|
||
|
||
export function FreeChatPage() {
|
||
const anonMode = useAnonymousMode();
|
||
const modelSlug = anonMode.isAnonymous ? anonMode.modelSlug : "";
|
||
const resetKey = anonMode.isAnonymous ? anonMode.resetKey : 0;
|
||
const webSearchEnabled = anonMode.isAnonymous ? anonMode.webSearchEnabled : true;
|
||
|
||
const [messages, setMessages] = useState<ThreadMessageLike[]>([]);
|
||
const [isRunning, setIsRunning] = useState(false);
|
||
const [tokenUsageStore] = useState(() => createTokenUsageStore());
|
||
const abortControllerRef = useRef<AbortController | null>(null);
|
||
// Mirror the latest messages into a ref so onNew stays a stable callback
|
||
// (it reads history on demand instead of depending on the array).
|
||
const messagesRef = useRef<ThreadMessageLike[]>([]);
|
||
messagesRef.current = messages;
|
||
|
||
// 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(async () => {
|
||
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" | undefined> => {
|
||
const reqBody: Record<string, unknown> = {
|
||
model_slug: modelSlug,
|
||
messages: messageHistory,
|
||
};
|
||
if (!webSearchEnabled) reqBody.disabled_tools = ["web_search"];
|
||
if (turnstileToken) reqBody.turnstile_token = turnstileToken;
|
||
|
||
const response = await fetch(buildBackendUrl("/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 toFreeChatHttpError(response.status, body);
|
||
}
|
||
|
||
const currentThinkingSteps = new Map<string, ThinkingStepData>();
|
||
const batcher = new FrameBatchedUpdater();
|
||
const contentPartsState: ContentPartsState = {
|
||
contentParts: [],
|
||
currentTextPartIndex: -1,
|
||
currentReasoningPartIndex: -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);
|
||
const forceFlush = () => {
|
||
scheduleFlush();
|
||
batcher.flush();
|
||
};
|
||
|
||
try {
|
||
for await (const parsed of readSSEStream(response)) {
|
||
switch (parsed.type) {
|
||
case "text-delta":
|
||
appendText(contentPartsState, parsed.delta);
|
||
scheduleFlush();
|
||
break;
|
||
|
||
case "reasoning-delta":
|
||
appendReasoning(contentPartsState, parsed.delta);
|
||
scheduleFlush();
|
||
break;
|
||
|
||
case "reasoning-end":
|
||
endReasoning(contentPartsState);
|
||
scheduleFlush();
|
||
break;
|
||
|
||
case "start-step":
|
||
addStepSeparator(contentPartsState);
|
||
scheduleFlush();
|
||
break;
|
||
|
||
case "finish-step":
|
||
break;
|
||
|
||
case "tool-input-start":
|
||
addToolCall(
|
||
contentPartsState,
|
||
TOOLS_WITH_UI,
|
||
parsed.toolCallId,
|
||
parsed.toolName,
|
||
{},
|
||
false,
|
||
parsed.langchainToolCallId,
|
||
parsed.metadata
|
||
);
|
||
forceFlush();
|
||
break;
|
||
|
||
case "tool-input-delta":
|
||
appendToolInputDelta(contentPartsState, parsed.toolCallId, parsed.inputTextDelta);
|
||
scheduleFlush();
|
||
break;
|
||
|
||
case "tool-input-available": {
|
||
const finalArgsText = JSON.stringify(parsed.input ?? {}, null, 2);
|
||
if (toolCallIndices.has(parsed.toolCallId)) {
|
||
updateToolCall(contentPartsState, parsed.toolCallId, {
|
||
args: parsed.input || {},
|
||
argsText: finalArgsText,
|
||
langchainToolCallId: parsed.langchainToolCallId,
|
||
metadata: parsed.metadata,
|
||
});
|
||
} else {
|
||
addToolCall(
|
||
contentPartsState,
|
||
TOOLS_WITH_UI,
|
||
parsed.toolCallId,
|
||
parsed.toolName,
|
||
parsed.input || {},
|
||
false,
|
||
parsed.langchainToolCallId,
|
||
parsed.metadata
|
||
);
|
||
updateToolCall(contentPartsState, parsed.toolCallId, {
|
||
argsText: finalArgsText,
|
||
});
|
||
}
|
||
forceFlush();
|
||
break;
|
||
}
|
||
|
||
case "tool-output-available":
|
||
updateToolCall(contentPartsState, parsed.toolCallId, {
|
||
result: parsed.output,
|
||
langchainToolCallId: parsed.langchainToolCallId,
|
||
metadata: parsed.metadata,
|
||
});
|
||
forceFlush();
|
||
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 Object.assign(new Error(parsed.errorText || "Server error"), {
|
||
errorCode: parsed.errorCode,
|
||
});
|
||
}
|
||
}
|
||
batcher.flush();
|
||
} catch (err) {
|
||
batcher.dispose();
|
||
throw err;
|
||
}
|
||
},
|
||
[modelSlug, tokenUsageStore, webSearchEnabled]
|
||
);
|
||
|
||
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;
|
||
|
||
trackAnonymousChatMessageSent({
|
||
modelSlug,
|
||
messageLength: userQuery.trim().length,
|
||
hasUploadedDoc: anonMode.isAnonymous && anonMode.uploadedDoc !== null ? true : false,
|
||
surface: "free_chat_page",
|
||
});
|
||
|
||
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 = messagesRef.current
|
||
.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 = normalizeFreeChatErrorMessage(error);
|
||
setMessages((prev) =>
|
||
prev.map((m) =>
|
||
m.id === assistantMsgId
|
||
? { ...m, content: [{ type: "text" as const, text: `Error: ${errorText}` }] }
|
||
: m
|
||
)
|
||
);
|
||
} finally {
|
||
setIsRunning(false);
|
||
abortControllerRef.current = null;
|
||
}
|
||
},
|
||
[modelSlug, anonMode, 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 = normalizeFreeChatErrorMessage(error);
|
||
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}>
|
||
<TimelineDataUI />
|
||
<StepSeparatorDataUI />
|
||
<div className="flex h-full flex-col overflow-hidden">
|
||
<RemoveAdsBanner />
|
||
|
||
{captchaRequired && TURNSTILE_SITE_KEY && (
|
||
<div className="flex justify-center border-b bg-muted/30 px-4 py-4">
|
||
<Alert className="w-auto max-w-md">
|
||
<ShieldCheck />
|
||
<AlertTitle>Quick verification to continue chatting</AlertTitle>
|
||
<AlertDescription>
|
||
<Turnstile
|
||
ref={turnstileRef}
|
||
siteKey={TURNSTILE_SITE_KEY}
|
||
onSuccess={handleTurnstileSuccess}
|
||
onError={() => turnstileRef.current?.reset()}
|
||
onExpire={() => turnstileRef.current?.reset()}
|
||
options={{ theme: "auto", size: "normal" }}
|
||
/>
|
||
</AlertDescription>
|
||
</Alert>
|
||
</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>
|
||
);
|
||
}
|