ColaFlow/docs/frontend/state-management-guide.md

State Management Guide - Zustand + TanStack Query

Document Overview

This guide explains how to effectively combine Zustand (client state) and TanStack Query v5 (server state) in ColaFlow's frontend architecture, with complete implementation examples.

Key Principle: Zustand manages authentication and UI state, TanStack Query manages all API data.

---

Table of Contents

1. Architecture Overview
2. Zustand Stores
3. TanStack Query Setup
4. Custom Hooks
5. State Synchronization
6. Best Practices
7. Performance Optimization

---

Architecture Overview

State Classification

State Type	Technology	Persistence	Example
Authentication	Zustand	Session Storage	User info, tenant info, access token
UI State	Zustand	Local Storage	Sidebar open/closed, theme, language
Server Data	TanStack Query	Memory (cache)	Projects, issues, MCP tokens
Form State	React Hook Form	Component local	Login form, create token form

Data Flow

User Action
    ↓
Component (UI Layer)
    ↓
    ├──→ Zustand Store (for auth/UI state)
    │    ├── setUser()
    │    ├── clearAuth()
    │    └── toggleSidebar()
    │
    └──→ TanStack Query Hook (for API data)
         ├── useLogin() → authService.login()
         ├── useMcpTokens() → mcpService.listTokens()
         └── useCreateMcpToken() → mcpService.createToken()
              ↓
         API Client (Axios)
              ↓
         Backend API (.NET 9)

---

Zustand Stores

1. Auth Store

File: stores/useAuthStore.ts

import { create } from 'zustand';
import { persist, createJSONStorage } from 'zustand/middleware';

// Types
export interface User {
  id: string;
  email: string;
  fullName: string;
  role: 'Admin' | 'Member' | 'Viewer';
  authProvider: 'Local' | 'AzureAD' | 'Google' | 'Okta';
  externalUserId?: string;
  avatarUrl?: string;
}

export interface Tenant {
  id: string;
  slug: string;
  name: string;
  plan: 'Free' | 'Starter' | 'Professional' | 'Enterprise';
  status: 'Active' | 'Trial' | 'Suspended' | 'Cancelled';
  ssoEnabled: boolean;
  ssoProvider?: 'AzureAD' | 'Google' | 'Okta' | 'GenericSaml';
}

interface AuthState {
  // State
  user: User | null;
  tenant: Tenant | null;
  accessToken: string | null;
  isAuthenticated: boolean;
  isLoading: boolean;

  // Actions
  setUser: (user: User, tenant: Tenant, accessToken: string) => void;
  setAccessToken: (token: string) => void;
  clearAuth: () => void;
  setLoading: (loading: boolean) => void;
}

export const useAuthStore = create<AuthState>()(
  persist(
    (set) => ({
      // Initial state
      user: null,
      tenant: null,
      accessToken: null,
      isAuthenticated: false,
      isLoading: false,

      // Actions
      setUser: (user, tenant, accessToken) => {
        set({
          user,
          tenant,
          accessToken,
          isAuthenticated: true,
          isLoading: false,
        });
      },

      setAccessToken: (accessToken) => {
        set({ accessToken });
      },

      clearAuth: () => {
        set({
          user: null,
          tenant: null,
          accessToken: null,
          isAuthenticated: false,
          isLoading: false,
        });
      },

      setLoading: (isLoading) => {
        set({ isLoading });
      },
    }),
    {
      name: 'colaflow-auth', // Key in storage
      storage: createJSONStorage(() => sessionStorage), // Use sessionStorage
      partialize: (state) => ({
        // Only persist user and tenant info, NOT the token
        user: state.user,
        tenant: state.tenant,
        // Do NOT persist accessToken (security)
      }),
    }
  )
);

// Selectors (for performance optimization)
export const useUser = () => useAuthStore((state) => state.user);
export const useTenant = () => useAuthStore((state) => state.tenant);
export const useIsAuthenticated = () => useAuthStore((state) => state.isAuthenticated);

