Skip to content

Commit b7fa794

Browse files
authored
fix(mothership): cancel desktop tools a signed-out turn starts late, guard the shown-once license key (#8757)
* 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. * 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. * test(mothership): assert the license key guard through tab and environment navigation * test(mothership): mock Chip for the license-key guard tests on top of #8761 #8761 already guards the shown-once license key (generatedKey in isDirty, cleared on discard), so this branch keeps staging's settings page and its tests now assert that guard. The page now renders Chip, which the shared emcn mock does not provide.
1 parent d8cb659 commit b7fa794

5 files changed

Lines changed: 267 additions & 28 deletions

File tree

Lines changed: 38 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,15 @@
11
import { describe, expect, it } from 'vitest'
22
import {
3-
leaseDesktopTool,
3+
desktopToolSession,
44
stopAllDesktopTools,
55
stopDesktopTools,
66
} from '@/app/workspace/[workspaceId]/home/hooks/desktop-tool-lifetimes'
77

88
describe('desktop tool leases', () => {
99
it('cancels every running tool of the stopped turn and no other turn', () => {
10-
const first = leaseDesktopTool('turn-a')
11-
const second = leaseDesktopTool('turn-a')
12-
const other = leaseDesktopTool('turn-b')
10+
const first = desktopToolSession().turn('turn-a').lease()
11+
const second = desktopToolSession().turn('turn-a').lease()
12+
const other = desktopToolSession().turn('turn-b').lease()
1313

1414
stopDesktopTools('turn-a', 'user_stop')
1515

@@ -21,9 +21,10 @@ describe('desktop tool leases', () => {
2121
})
2222

2323
it('keeps a turn reachable by Stop while any of its tools still runs', () => {
24-
const settled = leaseDesktopTool('turn-c')
25-
const running = leaseDesktopTool('turn-c')
26-
for (let turn = 0; turn < 500; turn++) leaseDesktopTool(`busy-${turn}`).release()
24+
const settled = desktopToolSession().turn('turn-c').lease()
25+
const running = desktopToolSession().turn('turn-c').lease()
26+
for (let turn = 0; turn < 500; turn++)
27+
desktopToolSession().turn(`busy-${turn}`).lease().release()
2728

2829
settled.release()
2930
settled.release()
@@ -33,11 +34,11 @@ describe('desktop tool leases', () => {
3334
})
3435

3536
it('gives a turn whose tools all settled a fresh lifetime for its next tool', () => {
36-
const settled = leaseDesktopTool('turn-d')
37+
const settled = desktopToolSession().turn('turn-d').lease()
3738
settled.release()
3839
stopDesktopTools('turn-d', 'user_stop')
3940

40-
const next = leaseDesktopTool('turn-d')
41+
const next = desktopToolSession().turn('turn-d').lease()
4142

4243
expect(settled.signal.aborted).toBe(false)
4344
expect(next.signal).not.toBe(settled.signal)
@@ -46,9 +47,9 @@ describe('desktop tool leases', () => {
4647
})
4748

4849
it('does not let a tool that settles after Stop release a newer lease on the turn', () => {
49-
const stopped = leaseDesktopTool('turn-e')
50+
const stopped = desktopToolSession().turn('turn-e').lease()
5051
stopDesktopTools('turn-e', 'user_stop')
51-
const next = leaseDesktopTool('turn-e')
52+
const next = desktopToolSession().turn('turn-e').lease()
5253

5354
stopped.release()
5455
stopDesktopTools('turn-e', 'user_stop')
@@ -57,17 +58,39 @@ describe('desktop tool leases', () => {
5758
})
5859

5960
it('cancels the running tools of every turn when the session ends', () => {
60-
const first = leaseDesktopTool('turn-f')
61-
const second = leaseDesktopTool('turn-g')
61+
const first = desktopToolSession().turn('turn-f').lease()
62+
const second = desktopToolSession().turn('turn-g').lease()
6263

6364
stopAllDesktopTools('signed_out')
64-
const next = leaseDesktopTool('turn-f')
6565

6666
expect(first.signal.aborted).toBe(true)
6767
expect(second.signal.aborted).toBe(true)
6868
expect(first.signal.reason).toBe('signed_out')
69-
expect(next.signal.aborted).toBe(false)
7069
first.release()
70+
second.release()
71+
})
72+
73+
it('cancels tools of a surface mounted before sign-out, even on a stream it reads later', () => {
74+
const surface = desktopToolSession()
75+
const running = surface.turn('turn-h')
76+
stopAllDesktopTools('signed_out')
77+
78+
const late = running.lease()
79+
const reconnected = surface.turn('turn-j').lease()
80+
81+
expect(late.signal.aborted).toBe(true)
82+
expect(late.signal.reason).toBe('signed_out')
83+
expect(reconnected.signal.aborted).toBe(true)
84+
late.release()
85+
reconnected.release()
86+
})
87+
88+
it('runs the tools of a surface mounted after the session ended', () => {
89+
stopAllDesktopTools('signed_out')
90+
91+
const next = desktopToolSession().turn('turn-i').lease()
92+
93+
expect(next.signal.aborted).toBe(false)
7194
next.release()
7295
})
7396
})

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

Lines changed: 35 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,9 @@ interface RunningTurnTools {
1515
*/
1616
const runningTurns = new Map<string, RunningTurnTools>()
1717

18+
/** Aborted by `stopAllDesktopTools`, then replaced, so each signed-in session has its own. */
19+
let session = new AbortController()
20+
1821
/** A running desktop tool's hold on its turn. */
1922
interface DesktopToolLease {
2023
/** Aborted only by the user's Stop of the turn, or by signing out. */
@@ -23,13 +26,40 @@ interface DesktopToolLease {
2326
release(): void
2427
}
2528

29+
/** A turn's desktop tools, in the session of the chat surface that runs the turn. */
30+
export interface DesktopToolTurn {
31+
/** Starts one desktop tool for the turn. */
32+
lease(): DesktopToolLease
33+
}
34+
35+
/** The desktop tools a chat surface starts, bound to the session the surface mounted in. */
36+
export interface DesktopToolSession {
37+
turn(streamId: string): DesktopToolTurn
38+
}
39+
40+
/**
41+
* Binds a chat surface to the current session. Take it once, when the surface mounts: a send or
42+
* reconnect still in flight at sign-out can deliver tool events after the stop, and each of them
43+
* then gets an already-aborted lease. Signing out leaves or reloads every chat surface, so a
44+
* surface mounted after sign-in binds to the new session.
45+
*/
46+
export function desktopToolSession(): DesktopToolSession {
47+
const startedIn = session.signal
48+
return {
49+
turn: (streamId) => ({
50+
lease: () =>
51+
startedIn.aborted ? { signal: startedIn, release() {} } : leaseDesktopTool(streamId),
52+
}),
53+
}
54+
}
55+
2656
/**
2757
* Starts a desktop tool (a browser action, a local file read or import) for a turn. Only the
2858
* user's Stop of that turn, or signing out (`stopAllDesktopTools`), cancels it: replacing the
2959
* stream reader, leaving the chat view, or stopping another chat's turn leaves it running to
3060
* finish and report its own result.
3161
*/
32-
export function leaseDesktopTool(streamId: string): DesktopToolLease {
62+
function leaseDesktopTool(streamId: string): DesktopToolLease {
3363
let turn = runningTurns.get(streamId)
3464
if (!turn) {
3565
turn = { stop: new AbortController(), running: 0 }
@@ -57,9 +87,12 @@ export function stopDesktopTools(streamId: string, reason: string): void {
5787

5888
/**
5989
* Cancels every leased desktop tool running in this tab (browser actions, local file reads and
60-
* imports), so none outlives the session that started it.
90+
* imports), and every one a turn of this session starts later, so none outlives the session
91+
* that started it.
6192
*/
6293
export function stopAllDesktopTools(reason: string): void {
94+
session.abort(reason)
95+
session = new AbortController()
6396
for (const turn of runningTurns.values()) turn.stop.abort(reason)
6497
runningTurns.clear()
6598
}

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

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -93,7 +93,9 @@ import { initTerminalTransport } from '@/lib/terminal/transport'
9393
import { getQueryClient } from '@/app/_shell/providers/get-query-client'
9494
import { chatUrl } from '@/app/workspace/[workspaceId]/home/hooks/chat-url'
9595
import {
96-
leaseDesktopTool,
96+
type DesktopToolSession,
97+
type DesktopToolTurn,
98+
desktopToolSession,
9799
stopDesktopTools,
98100
} from '@/app/workspace/[workspaceId]/home/hooks/desktop-tool-lifetimes'
99101
import { useFilePreviewController } from '@/app/workspace/[workspaceId]/home/hooks/preview'
@@ -518,10 +520,10 @@ function startClientBrowserTool(
518520
toolArgs: Record<string, unknown>,
519521
scopeId: string,
520522
eventTs?: string,
521-
turnStreamId?: string
523+
desktopTurn?: DesktopToolTurn
522524
): void {
523525
if (!isCurrentBrowserToolName(toolName)) return
524-
const lease = turnStreamId ? leaseDesktopTool(turnStreamId) : undefined
526+
const lease = desktopTurn?.lease()
525527
void executeBrowserToolOnClient(
526528
toolCallId,
527529
toolName,
@@ -1003,6 +1005,8 @@ export function useChat(
10031005
const chatIdRef = useRef<string | undefined>(initialChatId)
10041006
/** Cleared on unmount, so late async work cannot act on a surface the user left. */
10051007
const surfaceMountedRef = useRef(true)
1008+
const desktopToolsRef = useRef<DesktopToolSession | null>(null)
1009+
const desktopTools = (desktopToolsRef.current ??= desktopToolSession())
10061010
useEffect(() => {
10071011
surfaceMountedRef.current = true
10081012
return () => {
@@ -1654,7 +1658,7 @@ export function useChat(
16541658
toolCallId: string,
16551659
toolName: string,
16561660
toolArgs: Record<string, unknown>,
1657-
turnStreamId: string | undefined
1661+
desktopTurn: DesktopToolTurn | undefined
16581662
) => {
16591663
if (
16601664
!isNativeFileTool(toolName) &&
@@ -1666,7 +1670,7 @@ export function useChat(
16661670
return
16671671
}
16681672
handledClientLocalFilesystemToolIds.add(toolCallId)
1669-
const lease = turnStreamId ? leaseDesktopTool(turnStreamId) : undefined
1673+
const lease = desktopTurn?.lease()
16701674
const options = {
16711675
workspaceId,
16721676
chatId: chatIdRef.current ?? selectedChatIdRef.current,
@@ -2297,7 +2301,7 @@ export function useChat(
22972301
shouldContinue?: () => boolean
22982302
}
22992303
) => {
2300-
const turnStreamId = streamIdRef.current
2304+
const desktopTurn = streamIdRef.current ? desktopTools.turn(streamIdRef.current) : undefined
23012305
const activityTracker = getResourceActivityTracker(
23022306
expectedGen ?? streamGenRef.current,
23032307
options?.targetChatId
@@ -2318,7 +2322,7 @@ export function useChat(
23182322
eventTs?: string
23192323
) => {
23202324
const scopeId = activityScopeId()
2321-
startClientBrowserTool(toolCallId, toolName, toolArgs, scopeId, eventTs, turnStreamId)
2325+
startClientBrowserTool(toolCallId, toolName, toolArgs, scopeId, eventTs, desktopTurn)
23222326
}
23232327
const startClientTerminalToolForStream = (
23242328
toolCallId: string,
@@ -2351,7 +2355,7 @@ export function useChat(
23512355
removeResource,
23522356
startClientWorkflowTool,
23532357
startClientLocalFilesystemTool: (toolCallId, toolName, toolArgs) =>
2354-
startClientLocalFilesystemTool(toolCallId, toolName, toolArgs, turnStreamId),
2358+
startClientLocalFilesystemTool(toolCallId, toolName, toolArgs, desktopTurn),
23552359
startClientBrowserTool: startClientBrowserToolForStream,
23562360
startClientTerminalTool: startClientTerminalToolForStream,
23572361
startBrowserAgentRun: startBrowserAgentRunForStream,

0 commit comments

Comments
 (0)