🤖 Migrate to Zustand for simpler state management #243
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Implements Phase 1 of workspace state optimization by creating an external store pattern that enables selective re-renders via useSyncExternalStore. **Key Changes:** - Created WorkspaceStore: External store for workspace aggregators & streaming state - Created useWorkspaceStore hooks: useWorkspaceState, useWorkspaceAggregator, etc. - Migrated AIView to consume store hooks directly (no longer receives props) - Added WorkspaceStoreProvider at App level with lifecycle management - Store handles IPC subscriptions internally (cleaner than useEffect) **Benefits:** - AIView only re-renders when its specific workspace changes - Store lives outside React lifecycle (stable across renders) - Foundation for removing useWorkspaceAggregators hook **Testing:** - Unit tests verify store subscription, caching, and IPC integration - Types pass (make typecheck) **Next Steps:** - Migrate other hooks (useUnreadTracking, useResumeManager, etc.) - Migrate ProjectSidebar to use new hooks - Remove useWorkspaceAggregators once all consumers migrated
…spaceAggregators Completes Phase 1 by migrating all remaining components and hooks from useWorkspaceAggregators to the new WorkspaceStore with useSyncExternalStore. **Migrated Components & Hooks:** - AIView: Now uses useWorkspaceState(id) directly (previous commit) - ProjectSidebar: Uses useAllWorkspaceStates() for streaming indicators - useUnreadTracking: Consumes store directly, no longer needs workspaceStates param - useResumeManager: Consumes store directly - useAutoCompactContinue: Consumes store directly - App: Uses useWorkspaceRecency() and useAllWorkspaceStates() **Removed:** - useWorkspaceAggregators hook entirely deleted from codebase - All getWorkspaceState prop-drilling removed - workspaceStates/workspaceRecency passed as props removed **Architecture Benefits:** - Single external store manages all workspace state - Components only re-render when their subscribed data changes - Cleaner hook signatures (no prop drilling) - Store handles IPC subscriptions internally - Foundation for Phase 2 (GitStatusStore) **Testing:** - All 426 tests pass - Types pass (make typecheck) - Updated WorkspaceStore tests to match current WorkspaceMetadata schema
Simplifies external store architecture by:
1. Wrapping WorkspaceStore and GitStatusStore with Zustand
2. Eliminating useAll* anti-pattern hooks
3. Removing provider boilerplate
**Changes:**
1. **Added Zustand** (v5.0.8)
- Lightweight state management (~600 bytes gzipped)
- Handles subscription management automatically
- Built-in selector memoization
2. **Created Zustand wrappers**
- workspaceStoreZustand.ts - Wraps WorkspaceStore
- gitStatusStoreZustand.ts - Wraps GitStatusStore
- Preserved all existing store logic (IPC, aggregators, caching)
- Zustand only handles React integration layer
3. **Eliminated useAll* anti-pattern**
- Removed: useAllWorkspaceStates(), useAllGitStatuses()
- Problem: Caused all consumers to re-render on ANY workspace change
- Solution: Zustand selectors for fine-grained subscriptions
- Components now subscribe only to data they need
4. **Extracted WorkspaceListItem component**
- ProjectSidebar previously used useAllWorkspaceStates()
- Each workspace item now subscribes independently
- No re-renders when unrelated workspaces change
5. **Simplified App.tsx**
- Removed WorkspaceStoreProvider and GitStatusStoreProvider
- Removed manual store instantiation and cleanup
- Zustand manages global store instances automatically
- Reduced from ~30 lines of boilerplate to ~3 lines
6. **Removed debug logging**
- Cleaned up emoji logs from GitStatusStore
- Removed console statements from App.tsx
7. **Deleted obsolete files**
- src/hooks/useWorkspaceStore.tsx (108 LOC)
- src/hooks/useGitStatus.tsx (100 LOC)
**Benefits:**
- **Simpler architecture**: -61 LOC net (399 deleted, 338 added)
- **Better patterns**: Selector-based subscriptions, no useAll* hooks
- **Cleaner code**: No provider boilerplate, no manual subscription management
- **Preserved complexity**: Kept WorkspaceStore, GitStatusStore, CacheManager intact
**What we kept:**
✅ WorkspaceStore class (IPC, aggregators, caching logic)
✅ GitStatusStore class (polling, fetch with backoff)
✅ CacheManager (testable cache invalidation)
✅ All existing tests (452 pass)
**Architecture:**
```
React Components
↓
Zustand Selectors (auto-memoized, fine-grained)
↓
Zustand Wrappers (~80 LOC each)
↓
Core Store Classes (UNCHANGED)
├─ IPC subscriptions
├─ Message aggregators
├─ Cache management
└─ Git status polling
```
Tests: 452 pass, typecheck clean
_Generated with `cmux`_
- Move Zustand integration from separate files into WorkspaceStore.ts and GitStatusStore.ts - Delete unused useWorkspaceAggregators hook - Removes 124 LOC across 3 deleted files
- Move reference tests from separate files into WorkspaceStore.test.ts and GitStatusStore.test.ts - Keep only tests that verify public API (removed @ts-expect-error tests) - Deleted 3 test files (515 LOC)
- Wrap startRenaming, confirmRename, handleRenameKeyDown, and handleRemoveWorkspace with useCallback - Wrap WorkspaceListItem with React.memo - Prevents all workspace items from re-rendering when one workspace's state changes
- Compute streamingModels lazily when command palette sources are built - App.tsx no longer subscribes to getAllStates() on every render - Prevents App.tsx and ProjectSidebar from re-rendering on every stream update - Only CommandPalette data fetching is affected (when opened)
…reaming - WorkspaceListItem now subscribes only to sidebar-relevant fields (canInterrupt, currentModel, recencyTimestamp) - No longer re-renders on every message update during streaming - Zustand's shallow equality prevents re-renders when these 3 fields haven't changed - Sidebar should be completely still during active streams
- Import useShallow from zustand/react/shallow - Wrap selector with useShallow to enable shallow equality comparison - Prevents infinite re-render loop from returning new object references
Critical fix: - Move all hooks in AIView.tsx before early return (violates Rules of Hooks) - Made hooks resilient with optional chaining and safe defaults Lint fixes: - Remove unused imports (App.tsx, GitStatusStore.test.ts) - Fix import style in WorkspaceStore.ts (inline import -> type-only) - Change CacheManager any -> unknown for type safety - Remove duplicate getWorkspaceDisplayName in ProjectSidebar - Add type-safe getOnChatCallback helper in tests - Add eslint-disable comments for intentional dependency omissions - Remove orphaned JSDoc comment All static checks now pass: - ✅ ESLint (0 errors) - ✅ TypeScript (main + renderer) - ✅ Prettier formatting - ✅ 438 tests passing _Generated with `cmux`_
## Problem WorkspaceStore had 3 separate manual cache implementations: 1. Per-workspace state cache (Map<string, WorkspaceState>) 2. All-states cache with 16-line O(n) validation logic 3. Recency cache using JSON.stringify for hashing This was fragile, error-prone, and violated DRY principles. ## Changes ### 1. Centralized Caching (Fix #1 - HIGH) - Replaced 3 manual caches with single CacheManager instance - **Removed 58 lines** of manual cache validation logic - Single `cache.invalidateAll()` call in `emit()` - impossible to forget - Consistent with StreamingMessageAggregator pattern ### 2. Extract Recency Computation (Fix #2 - MEDIUM) - Created `src/utils/messages/recency.ts` - **Single array reverse** instead of two - 11 unit tests for edge cases ### 3. Extract Array Comparison (Fix #4 - LOW) - Created `src/utils/arrays.ts` - `arraysEqualByReference()` helper with 11 unit tests - Simplified `statesEqual()` from 34 lines to 10 ### 4. Enhanced Tests (Fix #3 - MEDIUM) - Added cache invalidation tests (4 new tests) - Added race condition tests (3 new tests) - Total: 41 new tests across 3 files ## Impact - **Net: -123 lines** in WorkspaceStore.ts (172 removed, 49 added) - **+502 lines total** (379 lines of new tests) - **All 467 tests pass** (19 WorkspaceStore, 11 recency, 11 arrays) - Zero type/lint errors ## What Wasn't Changed - Did NOT extract message handlers (148-line handleChatMessage) - Would be 4-6 hour effort with high integration risk - Deferred until handlers become more complex Generated with `cmux`
## Problem
GitStatusStore had similar issues to WorkspaceStore:
1. Two manual cache implementations (statusCache, allStatusCache)
2. 23 lines of manual cache validation in getAllStatuses()
3. Fragile path parsing to extract project name
4. compareGitStatus() method only used for manual cache comparison
## Changes
### 1. Centralized Caching
- Replaced 2 manual caches with single CacheManager instance
- **Removed 41 lines** of manual cache validation logic
- Single `cache.invalidateAll()` call in `emit()`
- Consistent with WorkspaceStore refactoring
### 2. Use projectName from Metadata
- Changed `groupWorkspacesByProject()` to use `m.projectName` directly
- Removed fragile path parsing: `m.workspacePath.split("/")[length-2]`
- More robust to path structure changes
### 3. Removed Unused compareGitStatus()
- Was only used for manual cache comparison
- CacheManager handles equality internally
- 10 lines removed
## Impact
- **Net: -41 lines** (53 removed, 12 added)
- **All 467 tests pass** (10 GitStatusStore tests)
- Zero type/lint errors
- No performance regression expected
## Before/After
### Before:
```typescript
// Two manual caches
private statusCache = new Map<string, GitStatus | null>();
private allStatusCache: Map<string, GitStatus | null> | null = null;
getAllStatuses(): Map<string, GitStatus | null> {
// 23 lines of manual validation logic
let hasChanges = false;
if (!this.allStatusCache || this.allStatusCache.size !== this.gitStatusMap.size) {
hasChanges = true;
} else {
for (const [id, status] of this.gitStatusMap) {
const cached = this.allStatusCache.get(id);
if (!this.compareGitStatus(cached ?? null, status ?? null)) {
hasChanges = true;
break;
}
}
}
// ... more validation
}
```
### After:
```typescript
// Single cache manager
private cache = new CacheManager();
getAllStatuses(): Map<string, GitStatus | null> {
return this.cache.get("all-statuses", () => {
return new Map(this.gitStatusMap);
});
}
```
## Summary
GitStatusStore now follows the same clean caching pattern as WorkspaceStore:
- Single CacheManager instance
- Automatic invalidation on emit()
- No manual cache validation
- More robust project grouping
Generated with `cmux`
ammar-agent
force-pushed
the
migrate-to-zustand
branch
from
c231611 to
6534e01
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Migrates workspace state and git status management to Zustand for simpler, more maintainable architecture.
What Changed
Replaced hooks with stores + Zustand:
useWorkspaceAggregators→WorkspaceStoreclass + Zustand wrapperGitStatusContext→GitStatusStoreclass + Zustand wrapperBenefits:
Component changes:
AIView: UsesuseWorkspaceState()anduseGitStatus()hooks instead of contextProjectSidebar: ExtractedWorkspaceListItemcomponent - each workspace subscribes independently instead of all workspaces re-rendering togetherTesting
make typecheck)