Fix mothership block req lifecycle

Author: Sg312

apps/sim/app/api/mothership/execute/route.ts

@@ -5,6 +5,7 @@ import { checkInternalAuth } from '@/lib/auth/hybrid'
 import { buildIntegrationToolSchemas } from '@/lib/copilot/chat/payload'
 import { generateWorkspaceContext } from '@/lib/copilot/chat/workspace-context'
 import { runHeadlessCopilotLifecycle } from '@/lib/copilot/request/lifecycle/headless'
+import { requestExplicitStreamAbort } from '@/lib/copilot/request/session/explicit-abort'
 import { generateId } from '@/lib/core/utils/uuid'
 import {
   assertActiveWorkspaceAccess,
@@ -26,8 +27,16 @@ const ExecuteRequestSchema = z.object({
   workspaceId: z.string().min(1, 'workspaceId is required'),
   userId: z.string().min(1, 'userId is required'),
   chatId: z.string().optional(),
+  messageId: z.string().optional(),
+  requestId: z.string().optional(),
+  workflowId: z.string().optional(),
+  executionId: z.string().optional(),
 })
 
+function isAbortError(error: unknown): boolean {
+  return error instanceof Error && error.name === 'AbortError'
+}
+
 /**
  * POST /api/mothership/execute
  *
@@ -37,6 +46,7 @@ const ExecuteRequestSchema = z.object({
  */
 export async function POST(req: NextRequest) {
   let messageId: string | undefined
+  let requestId: string | undefined
 
   try {
     const auth = await checkInternalAuth(req, { requireWorkflowId: false })
@@ -45,14 +55,29 @@ export async function POST(req: NextRequest) {
     }
 
     const body = await req.json()
-    const { messages, responseFormat, workspaceId, userId, chatId } =
-      ExecuteRequestSchema.parse(body)
+    const {
+      messages,
+      responseFormat,
+      workspaceId,
+      userId,
+      chatId,
+      messageId: providedMessageId,
+      requestId: providedRequestId,
+      workflowId,
+      executionId,
+    } = ExecuteRequestSchema.parse(body)
 
     await assertActiveWorkspaceAccess(workspaceId, userId)
 
     const effectiveChatId = chatId || generateId()
-    messageId = generateId()
-    const reqLogger = logger.withMetadata({ messageId })
+    messageId = providedMessageId || generateId()
+    requestId = providedRequestId || generateId()
+    const reqLogger = logger.withMetadata({
+      messageId,
+      requestId,
+      workflowId,
+      executionId,
+    })
     const [workspaceContext, integrationTools, userPermission] = await Promise.all([
       generateWorkspaceContext(workspaceId, userId),
       buildIntegrationToolSchemas(userId, messageId),
@@ -72,52 +97,96 @@ export async function POST(req: NextRequest) {
       ...(userPermission ? { userPermission } : {}),
     }
 
-    const result = await runHeadlessCopilotLifecycle(requestPayload, {
-      userId,
-      workspaceId,
-      chatId: effectiveChatId,
-      goRoute: '/api/mothership/execute',
-      autoExecuteTools: true,
-      interactive: false,
-    })
+    let allowExplicitAbort = true
+    let explicitAbortRequest: Promise<void> | undefined
+    const onAbort = () => {
+      if (!allowExplicitAbort || explicitAbortRequest || !messageId) {
+        return
+      }
 
-    if (!result.success) {
-      logger.error(
-        messageId
-          ? `Mothership execute failed [messageId:${messageId}]`
-          : 'Mothership execute failed',
-        {
-          error: result.error,
-          errors: result.errors,
-        }
-      )
-      return NextResponse.json(
-        {
-          error: result.error || 'Mothership execution failed',
-          content: result.content || '',
-        },
-        { status: 500 }
-      )
+      explicitAbortRequest = requestExplicitStreamAbort({
+        streamId: messageId,
+        userId,
+        chatId: effectiveChatId,
+      }).catch((error) => {
+        reqLogger.warn('Failed to send explicit abort for mothership execution', {
+          error: error instanceof Error ? error.message : String(error),
+        })
+      })
     }
 
-    const clientToolNames = new Set(integrationTools.map((t) => t.name))
-    const clientToolCalls = (result.toolCalls || []).filter(
-      (tc: { name: string }) => clientToolNames.has(tc.name) || tc.name.startsWith('mcp-')
-    )
+    if (req.signal.aborted) {
+      onAbort()
+    } else {
+      req.signal.addEventListener('abort', onAbort, { once: true })
+    }
+
+    try {
+      const result = await runHeadlessCopilotLifecycle(requestPayload, {
+        userId,
+        workspaceId,
+        chatId: effectiveChatId,
+        workflowId,
+        executionId,
+        simRequestId: requestId,
+        goRoute: '/api/mothership/execute',
+        autoExecuteTools: true,
+        interactive: false,
+        abortSignal: req.signal,
+      })
+
+      allowExplicitAbort = false
+
+      if (req.signal.aborted) {
+        reqLogger.info('Mothership execute aborted after lifecycle completion')
+        return NextResponse.json({ error: 'Mothership execution aborted' }, { status: 499 })
+      }
 
-    return NextResponse.json({
-      content: result.content,
-      model: 'mothership',
-      tokens: result.usage
-        ? {
-            prompt: result.usage.prompt,
-            completion: result.usage.completion,
-            total: (result.usage.prompt || 0) + (result.usage.completion || 0),
+      if (!result.success) {
+        logger.error(
+          messageId
+            ? `Mothership execute failed [messageId:${messageId}]`
+            : 'Mothership execute failed',
+          {
+            requestId,
+            workflowId,
+            executionId,
+            error: result.error,
+            errors: result.errors,
           }
-        : {},
-      cost: result.cost || undefined,
-      toolCalls: clientToolCalls,
-    })
+        )
+        return NextResponse.json(
+          {
+            error: result.error || 'Mothership execution failed',
+            content: result.content || '',
+          },
+          { status: 500 }
+        )
+      }
+
+      const clientToolNames = new Set(integrationTools.map((t) => t.name))
+      const clientToolCalls = (result.toolCalls || []).filter(
+        (tc: { name: string }) => clientToolNames.has(tc.name) || tc.name.startsWith('mcp-')
+      )
+
+      return NextResponse.json({
+        content: result.content,
+        model: 'mothership',
+        tokens: result.usage
+          ? {
+              prompt: result.usage.prompt,
+              completion: result.usage.completion,
+              total: (result.usage.prompt || 0) + (result.usage.completion || 0),
+            }
+          : {},
+        cost: result.cost || undefined,
+        toolCalls: clientToolCalls,
+      })
+    } finally {
+      allowExplicitAbort = false
+      req.signal.removeEventListener('abort', onAbort)
+      await explicitAbortRequest
+    }
   } catch (error) {
     if (error instanceof z.ZodError) {
       return NextResponse.json(
@@ -126,9 +195,23 @@ export async function POST(req: NextRequest) {
       )
     }
 
+    if (req.signal.aborted || isAbortError(error)) {
+      logger.info(
+        messageId
+          ? `Mothership execute aborted [messageId:${messageId}]`
+          : 'Mothership execute aborted',
+        {
+          requestId,
+        }
+      )
+
+      return NextResponse.json({ error: 'Mothership execution aborted' }, { status: 499 })
+    }
+
     logger.error(
       messageId ? `Mothership execute error [messageId:${messageId}]` : 'Mothership execute error',
       {
+        requestId,
         error: error instanceof Error ? error.message : 'Unknown error',
       }
     )