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

@@ -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()
+    }
+  })
+})

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

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<void> {
   })
 
   function configureRelayGraceTime(params: Record<string, unknown>): { 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) => {