refactor(cli): consolidate state management with Zustand stores

- Remove duplicate state from chat-store (streamingAgents, focusedAgentId, inputFocused, activeSubagents, isChainInProgress)
- Move streaming state to stream-store and UI state to ui-store
- Eliminate prop drilling by having components use stores directly
- Update MessageWithAgents to use useUiStore instead of prop drilling setFocusedAgentId
- Keep React contexts for derived values (ChatThemeProvider) and callbacks (MessageActionsProvider) only
- Add comprehensive state management architecture documentation to cli/knowledge.md

This ensures each piece of state has a single source of truth and follows the principle:
Use Zustand stores for state, React Context for derived values and callbacks.

Author: brandonkachen

cli/knowledge.md

@@ -1,5 +1,98 @@
 # CLI Package Knowledge
 
+## State Management Architecture
+
+The CLI uses **Zustand stores** for all state management, with React Context reserved exclusively for passing derived values and callbacks.
+
+### Zustand Stores
+
+All application state lives in specialized Zustand stores located in `cli/src/state/`:
+
+**State Stores:**
+- **auth-store.ts** - Authentication state (`isAuthenticated`, `user`)
+- **chat-store.ts** - Main chat state (`messages`, `inputValue`, `agentMode`, etc.)
+- **queue-store.ts** - Message queue and streaming status
+- **stream-store.ts** - Streaming agents and chain progress state
+- **ui-store.ts** - UI focus state (`focusedAgentId`, `inputFocused`)
+- **feedback-store.ts** - Feedback form state
+- **login-store.ts** - Login flow state
+
+**Key Principles:**
+- Each store manages a single domain of state
+- No duplicate state across stores
+- Components use `useShallow` selector for optimal re-render performance
+- Stores use `immer` middleware for immutable updates
+
+### React Context (Non-State Only)
+
+The CLI uses React Context **only** for passing derived values and callbacks, never for state:
+
+**Context Providers:**
+- **ChatThemeProvider** (`cli/src/contexts/chat-theme-context.tsx`) - Provides derived/computed values:
+  - `theme` (from hook)
+  - `markdownPalette` (memoized from theme)
+  - `availableWidth` (terminal dimension)
+  - `timerStartTime` (timer state)
+
+- **MessageActionsProvider** (`cli/src/contexts/message-actions-context.tsx`) - Provides callback functions:
+  - `onToggleCollapsed`
+  - `onBuildFast`
+  - `onBuildMax`
+  - `onFeedback`
+  - `onCloseFeedback`
+
+**Why This Works:**
+- Contexts don't manage state (no `useState` or state updates)
+- They only pass down **read-only derived values** or **stable callback references**
+- This avoids prop drilling without the performance issues of context-based state
+
+### Architecture Guidelines
+
+**✅ Use Zustand stores for:**
+- Any state that changes over time
+- State that needs to be shared across components
+- State that triggers re-renders
+
+**✅ Use React Context for:**
+- Derived/computed values (theme, dimensions, etc.)
+- Callback functions that don't change
+- Dependency injection of non-state values
+
+**❌ Never use React Context for:**
+- Managing state with `useState`
+- State that updates frequently
+- Anything that requires `setState` calls
+
+### Example: Accessing State
+
+```tsx
+// ✅ CORRECT: Use Zustand store for state
+import { useUiStore } from '../state/ui-store'
+
+const MyComponent = () => {
+  const { focusedAgentId, setFocusedAgentId } = useUiStore()
+  // Use state directly, no prop drilling needed
+}
+
+// ✅ CORRECT: Use context for callbacks/derived values
+import { useMessageActions } from '../contexts/message-actions-context'
+import { useChatTheme } from '../contexts/chat-theme-context'
+
+const MyComponent = () => {
+  const { onToggleCollapsed } = useMessageActions()
+  const { theme, markdownPalette } = useChatTheme()
+  // Use callbacks and derived values
+}
+```
+
+### Migration Notes
+
+Previously (before refactoring), some state was duplicated in `chat-store.ts`. The refactoring removed:
+- `streamingAgents`, `activeSubagents`, `isChainInProgress` (moved to `stream-store.ts`)
+- `focusedAgentId`, `inputFocused` (moved to `ui-store.ts`)
+
+This ensures each piece of state has a single source of truth.
+
 ## Test Naming Conventions
 
 **IMPORTANT**: Follow these naming patterns for automatic dependency detection:

cli/src/app.tsx

@@ -15,6 +15,7 @@ import { useTerminalDimensions } from './hooks/use-terminal-dimensions'
 import { useTheme } from './hooks/use-theme'
 import { getProjectRoot } from './project-files'
 import { useChatStore } from './state/chat-store'
+import { useUiStore } from './state/ui-store'
 import { createValidationErrorBlocks } from './utils/create-validation-error-blocks'
 import { openFileAtPath } from './utils/open-file'
 
