refactor(onboarding, model-connections): enhance onboarding logic and streamline model connection handling by integrating chat model checks and improving state management

surfsense_web/app/dashboard/[search_space_id]/client-layout.tsx

@@ -4,13 +4,12 @@ import { useAtomValue, useSetAtom } from "jotai";
 import { useParams, usePathname, useRouter } from "next/navigation";
 import { useTranslations } from "next-intl";
 import type React from "react";
-import { useCallback, useEffect, useMemo, useRef, useState } from "react";
-import { toast } from "sonner";
+import { useEffect, useState } from "react";
 import { pendingUserImageDataUrlsAtom } from "@/atoms/chat/pending-user-images.atom";
 import { myAccessAtom } from "@/atoms/members/members-query.atoms";
-import { updateModelRolesMutationAtom } from "@/atoms/model-connections/model-connections-mutation.atoms";
 import {
 	globalModelConnectionsAtom,
+	modelConnectionsAtom,
 	modelRolesAtom,
 } from "@/atoms/model-connections/model-connections-query.atoms";
 import { activeSearchSpaceIdAtom } from "@/atoms/search-spaces/search-space-query.atoms";
@@ -21,6 +20,7 @@ import { Card, CardContent, CardDescription, CardHeader, CardTitle } from "@/com
 import { useFolderSync } from "@/hooks/use-folder-sync";
 import { useGlobalLoadingEffect } from "@/hooks/use-global-loading";
 import { useElectronAPI } from "@/hooks/use-platform";
