Skip to content

Commit 2fc548f

Browse files
AlemTuzlakAlemTuzlak
committed
fix(react-core): enhance message handling by introducing isToolMessage utility and refactoring tests for clarity
1 parent 4a44b66 commit 2fc548f

2 files changed

Lines changed: 98 additions & 138 deletions

File tree

‎packages/react-core/src/v2/components/chat/CopilotChatMessageView.tsx‎

Lines changed: 14 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import {
88
AssistantMessage,
99
Message,
1010
ReasoningMessage,
11+
ToolMessage,
1112
UserMessage,
1213
} from "@ag-ui/core";
1314
import { twMerge } from "tailwind-merge";
@@ -16,6 +17,10 @@ import { getThreadClone } from "../../hooks/use-agent";
1617
import { useCopilotKit } from "../../providers/CopilotKitProvider";
1718
import { useCopilotChatConfiguration } from "../../providers/CopilotChatConfigurationProvider";
1819

20+
function isToolMessage(m: Message): m is ToolMessage {
21+
return m.role === "tool";
22+
}
23+
1924
/**
2025
* Memoized wrapper for assistant messages to prevent re-renders when other messages change.
2126
*/
@@ -69,22 +74,19 @@ const MemoizedAssistantMessage = React.memo(
6974
if (prevToolCalls && prevToolCalls.length > 0) {
7075
const toolCallIds = new Set(prevToolCalls.map((tc) => tc.id));
7176

72-
const prevToolResults = prevProps.messages.filter(
73-
(m) => m.role === "tool" && toolCallIds.has((m as any).toolCallId),
74-
);
75-
const nextToolResults = nextProps.messages.filter(
76-
(m) => m.role === "tool" && toolCallIds.has((m as any).toolCallId),
77-
);
77+
const prevToolResults = prevProps.messages
78+
.filter(isToolMessage)
79+
.filter((m) => toolCallIds.has(m.toolCallId));
80+
const nextToolResults = nextProps.messages
81+
.filter(isToolMessage)
82+
.filter((m) => toolCallIds.has(m.toolCallId));
7883

7984
// If number of tool results changed, re-render
8085
if (prevToolResults.length !== nextToolResults.length) return false;
8186

8287
// If any tool result content changed, re-render
8388
for (let i = 0; i < prevToolResults.length; i++) {
84-
if (
85-
(prevToolResults[i] as any).content !==
86-
(nextToolResults[i] as any).content
87-
)
89+
if (prevToolResults[i].content !== nextToolResults[i].content)
8890
return false;
8991
}
9092
}
@@ -371,13 +373,8 @@ export function CopilotChatMessageView({
371373
if (!existing) {
372374
messageAccumulator.set(message.id, message);
373375
} else if (message.role === "assistant" && existing.role === "assistant") {
374-
const content =
375-
(message as AssistantMessage).content ||
376-
(existing as AssistantMessage).content;
377-
messageAccumulator.set(message.id, {
378-
...message,
379-
content,
380-
} as AssistantMessage);
376+
const content = message.content || existing.content;
377+
messageAccumulator.set(message.id, { ...message, content });
381378
} else {
382379
messageAccumulator.set(message.id, message);
383380
}

‎packages/react-core/src/v2/components/chat/__tests__/CopilotChatMessageView.test.tsx‎

Lines changed: 84 additions & 121 deletions
Original file line numberDiff line numberDiff line change
@@ -5,42 +5,82 @@ import { vi } from "vitest";
55
import { CopilotKitProvider } from "../../../providers/CopilotKitProvider";
66
import { CopilotChatConfigurationProvider } from "../../../providers/CopilotChatConfigurationProvider";
77
import CopilotChatMessageView from "../CopilotChatMessageView";
8-
import {
8+
import type {
99
ActivityMessage,
1010
AssistantMessage,
1111
Message,
12+
ToolCall,
1213
UserMessage,
1314
} from "@ag-ui/core";
14-
import { ReactActivityMessageRenderer } from "../../../types";
15+
import type { ReactActivityMessageRenderer } from "../../../types";
16+
17+
// ---------------------------------------------------------------------------
18+
// Shared constants & helpers
19+
// ---------------------------------------------------------------------------
20+
21+
const AGENT_ID = "default";
22+
const THREAD_ID = "thread-test";
23+
24+
/** Typed factory — avoids `as UserMessage` casts everywhere. */
25+
function userMsg(id: string, content: string) {
26+
return { id, role: "user" as const, content };
27+
}
28+
29+
/** Typed factory — avoids `as AssistantMessage` casts everywhere. */
30+
function assistantMsg(
31+
id: string,
32+
content?: string,
33+
toolCalls?: ToolCall[],
34+
) {
35+
return { id, role: "assistant" as const, content, toolCalls };
36+
}
37+
38+
/** Typed factory — avoids `as ActivityMessage` casts everywhere. */
39+
function activityMsg(
40+
id: string,
41+
activityType: string,
42+
content: ActivityMessage["content"],
43+
) {
44+
return { id, role: "activity" as const, activityType, content };
45+
}
46+
47+
/** Typed factory — avoids `as any` casts on tool call objects. */
48+
function toolCall(id: string, name: string, args = "{}") {
49+
return {
50+
id,
51+
type: "function" as const,
52+
function: { name, arguments: args },
53+
};
54+
}
55+
56+
/**
57+
* Renders CopilotChatMessageView wrapped in the required providers.
58+
* Unified helper used by all describe blocks in this file.
59+
*/
60+
function renderMessageView({
61+
messages,
62+
renderActivityMessages,
63+
}: {
64+
messages: Message[];
65+
renderActivityMessages?: ReactActivityMessageRenderer<{ percent: number }>[];
66+
}) {
67+
return render(
68+
<CopilotKitProvider renderActivityMessages={renderActivityMessages}>
69+
<CopilotChatConfigurationProvider agentId={AGENT_ID} threadId={THREAD_ID}>
70+
<CopilotChatMessageView messages={messages} />
71+
</CopilotChatConfigurationProvider>
72+
</CopilotKitProvider>,
73+
);
74+
}
75+
76+
// ---------------------------------------------------------------------------
77+
// Tests
78+
// ---------------------------------------------------------------------------
1579

1680
describe("CopilotChatMessageView activity rendering", () => {
17-
const agentId = "default";
18-
const threadId = "thread-test";
19-
20-
function renderMessageView({
21-
messages,
22-
renderActivityMessages,
23-
}: {
24-
messages: Message[];
25-
renderActivityMessages?: ReactActivityMessageRenderer<any>[];
26-
}) {
27-
return render(
28-
<CopilotKitProvider renderActivityMessages={renderActivityMessages}>
29-
<CopilotChatConfigurationProvider agentId={agentId} threadId={threadId}>
30-
<CopilotChatMessageView messages={messages} />
31-
</CopilotChatConfigurationProvider>
32-
</CopilotKitProvider>,
33-
);
34-
}
35-
3681
it("renders activity messages via matching custom renderer", () => {
3782
const messages: Message[] = [
38-
{
39-
id: "act-1",
40-
role: "activity",
41-
activityType: "search-progress",
42-
content: { percent: 42 },
43-
} as ActivityMessage,
83+
activityMsg("act-1", "search-progress", { percent: 42 }),
4484
];
4585

4686
const renderers: ReactActivityMessageRenderer<{ percent: number }>[] = [
@@ -62,12 +102,7 @@ describe("CopilotChatMessageView activity rendering", () => {
62102

63103
it("skips rendering when no activity renderer matches", () => {
64104
const messages: Message[] = [
65-
{
66-
id: "act-2",
67-
role: "activity",
68-
activityType: "unknown-type",
69-
content: { message: "should not render" },
70-
} as ActivityMessage,
105+
activityMsg("act-2", "unknown-type", { message: "should not render" }),
71106
];
72107

73108
renderMessageView({ messages, renderActivityMessages: [] });
@@ -77,61 +112,17 @@ describe("CopilotChatMessageView activity rendering", () => {
77112
});
78113

79114
describe("CopilotChatMessageView duplicate message deduplication", () => {
80-
const agentId = "default";
81-
const threadId = "thread-test";
82-
83-
function renderMessageView({ messages }: { messages: Message[] }) {
84-
return render(
85-
<CopilotKitProvider>
86-
<CopilotChatConfigurationProvider agentId={agentId} threadId={threadId}>
87-
<CopilotChatMessageView messages={messages} />
88-
</CopilotChatConfigurationProvider>
89-
</CopilotKitProvider>,
90-
);
91-
}
92-
93115
it("preserves assistant text content when later duplicate has empty content (multi-tool-call scenario)", () => {
94116
const messages: Message[] = [
95-
{
96-
id: "user-1",
97-
role: "user",
98-
content: "Record a headache",
99-
} as UserMessage,
100-
{
101-
id: "assistant-1",
102-
role: "assistant",
103-
content: "Let me record that...",
104-
toolCalls: [],
105-
} as AssistantMessage,
106-
{
107-
id: "assistant-1",
108-
role: "assistant",
109-
content: "", // wiped by streaming tool-call update
110-
toolCalls: [
111-
{
112-
id: "tc-1",
113-
type: "function",
114-
function: { name: "captureData", arguments: "{}" },
115-
} as any,
116-
],
117-
} as AssistantMessage,
118-
{
119-
id: "assistant-1",
120-
role: "assistant",
121-
content: "", // still empty on second tool call
122-
toolCalls: [
123-
{
124-
id: "tc-1",
125-
type: "function",
126-
function: { name: "captureData", arguments: "{}" },
127-
} as any,
128-
{
129-
id: "tc-2",
130-
type: "function",
131-
function: { name: "updateMemory", arguments: "{}" },
132-
} as any,
133-
],
134-
} as AssistantMessage,
117+
userMsg("user-1", "Record a headache"),
118+
assistantMsg("assistant-1", "Let me record that..."),
119+
assistantMsg("assistant-1", "", [
120+
toolCall("tc-1", "captureData"),
121+
]),
122+
assistantMsg("assistant-1", "", [
123+
toolCall("tc-1", "captureData"),
124+
toolCall("tc-2", "updateMemory"),
125+
]),
135126
];
136127

137128
renderMessageView({ messages });
@@ -148,21 +139,9 @@ describe("CopilotChatMessageView duplicate message deduplication", () => {
148139

149140
it("deduplicates messages with the same id, keeping the last occurrence", () => {
150141
const messages: Message[] = [
151-
{
152-
id: "user-1",
153-
role: "user",
154-
content: "Hello",
155-
} as UserMessage,
156-
{
157-
id: "assistant-1",
158-
role: "assistant",
159-
content: "Partial response...",
160-
} as AssistantMessage,
161-
{
162-
id: "assistant-1",
163-
role: "assistant",
164-
content: "Full response from the assistant.",
165-
} as AssistantMessage,
142+
userMsg("user-1", "Hello"),
143+
assistantMsg("assistant-1", "Partial response..."),
144+
assistantMsg("assistant-1", "Full response from the assistant."),
166145
];
167146

168147
const consoleSpy = vi.spyOn(console, "error").mockImplementation(() => {});
@@ -191,26 +170,10 @@ describe("CopilotChatMessageView duplicate message deduplication", () => {
191170

192171
it("preserves order of unique messages", () => {
193172
const messages: Message[] = [
194-
{
195-
id: "user-1",
196-
role: "user",
197-
content: "First question",
198-
} as UserMessage,
199-
{
200-
id: "assistant-1",
201-
role: "assistant",
202-
content: "First answer",
203-
} as AssistantMessage,
204-
{
205-
id: "user-2",
206-
role: "user",
207-
content: "Second question",
208-
} as UserMessage,
209-
{
210-
id: "assistant-2",
211-
role: "assistant",
212-
content: "Second answer",
213-
} as AssistantMessage,
173+
userMsg("user-1", "First question"),
174+
assistantMsg("assistant-1", "First answer"),
175+
userMsg("user-2", "Second question"),
176+
assistantMsg("assistant-2", "Second answer"),
214177
];
215178

216179
renderMessageView({ messages });

0 commit comments

Comments
 (0)