Usage:

// In a component
import { useAuthStore, useUser, useTenant } from '@/stores/useAuthStore';

function UserProfile() {
  // Method 1: Select specific fields (recommended for performance)
  const user = useUser();
  const tenant = useTenant();

  // Method 2: Select multiple fields
  const { user, tenant } = useAuthStore((state) => ({
    user: state.user,
    tenant: state.tenant,
  }));

  // Method 3: Access actions
  const clearAuth = useAuthStore((state) => state.clearAuth);

  return (
    <div>
      <h1>{user?.fullName}</h1>
      <p>{tenant?.name}</p>
      <button onClick={clearAuth}>Logout</button>
    </div>
  );
}

---

2. UI Store

File: stores/useUiStore.ts

import { create } from 'zustand';
import { persist, createJSONStorage } from 'zustand/middleware';

interface UiState {
  // Sidebar
  sidebarOpen: boolean;
  toggleSidebar: () => void;
  setSidebarOpen: (open: boolean) => void;

  // Theme
  theme: 'light' | 'dark' | 'system';
  setTheme: (theme: 'light' | 'dark' | 'system') => void;

  // Language
  language: 'en' | 'zh-CN';
  setLanguage: (language: 'en' | 'zh-CN') => void;

  // Notifications
  notificationCount: number;
  setNotificationCount: (count: number) => void;
}

export const useUiStore = create<UiState>()(
  persist(
    (set) => ({
      // Sidebar
      sidebarOpen: true,
      toggleSidebar: () => set((state) => ({ sidebarOpen: !state.sidebarOpen })),
      setSidebarOpen: (open) => set({ sidebarOpen: open }),

      // Theme
      theme: 'system',
      setTheme: (theme) => set({ theme }),

      // Language
      language: 'en',
      setLanguage: (language) => set({ language }),

      // Notifications
      notificationCount: 0,
      setNotificationCount: (notificationCount) => set({ notificationCount }),
    }),
    {
      name: 'colaflow-ui',
      storage: createJSONStorage(() => localStorage), // Persist UI preferences
    }
  )
);

// Selectors
export const useSidebarOpen = () => useUiStore((state) => state.sidebarOpen);
export const useTheme = () => useUiStore((state) => state.theme);

---

TanStack Query Setup

Query Client Configuration

File: lib/query-client.ts

import { QueryClient, DefaultOptions } from '@tanstack/react-query';

const queryConfig: DefaultOptions = {
  queries: {
    // Stale time: Data is considered fresh for 5 minutes
    staleTime: 1000 * 60 * 5,

    // Garbage collection time: Unused data is garbage collected after 10 minutes
    gcTime: 1000 * 60 * 10, // Renamed from cacheTime in v5

    // Refetch on window focus
    refetchOnWindowFocus: false,

    // Retry failed requests
    retry: 1,

    // Retry delay
    retryDelay: (attemptIndex) => Math.min(1000 * 2 ** attemptIndex, 30000),
  },
  mutations: {
    // Retry failed mutations
    retry: 0,
  },
};

export const queryClient = new QueryClient({
  defaultOptions: queryConfig,
});

Provider Setup:

File: app/layout.tsx

'use client';

import { QueryClientProvider } from '@tanstack/react-query';
import { ReactQueryDevtools } from '@tanstack/react-query-devtools';
import { queryClient } from '@/lib/query-client';

export default function RootLayout({ children }: { children: React.ReactNode }) {
  return (
    <html lang="en">
      <body>
        <QueryClientProvider client={queryClient}>
          {children}

          {/* DevTools (only in development) */}
          {process.env.NODE_ENV === 'development' && (
            <ReactQueryDevtools initialIsOpen={false} />
          )}
        </QueryClientProvider>
      </body>
    </html>
  );
}

---

Custom Hooks

Authentication Hooks

useLogin

File: hooks/auth/useLogin.ts

