diff --git a/src/renderer/src/store/folder-workspaces/folder-workspace-mutations.ts b/src/renderer/src/store/folder-workspaces/folder-workspace-mutations.ts index b9eefe424a8..41b63ec5bf8 100644 --- a/src/renderer/src/store/folder-workspaces/folder-workspace-mutations.ts +++ b/src/renderer/src/store/folder-workspaces/folder-workspace-mutations.ts @@ -28,6 +28,10 @@ import { getFolderWorkspaceUpdateIdentity, reconcileFailedFolderWorkspaceUpdate } from './folder-workspace-catalog' +import { + captureWorkspaceChatDraftKeys, + deleteWorkspaceChatDrafts +} from '../slices/worktrees/teardown/removed-worktree-chat-drafts' export type FolderWorkspaceUpdateField = keyof FolderWorkspaceUpdates @@ -230,6 +234,9 @@ export function createFolderWorkspaceMutationActions( folderWorkspaceId, executionHostId ) + const workspaceKey = folderWorkspaceKey(folderWorkspaceId) + // Why before the host call: its announcement can start a refresh that drops these tabs. + const chatDraftKeys = captureWorkspaceChatDraftKeys(state, [workspaceKey]) try { // Why: deletion targets the folder's owner; focus may be on a different host. const target = getActiveRuntimeTarget({ activeRuntimeEnvironmentId: runtimeEnvironmentId }) @@ -247,7 +254,6 @@ export function createFolderWorkspaceMutationActions( if (!deleted) { return false } - const workspaceKey = folderWorkspaceKey(folderWorkspaceId) set((s) => ({ folderWorkspaces: s.folderWorkspaces.filter( (workspace) => @@ -261,6 +267,7 @@ export function createFolderWorkspaceMutationActions( // tear down Chromium guests before purging the remaining renderer state. await get().shutdownWorktreeBrowsers(workspaceKey) get().purgeWorktreeTerminalState([workspaceKey]) + deleteWorkspaceChatDrafts(chatDraftKeys) } return true } catch (err) { diff --git a/src/renderer/src/store/repos/repo-removal.ts b/src/renderer/src/store/repos/repo-removal.ts index 43525c3bba8..9a1f7b9e302 100644 --- a/src/renderer/src/store/repos/repo-removal.ts +++ b/src/renderer/src/store/repos/repo-removal.ts @@ -22,6 +22,10 @@ import type { RepoSlice } from './repo-state' import { ERROR_TOAST_DURATION } from './repo-state' import { mergeProjectCompatibilityForHostRepoChange } from './repo-catalog-identity' import { settingsForRepoOwner } from './owner-routing' +import { + captureWorkspaceChatDraftKeys, + deleteWorkspaceChatDrafts +} from '../slices/worktrees/teardown/removed-worktree-chat-drafts' export function worktreeBelongsToHost(worktree: { hostId?: string }, hostId: string): boolean { return (worktree.hostId ?? LOCAL_EXECUTION_HOST_ID) === hostId @@ -83,6 +87,11 @@ export function createRepoRemovalActions( const idExistsOnOtherHost = get().repos.some( (repo) => repo.id === projectId && getRepoExecutionHostId(repo) !== ownerHostId ) + // Why before the host call: its announcement can start a listing refresh that drops these tabs. + const chatDraftKeys = captureWorkspaceChatDraftKeys( + get(), + getKnownRepoWorktreeIds(get(), projectId, ownerHostId) + ) try { await (target.kind === 'local' ? idExistsOnOtherHost @@ -156,6 +165,7 @@ export function createRepoRemovalActions( // Why: use the canonical per-worktree purge to evict all worktree-scoped maps (hand-deletion leaked most); runs before the set() below so it still sees tabsByWorktree. get().purgeWorktreeTerminalState(purgeTargets) + deleteWorkspaceChatDrafts(chatDraftKeys) get().clearLocalDetectedAgentContextsForProjects(localAgentContextProjectIds) set((s) => { diff --git a/src/renderer/src/store/slices/workspace-chat-draft-removal.test.ts b/src/renderer/src/store/slices/workspace-chat-draft-removal.test.ts new file mode 100644 index 00000000000..6d56cded9cc --- /dev/null +++ b/src/renderer/src/store/slices/workspace-chat-draft-removal.test.ts @@ -0,0 +1,233 @@ +/** + * A workspace's unsent chat drafts die only with a user's delete, by keys read before the host + * round trip. The host announces a removal before it replies, and a listing purge that lands first + * must not hide them; a purge that is not a user's delete must not erase them. + */ +import { describe, it, expect, vi, beforeEach } from 'vitest' +import type * as AgentStatusModule from '@/lib/agent-status' +import type { FolderWorkspace } from '../../../../shared/folder-workspace-types' +import type { PublicKnownRuntimeEnvironment } from '../../../../shared/runtime-environments' +import { toRuntimeExecutionHostId } from '../../../../shared/execution-host' +import { folderWorkspaceKey } from '../../../../shared/workspace-scope' + +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(), remove: vi.fn(), updateMeta: vi.fn().mockResolvedValue({}) }, + repos: { remove: vi.fn() }, + folderWorkspaces: { delete: vi.fn() }, + pty: { kill: vi.fn() }, + 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, + makeTab, + makeTabGroup, + makeUnifiedTab, + TEST_REPO +} from './store-test-helpers' +import { + clearNativeChatComposerDraftsForTests, + readNativeChatComposerDraft, + structuredAgentSessionDraftScopeKey, + updateNativeChatComposerDraft +} from '@/components/native-chat/native-chat-composer-draft-store' + +const WT1 = 'repo1::/path/wt1' +const WT2 = 'repo1::/path/wt2' +const CHAT = structuredAgentSessionDraftScopeKey('session-1') +const PANE = 'tab-wt1:leaf-a' +const OTHER = 'tab-wt2:leaf-a' + +function seedWorkspace(store: ReturnType, workspaceId: string): void { + const chat = makeUnifiedTab({ + id: 'chat-tab', + entityId: 'session-1', + contentType: 'agent-session', + worktreeId: workspaceId, + groupId: 'group-1' + }) + seedStore(store, { + tabsByWorktree: { + [workspaceId]: [makeTab({ id: 'tab-wt1', worktreeId: workspaceId })], + [WT2]: [makeTab({ id: 'tab-wt2', worktreeId: WT2 })] + }, + unifiedTabsByWorktree: { [workspaceId]: [chat] }, + groupsByWorktree: { + [workspaceId]: [ + makeTabGroup({ + id: 'group-1', + worktreeId: workspaceId, + activeTabId: chat.id, + tabOrder: [chat.id] + }) + ] + } + }) + updateNativeChatComposerDraft(CHAT, { text: 'chat' }, 'immediate') + updateNativeChatComposerDraft(PANE, { text: 'pane' }, 'immediate') + updateNativeChatComposerDraft(OTHER, { text: 'other' }, 'immediate') +} + +function seedWorktree(store: ReturnType): void { + seedStore(store, { + worktreesByRepo: { + repo1: [ + makeWorktree({ id: WT1, repoId: 'repo1', path: '/path/wt1' }), + makeWorktree({ id: WT2, repoId: 'repo1', path: '/path/wt2' }) + ] + } + }) + seedWorkspace(store, WT1) +} + +function draftTexts(): string[] { + return [CHAT, PANE, OTHER].map((key) => readNativeChatComposerDraft(key).text) +} + +describe('workspace chat drafts on removal', () => { + beforeEach(() => { + vi.clearAllMocks() + mockApi.worktrees.remove.mockResolvedValue(undefined) + mockApi.repos.remove.mockResolvedValue(undefined) + clearNativeChatComposerDraftsForTests() + }) + + it('removing a worktree deletes its drafts even when a listing purge lands mid-removal', async () => { + const store = createTestStore() + seedWorktree(store) + mockApi.worktrees.remove.mockImplementation(async () => { + store.getState().purgeWorktreeTerminalState([WT1]) + }) + + const result = await store.getState().removeWorktree({ id: WT1, executionHostId: null }, true) + + expect(result).toEqual({ ok: true }) + expect(draftTexts()).toEqual(['', '', 'other']) + }) + + it('a removal the host refuses keeps the drafts', async () => { + vi.spyOn(console, 'warn').mockImplementation(() => {}) + const store = createTestStore() + seedWorktree(store) + mockApi.worktrees.remove.mockRejectedValue(new Error('Permission denied')) + + const result = await store.getState().removeWorktree({ id: WT1, executionHostId: null }, true) + + expect(result.ok).toBe(false) + expect(draftTexts()).toEqual(['chat', 'pane', 'other']) + }) + + it('a listing purge alone keeps the drafts', () => { + const store = createTestStore() + seedWorktree(store) + + store.getState().purgeWorktreeTerminalState([WT1]) + + expect(store.getState().unifiedTabsByWorktree[WT1]).toBeUndefined() + expect(draftTexts()).toEqual(['chat', 'pane', 'other']) + }) + + it('re-pairing a server under the same id keeps the drafts of chats it restores', () => { + const runtime = toRuntimeExecutionHostId('env-a') + const environment = (pairingRevision: number): PublicKnownRuntimeEnvironment => ({ + id: 'env-a', + name: 'Server', + createdAt: 1, + updatedAt: 1, + pairingRevision, + lastUsedAt: null, + runtimeId: null, + endpoints: [{ id: 'e', kind: 'websocket', label: 'Server', endpoint: 'ws://server' }], + preferredEndpointId: 'e' + }) + const store = createTestStore() + seedStore(store, { + runtimeEnvironments: [environment(1)], + repos: [TEST_REPO, { ...TEST_REPO, id: 'repoA', path: '/repoA', executionHostId: runtime }], + worktreesByRepo: { + repoA: [makeWorktree({ id: WT1, repoId: 'repoA', path: '/path/wt1', hostId: runtime })] + } + }) + seedWorkspace(store, WT1) + + store.getState().setRuntimeEnvironments([environment(2)]) + + expect(store.getState().unifiedTabsByWorktree[WT1]).toBeUndefined() + expect(draftTexts()).toEqual(['chat', 'pane', 'other']) + }) + + it('removing a project deletes its worktrees’ drafts only', async () => { + const store = createTestStore() + seedStore(store, { repos: [{ ...TEST_REPO, id: 'repo1' }] }) + seedWorktree(store) + seedStore(store, { + worktreesByRepo: { + repo1: [makeWorktree({ id: WT1, repoId: 'repo1', path: '/path/wt1' })], + repo2: [makeWorktree({ id: WT2, repoId: 'repo2', path: '/path/wt2' })] + } + }) + + await store.getState().removeProject('repo1') + + expect(draftTexts()).toEqual(['', '', 'other']) + }) + + it('deleting a folder workspace deletes its drafts', async () => { + const folder: FolderWorkspace = { + id: 'folder-1', + projectGroupId: 'group-a', + name: 'Folder', + folderPath: '/workspace/folder', + linkedTask: null, + comment: '', + isArchived: false, + isUnread: false, + isPinned: false, + sortOrder: 1, + lastActivityAt: 0, + createdAt: 1, + updatedAt: 1 + } + mockApi.folderWorkspaces.delete.mockResolvedValue(true) + const store = createTestStore() + store.setState({ + projectGroups: [ + { + id: 'group-a', + name: 'A', + parentPath: '/workspace', + parentGroupId: null, + createdFrom: 'manual', + tabOrder: 0, + isCollapsed: false, + color: null, + createdAt: 1, + updatedAt: 1, + executionHostId: 'local' + } + ], + folderWorkspaces: [folder] + }) + seedWorkspace(store, folderWorkspaceKey(folder.id)) + + await expect(store.getState().deleteFolderWorkspace(folder.id)).resolves.toBe(true) + + expect(draftTexts()).toEqual(['', '', 'other']) + }) +}) diff --git a/src/renderer/src/store/slices/worktrees/teardown/remove-worktree.ts b/src/renderer/src/store/slices/worktrees/teardown/remove-worktree.ts index 3ca041aeb75..02ba5b50a3b 100644 --- a/src/renderer/src/store/slices/worktrees/teardown/remove-worktree.ts +++ b/src/renderer/src/store/slices/worktrees/teardown/remove-worktree.ts @@ -41,6 +41,7 @@ import { recordRemovedWorktreeSnapshotPrune } from './removed-worktree-snapshot- import { clearSessionCommitDraftForWorktree } from '@/lib/source-control-commit-draft-session' import { dispatchWorktreeRemoval } from './dispatch-worktree-removal' import { tearDownRemovedWorktreeRendererState } from './removed-worktree-renderer-teardown' +import { captureWorktreeStateBeforeRemoval } from './worktree-state-before-removal' export function createRemoveWorktree( set: WorktreeSliceSet, @@ -108,9 +109,7 @@ export function createRemoveWorktree( worktreeId, requiredExecutionHostId ) - const terminalPtyIdsBeforeRemoval = (get().tabsByWorktree[worktreeId] ?? []).flatMap( - (tab) => get().ptyIdsByTabId[tab.id] ?? [] - ) + const beforeRemoval = captureWorktreeStateBeforeRemoval(get(), worktreeId) if (!forgetLocalOnly) { removalGenerationGuard?.assertCurrent() } @@ -252,7 +251,7 @@ export function createRemoveWorktree( worktreeId, hostId, requiredExecutionHostId, - terminalPtyIdsBeforeRemoval, + beforeRemoval, catalogVersion: removalResult?.catalogVersion }) // Why: Source Control may be unmounted during deletion, so it can't be the only stale-draft cleanup path. diff --git a/src/renderer/src/store/slices/worktrees/teardown/removed-worktree-chat-drafts.ts b/src/renderer/src/store/slices/worktrees/teardown/removed-worktree-chat-drafts.ts new file mode 100644 index 00000000000..d3ffd2d98ab --- /dev/null +++ b/src/renderer/src/store/slices/worktrees/teardown/removed-worktree-chat-drafts.ts @@ -0,0 +1,46 @@ +import type { AppState } from '../../../types' +import { + deleteNativeChatComposerDraft, + deleteNativeChatComposerDraftsForTab, + structuredAgentSessionDraftScopeKey +} from '@/components/native-chat/native-chat-composer-draft-store' + +/** Which unsent chat drafts a workspace's open tabs own: conversation draft keys and terminal tab + * ids (each tab's pane drafts). */ +export type WorkspaceChatDraftKeys = { + readonly conversations: readonly string[] + readonly terminalTabIds: readonly string[] +} + +/** + * Read before a user's delete goes to the host: the host announces the removal before it replies, + * and the listing refresh that starts can drop the tab lists these keys are found through. + */ +export function captureWorkspaceChatDraftKeys( + state: Pick, + workspaceIds: Iterable +): WorkspaceChatDraftKeys { + const conversations: string[] = [] + const terminalTabIds: string[] = [] + for (const workspaceId of workspaceIds) { + for (const tab of state.unifiedTabsByWorktree[workspaceId] ?? []) { + if (tab.contentType === 'agent-session') { + conversations.push(structuredAgentSessionDraftScopeKey(tab.entityId)) + } + } + for (const tab of state.tabsByWorktree[workspaceId] ?? []) { + terminalTabIds.push(tab.id) + } + } + return { conversations, terminalTabIds } +} + +/** Only after the delete succeeded: a refused or failed one keeps the drafts. */ +export function deleteWorkspaceChatDrafts(keys: WorkspaceChatDraftKeys): void { + for (const key of keys.conversations) { + deleteNativeChatComposerDraft(key) + } + for (const tabId of keys.terminalTabIds) { + deleteNativeChatComposerDraftsForTab(tabId) + } +} diff --git a/src/renderer/src/store/slices/worktrees/teardown/removed-worktree-renderer-teardown.ts b/src/renderer/src/store/slices/worktrees/teardown/removed-worktree-renderer-teardown.ts index 5bb297d05f3..4a4c0a4a2e6 100644 --- a/src/renderer/src/store/slices/worktrees/teardown/removed-worktree-renderer-teardown.ts +++ b/src/renderer/src/store/slices/worktrees/teardown/removed-worktree-renderer-teardown.ts @@ -11,11 +11,8 @@ import { disposeRemovedWorktreeParkedTerminalWatchers } from '../../../../compon import { detachedHeadAutoDerivedDisplayNames } from '../metadata/detached-head-display-name' import { applyRemoveWorktreeSuccessState } from './remove-worktree-store-cleanup' import { purgeOrphanedRuntimeSshProjects } from './orphaned-runtime-ssh-project-purge' -import { - deleteNativeChatComposerDraft, - deleteNativeChatComposerDraftsForTab, - structuredAgentSessionDraftScopeKey -} from '@/components/native-chat/native-chat-composer-draft-store' +import { deleteWorkspaceChatDrafts } from './removed-worktree-chat-drafts' +import type { WorktreeStateBeforeRemoval } from './worktree-state-before-removal' /** * Renderer-side teardown after the backend removal succeeded. @@ -30,12 +27,11 @@ export async function tearDownRemovedWorktreeRendererState(args: { worktreeId: string hostId: ExecutionHostId | undefined requiredExecutionHostId: ExecutionHostId | null - terminalPtyIdsBeforeRemoval: readonly string[] + beforeRemoval: WorktreeStateBeforeRemoval /** The catalog the host's removal produced, when the host stamps it. */ catalogVersion?: WorktreeCatalogVersion }): Promise { - const { set, get, worktreeId, hostId, requiredExecutionHostId, terminalPtyIdsBeforeRemoval } = - args + const { set, get, worktreeId, hostId, requiredExecutionHostId, beforeRemoval } = args // Why first: a listing scanned before this removal must not be applied after it and bring // the row back, so the version is on record before any await below yields. if (hostId && args.catalogVersion) { @@ -49,10 +45,11 @@ export async function tearDownRemovedWorktreeRendererState(args: { ) ) } - const structuredSessionIds: string[] = [] + // Why the captured keys: closing a structured chat keeps its conversation's draft, and a listing + // refresh may already have dropped these tabs. + deleteWorkspaceChatDrafts(beforeRemoval.chatDraftKeys) for (const tab of get().unifiedTabsByWorktree[worktreeId] ?? []) { if (tab.contentType === 'agent-session') { - structuredSessionIds.push(tab.entityId) get().closeUnifiedTab(tab.id, { preserveWorktreeSelection: true, recordInteraction: false @@ -86,20 +83,12 @@ export async function tearDownRemovedWorktreeRendererState(args: { detachedHeadAutoDerivedDisplayNames.delete(worktreeId) forgetForegroundTerminalTabs(tabIds) forgetAgentStartupDeliveriesForTabs(tabIds) - // Why: closing a structured chat keeps its conversation's draft, and terminal tabs skip - // closeTab, so both die with their worktree here. - for (const sessionId of structuredSessionIds) { - deleteNativeChatComposerDraft(structuredAgentSessionDraftScopeKey(sessionId)) - } - for (const tabId of tabIds) { - deleteNativeChatComposerDraftsForTab(tabId) - } // Why: snapshot the sidebar top-row anchor in the same tick we remove the row; recording at click time goes stale across the await. requestVirtualizedScrollAnchorRecord('[data-worktree-sidebar]') // Why: dispose parked terminal watchers only on explicit deletion; identity migration/remounts must keep buffered PTY state. - disposeRemovedWorktreeParkedTerminalWatchers(worktreeId, terminalPtyIdsBeforeRemoval) + disposeRemovedWorktreeParkedTerminalWatchers(worktreeId, beforeRemoval.terminalPtyIds) applyRemoveWorktreeSuccessState(set, worktreeId, tabIds, requiredExecutionHostId ?? hostId) get().removeWorkspaceSpaceWorktrees?.( hostId ? [{ id: worktreeId, executionHostId: hostId }] : [worktreeId] diff --git a/src/renderer/src/store/slices/worktrees/teardown/worktree-state-before-removal.ts b/src/renderer/src/store/slices/worktrees/teardown/worktree-state-before-removal.ts new file mode 100644 index 00000000000..c71a0bb3ee7 --- /dev/null +++ b/src/renderer/src/store/slices/worktrees/teardown/worktree-state-before-removal.ts @@ -0,0 +1,24 @@ +import type { AppState } from '../../../types' +import { + captureWorkspaceChatDraftKeys, + type WorkspaceChatDraftKeys +} from './removed-worktree-chat-drafts' + +/** What the teardown needs from the worktree's tab lists, which a listing refresh started during + * the removal round trip can drop before the teardown runs. */ +export type WorktreeStateBeforeRemoval = { + readonly terminalPtyIds: readonly string[] + readonly chatDraftKeys: WorkspaceChatDraftKeys +} + +export function captureWorktreeStateBeforeRemoval( + state: Pick, + worktreeId: string +): WorktreeStateBeforeRemoval { + return { + terminalPtyIds: (state.tabsByWorktree[worktreeId] ?? []).flatMap( + (tab) => state.ptyIdsByTabId[tab.id] ?? [] + ), + chatDraftKeys: captureWorkspaceChatDraftKeys(state, [worktreeId]) + } +}