From 690792633c9f1a188611d0c55e1bd3863321bd8e Mon Sep 17 00:00:00 2001 From: Mikey Date: Wed, 2 Sep 2026 15:03:04 -0700 Subject: [PATCH 1/2] Add shallow equality guards to MessageBlockStore to prevent unnecessary renders --- cli/src/chat.tsx | 6 +- .../__tests__/message-block-store.test.ts | 63 +++++++++++++++++++ cli/src/state/message-block-store.ts | 28 +++++++++ 3 files changed, 95 insertions(+), 2 deletions(-) create mode 100644 cli/src/state/__tests__/message-block-store.test.ts diff --git a/cli/src/chat.tsx b/cli/src/chat.tsx index 50bad5c951..a923034dd3 100644 --- a/cli/src/chat.tsx +++ b/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' @@ -110,6 +110,8 @@ import type { UseMutationResult } from '@tanstack/react-query' import type { SponsoredProposalContentBlock } from './types/chat' import type { Dispatch, SetStateAction } from 'react' +const EMPTY_RESPONSE_ADS: Record = {} + export const Chat = ({ consumeInitialPrompt, fileTree, @@ -1487,7 +1489,7 @@ export const Chat = ({ isWaitingForResponse, timerStartTime, availableWidth: messageAvailableWidth, - responseAds: showInlineAds ? responseAds : {}, + responseAds: showInlineAds ? responseAds : EMPTY_RESPONSE_ADS, }) }, [ theme, diff --git a/cli/src/state/__tests__/message-block-store.test.ts b/cli/src/state/__tests__/message-block-store.test.ts new file mode 100644 index 0000000000..903b8da17a --- /dev/null +++ b/cli/src/state/__tests__/message-block-store.test.ts @@ -0,0 +1,63 @@ +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() + }) +}) diff --git a/cli/src/state/message-block-store.ts b/cli/src/state/message-block-store.ts index dd2acdecdf..80b52a4477 100644 --- a/cli/src/state/message-block-store.ts +++ b/cli/src/state/message-block-store.ts @@ -154,11 +154,39 @@ export const useMessageBlockStore = create()( setContext: (updates) => set((state) => { + let changed = false + for (const key of Object.keys(updates) as Array) { + const prev = state.context[key] + const next = updates[key] + if (prev === next) continue + if ( + key === 'responseAds' && + prev && + next && + typeof prev === 'object' && + typeof next === 'object' && + Object.keys(prev).length === 0 && + Object.keys(next).length === 0 + ) { + continue + } + changed = true + break + } + if (!changed) return state.context = { ...state.context, ...updates } }), setCallbacks: (callbacks) => set((state) => { + let changed = false + for (const key of Object.keys(callbacks) as Array) { + if (state.callbacks[key] !== callbacks[key]) { + changed = true + break + } + } + if (!changed) return state.callbacks = callbacks }), From dcd9684960891f21779b317b2d3d86c271006fa7 Mon Sep 17 00:00:00 2001 From: Mikey Date: Wed, 2 Sep 2026 15:08:09 -0700 Subject: [PATCH 2/2] Refine MessageBlockStore equality guards and setCallbacks replacement semantics --- cli/src/chat.tsx | 7 +++-- .../__tests__/message-block-store.test.ts | 20 +++++++++++++ cli/src/state/message-block-store.ts | 30 ++++++++----------- 3 files changed, 37 insertions(+), 20 deletions(-) diff --git a/cli/src/chat.tsx b/cli/src/chat.tsx index a923034dd3..4aab29aba1 100644 --- a/cli/src/chat.tsx +++ b/cli/src/chat.tsx @@ -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' @@ -110,8 +113,6 @@ import type { UseMutationResult } from '@tanstack/react-query' import type { SponsoredProposalContentBlock } from './types/chat' import type { Dispatch, SetStateAction } from 'react' -const EMPTY_RESPONSE_ADS: Record = {} - export const Chat = ({ consumeInitialPrompt, fileTree, diff --git a/cli/src/state/__tests__/message-block-store.test.ts b/cli/src/state/__tests__/message-block-store.test.ts index 903b8da17a..8f0b35ba0d 100644 --- a/cli/src/state/__tests__/message-block-store.test.ts +++ b/cli/src/state/__tests__/message-block-store.test.ts @@ -60,4 +60,24 @@ describe('MessageBlockStore equality guards', () => { 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 + 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() + }) }) diff --git a/cli/src/state/message-block-store.ts b/cli/src/state/message-block-store.ts index 80b52a4477..03c9128ce8 100644 --- a/cli/src/state/message-block-store.ts +++ b/cli/src/state/message-block-store.ts @@ -118,6 +118,8 @@ type MessageBlockStore = MessageBlockStoreState & MessageBlockStoreActions const noop = () => {} const noopFeedback: MessageBlockCallbacks['onFeedback'] = () => {} +export const EMPTY_RESPONSE_ADS: Record = {} + const initialContext: MessageBlockContext = { theme: null, markdownPalette: null, @@ -125,7 +127,7 @@ const initialContext: MessageBlockContext = { isWaitingForResponse: false, timerStartTime: null, availableWidth: 80, - responseAds: {}, + responseAds: EMPTY_RESPONSE_ADS, } const initialCallbacks: MessageBlockCallbacks = { @@ -156,22 +158,10 @@ export const useMessageBlockStore = create()( set((state) => { let changed = false for (const key of Object.keys(updates) as Array) { - const prev = state.context[key] - const next = updates[key] - if (prev === next) continue - if ( - key === 'responseAds' && - prev && - next && - typeof prev === 'object' && - typeof next === 'object' && - Object.keys(prev).length === 0 && - Object.keys(next).length === 0 - ) { - continue + if (state.context[key] !== updates[key]) { + changed = true + break } - changed = true - break } if (!changed) return state.context = { ...state.context, ...updates } @@ -179,8 +169,14 @@ export const useMessageBlockStore = create()( 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 Object.keys(callbacks) as Array) { + for (const key of nextKeys as Array) { if (state.callbacks[key] !== callbacks[key]) { changed = true break