fix(mothership): cancel desktop tools a signed-out turn starts late, guard the shown-once license key

apps/sim/app/workspace/[workspaceId]/home/hooks/desktop-tool-lifetimes.test.ts

@@ -1,15 +1,15 @@
 import { describe, expect, it } from 'vitest'
 import {
-  leaseDesktopTool,
+  desktopToolSession,
   stopAllDesktopTools,
   stopDesktopTools,
 } from '@/app/workspace/[workspaceId]/home/hooks/desktop-tool-lifetimes'
 
 describe('desktop tool leases', () => {
   it('cancels every running tool of the stopped turn and no other turn', () => {
-    const first = leaseDesktopTool('turn-a')
-    const second = leaseDesktopTool('turn-a')
-    const other = leaseDesktopTool('turn-b')
+    const first = desktopToolSession().turn('turn-a').lease()
+    const second = desktopToolSession().turn('turn-a').lease()
+    const other = desktopToolSession().turn('turn-b').lease()
 
     stopDesktopTools('turn-a', 'user_stop')
 
@@ -21,9 +21,10 @@ describe('desktop tool leases', () => {
   })
 
   it('keeps a turn reachable by Stop while any of its tools still runs', () => {
-    const settled = leaseDesktopTool('turn-c')
-    const running = leaseDesktopTool('turn-c')
-    for (let turn = 0; turn < 500; turn++) leaseDesktopTool(`busy-${turn}`).release()
+    const settled = desktopToolSession().turn('turn-c').lease()
+    const running = desktopToolSession().turn('turn-c').lease()
+    for (let turn = 0; turn < 500; turn++)
+      desktopToolSession().turn(`busy-${turn}`).lease().release()
 
     settled.release()
     settled.release()
@@ -33,11 +34,11 @@ describe('desktop tool leases', () => {
   })
 
   it('gives a turn whose tools all settled a fresh lifetime for its next tool', () => {
-    const settled = leaseDesktopTool('turn-d')
+    const settled = desktopToolSession().turn('turn-d').lease()
     settled.release()
     stopDesktopTools('turn-d', 'user_stop')
 
-    const next = leaseDesktopTool('turn-d')
+    const next = desktopToolSession().turn('turn-d').lease()
 
     expect(settled.signal.aborted).toBe(false)
     expect(next.signal).not.toBe(settled.signal)
@@ -46,9 +47,9 @@ describe('desktop tool leases', () => {
   })
 
   it('does not let a tool that settles after Stop release a newer lease on the turn', () => {
-    const stopped = leaseDesktopTool('turn-e')
+    const stopped = desktopToolSession().turn('turn-e').lease()
     stopDesktopTools('turn-e', 'user_stop')
-    const next = leaseDesktopTool('turn-e')
+    const next = desktopToolSession().turn('turn-e').lease()
 
     stopped.release()
     stopDesktopTools('turn-e', 'user_stop')
@@ -57,17 +58,39 @@ describe('desktop tool leases', () => {
   })
 
   it('cancels the running tools of every turn when the session ends', () => {
-    const first = leaseDesktopTool('turn-f')
-    const second = leaseDesktopTool('turn-g')
+    const first = desktopToolSession().turn('turn-f').lease()
+    const second = desktopToolSession().turn('turn-g').lease()
 
     stopAllDesktopTools('signed_out')
-    const next = leaseDesktopTool('turn-f')
 
     expect(first.signal.aborted).toBe(true)
     expect(second.signal.aborted).toBe(true)
     expect(first.signal.reason).toBe('signed_out')
-    expect(next.signal.aborted).toBe(false)
     first.release()
+    second.release()
+  })
+
+  it('cancels tools of a surface mounted before sign-out, even on a stream it reads later', () => {
+    const surface = desktopToolSession()
+    const running = surface.turn('turn-h')
+    stopAllDesktopTools('signed_out')
+
+    const late = running.lease()
+    const reconnected = surface.turn('turn-j').lease()
+
+    expect(late.signal.aborted).toBe(true)
+    expect(late.signal.reason).toBe('signed_out')
+    expect(reconnected.signal.aborted).toBe(true)
+    late.release()
+    reconnected.release()
+  })
+
+  it('runs the tools of a surface mounted after the session ended', () => {
+    stopAllDesktopTools('signed_out')
+
+    const next = desktopToolSession().turn('turn-i').lease()
+
+    expect(next.signal.aborted).toBe(false)
     next.release()
   })
 })

apps/sim/app/workspace/[workspaceId]/home/hooks/desktop-tool-lifetimes.ts

@@ -15,6 +15,9 @@ interface RunningTurnTools {
  */
 const runningTurns = new Map<string, RunningTurnTools>()
 
+/** Aborted by `stopAllDesktopTools`, then replaced, so each signed-in session has its own. */
+let session = new AbortController()
+
 /** A running desktop tool's hold on its turn. */
 interface DesktopToolLease {
   /** Aborted only by the user's Stop of the turn, or by signing out. */
@@ -23,13 +26,40 @@ interface DesktopToolLease {
   release(): void
 }
 
+/** A turn's desktop tools, in the session of the chat surface that runs the turn. */
+export interface DesktopToolTurn {
+  /** Starts one desktop tool for the turn. */
+  lease(): DesktopToolLease
+}
+
+/** The desktop tools a chat surface starts, bound to the session the surface mounted in. */
+export interface DesktopToolSession {
+  turn(streamId: string): DesktopToolTurn
+}
+
+/**
+ * Binds a chat surface to the current session. Take it once, when the surface mounts: a send or
+ * reconnect still in flight at sign-out can deliver tool events after the stop, and each of them
+ * then gets an already-aborted lease. Signing out leaves or reloads every chat surface, so a
+ * surface mounted after sign-in binds to the new session.
+ */
+export function desktopToolSession(): DesktopToolSession {
+  const startedIn = session.signal
+  return {
+    turn: (streamId) => ({
+      lease: () =>
+        startedIn.aborted ? { signal: startedIn, release() {} } : leaseDesktopTool(streamId),
+    }),
+  }
+}
+
 /**
  * Starts a desktop tool (a browser action, a local file read or import) for a turn. Only the
  * user's Stop of that turn, or signing out (`stopAllDesktopTools`), cancels it: replacing the
  * stream reader, leaving the chat view, or stopping another chat's turn leaves it running to
  * finish and report its own result.
  */
-export function leaseDesktopTool(streamId: string): DesktopToolLease {
+function leaseDesktopTool(streamId: string): DesktopToolLease {
   let turn = runningTurns.get(streamId)
   if (!turn) {
     turn = { stop: new AbortController(), running: 0 }
@@ -57,9 +87,12 @@ export function stopDesktopTools(streamId: string, reason: string): void {
 
 /**
  * Cancels every leased desktop tool running in this tab (browser actions, local file reads and
- * imports), so none outlives the session that started it.
+ * imports), and every one a turn of this session starts later, so none outlives the session
+ * that started it.
  */
 export function stopAllDesktopTools(reason: string): void {
+  session.abort(reason)
+  session = new AbortController()
   for (const turn of runningTurns.values()) turn.stop.abort(reason)
   runningTurns.clear()
 }

apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts

@@ -93,7 +93,9 @@ import { initTerminalTransport } from '@/lib/terminal/transport'
 import { getQueryClient } from '@/app/_shell/providers/get-query-client'
 import { chatUrl } from '@/app/workspace/[workspaceId]/home/hooks/chat-url'
 import {
-  leaseDesktopTool,
+  type DesktopToolSession,
+  type DesktopToolTurn,
+  desktopToolSession,
   stopDesktopTools,
 } from '@/app/workspace/[workspaceId]/home/hooks/desktop-tool-lifetimes'
 import { useFilePreviewController } from '@/app/workspace/[workspaceId]/home/hooks/preview'
@@ -518,10 +520,10 @@ function startClientBrowserTool(
   toolArgs: Record<string, unknown>,
   scopeId: string,
   eventTs?: string,
-  turnStreamId?: string
+  desktopTurn?: DesktopToolTurn
 ): void {
   if (!isCurrentBrowserToolName(toolName)) return
-  const lease = turnStreamId ? leaseDesktopTool(turnStreamId) : undefined
+  const lease = desktopTurn?.lease()
   void executeBrowserToolOnClient(
     toolCallId,
     toolName,
@@ -1003,6 +1005,8 @@ export function useChat(
   const chatIdRef = useRef<string | undefined>(initialChatId)
   /** Cleared on unmount, so late async work cannot act on a surface the user left. */
   const surfaceMountedRef = useRef(true)
+  const desktopToolsRef = useRef<DesktopToolSession | null>(null)
+  const desktopTools = (desktopToolsRef.current ??= desktopToolSession())
   useEffect(() => {
     surfaceMountedRef.current = true
     return () => {
@@ -1654,7 +1658,7 @@ export function useChat(
       toolCallId: string,
       toolName: string,
       toolArgs: Record<string, unknown>,
-      turnStreamId: string | undefined
+      desktopTurn: DesktopToolTurn | undefined
     ) => {
       if (
         !isNativeFileTool(toolName) &&
@@ -1666,7 +1670,7 @@ export function useChat(
         return
       }
       handledClientLocalFilesystemToolIds.add(toolCallId)
-      const lease = turnStreamId ? leaseDesktopTool(turnStreamId) : undefined
+      const lease = desktopTurn?.lease()
       const options = {
         workspaceId,
         chatId: chatIdRef.current ?? selectedChatIdRef.current,
@@ -2297,7 +2301,7 @@ export function useChat(
         shouldContinue?: () => boolean
       }
     ) => {
-      const turnStreamId = streamIdRef.current
+      const desktopTurn = streamIdRef.current ? desktopTools.turn(streamIdRef.current) : undefined
       const activityTracker = getResourceActivityTracker(
         expectedGen ?? streamGenRef.current,
         options?.targetChatId
@@ -2318,7 +2322,7 @@ export function useChat(
         eventTs?: string
       ) => {
         const scopeId = activityScopeId()
-        startClientBrowserTool(toolCallId, toolName, toolArgs, scopeId, eventTs, turnStreamId)
+        startClientBrowserTool(toolCallId, toolName, toolArgs, scopeId, eventTs, desktopTurn)
       }
       const startClientTerminalToolForStream = (
         toolCallId: string,
@@ -2351,7 +2355,7 @@ export function useChat(
         removeResource,
         startClientWorkflowTool,
         startClientLocalFilesystemTool: (toolCallId, toolName, toolArgs) =>
-          startClientLocalFilesystemTool(toolCallId, toolName, toolArgs, turnStreamId),
+          startClientLocalFilesystemTool(toolCallId, toolName, toolArgs, desktopTurn),
         startClientBrowserTool: startClientBrowserToolForStream,
         startClientTerminalTool: startClientTerminalToolForStream,
         startBrowserAgentRun: startBrowserAgentRunForStream,