Skip to content

Commit f1a0d4c

Browse files
committed
improvement(copilot): make copilot_messages the sole transcript store, remove JSONB dual-write
Stop writing/reading the legacy copilot_chats.messages JSONB column now that reads are cut over to copilot_messages. Make appendCopilotChatMessages the primary write (throws on failure instead of swallowing), repoint peripheral readers (workspace VFS, chat cleanup, data drains, fork, superuser import) to copilot_messages, and persist the assistant turn inside finalizeAssistantTurn's transaction so it commits atomically with the stream-marker clear. The column itself is dropped in a follow-up migration after this bakes.
1 parent e8f6485 commit f1a0d4c

19 files changed

Lines changed: 496 additions & 591 deletions

File tree

apps/sim/app/api/copilot/chat/stop/route.test.ts

Lines changed: 54 additions & 143 deletions
Original file line numberDiff line numberDiff line change
@@ -1,79 +1,19 @@
11
/**
22
* @vitest-environment node
33
*/
4-
import { authMockFns } from '@sim/testing'
4+
import { authMockFns, dbChainMock, dbChainMockFns, resetDbChainMock } from '@sim/testing'
55
import { NextRequest } from 'next/server'
66
import { beforeEach, describe, expect, it, vi } from 'vitest'
77

8-
const {
9-
mockSelect,
10-
mockFrom,
11-
mockWhereSelect,
12-
mockLimit,
13-
mockForUpdate,
14-
mockUpdate,
15-
mockSet,
16-
mockWhereUpdate,
17-
mockReturning,
18-
mockPublishStatusChanged,
19-
mockSql,
20-
mockTransaction,
21-
} = vi.hoisted(() => {
22-
const mockSelect = vi.fn()
23-
const mockFrom = vi.fn()
24-
const mockWhereSelect = vi.fn()
25-
const mockLimit = vi.fn()
26-
const mockForUpdate = vi.fn()
27-
const mockUpdate = vi.fn()
28-
const mockSet = vi.fn()
29-
const mockWhereUpdate = vi.fn()
30-
const mockReturning = vi.fn()
31-
const mockPublishStatusChanged = vi.fn()
32-
const mockSql = vi.fn((strings: TemplateStringsArray, ...values: unknown[]) => ({
33-
strings,
34-
values,
35-
}))
36-
const mockTransaction = vi.fn(
37-
(callback: (tx: { select: typeof mockSelect; update: typeof mockUpdate }) => unknown) =>
38-
callback({ select: mockSelect, update: mockUpdate })
39-
)
40-
41-
return {
42-
mockSelect,
43-
mockFrom,
44-
mockWhereSelect,
45-
mockLimit,
46-
mockForUpdate,
47-
mockUpdate,
48-
mockSet,
49-
mockWhereUpdate,
50-
mockReturning,
51-
mockPublishStatusChanged,
52-
mockSql,
53-
mockTransaction,
54-
}
55-
})
8+
vi.mock('@sim/db', () => dbChainMock)
569

