diff --git a/src/gateway/server-methods/setup-admission.test.ts b/src/gateway/server-methods/setup-admission.test.ts index 3e721b1558f9..a9a52123998c 100644 --- a/src/gateway/server-methods/setup-admission.test.ts +++ b/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 | 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(() => { diff --git a/src/gateway/server-methods/setup-admission.ts b/src/gateway/server-methods/setup-admission.ts index 397bae4de5b8..aea72756c21e 100644 --- a/src/gateway/server-methods/setup-admission.ts +++ b/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(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) {