fix(mothership): time the lock lapse from its acquisition, type the re-run predicates as conditions, revoke access mid-recovery in the test

Author: waleedlatif1

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

@@ -24,6 +24,7 @@ import {
   like,
   lt,
   ne,
+  not,
   notInArray,
   or,
   type SQL,
@@ -733,12 +734,23 @@ export type ToolExecutionClaim =
   /** Already claimed or settled. */
   | { outcome: 'existing' }
 
+/**
+ * Every condition at once, typed as a condition: drizzle's `and` may return undefined, which a
+ * negation would render as `NOT (NULL)` and match nothing.
+ */
+function allOf(...conditions: SQL[]): SQL {
+  return sql`(${sql.join(
+    conditions.map((condition) => sql`(${condition})`),
+    sql` and `
+  )})`
+}
+
 /**
  * A Sim handler's execution of a call that has no result yet. Only a Sim handler sets
  * `executionStartedAt` without a browser workflow execution, so desktop and browser
  * executions, which settle their own lapses, never match.
  */
-const unfinishedSimExecution = and(
+const unfinishedSimExecution = allOf(
   eq(copilotAsyncToolCalls.status, ASYNC_TOOL_STATUS.running),
   isNotNull(copilotAsyncToolCalls.executionStartedAt),
   isNull(copilotAsyncToolCalls.executionSettledAt),
@@ -751,7 +763,7 @@ const unfinishedSimExecution = and(
  * owner's heartbeat never lets happen. A takeover runs it again until it has started
  * {@link SIM_TOOL_EXECUTION_MAX_ATTEMPTS} times.
  */
-const interruptedSimExecution = and(
+const interruptedSimExecution = allOf(
   unfinishedSimExecution,
   sql`${copilotAsyncToolCalls.executionLeaseExpiresAt} <= clock_timestamp()`,
   lt(copilotAsyncToolCalls.executionAttempts, SIM_TOOL_EXECUTION_MAX_ATTEMPTS)
@@ -761,7 +773,7 @@ const interruptedSimExecution = and(
  * An interrupted execution a takeover may still run again: its run still admits tools. Another
  * call's wait or workbench access must not revoke it as interrupted before that re-run.
  */
-const rerunnableSimExecution = and(
+const rerunnableSimExecution = allOf(
   interruptedSimExecution,
   sql`EXISTS (SELECT 1 FROM ${copilotRuns} r WHERE r.id = ${copilotAsyncToolCalls.runId} AND r.tool_admission_closed_at IS NULL AND r.status NOT IN (${sql.join(
     TERMINAL_RUN_STATUSES.map((status) => sql`${status}`),
@@ -1021,10 +1033,7 @@ export async function revokeExpiredSimToolExecutions(input: {
       tx,
       and(
         sql`EXISTS (SELECT 1 FROM ${copilotRuns} r WHERE r.id = ${copilotAsyncToolCalls.runId} AND r.id = ${input.runId}::uuid AND r.user_id = ${input.userId})`,
-        or(
-          eq(copilotAsyncToolCalls.toolCallId, input.toolCallId),
-          sql`NOT (${rerunnableSimExecution})`
-        )
+        or(eq(copilotAsyncToolCalls.toolCallId, input.toolCallId), not(rerunnableSimExecution))
       )
     )
   )
@@ -1387,7 +1396,7 @@ export async function prepareWorkbenchAccess(
       and(
         sql`EXISTS (SELECT 1 FROM ${copilotRuns} r WHERE r.id = ${copilotAsyncToolCalls.runId} AND r.chat_id = ${owner.chatId} AND r.user_id = ${input.userId})`,
         /** This run was just verified to admit tools; every older run's admission just closed. */
-        sql`NOT (${and(eq(copilotAsyncToolCalls.runId, input.runId), interruptedSimExecution)})`
+        not(allOf(eq(copilotAsyncToolCalls.runId, input.runId), interruptedSimExecution))
       )
     )
     const previous = await tx

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

@@ -57,16 +57,27 @@ const { redisUrl, inheritedEnv, worker } = await vi.hoisted(async () => {
 vi.mock('@/lib/auth', () => authMock)
 vi.mock('@/lib/core/config/env-flags', () => envFlagsMock)
 vi.mock('@/lib/core/config/feature-flags', () => featureFlagsMock)
-/** Recovery's own workspace re-check, so a test can lose access after authorization passed. */
-const { workspacePermission } = vi.hoisted(() => ({ workspacePermission: { revoked: false } }))
+/**
+ * Recovery's own workspace re-check runs only after authorization passed; `beforeRecheck` lets
+ * a test change real access at exactly that point.
+ */
+const { workspacePermission } = vi.hoisted(() => ({
+  workspacePermission: {
+    rechecks: 0,
+    beforeRecheck: undefined as (() => Promise<void>) | undefined,
+  },
+}))
 vi.mock('@/lib/workspaces/permissions/utils', async (importOriginal) => {
   const original = await importOriginal<typeof import('@/lib/workspaces/permissions/utils')>()
   return {
     ...original,
-    getUserEntityPermissions: (...args: Parameters<typeof original.getUserEntityPermissions>) =>
-      workspacePermission.revoked
-        ? Promise.resolve(null)
-        : original.getUserEntityPermissions(...args),
+    getUserEntityPermissions: async (
+      ...args: Parameters<typeof original.getUserEntityPermissions>
+    ) => {
+      workspacePermission.rechecks++
+      await workspacePermission.beforeRecheck?.()
+      return original.getUserEntityPermissions(...args)
+    },
   }
 })
 
@@ -344,7 +355,8 @@ describe.runIf(Boolean(redisUrl))('taking over a run whose controller died', ()
     worker.requests.length = 0
     resetEnvFlagsMock()
     featureFlagsMockFns.mockIsFeatureEnabled.mockResolvedValue(false)
-    workspacePermission.revoked = false
+    workspacePermission.rechecks = 0
+    workspacePermission.beforeRecheck = undefined
     await redis().del(ABANDONED_RUN_CURSOR_KEY)
     await db
       .insert(permissions)
@@ -722,11 +734,18 @@ describe.runIf(Boolean(redisUrl))('taking over a run whose controller died', ()
   })
 
   it('settles the run as revoked when access is lost after authorization passed', async () => {
-    workspacePermission.revoked = true
     const { streamId, runId, toolCallId } = await runAbandonedMidTool()
+    /** The user is removed from the workspace after authorization admitted them. */
+    workspacePermission.beforeRecheck = async () => {
+      await db
+        .delete(permissions)
+        .where(and(eq(permissions.userId, userId), eq(permissions.entityId, workspaceId)))
+    }
 
     await takeOverAbandonedRuns()
 
+    expect(workspacePermission.rechecks).toBe(1)
+
     const run = await storedRun(runId)
     expect(run.status).toBe('error')
     expect(run.error).toBe(REVOKED_RUN_ERROR)

apps/sim/lib/mothership/request/session/abort.test.ts

@@ -127,6 +127,33 @@ describe('startAbortPoller heartbeat', () => {
     }
   })
 
+  it('times the lapse from when the lock was taken, not from when polling started', async () => {
+    const controller = new AbortController()
+    const streamId = 'stream-slow-setup'
+    const chatId = 'chat-slow-setup'
+    redisConfigMockFns.mockExtendLock.mockRejectedValue(new Error('redis unreachable'))
+
+    const interval = startAbortPoller(streamId, controller, {
+      chatId,
+      lease: {
+        key: `copilot:chat-stream-lock:${chatId}`,
+        value: `${streamId}\ncontroller-1`,
+        acquiredAt: Date.now() - 50_000,
+      },
+    })
+
+    try {
+      await vi.advanceTimersByTimeAsync(9_000)
+      expect(controller.signal.aborted).toBe(false)
+
+      await vi.advanceTimersByTimeAsync(2_000)
+      expect(controller.signal.aborted).toBe(true)
+      expect(controller.signal.reason.name).toBe('StreamControllerSupersededError')
+    } finally {
+      clearInterval(interval)
+    }
+  })
+
   it('aborts the controller before clearing the marker so the marker is never observable as cleared while the signal is still unaborted', async () => {
     const controller = new AbortController()
     const streamId = 'stream-order-1'

apps/sim/lib/mothership/request/session/abort.ts

@@ -227,14 +227,19 @@ export async function acquirePendingChatStream(
     },
     async (span) => {
       const redis = getRedisClient()
-      const lease = { key: chatStreamLockKey(chatId), value: `${streamId}\n${generateId()}` }
+      const lease: ChatStreamLease = {
+        key: chatStreamLockKey(chatId),
+        value: `${streamId}\n${generateId()}`,
+      }
       span.setAttribute(TraceAttr.LockBackend, redis ? AbortBackend.Redis : AbortBackend.InProcess)
       if (redis) {
         const deadline = Date.now() + timeoutMs
         for (;;) {
           try {
+            const attemptedAt = Date.now()
             const acquired = await acquireLock(lease.key, lease.value, CHAT_STREAM_LOCK_TTL_SECONDS)
             if (acquired) {
+              lease.acquiredAt = attemptedAt
               registerPendingChatStream(chatId, streamId, lease)
               span.setAttribute(TraceAttr.LockAcquired, true)
               return true
@@ -357,7 +362,8 @@ export function startAbortPoller(
   const chatId = options?.chatId
   const lease = options?.lease ?? (chatId ? getLocalChatStreamLease(chatId, streamId) : undefined)
 
-  let lastHeartbeatAt = Date.now()
+  /** The lock's TTL runs from its last confirmed write: the acquisition until a renewal lands. */
+  let lastHeartbeatAt = lease?.acquiredAt ?? Date.now()
   let heartbeatOwnershipLost = false
   const stopController = (message: string) => {
     heartbeatOwnershipLost = true

apps/sim/lib/mothership/request/session/controller-lease.ts

@@ -4,6 +4,8 @@ import { getRedisClient } from '@/lib/core/config/redis'
 export interface ChatStreamLease {
   key: string
   value: string
+  /** When the lock was taken (ms epoch, read before the write), so its TTL is known to run from then. */
+  acquiredAt?: number
 }
 
 export class StreamControllerSupersededError extends Error {