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

- A turn's stream now binds to the session it started in. Its tool events can arrive after sign-out stops every desktop tool, and each one then gets an already-aborted lease instead of a fresh controller. A turn started after sign-in runs normally.
- The generated license key counts as an unsaved change, so leaving the Licenses tab asks first, and confirming drops it.

Author: waleedlatif1

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

@@ -1,15 +1,15 @@
 import { describe, expect, it } from 'vitest'
 import {
-  leaseDesktopTool,
+  desktopToolTurn,
   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 = desktopToolTurn('turn-a').lease()
+    const second = desktopToolTurn('turn-a').lease()
+    const other = desktopToolTurn('turn-b').lease()
 
     stopDesktopTools('turn-a', 'user_stop')
 
@@ -21,9 +21,9 @@ 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 = desktopToolTurn('turn-c').lease()
+    const running = desktopToolTurn('turn-c').lease()
+    for (let turn = 0; turn < 500; turn++) desktopToolTurn(`busy-${turn}`).lease().release()
 
     settled.release()
     settled.release()
@@ -33,11 +33,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 = desktopToolTurn('turn-d').lease()
     settled.release()
     stopDesktopTools('turn-d', 'user_stop')
 
-    const next = leaseDesktopTool('turn-d')
+    const next = desktopToolTurn('turn-d').lease()
 
     expect(settled.signal.aborted).toBe(false)
     expect(next.signal).not.toBe(settled.signal)
@@ -46,9 +46,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 = desktopToolTurn('turn-e').lease()
     stopDesktopTools('turn-e', 'user_stop')
-    const next = leaseDesktopTool('turn-e')
+    const next = desktopToolTurn('turn-e').lease()
 
     stopped.release()
     stopDesktopTools('turn-e', 'user_stop')
@@ -57,17 +57,35 @@ 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 = desktopToolTurn('turn-f').lease()
+    const second = desktopToolTurn('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 a tool a turn of the ended session starts after sign-out', () => {
+    const turn = desktopToolTurn('turn-h')
+    stopAllDesktopTools('signed_out')
+
+    const late = turn.lease()
+
+    expect(late.signal.aborted).toBe(true)
+    expect(late.signal.reason).toBe('signed_out')
+    late.release()
+  })
+
+  it('runs the tools of a turn started after the session ended', () => {
+    stopAllDesktopTools('signed_out')
+
+    const next = desktopToolTurn('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,31 @@ interface DesktopToolLease {
   release(): void
 }
 
+/** A turn's stream, bound to the session it started in. */
+export interface DesktopToolTurn {
+  /** Starts one desktop tool for the turn. */
+  lease(): DesktopToolLease
+}
+
+/**
+ * Binds a turn's stream to the current session. Take it once, when the stream starts: its tool
+ * events can still arrive after a sign-out, and each of them then gets an already-aborted lease.
+ */
+export function desktopToolTurn(streamId: string): DesktopToolTurn {
+  const startedIn = session.signal
+  return {
+    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 +78,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,8 @@ 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 DesktopToolTurn,
+  desktopToolTurn,
   stopDesktopTools,
 } from '@/app/workspace/[workspaceId]/home/hooks/desktop-tool-lifetimes'
 import { useFilePreviewController } from '@/app/workspace/[workspaceId]/home/hooks/preview'
@@ -518,10 +519,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,
@@ -1654,7 +1655,7 @@ export function useChat(
       toolCallId: string,
       toolName: string,
       toolArgs: Record<string, unknown>,
-      turnStreamId: string | undefined
+      desktopTurn: DesktopToolTurn | undefined
     ) => {
       if (
         !isNativeFileTool(toolName) &&
@@ -1666,7 +1667,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 +2298,7 @@ export function useChat(
         shouldContinue?: () => boolean
       }
     ) => {
-      const turnStreamId = streamIdRef.current
+      const desktopTurn = streamIdRef.current ? desktopToolTurn(streamIdRef.current) : undefined
       const activityTracker = getResourceActivityTracker(
         expectedGen ?? streamGenRef.current,
         options?.targetChatId
@@ -2318,7 +2319,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 +2352,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,

apps/sim/app/workspace/[workspaceId]/settings/components/mothership/mothership.test.tsx

@@ -0,0 +1,131 @@
+/**
+ * @vitest-environment jsdom
+ */
+import { act, type ChangeEventHandler, type ReactNode } from 'react'
+import { emcnMock } from '@sim/testing/mocks/emcn.mock'
+import { NuqsTestingAdapter } from 'nuqs/adapters/testing'
+import { createRoot, type Root } from 'react-dom/client'
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
+
+const { mockGenerate } = vi.hoisted(() => ({ mockGenerate: vi.fn() }))
+
+vi.mock('@sim/emcn', () => ({
+  ...emcnMock,
+  Badge: ({ children }: { children?: ReactNode }) => <span>{children}</span>,
+  Button: ({ children, ...props }: { children?: ReactNode }) => (
+    <button type='button' {...props}>
+      {children}
+    </button>
+  ),
+  ChipCopyInput: ({ value }: { value?: string }) => (
+    <input data-testid='license-key' readOnly value={value ?? ''} />
+  ),
+  ChipInput: ({
+    value,
+    onChange,
+    placeholder,
+  }: {
+    value?: string
+    onChange?: ChangeEventHandler<HTMLInputElement>
+    placeholder?: string
+  }) => <input placeholder={placeholder} value={value ?? ''} onChange={onChange} />,
+  ChipModalTabs: () => <div />,
+  ChipSelect: () => <div />,
+  Label: ({ children }: { children?: ReactNode }) => <span>{children}</span>,
+  Skeleton: () => <div />,
+}))
+
+vi.mock('@/app/workspace/[workspaceId]/settings/components/settings-panel', () => ({
+  SettingsPanel: ({ children }: { children?: ReactNode }) => <div>{children}</div>,
+}))
+
+vi.mock('@/app/workspace/[workspaceId]/settings/components/settings-empty-state', () => ({
+  SettingsEmptyState: () => null,
+}))
+
+vi.mock('@/hooks/queries/mothership-admin', () => ({
+  useGenerateLicense: () => ({ mutate: mockGenerate, isPending: false, error: null }),
+  useMothershipLicenses: () => ({ data: undefined, isLoading: false }),
+  useMothershipRequests: () => ({ data: undefined, isLoading: false }),
+  useMothershipUserBreakdown: () => ({ data: undefined, isLoading: false }),
+}))
+
+import { Mothership } from '@/app/workspace/[workspaceId]/settings/components/mothership/mothership'
+import { useSettingsDirtyStore } from '@/stores/settings/dirty/store'
+
+let container: HTMLDivElement
+let root: Root
+
+beforeEach(() => {
+  ;(globalThis as { IS_REACT_ACT_ENVIRONMENT?: boolean }).IS_REACT_ACT_ENVIRONMENT = true
+  useSettingsDirtyStore.getState().reset()
+  mockGenerate.mockImplementation(
+    (_input: unknown, options: { onSuccess: (result: { license_key: string }) => void }) =>
+      options.onSuccess({ license_key: 'sim_license_once' })
+  )
+  container = document.createElement('div')
+  document.body.appendChild(container)
+  root = createRoot(container)
+  act(() =>
+    root.render(
+      <NuqsTestingAdapter searchParams='?tab=licenses'>
+        <Mothership />
+      </NuqsTestingAdapter>
+    )
+  )
+})
+
+afterEach(() => {
+  act(() => root.unmount())
+  container.remove()
+})
+
+function type(placeholder: string, value: string) {
+  const input = container.querySelector<HTMLInputElement>(`input[placeholder="${placeholder}"]`)
+  expect(input).not.toBeNull()
+  act(() => {
+    Object.getOwnPropertyDescriptor(window.HTMLInputElement.prototype, 'value')?.set?.call(
+      input,
+      value
+    )
+    input?.dispatchEvent(new Event('input', { bubbles: true }))
+  })
+}
+
+function generateKey() {
+  type('e.g. Acme Corp', 'Acme')
+  type('Signed order form or written approval', 'order-1')
+  const generate = Array.from(container.querySelectorAll('button')).find(
+    (button) => !button.disabled
+  )
+  act(() => generate?.click())
+}
+
+function licenseKey() {
+  return container.querySelector<HTMLInputElement>('[data-testid="license-key"]')?.value
+}
+
+describe('Mothership license generation', () => {
+  it('asks before leaving while the shown-once license key is on screen', () => {
+    generateKey()
+    const leave = vi.fn()
+
+    const left = useSettingsDirtyStore.getState().requestLeave(leave)
+
+    expect(licenseKey()).toBe('sim_license_once')
+    expect(left).toBe(false)
+    expect(leave).not.toHaveBeenCalled()
+  })
+
+  it('drops the license key when the admin confirms leaving', () => {
+    generateKey()
+    const leave = vi.fn()
+    useSettingsDirtyStore.getState().requestLeave(leave)
+
+    act(() => useSettingsDirtyStore.getState().confirmLeave())
+
+    expect(leave).toHaveBeenCalledOnce()
+    expect(licenseKey()).toBeUndefined()
+    expect(useSettingsDirtyStore.getState().isDirty).toBe(false)
+  })
+})

apps/sim/app/workspace/[workspaceId]/settings/components/mothership/mothership.tsx

@@ -365,9 +365,10 @@ function LicensesTab({ environment }: { environment: MothershipEnv }) {
     setNewName('')
     setNewExpiry('')
     setApprovalReference('')
+    setGeneratedKey(null)
   }, [])
   useSettingsUnsavedGuard({
-    isDirty: Boolean(newName.trim() || newExpiry || approvalReference.trim()),
+    isDirty: Boolean(generatedKey || newName.trim() || newExpiry || approvalReference.trim()),
     navigationBlocked: generateLicense.isPending,
     onDiscard: discardDraft,
   })
@@ -382,8 +383,8 @@ function LicensesTab({ environment }: { environment: MothershipEnv }) {
       },
       {
         onSuccess: (result) => {
-          setGeneratedKey(result.license_key)
           discardDraft()
+          setGeneratedKey(result.license_key)
         },
       }
     )

apps/sim/stores/index.test.ts

@@ -13,7 +13,7 @@ vi.mock('@/stores/reset-all-stores', () => {
   return { resetAllStores: mockResetAllStores }
 })
 
-import { leaseDesktopTool } from '@/app/workspace/[workspaceId]/home/hooks/desktop-tool-lifetimes'
+import { desktopToolTurn } from '@/app/workspace/[workspaceId]/home/hooks/desktop-tool-lifetimes'
 import { clearUserData, RECENT_IMPERSONATIONS_STORAGE_KEY } from '@/stores'
 
 expect(mockModuleLoaded).not.toHaveBeenCalled()
@@ -98,8 +98,8 @@ describe('clearUserData', () => {
   })
 
   it('cancels desktop tools still running for the signed-out identity', async () => {
-    const localRead = leaseDesktopTool('turn-before-sign-out')
-    const browserAction = leaseDesktopTool('other-turn-before-sign-out')
+    const localRead = desktopToolTurn('turn-before-sign-out').lease()
+    const browserAction = desktopToolTurn('other-turn-before-sign-out').lease()
     mockResetAllStores.mockImplementationOnce(() => {
       throw new Error('Chunk unavailable')
     })