Skip to content

Commit f4d0c54

Browse files
authored
🤖 perf: reduce ChatInput rerenders during streaming (#1123)
_Generated with `mux`_ ## What Reduce `ChatInput` re-renders during streaming by: - adding an explicit `React.memo` boundary for `ChatInput` - stabilizing `AIView → ChatInput` props (memoized auto-compaction result, stable review callbacks) - removing handler churn tied to streaming updates by reading latest workspace state via a ref ## How to verify - Start a long streaming response and profile renders: `ChatInput` should no longer commit on each delta.
1 parent fa10fc7 commit f4d0c54

File tree

2 files changed

+58
-25
lines changed

2 files changed

+58
-25
lines changed

‎src/browser/components/AIView.tsx‎

Lines changed: 54 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -181,6 +181,13 @@ const AIViewInner: React.FC<AIViewProps> = ({
181181
const [expandedBashGroups, setExpandedBashGroups] = useState<Set<string>>(new Set());
182182

183183
// Extract state from workspace state
184+
185+
// Keep a ref to the latest workspace state so event handlers (passed to memoized children)
186+
// can stay referentially stable during streaming while still reading fresh data.
187+
const workspaceStateRef = useRef(workspaceState);
188+
useEffect(() => {
189+
workspaceStateRef.current = workspaceState;
190+
}, [workspaceState]);
184191
const { messages, canInterrupt, isCompacting, loading, currentModel } = workspaceState;
185192

186193
// Apply message transformations:
@@ -203,11 +210,9 @@ const AIViewInner: React.FC<AIViewProps> = ({
203210
// Get active stream message ID for token counting
204211
const activeStreamMessageId = aggregator?.getActiveStreamMessageId();
205212

206-
const autoCompactionResult = checkAutoCompaction(
207-
workspaceUsage,
208-
pendingModel,
209-
use1M,
210-
autoCompactionThreshold / 100
213+
const autoCompactionResult = useMemo(
214+
() => checkAutoCompaction(workspaceUsage, pendingModel, use1M, autoCompactionThreshold / 100),
215+
[workspaceUsage, pendingModel, use1M, autoCompactionThreshold]
211216
);
212217

213218
// Show warning when: shouldShowWarning flag is true AND not currently compacting
@@ -265,7 +270,16 @@ const AIViewInner: React.FC<AIViewProps> = ({
265270

266271
// Handler for review notes from Code Review tab - adds review (starts attached)
267272
// Depend only on addReview (not whole reviews object) to keep callback stable
268-
const { addReview } = reviews;
273+
const { addReview, checkReview } = reviews;
274+
275+
const handleCheckReviews = useCallback(
276+
(ids: string[]) => {
277+
for (const id of ids) {
278+
checkReview(id);
279+
}
280+
},
281+
[checkReview]
282+
);
269283
const handleReviewNote = useCallback(
270284
(data: ReviewNoteData) => {
271285
addReview(data);
@@ -310,31 +324,47 @@ const AIViewInner: React.FC<AIViewProps> = ({
310324
}, [api, workspaceId, workspaceState?.queuedMessage, workspaceState?.canInterrupt]);
311325

312326
const handleEditLastUserMessage = useCallback(async () => {
313-
if (!workspaceState) return;
327+
const current = workspaceStateRef.current;
328+
if (!current) return;
329+
330+
if (current.queuedMessage) {
331+
const queuedMessage = current.queuedMessage;
332+
333+
await api?.workspace.clearQueue({ workspaceId });
334+
chatInputAPI.current?.restoreText(queuedMessage.content);
314335

315-
if (workspaceState.queuedMessage) {
316-
await handleEditQueuedMessage();
336+
// Restore images if present
337+
if (queuedMessage.imageParts && queuedMessage.imageParts.length > 0) {
338+
chatInputAPI.current?.restoreImages(queuedMessage.imageParts);
339+
}
317340
return;
318341
}
319342

320343
// Otherwise, edit last user message
321-
const transformedMessages = mergeConsecutiveStreamErrors(workspaceState.messages);
344+
const transformedMessages = mergeConsecutiveStreamErrors(current.messages);
322345
const lastUserMessage = [...transformedMessages]
323346
.reverse()
324347
.find((msg): msg is Extract<DisplayedMessage, { type: "user" }> => msg.type === "user");
325-
if (lastUserMessage) {
326-
setEditingMessage({ id: lastUserMessage.historyId, content: lastUserMessage.content });
327-
setAutoScroll(false); // Show jump-to-bottom indicator
328-
329-
// Scroll to the message being edited
330-
requestAnimationFrame(() => {
331-
const element = contentRef.current?.querySelector(
332-
`[data-message-id="${lastUserMessage.historyId}"]`
333-
);
334-
element?.scrollIntoView({ behavior: "smooth", block: "center" });
335-
});
348+
349+
if (!lastUserMessage) {
350+
return;
336351
}
337-
}, [workspaceState, contentRef, setAutoScroll, handleEditQueuedMessage]);
352+
353+
setEditingMessage({ id: lastUserMessage.historyId, content: lastUserMessage.content });
354+
setAutoScroll(false); // Show jump-to-bottom indicator
355+
356+
// Scroll to the message being edited
357+
requestAnimationFrame(() => {
358+
const element = contentRef.current?.querySelector(
359+
`[data-message-id="${lastUserMessage.historyId}"]`
360+
);
361+
element?.scrollIntoView({ behavior: "smooth", block: "center" });
362+
});
363+
}, [api, workspaceId, chatInputAPI, contentRef, setAutoScroll]);
364+
365+
const handleEditLastUserMessageClick = useCallback(() => {
366+
void handleEditLastUserMessage();
367+
}, [handleEditLastUserMessage]);
338368

339369
const handleCancelEdit = useCallback(() => {
340370
setEditingMessage(undefined);
@@ -740,14 +770,14 @@ const AIViewInner: React.FC<AIViewProps> = ({
740770
isCompacting={isCompacting}
741771
editingMessage={editingMessage}
742772
onCancelEdit={handleCancelEdit}
743-
onEditLastUserMessage={() => void handleEditLastUserMessage()}
773+
onEditLastUserMessage={handleEditLastUserMessageClick}
744774
canInterrupt={canInterrupt}
745775
onReady={handleChatInputReady}
746776
autoCompactionCheck={autoCompactionResult}
747777
attachedReviews={reviews.attachedReviews}
748778
onDetachReview={reviews.detachReview}
749779
onDetachAllReviews={reviews.detachAllAttached}
750-
onCheckReviews={(ids) => ids.forEach((id) => reviews.checkReview(id))}
780+
onCheckReviews={handleCheckReviews}
751781
onUpdateReviewNote={reviews.updateReviewNote}
752782
/>
753783
</div>

‎src/browser/components/ChatInput/index.tsx‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -116,7 +116,7 @@ import type { ImagePart } from "@/common/orpc/types";
116116

117117
export type { ChatInputProps, ChatInputAPI };
118118

119-
export const ChatInput: React.FC<ChatInputProps> = (props) => {
119+
const ChatInputInner: React.FC<ChatInputProps> = (props) => {
120120
const { api } = useAPI();
121121
const { variant } = props;
122122

@@ -1632,6 +1632,9 @@ export const ChatInput: React.FC<ChatInputProps> = (props) => {
16321632
);
16331633
};
16341634

1635+
export const ChatInput = React.memo(ChatInputInner);
1636+
ChatInput.displayName = "ChatInput";
1637+
16351638
const TokenCountDisplay: React.FC<{ reader: TokenCountReader }> = ({ reader }) => {
16361639
const tokens = reader();
16371640
if (!tokens) {

0 commit comments

Comments
 (0)