Skip to content

Commit 6e08cba

Browse files
mmemme
authored
Merge pull request #46 from CopilotKit/mme/improve-rerenders
Fix re-rendering issues
2 parents 0fd1428 + 2397c56 commit 6e08cba

File tree

5 files changed

+2747
-101
lines changed

5 files changed

+2747
-101
lines changed

‎packages/react/src/components/chat/CopilotChatMessageView.tsx‎

Lines changed: 202 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,162 @@
1+
import React from "react";
12
import { WithSlots, renderSlot } from "@/lib/slots";
23
import CopilotChatAssistantMessage from "./CopilotChatAssistantMessage";
34
import CopilotChatUserMessage from "./CopilotChatUserMessage";
4-
import { Message } from "@ag-ui/core";
5+
import { ActivityMessage, AssistantMessage, Message, UserMessage } from "@ag-ui/core";
56
import { twMerge } from "tailwind-merge";
67
import { useRenderActivityMessage, useRenderCustomMessages } from "@/hooks";
78

9+
/**
10+
* Memoized wrapper for assistant messages to prevent re-renders when other messages change.
11+
*/
12+
const MemoizedAssistantMessage = React.memo(
13+
function MemoizedAssistantMessage({
14+
message,
15+
messages,
16+
isRunning,
17+
AssistantMessageComponent,
18+
}: {
19+
message: AssistantMessage;
20+
messages: Message[];
21+
isRunning: boolean;
22+
AssistantMessageComponent: typeof CopilotChatAssistantMessage;
23+
}) {
24+
return (
25+
<AssistantMessageComponent
26+
message={message}
27+
messages={messages}
28+
isRunning={isRunning}
29+
/>
30+
);
31+
},
32+
(prevProps, nextProps) => {
33+
// Only re-render if this specific message changed
34+
if (prevProps.message.id !== nextProps.message.id) return false;
35+
if (prevProps.message.content !== nextProps.message.content) return false;
36+
37+
// Compare tool calls if present
38+
const prevToolCalls = prevProps.message.toolCalls;
39+
const nextToolCalls = nextProps.message.toolCalls;
40+
if (prevToolCalls?.length !== nextToolCalls?.length) return false;
41+
if (prevToolCalls && nextToolCalls) {
42+
for (let i = 0; i < prevToolCalls.length; i++) {
43+
const prevTc = prevToolCalls[i];
44+
const nextTc = nextToolCalls[i];
45+
if (!prevTc || !nextTc) return false;
46+
if (prevTc.id !== nextTc.id) return false;
47+
if (prevTc.function.arguments !== nextTc.function.arguments) return false;
48+
}
49+
}
50+
51+
// Check if tool results changed for this message's tool calls
52+
// Tool results are separate messages with role="tool" that reference tool call IDs
53+
if (prevToolCalls && prevToolCalls.length > 0) {
54+
const toolCallIds = new Set(prevToolCalls.map(tc => tc.id));
55+
56+
const prevToolResults = prevProps.messages.filter(
57+
m => m.role === "tool" && toolCallIds.has((m as any).toolCallId)
58+
);
59+
const nextToolResults = nextProps.messages.filter(
60+
m => m.role === "tool" && toolCallIds.has((m as any).toolCallId)
61+
);
62+
63+
// If number of tool results changed, re-render
64+
if (prevToolResults.length !== nextToolResults.length) return false;
65+
66+
// If any tool result content changed, re-render
67+
for (let i = 0; i < prevToolResults.length; i++) {
68+
if ((prevToolResults[i] as any).content !== (nextToolResults[i] as any).content) return false;
69+
}
70+
}
71+
72+
// Only care about isRunning if this message is CURRENTLY the latest
73+
// (we don't need to re-render just because a message stopped being the latest)
74+
const nextIsLatest = nextProps.messages[nextProps.messages.length - 1]?.id === nextProps.message.id;
75+
if (nextIsLatest && prevProps.isRunning !== nextProps.isRunning) return false;
76+
77+
// Check if component reference changed
78+
if (prevProps.AssistantMessageComponent !== nextProps.AssistantMessageComponent) return false;
79+
80+
return true;
81+
}
82+
);
83+
84+
/**
85+
* Memoized wrapper for user messages to prevent re-renders when other messages change.
86+
*/
87+
const MemoizedUserMessage = React.memo(
88+
function MemoizedUserMessage({
89+
message,
90+
UserMessageComponent,
91+
}: {
92+
message: UserMessage;
93+
UserMessageComponent: typeof CopilotChatUserMessage;
94+
}) {
95+
return <UserMessageComponent message={message} />;
96+
},
97+
(prevProps, nextProps) => {
98+
// Only re-render if this specific message changed
99+
if (prevProps.message.id !== nextProps.message.id) return false;
100+
if (prevProps.message.content !== nextProps.message.content) return false;
101+
if (prevProps.UserMessageComponent !== nextProps.UserMessageComponent) return false;
102+
return true;
103+
}
104+
);
105+
106+
/**
107+
* Memoized wrapper for activity messages to prevent re-renders when other messages change.
108+
*/
109+
const MemoizedActivityMessage = React.memo(
110+
function MemoizedActivityMessage({
111+
message,
112+
renderActivityMessage,
113+
}: {
114+
message: ActivityMessage;
115+
renderActivityMessage: (message: ActivityMessage) => React.ReactElement | null;
116+
}) {
117+
return renderActivityMessage(message);
118+
},
119+
(prevProps, nextProps) => {
120+
// Only re-render if this specific activity message changed
121+
if (prevProps.message.id !== nextProps.message.id) return false;
122+
if (prevProps.message.activityType !== nextProps.message.activityType) return false;
123+
// Compare content - need to stringify since it's an object
124+
if (JSON.stringify(prevProps.message.content) !== JSON.stringify(nextProps.message.content)) return false;
125+
// Note: We don't compare renderActivityMessage function reference because it changes
126+
// frequently due to useCallback dependencies in useRenderActivityMessage.
127+
// The message content comparison is sufficient to determine if a re-render is needed.
128+
return true;
129+
}
130+
);
131+
132+
/**
133+
* Memoized wrapper for custom messages to prevent re-renders when other messages change.
134+
*/
135+
const MemoizedCustomMessage = React.memo(
136+
function MemoizedCustomMessage({
137+
message,
138+
position,
139+
renderCustomMessage,
140+
}: {
141+
message: Message;
142+
position: "before" | "after";
143+
renderCustomMessage: (params: { message: Message; position: "before" | "after" }) => React.ReactElement | null;
144+
}) {
145+
return renderCustomMessage({ message, position });
146+
},
147+
(prevProps, nextProps) => {
148+
// Only re-render if the message or position changed
149+
if (prevProps.message.id !== nextProps.message.id) return false;
150+
if (prevProps.position !== nextProps.position) return false;
151+
// Compare message content - for assistant messages this is a string, for others may differ
152+
if (prevProps.message.content !== nextProps.message.content) return false;
153+
if (prevProps.message.role !== nextProps.message.role) return false;
154+
// Note: We don't compare renderCustomMessage function reference because it changes
155+
// frequently. The message content comparison is sufficient to determine if a re-render is needed.
156+
return true;
157+
}
158+
);
159+
8160
export type CopilotChatMessageViewProps = Omit<
9161
WithSlots<
10162
{
@@ -43,48 +195,71 @@ export function CopilotChatMessageView({
43195
.flatMap((message) => {
44196
const elements: (React.ReactElement | null | undefined)[] = [];
45197

46-
// Render custom message before
198+
// Render custom message before (using memoized wrapper)
47199
if (renderCustomMessage) {
48200
elements.push(
49-
renderCustomMessage({
50-
message,
51-
position: "before",
52-
}),
201+
<MemoizedCustomMessage
202+
key={`${message.id}-custom-before`}
203+
message={message}
204+
position="before"
205+
renderCustomMessage={renderCustomMessage}
206+
/>
53207
);
54208
}
55209

56-
// Render the main message
210+
// Render the main message using memoized wrappers to prevent unnecessary re-renders
57211
if (message.role === "assistant") {
212+
// Determine the component to use (custom slot or default)
213+
const AssistantComponent = (
214+
typeof assistantMessage === "function"
215+
? assistantMessage
216+
: CopilotChatAssistantMessage
217+
) as typeof CopilotChatAssistantMessage;
218+
58219
elements.push(
59-
renderSlot(assistantMessage, CopilotChatAssistantMessage, {
60-
key: message.id,
61-
message,
62-
messages,
63-
isRunning,
64-
}),
220+
<MemoizedAssistantMessage
221+
key={message.id}
222+
message={message as AssistantMessage}
223+
messages={messages}
224+
isRunning={isRunning}
225+
AssistantMessageComponent={AssistantComponent}
226+
/>
65227
);
66228
} else if (message.role === "user") {
229+
// Determine the component to use (custom slot or default)
230+
const UserComponent = (
231+
typeof userMessage === "function"
232+
? userMessage
233+
: CopilotChatUserMessage
234+
) as typeof CopilotChatUserMessage;
235+
67236
elements.push(
68-
renderSlot(userMessage, CopilotChatUserMessage, {
69-
key: message.id,
70-
message,
71-
}),
237+
<MemoizedUserMessage
238+
key={message.id}
239+
message={message as UserMessage}
240+
UserMessageComponent={UserComponent}
241+
/>
72242
);
73243
} else if (message.role === "activity") {
74-
const renderedActivity = renderActivityMessage(message);
75-
76-
if (renderedActivity) {
77-
elements.push(renderedActivity);
78-
}
244+
// Use memoized wrapper to prevent re-renders when other messages change
245+
elements.push(
246+
<MemoizedActivityMessage
247+
key={message.id}
248+
message={message as ActivityMessage}
249+
renderActivityMessage={renderActivityMessage}
250+
/>
251+
);
79252
}
80253

81-
// Render custom message after
254+
// Render custom message after (using memoized wrapper)
82255
if (renderCustomMessage) {
83256
elements.push(
84-
renderCustomMessage({
85-
message,
86-
position: "after",
87-
}),
257+
<MemoizedCustomMessage
258+
key={`${message.id}-custom-after`}
259+
message={message}
260+
position="after"
261+
renderCustomMessage={renderCustomMessage}
262+
/>
88263
);
89264
}
90265

‎packages/react/src/components/chat/CopilotChatView.tsx‎

Lines changed: 12 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import React, { useRef, useState, useEffect } from "react";
2-
import { WithSlots, renderSlot } from "@/lib/slots";
2+
import { WithSlots, SlotValue, renderSlot } from "@/lib/slots";
33
import CopilotChatMessageView from "./CopilotChatMessageView";
44
import CopilotChatInput, { CopilotChatInputProps } from "./CopilotChatInput";
55
import CopilotChatSuggestionView, { CopilotChatSuggestionViewProps } from "./CopilotChatSuggestionView";
@@ -113,20 +113,19 @@ export function CopilotChatView({
113113
});
114114

115115
const BoundInput = renderSlot(input, CopilotChatInput, (inputProps ?? {}) as CopilotChatInputProps);
116+
116117
const hasSuggestions = Array.isArray(suggestions) && suggestions.length > 0;
117118
const BoundSuggestionView = hasSuggestions
118-
? renderSlot<typeof CopilotChatSuggestionView, CopilotChatSuggestionViewProps>(
119-
suggestionView,
120-
CopilotChatSuggestionView,
121-
{
122-
suggestions,
123-
loadingIndexes: suggestionLoadingIndexes,
124-
onSelectSuggestion,
125-
className: "mb-3 lg:ml-4 lg:mr-4 ml-0 mr-0",
126-
},
127-
)
119+
? renderSlot(suggestionView, CopilotChatSuggestionView, {
120+
suggestions,
121+
loadingIndexes: suggestionLoadingIndexes,
122+
onSelectSuggestion,
123+
className: "mb-3 lg:ml-4 lg:mr-4 ml-0 mr-0",
124+
})
128125
: null;
126+
129127
const BoundFeather = renderSlot(feather, CopilotChatView.Feather, {});
128+
130129
const BoundScrollView = renderSlot(scrollView, CopilotChatView.ScrollView, {
131130
autoScroll,
132131
scrollToBottomButton,
@@ -187,7 +186,7 @@ export namespace CopilotChatView {
187186
// Inner component that has access to StickToBottom context
188187
const ScrollContent: React.FC<{
189188
children: React.ReactNode;
190-
scrollToBottomButton?: React.FC<React.ButtonHTMLAttributes<HTMLButtonElement>>;
189+
scrollToBottomButton?: SlotValue<React.FC<React.ButtonHTMLAttributes<HTMLButtonElement>>>;
191190
inputContainerHeight: number;
192191
isResizing: boolean;
193192
}> = ({ children, scrollToBottomButton, inputContainerHeight, isResizing }) => {
@@ -219,7 +218,7 @@ export namespace CopilotChatView {
219218
export const ScrollView: React.FC<
220219
React.HTMLAttributes<HTMLDivElement> & {
221220
autoScroll?: boolean;
222-
scrollToBottomButton?: React.FC<React.ButtonHTMLAttributes<HTMLButtonElement>>;
221+
scrollToBottomButton?: SlotValue<React.FC<React.ButtonHTMLAttributes<HTMLButtonElement>>>;
223222
inputContainerHeight?: number;
224223
isResizing?: boolean;
225224
}

0 commit comments

Comments
 (0)