import { useMutation, useQueryClient } from '@tanstack/react-query';
import { authService } from '@/services/auth.service';
import { useAuthStore } from '@/stores/useAuthStore';
import { useRouter } from 'next/navigation';
import { toast } from 'sonner';

export function useLogin() {
  const setUser = useAuthStore((state) => state.setUser);
  const router = useRouter();
  const queryClient = useQueryClient();

  return useMutation({
    mutationFn: authService.login,

    onSuccess: (data) => {
      // Store user, tenant, and token in Zustand
      setUser(data.user, data.tenant, data.accessToken);

      // Show success message
      toast.success(`Welcome back, ${data.user.fullName}!`);

      // Prefetch user data (optional)
      queryClient.prefetchQuery({
        queryKey: ['user', data.user.id],
        queryFn: () => authService.getCurrentUser(),
      });

      // Redirect to dashboard or original page
      const redirectUrl = sessionStorage.getItem('redirect_after_login') || '/dashboard';
      sessionStorage.removeItem('redirect_after_login');
      router.push(redirectUrl);
    },

    onError: (error: any) => {
      // Show error message
      const message = error.response?.data?.message || 'Login failed. Please try again.';
      toast.error(message);
    },
  });
}

Usage:

// In a login component
import { useLogin } from '@/hooks/auth/useLogin';

function LoginForm() {
  const login = useLogin();

  const handleSubmit = (data: { email: string; password: string }) => {
    login.mutate(data);
  };

  return (
    <form onSubmit={handleSubmit}>
      {/* Form fields */}
      <button type="submit" disabled={login.isPending}>
        {login.isPending ? 'Logging in...' : 'Login'}
      </button>
    </form>
  );
}

---

useLogout

File: hooks/auth/useLogout.ts

import { useMutation, useQueryClient } from '@tanstack/react-query';
import { authService } from '@/services/auth.service';
import { useAuthStore } from '@/stores/useAuthStore';
import { useRouter } from 'next/navigation';
import { toast } from 'sonner';

export function useLogout() {
  const clearAuth = useAuthStore((state) => state.clearAuth);
  const router = useRouter();
  const queryClient = useQueryClient();

  return useMutation({
    mutationFn: authService.logout,

    onSuccess: () => {
      // Clear auth state
      clearAuth();

      // Clear all query cache
      queryClient.clear();

      // Show success message
      toast.success('Logged out successfully');

      // Redirect to login page
      router.push('/login');
    },

    onError: () => {
      // Even if logout fails on backend, clear client state
      clearAuth();
      queryClient.clear();
      router.push('/login');
    },
  });
}

---

Tenant Hooks

useCheckSlug

File: hooks/tenants/useCheckSlug.ts

import { useQuery } from '@tanstack/react-query';
import { tenantService } from '@/services/tenant.service';
import { useMemo } from 'react';

export function useCheckSlug(slug: string, enabled: boolean = true) {
  // Trim and lowercase
  const normalizedSlug = slug.trim().toLowerCase();

  return useQuery({
    queryKey: ['check-slug', normalizedSlug],
    queryFn: () => tenantService.checkSlugAvailability(normalizedSlug),
    enabled: enabled && normalizedSlug.length >= 3,
    staleTime: 5000, // Cache for 5 seconds
    retry: false, // Don't retry slug checks
  });
}

Usage:

import { useCheckSlug } from '@/hooks/tenants/useCheckSlug';

function TenantSlugInput() {
  const [slug, setSlug] = useState('');
  const { data, isLoading, error } = useCheckSlug(slug, slug.length >= 3);

  return (
    <div>
      <input
        value={slug}
        onChange={(e) => setSlug(e.target.value)}
        placeholder="your-company"
      />

      {isLoading && <span>Checking...</span>}

      {data?.available === true && (
        <span className="text-green-600">✓ Available</span>
      )}

      {data?.available === false && (
        <span className="text-red-600">✗ Taken</span>
      )}
    </div>
  );
}

---

useSsoConfig

