Repository navigation
10. State Management
The AI Client application manages all global client state through a single Zustand store exposed as useAppStore. This store is composed from two domain-specific slices — ChatSlice and EntitySlice — each responsible for a distinct area of the application's data model. Zustand v5 is used alongside React 19.
The store operates with no middleware. There is no Immer integration, no devtools configuration, no persistence layer, and no subscribeWithSelector middleware. All state is ephemeral: it is rehydrated from PostgreSQL on each page load and discarded when the browser tab closes. This design deliberately avoids the complexity of client-side caching and stale-data reconciliation.
Only client components consume the store. React Server Components never access useAppStore; they fetch data directly from the database via Server Actions and pass it down as props or trigger client-side hydration after authentication.
flowchart TD
STORE["useAppStore"]
CS["ChatSlice"]
ES["EntitySlice"]
COMP["Components"]
HOOKS["Hooks"]
CS -->|spread into| STORE
ES -->|spread into| STORE
STORE -->|selectors| COMP
STORE -->|selectors| HOOKS
The store is created using Zustand's slice pattern. Both slices are defined as StateCreator functions that receive set, get, and the combined store type. Their return values are spread together into a single create call, producing a unified AppState type defined as the intersection of both slices.
This approach keeps each slice self-contained and independently testable whilst still allowing cross-slice reads via get().
| Aspect | Detail |
|---|---|
| Creation | create<AppState>()(...) |
| Pattern | Slice pattern (spread composition) |
| Middleware | None |
| Persistence | None (ephemeral, rehydrated from DB) |
| Devtools | Not configured |
The absence of middleware is a conscious decision. Since state is always backed by PostgreSQL, there is no value in persisting it to localStorage or sessionStorage. Devtools are omitted in favour of keeping the bundle lean; debugging relies on React DevTools and network inspection of Server Action calls.
The ChatSlice manages conversation state as a normalised hash map keyed by chat ID: chats: Record<string, Chat>. Each Chat object contains a flat Record<string, Message> that forms an adjacency-list tree through parentId and childrenIds references on each message. This tree structure supports branching conversations — edits and regenerations create sibling nodes rather than overwriting history.
The active branch is tracked by currentLeafId, which points to the deepest node in the currently visible conversation path. Walking the tree from root to currentLeafId reconstructs the active thread.
flowchart TD
ROOT["Root Message"]
A["Message A"]
B["Message B — active branch"]
C["Message C — sibling branch"]
D["Message D"]
ROOT --> A
A --> B
A --> C
B --> D
style B fill:#2563eb,color:#fff
style D fill:#2563eb,color:#fff
Messages are never stored as a flat ordered list. The tree model enables non-linear conversation history: users can branch at any point, switch between branches, and navigate back to earlier forks. The childrenIds array on each message preserves sibling order.
The slice exposes 16 actions split into two categories:
-
Pure local mutations — modify store state only, with no server interaction. Examples include
addMessage,deleteMessage,setCurrentLeaf,deleteChat,updateMessageAttachments, andupsertChat. These are used when the caller manages persistence separately. -
Persistent actions — most are suffixed with
Db(e.g.deleteMessageDb,renameChatDb). These wrap the corresponding local mutation inwithOptimisticUpdate, snapshotting state before the mutation, applying it immediately, and rolling back on failure. The exception iscreateChatDb, which awaits the server action first and then adds the returned chat to the store (persist-then-add, not optimistic).setKnowledgebasealso persists viawithOptimisticUpdatedespite lacking theDbsuffix.
sequenceDiagram
participant U as User
participant HOOK as useStreamResponse
participant S as Store
participant SA as persistMessage
participant API as /api/chat
participant LS as Local State
participant DB as Database
U->>HOOK: streamResponse(content, parentId, ...)
HOOK->>S: addMessage (user message + sets currentLeafId)
HOOK->>SA: persistMessage (server action)
SA->>DB: INSERT user message
HOOK->>API: POST /api/chat (SSE stream begins)
loop Streaming tokens
API-->>LS: text / reasoning / tool-call events
Note over LS: Accumulated in React useState
end
API-->>DB: INSERT assistant message (server-side)
API-->>LS: "done" event (assistantMessageId + metadata)
LS->>S: addMessage (assistant message)
Note over S: currentLeafId updated atomically
During streaming, the assistant's response is not stored in the Zustand store. It accumulates in React local state (useState) within the streaming hook. Only when the "done" SSE event arrives is the complete assistant message — including tool-call metadata, model identifier, and reasoning tokens — committed to the store and persisted to the database. This prevents partial or corrupted messages from entering the global state.
The EntitySlice manages eight collections of application entities as flat arrays. Unlike the chat slice, it exposes load-only actions — there are no in-slice CRUD mutation methods.
| Entity | State Key | Loader Action |
|---|---|---|
| Projects | projects |
loadProjects() |
| Assistants | assistants |
loadAssistants() |
| Prompts | prompts |
loadPrompts() |
| MCP Servers | mcpServers |
loadMcpServers() |
| Public MCP Servers | publicMcpServers |
loadPublicMcpServers() |
| Knowledgebases | knowledgebases |
loadKnowledgebases() |
| Transform Agents | transformAgents |
loadTransformAgents() |
| MCP Prompts | mcpPrompts |
loadMcpPrompts() |
Seven of the eight loaders are generated by createEntityLoader, a generic factory function that encapsulates the standard fetch → map → set pattern. Given a Server Action fetcher and a row-to-entity mapper, it returns a loader function that fetches the full list from the database, transforms each row, and replaces the corresponding array in the store. The exception is loadMcpPrompts, which calls discoverAllPrompts() directly and sets the result without using the factory, as MCP prompts are ephemeral discovery results rather than database rows.
Entity mutations follow a load-after-mutate strategy:
- The component calls a Server Action directly (e.g.
createProject,deleteAssistant). - The Server Action persists the change to PostgreSQL.
- The component then calls the corresponding loader action (e.g.
loadProjects()). - The loader fetches the full list, maps rows, and replaces the store array entirely.
This pattern trades UI responsiveness for simplicity — there are no optimistic updates for entity operations, and the user may experience a brief loading state while the round-trip completes. There is also no cross-slice cascading: deleting a project does not automatically unlink associated chats in the store.
sequenceDiagram
participant C as Component
participant SA as Server Action
participant DB as Database
participant L as Loader
participant S as Store
C->>SA: createProject(data)
SA->>DB: INSERT INTO project
SA-->>C: Success
C->>L: loadProjects()
L->>DB: SELECT * FROM project
DB-->>L: Rows
L->>L: Map rows to store types
L->>S: Replace projects array
Chat operations use a snapshot-and-rollback optimistic update pattern implemented through withOptimisticUpdate. The lifecycle proceeds as follows:
- Snapshot the entire store state before any mutation.
- Apply the local mutation immediately, providing instant UI feedback.
- Await the corresponding Server Action.
- On success — the optimistic state is already correct; no further action needed.
- On failure — restore the full snapshot, reverting all state changes.
flowchart TD
START["Action Triggered"] --> SNAP["Snapshot Entire Store"]
SNAP --> MUTATE["Apply Local Mutation"]
MUTATE --> AWAIT["Await Server Action"]
AWAIT -->|Success| DONE["Keep Optimistic State"]
AWAIT -->|Failure| ROLLBACK["Restore Snapshot"]
ROLLBACK --> DONE2["State Reverted"]
The rollback is coarse-grained: it restores the entire store state, not just the affected entity. This simplifies the implementation — there is no need to track individual field changes — but means that concurrent mutations from other parts of the UI may also be reverted on failure.
Entity operations use no optimistic updates. The component waits for the Server Action to complete, then triggers a full reload. This produces a brief loading state but guarantees that the store always reflects the database truth. The full array is replaced on each load, avoiding merge conflicts entirely.
The streaming pipeline carefully separates concerns between the Zustand store and React local state:
-
User message — added to the store immediately via
addMessage(a pure local mutation). Persistence happens separately via apersistMessage()server action call — there is no snapshot-rollback for the user message path. -
Streaming text — accumulated in a React
useStatevariable within the streaming hook. The store is not updated during streaming. -
Tool calls — tracked as a
ToolCallState[]array in React local state. On stream completion, tool-call data is serialised into the message's metadata JSON. -
Completion — the assistant message is persisted to the database server-side in the API route before the
"done"SSE event is emitted. When the client receives the"done"event (which carries the persisted message ID and metadata), the full assistant message is committed to the Zustand store. -
Artifacts — derived via
useMemoinside the artifact panel hook. Artifacts are never stored in Zustand; they are computed from message metadata and Mermaid fenced code blocks on every render.
flowchart LR
UM["User Message"] -->|addMessage| STORE["Zustand Store"]
UM -->|persistMessage| DB["Database"]
ST["Stream Text"] -->|accumulate| RLS["React Local State"]
TC["Tool Calls"] -->|track| RLS
API["API Route"] -->|persist assistant msg| DB
API -->|done event| RLS
RLS -->|addMessage| STORE
STORE -->|useMemo| ART["Artifacts — computed, never stored"]
Initial state loading is centralised in the useResourceHydration hook, which orchestrates the loading of entity collections after authentication. Deduplication works per hook instance via a hydrationAttempted ref that prevents React StrictMode double-fire. Cross-component deduplication relies on checking whether data is already present in the store; if two components mount simultaneously with an empty store, both may start fetching. The hook tracks a loadingResources set as local state per instance, not a global flag.
Chat hydration follows an independent path — useResourceHydration does not trigger loadChats(). Chat loading is initiated separately (e.g. by the sidebar or by page-level components) and uses a two-pass strategy:
- First pass — creates all message objects in the store without linking children.
-
Second pass — populates
childrenIdson each message by scanningparentIdreferences.
This two-pass approach avoids ordering dependencies: messages can arrive from the database in any order without causing dangling childrenIds references.
sequenceDiagram
participant AUTH as Auth Layer
participant HYD as useResourceHydration
participant EL as Entity Loaders
participant CL as Chat Loader
participant DB as Database
participant S as Store
AUTH->>HYD: Authentication confirmed
HYD->>EL: Trigger entity loaders (parallel)
EL->>DB: Fetch projects, assistants, prompts, etc.
DB-->>EL: Entity rows
EL->>S: Replace entity arrays
Note over HYD: Chat loading is independent
AUTH->>CL: Sidebar / page triggers loadChats()
CL->>DB: Fetch chats + messages
DB-->>CL: Chat/message rows
CL->>CL: Pass 1 — create messages
CL->>CL: Pass 2 — link childrenIds
CL->>S: Set chats record
There is no stale data detection. Every mount triggers a full re-fetch from the database. This is acceptable given the single-user nature of the application and the relatively small dataset sizes involved.
Zustand re-renders a subscribing component only when its selector return value changes, using referential equality by default. The application uses several selector approaches with varying re-render characteristics.
| Pattern | Usage | Re-render Risk |
|---|---|---|
| Atomic selector | ~90% of usages | Low — primitive values or stable references |
useShallow |
~3 places | Low — shallow-compared objects/arrays |
| Full store subscription | ~5 places | High — re-renders on any state change |
| Computed/derived selector | Rare | Varies — depends on reference stability |
The dominant pattern is inline atomic selectors: useAppStore((state) => state.X). These extract a single primitive value or a stable reference, resulting in minimal unnecessary re-renders. Approximately 90% of all store subscriptions follow this pattern.
useShallow from zustand/react/shallow is used in only three places, where components need multiple values from the store compared by shallow equality rather than referential equality.
Five call sites subscribe to the full store without a selector — an anti-pattern that causes re-renders on every state change. These are candidates for refactoring.
The application does not use auto-generated selectors or a memoised selector library such as Reselect. Derived state is computed inline via useMemo within hooks and components.
Eight mapper functions form a clean boundary between database row types and store types. These mappers are applied during the hydration and loading phase, ensuring that components never encounter raw database representations.
The mappers handle four categories of transformation:
-
Null-to-undefined coercion — database columns that may be
nullare converted toundefinedfor cleaner TypeScript consumption. -
Null-to-empty-string — text fields that should never be
undefinedin the UI receive empty string defaults. -
ISO string-to-Date conversion — raw timestamp strings from the database are parsed into JavaScript
Dateobjects. - JSON string parsing — serialised JSON columns (message metadata, transform agent steps) are parsed into typed objects.
flowchart LR
ROW["DB Row — nulls, ISO strings, JSON text"]
MAP["Mapper Function"]
TYPE["Store Type — defaults, Dates, parsed objects"]
COMP["Component"]
ROW --> MAP --> TYPE --> COMP
This mapping layer ensures that no component logic needs to handle null coercion, date parsing, or JSON deserialisation. All such concerns are resolved once at the data boundary.
| Characteristic | Implementation |
|---|---|
| Store count | Single global store |
| Composition | Slice pattern (2 slices) |
| Middleware | None |
| Persistence | Ephemeral (DB-backed) |
| Chat state model | Normalised Record with tree adjacency list |
| Entity state model | Flat arrays, load-only |
| Optimistic updates | Chat operations only (snapshot rollback) |
| Streaming state | React local state, committed on completion |
| Derived state | Computed in hooks via useMemo
|
| Hydration | Post-auth, deduplicated, full re-fetch |
| Selector convention | Inline atomic selectors |
The architecture prioritises simplicity over sophistication. By avoiding middleware, persistence layers, and normalised entity caches, the state management remains straightforward to reason about. The trade-offs — brief loading states for entity mutations, coarse-grained rollbacks, and full re-fetches on mount — are acceptable given the application's single-user context and moderate data volumes.