diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/desktop-tool-lifetimes.test.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/desktop-tool-lifetimes.test.ts index 9cf43b508c9..917314c788e 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/desktop-tool-lifetimes.test.ts +++ b/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() }) }) diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/desktop-tool-lifetimes.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/desktop-tool-lifetimes.ts index 9e2f49aeee1..43e8cde05bf 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/desktop-tool-lifetimes.ts +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/desktop-tool-lifetimes.ts @@ -15,6 +15,9 @@ interface RunningTurnTools { */ const runningTurns = new Map() +/** 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() } diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts index 627272e31ae..8114ddee989 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts +++ b/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, 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(initialChatId) /** Cleared on unmount, so late async work cannot act on a surface the user left. */ const surfaceMountedRef = useRef(true) + const desktopToolsRef = useRef(null) + const desktopTools = (desktopToolsRef.current ??= desktopToolSession()) useEffect(() => { surfaceMountedRef.current = true return () => { @@ -1654,7 +1658,7 @@ export function useChat( toolCallId: string, toolName: string, toolArgs: Record, - 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, diff --git a/apps/sim/app/workspace/[workspaceId]/settings/components/mothership/mothership.test.tsx b/apps/sim/app/workspace/[workspaceId]/settings/components/mothership/mothership.test.tsx new file mode 100644 index 00000000000..d027a0cd1eb --- /dev/null +++ b/apps/sim/app/workspace/[workspaceId]/settings/components/mothership/mothership.test.tsx @@ -0,0 +1,179 @@ +/** + * @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 }) => {children}, + Chip: ({ children }: { children?: ReactNode }) => {children}, + Button: ({ children, ...props }: { children?: ReactNode }) => ( + + ), + ChipCopyInput: ({ value }: { value?: string }) => ( + + ), + ChipInput: ({ + value, + onChange, + placeholder, + }: { + value?: string + onChange?: ChangeEventHandler + placeholder?: string + }) => , + ChipModalTabs: ({ + tabs, + onChange, + }: { + tabs: Array<{ label: string; value: string }> + onChange: (value: string) => void + }) => ( +
+ {tabs.map((tab) => ( + + ))} +
+ ), + ChipSelect: ({ + options, + onChange, + }: { + options: Array<{ label: string; value: string }> + onChange: (value: string) => void + }) => ( +
+ {options.map((option) => ( + + ))} +
+ ), + Label: ({ children }: { children?: ReactNode }) => {children}, + Skeleton: () =>
, +})) + +vi.mock('@/app/workspace/[workspaceId]/settings/components/settings-panel', () => ({ + SettingsPanel: ({ children }: { children?: ReactNode }) =>
{children}
, +})) + +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 +let urlSearch: string + +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) + urlSearch = '?tab=licenses' + act(() => + root.render( + { + urlSearch = event.queryString + }} + > + + + ) + ) +}) + +afterEach(() => { + act(() => root.unmount()) + container.remove() +}) + +function type(placeholder: string, value: string) { + const input = container.querySelector(`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.dataset.tab && !button.dataset.env && !button.disabled + ) + act(() => generate?.click()) +} + +function licenseKey() { + return container.querySelector('[data-testid="license-key"]')?.value +} + +function switchToByok() { + act(() => container.querySelector('[data-tab="byok"]')?.click()) +} + +describe('Mothership license generation', () => { + it('stays on the tab while the shown-once license key is on screen', () => { + generateKey() + + switchToByok() + + expect(licenseKey()).toBe('sim_license_once') + expect(urlSearch).not.toContain('byok') + expect(useSettingsDirtyStore.getState().pendingLeave).not.toBeNull() + }) + + it('drops the license key when the admin confirms switching environments', async () => { + generateKey() + act(() => container.querySelector('[data-env="prod"]')?.click()) + + await act(async () => useSettingsDirtyStore.getState().confirmLeave()) + + await vi.waitFor(() => expect(urlSearch).toContain('env=prod')) + expect(licenseKey()).toBeUndefined() + expect(useSettingsDirtyStore.getState().isDirty).toBe(false) + }) +}) diff --git a/apps/sim/stores/index.test.ts b/apps/sim/stores/index.test.ts index bf4be01dc6e..b7f95fd28ae 100644 --- a/apps/sim/stores/index.test.ts +++ b/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 { desktopToolSession } from '@/app/workspace/[workspaceId]/home/hooks/desktop-tool-lifetimes' import { clearUserData, RECENT_IMPERSONATIONS_STORAGE_KEY } from '@/stores' import { useSettingsDirtyStore } from '@/stores/settings/dirty/store' @@ -113,8 +113,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 = desktopToolSession().turn('turn-before-sign-out').lease() + const browserAction = desktopToolSession().turn('other-turn-before-sign-out').lease() mockResetAllStores.mockImplementationOnce(() => { throw new Error('Chunk unavailable') })