fix(gateway): retain admission across setup sessions (#123418)

Co-authored-by: Peter Steinberger <steipete@gmail.com>
Co-authored-by: Colin Johnson <colin@solvely.net>

src/gateway/server-methods/setup-admission.test.ts

@@ -1,6 +1,13 @@
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import { createDeferred } from "../../../test/helpers/promise.js";
 import { useAutoCleanupTempDirTracker } from "../../../test/helpers/temp-dir.js";
+import { enqueueCommandInLane } from "../../process/command-queue.js";
+import { resetCommandQueueStateForTest } from "../../process/command-queue.test-support.js";
+import {
+  getActiveGatewayRootWorkCount,
+  resetGatewayWorkAdmission,
+  runWithGatewayIndependentRootWorkAdmission,
+} from "../../process/gateway-work-admission.js";
 import { withSetupMigrationTargetLock } from "../../wizard/setup.migration-snapshot.js";
 
 const tempDirs = useAutoCleanupTempDirTracker(afterEach);
@@ -23,6 +30,11 @@ describe("setup admission", () => {
     mocks.stateDir = tempDirs.make("openclaw-setup-admission-");
   });
 
+  afterEach(() => {
+    resetCommandQueueStateForTest();
+    resetGatewayWorkAdmission();
+  });
+
   it("rejects concurrent work instead of queueing it", async () => {
     const firstStarted = createDeferred();
     const releaseFirst = createDeferred();
@@ -86,6 +98,27 @@ describe("setup admission", () => {
     await whenAdmittedWizardSessionSettled(next!);
   });
 
+  it("retains root work for post-start session continuations", async () => {
+    const continueAfterStart = createDeferred();
+    let runner: Promise<void> | undefined;
+
+    await runWithGatewayIndependentRootWorkAdmission(async () => {
+      await createAdmittedWizardSession(() => {
+        runner = (async () => {
+          await continueAfterStart.promise;
+          await enqueueCommandInLane("setup-post-start-proof", async () => undefined);
+        })();
+        return { whenSettled: () => runner! };
+      });
+    });
+
+    const activeAfterStart = getActiveGatewayRootWorkCount();
+    continueAfterStart.resolve();
+    await expect(runner).resolves.toBeUndefined();
+    expect(activeAfterStart).toBe(1);
+    await vi.waitFor(() => expect(getActiveGatewayRootWorkCount()).toBe(0));
+  });
+
   it("releases an admitted session lease when construction fails", async () => {
     await expect(
       createAdmittedWizardSession(() => {

src/gateway/server-methods/setup-admission.ts

@@ -1,4 +1,5 @@
 import { resolveStateDir } from "../../config/paths.js";
+import { retainGatewayRootWorkAdmissionContinuation } from "../../process/gateway-work-admission.js";
 import {
   SetupTargetLockedError,
   withSetupMigrationTargetLock,
@@ -62,6 +63,12 @@ export async function createAdmittedWizardSession<T extends { whenSettled(): Pro
       : createSession();
     const settled = admissionSettled ?? session.whenSettled();
     wizardSessionAdmissionSettlements.set(session, settled);
+    // The runner outlives its start RPC and inherits that request's admission.
+    // Keep the root live so later prompts and post-auth probes remain subordinate work.
+    const releaseGatewayWork = retainGatewayRootWorkAdmissionContinuation();
+    if (releaseGatewayWork) {
+      void settled.then(releaseGatewayWork, releaseGatewayWork);
+    }
     void settled.then(releaseSession, releaseSession);
     return session;
   } catch (error) {

src/gateway/server-methods/wizard.ts

@@ -12,7 +12,6 @@ import {
   validateWizardStatusParams,
 } from "../../../packages/gateway-protocol/src/index.js";
 import type { OnboardOptions } from "../../commands/onboard-types.js";
-import { retainGatewayRootWorkAdmissionContinuation } from "../../process/gateway-work-admission.js";
 import { createNonExitingRuntime, ExitError, type RuntimeEnv } from "../../runtime.js";
 import type { WizardPrompter } from "../../wizard/prompts.js";
 import {
@@ -79,15 +78,6 @@ function sanitizeWizardResultForClient<T extends { step?: WizardStep }>(result:
   return result.step ? { ...result, step: sanitizeWizardStepForClient(result.step) } : result;
 }
 
-function retainGatewayWorkUntilSettled(session: WizardSession): void {
-  // Hosted wizard state spans RPC requests. Keep restart/suspend admission
-  // active between steps or a config reload can erase the process-local session.
-  const release = retainGatewayRootWorkAdmissionContinuation();
-  if (release) {
-    void whenAdmittedWizardSessionSettled(session).then(release, release);
-  }
-}
-
 /** Resolves a live wizard session or sends the public not-found error. */
 function findWizardSessionOrRespond(params: {
   context: GatewayRequestContext;
@@ -155,7 +145,6 @@ export const wizardHandlers: GatewayRequestHandlers = {
       );
       return;
     }
-    retainGatewayWorkUntilSettled(session);
     context.wizardSessions.set(sessionId, session);
     const result = await session.next();
     if (result.done) {