@@ -53,13 +54,14 @@ export const App = ({
 
   const [isAgentListCollapsed, setIsAgentListCollapsed] = useState(true)
   const inputRef = useRef<MultilineInputHandle | null>(null)
-  const { setInputFocused, resetChatStore } = useChatStore(
+  const { resetChatStore } = useChatStore(
     useShallow((store) => ({
-      setInputFocused: store.setInputFocused,
       resetChatStore: store.reset,
     })),
   )
 
+  const { setInputFocused } = useUiStore()
+
   const { isAuthenticated, handleLoginSuccess, logout } = useAuthState({
     requireAuth,
     hasInvalidCredentials,

cli/src/chat.tsx

@@ -247,8 +247,6 @@ export const Chat = ({
     isUserCollapsing,
   )
 
-
-
   const inertialScrollAcceleration = useMemo(
     () => createChatScrollAcceleration(),
     [],
@@ -761,7 +759,6 @@ export const Chat = ({
             streamingAgents={streamingAgents}
             messageTree={messageTree}
             messages={messages}
-            setFocusedAgentId={setFocusedAgentId}
             isWaitingForResponse={isWaitingForResponse}
           />
 

cli/src/components/chat-scroll-area.tsx

@@ -3,7 +3,6 @@ import React from 'react'
 import { MessageWithAgents } from './message-with-agents'
 import type { ChatMessage } from '../types/chat'
 import type { ScrollBoxRenderable } from '@opentui/core'
-import type { Dispatch, SetStateAction } from 'react'
 
 interface ChatScrollAreaProps {
   scrollRef: React.MutableRefObject<ScrollBoxRenderable | null>
@@ -14,7 +13,6 @@ interface ChatScrollAreaProps {
   streamingAgents: Set<string>
   messageTree: Map<string, ChatMessage[]>
   messages: ChatMessage[]
-  setFocusedAgentId: Dispatch<SetStateAction<string | null>>
   isWaitingForResponse: boolean
 }
 
@@ -27,7 +25,6 @@ export const ChatScrollArea = ({
   streamingAgents,
   messageTree,
   messages,
-  setFocusedAgentId,
   isWaitingForResponse,
 }: ChatScrollAreaProps) => {
   return (
@@ -78,7 +75,6 @@ export const ChatScrollArea = ({
             streamingAgents={streamingAgents}
             messageTree={messageTree}
             messages={messages}
-            setFocusedAgentId={setFocusedAgentId}
             isWaitingForResponse={isWaitingForResponse}
           />
         )

cli/src/components/message-with-agents.tsx

@@ -12,6 +12,7 @@ import {
 } from '../utils/markdown-renderer'
 import { useMessageActions } from '../contexts/message-actions-context'
 import { useChatTheme } from '../contexts/chat-theme-context'
+import { useUiStore } from '../state/ui-store'
 
 import type { ChatMessage } from '../types/chat'
 import type { ChatTheme } from '../types/theme-system'
@@ -23,7 +24,6 @@ interface MessageWithAgentsProps {
   streamingAgents: Set<string>
   messageTree: Map<string, ChatMessage[]>
   messages: ChatMessage[]
-  setFocusedAgentId: React.Dispatch<React.SetStateAction<string | null>>
   isWaitingForResponse: boolean
 }
 
@@ -35,7 +35,6 @@ export const MessageWithAgents = memo(
     streamingAgents,
     messageTree,
     messages,
-    setFocusedAgentId,
     isWaitingForResponse,
   }: MessageWithAgentsProps): ReactNode => {
     const { theme, markdownPalette, availableWidth, timerStartTime } = useChatTheme()
@@ -44,15 +43,14 @@ export const MessageWithAgents = memo(
     const isAgent = message.variant === 'agent'
 
     if (isAgent) {
-      return (            <AgentMessage
-              message={message}
-              depth={depth}
-              streamingAgents={streamingAgents}
-              messageTree={messageTree}
-              messages={messages}
-              setFocusedAgentId={setFocusedAgentId}
-              isWaitingForResponse={isWaitingForResponse}
+      return (            <AgentMessage            message={message}
+            depth={depth}
+            streamingAgents={streamingAgents}
+            messageTree={messageTree}
+            messages={messages}
+            isWaitingForResponse={isWaitingForResponse}
             />
+
       )
     }
 
@@ -214,7 +212,6 @@ export const MessageWithAgents = memo(
                 streamingAgents={streamingAgents}
                 messageTree={messageTree}
                 messages={messages}
-                setFocusedAgentId={setFocusedAgentId}
                 isWaitingForResponse={isWaitingForResponse}
               />
               </box>
@@ -232,7 +229,6 @@ interface AgentMessageProps {
   streamingAgents: Set<string>
   messageTree: Map<string, ChatMessage[]>
   messages: ChatMessage[]
-  setFocusedAgentId: React.Dispatch<React.SetStateAction<string | null>>
   isWaitingForResponse: boolean
 }
 
@@ -243,11 +239,11 @@ const AgentMessage = memo(
     streamingAgents,
     messageTree,
     messages,
-    setFocusedAgentId,
     isWaitingForResponse,
   }: AgentMessageProps): ReactNode => {
     const { theme, markdownPalette, availableWidth } = useChatTheme()
     const { onToggleCollapsed } = useMessageActions()
+    const { setFocusedAgentId } = useUiStore()
     const agentInfo = message.agent!
 
     // Get or initialize collapse state from message metadata
@@ -399,7 +395,6 @@ const AgentMessage = memo(
                   streamingAgents={streamingAgents}
                   messageTree={messageTree}
                   messages={messages}
-                  setFocusedAgentId={setFocusedAgentId}
                   isWaitingForResponse={isWaitingForResponse}
                 />
               </box>

cli/src/state/chat-store.ts

@@ -16,14 +16,9 @@ export type InputValue = {
 
 export type ChatStoreState = {
   messages: ChatMessage[]
-  streamingAgents: Set<string>
-  focusedAgentId: string | null
   inputValue: string
   cursorPosition: number
   lastEditDueToNav: boolean
-  inputFocused: boolean
-  activeSubagents: Set<string>
-  isChainInProgress: boolean
   slashSelectedIndex: number
   agentSelectedIndex: number
   agentMode: AgentMode
@@ -37,20 +32,9 @@ type ChatStoreActions = {
   setMessages: (
     value: ChatMessage[] | ((prev: ChatMessage[]) => ChatMessage[]),
   ) => void
-  setStreamingAgents: (
-    value: Set<string> | ((prev: Set<string>) => Set<string>),
-  ) => void
-  setFocusedAgentId: (
-    value: string | null | ((prev: string | null) => string | null),
-  ) => void
   setInputValue: (
     value: InputValue | ((prev: InputValue) => InputValue),
   ) => void
-  setInputFocused: (focused: boolean) => void
-  setActiveSubagents: (
-    value: Set<string> | ((prev: Set<string>) => Set<string>),
-  ) => void
-  setIsChainInProgress: (active: boolean) => void
   setSlashSelectedIndex: (value: number | ((prev: number) => number)) => void
   setAgentSelectedIndex: (value: number | ((prev: number) => number)) => void
   setAgentMode: (mode: AgentMode) => void
@@ -66,14 +50,9 @@ type ChatStore = ChatStoreState & ChatStoreActions
 
 const initialState: ChatStoreState = {
   messages: [],
-  streamingAgents: new Set<string>(),
-  focusedAgentId: null,
   inputValue: '',
   cursorPosition: 0,
   lastEditDueToNav: false,
-  inputFocused: true,
-  activeSubagents: new Set<string>(),
-  isChainInProgress: false,
   slashSelectedIndex: 0,
   agentSelectedIndex: 0,
   agentMode: 'DEFAULT',
@@ -93,18 +72,6 @@ export const useChatStore = create<ChatStore>()(
           typeof value === 'function' ? value(state.messages) : value
       }),
 
-    setStreamingAgents: (value) =>
-      set((state) => {
-        state.streamingAgents =
-          typeof value === 'function' ? value(state.streamingAgents) : value
-      }),
-
-    setFocusedAgentId: (value) =>
-      set((state) => {
-        state.focusedAgentId =
-          typeof value === 'function' ? value(state.focusedAgentId) : value
-      }),
-
     setInputValue: (value) =>
       set((state) => {
         const { text, cursorPosition, lastEditDueToNav } =
@@ -120,22 +87,6 @@ export const useChatStore = create<ChatStore>()(
         state.lastEditDueToNav = lastEditDueToNav
       }),
 
-    setInputFocused: (focused) =>
-      set((state) => {
-        state.inputFocused = focused
-      }),
-
-    setActiveSubagents: (value) =>
-      set((state) => {
-        state.activeSubagents =
-          typeof value === 'function' ? value(state.activeSubagents) : value
-      }),
-
-    setIsChainInProgress: (active) =>
-      set((state) => {
-        state.isChainInProgress = active
-      }),
-
     setSlashSelectedIndex: (value) =>
       set((state) => {
         state.slashSelectedIndex =
@@ -187,14 +138,9 @@ export const useChatStore = create<ChatStore>()(
     reset: () =>
       set((state) => {
         state.messages = initialState.messages.slice()
-        state.streamingAgents = new Set(initialState.streamingAgents)
-        state.focusedAgentId = initialState.focusedAgentId
         state.inputValue = initialState.inputValue
         state.cursorPosition = initialState.cursorPosition
         state.lastEditDueToNav = initialState.lastEditDueToNav
-        state.inputFocused = initialState.inputFocused
-        state.activeSubagents = new Set(initialState.activeSubagents)
-        state.isChainInProgress = initialState.isChainInProgress
         state.slashSelectedIndex = initialState.slashSelectedIndex
         state.agentSelectedIndex = initialState.agentSelectedIndex
         state.agentMode = initialState.agentMode