refactor: consolidate store improvements and add runState management

- Consolidate enableMapSet() to single call at app initialization
- Add runState to chat store for centralized state management
- Simplify feedback-container with single useFeedbackStore call
- Update use-send-message to save runState to store
- Include runState in feedback submission logging

These changes improve code organization and set foundation for future
simplifications where components can access runState directly from the
store instead of extracting from message metadata.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

Author: brandonkachen

cli/src/chat.tsx

@@ -128,6 +128,7 @@ export const Chat = ({
     addSessionCredits,
     resetChatStore,
     sessionCreditsUsed,
+    setRunState,
   } = useChatStore(
     useShallow((store) => ({
       inputValue: store.inputValue,
@@ -164,6 +165,7 @@ export const Chat = ({
       addSessionCredits: store.addSessionCredits,
       resetChatStore: store.reset,
       sessionCreditsUsed: store.sessionCreditsUsed,
+      setRunState: store.setRunState,
     })),
   )
 
@@ -428,6 +430,7 @@ export const Chat = ({
     lastMessageMode,
     setLastMessageMode,
     addSessionCredits,
+    setRunState,
     isQueuePausedRef,
     resumeQueue,
     continueChat,

cli/src/components/feedback-container.tsx

@@ -49,6 +49,7 @@ export const FeedbackContainer: React.FC<FeedbackContainerProps> = ({
   const messages = useChatStore((state) => state.messages)
   const agentMode = useChatStore((state) => state.agentMode)
   const sessionCreditsUsed = useChatStore((state) => state.sessionCreditsUsed)
+  const runState = useChatStore((state) => state.runState)
 
   const previousFeedbackModeRef = useRef(feedbackMode)
 
@@ -86,6 +87,7 @@ export const FeedbackContainer: React.FC<FeedbackContainerProps> = ({
         category: feedbackCategory,
         type: feedbackMessageId ? 'message' : 'general',
       },
+      ...(runState && { runState }),
     })
 
     if (feedbackMessageId) {
@@ -106,6 +108,7 @@ export const FeedbackContainer: React.FC<FeedbackContainerProps> = ({
     messages,
     agentMode,
     sessionCreditsUsed,
+    runState,
     markMessageFeedbackSubmitted,
     resetFeedbackForm,
     closeFeedback,

cli/src/hooks/use-send-message.ts

@@ -210,6 +210,7 @@ interface UseSendMessageOptions {
   lastMessageMode: AgentMode | null
   setLastMessageMode: (mode: AgentMode | null) => void
   addSessionCredits: (credits: number) => void
+  setRunState: (runState: any | null) => void
   isQueuePausedRef?: React.MutableRefObject<boolean>
   resumeQueue?: () => void
   continueChat: boolean
@@ -244,6 +245,7 @@ export const useSendMessage = ({
   lastMessageMode,
   setLastMessageMode,
   addSessionCredits,
+  setRunState,
   isQueuePausedRef,
   resumeQueue,
   continueChat,
@@ -259,6 +261,7 @@ export const useSendMessage = ({
       const loadedState = loadMostRecentChatState()
       if (loadedState) {
         previousRunStateRef.current = loadedState.runState
+        setRunState(loadedState.runState)
         setMessages(loadedState.messages)
 
         // Collapse all subagents and tools by default when continuing
@@ -275,7 +278,7 @@ export const useSendMessage = ({
         logger.info('No previous chat state found to continue from')
       }
     }
-  }, [continueChat, setMessages, setCollapsedAgents])
+  }, [continueChat, setMessages, setCollapsedAgents, setRunState])
   const spawnAgentsMapRef = useRef<
     Map<string, { index: number; agentType: string }>
   >(new Map())
@@ -1561,7 +1564,8 @@ export const useSendMessage = ({
         })
 
         previousRunStateRef.current = runState
-        
+        setRunState(runState)
+
         // Save both runState and current messages
         applyMessageUpdate((currentMessages) => {
           saveChatState(runState, currentMessages)

cli/src/index.tsx

@@ -9,8 +9,12 @@ import { validateAgents } from '@codebuff/sdk'
 import { render } from '@opentui/react'
 import { QueryClient, QueryClientProvider } from '@tanstack/react-query'
 import { Command } from 'commander'
+import { enableMapSet } from 'immer'
 import React from 'react'
 
+// Enable Map and Set support in Immer globally (once at app initialization)
+enableMapSet()
+
 import { App } from './app'
 import { initializeThemeStore } from './hooks/use-theme'
 import { getProjectRoot } from './project-files'

cli/src/state/chat-store.ts

@@ -1,4 +1,3 @@
-import { enableMapSet } from 'immer'
 import { create } from 'zustand'
 import { immer } from 'zustand/middleware/immer'
 
@@ -31,6 +30,7 @@ export type ChatStoreState = {
   hasReceivedPlanResponse: boolean
   lastMessageMode: AgentMode | null
   sessionCreditsUsed: number
+  runState: any | null // Using any to avoid Immer type issues with complex RunState type
 }
 
 type ChatStoreActions = {
@@ -64,13 +64,12 @@ type ChatStoreActions = {
   setHasReceivedPlanResponse: (value: boolean) => void
   setLastMessageMode: (mode: AgentMode | null) => void
   addSessionCredits: (credits: number) => void
+  setRunState: (runState: any | null) => void
   reset: () => void
 }
 
 type ChatStore = ChatStoreState & ChatStoreActions
 
-enableMapSet()
-
 const initialState: ChatStoreState = {
   messages: [],
   streamingAgents: new Set<string>(),
@@ -89,6 +88,7 @@ const initialState: ChatStoreState = {
   hasReceivedPlanResponse: false,
   lastMessageMode: null,
   sessionCreditsUsed: 0,
+  runState: null,
 }
 
 export const useChatStore = create<ChatStore>()(
@@ -197,6 +197,11 @@ export const useChatStore = create<ChatStore>()(
         state.sessionCreditsUsed += credits
       }),
 
+    setRunState: (runState) =>
+      set((state) => {
+        state.runState = runState
+      }),
+
     reset: () =>
       set((state) => {
         state.messages = initialState.messages.slice()
@@ -215,6 +220,7 @@ export const useChatStore = create<ChatStore>()(
         state.hasReceivedPlanResponse = initialState.hasReceivedPlanResponse
         state.lastMessageMode = initialState.lastMessageMode
         state.sessionCreditsUsed = initialState.sessionCreditsUsed
+        state.runState = initialState.runState
       }),
   })),
 )

cli/src/state/feedback-store.ts

@@ -1,8 +1,5 @@
 import { create } from 'zustand'
 import { immer } from 'zustand/middleware/immer'
-import { enableMapSet } from 'immer'
-
-enableMapSet()
 
 interface FeedbackState {
   feedbackMessageId: string | null

cli/src/state/login-store.ts

@@ -1,4 +1,3 @@
-import { enableMapSet } from 'immer'
 import { create } from 'zustand'
 import { immer } from 'zustand/middleware/immer'
 
@@ -43,8 +42,6 @@ type LoginStoreActions = {
 
 type LoginStore = LoginStoreState & LoginStoreActions
 
-enableMapSet()
-
 const initialState: LoginStoreState = {
   loginUrl: null,
   loading: false,