improvement(mothership): abort path race preventing persistence

Author: icecrasher321

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

@@ -10,24 +10,49 @@ const {
   mockFrom,
   mockWhereSelect,
   mockLimit,
+  mockForUpdate,
   mockUpdate,
   mockSet,
   mockWhereUpdate,
   mockReturning,
   mockPublishStatusChanged,
   mockSql,
-} = vi.hoisted(() => ({
-  mockSelect: vi.fn(),
-  mockFrom: vi.fn(),
-  mockWhereSelect: vi.fn(),
-  mockLimit: vi.fn(),
-  mockUpdate: vi.fn(),
-  mockSet: vi.fn(),
-  mockWhereUpdate: vi.fn(),
-  mockReturning: vi.fn(),
-  mockPublishStatusChanged: vi.fn(),
-  mockSql: vi.fn((strings: TemplateStringsArray, ...values: unknown[]) => ({ strings, values })),
-}))
+  mockTransaction,
+} = vi.hoisted(() => {
+  const mockSelect = vi.fn()
+  const mockFrom = vi.fn()
+  const mockWhereSelect = vi.fn()
+  const mockLimit = vi.fn()
+  const mockForUpdate = vi.fn()
+  const mockUpdate = vi.fn()
+  const mockSet = vi.fn()
+  const mockWhereUpdate = vi.fn()
+  const mockReturning = vi.fn()
+  const mockPublishStatusChanged = vi.fn()
+  const mockSql = vi.fn((strings: TemplateStringsArray, ...values: unknown[]) => ({
+    strings,
+    values,
+  }))
+  const mockTransaction = vi.fn(
+    (callback: (tx: { select: typeof mockSelect; update: typeof mockUpdate }) => unknown) =>
+      callback({ select: mockSelect, update: mockUpdate })
+  )
+
+  return {
+    mockSelect,
+    mockFrom,
+    mockWhereSelect,
+    mockLimit,
+    mockForUpdate,
+    mockUpdate,
+    mockSet,
+    mockWhereUpdate,
+    mockReturning,
+    mockPublishStatusChanged,
+    mockSql,
+    mockTransaction,
+  }
+})
 
 vi.mock('@sim/db/schema', () => ({
   copilotChats: {
@@ -41,8 +66,7 @@ vi.mock('@sim/db/schema', () => ({
 
 vi.mock('@sim/db', () => ({
   db: {
-    select: mockSelect,
-    update: mockUpdate,
+    transaction: mockTransaction,
   },
 }))
 
@@ -78,9 +102,11 @@ describe('copilot chat stop route', () => {
       {
         workspaceId: 'ws-1',
         messages: [{ id: 'stream-1', role: 'user', content: 'hello' }],
+        conversationId: 'stream-1',
       },
     ])
-    mockWhereSelect.mockReturnValue({ limit: mockLimit })
+    mockForUpdate.mockReturnValue({ limit: mockLimit })
+    mockWhereSelect.mockReturnValue({ for: mockForUpdate })
     mockFrom.mockReturnValue({ where: mockWhereSelect })
     mockSelect.mockReturnValue({ from: mockFrom })
 
@@ -153,4 +179,33 @@ describe('copilot chat stop route', () => {
       streamId: 'stream-1',
     })
   })
+
+  it('appends a stopped assistant message if the stream marker was already cleared', async () => {
+    mockLimit.mockResolvedValueOnce([
+      {
+        workspaceId: 'ws-1',
+        messages: [{ id: 'stream-1', role: 'user', content: 'hello' }],
+        conversationId: null,
+      },
+    ])
+
+    const response = await POST(
+      createRequest({
+        chatId: 'chat-1',
+        streamId: 'stream-1',
+        content: 'partial',
+      })
+    )
+
+    expect(response.status).toBe(200)
+    expect(await response.json()).toEqual({ success: true })
+
+    const setArg = mockSet.mock.calls[0]?.[0]
+    expect(setArg.messages).toBeTruthy()
+    const appendedPayload = JSON.parse(setArg.messages.values[1] as string)
+    expect(appendedPayload[0]).toMatchObject({
+      role: 'assistant',
+      content: 'partial',
+    })
+  })
 })

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

@@ -1,13 +1,11 @@
-import { db } from '@sim/db'
-import { copilotChats } from '@sim/db/schema'
 import { createLogger } from '@sim/logger'
 import { generateId } from '@sim/utils/id'
-import { and, eq, sql } from 'drizzle-orm'
 import { type NextRequest, NextResponse } from 'next/server'
 import { copilotChatStopContract } from '@/lib/api/contracts/copilot'
 import { parseRequest } from '@/lib/api/server'
 import { getSession } from '@/lib/auth'
 import { normalizeMessage, type PersistedMessage } from '@/lib/copilot/chat/persisted-message'
+import { finalizeAssistantTurn } from '@/lib/copilot/chat/terminal-state'
 import { CopilotStopOutcome } from '@/lib/copilot/generated/trace-attribute-values-v1'
 import { TraceAttr } from '@/lib/copilot/generated/trace-attributes-v1'
 import { TraceSpan } from '@/lib/copilot/generated/trace-spans-v1'
@@ -44,71 +42,33 @@ export const POST = withRouteHandler((req: NextRequest) =>
         ...(requestId ? { [TraceAttr.RequestId]: requestId } : {}),
       })
 
-      const [row] = await db
-        .select({
-          workspaceId: copilotChats.workspaceId,
-          messages: copilotChats.messages,
-        })
-        .from(copilotChats)
-        .where(and(eq(copilotChats.id, chatId), eq(copilotChats.userId, session.user.id)))
-        .limit(1)
-
-      if (!row) {
-        span.setAttribute(TraceAttr.CopilotStopOutcome, CopilotStopOutcome.ChatNotFound)
-        return NextResponse.json({ success: true })
-      }
-
-      const messages: Record<string, unknown>[] = Array.isArray(row.messages) ? row.messages : []
-      const userIdx = messages.findIndex((message) => message.id === streamId)
-      const alreadyHasResponse =
-        userIdx >= 0 &&
-        userIdx + 1 < messages.length &&
-        (messages[userIdx + 1] as Record<string, unknown>)?.role === 'assistant'
-      const canAppendAssistant =
-        userIdx >= 0 && userIdx === messages.length - 1 && !alreadyHasResponse
-
-      const updateWhere = and(
-        eq(copilotChats.id, chatId),
-        eq(copilotChats.userId, session.user.id),
-        eq(copilotChats.conversationId, streamId)
-      )
-
-      const setClause: Record<string, unknown> = {
-        conversationId: null,
-        updatedAt: new Date(),
-      }
-
       const hasContent = content.trim().length > 0
       const hasBlocks = Array.isArray(contentBlocks) && contentBlocks.length > 0
       const synthesizedStoppedBlocks = hasBlocks
         ? contentBlocks
         : hasContent
           ? [{ type: 'text', channel: 'assistant', content }, { type: 'stopped' }]
           : [{ type: 'stopped' }]
-      if (canAppendAssistant) {
-        const normalized = normalizeMessage({
-          id: generateId(),
-          role: 'assistant',
-          content,
-          timestamp: new Date().toISOString(),
-          contentBlocks: synthesizedStoppedBlocks,
-          // Persist so the UI copy-request-id button survives refetch.
-          ...(requestId ? { requestId } : {}),
-        })
-        const assistantMessage: PersistedMessage = normalized
-        setClause.messages = sql`${copilotChats.messages} || ${JSON.stringify([assistantMessage])}::jsonb`
-      }
-      span.setAttribute(TraceAttr.CopilotStopAppendedAssistant, canAppendAssistant)
-
-      const [updated] = await db
-        .update(copilotChats)
-        .set(setClause)
-        .where(updateWhere)
-        .returning({ workspaceId: copilotChats.workspaceId })
+      const assistantMessage: PersistedMessage = normalizeMessage({
+        id: generateId(),
+        role: 'assistant',
+        content,
+        timestamp: new Date().toISOString(),
+        contentBlocks: synthesizedStoppedBlocks,
+        ...(requestId ? { requestId } : {}),
+      })
+      const result = await finalizeAssistantTurn({
+        chatId,
+        userId: session.user.id,
+        userMessageId: streamId,
+        assistantMessage,
+        streamMarkerPolicy: 'active-or-cleared',
+      })
+      span.setAttribute(TraceAttr.CopilotStopAppendedAssistant, result.appendedAssistant)
 
-      if (updated?.workspaceId) {
+      if (result.updated && result.workspaceId) {
         taskPubSub?.publishStatusChanged({
-          workspaceId: updated.workspaceId,
+          workspaceId: result.workspaceId,
           chatId,
           type: 'completed',
           streamId,
@@ -117,7 +77,11 @@ export const POST = withRouteHandler((req: NextRequest) =>
 
       span.setAttribute(
         TraceAttr.CopilotStopOutcome,
-        updated ? CopilotStopOutcome.Persisted : CopilotStopOutcome.NoMatchingRow
+        result.found
+          ? result.updated
+            ? CopilotStopOutcome.Persisted
+            : CopilotStopOutcome.NoMatchingRow
+          : CopilotStopOutcome.ChatNotFound
       )
       return NextResponse.json({ success: true })
     } catch (error) {