From 669940ccf7ff63a8c50a88cea54fde58ccdd25fb Mon Sep 17 00:00:00 2001 From: Jinjing <6427696+AmethystLiang@users.noreply.github.com> Date: Sat, 1 Aug 2026 15:45:19 -0700 Subject: [PATCH] refactor(relay): extract grace reconfiguration to testable function Extract `applyRelayGraceTimeConfiguration` so the grace-time re-arming logic can be tested independently. relay.ts runs `main()` on import and exports nothing, making the call site (including retryDeferredShutdown hand-off into startGrace) otherwise untestable. --- src/relay/relay-grace-branch.test.ts | 106 ++++++++++++++++++++++++++- src/relay/relay-grace-branch.ts | 42 +++++++++++ src/relay/relay.ts | 29 +++----- 3 files changed, 156 insertions(+), 21 deletions(-) diff --git a/src/relay/relay-grace-branch.test.ts b/src/relay/relay-grace-branch.test.ts index 8e872cb5301..e7fee583239 100644 --- a/src/relay/relay-grace-branch.test.ts +++ b/src/relay/relay-grace-branch.test.ts @@ -1,7 +1,9 @@ -import { describe, expect, it } from 'vitest' +import { describe, expect, it, vi } from 'vitest' import { + applyRelayGraceTimeConfiguration, decideRelayGrace, decideRelayGraceReconfigure, + type RelayGraceBranch, type RelayGraceDecisionInput, type RelayGraceReconfigureInput } from './relay-grace-branch' @@ -148,3 +150,105 @@ describe('decideRelayGraceReconfigure', () => { expect(reconfigure({ shutdownInFlight: true })).toEqual({ rearm: false }) }) }) + +/** + * Models relay.ts's grace state around the configureGraceTime call site: `startGrace` re-decides the + * branch exactly as the relay does, so a dropped `retryDeferredShutdown` shows up as a branch downgrade. + */ +function relayGraceHost( + overrides: { + configuredGraceMs?: number + graceTimerArmed?: boolean + shutdownInFlight?: boolean + graceBranch?: RelayGraceBranch | null + relayIdle?: boolean + activePtyCount?: number + } = {} +) { + let configuredGraceMs = overrides.configuredGraceMs ?? 10_000 + let graceBranch: RelayGraceBranch | null = overrides.graceBranch ?? 'configured' + const startGrace = vi.fn( + (_reason: string, options?: { retryDeferredShutdown?: boolean }): void => { + graceBranch = decide({ + configuredGraceMs, + relayIdle: overrides.relayIdle ?? true, + activePtyCount: overrides.activePtyCount ?? 0, + retryDeferredShutdown: options?.retryDeferredShutdown === true + }).branch + } + ) + + return { + startGrace, + branch: () => graceBranch, + configuredGraceMs: () => configuredGraceMs, + apply: (graceTimeSeconds: unknown) => + applyRelayGraceTimeConfiguration(graceTimeSeconds, { + readConfiguredGraceMs: () => configuredGraceMs, + // Mirrors PtyHandler.setGraceTimeMs's clamp so the decision sees the stored value. + writeConfiguredGraceMs: (graceMs) => { + configuredGraceMs = Math.max(0, Math.floor(graceMs)) + }, + isGraceTimerArmed: () => overrides.graceTimerArmed ?? true, + isShutdownInFlight: () => overrides.shutdownInFlight ?? false, + readGraceBranch: () => graceBranch, + startGrace + }) + } +} + +describe('applyRelayGraceTimeConfiguration', () => { + it('re-arms a deferred shutdown without downgrading it to an ordinary configured window', () => { + // Why: the call site the branch selector alone cannot cover — dropping the startGrace option here + // would leave a refused kill with nothing left to retry it. + const host = relayGraceHost({ graceBranch: 'shutdown-deferred', configuredGraceMs: 0 }) + + expect(host.apply(86_400)).toEqual({ graceTimeMs: 86_400_000 }) + expect(host.startGrace).toHaveBeenCalledWith('grace reconfigured', { + retryDeferredShutdown: true + }) + expect(host.branch()).toBe('shutdown-deferred') + }) + + it('re-arms an ordinary window at the raised grace', () => { + const host = relayGraceHost() + + expect(host.apply(86_400)).toEqual({ graceTimeMs: 86_400_000 }) + expect(host.startGrace).toHaveBeenCalledWith('grace reconfigured', { + retryDeferredShutdown: false + }) + expect(host.branch()).toBe('configured') + }) + + it('stores the host-sleep zero and lets the idle cap own the window', () => { + const host = relayGraceHost() + + expect(host.apply(0)).toEqual({ graceTimeMs: 0 }) + expect(host.startGrace).toHaveBeenCalledTimes(1) + expect(host.branch()).toBe('idle-no-ptys') + }) + + it('does not re-arm on a re-asserted grace', () => { + const host = relayGraceHost({ configuredGraceMs: 86_400_000 }) + + expect(host.apply(86_400)).toEqual({ graceTimeMs: 86_400_000 }) + expect(host.startGrace).not.toHaveBeenCalled() + }) + + it('does not arm a window that was never running', () => { + const host = relayGraceHost({ graceTimerArmed: false }) + + expect(host.apply(86_400)).toEqual({ graceTimeMs: 86_400_000 }) + expect(host.startGrace).not.toHaveBeenCalled() + }) + + it('ignores a malformed or negative payload without touching the stored grace', () => { + for (const payload of [undefined, 'later', Number.NaN, -1]) { + const host = relayGraceHost() + + expect(host.apply(payload)).toEqual({ graceTimeMs: 10_000 }) + expect(host.configuredGraceMs()).toBe(10_000) + expect(host.startGrace).not.toHaveBeenCalled() + } + }) +}) diff --git a/src/relay/relay-grace-branch.ts b/src/relay/relay-grace-branch.ts index 0c57f00f82e..68ba3e2ec9d 100644 --- a/src/relay/relay-grace-branch.ts +++ b/src/relay/relay-grace-branch.ts @@ -36,6 +36,48 @@ export type RelayGraceReconfigureDecision = | { rearm: false } | { rearm: true; retryDeferredShutdown: boolean } +/** Live relay grace state, read through accessors so the caller's `let` bindings stay authoritative. */ +export type RelayGraceTimeConfigurationInput = { + readConfiguredGraceMs: () => number + /** Writes the new grace; the value is read back so a normalizing setter wins. */ + writeConfiguredGraceMs: (graceMs: number) => void + isGraceTimerArmed: () => boolean + isShutdownInFlight: () => boolean + readGraceBranch: () => RelayGraceBranch | null + startGrace: (reason: string, options?: { retryDeferredShutdown?: boolean }) => void +} + +/** + * Applies a `relay.configureGraceTime` payload, re-arming the running window when the value changed. + * + * Why extracted from relay.ts: that file runs `main()` on import and exports nothing, so the call site + * — including the `retryDeferredShutdown` hand-off into `startGrace` — is otherwise untestable. + */ +export function applyRelayGraceTimeConfiguration( + graceTimeSeconds: unknown, + state: RelayGraceTimeConfigurationInput +): { graceTimeMs: number } { + const seconds = Number(graceTimeSeconds) + if (Number.isFinite(seconds) && seconds >= 0) { + const previousConfiguredGraceMs = state.readConfiguredGraceMs() + // Why: the host sends 0 before system sleep so live remote PTYs survive longer than the ordinary grace window. + state.writeConfiguredGraceMs(Math.floor(seconds) * 1000) + const reconfigure = decideRelayGraceReconfigure({ + previousConfiguredGraceMs, + nextConfiguredGraceMs: state.readConfiguredGraceMs(), + graceTimerArmed: state.isGraceTimerArmed(), + shutdownInFlight: state.isShutdownInFlight(), + currentBranch: state.readGraceBranch() + }) + if (reconfigure.rearm) { + state.startGrace('grace reconfigured', { + retryDeferredShutdown: reconfigure.retryDeferredShutdown + }) + } + } + return { graceTimeMs: state.readConfiguredGraceMs() } +} + /** * Decides whether a live grace change must re-arm the running grace timer. * diff --git a/src/relay/relay.ts b/src/relay/relay.ts index 8bdc072a0cf..4bfc8bf3e92 100644 --- a/src/relay/relay.ts +++ b/src/relay/relay.ts @@ -52,8 +52,8 @@ import { import { resolveSetupAgentSequenceLaunchCommand } from '../shared/setup-agent-sequencing' import { pickRemoteCliEnv } from './remote-cli-env' import { + applyRelayGraceTimeConfiguration, decideRelayGrace, - decideRelayGraceReconfigure, type RelayGraceBranch } from './relay-grace-branch' import { relayLogLine } from './relay-diagnostic-log' @@ -685,25 +685,14 @@ async function main(): Promise { }) function configureRelayGraceTime(params: Record): { graceTimeMs: number } { - const seconds = Number(params.graceTimeSeconds) - if (Number.isFinite(seconds) && seconds >= 0) { - const previousGraceMs = ptyHandler.configuredGraceTimeMs - // Why: the host sends 0 before system sleep so live remote PTYs survive longer than the ordinary grace window. - ptyHandler.setGraceTimeMs(Math.floor(seconds) * 1000) - const reconfigure = decideRelayGraceReconfigure({ - previousConfiguredGraceMs: previousGraceMs, - nextConfiguredGraceMs: ptyHandler.configuredGraceTimeMs, - graceTimerArmed: graceDeadlineAt !== null && graceReason !== null, - shutdownInFlight, - currentBranch: graceBranch - }) - if (reconfigure.rearm) { - startGrace('grace reconfigured', { - retryDeferredShutdown: reconfigure.retryDeferredShutdown - }) - } - } - return { graceTimeMs: ptyHandler.configuredGraceTimeMs } + return applyRelayGraceTimeConfiguration(params.graceTimeSeconds, { + readConfiguredGraceMs: () => ptyHandler.configuredGraceTimeMs, + writeConfiguredGraceMs: (graceMs) => ptyHandler.setGraceTimeMs(graceMs), + isGraceTimerArmed: () => graceDeadlineAt !== null && graceReason !== null, + isShutdownInFlight: () => shutdownInFlight, + readGraceBranch: () => graceBranch, + startGrace + }) } dispatcher.onNotification(SSH_RELAY_CONFIGURE_GRACE_TIME_METHOD, (params) => {