Skip to content

Commit 5b056ef

Browse files
Ignore empty chat assistant placeholders (#5620)
Skip empty finished assistant messages during chat persistence and loading so stale placeholder rows cannot make existing chats appear stuck.
1 parent 33726ef commit 5b056ef

9 files changed

Lines changed: 399 additions & 23 deletions

File tree

Lines changed: 235 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,235 @@
1+
import { cleanup, fireEvent, render, screen } from "@testing-library/react";
2+
import { beforeEach, describe, expect, it, vi } from "vitest";
3+
4+
const mocks = vi.hoisted(() => ({
5+
chatRegenerate: vi.fn(),
6+
chatSendMessage: vi.fn(),
7+
chatSetMessages: vi.fn(),
8+
chatStop: vi.fn(),
9+
messages: [] as unknown[],
10+
store: null as unknown,
11+
}));
12+
13+
vi.mock("@ai-sdk/react", () => ({
14+
useChat: () => ({
15+
messages: mocks.messages,
16+
sendMessage: mocks.chatSendMessage,
17+
regenerate: mocks.chatRegenerate,
18+
stop: mocks.chatStop,
19+
status: "ready",
20+
error: undefined,
21+
setMessages: mocks.chatSetMessages,
22+
}),
23+
}));
24+
25+
vi.mock("~/chat/context/use-chat-context-pipeline", () => ({
26+
useChatContextPipeline: () => ({
27+
contextEntities: [],
28+
pendingRefs: [],
29+
}),
30+
}));
31+
32+
vi.mock("~/chat/transport/use-transport", () => ({
33+
useTransport: () => ({
34+
transport: {},
35+
isSystemPromptReady: true,
36+
}),
37+
}));
38+
39+
vi.mock("~/store/tinybase/store/main", () => ({
40+
STORE_ID: "main",
41+
UI: {
42+
useStore: () => mocks.store,
43+
useValues: () => ({ user_id: "user-1" }),
44+
},
45+
}));
46+
47+
import { ChatSession } from "./session-provider";
48+
49+
import { buildPersistedChatMessageRow } from "~/chat/store/persisted-messages";
50+
import type { HyprUIMessage } from "~/chat/types";
51+
52+
type FakeStore = ReturnType<typeof createStore>;
53+
54+
function createStore(rows: Record<string, Record<string, unknown>>) {
55+
const table = new Map(Object.entries(rows));
56+
57+
return {
58+
delRow: vi.fn((tableName: string, rowId: string) => {
59+
if (tableName === "chat_messages") {
60+
table.delete(rowId);
61+
}
62+
}),
63+
forEachRow: vi.fn(
64+
(
65+
tableName: string,
66+
callback: (rowId: string, forEachCell: () => void) => void,
67+
) => {
68+
if (tableName !== "chat_messages") {
69+
return;
70+
}
71+
table.forEach((_row, rowId) => callback(rowId, () => {}));
72+
},
73+
),
74+
getRow: vi.fn((tableName: string, rowId: string) => {
75+
if (tableName !== "chat_messages") {
76+
return undefined;
77+
}
78+
return table.get(rowId);
79+
}),
80+
setRow: vi.fn(
81+
(tableName: string, rowId: string, row: Record<string, unknown>) => {
82+
if (tableName === "chat_messages") {
83+
table.set(rowId, row);
84+
}
85+
},
86+
),
87+
transaction: vi.fn((callback: () => void) => callback()),
88+
};
89+
}
90+
91+
function persistedAssistantRow(message: HyprUIMessage) {
92+
return buildPersistedChatMessageRow({
93+
message,
94+
chatGroupId: "group-1",
95+
userId: "user-1",
96+
status: "ready",
97+
}) as unknown as Record<string, unknown>;
98+
}
99+
100+
function renderSession() {
101+
render(
102+
<ChatSession chatGroupId="group-1" sessionId="session-1">
103+
{({ regenerate }) => (
104+
<button type="button" onClick={regenerate}>
105+
Regenerate
106+
</button>
107+
)}
108+
</ChatSession>,
109+
);
110+
}
111+
112+
describe("ChatSession", () => {
113+
beforeEach(() => {
114+
cleanup();
115+
mocks.chatRegenerate.mockClear();
116+
mocks.chatSendMessage.mockClear();
117+
mocks.chatSetMessages.mockClear();
118+
mocks.chatStop.mockClear();
119+
mocks.messages = [];
120+
mocks.store = createStore({});
121+
});
122+
123+
it("does not delete the previous persisted assistant when retrying an unpersisted empty assistant", () => {
124+
const previousAssistant: HyprUIMessage = {
125+
id: "assistant-previous",
126+
role: "assistant",
127+
parts: [{ type: "text", text: "Previous answer" }],
128+
metadata: { createdAt: Date.parse("2024-01-01T00:00:01Z") },
129+
};
130+
const store = createStore({
131+
"assistant-previous": persistedAssistantRow(previousAssistant),
132+
});
133+
mocks.store = store;
134+
mocks.messages = [
135+
{
136+
id: "user-1",
137+
role: "user",
138+
parts: [{ type: "text", text: "First question" }],
139+
},
140+
previousAssistant,
141+
{
142+
id: "user-2",
143+
role: "user",
144+
parts: [{ type: "text", text: "Retry question" }],
145+
},
146+
{
147+
id: "assistant-empty",
148+
role: "assistant",
149+
parts: [],
150+
},
151+
];
152+
153+
renderSession();
154+
155+
fireEvent.click(screen.getByRole("button", { name: "Regenerate" }));
156+
157+
expect(store.delRow).toHaveBeenCalledWith(
158+
"chat_messages",
159+
"assistant-empty",
160+
);
161+
expect(
162+
(store as FakeStore).getRow("chat_messages", "assistant-previous"),
163+
).toBeDefined();
164+
expect(mocks.chatRegenerate).toHaveBeenCalledTimes(1);
165+
});
166+
167+
it("deletes the persisted row for the in-memory assistant being regenerated", () => {
168+
const assistant: HyprUIMessage = {
169+
id: "assistant-current",
170+
role: "assistant",
171+
parts: [{ type: "text", text: "Current answer" }],
172+
metadata: { createdAt: Date.parse("2024-01-01T00:00:01Z") },
173+
};
174+
const store = createStore({
175+
"assistant-current": persistedAssistantRow(assistant),
176+
});
177+
mocks.store = store;
178+
mocks.messages = [
179+
{
180+
id: "user-1",
181+
role: "user",
182+
parts: [{ type: "text", text: "Question" }],
183+
},
184+
assistant,
185+
];
186+
187+
renderSession();
188+
189+
fireEvent.click(screen.getByRole("button", { name: "Regenerate" }));
190+
191+
expect(store.delRow).toHaveBeenCalledWith(
192+
"chat_messages",
193+
"assistant-current",
194+
);
195+
expect(store.getRow("chat_messages", "assistant-current")).toBeUndefined();
196+
expect(mocks.chatRegenerate).toHaveBeenCalledTimes(1);
197+
});
198+
199+
it("deletes the last assistant row when a trailing user message is present", () => {
200+
const assistant: HyprUIMessage = {
201+
id: "assistant-current",
202+
role: "assistant",
203+
parts: [{ type: "text", text: "Current answer" }],
204+
metadata: { createdAt: Date.parse("2024-01-01T00:00:01Z") },
205+
};
206+
const store = createStore({
207+
"assistant-current": persistedAssistantRow(assistant),
208+
});
209+
mocks.store = store;
210+
mocks.messages = [
211+
{
212+
id: "user-1",
213+
role: "user",
214+
parts: [{ type: "text", text: "Question" }],
215+
},
216+
assistant,
217+
{
218+
id: "user-2",
219+
role: "user",
220+
parts: [{ type: "text", text: "Follow-up" }],
221+
},
222+
];
223+
224+
renderSession();
225+
226+
fireEvent.click(screen.getByRole("button", { name: "Regenerate" }));
227+
228+
expect(store.delRow).toHaveBeenCalledWith(
229+
"chat_messages",
230+
"assistant-current",
231+
);
232+
expect(store.getRow("chat_messages", "assistant-current")).toBeUndefined();
233+
expect(mocks.chatRegenerate).toHaveBeenCalledTimes(1);
234+
});
235+
});

apps/desktop/src/chat/components/session-provider.tsx

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ import {
1818
buildPersistedChatMessageRow,
1919
getPersistedChatMessages,
2020
getVisibleChatMessages,
21+
shouldPersistFinishedMessage,
2122
} from "~/chat/store/persisted-messages";
2223
import { stripEphemeralToolContext } from "~/chat/tools/strip-ephemeral-tool-context";
2324
import { useTransport } from "~/chat/transport/use-transport";
@@ -117,6 +118,10 @@ export function ChatSession({
117118
sanitizedParts === message.parts
118119
? message
119120
: { ...message, parts: sanitizedParts };
121+
if (!shouldPersistFinishedMessage(sanitizedMessage)) {
122+
store.delRow("chat_messages", sanitizedMessage.id);
123+
return;
124+
}
120125
store.setRow(
121126
"chat_messages",
122127
sanitizedMessage.id,
@@ -142,12 +147,16 @@ export function ChatSession({
142147

143148
const regenerate = useCallback(() => {
144149
if (!store || !chatGroupId) return;
145-
const last = [...getPersistedChatMessages(store, chatGroupId)]
146-
.reverse()
147-
.find((m) => m.message.role === "assistant");
148-
if (last) store.delRow("chat_messages", last.id);
150+
for (let i = messages.length - 1; i >= 0; i--) {
151+
if (messages[i].role !== "assistant") {
152+
continue;
153+
}
154+
155+
store.delRow("chat_messages", messages[i].id);
156+
break;
157+
}
149158
void chatRegenerate();
150-
}, [store, chatGroupId, chatRegenerate]);
159+
}, [store, chatGroupId, messages, chatRegenerate]);
151160

152161
const setMessages = useCallback(
153162
(next: HyprUIMessage[] | ((prev: HyprUIMessage[]) => HyprUIMessage[])) => {

apps/desktop/src/chat/components/shared.test.ts

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,4 +22,14 @@ describe("hasRenderableContent", () => {
2222
}),
2323
).toBe(true);
2424
});
25+
26+
test("returns false for blank text-only messages", () => {
27+
expect(
28+
hasRenderableContent({
29+
id: "message-3",
30+
role: "assistant",
31+
parts: [{ type: "text", text: " " }],
32+
}),
33+
).toBe(false);
34+
});
2535
});
Lines changed: 1 addition & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1 @@
1-
import type { HyprUIMessage } from "~/chat/types";
2-
3-
export function hasRenderableContent(message: HyprUIMessage): boolean {
4-
return message.parts.some((part) => {
5-
if (part.type === "step-start") {
6-
return false;
7-
}
8-
9-
if (part.type === "reasoning") {
10-
return part.text.trim().length > 0;
11-
}
12-
13-
return true;
14-
});
15-
}
1+
export { hasRenderableContent } from "~/chat/message-content";
Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
import type { HyprUIMessage } from "~/chat/types";
2+
3+
function isRecord(value: unknown): value is Record<string, unknown> {
4+
return typeof value === "object" && value !== null;
5+
}
6+
7+
export function hasRenderableParts(parts: unknown): boolean {
8+
if (!Array.isArray(parts)) {
9+
return false;
10+
}
11+
12+
return parts.some((part) => {
13+
if (!isRecord(part) || typeof part.type !== "string") {
14+
return false;
15+
}
16+
17+
if (part.type === "step-start") {
18+
return false;
19+
}
20+
21+
if (part.type === "reasoning" || part.type === "text") {
22+
return typeof part.text === "string" && part.text.trim().length > 0;
23+
}
24+
25+
return true;
26+
});
27+
}
28+
29+
export function hasRenderableContent(message: HyprUIMessage) {
30+
return hasRenderableParts(message.parts);
31+
}

0 commit comments

Comments
 (0)