State Management for Agent UIs: React Context, Zustand, and Server State with TanStack Query

The State Management Challenge in Agent UIs

Agent interfaces manage three distinct categories of state: UI state (sidebar open, selected conversation, theme), client state (message drafts, optimistic messages, local preferences), and server state (conversation history, agent configuration, user profile). Using a single approach for all three creates unnecessary complexity. The modern pattern separates these concerns: React Context for UI state, Zustand for client state, and TanStack Query for server state.

React Context for UI State

UI state is lightweight, changes infrequently, and affects the visual layout. React Context handles this well without any external library.

import {
  createContext,
  useContext,
  useState,
  ReactNode,
} from "react";

interface UIState {
  sidebarOpen: boolean;
  activeConversationId: string | null;
  theme: "light" | "dark";
}

interface UIActions {
  toggleSidebar: () => void;
  setActiveConversation: (id: string | null) => void;
  setTheme: (theme: "light" | "dark") => void;
}

const UIContext = createContext<
  (UIState & UIActions) | null
>(null);

export function UIProvider({ children }: { children: ReactNode }) {
  const [state, setState] = useState<UIState>({
    sidebarOpen: true,
    activeConversationId: null,
    theme: "light",
  });

  const actions: UIActions = {
    toggleSidebar: () =>
      setState((s) => ({ ...s, sidebarOpen: !s.sidebarOpen })),
    setActiveConversation: (id) =>
      setState((s) => ({ ...s, activeConversationId: id })),
    setTheme: (theme) =>
      setState((s) => ({ ...s, theme })),
  };

  return (
    <UIContext.Provider value={{ ...state, ...actions }}>
      {children}
    </UIContext.Provider>
  );
}

export function useUI() {
  const ctx = useContext(UIContext);
  if (!ctx) throw new Error("useUI must be inside UIProvider");
  return ctx;
}

Context is the right tool here because UI state changes are infrequent (toggling a sidebar, switching conversations) and the provider sits near the top of the tree. The common criticism that Context causes excessive re-renders applies when state changes rapidly, which UI state does not.

Zustand for Client-Side Message State

Message state changes frequently (every streamed token, every optimistic update) and is complex (multiple messages, status transitions, ordering). Zustand provides a lightweight store that avoids the re-render issues of Context.

import { create } from "zustand";

interface ChatMessage {
  id: string;
  role: "user" | "assistant";
  content: string;
  status: "optimistic" | "streaming" | "complete" | "error";
  conversationId: string;
}

interface MessageStore {
  messages: Map<string, ChatMessage[]>;
  addMessage: (convId: string, msg: ChatMessage) => void;
  appendToken: (convId: string, msgId: string, token: string) => void;
  updateStatus: (
    convId: string,
    msgId: string,
    status: ChatMessage["status"]
  ) => void;
  getConversationMessages: (convId: string) => ChatMessage[];
}

export const useMessageStore = create<MessageStore>(
  (set, get) => ({
    messages: new Map(),

    addMessage: (convId, msg) =>
      set((state) => {
        const newMap = new Map(state.messages);
        const existing = newMap.get(convId) || [];
        newMap.set(convId, [...existing, msg]);
        return { messages: newMap };
      }),

    appendToken: (convId, msgId, token) =>
      set((state) => {
        const newMap = new Map(state.messages);
        const msgs = (newMap.get(convId) || []).map((m) =>
          m.id === msgId
            ? { ...m, content: m.content + token }
            : m
        );
        newMap.set(convId, msgs);
        return { messages: newMap };
      }),

    updateStatus: (convId, msgId, status) =>
      set((state) => {
        const newMap = new Map(state.messages);
        const msgs = (newMap.get(convId) || []).map((m) =>
          m.id === msgId ? { ...m, status } : m
        );
        newMap.set(convId, msgs);
        return { messages: newMap };
      }),

    getConversationMessages: (convId) =>
      get().messages.get(convId) || [],
  })
);

Zustand shines here because components can subscribe to slices of the store. A component that only reads messages for one conversation will not re-render when messages in another conversation change.

