fix(mothership): page past live controllers, stop a scan on shutdown, dispose chat pubsub in tests

Author: waleedlatif1

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

@@ -497,29 +497,38 @@ const ABANDONED_RUN_IDLE_MS = 60_000
  * no run is provably abandoned.
  */
 export async function findAbandonedRecoverableRuns(limit: number): Promise<UnownedRun[]> {
-  const candidates: 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)
+  const abandoned: UnownedRun[] = []
+  let after: string | undefined
+  /** Pages past runs whose controller is alive, so they never crowd out an abandoned one. */
+  while (abandoned.length < limit) {
+    const candidates: 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),
+          after ? gt(copilotRuns.id, after) : undefined
+        )
       )
-    )
-    .orderBy(asc(copilotRuns.updatedAt))
-    .limit(limit)
-  if (candidates.length === 0) return []
-  try {
-    const locked = await findStreamsHoldingChatLock(candidates)
-    return candidates.filter((run) => !locked.has(run.streamId))
-  } catch (error) {
-    logger.warn('Chat stream ownership is unreadable; no run is taken over this round', {
-      error: getErrorMessage(error),
-    })
-    return []
+      .orderBy(asc(copilotRuns.id))
+      .limit(SWEEP_BATCH_SIZE)
+    if (candidates.length === 0) break
+    after = candidates[candidates.length - 1].id
+    try {
+      const locked = await findStreamsHoldingChatLock(candidates)
+      abandoned.push(...candidates.filter((run) => !locked.has(run.streamId)))
+    } catch (error) {
+      logger.warn('Chat stream ownership is unreadable; no run is taken over this round', {
+        error: getErrorMessage(error),
+      })
+      return []
+    }
+    if (candidates.length < SWEEP_BATCH_SIZE) break
   }
+  return abandoned.slice(0, limit)
 }

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

@@ -81,6 +81,7 @@ import {
   revokeExhaustedSimToolExecutions,
   upsertAsyncToolCall,
 } from '@/lib/mothership/async-runs/repository'
+import { chatPubSub } from '@/lib/mothership/chat-status'
 import { takeOverAbandonedRuns } from '@/lib/mothership/request/application/abandoned-runs'
 import {
   createProviderToolCallIdentity,
@@ -324,6 +325,7 @@ describe.runIf(Boolean(redisUrl))('taking over a run whose controller died', ()
     await db.delete(permissions).where(eq(permissions.entityId, workspaceId))
     await deleteWorkspaceFixture(db, eq(workspace.id, workspaceId))
     for (const id of [ownerId, userId]) await db.delete(user).where(eq(user.id, id))
+    chatPubSub?.dispose()
     await closeRedisConnection()
   })
 

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

@@ -26,6 +26,7 @@ const DELEGATION_TTL_MS = 60_000
 
 let timer: ReturnType<typeof setInterval> | undefined
 let running = false
+let stopped = false
 
 /**
  * Resumes the run as its user, re-checking that user's current access. A user who may no
@@ -77,6 +78,7 @@ export async function takeOverAbandonedRuns(): Promise<void> {
   running = true
   try {
     for (const run of await findAbandonedRecoverableRuns(TAKEOVER_BATCH_SIZE)) {
+      if (stopped) return
       await takeOver(run).catch((error) =>
         logger.warn('Abandoned run could not be taken over; a later round retries', {
           runId: run.id,
@@ -102,7 +104,10 @@ export function startAbandonedRunTakeover(): void {
     )
   }, TAKEOVER_INTERVAL_MS)
   timer.unref()
-  const stop = () => clearInterval(timer)
+  const stop = () => {
+    stopped = true
+    clearInterval(timer)
+  }
   process.once('SIGTERM', stop)
   process.once('SIGINT', stop)
 }