fix(ssh): handle owner displacement and graceful shutdown (#12367)

* fix(ssh): handle owner displacement and graceful shutdown

SSH connections can reconnect with valid session proof after network loss or
device sleep. When the incumbent owner is still half-open, allow the
reconnecting client to displace it outright rather than wait for socket closure
— a window that may never close. Retain displaced deliveries for the new owner
to rotate. During app shutdown, drain SSH sessions without terminating recovery
operations, and retry pending owner grants in case a replacement commits mid-drain.

* 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.

* fix(ssh): fix flaky test sync and add error code safety check

Test was using tick-based Promise.resolve() loops which don't guarantee
the async operation has started. Replace with signal-based synchronization
that waits for the actual lease flush. Also add nullish-coalescing to
error code check to prevent crashes if error is null or undefined.

* fix(ssh): fence reset transport opens during shutdown

* fix(ssh): keep recovery leases stable across reconnects

* fix(ssh): close transports owned by cancelled connect attempts

When a connect is cancelled after its transport has opened, that cancelled
attempt still owns the transport and must close it — otherwise it leaks. Add
disconnectConnection() to close by identity (not by target ID) so a cancelled
attempt closes only the transport it minted, without tearing down its
replacement's live transport. Track priorConnection to detect whether this
attempt opened a new transport or reused an existing one, and close only on
abandonment if this attempt owns the session.

* fix(ssh): fence old owner proofs and close superseded transports

When an owner reconnects with a new proof while an old one is still live,
the old proof is now fenced with SUPERSEDED_ERROR instead of retrying
indefinitely. The relay also closes stale transports to signal that their
recovery generation has been overtaken by a newer one.

This ensures overlapping reconnect scenarios complete with the newest proof
rather than getting blocked by stale recovery attempts.

* fix(relay): re-pin stdin/stdout fds after closing to prevent recycling

When the relay closes stdin/stdout to signal EOF to the SSH peer, the OS
can recycle those fds (0 and 1) for new sockets or files. If Node still
treats process.stdin/stdout as those numbers, subsequent operations
corrupt socket clients and trigger shutdown errors. Re-pin the fds by
opening /dev/null to keep them occupied and prevent recycling.
This commit is contained in:
Jinjing
2026-08-04 01:44:59 -07:00
parent d759c5094b
commit 72f534ba21
19 changed files with 1627 additions and 116 deletions
+4 -2
View File
@@ -250,7 +250,8 @@ 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'
import { getPluginsDataDir } from './plugins/plugin-discovery'
@@ -2940,7 +2941,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
@@ -2992,6 +2992,7 @@ app.on('will-quit', (e) => {
runtime?.getOffscreenBrowserBackend()?.destroyAll?.()
browserManager.setBrowserGuestStateChangedListener(null)
const emulatorShutdown = runtime?.getEmulatorBridge()?.destroyAllSessions() ?? Promise.resolve()
const sshShutdown = detachAllSshSessionsForShutdown()
killAllPty()
const watcherShutdown = shutdownWatchersOnce()
const storeFlush = store?.flushAsync() ?? Promise.resolve()
@@ -3034,6 +3035,7 @@ app.on('will-quit', (e) => {
{ name: 'runtime-rpc', promise: rpcStopAndClear },
{ name: 'watchers', promise: watcherShutdown },
{ name: 'emulator', promise: emulatorShutdown },
{ name: 'ssh', promise: sshShutdown },
{ name: 'plugin-hosts', promise: pluginHostShutdown },
{ name: 'usage-cache', promise: usageCacheFlush },
{ name: 'stats', promise: statsFlush },
+295 -1
View File
@@ -42,6 +42,7 @@ const {
mockConnectionManager: {
connect: vi.fn(),
disconnect: vi.fn(),
disconnectConnection: vi.fn(),
reconnect: vi.fn(),
getConnection: vi.fn(),
getState: vi.fn(),
@@ -241,7 +242,13 @@ vi.mock('../ssh/ssh-port-scanner', () => ({
}
}))
import { getSshConnectionManager, registerSshHandlers, resetSshHandlerStateForTests } from './ssh'
import {
detachAllSshSessionsForShutdown,
getActiveMultiplexer,
getSshConnectionManager,
registerSshHandlers,
resetSshHandlerStateForTests
} from './ssh'
import { RelayVersionMismatchError } from '../ssh/ssh-relay-version-mismatch-error'
import {
SSH_RELAY_CONFIGURE_GRACE_TIME_METHOD,
@@ -260,6 +267,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'
@@ -296,6 +304,7 @@ describe('SSH IPC handlers', () => {
const createConnectionManagerMock = () => ({
connect: vi.fn(),
disconnect: vi.fn(),
disconnectConnection: vi.fn(),
reconnect: vi.fn(),
getConnection: vi.fn(),
getState: vi.fn(),
@@ -365,6 +374,7 @@ describe('SSH IPC handlers', () => {
mockConnectionManager.connect.mockReset()
mockConnectionManager.disconnect.mockReset()
mockConnectionManager.disconnectConnection.mockReset().mockResolvedValue(undefined)
mockConnectionManager.reconnect.mockReset()
mockConnectionManager.getConnection.mockReset()
mockConnectionManager.getState.mockReset()
@@ -539,6 +549,145 @@ describe('SSH IPC handlers', () => {
expect(mockSshStore.removeTarget).toHaveBeenCalledWith('ssh-1')
})
it('detaches active SSH sessions during app shutdown without terminating recovery', async () => {
const target: SshTarget = {
id: 'ssh-1',
label: 'Server',
host: 'example.com',
port: 22,
username: 'deploy'
}
mockSshStore.getTarget.mockReturnValue(target)
mockConnectionManager.connect.mockResolvedValue({})
mockConnectionManager.getState.mockReturnValue({
targetId: 'ssh-1',
status: 'connected',
error: null,
reconnectAttempt: 0
})
await handlers.get('ssh:connect')!(null, { targetId: 'ssh-1' })
mockMux.dispose.mockClear()
mockConnectionManager.disconnectAll.mockClear().mockResolvedValue(undefined)
await detachAllSshSessionsForShutdown()
expect(mockPortForwardManager.removeAllForwards).toHaveBeenCalledWith('ssh-1')
expect(mockMux.dispose).toHaveBeenCalledWith('connection_lost')
expect(mockStore.markSshRemotePtyLeasesAsync).toHaveBeenCalledWith('ssh-1', 'detached')
expect(mockStore.removeSshPtyConsumerRecovery).not.toHaveBeenCalled()
expect(mockConnectionManager.disconnectAll).toHaveBeenCalled()
})
it('refuses the replacement session a paused connect would publish after shutdown began', async () => {
const target: SshTarget = {
id: 'ssh-1',
label: 'Server',
host: 'example.com',
port: 22,
username: 'deploy'
}
mockSshStore.getTarget.mockReturnValue(target)
mockConnectionManager.connect.mockResolvedValue({})
mockConnectionManager.getState.mockReturnValue({
targetId: 'ssh-1',
status: 'connected',
error: null,
reconnectAttempt: 0
})
mockConnectionManager.disconnectAll.mockResolvedValue(undefined)
await handlers.get('ssh:connect')!(null, { targetId: 'ssh-1' })
// Why here: the old session's lease flush is the one window where doConnect resumes straight into
// publishing its replacement session without re-checking authority.
let releaseDetach = (): void => {}
const detachFlush = new Promise<void>((resolve) => {
releaseDetach = resolve
})
let signalEnteredDetach = (): void => {}
const enteredDetach = new Promise<void>((resolve) => {
signalEnteredDetach = resolve
})
mockStore.markSshRemotePtyLeasesAsync.mockImplementationOnce(() => {
signalEnteredDetach()
return detachFlush
})
const replacement = handlers.get('ssh:connect')!(null, { targetId: 'ssh-1' })
// Why await the flush entry and not ticks: shutdown must not snapshot what to drain until the
// replacement connect is parked in the lease flush.
await enteredDetach
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('closed for app shutdown')
await shutdown
// 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(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(
'closed for app shutdown'
)
})
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: '' }
@@ -1839,6 +1988,106 @@ describe('SSH IPC handlers', () => {
expect(mockDeployAndLaunchRelay).toHaveBeenCalledTimes(1)
})
it('closes the transport a cancelled connect opened after the disconnect finished', async () => {
const target: SshTarget = {
id: 'ssh-1',
label: 'Server',
host: 'example.com',
port: 22,
username: 'deploy'
}
const lateConn = { id: 'late-transport' }
let resolveStaleConnect!: (connection: unknown) => void
mockSshStore.getTarget.mockReturnValue(target)
mockConnectionManager.getConnection.mockReturnValue(undefined)
mockConnectionManager.connect.mockReturnValueOnce(
new Promise((resolve) => {
resolveStaleConnect = resolve
})
)
mockConnectionManager.disconnect.mockResolvedValue(undefined)
const staleConnect = handlers.get('ssh:connect')!(null, {
targetId: 'ssh-1'
}) as Promise<SshConnectionState>
await vi.waitFor(() => expect(mockConnectionManager.connect).toHaveBeenCalledTimes(1))
// Why await the whole disconnect: the leak only exists once its teardown has already run, so
// nothing else is left to close the transport this attempt opens afterwards.
await handlers.get('ssh:disconnect')!(null, { targetId: 'ssh-1' })
resolveStaleConnect(lateConn)
await expect(staleConnect).rejects.toThrow('SSH connection attempt was cancelled')
expect(mockConnectionManager.disconnectConnection).toHaveBeenCalledWith('ssh-1', lateConn)
expect(getActiveMultiplexer('ssh-1')).toBeUndefined()
})
it('closes the transport when establish resumes after the connect was invalidated', async () => {
const target: SshTarget = {
id: 'ssh-1',
label: 'Server',
host: 'example.com',
port: 22,
username: 'deploy'
}
const conn = { id: 'establishing-transport' }
let releaseRelayLaunch = (): void => {}
mockSshStore.getTarget.mockReturnValue(target)
mockConnectionManager.getConnection.mockReturnValue(undefined)
mockConnectionManager.connect.mockResolvedValue(conn)
mockConnectionManager.disconnect.mockResolvedValue(undefined)
mockDeployAndLaunchRelay.mockImplementationOnce(
() =>
new Promise((resolve) => {
releaseRelayLaunch = () => resolve(createRelayLaunchResult())
})
)
const connect = handlers.get('ssh:connect')!(null, {
targetId: 'ssh-1'
}) as Promise<SshConnectionState>
await vi.waitFor(() => expect(mockDeployAndLaunchRelay).toHaveBeenCalledTimes(1))
await handlers.get('ssh:disconnect')!(null, { targetId: 'ssh-1' })
releaseRelayLaunch()
await expect(connect).rejects.toThrow('SSH connection attempt was cancelled')
expect(mockConnectionManager.disconnectConnection).toHaveBeenCalledWith('ssh-1', conn)
expect(getActiveMultiplexer('ssh-1')).toBeUndefined()
})
it('leaves a reused transport to its replacement when the connect is cancelled', async () => {
const target: SshTarget = {
id: 'ssh-1',
label: 'Server',
host: 'example.com',
port: 22,
username: 'deploy'
}
// Why: connect() hands back the already-open transport, so this attempt never owned it.
const sharedConn = { id: 'shared-transport' }
let resolveStaleConnect!: (connection: unknown) => void
mockSshStore.getTarget.mockReturnValue(target)
mockConnectionManager.getConnection.mockReturnValue(sharedConn)
mockConnectionManager.connect.mockReturnValueOnce(
new Promise((resolve) => {
resolveStaleConnect = resolve
})
)
mockConnectionManager.disconnect.mockResolvedValue(undefined)
const staleConnect = handlers.get('ssh:connect')!(null, {
targetId: 'ssh-1'
}) as Promise<SshConnectionState>
await vi.waitFor(() => expect(mockConnectionManager.connect).toHaveBeenCalledTimes(1))
await handlers.get('ssh:disconnect')!(null, { targetId: 'ssh-1' })
resolveStaleConnect(sharedConn)
await expect(staleConnect).rejects.toThrow('SSH connection attempt was cancelled')
expect(mockConnectionManager.disconnectConnection).not.toHaveBeenCalled()
})
it('keeps reconnect behind transport disconnect when forward teardown fails', async () => {
const target: SshTarget = {
id: 'ssh-1',
@@ -2327,6 +2576,51 @@ describe('SSH IPC handlers', () => {
expect(mockConnectionManager.disconnect).toHaveBeenCalledWith('ssh-1')
})
it('ssh:resetRelay does not open a transport when shutdown starts while it waits for a connect', async () => {
const target: SshTarget = {
id: 'ssh-1',
label: 'Server',
host: 'example.com',
port: 22,
username: 'deploy'
}
let failConnect!: (error: unknown) => void
const connectResult = new Promise((_resolve, reject) => {
failConnect = reject
})
mockSshStore.getTarget.mockReturnValue(target)
mockConnectionManager.connect.mockReturnValue(connectResult)
// Why undefined: reset must fall through to opening its own transport, which is the call under test.
mockConnectionManager.getConnection.mockReturnValue(undefined)
mockConnectionManager.getState.mockReturnValue({
targetId: 'ssh-1',
status: 'connected',
error: null,
reconnectAttempt: 0
})
const connectPromise = handlers.get('ssh:connect')!(null, {
targetId: 'ssh-1'
}) as Promise<unknown>
await vi.waitFor(() => expect(mockConnectionManager.connect).toHaveBeenCalledTimes(1))
// Why admitted first: the gate latches only after reset is already parked behind the connect, so
// the entry fence cannot be what stops it.
const resetPromise = handlers.get('ssh:resetRelay')!(null, {
targetId: 'ssh-1'
}) as Promise<void>
await Promise.resolve()
quitTeardownStartGate.tryStart({ preventDefault() {} })
failConnect(new Error('transport lost'))
await expect(connectPromise).rejects.toThrow('transport lost')
await expect(resetPromise).rejects.toThrow('closed for app shutdown')
// Why once: the resumed reset must not open a second transport that would outlive the drain.
expect(mockConnectionManager.connect).toHaveBeenCalledTimes(1)
expect(mockForceStopRelayForTarget).not.toHaveBeenCalled()
})
it('ssh:connect waits for an in-flight reset before starting a new connection', async () => {
const target: SshTarget = {
id: 'ssh-1',
+132 -4
View File
@@ -3,7 +3,7 @@ import { ipcMain, powerMonitor, type BrowserWindow } from 'electron'
import { appendFileSync } from 'node:fs'
import type { Store } from '../persistence'
import { SshConnectionStore } from '../ssh/ssh-connection-store'
import type { SshConnectionCallbacks } from '../ssh/ssh-connection'
import type { SshConnection, SshConnectionCallbacks } from '../ssh/ssh-connection'
import { SshConnectionManager } from '../ssh/ssh-connection-manager'
import type { SshChannelMultiplexer } from '../ssh/ssh-channel-multiplexer'
import { SshRelaySession, type SshRelayAiVaultHostInfo } from '../ssh/ssh-relay-session'
@@ -19,6 +19,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'
@@ -269,6 +270,33 @@ async function abandonFailedSshSession(targetId: string, session: SshRelaySessio
}
}
// Why: a connect cancelled after its transport already opened still owns that transport and its
// unpublished session — nothing else will reach them, so it has to close them itself. Why only the
// connection it minted, and by identity: disconnecting by target id (or closing a transport it merely
// reused) would tear down the replacement connect's live transport.
async function abandonCancelledConnectAttempt(
targetId: string,
session: SshRelaySession,
mintedConnection: SshConnection | null
): Promise<void> {
// Why the identity guard: every path that removes a session from activeSessions detaches it first,
// so re-detaching a superseded session would only clobber the replacement's lease state.
if (activeSessions.get(targetId) === session) {
await abandonFailedSshSession(targetId, session)
}
if (!mintedConnection) {
return
}
try {
await connectionManager!.disconnectConnection(targetId, mintedConnection)
} catch (error) {
// Why: the caller is about to throw the cancellation; a teardown throw must not replace it.
console.warn(
`[ssh] Failed to disconnect cancelled connect transport for ${targetId}: ${error instanceof Error ? error.message : String(error)}`
)
}
}
function relayGracePeriodForTarget(target: SshTarget | null | undefined): number | undefined {
return target?.relayGracePeriodSeconds
}
@@ -283,6 +311,15 @@ type ConnectAttempt = {
const connectInFlight = new Map<string, ConnectAttempt>()
const pendingTransportReconnects = new Set<string>()
// 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 (quitTeardownStartGate.hasStarted()) {
throw new Error('SSH connects are closed for app shutdown')
}
}
function invalidateConnectAttempt(targetId: string): void {
rotateSshProviderAuthority(targetId)
pendingTransportReconnects.delete(targetId)
@@ -299,6 +336,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 = {
@@ -1004,6 +1042,9 @@ export function registerSshHandlers(
if (!isCurrentSshProviderAuthority(observedAuthority)) {
throw createCancelledConnectAttemptError()
}
// Why: the shutdown drain fences and snapshots synchronously, so a connect either registers in
// connectInFlight below (and gets joined) or fails here — it can never slip between the two.
assertSshConnectsNotFenced()
pendingTransportReconnects.delete(targetId)
const promise = doConnect(targetId, replacePendingTransport)
@@ -1092,6 +1133,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,
@@ -1106,6 +1152,12 @@ export function registerSshHandlers(
const ownsSession = (): boolean =>
isCurrentConnectAttempt(targetId, authority) && activeSessions.get(targetId) === session
// Why captured here and not with existingState: connect() reuses an already-connected transport,
// and only a transport this attempt opened is this attempt's to close when it loses the race.
const priorConnection = connectionManager!.getConnection(targetId)
const mintedConnection = (): SshConnection | null =>
conn && conn !== priorConnection ? conn : null
try {
conn = await connectionManager!.connect(target)
if (!ownsSession()) {
@@ -1116,6 +1168,7 @@ export function registerSshHandlers(
const errObj = err instanceof Error ? err : new Error(String(err))
const status: SshConnectionStatus = isAuthError(errObj) ? 'auth-failed' : 'error'
if (!ownsSession()) {
await abandonCancelledConnectAttempt(targetId, session, mintedConnection())
throw createCancelledConnectAttemptError()
}
// Why: clear this failed connect's flag so a later non-prompting connect isn't deferred.
@@ -1156,6 +1209,7 @@ export function registerSshHandlers(
})
} catch (err) {
if (!ownsSession()) {
await abandonCancelledConnectAttempt(targetId, session, mintedConnection())
throw createCancelledConnectAttemptError()
}
await abandonFailedSshSession(targetId, session)
@@ -1262,7 +1316,13 @@ export function registerSshHandlers(
}
const existingConn = connectionManager!.getConnection(targetId)
const conn = existingConn ?? (await connectionManager!.connect(target))
let conn = existingConn
if (!conn) {
// Why re-check: admission fenced this reset before it parked on the in-flight connect, so shutdown
// may have started (and drained) while we waited — opening a transport now would outlive the drain.
assertSshConnectsNotFenced()
conn = await connectionManager!.connect(target)
}
try {
await forceStopRelayForTarget(conn, targetId)
} finally {
@@ -1295,6 +1355,8 @@ export function registerSshHandlers(
if (!target) {
throw new Error(`SSH target "${args.targetId}" not found`)
}
// Why: reset opens its own transport, so it must be fenced by shutdown the same way connect is.
assertSshConnectsNotFenced()
let resetPromise: Promise<void>
resetPromise = runTargetLifecycle(args.targetId, () =>
@@ -1354,17 +1416,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)
@@ -1472,6 +1543,61 @@ export function getSshConnectionManager(): SshConnectionManager | null {
return connectionManager
}
// Why: an invalidated connect only observes its cancellation at the next checkpoint, and one blocked
// in the transport handshake can sit there for the whole SSH timeout. Bound the join so the final
// drain always runs well inside the quit deadline instead of racing it.
const SSH_SHUTDOWN_INFLIGHT_JOIN_MS = 2_000
async function settleWithinMs(work: readonly Promise<unknown>[], timeoutMs: number): Promise<void> {
if (work.length === 0) {
return
}
await Promise.race([
Promise.allSettled(work),
new Promise<void>((resolve) => {
const timer = setTimeout(resolve, timeoutMs)
timer.unref?.()
})
])
}
// 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> {
const inFlightWork: Promise<unknown>[] = [
...[...connectInFlight.values()].map((attempt) => attempt.promise),
...resetRelayInFlight.values(),
...testConnectionProbes
]
for (const targetId of Array.from(connectInFlight.keys())) {
invalidateConnectAttempt(targetId)
}
const errors: unknown[] = []
const drain = async (): Promise<void> => {
const results = await Promise.allSettled([
...[...activeSessions.keys()].map((targetId) =>
teardownActiveSshSession(targetId, (session) => session.detachAndPersist())
),
connectionManager?.disconnectAll() ?? Promise.resolve()
])
for (const result of results) {
if (result.status === 'rejected') {
errors.push(result.reason as unknown)
}
}
}
await drain()
// Why: 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.
await settleWithinMs(inFlightWork, SSH_SHUTDOWN_INFLIGHT_JOIN_MS)
await drain()
if (errors.length > 0) {
throw new AggregateError(errors, 'SSH shutdown detach failed')
}
}
export async function resetSshHandlerStateForTests(): Promise<void> {
advertisedUrlWatcherUnsubscribe?.()
advertisedUrlWatcherUnsubscribe = null
@@ -1498,7 +1624,9 @@ export async function resetSshHandlerStateForTests(): Promise<void> {
resetSshProviderAuthorities()
resetRelayInFlight.clear()
testingTargets.clear()
testConnectionProbes.clear()
credentialRequestedForTarget.clear()
quitTeardownStartGate.resetForTests()
await connectionManager?.disconnectAll()
portForwardManager?.dispose()
+12
View File
@@ -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()
@@ -94,4 +94,36 @@ describe('SshConnectionManager', () => {
expect(await manager.connect(target)).toBe(replacement)
expect(manager.getConnection(target.id)).toBe(replacement)
})
it('closes a late-resolved connection without evicting the replacement', async () => {
let resolveFirst!: () => void
mockState.connectResults.push(
new Promise<void>((resolve) => {
resolveFirst = resolve
}),
Promise.resolve()
)
const manager = new SshConnectionManager({ onStateChange: vi.fn() })
const firstConnect = manager.connect(target)
await manager.disconnect(target.id)
const replacement = await manager.connect(target)
resolveFirst()
const late = await firstConnect
await manager.disconnectConnection(target.id, late)
expect(mockState.instances[0].disconnect).toHaveBeenCalledTimes(2)
expect(manager.getConnection(target.id)).toBe(replacement)
})
it('clears the pool entry when the closed connection is still the registered one', async () => {
const manager = new SshConnectionManager({ onStateChange: vi.fn() })
const conn = await manager.connect(target)
await manager.disconnectConnection(target.id, conn)
expect(mockState.instances[0].disconnect).toHaveBeenCalledOnce()
expect(manager.getConnection(target.id)).toBeUndefined()
})
})
+12
View File
@@ -76,6 +76,18 @@ export class SshConnectionManager {
}
}
/**
* Close one specific connection, clearing the pool entry only when it is still the registered one.
* 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<void> {
await conn.disconnect()
if (this.connections.get(targetId) === conn) {
this.connections.delete(targetId)
}
}
async reconnect(targetId: string): Promise<void> {
const conn = this.connections.get(targetId)
if (!conn) {
@@ -0,0 +1,106 @@
import { afterEach, describe, expect, it, vi } from 'vitest'
import {
PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR,
PTY_CONSUMER_OWNER_RECOVERY_SUPERSEDED_ERROR
} from '../../shared/pty-consumer-session'
import { retrySshOwnerRecoveryWhileBlocked } from './ssh-owner-recovery-retry'
function publicationPendingError(): Error & { code: number } {
return Object.assign(new Error('Owner grant publication is still pending'), {
code: PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR
})
}
function supersededError(): Error & { code: number } {
return Object.assign(new Error('Owner recovery generation was superseded'), {
code: PTY_CONSUMER_OWNER_RECOVERY_SUPERSEDED_ERROR
})
}
function openGate() {
return {
isCurrent: () => true,
onClosed: () => () => {}
}
}
describe('SSH owner recovery retry', () => {
afterEach(() => {
vi.useRealTimers()
})
it('recovers once the incumbent grant publication settles', async () => {
vi.useFakeTimers()
const attempt = vi
.fn<() => Promise<string>>()
.mockRejectedValueOnce(publicationPendingError())
.mockRejectedValueOnce(publicationPendingError())
.mockResolvedValue('recovered')
const recovery = retrySshOwnerRecoveryWhileBlocked(attempt, openGate())
await vi.advanceTimersByTimeAsync(75)
await expect(recovery).resolves.toBe('recovered')
expect(attempt).toHaveBeenCalledTimes(3)
})
it('stops at the publication-settlement deadline', async () => {
vi.useFakeTimers()
const error = publicationPendingError()
const attempt = vi.fn<() => Promise<never>>().mockRejectedValue(error)
const recovery = retrySshOwnerRecoveryWhileBlocked(attempt, openGate(), 60)
const rejection = expect(recovery).rejects.toBe(error)
await vi.advanceTimersByTimeAsync(60)
await rejection
expect(attempt).toHaveBeenCalledTimes(3)
})
it('does not retry unrelated failures', async () => {
const error = Object.assign(new Error('stale owner'), { code: -32041 })
const attempt = vi.fn<() => Promise<never>>().mockRejectedValue(error)
await expect(retrySshOwnerRecoveryWhileBlocked(attempt, openGate())).rejects.toBe(error)
expect(attempt).toHaveBeenCalledOnce()
})
it('retries while a superseded transport is closing', async () => {
vi.useFakeTimers()
const attempt = vi
.fn<() => Promise<string>>()
.mockRejectedValueOnce(supersededError())
.mockResolvedValue('recovered')
const recovery = retrySshOwnerRecoveryWhileBlocked(attempt, openGate())
await vi.advanceTimersByTimeAsync(25)
await expect(recovery).resolves.toBe('recovered')
expect(attempt).toHaveBeenCalledTimes(2)
})
it('stops waiting when the relay channel closes', async () => {
vi.useFakeTimers()
let current = true
let close: (() => void) | undefined
const error = publicationPendingError()
const attempt = vi.fn<() => Promise<never>>().mockRejectedValue(error)
const recovery = retrySshOwnerRecoveryWhileBlocked(attempt, {
isCurrent: () => current,
onClosed: (listener) => {
close = listener
return () => {
close = undefined
}
}
})
await vi.advanceTimersByTimeAsync(0)
expect(close).toBeTypeOf('function')
current = false
close?.()
await expect(recovery).rejects.toBe(error)
expect(attempt).toHaveBeenCalledOnce()
})
})
+69
View File
@@ -0,0 +1,69 @@
import {
PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR,
PTY_CONSUMER_OWNER_RECOVERY_SUPERSEDED_ERROR
} from '../../shared/pty-consumer-session'
// Why: bound polling when publication is settling or a superseded attempt is closing its transport.
export const SSH_OWNER_RECOVERY_WAIT_MS = 3_000
const SSH_OWNER_RECOVERY_INITIAL_DELAY_MS = 25
const SSH_OWNER_RECOVERY_MAX_DELAY_MS = 250
type SshOwnerRecoveryRetryGate = {
isCurrent: () => boolean
onClosed: (listener: () => void) => () => void
}
function waitForRetry(delayMs: number, gate: SshOwnerRecoveryRetryGate): Promise<void> {
return new Promise((resolve) => {
let settled = false
let timer: ReturnType<typeof setTimeout> | undefined
let unsubscribe = (): void => {}
const finish = (): void => {
if (settled) {
return
}
settled = true
clearTimeout(timer)
unsubscribe()
resolve()
}
unsubscribe = gate.onClosed(finish)
timer = setTimeout(finish, delayMs)
timer.unref?.()
if (!gate.isCurrent()) {
finish()
}
})
}
export async function retrySshOwnerRecoveryWhileBlocked<T>(
attempt: () => Promise<T>,
gate: SshOwnerRecoveryRetryGate,
waitMs: number = SSH_OWNER_RECOVERY_WAIT_MS
): Promise<T> {
const deadline = Date.now() + waitMs
let delayMs = SSH_OWNER_RECOVERY_INITIAL_DELAY_MS
while (true) {
try {
return await attempt()
} catch (error) {
const code = (error as { code?: unknown } | null | undefined)?.code
if (
(code !== PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR &&
code !== PTY_CONSUMER_OWNER_RECOVERY_SUPERSEDED_ERROR) ||
!gate.isCurrent()
) {
throw error
}
const remainingMs = deadline - Date.now()
if (remainingMs <= 0) {
throw error
}
await waitForRetry(Math.min(delayMs, remainingMs), gate)
if (!gate.isCurrent()) {
throw error
}
delayMs = Math.min(delayMs * 2, SSH_OWNER_RECOVERY_MAX_DELAY_MS)
}
}
}
@@ -1,4 +1,8 @@
import { beforeEach, describe, expect, it, vi } from 'vitest'
import {
PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR,
PTY_CONSUMER_OWNER_RECOVERY_SUPERSEDED_ERROR
} from '../../shared/pty-consumer-session'
import { SshRelaySession } from './ssh-relay-session'
import {
createMismatchedOwnerRecoveryError,
@@ -154,6 +158,82 @@ describe('SshRelaySession consumer recovery durability', () => {
session.dispose()
})
it('retries a pending incumbent publication before recovering its persisted lease', async () => {
vi.useFakeTimers()
try {
const targetId = 'target-owner-publication-pending'
const deps = createMockDeps()
vi.mocked(deps.mockStore.getSshPtyConsumerRecovery).mockReturnValue({
targetId,
clientInstanceId: 'persisted-client',
serverBuildId: 'test-relay-build',
clientGeneration: 1,
ownerGeneration: 1,
ownerLease: 'persisted-owner'
})
openConsumerSessionMock.mockRejectedValueOnce(
Object.assign(new Error('Owner grant publication is still pending'), {
code: PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR
})
)
const session = new SshRelaySession(
targetId,
deps.getMainWindow,
deps.mockStore,
deps.mockPortForward
)
const establishing = session.establish(deps.mockConn)
await vi.advanceTimersByTimeAsync(25)
await establishing
expect(openConsumerSessionMock).toHaveBeenCalledTimes(2)
expect(openConsumerSessionMock.mock.calls[0]?.[1]).toMatchObject({
clientInstanceId: 'persisted-client',
resume: { ownerGeneration: 1, ownerLease: 'persisted-owner' }
})
session.dispose()
} finally {
vi.useRealTimers()
}
})
it('preserves persisted recovery while a superseding transport is still live', async () => {
vi.useFakeTimers()
try {
const targetId = 'target-owner-generation-superseded'
const deps = createMockDeps()
vi.mocked(deps.mockStore.getSshPtyConsumerRecovery).mockReturnValue({
targetId,
clientInstanceId: 'persisted-client',
serverBuildId: 'test-relay-build',
clientGeneration: 1,
ownerGeneration: 1,
ownerLease: 'persisted-owner'
})
const superseded = Object.assign(new Error('Owner recovery generation was superseded'), {
code: PTY_CONSUMER_OWNER_RECOVERY_SUPERSEDED_ERROR
})
openConsumerSessionMock.mockRejectedValue(superseded)
const session = new SshRelaySession(
targetId,
deps.getMainWindow,
deps.mockStore,
deps.mockPortForward
)
const failed = expect(session.establish(deps.mockConn)).rejects.toBe(superseded)
await vi.advanceTimersByTimeAsync(3_000)
await failed
expect(openConsumerSessionMock.mock.calls.length).toBeGreaterThan(1)
expect(deps.mockStore.removeSshPtyConsumerRecovery).not.toHaveBeenCalled()
session.dispose()
} finally {
vi.useRealTimers()
}
})
it('keeps destructive disposal pending until consumer recovery is removed', async () => {
const { mockStore, mockPortForward, getMainWindow } = createMockDeps()
vi.mocked(mockStore.getSshPtyConsumerRecovery).mockReturnValue({
@@ -232,6 +312,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
+46 -28
View File
@@ -80,6 +80,7 @@ import type { Store } from '../persistence'
import type { OrcaRuntimeService } from '../runtime/orca-runtime'
import { DEFAULT_PTY_SOURCE_WINDOW_SU } from '../../shared/pty-source-credit-contract'
import { PTY_CONSUMER_STALE_OWNER_RECOVERY_ERROR } from '../../shared/pty-consumer-session'
import { retrySshOwnerRecoveryWhileBlocked } from './ssh-owner-recovery-retry'
import { runRemoteOrcaCli } from './ssh-remote-orca-cli'
import {
acknowledgeRemoteOrcaCliPostOutput,
@@ -991,17 +992,24 @@ export class SshRelaySession {
outputFlowControl: { requestedWindowSu: DEFAULT_PTY_SOURCE_WINDOW_SU }
}
try {
return await openSshPtyConsumerSession(mux, {
...options,
...(previousOwner
? {
resume: {
ownerGeneration: previousOwner.ownerGeneration,
ownerLease: previousOwner.ownerLease
}
}
: {})
})
return await retrySshOwnerRecoveryWhileBlocked(
() =>
openSshPtyConsumerSession(mux, {
...options,
...(previousOwner
? {
resume: {
ownerGeneration: previousOwner.ownerGeneration,
ownerLease: previousOwner.ownerLease
}
}
: {})
}),
{
isCurrent: () => ownsAttempt() && !mux.isDisposed(),
onClosed: (listener) => mux.onDispose(listener)
}
)
} catch (error) {
if (
!previousOwner ||
@@ -1010,26 +1018,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')
}
+10
View File
@@ -161,6 +161,16 @@ export class RelayDispatcher {
this.closeClient(client, new Error('Relay client detached'), true)
}
// Why: a displaced owner must lose its transport whichever client holds it, and the launch channel is
// often the half-open one. The primary keeps its id for setWrite() revival, so invalidate it instead.
releaseDisplacedClient(clientId: number): void {
if (clientId === this.primaryClient.id) {
this.invalidateClient()
return
}
this.detachClient(clientId)
}
feedClient(clientId: number, data: Buffer): void {
const client = this.clients.get(clientId)
if (!client) {
@@ -0,0 +1,368 @@
import { afterEach, describe, expect, it, vi } from 'vitest'
import {
RelayDispatcher,
type RelayClientSessionIdentity,
type SinkWriteSettlement
} from './dispatcher'
import { encodeJsonRpcFrame, MessageType } from './protocol'
import { SshPtyConsumerSessionAdapter } from './ssh-pty-consumer-session-adapter'
import {
PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR,
PTY_CONSUMER_OWNER_RECOVERY_SUPERSEDED_ERROR
} from '../shared/pty-consumer-session'
const endpointIdentity: RelayClientSessionIdentity = {
principal: 'endpoint-principal',
authenticated: true,
allowSessionOwner: true,
authenticationKind: 'endpoint-credential'
}
function requestFrame(id: number, method: string, params: Record<string, unknown>): Buffer {
return encodeJsonRpcFrame({ jsonrpc: '2.0', id, method, params }, id, 0)
}
function message(buffer: Buffer): Record<string, unknown> | null {
if (buffer[0] !== MessageType.Regular) {
return null
}
const length = buffer.readUInt32BE(9)
return JSON.parse(buffer.subarray(13, 13 + length).toString('utf8'))
}
function response(writes: readonly Buffer[], id: number): Record<string, unknown> | undefined {
return writes.map(message).find((entry) => entry?.id === id) as
| Record<string, unknown>
| undefined
}
function ownerHelloParams(resume?: Record<string, unknown>): Record<string, unknown> {
return {
protocolVersion: 1,
clientInstanceId: 'client-1',
requestedRole: 'session-owner',
capabilities: { outputFlowControl: { versions: [1], requestedWindowSu: 8 } },
...(resume ? { resume } : {})
}
}
async function flushRequests(): Promise<void> {
await new Promise((resolve) => setImmediate(resolve))
}
describe('relay PTY consumer owner displacement', () => {
let dispatcher: RelayDispatcher | null = null
afterEach(() => {
dispatcher?.dispose()
dispatcher = null
})
it('displaces a half-open owner and retains its deliveries for the reconnected client', async () => {
// Why no settle deferral: a half-open peer is indistinguishable from a live one at the sink — writes
// still "succeed" locally, so the relay never learns the owner is gone.
const staleWrites: Buffer[] = []
const closeStaleTransport = vi.fn()
dispatcher = new RelayDispatcher(
(data, settle) => {
staleWrites.push(Buffer.from(data))
settle({ ok: true })
return true
},
{ supportsWriteCallback: true, close: closeStaleTransport },
endpointIdentity
)
const detached: number[] = []
dispatcher.onClientDetached((clientId) => detached.push(clientId))
const adapter = new SshPtyConsumerSessionAdapter(dispatcher, 'build-a')
dispatcher.feed(requestFrame(1, 'pty.openClient', ownerHelloParams()))
await flushRequests()
const staleGrant = response(staleWrites, 1)!.result as Record<string, unknown>
expect(adapter.deliveryMode(1)).toBe('source-owner')
const staleDelivery = adapter.openDelivery(1, 'pty-1', 'incarnation-1')!
adapter.appendSource(staleDelivery, {
spanId: 'span-1',
data: 'tail',
displayStart: 0,
displayEnd: 4,
splittable: true,
transform: { transformed: false, rawLengthSu: 4, scalarSafe: true }
})
const reconnectWrites: Buffer[] = []
const reconnectClientId = dispatcher.attachClient(
(data, settle) => {
reconnectWrites.push(Buffer.from(data))
settle({ ok: true })
return true
},
{ supportsWriteCallback: true },
endpointIdentity
)
dispatcher.feedClient(
reconnectClientId,
requestFrame(
2,
'pty.openClient',
ownerHelloParams({
ownerGeneration: staleGrant.ownerGeneration,
ownerLease: staleGrant.ownerLease
})
)
)
await flushRequests()
expect(response(reconnectWrites, 2)!.result).toMatchObject({
role: 'session-owner',
ownerGeneration: 2
})
expect(adapter.deliveryMode(reconnectClientId)).toBe('source-owner')
// Why: revocation is what bounds the takeover — the stale owner cannot open or drive deliveries again
// even though its socket never closed.
expect(adapter.deliveryMode(1)).toBe('subscriber')
expect(adapter.openDelivery(1, 'pty-2', 'incarnation-1')).toBeNull()
expect(detached).toContain(1)
expect(closeStaleTransport).toHaveBeenCalledOnce()
// Why: retained, not closed — the reconnected owner still has to rotate the stale delivery forward.
expect(adapter.getDebugSnapshot()).toMatchObject({ deliveryTokens: 1, graceTimers: 1 })
const rotation = adapter.rotateDelivery(staleDelivery, reconnectClientId, 0)
expect(rotation.identity).toMatchObject({ ownerGeneration: 2, clientGeneration: 2 })
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(
(data, settle) => {
incumbentWrites.push(Buffer.from(data))
settle({ ok: true })
return true
},
{ supportsWriteCallback: true },
endpointIdentity
)
const detached: number[] = []
dispatcher.onClientDetached((clientId) => detached.push(clientId))
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 reconnectWrites: Buffer[] = []
let grantSettlement: ((result: SinkWriteSettlement) => void) | undefined
const reconnectClientId = dispatcher.attachClient(
(data, settle) => {
reconnectWrites.push(Buffer.from(data))
// Why: the takeover must not commit until the replacement grant is on the wire.
grantSettlement = settle
return true
},
{ supportsWriteCallback: true },
endpointIdentity
)
dispatcher.feedClient(
reconnectClientId,
requestFrame(
2,
'pty.openClient',
ownerHelloParams({
ownerGeneration: incumbentGrant.ownerGeneration,
ownerLease: incumbentGrant.ownerLease
})
)
)
await flushRequests()
expect(adapter.deliveryMode(1)).toBe('source-owner')
expect(detached).not.toContain(1)
grantSettlement!({ ok: false, error: new Error('reconnect socket closed mid-response') })
expect(adapter.deliveryMode(1)).toBe('source-owner')
expect(adapter.deliveryMode(reconnectClientId)).toBe('subscriber')
// Why: the rolled-back attempt must leave the incumbent's lease reclaimable by the next reconnect.
const retryWrites: Buffer[] = []
const retryClientId = dispatcher.attachClient(
(data, settle) => {
retryWrites.push(Buffer.from(data))
settle({ ok: true })
return true
},
{ supportsWriteCallback: true },
endpointIdentity
)
dispatcher.feedClient(
retryClientId,
requestFrame(
3,
'pty.openClient',
ownerHelloParams({
ownerGeneration: incumbentGrant.ownerGeneration,
ownerLease: incumbentGrant.ownerLease
})
)
)
await flushRequests()
// Why generation 3: the rolled-back attempt consumed generation 2, but not the incumbent's lease.
expect(response(retryWrites, 3)!.result).toMatchObject({
role: 'session-owner',
ownerGeneration: 3
})
expect(adapter.deliveryMode(1)).toBe('subscriber')
expect(detached).toContain(1)
})
it('fences an old proof while its replacement is live', 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 resume = {
ownerGeneration: incumbentGrant.ownerGeneration,
ownerLease: incumbentGrant.ownerLease
}
let firstReconnectSettlement: ((result: SinkWriteSettlement) => void) | undefined
const firstReconnectWrites: Buffer[] = []
const firstReconnectClientId = dispatcher.attachClient(
(data, settle) => {
firstReconnectWrites.push(Buffer.from(data))
// Why ??=: the test settles the grant response, not whatever frame the dispatcher wrote last.
firstReconnectSettlement ??= settle
return true
},
{ supportsWriteCallback: true },
endpointIdentity
)
dispatcher.feedClient(
firstReconnectClientId,
requestFrame(2, 'pty.openClient', ownerHelloParams(resume))
)
await flushRequests()
const retryWrites: Buffer[] = []
const retryClientId = dispatcher.attachClient(
(data, settle) => {
retryWrites.push(Buffer.from(data))
settle({ ok: true })
return true
},
{ supportsWriteCallback: true },
endpointIdentity
)
dispatcher.feedClient(
retryClientId,
requestFrame(3, 'pty.openClient', ownerHelloParams(resume))
)
await flushRequests()
expect(response(retryWrites, 3)!.error).toMatchObject({
code: PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR
})
firstReconnectSettlement!({ ok: true })
dispatcher.feedClient(
retryClientId,
requestFrame(4, 'pty.openClient', ownerHelloParams(resume))
)
await flushRequests()
expect(response(retryWrites, 4)!.error).toMatchObject({
code: PTY_CONSUMER_OWNER_RECOVERY_SUPERSEDED_ERROR
})
expect(adapter.deliveryMode(firstReconnectClientId)).toBe('source-owner')
expect(adapter.deliveryMode(retryClientId)).toBe('subscriber')
dispatcher.detachClient(firstReconnectClientId)
dispatcher.feedClient(
retryClientId,
requestFrame(5, 'pty.openClient', ownerHelloParams(resume))
)
await flushRequests()
expect(response(retryWrites, 5)!.result).toMatchObject({
role: 'session-owner',
ownerGeneration: 3,
ownerLease: incumbentGrant.ownerLease
})
expect(adapter.deliveryMode(retryClientId)).toBe('source-owner')
})
})
+34 -1
View File
@@ -10,7 +10,15 @@
import { createServer, createConnection, type Socket, type Server } from 'node:net'
import { join } from 'node:path'
import { unlinkSync, existsSync, statSync, readFileSync, chmodSync } from 'node:fs'
import {
unlinkSync,
existsSync,
statSync,
readFileSync,
chmodSync,
closeSync,
openSync
} from 'node:fs'
import {
RELAY_SENTINEL,
FrameDecoder,
@@ -596,6 +604,31 @@ async function main(): Promise<void> {
}
stdoutDrainWaiters.add(cb)
return () => stdoutDrainWaiters.delete(cb)
},
close: () => {
stdoutAlive = false
flushStdoutDrainWaiters()
// Why close then re-pin: the SSH peer must see EOF, but a long-lived daemon that
// frees fds 0/1 lets accept()/open() recycle them while Node still treats
// process.stdin/stdout as those numbers — corrupting socket clients and shutdown.
for (const fd of [process.stdin.fd, process.stdout.fd]) {
try {
closeSync(fd)
} catch {
// Already closed by the peer.
}
}
const devNull = process.platform === 'win32' ? 'NUL' : '/dev/null'
try {
openSync(devNull, 'r')
} catch {
/* best-effort pin of the lowest free fd (normally 0) */
}
try {
openSync(devNull, 'w')
} catch {
/* best-effort pin of the next free fd (normally 1) */
}
}
},
undefined,
+20 -4
View File
@@ -1,6 +1,7 @@
import {
PTY_CONSUMER_SESSION_PROTOCOL_VERSION,
PtyConsumerSession,
type PtyConsumerSessionAdmission,
type PtyConsumerSessionGrant,
type PtyConsumerSessionHello
} from '../shared/pty-consumer-session'
@@ -74,7 +75,7 @@ export class SshPtyConsumerSessionAdapter {
private readonly pausedDeliveryByPty = new Map<string, PtySourceDeliveryIdentity>()
constructor(
dispatcher: RelayDispatcher,
private readonly dispatcher: RelayDispatcher,
serverBuildId: string,
private readonly setDeliveryPaused?: (id: string, paused: boolean) => void,
onSourceCreditAvailable?: (id: string) => void
@@ -268,15 +269,30 @@ export class SshPtyConsumerSessionAdapter {
throw new Error('SSH PTY consumer response publication fence is unavailable')
}
context.onResponseSettled((result) => {
if (result.ok) {
admission.commitPublication()
} else {
if (!result.ok) {
admission.rollbackPublication()
return
}
admission.commitPublication()
this.closeDisplacedOwner(admission.displacedOwner)
})
return admission.grant
}
// Why: only after the replacement grant is published — until then the admission can still roll back
// onto the incumbent. Its deliveries are retained for the new owner to rotate, exactly as on detach.
private closeDisplacedOwner(displaced: PtyConsumerSessionAdmission['displacedOwner']): void {
if (!displaced) {
return
}
this.clearPausedForGrant(displaced.grant)
this.sourceCredit.retainOrCloseOnDetach(displaced.grant)
const clientId = Number(displaced.connectionId)
if (Number.isSafeInteger(clientId)) {
this.dispatcher.releaseDisplacedClient(clientId)
}
}
private clearPausedIdentity(identity: PtySourceDeliveryIdentity): void {
if (this.pausedDeliveryByPty.get(identity.id) === identity) {
this.pausedDeliveryByPty.delete(identity.id)
+63
View File
@@ -0,0 +1,63 @@
import {
PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR,
PTY_CONSUMER_OWNER_RECOVERY_SUPERSEDED_ERROR,
PTY_CONSUMER_STALE_OWNER_RECOVERY_ERROR,
type PtyConsumerAuthentication,
type PtyConsumerSessionHello
} from './pty-consumer-session-contract'
type IncumbentOwner = {
principal: string
clientInstanceId: string
generation: number
lease: string
state: 'pending' | 'active' | 'disconnected'
replaces?: { generation: number }
}
function throwRecoveryError(message: string, code: number): never {
throw Object.assign(new Error(message), { code })
}
export function assertPtyConsumerOwnerRecovery(
hello: PtyConsumerSessionHello,
authentication: PtyConsumerAuthentication,
current: IncumbentOwner
): void {
const resume = hello.resume
if (!resume) {
throw new Error('Owner recovery proof is required')
}
if (
resume.ownerLease !== current.lease ||
hello.clientInstanceId !== current.clientInstanceId ||
authentication.principal !== current.principal
) {
throwRecoveryError(
'Owner recovery lease is stale or belongs to another principal',
PTY_CONSUMER_STALE_OWNER_RECOVERY_ERROR
)
}
if (current.state === 'active' && resume.ownerGeneration < current.generation) {
throwRecoveryError(
'Owner recovery generation was superseded',
PTY_CONSUMER_OWNER_RECOVERY_SUPERSEDED_ERROR
)
}
const generationMatches =
resume.ownerGeneration === current.generation ||
(current.state === 'pending' && resume.ownerGeneration === current.replaces?.generation) ||
(current.state === 'disconnected' && resume.ownerGeneration < current.generation)
if (!generationMatches) {
throwRecoveryError(
'Owner recovery generation is stale',
PTY_CONSUMER_STALE_OWNER_RECOVERY_ERROR
)
}
if (current.state === 'pending') {
throwRecoveryError(
'Owner grant publication is still pending',
PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR
)
}
}
@@ -1,6 +1,10 @@
export const PTY_CONSUMER_SESSION_PROTOCOL_VERSION = 1
export const PTY_CONSUMER_OWNER_GRACE_MS = 30_000
export const PTY_CONSUMER_STALE_OWNER_RECOVERY_ERROR = -32041
// Why: recovery is blocked only while the incumbent owner's grant publication is still settling — a
// window bounded by one response write, so the client may retry within a short budget.
export const PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR = -32042
export const PTY_CONSUMER_OWNER_RECOVERY_SUPERSEDED_ERROR = -32043
export type PtyConsumerRole = 'session-owner' | 'subscriber'
@@ -41,8 +45,16 @@ export type PtyConsumerAuthentication = {
allowSessionOwner: boolean
}
export type PtyConsumerDisplacedOwner = {
connectionId: string
grant: Readonly<PtyConsumerSessionGrant>
}
export type PtyConsumerSessionAdmission = {
grant: Readonly<PtyConsumerSessionGrant>
// Why: set when this admission takes over a still-attached owner. The transport layer owns closing
// that connection and releasing its deliveries — do it only once the new grant has been published.
displacedOwner?: Readonly<PtyConsumerDisplacedOwner>
commitPublication: () => void
rollbackPublication: () => void
}
+52
View File
@@ -0,0 +1,52 @@
import type { PtyConsumerSessionHello } from './pty-consumer-session-contract'
export const MAX_CAPABILITY_VERSIONS = 8
export function assertNonEmptyString(value: unknown, name: string): asserts value is string {
if (typeof value !== 'string' || value.length === 0 || value.length > 512) {
throw new Error(`${name} must be a non-empty string of at most 512 characters`)
}
}
export function validateHello(hello: PtyConsumerSessionHello): void {
assertNonEmptyString(hello.clientInstanceId, 'clientInstanceId')
if (hello.requestedRole !== 'session-owner' && hello.requestedRole !== 'subscriber') {
throw new Error('requestedRole must be session-owner or subscriber')
}
if (hello.resume) {
if (!Number.isSafeInteger(hello.resume.ownerGeneration) || hello.resume.ownerGeneration <= 0) {
throw new Error('resume.ownerGeneration must be a positive safe integer')
}
assertNonEmptyString(hello.resume.ownerLease, 'resume.ownerLease')
}
const flow = hello.capabilities?.outputFlowControl
if (!flow) {
return
}
if (
!Array.isArray(flow.versions) ||
flow.versions.length > MAX_CAPABILITY_VERSIONS ||
flow.versions.some((version) => !Number.isSafeInteger(version) || version <= 0)
) {
throw new Error('outputFlowControl.versions must contain positive safe integers')
}
if (!Number.isSafeInteger(flow.requestedWindowSu) || flow.requestedWindowSu <= 0) {
throw new Error('outputFlowControl.requestedWindowSu must be a positive safe integer')
}
}
// Why: a duplicate open on one connection is only idempotent if it asks for exactly the same thing.
export function helloFingerprint(hello: PtyConsumerSessionHello): string {
const flow = hello.capabilities?.outputFlowControl
return JSON.stringify({
clientInstanceId: hello.clientInstanceId,
requestedRole: hello.requestedRole,
resume: hello.resume,
outputFlowControl: flow
? {
versions: [...flow.versions].sort((a, b) => a - b),
requestedWindowSu: flow.requestedWindowSu
}
: undefined
})
}
+181 -5
View File
@@ -1,5 +1,7 @@
import { describe, expect, it } from 'vitest'
import {
PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR,
PTY_CONSUMER_OWNER_RECOVERY_SUPERSEDED_ERROR,
PTY_CONSUMER_STALE_OWNER_RECOVERY_ERROR,
PtyConsumerSession,
type PtyConsumerAuthentication,
@@ -121,7 +123,7 @@ describe('PtyConsumerSession', () => {
).toThrow('authentication required')
})
it('rotates the lease and increments owner generation on valid recovery', () => {
it('keeps the lease stable and increments owner generation on valid recovery', () => {
const session = createSession()
const first = session.admit(ownerHello(), auth('connection-1'))
first.commitPublication()
@@ -141,15 +143,97 @@ describe('PtyConsumerSession', () => {
expect(recovered.grant).toMatchObject({
role: 'session-owner',
ownerGeneration: 2,
ownerLease: 'lease-2'
ownerLease: 'lease-1'
})
})
it('rejects recovery while the current owner is still active', () => {
it('displaces a still-attached owner that a matching resume proof reclaims', () => {
const session = createSession()
const first = session.admit(ownerHello(), auth('connection-1'))
first.commitPublication()
const recovered = session.admit(
ownerHello({
resume: {
ownerGeneration: first.grant.ownerGeneration!,
ownerLease: first.grant.ownerLease!
}
}),
auth('connection-2')
)
expect(recovered.displacedOwner).toEqual({
connectionId: 'connection-1',
grant: first.grant
})
// Why: the incumbent keeps authority until the replacement grant is actually published.
expect(session.activeGrant('connection-1')).toBe(first.grant)
recovered.commitPublication()
expect(recovered.grant).toMatchObject({ role: 'session-owner', ownerGeneration: 2 })
expect(session.activeGrant('connection-1')).toBeNull()
expect(session.activeGrant('connection-2')).toBe(recovered.grant)
})
it('restores the displaced owner when the replacement publication rolls back', () => {
const session = createSession()
const first = session.admit(ownerHello(), auth('connection-1'))
first.commitPublication()
const recovered = session.admit(
ownerHello({
resume: {
ownerGeneration: first.grant.ownerGeneration!,
ownerLease: first.grant.ownerLease!
}
}),
auth('connection-2')
)
recovered.rollbackPublication()
expect(session.activeGrant('connection-1')).toBe(first.grant)
expect(session.activeGrant('connection-2')).toBeNull()
// Why: the restored incumbent must still hold the lease it was admitted with.
const reclaimed = session.admit(
ownerHello({
resume: {
ownerGeneration: first.grant.ownerGeneration!,
ownerLease: first.grant.ownerLease!
}
}),
auth('connection-3')
)
expect(reclaimed.displacedOwner?.connectionId).toBe('connection-1')
})
it('expires a displaced owner restored after its connection already closed', () => {
let now = 10
const session = createSession({ now: () => now })
const first = session.admit(ownerHello(), auth('connection-1'))
first.commitPublication()
const recovered = session.admit(
ownerHello({
resume: {
ownerGeneration: first.grant.ownerGeneration!,
ownerLease: first.grant.ownerLease!
}
}),
auth('connection-2')
)
session.close('connection-1')
recovered.rollbackPublication()
now += 30_000
session.sweepExpired()
const fresh = session.admit(ownerHello(), auth('connection-4'))
expect(fresh.grant.role).toBe('session-owner')
})
it('refuses recovery while the incumbent grant publication is still settling', () => {
const session = createSession()
const first = session.admit(ownerHello(), auth('connection-1'))
expect(() =>
session.admit(
ownerHello({
@@ -160,7 +244,99 @@ describe('PtyConsumerSession', () => {
}),
auth('connection-2')
)
).toThrow('Active owner')
).toThrow(
expect.objectContaining({
code: PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR,
message: expect.stringContaining('still pending')
})
)
})
it('fences an old recovery generation after its replacement commits', () => {
const session = createSession()
const first = session.admit(ownerHello(), auth('connection-1'))
first.commitPublication()
const resume = {
ownerGeneration: first.grant.ownerGeneration!,
ownerLease: first.grant.ownerLease!
}
const replacement = session.admit(ownerHello({ resume }), auth('connection-2'))
expect(() => session.admit(ownerHello({ resume }), auth('connection-3'))).toThrow(
expect.objectContaining({ code: PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR })
)
replacement.commitPublication()
expect(() => session.admit(ownerHello({ resume }), auth('connection-3'))).toThrow(
expect.objectContaining({ code: PTY_CONSUMER_OWNER_RECOVERY_SUPERSEDED_ERROR })
)
const retry = session.admit(
ownerHello({
resume: {
ownerGeneration: replacement.grant.ownerGeneration!,
ownerLease: replacement.grant.ownerLease!
}
}),
auth('connection-3')
)
expect(retry.grant).toMatchObject({
role: 'session-owner',
ownerGeneration: 3,
ownerLease: first.grant.ownerLease
})
expect(retry.displacedOwner?.connectionId).toBe('connection-2')
// Why: the committed replacement already retired the incumbent, so only connection-2 is left to displace.
expect(session.activeGrant('connection-1')).toBeNull()
})
it('accepts an older stable proof after its replacement disconnects', () => {
const session = createSession()
const first = session.admit(ownerHello(), auth('connection-1'))
first.commitPublication()
const resume = {
ownerGeneration: first.grant.ownerGeneration!,
ownerLease: first.grant.ownerLease!
}
const replacement = session.admit(ownerHello({ resume }), auth('connection-2'))
replacement.commitPublication()
session.close('connection-2')
const retry = session.admit(ownerHello({ resume }), auth('connection-3'))
expect(retry.grant).toMatchObject({
role: 'session-owner',
ownerGeneration: 3,
ownerLease: first.grant.ownerLease
})
})
it('retries overlapping recovery against the incumbent after publication rolls back', () => {
const session = createSession()
const first = session.admit(ownerHello(), auth('connection-1'))
first.commitPublication()
const resume = {
ownerGeneration: first.grant.ownerGeneration!,
ownerLease: first.grant.ownerLease!
}
const replacement = session.admit(ownerHello({ resume }), auth('connection-2'))
expect(() => session.admit(ownerHello({ resume }), auth('connection-3'))).toThrow(
expect.objectContaining({ code: PTY_CONSUMER_OWNER_RECOVERY_PENDING_ERROR })
)
replacement.rollbackPublication()
const retry = session.admit(ownerHello({ resume }), auth('connection-3'))
expect(retry.grant).toMatchObject({
role: 'session-owner',
ownerGeneration: 3,
ownerLease: first.grant.ownerLease
})
expect(retry.displacedOwner?.connectionId).toBe('connection-1')
// Why: the rolled-back replacement never published, so it holds no grant the retry could displace.
expect(session.activeGrant('connection-2')).toBeNull()
})
it('types mismatched recovery without disturbing principal or lease ownership', () => {
@@ -201,7 +377,7 @@ describe('PtyConsumerSession', () => {
expect(recovered.grant).toMatchObject({
role: 'session-owner',
ownerGeneration: 2,
ownerLease: 'lease-2'
ownerLease: 'lease-1'
})
})
+59 -71
View File
@@ -4,22 +4,28 @@ import {
PTY_CONSUMER_SESSION_PROTOCOL_VERSION,
PTY_CONSUMER_STALE_OWNER_RECOVERY_ERROR,
type PtyConsumerAuthentication,
type PtyConsumerDisplacedOwner,
type PtyConsumerSessionAdmission,
type PtyConsumerSessionGrant,
type PtyConsumerSessionHello,
type PtyConsumerSessionOptions
} from './pty-consumer-session-contract'
import {
assertNonEmptyString,
helloFingerprint,
MAX_CAPABILITY_VERSIONS,
validateHello
} from './pty-consumer-session-hello'
import { assertPtyConsumerOwnerRecovery } from './pty-consumer-owner-recovery'
export * from './pty-consumer-session-contract'
const MAX_CAPABILITY_VERSIONS = 8
type ClientRecord = {
fingerprint: string
principal: string
clientInstanceId: string
grant: Readonly<PtyConsumerSessionGrant>
state: 'pending' | 'active'
state: 'pending' | 'active' | 'displaced'
publicationState: 'pending' | 'committed' | 'rolled-back'
}
@@ -34,54 +40,6 @@ type OwnerRecord = {
replaces?: OwnerRecord
}
function assertNonEmptyString(value: unknown, name: string): asserts value is string {
if (typeof value !== 'string' || value.length === 0 || value.length > 512) {
throw new Error(`${name} must be a non-empty string of at most 512 characters`)
}
}
function validateHello(hello: PtyConsumerSessionHello): void {
assertNonEmptyString(hello.clientInstanceId, 'clientInstanceId')
if (hello.requestedRole !== 'session-owner' && hello.requestedRole !== 'subscriber') {
throw new Error('requestedRole must be session-owner or subscriber')
}
if (hello.resume) {
if (!Number.isSafeInteger(hello.resume.ownerGeneration) || hello.resume.ownerGeneration <= 0) {
throw new Error('resume.ownerGeneration must be a positive safe integer')
}
assertNonEmptyString(hello.resume.ownerLease, 'resume.ownerLease')
}
const flow = hello.capabilities?.outputFlowControl
if (!flow) {
return
}
if (
!Array.isArray(flow.versions) ||
flow.versions.length > MAX_CAPABILITY_VERSIONS ||
flow.versions.some((version) => !Number.isSafeInteger(version) || version <= 0)
) {
throw new Error('outputFlowControl.versions must contain positive safe integers')
}
if (!Number.isSafeInteger(flow.requestedWindowSu) || flow.requestedWindowSu <= 0) {
throw new Error('outputFlowControl.requestedWindowSu must be a positive safe integer')
}
}
function helloFingerprint(hello: PtyConsumerSessionHello): string {
const flow = hello.capabilities?.outputFlowControl
return JSON.stringify({
clientInstanceId: hello.clientInstanceId,
requestedRole: hello.requestedRole,
resume: hello.resume,
outputFlowControl: flow
? {
versions: [...flow.versions].sort((a, b) => a - b),
requestedWindowSu: flow.requestedWindowSu
}
: undefined
})
}
export class PtyConsumerSession {
private readonly clients = new Map<string, ClientRecord>()
private readonly now: () => number
@@ -160,7 +118,7 @@ export class PtyConsumerSession {
if (owner) {
this.owner = owner
}
return this.admissionFor(client)
return this.admissionFor(client, this.displacedOwnerFor(owner))
}
close(connectionId: string): void {
@@ -170,6 +128,17 @@ export class PtyConsumerSession {
}
this.clients.delete(connectionId)
if (this.owner?.connectionId !== connectionId) {
// Why: a pending replacement can still roll back onto the owner it is displacing; restoring an
// 'active' record whose connection has since closed would wedge an owner that can never expire.
if (
this.owner?.replaces?.connectionId === connectionId &&
this.owner.replaces.state === 'active'
) {
this.owner = {
...this.owner,
replaces: { ...this.owner.replaces, state: 'disconnected', disconnectedAt: this.now() }
}
}
return
}
if (this.owner.state === 'pending') {
@@ -192,9 +161,13 @@ export class PtyConsumerSession {
return client?.state === 'active' ? client.grant : null
}
private admissionFor(client: ClientRecord): PtyConsumerSessionAdmission {
private admissionFor(
client: ClientRecord,
displacedOwner?: Readonly<PtyConsumerDisplacedOwner>
): PtyConsumerSessionAdmission {
return {
grant: client.grant,
...(displacedOwner ? { displacedOwner } : {}),
commitPublication: () => {
if (client.publicationState !== 'pending') {
return
@@ -206,6 +179,7 @@ export class PtyConsumerSession {
client.state = 'active'
const owner = this.owner
if (owner?.connectionId === this.connectionIdFor(client) && owner.state === 'pending') {
this.retireDisplacedOwner(owner.replaces)
this.owner = { ...owner, state: 'active', replaces: undefined }
}
},
@@ -254,32 +228,46 @@ export class PtyConsumerSession {
if (!hello.resume) {
return null
}
const recoveryMatches =
hello.resume.ownerGeneration === current.generation &&
hello.resume.ownerLease === current.lease &&
hello.clientInstanceId === current.clientInstanceId &&
authentication.principal === current.principal
if (!recoveryMatches) {
throw Object.assign(
new Error('Owner recovery lease is stale or belongs to another principal'),
{ code: PTY_CONSUMER_STALE_OWNER_RECOVERY_ERROR }
)
}
if (current.state === 'pending') {
throw new Error('Owner grant publication is still pending')
}
if (current.state !== 'disconnected') {
throw new Error('Active owner cannot be replaced by recovery')
}
assertPtyConsumerOwnerRecovery(hello, authentication, current)
// Why an active owner is displaced rather than refused: the resume proof matched this owner's
// generation, lease, client instance, and principal on a *different* transport, so the requester is
// the same logical owner reconnecting. Waiting for the incumbent's socket to close is unbounded —
// a half-open connection after sleep/resume or NAT loss never gets there.
return this.newOwner(hello, authentication, current)
}
private displacedOwnerFor(
owner: OwnerRecord | null
): Readonly<PtyConsumerDisplacedOwner> | undefined {
const replaced = owner?.replaces
if (replaced?.state !== 'active') {
return undefined
}
const client = this.clients.get(replaced.connectionId)
if (client?.state !== 'active') {
return undefined
}
return Object.freeze({ connectionId: replaced.connectionId, grant: client.grant })
}
// Why: the displaced connection may still be writable (half-open), so revoke its grant the moment the
// replacement is published — a stale owner must not keep driving deliveries under the old generation.
private retireDisplacedOwner(replaced: OwnerRecord | undefined): void {
if (replaced?.state !== 'active') {
return
}
const client = this.clients.get(replaced.connectionId)
if (client?.state === 'active') {
client.state = 'displaced'
}
}
private newOwner(
hello: PtyConsumerSessionHello,
authentication: PtyConsumerAuthentication,
replaces: OwnerRecord | null
): OwnerRecord {
const lease = this.createLease()
const lease = replaces?.lease ?? this.createLease()
assertNonEmptyString(lease, 'ownerLease')
return {
connectionId: authentication.connectionId,