Skip to content

Commit 5823d8f

Browse files
kfaracikclaude
andcommitted
Merge branch 'main' into feat/221-rag-sources
Reconciles the RAG sources/citations stack with the conversation forking feature (#227): MessageItem carries both the Sources button and the Copy/Fork action row, Messages renders branch markers next to RAG-annotated messages, and the completed assistant message in llmStore now keeps its persisted id together with the cited sources. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2 parents ba68d59 + ec8fe09 commit 5823d8f

32 files changed

Lines changed: 1774 additions & 153 deletions

__tests__/DrawerMenu.test.tsx

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,10 +46,12 @@ const mockChats = [
4646

4747
const mockRenameChat = jest.fn();
4848
const mockDeleteChat = jest.fn();
49+
let mockPhantomChat: { id: number } | null = null;
4950
jest.mock('../store/chatStore', () => ({
5051
useChatStore: jest.fn(() => ({
5152
chats: mockChats,
52-
phantomChat: null,
53+
phantomChat: mockPhantomChat,
54+
getChatById: (id: number) => mockChats.find((chat) => chat.id === id),
5355
renameChat: mockRenameChat,
5456
deleteChat: mockDeleteChat,
5557
})),
@@ -90,6 +92,7 @@ const renderMenu = (props: Partial<MenuProps> = {}) =>
9092
beforeEach(() => {
9193
jest.clearAllMocks();
9294
mockPathname = '/';
95+
mockPhantomChat = null;
9396
setPlatform('ios');
9497
});
9598

@@ -177,6 +180,28 @@ describe('DrawerMenu — collapsed', () => {
177180
expect(onNavigate).toHaveBeenCalled();
178181
});
179182

183+
it('only closes the drawer when already on the new chat screen', () => {
184+
mockPhantomChat = { id: 4 };
185+
mockPathname = '/chat/4';
186+
const onNavigate = jest.fn();
187+
renderMenu({ onNavigate });
188+
189+
fireEvent.press(screen.getByTestId('drawer-new-chat'));
190+
191+
expect(mockStartPhantomChat).not.toHaveBeenCalled();
192+
expect(onNavigate).toHaveBeenCalled();
193+
});
194+
195+
it('starts a phantom chat when a forked chat claimed the phantom id', () => {
196+
mockPhantomChat = { id: 3 };
197+
mockPathname = '/chat/3';
198+
renderMenu();
199+
200+
fireEvent.press(screen.getByTestId('drawer-new-chat'));
201+
202+
expect(mockStartPhantomChat).toHaveBeenCalledWith({}, 'replace');
203+
});
204+
180205
it('navigates to the chat when an item is pressed', () => {
181206
renderMenu();
182207

__tests__/MessageItem.test.tsx

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,11 +96,20 @@ const setLLMState = (state: MockLLMState) =>
9696
selector ? selector(state) : state
9797
);
9898

99+
const baseMessage = {
100+
id: 1,
101+
role: 'assistant',
102+
content: 'Hello world',
103+
chatId: 1,
104+
timestamp: 0,
105+
} as React.ComponentProps<typeof MessageItem>['message'];
106+
99107
const renderItem = (
100108
props: Partial<React.ComponentProps<typeof MessageItem>> = {}
101109
) =>
102110
render(
103111
<MessageItem
112+
message={baseMessage}
104113
content="Hello world"
105114
role="assistant"
106115
isLastMessage={false}

__tests__/chatRepository.test.ts

Lines changed: 238 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,10 @@
1-
import { persistMessage } from '../database/chatRepository';
1+
import { forkChat, persistMessage } from '../database/chatRepository';
22
import type { SQLiteDatabase } from 'expo-sqlite';
33

4+
type TransactionCallback = Parameters<
5+
SQLiteDatabase['withTransactionAsync']
6+
>[0];
7+
48
jest.mock('expo-sqlite', () => ({
59
useSQLiteContext: jest.fn(() => ({})),
610
}));
@@ -67,3 +71,236 @@ describe('persistMessage with imagePath', () => {
6771
);
6872
});
6973
});
74+
75+
describe('forkChat', () => {
76+
it('creates a branch chat and copies messages up to the target message', async () => {
77+
const runAsync = jest
78+
.fn()
79+
.mockResolvedValueOnce({ lastInsertRowId: 10 })
80+
.mockResolvedValueOnce({ lastInsertRowId: 101 })
81+
.mockResolvedValueOnce({ lastInsertRowId: 102 })
82+
.mockResolvedValue({ lastInsertRowId: 0 });
83+
const getFirstAsync = jest.fn().mockResolvedValue({
84+
id: 1,
85+
title: 'Original',
86+
modelId: 7,
87+
lastUsed: 1,
88+
});
89+
const getAllAsync = jest
90+
.fn()
91+
.mockResolvedValueOnce([
92+
{
93+
id: 1,
94+
chatId: 1,
95+
role: 'user',
96+
content: 'one',
97+
timestamp: 100,
98+
},
99+
{
100+
id: 2,
101+
chatId: 1,
102+
role: 'assistant',
103+
content: 'two',
104+
timestamp: 200,
105+
modelName: 'model',
106+
},
107+
{
108+
id: 3,
109+
chatId: 1,
110+
role: 'user',
111+
content: 'three',
112+
timestamp: 300,
113+
},
114+
])
115+
.mockResolvedValueOnce([]);
116+
const mockDb = {
117+
runAsync,
118+
getFirstAsync,
119+
getAllAsync,
120+
withTransactionAsync: async (callback: TransactionCallback) => callback(),
121+
} as SQLiteDatabase;
122+
123+
const newChatId = await forkChat(mockDb, 1, 2);
124+
125+
expect(newChatId).toBe(10);
126+
expect(runAsync).toHaveBeenCalledWith(
127+
`INSERT INTO chats (title, modelId, lastUsed) VALUES (?, ?, ?)`,
128+
['Original', 7, expect.any(Number)]
129+
);
130+
expect(runAsync).toHaveBeenCalledWith(
131+
expect.stringContaining('INSERT INTO messages'),
132+
[10, 'user', 'one', 100, '', 0, 0, null, null]
133+
);
134+
expect(runAsync).toHaveBeenCalledWith(
135+
expect.stringContaining('INSERT INTO messages'),
136+
[10, 'assistant', 'two', 200, 'model', 0, 0, null, null]
137+
);
138+
expect(runAsync).toHaveBeenCalledWith(
139+
expect.stringContaining('INSERT INTO chatBranches'),
140+
[10, 102, 1, 2, 'Original', 'two']
141+
);
142+
expect(runAsync).toHaveBeenCalledWith(
143+
expect.stringContaining('INSERT INTO chatSettings'),
144+
[10, 1]
145+
);
146+
expect(runAsync).toHaveBeenCalledWith(
147+
expect.stringContaining('INSERT INTO chatSources'),
148+
[10, 1]
149+
);
150+
});
151+
152+
it('throws when the target message is not in the original chat', async () => {
153+
const mockDb = {
154+
runAsync: jest.fn(),
155+
getFirstAsync: jest.fn().mockResolvedValue({
156+
id: 1,
157+
title: 'Original',
158+
modelId: 7,
159+
lastUsed: 1,
160+
}),
161+
getAllAsync: jest.fn().mockResolvedValueOnce([
162+
{
163+
id: 1,
164+
chatId: 1,
165+
role: 'user',
166+
content: 'one',
167+
timestamp: 100,
168+
},
169+
]),
170+
withTransactionAsync: async (callback: TransactionCallback) => callback(),
171+
} as SQLiteDatabase;
172+
173+
await expect(forkChat(mockDb, 1, 999)).rejects.toThrow(
174+
'Message 999 not found in chat 1'
175+
);
176+
expect(mockDb.runAsync).not.toHaveBeenCalled();
177+
});
178+
179+
it('preserves existing branch markers when forking from a branch', async () => {
180+
const runAsync = jest
181+
.fn()
182+
.mockResolvedValueOnce({ lastInsertRowId: 10 })
183+
.mockResolvedValueOnce({ lastInsertRowId: 101 })
184+
.mockResolvedValueOnce({ lastInsertRowId: 102 })
185+
.mockResolvedValueOnce({ lastInsertRowId: 103 })
186+
.mockResolvedValue({ lastInsertRowId: 0 });
187+
const mockDb = {
188+
runAsync,
189+
getFirstAsync: jest.fn().mockResolvedValue({
190+
id: 1,
191+
title: 'Fork 1',
192+
modelId: 7,
193+
lastUsed: 1,
194+
}),
195+
getAllAsync: jest
196+
.fn()
197+
.mockResolvedValueOnce([
198+
{
199+
id: 1,
200+
chatId: 1,
201+
role: 'user',
202+
content: 'one',
203+
timestamp: 100,
204+
},
205+
{
206+
id: 2,
207+
chatId: 1,
208+
role: 'assistant',
209+
content: 'two',
210+
timestamp: 200,
211+
},
212+
{
213+
id: 3,
214+
chatId: 1,
215+
role: 'assistant',
216+
content: 'three',
217+
timestamp: 300,
218+
},
219+
])
220+
.mockResolvedValueOnce([
221+
{
222+
id: 1,
223+
chatId: 1,
224+
afterMessageId: 2,
225+
sourceChatId: 8,
226+
sourceMessageId: 20,
227+
sourceChatTitle: 'Original',
228+
sourceMessagePreview: 'two',
229+
createdAt: 1,
230+
},
231+
]),
232+
withTransactionAsync: async (callback: TransactionCallback) => callback(),
233+
} as SQLiteDatabase;
234+
235+
await forkChat(mockDb, 1, 3);
236+
237+
expect(runAsync).toHaveBeenCalledWith(
238+
expect.stringContaining('INSERT INTO chatBranches'),
239+
[10, 102, 8, 20, 'Original', 'two']
240+
);
241+
expect(runAsync).toHaveBeenCalledWith(
242+
expect.stringContaining('INSERT INTO chatBranches'),
243+
[10, 103, 1, 3, 'Fork 1', 'three']
244+
);
245+
});
246+
247+
it('replaces an existing branch marker at the target message with the latest branch marker', async () => {
248+
const runAsync = jest
249+
.fn()
250+
.mockResolvedValueOnce({ lastInsertRowId: 10 })
251+
.mockResolvedValueOnce({ lastInsertRowId: 101 })
252+
.mockResolvedValueOnce({ lastInsertRowId: 102 })
253+
.mockResolvedValue({ lastInsertRowId: 0 });
254+
const mockDb = {
255+
runAsync,
256+
getFirstAsync: jest.fn().mockResolvedValue({
257+
id: 1,
258+
title: 'Fork 1',
259+
modelId: 7,
260+
lastUsed: 1,
261+
}),
262+
getAllAsync: jest
263+
.fn()
264+
.mockResolvedValueOnce([
265+
{
266+
id: 1,
267+
chatId: 1,
268+
role: 'user',
269+
content: 'one',
270+
timestamp: 100,
271+
},
272+
{
273+
id: 2,
274+
chatId: 1,
275+
role: 'assistant',
276+
content: 'two',
277+
timestamp: 200,
278+
},
279+
])
280+
.mockResolvedValueOnce([
281+
{
282+
id: 1,
283+
chatId: 1,
284+
afterMessageId: 2,
285+
sourceChatId: 8,
286+
sourceMessageId: 20,
287+
sourceChatTitle: 'Original',
288+
sourceMessagePreview: 'old marker',
289+
createdAt: 1,
290+
},
291+
]),
292+
withTransactionAsync: async (callback: TransactionCallback) => callback(),
293+
} as SQLiteDatabase;
294+
295+
await forkChat(mockDb, 1, 2);
296+
297+
expect(runAsync).not.toHaveBeenCalledWith(
298+
expect.stringContaining('INSERT INTO chatBranches'),
299+
[10, 102, 8, 20, 'Original', 'old marker']
300+
);
301+
expect(runAsync).toHaveBeenCalledWith(
302+
expect.stringContaining('INSERT INTO chatBranches'),
303+
[10, 102, 1, 2, 'Fork 1', 'two']
304+
);
305+
});
306+
});

__tests__/chatStore.test.ts

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -190,6 +190,32 @@ describe('setChatModel', () => {
190190
});
191191
});
192192

193+
describe('forkChat', () => {
194+
it('returns the new chat id and reloads chats', async () => {
195+
(chatRepository.forkChat as jest.Mock).mockResolvedValue(10);
196+
(chatRepository.getAllChats as jest.Mock).mockResolvedValue([
197+
mockChat(10),
198+
mockChat(1),
199+
]);
200+
201+
const newChatId = await useChatStore.getState().forkChat(1, 2);
202+
203+
expect(newChatId).toBe(10);
204+
expect(chatRepository.forkChat).toHaveBeenCalledWith(mockDb, 1, 2);
205+
expect(chatRepository.getAllChats).toHaveBeenCalledWith(mockDb);
206+
expect(useChatStore.getState().chats[0].id).toBe(10);
207+
});
208+
209+
it('returns undefined when db is not set', async () => {
210+
useChatStore.setState({ db: null });
211+
212+
const newChatId = await useChatStore.getState().forkChat(1, 2);
213+
214+
expect(newChatId).toBeUndefined();
215+
expect(chatRepository.forkChat).not.toHaveBeenCalled();
216+
});
217+
});
218+
193219
describe('enableSource', () => {
194220
it('adds sourceId to phantomChat without hitting the db', async () => {
195221
useChatStore.setState({

__tests__/llmStore.test.ts

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -405,6 +405,30 @@ describe('sendChatMessage', () => {
405405
expect(messagesBeforeGenerate[1].content).toBe('');
406406
});
407407

408+
it('replaces assistant placeholder with persisted message id after generation', async () => {
409+
mockPersistMessage.mockResolvedValueOnce(41).mockResolvedValueOnce(42);
410+
useLLMStore.setState({
411+
model: baseModel,
412+
activeChatId: 1,
413+
activeChatMessages: [],
414+
});
415+
416+
await useLLMStore
417+
.getState()
418+
.sendChatMessage('ping', 1, noSources, settings);
419+
420+
const messages = useLLMStore.getState().activeChatMessages;
421+
expect(messages).toHaveLength(2);
422+
expect(messages[0].id).toBe(41);
423+
expect(messages[1]).toEqual(
424+
expect.objectContaining({
425+
id: 42,
426+
role: 'assistant',
427+
content: 'The answer is 42.',
428+
})
429+
);
430+
});
431+
408432
it('recovers gracefully when generation returns null', async () => {
409433
mockInstance.generate.mockResolvedValue(null);
410434
useLLMStore.setState({

0 commit comments

Comments
 (0)