diff --git a/docs/reference/commands.mdx b/docs/reference/commands.mdx index ad3c75cef5..cc37b748fb 100644 --- a/docs/reference/commands.mdx +++ b/docs/reference/commands.mdx @@ -241,7 +241,7 @@ If the session lacks that registration receipt, the provider is missing, or its Completed onboarding sessions are not resumable. -Use `--resume` only for interrupted `in_progress` sessions, not to change provider, model, agent, or sandbox recreation settings after onboarding has completed. +Use `--resume` only for resumable interrupted or failed sessions, not to change provider, model, agent, or sandbox recreation settings after onboarding has completed. During resume, NemoClaw reruns preflight, gateway, provider, and sandbox repair checks even when the saved session has already reached a later nonterminal onboarding phase. If the recorded session conflicts with flags you pass on the recovery run, NemoClaw exits and tells you to either rerun with the original settings or start over. diff --git a/scripts/install.sh b/scripts/install.sh index 0eb66b3494..1a9ba56357 100755 --- a/scripts/install.sh +++ b/scripts/install.sh @@ -3088,6 +3088,7 @@ activate_express_install() { export NEMOCLAW_NON_INTERACTIVE_SUDO_MODE=prompt export NEMOCLAW_YES=1 export NEMOCLAW_POLICY_MODE=suggested + unset NEMOCLAW_STATION_EXPRESS case "$platform" in "DGX Spark") export NEMOCLAW_SANDBOX_NAME="${NEMOCLAW_SANDBOX_NAME:-my-assistant}" @@ -3097,6 +3098,7 @@ activate_express_install() { fi ;; "DGX Station") + export NEMOCLAW_STATION_EXPRESS=1 export NEMOCLAW_SANDBOX_NAME="${NEMOCLAW_SANDBOX_NAME:-my-assistant}" export NEMOCLAW_PROVIDER=install-vllm configure_station_express_model diff --git a/src/lib/onboard.ts b/src/lib/onboard.ts index 035d7686d5..2337ad4913 100644 --- a/src/lib/onboard.ts +++ b/src/lib/onboard.ts @@ -4000,7 +4000,7 @@ async function preflightAuthoritativeRebuildTarget( } // ── Main ───────────────────────────────────────────────────────── -const onboard = onboardEntryOptions.withNonInteractiveEnvironment(runOnboard); +const onboard = onboardEntryOptions.wrapOnboard(runOnboard, onboardSession.loadSession); async function runOnboard(opts: OnboardOptions = {}): Promise { setupInferenceFactory.assertNoOpenShellGatewayEndpointOverride(); const runtimeControlRequests = runtimeControlFlow.applyOnboardRuntimeControlRequests(opts); @@ -4055,7 +4055,7 @@ async function runOnboard(opts: OnboardOptions = {}): Promise { } // Validate provider/model hints before preflight so configuration errors are not reported as Docker failures. // biome-ignore format: keep src/lib/onboard.ts net-neutral for growth guardrail. - resumeConfig.preflightEarlyOnboardEnvForResume(isNonInteractive(), opts.authoritativeResumeConfig === true); + const stationSessionInput = onboardEntryOptions.prepareSessionInput(runtimeControlRequests, requestedSandboxName, resume, () => resumeConfig.preflightEarlyOnboardEnvForResume(isNonInteractive(), opts.authoritativeResumeConfig === true)); const ownsOnboardLock = opts.onboardLockAlreadyHeld !== true; const lockResult = ownsOnboardLock ? onboardSession.acquireOnboardLock( @@ -4148,7 +4148,7 @@ async function runOnboard(opts: OnboardOptions = {}): Promise { authoritativeResumeConfig: opts.authoritativeResumeConfig === true, agentFlag: opts.agent || null, envAgent: process.env.NEMOCLAW_AGENT || null, - ...runtimeControlRequests, + ...stationSessionInput, }, { loadSession: onboardSession.loadSession, diff --git a/src/lib/onboard/entry-options.ts b/src/lib/onboard/entry-options.ts index 97c80328fd..1ca1e79467 100644 --- a/src/lib/onboard/entry-options.ts +++ b/src/lib/onboard/entry-options.ts @@ -1,6 +1,12 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +import { + requireStationExpressResumeIntent, + type StationExpressSessionLike, + wrapOnboard as wrapStationExpressOnboard, +} from "./station-express-resume"; + export interface OnboardEntryOptionsInput { opts: { resume?: boolean; @@ -44,6 +50,7 @@ export interface ResolvedOnboardEntryOptions { } type NonInteractiveEntryOptions = { nonInteractive?: boolean }; +type ResumableEntryOptions = NonInteractiveEntryOptions & { resume?: boolean; fresh?: boolean }; /** Scope the CLI flag to helpers that still read the compatibility environment variable. */ export function withNonInteractiveEnvironment( @@ -64,6 +71,26 @@ export function withNonInteractiveEnvironment( + run: (options?: Options) => Promise, + loadSession: () => StationExpressSessionLike | null, +): (options?: Options) => Promise { + return wrapStationExpressOnboard(withNonInteractiveEnvironment(run), loadSession); +} + +export function prepareSessionInput( + runtimeControlRequests: RuntimeControlRequests, + sandboxName: string | null, + resume: boolean, + preflight: () => void, +) { + preflight(); + return { + ...runtimeControlRequests, + stationExpressIntent: requireStationExpressResumeIntent(process.env, sandboxName, resume), + }; +} + export function resolveOnboardEntryOptions( input: OnboardEntryOptionsInput, deps: OnboardEntryOptionsDeps, diff --git a/src/lib/onboard/machine/handlers/provider-inference.test.ts b/src/lib/onboard/machine/handlers/provider-inference.test.ts index 9c6b857a27..07f87d481e 100644 --- a/src/lib/onboard/machine/handlers/provider-inference.test.ts +++ b/src/lib/onboard/machine/handlers/provider-inference.test.ts @@ -419,6 +419,34 @@ describe("handleProviderInferenceState", () => { expect(calls.deleteEnv).toHaveBeenCalledWith("COMPATIBLE_API_KEY"); }); + it("retains Station Express intent without committing a failed managed provider selection", async () => { + const setupNim = vi.fn(async () => { + throw new Error("injected managed vLLM download failure"); + }); + const { deps, calls } = createDeps({ setupNim }); + const session = createSession({ + mode: "non-interactive", + stationExpressIntent: { + version: 1, + model: "nemotron-3-ultra-550b-a55b", + sandboxName: "my-assistant", + }, + }); + + await expect(handleProviderInferenceState(baseOptions(deps, session))).rejects.toThrow( + "injected managed vLLM download failure", + ); + + expect(session.stationExpressIntent).toEqual({ + version: 1, + model: "nemotron-3-ultra-550b-a55b", + sandboxName: "my-assistant", + }); + expect(session.provider).toBeNull(); + expect(session.model).toBeNull(); + expect(calls.complete).not.toHaveBeenCalledWith("provider_selection", expect.anything()); + }); + it("exits through the injected CLI boundary when provider selection is incomplete", async () => { const setupNim = vi.fn(async () => ({ ...baseSelection, model: null })); const { deps, calls } = createDeps({ setupNim }); diff --git a/src/lib/onboard/session-bootstrap.test.ts b/src/lib/onboard/session-bootstrap.test.ts index 3b64d57f25..244765fb07 100644 --- a/src/lib/onboard/session-bootstrap.test.ts +++ b/src/lib/onboard/session-bootstrap.test.ts @@ -87,6 +87,32 @@ describe("prepareOnboardSession", () => { expect(getSession()?.sessionId).not.toBe("old-session"); }); + it("checkpoints Station Express choices before managed vLLM setup", async () => { + const { deps } = createDeps(); + const stationExpress = { + version: 1 as const, + model: "nemotron-3-ultra-550b-a55b", + sandboxName: "my-assistant", + }; + + const result = await prepareOnboardSession( + { + resume: false, + fresh: false, + requestedFromDockerfile: null, + requestedSandboxName: "my-assistant", + cannotPrompt: true, + nonInteractive: true, + stationExpressIntent: stationExpress, + }, + deps, + ); + + expect(result.session?.stationExpressIntent).toEqual(stationExpress); + expect(result.session?.provider).toBeNull(); + expect(result.session?.model).toBeNull(); + }); + it("defaults a fresh session to progressive disclosure", async () => { const { deps } = createDeps(); const result = await prepareOnboardSession( diff --git a/src/lib/onboard/session-bootstrap.ts b/src/lib/onboard/session-bootstrap.ts index ca776f7181..d67abbabf4 100644 --- a/src/lib/onboard/session-bootstrap.ts +++ b/src/lib/onboard/session-bootstrap.ts @@ -4,6 +4,7 @@ import type { Session } from "../state/onboard-session"; import { DEFAULT_TOOL_DISCLOSURE, type ToolDisclosure } from "../tool-disclosure"; import type { ResumeConfigConflict } from "./resume-config"; +import type { StationExpressResumeIntent } from "./station-express-resume"; export interface OnboardSessionBootstrapInput { resume: boolean; @@ -17,6 +18,7 @@ export interface OnboardSessionBootstrapInput { envAgent?: string | null; requestedToolDisclosure?: ToolDisclosure | null; requestedObservabilityEnabled?: boolean | null; + stationExpressIntent?: StationExpressResumeIntent | null; } export interface OnboardSessionBootstrapDeps { @@ -225,6 +227,7 @@ function prepareFreshSession( toolDisclosure: input.requestedToolDisclosure ?? DEFAULT_TOOL_DISCLOSURE, observabilityEnabled: input.requestedObservabilityEnabled === true, observabilityRequestedExplicitly: typeof input.requestedObservabilityEnabled === "boolean", + stationExpressIntent: input.stationExpressIntent ?? null, metadata: { gatewayName: "nemoclaw", fromDockerfile: fromDockerfile || null }, }), ); diff --git a/src/lib/onboard/station-express-resume.test.ts b/src/lib/onboard/station-express-resume.test.ts new file mode 100644 index 0000000000..f937fe4cc6 --- /dev/null +++ b/src/lib/onboard/station-express-resume.test.ts @@ -0,0 +1,164 @@ +// 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 { createSession } from "../state/onboard-session"; +import { + getStationExpressResumeIntent, + parseStationExpressResumeIntent, + STATION_EXPRESS_ENV, + withStationExpressResumeEnvironment, +} from "./station-express-resume"; + +const ultraIntent = { + version: 1 as const, + model: "nemotron-3-ultra-550b-a55b", + sandboxName: "my-assistant", +}; + +function expressEnv(): NodeJS.ProcessEnv { + return { + [STATION_EXPRESS_ENV]: "1", + NEMOCLAW_NON_INTERACTIVE: "1", + NEMOCLAW_YES: "1", + NEMOCLAW_POLICY_MODE: "suggested", + NEMOCLAW_SANDBOX_NAME: "my-assistant", + NEMOCLAW_PROVIDER: "install-vllm", + NEMOCLAW_VLLM_MODEL: "nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-NVFP4", + NEMOCLAW_MODEL: "nvidia/nemotron-3-ultra-550b-a55b", + }; +} + +function resumeDeps( + session = createSession({ mode: "non-interactive", stationExpressIntent: ultraIntent }), +) { + return { + loadSession: vi.fn(() => session), + error: vi.fn(), + exitProcess: vi.fn((code: number): never => { + throw new Error(`exit ${String(code)}`); + }), + }; +} + +describe("DGX Station Express resume", () => { + it("captures a canonical secret-free intent from the installer environment", () => { + expect(getStationExpressResumeIntent(expressEnv(), "my-assistant")).toEqual({ + ok: true, + intent: ultraIntent, + }); + }); + + it("ignores ordinary onboarding without the Station Express marker", () => { + expect(getStationExpressResumeIntent({}, null)).toEqual({ ok: true, intent: null }); + }); + + it("rejects malformed or expanded persisted intent", () => { + expect( + parseStationExpressResumeIntent({ ...ultraIntent, token: "must-not-persist" }), + ).toBeNull(); + expect( + parseStationExpressResumeIntent({ ...ultraIntent, model: "qwen3.6-35b-a3b-nvfp4" }), + ).toBeNull(); + }); + + it("restores the saved provider and model for a plain failed-session resume", async () => { + const env: NodeJS.ProcessEnv = { NEMOCLAW_PROVIDER: "" }; + const deps = resumeDeps( + createSession({ + mode: "non-interactive", + status: "failed", + stationExpressIntent: ultraIntent, + }), + ); + const run = vi.fn(async () => { + expect(env).toMatchObject({ + NEMOCLAW_STATION_EXPRESS: "1", + NEMOCLAW_NON_INTERACTIVE: "1", + NEMOCLAW_YES: "1", + NEMOCLAW_POLICY_MODE: "suggested", + NEMOCLAW_SANDBOX_NAME: "my-assistant", + NEMOCLAW_PROVIDER: "install-vllm", + NEMOCLAW_VLLM_MODEL: "nemotron-3-ultra-550b-a55b", + NEMOCLAW_MODEL: "nvidia/nemotron-3-ultra-550b-a55b", + }); + }); + + await withStationExpressResumeEnvironment(run, deps, env)({ resume: true }); + + expect(run).toHaveBeenCalledTimes(1); + expect(env).toEqual({ NEMOCLAW_PROVIDER: "" }); + }); + + it("also restores an automatically resumed in-progress Express session", async () => { + const env: NodeJS.ProcessEnv = {}; + const deps = resumeDeps(); + const run = vi.fn(async () => { + expect(env.NEMOCLAW_PROVIDER).toBe("install-vllm"); + }); + + await withStationExpressResumeEnvironment(run, deps, env)({}); + + expect(run).toHaveBeenCalledTimes(1); + expect(env).toEqual({}); + }); + + it("reuses a completed provider selection without replaying managed installation", async () => { + const completeProviderStep = { + status: "complete" as const, + startedAt: "2026-07-16T00:00:00.000Z", + completedAt: "2026-07-16T00:01:00.000Z", + error: null, + }; + const session = createSession({ + mode: "non-interactive", + status: "failed", + stationExpressIntent: ultraIntent, + provider: "vllm-local", + model: "nvidia/nemotron-3-ultra-550b-a55b", + steps: { + provider_selection: completeProviderStep, + }, + }); + const env: NodeJS.ProcessEnv = {}; + const deps = resumeDeps(session); + const run = vi.fn(async () => { + expect(env.NEMOCLAW_NON_INTERACTIVE).toBe("1"); + expect(env.NEMOCLAW_POLICY_MODE).toBe("suggested"); + expect(env.NEMOCLAW_PROVIDER).toBeUndefined(); + expect(env.NEMOCLAW_VLLM_MODEL).toBeUndefined(); + expect(env.NEMOCLAW_MODEL).toBeUndefined(); + }); + + await withStationExpressResumeEnvironment(run, deps, env)({ resume: true }); + + expect(run).toHaveBeenCalledTimes(1); + expect(env).toEqual({}); + }); + + it("does not restore discarded intent for --fresh", async () => { + const env: NodeJS.ProcessEnv = {}; + const deps = resumeDeps(); + const run = vi.fn(async () => { + expect(env.NEMOCLAW_PROVIDER).toBeUndefined(); + }); + + await withStationExpressResumeEnvironment(run, deps, env)({ fresh: true }); + + expect(run).toHaveBeenCalledTimes(1); + }); + + it("fails closed when an explicit resume override selects another model", async () => { + const env: NodeJS.ProcessEnv = { NEMOCLAW_VLLM_MODEL: "deepseek-v4-flash" }; + const deps = resumeDeps(); + const run = vi.fn(async () => undefined); + + await expect( + withStationExpressResumeEnvironment(run, deps, env)({ resume: true }), + ).rejects.toThrow("exit 1"); + + expect(run).not.toHaveBeenCalled(); + expect(deps.error).toHaveBeenCalledWith(expect.stringContaining("NEMOCLAW_VLLM_MODEL")); + }); +}); diff --git a/src/lib/onboard/station-express-resume.ts b/src/lib/onboard/station-express-resume.ts new file mode 100644 index 0000000000..c6087b987e --- /dev/null +++ b/src/lib/onboard/station-express-resume.ts @@ -0,0 +1,276 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { selectVllmModelFromEnv, type VllmModelDef } from "../inference/vllm-models"; +import { NAME_MAX_LENGTH, NAME_VALID_PATTERN } from "../name-validation"; + +export const STATION_EXPRESS_ENV = "NEMOCLAW_STATION_EXPRESS"; +export const STATION_EXPRESS_INTENT_VERSION = 1; + +export interface StationExpressResumeIntent { + version: typeof STATION_EXPRESS_INTENT_VERSION; + model: string; + sandboxName: string; +} + +export interface StationExpressSessionLike { + resumable?: boolean; + status?: string; + mode?: string; + stationExpressIntent?: StationExpressResumeIntent | null; + steps?: { provider_selection?: { status?: string | null } | null } | null; +} + +interface ResumeOptionsLike { + resume?: boolean; + fresh?: boolean; +} + +interface StationExpressResumeDeps { + loadSession(): StationExpressSessionLike | null; + error(message: string): void; + exitProcess(code: number): never; +} + +type StationExpressFailureDeps = Pick; + +type IntentResult = + | { ok: true; intent: StationExpressResumeIntent | null } + | { ok: false; message: string }; + +const RESUME_ENV = [ + STATION_EXPRESS_ENV, + "NEMOCLAW_NON_INTERACTIVE", + "NEMOCLAW_YES", + "NEMOCLAW_POLICY_MODE", + "NEMOCLAW_SANDBOX_NAME", + "NEMOCLAW_PROVIDER", + "NEMOCLAW_VLLM_MODEL", + "NEMOCLAW_MODEL", +] as const; + +function isObject(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +function stationModel(value: unknown): VllmModelDef | null { + if (typeof value !== "string") return null; + try { + const model = selectVllmModelFromEnv({ NEMOCLAW_VLLM_MODEL: value }); + return model?.platforms.includes("station") ? model : null; + } catch { + return null; + } +} + +function servedModel(model: VllmModelDef): string { + return model.servedModelId ?? model.id; +} + +function validSandboxName(value: unknown): value is string { + return ( + typeof value === "string" && value.length <= NAME_MAX_LENGTH && NAME_VALID_PATTERN.test(value) + ); +} + +export function parseStationExpressResumeIntent(value: unknown): StationExpressResumeIntent | null { + if (!isObject(value)) return null; + const keys = Object.keys(value).sort(); + if (keys.join(",") !== "model,sandboxName,version") return null; + if (value.version !== STATION_EXPRESS_INTENT_VERSION) return null; + const model = stationModel(value.model); + if (!model || value.model !== model.envValue || !validSandboxName(value.sandboxName)) return null; + return { + version: STATION_EXPRESS_INTENT_VERSION, + model: model.envValue, + sandboxName: value.sandboxName, + }; +} + +function expectedEnvironment( + intent: StationExpressResumeIntent, + includeProviderSelection = true, +): Partial> | null { + const model = stationModel(intent.model); + if (!model) return null; + const expected: Partial> = { + [STATION_EXPRESS_ENV]: "1", + NEMOCLAW_NON_INTERACTIVE: "1", + NEMOCLAW_YES: "1", + NEMOCLAW_POLICY_MODE: "suggested", + NEMOCLAW_SANDBOX_NAME: intent.sandboxName, + }; + if (includeProviderSelection) { + expected.NEMOCLAW_PROVIDER = "install-vllm"; + expected.NEMOCLAW_VLLM_MODEL = model.envValue; + expected.NEMOCLAW_MODEL = servedModel(model); + } + return expected; +} + +function equivalentEnvironmentValue( + name: (typeof RESUME_ENV)[number], + actual: string, + expected: string, +): boolean { + if (name === "NEMOCLAW_VLLM_MODEL") { + return stationModel(actual)?.envValue === expected; + } + if (name === "NEMOCLAW_SANDBOX_NAME") { + return actual.trim().toLowerCase() === expected; + } + return actual.trim().toLowerCase() === expected.toLowerCase(); +} + +function validateExpectedEnvironment( + env: NodeJS.ProcessEnv, + expected: Partial>, +): string | null { + for (const name of RESUME_ENV) { + const expectedValue = expected[name]; + if (expectedValue === undefined) continue; + const actual = env[name]; + if (typeof actual !== "string" || actual.trim().length === 0) continue; + if (!equivalentEnvironmentValue(name, actual, expectedValue)) return name; + } + return null; +} + +export function getStationExpressResumeIntent( + env: NodeJS.ProcessEnv, + sandboxName: string | null, +): IntentResult { + const marker = String(env[STATION_EXPRESS_ENV] ?? "").trim(); + if (!marker) return { ok: true, intent: null }; + if (marker !== "1") { + return { ok: false, message: `${STATION_EXPRESS_ENV} must be 1 when set.` }; + } + + const model = stationModel(env.NEMOCLAW_VLLM_MODEL); + if (!model || !sandboxName || !validSandboxName(sandboxName)) { + return { + ok: false, + message: "DGX Station Express requires a registered Station vLLM model and sandbox name.", + }; + } + const intent: StationExpressResumeIntent = { + version: STATION_EXPRESS_INTENT_VERSION, + model: model.envValue, + sandboxName, + }; + const expected = expectedEnvironment(intent); + if (!expected) { + return { ok: false, message: "DGX Station Express model state is invalid." }; + } + for (const name of [ + STATION_EXPRESS_ENV, + "NEMOCLAW_NON_INTERACTIVE", + "NEMOCLAW_YES", + "NEMOCLAW_POLICY_MODE", + "NEMOCLAW_PROVIDER", + ] as const) { + const actual = String(env[name] ?? ""); + const expectedValue = expected[name]; + if (!expectedValue || !equivalentEnvironmentValue(name, actual, expectedValue)) { + return { ok: false, message: `DGX Station Express requires ${name}=${expectedValue}.` }; + } + } + const conflict = validateExpectedEnvironment(env, expected); + if (conflict) { + return { ok: false, message: `DGX Station Express has a conflicting ${conflict} value.` }; + } + return { ok: true, intent }; +} + +/** Validate initial Station Express intent before onboarding acquires its session lock. */ +export function requireStationExpressResumeIntent( + env: NodeJS.ProcessEnv, + sandboxName: string | null, + resume: boolean, + deps: StationExpressFailureDeps = { + error: (message) => console.error(message), + exitProcess: (code) => process.exit(code), + }, +): StationExpressResumeIntent | null { + if (resume) return null; + const result = getStationExpressResumeIntent(env, sandboxName); + if (!result.ok) { + deps.error(` ${result.message}`); + deps.exitProcess(1); + } + return result.intent; +} + +function shouldRestoreStationExpress( + options: ResumeOptionsLike | undefined, + session: StationExpressSessionLike | null, +): session is StationExpressSessionLike & { stationExpressIntent: StationExpressResumeIntent } { + if (options?.fresh === true || !session?.stationExpressIntent || session.resumable === false) + return false; + return options?.resume === true || session.status === "in_progress"; +} + +export function withStationExpressResumeEnvironment( + run: (options?: Options) => Promise, + deps: StationExpressResumeDeps, + env: NodeJS.ProcessEnv = process.env, +): (options?: Options) => Promise { + return async (options) => { + const session = deps.loadSession(); + if (!shouldRestoreStationExpress(options, session)) return run(options); + const intent = parseStationExpressResumeIntent(session.stationExpressIntent); + if (session.mode !== "non-interactive" || !intent) { + deps.error( + " DGX Station Express resume state is invalid. Run nemoclaw onboard --fresh to start again.", + ); + deps.exitProcess(1); + } + + const expected = expectedEnvironment( + intent, + session.steps?.provider_selection?.status !== "complete", + ); + if (!expected) { + deps.error( + " DGX Station Express resume model is no longer supported. Run nemoclaw onboard --fresh to start again.", + ); + deps.exitProcess(1); + } + const conflict = validateExpectedEnvironment(env, expected); + if (conflict) { + deps.error( + ` DGX Station Express resume conflicts with ${conflict}. Unset ${conflict} and rerun nemoclaw onboard --resume, or run nemoclaw onboard --fresh to start again.`, + ); + deps.exitProcess(1); + } + + const previous = new Map<(typeof RESUME_ENV)[number], string | undefined>(); + for (const name of RESUME_ENV) { + const expectedValue = expected[name]; + if (expectedValue === undefined) continue; + previous.set(name, env[name]); + env[name] = expectedValue; + } + try { + await run(options); + } finally { + for (const name of RESUME_ENV) { + if (!previous.has(name)) continue; + const value = previous.get(name); + if (value === undefined) delete env[name]; + else env[name] = value; + } + } + }; +} + +export function wrapOnboard( + run: (options?: Options) => Promise, + loadSession: StationExpressResumeDeps["loadSession"], +): (options?: Options) => Promise { + return withStationExpressResumeEnvironment(run, { + loadSession, + error: (message) => console.error(message), + exitProcess: (code) => process.exit(code), + }); +} diff --git a/src/lib/state/onboard-session-station-express.test.ts b/src/lib/state/onboard-session-station-express.test.ts new file mode 100644 index 0000000000..eb1e900f75 --- /dev/null +++ b/src/lib/state/onboard-session-station-express.test.ts @@ -0,0 +1,89 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +type OnboardSessionModule = typeof import("./onboard-session"); +type LoadedSession = NonNullable>; +let session: OnboardSessionModule; +let tmpDir: string; + +function requireLoadedSession( + loaded: ReturnType, +): LoadedSession { + expect(loaded).not.toBeNull(); + return loaded as LoadedSession; +} + +beforeEach(async () => { + tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), "nemoclaw-station-express-session-")); + vi.stubEnv("HOME", tmpDir); + vi.resetModules(); + session = await import("./onboard-session"); + session.clearSession(); + session.releaseOnboardLock(); +}); + +afterEach(() => { + session.clearSession(); + session.releaseOnboardLock(); + fs.rmSync(tmpDir, { recursive: true, force: true }); + vi.unstubAllEnvs(); +}); + +describe("Station Express onboarding session state", () => { + it("round-trips only canonical secret-free resume state", () => { + const stationExpress = { + version: 1 as const, + model: "nemotron-3-ultra-550b-a55b", + sandboxName: "my-assistant", + }; + session.saveSession( + session.createSession({ mode: "non-interactive", stationExpressIntent: stationExpress }), + ); + + expect(requireLoadedSession(session.loadSession()).stationExpressIntent).toEqual( + stationExpress, + ); + expect(fs.readFileSync(session.SESSION_FILE, "utf8")).not.toContain("token"); + }); + + it("accepts legacy sessions without resume state and rejects malformed state", () => { + const legacy = session.createSession() as unknown as Record; + delete legacy.stationExpressIntent; + expect( + requireLoadedSession(session.normalizeSession(legacy as never)).stationExpressIntent, + ).toBeNull(); + + const malformed = { + ...session.createSession({ mode: "non-interactive" }), + stationExpressIntent: { + version: 1, + model: "nemotron-3-ultra-550b-a55b", + sandboxName: "my-assistant", + HF_TOKEN: "must-not-persist", + }, + }; + expect(session.normalizeSession(malformed as never)).toBeNull(); + }); + + it("clears resume intent only after successful completion", () => { + session.saveSession( + session.createSession({ + mode: "non-interactive", + stationExpressIntent: { + version: 1, + model: "nemotron-3-ultra-550b-a55b", + sandboxName: "my-assistant", + }, + }), + ); + + session.completeSession(); + + expect(requireLoadedSession(session.loadSession()).stationExpressIntent).toBeNull(); + }); +}); diff --git a/src/lib/state/onboard-session.ts b/src/lib/state/onboard-session.ts index be7589bbdc..a4add0a792 100644 --- a/src/lib/state/onboard-session.ts +++ b/src/lib/state/onboard-session.ts @@ -30,6 +30,10 @@ import { isTerminalOnboardMachineState, } from "../onboard/machine/transitions"; import type { OnboardMachineState, OnboardNonTerminalMachineState } from "../onboard/machine/types"; +import { + parseStationExpressResumeIntent, + type StationExpressResumeIntent, +} from "../onboard/station-express-resume"; import { redactSensitiveText, redactUrl } from "../security/redact"; import { assignSafeToolDisclosureUpdate, @@ -153,6 +157,8 @@ export interface Session { sandboxName: string | null; provider: string | null; model: string | null; + /** Secret-free installer choices needed to retry an interrupted DGX Station Express run. */ + stationExpressIntent: StationExpressResumeIntent | null; endpointUrl: string | null; credentialEnv: string | null; hermesAuthMethod: HermesAuthMethod | null; @@ -631,6 +637,7 @@ export function createSession(overrides: Partial = {}): Session { sandboxName: overrides.sandboxName ?? null, provider: overrides.provider ?? null, model: overrides.model ?? null, + stationExpressIntent: parseStationExpressResumeIntent(overrides.stationExpressIntent), endpointUrl: overrides.endpointUrl ?? null, credentialEnv: overrides.credentialEnv ?? null, hermesAuthMethod: overrides.hermesAuthMethod ?? null, @@ -673,6 +680,13 @@ export function createSession(overrides: Partial = {}): Session { export function normalizeSession(data: Session | SessionJsonValue | undefined): Session | null { if (!isObject(data) || data.version !== SESSION_VERSION) return null; + const stationExpressIntent = parseStationExpressResumeIntent(data.stationExpressIntent); + if ( + hasOwn(data, "stationExpressIntent") && + data.stationExpressIntent !== null && + !stationExpressIntent + ) + return null; const normalized = createSession({ sessionId: readString(data.sessionId) ?? undefined, @@ -683,6 +697,7 @@ export function normalizeSession(data: Session | SessionJsonValue | undefined): sandboxName: readString(data.sandboxName), provider: readString(data.provider), model: readString(data.model), + stationExpressIntent, endpointUrl: typeof data.endpointUrl === "string" ? redactUrl(data.endpointUrl) : null, credentialEnv: readString(data.credentialEnv), hermesAuthMethod: readHermesAuthMethod(data.hermesAuthMethod), @@ -712,6 +727,14 @@ export function normalizeSession(data: Session | SessionJsonValue | undefined): }); normalized.resumable = data.resumable !== false; normalized.status = readString(data.status) ?? normalized.status; + if ( + normalized.stationExpressIntent && + (normalized.mode !== "non-interactive" || + normalized.resumable === false || + normalized.status === "complete") + ) { + return null; + } if (isObject(data.steps)) { for (const [name, step] of Object.entries(data.steps)) { @@ -1504,6 +1527,7 @@ export function completeSession(updates: SessionUpdates = {}): Session { Object.assign(session, safeUpdates); session.status = "complete"; session.resumable = false; + session.stationExpressIntent = null; session.failure = null; transitionMachineSnapshot(session, "complete", now); return session; diff --git a/test/install-express-prompt.test.ts b/test/install-express-prompt.test.ts index 74f3600ac3..4e66742e44 100644 --- a/test/install-express-prompt.test.ts +++ b/test/install-express-prompt.test.ts @@ -46,9 +46,10 @@ ensure_openshell_build_deps() { :; } # Stop immediately after the real express prompt configures the DeepSeek # recipe, before setup-jetson.sh or any installation side effect can run. bash() { - printf "RESULT NON_INTERACTIVE=%s SUDO_MODE=%s PROVIDER=%s MODEL=%s VLLM_MODEL=%s POLICY=%s YES=%s SANDBOX=%s\\n" \ + printf "RESULT NON_INTERACTIVE=%s SUDO_MODE=%s PROVIDER=%s MODEL=%s VLLM_MODEL=%s POLICY=%s YES=%s SANDBOX=%s STATION_EXPRESS=%s\\n" \ "\${NON_INTERACTIVE:-}" "\${NEMOCLAW_NON_INTERACTIVE_SUDO_MODE:-}" "\${NEMOCLAW_PROVIDER:-}" "\${NEMOCLAW_MODEL:-}" \ - "\${NEMOCLAW_VLLM_MODEL:-}" "\${NEMOCLAW_POLICY_MODE:-}" "\${NEMOCLAW_YES:-}" "\${NEMOCLAW_SANDBOX_NAME:-}" + "\${NEMOCLAW_VLLM_MODEL:-}" "\${NEMOCLAW_POLICY_MODE:-}" "\${NEMOCLAW_YES:-}" "\${NEMOCLAW_SANDBOX_NAME:-}" \ + "\${NEMOCLAW_STATION_EXPRESS:-}" exit 0 } main "$@" @@ -61,9 +62,10 @@ NON_INTERACTIVE="\${NON_INTERACTIVE:-}" NEMOCLAW_PROVIDER="\${NEMOCLAW_PROVIDER:-}" NEMOCLAW_NO_EXPRESS="\${NEMOCLAW_NO_EXPRESS:-}" maybe_offer_express_install -printf "RESULT NON_INTERACTIVE=%s SUDO_MODE=%s PROVIDER=%s MODEL=%s VLLM_MODEL=%s POLICY=%s YES=%s SANDBOX=%s\\n" \\ +printf "RESULT NON_INTERACTIVE=%s SUDO_MODE=%s PROVIDER=%s MODEL=%s VLLM_MODEL=%s POLICY=%s YES=%s SANDBOX=%s STATION_EXPRESS=%s\\n" \\ "\${NON_INTERACTIVE:-}" "\${NEMOCLAW_NON_INTERACTIVE_SUDO_MODE:-}" "\${NEMOCLAW_PROVIDER:-}" "\${NEMOCLAW_MODEL:-}" \\ - "\${NEMOCLAW_VLLM_MODEL:-}" "\${NEMOCLAW_POLICY_MODE:-}" "\${NEMOCLAW_YES:-}" "\${NEMOCLAW_SANDBOX_NAME:-}" + "\${NEMOCLAW_VLLM_MODEL:-}" "\${NEMOCLAW_POLICY_MODE:-}" "\${NEMOCLAW_YES:-}" "\${NEMOCLAW_SANDBOX_NAME:-}" \\ + "\${NEMOCLAW_STATION_EXPRESS:-}" ''' env = dict(os.environ) env["INSTALLER_UNDER_TEST"] = installer @@ -211,6 +213,7 @@ detect_express_platform expect(output).toMatch( /RESULT NON_INTERACTIVE=1 SUDO_MODE=prompt PROVIDER=install-vllm MODEL= VLLM_MODEL= POLICY=suggested YES=1 SANDBOX=my-assistant/, ); + expect(output).toMatch(/STATION_EXPRESS=\s/); }); it("preserves a preset Spark vLLM model in the prompt and exported env", () => { @@ -261,6 +264,7 @@ detect_express_platform expect(output).toMatch( /RESULT NON_INTERACTIVE=1 SUDO_MODE=prompt PROVIDER=install-vllm MODEL=nvidia\/nemotron-3-ultra-550b-a55b VLLM_MODEL=nemotron-3-ultra-550b-a55b POLICY=suggested YES=1 SANDBOX=my-assistant/, ); + expect(output).toMatch(/STATION_EXPRESS=1/); }); it("normalizes the canonical Ultra served alias to the registered model slug", () => { diff --git a/test/install-station-host-preparation.test.ts b/test/install-station-host-preparation.test.ts index 67228d531e..e6aec2ed66 100644 --- a/test/install-station-host-preparation.test.ts +++ b/test/install-station-host-preparation.test.ts @@ -1156,8 +1156,8 @@ NON_INTERACTIVE='' NEMOCLAW_PROVIDER='' NEMOCLAW_NO_EXPRESS='' maybe_offer_express_install -printf 'RESULT PLATFORM=%s PROVIDER=%s MODEL=%s VLLM_MODEL=%s\n' \ - "$_SELECTED_EXPRESS_PLATFORM" "$NEMOCLAW_PROVIDER" "\${NEMOCLAW_MODEL:-}" "$NEMOCLAW_VLLM_MODEL" +printf 'RESULT PLATFORM=%s PROVIDER=%s MODEL=%s VLLM_MODEL=%s STATION_EXPRESS=%s\n' \ + "$_SELECTED_EXPRESS_PLATFORM" "$NEMOCLAW_PROVIDER" "\${NEMOCLAW_MODEL:-}" "$NEMOCLAW_VLLM_MODEL" "$NEMOCLAW_STATION_EXPRESS" `, ], { @@ -1178,7 +1178,7 @@ printf 'RESULT PLATFORM=%s PROVIDER=%s MODEL=%s VLLM_MODEL=%s\n' \ expect(output).toMatch(/Resuming the accepted express install/); expect(output).not.toMatch(/Run express install with these settings/); expect(output).toMatch( - /RESULT PLATFORM=DGX Station PROVIDER=install-vllm MODEL=nvidia\/nemotron-3-ultra-550b-a55b VLLM_MODEL=nemotron-3-ultra-550b-a55b/, + /RESULT PLATFORM=DGX Station PROVIDER=install-vllm MODEL=nvidia\/nemotron-3-ultra-550b-a55b VLLM_MODEL=nemotron-3-ultra-550b-a55b STATION_EXPRESS=1/, ); });