feat(ssh): track connection-manager drains, test probes and provider continuations (#16741 T2 P3+P8a)

src/main/ipc/filesystem-import-ssh-remote-existence.ts

@@ -0,0 +1,29 @@
+import type { IFilesystemProvider } from '../providers/types'
+
+/** Whether a remote path exists; only a definite "missing" answer reads as false. */
+export async function remotePathExists(
+  provider: IFilesystemProvider,
+  remotePath: string
+): Promise<boolean> {
+  try {
+    await provider.stat(remotePath)
+    return true
+  } catch (error) {
+    if (isRemoteMissingError(error)) {
+      return false
+    }
+    throw error
+  }
+}
+
+function isRemoteMissingError(error: unknown): boolean {
+  if (!(error instanceof Error)) {
+    return false
+  }
+  return (
+    ('code' in error && error.code === 'ENOENT') ||
+    /\b(ENOENT|ENOTDIR)\b|no such file or directory|cannot find (?:the )?(?:file|path)|(?:file|path) not found/i.test(
+      error.message
+    )
+  )
+}

src/main/ipc/filesystem-import-ssh.ts

@@ -8,6 +8,8 @@ import type { FileUploadSession, IFilesystemProvider } from '../providers/types'
 import type { ImportItemResult } from '../../shared/filesystem-import-result-types'
 import { assertSafeRemotePathSegment, type RemotePathFlavor } from '../ssh/ssh-remote-platform'
 import { isWindowsAbsolutePathLike } from '../../shared/cross-platform-path'
+import { runSshProviderContinuation } from '../ssh/ssh-provider-continuations'
+import { remotePathExists } from './filesystem-import-ssh-remote-existence'
 import {
   captureLocalUploadRoot,
   preScanSshImportDirectory,
@@ -17,6 +19,12 @@ import {
 // Why: the SSH import path uses SshFilesystemProvider instead of direct SFTP so
 // system-SSH transports (ProxyCommand/ProxyJump/FIDO2) get the same workflows.
 export async function importExternalPathsSsh(
+  ...args: Parameters<typeof importExternalPathsSshTracked>
+): Promise<{ results: ImportItemResult[] }> {
+  return runSshProviderContinuation(args[2], () => importExternalPathsSshTracked(...args))
+}
+
+async function importExternalPathsSshTracked(
   sourcePaths: string[],
   destDir: string,
   connectionId: string,
@@ -269,31 +277,3 @@ async function ensureDropStagingDir(
   assertCurrent?.()
   await provider.createDir(destDir)
 }
-
-async function remotePathExists(
-  provider: IFilesystemProvider,
-  remotePath: string
-): Promise<boolean> {
-  try {
-    await provider.stat(remotePath)
-    return true
-  } catch (error) {
-    if (isRemoteMissingError(error)) {
-      return false
-    }
-    throw error
-  }
-}
-
-function isRemoteMissingError(error: unknown): boolean {
-  if (!(error instanceof Error)) {
-    return false
-  }
-  const code = (error as NodeJS.ErrnoException).code
-  return (
-    code === 'ENOENT' ||
-    /\b(ENOENT|ENOTDIR)\b|no such file or directory|cannot find (?:the )?(?:file|path)|(?:file|path) not found/i.test(
-      error.message
-    )
-  )
-}

src/main/ipc/filesystem-provider-continuations.test.ts

@@ -0,0 +1,238 @@
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+import type { IFilesystemProvider } from '../providers/types'
+import type { FilesystemHandlerContext } from './filesystem/filesystem-handler-context'
+import { hasSshProviderContinuations } from '../ssh/ssh-provider-continuations'
+import {
+  registerSshFilesystemProvider,
+  unregisterSshFilesystemProvider
+} from '../providers/ssh-filesystem-dispatch'
+import {
+  advanceSshConnectionGeneration,
+  resetSshConnectionGenerations
+} from '../ssh/ssh-connection-generation'
+import { importExternalPathsSsh } from './filesystem-import-ssh'
+import { captureLocalUploadRoot, uploadSshImportDirectory } from './filesystem-import-ssh-directory'
+import { registerFilesystemWriteHandlers } from './filesystem/filesystem-write-handlers'
+
+const mocks = vi.hoisted(() => ({
+  handle: vi.fn(),
+  lstat: vi.fn(),
+  writeFile: vi.fn(),
+  trashItem: vi.fn(),
+  resolveAuthorizedPath: vi.fn(),
+  tryDeleteWslUncPath: vi.fn()
+}))
+vi.mock('electron', () => ({ ipcMain: { handle: mocks.handle }, shell: mocks }))
+vi.mock('node:fs/promises', () => ({ lstat: mocks.lstat, writeFile: mocks.writeFile }))
+vi.mock('./filesystem-auth', () => ({
+  authorizeExternalPath: vi.fn(),
+  resolveAuthorizedPath: mocks.resolveAuthorizedPath
+}))
+vi.mock('./filesystem-mutations', () => ({ registerFilesystemMutationHandlers: vi.fn() }))
+vi.mock('../wsl-unc-delete', () => ({ tryDeleteWslUncPath: mocks.tryDeleteWslUncPath }))
+vi.mock('./ssh', () => ({
+  getSshConnectionManager: () => ({
+    getConnection: () => ({ getState: () => ({ status: 'connected' }) })
+  })
+}))
+vi.mock('./filesystem-import-ssh-directory', () => ({
+  captureLocalUploadRoot: vi.fn(),
+  preScanSshImportDirectory: vi.fn(),
+  uploadSshImportDirectory: vi.fn()
+}))
+
+const targetId = 'filesystem-continuation-target'
+// The IPC payload fields these handlers read; each test passes the subset its channel needs.
+type FilesystemHandlerArgs = Record<string, unknown>
+const handlers = new Map<
+  string,
+  (_event: unknown, args: FilesystemHandlerArgs) => Promise<unknown>
+>()
+const fileStat = { isFile: () => true, isDirectory: () => false, isSymbolicLink: () => false }
+
+beforeEach(() => {
+  vi.resetAllMocks()
+  resetSshConnectionGenerations()
+  unregisterSshFilesystemProvider(targetId)
+  expect(hasSshProviderContinuations(targetId)).toBe(false)
+  handlers.clear()
+  mocks.handle.mockImplementation((channel, handler) => handlers.set(channel, handler))
+  mocks.lstat.mockResolvedValue(fileStat)
+  mocks.resolveAuthorizedPath.mockImplementation(async (path) => path)
+  mocks.tryDeleteWslUncPath.mockResolvedValue(false)
+  // oxlint-disable-next-line typescript/consistent-type-assertions -- SAFETY: the write handlers read no store field on the SSH branch under test.
+  registerFilesystemWriteHandlers({ store: {} } as FilesystemHandlerContext)
+})
+
+describe('SSH filesystem mutation continuation settlement', () => {
+  it.each(['fs:writeFile', 'fs:deletePath'])(
+    '%s retains removed providers until settlement',
+    async (channel) => {
+      const pending = Promise.withResolvers<void>()
+      const operation = vi.fn(() => {
+        expect(hasSshProviderContinuations(targetId)).toBe(true)
+        return pending.promise
+      })
+      // oxlint-disable-next-line typescript/consistent-type-assertions -- SAFETY: the stub implements only the provider methods these handlers call.
+      registerSshFilesystemProvider(targetId, {
+        writeFile: operation,
+        deletePath: operation
+      } as unknown as IFilesystemProvider)
+      const result = handlers.get(channel)!(null, {
+        connectionId: targetId,
+        expectedSshTargetId: targetId,
+        expectedSshConnectionGeneration: 0,
+        filePath: '/remote/file',
+        content: 'text',
+        targetPath: '/remote/file',
+        recursive: true
+      })
+      unregisterSshFilesystemProvider(targetId)
+      advanceSshConnectionGeneration(targetId)
+      expect(hasSshProviderContinuations(targetId)).toBe(true)
+      expect(hasSshProviderContinuations('other-target')).toBe(false)
+      pending.resolve()
+      await result
+      expect(hasSshProviderContinuations(targetId)).toBe(false)
+      expect(operation).toHaveBeenCalledWith(
+        ...(channel === 'fs:writeFile' ? ['/remote/file', 'text'] : ['/remote/file', true])
+      )
+    }
+  )
+
+  it.each(['fs:writeFile', 'fs:deletePath'])(
+    '%s releases after provider rejection',
+    async (channel) => {
+      const pending = Promise.withResolvers<void>()
+      // oxlint-disable-next-line typescript/consistent-type-assertions -- SAFETY: the stub implements only the provider methods these handlers call.
+      registerSshFilesystemProvider(targetId, {
+        writeFile: () => pending.promise,
+        deletePath: () => pending.promise
+      } as unknown as IFilesystemProvider)
+      const result = handlers.get(channel)!(null, {
+        connectionId: targetId,
+        expectedSshTargetId: targetId,
+        expectedSshConnectionGeneration: 0,
+        filePath: '/remote/file',
+        targetPath: '/remote/file',
+        content: ''
+      })
+      const failure = expect(result).rejects.toThrow('late failure')
+      expect(hasSshProviderContinuations(targetId)).toBe(true)
+      pending.reject(new Error('late failure'))
+      await failure
+      expect(hasSshProviderContinuations(targetId)).toBe(false)
+    }
+  )
+
+  it.each(['fs:writeFile', 'fs:deletePath'])(
+    '%s leaves local handling untracked',
+    async (channel) => {
+      const pending = Promise.withResolvers<void>()
+      mocks.writeFile.mockReturnValue(pending.promise)
+      mocks.trashItem.mockReturnValue(pending.promise)
+      const result = handlers.get(channel)!(null, {
+        filePath: '/local/file',
+        targetPath: '/local/file',
+        content: 'text'
+      })
+      expect(hasSshProviderContinuations(targetId)).toBe(false)
+      pending.resolve()
+      await result
+      expect(channel === 'fs:writeFile' ? mocks.writeFile : mocks.trashItem).toHaveBeenCalled()
+      expect(hasSshProviderContinuations(targetId)).toBe(false)
+    }
+  )
+})
+
+describe('complete SSH import continuation settlement', () => {
+  it('retains a failed directory import through pending rollback', async () => {
+    const rollback = Promise.withResolvers<void>()
+    const rollingBack = Promise.withResolvers<void>()
+    mocks.lstat.mockResolvedValue({ ...fileStat, isDirectory: () => true, isFile: () => false })
+    vi.mocked(captureLocalUploadRoot).mockResolvedValue('/source/directory')
+    vi.mocked(uploadSshImportDirectory).mockRejectedValue(new Error('upload failed'))
+    const close = vi.fn(() => expect(hasSshProviderContinuations(targetId)).toBe(true))
+    const deletePath = vi.fn(() => {
+      rollingBack.resolve()
+      return rollback.promise
+    })
+    // oxlint-disable-next-line typescript/consistent-type-assertions -- SAFETY: the stub implements only the provider methods these handlers call.
+    registerSshFilesystemProvider(targetId, {
+      openFileUploadSession: async () => ({ close }),
+      stat: async () => {
+        throw Object.assign(new Error('missing'), { code: 'ENOENT' })
+      },
+      createDirNoClobber: async () => {},
+      deletePath
+    } as unknown as IFilesystemProvider)
+    const result = importExternalPathsSsh(['/source/directory'], '/remote', targetId)
+    await rollingBack.promise
+    unregisterSshFilesystemProvider(targetId)
+    expect(hasSshProviderContinuations(targetId)).toBe(true)
+    expect(close).not.toHaveBeenCalled()
+    rollback.resolve()
+    expect((await result).results[0]).toMatchObject({ status: 'failed', reason: 'upload failed' })
+    expect(deletePath).toHaveBeenCalledWith('/remote/directory', true)
+    expect(close).toHaveBeenCalledOnce()
+    expect(hasSshProviderContinuations(targetId)).toBe(false)
+  })
+
+  it('retains tracking across local inspection, provider removal and session close', async () => {
+    const inspection = Promise.withResolvers<typeof fileStat>()
+    const inspected = Promise.withResolvers<void>()
+    const uploaded = Promise.withResolvers<void>()
+    const upload = Promise.withResolvers<void>()
+    mocks.lstat.mockImplementation(() => {
+      inspected.resolve()
+      return inspection.promise
+    })
+    const close = vi.fn(() => expect(hasSshProviderContinuations(targetId)).toBe(true))
+    // oxlint-disable-next-line typescript/consistent-type-assertions -- SAFETY: the stub implements only the provider methods these handlers call.
+    registerSshFilesystemProvider(targetId, {
+      openFileUploadSession: async () => ({
+        close,
+        uploadFile: () => {
+          uploaded.resolve()
+          return upload.promise
+        }
+      }),
+      stat: async () => {
+        throw Object.assign(new Error('missing'), { code: 'ENOENT' })
+      }
+    } as unknown as IFilesystemProvider)
+    const result = importExternalPathsSsh(['/source/file'], '/remote', targetId)
+    expect(hasSshProviderContinuations(targetId)).toBe(true)
+    await inspected.promise
+    unregisterSshFilesystemProvider(targetId)
+    advanceSshConnectionGeneration(targetId)
+    expect(hasSshProviderContinuations(targetId)).toBe(true)
+    inspection.resolve(fileStat)
+    await uploaded.promise
+    expect(hasSshProviderContinuations(targetId)).toBe(true)
+    upload.resolve()
+    expect((await result).results[0].status).toBe('imported')
+    expect(close).toHaveBeenCalledOnce()
+    expect(hasSshProviderContinuations(targetId)).toBe(false)
+  })
+
+  it('includes staging preparation and releases on its failure', async () => {
+    const staging = Promise.withResolvers<void>()
+    // oxlint-disable-next-line typescript/consistent-type-assertions -- SAFETY: the stub implements only the provider methods these handlers call.
+    registerSshFilesystemProvider(targetId, {
+      createDir: () => {
+        expect(hasSshProviderContinuations(targetId)).toBe(true)
+        return staging.promise
+      }
+    } as unknown as IFilesystemProvider)
+    const result = importExternalPathsSsh(['/source/file'], '/remote/.orca/drops', targetId, {
+      ensureDir: true
+    })
+    const failure = expect(result).rejects.toThrow('staging failed')
+    unregisterSshFilesystemProvider(targetId)
+    expect(hasSshProviderContinuations(targetId)).toBe(true)
+    staging.reject(new Error('staging failed'))
+    await failure
+    expect(hasSshProviderContinuations(targetId)).toBe(false)
+  })
+})

src/main/ipc/filesystem/filesystem-write-handlers.ts

@@ -2,6 +2,7 @@ import { ipcMain, shell } from 'electron'
 import { lstat, writeFile } from 'node:fs/promises'
 import type { SshMutationExpectation } from '../../../shared/ssh-types'
 import { assertSshMutationExpectation } from '../../ssh/ssh-connection-generation'
+import { runSshProviderContinuation } from '../../ssh/ssh-provider-continuations'
 import { requireSshFilesystemProvider } from '../../providers/ssh-filesystem-dispatch'
 import { tryDeleteWslUncPath } from '../../wsl-unc-delete'
 import { authorizeExternalPath, resolveAuthorizedPath } from '../filesystem-auth'
@@ -26,7 +27,9 @@ export function registerFilesystemWriteHandlers(context: FilesystemHandlerContex
       )
       if (args.connectionId) {
         const provider = requireSshFilesystemProvider(args.connectionId)
-        return provider.writeFile(args.filePath, args.content)
+        return runSshProviderContinuation(args.connectionId, () =>
+          provider.writeFile(args.filePath, args.content)
+        )
       }
       const filePath = await resolveAuthorizedPath(args.filePath, store)
       try {
@@ -61,7 +64,9 @@ export function registerFilesystemWriteHandlers(context: FilesystemHandlerContex
       )
       if (args.connectionId) {
         const provider = requireSshFilesystemProvider(args.connectionId)
-        return provider.deletePath(args.targetPath, args.recursive)
+        return runSshProviderContinuation(args.connectionId, () =>
+          provider.deletePath(args.targetPath, args.recursive)
+        )
       }
       // Why: preserve the symlink so we delete the link, not its target (realpath would trash the real file, possibly outside all roots).
       const targetPath = await resolveAuthorizedPath(args.targetPath, store, {

src/main/ipc/ssh-active-relay-sessions.ts

@@ -1,5 +1,8 @@
 import type { SshRelaySession } from '../ssh/ssh-relay-session'
-import { setSshActiveMultiplexerResolver } from '../ssh/ssh-target-registry'
+import {
+  setDirectSshAuthorityResolver,
+  setSshActiveMultiplexerResolver
+} from '../ssh/ssh-target-registry'
 import { setWorktreeRemovalSshHostHomeResolver } from '../worktree-removal-execution-host-route'
 
 // One session per SSH target owns the whole relay lifecycle (mux, providers, abort controller, state machine).
@@ -22,3 +25,5 @@ export function getActiveSshHostHomeDirectory(targetId: string): string | null {
 }
 
 setWorktreeRemovalSshHostHomeResolver(getActiveSshHostHomeDirectory)
+
+setDirectSshAuthorityResolver((targetId) => activeSessions.has(targetId))