diff --git a/src/lib/onboard/machine/core-flow-phases.ts b/src/lib/onboard/machine/core-flow-phases.ts index d86848663f..91011a26a8 100644 --- a/src/lib/onboard/machine/core-flow-phases.ts +++ b/src/lib/onboard/machine/core-flow-phases.ts @@ -2,7 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 import type { WebSearchConfig } from "../../inference/web-search"; -import type { OnboardFlowContext } from "./flow-context"; +import { assertProviderSelectedContext, type OnboardFlowContext } from "./flow-context"; import { runCoreOnboardFlowSequence } from "./flow-slices"; import { handleProviderInferenceState, @@ -97,9 +97,7 @@ export function createCoreOnboardFlowPhases< const sandboxPhase: OnboardSequencePhase = { state: "sandbox", async run(context) { - if (!context.model || !context.provider || !context.sandboxGpuConfig) { - throw new Error("Onboarding state is incomplete before sandbox setup."); - } + assertProviderSelectedContext(context, "sandbox setup"); const sandboxStateResult = await handleSandboxState({ resume: context.resume, fresh: context.fresh, diff --git a/src/lib/onboard/machine/final-flow-phases.ts b/src/lib/onboard/machine/final-flow-phases.ts index 7ed2685f03..10d33d5266 100644 --- a/src/lib/onboard/machine/final-flow-phases.ts +++ b/src/lib/onboard/machine/final-flow-phases.ts @@ -2,7 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 import type { WebSearchConfig } from "../../inference/web-search"; -import type { OnboardFlowContext } from "./flow-context"; +import { assertSandboxCreatedContext, type OnboardFlowContext } from "./flow-context"; import { createAgentSetupPhase, createFinalizationPhase, @@ -38,15 +38,6 @@ export interface FinalOnboardFlowPhaseOptions< >["deps"]; } -function requireFinalContext( - context: Context, - stepName: string, -): asserts context is Context & { sandboxName: string; model: string; provider: string } { - if (!context.sandboxName || !context.model || !context.provider) { - throw new Error(`Onboarding state is incomplete before ${stepName}.`); - } -} - export function createFinalOnboardFlowPhases< Context extends OnboardFlowContext, VerifyChain = unknown, @@ -57,7 +48,7 @@ export function createFinalOnboardFlowPhases< const createBranchPhase = options.branchState === "agent_setup" ? createAgentSetupPhase : createOpenclawSetupPhase; const branchSetupPhase = createBranchPhase(async (context) => { - requireFinalContext(context, "agent setup"); + assertSandboxCreatedContext(context, "agent setup"); const agentSetupResult = await handleAgentSetupState({ agent: context.agent, sandboxName: context.sandboxName, @@ -76,7 +67,7 @@ export function createFinalOnboardFlowPhases< }); const policiesPhase = createPoliciesPhase(async (context) => { - requireFinalContext(context, "policies"); + assertSandboxCreatedContext(context, "policies"); const policiesResult = await handlePoliciesState({ resume: context.resume, sandboxName: context.sandboxName, @@ -101,7 +92,7 @@ export function createFinalOnboardFlowPhases< }); const finalizationPhase = createFinalizationPhase(async (context) => { - requireFinalContext(context, "finalization"); + assertSandboxCreatedContext(context, "finalization"); const finalizationResult = await handleFinalizationState({ sandboxName: context.sandboxName, model: context.model, diff --git a/src/lib/onboard/machine/flow-context.test.ts b/src/lib/onboard/machine/flow-context.test.ts index c1f98d62fc..dbd7eafcdd 100644 --- a/src/lib/onboard/machine/flow-context.test.ts +++ b/src/lib/onboard/machine/flow-context.test.ts @@ -4,12 +4,14 @@ import { describe, expect, it } from "vitest"; import { createSession } from "../../state/onboard-session"; -import { advanceTo } from "./result"; import { + assertProviderSelectedContext, + assertSandboxCreatedContext, mergeOnboardFlowContext, - onboardFlowPhaseResult, type OnboardFlowContext, + onboardFlowPhaseResult, } from "./flow-context"; +import { advanceTo } from "./result"; function baseContext(): OnboardFlowContext { return { @@ -61,4 +63,41 @@ describe("onboard flow context helpers", () => { expect(result.context.provider).toBe("nvidia-prod"); expect(result.result).toMatchObject({ next: "gateway", transitionKind: "advance" }); }); + + it("asserts provider-selected context before sandbox setup", () => { + const context = mergeOnboardFlowContext(baseContext(), { + provider: "nvidia-prod", + model: "model", + }); + + expect(() => assertProviderSelectedContext(context, "sandbox setup")).not.toThrow(); + }); + + it("rejects missing provider-selected context fields", () => { + expect(() => assertProviderSelectedContext(baseContext(), "sandbox setup")).toThrow( + /Onboarding state is incomplete before sandbox setup\./, + ); + }); + + it("asserts sandbox-created context before final phases", () => { + const context = mergeOnboardFlowContext(baseContext(), { + sandboxName: "my-assistant", + provider: "nvidia-prod", + model: "model", + sandboxGpuConfig: null, + }); + + expect(() => assertSandboxCreatedContext(context, "policies")).not.toThrow(); + }); + + it("rejects missing sandbox name before final phases", () => { + const context = mergeOnboardFlowContext(baseContext(), { + provider: "nvidia-prod", + model: "model", + }); + + expect(() => assertSandboxCreatedContext(context, "policies")).toThrow( + /Onboarding state is incomplete before policies\./, + ); + }); }); diff --git a/src/lib/onboard/machine/flow-context.ts b/src/lib/onboard/machine/flow-context.ts index 0d36fa1f3f..761e71681b 100644 --- a/src/lib/onboard/machine/flow-context.ts +++ b/src/lib/onboard/machine/flow-context.ts @@ -30,11 +30,44 @@ export interface OnboardFlowContext = Context & { + model: string; + provider: string; + sandboxGpuConfig: NonNullable; +}; + +export type SandboxCreatedOnboardFlowContext = Context & { + sandboxName: string; + model: string; + provider: string; +}; + +export type FinalOnboardFlowContext = + SandboxCreatedOnboardFlowContext; + export interface OnboardFlowPhaseResult { context: Context; result: OnboardStateHandlerResult; } +export function assertProviderSelectedContext( + context: Context, + stepName: string, +): asserts context is ProviderSelectedOnboardFlowContext { + if (!context.model || !context.provider || !context.sandboxGpuConfig) { + throw new Error(`Onboarding state is incomplete before ${stepName}.`); + } +} + +export function assertSandboxCreatedContext( + context: Context, + stepName: string, +): asserts context is SandboxCreatedOnboardFlowContext { + if (!context.sandboxName || !context.model || !context.provider) { + throw new Error(`Onboarding state is incomplete before ${stepName}.`); + } +} + export function mergeOnboardFlowContext( context: Context, patch: Partial,