diff --git a/src/gateway/server-methods/system-agent.test.ts b/src/gateway/server-methods/system-agent.test.ts index 0bae50d0dfa9..32b2a2fca9c6 100644 --- a/src/gateway/server-methods/system-agent.test.ts +++ b/src/gateway/server-methods/system-agent.test.ts @@ -14,7 +14,10 @@ import { resetCommandQueueStateForTest } from "../../process/command-queue.test- import { getActiveGatewayRootWorkCount } from "../../process/gateway-work-admission.js"; import { CommandLane } from "../../process/lanes.js"; import { defaultRuntime } from "../../runtime.js"; -import { SystemAgentChatEngine } from "../../system-agent/chat-engine.js"; +import { + SystemAgentChatEngine, + SystemAgentWizardAnswerError, +} from "../../system-agent/chat-engine.js"; import { SYSTEM_AGENT_HOSTED_WIZARD_TIMEOUT_MS } from "../../system-agent/chat-wizard-host.js"; import { createSystemAgentVerifiedInferenceTestFixture, @@ -1129,6 +1132,9 @@ describe("openclaw.chat", () => { await expect(qrEngine.pollStep(stepId)).resolves.toMatchObject({ wizardSettling: true }); await auditStarted.promise; + await expect(qrEngine.pollStep("stale-step")).rejects.toBeInstanceOf( + SystemAgentWizardAnswerError, + ); const admission = callChat(makeContext(sessions), { sessionId: "new-session" }); await waitOneTask(); diff --git a/src/system-agent/chat-engine.ts b/src/system-agent/chat-engine.ts index d8a42ff4e43c..2e2b2835e8b9 100644 --- a/src/system-agent/chat-engine.ts +++ b/src/system-agent/chat-engine.ts @@ -81,7 +81,7 @@ export class SystemAgentChatEngine { private disposal: Promise | null = null; private persistentApplySettlement: Promise | null = null; private retainedPollReplies = new Map(); - private passivePollsInFlight = 0; + private passivePollsInFlight = new Map(); constructor( private readonly options: SystemAgentChatEngineOptions, @@ -132,7 +132,7 @@ export class SystemAgentChatEngine { this.pruneExpiredPollReplies(); return ( this.wizard.hasPendingQrCode() || - this.passivePollsInFlight > 0 || + this.passivePollsInFlight.size > 0 || this.retainedPollReplies.size > 0 ); } @@ -223,7 +223,10 @@ export class SystemAgentChatEngine { } return { ...retained.reply }; } - this.passivePollsInFlight += 1; + if (!this.passivePollsInFlight.has(stepId)) { + this.wizard.assertPollableStep(stepId); + } + this.passivePollsInFlight.set(stepId, (this.passivePollsInFlight.get(stepId) ?? 0) + 1); const observation = this.turnQueue .then(async () => { this.assertActive(); @@ -250,7 +253,12 @@ export class SystemAgentChatEngine { return reply; }) .finally(() => { - this.passivePollsInFlight -= 1; + const remaining = (this.passivePollsInFlight.get(stepId) ?? 1) - 1; + if (remaining === 0) { + this.passivePollsInFlight.delete(stepId); + } else { + this.passivePollsInFlight.set(stepId, remaining); + } }); this.turnQueue = observation.catch(() => undefined); let cancelTimer: (() => void) | undefined; diff --git a/src/system-agent/chat-wizard-host.ts b/src/system-agent/chat-wizard-host.ts index b347c99c7ff8..4e46c6f77966 100644 --- a/src/system-agent/chat-wizard-host.ts +++ b/src/system-agent/chat-wizard-host.ts @@ -239,8 +239,7 @@ export class ChatWizardHost { /** Observe a QR-owned step without answering the dependency-owned prompt. */ async pollStep(stepId: string): Promise { - this.expireActiveQrIfNeeded(); - this.pruneExpiredPassiveQrRetention(); + this.assertPollableStep(stepId); const bridge = this.bridge; if (!bridge) { throw new SystemAgentWizardAnswerError("The hosted wizard step is no longer active."); @@ -248,9 +247,6 @@ export class ChatWizardHost { if (bridge.step?.id === stepId) { return { text: renderWizardStep(bridge.step), configWritten: false }; } - if (bridge.passiveQrStepId !== stepId) { - throw new SystemAgentWizardAnswerError("The hosted wizard poll targets a stale step."); - } const result = this.renderPendingQrOwner(bridge) ?? (await this.pump()); return bridge.passiveQrRetentionExpiresAtMs === undefined ? result @@ -260,6 +256,21 @@ export class ChatWizardHost { }; } + assertPollableStep(stepId: string): void { + this.expireActiveQrIfNeeded(); + this.pruneExpiredPassiveQrRetention(); + const bridge = this.bridge; + if (!bridge) { + throw new SystemAgentWizardAnswerError("The hosted wizard step is no longer active."); + } + if (bridge.step?.id === stepId) { + return; + } + if (bridge.passiveQrStepId !== stepId) { + throw new SystemAgentWizardAnswerError("The hosted wizard poll targets a stale step."); + } + } + async resolveReply(text: string): Promise { const bridge = this.bridge; if (!bridge) {