From 7e3df6433cd68378d241f7e653d109848f3d865e Mon Sep 17 00:00:00 2001 From: m4air Date: Thu, 1 Oct 2026 09:08:43 -0700 Subject: [PATCH 1/2] feat(ssh): track connection-manager drains, test probes and provider continuations (#16741 T2 P3+P8a) The SSH connection manager now registers each connection it allocates against the connection's transport-closure notice, and tracks every connect, disconnect, reconnect and teardown per target, so later slices can tell when a target's local work and transports are really gone. Ordinary connect, disconnect and quit behavior is unchanged apart from cleanup. Manager (ssh-connection-manager.ts): - registerConnection counts every connection a target allocated, retired pool entries included, until its transport-closure notice arrives. - connect/disconnect/reconnect/disconnectConnection/disconnectAll run through trackTargetOperation; a failed teardown is remembered as unconfirmed. - A failed startup is now disconnected, not just dropped from the pool. - disconnectAndDrain(targetId, signal) and disconnectConnection(..., drain), bounded by a 10s timeout. - disconnectAll(shouldDisconnect) only tears down targets the filter allows. IPC: - ssh-connect-attempt-registry: runSshTestConnectionProbe publishes a probe per target before it starts and clears testingTargets / credential flags only when the target's last probe settles. ssh:testConnection uses it, so quit still joins an in-flight probe within the shutdown budget. - ssh-target-lifecycle-queue: runTargetLifecycle returns the operation's value. - ssh-shutdown-drain: a mayDetach predicate (default: every target) is plumbed through detach, invalidation and disconnectAll. - ssh-renderer-broadcast: targets owned by a runtime are hidden by owner too, not only by id prefix. - ssh-target-registry: direct-authority resolver, installed by ssh-active-relay-sessions. - Provider dispatch: unregister*IfCurrent for git and filesystem providers. P8a, SSH provider continuations: ssh-provider-continuations tracks local settlement of SSH filesystem writes/deletes, imports, detected-worktree listings and worktree/folder removals per target. It records only local settlement, never remote execution or exit. Porting note (source: #16741 head a68b6f3531, merge-base 277c289bd4): - Left for later slices: connectExclusive and assertTargetTransportsClosed (T8), hasTargetActivity (T3), assertProfileLifetimeAdmission (P8b, stubbed out of the manager and continuations), the reset admission checks in ssh-connection-handlers and the reset predicate in the shutdown drain (T3), the network-tunnel resolver (T4), and the managed-orcad provisioning test hunks in ssh-target-registry.test.ts (T6). - Tests read the manager's bookkeeping through ssh-connection-manager-test-probes until T3/T8 expose readers. Two of the four deferred disconnect-drain cases are restored; the two that need connectExclusive, two closure cases that need it and three profile-admission cases stay deferred. - The removal and detected-worktree call sites were rebased onto main's background removal and execution-host routing. filesystem-import-ssh moved its remote-existence helpers out to stay under 300 lines. --- .../filesystem-import-ssh-remote-existence.ts | 29 +++ src/main/ipc/filesystem-import-ssh.ts | 36 +-- .../filesystem-provider-continuations.test.ts | 233 ++++++++++++++++++ .../filesystem/filesystem-write-handlers.ts | 9 +- src/main/ipc/ssh-active-relay-sessions.ts | 7 +- src/main/ipc/ssh-connect-attempt-registry.ts | 41 +++ src/main/ipc/ssh-connection-handlers.ts | 14 +- src/main/ipc/ssh-ipc-module-mocks.ts | 1 + src/main/ipc/ssh-renderer-broadcast.ts | 7 +- src/main/ipc/ssh-shutdown-drain.ts | 39 ++- .../ipc/ssh-state-broadcast-fanout.test.ts | 35 +++ src/main/ipc/ssh-target-lifecycle-queue.ts | 24 +- ...ssh-test-connection-probe-registry.test.ts | 98 ++++++++ src/main/ipc/ssh-test-probe-presence.test.ts | 36 +++ .../worktrees-ssh-provider-authority.test.ts | 94 +++++++ .../register-detected-worktree-handlers.ts | 50 ++-- .../register-worktree-removal-handlers.ts | 13 +- .../worktree-removal-continuations.test.ts | 98 ++++++++ src/main/providers/ssh-filesystem-dispatch.ts | 10 + src/main/providers/ssh-git-dispatch.ts | 15 ++ .../ssh-connection-disconnect-drain.test.ts | 26 ++ .../ssh-connection-manager-closure.test.ts | 108 ++++++++ .../ssh-connection-manager-registry.test.ts | 53 ++++ .../ssh/ssh-connection-manager-test-probes.ts | 43 ++++ src/main/ssh/ssh-connection-manager.test.ts | 50 ++++ src/main/ssh/ssh-connection-manager.ts | 137 +++++++++- .../ssh/ssh-connection-store-test-fixture.ts | 107 ++++++++ src/main/ssh/ssh-connection-store.test.ts | 69 +----- src/main/ssh/ssh-provider-continuations.ts | 24 ++ src/main/ssh/ssh-target-registry.ts | 12 + 30 files changed, 1350 insertions(+), 168 deletions(-) create mode 100644 src/main/ipc/filesystem-import-ssh-remote-existence.ts create mode 100644 src/main/ipc/filesystem-provider-continuations.test.ts create mode 100644 src/main/ipc/ssh-test-connection-probe-registry.test.ts create mode 100644 src/main/ipc/ssh-test-probe-presence.test.ts create mode 100644 src/main/ipc/worktrees/removal/worktree-removal-continuations.test.ts create mode 100644 src/main/ssh/ssh-connection-manager-closure.test.ts create mode 100644 src/main/ssh/ssh-connection-manager-test-probes.ts create mode 100644 src/main/ssh/ssh-connection-store-test-fixture.ts create mode 100644 src/main/ssh/ssh-provider-continuations.ts diff --git a/src/main/ipc/filesystem-import-ssh-remote-existence.ts b/src/main/ipc/filesystem-import-ssh-remote-existence.ts new file mode 100644 index 000000000000..e853f291f803 --- /dev/null +++ b/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 { + 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 + ) + ) +} diff --git a/src/main/ipc/filesystem-import-ssh.ts b/src/main/ipc/filesystem-import-ssh.ts index 9268c6ece247..39b2350bb44e 100644 --- a/src/main/ipc/filesystem-import-ssh.ts +++ b/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 +): 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 { - 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 - ) - ) -} diff --git a/src/main/ipc/filesystem-provider-continuations.test.ts b/src/main/ipc/filesystem-provider-continuations.test.ts new file mode 100644 index 000000000000..be7a4b25e086 --- /dev/null +++ b/src/main/ipc/filesystem-provider-continuations.test.ts @@ -0,0 +1,233 @@ +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' +const handlers = new Map Promise>() +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() + 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() + // 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() + 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() + const rollingBack = Promise.withResolvers() + 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() + const inspected = Promise.withResolvers() + const uploaded = Promise.withResolvers() + const upload = Promise.withResolvers() + 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() + // 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) + }) +}) diff --git a/src/main/ipc/filesystem/filesystem-write-handlers.ts b/src/main/ipc/filesystem/filesystem-write-handlers.ts index 07d4de5b56a9..0cc5a921b451 100644 --- a/src/main/ipc/filesystem/filesystem-write-handlers.ts +++ b/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, { diff --git a/src/main/ipc/ssh-active-relay-sessions.ts b/src/main/ipc/ssh-active-relay-sessions.ts index 69b781be41d4..3c12496f7d4f 100644 --- a/src/main/ipc/ssh-active-relay-sessions.ts +++ b/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)) diff --git a/src/main/ipc/ssh-connect-attempt-registry.ts b/src/main/ipc/ssh-connect-attempt-registry.ts index 5a1e36329824..5aba526d6832 100644 --- a/src/main/ipc/ssh-connect-attempt-registry.ts +++ b/src/main/ipc/ssh-connect-attempt-registry.ts @@ -44,3 +44,44 @@ export const resetRelayInFlight = new Map>() // Why: ssh:testConnection connects then disconnects; suppressing broadcasts during the test avoids worktree cards flashing connected → disconnected. export const testingTargets = new Set() export const testConnectionProbes = new Set>() +const testConnectionProbesByTarget = new Map>>() + +export function hasSshTestConnectionProbes(targetId: string): boolean { + return (testConnectionProbesByTarget.get(targetId)?.size ?? 0) > 0 +} + +export function runSshTestConnectionProbe( + targetId: string, + operation: () => Promise +): Promise { + const probes = testConnectionProbesByTarget.get(targetId) ?? new Set>() + let probe!: Promise + // Publish before connection callbacks can start a reset or shutdown drain. + probe = Promise.resolve() + .then(operation) + .finally(() => { + testConnectionProbes.delete(probe) + probes.delete(probe) + if (probes.size === 0 && testConnectionProbesByTarget.get(targetId) === probes) { + testConnectionProbesByTarget.delete(targetId) + testingTargets.delete(targetId) + credentialRequestedForTarget.delete(targetId) + } + }) + probes.add(probe) + testConnectionProbesByTarget.set(targetId, probes) + testConnectionProbes.add(probe) + testingTargets.add(targetId) + return probe +} + +/** Reserve target admission before joining; failed probes still own cleanup until settled. */ +export async function awaitSshTestConnectionProbes(targetId: string): Promise { + while (true) { + const probes = testConnectionProbesByTarget.get(targetId) + if (!probes?.size) { + return + } + await Promise.allSettled(probes) + } +} diff --git a/src/main/ipc/ssh-connection-handlers.ts b/src/main/ipc/ssh-connection-handlers.ts index 9ea533c77c30..946dfed7bf9f 100644 --- a/src/main/ipc/ssh-connection-handlers.ts +++ b/src/main/ipc/ssh-connection-handlers.ts @@ -23,8 +23,7 @@ import { credentialRequestedForTarget, invalidateConnectAttempt, resetRelayInFlight, - testConnectionProbes, - testingTargets + runSshTestConnectionProbe } from './ssh-connect-attempt-registry' import { connectTarget } from './ssh-connect-flow' import { connectionManager, persistedStore } from './ssh-ipc-context' @@ -266,18 +265,16 @@ export function registerSshConnectionHandlers(): void { } } - testingTargets.add(args.targetId) // Why a tracked promise and not just the id: a probe holds a real transport that no session owns, // so shutdown has to be able to join it before the final drain disconnects what is left. - const probe = (async () => { + const probe = runSshTestConnectionProbe(args.targetId, async () => { // Why: a probe transport opened after the shutdown drain would outlive orderly teardown. assertSshConnectsNotFenced() const conn = await connectionManager!.connect(target) const state = conn.getState() await connectionManager!.disconnect(args.targetId) return state - })() - testConnectionProbes.add(probe) + }) try { return { success: true, state: await probe } } catch (err) { @@ -285,11 +282,6 @@ export function registerSshConnectionHandlers(): void { success: false, error: err instanceof Error ? err.message : String(err) } - } finally { - testConnectionProbes.delete(probe) - testingTargets.delete(args.targetId) - // Why: clear so a test's credential prompt doesn't leave lastRequiredPassphrase=true and defer this target at startup. - credentialRequestedForTarget.delete(args.targetId) } }) } diff --git a/src/main/ipc/ssh-ipc-module-mocks.ts b/src/main/ipc/ssh-ipc-module-mocks.ts index 0855a739fb2b..af05f0b12049 100644 --- a/src/main/ipc/ssh-ipc-module-mocks.ts +++ b/src/main/ipc/ssh-ipc-module-mocks.ts @@ -144,6 +144,7 @@ export function createSshIpcMocks(): SshIpcMocks { installSshPtySourceCancellationPublisher: vi.fn().mockReturnValue(() => {}) }, sshConnectionStore: { + isRuntimeOwnedSshTarget: (target: { owner?: unknown }) => target.owner !== undefined, SshConnectionStore: class MockSshConnectionStore { constructor() { return mockSshStore diff --git a/src/main/ipc/ssh-renderer-broadcast.ts b/src/main/ipc/ssh-renderer-broadcast.ts index 999b8e9a3462..8dbc1cfede16 100644 --- a/src/main/ipc/ssh-renderer-broadcast.ts +++ b/src/main/ipc/ssh-renderer-broadcast.ts @@ -12,8 +12,10 @@ import { enrichSshForwardEntries, getWorktreeIdsForConnection } from '../ports/ssh-advertised-url-enrichment' +import { isRuntimeOwnedSshTarget } from '../ssh/ssh-connection-store' import { getSshProviderAuthority } from '../ssh/ssh-provider-authority' import { getSshPlainSshMode } from '../ssh/ssh-plain-ssh-mode' +import { getSshTargetRegistryStore } from '../ssh/ssh-target-registry' import { activeSessions } from './ssh-active-relay-sessions' import { connectionManager, @@ -30,8 +32,9 @@ export function broadcastSshState( targetId: string, state: SshConnectionState ): void { - // Why: runtime-owned (ephemeral-VM) targets are hidden from the renderer, so broadcasting their state only triggers wasted listTargets() lookups. - if (isRuntimeOwnedSshTargetId(targetId)) { + const target = getSshTargetRegistryStore()?.getTarget(targetId) + // Why: owned targets are hidden from clients; broadcasting them leaks internal transport state into persisted reconnect hints. + if (isRuntimeOwnedSshTargetId(targetId) || (target && isRuntimeOwnedSshTarget(target))) { currentRuntime?.invalidateSshWorktreeScanCache?.(targetId) return } diff --git a/src/main/ipc/ssh-shutdown-drain.ts b/src/main/ipc/ssh-shutdown-drain.ts index f66db066c011..93a25071bfc6 100644 --- a/src/main/ipc/ssh-shutdown-drain.ts +++ b/src/main/ipc/ssh-shutdown-drain.ts @@ -61,22 +61,34 @@ async function settleTasksWithinMs( return { timedOut: [...pending], errors } } -function sshShutdownTasks(targetIds: readonly string[]): SshShutdownTask[] { +/** Which targets the quit drain may detach; a target it may not keeps its session and transport. */ +export type SshShutdownDetachPredicate = (targetId: string) => boolean + +const detachEveryTarget: SshShutdownDetachPredicate = () => true + +function sshShutdownTasks( + targetIds: readonly string[], + mayDetach: SshShutdownDetachPredicate +): SshShutdownTask[] { return [ ...targetIds - .filter((targetId) => activeSessions.has(targetId)) + .filter((targetId) => activeSessions.has(targetId) && mayDetach(targetId)) .map((targetId) => ({ targetId, promise: teardownActiveSshSession(targetId, (session) => session.detachAndPersist()) })), - { targetId: '*transports', promise: connectionManager?.disconnectAll() ?? Promise.resolve() } + { + targetId: '*transports', + promise: connectionManager?.disconnectAll(mayDetach) ?? Promise.resolve() + } ] } async function drainSshShutdown( targetIds: readonly string[], inFlight: readonly SshShutdownTask[], - detachErrors: readonly unknown[] = [] + detachErrors: readonly unknown[] = [], + mayDetach: SshShutdownDetachPredicate = detachEveryTarget ): Promise { const deadline = Date.now() + SSH_SHUTDOWN_BUDGET_MS const unfinished: SshShutdownUnfinished[] = [] @@ -96,12 +108,12 @@ async function drainSshShutdown( return settled.timedOut.length === 0 } - await runPhase('drain', sshShutdownTasks(targetIds)) + await runPhase('drain', sshShutdownTasks(targetIds, mayDetach)) // Why a second drain after the join: a connect paused in old-session teardown still publishes its // replacement session and opens a transport before it reaches the cancellation checkpoint, so the // first drain can miss both. if (await runPhase('in-flight-join', inFlight)) { - await runPhase('final-drain', sshShutdownTasks([...activeSessions.keys()])) + await runPhase('final-drain', sshShutdownTasks([...activeSessions.keys()], mayDetach)) } if (errors.length > 0 || unfinished.length > 0) { @@ -123,7 +135,9 @@ async function drainSshShutdown( // // Why no fence latch here: the committed quit path sets it before calling this, so it is already on // for the snapshot below. Called without that gate (tests), this degrades to a plain drain. -export function beginSshShutdown(): Promise { +export function beginSshShutdown( + mayDetach: SshShutdownDetachPredicate = detachEveryTarget +): Promise { if (sshShutdownDrain) { return sshShutdownDrain } @@ -136,13 +150,18 @@ export function beginSshShutdown(): Promise { ...[...testConnectionProbes].map((promise) => ({ targetId: '*probe', promise })) ] for (const targetId of Array.from(connectInFlight.keys())) { - invalidateConnectAttempt(targetId) + if (mayDetach(targetId)) { + invalidateConnectAttempt(targetId) + } } const targetIds = [...activeSessions.keys()] // Why before any await: this is the whole point of the split. Each session marks its recovery lease // detached in memory now, and the final flush persists it — the remote PTYs keep running. const detachErrors: unknown[] = [] - for (const session of activeSessions.values()) { + for (const [targetId, session] of activeSessions) { + if (!mayDetach(targetId)) { + continue + } // Why per-session: this runs synchronously inside a non-async will-quit listener, so one throw // (teardownProviders -> webContents.send on a destroyed renderer, routine on quit) would escape // it and skip every later session, the drain assignment, and the store flush that persists all @@ -153,6 +172,6 @@ export function beginSshShutdown(): Promise { detachErrors.push(error) } } - sshShutdownDrain = drainSshShutdown(targetIds, inFlight, detachErrors) + sshShutdownDrain = drainSshShutdown(targetIds, inFlight, detachErrors, mayDetach) return sshShutdownDrain } diff --git a/src/main/ipc/ssh-state-broadcast-fanout.test.ts b/src/main/ipc/ssh-state-broadcast-fanout.test.ts index e7f6480c261c..9c9ae661237d 100644 --- a/src/main/ipc/ssh-state-broadcast-fanout.test.ts +++ b/src/main/ipc/ssh-state-broadcast-fanout.test.ts @@ -175,6 +175,41 @@ describe('SSH IPC handlers', () => { expect(runtime.notifySshStateChanged).not.toHaveBeenCalled() }) + it('keeps claimed managed-orcad SSH state off clients under the original target id', () => { + const runtime = { + onPtyData: vi.fn(), + onPtyExit: vi.fn(), + invalidateSshWorktreeScanCache: vi.fn(), + notifySshStateChanged: vi.fn() + } + registerSshHandlers(mockStore as never, () => mockWindow as never, runtime as never) + mockSshStore.getTarget.mockReturnValue({ + id: 'ssh-1', + label: 'Managed orcad host', + host: 'example.com', + port: 22, + username: 'deploy', + owner: { type: 'on-demand-runtime', runtimeId: 'managed-orcad:environment-1' } + } satisfies SshTarget) + const callbacks = mockConnectionManager.callbacksRef.current as { + onStateChange: (targetId: string, state: SshConnectionState) => void + } + + callbacks.onStateChange('ssh-1', { + targetId: 'ssh-1', + status: 'connected', + error: null, + reconnectAttempt: 0 + }) + + expect(runtime.invalidateSshWorktreeScanCache).toHaveBeenCalledWith('ssh-1') + expect(runtime.notifySshStateChanged).not.toHaveBeenCalled() + expect(mockWindow.webContents.send).not.toHaveBeenCalledWith( + 'ssh:state-changed', + expect.anything() + ) + }) + it('invalidates runtime scans from hidden SSH state broadcasts', () => { const runtime = { onPtyData: vi.fn(), diff --git a/src/main/ipc/ssh-target-lifecycle-queue.ts b/src/main/ipc/ssh-target-lifecycle-queue.ts index 014a5037175d..f6c5221aa5c0 100644 --- a/src/main/ipc/ssh-target-lifecycle-queue.ts +++ b/src/main/ipc/ssh-target-lifecycle-queue.ts @@ -1,25 +1,27 @@ // Serializes disconnect/remove/terminate/reset for a single SSH target so they cannot interleave. export const targetLifecycleInFlight = new Map>() -export function runTargetLifecycle( - targetId: string, - operation: () => Promise -): Promise { +export function runTargetLifecycle(targetId: string, operation: () => Promise): Promise { const prior = targetLifecycleInFlight.get(targetId) const operationPromise = (async () => { if (prior) { await prior.catch(() => undefined) } - await operation() + return operation() })() let trackedPromise!: Promise - trackedPromise = operationPromise.finally(() => { - if (targetLifecycleInFlight.get(targetId) === trackedPromise) { - targetLifecycleInFlight.delete(targetId) - } - }) + trackedPromise = operationPromise + .then( + () => undefined, + () => undefined + ) + .finally(() => { + if (targetLifecycleInFlight.get(targetId) === trackedPromise) { + targetLifecycleInFlight.delete(targetId) + } + }) targetLifecycleInFlight.set(targetId, trackedPromise) - return trackedPromise + return operationPromise } export async function awaitTargetLifecycle(targetId: string): Promise { diff --git a/src/main/ipc/ssh-test-connection-probe-registry.test.ts b/src/main/ipc/ssh-test-connection-probe-registry.test.ts new file mode 100644 index 000000000000..c8b75e1c1ea0 --- /dev/null +++ b/src/main/ipc/ssh-test-connection-probe-registry.test.ts @@ -0,0 +1,98 @@ +import { expect, it, vi } from 'vitest' +import { + awaitSshTestConnectionProbes, + credentialRequestedForTarget, + runSshTestConnectionProbe, + testConnectionProbes, + testingTargets +} from './ssh-connect-attempt-registry' + +function deferred() { + let resolve!: () => void + const promise = new Promise((done) => { + resolve = done + }) + return { promise, resolve } +} + +it('publishes global and target tracking before callback-capable work', async () => { + const close = deferred() + let joined!: Promise + const finished = vi.fn() + const operation = vi.fn(async () => { + expect(testConnectionProbes.has(probe)).toBe(true) + expect(testingTargets.has('published')).toBe(true) + joined = awaitSshTestConnectionProbes('published').then(finished) + await close.promise + return 'connected' + }) + const probe = runSshTestConnectionProbe('published', operation) + expect(operation).not.toHaveBeenCalled() + expect(testConnectionProbes.has(probe)).toBe(true) + await Promise.resolve() + expect(finished).not.toHaveBeenCalled() + close.resolve() + await expect(probe).resolves.toBe('connected') + await joined + expect(finished).toHaveBeenCalledTimes(1) + expect(testConnectionProbes.has(probe)).toBe(false) + expect(testingTargets.has('published')).toBe(false) +}) + +it('joins every admitted probe and retains callback suppression until the final settlement', async () => { + const first = deferred() + const second = deferred() + const a = runSshTestConnectionProbe('multiple', () => first.promise) + const b = runSshTestConnectionProbe('multiple', () => second.promise) + credentialRequestedForTarget.add('multiple') + const finished = vi.fn() + const joined = awaitSshTestConnectionProbes('multiple').then(finished) + first.resolve() + await a + expect(testingTargets.has('multiple')).toBe(true) + expect(credentialRequestedForTarget.has('multiple')).toBe(true) + expect(finished).not.toHaveBeenCalled() + second.resolve() + await b + await joined + expect(testingTargets.has('multiple')).toBe(false) + expect(credentialRequestedForTarget.has('multiple')).toBe(false) +}) + +it('joins failed probes without mistaking rejection for unfinished work', async () => { + const failure = new Error('connection failed') + const probe = runSshTestConnectionProbe('failed', async () => { + throw failure + }) + const result = expect(probe).rejects.toBe(failure) + await expect(awaitSshTestConnectionProbes('failed')).resolves.toBeUndefined() + await result + expect(testConnectionProbes.has(probe)).toBe(false) + expect(testingTargets.has('failed')).toBe(false) +}) + +it('does not wait for another target', async () => { + const close = deferred() + const probe = runSshTestConnectionProbe('other', () => close.promise) + await awaitSshTestConnectionProbes('absent') + expect(testConnectionProbes.has(probe)).toBe(true) + close.resolve() + await probe +}) + +it('rechecks the target after a joined snapshot settles', async () => { + const first = deferred() + const second = deferred() + const a = runSshTestConnectionProbe('later', () => first.promise) + const finished = vi.fn() + const joined = awaitSshTestConnectionProbes('later').then(finished) + const b = runSshTestConnectionProbe('later', () => second.promise) + first.resolve() + await a + await Promise.resolve() + expect(finished).not.toHaveBeenCalled() + second.resolve() + await b + await joined + expect(finished).toHaveBeenCalledTimes(1) +}) diff --git a/src/main/ipc/ssh-test-probe-presence.test.ts b/src/main/ipc/ssh-test-probe-presence.test.ts new file mode 100644 index 000000000000..73287c4798f0 --- /dev/null +++ b/src/main/ipc/ssh-test-probe-presence.test.ts @@ -0,0 +1,36 @@ +import { expect, it } from 'vitest' +import { + hasSshTestConnectionProbes, + runSshTestConnectionProbe, + testingTargets +} from './ssh-connect-attempt-registry' + +it('publishes target-scoped presence before callbacks and retains it until settlement', async () => { + const target = 'probe-presence-success' + const work = Promise.withResolvers() + const probe = runSshTestConnectionProbe(target, async () => { + expect(hasSshTestConnectionProbes(target)).toBe(true) + await work.promise + }) + expect(hasSshTestConnectionProbes(target)).toBe(true) + expect(hasSshTestConnectionProbes('unrelated')).toBe(false) + testingTargets.delete(target) + expect(hasSshTestConnectionProbes(target)).toBe(true) + work.resolve() + await probe + expect(hasSshTestConnectionProbes(target)).toBe(false) +}) + +it('keeps other probes visible after a failed probe settles', async () => { + const target = 'probe-presence-failure' + const work = Promise.withResolvers() + const pending = runSshTestConnectionProbe(target, () => work.promise) + const failed = runSshTestConnectionProbe(target, async () => { + throw new Error('failed') + }) + await expect(failed).rejects.toThrow('failed') + expect(hasSshTestConnectionProbes(target)).toBe(true) + work.resolve() + await pending + expect(hasSshTestConnectionProbes(target)).toBe(false) +}) diff --git a/src/main/ipc/worktrees-ssh-provider-authority.test.ts b/src/main/ipc/worktrees-ssh-provider-authority.test.ts index 8d14755ab18d..8eb025da52f6 100644 --- a/src/main/ipc/worktrees-ssh-provider-authority.test.ts +++ b/src/main/ipc/worktrees-ssh-provider-authority.test.ts @@ -5,6 +5,7 @@ import type { ProviderRequestId } from '../../shared/detected-worktree-provider- import { toSshExecutionHostId } from '../../shared/execution-host' import { DETECTED_WORKTREE_PROVIDER_TIMEOUT_MS } from './worktrees/listing/register-detected-worktree-handlers' import { getSshProviderAuthority, rotateSshProviderAuthority } from '../ssh/ssh-provider-authority' +import { hasSshProviderContinuations } from '../ssh/ssh-provider-continuations' import { getSshGitProviderMock } from './worktrees-test-module-mocks' import { handlers, ipcEvent, setupWorktreeHandlers, store } from './worktrees-test-harness' @@ -95,6 +96,99 @@ describe('registerWorktreeHandlers', () => { setupWorktreeHandlers() }) + it.each(['resolve', 'reject'] as const)( + 'retains canceled SSH provider continuations until the provider settles: %s', + async (outcome) => { + const targetId = `continuation-cancel-${outcome}` + const repo = { + id: `repo-${targetId}`, + path: '/remote/repo', + displayName: 'repo', + badgeColor: '#000', + addedAt: 0, + connectionId: targetId + } + let resolveProvider!: (value: GitWorktreeInfo[]) => void + let rejectProvider!: (error: Error) => void + const provider = { + listWorktrees: vi.fn( + () => + new Promise((resolve, reject) => { + resolveProvider = resolve + rejectProvider = reject + }) + ) + } + store.getRepos.mockReturnValue([repo]) + getSshGitProviderMock.mockReturnValue(provider) + // oxlint-disable-next-line typescript/consistent-type-assertions -- SAFETY: request ids are minted by the renderer; any unique string is a valid id here. + const providerRequestId = `request-${targetId}` as ProviderRequestId + const pending = handlers['worktrees:listDetected'](ipcEvent, { + providerRequestId, + repoId: repo.id, + executionHostId: toSshExecutionHostId(targetId), + expectedAuthority: getSshProviderAuthority(targetId) + }) + await Promise.resolve() + expect(provider.listWorktrees).toHaveBeenCalledOnce() + expect(hasSshProviderContinuations(targetId)).toBe(true) + handlers['worktrees:cancelListDetected'](ipcEvent, { providerRequestId }) + await expect(pending).resolves.toMatchObject({ status: 'canceled' }) + expect(hasSshProviderContinuations(targetId)).toBe(true) + expect(hasSshProviderContinuations(`${targetId}-other`)).toBe(false) + if (outcome === 'resolve') { + resolveProvider([]) + } else { + rejectProvider(new Error('late provider failure')) + } + await vi.waitFor(() => expect(hasSshProviderContinuations(targetId)).toBe(false)) + expect(store.setWorktreeMeta).not.toHaveBeenCalled() + expect(store.removeWorktreeLineage).not.toHaveBeenCalled() + } + ) + + it('tracks legacy SSH listing until settlement without tracking a local folder listing', async () => { + const targetId = 'continuation-legacy' + const repo = { + id: 'repo-legacy', + path: '/remote/repo', + displayName: 'repo', + badgeColor: '#000', + addedAt: 0, + connectionId: targetId + } + const localRepo: Repo = { + id: 'repo-local', + path: '/local/repo', + displayName: 'local', + badgeColor: '#000', + addedAt: 0, + kind: 'folder' + } + let resolveProvider!: (value: GitWorktreeInfo[]) => void + const provider = { + listWorktrees: vi.fn( + () => + new Promise((resolve) => { + resolveProvider = resolve + }) + ) + } + store.getRepos.mockReturnValue([repo, localRepo]) + getSshGitProviderMock.mockReturnValue(provider) + const pending = handlers['worktrees:listDetected'](ipcEvent, { repoId: repo.id }) + expect(hasSshProviderContinuations(targetId)).toBe(true) + await expect( + handlers['worktrees:listDetected'](ipcEvent, { repoId: localRepo.id }) + ).resolves.toMatchObject({ repoId: localRepo.id, authoritative: true }) + expect(hasSshProviderContinuations(localRepo.id)).toBe(false) + expect(hasSshProviderContinuations(targetId)).toBe(true) + expect(provider.listWorktrees).toHaveBeenCalledOnce() + resolveProvider([]) + await expect(pending).resolves.toMatchObject({ repoId: repo.id, authoritative: true }) + expect(hasSshProviderContinuations(targetId)).toBe(false) + }) + it.each([ ['malformed', 'ssh:%'], ['contradictory', toSshExecutionHostId('target-b')] diff --git a/src/main/ipc/worktrees/listing/register-detected-worktree-handlers.ts b/src/main/ipc/worktrees/listing/register-detected-worktree-handlers.ts index 442a70985c8d..de466ff519e3 100644 --- a/src/main/ipc/worktrees/listing/register-detected-worktree-handlers.ts +++ b/src/main/ipc/worktrees/listing/register-detected-worktree-handlers.ts @@ -1,4 +1,5 @@ import { ipcMain } from 'electron' +import { runSshProviderContinuation } from '../../../ssh/ssh-provider-continuations' import type { DetectedWorktreeListResult } from '../../../../shared/worktree/types' import type { HostQualifiedDetectedWorktreeResult, @@ -72,16 +73,21 @@ export function registerDetectedWorktreeHandlers(context: WorktreeIpcContext): v }, DETECTED_WORKTREE_PROVIDER_TIMEOUT_MS) : undefined try { - const providerResult = listHostQualifiedDetectedWorktrees( - store, - args, - controller - ? { - signal: controller.signal, - status: () => (timedOut ? 'timed-out' : 'canceled') - } - : undefined - ) + const list = () => + listHostQualifiedDetectedWorktrees( + store, + args, + controller + ? { + signal: controller.signal, + status: () => (timedOut ? 'timed-out' : 'canceled') + } + : undefined + ) + const providerResult = + parsedHost?.kind === 'ssh' + ? runSshProviderContinuation(parsedHost.targetId, list) + : list() return abortedResult ? await Promise.race([providerResult, abortedResult]) : await providerResult @@ -106,17 +112,19 @@ export function registerDetectedWorktreeHandlers(context: WorktreeIpcContext): v const connectionId = getSshTargetIdForExecutionHost(getRepoExecutionHostId(repo)) const provider = connectionId ? getSshGitProvider(connectionId) : undefined const authority = connectionId ? { ...getSshProviderAuthority(connectionId) } : undefined - const result = await listDetectedWorktreesForCapturedRepo( - store, - repo, - () => - isCapturedRepoCurrent(store, repo) && - (!connectionId || - (getSshGitProvider(connectionId) === provider && - authority !== undefined && - isCurrentSshProviderAuthority(authority))), - provider - ) + const list = () => + listDetectedWorktreesForCapturedRepo( + store, + repo, + () => + isCapturedRepoCurrent(store, repo) && + (!connectionId || + (getSshGitProvider(connectionId) === provider && + authority !== undefined && + isCurrentSshProviderAuthority(authority))), + provider + ) + const result = await (connectionId ? runSshProviderContinuation(connectionId, list) : list()) return result && !('providerAbortStatus' in result) ? result : { diff --git a/src/main/ipc/worktrees/removal/register-worktree-removal-handlers.ts b/src/main/ipc/worktrees/removal/register-worktree-removal-handlers.ts index d91743b01850..d5e09d2f8f28 100644 --- a/src/main/ipc/worktrees/removal/register-worktree-removal-handlers.ts +++ b/src/main/ipc/worktrees/removal/register-worktree-removal-handlers.ts @@ -1,7 +1,10 @@ import { ipcMain } from 'electron' import { getLocalWorktreeCatalogVersion } from '../../../local-worktree-scan-generation' import type { RemoveWorktreeResult } from '../../../../shared/worktree/create-types' -import { getRepoExecutionHostId } from '../../../../shared/execution-host' +import { + getRepoExecutionHostId, + getSshTargetIdForExecutionHost +} from '../../../../shared/execution-host' import { withWorktreeSpan } from '../../../observability/instrumentation' import { parseWorktreeId } from '../../worktree-logic' import type { RemoveWorktreeArgs } from '../ipc-context-schemas' @@ -17,6 +20,7 @@ import { waitForPendingWorktreeRemoval } from '../../../worktree-background-removal' import { runSerializedWorktreeRemovalAcceptance } from '../../../worktree-removal-acceptance-queue' +import { runSshProviderContinuation } from '../../../ssh/ssh-provider-continuations' export function registerWorktreeRemovalHandlers(context: WorktreeIpcContext): void { const { store, options, worktreeRemovalsInFlight } = context @@ -36,6 +40,7 @@ export function registerWorktreeRemovalHandlers(context: WorktreeIpcContext): vo if (pending) { return { ...(await pending), catalogVersion: getLocalWorktreeCatalogVersion(repoId) } } + const targetId = getSshTargetIdForExecutionHost(removalHostId) ?? repo.connectionId const inFlightKey = getWorktreeRemovalInFlightKey(args.worktreeId, removalHostId) const optionsKey = getWorktreeRemovalOptionsKey(args) const inFlightRemoval = worktreeRemovalsInFlight.get(inFlightKey) @@ -48,11 +53,15 @@ export function registerWorktreeRemovalHandlers(context: WorktreeIpcContext): vo // Why: concurrent stale-toast/double-click/sidebar races can hit the same worktree; share the op so only one path touches Git and disk. const removal = withWorktreeSpan({ stage: 'remove', path: worktreePath }, async () => { + const execute = () => + executeWorktreeRemoval(context, args, repo, repoId, worktreePath, removalHostId) const accept = async (): Promise => // Why: another client's removal of this worktree may have been accepted during the wait. waitForPendingWorktreeRemoval(args.worktreeId, removalHostId) ? { removing: true } - : executeWorktreeRemoval(context, args, repo, repoId, worktreePath, removalHostId) + : targetId + ? runSshProviderContinuation(targetId, execute) + : execute() const accepted = await (repo.connectionId ? accept() : runSerializedWorktreeRemovalAcceptance(repo.path, accept)) diff --git a/src/main/ipc/worktrees/removal/worktree-removal-continuations.test.ts b/src/main/ipc/worktrees/removal/worktree-removal-continuations.test.ts new file mode 100644 index 000000000000..a57bd68d594f --- /dev/null +++ b/src/main/ipc/worktrees/removal/worktree-removal-continuations.test.ts @@ -0,0 +1,98 @@ +import { beforeEach, expect, it, vi } from 'vitest' +import { hasSshProviderContinuations } from '../../../ssh/ssh-provider-continuations' +import type { WorktreeIpcContext } from '../worktree-ipc-context' +import { registerWorktreeRemovalHandlers } from './register-worktree-removal-handlers' + +const mocks = vi.hoisted(() => ({ + handle: vi.fn(), + execute: vi.fn(), + resolveRepo: vi.fn() +})) +vi.mock('electron', () => ({ ipcMain: { handle: mocks.handle } })) +vi.mock('../../../observability/instrumentation', () => ({ + withWorktreeSpan: (_options: unknown, operation: () => unknown) => operation() +})) +vi.mock('../../worktree-logic', () => ({ + parseWorktreeId: () => ({ repoId: 'repo', worktreePath: '/work/feature' }) +})) +vi.mock('../repo-host-ownership', () => ({ resolveRepoForExecutionHost: mocks.resolveRepo })) +vi.mock('./execute-worktree-removal', () => ({ executeWorktreeRemoval: mocks.execute })) + +beforeEach(() => vi.clearAllMocks()) + +function setup(connectionId?: string, kind = 'git') { + mocks.resolveRepo.mockReturnValue({ id: 'repo', connectionId, kind }) + const worktreeRemovalsInFlight = new Map() + // oxlint-disable-next-line typescript/consistent-type-assertions -- SAFETY: removal reads only store and worktreeRemovalsInFlight from the context. + registerWorktreeRemovalHandlers({ + store: {}, + worktreeRemovalsInFlight + } as unknown as WorktreeIpcContext) + const handler = mocks.handle.mock.calls[0]![1] + return { + remove: (): Promise => handler({}, { worktreeId: 'repo:/work/feature' }), + worktreeRemovalsInFlight + } +} + +it.each(['git', 'folder'])( + 'tracks the full SSH %s removal before invocation and through cleanup', + async (kind) => { + const state = setup('target', kind) + const removed = Promise.withResolvers() + const cleaned = Promise.withResolvers() + mocks.execute.mockImplementation(async () => { + expect(hasSshProviderContinuations('target')).toBe(true) + await removed.promise + await cleaned.promise + return { success: true } + }) + const first = state.remove() + const joined = state.remove() + expect(mocks.execute).toHaveBeenCalledOnce() + expect(hasSshProviderContinuations('elsewhere')).toBe(false) + removed.resolve() + await Promise.resolve() + expect(hasSshProviderContinuations('target')).toBe(true) + cleaned.resolve() + await Promise.all([first, joined]) + expect(hasSshProviderContinuations('target')).toBe(false) + expect(state.worktreeRemovalsInFlight.size).toBe(0) + } +) + +it('retains a failed removal until its underlying operation rejects', async () => { + const state = setup('target') + const pending = Promise.withResolvers() + mocks.execute.mockReturnValue(pending.promise) + const result = state.remove() + expect(hasSshProviderContinuations('target')).toBe(true) + pending.reject(new Error('cleanup failed')) + await expect(result).rejects.toThrow('cleanup failed') + expect(hasSshProviderContinuations('target')).toBe(false) + expect(state.worktreeRemovalsInFlight.size).toBe(0) +}) + +it('does not register local removal as SSH work', async () => { + const state = setup() + const pending = Promise.withResolvers() + mocks.execute.mockReturnValue(pending.promise) + const result = state.remove() + await vi.waitFor(() => expect(mocks.execute).toHaveBeenCalledOnce()) + expect(hasSshProviderContinuations('target')).toBe(false) + pending.resolve({}) + await result +}) + +it('tracks a folder owner expressed only by its unified execution host', async () => { + const state = setup() + mocks.resolveRepo.mockReturnValue({ id: 'repo', kind: 'folder', executionHostId: 'ssh:target' }) + const pending = Promise.withResolvers() + mocks.execute.mockReturnValue(pending.promise) + const result = state.remove() + // Local acceptance is serialized per repo path first, so the continuation starts after a turn. + await vi.waitFor(() => expect(hasSshProviderContinuations('target')).toBe(true)) + pending.resolve({}) + await result + expect(hasSshProviderContinuations('target')).toBe(false) +}) diff --git a/src/main/providers/ssh-filesystem-dispatch.ts b/src/main/providers/ssh-filesystem-dispatch.ts index 04ef209bef15..b7171f75dc05 100644 --- a/src/main/providers/ssh-filesystem-dispatch.ts +++ b/src/main/providers/ssh-filesystem-dispatch.ts @@ -37,6 +37,16 @@ export function unregisterSshFilesystemProvider(connectionId: string): void { sshProviders.delete(connectionId) } +export function unregisterSshFilesystemProviderIfCurrent( + connectionId: string, + expected: IFilesystemProvider +): boolean { + if (sshProviders.get(connectionId) !== expected) { + return false + } + return sshProviders.delete(connectionId) +} + export function getSshFilesystemProvider(connectionId: string): IFilesystemProvider | undefined { return sshProviders.get(connectionId) } diff --git a/src/main/providers/ssh-git-dispatch.ts b/src/main/providers/ssh-git-dispatch.ts index 2a73d57f9c76..73f120bcb136 100644 --- a/src/main/providers/ssh-git-dispatch.ts +++ b/src/main/providers/ssh-git-dispatch.ts @@ -49,6 +49,21 @@ export function _getSshGitProviderGenerationCacheSize(): number { return sshProviderGenerations.size } +export function unregisterSshGitProviderIfCurrent( + connectionId: string, + expected: SshGitProvider, + expectedGeneration: number +): boolean { + if ( + sshProviders.get(connectionId) !== expected || + getSshGitProviderGeneration(connectionId) !== expectedGeneration + ) { + return false + } + unregisterSshGitProvider(connectionId) + return true +} + export function getSshGitProvider(connectionId: string): SshGitProvider | undefined { return sshProviders.get(connectionId) } diff --git a/src/main/ssh/ssh-connection-disconnect-drain.test.ts b/src/main/ssh/ssh-connection-disconnect-drain.test.ts index b5244ff5f327..4b04eae0e9a4 100644 --- a/src/main/ssh/ssh-connection-disconnect-drain.test.ts +++ b/src/main/ssh/ssh-connection-disconnect-drain.test.ts @@ -1,6 +1,8 @@ import { EventEmitter } from 'node:events' import { afterEach, beforeEach, expect, it, vi } from 'vitest' import { SshConnection } from './ssh-connection' +import { SshConnectionManager } from './ssh-connection-manager' +import { assertManagerTargetTransportsClosed } from './ssh-connection-manager-test-probes' import { createCallbacks, createResolvedConfig, createTarget } from './ssh-connection-test-fixtures' import { clientInstances, resetSshConnectionMocks, ssh2Mock } from './ssh-connection-test-harness' import { resolveWithSshG } from './ssh-config-parser' @@ -233,6 +235,30 @@ it('refuses a disconnected connection without claiming drain proof', async () => await expect(conn.disconnectAndDrain(new AbortController().signal)).rejects.toThrow() }) +it('retains ordinary manager transport debt after pool removal until physical close', async () => { + const manager = new SshConnectionManager(createCallbacks()) + const target = createTarget() + const conn = await manager.connect(target) + const client = conn.getClient()! + await manager.disconnect(target.id) + expect(manager.getConnection(target.id)).toBeUndefined() + expect(() => assertManagerTargetTransportsClosed(manager, target.id)).toThrow('closure_unproven') + client.emit('close') + expect(() => assertManagerTargetTransportsClosed(manager, target.id)).not.toThrow() +}) + +it('cleans failed ordinary startup without mistaking destroy for physical close', async () => { + const manager = new SshConnectionManager(createCallbacks()) + const target = createTarget() + ssh2Mock.connectSequence = [new Error('startup refused')] + ssh2Mock.destroyErrorMessage = 'socket still closing' + await expect(manager.connect(target)).rejects.toThrow('startup refused') + expect(manager.getConnection(target.id)).toBeUndefined() + expect(() => assertManagerTargetTransportsClosed(manager, target.id)).toThrow('closure_unproven') + clientInstances[0]!.emit('close') + expect(() => assertManagerTargetTransportsClosed(manager, target.id)).not.toThrow() +}) + it('closes owned system SSH but refuses unproven drain', async () => { vi.mocked(resolveWithSshG).mockResolvedValueOnce(createResolvedConfig()) const conn = new SshConnection(createTarget({ configHost: 'fdpass-host' }), createCallbacks()) diff --git a/src/main/ssh/ssh-connection-manager-closure.test.ts b/src/main/ssh/ssh-connection-manager-closure.test.ts new file mode 100644 index 000000000000..79e456b16533 --- /dev/null +++ b/src/main/ssh/ssh-connection-manager-closure.test.ts @@ -0,0 +1,108 @@ +import { + assertManagerTargetTransportsClosed, + managerTargetActive, + managerUnclosedTransportTargets +} from './ssh-connection-manager-test-probes' +import { beforeEach, expect, it, vi } from 'vitest' +import type { SshTarget } from '../../shared/ssh-types' + +type ClosureMockState = { + clients: { close: () => void; status: string }[] + connectError?: Error + disconnectError?: Error +} + +const state = vi.hoisted((): ClosureMockState => ({ clients: [] })) +vi.mock('./ssh-connection', () => ({ + SshConnection: class { + status = 'connected' + close = () => {} + constructor() { + state.clients.push(this) + } + subscribeTransportClosure(callback: () => void) { + this.close = callback + return () => {} + } + async connect() { + if (state.connectError) { + throw state.connectError + } + } + async disconnect() { + if (state.disconnectError) { + throw state.disconnectError + } + this.status = 'disconnected' + } + async disconnectAndDrain() { + await this.disconnect() + this.close() + } + getState() { + return { status: this.status } + } + setCallbacks() {} + } +})) +import { SshConnectionManager } from './ssh-connection-manager' + +const target: SshTarget = { + id: 'owned', + label: 'Owned', + host: 'example.test', + port: 22, + username: 'deploy' +} +beforeEach(() => { + state.clients.length = 0 + state.connectError = undefined + state.disconnectError = undefined +}) + +it('retains failed startup transport evidence until its disposed connection physically closes', async () => { + const manager = new SshConnectionManager({ onStateChange: vi.fn() }) + state.connectError = new Error('authentication rejected') + await expect(manager.connect(target)).rejects.toThrow('authentication rejected') + expect(state.clients[0].status).toBe('disconnected') + expect(manager.getConnection(target.id)).toBeUndefined() + expect(managerTargetActive(manager, target.id)).toBe(false) + expect(() => assertManagerTargetTransportsClosed(manager, target.id)).toThrow('closure_unproven') + state.clients[0].close() + expect(() => assertManagerTargetTransportsClosed(manager, target.id)).not.toThrow() +}) + +it('requires every allocation and ignores duplicate close notifications', async () => { + const manager = new SshConnectionManager({ onStateChange: vi.fn() }) + await manager.connect(target) + await manager.disconnect(target.id) + await manager.connect(target) + await manager.disconnect(target.id) + state.clients[0].close() + state.clients[0].close() + expect(() => assertManagerTargetTransportsClosed(manager, target.id)).toThrow('closure_unproven') + state.clients[1].close() + expect(() => assertManagerTargetTransportsClosed(manager, target.id)).not.toThrow() +}) + +it('does not let an old close remove a replacement or taint an unrelated target', async () => { + const manager = new SshConnectionManager({ onStateChange: vi.fn() }) + await manager.connect(target) + await manager.disconnect(target.id) + const replacement = await manager.connect(target) + state.clients[0].close() + expect(manager.getConnection(target.id)).toBe(replacement) + expect(() => assertManagerTargetTransportsClosed(manager, target.id)).toThrow() + expect(() => assertManagerTargetTransportsClosed(manager, 'unrelated')).not.toThrow() +}) + +it('releases bookkeeping after repeated observed closure', async () => { + const manager = new SshConnectionManager({ onStateChange: vi.fn() }) + for (let index = 0; index < 100; index++) { + await manager.connect(target) + await manager.disconnect(target.id) + state.clients[index].close() + expect(() => assertManagerTargetTransportsClosed(manager, target.id)).not.toThrow() + } + expect(managerUnclosedTransportTargets(manager)).toBe(0) +}) diff --git a/src/main/ssh/ssh-connection-manager-registry.test.ts b/src/main/ssh/ssh-connection-manager-registry.test.ts index e6a6bc24843d..a7b2b06e8059 100644 --- a/src/main/ssh/ssh-connection-manager-registry.test.ts +++ b/src/main/ssh/ssh-connection-manager-registry.test.ts @@ -88,4 +88,57 @@ describe('SshConnectionManager', () => { expect(mgr.getConnection('a')).toBeUndefined() expect(mgr.getConnection('b')).toBeUndefined() }) + + it('disconnectAll preserves excluded reset connections in the registry', async () => { + const mgr = new SshConnectionManager(createCallbacks()) + const retained = await mgr.connect(createTarget({ id: 'reset' })) + await mgr.connect(createTarget({ id: 'ordinary' })) + const disconnect = vi.spyOn(retained, 'disconnect') + await mgr.disconnectAll((id) => id !== 'reset') + expect(disconnect).not.toHaveBeenCalled() + expect(mgr.getConnection('reset')).toBe(retained) + expect(mgr.getConnection('ordinary')).toBeUndefined() + await mgr.disconnectAll() + }) + + it('checks disconnect admission immediately before each transport effect', async () => { + const mgr = new SshConnectionManager(createCallbacks()) + const first = await mgr.connect(createTarget({ id: 'first' })) + const second = await mgr.connect(createTarget({ id: 'second' })) + const disconnectSecond = vi.spyOn(second, 'disconnect') + let reserved = false + const originalDisconnect = first.disconnect.bind(first) + vi.spyOn(first, 'disconnect').mockImplementation(() => { + reserved = true + return originalDisconnect() + }) + await mgr.disconnectAll((id) => id !== 'second' || !reserved) + expect(disconnectSecond).not.toHaveBeenCalled() + expect(mgr.getConnection('second')).toBe(second) + await mgr.disconnectAll() + }) + + it('does not erase a replacement registered while an admitted disconnect settles', async () => { + const mgr = new SshConnectionManager(createCallbacks()) + const target = createTarget({ id: 'target' }) + const old = await mgr.connect(target) + const state = old.getState() + vi.spyOn(old, 'getState').mockReturnValue({ ...state, status: 'disconnected' }) + let finish!: () => void + vi.spyOn(old, 'disconnect') + .mockImplementationOnce( + () => + new Promise((resolve) => { + finish = resolve + }) + ) + .mockResolvedValue(undefined) + const drain = mgr.disconnectAll() + const replacement = await mgr.connect(target) + expect(replacement).not.toBe(old) + finish() + await drain + expect(mgr.getConnection('target')).toBe(replacement) + await mgr.disconnectAll() + }) }) diff --git a/src/main/ssh/ssh-connection-manager-test-probes.ts b/src/main/ssh/ssh-connection-manager-test-probes.ts new file mode 100644 index 000000000000..3c85210d5968 --- /dev/null +++ b/src/main/ssh/ssh-connection-manager-test-probes.ts @@ -0,0 +1,43 @@ +import type { SshConnectionManager } from './ssh-connection-manager' + +type ManagerBookkeeping = { + connections: Map + connectingTargets: Map + pendingTargetOperations: Map + unconfirmedTargetTeardowns: Set + unclosedTransportsByTarget: Map +} + +function bookkeeping(manager: SshConnectionManager): ManagerBookkeeping { + // oxlint-disable-next-line typescript/consistent-type-assertions -- SAFETY: these are SshConnectionManager's own private fields; tests read them until T3/T8 expose readers. + return manager as unknown as ManagerBookkeeping +} + +/** Test-side reading of the target activity T3's hasTargetActivity will expose. */ +export function managerTargetActive(manager: SshConnectionManager, targetId: string): boolean { + const state = bookkeeping(manager) + return ( + state.connections.has(targetId) || + state.connectingTargets.has(targetId) || + state.unconfirmedTargetTeardowns.has(targetId) || + state.pendingTargetOperations.has(targetId) + ) +} + +/** Test-side reading of the closure proof T8's assertTargetTransportsClosed will expose. */ +export function assertManagerTargetTransportsClosed( + manager: SshConnectionManager, + targetId: string +): void { + if ( + managerTargetActive(manager, targetId) || + bookkeeping(manager).unclosedTransportsByTarget.has(targetId) + ) { + throw new Error('ssh_target_transport_closure_unproven') + } +} + +/** How many targets still owe a physical transport close. */ +export function managerUnclosedTransportTargets(manager: SshConnectionManager): number { + return bookkeeping(manager).unclosedTransportsByTarget.size +} diff --git a/src/main/ssh/ssh-connection-manager.test.ts b/src/main/ssh/ssh-connection-manager.test.ts index b2de497a006d..b5c61574feca 100644 --- a/src/main/ssh/ssh-connection-manager.test.ts +++ b/src/main/ssh/ssh-connection-manager.test.ts @@ -1,3 +1,4 @@ +import { managerTargetActive } from './ssh-connection-manager-test-probes' import { beforeEach, describe, expect, it, vi } from 'vitest' import type { SshTarget } from '../../shared/ssh-types' @@ -12,6 +13,7 @@ const mockState = vi.hoisted(() => ({ vi.mock('./ssh-connection', () => ({ SshConnection: class MockSshConnection { + subscribeTransportClosure = vi.fn(() => () => {}) status: 'connecting' | 'connected' | 'disconnected' = 'connecting' connect = vi.fn(async () => { await (mockState.connectResults.shift() ?? Promise.resolve()) @@ -20,6 +22,7 @@ vi.mock('./ssh-connection', () => ({ disconnect = vi.fn(async () => { this.status = 'disconnected' }) + reconnect = vi.fn(async () => {}) constructor() { mockState.instances.push(this) @@ -45,6 +48,53 @@ const target = { } as SshTarget describe('SshConnectionManager', () => { + it('reports invalidated pending attempts even after disconnect removes the registration', async () => { + let reject!: (error: Error) => void + mockState.connectResults.push( + new Promise((_resolve, fail) => { + reject = fail + }) + ) + const manager = new SshConnectionManager({ onStateChange: vi.fn() }) + const pending = manager.connect(target) + const rejected = expect(pending).rejects.toThrow('cancelled') + await manager.disconnect(target.id) + expect(manager.getConnection(target.id)).toBeUndefined() + expect(managerTargetActive(manager, target.id)).toBe(true) + expect(managerTargetActive(manager, 'unrelated')).toBe(false) + reject(new Error('cancelled')) + await rejected + expect(managerTargetActive(manager, target.id)).toBe(false) + }) + + it('tracks detached exact-connection teardown until its promise settles', async () => { + const manager = new SshConnectionManager({ onStateChange: vi.fn() }) + const conn = await manager.connect(target) + await manager.disconnect(target.id) + let finish!: () => void + mockState.instances[0].disconnect.mockImplementationOnce( + () => + new Promise((resolve) => { + finish = resolve + }) + ) + const pending = manager.disconnectConnection(target.id, conn) + expect(managerTargetActive(manager, target.id)).toBe(true) + finish() + await pending + expect(managerTargetActive(manager, target.id)).toBe(false) + }) + + it('retains uncertainty after bulk teardown fails and removes the registration', async () => { + const manager = new SshConnectionManager({ onStateChange: vi.fn() }) + await manager.connect(target) + mockState.instances[0].disconnect.mockRejectedValueOnce(new Error('close unconfirmed')) + await manager.disconnectAll() + expect(manager.getConnection(target.id)).toBeUndefined() + expect(managerTargetActive(manager, target.id)).toBe(true) + expect(managerTargetActive(manager, 'unrelated')).toBe(false) + }) + beforeEach(() => { mockState.connectResults.length = 0 mockState.instances.length = 0 diff --git a/src/main/ssh/ssh-connection-manager.ts b/src/main/ssh/ssh-connection-manager.ts index e553bf906b59..aad5bdde5613 100644 --- a/src/main/ssh/ssh-connection-manager.ts +++ b/src/main/ssh/ssh-connection-manager.ts @@ -6,17 +6,46 @@ import { SshConnection, type SshConnectionCallbacks } from './ssh-connection' // 300-line oxlint max-lines threshold while preserving a clear // single-responsibility boundary (connection lifecycle vs. pool management). +// Bounds an owned drain so a host that never confirms close cannot hang its caller. +const OWNED_CONNECTION_DRAIN_TIMEOUT_MS = 10_000 + export class SshConnectionManager { private connections = new Map() private callbacks: SshConnectionCallbacks // Why: attempt identity lets disconnect unblock a replacement without the // cancelled attempt later clearing the replacement's state. private connectingTargets = new Map() + // Local operation and teardown bookkeeping; never a verdict on remote process exit. + private pendingTargetOperations = new Map() + private unconfirmedTargetTeardowns = new Set() + // Counts every connection a target allocated, retired pool entries included, until it closes. + private readonly unclosedTransportsByTarget = new Map() constructor(callbacks: SshConnectionCallbacks) { this.callbacks = callbacks } + private registerConnection(targetId: string, connection: SshConnection): void { + this.unclosedTransportsByTarget.set( + targetId, + (this.unclosedTransportsByTarget.get(targetId) ?? 0) + 1 + ) + let observed = false + connection.subscribeTransportClosure(() => { + if (observed) { + return + } + observed = true + const remaining = (this.unclosedTransportsByTarget.get(targetId) ?? 1) - 1 + if (remaining === 0) { + this.unclosedTransportsByTarget.delete(targetId) + } else { + this.unclosedTransportsByTarget.set(targetId, remaining) + } + }) + this.connections.set(targetId, connection) + } + setCallbacks(callbacks: SshConnectionCallbacks): void { this.callbacks = callbacks for (const connection of this.connections.values()) { @@ -25,6 +54,10 @@ export class SshConnectionManager { } async connect(target: SshTarget): Promise { + return this.trackTargetOperation(target.id, () => this.connectTarget(target)) + } + + private async connectTarget(target: SshTarget): Promise { const existing = this.connections.get(target.id) if (existing?.getState().status === 'connected') { return existing @@ -43,13 +76,16 @@ export class SshConnectionManager { } const conn = new SshConnection(target, this.callbacks) - this.connections.set(target.id, conn) + this.registerConnection(target.id, conn) try { await conn.connect() } catch (err) { - if (this.connections.get(target.id) === conn) { - this.connections.delete(target.id) + // Why: a failed startup can still hold sockets, so it is disconnected, not just forgotten. + try { + await this.disconnectConnection(target.id, conn) + } catch (cleanupError) { + throw new AggregateError([err, cleanupError], 'ssh_connection_startup_cleanup_failed') } throw err } @@ -63,6 +99,29 @@ export class SshConnectionManager { } async disconnect(targetId: string): Promise { + return this.trackTargetOperation(targetId, () => this.disconnectTarget(targetId), true) + } + + /** Drains only the registered connection for this target; local closure is not remote exit. */ + async disconnectAndDrain(targetId: string, signal: AbortSignal): Promise { + return this.trackTargetOperation( + targetId, + async () => { + this.connectingTargets.delete(targetId) + const conn = this.connections.get(targetId) + if (!conn) { + return + } + await conn.disconnectAndDrain(signal) + if (this.connections.get(targetId) === conn) { + this.connections.delete(targetId) + } + }, + true + ) + } + + private async disconnectTarget(targetId: string): Promise { // Why: disconnect invalidates the old attempt immediately so a reconnect // need not wait for the cancelled socket's late completion. this.connectingTargets.delete(targetId) @@ -81,11 +140,19 @@ export class SshConnectionManager { * Why: a cancelled connect whose transport opened late owns that exact connection — disconnecting * by target id would tear down the replacement's live transport instead. */ - async disconnectConnection(targetId: string, conn: SshConnection): Promise { - await conn.disconnect() - if (this.connections.get(targetId) === conn) { - this.connections.delete(targetId) - } + async disconnectConnection(targetId: string, conn: SshConnection, drain = false): Promise { + return this.trackTargetOperation( + targetId, + async () => { + await (drain + ? conn.disconnectAndDrain(AbortSignal.timeout(OWNED_CONNECTION_DRAIN_TIMEOUT_MS)) + : conn.disconnect()) + if (this.connections.get(targetId) === conn) { + this.connections.delete(targetId) + } + }, + true + ) } async reconnect(targetId: string): Promise { @@ -93,7 +160,33 @@ export class SshConnectionManager { if (!conn) { return } - await conn.reconnect() + await this.trackTargetOperation(targetId, () => conn.reconnect()) + } + + private async trackTargetOperation( + targetId: string, + operation: () => Promise, + retainFailure = false + ): Promise { + this.pendingTargetOperations.set( + targetId, + (this.pendingTargetOperations.get(targetId) ?? 0) + 1 + ) + try { + return await operation() + } catch (error) { + if (retainFailure) { + this.unconfirmedTargetTeardowns.add(targetId) + } + throw error + } finally { + const remaining = (this.pendingTargetOperations.get(targetId) ?? 1) - 1 + if (remaining === 0) { + this.pendingTargetOperations.delete(targetId) + } else { + this.pendingTargetOperations.set(targetId, remaining) + } + } } getConnection(targetId: string): SshConnection | undefined { @@ -112,9 +205,27 @@ export class SshConnectionManager { return states } - async disconnectAll(): Promise { - const disconnects = Array.from(this.connections.values()).map((c) => c.disconnect()) - await Promise.allSettled(disconnects) - this.connections.clear() + async disconnectAll(shouldDisconnect: (targetId: string) => boolean = () => true): Promise { + await Promise.allSettled( + Array.from(this.connections).map(async ([targetId, connection]) => { + if (!shouldDisconnect(targetId)) { + return + } + await this.trackTargetOperation( + targetId, + async () => { + try { + await connection.disconnect() + } finally { + // A later registration or an excluded target is not this drain's to remove. + if (this.connections.get(targetId) === connection) { + this.connections.delete(targetId) + } + } + }, + true + ) + }) + ) } } diff --git a/src/main/ssh/ssh-connection-store-test-fixture.ts b/src/main/ssh/ssh-connection-store-test-fixture.ts new file mode 100644 index 000000000000..1908112a8ed8 --- /dev/null +++ b/src/main/ssh/ssh-connection-store-test-fixture.ts @@ -0,0 +1,107 @@ +import { vi } from 'vitest' +import type { FolderWorkspace } from '../../shared/folder-workspace-types' +import type { Repo } from '../../shared/repo-types' +import type { RemovedSshTargetTombstone, SshTarget } from '../../shared/ssh-types' + +export function createMockStore() { + const targets: SshTarget[] = [] + const repos: Repo[] = [] + const folderWorkspaces: FolderWorkspace[] = [] + const projectGroups: { id: string; name: string; connectionId?: string | null }[] = [] + const leases: { + state: 'attached' | 'detached' | 'terminated' | 'expired' + ptyId?: string + worktreeId?: string + tabId?: string + leafId?: string + updatedAt?: number + }[] = [] + let deletedAliases: string[] = [] + const removedTombstones: RemovedSshTargetTombstone[] = [] + const reassignments: { oldTargetId: string; newTargetId: string }[] = [] + let generationCounter = 0 + + const dropTombstone = (oldTargetId: string) => { + const kept = removedTombstones.filter((t) => t.oldTargetId !== oldTargetId) + removedTombstones.length = 0 + removedTombstones.push(...kept) + } + + return { + allocateSshTargetGeneration: vi.fn(() => { + generationCounter += 1 + return generationCounter + }), + getSshTargets: vi.fn(() => [...targets]), + getSshTarget: vi.fn((id: string) => targets.find((t) => t.id === id)), + getRepos: vi.fn(() => [...repos]), + getProjectGroups: vi.fn(() => [...projectGroups]), + getFolderWorkspaces: vi.fn(() => [...folderWorkspaces]), + getSshRemotePtyLeases: vi.fn(() => [...leases]), + inspectOrcadMigrationSourceDependencies: vi.fn(() => ({ + totalCount: 0, + counts: { + automation: 0, + 'automation-run': 0, + 'mobile-tab-selection': 0, + 'retired-worktree-name': 0, + 'saved-port-forward': 0, + 'sparse-preset': 0, + 'terminal-lease': 0, + 'terminal-recovery': 0, + 'ui-routing': 0, + 'workspace-lineage': 0, + 'workspace-session': 0, + 'worktree-lineage': 0, + 'worktree-metadata': 0 + } + })), + addSshTarget: vi.fn((target: SshTarget) => targets.push(target)), + updateSshTarget: vi.fn((id: string, updates: Partial>) => { + const target = targets.find((t) => t.id === id) + if (!target) { + return null + } + Object.assign(target, updates) + return { ...target } + }), + removeSshTarget: vi.fn((id: string) => { + const idx = targets.findIndex((t) => t.id === id) + if (idx !== -1) { + targets.splice(idx, 1) + } + }), + getDeletedSshConfigAliases: vi.fn(() => [...deletedAliases]), + addDeletedSshConfigAlias: vi.fn((alias: string) => { + if (!deletedAliases.includes(alias)) { + deletedAliases.push(alias) + } + }), + removeDeletedSshConfigAlias: vi.fn((alias: string) => { + deletedAliases = deletedAliases.filter((entry) => entry !== alias) + }), + clearDeletedSshConfigAliases: vi.fn(() => { + deletedAliases = [] + }), + removedTombstones, + reassignments, + repos, + projectGroups, + folderWorkspaces, + leases, + getRemovedSshTargetTombstones: vi.fn(() => [...removedTombstones]), + addRemovedSshTargetTombstone: vi.fn((tombstone: RemovedSshTargetTombstone) => { + const filtered = removedTombstones.filter((t) => t.oldTargetId !== tombstone.oldTargetId) + removedTombstones.length = 0 + removedTombstones.push(...filtered, tombstone) + }), + removeRemovedSshTargetTombstone: vi.fn(dropTombstone), + // Nothing in this mock stores automations, so releasing always drops. + releaseRemovedSshTargetTombstone: vi.fn(dropTombstone), + reassignSshTargetId: vi.fn((oldTargetId: string, newTargetId: string) => { + reassignments.push({ oldTargetId, newTargetId }) + // Pretend one repo referenced the old id. + return ['repo-1'] + }) + } +} diff --git a/src/main/ssh/ssh-connection-store.test.ts b/src/main/ssh/ssh-connection-store.test.ts index 24845d2d3736..fc2b39e88af9 100644 --- a/src/main/ssh/ssh-connection-store.test.ts +++ b/src/main/ssh/ssh-connection-store.test.ts @@ -1,6 +1,7 @@ import { describe, expect, it, vi, beforeEach } from 'vitest' import { SshConnectionStore } from './ssh-connection-store' -import type { RemovedSshTargetTombstone, SshTarget } from '../../shared/ssh-types' +import { createMockStore } from './ssh-connection-store-test-fixture' +import type { SshTarget } from '../../shared/ssh-types' const { loadUserSshConfigMock, sshConfigHostsToTargetsMock } = vi.hoisted(() => ({ loadUserSshConfigMock: vi.fn(), @@ -12,72 +13,6 @@ vi.mock('./ssh-config-parser', () => ({ sshConfigHostsToTargets: sshConfigHostsToTargetsMock })) -function createMockStore() { - const targets: SshTarget[] = [] - let deletedAliases: string[] = [] - const removedTombstones: RemovedSshTargetTombstone[] = [] - const reassignments: { oldTargetId: string; newTargetId: string }[] = [] - let generationCounter = 0 - - const dropTombstone = (oldTargetId: string) => { - const kept = removedTombstones.filter((t) => t.oldTargetId !== oldTargetId) - removedTombstones.length = 0 - removedTombstones.push(...kept) - } - - return { - allocateSshTargetGeneration: vi.fn(() => { - generationCounter += 1 - return generationCounter - }), - getSshTargets: vi.fn(() => [...targets]), - getSshTarget: vi.fn((id: string) => targets.find((t) => t.id === id)), - addSshTarget: vi.fn((target: SshTarget) => targets.push(target)), - updateSshTarget: vi.fn((id: string, updates: Partial>) => { - const target = targets.find((t) => t.id === id) - if (!target) { - return null - } - Object.assign(target, updates) - return { ...target } - }), - removeSshTarget: vi.fn((id: string) => { - const idx = targets.findIndex((t) => t.id === id) - if (idx !== -1) { - targets.splice(idx, 1) - } - }), - getDeletedSshConfigAliases: vi.fn(() => [...deletedAliases]), - addDeletedSshConfigAlias: vi.fn((alias: string) => { - if (!deletedAliases.includes(alias)) { - deletedAliases.push(alias) - } - }), - removeDeletedSshConfigAlias: vi.fn((alias: string) => { - deletedAliases = deletedAliases.filter((entry) => entry !== alias) - }), - clearDeletedSshConfigAliases: vi.fn(() => { - deletedAliases = [] - }), - removedTombstones, - reassignments, - getRemovedSshTargetTombstones: vi.fn(() => [...removedTombstones]), - addRemovedSshTargetTombstone: vi.fn((tombstone: RemovedSshTargetTombstone) => { - const filtered = removedTombstones.filter((t) => t.oldTargetId !== tombstone.oldTargetId) - removedTombstones.length = 0 - removedTombstones.push(...filtered, tombstone) - }), - removeRemovedSshTargetTombstone: vi.fn(dropTombstone), - // Nothing in this mock stores automations, so releasing always drops. - releaseRemovedSshTargetTombstone: vi.fn(dropTombstone), - reassignSshTargetId: vi.fn((oldTargetId: string, newTargetId: string) => { - reassignments.push({ oldTargetId, newTargetId }) - // Pretend one repo referenced the old id. - return ['repo-1'] - }) - } -} - describe('SshConnectionStore', () => { let mockStore: ReturnType let sshStore: SshConnectionStore diff --git a/src/main/ssh/ssh-provider-continuations.ts b/src/main/ssh/ssh-provider-continuations.ts new file mode 100644 index 000000000000..418c1d896767 --- /dev/null +++ b/src/main/ssh/ssh-provider-continuations.ts @@ -0,0 +1,24 @@ +const pendingByTarget = new Map>() + +/** Tracks local continuation settlement, never host execution or cancellation acknowledgment. */ +export async function runSshProviderContinuation( + targetId: string, + operation: () => Promise +): Promise { + const pending = pendingByTarget.get(targetId) ?? new Set() + const token = {} + pendingByTarget.set(targetId, pending) + pending.add(token) + try { + return await operation() + } finally { + pending.delete(token) + if (pending.size === 0 && pendingByTarget.get(targetId) === pending) { + pendingByTarget.delete(targetId) + } + } +} + +export function hasSshProviderContinuations(targetId: string): boolean { + return (pendingByTarget.get(targetId)?.size ?? 0) > 0 +} diff --git a/src/main/ssh/ssh-target-registry.ts b/src/main/ssh/ssh-target-registry.ts index 9959e8072a6d..aef003025ccc 100644 --- a/src/main/ssh/ssh-target-registry.ts +++ b/src/main/ssh/ssh-target-registry.ts @@ -90,6 +90,7 @@ export function getActiveMultiplexer(connectionId: string): SshChannelMultiplexe } let registeredGetSshConnectionManager: (() => SshConnectionManager | null) | null = null +let registeredHasDirectSshAuthority: ((targetId: string) => boolean) | null = null export function setSshConnectionManagerResolver( resolve: (() => SshConnectionManager | null) | null @@ -108,3 +109,14 @@ export function setSshConnectionManagerResolver( export function getSshConnectionManager(): SshConnectionManager | null { return registeredGetSshConnectionManager?.() ?? null } + +export function setDirectSshAuthorityResolver( + resolve: ((targetId: string) => boolean) | null +): void { + registeredHasDirectSshAuthority = resolve +} + +/** Whether this process currently holds the target's relay session; false when unregistered. */ +export function hasRegisteredDirectSshAuthority(targetId: string): boolean { + return registeredHasDirectSshAuthority?.(targetId) ?? false +} From b6f9c64a39812fa2b738f78435c5414d9eaac77b Mon Sep 17 00:00:00 2001 From: m4air Date: Thu, 1 Oct 2026 10:40:20 -0700 Subject: [PATCH 2/2] test(ssh): name the filesystem continuation handler payload type --- src/main/ipc/filesystem-provider-continuations.test.ts | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/main/ipc/filesystem-provider-continuations.test.ts b/src/main/ipc/filesystem-provider-continuations.test.ts index be7a4b25e086..c585d7ac6f67 100644 --- a/src/main/ipc/filesystem-provider-continuations.test.ts +++ b/src/main/ipc/filesystem-provider-continuations.test.ts @@ -42,7 +42,12 @@ vi.mock('./filesystem-import-ssh-directory', () => ({ })) const targetId = 'filesystem-continuation-target' -const handlers = new Map Promise>() +// The IPC payload fields these handlers read; each test passes the subset its channel needs. +type FilesystemHandlerArgs = Record +const handlers = new Map< + string, + (_event: unknown, args: FilesystemHandlerArgs) => Promise +>() const fileStat = { isFile: () => true, isDirectory: () => false, isSymbolicLink: () => false } beforeEach(() => {