fix(daemon): keep idle reap closed during shutdown

packages/daemon/src/server/http/daemon-runtime-quiesce.test.ts

@@ -22,12 +22,15 @@ const probe = vi.hoisted(() => ({
   held: undefined as Promise<void> | 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<typeof import('@agent-device/host-kit/diagnostics')>();
   return {
     ...actual,
     emitDiagnostic: (...args: Parameters<typeof actual.emitDiagnostic>) => {
+      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<typeof import('./transport.ts')>();
+  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<typeof import('@agent-device/daemon-contracts/device-claims')>();
@@ -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<void>((resolve) => {
+    probe.entered = resolve;
+  });
+  probe.held = new Promise<void>((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<void>((resolve) => {
+    probe.entered = resolve;
+  });
+  probe.held = new Promise<void>((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 {

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<ReturnType<DaemonInvokeFn>>();
   // 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<ShutdownReport> => {
-    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.

packages/daemon/src/sessions/server/daemon-idle-reap.ts

@@ -100,6 +100,8 @@
    */
   noteActivity: () => void;
   cancel: () => void;
+  /** Permanently stops idle reap, including activity that settles during shutdown. */
+  close: () => void;
   readonly idleMs: number;
 };
 
@@ -123,6 +125,7 @@
   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 @@
 
   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 @@
   };
 
   const noteActivity = (): void => {
+    if (phase === 'closed') return;
     if (isIdleNow()) schedule();
     else cancel();
   };
 
-  return { noteActivity, cancel, idleMs };
+  return {
+    noteActivity,
+    cancel,
+    close: () => {
+      phase = 'closed';
+      cancel();
+    },
+    idleMs,
+  };
 }
 
 /**
@@ -169,9 +181,11 @@
  * Admission and the work census change in one synchronous step to avoid racing new work.
  * 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
    * lifecycle request has closed admission.
    */
@@ -256,6 +270,9 @@
   };
 
   return {
+    close: () => {
+      admitting = false;
+    },
     admit: (handler) => async (req) => {
       if (req.command === DAEMON_QUIESCE_COMMAND) return leave(req);
       return admitting ? await handler(req) : daemonShuttingDownResponse();