fix(gateway): settle setup admission with wizard sessions

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

@@ -15,6 +15,7 @@ vi.mock("../../config/paths.js", async () => ({
 import {
   createAdmittedWizardSession,
   runExclusiveSystemAgentSetupActivation,
+  whenAdmittedWizardSessionSettled,
 } from "./setup-admission.js";
 
 describe("setup admission", () => {
@@ -69,7 +70,7 @@ describe("setup admission", () => {
 
   it("holds an admitted session lease until its runner settles", async () => {
     const settled = createDeferred();
-    await createAdmittedWizardSession(() => ({
+    const session = await createAdmittedWizardSession(() => ({
       whenSettled: () => settled.promise,
     }));
 
@@ -77,14 +78,12 @@ describe("setup admission", () => {
       createAdmittedWizardSession(() => ({ whenSettled: () => Promise.resolve() })),
     ).resolves.toBeUndefined();
     settled.resolve();
-    await settled.promise;
-    await vi.waitFor(async () => {
-      const next = await createAdmittedWizardSession(() => ({
-        whenSettled: () => Promise.resolve(),
-      }));
-      expect(next).toBeDefined();
-      await next?.whenSettled();
-    });
+    await whenAdmittedWizardSessionSettled(session!);
+    const next = await createAdmittedWizardSession(() => ({
+      whenSettled: () => Promise.resolve(),
+    }));
+    expect(next).toBeDefined();
+    await whenAdmittedWizardSessionSettled(next!);
   });
 
   it("releases an admitted session lease when construction fails", async () => {
@@ -93,9 +92,11 @@ describe("setup admission", () => {
         throw new Error("construction failed");
       }),
     ).rejects.toThrow("construction failed");
-    await expect(
-      createAdmittedWizardSession(() => ({ whenSettled: () => Promise.resolve() })),
-    ).resolves.toBeDefined();
+    const recovered = await createAdmittedWizardSession(() => ({
+      whenSettled: () => Promise.resolve(),
+    }));
+    expect(recovered).toBeDefined();
+    await whenAdmittedWizardSessionSettled(recovered!);
   });
 
   it("reserves wizard admission while setup waits to acquire its target lock", async () => {

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

@@ -6,6 +6,7 @@ export const SETUP_ADMISSION_BUSY_MESSAGE =
   "OpenClaw setup is already in progress; try again when it finishes.";
 
 let wizardSessionInProgress = false;
+const wizardSessionAdmissionSettlements = new WeakMap<object, Promise<unknown>>();
 
 export class SetupAdmissionBusyError extends Error {}
 
@@ -27,6 +28,12 @@ export async function runExclusiveSystemAgentSetupActivation<T>(
   }
 }
 
+export function whenAdmittedWizardSessionSettled<T extends { whenSettled(): Promise<unknown> }>(
+  session: T,
+): Promise<unknown> {
+  return wizardSessionAdmissionSettlements.get(session) ?? session.whenSettled();
+}
+
 export async function createAdmittedWizardSession<T extends { whenSettled(): Promise<unknown> }>(
   createSession: () => T,
   lockSetupTarget = true,
@@ -39,16 +46,20 @@ export async function createAdmittedWizardSession<T extends { whenSettled(): Pro
     wizardSessionInProgress = false;
   };
   try {
+    let admissionSettled: Promise<unknown> | undefined;
     const session = lockSetupTarget
       ? await new Promise<T>((resolve, reject) => {
-          void runExclusiveSystemAgentSetupActivation(async () => {
+          admissionSettled = runExclusiveSystemAgentSetupActivation(async () => {
             const createdSession = createSession();
             resolve(createdSession);
             await createdSession.whenSettled();
-          }).catch(reject);
+          });
+          void admissionSettled.catch(reject);
         })
       : createSession();
-    void session.whenSettled().then(releaseSession, releaseSession);
+    const settled = admissionSettled ?? session.whenSettled();
+    wizardSessionAdmissionSettlements.set(session, settled);
+    void settled.then(releaseSession, releaseSession);
     return session;
   } catch (error) {
     releaseSession();

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

@@ -10,7 +10,10 @@ import {
 import type { RuntimeEnv } from "../../runtime.js";
 import type { WizardPrompter } from "../../wizard/prompts.js";
 import { createWizardSessionTracker } from "../server-wizard-sessions.js";
-import { runExclusiveSystemAgentSetupActivation } from "./setup-admission.js";
+import {
+  runExclusiveSystemAgentSetupActivation,
+  whenAdmittedWizardSessionSettled,
+} from "./setup-admission.js";
 import { systemAgentHandlers } from "./system-agent.js";
 import type { GatewayRequestHandlerOptions } from "./types.js";
 import { type SetupWizardRunner, wizardHandlers } from "./wizard.js";
@@ -51,7 +54,7 @@ async function cancelWizardSessions(
 ) {
   for (const session of sessions.values()) {
     session.cancel();
-    await session.whenSettled();
+    await whenAdmittedWizardSessionSettled(session);
   }
 }
 
@@ -201,7 +204,7 @@ describe("wizard setup ownership", () => {
       "admitted classic setup session",
     );
     session.cancel();
-    await session.whenSettled();
+    await whenAdmittedWizardSessionSettled(session);
   });
 
   it("makes structured setup retry while a classic runner owns admission, then releases", async () => {
@@ -241,7 +244,7 @@ describe("wizard setup ownership", () => {
       [...tracker.wizardSessions.values()][0],
       "active classic setup session",
     );
-    await session.whenSettled();
+    await whenAdmittedWizardSessionSettled(session);
     const structuredTask = vi.fn(async () => "ok");
     await expect(runExclusiveSystemAgentSetupActivation(structuredTask)).resolves.toBe("ok");
     expect(structuredTask).toHaveBeenCalledOnce();

src/gateway/server-methods/wizard.ts

@@ -21,7 +21,11 @@ import {
   type WizardStep,
 } from "../../wizard/session.js";
 import { formatForLog } from "../ws-log.js";
-import { createAdmittedWizardSession, SETUP_ADMISSION_BUSY_MESSAGE } from "./setup-admission.js";
+import {
+  createAdmittedWizardSession,
+  SETUP_ADMISSION_BUSY_MESSAGE,
+  whenAdmittedWizardSessionSettled,
+} from "./setup-admission.js";
 import type { GatewayRequestContext, GatewayRequestHandlers, RespondFn } from "./types.js";
 import { assertValidParams } from "./validation.js";
 
@@ -80,7 +84,7 @@ function retainGatewayWorkUntilSettled(session: WizardSession): void {
   // active between steps or a config reload can erase the process-local session.
   const release = retainGatewayRootWorkAdmissionContinuation();
   if (release) {
-    void session.whenSettled().then(release);
+    void whenAdmittedWizardSessionSettled(session).then(release);
   }
 }
 
@@ -157,7 +161,7 @@ export const wizardHandlers: GatewayRequestHandlers = {
     if (result.done) {
       // Let the runner release setup admission before the terminal response,
       // so an immediate replacement wizard is not rejected as still busy.
-      await session.whenSettled();
+      await whenAdmittedWizardSessionSettled(session);
       context.purgeWizardSession(sessionId);
     }
     respond(true, { sessionId, ...sanitizeWizardResultForClient(result) }, undefined);
@@ -198,7 +202,7 @@ export const wizardHandlers: GatewayRequestHandlers = {
     const result = await session.next();
     if (result.done) {
       // Keep terminal response ordering identical to wizard.start.
-      await session.whenSettled();
+      await whenAdmittedWizardSessionSettled(session);
       context.purgeWizardSession(sessionId);
     }
     respond(true, sanitizeWizardResultForClient(result), undefined);
@@ -215,7 +219,9 @@ export const wizardHandlers: GatewayRequestHandlers = {
     const cancelled = session.cancel();
     const status = readWizardStatus(session);
     if (cancelled) {
-      void session.whenSettled().then(() => context.purgeWizardSession(sessionId));
+      void whenAdmittedWizardSessionSettled(session).then(() =>
+        context.purgeWizardSession(sessionId),
+      );
     } else {
       context.purgeWizardSession(sessionId);
     }