Skip to content

Commit d6d10ec

Browse files
brandonkachenbrandonkachen
committed
switch to zustand
1 parent 4b93d15 commit d6d10ec

File tree

3 files changed

+77
-166
lines changed

3 files changed

+77
-166
lines changed

‎bun.lock‎

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎cli/package.json‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,10 +34,11 @@
3434
"@opentui/core-darwin-arm64": "0.1.26",
3535
"@opentui/react": "github:CodebuffAI/opentui#codebuff/custom",
3636
"react": "^19.0.0",
37+
"react-reconciler": "^0.32.0",
3738
"remark-parse": "^11.0.0",
3839
"unified": "^11.0.0",
39-
"react-reconciler": "^0.32.0",
40-
"yoga-layout": "^3.2.1"
40+
"yoga-layout": "^3.2.1",
41+
"zustand": "^5.0.8"
4142
},
4243
"devDependencies": {
4344
"@types/bun": "^1.2.11",

‎cli/src/state/chat-store.ts‎

Lines changed: 73 additions & 164 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,9 @@
1-
import { useCallback, useRef, useSyncExternalStore } from 'react'
1+
import { create } from 'zustand'
22

33
import { formatTimestamp } from '../utils/helpers'
44

55
import type { ChatMessage } from '../chat'
66

7-
type Listener = () => void
8-
9-
type StateSetter<T> = (value: T | ((prev: T) => T)) => void
10-
117
export type ChatStoreState = {
128
messages: ChatMessage[]
139
streamingAgents: Set<string>
@@ -20,17 +16,19 @@ export type ChatStoreState = {
2016
}
2117

2218
type ChatStoreActions = {
23-
setMessages: StateSetter<ChatMessage[]>
24-
setStreamingAgents: StateSetter<Set<string>>
25-
setCollapsedAgents: StateSetter<Set<string>>
26-
setFocusedAgentId: StateSetter<string | null>
27-
setInputValue: StateSetter<string>
19+
setMessages: (value: ChatMessage[] | ((prev: ChatMessage[]) => ChatMessage[])) => void
20+
setStreamingAgents: (value: Set<string> | ((prev: Set<string>) => Set<string>)) => void
21+
setCollapsedAgents: (value: Set<string> | ((prev: Set<string>) => Set<string>)) => void
22+
setFocusedAgentId: (value: string | null | ((prev: string | null) => string | null)) => void
23+
setInputValue: (value: string | ((prev: string) => string)) => void
2824
setInputFocused: (focused: boolean) => void
29-
setActiveSubagents: StateSetter<Set<string>>
25+
setActiveSubagents: (value: Set<string> | ((prev: Set<string>) => Set<string>)) => void
3026
setIsChainInProgress: (active: boolean) => void
3127
reset: () => void
3228
}
3329

30+
type ChatStore = ChatStoreState & ChatStoreActions
31+
3432
const initialState: ChatStoreState = {
3533
messages: [
3634
{
@@ -50,113 +48,49 @@ const initialState: ChatStoreState = {
5048
isChainInProgress: false,
5149
}
5250

53-
let state: ChatStoreState = initialState
54-
const listeners = new Set<Listener>()
55-
56-
const notify = (): void => {
57-
for (const listener of listeners) {
58-
listener()
59-
}
60-
}
61-
62-
const resolveState = <T>(
63-
update: T | ((prev: T) => T),
64-
prev: T,
65-
): T => {
66-
if (typeof update === 'function') {
67-
return (update as (value: T) => T)(prev)
68-
}
69-
return update
70-
}
71-
72-
const assignState = (next: ChatStoreState): void => {
73-
state = next
74-
notify()
75-
}
76-
77-
const setPartialState = (
78-
updater: (current: ChatStoreState) => ChatStoreState,
79-
): void => {
80-
const next = updater(state)
81-
if (next === state) {
82-
return
83-
}
84-
assignState(next)
85-
}
86-
87-
const actions: ChatStoreActions = {
88-
setMessages: (update) => {
89-
setPartialState((current) => {
90-
const nextMessages = resolveState(update, current.messages)
91-
if (nextMessages === current.messages) {
92-
return current
93-
}
94-
return { ...current, messages: nextMessages }
95-
})
96-
},
97-
setStreamingAgents: (update) => {
98-
setPartialState((current) => {
99-
const nextAgents = resolveState(update, current.streamingAgents)
100-
if (nextAgents === current.streamingAgents) {
101-
return current
102-
}
103-
return { ...current, streamingAgents: nextAgents }
104-
})
105-
},
106-
setCollapsedAgents: (update) => {
107-
setPartialState((current) => {
108-
const nextCollapsed = resolveState(update, current.collapsedAgents)
109-
if (nextCollapsed === current.collapsedAgents) {
110-
return current
111-
}
112-
return { ...current, collapsedAgents: nextCollapsed }
113-
})
114-
},
115-
setFocusedAgentId: (update) => {
116-
setPartialState((current) => {
117-
const nextFocused = resolveState(update, current.focusedAgentId)
118-
if (current.focusedAgentId === nextFocused) {
119-
return current
120-
}
121-
return { ...current, focusedAgentId: nextFocused }
122-
})
123-
},
124-
setInputValue: (update) => {
125-
setPartialState((current) => {
126-
const nextValue = resolveState(update, current.inputValue)
127-
if (nextValue === current.inputValue) {
128-
return current
129-
}
130-
return { ...current, inputValue: nextValue }
131-
})
132-
},
133-
setInputFocused: (focused) => {
134-
setPartialState((current) => {
135-
if (current.inputFocused === focused) {
136-
return current
137-
}
138-
return { ...current, inputFocused: focused }
139-
})
140-
},
141-
setActiveSubagents: (update) => {
142-
setPartialState((current) => {
143-
const nextSubagents = resolveState(update, current.activeSubagents)
144-
if (nextSubagents === current.activeSubagents) {
145-
return current
146-
}
147-
return { ...current, activeSubagents: nextSubagents }
148-
})
149-
},
150-
setIsChainInProgress: (active) => {
151-
setPartialState((current) => {
152-
if (current.isChainInProgress === active) {
153-
return current
154-
}
155-
return { ...current, isChainInProgress: active }
156-
})
157-
},
158-
reset: () => {
159-
assignState({
51+
export const useChatStore = create<ChatStore>((set) => ({
52+
...initialState,
53+
54+
setMessages: (value) =>
55+
set((state) => ({
56+
messages: typeof value === 'function' ? value(state.messages) : value,
57+
})),
58+
59+
setStreamingAgents: (value) =>
60+
set((state) => ({
61+
streamingAgents:
62+
typeof value === 'function' ? value(state.streamingAgents) : value,
63+
})),
64+
65+
setCollapsedAgents: (value) =>
66+
set((state) => ({
67+
collapsedAgents:
68+
typeof value === 'function' ? value(state.collapsedAgents) : value,
69+
})),
70+
71+
setFocusedAgentId: (value) =>
72+
set((state) => ({
73+
focusedAgentId:
74+
typeof value === 'function' ? value(state.focusedAgentId) : value,
75+
})),
76+
77+
setInputValue: (value) =>
78+
set((state) => ({
79+
inputValue: typeof value === 'function' ? value(state.inputValue) : value,
80+
})),
81+
82+
setInputFocused: (focused) => set({ inputFocused: focused }),
83+
84+
setActiveSubagents: (value) =>
85+
set((state) => ({
86+
activeSubagents:
87+
typeof value === 'function' ? value(state.activeSubagents) : value,
88+
})),
89+
90+
setIsChainInProgress: (active) => set({ isChainInProgress: active }),
91+
92+
reset: () =>
93+
set({
16094
messages: initialState.messages.slice(),
16195
streamingAgents: new Set(initialState.streamingAgents),
16296
collapsedAgents: new Set(initialState.collapsedAgents),
@@ -165,52 +99,27 @@ const actions: ChatStoreActions = {
16599
inputFocused: initialState.inputFocused,
166100
activeSubagents: new Set(initialState.activeSubagents),
167101
isChainInProgress: initialState.isChainInProgress,
168-
})
169-
},
170-
}
102+
}),
103+
}))
171104

105+
// For backwards compatibility with non-hook usage
172106
export const chatStore = {
173-
subscribe(listener: Listener): (() => void) {
174-
listeners.add(listener)
175-
return () => {
176-
listeners.delete(listener)
177-
}
178-
},
179-
getState(): ChatStoreState {
180-
return state
181-
},
182-
...actions,
183-
}
184-
185-
type ChatStoreSnapshot = ChatStoreState & ChatStoreActions
186-
187-
const getSnapshot = (): ChatStoreSnapshot => ({
188-
...state,
189-
...actions,
190-
})
191-
192-
export const useChatStore = <T>(
193-
selector: (snapshot: ChatStoreSnapshot) => T,
194-
isEqual: (a: T, b: T) => boolean = Object.is,
195-
): T => {
196-
const selectorRef = useRef(selector)
197-
selectorRef.current = selector
198-
199-
const lastSelectionRef = useRef<T>()
200-
201-
const getSelectedSnapshot = useCallback(() => {
202-
const selection = selectorRef.current(getSnapshot())
203-
const last = lastSelectionRef.current
204-
if (last !== undefined && isEqual(selection, last)) {
205-
return last
206-
}
207-
lastSelectionRef.current = selection
208-
return selection
209-
}, [isEqual])
210-
211-
return useSyncExternalStore(
212-
chatStore.subscribe,
213-
getSelectedSnapshot,
214-
getSelectedSnapshot,
215-
)
107+
subscribe: useChatStore.subscribe,
108+
getState: useChatStore.getState,
109+
setMessages: (value: ChatMessage[] | ((prev: ChatMessage[]) => ChatMessage[])) =>
110+
useChatStore.getState().setMessages(value),
111+
setStreamingAgents: (value: Set<string> | ((prev: Set<string>) => Set<string>)) =>
112+
useChatStore.getState().setStreamingAgents(value),
113+
setCollapsedAgents: (value: Set<string> | ((prev: Set<string>) => Set<string>)) =>
114+
useChatStore.getState().setCollapsedAgents(value),
115+
setFocusedAgentId: (value: string | null | ((prev: string | null) => string | null)) =>
116+
useChatStore.getState().setFocusedAgentId(value),
117+
setInputValue: (value: string | ((prev: string) => string)) =>
118+
useChatStore.getState().setInputValue(value),
119+
setInputFocused: (focused: boolean) => useChatStore.getState().setInputFocused(focused),
120+
setActiveSubagents: (value: Set<string> | ((prev: Set<string>) => Set<string>)) =>
121+
useChatStore.getState().setActiveSubagents(value),
122+
setIsChainInProgress: (active: boolean) =>
123+
useChatStore.getState().setIsChainInProgress(active),
124+
reset: () => useChatStore.getState().reset(),
216125
}

0 commit comments

Comments
 (0)