diff --git a/src/lib/onboard.ts b/src/lib/onboard.ts index 3b08dff5f7..7281ed598f 100644 --- a/src/lib/onboard.ts +++ b/src/lib/onboard.ts @@ -6411,7 +6411,7 @@ async function onboard(opts: OnboardOptions = {}): Promise { }, }, }); - session = (await onboardRuntimeBoundary.recordStateResultsWithStepCompatibility([...providerInferenceResult.retryStateResults, providerInferenceResult.stateResult]), providerInferenceResult.session); + await onboardRuntimeBoundary.recordStateResultsWithStepCompatibility(providerInferenceResult.stateResults); sandboxName = providerInferenceResult.sandboxName; const { model, @@ -6429,7 +6429,7 @@ async function onboard(opts: OnboardOptions = {}): Promise { resume, fresh, resumeAgentChanged, - session, + session: providerInferenceResult.session, sandboxName, model, provider, diff --git a/src/lib/onboard/machine/handlers/provider-inference.test.ts b/src/lib/onboard/machine/handlers/provider-inference.test.ts index ff16550458..a2ebd78cb7 100644 --- a/src/lib/onboard/machine/handlers/provider-inference.test.ts +++ b/src/lib/onboard/machine/handlers/provider-inference.test.ts @@ -172,6 +172,16 @@ describe("handleProviderInferenceState", () => { metadata: { state: "inference", provider: "nvidia-prod", model: "nvidia/test" }, }); expect(result.retryStateResults).toEqual([]); + expect(result.stateResults).toEqual([ + { + type: "transition", + next: "inference", + transitionKind: "advance", + updates: undefined, + metadata: { state: "provider_selection", provider: "nvidia-prod", model: "nvidia/test" }, + }, + result.stateResult, + ]); }); it("clears non-NVIDIA provider credentials when inference setup fails", async () => { @@ -435,6 +445,12 @@ describe("handleProviderInferenceState", () => { }, ]); expect(result.stateResult).toMatchObject({ next: "sandbox", transitionKind: "advance" }); + expect(result.stateResults.map((stateResult) => [stateResult.next, stateResult.transitionKind])).toEqual([ + ["inference", "advance"], + ["provider_selection", "retry"], + ["inference", "advance"], + ["sandbox", "advance"], + ]); }); it("aborts before inference setup when the configuration summary is rejected", async () => { diff --git a/src/lib/onboard/machine/handlers/provider-inference.ts b/src/lib/onboard/machine/handlers/provider-inference.ts index fe5177cc42..52e930cdf3 100644 --- a/src/lib/onboard/machine/handlers/provider-inference.ts +++ b/src/lib/onboard/machine/handlers/provider-inference.ts @@ -4,7 +4,7 @@ import type { WebSearchConfig } from "../../../inference/web-search"; import type { Session, SessionUpdates } from "../../../state/onboard-session"; import { withInferenceTrace, withProviderSelectionTrace } from "../../tracing"; -import { advanceTo, retryTo, type OnboardStateTransitionResult } from "../result"; +import { advanceTo, type OnboardStateTransitionResult, retryTo } from "../result"; export type ProviderInferenceRetry = { retry: "selection" } | { ok: true; retry?: undefined }; @@ -127,6 +127,7 @@ export interface ProviderInferenceStateResult { webSearchConfig: WebSearchConfig | null; session: Session | null; stateResult: OnboardStateTransitionResult; + stateResults: OnboardStateTransitionResult[]; retryStateResults: OnboardStateTransitionResult[]; } @@ -177,6 +178,7 @@ export async function handleProviderInferenceState({ const webSearchConfig = initial.webSearchConfig; let forceProviderSelection = initialForceProviderSelection; let allowToolsIncompatible = false; + const stateResults: OnboardStateTransitionResult[] = []; const retryStateResults: OnboardStateTransitionResult[] = []; while (true) { @@ -261,6 +263,11 @@ export async function handleProviderInferenceState({ }), ); } + stateResults.push( + advanceTo("inference", { + metadata: { state: "provider_selection", provider, model }, + }), + ); env.NEMOCLAW_OPENSHELL_BIN = deps.getOpenshellBinary(); const needsBedrockRuntimeAdapter = deps.needsBedrockRuntimeAdapter(provider, endpointUrl); const resumeInference = @@ -297,11 +304,11 @@ export async function handleProviderInferenceState({ clearStagedCredentialEnv(deps, credentialEnv); } if (inferenceResult?.retry === "selection") { - retryStateResults.push( - retryTo("provider_selection", { - metadata: { state: "inference", provider, model, reason: "selection_retry" }, - }), - ); + const retryStateResult = retryTo("provider_selection", { + metadata: { state: "inference", provider, model, reason: "selection_retry" }, + }); + retryStateResults.push(retryStateResult); + stateResults.push(retryStateResult); forceProviderSelection = true; continue; } @@ -395,11 +402,11 @@ export async function handleProviderInferenceState({ clearStagedCredentialEnv(deps, credentialEnv); } if (inferenceResult?.retry === "selection") { - retryStateResults.push( - retryTo("provider_selection", { - metadata: { state: "inference", provider, model, reason: "selection_retry" }, - }), - ); + const retryStateResult = retryTo("provider_selection", { + metadata: { state: "inference", provider, model, reason: "selection_retry" }, + }); + retryStateResults.push(retryStateResult); + stateResults.push(retryStateResult); forceProviderSelection = true; continue; } @@ -411,6 +418,11 @@ export async function handleProviderInferenceState({ break; } + const stateResult = advanceTo("sandbox", { + metadata: { state: "inference", provider, model }, + }); + stateResults.push(stateResult); + return { sandboxName, model, @@ -423,9 +435,8 @@ export async function handleProviderInferenceState({ nimContainer, webSearchConfig, session, - stateResult: advanceTo("sandbox", { - metadata: { state: "inference", provider, model }, - }), + stateResult, + stateResults, retryStateResults, }; }