refactor: update LLM role management logic and enhance UI feedback

surfsense_web/atoms/new-llm-config/new-llm-config-mutation.atoms.ts

@@ -109,10 +109,11 @@ export const updateLLMPreferencesMutationAtom = atomWithMutation((get) => {
 		mutationFn: async (request: UpdateLLMPreferencesRequest) => {
 			return newLLMConfigApiService.updateLLMPreferences(request);
 		},
-		onSuccess: () => {
+		onSuccess: (_data, request: UpdateLLMPreferencesRequest) => {
-			queryClient.invalidateQueries({
+			queryClient.setQueryData(
-				queryKey: cacheKeys.newLLMConfigs.preferences(Number(searchSpaceId)),
+				cacheKeys.newLLMConfigs.preferences(Number(searchSpaceId)),
-			});
+				(old: Record<string, unknown> | undefined) => ({ ...old, ...request.data })
+			);
 		},
 		onError: (error: Error) => {
 			toast.error(error.message || "Failed to update LLM preferences");

surfsense_web/components/settings/llm-role-manager.tsx

@@ -4,16 +4,14 @@ import { useAtomValue } from "jotai";
 import {
 	AlertCircle,
 	Bot,
-	CheckCircle,
+	CircleCheck,
 	CircleDashed,
 	FileText,
 	ImageIcon,
 	RefreshCw,
-	RotateCcw,
-	Save,
 	Shuffle,
 } from "lucide-react";
-import { useEffect, useState } from "react";
+import { useCallback, useEffect, useRef, useState } from "react";
 import { toast } from "sonner";
 import {
 	globalImageGenConfigsAtom,
@@ -40,6 +38,7 @@ import {
 	SelectValue,
 } from "@/components/ui/select";
 import { Skeleton } from "@/components/ui/skeleton";
+import { Spinner } from "@/components/ui/spinner";
 import { getProviderIcon } from "@/lib/provider-icons";
 import { cn } from "@/lib/utils";
 
@@ -48,8 +47,8 @@ const ROLE_DESCRIPTIONS = {
 		icon: Bot,
 		title: "Agent LLM",
 		description: "Primary LLM for chat interactions and agent operations",
-		color: "text-blue-600 dark:text-blue-400",
+		color: "text-muted-foreground",
-		bgColor: "bg-blue-500/10",
+		bgColor: "bg-muted",
 		prefKey: "agent_llm_id" as const,
 		configType: "llm" as const,
 	},
@@ -57,8 +56,8 @@ const ROLE_DESCRIPTIONS = {
 		icon: FileText,
 		title: "Document Summary LLM",
 		description: "Handles document summarization and research synthesis",
-		color: "text-purple-600 dark:text-purple-400",
+		color: "text-muted-foreground",
-		bgColor: "bg-purple-500/10",
+		bgColor: "bg-muted",
 		prefKey: "document_summary_llm_id" as const,
 		configType: "llm" as const,
 	},
@@ -66,8 +65,8 @@ const ROLE_DESCRIPTIONS = {
 		icon: ImageIcon,
 		title: "Image Generation Model",
 		description: "Model used for AI image generation (DALL-E, GPT Image, etc.)",
-		color: "text-teal-600 dark:text-teal-400",
+		color: "text-muted-foreground",
-		bgColor: "bg-teal-500/10",
+		bgColor: "bg-muted",
 		prefKey: "image_generation_config_id" as const,
 		configType: "image" as const,
 	},
@@ -118,88 +117,41 @@ export function LLMRoleManager({ searchSpaceId }: LLMRoleManagerProps) {
 		image_generation_config_id: preferences.image_generation_config_id ?? "",
 	}));
 
-	const [hasChanges, setHasChanges] = useState(false);
+	const [savingRole, setSavingRole] = useState<string | null>(null);
-	const [isSaving, setIsSaving] = useState(false);
+	const savingRef = useRef(false);
 
 	useEffect(() => {
-		const newAssignments = {
+		if (!savingRef.current) {
-			agent_llm_id: preferences.agent_llm_id ?? "",
+			setAssignments({
-			document_summary_llm_id: preferences.document_summary_llm_id ?? "",
+				agent_llm_id: preferences.agent_llm_id ?? "",
-			image_generation_config_id: preferences.image_generation_config_id ?? "",
+				document_summary_llm_id: preferences.document_summary_llm_id ?? "",
-		};
+				image_generation_config_id: preferences.image_generation_config_id ?? "",
-		setAssignments(newAssignments);
+			});
-		setHasChanges(false);
+		}
 	}, [
 		preferences?.agent_llm_id,
 		preferences?.document_summary_llm_id,
 		preferences?.image_generation_config_id,
 	]);
 
-	const handleRoleAssignment = (prefKey: string, configId: string) => {
+	const handleRoleAssignment = useCallback(async (prefKey: string, configId: string) => {
-		const newAssignments = {
+		const value = configId === "unassigned" ? "" : parseInt(configId);
-			...assignments,
-			[prefKey]: configId === "unassigned" ? "" : parseInt(configId),
-		};
 
-		setAssignments(newAssignments);
+		setAssignments((prev) => ({ ...prev, [prefKey]: value }));
+		setSavingRole(prefKey);
+		savingRef.current = true;
 
-		const currentPrefs = {
+		try {
-			agent_llm_id: preferences.agent_llm_id ?? "",
+			await updatePreferences({
-			document_summary_llm_id: preferences.document_summary_llm_id ?? "",
+				search_space_id: searchSpaceId,
-			image_generation_config_id: preferences.image_generation_config_id ?? "",
+				data: { [prefKey]: value || undefined },
-		};
+			});
-
+			toast.success("Role assignment updated");
-		const hasChangesNow = Object.keys(newAssignments).some(
+		} finally {
-			(key) =>
+			setSavingRole(null);
-				newAssignments[key as keyof typeof newAssignments] !==
+			savingRef.current = false;
-				currentPrefs[key as keyof typeof currentPrefs]
+		}
-		);
+	}, [updatePreferences, searchSpaceId]);
-
-		setHasChanges(hasChangesNow);
-	};
-
-	const handleSave = async () => {
-		setIsSaving(true);
-
-		const toNumericOrUndefined = (val: string | number) =>
-			typeof val === "string" ? (val ? parseInt(val) : undefined) : val;
-
-		const numericAssignments = {
-			agent_llm_id: toNumericOrUndefined(assignments.agent_llm_id),
-			document_summary_llm_id: toNumericOrUndefined(assignments.document_summary_llm_id),
-			image_generation_config_id: toNumericOrUndefined(assignments.image_generation_config_id),
-		};
-
-		await updatePreferences({
-			search_space_id: searchSpaceId,
-			data: numericAssignments,
-		});
-
-		setHasChanges(false);
-		toast.success("Role assignments saved successfully!");
-
-		setIsSaving(false);
-	};
-
-	const handleReset = () => {
-		setAssignments({
-			agent_llm_id: preferences.agent_llm_id ?? "",
-			document_summary_llm_id: preferences.document_summary_llm_id ?? "",
-			image_generation_config_id: preferences.image_generation_config_id ?? "",
-		});
-		setHasChanges(false);
-	};
-
-	const isAssignmentComplete =
-		assignments.agent_llm_id !== "" &&
-		assignments.agent_llm_id !== null &&
-		assignments.agent_llm_id !== undefined &&
-		assignments.document_summary_llm_id !== "" &&
-		assignments.document_summary_llm_id !== null &&
-		assignments.document_summary_llm_id !== undefined &&
-		assignments.image_generation_config_id !== "" &&
-		assignments.image_generation_config_id !== null &&
-		assignments.image_generation_config_id !== undefined;
 
 	// Combine global and custom LLM configs
 	const allLLMConfigs = [
@@ -213,6 +165,11 @@ export function LLMRoleManager({ searchSpaceId }: LLMRoleManagerProps) {
 		...(userImageConfigs ?? []).filter((config) => config.id && config.id.toString().trim() !== ""),
 	];
 
+	const isAssignmentComplete =
+		allLLMConfigs.some((c) => c.id === assignments.agent_llm_id) &&
+		allLLMConfigs.some((c) => c.id === assignments.document_summary_llm_id) &&
+		allImageConfigs.some((c) => c.id === assignments.image_generation_config_id);
+
 	const isLoading =
 		configsLoading ||
 		preferencesLoading ||
@@ -244,9 +201,9 @@ export function LLMRoleManager({ searchSpaceId }: LLMRoleManagerProps) {
 				{isAssignmentComplete && !isLoading && !hasError && (
 					<Badge
 						variant="outline"
-						className="text-xs gap-1.5 border-emerald-500/30 text-emerald-700 dark:text-emerald-300 bg-emerald-500/5"
+						className="text-xs gap-1.5 text-muted-foreground"
 					>
-						<CheckCircle className="h-3 w-3" />
+						<CircleCheck className="h-3 w-3" />
 						All roles assigned
 					</Badge>
 				)}
@@ -332,10 +289,7 @@ export function LLMRoleManager({ searchSpaceId }: LLMRoleManagerProps) {
 						const roleAllConfigs = isImageRole ? allImageConfigs : allLLMConfigs;
 
 						const assignedConfig = roleAllConfigs.find((config) => config.id === currentAssignment);
-						const isAssigned =
+						const isAssigned = !!assignedConfig;
-							currentAssignment !== "" &&
-							currentAssignment !== null &&
-							currentAssignment !== undefined;
 						const isAutoMode =
 							assignedConfig && "is_auto_mode" in assignedConfig && assignedConfig.is_auto_mode;
 
@@ -361,8 +315,10 @@ export function LLMRoleManager({ searchSpaceId }: LLMRoleManagerProps) {
 													</p>
 												</div>
 											</div>
-											{isAssigned ? (
+											{savingRole === role.prefKey ? (
-												<CheckCircle className="w-4 h-4 text-emerald-500 shrink-0 mt-0.5" />
+												<Spinner size="sm" className="shrink-0 mt-0.5 text-muted-foreground" />
+											) : isAssigned ? (
+												<CircleCheck className="w-4 h-4 text-muted-foreground/40 shrink-0 mt-0.5" />
 											) : (
 												<CircleDashed className="w-4 h-4 text-muted-foreground/40 shrink-0 mt-0.5" />
 											)}
@@ -374,7 +330,7 @@ export function LLMRoleManager({ searchSpaceId }: LLMRoleManagerProps) {
 												Configuration
 											</Label>
 											<Select
-												value={currentAssignment?.toString() || "unassigned"}
+												value={isAssigned ? currentAssignment.toString() : "unassigned"}
 												onValueChange={(value) => handleRoleAssignment(role.prefKey, value)}
 											>
 												<SelectTrigger className="w-full h-9 md:h-10 text-xs md:text-sm">
@@ -534,33 +490,6 @@ export function LLMRoleManager({ searchSpaceId }: LLMRoleManagerProps) {
 				</div>
 			)}
 
-			{/* Save / Reset Bar */}
-			{hasChanges && (
-				<div className="flex items-center justify-between gap-3 rounded-lg border border-border bg-muted/50 p-3 md:p-4">
-					<p className="text-xs md:text-sm text-muted-foreground">You have unsaved changes</p>
-					<div className="flex items-center gap-2">
-						<Button
-							variant="outline"
-							size="sm"
-							onClick={handleReset}
-							disabled={isSaving}
-							className="h-8 text-xs gap-1.5"
-						>
-							<RotateCcw className="w-3 h-3" />
-							Reset
-						</Button>
-						<Button
-							size="sm"
-							onClick={handleSave}
-							disabled={isSaving}
-							className="h-8 text-xs gap-1.5"
-						>
-							<Save className="w-3 h-3" />
-							{isSaving ? "Saving…" : "Save Changes"}
-						</Button>
-					</div>
-				</div>
-			)}
 		</div>
 	);
 }