diff --git a/src/lib/onboard/machine/flow-slices.test.ts b/src/lib/onboard/machine/flow-slices.test.ts new file mode 100644 index 0000000000..5afd843e4a --- /dev/null +++ b/src/lib/onboard/machine/flow-slices.test.ts @@ -0,0 +1,135 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { describe, expect, it } from "vitest"; + +import { + createSession, + filterSafeUpdates, + MACHINE_SNAPSHOT_VERSION, + normalizeSession, + type Session, + type SessionUpdates, +} from "../../state/onboard-session"; +import type { OnboardFlowContext } from "./flow-context"; +import { advanceTo } from "./result"; +import { OnboardRuntime, type OnboardRuntimeDeps } from "./runtime"; +import type { OnboardSequencePhase } from "./sequence-runner"; +import { initialOnboardFlowPhases, runInitialOnboardFlowSequence } from "./flow-slices"; + +function cloneSession(session: Session): Session { + return normalizeSession(JSON.parse(JSON.stringify(session))) ?? session; +} + +function runtime(initialSession: Session = createSession()) { + let session = cloneSession(initialSession); + const updateSession = (mutator: (value: Session) => Session | void): Session => { + session = cloneSession(mutator(cloneSession(session)) ?? session); + 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))), + markStepCompleteRecordOnly: (_stepName, updates: SessionUpdates = {}) => + updateSession((current) => Object.assign(current, filterSafeUpdates(updates))), + markStepSkipped: () => cloneSession(session), + markStepFailed: () => cloneSession(session), + markStepFailedRecordOnly: () => cloneSession(session), + completeSession: () => cloneSession(session), + filterSafeUpdates, + emitEvent: () => undefined, + now: () => "2026-05-29T00:00:00.000Z", + }; + return new OnboardRuntime(deps); +} + +function context(): OnboardFlowContext { + return { + resume: false, + fresh: false, + session: createSession(), + agent: null, + recordedSandboxName: null, + requestedSandboxName: null, + sandboxName: null, + fromDockerfile: null, + model: null, + provider: null, + endpointUrl: null, + credentialEnv: null, + hermesAuthMethod: null, + hermesToolGateways: [], + preferredInferenceApi: null, + nimContainer: null, + webSearchConfig: null, + webSearchSupported: false, + selectedMessagingChannels: [], + gpu: null, + sandboxGpuConfig: null, + gpuPassthrough: false, + }; +} + +function phase( + state: OnboardSequencePhase["state"], + next: ReturnType["next"], +): OnboardSequencePhase { + return { state, run: (ctx) => ({ context: ctx, result: advanceTo(next) }) }; +} + +describe("onboard flow slices", () => { + it("selects only initial preflight/gateway phases", () => { + expect( + initialOnboardFlowPhases([ + phase("preflight", "gateway"), + phase("gateway", "provider_selection"), + phase("provider_selection", "inference"), + ]).map((entry) => entry.state), + ).toEqual(["init", "preflight", "gateway"]); + }); + + it("runs the initial slice from a default session and stops at provider selection", async () => { + const result = await runInitialOnboardFlowSequence({ + context: context(), + runtime: runtime(), + phases: [ + phase("preflight", "gateway"), + phase("gateway", "provider_selection"), + phase("provider_selection", "inference"), + ], + }); + + expect(result.session.machine.state).toBe("provider_selection"); + }); + + it("runs the initial slice from preflight and stops at provider selection", async () => { + const result = await runInitialOnboardFlowSequence({ + context: context(), + runtime: runtime( + createSession({ + machine: { + version: MACHINE_SNAPSHOT_VERSION, + state: "preflight", + stateEnteredAt: "2026-05-29T00:00:00.000Z", + revision: 0, + }, + }), + ), + phases: [ + phase("preflight", "gateway"), + phase("gateway", "provider_selection"), + phase("provider_selection", "inference"), + ], + }); + + expect(result.session.machine.state).toBe("provider_selection"); + }); +}); diff --git a/src/lib/onboard/machine/flow-slices.ts b/src/lib/onboard/machine/flow-slices.ts new file mode 100644 index 0000000000..bf88e7e75b --- /dev/null +++ b/src/lib/onboard/machine/flow-slices.ts @@ -0,0 +1,33 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import type { OnboardFlowContext } from "./flow-context"; +import { onboardFlowPhaseResult } from "./flow-context"; +import { advanceTo } from "./result"; +import type { OnboardMachineRunnerRuntime } from "./runner"; +import type { OnboardSequencePhase } from "./sequence-runner"; +import { runOnboardSequenceWithRunner } from "./sequence-runner"; + +export function initialOnboardFlowPhases( + phases: readonly OnboardSequencePhase[], +): OnboardSequencePhase[] { + return [ + { + state: "init", + run: (context) => onboardFlowPhaseResult(context, advanceTo("preflight")), + }, + ...phases.filter((phase) => phase.state === "preflight" || phase.state === "gateway"), + ]; +} + +export async function runInitialOnboardFlowSequence(options: { + context: Context; + runtime: OnboardMachineRunnerRuntime; + phases: readonly OnboardSequencePhase[]; +}) { + return runOnboardSequenceWithRunner({ + ...options, + phases: initialOnboardFlowPhases(options.phases), + stopStates: ["provider_selection"], + }); +} diff --git a/src/lib/onboard/machine/sequence-runner.ts b/src/lib/onboard/machine/sequence-runner.ts index 6a9f5e86e4..b2d71a18e6 100644 --- a/src/lib/onboard/machine/sequence-runner.ts +++ b/src/lib/onboard/machine/sequence-runner.ts @@ -27,6 +27,7 @@ export interface OnboardSequenceRunnerOptions { phases: readonly OnboardSequencePhase[]; maxTransitions?: OnboardMachineRunnerOptions["maxTransitions"]; sequenceOwnership?: OnboardMachineRunnerOptions["sequenceOwnership"]; + stopStates?: OnboardMachineRunnerOptions["stopStates"]; } export class DuplicateOnboardSequencePhaseError extends Error { @@ -69,6 +70,7 @@ export async function runOnboardSequenceWithRunner({ phases, maxTransitions, sequenceOwnership, + stopStates, }: OnboardSequenceRunnerOptions) { let pendingContext = initialContext; return runOnboardMachine({ @@ -76,6 +78,7 @@ export async function runOnboardSequenceWithRunner({ runtime, maxTransitions, sequenceOwnership, + stopStates, handlers: buildOnboardSequenceHandlers(phases, (context) => { pendingContext = context; }),