fix(lab): ensure manual Lab runs are cancellable during shutdown

src/lab/automation/orchestrator.ts

@@ -1,6 +1,9 @@
 import { readConfigDiagnostics } from "../../config";
 import { registerCurrentServerResourceCleanup } from "../../lib/server-resource-ownership";
-import { registerOptionalShutdownHook } from "../../lib/optional-shutdown-hooks";
+import {
+  didRunOptionalShutdownHooks,
+  registerOptionalShutdownHook,
+} from "../../lib/optional-shutdown-hooks";
 import { queryLabStatus } from "../query";
 import { rebuildLabProjection } from "../projection/rebuild";
 import { planLabAutomationRuns } from "./planner";
@@ -403,6 +406,13 @@ export async function runLabAutomationTick(configDir?: string): Promise<void> {
 
 export function startLabAutomationScheduler(configDir?: string): void {
   const key = configKey(configDir);
+  // Same late-request race as enqueueManualLabRun: a policy PUT or CLI enable can resume
+  // after the shutdown sweep already ran. Starting here would reset the latch and leave a
+  // live interval dispatching Lab work outside the completed sweep.
+  if (didRunOptionalShutdownHooks()) {
+    requestLabAutomationShutdown();
+    return;
+  }
   const currentOwner = dispatchDepsByConfigDir.get(key)?.token;
   const existing = schedulerTimers.get(key);
   if (existing) {
@@ -418,6 +428,14 @@ export function startLabAutomationScheduler(configDir?: string): void {
     requestLabAutomationShutdown();
     stopLabAutomationScheduler(configDir);
   });
+  // The sweep may have run in the gap between the entry check and this registration; it
+  // snapshots the registry once, so a hook that landed afterwards is orphaned. Keep the
+  // latch set and never start the timer — the registration stays live so a repeat sweep
+  // still tears this scheduler down.
+  if (didRunOptionalShutdownHooks()) {
+    requestLabAutomationShutdown();
+    return;
+  }
   shutdownRequested = false;
   const { policy, routes } = loadLabAutomationConfig(configDir);
   const now = Date.now();
@@ -464,6 +482,14 @@ export async function enqueueManualLabRun(
   configDir?: string,
   abortSignal?: AbortSignal,
 ): Promise<LabAutomationRunRecordV1 | null> {
+  // The management route is not covered by the data-plane drain gate, so a request accepted
+  // before listener teardown can reach this point after the shutdown sweep already ran. A
+  // hook registered then is never invoked; run its teardown inline instead of letting the
+  // dispatch escape shutdown.
+  if (didRunOptionalShutdownHooks()) {
+    requestLabAutomationShutdown();
+    return null;
+  }
   const now = Date.now();
   const created = mutateLabAutomationState(configDir, (state) => {
     const next = enqueuePlannedRuns(state, [planned], "manual", now);
@@ -473,8 +499,21 @@ export async function enqueueManualLabRun(
     return { state: next, value: run };
   });
   if (!created) return null;
-  // Manual execution is independent of automation enablement/layer toggles.
-  await runDispatchBatch(configDir, { manualRunId: created.runId, abortSignal });
+  // Manual execution can run without activation or a scheduler, so it must own a shutdown
+  // hook for the lifetime of its dispatch instead of relying on either of those paths.
+  const detachShutdownHook = registerOptionalShutdownHook(
+    `lab-automation-manual:${created.runId}`,
+    requestLabAutomationShutdown,
+  );
+  // The sweep may have run in the gap between the entry check and this registration; it
+  // snapshots the registry once, so a hook that landed afterwards is orphaned.
+  if (didRunOptionalShutdownHooks()) requestLabAutomationShutdown();
+  try {
+    // Manual execution is independent of automation enablement/layer toggles.
+    await runDispatchBatch(configDir, { manualRunId: created.runId, abortSignal });
+  } finally {
+    detachShutdownHook();
+  }
   return loadLabAutomationState(configDir).runs.find((row) => row.runId === created.runId) ?? created;
 }
 
@@ -496,4 +535,4 @@ export function cancelLabAutomationRun(runId: string, configDir?: string): boole
     }
     return { state: next, value: true };
   });
-}
+}

src/lib/optional-shutdown-hooks.ts

@@ -22,12 +22,22 @@
 type ShutdownHook = () => void;
 
 const hooks = new Map<string, ShutdownHook>();
+let hooksRan = false;
 
 /**
  * Register (or replace) the teardown for one optional subsystem.
  *
  * Keyed so repeated activation of the same subsystem cannot accumulate duplicate hooks.
  * Returns a detach function so an owner-scoped lease can release its registration.
+ *
+ * Registration after a sweep is NOT retro-applied: the hook waits for the next
+ * `runOptionalShutdownHooks`, which a draining process never reaches. Callers whose work
+ * must not outlive the sweep should gate on `didRunOptionalShutdownHooks`.
+ *
+ * The `hooksRan` latch is process-lifetime: every production caller runs the sweep inside
+ * `drainAndShutdown`, whose callers then exit or hand off to a newly spawned process —
+ * there is no in-process restart after a sweep. `resetOptionalShutdownHooksForTests`
+ * models that fresh process; it is the only way a post-sweep subsystem may start again.
  */
 export function registerOptionalShutdownHook(key: string, hook: ShutdownHook): () => void {
   hooks.set(key, hook);
@@ -37,8 +47,14 @@ export function registerOptionalShutdownHook(key: string, hook: ShutdownHook): (
   };
 }
 
