fix(gateway): retain work admission across hosted wizard steps (#120582)

src/gateway/server-methods/wizard.test.ts

@@ -1,6 +1,12 @@
 // Wizard server-method tests cover stable lifecycle errors for process-local sessions.
 import { expectDefined } from "@openclaw/normalization-core";
 import { describe, expect, it, vi } from "vitest";
+import { createSafeGatewayRestartPreflight } from "../../infra/restart-coordinator.js";
+import {
+  getActiveGatewayRootWorkCount,
+  resetGatewayWorkAdmission,
+  runWithGatewayIndependentRootWorkAdmission,
+} from "../../process/gateway-work-admission.js";
 import type { RuntimeEnv } from "../../runtime.js";
 import { createDeferred } from "../../shared/deferred.js";
 import type { WizardPrompter } from "../../wizard/prompts.js";
@@ -138,6 +144,44 @@ describe("hosted wizard runtime isolation", () => {
 });
 
 describe("wizard setup ownership", () => {
+  it("retains gateway work admission between requests until the wizard settles", async () => {
+    resetGatewayWorkAdmission();
+    const runnerSettled = createDeferred();
+    const tracker = createWizardSessionTracker();
+    const context = {
+      ...tracker,
+      wizardRunner: async (_opts: unknown, _runtime: RuntimeEnv, prompter: WizardPrompter) => {
+        prompter.progress("working");
+        await runnerSettled.promise;
+      },
+    };
+
+    try {
+      await runWithGatewayIndependentRootWorkAdmission(async () => {
+        const respond = vi.fn();
+        await expectDefined(
+          wizardHandlers["wizard.start"],
+          "wizard.start test invariant",
+        )({ params: { mode: "local" }, respond, context } as never);
+        expect(respond.mock.calls[0]?.[1]).toMatchObject({ status: "running" });
+      });
+
+      expect(getActiveGatewayRootWorkCount()).toBe(1);
+      expect(createSafeGatewayRestartPreflight()).toMatchObject({
+        safe: false,
+        blockers: [expect.objectContaining({ kind: "root-request", count: 1 })],
+      });
+      runnerSettled.resolve();
+      await vi.waitFor(() => {
+        expect(getActiveGatewayRootWorkCount()).toBe(0);
+      });
+      expect(createSafeGatewayRestartPreflight().safe).toBe(true);
+    } finally {
+      runnerSettled.resolve();
+      resetGatewayWorkAdmission();
+    }
+  });
+
   it("blocks a replacement wizard until the cancelled runner settles", async () => {
     const runnerSettled = createDeferred();
     const tracker = createWizardSessionTracker();

src/gateway/server-methods/wizard.ts

@@ -12,6 +12,7 @@ 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 {
@@ -73,6 +74,15 @@ 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 session.whenSettled().then(release);
+  }
+}
+
 /** Resolves a live wizard session or sends the public not-found error. */
 function findWizardSessionOrRespond(params: {
   context: GatewayRequestContext;
@@ -136,6 +146,7 @@ export const wizardHandlers: GatewayRequestHandlers = {
               ),
             ),
           );
+    retainGatewayWorkUntilSettled(session);
     context.wizardSessions.set(sessionId, session);
     const result = await session.next();
     if (result.done) {