File: hooks/tenants/useSsoConfig.ts

import { useQuery, useMutation, useQueryClient } from '@tanstack/react-query';
import { tenantService } from '@/services/tenant.service';
import { useTenant } from '@/stores/useAuthStore';
import { toast } from 'sonner';

// Get SSO config
export function useSsoConfig() {
  const tenant = useTenant();

  return useQuery({
    queryKey: ['sso-config', tenant?.id],
    queryFn: () => tenantService.getSsoConfig(tenant!.id),
    enabled: !!tenant,
    staleTime: 1000 * 60 * 10, // 10 minutes
  });
}

// Update SSO config
export function useUpdateSsoConfig() {
  const tenant = useTenant();
  const queryClient = useQueryClient();

  return useMutation({
    mutationFn: (config: any) => tenantService.updateSsoConfig(tenant!.id, config),

    onSuccess: () => {
      // Invalidate SSO config query
      queryClient.invalidateQueries({ queryKey: ['sso-config', tenant?.id] });

      toast.success('SSO configuration updated successfully');
    },

    onError: (error: any) => {
      const message = error.response?.data?.message || 'Failed to update SSO configuration';
      toast.error(message);
    },
  });
}

// Test SSO connection
export function useTestSsoConnection() {
  const tenant = useTenant();

  return useMutation({
    mutationFn: () => tenantService.testSsoConnection(tenant!.id),

    onSuccess: (data) => {
      if (data.success) {
        toast.success('SSO connection successful!');
      } else {
        toast.error(`Connection failed: ${data.message}`);
      }
    },

    onError: () => {
      toast.error('Failed to test SSO connection');
    },
  });
}

---

MCP Token Hooks

useMcpTokens

File: hooks/mcp/useMcpTokens.ts

import { useQuery, useMutation, useQueryClient } from '@tanstack/react-query';
import { mcpService } from '@/services/mcp.service';
import { toast } from 'sonner';

// List all tokens
export function useMcpTokens() {
  return useQuery({
    queryKey: ['mcp-tokens'],
    queryFn: mcpService.listTokens,
    staleTime: 1000 * 60 * 5, // 5 minutes
  });
}

// Get single token details
export function useMcpToken(tokenId: string) {
  return useQuery({
    queryKey: ['mcp-token', tokenId],
    queryFn: () => mcpService.getToken(tokenId),
    enabled: !!tokenId,
  });
}

// Create new token
export function useCreateMcpToken() {
  const queryClient = useQueryClient();

  return useMutation({
    mutationFn: mcpService.createToken,

    onSuccess: () => {
      // Invalidate tokens list
      queryClient.invalidateQueries({ queryKey: ['mcp-tokens'] });

      toast.success('MCP Token created successfully');
    },

    onError: (error: any) => {
      const message = error.response?.data?.message || 'Failed to create token';
      toast.error(message);
    },
  });
}

// Revoke token
export function useRevokeMcpToken() {
  const queryClient = useQueryClient();

  return useMutation({
    mutationFn: ({ tokenId, reason }: { tokenId: string; reason?: string }) =>
      mcpService.revokeToken(tokenId, reason),

    onSuccess: () => {
      // Invalidate tokens list
      queryClient.invalidateQueries({ queryKey: ['mcp-tokens'] });

      toast.success('Token revoked successfully');
    },

    onError: () => {
      toast.error('Failed to revoke token');
    },
  });
}

// Get audit logs
export function useMcpAuditLogs(tokenId: string, filters?: any) {
  return useQuery({
    queryKey: ['mcp-audit-logs', tokenId, filters],
    queryFn: () => mcpService.getAuditLogs(tokenId, filters),
    enabled: !!tokenId,
    staleTime: 1000 * 60, // 1 minute
  });
}

---

State Synchronization

Keeping AuthStore and Query Cache in Sync

Problem: When user logs out, TanStack Query cache may still contain user-specific data.

Solution: Clear query cache on logout.

