Chapter 4: Performance Optimization: State Management (Zustand) and Caching for Complex Chat History
When your AI chat evolves beyond the simple "one-question-one-answer" model from the previous chapter to remembering context and supporting multiple "chat sessions," the basic useState approach will quickly fall apart.
If you cram all conversation data, loading states, and history into a single component, every state change will trigger massive re-renders across the entire page, potentially causing your typewriter effect to stutter.
This chapter will guide you through using Zustand, React's current most popular lightweight global state management library, to build a high-performance chat history management system.
Why Not Redux or Context API?
Traditionally, we had two options for cross-component state:
- Context API: Prone to severe "render hell." When a Provider's value changes, all child components blindly re-render, creating performance nightmares for scenarios like streaming that update state every 10ms.
- Redux: Too heavy. Requires writing tons of Action, Reducer, and Dispatch boilerplate code, drastically slowing down Vibe Coding development speed.
Zustand excels with: extreme simplicity, no Context Provider dependency, and precise "Selectors" to prevent unnecessary renders.
Step 1: Define the Store
First install Zustand: npm install zustand
Create a src/store/useChatStore.ts file. This Store will be our chat room's brain:
import { create } from 'zustand'
export type Message = {
id: string
role: 'user' | 'assistant'
content: string
}
export type ChatSession = {
id: string
title: string
messages: Message[]
}
interface ChatState {
sessions: ChatSession[]
activeSessionId: string | null
// Actions
setActiveSession: (id: string) => void
addMessage: (sessionId: string, msg: Message) => void
updateLastMessage: (sessionId: string, chunk: string) => void
createNewSession: () => string
}
export const useChatStore = create<ChatState>((set, get) => ({
// Initial state
sessions: [],
activeSessionId: null,
// Switch active chat window
setActiveSession: (id) => set({ activeSessionId: id }),
// Create new session
createNewSession: () => {
const newId = crypto.randomUUID()
set((state) => ({
// Prepend new session to array
sessions: [{ id: newId, title: 'New Exploration', messages: [] }, ...state.sessions],
activeSessionId: newId
}))
return newId
},
// Add complete new message
addMessage: (sessionId, msg) => set((state) => ({
sessions: state.sessions.map(session =>
session.id === sessionId
? { ...session, messages: [...session.messages, msg] }
: session
)
})),
// Update last message (optimized for streaming typewriter effect!)
updateLastMessage: (sessionId, chunk) => set((state) => ({
sessions: state.sessions.map(session => {
if (session.id !== sessionId) return session;
const lastMsg = session.messages[session.messages.length - 1];
if (!lastMsg || lastMsg.role !== 'assistant') return session;
// Clone array and append new chunk to last message
const newMessages = [...session.messages];
newMessages[newMessages.length - 1] = {
...lastMsg,
content: lastMsg.content + chunk
};
return { ...session, messages: newMessages };
})
})),
}))
Step 2: Precise State Usage in Components
With our Store ready, let's see how to use it in UI. Zustand's power lies in "taking only what you need," preventing unnecessary re-renders when unrelated state changes.
For example, here's a sidebar component showing "history list":
"use client"
import { useChatStore } from "@/store/useChatStore"
import { PlusCircle, MessageSquare } from "lucide-react"
export function Sidebar() {
// Precisely extract needed state/actions
// When messages update for typewriter effect, Sidebar won't re-render!
const sessions = useChatStore(state => state.sessions)
const activeSessionId = useChatStore(state => state.activeSessionId)
const setActiveSession = useChatStore(state => state.setActiveSession)
const createNewSession = useChatStore(state => state.createNewSession)
return (
<div className="w-72 bg-gray-900 text-white h-screen p-4 flex flex-col gap-4 shadow-2xl z-10">
<button
onClick={createNewSession}
className="flex items-center gap-2 bg-white/10 hover:bg-white/20 text-white font-bold py-3 px-4 rounded-xl transition-all border border-white/5"
>
<PlusCircle className="w-5 h-5" />
New Chat
</button>
<div className="flex-1 overflow-y-auto mt-2 space-y-2 pr-2 custom-scrollbar">
{sessions.map(s => (
<div
key={s.id}
onClick={() => setActiveSession(s.id)}
className={`flex items-center gap-3 p-3 rounded-xl cursor-pointer transition-all ${
activeSessionId === s.id
? 'bg-blue-600 font-bold shadow-lg shadow-blue-500/30'
: 'hover:bg-white/5 text-white/70'
}`}
>
<MessageSquare className="w-4 h-4 opacity-50" />
<span className="truncate flex-1">{s.title}</span>
</div>
))}
</div>
</div>
)
}
The brilliance here: When streaming typewriter effects trigger updateLastMessage in the chat window, the Sidebar won't re-render at all!
Add Persistence with localStorage
Zustand's powerful middleware makes this trivial. Imagine a user having a valuable AI conversation who accidentally hits F5 - they'd be furious if history disappeared. Just wrap our store with persist:
import { create } from 'zustand'
import { persist } from 'zustand/middleware'
export const useChatStore = create<ChatState>()(
persist(
(set, get) => ({
// ... all previous state and logic ...
}),
{
// localStorage key name
name: 'vibe-tutor-chat-storage',
// Optional: Filter out activeSessionId if needed
partialize: (state) => ({ sessions: state.sessions }),
}
)
)
Just five extra lines! Now, no matter how users refresh or close tabs, their "entire chat history" and "sessions" will magically reappear.
With Zustand's global state management and Persist caching, we've built an ultra-smooth, locally-persistent AI chat core without heavy database operations - dramatically reducing server costs while delivering lightning-fast UX.
Zustand Chat Store
import { create } from 'zustand';
import { persist, createJSONStorage } from 'zustand/middleware';
interface ChatSession {
id: string;
title: string;
messages: Message[];
createdAt: number;
}
interface ChatStore {
sessions: ChatSession[];
activeSessionId: string | null;
createSession: () => string;
deleteSession: (id: string) => void;
setActiveSession: (id: string) => void;
addMessage: (sessionId: string, message: Message) => void;
getActiveSession: () => ChatSession | undefined;
}
export const useChatStore = create<ChatStore>()(
persist(
(set, get) => ({
sessions: [],
activeSessionId: null,
createSession: () => {
const id = crypto.randomUUID();
const session: ChatSession = {
id,
title: 'New Chat',
messages: [],
createdAt: Date.now(),
};
set((state) => ({
sessions: [...state.sessions, session],
activeSessionId: id,
}));
return id;
},
deleteSession: (id) => {
set((state) => ({
sessions: state.sessions.filter((s) => s.id !== id),
activeSessionId:
state.activeSessionId === id ? null : state.activeSessionId,
}));
},
setActiveSession: (id) => set({ activeSessionId: id }),
addMessage: (sessionId, message) => {
set((state) => ({
sessions: state.sessions.map((s) =>
s.id === sessionId
? { ...s, messages: [...s.messages, message] }
: s
),
}));
},
getActiveSession: () => {
const { sessions, activeSessionId } = get();
return sessions.find((s) => s.id === activeSessionId);
},
}),
{
name: 'chat-sessions',
storage: createJSONStorage(() => localStorage),
}
)
);
Using the Store
export default function ChatApp() {
const {
sessions,
activeSessionId,
createSession,
setActiveSession,
addMessage,
} = useChatStore();
const activeSession = useChatStore((s) =>
s.sessions.find((sess) => sess.id === s.activeSessionId)
);
const handleNewChat = () => {
createSession();
};
const handleSend = async (text: string) => {
if (!activeSessionId) return;
const userMsg: Message = { role: 'user', content: text, id: Date.now().toString() };
addMessage(activeSessionId, userMsg);
// Call AI API
const response = await fetch('/api/chat', {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({
messages: [...(activeSession?.messages || []), userMsg],
sessionId: activeSessionId,
}),
});
// Handle streaming response...
};
return (
<div className="flex h-screen">
{/* Sidebar */}
<div className="w-64 bg-gray-100 p-4">
<button onClick={handleNewChat} className="btn-primary">
+ New Chat
</button>
<div className="mt-4 space-y-2">
{sessions.map((s) => (
<div
key={s.id}
className={`cursor-pointer p-2 rounded ${
s.id === activeSessionId ? 'bg-blue-100' : 'hover:bg-gray-200'
}`}
onClick={() => setActiveSession(s.id)}
>
{s.title}
</div>
))}
</div>
</div>
{/* Chat Area */}
<div className="flex-1 flex flex-col">
{activeSession ? (
<>
<MessageList messages={activeSession.messages} />
<ChatInput onSend={handleSend} />
</>
) : (
<EmptyState onCreateChat={handleNewChat} />
)}
</div>
</div>
);
}
Summary
Zustand with persist middleware provides fast, persistent chat state management. Users can create multiple sessions, switch between them, and their chat history survives page refreshes โ all without a database.
Key takeaways:
- Zustand: lightweight global state management |
- Persist middleware: saves state to localStorage automatically |
- Chat sessions: create, delete, switch with Zustand selectors |
- Active session tracking: single
activeSessionIdstate | - Messages stored per session in the Zustand store |
- No database needed for history โ localStorage is sufficient |
- Sidebar UI for session list, main area for active chat |
You've completed this course! You can now build real-time AI chat apps.