diff --git a/src/renderer/src/store/slices/worktree-removal-maps-leak.test.ts b/src/renderer/src/store/slices/worktree-removal-maps-leak.test.ts new file mode 100644 index 000000000..d765da09e --- /dev/null +++ b/src/renderer/src/store/slices/worktree-removal-maps-leak.test.ts @@ -0,0 +1,202 @@ +/** + * Memory-leak regression: per-worktree and per-page maps that the worktree-removal + * paths previously missed must be evicted on removal. + * + * Two families of entity-keyed Records grew monotonically over a session: + * + * - Worktree-keyed: `remoteStatusesByWorktree`, `recentlyClosedEditorTabsByWorktree` + * and `defaultTerminalTabsAppliedByWorktreeId` were re-keyed on rename but absent + * from BOTH removal paths (the bulk `buildWorktreePurgeState` and the single + * `removeWorktree` reducer), while their siblings (gitStatusByWorktree, + * recentlyClosedBrowserTabsByWorktree, …) were purged. + * - Page/workspace-keyed browser maps: `browserAnnotationsByPageId`, + * `remoteBrowserPageHandlesByPageId`, `pendingAddressBarFocusByPageId`, + * `pendingAddressBarFocusByTabId` and `recentlyClosedBrowserPagesByWorkspace` + * were cleaned only on the single-worktree path (via shutdownWorktreeBrowsers → + * closeBrowserTab); the bulk reconcile path (CLI/SSH/other-window git worktree + * removal) skipped them, orphaning entries permanently. + * + * worktreeId, browser workspace id and browser page id are all unbounded, ephemeral + * key spaces (fresh UUIDs / path-derived ids, never reused). + */ +import { describe, it, expect, vi, beforeEach } from 'vitest' +import type * as AgentStatusModule from '@/lib/agent-status' +import type { BrowserPage, BrowserWorkspace } from '../../../../shared/types' + +vi.mock('sonner', () => ({ + toast: { info: vi.fn(), success: vi.fn(), error: vi.fn(), warning: vi.fn() } +})) + +vi.mock('@/components/terminal-pane/pty-dispatcher', () => ({ + restorePtyDataHandlersAfterFailedShutdown: vi.fn(), + unregisterPtyDataHandlers: vi.fn() +})) + +vi.mock('@/lib/agent-status', async (importOriginal) => { + const actual = await importOriginal() + return { ...actual, detectAgentStatusFromTitle: vi.fn().mockReturnValue(null) } +}) + +const mockApi = { + worktrees: { + list: vi.fn().mockResolvedValue([]), + remove: vi.fn().mockResolvedValue(undefined), + forceDeletePreservedBranch: vi.fn().mockResolvedValue({ deleted: true }), + updateMeta: vi.fn().mockResolvedValue({}) + }, + pty: { kill: vi.fn().mockResolvedValue(undefined) }, + runtimeEnvironments: { call: vi.fn().mockResolvedValue({ ok: true, result: {} }) } +} + +// @ts-expect-error -- minimal window.api stub for the store under test +globalThis.window = { api: mockApi } + +import { createTestStore, seedStore, makeWorktree, makeOpenFile } from './store-test-helpers' + +const WT1 = 'repo1::/path/wt1' +const WT2 = 'repo1::/path/wt2' + +function makeWorkspace(id: string, worktreeId: string): BrowserWorkspace { + return { + id, + worktreeId, + url: 'about:blank', + title: '', + loading: false, + faviconUrl: null, + canGoBack: false, + canGoForward: false, + loadError: null, + createdAt: 0 + } +} + +function makePage(id: string, workspaceId: string, worktreeId: string): BrowserPage { + return { + id, + workspaceId, + worktreeId, + url: 'about:blank', + title: '', + loading: false, + faviconUrl: null, + canGoBack: false, + canGoForward: false, + loadError: null, + createdAt: 0 + } +} + +describe('worktree removal evicts the per-worktree + per-page maps it previously missed', () => { + beforeEach(() => { + vi.clearAllMocks() + mockApi.worktrees.remove.mockResolvedValue(undefined) + }) + + function seedWorktreeKeyedMaps(store: ReturnType): void { + seedStore(store, { + worktreesByRepo: { + repo1: [ + makeWorktree({ id: WT1, repoId: 'repo1', path: '/path/wt1' }), + makeWorktree({ id: WT2, repoId: 'repo1', path: '/path/wt2' }) + ] + }, + remoteStatusesByWorktree: { + [WT1]: { hasUpstream: true, ahead: 1, behind: 0 }, + [WT2]: { hasUpstream: true, ahead: 0, behind: 2 } + }, + recentlyClosedEditorTabsByWorktree: { + [WT1]: [makeOpenFile({ id: '/path/wt1/closed.ts', worktreeId: WT1 })], + [WT2]: [makeOpenFile({ id: '/path/wt2/closed.ts', worktreeId: WT2 })] + }, + defaultTerminalTabsAppliedByWorktreeId: { [WT1]: true, [WT2]: true } + }) + } + + it('bulk purgeWorktreeTerminalState drops worktree-keyed maps for the removed worktree only', () => { + const store = createTestStore() + seedWorktreeKeyedMaps(store) + + store.getState().purgeWorktreeTerminalState([WT1]) + + const s = store.getState() + // Evicted for the removed worktree. + expect(s.remoteStatusesByWorktree[WT1]).toBeUndefined() + expect(s.recentlyClosedEditorTabsByWorktree[WT1]).toBeUndefined() + expect(s.defaultTerminalTabsAppliedByWorktreeId[WT1]).toBeUndefined() + // Retained for the surviving worktree (guard over-eviction). + expect(s.remoteStatusesByWorktree[WT2]).toBeDefined() + expect(s.recentlyClosedEditorTabsByWorktree[WT2]).toBeDefined() + expect(s.defaultTerminalTabsAppliedByWorktreeId[WT2]).toBe(true) + }) + + it('single removeWorktree drops worktree-keyed maps for the removed worktree only', async () => { + const store = createTestStore() + seedWorktreeKeyedMaps(store) + + const result = await store.getState().removeWorktree(WT1) + expect(result).toEqual({ ok: true }) + + const s = store.getState() + expect(s.remoteStatusesByWorktree[WT1]).toBeUndefined() + expect(s.recentlyClosedEditorTabsByWorktree[WT1]).toBeUndefined() + expect(s.defaultTerminalTabsAppliedByWorktreeId[WT1]).toBeUndefined() + expect(s.remoteStatusesByWorktree[WT2]).toBeDefined() + expect(s.recentlyClosedEditorTabsByWorktree[WT2]).toBeDefined() + expect(s.defaultTerminalTabsAppliedByWorktreeId[WT2]).toBe(true) + }) + + it('bulk purgeWorktreeTerminalState drops page/workspace-keyed browser maps for the removed worktree only', () => { + const store = createTestStore() + const WS1 = 'ws-1' + const WS2 = 'ws-2' + const P1 = 'page-1' + const P2 = 'page-2' + seedStore(store, { + worktreesByRepo: { + repo1: [ + makeWorktree({ id: WT1, repoId: 'repo1', path: '/path/wt1' }), + makeWorktree({ id: WT2, repoId: 'repo1', path: '/path/wt2' }) + ] + }, + browserTabsByWorktree: { + [WT1]: [makeWorkspace(WS1, WT1)], + [WT2]: [makeWorkspace(WS2, WT2)] + }, + browserPagesByWorkspace: { + [WS1]: [makePage(P1, WS1, WT1)], + [WS2]: [makePage(P2, WS2, WT2)] + }, + browserAnnotationsByPageId: { [P1]: [], [P2]: [] }, + remoteBrowserPageHandlesByPageId: { + [P1]: { environmentId: 'env-1', remotePageId: 'r-1' }, + [P2]: { environmentId: 'env-2', remotePageId: 'r-2' } + }, + pendingAddressBarFocusByPageId: { [P1]: true, [P2]: true }, + // createBrowserTab writes BOTH the workspace id and the page id here. + pendingAddressBarFocusByTabId: { [WS1]: true, [P1]: true, [WS2]: true, [P2]: true }, + recentlyClosedBrowserPagesByWorkspace: { + [WS1]: [makePage(P1, WS1, WT1)], + [WS2]: [makePage(P2, WS2, WT2)] + } + }) + + store.getState().purgeWorktreeTerminalState([WT1]) + + const s = store.getState() + // Removed worktree's workspace + page entries are gone. + expect(s.browserAnnotationsByPageId[P1]).toBeUndefined() + expect(s.remoteBrowserPageHandlesByPageId[P1]).toBeUndefined() + expect(s.pendingAddressBarFocusByPageId[P1]).toBeUndefined() + expect(s.pendingAddressBarFocusByTabId[WS1]).toBeUndefined() + expect(s.pendingAddressBarFocusByTabId[P1]).toBeUndefined() + expect(s.recentlyClosedBrowserPagesByWorkspace[WS1]).toBeUndefined() + // Surviving worktree's entries remain (guard over-eviction). + expect(s.browserAnnotationsByPageId[P2]).toBeDefined() + expect(s.remoteBrowserPageHandlesByPageId[P2]).toBeDefined() + expect(s.pendingAddressBarFocusByPageId[P2]).toBe(true) + expect(s.pendingAddressBarFocusByTabId[WS2]).toBe(true) + expect(s.pendingAddressBarFocusByTabId[P2]).toBe(true) + expect(s.recentlyClosedBrowserPagesByWorkspace[WS2]).toBeDefined() + }) +}) diff --git a/src/renderer/src/store/slices/worktrees.ts b/src/renderer/src/store/slices/worktrees.ts index e1dbca326..2486ad963 100644 --- a/src/renderer/src/store/slices/worktrees.ts +++ b/src/renderer/src/store/slices/worktrees.ts @@ -936,8 +936,9 @@ function buildWorktreeRenameState( for (const key of WORKTREE_ID_KEYED_MAP_KEYS) { renameKey(key, renameValueByKey[key] as ((value: unknown) => unknown) | undefined) } - // Not in the shared purge list (purge drops them only via single removeWorktree); - // re-key them here so a renamed worktree keeps its editor-undo + push/pull state. + // Re-key these on rename so a renamed worktree keeps its editor-undo + push/pull + // state. (Both removal paths — buildWorktreePurgeState and the single + // removeWorktree reducer — now also purge them on removal.) renameKey('recentlyClosedEditorTabsByWorktree', (files: { worktreeId: string }[]) => files.map(withNewWorktreeId) ) @@ -1024,6 +1025,7 @@ function buildWorktreePurgeState(s: AppState, worktreeIds: string[]): Partial() const doomedBrowserWorkspaceIds = new Set() + const doomedPageIds = new Set() const removedFileIds = new Set() for (const id of worktreeIdSet) { for (const tab of s.tabsByWorktree[id] ?? []) { @@ -1033,6 +1035,15 @@ function buildWorktreePurgeState(s: AppState, worktreeIds: string[]): Partial(obj: Record): Record => { + let changed = false + const out = { ...obj } + for (const pageId of doomedPageIds) { + if (pageId in out) { + delete out[pageId] + changed = true + } + } + return changed ? out : obj + } const omitByFileId = (obj: Record): Record => { let changed = false const out = { ...obj } @@ -1166,6 +1188,20 @@ function buildWorktreePurgeState(s: AppState, worktreeIds: string[]): Partial { if (s.activeWorkspaceKey && worktreeIdSet.has(s.activeWorkspaceKey)) { @@ -2085,6 +2132,23 @@ export const createWorktreeSlice: StateCreator delete nextSearch[worktreeId] return nextSearch })(), + // Why: these worktree-keyed maps are re-keyed on rename but were missed + // by both removal paths, leaking one entry per removed worktree. + remoteStatusesByWorktree: (() => { + const next = { ...s.remoteStatusesByWorktree } + delete next[worktreeId] + return next + })(), + recentlyClosedEditorTabsByWorktree: (() => { + const next = { ...s.recentlyClosedEditorTabsByWorktree } + delete next[worktreeId] + return next + })(), + defaultTerminalTabsAppliedByWorktreeId: (() => { + const next = { ...s.defaultTerminalTabsAppliedByWorktreeId } + delete next[worktreeId] + return next + })(), activeWorktreeId: removedActiveWorktree ? null : s.activeWorktreeId, activeTabId: s.activeTabId && tabIds.has(s.activeTabId) ? null : s.activeTabId, openFiles: newOpenFiles,