Fix SSH sessions after host sleep

src/main/ipc/ssh.test.ts

@@ -5,6 +5,8 @@ import { describe, expect, it, vi, beforeEach } from 'vitest'
 
 const {
   handleMock,
+  powerMonitorOffMock,
+  powerMonitorOnMock,
   mockSshStore,
   mockConnectionManager,
   mockDeployAndLaunchRelay,
@@ -16,6 +18,8 @@ const {
   mockPortForwardManager
 } = vi.hoisted(() => ({
   handleMock: vi.fn(),
+  powerMonitorOffMock: vi.fn(),
+  powerMonitorOnMock: vi.fn(),
   mockSshStore: {
     listTargets: vi.fn().mockReturnValue([]),
     getTarget: vi.fn(),
@@ -27,6 +31,7 @@ const {
   mockConnectionManager: {
     connect: vi.fn(),
     disconnect: vi.fn(),
+    reconnect: vi.fn(),
     getConnection: vi.fn(),
     getState: vi.fn(),
     disconnectAll: vi.fn()
@@ -66,6 +71,10 @@ vi.mock('electron', () => ({
     once: vi.fn(),
     removeHandler: vi.fn(),
     removeAllListeners: vi.fn()
+  },
+  powerMonitor: {
+    on: powerMonitorOnMock,
+    off: powerMonitorOffMock
   }
 }))
 
@@ -158,7 +167,7 @@ vi.mock('../ssh/ssh-port-forward', () => ({
 }))
 
 import { registerSshHandlers } from './ssh'
-import type { SshTarget } from '../../shared/ssh-types'
+import { SSH_RELAY_CONFIGURE_GRACE_TIME_METHOD, type SshTarget } from '../../shared/ssh-types'
 import {
   clearProviderPtyState,
   deletePtyOwnership,
@@ -200,6 +209,7 @@ describe('SSH IPC handlers', () => {
 
     mockConnectionManager.connect.mockReset()
     mockConnectionManager.disconnect.mockReset()
+    mockConnectionManager.reconnect.mockReset()
     mockConnectionManager.getConnection.mockReset()
     mockConnectionManager.getState.mockReset()
     mockConnectionManager.disconnectAll.mockReset()
@@ -222,6 +232,8 @@ describe('SSH IPC handlers', () => {
     mockPortForwardManager.listForwards.mockReset().mockReturnValue([])
     mockPortForwardManager.removeAllForwards.mockReset()
     mockPortForwardManager.dispose.mockReset()
+    powerMonitorOnMock.mockReset()
+    powerMonitorOffMock.mockReset()
     vi.mocked(getSshPtyProvider).mockReset()
     vi.mocked(getPtyIdsForConnection).mockReset().mockReturnValue([])
     vi.mocked(clearProviderPtyState).mockReset()
@@ -752,6 +764,69 @@ describe('SSH IPC handlers', () => {
     expect(mockConnectionManager.disconnect).toHaveBeenCalledWith('ssh-1')
   })
 
+  it('forces active SSH sessions to reconnect when the system resumes from sleep', async () => {
+    const target: SshTarget = {
+      id: 'ssh-1',
+      label: 'Server',
+      host: 'example.com',
+      port: 22,
+      username: 'deploy'
+    }
+    const conn = {}
+    mockSshStore.getTarget.mockReturnValue(target)
+    mockConnectionManager.connect.mockResolvedValue(conn)
+    mockConnectionManager.getConnection.mockReturnValue(conn)
+    mockConnectionManager.getState.mockReturnValue({
+      targetId: 'ssh-1',
+      status: 'connected',
+      error: null,
+      reconnectAttempt: 0
+    })
+
+    await handlers.get('ssh:connect')!(null, { targetId: 'ssh-1' })
+
+    const resumeListener = powerMonitorOnMock.mock.calls.find(([event]) => event === 'resume')?.[1]
+    expect(resumeListener).toBeTypeOf('function')
+
+    resumeListener()
+
+    expect(mockConnectionManager.reconnect).toHaveBeenCalledWith('ssh-1')
+  })
+
+  it('extends active relay grace while the system is suspending', async () => {
+    const target: SshTarget = {
+      id: 'ssh-1',
+      label: 'Server',
+      host: 'example.com',
+      port: 22,
+      username: 'deploy'
+    }
+    const conn = {}
+    mockSshStore.getTarget.mockReturnValue(target)
+    mockConnectionManager.connect.mockResolvedValue(conn)
+    mockConnectionManager.getConnection.mockReturnValue(conn)
+    mockConnectionManager.getState.mockReturnValue({
+      targetId: 'ssh-1',
+      status: 'connected',
+      error: null,
+      reconnectAttempt: 0
+    })
+
+    await handlers.get('ssh:connect')!(null, { targetId: 'ssh-1' })
+    mockMux.notify.mockClear()
+
+    const suspendListener = powerMonitorOnMock.mock.calls.find(
+      ([event]) => event === 'suspend'
+    )?.[1]
+    expect(suspendListener).toBeTypeOf('function')
+
+    suspendListener()
+
+    expect(mockMux.notify).toHaveBeenCalledWith(SSH_RELAY_CONFIGURE_GRACE_TIME_METHOD, {
+      graceTimeSeconds: 0
+    })
+  })
+
   it('ssh:resetRelay expires active-session leases instead of marking them terminated', async () => {
     const target: SshTarget = {
       id: 'ssh-1',

src/main/ipc/ssh.ts

@@ -1,6 +1,6 @@
 /* oxlint-disable max-lines -- Why: co-locates SSH IPC handlers, port-forward
 broadcasting, and session lifecycle in one file to keep the data flow obvious. */
-import { ipcMain, type BrowserWindow } from 'electron'
+import { ipcMain, powerMonitor, type BrowserWindow } from 'electron'
 import type { Store } from '../persistence'
 import { SshConnectionStore } from '../ssh/ssh-connection-store'
 import { SshConnectionManager, type SshConnectionCallbacks } from '../ssh/ssh-connection'
@@ -44,6 +44,7 @@ let registeredConnectSshTarget: ((targetId: string) => Promise<SshConnectionStat
 let registeredGetSshState: ((targetId: string) => SshConnectionState | undefined) | null = null
 let persistedStore: Store | null = null
 let advertisedUrlWatcherUnsubscribe: (() => void) | null = null
+let powerMonitorUnsubscribe: (() => void) | null = null
 
 export async function connectRegisteredSshTarget(targetId: string): Promise<SshConnectionState> {
   if (!registeredConnectSshTarget) {
@@ -308,6 +309,37 @@ function registerAdvertisedUrlRefresh(getMainWindow: () => BrowserWindow | null)
   })
 }
 
+function registerPowerMonitorReconnect(): void {
+  powerMonitorUnsubscribe?.()
+  const onSuspend = (): void => {
+    for (const session of activeSessions.values()) {
+      session.prepareForHostSleep()
+    }
+  }
+  const onResume = (): void => {
+    for (const targetId of activeSessions.keys()) {
+      const conn = connectionManager?.getConnection(targetId)
+      if (!conn) {
+        continue
+      }
+      const reconnect = connectionManager?.reconnect(targetId)
+      void reconnect?.catch((err) => {
+        console.warn(
+          `[ssh] Failed to reconnect ${targetId} after system resume: ${
+            err instanceof Error ? err.message : String(err)
+          }`
+        )
+      })
+    }
+  }
+  powerMonitor.on('suspend', onSuspend)
+  powerMonitor.on('resume', onResume)
+  powerMonitorUnsubscribe = () => {
+    powerMonitor.off('suspend', onSuspend)
+    powerMonitor.off('resume', onResume)
+  }
+}
+
 export function registerSshHandlers(
   store: Store,
   getMainWindow: () => BrowserWindow | null,
@@ -405,6 +437,7 @@ export function registerSshHandlers(
 
   connectionManager = new SshConnectionManager(callbacks)
   portForwardManager = new SshPortForwardManager()
+  registerPowerMonitorReconnect()
   registerSshBrowseHandler(() => connectionManager)
 
   // ── Target CRUD ────────────────────────────────────────────────────

src/main/ssh/ssh-connection-manager.ts

@@ -60,6 +60,14 @@ export class SshConnectionManager {
     this.connections.delete(targetId)
   }
 
+  async reconnect(targetId: string): Promise<void> {
+    const conn = this.connections.get(targetId)
+    if (!conn) {
+      return
+    }
+    await conn.reconnect()
+  }
+
   getConnection(targetId: string): SshConnection | undefined {
     return this.connections.get(targetId)
   }

src/main/ssh/ssh-connection.test.ts

@@ -201,6 +201,23 @@ describe('SshConnection', () => {
     expect(clientInstances[1].setNoDelay).toHaveBeenCalledWith(true)
   })
 
+  it('forces a fresh SSH connection for an explicit reconnect', async () => {
+    const states: string[] = []
+    const conn = new SshConnection(
+      createTarget(),
+      createCallbacks({
+        onStateChange: vi.fn((_id, state) => states.push(state.status))
+      })
+    )
+    await conn.connect()
+
+    await conn.reconnect()
+
+    expect(clientInstances).toHaveLength(2)
+    expect(states).toEqual(['connecting', 'connected', 'reconnecting', 'connecting', 'connected'])
+    expect(conn.getState().status).toBe('connected')
+  })
+
   it('transitions through connecting → connected states', async () => {
     const states: string[] = []
     const callbacks = createCallbacks({

src/main/ssh/ssh-connection.ts

@@ -339,6 +339,23 @@ export class SshConnection {
     }
   }
 
+  async reconnect(): Promise<void> {
+    if (this.disposed || this.state.status === 'connecting') {
+      return
+    }
+    if (this.reconnectTimer) {
+      clearTimeout(this.reconnectTimer)
+      this.reconnectTimer = null
+    }
+    // Why: OS sleep/wake can leave ssh2 thinking a dead TCP socket is still
+    // connected. Tear down the local transport and run the normal reconnect
+    // path so the relay session can reattach remote PTYs after wake.
+    this.closeTransportsForReconnect()
+    this.state.reconnectAttempt = 0
+    this.setState('reconnecting')
+    await this.runReconnectAttempt(0)
+  }
+
   private async doSystemSshProbe(connectGeneration: number): Promise<void> {
     this.useSystemSshTransport = true
     this.client = null
@@ -515,30 +532,57 @@ export class SshConnection {
       if (this.disposed) {
         return
       }
-      try {
-        // Why: reset reconnectAttempt before attemptConnect so setState('connected')
-        // broadcasts reconnectAttempt=0, which ssh.ts uses to trigger relay re-establishment.
-        this.state.reconnectAttempt = 0
-        await this.attemptConnect()
-      } catch (err) {
-        if (this.disposed) {
-          return
-        }
-        const error = err instanceof Error ? err : new Error(String(err))
-        if (isAuthError(error) || isPassphraseError(error)) {
-          this.setState('auth-failed', error.message)
-          return
-        }
-        if (!isTransientError(error)) {
-          this.setState('error', error.message)
-          return
-        }
-        this.state.reconnectAttempt = attempt + 1
-        this.scheduleReconnect()
-      }
+      await this.runReconnectAttempt(attempt)
     }, RECONNECT_BACKOFF_MS[attempt])
   }
 
+  private async runReconnectAttempt(attempt: number): Promise<void> {
+    try {
+      // Why: reset reconnectAttempt before attemptConnect so setState('connected')
+      // broadcasts reconnectAttempt=0, which ssh.ts uses to trigger relay re-establishment.
+      this.state.reconnectAttempt = 0
+      await this.attemptConnect()
+    } catch (err) {
+      if (this.disposed) {
+        return
+      }
+      const error = err instanceof Error ? err : new Error(String(err))
+      if (isAuthError(error) || isPassphraseError(error)) {
+        this.setState('auth-failed', error.message)
+        return
+      }
+      if (!isTransientError(error)) {
+        this.setState('error', error.message)
+        return
+      }
+      this.state.reconnectAttempt = attempt + 1
+      this.scheduleReconnect()
+    }
+  }
+
+  private closeTransportsForReconnect(): void {
+    this.connectGeneration += 1
+    const client = this.client
+    this.client = null
+    try {
+      client?.end()
+      client?.destroy()
+    } catch {
+      /* best-effort transport teardown */
+    }
+    this.proxyProcess?.kill()
+    this.proxyProcess = null
+    this.systemOperationAbortController.abort()
+    this.systemOperationAbortController = new AbortController()
+    for (const channel of this.systemCommandChannels) {
+      channel.close()
+    }
+    this.systemCommandChannels.clear()
+    this.systemSsh?.kill()
+    this.systemSsh = null
+    this.useSystemSshTransport = false
+  }
+
   async connectViaSystemSsh(): Promise<SystemSshProcess> {
     if (this.disposed) {
       throw new Error('Connection disposed')

src/main/ssh/ssh-relay-session.test.ts

@@ -7,6 +7,7 @@ import type { SshConnection } from './ssh-connection'
 import type { Store } from '../persistence'
 import type { SshPortForwardManager } from './ssh-port-forward'
 import { AGENT_HOOK_INSTALL_PLUGINS_METHOD } from '../../shared/agent-hook-relay'
+import { SSH_RELAY_CONFIGURE_GRACE_TIME_METHOD } from '../../shared/ssh-types'
 
 const { muxRequestMock, installRemoteManagedAgentHooksMock } = vi.hoisted(() => ({
   muxRequestMock: vi.fn(),
@@ -524,6 +525,30 @@ describe('SshRelaySession', () => {
     expect(deployAndLaunchRelay).toHaveBeenCalledWith(mockConn, undefined, 600, 'target-1')
   })
 
+  it('restores the configured relay grace after establish', async () => {
+    const { mockConn, mockStore, mockPortForward, getMainWindow } = createMockDeps()
+    const session = new SshRelaySession('target-1', getMainWindow, mockStore, mockPortForward)
+
+    await session.establish(mockConn, 600)
+
+    expect(session.getMux()?.notify).toHaveBeenCalledWith(SSH_RELAY_CONFIGURE_GRACE_TIME_METHOD, {
+      graceTimeSeconds: 600
+    })
+  })
+
+  it('sets relay grace to unlimited before host sleep', async () => {
+    const { mockConn, mockStore, mockPortForward, getMainWindow } = createMockDeps()
+    const session = new SshRelaySession('target-1', getMainWindow, mockStore, mockPortForward)
+    await session.establish(mockConn)
+    vi.mocked(session.getMux()!.notify).mockClear()
+
+    session.prepareForHostSleep()
+
+    expect(session.getMux()?.notify).toHaveBeenCalledWith(SSH_RELAY_CONFIGURE_GRACE_TIME_METHOD, {
+      graceTimeSeconds: 0
+    })
+  })
+
   it('cleans up port forwards on dispose', async () => {
     const { mockConn, mockStore, mockPortForward, getMainWindow } = createMockDeps()
     const session = new SshRelaySession('target-1', getMainWindow, mockStore, mockPortForward)

src/main/ssh/ssh-relay-session.ts

@@ -48,12 +48,29 @@ import { notifyRemoteWorkspaceHandlers } from '../ipc/remote-workspace-events'
 import { PortScanner } from './ssh-port-scanner'
 import type { SshPortForwardManager } from './ssh-port-forward'
 import type { SshConnection } from './ssh-connection'
-import type { DetectedPort } from '../../shared/ssh-types'
+import {
+  DEFAULT_SSH_RELAY_GRACE_PERIOD_SECONDS,
+  type DetectedPort,
+  MAX_SSH_RELAY_GRACE_PERIOD_SECONDS,
+  MIN_SSH_RELAY_GRACE_PERIOD_SECONDS,
+  SSH_RELAY_CONFIGURE_GRACE_TIME_METHOD
+} from '../../shared/ssh-types'
 import type { Store } from '../persistence'
 import type { OrcaRuntimeService } from '../runtime/orca-runtime'
 
 export type RelaySessionState = 'idle' | 'deploying' | 'ready' | 'reconnecting' | 'disposed'
 
+function normalizeRelayGracePeriodSeconds(graceTimeSeconds: number | undefined): number {
+  const raw = graceTimeSeconds ?? DEFAULT_SSH_RELAY_GRACE_PERIOD_SECONDS
+  const requested = Number.isFinite(raw) ? Math.floor(raw) : DEFAULT_SSH_RELAY_GRACE_PERIOD_SECONDS
+  return requested === 0
+    ? 0
+    : Math.max(
+        MIN_SSH_RELAY_GRACE_PERIOD_SECONDS,
+        Math.min(MAX_SSH_RELAY_GRACE_PERIOD_SECONDS, requested)
+      )
+}
+
 export class SshRelaySession {
   private _state: RelaySessionState = 'idle'
   private mux: SshChannelMultiplexer | null = null
@@ -133,6 +150,14 @@ export class SshRelaySession {
     return this.portScanner
   }
 
+  prepareForHostSleep(): void {
+    const mux = this.mux
+    if (!mux || mux.isDisposed() || this.isDisposed()) {
+      return
+    }
+    mux.notify(SSH_RELAY_CONFIGURE_GRACE_TIME_METHOD, { graceTimeSeconds: 0 })
+  }
+
   // Why: single entry point for relay setup — used by both initial connect
   // and app-restart reconnect. Having one path eliminates the risk of
   // forgetting a registration step.
@@ -203,6 +228,7 @@ export class SshRelaySession {
         throw new Error('Session disposed during establish')
       }
 
+      this.configureRelayGraceTime(mux, graceTimeSeconds)
       this.watchMuxForRelayLoss(mux)
       this._state = 'ready'
       this.startPortScanning()
@@ -328,6 +354,7 @@ export class SshRelaySession {
         return
       }
 
+      this.configureRelayGraceTime(mux, graceTimeSeconds)
       this.watchMuxForRelayLoss(mux)
       this._state = 'ready'
       this.startPortScanning()
@@ -462,6 +489,15 @@ export class SshRelaySession {
     return true
   }
 
+  private configureRelayGraceTime(
+    mux: SshChannelMultiplexer,
+    graceTimeSeconds: number | undefined
+  ): void {
+    mux.notify(SSH_RELAY_CONFIGURE_GRACE_TIME_METHOD, {
+      graceTimeSeconds: normalizeRelayGracePeriodSeconds(graceTimeSeconds)
+    })
+  }
+
   // Why: the relay can inject ORCA_AGENT_HOOK_* env into SSH PTYs, but
   // hook-script agents (Claude/Codex/Gemini/etc.) still need their config
   // files on the remote host to call Orca's managed script. Install those

src/relay/pty-handler.test.ts

@@ -146,6 +146,18 @@ describe('PtyHandler', () => {
     expect(onExpire).not.toHaveBeenCalled()
   })
 
+  it('uses the configured grace time for future disconnect timers', () => {
+    const onExpire = vi.fn()
+
+    handler.setGraceTimeMs(250)
+    handler.startGraceTimer(onExpire)
+
+    vi.advanceTimersByTime(249)
+    expect(onExpire).not.toHaveBeenCalled()
+    vi.advanceTimersByTime(1)
+    expect(onExpire).toHaveBeenCalledTimes(1)
+  })
+
   it('spawns a PTY and returns an id', async () => {
     const result = await dispatcher.callRequest('pty.spawn', { cols: 80, rows: 24 })
     expect(result).toEqual({ id: 'pty-1' })

src/relay/pty-handler.ts

@@ -150,6 +150,14 @@ export class PtyHandler {
     this.registerHandlers()
   }
 
+  setGraceTimeMs(graceTimeMs: number): void {
+    this.graceTimeMs = Math.max(0, Math.floor(graceTimeMs))
+  }
+
+  get configuredGraceTimeMs(): number {
+    return this.graceTimeMs
+  }
+
   /** Subscribe to PTY-exit events. Used by the relay-hook server to evict
    *  per-paneKey cached payloads when the backing PTY ends. */
   setExitListener(listener: PtyExitListener | null): void {

src/relay/relay.ts

@@ -42,7 +42,10 @@ import {
   AGENT_HOOK_NOTIFICATION_METHOD,
   AGENT_HOOK_REQUEST_REPLAY_METHOD
 } from '../shared/agent-hook-relay'
-import { DEFAULT_SSH_RELAY_GRACE_PERIOD_SECONDS } from '../shared/ssh-types'
+import {
+  DEFAULT_SSH_RELAY_GRACE_PERIOD_SECONDS,
+  SSH_RELAY_CONFIGURE_GRACE_TIME_METHOD
+} from '../shared/ssh-types'
 import { assertPluginSourceUnderByteCap } from './plugin-source-limit'
 import { resolveOpenCodeSourceConfigDir, resolvePiSourceAgentDir } from './plugin-overlay-env'
 import { detectPiAgentKindFromCommand } from '../shared/pi-agent-kind'
@@ -303,6 +306,23 @@ async function main(): Promise<void> {
   const _workspaceSessionHandler = new WorkspaceSessionHandler(dispatcher)
   void _workspaceSessionHandler
 
+  function configureRelayGraceTime(params: Record<string, unknown>): { graceTimeMs: number } {
+    const seconds = Number(params.graceTimeSeconds)
+    if (Number.isFinite(seconds) && seconds >= 0) {
+      // Why: the host sends 0 before system sleep so live remote PTYs survive
+      // longer than the ordinary disconnect grace window.
+      ptyHandler.setGraceTimeMs(Math.floor(seconds) * 1000)
+    }
+    return { graceTimeMs: ptyHandler.configuredGraceTimeMs }
+  }
+
+  dispatcher.onNotification(SSH_RELAY_CONFIGURE_GRACE_TIME_METHOD, (params) => {
+    configureRelayGraceTime(params)
+  })
+  dispatcher.onRequest(SSH_RELAY_CONFIGURE_GRACE_TIME_METHOD, async (params) =>
+    configureRelayGraceTime(params)
+  )
+
   // ── Agent-hook server ─────────────────────────────────────────────
   // Why: hosts a loopback HTTP receiver inside the relay process so agent
   // CLIs running in remote PTYs can post hook events without leaving the

src/shared/ssh-types.ts

@@ -3,6 +3,7 @@
 export const MIN_SSH_RELAY_GRACE_PERIOD_SECONDS = 60
 export const MAX_SSH_RELAY_GRACE_PERIOD_SECONDS = 7 * 24 * 60 * 60
 export const DEFAULT_SSH_RELAY_GRACE_PERIOD_SECONDS = 3 * 60 * 60
+export const SSH_RELAY_CONFIGURE_GRACE_TIME_METHOD = 'relay.configureGraceTime'
 
 export type SshTarget = {
   id: string