From eb579a0e542e738958dabb33b88920299d970c71 Mon Sep 17 00:00:00 2001 From: Neil <4138956+nwparker@users.noreply.github.com> Date: Tue, 26 May 2026 14:53:16 -0700 Subject: [PATCH] Fix SSH sessions after host sleep --- src/main/ipc/ssh.test.ts | 77 ++++++++++++++++++++++- src/main/ipc/ssh.ts | 35 ++++++++++- src/main/ssh/ssh-connection-manager.ts | 8 +++ src/main/ssh/ssh-connection.test.ts | 17 +++++ src/main/ssh/ssh-connection.ts | 86 +++++++++++++++++++------- src/main/ssh/ssh-relay-session.test.ts | 25 ++++++++ src/main/ssh/ssh-relay-session.ts | 38 +++++++++++- src/relay/pty-handler.test.ts | 12 ++++ src/relay/pty-handler.ts | 8 +++ src/relay/relay.ts | 22 ++++++- src/shared/ssh-types.ts | 1 + 11 files changed, 304 insertions(+), 25 deletions(-) diff --git a/src/main/ipc/ssh.test.ts b/src/main/ipc/ssh.test.ts index e526005b3..570ff42f1 100644 --- a/src/main/ipc/ssh.test.ts +++ b/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', diff --git a/src/main/ipc/ssh.ts b/src/main/ipc/ssh.ts index 4e2f03037..ecf88092d 100644 --- a/src/main/ipc/ssh.ts +++ b/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 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 { 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 ──────────────────────────────────────────────────── diff --git a/src/main/ssh/ssh-connection-manager.ts b/src/main/ssh/ssh-connection-manager.ts index 6e87cee1e..b5eb04bf2 100644 --- a/src/main/ssh/ssh-connection-manager.ts +++ b/src/main/ssh/ssh-connection-manager.ts @@ -60,6 +60,14 @@ export class SshConnectionManager { this.connections.delete(targetId) } + async reconnect(targetId: string): Promise { + const conn = this.connections.get(targetId) + if (!conn) { + return + } + await conn.reconnect() + } + getConnection(targetId: string): SshConnection | undefined { return this.connections.get(targetId) } diff --git a/src/main/ssh/ssh-connection.test.ts b/src/main/ssh/ssh-connection.test.ts index 55e100953..ddb9d175d 100644 --- a/src/main/ssh/ssh-connection.test.ts +++ b/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({ diff --git a/src/main/ssh/ssh-connection.ts b/src/main/ssh/ssh-connection.ts index 422a1f2b7..8d6c0eb08 100644 --- a/src/main/ssh/ssh-connection.ts +++ b/src/main/ssh/ssh-connection.ts @@ -339,6 +339,23 @@ export class SshConnection { } } + async reconnect(): Promise { + 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 { 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 { + 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 { if (this.disposed) { throw new Error('Connection disposed') diff --git a/src/main/ssh/ssh-relay-session.test.ts b/src/main/ssh/ssh-relay-session.test.ts index de52e2c7e..ee4fbb8bb 100644 --- a/src/main/ssh/ssh-relay-session.test.ts +++ b/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) diff --git a/src/main/ssh/ssh-relay-session.ts b/src/main/ssh/ssh-relay-session.ts index ba35666c8..1282dfdce 100644 --- a/src/main/ssh/ssh-relay-session.ts +++ b/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 diff --git a/src/relay/pty-handler.test.ts b/src/relay/pty-handler.test.ts index 27fcf52ab..a31d35259 100644 --- a/src/relay/pty-handler.test.ts +++ b/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' }) diff --git a/src/relay/pty-handler.ts b/src/relay/pty-handler.ts index abc5ba208..c6c6ae7e6 100644 --- a/src/relay/pty-handler.ts +++ b/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 { diff --git a/src/relay/relay.ts b/src/relay/relay.ts index ef25635d9..213d71623 100644 --- a/src/relay/relay.ts +++ b/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 { const _workspaceSessionHandler = new WorkspaceSessionHandler(dispatcher) void _workspaceSessionHandler + function configureRelayGraceTime(params: Record): { 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 diff --git a/src/shared/ssh-types.ts b/src/shared/ssh-types.ts index dd23298c2..b1d4868a1 100644 --- a/src/shared/ssh-types.ts +++ b/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