diff --git a/src/gateway/server-methods/setup-admission.test.ts b/src/gateway/server-methods/setup-admission.test.ts index cf0a60627c7d..3e721b1558f9 100644 --- a/src/gateway/server-methods/setup-admission.test.ts +++ b/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 () => { diff --git a/src/gateway/server-methods/setup-admission.ts b/src/gateway/server-methods/setup-admission.ts index 9acbeb98ea6f..a03944b9f7cd 100644 --- a/src/gateway/server-methods/setup-admission.ts +++ b/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>(); export class SetupAdmissionBusyError extends Error {} @@ -27,6 +28,12 @@ export async function runExclusiveSystemAgentSetupActivation( } } +export function whenAdmittedWizardSessionSettled }>( + session: T, +): Promise { + return wizardSessionAdmissionSettlements.get(session) ?? session.whenSettled(); +} + export async function createAdmittedWizardSession }>( createSession: () => T, lockSetupTarget = true, @@ -39,16 +46,20 @@ export async function createAdmittedWizardSession | undefined; const session = lockSetupTarget ? await new Promise((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(); diff --git a/src/gateway/server-methods/wizard.test.ts b/src/gateway/server-methods/wizard.test.ts index 739a0e57f3da..91d964f306a0 100644 --- a/src/gateway/server-methods/wizard.test.ts +++ b/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(); diff --git a/src/gateway/server-methods/wizard.ts b/src/gateway/server-methods/wizard.ts index bc9bfcee1acf..a2d14530d76a 100644 --- a/src/gateway/server-methods/wizard.ts +++ b/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); }