From e1beab798a2c8dc671346e8cb0d84e6a80f9945d Mon Sep 17 00:00:00 2001
From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com>
Date: Wed, 24 Jun 2026 12:24:27 +0530
Subject: [PATCH] fix(zero): enhance zero context fetching with desktop
 authentication

---
 .../components/providers/ZeroProvider.tsx     | 88 +++++++++++++++++--
 1 file changed, 79 insertions(+), 9 deletions(-)

diff --git a/surfsense_web/components/providers/ZeroProvider.tsx b/surfsense_web/components/providers/ZeroProvider.tsx
index 4511fe842..530e9c958 100644
--- a/surfsense_web/components/providers/ZeroProvider.tsx
+++ b/surfsense_web/components/providers/ZeroProvider.tsx
@@ -5,17 +5,22 @@ import {
 	useZero,
 	ZeroProvider as ZeroReactProvider,
 } from "@rocicorp/zero/react";
-import { useAtomValue } from "jotai";
 import { usePathname } from "next/navigation";
 import { useEffect, useMemo, useState } from "react";
-import { currentUserAtom } from "@/atoms/user/user-query.atoms";
 import { useSession } from "@/hooks/use-session";
 import { getDesktopAccessToken } from "@/lib/auth-fetch";
 import { handleUnauthorized, isPublicRoute, refreshSession } from "@/lib/auth-utils";
+import { buildBackendUrl } from "@/lib/env-config";
+import type { Context } from "@/types/zero";
 import { queries } from "@/zero/queries";
 import { schema } from "@/zero/schema";
 
 const configuredCacheURL = process.env.NEXT_PUBLIC_ZERO_CACHE_URL;
+type ZeroContext = Exclude<Context, undefined>;
+type LoadedZeroContext = {
+	context: ZeroContext;
+	desktopAuth?: string;
+};
 
 function getCacheURL() {
 	if (configuredCacheURL) return configuredCacheURL;
@@ -25,6 +30,30 @@ function getCacheURL() {
 	return "http://localhost:4848";
 }
 
+async function fetchZeroContext(isDesktop: boolean): Promise<LoadedZeroContext | null> {
+	const headers: HeadersInit = {};
+	let desktopAuth: string | undefined;
+
+	if (isDesktop) {
+		const token = await getDesktopAccessToken();
+		if (!token) return null;
+		desktopAuth = token;
+		headers.Authorization = `Bearer ${token}`;
+	}
+
+	const response = await fetch(buildBackendUrl("/zero/context"), {
+		credentials: "include",
+		headers,
+	});
+
+	if (!response.ok) return null;
+
+	return {
+		context: (await response.json()) as ZeroContext,
+		desktopAuth,
+	};
+}
+
 function ZeroAuthSync({ isDesktop }: { isDesktop: boolean }) {
 	const zero = useZero();
 	const connectionState = useConnectionState();
@@ -70,17 +99,51 @@ function AuthenticatedZeroProvider({
 	children: React.ReactNode;
 	isDesktop: boolean;
 }) {
-	const { data: user, isLoading } = useAtomValue(currentUserAtom);
+	const [loadedContext, setLoadedContext] = useState<LoadedZeroContext | null>(null);
 
-	const userId = user?.id;
-	const userID = userId ? String(userId) : undefined;
+	useEffect(() => {
+		let isMounted = true;
 
-	if (isLoading || !userID) {
+		const load = async () => {
+			const nextContext = await fetchZeroContext(isDesktop);
+			if (isMounted) {
+				setLoadedContext(nextContext);
+			}
+		};
+
+		void load();
+
+		if (!isDesktop || typeof window === "undefined" || !window.electronAPI?.onAuthChanged) {
+			return () => {
+				isMounted = false;
+			};
+		}
+
+		const unsubscribe = window.electronAPI.onAuthChanged(({ accessToken }) => {
+			if (!accessToken) {
+				setLoadedContext(null);
+				return;
+			}
+			void load();
+		});
+
+		return () => {
+			isMounted = false;
+			unsubscribe();
+		};
+	}, [isDesktop]);
+
+	if (!loadedContext) {
 		return <>{children}</>;
 	}
 
 	return (
-		<ZeroClientProvider userID={userID} isDesktop={isDesktop}>
+		<ZeroClientProvider
+			userID={loadedContext.context.userId}
+			context={loadedContext.context}
+			isDesktop={isDesktop}
+			initialDesktopAuth={loadedContext.desktopAuth}
+		>
 			{children}
 		</ZeroClientProvider>
 	);
@@ -89,15 +152,22 @@ function AuthenticatedZeroProvider({
 function ZeroClientProvider({
 	children,
 	userID,
+	context,
 	isDesktop,
+	initialDesktopAuth,
 }: {
 	children: React.ReactNode;
 	userID: string;
+	context: ZeroContext;
 	isDesktop: boolean;
+	initialDesktopAuth?: string;
 }) {
 	const cacheURL = useMemo(() => getCacheURL(), []);
-	const [desktopAuth, setDesktopAuth] = useState<string | undefined>(undefined);
-	const context = useMemo(() => ({ userId: userID }), [userID]);
+	const [desktopAuth, setDesktopAuth] = useState<string | undefined>(initialDesktopAuth);
+
+	useEffect(() => {
+		setDesktopAuth(initialDesktopAuth);
+	}, [initialDesktopAuth]);
 
 	useEffect(() => {
 		if (!isDesktop) return;