diff --git a/src/lib/onboard/machine/runner.test.ts b/src/lib/onboard/machine/runner.test.ts new file mode 100644 index 0000000000..e0978813cd --- /dev/null +++ b/src/lib/onboard/machine/runner.test.ts @@ -0,0 +1,219 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { describe, expect, it, vi } from "vitest"; + +import { + MACHINE_SNAPSHOT_VERSION, + createSession, + filterSafeUpdates, + normalizeSession, + sanitizeFailure, + type Session, + type SessionUpdates, +} from "../../state/onboard-session"; +import { advanceTo, branchTo, completeOnboardMachine, failOnboardMachine, retryTo } from "./result"; +import { OnboardRuntime, type OnboardRuntimeDeps } from "./runtime"; +import { + MissingOnboardStateHandlerError, + OnboardMachineTransitionLimitError, + runOnboardMachine, + type OnboardStateHandlers, +} from "./runner"; + +interface RunnerContext { + attempts: number; + visited: string[]; +} + +function cloneSession(session: Session): Session { + return normalizeSession(JSON.parse(JSON.stringify(session))) ?? session; +} + +function createRuntime(initialSession: Session = createSession()) { + let session = cloneSession(initialSession); + const updateSession = (mutator: (value: Session) => Session | void): Session => { + const next = mutator(cloneSession(session)) ?? session; + session = cloneSession(next); + return cloneSession(session); + }; + const deps: OnboardRuntimeDeps = { + loadSession: () => cloneSession(session), + createSession, + saveSession: (next) => { + session = cloneSession(next); + return cloneSession(session); + }, + updateSession, + markStepStarted: () => cloneSession(session), + markStepComplete: (_stepName, updates: SessionUpdates = {}) => + updateSession((current) => { + Object.assign(current, filterSafeUpdates(updates)); + return current; + }), + markStepSkipped: () => cloneSession(session), + markStepFailed: (_stepName, message) => + updateSession((current) => { + current.status = "failed"; + current.failure = sanitizeFailure({ step: _stepName, message, recordedAt: "now" }); + return current; + }), + completeSession: (updates: SessionUpdates = {}) => + updateSession((current) => { + Object.assign(current, filterSafeUpdates(updates)); + current.status = "complete"; + current.resumable = false; + return current; + }), + filterSafeUpdates, + emitEvent: () => undefined, + now: () => "2026-05-28T00:00:00.000Z", + }; + return new OnboardRuntime(deps); +} + +describe("runOnboardMachine", () => { + it("runs handlers until completion while applying retry and branch transitions", async () => { + const runtime = createRuntime(); + const calls: string[] = []; + const handlers: OnboardStateHandlers = { + init: () => advanceTo("preflight"), + preflight: () => advanceTo("gateway"), + gateway: () => advanceTo("provider_selection"), + provider_selection: () => advanceTo("inference"), + inference: (context) => { + calls.push(`inference:${context.attempts}`); + return context.attempts === 0 ? retryTo("provider_selection") : advanceTo("sandbox"); + }, + sandbox: () => branchTo("openclaw"), + openclaw: () => advanceTo("policies"), + policies: () => advanceTo("finalizing"), + finalizing: () => advanceTo("post_verify"), + post_verify: () => completeOnboardMachine({ sandboxName: "my-assistant" }), + }; + + const result = await runOnboardMachine({ + context: { attempts: 0, visited: [] } as RunnerContext, + runtime, + handlers, + updateContext: ({ context, state }) => ({ + attempts: state === "inference" ? context.attempts + 1 : context.attempts, + visited: [...context.visited, state], + }), + }); + + expect(result.session).toMatchObject({ + status: "complete", + sandboxName: "my-assistant", + machine: { state: "complete" }, + }); + expect(calls).toEqual(["inference:0", "inference:1"]); + expect(result.context.visited).toEqual([ + "init", + "preflight", + "gateway", + "provider_selection", + "inference", + "provider_selection", + "inference", + "sandbox", + "openclaw", + "policies", + "finalizing", + "post_verify", + ]); + }); + + it("stops on failed terminal results", async () => { + const runtime = createRuntime(); + const policies = vi.fn(() => advanceTo("finalizing")); + + const result = await runOnboardMachine({ + context: { attempts: 0, visited: [] } as RunnerContext, + runtime, + handlers: { + init: () => advanceTo("preflight"), + preflight: () => failOnboardMachine("preflight failed", { step: "preflight" }), + policies, + }, + }); + + expect(result.session).toMatchObject({ + status: "failed", + failure: { step: "preflight", message: "preflight failed" }, + machine: { state: "failed" }, + }); + expect(policies).not.toHaveBeenCalled(); + }); + + it("returns immediately for terminal sessions", async () => { + const startedAt = "2026-05-28T00:00:00.000Z"; + const completeSession = createSession({ + resumable: false, + machine: { + version: MACHINE_SNAPSHOT_VERSION, + state: "complete", + stateEnteredAt: startedAt, + revision: 1, + }, + }); + completeSession.status = "complete"; + const runtime = createRuntime(completeSession); + const init = vi.fn(() => advanceTo("preflight")); + + const result = await runOnboardMachine({ + context: { attempts: 0, visited: [] } as RunnerContext, + runtime, + handlers: { init }, + }); + + expect(result.session).toMatchObject({ status: "complete", machine: { state: "complete" } }); + expect(init).not.toHaveBeenCalled(); + }); + + it("propagates runtime transition errors without updating context", async () => { + const runtime = createRuntime(); + const updateContext = vi.fn(({ context }) => context); + + await expect( + runOnboardMachine({ + context: { attempts: 0, visited: [] } as RunnerContext, + runtime, + handlers: { init: () => advanceTo("sandbox") }, + updateContext, + }), + ).rejects.toThrow("Invalid onboarding machine transition"); + expect(updateContext).not.toHaveBeenCalled(); + }); + + it("throws when a non-terminal state has no handler", async () => { + const runtime = createRuntime(); + + await expect( + runOnboardMachine({ + context: { attempts: 0, visited: [] } as RunnerContext, + runtime, + handlers: {}, + }), + ).rejects.toThrow(MissingOnboardStateHandlerError); + }); + + it("throws when retry-capable handlers exceed the transition limit", async () => { + const runtime = createRuntime(); + + await expect( + runOnboardMachine({ + context: { attempts: 0, visited: [] } as RunnerContext, + runtime, + handlers: { + init: () => advanceTo("preflight"), + preflight: () => advanceTo("gateway"), + gateway: () => advanceTo("provider_selection"), + provider_selection: () => advanceTo("inference"), + inference: () => retryTo("provider_selection"), + }, + maxTransitions: 5, + }), + ).rejects.toThrow(OnboardMachineTransitionLimitError); + }); +}); diff --git a/src/lib/onboard/machine/runner.ts b/src/lib/onboard/machine/runner.ts new file mode 100644 index 0000000000..126a9309ce --- /dev/null +++ b/src/lib/onboard/machine/runner.ts @@ -0,0 +1,100 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import type { Session } from "../../state/onboard-session"; +import type { OnboardStateResult } from "./result"; +import { isTerminalOnboardMachineState } from "./transitions"; +import type { OnboardMachineState, OnboardNonTerminalMachineState } from "./types"; + +export type OnboardStateHandler = ( + context: Context, +) => Promise | OnboardStateResult; + +export type OnboardStateHandlers = Partial< + Record> +>; + +export interface OnboardMachineRunnerRuntime { + session(): Promise; + applyResult(result: OnboardStateResult): Promise; +} + +export interface OnboardMachineRunnerOptions { + context: Context; + runtime: OnboardMachineRunnerRuntime; + handlers: OnboardStateHandlers; + /** + * Safety valve for retry-capable handlers. Handlers should bound their own + * retry loops, but the runner refuses to apply unbounded transitions. + */ + maxTransitions?: number; + updateContext?(input: { + context: Context; + state: OnboardMachineState; + result: OnboardStateResult; + session: Session; + }): Context | Promise; +} + +export interface OnboardMachineRunnerResult { + context: Context; + session: Session; +} + +export class MissingOnboardStateHandlerError extends Error { + readonly state: OnboardNonTerminalMachineState; + + constructor(state: OnboardNonTerminalMachineState) { + super(`Missing onboarding machine handler for state: ${state}`); + this.name = "MissingOnboardStateHandlerError"; + this.state = state; + } +} + +export class OnboardMachineTransitionLimitError extends Error { + readonly maxTransitions: number; + + constructor(maxTransitions: number) { + super(`Onboarding machine exceeded transition limit: ${maxTransitions}`); + this.name = "OnboardMachineTransitionLimitError"; + this.maxTransitions = maxTransitions; + } +} + +const DEFAULT_MAX_TRANSITIONS = 100; + +function normalizeMaxTransitions(value: number | undefined): number { + if (value === undefined) return DEFAULT_MAX_TRANSITIONS; + return Math.max(1, Math.trunc(value)); +} + +export async function runOnboardMachine({ + context: initialContext, + runtime, + handlers, + maxTransitions, + updateContext, +}: OnboardMachineRunnerOptions): Promise> { + let context = initialContext; + let session = await runtime.session(); + let transitions = 0; + const transitionLimit = normalizeMaxTransitions(maxTransitions); + + while (!isTerminalOnboardMachineState(session.machine.state)) { + if (transitions >= transitionLimit) { + throw new OnboardMachineTransitionLimitError(transitionLimit); + } + const state = session.machine.state; + const handler = handlers[state as OnboardNonTerminalMachineState]; + if (!handler) throw new MissingOnboardStateHandlerError(state as OnboardNonTerminalMachineState); + + const result = await handler(context); + session = await runtime.applyResult(result); + transitions += 1; + context = updateContext + ? await updateContext({ context, state, result, session }) + : context; + } + + return { context, session }; +}