fix(mothership): bind desktop tools to the session the chat surface mounted in

A send or reconnect still in flight at sign-out reaches the stream reader after the stop, so a per-reader capture took the new session. The surface now takes its session once at mount; signing out leaves or reloads every chat surface.

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 {
-  desktopToolTurn,
+  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 = desktopToolTurn('turn-a').lease()
-    const second = desktopToolTurn('turn-a').lease()
-    const other = desktopToolTurn('turn-b').lease()
+    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 = desktopToolTurn('turn-c').lease()
-    const running = desktopToolTurn('turn-c').lease()
-    for (let turn = 0; turn < 500; turn++) desktopToolTurn(`busy-${turn}`).lease().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 = desktopToolTurn('turn-d').lease()
+    const settled = desktopToolSession().turn('turn-d').lease()
     settled.release()
     stopDesktopTools('turn-d', 'user_stop')
 
-    const next = desktopToolTurn('turn-d').lease()
+    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 = desktopToolTurn('turn-e').lease()
+    const stopped = desktopToolSession().turn('turn-e').lease()
     stopDesktopTools('turn-e', 'user_stop')
-    const next = desktopToolTurn('turn-e').lease()
+    const next = desktopToolSession().turn('turn-e').lease()
 
     stopped.release()
     stopDesktopTools('turn-e', 'user_stop')
@@ -57,8 +58,8 @@ describe('desktop tool leases', () => {
   })
 
   it('cancels the running tools of every turn when the session ends', () => {
-    const first = desktopToolTurn('turn-f').lease()
-    const second = desktopToolTurn('turn-g').lease()
+    const first = desktopToolSession().turn('turn-f').lease()
+    const second = desktopToolSession().turn('turn-g').lease()
 
     stopAllDesktopTools('signed_out')
 
@@ -69,21 +70,25 @@ describe('desktop tool leases', () => {
     second.release()
   })
 
-  it('cancels a tool a turn of the ended session starts after sign-out', () => {
-    const turn = desktopToolTurn('turn-h')
+  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 = turn.lease()
+    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 turn started after the session ended', () => {
+  it('runs the tools of a surface mounted after the session ended', () => {
     stopAllDesktopTools('signed_out')
 
-    const next = desktopToolTurn('turn-i').lease()
+    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

@@ -26,21 +26,30 @@ interface DesktopToolLease {
   release(): void
 }
 
-/** A turn's stream, bound to the session it started in. */
+/** 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 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.
+ * 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 desktopToolTurn(streamId: string): DesktopToolTurn {
+export function desktopToolSession(): DesktopToolSession {
   const startedIn = session.signal
   return {
-    lease: () =>
-      startedIn.aborted ? { signal: startedIn, release() {} } : leaseDesktopTool(streamId),
+    turn: (streamId) => ({
+      lease: () =>
+        startedIn.aborted ? { signal: startedIn, release() {} } : leaseDesktopTool(streamId),
+    }),
   }
 }
 

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

@@ -93,8 +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 {
+  type DesktopToolSession,
   type DesktopToolTurn,
-  desktopToolTurn,
+  desktopToolSession,
   stopDesktopTools,
 } from '@/app/workspace/[workspaceId]/home/hooks/desktop-tool-lifetimes'
 import { useFilePreviewController } from '@/app/workspace/[workspaceId]/home/hooks/preview'
@@ -1004,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 () => {
@@ -2298,7 +2301,7 @@ export function useChat(
         shouldContinue?: () => boolean
       }
     ) => {
-      const desktopTurn = streamIdRef.current ? desktopToolTurn(streamIdRef.current) : undefined
+      const desktopTurn = streamIdRef.current ? desktopTools.turn(streamIdRef.current) : undefined
       const activityTracker = getResourceActivityTracker(
         expectedGen ?? streamGenRef.current,
         options?.targetChatId

apps/sim/stores/index.test.ts

@@ -13,7 +13,7 @@ vi.mock('@/stores/reset-all-stores', () => {
   return { resetAllStores: mockResetAllStores }
 })
 
-import { desktopToolTurn } from '@/app/workspace/[workspaceId]/home/hooks/desktop-tool-lifetimes'
+import { desktopToolSession } 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 = desktopToolTurn('turn-before-sign-out').lease()
-    const browserAction = desktopToolTurn('other-turn-before-sign-out').lease()
+    const localRead = desktopToolSession().turn('turn-before-sign-out').lease()
+    const browserAction = desktopToolSession().turn('other-turn-before-sign-out').lease()
     mockResetAllStores.mockImplementationOnce(() => {
       throw new Error('Chunk unavailable')
     })