From b4e74e66f6ce37de044b1ac26965c6207de0ac57 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Micha=C5=82=20Pierzcha=C5=82a?= Date: Mon, 2 Nov 2026 10:22:28 +0100 Subject: [PATCH] fix(daemon): keep idle reap closed during shutdown --- .../http/daemon-runtime-quiesce.test.ts | 94 +++++++++++++++++++ .../daemon/src/server/http/daemon-runtime.ts | 9 +- .../src/sessions/server/daemon-idle-reap.ts | 23 ++++- 3 files changed, 117 insertions(+), 9 deletions(-) diff --git a/packages/daemon/src/server/http/daemon-runtime-quiesce.test.ts b/packages/daemon/src/server/http/daemon-runtime-quiesce.test.ts index f6816220cc..7ed9248786 100644 --- a/packages/daemon/src/server/http/daemon-runtime-quiesce.test.ts +++ b/packages/daemon/src/server/http/daemon-runtime-quiesce.test.ts @@ -22,12 +22,15 @@ const probe = vi.hoisted(() => ({ held: undefined as Promise | undefined, entered: undefined as (() => void) | undefined, diagnosticFailure: undefined as Error | undefined, + handleRequest: undefined as import('../../core/daemon-request.ts').DaemonInvokeFn | undefined, + idleReaps: 0, })); vi.mock('@agent-device/host-kit/diagnostics', async (importOriginal) => { const actual = await importOriginal(); return { ...actual, emitDiagnostic: (...args: Parameters) => { + if (args[0].phase === 'daemon_idle_reap') probe.idleReaps += 1; if (probe.diagnosticFailure) { if (args[0].phase === 'daemon_fatal') throw probe.diagnosticFailure; throw new Error('secondary diagnostic failure'); @@ -37,6 +40,17 @@ vi.mock('@agent-device/host-kit/diagnostics', async (importOriginal) => { }; }); +vi.mock('./transport.ts', async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + createSocketServer: (handler: import('../../core/daemon-request.ts').DaemonInvokeFn) => { + probe.handleRequest = handler; + return actual.createSocketServer(handler); + }, + }; +}); + vi.mock('@agent-device/daemon-contracts/device-claims', async (importOriginal) => { const actual = await importOriginal(); @@ -104,6 +118,9 @@ afterEach(() => { probe.store = undefined; probe.registry = undefined; probe.diagnosticFailure = undefined; + probe.handleRequest = undefined; + probe.idleReaps = 0; + vi.useRealTimers(); vi.restoreAllMocks(); }); @@ -255,6 +272,83 @@ test('a daemon with a request in flight refuses to quiesce and keeps serving', a } }); +test('request completion after close never re-arms idle reap', async () => { + vi.useFakeTimers(); + const daemon = await bootDaemon({ idleTimeoutMs: 80 }); + let release: () => void = () => {}; + const entered = new Promise((resolve) => { + probe.entered = resolve; + }); + probe.held = new Promise((resolve) => { + release = resolve; + }); + const running = daemon.runtime.run(() => + probe.handleRequest!({ + token: daemon.runtime.token, + session: 'default', + command: 'session_list', + positionals: [], + }), + ); + await entered; + + try { + const closing = daemon.runtime.close(); + release(); + expect((await running).ok).toBe(true); + await vi.runOnlyPendingTimersAsync(); + expect((await closing).requests).toEqual({ drained: 1, abandoned: 0 }); + await vi.advanceTimersByTimeAsync(81); + expect(probe.idleReaps).toBe(0); + } finally { + release(); + await running; + await vi.runOnlyPendingTimersAsync(); + await daemon.runtime.close(); + } +}); + +test('quiesce during an external close reports draining while a request is still running', async () => { + const daemon = await bootDaemon(); + let release: () => void = () => {}; + const entered = new Promise((resolve) => { + probe.entered = resolve; + }); + probe.held = new Promise((resolve) => { + release = resolve; + }); + const running = daemon.runtime.run(() => + probe.handleRequest!({ + token: daemon.runtime.token, + session: 'default', + command: 'session_list', + positionals: [], + }), + ); + await entered; + const closing = daemon.runtime.close(); + + try { + const response = await daemon.runtime.run(() => + probe.handleRequest!( + buildDaemonLifecycleRequest({ + token: daemon.runtime.token, + fence: { + pid: daemon.runtime.identity.pid, + processStartTime: daemon.runtime.identity.startTime, + }, + requestId: 'quiesce-during-external-close', + }), + ), + ); + expect(response).toEqual({ ok: true, data: { outcome: 'draining' } }); + } finally { + release(); + await running; + await closing; + } +}); + test('a daemon holding a session, a repair session or a live lease refuses to quiesce', async () => { const daemon = await bootDaemon(); try { diff --git a/packages/daemon/src/server/http/daemon-runtime.ts b/packages/daemon/src/server/http/daemon-runtime.ts index 1aa1fce621..ed14188bf8 100644 --- a/packages/daemon/src/server/http/daemon-runtime.ts +++ b/packages/daemon/src/server/http/daemon-runtime.ts @@ -52,7 +52,6 @@ import type { SessionRef } from '../../core/session-state.ts'; import { createDaemonIdleReap, createDaemonLifecycleAdmission, - daemonShuttingDownResponse, } from '../../sessions/server/daemon-idle-reap.ts'; import { withDaemonDiagnosticsScope } from '@agent-device/daemon-contracts/daemon-diagnostics-scope'; import { createSessionIdleExpiry } from './daemon-session-idle-expiry.ts'; @@ -451,7 +450,6 @@ async function startDaemonRuntimeOnHost( // in-flight requests, no active recording) past AGENT_DEVICE_DAEMON_IDLE_TIMEOUT_MS. // `close` is defined below but only invoked asynchronously by the timer, // well after this closure captures it. - let shuttingDown = false; const inFlightRequests = new Set>(); // Who owns this process decides what it publishes and which lifecycle request it answers. A host // that starts the runtime passes `embedded`; the daemon executable leaves it detached. @@ -468,7 +466,7 @@ async function startDaemonRuntimeOnHost( }); rollback.onRollback(() => { - idleReap.cancel(); + idleReap.close(); sessionIdleExpiry.cancel(); leaseSessionEnd.cancel(); }); @@ -491,7 +489,6 @@ async function startDaemonRuntimeOnHost( }); const handleRequest: DaemonInvokeFn = lifecycleAdmission.admit(async (req) => { - if (shuttingDown) return daemonShuttingDownResponse(); idleReap.cancel(); const dispatch = runtimeState.run(() => dispatchRequest(req)); inFlightRequests.add(dispatch); @@ -727,11 +724,11 @@ async function startDaemonRuntimeOnHost( }; const closeSteps = async (closeOptions: { cause?: unknown }): Promise => { - idleReap.cancel(); + lifecycleAdmission.close(); + idleReap.close(); sessionIdleExpiry.cancel(); leaseSessionEnd.cancel(); clearInterval(leaseExpirySweep); - shuttingDown = true; stopMetadataLossWatch(); // Cancellation and the listener close run together: a connection held by a request that ignores // cancellation is force-closed at the same budget the drain gives it. diff --git a/packages/daemon/src/sessions/server/daemon-idle-reap.ts b/packages/daemon/src/sessions/server/daemon-idle-reap.ts index a63abc271e..b1a763cc34 100644 --- a/packages/daemon/src/sessions/server/daemon-idle-reap.ts +++ b/packages/daemon/src/sessions/server/daemon-idle-reap.ts @@ -100,6 +100,8 @@ export type DaemonIdleReapController = { */ noteActivity: () => void; cancel: () => void; + /** Permanently stops idle reap, including activity that settles during shutdown. */ + close: () => void; readonly idleMs: number; }; @@ -123,6 +125,7 @@ export function createDaemonIdleReap(params: { ownership?: DaemonRuntimeOwnership; }): DaemonIdleReapController { const idleMs = resolveDaemonIdleReapMs(params.env, params.ownership); + let phase: 'open' | 'closed' = 'open'; let timer: NodeJS.Timeout | undefined; const isIdleNow = (): boolean => @@ -139,13 +142,13 @@ export function createDaemonIdleReap(params: { const schedule = (): void => { cancel(); - if (idleMs <= 0) return; + if (phase === 'closed' || idleMs <= 0) return; timer = setTimeout(() => { timer = undefined; // Re-check at fire time instead of trusting the scheduling snapshot: a // session open or new request must never lose a race against a // previously scheduled reap. - if (!isIdleNow()) return; + if (phase === 'closed' || !isIdleNow()) return; if (params.hasRetainedLeases()) { schedule(); return; @@ -157,11 +160,20 @@ export function createDaemonIdleReap(params: { }; const noteActivity = (): void => { + if (phase === 'closed') return; if (isIdleNow()) schedule(); else cancel(); }; - return { noteActivity, cancel, idleMs }; + return { + noteActivity, + cancel, + close: () => { + phase = 'closed'; + cancel(); + }, + idleMs, + }; } /** @@ -170,6 +182,8 @@ export function createDaemonIdleReap(params: { * Embedded runtimes refuse replacement because their host owns close(). */ export type DaemonLifecycleAdmission = Readonly<{ + /** Closes admission when any owner starts draining the runtime. */ + close(): void; /** * Puts the lifecycle request and the closed admission it leaves behind in front of `handler`: a * lifecycle request gets its verdict, and any other request gets the shutting-down refusal once a @@ -256,6 +270,9 @@ export function createDaemonLifecycleAdmission(params: { }; return { + close: () => { + admitting = false; + }, admit: (handler) => async (req) => { if (req.command === DAEMON_QUIESCE_COMMAND) return leave(req); return admitting ? await handler(req) : daemonShuttingDownResponse();