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](#architecture-overview)
2. [Zustand Stores](#zustand-stores)
3. [TanStack Query Setup](#tanstack-query-setup)
4. [Custom Hooks](#custom-hooks)
5. [State Synchronization](#state-synchronization)
6. [Best Practices](#best-practices)
7. [Performance Optimization](#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`

```typescript
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**:

```tsx
// 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`

```typescript
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`

```typescript
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`

```typescript
'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`

```typescript
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**:

```tsx
// 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`

```typescript
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`

```typescript
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**:

```tsx
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`

```typescript
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`

```typescript
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.

```typescript
// 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

```typescript
// 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**:
```typescript
// 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**:
```typescript
// Inconsistent keys
queryKey: ['tokens']
queryKey: ['mcpToken', tokenId]
queryKey: ['audit', tokenId]
```

---

### 2. Use Selectors for Performance

**Good** (only re-renders when `user` changes):
```tsx
const user = useAuthStore((state) => state.user);
```

**Bad** (re-renders when ANY auth state changes):
```tsx
const { user } = useAuthStore();
```

**Even Better** (dedicated selector):
```tsx
// stores/useAuthStore.ts
export const useUser = () => useAuthStore((state) => state.user);

// In component
const user = useUser();
```

---

### 3. Prefetch Critical Data

```typescript
// 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

```tsx
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:

```typescript
// ✅ 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:

```tsx
// 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:

```typescript
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

```typescript
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**:
```tsx
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:

```typescript
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**:
```tsx
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)