diff --git a/src/gateway/server-methods/wizard.test.ts b/src/gateway/server-methods/wizard.test.ts index 4b5f0346ddb0..79bee0cd0e27 100644 --- a/src/gateway/server-methods/wizard.test.ts +++ b/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(); diff --git a/src/gateway/server-methods/wizard.ts b/src/gateway/server-methods/wizard.ts index ba2561e885a8..0b0944cbde84 100644 --- a/src/gateway/server-methods/wizard.ts +++ b/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(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) {