57-
vi.mock('@sim/db/schema', () => ({
58-
copilotChats: {
59-
id: 'copilotChats.id',
60-
userId: 'copilotChats.userId',
61-
workspaceId: 'copilotChats.workspaceId',
62-
messages: 'copilotChats.messages',
63-
conversationId: 'copilotChats.conversationId',
64-
},
65-
}))
66-
67-
vi.mock('@sim/db', () => ({
68-
db: {
69-
transaction: mockTransaction,
70-
},
10+
const { mockAppendCopilotChatMessages, mockPublishStatusChanged } = vi.hoisted(() => ({
11+
mockAppendCopilotChatMessages: vi.fn(),
12+
mockPublishStatusChanged: vi.fn(),
7113
}))
7214

73-
vi.mock('drizzle-orm', () => ({
74-
and: vi.fn((...conditions: unknown[]) => ({ conditions, type: 'and' })),
75-
eq: vi.fn((field: unknown, value: unknown) => ({ field, value, type: 'eq' })),
76-
sql: mockSql,
15+
vi.mock('@/lib/copilot/chat/messages-store', () => ({
16+
appendCopilotChatMessages: mockAppendCopilotChatMessages,
7717
}))
7818

7919
vi.mock('@/lib/copilot/tasks', () => ({
@@ -92,81 +32,72 @@ function createRequest(body: Record<string, unknown>) {
9232
})
9333
}
9434

35+
/**
36+
* Sequence the two in-tx reads `finalizeAssistantTurn` issues: the chat row
37+
* (`FOR UPDATE ... LIMIT 1`) and the assistant existence check (both terminate
38+
* on `.limit(1)`).
39+
*/
40+
function mockReads(opts: { chat: Record<string, unknown> | null; assistantExists?: boolean }) {
41+
dbChainMockFns.limit.mockResolvedValueOnce(opts.chat ? [opts.chat] : [])
42+
dbChainMockFns.limit.mockResolvedValueOnce(opts.assistantExists ? [{ id: 'cm-asst' }] : [])
43+
}
44+
9545
describe('copilot chat stop route', () => {
9646
beforeEach(() => {
9747
vi.clearAllMocks()
98-
48+
// Drain any `mockResolvedValueOnce` queue left by an early-return test
49+
// (clearAllMocks/resetDbChainMock don't clear the once-queue), then restore
50+
// the default chain implementations.
51+
dbChainMockFns.limit.mockReset()
52+
resetDbChainMock()
9953
authMockFns.mockGetSession.mockResolvedValue({ user: { id: 'user-1' } })
100-
101-
mockLimit.mockResolvedValue([
102-
{
103-
workspaceId: 'ws-1',
104-
messages: [{ id: 'stream-1', role: 'user', content: 'hello' }],
105-
conversationId: 'stream-1',
106-
},
107-
])
108-
mockForUpdate.mockReturnValue({ limit: mockLimit })
109-
mockWhereSelect.mockReturnValue({ for: mockForUpdate })
110-
mockFrom.mockReturnValue({ where: mockWhereSelect })
111-
mockSelect.mockReturnValue({ from: mockFrom })
112-
113-
mockReturning.mockResolvedValue([{ workspaceId: 'ws-1' }])
114-
mockWhereUpdate.mockReturnValue({ returning: mockReturning })
115-
mockSet.mockReturnValue({ where: mockWhereUpdate })
116-
mockUpdate.mockReturnValue({ set: mockSet })
11754
})
11855

11956
it('returns 401 when unauthenticated', async () => {
12057
authMockFns.mockGetSession.mockResolvedValueOnce(null)
12158

12259
const response = await POST(
123-
createRequest({
124-
chatId: 'chat-1',
125-
streamId: 'stream-1',
126-
content: '',
127-
})
60+
createRequest({ chatId: 'chat-1', streamId: 'stream-1', content: '' })
12861
)
12962

13063
expect(response.status).toBe(401)
13164
expect(await response.json()).toEqual({ error: 'Unauthorized' })
13265
})
13366

13467
it('is a no-op when the chat is missing', async () => {
135-
mockLimit.mockResolvedValueOnce([])
68+
mockReads({ chat: null })
13669

13770
const response = await POST(
138-
createRequest({
139-
chatId: 'missing-chat',
140-
streamId: 'stream-1',
141-
content: '',
142-
})
71+
createRequest({ chatId: 'missing-chat', streamId: 'stream-1', content: '' })
14372
)
14473

14574
expect(response.status).toBe(200)
14675
expect(await response.json()).toEqual({ success: true })
147-
expect(mockUpdate).not.toHaveBeenCalled()
76+
expect(mockAppendCopilotChatMessages).not.toHaveBeenCalled()
14877
})
14978

15079
it('appends a stopped assistant message even with no content', async () => {
80+
mockReads({
81+
chat: { workspaceId: 'ws-1', conversationId: 'stream-1', model: null },
82+
assistantExists: false,
83+
})
84+
15185
const response = await POST(
152-
createRequest({
153-
chatId: 'chat-1',
154-
streamId: 'stream-1',
155-
content: '',
156-
})
86+
createRequest({ chatId: 'chat-1', streamId: 'stream-1', content: '' })
15787
)
15888

15989
expect(response.status).toBe(200)
16090
expect(await response.json()).toEqual({ success: true })
16191

162-
const setArg = mockSet.mock.calls[0]?.[0]
163-
expect(setArg).toBeTruthy()
92+
// The stream marker is cleared and nothing is written to the JSONB column.
93+
const setArg = dbChainMockFns.set.mock.calls[0]?.[0] as Record<string, unknown>
16494
expect(setArg.conversationId).toBeNull()
165-
expect(setArg.messages).toBeTruthy()
95+
expect(Object.hasOwn(setArg, 'messages')).toBe(false)
16696

167-
const appendedPayload = JSON.parse(setArg.messages.values[1] as string)
168-
expect(appendedPayload).toHaveLength(1)
169-
expect(appendedPayload[0]).toMatchObject({
97+
// The stopped assistant turn is persisted to copilot_messages.
98+
expect(mockAppendCopilotChatMessages).toHaveBeenCalledTimes(1)
99+
const [, appended] = mockAppendCopilotChatMessages.mock.calls[0]
100+
expect(appended[0]).toMatchObject({
170101
role: 'assistant',
171102
content: '',
172103
contentBlocks: [{ type: 'complete', status: 'cancelled' }],
@@ -181,32 +112,21 @@ describe('copilot chat stop route', () => {
181112
})
182113

183114
it('appends a stopped assistant message if the stream marker was already cleared', async () => {
184-
mockLimit.mockResolvedValueOnce([
185-
{
186-
workspaceId: 'ws-1',
187-
messages: [{ id: 'stream-1', role: 'user', content: 'hello' }],
188-
conversationId: null,
189-
},
190-
])
115+
mockReads({
116+
chat: { workspaceId: 'ws-1', conversationId: null, model: null },
117+
assistantExists: false,
118+
})
191119

192120
const response = await POST(
193-
createRequest({
194-
chatId: 'chat-1',
195-
streamId: 'stream-1',
196-
content: 'partial',
197-
})
121+
createRequest({ chatId: 'chat-1', streamId: 'stream-1', content: 'partial' })
198122
)
199123

200124
expect(response.status).toBe(200)
201125
expect(await response.json()).toEqual({ success: true })
202126

203-
const setArg = mockSet.mock.calls[0]?.[0]
204-
expect(setArg.messages).toBeTruthy()
205-
const appendedPayload = JSON.parse(setArg.messages.values[1] as string)
206-
expect(appendedPayload[0]).toMatchObject({
207-
role: 'assistant',
208-
content: 'partial',
209-
})
127+
expect(mockAppendCopilotChatMessages).toHaveBeenCalledTimes(1)
128+
const [, appended] = mockAppendCopilotChatMessages.mock.calls[0]
129+
expect(appended[0]).toMatchObject({ role: 'assistant', content: 'partial' })
210130

211131
expect(mockPublishStatusChanged).toHaveBeenCalledWith({
212132
workspaceId: 'ws-1',
@@ -217,28 +137,19 @@ describe('copilot chat stop route', () => {
217137
})
218138

219139
it('republishes completed status when the assistant was already persisted', async () => {
220-
mockLimit.mockResolvedValueOnce([
221-
{
222-
workspaceId: 'ws-1',
223-
messages: [
224-
{ id: 'stream-1', role: 'user', content: 'hello' },
225-
{ id: 'assistant-1', role: 'assistant', content: 'partial' },
226-
],
227-
conversationId: null,
228-
},
229-
])
140+
mockReads({
141+
chat: { workspaceId: 'ws-1', conversationId: null, model: null },
142+
assistantExists: true,
143+
})
230144

231145
const response = await POST(
232-
createRequest({
233-
chatId: 'chat-1',
234-
streamId: 'stream-1',
235-
content: 'partial',
236-
})
146+
createRequest({ chatId: 'chat-1', streamId: 'stream-1', content: 'partial' })
237147
)
238148

239149
expect(response.status).toBe(200)
240150
expect(await response.json()).toEqual({ success: true })
241-
expect(mockUpdate).not.toHaveBeenCalled()
151+
expect(mockAppendCopilotChatMessages).not.toHaveBeenCalled()
152+
expect(dbChainMockFns.set).not.toHaveBeenCalled()
242153
expect(mockPublishStatusChanged).toHaveBeenCalledWith({
243154
workspaceId: 'ws-1',
244155
chatId: 'chat-1',

apps/sim/app/api/copilot/chat/update-messages/route.test.ts

Lines changed: 24 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ const {
1616
mockSet,
1717
mockUpdateWhere,
1818
mockReturning,
19+
mockReplaceCopilotChatMessages,
1920
} = vi.hoisted(() => ({
2021
mockSelect: vi.fn(),
2122
mockFrom: vi.fn(),
@@ -25,6 +26,7 @@ const {
2526
mockSet: vi.fn(),
2627
mockUpdateWhere: vi.fn(),
2728
mockReturning: vi.fn(),
29+
mockReplaceCopilotChatMessages: vi.fn(),
2830
}))
2931

3032
vi.mock('@sim/db', () => ({
@@ -34,6 +36,10 @@ vi.mock('@sim/db', () => ({
3436
},
3537
}))
3638

39+
vi.mock('@/lib/copilot/chat/messages-store', () => ({
40+
replaceCopilotChatMessages: mockReplaceCopilotChatMessages,
41+
}))
42+
3743
vi.mock('drizzle-orm', () => ({
3844
and: vi.fn((...conditions: unknown[]) => ({ conditions, type: 'and' })),
3945
eq: vi.fn((field: unknown, value: unknown) => ({ field, value, type: 'eq' })),
@@ -257,9 +263,10 @@ describe('Copilot Chat Update Messages API Route', () => {
257263

258264
expect(mockSelect).toHaveBeenCalled()
259265
expect(mockUpdate).toHaveBeenCalled()
260-
expect(mockSet).toHaveBeenCalledWith({
261-
messages,
262-
updatedAt: expect.any(Date),
266+
// The transcript is written to copilot_messages, not the chat row.
267+
expect(mockSet).toHaveBeenCalledWith({ updatedAt: expect.any(Date) })
268+
expect(mockReplaceCopilotChatMessages).toHaveBeenCalledWith('chat-123', messages, {
269+
chatModel: 'gpt-4',
263270
})
264271
})
265272

@@ -315,8 +322,12 @@ describe('Copilot Chat Update Messages API Route', () => {
315322
messageCount: 2,
316323
})
317324

318-
expect(mockSet).toHaveBeenCalledWith({
319-
messages: [
325+
expect(mockSet).toHaveBeenCalledWith({ updatedAt: expect.any(Date) })
326+
// Messages are normalized (toolCalls folded into contentBlocks) before
327+
// being persisted to copilot_messages.
328+
expect(mockReplaceCopilotChatMessages).toHaveBeenCalledWith(
329+
'chat-456',
330+
[
320331
{
321332
id: 'msg-1',
322333
role: 'user',
@@ -345,8 +356,8 @@ describe('Copilot Chat Update Messages API Route', () => {
345356
],
346357
},
347358
],
348-
updatedAt: expect.any(Date),
349-
})
359+
{ chatModel: 'gpt-4' }
360+
)
350361
})
351362

352363
it('should handle empty messages array', async () => {
@@ -373,9 +384,9 @@ describe('Copilot Chat Update Messages API Route', () => {
373384
messageCount: 0,
374385
})
375386

376-
expect(mockSet).toHaveBeenCalledWith({
377-
messages: [],
378-
updatedAt: expect.any(Date),
387+
expect(mockSet).toHaveBeenCalledWith({ updatedAt: expect.any(Date) })
388+
expect(mockReplaceCopilotChatMessages).toHaveBeenCalledWith('chat-789', [], {
389+
chatModel: 'gpt-4',
379390
})
380391
})
381392

@@ -485,9 +496,9 @@ describe('Copilot Chat Update Messages API Route', () => {
485496
messageCount: 100,
486497
})
487498

488-
expect(mockSet).toHaveBeenCalledWith({
489-
messages,
490-
updatedAt: expect.any(Date),
499+
expect(mockSet).toHaveBeenCalledWith({ updatedAt: expect.any(Date) })
500+
expect(mockReplaceCopilotChatMessages).toHaveBeenCalledWith('chat-large', messages, {
501+
chatModel: 'gpt-4',
491502
})
492503
})
493504

apps/sim/app/api/copilot/chat/update-messages/route.ts

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ import { type NextRequest, NextResponse } from 'next/server'
66
import { updateCopilotMessagesContract } from '@/lib/api/contracts/copilot'
77
import { parseRequest } from '@/lib/api/server'
88
import { getAccessibleCopilotChatAuth } from '@/lib/copilot/chat/lifecycle'
9-
import { replaceCopilotChatMessages } from '@/lib/copilot/chat/messages-dual-write'
9+
import { replaceCopilotChatMessages } from '@/lib/copilot/chat/messages-store'
1010
import { normalizeMessage, type PersistedMessage } from '@/lib/copilot/chat/persisted-message'
1111
import {
1212
authenticateCopilotRequestSessionOnly,
@@ -73,9 +73,9 @@ export const POST = withRouteHandler(async (req: NextRequest) => {
7373
return createNotFoundResponse('Chat not found or unauthorized')
7474
}
7575

76-
// Update chat with new messages, plan artifact, and config
76+
// Update chat metadata (plan artifact / config); the transcript itself is
77+
// written to copilot_messages via replaceCopilotChatMessages below.
7778
const updateData: Record<string, unknown> = {
78-
messages: normalizedMessages,
7979
updatedAt: new Date(),
8080
}
8181

0 commit comments

Comments
 (0)