From 5efd4291bfe54c7815f052be953ebcd6195be9e8 Mon Sep 17 00:00:00 2001
From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com>
Date: Sun, 12 Jul 2026 20:29:59 +0530
Subject: [PATCH] feat(model-connections): add llmSetupStatusAtomFamily and
 invalidate on mutations

---
 .../model-connections-mutation.atoms.ts             |  6 ++++++
 .../model-connections-query.atoms.ts                | 13 +++++++++++++
 2 files changed, 19 insertions(+)

diff --git a/surfsense_web/atoms/model-connections/model-connections-mutation.atoms.ts b/surfsense_web/atoms/model-connections/model-connections-mutation.atoms.ts
index 5228833e2..c30d65587 100644
--- a/surfsense_web/atoms/model-connections/model-connections-mutation.atoms.ts
+++ b/surfsense_web/atoms/model-connections/model-connections-mutation.atoms.ts
@@ -25,6 +25,9 @@ function invalidateModelConnections(workspaceId: number) {
 	queryClient.invalidateQueries({
 		queryKey: cacheKeys.modelConnections.roles(workspaceId),
 	});
+	queryClient.invalidateQueries({
+		queryKey: cacheKeys.modelConnections.setupStatus(workspaceId),
+	});
 }
 
 function upsertModelConnection(workspaceId: number, connection: ConnectionRead) {
@@ -208,6 +211,9 @@ export const updateModelRolesMutationAtom = atomWithMutation((get) => {
 			queryClient.invalidateQueries({
 				queryKey: cacheKeys.modelConnections.roles(workspaceId),
 			});
+			queryClient.invalidateQueries({
+				queryKey: cacheKeys.modelConnections.setupStatus(workspaceId),
+			});
 		},
 		onError: (error: Error) => toast.error(error.message || "Failed to update model roles"),
 	};
diff --git a/surfsense_web/atoms/model-connections/model-connections-query.atoms.ts b/surfsense_web/atoms/model-connections/model-connections-query.atoms.ts
index 36cc6ae4e..f73fc3f18 100644
--- a/surfsense_web/atoms/model-connections/model-connections-query.atoms.ts
+++ b/surfsense_web/atoms/model-connections/model-connections-query.atoms.ts
@@ -1,4 +1,5 @@
 import { atomWithQuery } from "jotai-tanstack-query";
+import { atomFamily } from "jotai/utils";
 import { modelConnectionsApiService } from "@/lib/apis/model-connections-api.service";
 import { isAuthenticated } from "@/lib/auth-utils";
 import { cacheKeys } from "@/lib/query-client/cache-keys";
@@ -44,3 +45,15 @@ export const modelRolesAtom = atomWithQuery((get) => {
 		queryFn: () => modelConnectionsApiService.getModelRoles(workspaceId),
 	};
 });
+
+// Keyed by the route workspaceId (not activeWorkspaceIdAtom) so the onboarding
+// gate's verdict can never be computed against a stale workspace during a
+// workspace switch.
+export const llmSetupStatusAtomFamily = atomFamily((workspaceId: number) =>
+	atomWithQuery(() => ({
+		queryKey: cacheKeys.modelConnections.setupStatus(workspaceId),
+		enabled: workspaceId > 0 && isAuthenticated(),
+		staleTime: 5 * 60 * 1000,
+		queryFn: () => modelConnectionsApiService.getLlmSetupStatus(workspaceId),
+	}))
+);