fix(mothership): release the wake's own chat lease when the run lookup fails

Author: waleedlatif1

apps/sim/lib/mothership/tasks/application/prepare-wake.integration.ts

@@ -12,7 +12,7 @@ const { redisUrl, inheritedRedisUrl } = await vi.hoisted(async () => {
   const url = readTestRedisUrl()
   const inheritedRedisUrl = process.env.REDIS_URL
   /** The real Redis module reads this at import. */
-  process.env.REDIS_URL = url
+  if (url) process.env.REDIS_URL = url
   return { redisUrl: url, inheritedRedisUrl }
 })
 
@@ -31,7 +31,7 @@ import { copilotChats, copilotRuns, permissions, user, workspace } from '@sim/db
 import { generateId } from '@sim/utils/id'
 import { eq, inArray } from 'drizzle-orm'
 import { NextRequest } from 'next/server'
-import { closeRedisConnection } from '@/lib/core/config/redis'
+import { closeRedisConnection, getRedisClient } from '@/lib/core/config/redis'
 import {
   createRunSegment,
   getLatestRunForStream,
@@ -40,14 +40,15 @@ import {
 import { chatPubSub } from '@/lib/mothership/chat-status'
 import {
   acquirePendingChatStream,
+  getLocalChatStreamLease,
   releasePendingChatStream,
 } from '@/lib/mothership/request/session/abort'
 import { POST as wakeRoute } from '@/app/api/mothership/wake/route'
 
 afterAll(async () => {
   chatPubSub?.dispose()
   await closeRedisConnection()
-  if (inheritedRedisUrl === undefined) process.env.REDIS_URL = undefined
+  if (inheritedRedisUrl === undefined) Reflect.deleteProperty(process.env, 'REDIS_URL')
   else process.env.REDIS_URL = inheritedRedisUrl
 })
 
@@ -174,4 +175,36 @@ describe.runIf(Boolean(redisUrl))('task wake retries', () => {
     expect(await acquirePendingChatStream(chatId, nextTurn, 0)).toBe(true)
     await releasePendingChatStream(chatId, nextTurn)
   })
+
+  it("keeps a retry's chat lock when an earlier wake's slow lookup fails after its own lock expired", async () => {
+    const chatId = await idleChat()
+    const runId = generateId()
+    let lookupStarted!: () => void
+    const started = new Promise<void>((resolve) => {
+      lookupStarted = resolve
+    })
+    let failLookup!: () => void
+    const failed = new Promise<void>((resolve) => {
+      failLookup = resolve
+    })
+    vi.mocked(getLatestRunForStream).mockImplementationOnce(async () => {
+      lookupStarted()
+      await failed
+      throw new Error('statement timeout')
+    })
+
+    const slowWake = wake(chatId, runId)
+    await started
+    /** The first wake's lock outlives its TTL while the lookup hangs. */
+    const firstLease = getLocalChatStreamLease(chatId, runId)
+    await getRedisClient()?.del(firstLease?.key ?? '')
+    expect((await wake(chatId, runId)).status).toBe(202)
+
+    failLookup()
+    expect((await slowWake).status).toBe(500)
+
+    const nextTurn = generateId()
+    expect(await acquirePendingChatStream(chatId, nextTurn, 0)).toBe(false)
+    await releasePendingChatStream(chatId, runId)
+  }, 15_000)
 })

apps/sim/lib/mothership/tasks/application/prepare-wake.ts

@@ -44,24 +44,21 @@ export const prepareTaskWake = defineAuthorizedChatUseCase({
     if (!(await acquirePendingChatStream(input.chatId, input.runId))) {
       throw new OrchestrationError('conflict', 'Another stream holds this chat; retry the wake')
     }
+    const lease = getLocalChatStreamLease(input.chatId, input.runId)
     /**
      * The worker retries a wake under the same run ID until its own run appears. A turn sim
      * already ran under that ID without reaching the worker (a usage-limit refusal) can never
      * open again, so answer not-found: the worker dismisses the notification instead of
      * retrying forever. Checked under the chat lock, so an in-flight turn still answers busy.
-     * Any throw here releases the lock just taken, since the wake turn that would release it
-     * never starts.
+     * Any throw here releases the lock just taken, by its own lease: a slow lookup can outlive
+     * the lock, and a retry under the same run ID may hold the chat by then.
      */
     try {
       if (await getLatestRunForStream(input.runId)) {
         throw new OrchestrationError('not_found', 'This wake already ran')
       }
     } catch (error) {
-      await releasePendingChatStream(
-        input.chatId,
-        input.runId,
-        getLocalChatStreamLease(input.chatId, input.runId)
-      )
+      await releasePendingChatStream(input.chatId, input.runId, lease)
       throw error
     }
     return { accepted: true } as const