Skip to content

Commit 448c531

Browse files
brandonkachencodebuff-team
brandonkachen
and
codebuff-team
committed
Refactor: Modernize Zustand store implementation
- Add immer middleware to reduce boilerplate in setters - Use useShallow for efficient multi-value selection in chat.tsx - Remove unused chatStore non-hook wrapper export - Move slashSelectedIndex and agentSelectedIndex to store Benefits: - Cleaner setter syntax with direct state mutations - Single useChatStore call with useShallow instead of multiple calls - Reduced boilerplate while maintaining type safety 🤖 Generated with Codebuff Co-Authored-By: Codebuff <noreply@codebuff.com>
1 parent d6d10ec commit 448c531

File tree

2 files changed

+102
-86
lines changed

2 files changed

+102
-86
lines changed

‎cli/src/chat.tsx‎

Lines changed: 46 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ import { logger } from './utils/logger'
2525
import { buildMessageTree } from './utils/message-tree-utils'
2626
import { chatThemes, createMarkdownPalette } from './utils/theme-system'
2727
import { useChatStore } from './state/chat-store'
28+
import { useShallow } from 'zustand/react/shallow'
2829

2930
import type { ToolName } from '@codebuff/sdk'
3031
import type { InputRenderable, ScrollBoxRenderable } from '@opentui/core'
@@ -85,43 +86,59 @@ export const App = ({ initialPrompt }: { initialPrompt?: string } = {}) => {
8586
const theme = chatThemes[themeName]
8687
const markdownPalette = useMemo(() => createMarkdownPalette(theme), [theme])
8788

88-
const inputValue = useChatStore((store) => store.inputValue)
89-
const setInputValue = useChatStore((store) => store.setInputValue)
90-
const inputFocused = useChatStore((store) => store.inputFocused)
91-
const setInputFocused = useChatStore((store) => store.setInputFocused)
92-
const [slashSelectedIndex, setSlashSelectedIndex] = useState<number>(0)
93-
const [agentSelectedIndex, setAgentSelectedIndex] = useState<number>(0)
89+
const {
90+
inputValue,
91+
setInputValue,
92+
inputFocused,
93+
setInputFocused,
94+
slashSelectedIndex,
95+
setSlashSelectedIndex,
96+
agentSelectedIndex,
97+
setAgentSelectedIndex,
98+
collapsedAgents,
99+
setCollapsedAgents,
100+
streamingAgents,
101+
setStreamingAgents,
102+
focusedAgentId,
103+
setFocusedAgentId,
104+
messages,
105+
setMessages,
106+
activeSubagents,
107+
setActiveSubagents,
108+
isChainInProgress,
109+
setIsChainInProgress,
110+
} = useChatStore(
111+
useShallow((store) => ({
112+
inputValue: store.inputValue,
113+
setInputValue: store.setInputValue,
114+
inputFocused: store.inputFocused,
115+
setInputFocused: store.setInputFocused,
116+
slashSelectedIndex: store.slashSelectedIndex,
117+
setSlashSelectedIndex: store.setSlashSelectedIndex,
118+
agentSelectedIndex: store.agentSelectedIndex,
119+
setAgentSelectedIndex: store.setAgentSelectedIndex,
120+
collapsedAgents: store.collapsedAgents,
121+
setCollapsedAgents: store.setCollapsedAgents,
122+
streamingAgents: store.streamingAgents,
123+
setStreamingAgents: store.setStreamingAgents,
124+
focusedAgentId: store.focusedAgentId,
125+
setFocusedAgentId: store.setFocusedAgentId,
126+
messages: store.messages,
127+
setMessages: store.setMessages,
128+
activeSubagents: store.activeSubagents,
129+
setActiveSubagents: store.setActiveSubagents,
130+
isChainInProgress: store.isChainInProgress,
131+
setIsChainInProgress: store.setIsChainInProgress,
132+
})),
133+
)
94134

95135
const activeAgentStreamsRef = useRef<number>(0)
96-
const isChainInProgress = useChatStore((store) => store.isChainInProgress)
97-
const setIsChainInProgress = useChatStore(
98-
(store) => store.setIsChainInProgress,
99-
)
100136
const isChainInProgressRef = useRef<boolean>(isChainInProgress)
101137

102138
const { clipboardMessage } = useClipboard()
103139

104-
const collapsedAgents = useChatStore((store) => store.collapsedAgents)
105-
const setCollapsedAgents = useChatStore(
106-
(store) => store.setCollapsedAgents,
107-
)
108-
const streamingAgents = useChatStore((store) => store.streamingAgents)
109-
const setStreamingAgents = useChatStore(
110-
(store) => store.setStreamingAgents,
111-
)
112-
const focusedAgentId = useChatStore((store) => store.focusedAgentId)
113-
const setFocusedAgentId = useChatStore(
114-
(store) => store.setFocusedAgentId,
115-
)
116140
const agentRefsMap = useRef<Map<string, any>>(new Map())
117-
118-
const messages = useChatStore((store) => store.messages)
119-
const setMessages = useChatStore((store) => store.setMessages)
120141
const hasAutoSubmittedRef = useRef(false)
121-
const activeSubagents = useChatStore((store) => store.activeSubagents)
122-
const setActiveSubagents = useChatStore(
123-
(store) => store.setActiveSubagents,
124-
)
125142
const activeSubagentsRef = useRef<Set<string>>(activeSubagents)
126143

127144
useEffect(() => {

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

Lines changed: 56 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import { create } from 'zustand'
2+
import { immer } from 'zustand/middleware/immer'
23

34
import { formatTimestamp } from '../utils/helpers'
45

@@ -13,6 +14,8 @@ export type ChatStoreState = {
1314
inputFocused: boolean
1415
activeSubagents: Set<string>
1516
isChainInProgress: boolean
17+
slashSelectedIndex: number
18+
agentSelectedIndex: number
1619
}
1720

1821
type ChatStoreActions = {
@@ -24,6 +27,8 @@ type ChatStoreActions = {
2427
setInputFocused: (focused: boolean) => void
2528
setActiveSubagents: (value: Set<string> | ((prev: Set<string>) => Set<string>)) => void
2629
setIsChainInProgress: (active: boolean) => void
30+
setSlashSelectedIndex: (value: number | ((prev: number) => number)) => void
31+
setAgentSelectedIndex: (value: number | ((prev: number) => number)) => void
2732
reset: () => void
2833
}
2934

@@ -46,80 +51,74 @@ const initialState: ChatStoreState = {
4651
inputFocused: true,
4752
activeSubagents: new Set<string>(),
4853
isChainInProgress: false,
54+
slashSelectedIndex: 0,
55+
agentSelectedIndex: 0,
4956
}
5057

51-
export const useChatStore = create<ChatStore>((set) => ({
58+
export const useChatStore = create<ChatStore>()(immer((set) => ({
5259
...initialState,
5360

5461
setMessages: (value) =>
55-
set((state) => ({
56-
messages: typeof value === 'function' ? value(state.messages) : value,
57-
})),
62+
set((state) => {
63+
state.messages = typeof value === 'function' ? value(state.messages) : value
64+
}),
5865

5966
setStreamingAgents: (value) =>
60-
set((state) => ({
61-
streamingAgents:
62-
typeof value === 'function' ? value(state.streamingAgents) : value,
63-
})),
67+
set((state) => {
68+
state.streamingAgents = typeof value === 'function' ? value(state.streamingAgents) : value
69+
}),
6470

6571
setCollapsedAgents: (value) =>
66-
set((state) => ({
67-
collapsedAgents:
68-
typeof value === 'function' ? value(state.collapsedAgents) : value,
69-
})),
72+
set((state) => {
73+
state.collapsedAgents = typeof value === 'function' ? value(state.collapsedAgents) : value
74+
}),
7075

7176
setFocusedAgentId: (value) =>
72-
set((state) => ({
73-
focusedAgentId:
74-
typeof value === 'function' ? value(state.focusedAgentId) : value,
75-
})),
77+
set((state) => {
78+
state.focusedAgentId = typeof value === 'function' ? value(state.focusedAgentId) : value
79+
}),
7680

7781
setInputValue: (value) =>
78-
set((state) => ({
79-
inputValue: typeof value === 'function' ? value(state.inputValue) : value,
80-
})),
82+
set((state) => {
83+
state.inputValue = typeof value === 'function' ? value(state.inputValue) : value
84+
}),
8185

82-
setInputFocused: (focused) => set({ inputFocused: focused }),
86+
setInputFocused: (focused) =>
87+
set((state) => {
88+
state.inputFocused = focused
89+
}),
8390

8491
setActiveSubagents: (value) =>
85-
set((state) => ({
86-
activeSubagents:
87-
typeof value === 'function' ? value(state.activeSubagents) : value,
88-
})),
92+
set((state) => {
93+
state.activeSubagents = typeof value === 'function' ? value(state.activeSubagents) : value
94+
}),
8995

90-
setIsChainInProgress: (active) => set({ isChainInProgress: active }),
96+
setIsChainInProgress: (active) =>
97+
set((state) => {
98+
state.isChainInProgress = active
99+
}),
100+
101+
setSlashSelectedIndex: (value) =>
102+
set((state) => {
103+
state.slashSelectedIndex = typeof value === 'function' ? value(state.slashSelectedIndex) : value
104+
}),
105+
106+
setAgentSelectedIndex: (value) =>
107+
set((state) => {
108+
state.agentSelectedIndex = typeof value === 'function' ? value(state.agentSelectedIndex) : value
109+
}),
91110

92111
reset: () =>
93-
set({
94-
messages: initialState.messages.slice(),
95-
streamingAgents: new Set(initialState.streamingAgents),
96-
collapsedAgents: new Set(initialState.collapsedAgents),
97-
focusedAgentId: initialState.focusedAgentId,
98-
inputValue: initialState.inputValue,
99-
inputFocused: initialState.inputFocused,
100-
activeSubagents: new Set(initialState.activeSubagents),
101-
isChainInProgress: initialState.isChainInProgress,
112+
set((state) => {
113+
state.messages = initialState.messages.slice()
114+
state.streamingAgents = new Set(initialState.streamingAgents)
115+
state.collapsedAgents = new Set(initialState.collapsedAgents)
116+
state.focusedAgentId = initialState.focusedAgentId
117+
state.inputValue = initialState.inputValue
118+
state.inputFocused = initialState.inputFocused
119+
state.activeSubagents = new Set(initialState.activeSubagents)
120+
state.isChainInProgress = initialState.isChainInProgress
121+
state.slashSelectedIndex = initialState.slashSelectedIndex
122+
state.agentSelectedIndex = initialState.agentSelectedIndex
102123
}),
103-
}))
104-
105-
// For backwards compatibility with non-hook usage
106-
export const chatStore = {
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(),
125-
}
124+
})))

0 commit comments

Comments
 (0)