fix(mothership): take over runs whose controller died and re-run their interrupted tool

apps/sim/instrumentation-node.ts

@@ -414,4 +414,9 @@ export async function register() {
 
   const { startServiceUsageReplay } = await import('./lib/mothership/billing/service-delivery')
   startServiceUsageReplay()
+
+  const { startAbandonedRunTakeover } = await import(
+    '@/lib/mothership/request/application/abandoned-runs'
+  )
+  startAbandonedRunTakeover()
 }

apps/sim/lib/core/config/feature-flags.ts

@@ -120,6 +120,13 @@ const FEATURE_FLAGS = {
       'targeting. Off-AppConfig falls back to KNOWLEDGE_MEMBER_ACCESS.',
     fallback: 'KNOWLEDGE_MEMBER_ACCESS',
   },
+  'mothership-abandoned-run-takeover-disabled': {
+    description:
+      'Kill switch for Sim taking over Chat runs whose controller died (a deploy or a crash) and ' +
+      're-running their interrupted tool. Global on/off only; on stops every takeover scan. ' +
+      'Off by default, so takeovers run unless this is turned on.',
+    fallback: null,
+  },
 } satisfies Record<string, FeatureFlagDefinition>
 
 /**

apps/sim/lib/desktop/executor/bound-turn.integration.ts

@@ -511,7 +511,7 @@ describe.runIf(Boolean(redisUrl))("a turn bound to a desktop's background execut
 
       await lapse(toolCallId, 'lease')
       /** A Sim tool waiting on the same run sweeps its expired executions; this one is not its. */
-      await revokeExpiredSimToolExecutions({ runId: run.runId, userId })
+      await revokeExpiredSimToolExecutions({ runId: run.runId, userId, toolCallId: 'waiting-call' })
       expect((await storedCall(toolCallId)).status).toBe('running')
       await answer
       const late = await desktop.complete(toolCallId, executionToken, { output: 'built' })

apps/sim/lib/mothership/async-runs/browser-download-claim.integration.ts

