fix(mothership): renew a running code call's sandbox CLI scope past its expiry

Author: waleedlatif1

apps/sim/lib/mothership/tools/sandbox-resources.test.ts

@@ -11,19 +11,33 @@ vi.mock('@/lib/core/config/feature-flags', () => featureFlagsMock)
 
 const { values, lists, seen, persist, redis } = vi.hoisted(() => {
   const values = new Map<string, string>()
+  const expiresAt = new Map<string, number>()
   const lists = new Map<string, string[]>()
   const seen = new Map<string, Set<string>>()
+  const live = (key: string) => {
+    if ((expiresAt.get(key) ?? Number.POSITIVE_INFINITY) <= Date.now()) {
+      values.delete(key)
+      expiresAt.delete(key)
+    }
+    return values.has(key)
+  }
   return {
     values,
     lists,
     seen,
     persist: vi.fn(async () => {}),
     redis: {
-      set: vi.fn(async (key: string, value: string) => {
+      set: vi.fn(async (key: string, value: string, _mode?: 'EX', ttlSeconds?: number) => {
         values.set(key, value)
+        if (ttlSeconds) expiresAt.set(key, Date.now() + ttlSeconds * 1000)
+      }),
+      expire: vi.fn(async (key: string, ttlSeconds: number) => {
+        if (!live(key)) return 0
+        expiresAt.set(key, Date.now() + ttlSeconds * 1000)
+        return 1
       }),
-      get: vi.fn(async (key: string) => values.get(key) ?? null),
-      exists: vi.fn(async (key: string) => (values.has(key) ? 1 : 0)),
+      get: vi.fn(async (key: string) => (live(key) ? (values.get(key) ?? null) : null)),
+      exists: vi.fn(async (key: string) => (live(key) ? 1 : 0)),
       del: vi.fn(async (key: string) => {
         values.delete(key)
         lists.delete(key)
@@ -183,6 +197,33 @@ describe('sandbox resource ownership and publication', () => {
     })
   })
 
+  it('keeps the callback scope live for a call that outlasts its expiry, and revokes it after', async () => {
+    vi.useFakeTimers()
+    try {
+      const finish = Promise.withResolvers<void>()
+      let token = ''
+      const call = withSandboxResourceScope(
+        identity,
+        new AbortController().signal,
+        undefined,
+        async () => {
+          token =
+            (await sandboxResourceEndpoint('https://sim.test', args, 'secret')).split('/').at(-1) ??
+            ''
+          await finish.promise
+        }
+      )
+      await vi.advanceTimersByTimeAsync(15 * 60_000)
+      expect(await readSandboxResourceScope(token, 'secret')).not.toBeNull()
+      finish.resolve()
+      await vi.advanceTimersByTimeAsync(1_000)
+      await call
+      expect(await readSandboxResourceScope(token, 'secret')).toBeNull()
+    } finally {
+      vi.useRealTimers()
+    }
+  })
+
   it('revokes on cancellation and never emits an abandoned callback', async () => {
     const controller = new AbortController()
     const emit = vi.fn()

apps/sim/lib/mothership/tools/sandbox-resources.ts

@@ -16,6 +16,8 @@ import { isProjectsEnabled } from '@/lib/projects/rollout.server'
 
 const logger = createLogger('MothershipSandboxResources')
 const CONTEXT_TTL_SECONDS = 600
+/** A running call renews its scope well inside the TTL; the TTL only bounds a crashed owner. */
+const CONTEXT_RENEW_MS = (CONTEXT_TTL_SECONDS * 1000) / 3
 const POLL_MS = 200
 const scopeSchema = z
   .object({
@@ -233,10 +235,23 @@ export async function withSandboxResourceScope<T>(
     () => {},
     () => {}
   )
+  const renew = async () => {
+    if (!active.token || signal.aborted) return
+    try {
+      await redisClient().expire(contextKey(active.token), CONTEXT_TTL_SECONDS)
+    } catch {
+      logger.warn('Sandbox resource scope renewal deferred', { toolCallId: identity.toolCallId })
+    }
+  }
+  let renewedAt = Date.now()
   try {
     while (running) {
       await Promise.race([settled, sleep(POLL_MS)])
       await drain()
+      if (Date.now() - renewedAt >= CONTEXT_RENEW_MS) {
+        renewedAt = Date.now()
+        await renew()
+      }
     }
     return await execution
   } finally {