Selectors for Performance

Use Zustand selectors to minimize re-renders. Components that only need to know whether a conversation has unread messages should not re-render when message content changes.

function useConversationMessages(convId: string) {
  return useMessageStore(
    (state) => state.messages.get(convId) || []
  );
}

function useIsStreaming(convId: string) {
  return useMessageStore((state) => {
    const msgs = state.messages.get(convId) || [];
    return msgs.some((m) => m.status === "streaming");
  });
}

function useMessageCount(convId: string) {
  return useMessageStore(
    (state) => (state.messages.get(convId) || []).length
  );
}

Each selector creates a subscription that only triggers re-renders when its return value changes. This is far more efficient than subscribing to the entire store.

TanStack Query for Server State

Server state — conversation history, agent configuration, user profile — lives on the backend and should be fetched, cached, and synchronized. TanStack Query handles this with automatic caching, background refetching, and stale-while-revalidate patterns.

import {
  useQuery,
  useMutation,
  useQueryClient,
} from "@tanstack/react-query";

interface Conversation {
  id: string;
  title: string;
  createdAt: string;
  messageCount: number;
}

function useConversations() {
  return useQuery<Conversation[]>({
    queryKey: ["conversations"],
    queryFn: () =>
      fetch("/api/conversations").then((r) => r.json()),
    staleTime: 60_000,
  });
}

function useCreateConversation() {
  const queryClient = useQueryClient();

  return useMutation({
    mutationFn: (title: string) =>
      fetch("/api/conversations", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ title }),
      }).then((r) => r.json()),

    onSuccess: () => {
      queryClient.invalidateQueries({
        queryKey: ["conversations"],
      });
    },
  });
}

The staleTime: 60_000 tells TanStack Query that the conversation list is fresh for 60 seconds. During that window, navigating away and back will show cached data instantly without a loading spinner.

Hydrating Zustand from Server State

When the user opens a conversation, fetch the history from the server and populate the Zustand store. This bridges server and client state.

function useLoadConversation(convId: string) {
  const addMessage = useMessageStore((s) => s.addMessage);

  return useQuery({
    queryKey: ["conversation-history", convId],
    queryFn: async () => {
      const res = await fetch(`/api/conversations/${convId}/messages`);
      const messages: ChatMessage[] = await res.json();
      messages.forEach((msg) => addMessage(convId, msg));
      return messages;
    },
    staleTime: Infinity, // Only fetch once per session
  });
}

Setting staleTime: Infinity ensures the history is fetched once when the conversation opens and not re-fetched on window focus or component remount. New messages are added through the Zustand store directly from the streaming hook.

When to Use Which Pattern

The decision tree is straightforward. If the state affects layout or visual mode and changes infrequently, use React Context. If the state is client-only, changes frequently, and multiple components need it, use Zustand. If the state comes from the server and needs caching, refetching, and synchronization, use TanStack Query.

FAQ

Can I use just Zustand for everything instead of three separate tools?

You can, but you lose the automatic caching and background refetching of TanStack Query. You would need to manually implement stale-while-revalidate, deduplication of in-flight requests, and cache invalidation. For simple apps with few API calls, an all-Zustand approach works. For production agent interfaces with many endpoints, the combination is worth the added dependency.

How do I persist Zustand state across page refreshes?

Zustand provides a persist middleware that serializes state to localStorage or sessionStorage. Wrap your store creation with persist and specify a storage key. Be selective about what you persist — message content should come from the server on refresh, but user preferences like theme and sidebar state are good candidates for local persistence.

How do I share state between the chat component and a separate analytics panel?

Both components can subscribe to the same Zustand store using different selectors. The chat component subscribes to messages for the active conversation. The analytics panel subscribes to aggregated metrics derived from the same store. Since Zustand stores are global singletons, both components automatically share the same data without prop drilling or context nesting.

---

#StateManagement #Zustand #TanStackQuery #ReactContext #TypeScript #AgenticAI #LearnAI #AIEngineering