// hooks/auth/useLogout.ts
export function useLogout() {
  const clearAuth = useAuthStore((state) => state.clearAuth);
  const queryClient = useQueryClient();

  return useMutation({
    mutationFn: authService.logout,

    onSuccess: () => {
      // 1. Clear auth state
      clearAuth();

      // 2. Clear ALL query cache
      queryClient.clear();

      // Or selectively clear specific queries:
      // queryClient.removeQueries({ queryKey: ['projects'] });
      // queryClient.removeQueries({ queryKey: ['issues'] });

      // 3. Redirect
      router.push('/login');
    },
  });
}

---

Optimistic Updates

Use Case: Immediately update UI before API call completes.

Example: Revoke MCP token

// hooks/mcp/useMcpTokens.ts
export function useRevokeMcpToken() {
  const queryClient = useQueryClient();

  return useMutation({
    mutationFn: ({ tokenId, reason }: { tokenId: string; reason?: string }) =>
      mcpService.revokeToken(tokenId, reason),

    // Optimistic update
    onMutate: async ({ tokenId }) => {
      // Cancel any outgoing refetches
      await queryClient.cancelQueries({ queryKey: ['mcp-tokens'] });

      // Snapshot the previous value
      const previousTokens = queryClient.getQueryData(['mcp-tokens']);

      // Optimistically update to the new value
      queryClient.setQueryData(['mcp-tokens'], (old: any) => ({
        ...old,
        tokens: old?.tokens.map((token: any) =>
          token.id === tokenId ? { ...token, status: 'Revoked' } : token
        ),
      }));

      // Return context with previous value
      return { previousTokens };
    },

    // On error, roll back to previous value
    onError: (_err, _variables, context) => {
      if (context?.previousTokens) {
        queryClient.setQueryData(['mcp-tokens'], context.previousTokens);
      }

      toast.error('Failed to revoke token');
    },

    // Always refetch after success or error
    onSettled: () => {
      queryClient.invalidateQueries({ queryKey: ['mcp-tokens'] });
    },
  });
}

---

Best Practices

1. Use Query Keys Consistently

Good:

// All related queries use the same base key
queryKey: ['mcp-tokens']
queryKey: ['mcp-token', tokenId]
queryKey: ['mcp-audit-logs', tokenId]

// Invalidate all MCP token queries:
queryClient.invalidateQueries({ queryKey: ['mcp-tokens'] });
queryClient.invalidateQueries({ queryKey: ['mcp-token'] });

Bad:

// Inconsistent keys
queryKey: ['tokens']
queryKey: ['mcpToken', tokenId]
queryKey: ['audit', tokenId]

---

2. Use Selectors for Performance

Good (only re-renders when user changes):

const user = useAuthStore((state) => state.user);

Bad (re-renders when ANY auth state changes):

const { user } = useAuthStore();

Even Better (dedicated selector):

// stores/useAuthStore.ts
export const useUser = () => useAuthStore((state) => state.user);

// In component
const user = useUser();

---

3. Prefetch Critical Data

// app/(dashboard)/layout.tsx
export default function DashboardLayout({ children }: { children: React.ReactNode }) {
  const queryClient = useQueryClient();
  const user = useUser();

  useEffect(() => {
    if (user) {
      // Prefetch projects
      queryClient.prefetchQuery({
        queryKey: ['projects'],
        queryFn: projectService.getAll,
      });

      // Prefetch issues
      queryClient.prefetchQuery({
        queryKey: ['issues'],
        queryFn: issueService.getAll,
      });
    }
  }, [user, queryClient]);

  return <div>{children}</div>;
}

---

4. Handle Loading and Error States

function ProjectList() {
  const { data, isLoading, error } = useProjects();

  if (isLoading) {
    return <LoadingSpinner />;
  }

  if (error) {
    return (
      <ErrorMessage
        message="Failed to load projects"
        retry={() => refetch()}
      />
    );
  }

  if (!data || data.length === 0) {
    return <EmptyState message="No projects yet" />;
  }

  return (
    <div>
      {data.map((project) => (
        <ProjectCard key={project.id} project={project} />
      ))}
    </div>
  );
}

