address comments

Author: icecrasher321

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

@@ -207,5 +207,43 @@ describe('copilot chat stop route', () => {
       role: 'assistant',
       content: 'partial',
     })
+
+    expect(mockPublishStatusChanged).toHaveBeenCalledWith({
+      workspaceId: 'ws-1',
+      chatId: 'chat-1',
+      type: 'completed',
+      streamId: 'stream-1',
+    })
+  })
+
+  it('republishes completed status when the assistant was already persisted', async () => {
+    mockLimit.mockResolvedValueOnce([
+      {
+        workspaceId: 'ws-1',
+        messages: [
+          { id: 'stream-1', role: 'user', content: 'hello' },
+          { id: 'assistant-1', role: 'assistant', content: 'partial' },
+        ],
+        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 })
+    expect(mockUpdate).not.toHaveBeenCalled()
+    expect(mockPublishStatusChanged).toHaveBeenCalledWith({
+      workspaceId: 'ws-1',
+      chatId: 'chat-1',
+      type: 'completed',
+      streamId: 'stream-1',
+    })
   })
 })

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

@@ -4,9 +4,16 @@ 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 {
+  normalizeMessage,
+  type PersistedMessage,
+  withStoppedContentBlock,
+} 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 {
+  CopilotChatFinalizeOutcome,
+  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'
 import { withIncomingGoSpan } from '@/lib/copilot/request/otel'
@@ -49,14 +56,16 @@ export const POST = withRouteHandler((req: NextRequest) =>
         : hasContent
           ? [{ type: 'text', channel: 'assistant', content }, { type: 'stopped' }]
           : [{ type: 'stopped' }]
-      const assistantMessage: PersistedMessage = normalizeMessage({
-        id: generateId(),
-        role: 'assistant',
-        content,
-        timestamp: new Date().toISOString(),
-        contentBlocks: synthesizedStoppedBlocks,
-        ...(requestId ? { requestId } : {}),
-      })
+      const assistantMessage: PersistedMessage = withStoppedContentBlock(
+        normalizeMessage({
+          id: generateId(),
+          role: 'assistant',
+          content,
+          timestamp: new Date().toISOString(),
+          contentBlocks: synthesizedStoppedBlocks,
+          ...(requestId ? { requestId } : {}),
+        })
+      )
       const result = await finalizeAssistantTurn({
         chatId,
         userId: session.user.id,
@@ -65,8 +74,15 @@ export const POST = withRouteHandler((req: NextRequest) =>
         streamMarkerPolicy: 'active-or-cleared',
       })
       span.setAttribute(TraceAttr.CopilotStopAppendedAssistant, result.appendedAssistant)
+      const stopOutcome = !result.found
+        ? CopilotStopOutcome.ChatNotFound
+        : result.updated || result.outcome === CopilotChatFinalizeOutcome.AssistantAlreadyPersisted
+          ? CopilotStopOutcome.Persisted
+          : CopilotStopOutcome.NoMatchingRow
+      const shouldPublishCompleted =
+        result.updated || result.outcome === CopilotChatFinalizeOutcome.AssistantAlreadyPersisted
 
-      if (result.updated && result.workspaceId) {
+      if (shouldPublishCompleted && result.workspaceId) {
         taskPubSub?.publishStatusChanged({
           workspaceId: result.workspaceId,
           chatId,
@@ -75,14 +91,7 @@ export const POST = withRouteHandler((req: NextRequest) =>
         })
       }
 
-      span.setAttribute(
-        TraceAttr.CopilotStopOutcome,
-        result.found
-          ? result.updated
-            ? CopilotStopOutcome.Persisted
-            : CopilotStopOutcome.NoMatchingRow
-          : CopilotStopOutcome.ChatNotFound
-      )
+      span.setAttribute(TraceAttr.CopilotStopOutcome, stopOutcome)
       return NextResponse.json({ success: true })
     } catch (error) {
       logger.error('Error stopping chat stream:', error)

apps/sim/lib/copilot/chat/persisted-message.ts

@@ -224,6 +224,45 @@ export function buildPersistedAssistantMessage(
   return message
 }
 
+export function withStoppedContentBlock(message: PersistedMessage): PersistedMessage {
+  const contentBlocks = message.contentBlocks ?? []
+  const hasAssistantText = contentBlocks.some(
+    (block) =>
+      block.type === MothershipStreamV1EventType.text &&
+      block.channel !== MothershipStreamV1TextChannel.thinking &&
+      block.content?.trim()
+  )
+  if (
+    contentBlocks.some(
+      (block) =>
+        block.type === MothershipStreamV1EventType.complete &&
+        block.status === MothershipStreamV1CompletionStatus.cancelled
+    )
+  ) {
+    return message
+  }
+
+  return normalizeMessage({
+    ...message,
+    contentBlocks: [
+      ...(hasAssistantText || !message.content.trim()
+        ? []
+        : [
+            {
+              type: MothershipStreamV1EventType.text,
+              channel: MothershipStreamV1TextChannel.assistant,
+              content: message.content,
+            },
+          ]),
+      ...contentBlocks,
+      {
+        type: MothershipStreamV1EventType.complete,
+        status: MothershipStreamV1CompletionStatus.cancelled,
+      },
+    ],
+  })
+}
+
 export interface UserMessageParams {
   id: string
   content: string

apps/sim/lib/copilot/chat/post.test.ts

@@ -27,6 +27,8 @@ const {
   getPendingChatStreamId,
   releasePendingChatStream,
   resolveOrCreateChat,
+  finalizeAssistantTurn,
+  mockPublishStatusChanged,
 } = vi.hoisted(() => ({
   getEffectiveDecryptedEnv: vi.fn(),
   generateWorkspaceContext: vi.fn(),
@@ -38,6 +40,8 @@ const {
   getPendingChatStreamId: vi.fn(),
   releasePendingChatStream: vi.fn(),
   resolveOrCreateChat: vi.fn(),
+  finalizeAssistantTurn: vi.fn(),
+  mockPublishStatusChanged: vi.fn(),
 }))
 
 const getSession = authMockFns.mockGetSession
@@ -78,9 +82,13 @@ vi.mock('@/lib/copilot/chat/lifecycle', () => ({
   resolveOrCreateChat,
 }))
 
+vi.mock('@/lib/copilot/chat/terminal-state', () => ({
+  finalizeAssistantTurn,
+}))
+
 vi.mock('@/lib/copilot/tasks', () => ({
   taskPubSub: {
-    publishStatusChanged: vi.fn(),
+    publishStatusChanged: mockPublishStatusChanged,
   },
 }))
 
@@ -137,6 +145,13 @@ describe('handleUnifiedChatPost', () => {
       conversationHistory: [],
       isNew: true,
     })
+    finalizeAssistantTurn.mockResolvedValue({
+      found: true,
+      updated: true,
+      appendedAssistant: true,
+      workspaceId: 'ws-1',
+      outcome: 'appended_assistant',
+    })
   })
 
   it('routes workflow-attached chat requests through the copilot backend path', async () => {
@@ -176,6 +191,7 @@ describe('handleUnifiedChatPost', () => {
         body: JSON.stringify({
           message: 'Hello',
           workspaceId: 'ws-1',
+          createNewChat: true,
         }),
       })
     )
@@ -205,6 +221,90 @@ describe('handleUnifiedChatPost', () => {
     )
   })
 
+  it('persists cancelled partial responses from the server lifecycle', async () => {
+    await handleUnifiedChatPost(
+      new NextRequest('http://localhost/api/copilot/chat', {
+        method: 'POST',
+        body: JSON.stringify({
+          message: 'Hello',
+          workspaceId: 'ws-1',
+          createNewChat: true,
+        }),
+      })
+    )
+
+    const streamArgs = createSSEStream.mock.calls[0]?.[0]
+    const onComplete = streamArgs?.orchestrateOptions?.onComplete
+    expect(onComplete).toBeTypeOf('function')
+
+    await onComplete({
+      success: false,
+      cancelled: true,
+      content: 'partial answer',
+      contentBlocks: [],
+      toolCalls: [],
+      chatId: 'chat-1',
+      requestId: 'request-1',
+    })
+
+    expect(finalizeAssistantTurn).toHaveBeenCalledWith(
+      expect.objectContaining({
+        chatId: 'chat-1',
+        userMessageId: expect.any(String),
+        streamMarkerPolicy: 'active-or-cleared',
+        assistantMessage: expect.objectContaining({
+          role: 'assistant',
+          content: 'partial answer',
+          contentBlocks: expect.arrayContaining([
+            expect.objectContaining({ type: 'complete', status: 'cancelled' }),
+          ]),
+        }),
+      })
+    )
+  })
+
+  it('republishes completed status when cancelled lifecycle persistence already ran', async () => {
+    await handleUnifiedChatPost(
+      new NextRequest('http://localhost/api/copilot/chat', {
+        method: 'POST',
+        body: JSON.stringify({
+          message: 'Hello',
+          workspaceId: 'ws-1',
+          createNewChat: true,
+        }),
+      })
+    )
+
+    const streamArgs = createSSEStream.mock.calls[0]?.[0]
+    const onComplete = streamArgs?.orchestrateOptions?.onComplete
+    expect(onComplete).toBeTypeOf('function')
+
+    finalizeAssistantTurn.mockResolvedValueOnce({
+      found: true,
+      updated: false,
+      appendedAssistant: false,
+      workspaceId: 'ws-1',
+      outcome: 'assistant_already_persisted',
+    })
+
+    await onComplete({
+      success: false,
+      cancelled: true,
+      content: 'partial answer',
+      contentBlocks: [],
+      toolCalls: [],
+      chatId: 'chat-1',
+      requestId: 'request-1',
+    })
+
+    expect(mockPublishStatusChanged).toHaveBeenCalledWith({
+      workspaceId: 'ws-1',
+      chatId: 'chat-1',
+      type: 'completed',
+      streamId: streamArgs?.streamId,
+    })
+  })
+
   it('rejects requests that have neither workflow nor workspace attachment', async () => {
     const response = await handleUnifiedChatPost(
       new NextRequest('http://localhost/api/copilot/chat', {

apps/sim/lib/copilot/chat/post.ts

@@ -14,6 +14,7 @@ import { buildCopilotRequestPayload } from '@/lib/copilot/chat/payload'
 import {
   buildPersistedAssistantMessage,
   buildPersistedUserMessage,
+  withStoppedContentBlock,
 } from '@/lib/copilot/chat/persisted-message'
 import {
   processContextsServer,
@@ -23,6 +24,7 @@ import { finalizeAssistantTurn } from '@/lib/copilot/chat/terminal-state'
 import { generateWorkspaceContext } from '@/lib/copilot/chat/workspace-context'
 import { COPILOT_REQUEST_MODES } from '@/lib/copilot/constants'
 import {
+  CopilotChatFinalizeOutcome,
   CopilotChatPersistOutcome,
   CopilotTransport,
 } from '@/lib/copilot/generated/trace-attribute-values-v1'
@@ -425,13 +427,31 @@ function buildOnComplete(params: {
 
     if (!chatId) return
 
-    // On cancel, /chat/stop is the sole DB writer — it persists
-    // partial content AND clears conversationId in one UPDATE. If we
-    // finalize here first the filter misses and content vanishes.
-    // Real errors still finalize so the stream marker clears.
-    if (result.cancelled) return
-
     try {
+      if (result.cancelled) {
+        const finalization = await finalizeAssistantTurn({
+          chatId,
+          userMessageId,
+          assistantMessage: withStoppedContentBlock(
+            buildPersistedAssistantMessage(result, requestId)
+          ),
+          streamMarkerPolicy: 'active-or-cleared',
+        })
+        const shouldPublishCompletion =
+          finalization.updated ||
+          finalization.outcome === CopilotChatFinalizeOutcome.AssistantAlreadyPersisted
+
+        if (notifyWorkspaceStatus && workspaceId && shouldPublishCompletion) {
+          taskPubSub?.publishStatusChanged({
+            workspaceId,
+            chatId,
+            type: 'completed',
+            streamId: userMessageId,
+          })
+        }
+        return
+      }
+
       await finalizeAssistantTurn({
         chatId,
         userMessageId,