11/**
22 * @vitest -environment node
33 */
4- import { authMockFns } from '@sim/testing'
4+ import { authMockFns , dbChainMock , dbChainMockFns , resetDbChainMock } from '@sim/testing'
55import { NextRequest } from 'next/server'
66import { 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
7919vi . 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+
9545describe ( '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' ,
0 commit comments