Add shallow equality guards to MessageBlockStore to prevent unnecessary renders

cli/src/chat.tsx

@@ -49,7 +49,7 @@ import { useChatStreaming } from './hooks/use-chat-streaming'
 import { useChatUI } from './hooks/use-chat-ui'
 import { useClipboard } from './hooks/use-clipboard'
 import { useEvent } from './hooks/use-event'
-import { useGravityAd } from './hooks/use-gravity-ad'
+import { useGravityAd, type AdResponse } from './hooks/use-gravity-ad'
 import { useInputHistory } from './hooks/use-input-history'
 import { usePublishMutation } from './hooks/use-publish-mutation'
 import { useSuggestionEngine } from './hooks/use-suggestion-engine'
@@ -61,7 +61,10 @@ import { useChatStore } from './state/chat-store'
 import { useQueuePanelStore } from './state/queue-panel-store'
 import { useReviewStore } from './state/review-store'
 import { useFeedbackStore } from './state/feedback-store'
-import { useMessageBlockStore } from './state/message-block-store'
+import {
+  useMessageBlockStore,
+  EMPTY_RESPONSE_ADS,
+} from './state/message-block-store'
 import { usePublishStore } from './state/publish-store'
 import { reportActivity } from './utils/activity-tracker'
 import { stopActiveRun } from './utils/active-run'
@@ -1487,7 +1490,7 @@ export const Chat = ({
       isWaitingForResponse,
       timerStartTime,
       availableWidth: messageAvailableWidth,
-      responseAds: showInlineAds ? responseAds : {},
+      responseAds: showInlineAds ? responseAds : EMPTY_RESPONSE_ADS,
     })
   }, [
     theme,

cli/src/state/__tests__/message-block-store.test.ts

@@ -0,0 +1,83 @@
+import { describe, test, expect, beforeEach } from 'bun:test'
+
+import { useMessageBlockStore } from '../message-block-store'
+
+describe('MessageBlockStore equality guards', () => {
+  beforeEach(() => {
+    useMessageBlockStore.getState().reset()
+  })
+
+  test('notifies subscribers when context property actually changes', () => {
+    let notifications = 0
+    const unsub = useMessageBlockStore.subscribe(() => {
+      notifications++
+    })
+
+    useMessageBlockStore.getState().setContext({ availableWidth: 120 })
+    expect(notifications).toBe(1)
+    expect(useMessageBlockStore.getState().context.availableWidth).toBe(120)
+
+    unsub()
+  })
+
+  test('does NOT notify subscribers when setContext is called with identical values', () => {
+    let notifications = 0
+    const unsub = useMessageBlockStore.subscribe(() => {
+      notifications++
+    })
+
+    // availableWidth is already 80 in initialContext
+    useMessageBlockStore.getState().setContext({ availableWidth: 80 })
+    expect(notifications).toBe(0)
+
+    // Call with existing isWaitingForResponse: false
+    useMessageBlockStore.getState().setContext({ isWaitingForResponse: false })
+    expect(notifications).toBe(0)
+
+    unsub()
+  })
+
+  test('does NOT notify subscribers when setCallbacks is called with identical callbacks', () => {
+    let notifications = 0
+    const currentCallbacks = useMessageBlockStore.getState().callbacks
+
+    const unsub = useMessageBlockStore.subscribe(() => {
+      notifications++
+    })
+
+    // Passing identical callbacks reference dictionary
+    useMessageBlockStore.getState().setCallbacks({ ...currentCallbacks })
+    expect(notifications).toBe(0)
+
+    // Modifying one callback triggers notification
+    const newFn = () => {}
+    useMessageBlockStore.getState().setCallbacks({
+      ...currentCallbacks,
+      onBuildFast: newFn,
+    })
+    expect(notifications).toBe(1)
+    expect(useMessageBlockStore.getState().callbacks.onBuildFast).toBe(newFn)
+
+    unsub()
+  })
+
+  test('notifies subscribers when setCallbacks has fewer or different keys (full replacement semantics)', () => {
+    let notifications = 0
+    const currentCallbacks = useMessageBlockStore.getState().callbacks
+
+    const unsub = useMessageBlockStore.subscribe(() => {
+      notifications++
+    })
+
+    // Create a copy omitting one key
+    const subset = { ...currentCallbacks } as Partial<typeof currentCallbacks>
+    delete subset.onBuildFast
+
+    // Passing subset must count as changed because full replacement would remove onBuildFast
+    useMessageBlockStore.getState().setCallbacks(subset as any)
+    expect(notifications).toBe(1)
+    expect('onBuildFast' in useMessageBlockStore.getState().callbacks).toBe(false)
+
+    unsub()
+  })
+})

cli/src/state/message-block-store.ts

@@ -118,14 +118,16 @@ type MessageBlockStore = MessageBlockStoreState & MessageBlockStoreActions
 const noop = () => {}
 const noopFeedback: MessageBlockCallbacks['onFeedback'] = () => {}
 
+export const EMPTY_RESPONSE_ADS: Record<string, AdResponse[]> = {}
+
 const initialContext: MessageBlockContext = {
   theme: null,
   markdownPalette: null,
   messageTree: null,
   isWaitingForResponse: false,
   timerStartTime: null,
   availableWidth: 80,
-  responseAds: {},
+  responseAds: EMPTY_RESPONSE_ADS,
 }
 
 const initialCallbacks: MessageBlockCallbacks = {
@@ -154,11 +156,33 @@ export const useMessageBlockStore = create<MessageBlockStore>()(
 
     setContext: (updates) =>
       set((state) => {
+        let changed = false
+        for (const key of Object.keys(updates) as Array<keyof MessageBlockContext>) {
+          if (state.context[key] !== updates[key]) {
+            changed = true
+            break
+          }
+        }
+        if (!changed) return
         state.context = { ...state.context, ...updates }
       }),
 
     setCallbacks: (callbacks) =>
       set((state) => {
+        const prevKeys = Object.keys(state.callbacks)
+        const nextKeys = Object.keys(callbacks)
+        if (prevKeys.length !== nextKeys.length) {
+          state.callbacks = callbacks
+          return
+        }
+        let changed = false
+        for (const key of nextKeys as Array<keyof MessageBlockCallbacks>) {
+          if (state.callbacks[key] !== callbacks[key]) {
+            changed = true
+            break
+          }
+        }
+        if (!changed) return
         state.callbacks = callbacks
       }),