+/** Whether `runOptionalShutdownHooks` has run at least once since the last test reset. */
+export function didRunOptionalShutdownHooks(): boolean {
+  return hooksRan;
+}
+
 /** Run every registered teardown. Never throws. */
 export function runOptionalShutdownHooks(): void {
+  hooksRan = true;
   for (const [key, hook] of [...hooks]) {
     try {
       hook();
@@ -54,4 +70,5 @@ export function runOptionalShutdownHooks(): void {
 /** Test-only reset so an isolated lifecycle test does not inherit registrations. */
 export function resetOptionalShutdownHooksForTests(): void {
   hooks.clear();
+  hooksRan = false;
 }

tests/lab/lab-automation-review-regressions.test.ts

@@ -16,10 +16,13 @@ import { enqueuePlannedRuns } from "../../src/lab/automation/queue";
 import { rollBudgetWindow, runBudgetRemaining } from "../../src/lab/automation/budgets";
 import {
   enqueueManualLabRun,
+  isLabAutomationSchedulerRunning,
+  reconcileLabAutomationQueue,
   requestLabAutomationShutdown,
   resetLabAutomationSchedulerStateForTests,
   runLabAutomationTick,
   setLabAutomationDispatchDeps,
+  startLabAutomationScheduler,
   stopLabAutomationScheduler,
 } from "../../src/lab/automation/orchestrator";
 import { dispatchLabAutomationRun } from "../../src/lab/automation/dispatch";
@@ -31,6 +34,10 @@ import type {
 import { LabAutomationError } from "../../src/lab/automation/types";
 import { LAB_AUTOMATION_HARD_MAX } from "../../src/lab/automation/constants";
 import { createHostIssuedLabRouteExecutor } from "../../src/lib/lab-live-host";
+import {
+  resetOptionalShutdownHooksForTests,
+  runOptionalShutdownHooks,
+} from "../../src/lib/optional-shutdown-hooks";
 import { readInstallationSalt } from "../../src/lab/subject/installation-salt";
 import type { NormalizedObservation } from "../../src/lab/conformance/types";
 import {
@@ -150,6 +157,7 @@ afterEach(() => {
   requestLabAutomationShutdown();
   stopLabAutomationScheduler();
   resetLabAutomationSchedulerStateForTests();
+  resetOptionalShutdownHooksForTests();
   setLabAutomationDispatchDeps({});
   resetCompatibilityVersionCacheForTests();
   delete process.env.OPENCODEX_HOME;
@@ -240,6 +248,121 @@ describe("CL-08 independent review regressions", () => {
     expect(result?.state).not.toBe("running");
   });
 
+  test("manual run remains cancellable by shutdown after its activation lease is released", async () => {
+    const home = tempHome();
+    prepareHome(home);
+    const config = providerConfig();
+    const plan = planManualLabRun({
+      evidenceLayer: "live_route_compatibility",
+      scenarioId: "responses-core.live.basic-turn",
+      providerName: "fixture-provider",
+      modelId: "fixture-model",
+      config,
+      configDir: home,
+    });
+    let startedResolve!: () => void;
+    const started = new Promise<void>((resolve) => { startedResolve = resolve; });
+    const release = setLabAutomationDispatchDeps({
+      configDir: home,
+      loadConfig: () => config,
+      resolve: fixtureDnsResolve(),
+      routeExecutor: createHostIssuedLabRouteExecutor(async (input) => {
+        startedResolve();
+        if (!input.signal.aborted) {
+          await new Promise<void>((resolve) => input.signal.addEventListener("abort", () => resolve(), { once: true }));
+        }
+        return passObservation();
+      }),
+    });
+
+    const manualRun = enqueueManualLabRun(plan, home);
+    await started;
+    release();
+    runOptionalShutdownHooks();
+
+    const result = await manualRun;
+    expect(result?.state).toBe("cancelled");
+    expect(result?.terminalCode).toBe("cancelled");
+  });
+
+  // Regression: the management route is outside the data-plane drain gate, so a manual run
+  // accepted while drainAndShutdown is in flight could register its per-run hook AFTER the
+  // hooks snapshot ran — orphaned, never invoked, and free to dispatch past shutdown (the
+  // Lab was never activated in that scenario, so no scheduler hook had set
+  // shutdownRequested either). The enqueue path now treats a completed sweep as already
+  // fired and refuses to queue or dispatch.
+  test("a manual run arriving after the shutdown sweep does not queue or dispatch", async () => {
+    const home = tempHome();
+    prepareHome(home);
+    saveLabAutomationPolicy(defaultLabAutomationPolicyV1(), home);
+    const plan = planManualLabRun({
+      evidenceLayer: "protocol_conformance",
+      scenarioId: "responses-core.protocol.request-shape",
+      configDir: home,
+    });
+    runOptionalShutdownHooks();
+    const result = await enqueueManualLabRun(plan, home);
+    expect(result).toBeNull();
+    expect(loadLabAutomationState(home).runs).toHaveLength(0);
+  });
+
+  // Regression: a policy PUT shares the same post-sweep race as the manual-run POST — the
+  // route resumes past the snapshot and applySchedulerPolicy used to start a scheduler
+  // whose startup cleared `shutdownRequested`, leaving a live interval dispatching Lab
+  // work outside the completed sweep. The start is now refused and the latch stays set.
+  test("a policy update landing after the shutdown sweep cannot restart Lab automation", async () => {
+    const home = tempHome();
+    prepareHome(home);
+    const config = providerConfig();
+    const policy = livePolicy();
+    saveLabAutomationPolicy(policy, home);
+    saveLabAutomationRoutes({
+      schemaVersion: 1,
+      routes: [{ providerName: "fixture-provider", modelId: "fixture-model" }],
+    }, home);
+    let invokes = 0;
+    setLabAutomationDispatchDeps({
+      configDir: home,
+      loadConfig: () => config,
+      resolve: fixtureDnsResolve(),
+      routeExecutor: createHostIssuedLabRouteExecutor(async () => {
+        invokes += 1;
+        return passObservation();
+      }),
+    });
+    runOptionalShutdownHooks();
+
+    // The late PUT path: applySchedulerPolicy reconciles then starts the scheduler.
+    reconcileLabAutomationQueue(home);
+    startLabAutomationScheduler(home);
+    expect(isLabAutomationSchedulerRunning(home)).toBe(false);
+    await runLabAutomationTick(home);
+    expect(invokes).toBe(0);
+  });
+
+  // `hooksRan` is process-lifetime (drainAndShutdown callers exit; restarts are new
+  // processes), so a later manual run must stay rejected. The test reset models the
+  // fresh process: with both latches cleared, the manual path dispatches again.
+  test("manual runs stay rejected after the sweep until a test reset re-arms them", async () => {
+    const home = tempHome();
+    prepareHome(home);
+    saveLabAutomationPolicy(defaultLabAutomationPolicyV1(), home);
+    const plan = planManualLabRun({
+      evidenceLayer: "protocol_conformance",
+      scenarioId: "responses-core.protocol.request-shape",
+      configDir: home,
+    });
+    runOptionalShutdownHooks();
+    expect(await enqueueManualLabRun(plan, home)).toBeNull();
+
+    resetOptionalShutdownHooksForTests();
+    resetLabAutomationSchedulerStateForTests();
+    const result = await enqueueManualLabRun(plan, home);
+    expect(result).not.toBeNull();
+    expect(result?.state).not.toBe("queued");
+    expect(result?.state).not.toBe("running");
+  });
+
   test("disabling the live layer cancels previously queued scheduled live work", async () => {
     const home = tempHome();
     prepareHome(home);

tests/lib/optional-shutdown-hooks.test.ts

@@ -188,14 +188,31 @@ describe("scheduler hook keying", () => {
     }
   });
 
-  // An in-process restart (service restart, test suite) must re-arm the hook.
-  test("a scheduler restarted after shutdown is stoppable again", () => {
+  // The sweep is process-lifetime: drainAndShutdown callers exit afterwards, so a
+  // scheduler start arriving late (a policy PUT resuming past the snapshot) must be
+  // refused — not re-armed — or its timer would dispatch outside the completed sweep.
+  test("a scheduler start after the shutdown sweep is refused", () => {
+    resetOptionalShutdownHooksForTests();
+    const configDir = mkdtempSync(join(tmpdir(), "ocx-shutdown-late-start-"));
+    try {
+      runOptionalShutdownHooks();
+      startLabAutomationScheduler(configDir);
+      expect(isLabAutomationSchedulerRunning(configDir)).toBe(false);
+    } finally {
+      stopLabAutomationScheduler(configDir);
+    }
+  });
+
+  // An in-process restart exists only in tests, where the reset models a fresh process;
+  // with the latch cleared a new scheduler must re-arm its hook for the next sweep.
+  test("a scheduler restarted after a test reset is stoppable again", () => {
     resetOptionalShutdownHooksForTests();
     const configDir = mkdtempSync(join(tmpdir(), "ocx-shutdown-restart-"));
     try {
       startLabAutomationScheduler(configDir);
       runOptionalShutdownHooks();
       expect(isLabAutomationSchedulerRunning(configDir)).toBe(false);
+      resetOptionalShutdownHooksForTests();
       startLabAutomationScheduler(configDir);
       expect(isLabAutomationSchedulerRunning(configDir)).toBe(true);
       runOptionalShutdownHooks();