fix(desktop): serialize a device result with Stop through the run row

A result now locks the run row before checking tool admission, the same
serialization the claim and Stop use, so it can no longer commit after a
concurrent Stop. The schema declares the run-device index concurrent, as
the migration builds it.

Author: waleedlatif1

apps/sim/lib/desktop/application/executor.integration.ts

@@ -28,6 +28,7 @@ import {
   user,
   workspace,
 } from '@sim/db/schema'
+import { createDeferred } from '@sim/testing/helpers/deferred'
 import { featureFlagsMock, featureFlagsMockFns } from '@sim/testing/mocks/feature-flags.mock'
 import { sleep } from '@sim/utils/helpers'
 import { generateId } from '@sim/utils/id'
@@ -544,6 +545,37 @@ describe.runIf(Boolean(redisUrl))('desktop background executor protocol', () =>
   })
 
   describe('completion', () => {
+    it('serializes a result with a concurrent Stop, so it never lands after the Stop commits', async () => {
+      const desktop = await signedInDesktop()
+      const run = await boundRun(desktop)
+      const toolCallId = await pendingCall(run.runId)
+      await offer(toolCallId)
+      const { executionToken } = await claim(desktop.principal, desktop.deviceId, toolCallId)
+
+      /** Stop closes admission on the run row; hold its transaction open while the result arrives. */
+      const stopHeld = createDeferred<void>()
+      const releaseStop = createDeferred<void>()
+      const stopping = db.transaction(async (tx) => {
+        await tx
+          .update(copilotRuns)
+          .set({ toolAdmissionClosedAt: sql`now()` })
+          .where(eq(copilotRuns.id, run.runId))
+        stopHeld.resolve()
+        await releaseStop.promise
+      })
+      await stopHeld.promise
+      const completing = completeDesktopTool.execute({
+        principal: desktop.principal,
+        input: { deviceId: desktop.deviceId, toolCallId, executionToken, status: 'success' },
+      })
+      await sleep(300)
+      releaseStop.resolve()
+      await stopping
+
+      await expect(completing).resolves.toEqual({ outcome: 'superseded', status: 'cancelled' })
+      expect((await row(toolCallId)).status).toBe('cancelled')
+    })
+
     it('records a result once, wakes the waiting run, and answers a retry as a duplicate', async () => {
       const desktop = await signedInDesktop()
       const run = await boundRun(desktop)
@@ -713,9 +745,17 @@ describe.runIf(Boolean(redisUrl))('desktop background executor protocol', () =>
         principal: desktop.principal,
         input: { deviceId: desktop.deviceId },
       })
-      await expect(stream.revalidate()).resolves.toBeUndefined()
-      await db.delete(session).where(eq(session.id, desktop.sessionId))
-      await expect(stream.revalidate()).rejects.toBeInstanceOf(DesktopDeviceUnrecognizedError)
+      const close = stream.subscribe(() => {})
+      try {
+        await expect.poll(() => isDesktopPresent(desktop.deviceId)).toBe(true)
+        await expect(stream.revalidate()).resolves.toBeUndefined()
+        await db.delete(session).where(eq(session.id, desktop.sessionId))
+        /** The SSE transport closes a stream whose revalidation rejects; closing releases presence. */
+        await expect(stream.revalidate()).rejects.toBeInstanceOf(DesktopDeviceUnrecognizedError)
+      } finally {
+        close()
+      }
+      await expect.poll(() => isDesktopPresent(desktop.deviceId)).toBe(false)
     })
   })
 })

apps/sim/lib/desktop/executor/repository.ts

@@ -301,31 +301,38 @@ export async function recordDesktopCallResult(input: {
   result: AsyncCompletionData
   error: string | null
 }): Promise<boolean> {
-  const [row] = await db
-    .update(copilotAsyncToolCalls)
-    .set({
-      status: input.status,
-      result: sanitizeValueForJsonb(input.result),
-      error: input.error,
-      claimedBy: null,
-      claimedAt: null,
-      completedAt: sql`now()`,
-      executionSettledAt: sql`now()`,
-      updatedAt: sql`now()`,
-    })
-    .where(
-      and(
-        eq(copilotAsyncToolCalls.toolCallId, input.toolCallId),
-        eq(copilotAsyncToolCalls.runId, input.runId),
-        eq(copilotAsyncToolCalls.status, ASYNC_TOOL_STATUS.running),
-        eq(copilotAsyncToolCalls.executionOwnerToken, input.ownerToken),
-        isNull(copilotAsyncToolCalls.executionRevokedAt),
-        sql`EXISTS (SELECT 1 FROM ${copilotRuns} r WHERE r.id = ${copilotAsyncToolCalls.runId}
-          AND r.tool_admission_closed_at IS NULL)`
+  return db.transaction(async (tx) => {
+    /** Serialized with Stop through the run row, exactly like the claim. */
+    const [run] = await tx
+      .select({ toolAdmissionClosedAt: copilotRuns.toolAdmissionClosedAt })
+      .from(copilotRuns)
+      .where(eq(copilotRuns.id, input.runId))
+      .for('update')
+    if (!run || run.toolAdmissionClosedAt) return false
+    const [row] = await tx
+      .update(copilotAsyncToolCalls)
+      .set({
+        status: input.status,
+        result: sanitizeValueForJsonb(input.result),
+        error: input.error,
+        claimedBy: null,
+        claimedAt: null,
+        completedAt: sql`now()`,
+        executionSettledAt: sql`now()`,
+        updatedAt: sql`now()`,
+      })
+      .where(
+        and(
+          eq(copilotAsyncToolCalls.toolCallId, input.toolCallId),
+          eq(copilotAsyncToolCalls.runId, input.runId),
+          eq(copilotAsyncToolCalls.status, ASYNC_TOOL_STATUS.running),
+          eq(copilotAsyncToolCalls.executionOwnerToken, input.ownerToken),
+          isNull(copilotAsyncToolCalls.executionRevokedAt)
+        )
       )
-    )
-    .returning({ toolCallId: copilotAsyncToolCalls.toolCallId })
-  return Boolean(row)
+      .returning({ toolCallId: copilotAsyncToolCalls.toolCallId })
+    return Boolean(row)
+  })
 }
 
 export type DesktopCallAcknowledgement =

packages/db/migrations/meta/0395_snapshot.json

@@ -3165,7 +3165,7 @@
           ],
           "isUnique": false,
           "where": "\"copilot_runs\".\"desktop_device_id\" IS NOT NULL",
-          "concurrently": false,
+          "concurrently": true,
           "method": "btree",
           "with": {}
         }

packages/db/schema.ts

@@ -4420,7 +4420,8 @@ export const copilotRuns = pgTable(
     streamIdUnique: uniqueIndex('copilot_runs_stream_id_unique').on(table.streamId),
     desktopDeviceStartedAtIdx: index('copilot_runs_desktop_device_started_at_idx')
       .on(table.desktopDeviceId, table.startedAt)
-      .where(sql`${table.desktopDeviceId} IS NOT NULL`),
+      .where(sql`${table.desktopDeviceId} IS NOT NULL`)
+      .concurrently(),
   })
 )