@@ -52,6 +52,7 @@ describe('browser download admission with PostgreSQL', () => {
       claimed_at timestamp, claimed_by text, execution_started_at timestamp,
       execution_settled_at timestamp, execution_owner_token text,
       execution_lease_expires_at timestamptz, execution_revoked_at timestamptz,
+      execution_attempts integer NOT NULL DEFAULT 0,
       client_workflow_execution_id text, sandbox_processes jsonb NOT NULL DEFAULT '{}',
       completed_at timestamp, created_at timestamp NOT NULL DEFAULT now(),
       updated_at timestamp NOT NULL DEFAULT now()

apps/sim/lib/mothership/async-runs/execution-lease.ts

@@ -9,9 +9,18 @@ export interface SimToolExecutionOwner {
   ownerToken: string
 }
 
+/**
+ * Sim executions one call may start: the first, and a re-run each time a takeover finds the
+ * previous owner died before recording a result. A tool that kills its executor every time
+ * ends with {@link EXHAUSTED_SIM_TOOL_MESSAGE} instead of re-running forever.
+ */
+export const SIM_TOOL_EXECUTION_MAX_ATTEMPTS = 3
+
 export const INTERRUPTED_SIM_TOOL_MESSAGE =
   'Tool execution lost its owner before a result was recorded. Its outcome is unknown; inspect current state before retrying a mutation.'
 
+export const EXHAUSTED_SIM_TOOL_MESSAGE = `Tool execution was interrupted ${SIM_TOOL_EXECUTION_MAX_ATTEMPTS} times because the Sim server running it stopped, so it was not run again. Its outcome is unknown; inspect current state before retrying a mutation.`
+
 export class SimToolExecutionLeaseLostError extends Error {
   constructor() {
     super(INTERRUPTED_SIM_TOOL_MESSAGE)

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

@@ -73,6 +73,8 @@ export const ORPHANED_RUN_GRACE_MS = 60 * 60 * 1000
 export const LEGACY_RUN_GRACE_MS = 24 * 60 * 60 * 1000
 
 export const ORPHANED_RUN_ERROR = 'This response was interrupted before it finished.'
+export const REVOKED_RUN_ERROR =
+  'This response was stopped because you no longer have access to continue it.'
 export const LEGACY_RUN_ERROR = 'Run was never finalized (pre-lease run).'
 
 const SWEEP_BATCH_SIZE = 500
@@ -124,7 +126,7 @@ const stopRequested = sql`(EXISTS (SELECT 1 FROM ${copilotRequestStops} s
     WHERE s.user_id = ${copilotRuns.userId} AND s.organization_id = ${copilotRuns.organizationId}
       AND s.stream_id = ${copilotRuns.streamId}))`
 
-interface UnownedRun {
+export interface UnownedRun {
   id: string
   chatId: string
   streamId: string
@@ -144,9 +146,15 @@ const unownedRunColumns = {
   controllerToken,
 }
 
+const SETTLED_RUN_ERRORS = {
+  orphaned: ORPHANED_RUN_ERROR,
+  revoked: REVOKED_RUN_ERROR,
+  legacy: LEGACY_RUN_ERROR,
+} as const
+
 /** A run ends cancelled only if its user pressed Stop, whoever settles it. */
-function terminalValues(reason: 'orphaned' | 'legacy') {
-  const error = reason === 'orphaned' ? ORPHANED_RUN_ERROR : LEGACY_RUN_ERROR
+function terminalValues(reason: keyof typeof SETTLED_RUN_ERRORS) {
+  const error = SETTLED_RUN_ERRORS[reason]
   return {
     status: sql`(CASE WHEN ${stopRequested} THEN 'cancelled' ELSE 'error' END)::copilot_run_status`,
     error: sql`CASE WHEN ${stopRequested} THEN NULL ELSE ${error}::text END`,
@@ -171,7 +179,8 @@ function terminalValues(reason: 'orphaned' | 'legacy') {
 async function settleRuns(
   tx: DbTransaction,
   runs: UnownedRun[],
-  guard: SQL | undefined
+  guard: SQL | undefined,
+  reason: Exclude<keyof typeof SETTLED_RUN_ERRORS, 'legacy'> = 'orphaned'
 ): Promise<{ settled: UnownedRun[]; released: UnownedRun[] }> {
   if (runs.length === 0) return { settled: [], released: [] }
   const chatIds = [...new Set(runs.map((run) => run.chatId))]
@@ -235,7 +244,7 @@ async function settleRuns(
   for (const run of runs) {
     if (run.controllerToken === null) continue
     await apply([run], eq(controllerToken, run.controllerToken), {
-      ...terminalValues('orphaned'),
+      ...terminalValues(reason),
       completedAt: sql`now()`,
       updatedAt: sql`now()`,
     })
@@ -337,9 +346,9 @@ async function releaseChatLocks(fences: ChatLockFence[]): Promise<void> {
   )
 }
 
-async function readSweepCursor(): Promise<string | undefined> {
+async function readSweepCursor(key = SWEEP_CURSOR_KEY): Promise<string | undefined> {
   try {
-    return (await getRedisClient()?.get(SWEEP_CURSOR_KEY)) ?? undefined
+    return (await getRedisClient()?.get(key)) ?? undefined
   } catch (error) {
     logger.warn('Orphaned-run sweep cursor is unreadable; starting from the first run', {
       error: getErrorMessage(error),
@@ -348,12 +357,12 @@ async function readSweepCursor(): Promise<string | undefined> {
   }
 }
 
-async function writeSweepCursor(cursor: string | undefined): Promise<void> {
+async function writeSweepCursor(cursor: string | undefined, key = SWEEP_CURSOR_KEY): Promise<void> {
   try {
     const redis = getRedisClient()
     if (!redis) return
-    if (cursor) await redis.set(SWEEP_CURSOR_KEY, cursor, 'EX', SWEEP_CURSOR_TTL_SECONDS)
-    else await redis.del(SWEEP_CURSOR_KEY)
+    if (cursor) await redis.set(key, cursor, 'EX', SWEEP_CURSOR_TTL_SECONDS)
+    else await redis.del(key)
   } catch (error) {
     logger.warn('Orphaned-run sweep cursor could not be saved', { error: getErrorMessage(error) })
   }
@@ -453,3 +462,96 @@ export async function settleStoppedRunWithoutController(runId: string): Promise<
   announceReleased(released)
   return settled.length > 0
 }
+
+/**
+ * Settles a run its controller abandoned as an error when its user may no longer continue
+ * it. The run's chat lock is fenced as a takeover's would be, so a reconnect cannot claim
+ * the run while it settles; a run whose lock is held, or whose controller changed, is left.
+ */
+export async function settleRunWithoutAccess(run: UnownedRun): Promise<boolean> {
+  if (!run.controllerToken) return false
+  const fences = await fenceChatLocks([run])
+  if (fences.length === 0) return false
+  try {
+    const { settled, released } = await db.transaction((tx) =>
+      settleRuns(tx, [run], undefined, 'revoked')
+    )
+    announceReleased(released)
+    return settled.length > 0
+  } finally {
+    await releaseChatLocks(fences)
+  }
+}
+
+/**
+ * A run its controller is still driving writes its status at least once per leg, and a dead
+ * controller's chat lock lapses within this long, so a run written to more recently is not
+ * worth checking yet.
+ */
+const ABANDONED_RUN_IDLE_MS = 60_000
+/** Where the last takeover scan stopped, so live runs it skipped never hide the runs after them. */
+export const ABANDONED_RUN_CURSOR_KEY = 'copilot:abandoned-runs:scan-cursor'
+
+/**
+ * Runs an interactive controller abandoned that a takeover can still resume: unfinished and
+ * not stopped, saved with their recovery intent, idle for a minute but within the recovery
+ * window, and with no controller holding their chat lock. Each scan examines at most
+ * `maxExamined` candidates, resuming after the last one the previous scan examined and wrapping
+ * to the first, so runs whose controller is alive cost one examination per pass and never
+ * crowd out an abandoned run. Empty while the locks cannot be read, since then no run is
+ * provably abandoned.
+ */
+export async function findAbandonedRecoverableRuns(
+  limit: number,
+  { pageSize = SWEEP_BATCH_SIZE, maxExamined = SWEEP_BATCH_SIZE } = {}
+): Promise<UnownedRun[]> {
+  const abandoned: UnownedRun[] = []
+  const start = await readSweepCursor(ABANDONED_RUN_CURSOR_KEY)
+  let cursor = start
+  let wrapped = start === undefined
+  let examined = 0
+  while (abandoned.length < limit && examined < maxExamined) {
+    const size = Math.min(pageSize, maxExamined - examined)
+    const page: UnownedRun[] = await db
+      .select(unownedRunColumns)
+      .from(copilotRuns)
+      .where(
+        and(
+          inArray(copilotRuns.status, UNFINISHED_RUN_STATUSES),
+          isNotNull(controllerToken),
+          sql`${copilotRuns.requestContext}->'recovery'->>'kind' = 'interactive_stream'`,
+          idleFor(ABANDONED_RUN_IDLE_MS),
+          not(idleFor(ORPHANED_RUN_GRACE_MS)),
+          not(stopRequested),
+          cursor ? gt(copilotRuns.id, cursor) : undefined,
+          wrapped && start ? lte(copilotRuns.id, start) : undefined
+        )
+      )
+      .orderBy(asc(copilotRuns.id))
+      .limit(size)
+    let locked: Set<string>
+    try {
+      locked = page.length > 0 ? await findStreamsHoldingChatLock(page) : new Set()
+    } catch (error) {
+      logger.warn('Chat stream ownership is unreadable; no run is taken over this round', {
+        error: getErrorMessage(error),
+      })
+      return []
+    }
+    for (const run of page) {
+      examined++
+      cursor = run.id
+      if (locked.has(run.streamId)) continue
+      abandoned.push(run)
+      if (abandoned.length === limit) break
+    }
+    if (abandoned.length < limit && page.length < size) {
+      /** The end of the table: wrap once to cover the runs before the starting point. */
+      cursor = undefined
+      if (wrapped) break
+      wrapped = true
+    }
+  }
+  await writeSweepCursor(cursor, ABANDONED_RUN_CURSOR_KEY)
+  return abandoned
+}

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

@@ -5,7 +5,13 @@ import {
   copilotRequestStops,
   copilotRuns,
 } from '@sim/db/schema'
-import { dbChainMockFns, hasMockCondition, queueTableRows, resetDbChainMock } from '@sim/testing'
+import {
+  dbChainMockFns,
+  flattenMockConditions,
+  hasMockCondition,
+  queueTableRows,
+  resetDbChainMock,
+} from '@sim/testing'
 import { beforeEach, describe, expect, it } from 'vitest'
 import {
   areStreamToolExecutionsSettled,
@@ -185,6 +191,21 @@ describe('run admission and early Stop', () => {
   })
 })
 
+/** Finds a condition anywhere in a mocked expression, including inside `sql` and `sql.join` fragments. */
+function someNestedCondition(
+  node: unknown,
+  predicate: (condition: Record<string, unknown>) => boolean
+): boolean {
+  if (!node || typeof node !== 'object') return false
+  const record = node as Record<string, unknown>
+  if (predicate(record)) return true
+  return ['conditions', 'values', 'fragments'].some(
+    (key) =>
+      Array.isArray(record[key]) &&
+      (record[key] as unknown[]).some((child) => someNestedCondition(child, predicate))
+  )
+}
+
 describe('durable Sim tool ownership', () => {
   beforeEach(() => {
     resetDbChainMock()
@@ -240,14 +261,26 @@ describe('durable Sim tool ownership', () => {
         (condition) => condition.type === 'eq' && condition.right === input.userId
       )
     ).toBe(true)
+    /** A first execution, or a re-run of one whose owner died before it recorded a result. */
+    const claimable = flattenMockConditions(dbChainMockFns.where.mock.calls[1]?.[0]).find(
+      (condition) => condition.type === 'or'
+    )
+    const [firstExecution, interruptedExecution] = (claimable?.conditions ?? []) as unknown[]
     expect(
       hasMockCondition(
-        dbChainMockFns.where.mock.calls[1]?.[0],
+        firstExecution,
         (condition) =>
           condition.type === 'isNull' &&
           condition.column === copilotAsyncToolCalls.executionStartedAt
       )
     ).toBe(true)
+    expect(
+      someNestedCondition(
+        interruptedExecution,
+        (condition) =>
+          condition.type === 'lt' && condition.left === copilotAsyncToolCalls.executionAttempts
+      )
+    ).toBe(true)
     expect(dbChainMockFns.set).toHaveBeenCalledWith(
       expect.objectContaining({ executionStartedAt: expect.any(Date), status: 'running' })
     )