refactor(frontend): migrate selected org ID to Zustand (#12153)

frontend/__tests__/stores/selected-organization-store.test.ts

@@ -0,0 +1,51 @@
+import { act, renderHook } from "@testing-library/react";
+import { describe, expect, it } from "vitest";
+import { useSelectedOrganizationStore } from "#/stores/selected-organization-store";
+
+describe("useSelectedOrganizationStore", () => {
+  it("should have null as initial orgId", () => {
+    const { result } = renderHook(() => useSelectedOrganizationStore());
+    expect(result.current.orgId).toBeNull();
+  });
+
+  it("should update orgId when setOrgId is called", () => {
+    const { result } = renderHook(() => useSelectedOrganizationStore());
+
+    act(() => {
+      result.current.setOrgId("org-123");
+    });
+
+    expect(result.current.orgId).toBe("org-123");
+  });
+
+  it("should allow setting orgId to null", () => {
+    const { result } = renderHook(() => useSelectedOrganizationStore());
+
+    act(() => {
+      result.current.setOrgId("org-123");
+    });
+
+    expect(result.current.orgId).toBe("org-123");
+
+    act(() => {
+      result.current.setOrgId(null);
+    });
+
+    expect(result.current.orgId).toBeNull();
+  });
+
+  it("should share state across multiple hook instances", () => {
+    const { result: result1 } = renderHook(() =>
+      useSelectedOrganizationStore(),
+    );
+    const { result: result2 } = renderHook(() =>
+      useSelectedOrganizationStore(),
+    );
+
+    act(() => {
+      result1.current.setOrgId("shared-org");
+    });
+
+    expect(result2.current.orgId).toBe("shared-org");
+  });
+});

frontend/src/context/use-selected-organization.ts

@@ -1,12 +1,16 @@
-import { useSelectedOrgId } from "#/hooks/query/use-selected-organization-id";
-import { useUpdateSelectedOrganizationId } from "#/hooks/mutation/use-update-selected-organization-id";
+import { useRevalidator } from "react-router";
+import { useSelectedOrganizationStore } from "#/stores/selected-organization-store";
 
 export const useSelectedOrganizationId = () => {
-  const { data: orgId } = useSelectedOrgId();
-  const updateState = useUpdateSelectedOrganizationId();
+  const revalidator = useRevalidator();
+  const { orgId, setOrgId: setOrgIdStore } = useSelectedOrganizationStore();
 
-  return {
-    orgId,
-    setOrgId: (newValue: string | null) => updateState.mutate(newValue),
+  const setOrgId = (newOrgId: string | null) => {
+    setOrgIdStore(newOrgId);
+    // Revalidate route to ensure the latest orgId is used.
+    // This is useful for redirecting the user away from admin-only org pages.
+    revalidator.revalidate();
   };
+
+  return { orgId, setOrgId };
 };

frontend/src/hooks/mutation/use-update-selected-organization-id.ts

@@ -1,38 +0,0 @@
-import { useMutation, useQueryClient } from "@tanstack/react-query";
-import { SELECTED_ORGANIZATION_QUERY_KEY } from "#/hooks/query/use-selected-organization-id";
-
-export function useUpdateSelectedOrganizationId() {
-  const queryClient = useQueryClient();
-
-  return useMutation({
-    mutationFn: async (newValue: string | null) => {
-      queryClient.setQueryData([SELECTED_ORGANIZATION_QUERY_KEY], newValue);
-      return newValue;
-    },
-    onMutate: async (newValue) => {
-      await queryClient.cancelQueries({
-        queryKey: [SELECTED_ORGANIZATION_QUERY_KEY],
-      });
-
-      // Snapshot the previous value
-      const previousValue = queryClient.getQueryData([
-        SELECTED_ORGANIZATION_QUERY_KEY,
-      ]);
-      queryClient.setQueryData([SELECTED_ORGANIZATION_QUERY_KEY], newValue);
-
-      return { previousValue };
-    },
-    onError: (_, __, context) => {
-      queryClient.setQueryData(
-        [SELECTED_ORGANIZATION_QUERY_KEY],
-        context?.previousValue,
-      );
-    },
-    // Always refetch after error or success
-    onSettled: () => {
-      queryClient.invalidateQueries({
-        queryKey: [SELECTED_ORGANIZATION_QUERY_KEY],
-      });
-    },
-  });
-}

frontend/src/hooks/query/use-selected-organization-id.ts

@@ -1,23 +0,0 @@
-import { useQuery, useQueryClient } from "@tanstack/react-query";
-import { useRevalidator } from "react-router";
-
-export const SELECTED_ORGANIZATION_QUERY_KEY = "selected_organization";
-
-export function useSelectedOrgId() {
-  const queryClient = useQueryClient();
-  const revalidator = useRevalidator();
-
-  return useQuery({
-    queryKey: [SELECTED_ORGANIZATION_QUERY_KEY],
-    initialData: null as string | null,
-    queryFn: () => {
-      const storedOrgId = queryClient.getQueryData<string>([
-        SELECTED_ORGANIZATION_QUERY_KEY,
-      ]);
-      // Revalidate route clientLoader to ensure the latest orgId is used.
-      // This is useful for redirecting the user away from admin-only org pages.
-      revalidator.revalidate();
-      return storedOrgId || null; // Return null if no org ID is set
-    },
-  });
-}

frontend/src/routes/manage-org.tsx

@@ -12,10 +12,8 @@ import { BrandButton } from "#/components/features/settings/brand-button";
 import { useMe } from "#/hooks/query/use-me";
 import { useConfig } from "#/hooks/query/use-config";
 import { rolePermissions } from "#/utils/org/permissions";
-import {
-  getSelectedOrgFromQueryClient,
-  getMeFromQueryClient,
-} from "#/utils/query-client-getters";
+import { getSelectedOrganizationIdFromStore } from "#/stores/selected-organization-store";
+import { getMeFromQueryClient } from "#/utils/query-client-getters";
 import { queryClient } from "#/query-client-config";
 import { I18nKey } from "#/i18n/declaration";
 import { amountIsValid } from "#/utils/amount-is-valid";
@@ -212,7 +210,7 @@ function AddCreditsModal({ onClose }: AddCreditsModalProps) {
 }
 
 export const clientLoader = async () => {
-  const selectedOrgId = getSelectedOrgFromQueryClient();
+  const selectedOrgId = getSelectedOrganizationIdFromStore();
   let me = getMeFromQueryClient(selectedOrgId);
 
   if (!me && selectedOrgId) {

frontend/src/routes/manage-organization-members.tsx

@@ -14,14 +14,12 @@ import { BrandButton } from "#/components/features/settings/brand-button";
 import { rolePermissions } from "#/utils/org/permissions";
 import { organizationService } from "#/api/organization-service/organization-service.api";
 import { queryClient } from "#/query-client-config";
-import {
-  getSelectedOrgFromQueryClient,
-  getMeFromQueryClient,
-} from "#/utils/query-client-getters";
+import { getSelectedOrganizationIdFromStore } from "#/stores/selected-organization-store";
+import { getMeFromQueryClient } from "#/utils/query-client-getters";
 import { I18nKey } from "#/i18n/declaration";
 
 export const clientLoader = async () => {
-  const selectedOrgId = getSelectedOrgFromQueryClient();
+  const selectedOrgId = getSelectedOrganizationIdFromStore();
   let me = getMeFromQueryClient(selectedOrgId);
 
   if (!me && selectedOrgId) {

frontend/src/stores/selected-organization-store.ts

@@ -0,0 +1,30 @@
+import { create } from "zustand";
+import { devtools } from "zustand/middleware";
+
+interface SelectedOrganizationState {
+  orgId: string | null;
+}
+
+interface SelectedOrganizationActions {
+  setOrgId: (orgId: string | null) => void;
+}
+
+type SelectedOrganizationStore = SelectedOrganizationState &
+  SelectedOrganizationActions;
+
+const initialState: SelectedOrganizationState = {
+  orgId: null,
+};
+
+export const useSelectedOrganizationStore = create<SelectedOrganizationStore>()(
+  devtools(
+    (set) => ({
+      ...initialState,
+      setOrgId: (orgId) => set({ orgId }),
+    }),
+    { name: "SelectedOrganizationStore" },
+  ),
+);
+
+export const getSelectedOrganizationIdFromStore = (): string | null =>
+  useSelectedOrganizationStore.getState().orgId;

frontend/src/utils/query-client-getters.ts

@@ -1,9 +1,5 @@
 import { queryClient } from "#/query-client-config";
 import { OrganizationMember } from "#/types/org";
-import { SELECTED_ORGANIZATION_QUERY_KEY } from "#/hooks/query/use-selected-organization-id";
 
-export const getMeFromQueryClient = (orgId: string | undefined) =>
+export const getMeFromQueryClient = (orgId: string | null) =>
   queryClient.getQueryData<OrganizationMember>(["organizations", orgId, "me"]);
-
-export const getSelectedOrgFromQueryClient = () =>
-  queryClient.getQueryData<string>([SELECTED_ORGANIZATION_QUERY_KEY]);