Skip to content
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 3 additions & 9 deletions src/lib/onboard.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4417,15 +4417,8 @@ async function runOnboard(opts: OnboardOptions = {}): Promise<void> {
process.exit(1);
}

const coreFlowContext: InitialOnboardFlowContext = {
...initialContext,
session,
sandboxName,
selectedMessagingChannels,
gpu,
sandboxGpuConfig,
gpuPassthrough,
};
// biome-ignore format: keep src/lib/onboard.ts net-neutral for growth guardrail.
const coreFlowContext: InitialOnboardFlowContext = { ...initialContext, session, sandboxName, selectedMessagingChannels, gpu, sandboxGpuConfig, gpuPassthrough };
// biome-ignore format: keep src/lib/onboard.ts net-neutral for growth guardrail.
const runCoreGatewayOpenshell = setupInferenceFactory.createGatewayScopedOpenshellRunner(runOpenshell, GATEWAY_NAME);
const [providerInferencePhase, sandboxPhase] =
Expand All @@ -4442,6 +4435,7 @@ async function runOnboard(opts: OnboardOptions = {}): Promise<void> {
providerDeps: {
checkGatewayRouteCompatibility,
preflightGatewayRouteDiscovery,
hasRegisteredSandbox: (name) => Boolean(registry.getSandbox(name)),
withGatewayRouteMutationLock: gatewayRouteMutationLock.withGatewayRouteMutationLock,
normalizeHermesAuthMethod,
setupNim: (g, s, a, recover, gateway, assertRouteCompatible, canProbeRoute) =>
Expand Down
1 change: 1 addition & 0 deletions src/lib/onboard/machine/core-flow-phases.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,7 @@ function createPhases(
requiredEndpointUrl: null,
requiredInferenceApi: null,
}),
hasRegisteredSandbox: () => false,
withGatewayRouteMutationLock: async <T>(
_gatewayName: string,
operation: () => Promise<T> | T,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,7 @@ function createDeps() {
const deps: Options["deps"] = {
checkGatewayRouteCompatibility: calls.checkGatewayRouteCompatibility,
preflightGatewayRouteDiscovery: calls.preflightGatewayRouteDiscovery,
hasRegisteredSandbox: () => false,
withGatewayRouteMutationLock: async (_gatewayName, operation) => await operation(),
normalizeHermesAuthMethod: () => null,
setupNim: calls.setupNim,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,7 @@ export function createDeps(
deps: {
checkGatewayRouteCompatibility: calls.checkGatewayRouteCompatibility,
preflightGatewayRouteDiscovery: calls.preflightGatewayRouteDiscovery,
hasRegisteredSandbox: () => false,
withGatewayRouteMutationLock: async <T>(
_gatewayName: string,
operation: () => Promise<T> | T,
Expand Down
20 changes: 20 additions & 0 deletions src/lib/onboard/machine/handlers/provider-inference.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -313,6 +313,26 @@ describe("handleProviderInferenceState", () => {
);
});

it("disables recorded provider recovery for a brand-new sandbox identity (#6630)", async () => {
const { deps, calls } = createDeps();

await handleProviderInferenceState({
...baseOptions(deps),
fresh: false,
sandboxName: "dc-after",
});

expect(calls.setupNim).toHaveBeenCalledWith(
{ type: "nvidia" },
"dc-after",
null,
false,
"nemoclaw",
expect.any(Function),
expect.any(Function),
);
});

it("does not use resume shortcuts when fresh is also set", async () => {
const session = createSession({ provider: "ollama-local", model: "llama3.1" });
session.steps.provider_selection.status = "complete";
Expand Down
11 changes: 10 additions & 1 deletion src/lib/onboard/machine/handlers/provider-inference.ts
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import type {
} from "../../../inference/gateway-route-compatibility";
import type { WebSearchConfig } from "../../../inference/web-search";
import type { HermesAuthMethod, Session, SessionUpdates } from "../../../state/onboard-session";
import { shouldRecoverRecordedProvider } from "../../provider-recovery";
import { withInferenceTrace, withProviderSelectionTrace } from "../../tracing";
import { advanceTo, type OnboardStateTransitionResult, retryTo } from "../result";
import {
Expand Down Expand Up @@ -87,6 +88,7 @@ export interface ProviderInferenceStateOptions<Gpu, Agent, Host> {
deps: {
checkGatewayRouteCompatibility: CurrentGatewayRouteCompatibilityCheck;
preflightGatewayRouteDiscovery: CurrentGatewayRouteDiscoveryPreflight;
hasRegisteredSandbox(sandboxName: string): boolean;
withGatewayRouteMutationLock<T>(
gatewayName: string,
operation: () => Promise<T> | T,
Expand Down Expand Up @@ -319,6 +321,13 @@ export async function handleProviderInferenceState<Gpu, Agent, Host>({
let reuseGatewayCredentialWithoutLocalKey = false;
let endpointPinnedAddresses: string[] | undefined;
const effectiveResume = resume && !fresh;
const recoverRecordedProvider = shouldRecoverRecordedProvider({
fresh,
resume,
sandboxName,
hasRegisteredSandbox: Boolean(sandboxName && deps.hasRegisteredSandbox(sandboxName)),
sessionSandboxName: session?.sandboxName ?? null,
});
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Outdated
const stateResults: OnboardStateTransitionResult[] = [];
const retryStateResults: OnboardStateTransitionResult[] = [];

Expand Down Expand Up @@ -447,7 +456,7 @@ export async function handleProviderInferenceState<Gpu, Agent, Host>({
gpu,
sandboxName,
agent,
!fresh,
recoverRecordedProvider,
gatewayName,
(route) => guardProviderInferenceRouteSelection(deps, gatewayName, sandboxName, route),
(provider) =>
Expand Down
82 changes: 81 additions & 1 deletion src/lib/onboard/provider-recovery.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,11 @@ import { afterEach, describe, expect, it, vi } from "vitest";

import * as onboardSession from "../state/onboard-session";
import * as registry from "../state/registry";
import { createProviderRecoveryHelpers, validateLiveGatewayInference } from "./provider-recovery";
import {
createProviderRecoveryHelpers,
shouldRecoverRecordedProvider,
validateLiveGatewayInference,
} from "./provider-recovery";

afterEach(() => {
vi.restoreAllMocks();
Expand Down Expand Up @@ -33,6 +37,82 @@ describe("validateLiveGatewayInference", () => {
});
});

describe("shouldRecoverRecordedProvider", () => {
it.each([
{
label: "rejects gateway recovery for a brand-new sandbox",
fresh: false,
resume: false,
sandboxName: "dc-after",
hasRegisteredSandbox: false,
sessionSandboxName: null,
expected: false,
},
{
label: "allows gateway recovery before an interactive sandbox name is selected",
fresh: false,
resume: false,
sandboxName: null,
hasRegisteredSandbox: false,
sessionSandboxName: null,
expected: true,
},
{
label: "allows gateway recovery while resuming",
fresh: false,
resume: true,
sandboxName: "dc-after",
hasRegisteredSandbox: false,
sessionSandboxName: null,
expected: true,
},
{
label: "allows gateway recovery for a registered sandbox",
fresh: false,
resume: false,
sandboxName: "dc-after",
hasRegisteredSandbox: true,
sessionSandboxName: null,
expected: true,
},
{
label: "allows gateway recovery for a matching session",
fresh: false,
resume: false,
sandboxName: "dc-after",
hasRegisteredSandbox: false,
sessionSandboxName: "dc-after",
expected: true,
},
{
label: "rejects gateway recovery when fresh overrides existing identity",
fresh: true,
resume: true,
sandboxName: "dc-after",
hasRegisteredSandbox: true,
sessionSandboxName: "dc-after",
expected: false,
},
])("$label", ({
fresh,
resume,
sandboxName,
hasRegisteredSandbox,
sessionSandboxName,
expected,
}) => {
expect(
shouldRecoverRecordedProvider({
fresh,
resume,
sandboxName,
hasRegisteredSandbox,
sessionSandboxName,
}),
).toBe(expected);
});
});

describe("provider recovery persisted routing state", () => {
function helpers() {
return createProviderRecoveryHelpers({
Expand Down
18 changes: 18 additions & 0 deletions src/lib/onboard/provider-recovery.ts
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,24 @@ const MAX_LIVE_PROVIDER_LENGTH = 128;
const MAX_LIVE_MODEL_LENGTH = 512;
const SAFE_LIVE_PROVIDER = /^[A-Za-z0-9._:-]+$/;

export function shouldRecoverRecordedProvider(input: {
fresh: boolean;
resume: boolean;
sandboxName: string | null;
hasRegisteredSandbox: boolean;
sessionSandboxName: string | null;
}): boolean {
return (
!input.fresh &&
(input.resume ||
!input.sandboxName ||
Boolean(
input.sandboxName &&
(input.hasRegisteredSandbox || input.sessionSandboxName === input.sandboxName),
))
);
}

export function validateLiveGatewayInference(
value: { provider: string | null; model: string | null } | null,
): { provider: string; model: string } | null {
Expand Down
Loading