fix(ssh): handle owner displacement and graceful shutdown

Make QuitTeardownStartGate a shared singleton so SSH connects use the
same shutdown fence as the main quit path. Track test-connection probes
to ensure they complete before final teardown. Guard owner displacement
to prevent stale owners from clearing recovery state claimed by newer
owners.

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

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: '' }

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<string, ConnectAttempt>()
 const pendingTransportReconnects = new Set<string>()
 
-// 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<string, Promise<void>>()
 
 // Why: ssh:testConnection connects then disconnects; suppressing broadcasts during the test avoids worktree cards flashing connected → disconnected.
 const testingTargets = new Set<string>()
+const testConnectionProbes = new Set<Promise<unknown>>()
 
 // 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<unknown>[], 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<void> {
-  // 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<unknown>[] = [
     ...[...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<void> {
   resetSshProviderAuthorities()
   resetRelayInFlight.clear()
   testingTargets.clear()
+  testConnectionProbes.clear()
   credentialRequestedForTarget.clear()
-  sshShutdownFenced = false
+  quitTeardownStartGate.resetForTests()
 
   await connectionManager?.disconnectAll()
   portForwardManager?.dispose()

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

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

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')
       }

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<string, unknown>
+    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(