fix(mothership): let the server re-claim a workflow call its dead claimant never started

Author: waleedlatif1

apps/sim/lib/mothership/async-runs/repository.ts

@@ -21,6 +21,7 @@ import {
   inArray,
   isNotNull,
   isNull,
+  like,
   lt,
   ne,
   notInArray,
@@ -1591,7 +1592,23 @@ export async function claimWorkflowToolExecution(
                   )
                 ),
                 /** Only the server re-runs a workflow whose Sim executor died; no browser resumes it. */
-                executor === 'sim' ? interruptedSimExecution : undefined
+                executor === 'sim' ? interruptedSimExecution : undefined,
+                /**
+                 * A server claim whose claimant died before starting the execution. Starting is
+                 * fenced by `claimToolExecution`, so a claimant still alive and this one cannot
+                 * both run it: the second to start observes the first's result.
+                 */
+                executor === 'sim'
+                  ? and(
+                      like(copilotAsyncToolCalls.claimedBy, `${WORKFLOW_EXECUTION_CLAIM_PREFIX}%`),
+                      isNull(copilotAsyncToolCalls.executionStartedAt),
+                      isNull(copilotAsyncToolCalls.clientWorkflowExecutionId),
+                      inArray(copilotAsyncToolCalls.status, [
+                        ASYNC_TOOL_STATUS.running,
+                        ASYNC_TOOL_STATUS.delivered,
+                      ])
+                    )
+                  : undefined
               )
             )
           )

apps/sim/lib/mothership/request/application/abandoned-runs.integration.ts

@@ -588,4 +588,47 @@ describe.runIf(Boolean(redisUrl))('taking over a run whose controller died', ()
       await db.update(copilotRuns).set({ status: 'cancelled' }).where(eq(copilotRuns.id, runId))
     }
   })
+
+  it('re-claims a server workflow claim whose claimant died before starting it, once', async () => {
+    const chatId = generateId()
+    chatIds.push(chatId)
+    const runId = generateId()
+    const streamId = generateId()
+    await db.insert(copilotChats).values({ id: chatId, userId, workspaceId, type: 'mothership' })
+    await db.insert(copilotRuns).values({
+      id: runId,
+      executionId: generateId(),
+      chatId,
+      userId,
+      workspaceId,
+      streamId,
+      status: 'paused_waiting_for_tool',
+      toolExecutionVersion: 2,
+    })
+    const toolCallId = scopeProviderToolCallId(
+      'workflow-call',
+      createProviderToolCallIdentity(runId)
+    )
+    await upsertAsyncToolCall({ runId, toolCallId, toolName: 'run_workflow', args: {} })
+    await db
+      .update(copilotAsyncToolCalls)
+      .set({ status: 'running' })
+      .where(eq(copilotAsyncToolCalls.toolCallId, toolCallId))
+    try {
+      /** The dead controller's server fallback claimed the call, then died before starting it. */
+      expect(await claimWorkflowToolExecution(toolCallId, generateId(), 'sim')).not.toBeNull()
+
+      expect(await claimWorkflowToolExecution(toolCallId, generateId(), 'client')).toBeNull()
+      expect(await claimWorkflowToolExecution(toolCallId, generateId(), 'sim')).not.toBeNull()
+      const successor = { toolCallId, runId, userId, ownerToken: 'successor' }
+      expect(await claimToolExecution(successor)).toEqual({ outcome: 'claimed' })
+      expect(await claimToolExecution({ ...successor, ownerToken: 'late-claimant' })).toEqual({
+        outcome: 'existing',
+      })
+      expect(await claimWorkflowToolExecution(toolCallId, generateId(), 'sim')).toBeNull()
+      expect(await storedTool(toolCallId)).toMatchObject({ attempts: 1, ownerToken: 'successor' })
+    } finally {
+      await db.update(copilotRuns).set({ status: 'cancelled' }).where(eq(copilotRuns.id, runId))
+    }
+  })
 })