From 50f37a479cf6fcd49650e53d7201da58de12463a Mon Sep 17 00:00:00 2001 From: jesse-merhi <79823012+jesse-merhi@users.noreply.github.com> Date: Wed, 12 Aug 2026 22:19:48 +1000 Subject: [PATCH] feat(gateway): recover Custodian wizard sessions after reload --- .../OpenClawProtocol/GatewayModels.swift | 12 +- packages/gateway-protocol/CHANGELOG.md | 1 + .../gateway-protocol/src/schema/frames.ts | 1 + .../src/schema/openclaw.test.ts | 19 ++ .../gateway-protocol/src/schema/openclaw.ts | 8 + .../system-agent-chat-history.test.ts | 257 ++++++++++++++++++ .../system-agent-chat-history.ts | 120 ++++++++ .../system-agent-gateway-queue.ts | 18 ++ .../system-agent-reset-boundary.test.ts | 140 +++++++++- .../system-agent-session-ownership.test.ts | 2 + .../system-agent-session-queue.ts | 16 ++ src/gateway/server-methods/system-agent.ts | 119 +++----- .../server.auth.default-token.suite.ts | 3 + .../server/ws-connection/connect-hello.ts | 1 + src/system-agent/chat-engine.ts | 12 + src/system-agent/chat-wizard-host.test.ts | 26 ++ src/system-agent/chat-wizard-host.ts | 7 +- 17 files changed, 671 insertions(+), 91 deletions(-) create mode 100644 src/gateway/server-methods/system-agent-chat-history.test.ts create mode 100644 src/gateway/server-methods/system-agent-chat-history.ts create mode 100644 src/gateway/server-methods/system-agent-gateway-queue.ts create mode 100644 src/gateway/server-methods/system-agent-session-queue.ts diff --git a/apps/shared/OpenClawKit/Sources/OpenClawProtocol/GatewayModels.swift b/apps/shared/OpenClawKit/Sources/OpenClawProtocol/GatewayModels.swift index 96c95b68bf4d..a5e1fffd0370 100644 --- a/apps/shared/OpenClawKit/Sources/OpenClawProtocol/GatewayModels.swift +++ b/apps/shared/OpenClawKit/Sources/OpenClawProtocol/GatewayModels.swift @@ -10923,15 +10923,19 @@ public struct SystemAgentChatResult: Codable, Sendable { public struct SystemAgentChatHistoryParams: Codable, Sendable { public let limit: Int? + public let sessionid: String? public init( - limit: Int? = nil) + limit: Int? = nil, + sessionid: String? = nil) { self.limit = limit + self.sessionid = sessionid } private enum CodingKeys: String, CodingKey { case limit + case sessionid = "sessionId" } } @@ -10959,15 +10963,19 @@ public struct SystemAgentChatHistoryTurn: Codable, Sendable { public struct SystemAgentChatHistoryResult: Codable, Sendable { public let turns: [SystemAgentChatHistoryTurn] + public let activewizard: [String: AnyCodable]? public init( - turns: [SystemAgentChatHistoryTurn]) + turns: [SystemAgentChatHistoryTurn], + activewizard: [String: AnyCodable]? = nil) { self.turns = turns + self.activewizard = activewizard } private enum CodingKeys: String, CodingKey { case turns + case activewizard = "activeWizard" } } diff --git a/packages/gateway-protocol/CHANGELOG.md b/packages/gateway-protocol/CHANGELOG.md index fb5b6cd63328..0250212f2ec9 100644 --- a/packages/gateway-protocol/CHANGELOG.md +++ b/packages/gateway-protocol/CHANGELOG.md @@ -14,6 +14,7 @@ version and the additive schema surface. Dates are authoring dates (2026). - Slim worker and session-catalog payloads to the active wire contract. - Remove dead protocol surfaces and add since-vintage metadata to retained schemas and methods. - Add optional `step` on `SystemAgentChatResult` carrying the full awaited wizard step. +- Add owner-bound active-wizard snapshots to Custodian chat history for reload recovery. ## Protocol v4 (current) diff --git a/packages/gateway-protocol/src/schema/frames.ts b/packages/gateway-protocol/src/schema/frames.ts index bbd3c7002bc4..9bb01aecf996 100644 --- a/packages/gateway-protocol/src/schema/frames.ts +++ b/packages/gateway-protocol/src/schema/frames.ts @@ -10,6 +10,7 @@ export const GATEWAY_SERVER_CAPS = { BOARD_WIDGET_PUT_CANVAS_DOC: "board-widget-put-canvas-doc", CHAT_SEND_ROUTING_CONTRACT: "chat-send-routing-contract", SYSTEM_AGENT_WIZARD_CANCEL: "openclaw-chat-wizard-cancel", + SYSTEM_AGENT_CHAT_HISTORY_SESSION_RECOVERY: "openclaw-chat-history-session-recovery", SYSTEM_AGENT_SETUP_MODEL_REF: "openclaw-setup-model-ref", TASK_SUGGESTIONS_ACCEPT_MODES: "taskSuggestions.acceptModes", } as const; diff --git a/packages/gateway-protocol/src/schema/openclaw.test.ts b/packages/gateway-protocol/src/schema/openclaw.test.ts index 7bfff7c538bc..4f83263951e1 100644 --- a/packages/gateway-protocol/src/schema/openclaw.test.ts +++ b/packages/gateway-protocol/src/schema/openclaw.test.ts @@ -91,6 +91,8 @@ describe("OpenClaw chat question protocol", () => { describe("OpenClaw chat history protocol", () => { it("accepts the default request and bounds explicit limits", () => { expect(validateSystemAgentChatHistoryParams({})).toBe(true); + expect(validateSystemAgentChatHistoryParams({ sessionId: "recover-session" })).toBe(true); + expect(validateSystemAgentChatHistoryParams({ sessionId: "" })).toBe(false); expect(validateSystemAgentChatHistoryParams({ limit: 1 })).toBe(true); expect(validateSystemAgentChatHistoryParams({ limit: 500 })).toBe(true); expect(validateSystemAgentChatHistoryParams({ limit: 0 })).toBe(false); @@ -112,6 +114,23 @@ describe("OpenClaw chat history protocol", () => { }), ).toBe(false); }); + + it("accepts an optional sanitized active wizard snapshot", () => { + expect( + Value.Check(SystemAgentChatHistoryResultSchema, { + turns: [], + activeWizard: { + sessionId: "recover-session", + step: { + id: "secret", + type: "text", + message: "Twitch client secret", + sensitive: true, + }, + }, + }), + ).toBe(true); + }); }); describe("OpenClaw setup detection protocol", () => { diff --git a/packages/gateway-protocol/src/schema/openclaw.ts b/packages/gateway-protocol/src/schema/openclaw.ts index 3d8ba4fd288b..791ed07d670f 100644 --- a/packages/gateway-protocol/src/schema/openclaw.ts +++ b/packages/gateway-protocol/src/schema/openclaw.ts @@ -109,6 +109,8 @@ export const SystemAgentChatResultSchema = closedObject({ export const SystemAgentChatHistoryParamsSchema = closedObject({ limit: Type.Optional(Type.Integer({ minimum: 1, maximum: 500, default: 100 })), + /** Exact volatile chat session to inspect for a reload-recoverable active wizard. */ + sessionId: Type.Optional(NonEmptyString), }); export const SystemAgentChatHistoryTurnSchema = closedObject({ @@ -119,6 +121,12 @@ export const SystemAgentChatHistoryTurnSchema = closedObject({ export const SystemAgentChatHistoryResultSchema = closedObject({ turns: Type.Array(SystemAgentChatHistoryTurnSchema), + activeWizard: Type.Optional( + closedObject({ + sessionId: NonEmptyString, + step: WizardStepSchema, + }), + ), }); export const SystemChangeKindSchema = Type.Union([ diff --git a/src/gateway/server-methods/system-agent-chat-history.test.ts b/src/gateway/server-methods/system-agent-chat-history.test.ts new file mode 100644 index 000000000000..6785c14b682c --- /dev/null +++ b/src/gateway/server-methods/system-agent-chat-history.test.ts @@ -0,0 +1,257 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { createDeferred } from "../../../test/helpers/promise.js"; +import { + appendSystemAgentRecoveryHistory, + setSystemAgentRecoveryHistory, + systemAgentChatHistoryHandler, +} from "./system-agent-chat-history.js"; +import { runSystemAgentGatewayTask } from "./system-agent-gateway-queue.js"; +import { getSystemAgentSessionQueue } from "./system-agent-session-queue.js"; +import type { GatewayClient } from "./types.js"; + +const turns = [ + { role: "user" as const, text: "one", at: 1 }, + { role: "assistant" as const, text: "two", at: 2 }, +]; + +const transcriptStoreMocks = vi.hoisted(() => ({ + readTranscriptTail: vi.fn(), +})); + +vi.mock("../../system-agent/transcript-store.js", () => ({ + readTranscriptTail: transcriptStoreMocks.readTranscriptTail, +})); + +const ownerClient = { + connId: "conn-owner", + connect: { device: { id: "device-owner" } }, +} as GatewayClient; + +function makeInvocation(params: { + sessionId?: string; + limit?: number; + client?: GatewayClient; + activeWizardStep?: ReturnType; +}) { + const calls: Array<{ ok: boolean; payload?: unknown; error?: unknown }> = []; + const activeWizardStep = params.activeWizardStep ?? vi.fn().mockResolvedValue(undefined); + const session = { + ownerKey: "device:device-owner", + engine: { + activeWizardStep, + }, + lastUsedAt: 1, + }; + setSystemAgentRecoveryHistory(session.engine, turns); + const context = { + systemAgentSessions: new Map(params.sessionId ? [[params.sessionId, session]] : []), + }; + const options = { + params: { + ...(params.sessionId ? { sessionId: params.sessionId } : {}), + ...(params.limit ? { limit: params.limit } : {}), + }, + client: params.client ?? ownerClient, + context, + respond: (ok: boolean, payload?: unknown, error?: unknown) => { + calls.push({ ok, payload, error }); + }, + } as never; + return { activeWizardStep, calls, context, options, session }; +} + +describe("openclaw.chat.history wizard recovery", () => { + beforeEach(() => { + transcriptStoreMocks.readTranscriptTail.mockReset().mockReturnValue(turns); + }); + + it("returns an active wizard only to its bound owner", async () => { + const activeWizardStep = vi.fn().mockResolvedValue({ + id: "secret", + type: "text", + message: "Bot token", + sensitive: true, + }); + const owner = makeInvocation({ sessionId: "recover-session", activeWizardStep }); + + await systemAgentChatHistoryHandler(owner.options); + + expect(owner.calls).toEqual([ + { + ok: true, + payload: { + turns, + activeWizard: { + sessionId: "recover-session", + step: { + id: "secret", + type: "text", + message: "Bot token", + sensitive: true, + }, + }, + }, + error: undefined, + }, + ]); + expect(activeWizardStep).toHaveBeenCalledOnce(); + expect(owner.session.lastUsedAt).toBeGreaterThan(1); + const foreign = makeInvocation({ + sessionId: "recover-session", + client: { + connId: "conn-foreign", + connect: { device: { id: "device-foreign" } }, + } as GatewayClient, + activeWizardStep, + }); + + await systemAgentChatHistoryHandler(foreign.options); + + expect(foreign.calls).toEqual([ + { + ok: true, + payload: { turns }, + error: undefined, + }, + ]); + expect(activeWizardStep).toHaveBeenCalledOnce(); + expect(foreign.session.lastUsedAt).toBe(1); + }); + + it("falls back to the global audit history after a Gateway reload", async () => { + const invocation = makeInvocation({ sessionId: "recover-session" }); + invocation.context.systemAgentSessions.clear(); + + await systemAgentChatHistoryHandler(invocation.options); + + expect(transcriptStoreMocks.readTranscriptTail).toHaveBeenCalledWith(100); + expect(invocation.activeWizardStep).not.toHaveBeenCalled(); + expect(invocation.calls).toEqual([ + { + ok: true, + payload: { turns }, + error: undefined, + }, + ]); + }); + + it("bounds live recovery turns to the history protocol maximum", async () => { + const invocation = makeInvocation({ sessionId: "recover-session", limit: 500 }); + setSystemAgentRecoveryHistory( + invocation.session.engine, + Array.from({ length: 500 }, (_, index) => ({ + role: "assistant" as const, + text: `old-${index}`, + at: index, + })), + ); + appendSystemAgentRecoveryHistory(invocation.session.engine, [ + { role: "user", text: "new question", at: 500 }, + { role: "assistant", text: "new reply", at: 501 }, + ]); + + await systemAgentChatHistoryHandler(invocation.options); + + const recovered = ( + invocation.calls[0]?.payload as { turns?: Array<{ text: string }> } | undefined + )?.turns; + expect(recovered).toHaveLength(500); + expect(recovered?.[0]?.text).toBe("old-2"); + expect(recovered?.slice(-2).map((turn) => turn.text)).toEqual(["new question", "new reply"]); + }); + + it("waits for the session queue before reading the recovery transcript", async () => { + const turnStarted = createDeferred(); + const releaseTurn = createDeferred(); + const invocation = makeInvocation({ sessionId: "recover-session" }); + const turn = getSystemAgentSessionQueue(invocation.context.systemAgentSessions).enqueue( + "recover-session", + async () => { + turnStarted.resolve(); + await releaseTurn.promise; + setSystemAgentRecoveryHistory(invocation.session.engine, [ + { role: "user", text: "committed question", at: 2 }, + { role: "assistant", text: "committed reply", at: 3 }, + ]); + }, + ); + await turnStarted.promise; + + const history = systemAgentChatHistoryHandler(invocation.options); + await Promise.resolve(); + const callsBeforeRelease = [...invocation.calls]; + releaseTurn.resolve(); + await Promise.all([turn, history]); + + expect(callsBeforeRelease).toEqual([]); + expect(invocation.calls).toEqual([ + { + ok: true, + payload: { + turns: [ + { role: "user", text: "committed question", at: 2 }, + { role: "assistant", text: "committed reply", at: 3 }, + ], + }, + error: undefined, + }, + ]); + }); + + it("waits for the global Gateway queue before recovering a session", async () => { + const taskStarted = createDeferred(); + const releaseTask = createDeferred(); + const invocation = makeInvocation({ sessionId: "recover-session" }); + const globalTask = runSystemAgentGatewayTask(async () => { + taskStarted.resolve(); + await releaseTask.promise; + }); + await taskStarted.promise; + + const history = systemAgentChatHistoryHandler(invocation.options); + await Promise.resolve(); + expect(invocation.calls).toEqual([]); + + releaseTask.resolve(); + await Promise.all([globalTask, history]); + + expect(invocation.calls).toEqual([ + { + ok: true, + payload: { turns }, + error: undefined, + }, + ]); + }); + + it("does not read a predecessor replaced under the same session id", async () => { + const taskStarted = createDeferred(); + const releaseTask = createDeferred(); + const invocation = makeInvocation({ sessionId: "recover-session" }); + const globalTask = runSystemAgentGatewayTask(async () => { + taskStarted.resolve(); + await releaseTask.promise; + }); + await taskStarted.promise; + + const history = systemAgentChatHistoryHandler(invocation.options); + const replacementActiveWizardStep = vi.fn().mockResolvedValue(undefined); + invocation.context.systemAgentSessions.set("recover-session", { + ownerKey: invocation.session.ownerKey, + engine: { activeWizardStep: replacementActiveWizardStep }, + lastUsedAt: 1, + }); + releaseTask.resolve(); + await Promise.all([globalTask, history]); + + expect(invocation.activeWizardStep).not.toHaveBeenCalled(); + expect(replacementActiveWizardStep).not.toHaveBeenCalled(); + expect(invocation.calls).toEqual([ + { + ok: true, + payload: { turns }, + error: undefined, + }, + ]); + }); +}); diff --git a/src/gateway/server-methods/system-agent-chat-history.ts b/src/gateway/server-methods/system-agent-chat-history.ts new file mode 100644 index 000000000000..cd8b9a359c69 --- /dev/null +++ b/src/gateway/server-methods/system-agent-chat-history.ts @@ -0,0 +1,120 @@ +import { + validateSystemAgentChatHistoryParams, + type SystemAgentChatHistoryTurn, +} from "../../../packages/gateway-protocol/src/index.js"; +import type { SystemAgentChatEngine } from "../../system-agent/chat-engine.js"; +import { resolveSystemAgentDelegationKey } from "../../system-agent/delegation-session.js"; +import { readTranscriptTail } from "../../system-agent/transcript-store.js"; +import { runSystemAgentGatewayTask } from "./system-agent-gateway-queue.js"; +import { getSystemAgentSessionQueue } from "./system-agent-session-queue.js"; +import type { GatewayClient, GatewayRequestHandler } from "./types.js"; +import { assertValidParams } from "./validation.js"; + +const DEFAULT_SYSTEM_AGENT_HISTORY_LIMIT = 100; +const MAX_SYSTEM_AGENT_RECOVERY_TURNS = 500; +const recoveryTurnsByEngine = new WeakMap(); + +export function setSystemAgentRecoveryHistory( + engine: object, + turns: readonly SystemAgentChatHistoryTurn[], +): void { + recoveryTurnsByEngine.set(engine, turns.slice(-MAX_SYSTEM_AGENT_RECOVERY_TURNS)); +} + +export function appendSystemAgentRecoveryHistory( + engine: object, + turns: readonly SystemAgentChatHistoryTurn[], +): void { + const recoveryTurns = [...(recoveryTurnsByEngine.get(engine) ?? []), ...turns]; + recoveryTurnsByEngine.set(engine, recoveryTurns.slice(-MAX_SYSTEM_AGENT_RECOVERY_TURNS)); +} + +function readSystemAgentRecoveryHistory( + engine: object, + limit = DEFAULT_SYSTEM_AGENT_HISTORY_LIMIT, +): SystemAgentChatHistoryTurn[] { + return (recoveryTurnsByEngine.get(engine) ?? []).slice(-limit); +} + +export function resolveSystemAgentSessionOwnerKey(params: { + delegation?: { agentId?: string; sessionKey?: string }; + client: GatewayClient | null; +}): string | undefined { + const delegationKey = resolveSystemAgentDelegationKey(params.delegation); + if (delegationKey !== undefined) { + // Delegation is the host-only, cross-connection owner asserted by the regular-agent + // tool path. Keep its agent/session tuple authoritative across gateway reconnects. + return delegationKey; + } + // Authenticated users survive reconnects and may span paired devices. Otherwise + // bind to the verified device, with the server-issued connection as a last resort. + const userId = params.client?.authenticatedUserId?.trim(); + if (userId) { + return `user:${userId}`; + } + const deviceId = params.client?.connect.device?.id.trim(); + if (deviceId) { + return `device:${deviceId}`; + } + const connId = params.client?.connId?.trim(); + return connId ? `connection:${connId}` : undefined; +} + +export const systemAgentChatHistoryHandler: GatewayRequestHandler = async ({ + params, + respond, + client, + context, +}) => { + if ( + !assertValidParams( + params, + validateSystemAgentChatHistoryParams, + "openclaw.chat.history", + respond, + ) + ) { + return; + } + const requestedSessionId = params.sessionId; + const session = requestedSessionId + ? context.systemAgentSessions.get(requestedSessionId) + : undefined; + const ownerKey = resolveSystemAgentSessionOwnerKey({ client }); + const recovery = + requestedSessionId && session && ownerKey === session.ownerKey + ? await runSystemAgentGatewayTask( + async () => + await getSystemAgentSessionQueue(context.systemAgentSessions).enqueue( + requestedSessionId, + async () => { + if (context.systemAgentSessions.get(requestedSessionId) !== session) { + return undefined; + } + session.lastUsedAt = Date.now(); + const engine = session.engine as typeof session.engine & + Pick; + return { + turns: readSystemAgentRecoveryHistory( + engine, + params.limit ?? DEFAULT_SYSTEM_AGENT_HISTORY_LIMIT, + ), + step: await engine.activeWizardStep(), + }; + }, + ), + ) + : undefined; + const turns = + recovery?.turns ?? readTranscriptTail(params.limit ?? DEFAULT_SYSTEM_AGENT_HISTORY_LIMIT); + respond( + true, + { + turns, + ...(requestedSessionId && recovery?.step + ? { activeWizard: { sessionId: requestedSessionId, step: recovery.step } } + : {}), + }, + undefined, + ); +}; diff --git a/src/gateway/server-methods/system-agent-gateway-queue.ts b/src/gateway/server-methods/system-agent-gateway-queue.ts new file mode 100644 index 000000000000..23232daf4c86 --- /dev/null +++ b/src/gateway/server-methods/system-agent-gateway-queue.ts @@ -0,0 +1,18 @@ +import { KeyedAsyncQueue } from "../../plugin-sdk/keyed-async-queue.js"; +import { enqueueCommandInLane, setCommandLaneConcurrency } from "../../process/command-queue.js"; +import { CommandLane } from "../../process/lanes.js"; + +const SYSTEM_AGENT_GATEWAY_EXECUTION_KEY = "gateway"; +const systemAgentGatewayExecutionQueue = new KeyedAsyncQueue(); + +export async function runSystemAgentGatewayTask(task: () => Promise): Promise { + // Track every accepted RPC as active, never queued: restart draining snapshots + // active ids, so a queued OpenClaw request could otherwise outlive its socket. + setCommandLaneConcurrency(CommandLane.SystemAgent, Number.MAX_SAFE_INTEGER); + return await enqueueCommandInLane(CommandLane.SystemAgent, () => + // Bound expensive detection, activation, and agent turns without hiding + // accepted work from restart draining. This also makes session eviction and + // setup writes atomic with respect to other OpenClaw gateway requests. + systemAgentGatewayExecutionQueue.enqueue(SYSTEM_AGENT_GATEWAY_EXECUTION_KEY, task), + ); +} diff --git a/src/gateway/server-methods/system-agent-reset-boundary.test.ts b/src/gateway/server-methods/system-agent-reset-boundary.test.ts index 853f1ccf7892..79e5c6090909 100644 --- a/src/gateway/server-methods/system-agent-reset-boundary.test.ts +++ b/src/gateway/server-methods/system-agent-reset-boundary.test.ts @@ -10,8 +10,13 @@ import { closeOpenClawStateDatabase } from "../../state/openclaw-state-db.js"; import { SystemAgentChatEngine } from "../../system-agent/chat-engine.js"; import { SystemAgentInferenceUnavailableError } from "../../system-agent/inference-error.js"; import { createSystemAgentVerifiedInferenceTestFixture } from "../../system-agent/system-agent.test-helpers.js"; -import { appendTranscriptTurn, readTranscriptTail } from "../../system-agent/transcript-store.js"; +import { + appendTranscriptReset, + appendTranscriptTurn, + readTranscriptTail, +} from "../../system-agent/transcript-store.js"; import { withTestDir } from "../../test-helpers/temp-dir.js"; +import { setSystemAgentRecoveryHistory } from "./system-agent-chat-history.js"; import { systemAgentHandlers, type SystemAgentChatSession } from "./system-agent.js"; import type { GatewayClient, GatewayRequestContext } from "./types.js"; @@ -133,6 +138,139 @@ async function withTranscriptState(prefix: string, run: () => Promise): Pr } describe("openclaw.chat reset boundary", () => { + it("recovers the live session transcript without global pre-reset turns", async () => { + await withTranscriptState("openclaw-session-recovery-boundary-", async () => { + appendTranscriptTurn({ role: "user", text: "discarded setup request", at: 1 }); + appendTranscriptTurn({ role: "assistant", text: "discarded setup reply", at: 2 }); + appendTranscriptReset(); + const fixture = await createSystemAgentVerifiedInferenceTestFixture(verifiedConfig); + const engine = new SystemAgentChatEngine( + { + surface: "gateway", + verifiedInference: fixture.binding, + deps: { + ...fixture.deps, + readConfigFileSnapshot: async () => + ({ + exists: true, + valid: true, + path: "/tmp/openclaw.json", + hash: "verified-config", + config: verifiedConfig, + runtimeConfig: verifiedConfig, + sourceConfig: verifiedConfig, + issues: [], + }) as never, + }, + }, + { + wizardDependencies: { + runChannelSetupWizard: async (_channel, prompter) => { + await prompter.text({ message: "Bot token" }); + }, + }, + }, + ); + const sessions = new Map([ + [ + "recover-session", + { + engine, + welcome: "welcome text", + lastUsedAt: 1, + ownerKey: "device:device-test", + }, + ], + ]); + const context = { systemAgentSessions: sessions } as unknown as GatewayRequestContext; + const chatResponses: Array<{ ok: boolean; payload?: unknown; error?: unknown }> = []; + + await expectDefined( + systemAgentHandlers["openclaw.chat"], + 'systemAgentHandlers["openclaw.chat"] test invariant', + )({ + params: { sessionId: "recover-session", message: "connect telegram" }, + client, + context, + respond: (ok: boolean, payload?: unknown, error?: unknown) => + chatResponses.push({ ok, payload, error }), + } as never); + + expect(chatResponses).toEqual([ + { + ok: true, + payload: expect.objectContaining({ + wizardInputPending: true, + step: expect.objectContaining({ message: "Bot token" }), + }), + error: undefined, + }, + ]); + const historyResponses: Array<{ ok: boolean; payload?: unknown; error?: unknown }> = []; + await expectDefined( + systemAgentHandlers["openclaw.chat.history"], + 'systemAgentHandlers["openclaw.chat.history"] test invariant', + )({ + params: { sessionId: "recover-session" }, + client, + context, + respond: (ok: boolean, payload?: unknown, error?: unknown) => + historyResponses.push({ ok, payload, error }), + } as never); + + expect(historyResponses[0]).toMatchObject({ + ok: true, + payload: { + turns: [ + { role: "user", text: "connect telegram" }, + { role: "assistant", text: expect.any(String) }, + ], + activeWizard: expect.objectContaining({ sessionId: "recover-session" }), + }, + }); + expect(historyResponses[0]).not.toHaveProperty("payload.turns.0.sessionId"); + }); + }); + + it("keeps another live session's recovery turns when one session resets", async () => { + await withTranscriptState("openclaw-session-reset-isolation-", async () => { + const recoveryTurns = [ + { role: "user" as const, text: "other live question", at: 1 }, + { role: "assistant" as const, text: "other live answer", at: 2 }, + ]; + const sessions = discardableSessions(async () => undefined); + const engine = { + activeWizardStep: vi.fn(async () => undefined), + }; + sessions.set("s2", { + engine, + welcome: "welcome text", + lastUsedAt: 2, + ownerKey: "device:device-test", + } as unknown as SystemAgentChatSession); + setSystemAgentRecoveryHistory(engine, recoveryTurns); + inferenceFallbackMocks.verifySystemAgentInferenceWithFallback.mockResolvedValueOnce({ + ok: false, + status: "unavailable", + error: "no configured model", + }); + + await resetSession({ sessions }); + const responses: Array<{ ok: boolean; payload?: unknown }> = []; + await expectDefined( + systemAgentHandlers["openclaw.chat.history"], + 'systemAgentHandlers["openclaw.chat.history"] test invariant', + )({ + params: { sessionId: "s2" }, + client, + context: { systemAgentSessions: sessions } as unknown as GatewayRequestContext, + respond: (ok: boolean, payload?: unknown) => responses.push({ ok, payload }), + } as never); + + expect(responses).toEqual([{ ok: true, payload: { turns: recoveryTurns } }]); + }); + }); + // The reset discards the live session before initialization runs, so the // durable boundary has to survive a failed replacement. Otherwise the next // ordinary session seeds from the pre-reset transcript and undoes the reset. diff --git a/src/gateway/server-methods/system-agent-session-ownership.test.ts b/src/gateway/server-methods/system-agent-session-ownership.test.ts index 9113e3306b78..683c12532767 100644 --- a/src/gateway/server-methods/system-agent-session-ownership.test.ts +++ b/src/gateway/server-methods/system-agent-session-ownership.test.ts @@ -36,6 +36,7 @@ vi.mock("../../system-agent/greeting.js", () => ({ })); type FakeEngine = { + activeWizardStep: ReturnType; answerWizard: ReturnType; cancelWizard: ReturnType; handle: ReturnType; @@ -51,6 +52,7 @@ type FakeEngine = { function makeEngine(): FakeEngine { return { + activeWizardStep: vi.fn(async () => undefined), answerWizard: vi.fn(async () => { throw new SystemAgentWizardAnswerError("No hosted wizard is awaiting an answer."); }), diff --git a/src/gateway/server-methods/system-agent-session-queue.ts b/src/gateway/server-methods/system-agent-session-queue.ts new file mode 100644 index 000000000000..e89cce77e2c7 --- /dev/null +++ b/src/gateway/server-methods/system-agent-session-queue.ts @@ -0,0 +1,16 @@ +import { KeyedAsyncQueue } from "../../plugin-sdk/keyed-async-queue.js"; + +type SystemAgentSessionCollection = ReadonlyMap; + +const systemAgentSessionQueues = new WeakMap(); + +export function getSystemAgentSessionQueue( + sessions: SystemAgentSessionCollection, +): KeyedAsyncQueue { + let queue = systemAgentSessionQueues.get(sessions); + if (!queue) { + queue = new KeyedAsyncQueue(); + systemAgentSessionQueues.set(sessions, queue); + } + return queue; +} diff --git a/src/gateway/server-methods/system-agent.ts b/src/gateway/server-methods/system-agent.ts index 74ba6e699ddb..7c7a40dd01ce 100644 --- a/src/gateway/server-methods/system-agent.ts +++ b/src/gateway/server-methods/system-agent.ts @@ -6,11 +6,11 @@ import { ErrorCodes, errorShape, validateSystemAgentChatParams, - validateSystemAgentChatHistoryParams, validateSystemAgentSetupActivateParams, validateSystemAgentSetupAuthStartParams, validateSystemAgentSetupDetectParams, validateSystemAgentSetupVerifyParams, + type SystemAgentChatHistoryTurn, type SystemAgentChatQuestion, } from "../../../packages/gateway-protocol/src/index.js"; import { @@ -18,16 +18,12 @@ import { SYSTEM_AGENT_APPROVAL_TIMEOUT_MS, type SystemAgentApprovalRequestPayload, } from "../../infra/system-agent-approvals.js"; -import { KeyedAsyncQueue } from "../../plugin-sdk/keyed-async-queue.js"; -import { enqueueCommandInLane, setCommandLaneConcurrency } from "../../process/command-queue.js"; import { runWithGatewayIndependentRootWorkContinuation } from "../../process/gateway-work-admission.js"; -import { CommandLane } from "../../process/lanes.js"; import { defaultRuntime } from "../../runtime.js"; import { SystemAgentChatEngine, SystemAgentWizardAnswerError, } from "../../system-agent/chat-engine.js"; -import { resolveSystemAgentDelegationKey } from "../../system-agent/delegation-session.js"; import { acknowledgeSystemAgentGreetingDelivery, buildSystemAgentGreetingQuestion, @@ -56,14 +52,21 @@ import { SETUP_ADMISSION_BUSY_MESSAGE, SetupAdmissionBusyError, } from "./setup-admission.js"; +import { + appendSystemAgentRecoveryHistory, + resolveSystemAgentSessionOwnerKey, + setSystemAgentRecoveryHistory, + systemAgentChatHistoryHandler, +} from "./system-agent-chat-history.js"; import { sanitizeSystemAgentChatParams } from "./system-agent-chat-params.js"; import { buildSystemAgentChatResult, getSystemAgentChatInputError, runSystemAgentChatInput, } from "./system-agent-chat-turn.js"; -import type { GatewayClient, GatewayRequestContext, GatewayRequestHandlers } from "./types.js"; -import type { RespondFn } from "./types.js"; +import { runSystemAgentGatewayTask } from "./system-agent-gateway-queue.js"; +import { getSystemAgentSessionQueue } from "./system-agent-session-queue.js"; +import type { GatewayRequestContext, GatewayRequestHandlers, RespondFn } from "./types.js"; import { assertValidParams } from "./validation.js"; /** @@ -72,35 +75,17 @@ import { assertValidParams } from "./validation.js"; * the pre-inference phase; a new chat session starts only after a live model * turn succeeds. * - * The bounded session map owns only in-flight wizard and approval state. The - * sanitized conversation is a durable machine-wide logbook; `reset: true` - * replaces the in-memory session without deleting that transcript. + * The bounded session map owns in-flight wizard, approval, and recovery state. + * The sanitized conversation is also a durable machine-wide logbook; + * `reset: true` replaces the in-memory session without deleting that audit log. */ export type SystemAgentChatSession = GatewayRequestContext["systemAgentSessions"] extends Map ? Session : never; const MAX_SYSTEM_AGENT_SESSIONS = 8; const SYSTEM_AGENT_SEED_HISTORY_LIMIT = 30; -const DEFAULT_SYSTEM_AGENT_HISTORY_LIMIT = 100; const PROVIDER_AUTH_SESSION_TIMEOUT_MS = 25 * 60 * 1000; const PROVIDER_PREPARE_SESSION_TIMEOUT_MS = 2 * 60 * 60 * 1000; -const SYSTEM_AGENT_GATEWAY_EXECUTION_KEY = "gateway"; -const systemAgentGatewayExecutionQueue = new KeyedAsyncQueue(); -const systemAgentSessionQueues = new WeakMap< - Map, - KeyedAsyncQueue ->(); - -function getSystemAgentSessionQueue( - sessions: Map, -): KeyedAsyncQueue { - let queue = systemAgentSessionQueues.get(sessions); - if (!queue) { - queue = new KeyedAsyncQueue(); - systemAgentSessionQueues.set(sessions, queue); - } - return queue; -} function acknowledgeDeliveredSystemAgentWelcome(session: SystemAgentChatSession): void { const auditSequence = session.welcomeAuditSequence; @@ -111,42 +96,6 @@ function acknowledgeDeliveredSystemAgentWelcome(session: SystemAgentChatSession) delete session.welcomeAuditSequence; } -async function runSystemAgentGatewayTask(task: () => Promise): Promise { - // Track every accepted RPC as active, never queued: restart draining snapshots - // active ids, so a queued OpenClaw request could otherwise outlive its socket. - setCommandLaneConcurrency(CommandLane.SystemAgent, Number.MAX_SAFE_INTEGER); - return await enqueueCommandInLane(CommandLane.SystemAgent, () => - // Bound expensive detection, activation, and agent turns without hiding - // accepted work from restart draining. This also makes session eviction and - // setup writes atomic with respect to other OpenClaw gateway requests. - systemAgentGatewayExecutionQueue.enqueue(SYSTEM_AGENT_GATEWAY_EXECUTION_KEY, task), - ); -} - -function resolveSystemAgentSessionOwnerKey(params: { - delegation?: { agentId?: string; sessionKey?: string }; - client: GatewayClient | null; -}): string | undefined { - const delegationKey = resolveSystemAgentDelegationKey(params.delegation); - if (delegationKey !== undefined) { - // Delegation is the host-only, cross-connection owner asserted by the regular-agent - // tool path. Keep its agent/session tuple authoritative across gateway reconnects. - return delegationKey; - } - // Authenticated users survive reconnects and may span paired devices. Otherwise - // bind to the verified device, with the server-issued connection as a last resort. - const userId = params.client?.authenticatedUserId?.trim(); - if (userId) { - return `user:${userId}`; - } - const deviceId = params.client?.connect.device?.id.trim(); - if (deviceId) { - return `device:${deviceId}`; - } - const connId = params.client?.connId?.trim(); - return connId ? `connection:${connId}` : undefined; -} - async function evictOldestSession( sessions: Map, context: GatewayRequestContext, @@ -172,13 +121,18 @@ async function evictOldestSession( } } -function persistEngineHistory(engine: SystemAgentChatSession["engine"], startIndex: number): void { +function persistEngineHistory( + engine: SystemAgentChatSession["engine"], + startIndex: number, +): SystemAgentChatHistoryTurn[] { const at = Date.now(); - for (const turn of engine.historySince(startIndex)) { + const turns = engine.historySince(startIndex).map((turn) => ({ ...turn, at })); + for (const turn of turns) { // Engine history is authoritative here: sensitive user text has already // been replaced by the mask marker before it crosses this boundary. - appendTranscriptTurn({ ...turn, at }); + appendTranscriptTurn(turn); } + return turns; } function queueDelegatedApproval(params: { @@ -263,23 +217,7 @@ export const systemAgentHandlers: GatewayRequestHandlers = { undefined, ); }, - "openclaw.chat.history": ({ params, respond }) => { - if ( - !assertValidParams( - params, - validateSystemAgentChatHistoryParams, - "openclaw.chat.history", - respond, - ) - ) { - return; - } - respond( - true, - { turns: readTranscriptTail(params.limit ?? DEFAULT_SYSTEM_AGENT_HISTORY_LIMIT) }, - undefined, - ); - }, + "openclaw.chat.history": systemAgentChatHistoryHandler, /** Structured onboarding: list reusable AI access on this host. */ "openclaw.setup.detect": async ({ params, respond }) => { if ( @@ -636,7 +574,7 @@ export const systemAgentHandlers: GatewayRequestHandlers = { respond(false, undefined, errorShape(ErrorCodes.UNAVAILABLE, error.message)); return; } - persistEngineHistory(engine, welcomeHistoryStart); + const recoveryTurns = persistEngineHistory(engine, welcomeHistoryStart); await evictOldestSession(sessions, context); session = { engine, @@ -648,6 +586,7 @@ export const systemAgentHandlers: GatewayRequestHandlers = { lastUsedAt: Date.now(), ownerKey, }; + setSystemAgentRecoveryHistory(engine, recoveryTurns); sessions.set(sessionId, session); if (welcomeOnly) { respond( @@ -701,7 +640,10 @@ export const systemAgentHandlers: GatewayRequestHandlers = { } reply = turnReply; } catch (error) { - persistEngineHistory(session.engine, historyStart); + appendSystemAgentRecoveryHistory( + session.engine, + persistEngineHistory(session.engine, historyStart), + ); if (error instanceof SystemAgentWizardAnswerError) { respond(false, undefined, errorShape(ErrorCodes.INVALID_REQUEST, error.message)); return; @@ -730,7 +672,10 @@ export const systemAgentHandlers: GatewayRequestHandlers = { ); return; } - persistEngineHistory(session.engine, historyStart); + appendSystemAgentRecoveryHistory( + session.engine, + persistEngineHistory(session.engine, historyStart), + ); const delegation = params.delegation; let proposalId: string | undefined; if (delegation) { diff --git a/src/gateway/server.auth.default-token.suite.ts b/src/gateway/server.auth.default-token.suite.ts index 4f78aeaf3290..0cf77f20a87b 100644 --- a/src/gateway/server.auth.default-token.suite.ts +++ b/src/gateway/server.auth.default-token.suite.ts @@ -162,6 +162,9 @@ export function registerDefaultAuthTokenSuite(): void { expect(payload?.features?.capabilities).toContain( GATEWAY_SERVER_CAPS.SYSTEM_AGENT_WIZARD_CANCEL, ); + expect(payload?.features?.capabilities).toContain( + GATEWAY_SERVER_CAPS.SYSTEM_AGENT_CHAT_HISTORY_SESSION_RECOVERY, + ); expect(payload?.features?.capabilities).toContain( GATEWAY_SERVER_CAPS.SYSTEM_AGENT_SETUP_MODEL_REF, ); diff --git a/src/gateway/server/ws-connection/connect-hello.ts b/src/gateway/server/ws-connection/connect-hello.ts index c79666b2c133..04a5d2b4ee83 100644 --- a/src/gateway/server/ws-connection/connect-hello.ts +++ b/src/gateway/server/ws-connection/connect-hello.ts @@ -121,6 +121,7 @@ export async function sendGatewayHello( GATEWAY_SERVER_CAPS.BOARD_WIDGET_PUT_CANVAS_DOC, GATEWAY_SERVER_CAPS.CHAT_SEND_ROUTING_CONTRACT, GATEWAY_SERVER_CAPS.SYSTEM_AGENT_WIZARD_CANCEL, + GATEWAY_SERVER_CAPS.SYSTEM_AGENT_CHAT_HISTORY_SESSION_RECOVERY, GATEWAY_SERVER_CAPS.SYSTEM_AGENT_SETUP_MODEL_REF, GATEWAY_SERVER_CAPS.TASK_SUGGESTIONS_ACCEPT_MODES, ], diff --git a/src/system-agent/chat-engine.ts b/src/system-agent/chat-engine.ts index 3afdafeccaac..967ec1572b90 100644 --- a/src/system-agent/chat-engine.ts +++ b/src/system-agent/chat-engine.ts @@ -2,6 +2,7 @@ import type { SystemAgentWizardCancel, WizardAnswer, + WizardStep, } from "../../packages/gateway-protocol/src/index.js"; import type { RuntimeEnv } from "../runtime.js"; import { @@ -187,6 +188,17 @@ export class SystemAgentChatEngine { return await turn; } + async activeWizardStep(): Promise { + // Recovery snapshots share the turn queue so they cannot observe a step + // while its answer or cancellation mutation is still in flight. + const snapshot = this.turnQueue.then(() => this.wizard.clientStep); + this.turnQueue = snapshot.then( + () => undefined, + () => undefined, + ); + return await snapshot; + } + private completeTurn(reply: SystemAgentChatReply, userHistoryText: string): SystemAgentChatReply { const completed = this.wizard.decorateReply(reply); this.history.push({ role: "user", text: userHistoryText }); diff --git a/src/system-agent/chat-wizard-host.test.ts b/src/system-agent/chat-wizard-host.test.ts index 0c11624da75c..3cb19d65cd86 100644 --- a/src/system-agent/chat-wizard-host.test.ts +++ b/src/system-agent/chat-wizard-host.test.ts @@ -411,6 +411,32 @@ describe("SystemAgentChatEngine wizard", () => { expect(plain.step?.initialValue).toBe("123456:REAL-SECRET"); }); + it("snapshots only the sanitized active wizard step for reload recovery", async () => { + useTempStateDir(); + const engine = new SystemAgentChatEngine({ + surface: "gateway", + runAgentTurn: async () => null, + planWithAssistant: async () => null, + deps: { loadOverview: fakeOverviewLoader() }, + runChannelSetupWizard: async (_channel: string, prompter: WizardPrompter) => { + await prompter.text({ + message: "Bot token", + initialValue: "REAL-SECRET", + sensitive: true, + }); + }, + }); + + const prompt = await engine.handle("connect telegram"); + const snapshot = await engine.activeWizardStep(); + expect(snapshot).toEqual(prompt.step); + expect(snapshot).not.toHaveProperty("initialValue"); + + const stepId = expectDefined(snapshot?.id, "expected an active wizard step"); + await engine.cancelWizard({ stepId }); + await expect(engine.activeWizardStep()).resolves.toBeUndefined(); + }); + it("omits the wizard step outside an awaiting hosted wizard", async () => { useTempStateDir(); const engine = new SystemAgentChatEngine({ diff --git a/src/system-agent/chat-wizard-host.ts b/src/system-agent/chat-wizard-host.ts index b4689a7936f1..d5b7e51dfa51 100644 --- a/src/system-agent/chat-wizard-host.ts +++ b/src/system-agent/chat-wizard-host.ts @@ -269,6 +269,11 @@ export class ChatWizardHost { return this.bridge?.step?.sensitive === true; } + get clientStep(): WizardStep | undefined { + const step = this.bridge?.step; + return step ? sanitizeWizardStepForClient(step) : undefined; + } + dispose(): void { this.bridge?.session.cancel(); this.bridge = null; @@ -281,7 +286,7 @@ export class ChatWizardHost { ? { ...reply, text: `${reply.text}\n${WIZARD_CANCEL_HINT}` } : reply; const question = wizardStepChatQuestion(step); - const clientStep = step ? sanitizeWizardStepForClient(step) : null; + const clientStep = this.clientStep; return { ...completedReply, ...(step?.sensitive === true ? { sensitive: true } : {}),