From b5bab255eb282b5a70d609ec0a8d94b34df7a0af Mon Sep 17 00:00:00 2001 From: JonathanLab Date: Mon, 30 Mar 2026 16:11:18 +0200 Subject: [PATCH] refactor: hook up renderer auth logic to service --- .../src/main/services/auth/service.test.ts | 66 ++++++++ apps/code/src/main/services/auth/service.ts | 4 +- apps/code/src/renderer/App.tsx | 25 ++- .../components/ScopeReauthPrompt.test.tsx | 148 +++++++----------- .../renderer/components/ScopeReauthPrompt.tsx | 30 ++-- .../features/auth/components/AuthScreen.tsx | 42 +++-- .../auth/components/InviteCodeScreen.tsx | 29 ++-- .../features/auth/hooks/authClient.ts | 63 ++++++++ .../features/auth/hooks/authMutations.ts | 91 +++++++++++ .../features/auth/hooks/authQueries.ts | 90 +++++++++++ .../features/auth/hooks/useAuthSession.ts | 97 ++++++++++++ .../features/auth/stores/authUiStateStore.ts | 34 ++++ .../inbox/components/DataSourceSetup.tsx | 17 +- .../inbox/components/InboxSignalsTab.tsx | 6 +- .../inbox/hooks/useExternalDataSources.ts | 4 +- .../inbox/hooks/useSignalSourceConfigs.ts | 4 +- .../inbox/hooks/useSignalSourceManager.ts | 7 +- .../onboarding/components/BillingStep.tsx | 5 +- .../components/GitIntegrationStep.tsx | 45 +++--- .../onboarding/components/OnboardingFlow.tsx | 10 +- .../onboarding/components/OrgBillingStep.tsx | 50 +++--- .../onboarding/components/TutorialStep.tsx | 6 +- .../onboarding/hooks/useOnboardingFlow.ts | 8 +- .../hooks/useProjectsWithIntegrations.ts | 6 +- .../onboarding/stores/onboardingStore.ts | 68 ++++++++ .../features/projects/hooks/useProjects.tsx | 34 ++-- .../sessions/hooks/useChatTitleGenerator.ts | 4 +- .../features/sessions/service/service.test.ts | 98 +++++++++--- .../features/sessions/service/service.ts | 80 ++++------ .../components/sections/AccountSettings.tsx | 32 ++-- .../components/sections/AdvancedSettings.tsx | 6 +- .../components/sections/GeneralSettings.tsx | 6 +- .../sidebar/components/ProjectSwitcher.tsx | 16 +- .../task-detail/hooks/usePreviewSession.ts | 4 +- .../task-detail/hooks/useTaskCreation.ts | 6 +- .../features/task-detail/service/service.ts | 6 +- .../renderer/hooks/useAuthenticatedClient.ts | 10 +- .../hooks/useAuthenticatedInfiniteQuery.ts | 6 +- .../hooks/useAuthenticatedMutation.ts | 4 +- .../renderer/hooks/useAuthenticatedQuery.ts | 15 +- .../src/renderer/hooks/useOrganizations.ts | 15 +- .../src/renderer/hooks/useProjectQuery.ts | 4 +- .../src/renderer/hooks/useTaskDeepLink.ts | 6 +- apps/code/src/renderer/utils/generateTitle.ts | 6 +- 44 files changed, 926 insertions(+), 387 deletions(-) create mode 100644 apps/code/src/renderer/features/auth/hooks/authClient.ts create mode 100644 apps/code/src/renderer/features/auth/hooks/authMutations.ts create mode 100644 apps/code/src/renderer/features/auth/hooks/authQueries.ts create mode 100644 apps/code/src/renderer/features/auth/hooks/useAuthSession.ts create mode 100644 apps/code/src/renderer/features/auth/stores/authUiStateStore.ts create mode 100644 apps/code/src/renderer/features/onboarding/stores/onboardingStore.ts diff --git a/apps/code/src/main/services/auth/service.test.ts b/apps/code/src/main/services/auth/service.test.ts index c90c932c0..dc822fe71 100644 --- a/apps/code/src/main/services/auth/service.test.ts +++ b/apps/code/src/main/services/auth/service.test.ts @@ -171,4 +171,70 @@ describe("AuthService", () => { "rotated-refresh-token", ); }); + + it("preserves the selected project across logout and re-login for the same account", async () => { + vi.mocked(oauthService.startFlow) + .mockResolvedValueOnce({ + success: true, + data: { + access_token: "initial-access-token", + refresh_token: "initial-refresh-token", + expires_in: 3600, + token_type: "Bearer", + scope: "", + scoped_teams: [42, 84], + scoped_organizations: ["org-1"], + }, + }) + .mockResolvedValueOnce({ + success: true, + data: { + access_token: "second-access-token", + refresh_token: "second-refresh-token", + expires_in: 3600, + token_type: "Bearer", + scope: "", + scoped_teams: [42, 84], + scoped_organizations: ["org-1"], + }, + }); + vi.mocked(oauthService.refreshToken).mockResolvedValue({ + success: true, + data: { + access_token: "refreshed-access-token", + refresh_token: "refreshed-refresh-token", + expires_in: 3600, + token_type: "Bearer", + scope: "", + scoped_teams: [42, 84], + scoped_organizations: ["org-1"], + }, + }); + + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue({ + json: vi.fn().mockResolvedValue({ has_access: true }), + }) as unknown as typeof fetch, + ); + + await service.login("us"); + await service.selectProject(84); + await service.logout(); + + expect(service.getState()).toMatchObject({ + status: "anonymous", + cloudRegion: "us", + projectId: 84, + }); + + await service.login("us"); + + expect(service.getState()).toMatchObject({ + status: "authenticated", + cloudRegion: "us", + projectId: 84, + availableProjectIds: [42, 84], + }); + }); }); diff --git a/apps/code/src/main/services/auth/service.ts b/apps/code/src/main/services/auth/service.ts index 6472f5f1f..4b9808e43 100644 --- a/apps/code/src/main/services/auth/service.ts +++ b/apps/code/src/main/services/auth/service.ts @@ -219,9 +219,11 @@ export class AuthService extends TypedEventEmitter { } async logout(): Promise { + const { cloudRegion, projectId } = this.state; + this.authSessionRepository.clearCurrent(); this.session = null; - this.setAnonymousState(); + this.setAnonymousState({ cloudRegion, projectId }); return this.getState(); } diff --git a/apps/code/src/renderer/App.tsx b/apps/code/src/renderer/App.tsx index 3ea3e1d85..b0d2fbeba 100644 --- a/apps/code/src/renderer/App.tsx +++ b/apps/code/src/renderer/App.tsx @@ -5,8 +5,10 @@ import { ScopeReauthPrompt } from "@components/ScopeReauthPrompt"; import { UpdatePrompt } from "@components/UpdatePrompt"; import { AuthScreen } from "@features/auth/components/AuthScreen"; import { InviteCodeScreen } from "@features/auth/components/InviteCodeScreen"; -import { useAuthStore } from "@features/auth/stores/authStore"; +import { useAuthStateValue } from "@features/auth/hooks/authQueries"; +import { useAuthSession } from "@features/auth/hooks/useAuthSession"; import { OnboardingFlow } from "@features/onboarding/components/OnboardingFlow"; +import { useOnboardingStore } from "@features/onboarding/stores/onboardingStore"; import { Flex, Spinner, Text } from "@radix-ui/themes"; import { initializeConnectivityStore } from "@renderer/stores/connectivityStore"; import { useFocusStore } from "@renderer/stores/focusStore"; @@ -25,10 +27,14 @@ const log = logger.scope("app"); function App() { const trpcReact = useTRPC(); - const { isAuthenticated, hasCompletedOnboarding, hasCodeAccess } = - useAuthStore(); + const { isBootstrapped } = useAuthSession(); + const authState = useAuthStateValue((state) => state); + const hasCompletedOnboarding = useOnboardingStore( + (state) => state.hasCompletedOnboarding, + ); + const isAuthenticated = authState.status === "authenticated"; + const hasCodeAccess = authState.hasCodeAccess; const isDarkMode = useThemeStore((state) => state.isDarkMode); - const [isLoading, setIsLoading] = useState(true); const [showTransition, setShowTransition] = useState(false); const wasInMainApp = useRef(isAuthenticated && hasCompletedOnboarding); @@ -114,15 +120,6 @@ function App() { }), ); - // Initialize auth state from main process - useEffect(() => { - const initialize = async () => { - await useAuthStore.getState().initializeOAuth(); - setIsLoading(false); - }; - void initialize(); - }, []); - // Handle transition into main app — only show the dark overlay if dark mode is active useEffect(() => { const isInMainApp = isAuthenticated && hasCompletedOnboarding; @@ -136,7 +133,7 @@ function App() { setShowTransition(false); }; - if (isLoading) { + if (!isBootstrapped) { return ( diff --git a/apps/code/src/renderer/components/ScopeReauthPrompt.test.tsx b/apps/code/src/renderer/components/ScopeReauthPrompt.test.tsx index 1253f3476..0cb091af0 100644 --- a/apps/code/src/renderer/components/ScopeReauthPrompt.test.tsx +++ b/apps/code/src/renderer/components/ScopeReauthPrompt.test.tsx @@ -1,65 +1,40 @@ +import { Theme } from "@radix-ui/themes"; import { render, screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; +import type { ReactElement } from "react"; import { beforeEach, describe, expect, it, vi } from "vitest"; +import { ScopeReauthPrompt } from "./ScopeReauthPrompt"; -vi.mock("@renderer/trpc/client", () => ({ - trpcClient: { - auth: { - getState: { query: vi.fn() }, - onStateChanged: { subscribe: vi.fn(() => ({ unsubscribe: vi.fn() })) }, - getValidAccessToken: { - query: vi.fn().mockResolvedValue({ - accessToken: "token", - apiHost: "https://us.posthog.com", - }), - }, - refreshAccessToken: { - mutate: vi.fn().mockResolvedValue({ - accessToken: "token", - apiHost: "https://us.posthog.com", - }), - }, - login: { - mutate: vi.fn().mockResolvedValue({ - state: { - status: "authenticated", - bootstrapComplete: true, - cloudRegion: "us", - projectId: 1, - availableProjectIds: [1], - availableOrgIds: [], - hasCodeAccess: true, - needsScopeReauth: false, - }, - }), - }, - signup: { mutate: vi.fn() }, - selectProject: { mutate: vi.fn() }, - redeemInviteCode: { mutate: vi.fn() }, - logout: { - mutate: vi.fn().mockResolvedValue({ - status: "anonymous", - bootstrapComplete: true, - cloudRegion: null, - projectId: null, - availableProjectIds: [], - availableOrgIds: [], - hasCodeAccess: null, - needsScopeReauth: false, - }), - }, - }, - analytics: { - setUserId: { mutate: vi.fn().mockResolvedValue(undefined) }, - resetUser: { mutate: vi.fn().mockResolvedValue(undefined) }, - }, - }, +const authState = { + status: "anonymous" as const, + bootstrapComplete: true, + cloudRegion: null as "us" | "eu" | "dev" | null, + projectId: null, + availableProjectIds: [], + availableOrgIds: [], + hasCodeAccess: null, + needsScopeReauth: false, +}; + +const mockLoginMutateAsync = vi.fn(); +const mockLogoutMutate = vi.fn(() => { + authState.needsScopeReauth = false; + authState.cloudRegion = null; +}); + +vi.mock("@features/auth/hooks/authQueries", () => ({ + useAuthStateValue: (selector: (state: typeof authState) => unknown) => + selector(authState), })); -vi.mock("@utils/analytics", () => ({ - identifyUser: vi.fn(), - resetUser: vi.fn(), - track: vi.fn(), +vi.mock("@features/auth/hooks/authMutations", () => ({ + useLoginMutation: () => ({ + mutateAsync: mockLoginMutateAsync, + isPending: false, + }), + useLogoutMutation: () => ({ + mutate: mockLogoutMutate, + }), })); vi.mock("@utils/logger", () => ({ @@ -73,40 +48,18 @@ vi.mock("@utils/logger", () => ({ }, })); -vi.mock("@utils/queryClient", () => ({ - queryClient: { - clear: vi.fn(), - setQueryData: vi.fn(), - removeQueries: vi.fn(), - }, -})); - -vi.mock("@stores/navigationStore", () => ({ - useNavigationStore: { - getState: () => ({ navigateToTaskInput: vi.fn() }), - }, -})); - -import { - resetAuthStoreModuleStateForTest, - useAuthStore, -} from "@features/auth/stores/authStore"; -import { Theme } from "@radix-ui/themes"; -import type { ReactElement } from "react"; -import { ScopeReauthPrompt } from "./ScopeReauthPrompt"; - function renderWithTheme(ui: ReactElement) { return render({ui}); } describe("ScopeReauthPrompt", () => { beforeEach(() => { - localStorage.clear(); - resetAuthStoreModuleStateForTest(); - useAuthStore.setState({ - needsScopeReauth: false, - cloudRegion: null, - }); + vi.clearAllMocks(); + authState.status = "anonymous"; + authState.cloudRegion = null; + authState.projectId = null; + authState.hasCodeAccess = null; + authState.needsScopeReauth = false; }); it("does not render dialog when needsScopeReauth is false", () => { @@ -117,25 +70,34 @@ describe("ScopeReauthPrompt", () => { }); it("renders dialog when needsScopeReauth is true", () => { - useAuthStore.setState({ needsScopeReauth: true, cloudRegion: "us" }); + authState.needsScopeReauth = true; + authState.cloudRegion = "us"; + renderWithTheme(); + expect(screen.getByText("Re-authentication required")).toBeInTheDocument(); }); it("disables Sign in button when cloudRegion is null", () => { - useAuthStore.setState({ needsScopeReauth: true, cloudRegion: null }); + authState.needsScopeReauth = true; + renderWithTheme(); + expect(screen.getByRole("button", { name: "Sign in" })).toBeDisabled(); }); it("enables Sign in button when cloudRegion is set", () => { - useAuthStore.setState({ needsScopeReauth: true, cloudRegion: "us" }); + authState.needsScopeReauth = true; + authState.cloudRegion = "us"; + renderWithTheme(); + expect(screen.getByRole("button", { name: "Sign in" })).not.toBeDisabled(); }); it("shows Log out button as an escape hatch when cloudRegion is null", () => { - useAuthStore.setState({ needsScopeReauth: true, cloudRegion: null }); + authState.needsScopeReauth = true; + renderWithTheme(); const logoutButton = screen.getByRole("button", { name: "Log out" }); @@ -145,14 +107,14 @@ describe("ScopeReauthPrompt", () => { it("calls logout when Log out button is clicked", async () => { const user = userEvent.setup(); - useAuthStore.setState({ needsScopeReauth: true, cloudRegion: null }); + authState.needsScopeReauth = true; + renderWithTheme(); await user.click(screen.getByRole("button", { name: "Log out" })); - const state = useAuthStore.getState(); - expect(state.needsScopeReauth).toBe(false); - expect(state.isAuthenticated).toBe(false); - expect(state.cloudRegion).toBeNull(); + expect(mockLogoutMutate).toHaveBeenCalledTimes(1); + expect(authState.needsScopeReauth).toBe(false); + expect(authState.cloudRegion).toBeNull(); }); }); diff --git a/apps/code/src/renderer/components/ScopeReauthPrompt.tsx b/apps/code/src/renderer/components/ScopeReauthPrompt.tsx index 2fbcfdc8b..0bbedcfed 100644 --- a/apps/code/src/renderer/components/ScopeReauthPrompt.tsx +++ b/apps/code/src/renderer/components/ScopeReauthPrompt.tsx @@ -1,17 +1,19 @@ -import { useAuthStore } from "@features/auth/stores/authStore"; +import { + useLoginMutation, + useLogoutMutation, +} from "@features/auth/hooks/authMutations"; +import { useAuthStateValue } from "@features/auth/hooks/authQueries"; import { ShieldWarning } from "@phosphor-icons/react"; import { Button, Dialog, Flex, Text } from "@radix-ui/themes"; import { logger } from "@utils/logger"; -import { useState } from "react"; const log = logger.scope("scope-reauth-prompt"); export function ScopeReauthPrompt() { - const needsScopeReauth = useAuthStore((s) => s.needsScopeReauth); - const cloudRegion = useAuthStore((s) => s.cloudRegion); - const loginWithOAuth = useAuthStore((s) => s.loginWithOAuth); - const logout = useAuthStore((s) => s.logout); - const [isLoading, setIsLoading] = useState(false); + const needsScopeReauth = useAuthStateValue((state) => state.needsScopeReauth); + const cloudRegion = useAuthStateValue((state) => state.cloudRegion); + const loginMutation = useLoginMutation(); + const logoutMutation = useLogoutMutation(); const handleSignIn = async () => { if (!cloudRegion) { @@ -19,13 +21,10 @@ export function ScopeReauthPrompt() { return; } - setIsLoading(true); try { - await loginWithOAuth(cloudRegion); + await loginMutation.mutateAsync(cloudRegion); } catch (error) { log.error("Re-authentication failed", error); - } finally { - setIsLoading(false); } }; @@ -50,13 +49,18 @@ export function ScopeReauthPrompt() { - diff --git a/apps/code/src/renderer/features/onboarding/components/TutorialStep.tsx b/apps/code/src/renderer/features/onboarding/components/TutorialStep.tsx index c1cc7fe0a..8f4e5cb13 100644 --- a/apps/code/src/renderer/features/onboarding/components/TutorialStep.tsx +++ b/apps/code/src/renderer/features/onboarding/components/TutorialStep.tsx @@ -1,11 +1,11 @@ import { TourHighlight } from "@components/TourHighlight"; -import { useAuthStore } from "@features/auth/stores/authStore"; import { FolderPicker } from "@features/folder-picker/components/FolderPicker"; import { GitHubRepoPicker } from "@features/folder-picker/components/GitHubRepoPicker"; import { BranchSelector } from "@features/git-interaction/components/BranchSelector"; import type { MessageEditorHandle } from "@features/message-editor/components/MessageEditor"; import { ModeIndicatorInput } from "@features/message-editor/components/ModeIndicatorInput"; import { useDraftStore } from "@features/message-editor/stores/draftStore"; +import { useOnboardingStore } from "@features/onboarding/stores/onboardingStore"; import { getSessionService } from "@features/sessions/service/service"; import { cycleModeOption, @@ -61,7 +61,9 @@ interface TutorialStepProps { export function TutorialStep({ onComplete, onBack }: TutorialStepProps) { const { allowBypassPermissions } = useSettingsStore(); - const { completeOnboarding } = useAuthStore(); + const completeOnboarding = useOnboardingStore( + (state) => state.completeOnboarding, + ); // Tour state machine const { diff --git a/apps/code/src/renderer/features/onboarding/hooks/useOnboardingFlow.ts b/apps/code/src/renderer/features/onboarding/hooks/useOnboardingFlow.ts index 3fca32901..e411d5283 100644 --- a/apps/code/src/renderer/features/onboarding/hooks/useOnboardingFlow.ts +++ b/apps/code/src/renderer/features/onboarding/hooks/useOnboardingFlow.ts @@ -1,9 +1,11 @@ +import { useOnboardingStore } from "@features/onboarding/stores/onboardingStore"; import { useFeatureFlag } from "@hooks/useFeatureFlag"; -import { useEffect, useMemo, useState } from "react"; +import { useEffect, useMemo } from "react"; import { ONBOARDING_STEPS, type OnboardingStep } from "../types"; export function useOnboardingFlow() { - const [currentStep, setCurrentStep] = useState("welcome"); + const currentStep = useOnboardingStore((state) => state.currentStep); + const setCurrentStep = useOnboardingStore((state) => state.setCurrentStep); const billingEnabled = useFeatureFlag("twig-billing", false); // Show billing onboarding steps only when billing is enabled @@ -21,7 +23,7 @@ export function useOnboardingFlow() { if (!activeSteps.includes(currentStep)) { setCurrentStep(activeSteps[0]); } - }, [activeSteps, currentStep]); + }, [activeSteps, currentStep, setCurrentStep]); const currentIndex = activeSteps.indexOf(currentStep); const isFirstStep = currentIndex === 0; diff --git a/apps/code/src/renderer/features/onboarding/hooks/useProjectsWithIntegrations.ts b/apps/code/src/renderer/features/onboarding/hooks/useProjectsWithIntegrations.ts index 55b62e437..e441f606a 100644 --- a/apps/code/src/renderer/features/onboarding/hooks/useProjectsWithIntegrations.ts +++ b/apps/code/src/renderer/features/onboarding/hooks/useProjectsWithIntegrations.ts @@ -1,4 +1,5 @@ -import { useAuthStore } from "@features/auth/stores/authStore"; +import { useAuthenticatedClient } from "@features/auth/hooks/authClient"; +import { AUTH_SCOPED_QUERY_META } from "@features/auth/hooks/authQueries"; import type { Integration } from "@features/integrations/stores/integrationStore"; import { useProjects } from "@features/projects/hooks/useProjects"; import { useQueries } from "@tanstack/react-query"; @@ -14,7 +15,7 @@ export interface ProjectWithIntegrations { export function useProjectsWithIntegrations() { const { projects, isLoading: projectsLoading } = useProjects(); - const client = useAuthStore((s) => s.client); + const client = useAuthenticatedClient(); // Fetch integrations for each project in parallel const integrationQueries = useQueries({ @@ -26,6 +27,7 @@ export function useProjectsWithIntegrations() { }, enabled: !!client && projects.length > 0, staleTime: 60 * 1000, // 1 minute + meta: AUTH_SCOPED_QUERY_META, })), }); diff --git a/apps/code/src/renderer/features/onboarding/stores/onboardingStore.ts b/apps/code/src/renderer/features/onboarding/stores/onboardingStore.ts new file mode 100644 index 000000000..4d165dbd6 --- /dev/null +++ b/apps/code/src/renderer/features/onboarding/stores/onboardingStore.ts @@ -0,0 +1,68 @@ +import { create } from "zustand"; +import { persist } from "zustand/middleware"; +import type { OnboardingStep } from "../types"; + +interface OnboardingStoreState { + currentStep: OnboardingStep; + hasCompletedOnboarding: boolean; + isConnectingGithub: boolean; + selectedPlan: "free" | "pro" | null; + selectedOrgId: string | null; + selectedProjectId: number | null; +} + +interface OnboardingStoreActions { + setCurrentStep: (step: OnboardingStep) => void; + completeOnboarding: () => void; + resetOnboarding: () => void; + resetSelections: () => void; + setConnectingGithub: (isConnecting: boolean) => void; + selectPlan: (plan: "free" | "pro") => void; + selectOrg: (orgId: string) => void; + selectProjectId: (projectId: number | null) => void; +} + +type OnboardingStore = OnboardingStoreState & OnboardingStoreActions; + +const initialState: OnboardingStoreState = { + currentStep: "welcome", + hasCompletedOnboarding: false, + isConnectingGithub: false, + selectedPlan: null, + selectedOrgId: null, + selectedProjectId: null, +}; + +export const useOnboardingStore = create()( + persist( + (set) => ({ + ...initialState, + + setCurrentStep: (step) => set({ currentStep: step }), + completeOnboarding: () => set({ hasCompletedOnboarding: true }), + resetOnboarding: () => set({ ...initialState }), + resetSelections: () => + set({ + currentStep: "welcome", + isConnectingGithub: false, + selectedPlan: null, + selectedOrgId: null, + selectedProjectId: null, + }), + setConnectingGithub: (isConnectingGithub) => set({ isConnectingGithub }), + selectPlan: (plan) => set({ selectedPlan: plan }), + selectOrg: (orgId) => set({ selectedOrgId: orgId }), + selectProjectId: (selectedProjectId) => set({ selectedProjectId }), + }), + { + name: "onboarding-store", + partialize: (state) => ({ + currentStep: state.currentStep, + hasCompletedOnboarding: state.hasCompletedOnboarding, + selectedPlan: state.selectedPlan, + selectedOrgId: state.selectedOrgId, + selectedProjectId: state.selectedProjectId, + }), + }, + ), +); diff --git a/apps/code/src/renderer/features/projects/hooks/useProjects.tsx b/apps/code/src/renderer/features/projects/hooks/useProjects.tsx index 18f611ef4..e7406bcc0 100644 --- a/apps/code/src/renderer/features/projects/hooks/useProjects.tsx +++ b/apps/code/src/renderer/features/projects/hooks/useProjects.tsx @@ -1,5 +1,9 @@ -import { useAuthStore } from "@features/auth/stores/authStore"; -import { useQuery } from "@tanstack/react-query"; +import { useOptionalAuthenticatedClient } from "@features/auth/hooks/authClient"; +import { useSelectProjectMutation } from "@features/auth/hooks/authMutations"; +import { + useAuthStateValue, + useCurrentUser, +} from "@features/auth/hooks/authQueries"; import { logger } from "@utils/logger"; import { useEffect, useMemo } from "react"; @@ -36,20 +40,12 @@ export function groupProjectsByOrg(projects: ProjectInfo[]): GroupedProjects[] { } export function useProjects() { - const availableProjectIds = useAuthStore((s) => s.availableProjectIds); - const client = useAuthStore((s) => s.client); - const currentProjectId = useAuthStore((s) => s.projectId); - - const { - data: currentUser, - isLoading, - error, - } = useQuery({ - queryKey: ["currentUser"], - queryFn: () => client?.getCurrentUser(), - enabled: !!client, - staleTime: 5 * 60 * 1000, - }); + const availableProjectIds = useAuthStateValue( + (state) => state.availableProjectIds, + ); + const currentProjectId = useAuthStateValue((state) => state.projectId); + const client = useOptionalAuthenticatedClient(); + const { data: currentUser, isLoading, error } = useCurrentUser({ client }); const projects = useMemo(() => { if (!currentUser?.organization) return []; @@ -84,7 +80,7 @@ export function useProjects() { .filter((p): p is ProjectInfo => p !== null); }, [currentUser, availableProjectIds]); - const selectProject = useAuthStore((s) => s.selectProject); + const selectProjectMutation = useSelectProjectMutation(); const currentProject = projects.find((p) => p.id === currentProjectId); const groupedProjects = groupProjectsByOrg(projects); @@ -97,9 +93,9 @@ export function useProjects() { ? "no project selected" : "current project not found in list", }); - selectProject(projects[0].id); + selectProjectMutation.mutate(projects[0].id); } - }, [projects, currentProject, currentProjectId, selectProject]); + }, [currentProject, currentProjectId, projects, selectProjectMutation]); return { projects, diff --git a/apps/code/src/renderer/features/sessions/hooks/useChatTitleGenerator.ts b/apps/code/src/renderer/features/sessions/hooks/useChatTitleGenerator.ts index 1202fe781..b5a8d5e46 100644 --- a/apps/code/src/renderer/features/sessions/hooks/useChatTitleGenerator.ts +++ b/apps/code/src/renderer/features/sessions/hooks/useChatTitleGenerator.ts @@ -1,4 +1,4 @@ -import { useAuthStore } from "@features/auth/stores/authStore"; +import { getAuthenticatedClient } from "@features/auth/hooks/authClient"; import { getSessionService } from "@features/sessions/service/service"; import { useSessionStore } from "@features/sessions/stores/sessionStore"; import type { Task } from "@shared/types"; @@ -71,7 +71,7 @@ export function useChatTitleGenerator(taskId: string): void { const title = await generateTitle(content); if (title) { - const client = useAuthStore.getState().client; + const client = await getAuthenticatedClient(); if (client) { await client.updateTask(taskId, { title }); queryClient.setQueriesData( diff --git a/apps/code/src/renderer/features/sessions/service/service.test.ts b/apps/code/src/renderer/features/sessions/service/service.test.ts index dbe19f2db..1bf42f202 100644 --- a/apps/code/src/renderer/features/sessions/service/service.test.ts +++ b/apps/code/src/renderer/features/sessions/service/service.test.ts @@ -61,21 +61,49 @@ vi.mock("@features/sessions/stores/sessionStore", () => ({ mergeConfigOptions: vi.fn((live: unknown[], _persisted: unknown[]) => live), })); -const mockAuthStore = vi.hoisted(() => ({ - useAuthStore: { - getState: vi.fn(() => ({ - cloudRegion: "us", - projectId: 123, - client: { - createTaskRun: vi.fn(), - appendTaskRunLog: vi.fn(), - }, - })), - }, - setSessionResetCallback: vi.fn(), +const mockBuildAuthenticatedClient = vi.hoisted(() => + vi.fn< + () => { + createTaskRun: ReturnType; + appendTaskRunLog: ReturnType; + } | null + >(() => ({ + createTaskRun: vi.fn(), + appendTaskRunLog: vi.fn(), + })), +); + +const mockAuth = vi.hoisted(() => ({ + fetchAuthState: vi.fn<() => Promise>>(async () => ({ + status: "authenticated", + bootstrapComplete: true, + cloudRegion: "us", + projectId: 123, + availableProjectIds: [123], + availableOrgIds: [], + hasCodeAccess: true, + needsScopeReauth: false, + })), + getAuthenticatedClient: vi.fn<() => Promise | null>>( + async () => mockBuildAuthenticatedClient(), + ), + createAuthenticatedClient: vi.fn((authState: Record) => { + return authState.status === "authenticated" + ? mockBuildAuthenticatedClient() + : null; + }), })); -vi.mock("@features/auth/stores/authStore", () => mockAuthStore); +vi.mock("@features/auth/hooks/authQueries", () => ({ + AUTH_SCOPED_QUERY_META: { authScoped: true }, + clearAuthScopedQueries: vi.fn(), + getAuthIdentity: vi.fn(), + fetchAuthState: mockAuth.fetchAuthState, +})); +vi.mock("@features/auth/hooks/authClient", () => ({ + getAuthenticatedClient: mockAuth.getAuthenticatedClient, + createAuthenticatedClient: mockAuth.createAuthenticatedClient, +})); vi.mock("@features/sessions/stores/modelsStore", () => ({ useModelsStore: { @@ -280,13 +308,19 @@ describe("SessionService", () => { // Track how many times createTaskRun is called const createTaskRunMock = vi.fn().mockResolvedValue({ id: "run-123" }); - mockAuthStore.useAuthStore.getState.mockReturnValue({ + mockAuth.fetchAuthState.mockResolvedValue({ + status: "authenticated", + bootstrapComplete: true, cloudRegion: "us", projectId: 123, - client: { - createTaskRun: createTaskRunMock, - appendTaskRunLog: vi.fn(), - }, + availableProjectIds: [123], + availableOrgIds: [], + hasCodeAccess: true, + needsScopeReauth: false, + }); + mockBuildAuthenticatedClient.mockReturnValue({ + createTaskRun: createTaskRunMock, + appendTaskRunLog: vi.fn(), }); mockTrpcAgent.start.mutate.mockResolvedValue({ @@ -333,11 +367,17 @@ describe("SessionService", () => { it("creates error session when auth is missing", async () => { const service = getSessionService(); - mockAuthStore.useAuthStore.getState.mockReturnValue({ + mockAuth.fetchAuthState.mockResolvedValue({ + status: "anonymous", + bootstrapComplete: true, cloudRegion: null, projectId: null, - client: null, - } as unknown as ReturnType); + availableProjectIds: [], + availableOrgIds: [], + hasCodeAccess: null, + needsScopeReauth: false, + }); + mockBuildAuthenticatedClient.mockReturnValue(null); await service.connectToTask({ task: createMockTask(), @@ -414,13 +454,19 @@ describe("SessionService", () => { // Setup: create a task run to trigger subscription creation const createTaskRunMock = vi.fn().mockResolvedValue({ id: "run-456" }); - mockAuthStore.useAuthStore.getState.mockReturnValue({ + mockAuth.fetchAuthState.mockResolvedValue({ + status: "authenticated", + bootstrapComplete: true, cloudRegion: "us", projectId: 123, - client: { - createTaskRun: createTaskRunMock, - appendTaskRunLog: vi.fn(), - }, + availableProjectIds: [123], + availableOrgIds: [], + hasCodeAccess: true, + needsScopeReauth: false, + }); + mockBuildAuthenticatedClient.mockReturnValue({ + createTaskRun: createTaskRunMock, + appendTaskRunLog: vi.fn(), }); mockTrpcAgent.start.mutate.mockResolvedValue({ channel: "test-channel", diff --git a/apps/code/src/renderer/features/sessions/service/service.ts b/apps/code/src/renderer/features/sessions/service/service.ts index d6c2a20a4..5eb919f47 100644 --- a/apps/code/src/renderer/features/sessions/service/service.ts +++ b/apps/code/src/renderer/features/sessions/service/service.ts @@ -4,9 +4,10 @@ import type { SessionConfigOption, } from "@agentclientprotocol/sdk"; import { - setSessionResetCallback, - useAuthStore, -} from "@features/auth/stores/authStore"; + createAuthenticatedClient, + getAuthenticatedClient, +} from "@features/auth/hooks/authClient"; +import { fetchAuthState } from "@features/auth/hooks/authQueries"; import { useSessionAdapterStore } from "@features/sessions/stores/sessionAdapterStore"; import { getPersistedConfigOptions, @@ -65,7 +66,7 @@ export const PREVIEW_TASK_ID = "__preview__"; interface AuthCredentials { apiHost: string; projectId: number; - client: ReturnType["client"]; + client: Awaited>; } export interface ConnectParams { @@ -102,8 +103,6 @@ export function resetSessionService(): void { }); } -setSessionResetCallback(resetSessionService); - export class SessionService { private connectingTasks = new Map>(); private subscriptions = new Map< @@ -193,7 +192,7 @@ export class SessionService { const taskTitle = task.title || task.description || "Task"; try { - const auth = this.getAuthCredentials(); + const auth = await this.getAuthCredentials(); if (!auth) { log.error("Missing auth credentials"); const taskRunId = latestRun?.id ?? `error-${taskId}`; @@ -656,7 +655,7 @@ export class SessionService { await this.cleanupPreviewSession(); if (abort.signal.aborted) return; - const auth = this.getAuthCredentials(); + const auth = await this.getAuthCredentials(); if (!auth) { log.info("Skipping preview session - not authenticated"); return; @@ -1253,7 +1252,7 @@ export class SessionService { return { stopReason: "queued" }; } - const auth = this.getCloudCommandAuth(); + const auth = await this.getCloudCommandAuth(); if (!auth) { throw new Error("Authentication required for cloud commands"); } @@ -1384,7 +1383,7 @@ export class SessionService { session: AgentSession, promptText: string, ): Promise<{ stopReason: string }> { - const client = useAuthStore.getState().client; + const client = await getAuthenticatedClient(); if (!client) { throw new Error("Authentication required for cloud commands"); } @@ -1461,7 +1460,7 @@ export class SessionService { return false; } - const auth = this.getCloudCommandAuth(); + const auth = await this.getCloudCommandAuth(); if (!auth) { log.error("No auth for cloud cancel"); return false; @@ -1501,11 +1500,11 @@ export class SessionService { } } - private getCloudCommandAuth(): { + private async getCloudCommandAuth(): Promise<{ apiHost: string; teamId: number; - } | null { - const authState = useAuthStore.getState(); + } | null> { + const authState = await fetchAuthState(); if (!authState.cloudRegion || !authState.projectId) return null; return { apiHost: getCloudUrlFromRegion(authState.cloudRegion), @@ -1809,7 +1808,7 @@ export class SessionService { if (session?.initialPrompt?.length) { const { taskTitle, initialPrompt } = session; await this.teardownSession(session.taskRunId); - const auth = this.getAuthCredentials(); + const auth = await this.getAuthCredentials(); if (!auth) { throw new Error( "Unable to reach server. Please check your connection.", @@ -1865,7 +1864,7 @@ export class SessionService { } this.unsubscribeFromChannel(taskRunId); - const auth = this.getAuthCredentials(); + const auth = await this.getAuthCredentials(); if (!auth) { throw new Error("Unable to reach server. Please check your connection."); } @@ -1942,21 +1941,20 @@ export class SessionService { }); } - // Get auth for host info - const auth = useAuthStore.getState(); - if (!auth.projectId || !auth.cloudRegion) { - log.warn("No auth for cloud task watcher", { taskId }); - return () => {}; - } + void fetchAuthState() + .then((authState) => { + if (!authState.projectId || !authState.cloudRegion) { + log.warn("No auth for cloud task watcher", { taskId }); + return; + } - // Start main-process watcher - trpcClient.cloudTask.watch - .mutate({ - taskId, - runId, - apiHost: getCloudUrlFromRegion(auth.cloudRegion), - teamId: auth.projectId, - viewing, + return trpcClient.cloudTask.watch.mutate({ + taskId, + runId, + apiHost: getCloudUrlFromRegion(authState.cloudRegion), + teamId: authState.projectId, + viewing, + }); }) .catch((err: unknown) => log.warn("Failed to start cloud task watcher", { taskId, err }), @@ -2163,13 +2161,13 @@ export class SessionService { // --- Helper Methods --- - private getAuthCredentials(): AuthCredentials | null { - const authState = useAuthStore.getState(); + private async getAuthCredentials(): Promise { + const authState = await fetchAuthState(); const apiHost = authState.cloudRegion ? getCloudUrlFromRegion(authState.cloudRegion) : null; const projectId = authState.projectId; - const client = authState.client; + const client = createAuthenticatedClient(authState); if (!apiHost || !projectId || !client) return null; return { apiHost, projectId, client }; @@ -2288,23 +2286,13 @@ export class SessionService { // Don't update processedLineCount - it tracks S3 log lines, not local events sessionStoreSetters.appendEvents(session.taskRunId, [event]); - const auth = useAuthStore.getState(); - if (auth.client) { + const client = await getAuthenticatedClient(); + if (client) { try { - await auth.client.appendTaskRunLog(taskId, session.taskRunId, [ - storedEntry, - ]); + await client.appendTaskRunLog(taskId, session.taskRunId, [storedEntry]); } catch (error) { log.warn("Failed to persist event to logs", { error }); } } } } - -// Register callback when module loads (not during tests) -if (typeof window !== "undefined" && !import.meta.env?.VITEST) { - setSessionResetCallback(() => { - log.info("Auth triggered session reset"); - resetSessionService(); - }); -} diff --git a/apps/code/src/renderer/features/settings/components/sections/AccountSettings.tsx b/apps/code/src/renderer/features/settings/components/sections/AccountSettings.tsx index fc6783b53..5ce0c9956 100644 --- a/apps/code/src/renderer/features/settings/components/sections/AccountSettings.tsx +++ b/apps/code/src/renderer/features/settings/components/sections/AccountSettings.tsx @@ -1,26 +1,30 @@ -import { useAuthStore } from "@features/auth/stores/authStore"; +import { useOptionalAuthenticatedClient } from "@features/auth/hooks/authClient"; +import { useLogoutMutation } from "@features/auth/hooks/authMutations"; +import { + useAuthStateValue, + useCurrentUser, +} from "@features/auth/hooks/authQueries"; +import { useOnboardingStore } from "@features/onboarding/stores/onboardingStore"; import { SettingRow } from "@features/settings/components/SettingRow"; import { SignOut } from "@phosphor-icons/react"; import { Avatar, Badge, Button, Flex, Spinner, Text } from "@radix-ui/themes"; import { REGION_LABELS } from "@shared/constants/oauth"; -import { useQuery } from "@tanstack/react-query"; export function AccountSettings() { - const { client, isAuthenticated, selectedPlan, logout, cloudRegion } = - useAuthStore(); - - // Fetch current user from PostHog - const { data: user, isLoading } = useQuery({ - queryKey: ["currentUser"], - queryFn: async () => { - if (!client) return null; - return await client.getCurrentUser(); - }, - enabled: !!client && isAuthenticated, + const isAuthenticated = useAuthStateValue( + (state) => state.status === "authenticated", + ); + const cloudRegion = useAuthStateValue((state) => state.cloudRegion); + const selectedPlan = useOnboardingStore((state) => state.selectedPlan); + const logoutMutation = useLogoutMutation(); + const client = useOptionalAuthenticatedClient(); + const { data: user, isLoading } = useCurrentUser({ + client, + enabled: isAuthenticated, }); const handleLogout = () => { - logout(); + logoutMutation.mutate(); }; if (!isAuthenticated) { diff --git a/apps/code/src/renderer/features/settings/components/sections/AdvancedSettings.tsx b/apps/code/src/renderer/features/settings/components/sections/AdvancedSettings.tsx index b8a46809c..90dc47a07 100644 --- a/apps/code/src/renderer/features/settings/components/sections/AdvancedSettings.tsx +++ b/apps/code/src/renderer/features/settings/components/sections/AdvancedSettings.tsx @@ -1,4 +1,4 @@ -import { useAuthStore } from "@features/auth/stores/authStore"; +import { useOnboardingStore } from "@features/onboarding/stores/onboardingStore"; import { SettingRow } from "@features/settings/components/SettingRow"; import { useSettingsStore } from "@features/settings/stores/settingsStore"; import { useFeatureFlag } from "@hooks/useFeatureFlag"; @@ -22,9 +22,7 @@ export function AdvancedSettings() { diff --git a/apps/code/src/renderer/features/settings/components/sections/GeneralSettings.tsx b/apps/code/src/renderer/features/settings/components/sections/GeneralSettings.tsx index d85adf505..c8cbc2cc6 100644 --- a/apps/code/src/renderer/features/settings/components/sections/GeneralSettings.tsx +++ b/apps/code/src/renderer/features/settings/components/sections/GeneralSettings.tsx @@ -1,4 +1,4 @@ -import { useAuthStore } from "@features/auth/stores/authStore"; +import { useAuthStateValue } from "@features/auth/hooks/authQueries"; import { SettingRow } from "@features/settings/components/SettingRow"; import { type AutoConvertLongText, @@ -479,8 +479,8 @@ export function GeneralSettings() { } function HedgehogDescription() { - const cloudRegion = useAuthStore((s) => s.cloudRegion); - const projectId = useAuthStore((s) => s.projectId); + const cloudRegion = useAuthStateValue((state) => state.cloudRegion); + const projectId = useAuthStateValue((state) => state.projectId); const customizeUrl = cloudRegion && projectId diff --git a/apps/code/src/renderer/features/sidebar/components/ProjectSwitcher.tsx b/apps/code/src/renderer/features/sidebar/components/ProjectSwitcher.tsx index 1a79e44f1..7c0f12cfb 100644 --- a/apps/code/src/renderer/features/sidebar/components/ProjectSwitcher.tsx +++ b/apps/code/src/renderer/features/sidebar/components/ProjectSwitcher.tsx @@ -1,4 +1,8 @@ -import { useAuthStore } from "@features/auth/stores/authStore"; +import { + useLogoutMutation, + useSelectProjectMutation, +} from "@features/auth/hooks/authMutations"; +import { useAuthStateValue } from "@features/auth/hooks/authQueries"; import { Command } from "@features/command/components/Command"; import { useProjects } from "@features/projects/hooks/useProjects"; import { useSettingsDialogStore } from "@features/settings/stores/settingsDialogStore"; @@ -52,9 +56,9 @@ export function ProjectSwitcher() { setLearnMoreOpen(false); } }, [popoverOpen]); - const cloudRegion = useAuthStore((s) => s.cloudRegion); - const selectProject = useAuthStore((s) => s.selectProject); - const logout = useAuthStore((s) => s.logout); + const cloudRegion = useAuthStateValue((state) => state.cloudRegion); + const selectProjectMutation = useSelectProjectMutation(); + const logoutMutation = useLogoutMutation(); const { groupedProjects, currentProject, @@ -65,7 +69,7 @@ export function ProjectSwitcher() { const handleProjectSelect = (projectId: number) => { if (projectId !== currentProjectId) { - selectProject(projectId); + selectProjectMutation.mutate(projectId); } setPopoverOpen(false); setDialogOpen(false); @@ -112,7 +116,7 @@ export function ProjectSwitcher() { const handleLogout = () => { setPopoverOpen(false); - logout(); + logoutMutation.mutate(); }; return ( diff --git a/apps/code/src/renderer/features/task-detail/hooks/usePreviewSession.ts b/apps/code/src/renderer/features/task-detail/hooks/usePreviewSession.ts index 4752af8e3..90f07764c 100644 --- a/apps/code/src/renderer/features/task-detail/hooks/usePreviewSession.ts +++ b/apps/code/src/renderer/features/task-detail/hooks/usePreviewSession.ts @@ -1,5 +1,5 @@ import type { SessionConfigOption } from "@agentclientprotocol/sdk"; -import { useAuthStore } from "@features/auth/stores/authStore"; +import { useAuthStateValue } from "@features/auth/hooks/authQueries"; import { getSessionService, PREVIEW_TASK_ID, @@ -31,7 +31,7 @@ interface PreviewSessionResult { export function usePreviewSession( adapter: "claude" | "codex", ): PreviewSessionResult { - const projectId = useAuthStore((s) => s.projectId); + const projectId = useAuthStateValue((state) => state.projectId); useEffect(() => { if (!projectId) return; diff --git a/apps/code/src/renderer/features/task-detail/hooks/useTaskCreation.ts b/apps/code/src/renderer/features/task-detail/hooks/useTaskCreation.ts index 62c769a15..2b40713d0 100644 --- a/apps/code/src/renderer/features/task-detail/hooks/useTaskCreation.ts +++ b/apps/code/src/renderer/features/task-detail/hooks/useTaskCreation.ts @@ -1,4 +1,4 @@ -import { useAuthStore } from "@features/auth/stores/authStore"; +import { useAuthStateValue } from "@features/auth/hooks/authQueries"; import type { MessageEditorHandle } from "@features/message-editor/components/MessageEditor"; import { useTaskInputHistoryStore } from "@features/message-editor/stores/taskInputHistoryStore"; import { @@ -105,7 +105,9 @@ export function useTaskCreation({ }: UseTaskCreationOptions): UseTaskCreationReturn { const [isCreatingTask, setIsCreatingTask] = useState(false); const { navigateToTask } = useNavigationStore(); - const { isAuthenticated } = useAuthStore(); + const isAuthenticated = useAuthStateValue( + (state) => state.status === "authenticated", + ); const { invalidateTasks } = useCreateTask(); const { isOnline } = useConnectivity(); diff --git a/apps/code/src/renderer/features/task-detail/service/service.ts b/apps/code/src/renderer/features/task-detail/service/service.ts index a2b4283df..a754f7f9e 100644 --- a/apps/code/src/renderer/features/task-detail/service/service.ts +++ b/apps/code/src/renderer/features/task-detail/service/service.ts @@ -1,4 +1,4 @@ -import { useAuthStore } from "@features/auth/stores/authStore"; +import { getAuthenticatedClient } from "@features/auth/hooks/authClient"; import { useDraftStore } from "@features/message-editor/stores/draftStore"; import { useSettingsStore } from "@features/settings/stores/settingsStore"; import { workspaceApi } from "@features/workspace/hooks/useWorkspace"; @@ -48,7 +48,7 @@ export class TaskService { }; } - const posthogClient = useAuthStore.getState().client; + const posthogClient = await getAuthenticatedClient(); if (!posthogClient) { return { success: false, @@ -95,7 +95,7 @@ export class TaskService { ): Promise { log.info("Opening existing task", { taskId, taskRunId }); - const posthogClient = useAuthStore.getState().client; + const posthogClient = await getAuthenticatedClient(); if (!posthogClient) { return { success: false, diff --git a/apps/code/src/renderer/hooks/useAuthenticatedClient.ts b/apps/code/src/renderer/hooks/useAuthenticatedClient.ts index 777c686d7..b22d29a56 100644 --- a/apps/code/src/renderer/hooks/useAuthenticatedClient.ts +++ b/apps/code/src/renderer/hooks/useAuthenticatedClient.ts @@ -1,11 +1,5 @@ -import { useAuthStore } from "@features/auth/stores/authStore"; +import { useAuthenticatedClient as useClient } from "@features/auth/hooks/authClient"; export function useAuthenticatedClient() { - const client = useAuthStore((state) => state.client); - - if (!client) { - throw new Error("Not authenticated"); - } - - return client; + return useClient(); } diff --git a/apps/code/src/renderer/hooks/useAuthenticatedInfiniteQuery.ts b/apps/code/src/renderer/hooks/useAuthenticatedInfiniteQuery.ts index a1a486543..5bba77c02 100644 --- a/apps/code/src/renderer/hooks/useAuthenticatedInfiniteQuery.ts +++ b/apps/code/src/renderer/hooks/useAuthenticatedInfiniteQuery.ts @@ -1,4 +1,5 @@ -import { useAuthStore } from "@features/auth/stores/authStore"; +import { useOptionalAuthenticatedClient } from "@features/auth/hooks/authClient"; +import { AUTH_SCOPED_QUERY_META } from "@features/auth/hooks/authQueries"; import type { PostHogAPIClient } from "@renderer/api/posthogClient"; import type { QueryKey } from "@tanstack/react-query"; import { useInfiniteQuery } from "@tanstack/react-query"; @@ -33,7 +34,7 @@ export function useAuthenticatedInfiniteQuery< queryFn: AuthenticatedInfiniteQueryFn, options: UseAuthenticatedInfiniteQueryOptions, ) { - const client = useAuthStore((state) => state.client); + const client = useOptionalAuthenticatedClient(); return useInfiniteQuery({ queryKey, @@ -47,5 +48,6 @@ export function useAuthenticatedInfiniteQuery< refetchInterval: options.refetchInterval, refetchIntervalInBackground: options.refetchIntervalInBackground, staleTime: options.staleTime, + meta: AUTH_SCOPED_QUERY_META, }); } diff --git a/apps/code/src/renderer/hooks/useAuthenticatedMutation.ts b/apps/code/src/renderer/hooks/useAuthenticatedMutation.ts index 8a96b90c0..99d57e660 100644 --- a/apps/code/src/renderer/hooks/useAuthenticatedMutation.ts +++ b/apps/code/src/renderer/hooks/useAuthenticatedMutation.ts @@ -1,4 +1,4 @@ -import { useAuthStore } from "@features/auth/stores/authStore"; +import { useOptionalAuthenticatedClient } from "@features/auth/hooks/authClient"; import type { PostHogAPIClient } from "@renderer/api/posthogClient"; import type { UseMutationOptions, @@ -19,7 +19,7 @@ export function useAuthenticatedMutation< mutationFn: AuthenticatedMutationFn, options?: Omit, "mutationFn">, ): UseMutationResult { - const client = useAuthStore((state) => state.client); + const client = useOptionalAuthenticatedClient(); return useMutation({ mutationFn: async (variables: TVariables) => { diff --git a/apps/code/src/renderer/hooks/useAuthenticatedQuery.ts b/apps/code/src/renderer/hooks/useAuthenticatedQuery.ts index 8ea157749..2bb3636d3 100644 --- a/apps/code/src/renderer/hooks/useAuthenticatedQuery.ts +++ b/apps/code/src/renderer/hooks/useAuthenticatedQuery.ts @@ -1,4 +1,5 @@ -import { useAuthStore } from "@features/auth/stores/authStore"; +import { useOptionalAuthenticatedClient } from "@features/auth/hooks/authClient"; +import { AUTH_SCOPED_QUERY_META } from "@features/auth/hooks/authQueries"; import type { PostHogAPIClient } from "@renderer/api/posthogClient"; import type { QueryKey, @@ -21,7 +22,8 @@ export function useAuthenticatedQuery< "queryKey" | "queryFn" >, ): UseQueryResult { - const client = useAuthStore((state) => state.client); + const client = useOptionalAuthenticatedClient(); + const { meta, ...restOptions } = options ?? {}; return useQuery({ queryKey, @@ -30,7 +32,12 @@ export function useAuthenticatedQuery< return await queryFn(client); }, enabled: - !!client && (options?.enabled !== undefined ? options.enabled : true), - ...options, + !!client && + (restOptions.enabled !== undefined ? restOptions.enabled : true), + meta: { + ...AUTH_SCOPED_QUERY_META, + ...meta, + }, + ...restOptions, }); } diff --git a/apps/code/src/renderer/hooks/useOrganizations.ts b/apps/code/src/renderer/hooks/useOrganizations.ts index 92dffc15f..f5463bbcc 100644 --- a/apps/code/src/renderer/hooks/useOrganizations.ts +++ b/apps/code/src/renderer/hooks/useOrganizations.ts @@ -1,7 +1,8 @@ -import { useAuthStore } from "@features/auth/stores/authStore"; +import { useOptionalAuthenticatedClient } from "@features/auth/hooks/authClient"; +import { useCurrentUser } from "@features/auth/hooks/authQueries"; +import { useOnboardingStore } from "@features/onboarding/stores/onboardingStore"; import { useAuthenticatedQuery } from "@hooks/useAuthenticatedQuery"; import type { PostHogAPIClient } from "@renderer/api/posthogClient"; -import { useQueryClient } from "@tanstack/react-query"; import { useMemo } from "react"; export interface OrgWithBilling { @@ -52,8 +53,9 @@ async function fetchOrgsWithBilling( } export function useOrganizations() { - const selectedOrgId = useAuthStore((s) => s.selectedOrgId); - const queryClient = useQueryClient(); + const selectedOrgId = useOnboardingStore((state) => state.selectedOrgId); + const client = useOptionalAuthenticatedClient(); + const { data: currentUser } = useCurrentUser({ client }); const { data: orgsWithBilling, @@ -70,9 +72,6 @@ export function useOrganizations() { if (!orgsWithBilling?.length) return null; // Default to the user's currently active org in PostHog - const currentUser = queryClient.getQueryData<{ - organization?: { id: string }; - }>(["currentUser"]); const userCurrentOrgId = currentUser?.organization?.id; if ( userCurrentOrgId && @@ -85,7 +84,7 @@ export function useOrganizations() { (org) => org.has_active_subscription, ); return (withBilling ?? orgsWithBilling[0]).id; - }, [selectedOrgId, orgsWithBilling, queryClient]); + }, [currentUser?.organization?.id, orgsWithBilling, selectedOrgId]); const sortedOrgs = useMemo(() => { return [...(orgsWithBilling ?? [])].sort((a, b) => diff --git a/apps/code/src/renderer/hooks/useProjectQuery.ts b/apps/code/src/renderer/hooks/useProjectQuery.ts index 11090b776..a0137df38 100644 --- a/apps/code/src/renderer/hooks/useProjectQuery.ts +++ b/apps/code/src/renderer/hooks/useProjectQuery.ts @@ -1,8 +1,8 @@ -import { useAuthStore } from "@features/auth/stores/authStore"; +import { useAuthStateValue } from "@features/auth/hooks/authQueries"; import { useAuthenticatedQuery } from "./useAuthenticatedQuery"; export function useProjectQuery() { - const projectId = useAuthStore((state) => state.projectId); + const projectId = useAuthStateValue((state) => state.projectId); return useAuthenticatedQuery( ["project", projectId], diff --git a/apps/code/src/renderer/hooks/useTaskDeepLink.ts b/apps/code/src/renderer/hooks/useTaskDeepLink.ts index e97aeaa47..73c0b101d 100644 --- a/apps/code/src/renderer/hooks/useTaskDeepLink.ts +++ b/apps/code/src/renderer/hooks/useTaskDeepLink.ts @@ -1,4 +1,4 @@ -import { useAuthStore } from "@features/auth/stores/authStore"; +import { useAuthStateValue } from "@features/auth/hooks/authQueries"; import { useTaskViewed } from "@features/sidebar/hooks/useTaskViewed"; import type { TaskService } from "@features/task-detail/service/service"; import { get } from "@renderer/di/container"; @@ -30,7 +30,9 @@ export function useTaskDeepLink() { const navigateToTask = useNavigationStore((state) => state.navigateToTask); const { markAsViewed } = useTaskViewed(); const queryClient = useQueryClient(); - const isAuthenticated = useAuthStore((state) => state.isAuthenticated); + const isAuthenticated = useAuthStateValue( + (state) => state.status === "authenticated", + ); const hasFetchedPending = useRef(false); const handleOpenTask = useCallback( diff --git a/apps/code/src/renderer/utils/generateTitle.ts b/apps/code/src/renderer/utils/generateTitle.ts index dd3f33dcf..181b311ef 100644 --- a/apps/code/src/renderer/utils/generateTitle.ts +++ b/apps/code/src/renderer/utils/generateTitle.ts @@ -1,4 +1,4 @@ -import { useAuthStore } from "@features/auth/stores/authStore"; +import { fetchAuthState } from "@features/auth/hooks/authQueries"; import { trpcClient } from "@renderer/trpc"; import { logger } from "@utils/logger"; @@ -41,8 +41,8 @@ Never wrap the title in quotes.`; export async function generateTitle(content: string): Promise { try { - const authState = useAuthStore.getState(); - if (!authState.isAuthenticated) return null; + const authState = await fetchAuthState(); + if (authState.status !== "authenticated") return null; const result = await trpcClient.llmGateway.prompt.mutate({ system: SYSTEM_PROMPT,