diff --git a/src/main/index.ts b/src/main/index.ts index 7ffd94df50e..4375f5d3762 100644 --- a/src/main/index.ts +++ b/src/main/index.ts @@ -250,7 +250,7 @@ import { buildHeadlessAutomationWorktreeCreateArgs } from './automations/headles import { AgentAwakeService } from './agent-awake-service' import { registerSystemResumeBroadcast } from './system-resume-broadcast' import { settleTeardownWithinDeadline } from './quit-teardown-deadline' -import { QuitTeardownStartGate } from './quit-teardown-start-gate' +import { quitTeardownStartGate } from './quit-teardown-start-gate' import { detachAllSshSessionsForShutdown } from './ipc/ssh' import { PluginService } from './plugins/plugin-service' import { PluginKillListService } from './plugins/plugin-kill-list-service' @@ -2943,7 +2943,6 @@ app.on('before-quit', () => { // Why: will-quit fires twice — first pass preventDefaults and runs teardown; second pass exits. let daemonDisconnectDone = false -const quitTeardownStartGate = new QuitTeardownStartGate() app.on('will-quit', (e) => { // Why return instead of re-running teardown: the second pass is Electron re-firing after // our own app.quit(), so every step below already ran and every durable write already diff --git a/src/main/ipc/ssh.test.ts b/src/main/ipc/ssh.test.ts index 62f1ef26985..bea2164905d 100644 --- a/src/main/ipc/ssh.test.ts +++ b/src/main/ipc/ssh.test.ts @@ -282,6 +282,7 @@ import { } from './pty' import { assertSshMutationExpectation } from '../ssh/ssh-connection-generation' import { getSshPtyConsumerRecovery } from '../ssh/ssh-pty-consumer-recovery' +import { quitTeardownStartGate } from '../quit-teardown-start-gate' describe('SSH IPC handlers', () => { const relayBuildId = '0.1.0+ipc-test' @@ -592,7 +593,7 @@ describe('SSH IPC handlers', () => { expect(mockConnectionManager.disconnectAll).toHaveBeenCalled() }) - it('drains the replacement session a paused connect publishes after shutdown began', async () => { + it('refuses the replacement session a paused connect would publish after shutdown began', async () => { const target: SshTarget = { id: 'ssh-1', label: 'Server', @@ -627,18 +628,21 @@ describe('SSH IPC handlers', () => { expect(mockStore.markSshRemotePtyLeasesAsync).toHaveBeenCalledWith('ssh-1', 'detached') mockConnectionManager.disconnectAll.mockClear() + // Why latch the gate here: the committed quit path owns it, so the drain alone must not fence. + quitTeardownStartGate.tryStart({ preventDefault() {} }) const shutdown = detachAllSshSessionsForShutdown() releaseDetach() - await expect(replacement).rejects.toThrow() + await expect(replacement).rejects.toThrow('closed for app shutdown') await shutdown - // Why two detaches: the old session's, plus the replacement the paused connect published after the - // first drain had already snapshotted activeSessions. + // Why one detach: the fence sits at the publication point, so the resumed connect never registers a + // replacement session — only the old session it had already torn down was detached. const detaches = mockStore.markSshRemotePtyLeasesAsync.mock.calls.filter( (call) => call[1] === 'detached' ) - expect(detaches).toHaveLength(2) + expect(detaches).toHaveLength(1) + // Why twice: once for the drain's snapshot, once after joining the connect that was still in flight. expect(mockConnectionManager.disconnectAll).toHaveBeenCalledTimes(2) expect(getActiveMultiplexer('ssh-1')).toBeUndefined() await expect(handlers.get('ssh:connect')!(null, { targetId: 'ssh-1' })).rejects.toThrow( @@ -646,6 +650,53 @@ describe('SSH IPC handlers', () => { ) }) + it('joins an in-flight test-connection probe before the final shutdown disconnect', async () => { + const target: SshTarget = { + id: 'ssh-probe', + label: 'Server', + host: 'example.com', + port: 22, + username: 'deploy' + } + mockSshStore.getTarget.mockReturnValue(target) + mockConnectionManager.getState.mockReturnValue({ + targetId: 'ssh-probe', + status: 'connected', + error: null, + reconnectAttempt: 0 + }) + mockConnectionManager.disconnectAll.mockResolvedValue(undefined) + mockConnectionManager.disconnect.mockClear().mockResolvedValue(undefined) + + // Why: a probe holds a transport no session owns, so shutdown has to wait for it to hand it back. + const probeState = { + targetId: 'ssh-probe', + status: 'connected', + error: null, + reconnectAttempt: 0 + } + let openProbeTransport = (): void => {} + mockConnectionManager.connect.mockImplementationOnce( + () => + new Promise((resolve) => { + openProbeTransport = () => resolve({ getState: () => probeState }) + }) + ) + const probe = handlers.get('ssh:testConnection')!(null, { targetId: 'ssh-probe' }) + for (let tick = 0; tick < 5; tick++) { + await Promise.resolve() + } + expect(mockConnectionManager.connect).toHaveBeenCalledWith(target) + + quitTeardownStartGate.tryStart({ preventDefault() {} }) + const shutdown = detachAllSshSessionsForShutdown() + openProbeTransport() + await shutdown + + expect(await probe).toMatchObject({ success: true }) + expect(mockConnectionManager.disconnect).toHaveBeenCalledWith('ssh-probe') + }) + it('ssh:importConfig returns imported targets', async () => { const imported: SshTarget[] = [ { id: 'ssh-imp', label: 'staging', host: 'staging.com', port: 22, username: '' } diff --git a/src/main/ipc/ssh.ts b/src/main/ipc/ssh.ts index 75dfa1f7b9f..174edf024b5 100644 --- a/src/main/ipc/ssh.ts +++ b/src/main/ipc/ssh.ts @@ -25,6 +25,7 @@ import type { DirectSshAuthority } from '../../shared/ssh-types' import { SSH_TERMINATE_RECONNECT_REQUIRED } from '../../shared/constants' +import { quitTeardownStartGate } from '../quit-teardown-start-gate' import { isRuntimeOwnedSshTargetId } from '../../shared/execution-host' import { isAuthError } from '../ssh/ssh-connection-utils' import { createCancelledConnectAttemptError } from '../ssh/ssh-connect-attempt-cancellation' @@ -306,12 +307,11 @@ type ConnectAttempt = { const connectInFlight = new Map() const pendingTransportReconnects = new Set() -// Why: once the committed quit path starts draining SSH, a connect that has not yet published its -// session must not start one — the drain has already snapshotted what it will tear down. -let sshShutdownFenced = false - +// Why the quit gate rather than a local latch: "the committed quit has begun" already has an owner, +// and a private copy could be set by something that is not actually quitting — leaving SSH connects +// refused for the rest of the process lifetime. function assertSshConnectsNotFenced(): void { - if (sshShutdownFenced) { + if (quitTeardownStartGate.hasStarted()) { throw new Error('SSH connects are closed for app shutdown') } } @@ -332,6 +332,7 @@ const resetRelayInFlight = new Map>() // Why: ssh:testConnection connects then disconnects; suppressing broadcasts during the test avoids worktree cards flashing connected → disconnected. const testingTargets = new Set() +const testConnectionProbes = new Set>() // Why: without backoff, a relay channel that keeps dying reconnects as fast as the network allows, hammering local + remote sshd; track attempts and back off to end the loop recoverably. type RelayLostBackoffState = { @@ -1143,6 +1144,11 @@ export function registerSshHandlers( } } + // Why here and not only at entry: this is the publication point, and it is the last statement + // before the transport opens. Checking it in the same synchronous block as activeSessions.set + // means a connect either registers before the shutdown drain snapshots, or registers never and + // owns nothing to clean up. + assertSshConnectsNotFenced() // Why: create the session early so onStateChange sees it in 'deploying' and skips reconnect logic. const session = new SshRelaySession( targetId, @@ -1407,19 +1413,26 @@ export function registerSshHandlers( } testingTargets.add(args.targetId) - try { + // 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 () => { // 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 { success: true, state } + return state + })() + testConnectionProbes.add(probe) + try { + return { success: true, state: await probe } } catch (err) { return { 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) @@ -1545,13 +1558,13 @@ async function settleWithinMs(work: readonly Promise[], timeoutMs: numb ]) } +// 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 async function detachAllSshSessionsForShutdown(): Promise { - // Why synchronous, before the first await: connectTarget registers its attempt and checks the fence - // in one uninterrupted block, so fencing and snapshotting here leaves no gap for an unjoined connect. - sshShutdownFenced = true const inFlightWork: Promise[] = [ ...[...connectInFlight.values()].map((attempt) => attempt.promise), - ...resetRelayInFlight.values() + ...resetRelayInFlight.values(), + ...testConnectionProbes ] for (const targetId of Array.from(connectInFlight.keys())) { invalidateConnectAttempt(targetId) @@ -1608,8 +1621,9 @@ export async function resetSshHandlerStateForTests(): Promise { resetSshProviderAuthorities() resetRelayInFlight.clear() testingTargets.clear() + testConnectionProbes.clear() credentialRequestedForTarget.clear() - sshShutdownFenced = false + quitTeardownStartGate.resetForTests() await connectionManager?.disconnectAll() portForwardManager?.dispose() diff --git a/src/main/quit-teardown-start-gate.ts b/src/main/quit-teardown-start-gate.ts index be970495a4b..a96b7b3f38b 100644 --- a/src/main/quit-teardown-start-gate.ts +++ b/src/main/quit-teardown-start-gate.ts @@ -9,4 +9,16 @@ export class QuitTeardownStartGate { this.started = true return true } + + hasStarted(): boolean { + return this.started + } + + resetForTests(): void { + this.started = false + } } + +// Why shared, not per-caller: subsystems fence new work on "the committed quit has begun", and a +// second instance would answer no while the real teardown is already draining what it snapshotted. +export const quitTeardownStartGate = new QuitTeardownStartGate() diff --git a/src/main/ssh/ssh-relay-session-recovery-durability.test.ts b/src/main/ssh/ssh-relay-session-recovery-durability.test.ts index 9a28d491512..be5cc07f80a 100644 --- a/src/main/ssh/ssh-relay-session-recovery-durability.test.ts +++ b/src/main/ssh/ssh-relay-session-recovery-durability.test.ts @@ -273,6 +273,46 @@ describe('SshRelaySession consumer recovery durability', () => { expect(openConsumerSessionMock).toHaveBeenCalledTimes(1) }) + it('leaves recovery state alone when a newer owner already claimed the target record', async () => { + const targetId = 'target-stale-owner-loser' + const { mockConn, mockStore, mockPortForward, getMainWindow } = createMockDeps() + vi.mocked(mockStore.getSshPtyConsumerRecovery).mockReturnValue({ + targetId, + clientInstanceId: 'persisted-client', + serverBuildId: 'test-relay-build', + clientGeneration: 1, + ownerGeneration: 1, + ownerLease: 'stale-owner' + }) + const winner = { + mode: 'negotiated' as const, + clientInstanceId: 'persisted-client', + clientGeneration: 2, + ownerGeneration: 5, + ownerLease: 'winner-owner' + } + openConsumerSessionMock.mockImplementationOnce(() => { + // Why inside the rejection: the record is target-scoped, so the winner can land while this + // attempt is still unwinding its own resume. + const record = getSshPtyConsumerRecovery(targetId)! + record.owner = winner + record.checkpointsByAppPtyId.set('pty-1', { + id: 'pty-1' + } as unknown as never) + return Promise.reject(createMismatchedOwnerRecoveryError()) + }) + openConsumerSessionMock.mockRejectedValueOnce(new Error('fresh open failed')) + const session = new SshRelaySession(targetId, getMainWindow, mockStore, mockPortForward) + + await expect(session.establish(mockConn)).rejects.toThrow('fresh open failed') + + const record = getSshPtyConsumerRecovery(targetId)! + expect(record.owner).toBe(winner) + expect(record.checkpointsByAppPtyId.has('pty-1')).toBe(true) + expect(mockStore.removeSshPtyConsumerRecovery).not.toHaveBeenCalled() + session.dispose() + }) + it('does not remember a consumer opened after establish was disposed', async () => { const { mockConn, mockStore, mockPortForward, getMainWindow } = createMockDeps() let signalOpenStarted!: () => void diff --git a/src/main/ssh/ssh-relay-session.ts b/src/main/ssh/ssh-relay-session.ts index e6a821f74a8..ea293e5b8ca 100644 --- a/src/main/ssh/ssh-relay-session.ts +++ b/src/main/ssh/ssh-relay-session.ts @@ -1054,26 +1054,36 @@ export class SshRelaySession { throw error } const recovery = getSshPtyConsumerRecovery(this.targetId) - if (recovery) { - delete recovery.owner - recovery.checkpointsByAppPtyId.clear() - for (const [ptyId, migration] of recovery.modelMigrationsByAppPtyId) { - recovery.modelMigrationsByAppPtyId.set( - ptyId, - migration.then(() => - Object.freeze({ - status: 'checkpoint-unavailable' as const, - reason: 'completion-failed' as const - }) + // Why identity-guarded: the record is target-scoped and its clientInstanceId is shared by every + // session for that target, so only a record still describing the owner this attempt tried to + // resume is ours to drop — otherwise a loser wipes the winner's checkpoints. + const ownsRecoveryRecord = + ownsAttempt() && + (!recovery?.owner || + (recovery.owner.ownerGeneration === previousOwner.ownerGeneration && + recovery.owner.ownerLease === previousOwner.ownerLease)) + if (ownsRecoveryRecord) { + if (recovery) { + delete recovery.owner + recovery.checkpointsByAppPtyId.clear() + for (const [ptyId, migration] of recovery.modelMigrationsByAppPtyId) { + recovery.modelMigrationsByAppPtyId.set( + ptyId, + migration.then(() => + Object.freeze({ + status: 'checkpoint-unavailable' as const, + reason: 'completion-failed' as const + }) + ) ) - ) + } } + await removeSshPtyConsumerOwnerRecovery( + this.targetId, + this.ptyConsumerClientInstanceId, + this.store + ) } - await removeSshPtyConsumerOwnerRecovery( - this.targetId, - this.ptyConsumerClientInstanceId, - this.store - ) if (!ownsAttempt()) { throw new Error('Session disposed during establish') } diff --git a/src/relay/relay-pty-consumer-owner-displacement.test.ts b/src/relay/relay-pty-consumer-owner-displacement.test.ts index ede575333d5..e91d41d6852 100644 --- a/src/relay/relay-pty-consumer-owner-displacement.test.ts +++ b/src/relay/relay-pty-consumer-owner-displacement.test.ts @@ -127,6 +127,68 @@ describe('relay PTY consumer owner displacement', () => { expect(rotation.recovery.map((span) => span.data)).toEqual(['tail']) }) + it('retains the incumbent delivery once when its socket closes mid-publication', async () => { + const incumbentWrites: Buffer[] = [] + dispatcher = new RelayDispatcher( + (data, settle) => { + incumbentWrites.push(Buffer.from(data)) + settle({ ok: true }) + return true + }, + { supportsWriteCallback: true }, + endpointIdentity + ) + const adapter = new SshPtyConsumerSessionAdapter(dispatcher, 'build-a') + + dispatcher.feed(requestFrame(1, 'pty.openClient', ownerHelloParams())) + await flushRequests() + const incumbentGrant = response(incumbentWrites, 1)!.result as Record + const incumbentDelivery = adapter.openDelivery(1, 'pty-1', 'incarnation-1')! + adapter.appendSource(incumbentDelivery, { + spanId: 'span-1', + data: 'tail', + displayStart: 0, + displayEnd: 4, + splittable: true, + transform: { transformed: false, rawLengthSu: 4, scalarSafe: true } + }) + + const reconnectWrites: Buffer[] = [] + let grantSettlement: ((result: SinkWriteSettlement) => void) | undefined + const reconnectClientId = dispatcher.attachClient( + (data, settle) => { + reconnectWrites.push(Buffer.from(data)) + grantSettlement = settle + return true + }, + { supportsWriteCallback: true }, + endpointIdentity + ) + dispatcher.feedClient( + reconnectClientId, + requestFrame( + 2, + 'pty.openClient', + ownerHelloParams({ + ownerGeneration: incumbentGrant.ownerGeneration, + ownerLease: incumbentGrant.ownerLease + }) + ) + ) + await flushRequests() + + // Why this ordering: the incumbent's socket close and the replacement's grant write are independent + // events, so detach-then-commit has to leave exactly one retention, not two competing ones. + dispatcher.invalidateClient() + grantSettlement!({ ok: true }) + + expect(adapter.deliveryMode(reconnectClientId)).toBe('source-owner') + expect(adapter.getDebugSnapshot()).toMatchObject({ deliveryTokens: 1, graceTimers: 1 }) + const rotation = adapter.rotateDelivery(incumbentDelivery, reconnectClientId, 0) + expect(rotation.identity).toMatchObject({ ownerGeneration: 2, clientGeneration: 2 }) + expect(adapter.getDebugSnapshot()).toMatchObject({ graceTimers: 0 }) + }) + it('keeps the incumbent owner when the replacement grant fails to publish', async () => { const incumbentWrites: Buffer[] = [] dispatcher = new RelayDispatcher(