+import { isLlmOnboardingComplete } from "@/lib/onboarding";
 
 export function DashboardClientLayout({
 	children,
@@ -33,39 +33,32 @@ export function DashboardClientLayout({
 	const router = useRouter();
 	const pathname = usePathname();
 	const { search_space_id } = useParams();
+	const activeSearchSpaceId = useAtomValue(activeSearchSpaceIdAtom);
 	const setActiveSearchSpaceIdState = useSetAtom(activeSearchSpaceIdAtom);
 	const setPendingUserImageUrls = useSetAtom(pendingUserImageDataUrlsAtom);
 
 	const {
 		data: modelRoles = {},
-		isFetching: loading,
+		isLoading: loading,
 		error,
-		refetch: refetchModelRoles,
 	} = useAtomValue(modelRolesAtom);
-	const { data: globalConnections = [], isFetching: globalConfigsLoading } = useAtomValue(
+	const { data: globalConnections = [], isLoading: globalConfigsLoading } = useAtomValue(
 		globalModelConnectionsAtom
 	);
-	const { mutateAsync: updateModelRoles } = useAtomValue(updateModelRolesMutationAtom);
-
-	const firstGlobalChatModel = useMemo(() => {
-		for (const connection of globalConnections) {
-			const model = connection.models.find((item) => item.enabled && item.supports_chat);
-			if (model) return model;
-		}
-		return null;
-	}, [globalConnections]);
-
-	const isOnboardingComplete = useCallback(() => {
-		return (modelRoles.chat_model_id ?? 0) !== 0 || Boolean(firstGlobalChatModel);
-	}, [modelRoles.chat_model_id, firstGlobalChatModel]);
+	const { data: modelConnections = [], isLoading: modelConnectionsLoading } =
+		useAtomValue(modelConnectionsAtom);
 
 	const { data: access = null, isLoading: accessLoading } = useAtomValue(myAccessAtom);
 	const [hasCheckedOnboarding, setHasCheckedOnboarding] = useState(false);
-	const [isAutoConfiguring, setIsAutoConfiguring] = useState(false);
-	const hasAttemptedAutoConfig = useRef(false);
 
 	const isOnboardingPage = pathname?.includes("/onboard");
 	const isOwner = access?.is_owner ?? false;
+	const isSearchSpaceReady = activeSearchSpaceId === searchSpaceId;
+
+	useEffect(() => {
+		if (isSearchSpaceReady) return;
+		setHasCheckedOnboarding(false);
+	}, [isSearchSpaceReady]);
 
 	useEffect(() => {
 		if (isOnboardingPage) {
@@ -74,13 +67,18 @@ export function DashboardClientLayout({
 		}
 
 		if (
+			isSearchSpaceReady &&
 			!loading &&
 			!accessLoading &&
 			!globalConfigsLoading &&
-			!hasCheckedOnboarding &&
-			!isAutoConfiguring
+			!modelConnectionsLoading &&
+			!hasCheckedOnboarding
 		) {
-			const onboardingComplete = isOnboardingComplete();
+			const onboardingComplete = isLlmOnboardingComplete(
+				modelRoles.chat_model_id,
+				globalConnections,
+				modelConnections
+			);
 
 			if (onboardingComplete) {
 				setHasCheckedOnboarding(true);
@@ -92,50 +90,23 @@ export function DashboardClientLayout({
 				return;
 			}
 
-			if (firstGlobalChatModel && !hasAttemptedAutoConfig.current) {
-				hasAttemptedAutoConfig.current = true;
-				setIsAutoConfiguring(true);
-
-				const autoConfigureWithGlobal = async () => {
-					try {
-						await updateModelRoles({ chat_model_id: firstGlobalChatModel.id });
-
-						await refetchModelRoles();
-
-						toast.success("AI configured automatically!", {
-							description: `Using ${firstGlobalChatModel.display_name || firstGlobalChatModel.model_id}. Customize in Settings.`,
-						});
-
-						setHasCheckedOnboarding(true);
-					} catch (error) {
-						console.error("Auto-configuration failed:", error);
-						router.push(`/dashboard/${searchSpaceId}/onboard`);
-					} finally {
-						setIsAutoConfiguring(false);
-					}
-				};
-
-				autoConfigureWithGlobal();
-				return;
-			}
-
 			router.push(`/dashboard/${searchSpaceId}/onboard`);
 			setHasCheckedOnboarding(true);
 		}
 	}, [
+		isSearchSpaceReady,
 		loading,
 		accessLoading,
 		globalConfigsLoading,
-		isOnboardingComplete,
+		modelConnectionsLoading,
+		modelRoles.chat_model_id,
+		globalConnections,
+		modelConnections,
 		isOnboardingPage,
 		isOwner,
-		isAutoConfiguring,
-		firstGlobalChatModel,
 		router,
 		searchSpaceId,
 		hasCheckedOnboarding,
-		updateModelRoles,
-		refetchModelRoles,
 	]);
 
 	const electronAPI = useElectronAPI();
@@ -188,9 +159,12 @@ export function DashboardClientLayout({
 	// Determine if we should show loading
 	const shouldShowLoading =
 		(!hasCheckedOnboarding &&
-			(loading || accessLoading || globalConfigsLoading) &&
-			!isOnboardingPage) ||
-		isAutoConfiguring;
+			(!isSearchSpaceReady ||
+				loading ||
+				accessLoading ||
+				globalConfigsLoading ||
+				modelConnectionsLoading) &&
+			!isOnboardingPage);
 
 	// Use global loading screen - spinner animation won't reset
 	useGlobalLoadingEffect(shouldShowLoading);

surfsense_web/app/dashboard/[search_space_id]/onboard/page.tsx

@@ -2,9 +2,7 @@
 
 import { useAtomValue } from "jotai";
 import { useParams, useRouter } from "next/navigation";
-import { useEffect, useMemo, useRef, useState } from "react";
-import { toast } from "sonner";
-import { updateModelRolesMutationAtom } from "@/atoms/model-connections/model-connections-mutation.atoms";
+import { useEffect, useMemo } from "react";
 import {
 	globalModelConnectionsAtom,
 	modelConnectionsAtom,
@@ -12,85 +10,40 @@ import {
 } from "@/atoms/model-connections/model-connections-query.atoms";
 import { Logo } from "@/components/Logo";
 import { ModelProviderConnectionsPanel } from "@/components/settings/model-connections/model-provider-connections-panel";
-import { capability } from "@/components/settings/model-connections/model-utils";
 import { Button } from "@/components/ui/button";
 import { useGlobalLoadingEffect } from "@/hooks/use-global-loading";
 import { getBearerToken, redirectToLogin } from "@/lib/auth-utils";
+import { hasEnabledChatModel, isLlmOnboardingComplete } from "@/lib/onboarding";
 
 export default function OnboardPage() {
 	const router = useRouter();
 	const params = useParams();
 	const searchSpaceId = Number(params.search_space_id);
-	const { data: globalConnections = [], isFetching: globalLoading } = useAtomValue(
+	const { data: globalConnections = [], isLoading: globalLoading } = useAtomValue(
 		globalModelConnectionsAtom
 	);
-	const { data: connections = [], isFetching: connectionsLoading } =
-		useAtomValue(modelConnectionsAtom);
-	const { data: roles = {}, isFetching: rolesLoading } = useAtomValue(modelRolesAtom);
-	const { mutateAsync: updateRoles, isPending } = useAtomValue(updateModelRolesMutationAtom);
-	const [isAutoConfiguring, setIsAutoConfiguring] = useState(false);
-	const hasAttemptedAutoConfig = useRef(false);
+	const { data: connections = [] } = useAtomValue(modelConnectionsAtom);
+	const { data: roles = {}, isLoading: rolesLoading } = useAtomValue(modelRolesAtom);
 
 	useEffect(() => {
 		if (!getBearerToken()) redirectToLogin();
 	}, []);
 
-	const firstGlobalChatModel = useMemo(() => {
-		for (const connection of globalConnections) {
-			const model = connection.models.find((item) => item.enabled && item.supports_chat);
-			if (model) return model;
-		}
-		return null;
-	}, [globalConnections]);
-	const hasEnabledChatModel = useMemo(
-		() =>
-			connections.some(
-				(connection) =>
-					connection.enabled &&
-					connection.models.some((model) => model.enabled && capability(model, "chat"))
-			),
-		[connections]
+	const hasUsableChatModel = useMemo(
+		() => hasEnabledChatModel([...globalConnections, ...connections]),
+		[globalConnections, connections]
 	);
 
-	const isComplete = (roles.chat_model_id ?? 0) !== 0 || Boolean(firstGlobalChatModel);
-
-	useEffect(() => {
-		if (globalLoading || rolesLoading || hasAttemptedAutoConfig.current) return;
-		if ((roles.chat_model_id ?? 0) !== 0) {
-			router.push(`/dashboard/${searchSpaceId}/new-chat`);
-			return;
-		}
-		if (!firstGlobalChatModel) return;
-
-		hasAttemptedAutoConfig.current = true;
-		setIsAutoConfiguring(true);
-		updateRoles({ chat_model_id: firstGlobalChatModel.id })
-			.then(() => {
-				toast.success("AI configured automatically", {
-					description: `Using ${firstGlobalChatModel.display_name || firstGlobalChatModel.model_id}.`,
-				});
-				router.push(`/dashboard/${searchSpaceId}/new-chat`);
-			})
-			.catch((error) => {
-				console.error("Auto-configuration failed:", error);
-				toast.error("Auto-configuration failed. Add a connection manually.");
-				setIsAutoConfiguring(false);
-			});
-	}, [
-		firstGlobalChatModel,
-		globalLoading,
+	const onboardingComplete = isLlmOnboardingComplete(
 		roles.chat_model_id,
-		rolesLoading,
-		router,
-		searchSpaceId,
-		updateRoles,
-	]);
+		globalConnections,
+		connections
+	);
 
-	const isLoading =
-		globalLoading || connectionsLoading || rolesLoading || isAutoConfiguring || isPending;
+	const isLoading = globalLoading || rolesLoading;
 	useGlobalLoadingEffect(isLoading);
 
-	if (isLoading || isComplete) return null;
+	if (isLoading) return null;
 
 	return (
 		<div className="flex min-h-screen select-none flex-col items-center justify-center bg-main-panel p-4">
@@ -109,7 +62,7 @@ export default function OnboardPage() {
 					footerAction={
 						<Button
 							className="min-w-[112px]"
-							disabled={!hasEnabledChatModel}
+							disabled={!onboardingComplete || !hasUsableChatModel}
 							onClick={() => router.push(`/dashboard/${searchSpaceId}/new-chat`)}
 						>
 							Start

surfsense_web/atoms/model-connections/model-connections-mutation.atoms.ts

@@ -2,6 +2,7 @@ import { atomWithMutation } from "jotai-tanstack-query";
 import { toast } from "sonner";
 import type {
 	ConnectionCreateRequest,
+	ConnectionRead,
 	ConnectionUpdateRequest,
 	ModelCreateRequest,
 	ModelPreviewRead,
@@ -26,15 +27,33 @@ function invalidateModelConnections(searchSpaceId: number) {
 	});
 }
 
+function upsertModelConnection(searchSpaceId: number, connection: ConnectionRead) {
+	queryClient.setQueryData<ConnectionRead[]>(
+		cacheKeys.modelConnections.all(searchSpaceId),
+		(current = []) => {
+			if (current.some((item) => item.id === connection.id)) {
+				return current.map((item) => (item.id === connection.id ? connection : item));
+			}
+			return [...current, connection];
+		}
+	);
+}
+
 export const createModelConnectionMutationAtom = atomWithMutation((get) => {
 	const searchSpaceId = Number(get(activeSearchSpaceIdAtom));
 	return {
 		mutationKey: ["model-connections", "create"],
 		mutationFn: (request: ConnectionCreateRequest) =>
 			modelConnectionsApiService.createConnection(request),
-		onSuccess: () => {
+		onSuccess: (connection: ConnectionRead, request: ConnectionCreateRequest) => {
+			const resolvedSearchSpaceId = Number(
+				request.search_space_id ?? connection.search_space_id ?? searchSpaceId
+			);
 			toast.success("Connection created");
-			invalidateModelConnections(searchSpaceId);
+			if (resolvedSearchSpaceId > 0) {
+				upsertModelConnection(resolvedSearchSpaceId, connection);
+				invalidateModelConnections(resolvedSearchSpaceId);
+			}
 		},
 		onError: (error: Error) => toast.error(error.message || "Failed to create connection"),
 	};

surfsense_web/lib/onboarding.ts

@@ -1,8 +1,29 @@
-export function isLlmOnboardingComplete(
-	agentLlmId: number | null | undefined,
-	hasGlobalConfigs: boolean
-): boolean {
-	if (agentLlmId === null || agentLlmId === undefined) return false;
-	if (agentLlmId === 0) return hasGlobalConfigs;
-	return true;
+import type { ConnectionRead } from "@/contracts/types/model-connections.types";
+
+export function hasEnabledChatModel(connections: ConnectionRead[]): boolean {
+	return connections.some(
+		(connection) =>
+			connection.enabled &&
+			connection.models.some((model) => model.enabled && Boolean(model.supports_chat))
+	);
+}
+
+export function isLlmOnboardingComplete(
+	chatModelId: number | null | undefined,
+	globalConnections: ConnectionRead[],
+	searchSpaceConnections: ConnectionRead[]
+): boolean {
+	const connections = [...globalConnections, ...searchSpaceConnections];
+	const resolvedChatModelId = chatModelId ?? 0;
+
+	if (resolvedChatModelId === 0) {
+		return hasEnabledChatModel(connections);
+	}
+
+	return connections.some((connection) =>
+		connection.models.some(
+			(model) =>
+				model.id === resolvedChatModelId && model.enabled && Boolean(model.supports_chat)
+		)
+	);
 }