Merge remote-tracking branch 'origin/brennanb2025/chat-drafts-survive-reload' into brennanb2025/chat-send-settles-once

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) {

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) => {

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<typeof AgentStatusModule>()
+  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<typeof createTestStore>, 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<typeof createTestStore>): 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'])
+  })
+})

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.

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<AppState, 'tabsByWorktree' | 'unifiedTabsByWorktree'>,
+  workspaceIds: Iterable<string>
+): 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)
+  }
+}

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<void> {
-  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]

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<AppState, 'tabsByWorktree' | 'unifiedTabsByWorktree' | 'ptyIdsByTabId'>,
+  worktreeId: string
+): WorktreeStateBeforeRemoval {
+  return {
+    terminalPtyIds: (state.tabsByWorktree[worktreeId] ?? []).flatMap(
+      (tab) => state.ptyIdsByTabId[tab.id] ?? []
+    ),
+    chatDraftKeys: captureWorkspaceChatDraftKeys(state, [worktreeId])
+  }
+}