---

5. Use Mutations for Write Operations

Always use mutations (not queries) for creating, updating, deleting data:

// ✅ Good
const createProject = useMutation({
  mutationFn: projectService.create,
});

// ❌ Bad (don't use query for mutations)
const createProject = useQuery({
  queryKey: ['create-project'],
  queryFn: projectService.create,
});

---

Performance Optimization

1. Query Deduplication

TanStack Query automatically deduplicates identical queries:

// Both components will share the same query
function ComponentA() {
  const { data } = useMcpTokens();
  // ...
}

function ComponentB() {
  const { data } = useMcpTokens(); // Same query, shared data
  // ...
}

---

2. Background Refetching

Keep data fresh without blocking UI:

export function useMcpTokens() {
  return useQuery({
    queryKey: ['mcp-tokens'],
    queryFn: mcpService.listTokens,
    staleTime: 1000 * 60 * 5, // 5 minutes

    // Refetch in background when data is stale
    refetchOnMount: true,
    refetchOnReconnect: true,
    refetchOnWindowFocus: false, // Disable for better UX
  });
}

---

3. Pagination

export function useMcpAuditLogs(tokenId: string, page: number, limit: number) {
  return useQuery({
    queryKey: ['mcp-audit-logs', tokenId, page, limit],
    queryFn: () => mcpService.getAuditLogs(tokenId, { page, limit }),
    enabled: !!tokenId,

    // Keep previous data while fetching next page
    placeholderData: (previousData) => previousData,
  });
}

Usage:

function AuditLogTable({ tokenId }: { tokenId: string }) {
  const [page, setPage] = useState(1);
  const { data, isLoading, isPlaceholderData } = useMcpAuditLogs(tokenId, page, 50);

  return (
    <div>
      <Table data={data?.logs} loading={isPlaceholderData} />

      <Pagination
        currentPage={page}
        totalPages={data?.pagination.totalPages}
        onPageChange={setPage}
      />
    </div>
  );
}

---

4. Infinite Queries

For infinite scroll:

import { useInfiniteQuery } from '@tanstack/react-query';

export function useInfiniteAuditLogs(tokenId: string) {
  return useInfiniteQuery({
    queryKey: ['mcp-audit-logs-infinite', tokenId],
    queryFn: ({ pageParam = 1 }) =>
      mcpService.getAuditLogs(tokenId, { page: pageParam, limit: 50 }),
    getNextPageParam: (lastPage) => {
      const { page, totalPages } = lastPage.pagination;
      return page < totalPages ? page + 1 : undefined;
    },
    enabled: !!tokenId,
  });
}

Usage:

function InfiniteAuditLogList({ tokenId }: { tokenId: string }) {
  const {
    data,
    fetchNextPage,
    hasNextPage,
    isFetchingNextPage,
  } = useInfiniteAuditLogs(tokenId);

  return (
    <div>
      {data?.pages.map((page) =>
        page.logs.map((log) => <AuditLogItem key={log.id} log={log} />)
      )}

      {hasNextPage && (
        <button onClick={() => fetchNextPage()} disabled={isFetchingNextPage}>
          {isFetchingNextPage ? 'Loading...' : 'Load More'}
        </button>
      )}
    </div>
  );
}

---

Conclusion

This state management guide demonstrates how to effectively combine Zustand and TanStack Query in ColaFlow's frontend:

Key Takeaways:

• ✅ Use Zustand for authentication and UI state
• ✅ Use TanStack Query for all server data
• ✅ Clear query cache on logout
• ✅ Use optimistic updates for instant UX
• ✅ Use selectors for performance
• ✅ Prefetch critical data
• ✅ Handle loading/error states gracefully
• ✅ Use mutations for write operations

Next Document: component-library.md (Reusable components catalog)