Skip to content
Merged
2 changes: 2 additions & 0 deletions src/lib/actions/sandbox/rebuild-gpu-opt-out.ts
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ import type {
PreparedImageRebuildHandoff,
} from "../../onboard/prepared-dcode-rebuild";
import type {
ProviderRecoveryReceipt,
RebuildProviderReconfigureHandoff,
RebuildRouteHandoff,
} from "../../onboard/rebuild-route-handoff";
Expand Down Expand Up @@ -106,6 +107,7 @@ export type RebuildRecreateOnboardOpts = {
preparedDcodeRebuild?: PreparedDcodeRebuildHandoff;
rebuildRegistryInferenceRoute?: RebuildRouteHandoff;
rebuildProviderReconfigure?: RebuildProviderReconfigureHandoff;
providerRecoveryReceipt?: ProviderRecoveryReceipt;
preparedImageRebuild?: PreparedImageRebuildHandoff;
autoYes: boolean;
toolDisclosure: ToolDisclosure;
Expand Down
28 changes: 27 additions & 1 deletion src/lib/actions/sandbox/rebuild-preflight-target-phase.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,13 @@ import { CLI_NAME } from "../../cli/branding";
import type { SandboxMessagingPlan } from "../../messaging";
import { isSandboxBaseImageRefreshRequested } from "../../onboard/base-image-resolution-flow";
import type { DcodeAutoApprovalMode } from "../../onboard/dcode-auto-approval";
import { createRebuildProviderReconfigureHandoff } from "../../onboard/rebuild-route-handoff";
import { randomUUID } from "node:crypto";

import {
createRebuildProviderReconfigureHandoff,
mintProviderRecoveryReceipt,
type RegistryInferenceRoute,
} from "../../onboard/rebuild-route-handoff";
import { readSandboxBaseImageResolutionMetadata } from "../../sandbox-base-image";
import * as registry from "../../state/registry";
import type { ToolDisclosure } from "../../tool-disclosure";
Expand Down Expand Up @@ -36,6 +42,9 @@ import {
stageRebuildHermesDashboardConfig,
} from "./rebuild-target-preflight";

/** Upper bound on how long a minted provider-recovery receipt stays valid. */
const PROVIDER_RECOVERY_RECEIPT_TTL_MS = 60 * 60 * 1000;

export interface RebuildPreparedTarget {
targetConfig: RebuildTargetConfig;
recreateOptions: RebuildRecreateOnboardOpts;
Expand Down Expand Up @@ -150,6 +159,23 @@ export async function prepareRebuildTargetPreflights(args: {
) {
return null;
}
const recoveryRoute: RegistryInferenceRoute = resumeConfig.registryInferenceRoute ?? {
provider: resumeConfig.provider,
model: resumeConfig.model,
endpointUrl: resumeConfig.endpointUrl,
preferredInferenceApi: resumeConfig.preferredInferenceApi ?? "",
source: "registry",
};
recreateOptions.providerRecoveryReceipt = mintProviderRecoveryReceipt(
{
sandboxName,
gatewayName: recreateOptions.targetGatewayName,
provider: resumeConfig.provider,
model: resumeConfig.model,
route: recoveryRoute,
},
{ nonce: randomUUID(), expiresAtMs: Date.now() + PROVIDER_RECOVERY_RECEIPT_TTL_MS },
);
if (!(await ensureRebuildTargetGatewaySelected(sandboxName, sandboxEntry, log, bail)))
return null;
if (!checkRebuildGatewaySchemaPreflight(sandboxName, sandboxEntry, bail)) return null;
Expand Down
51 changes: 48 additions & 3 deletions src/lib/onboard/authoritative-rebuild-target.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,10 @@ import {
rebuildProviderFlowOptions,
resolveAuthoritativeOnboardGatewayBinding,
} from "./authoritative-rebuild-target";
import {
mintProviderRecoveryReceipt,
type ProviderRecoveryReceiptTarget,
} from "./rebuild-route-handoff";

const target = {
sandboxName: "alpha",
Expand Down Expand Up @@ -107,19 +111,19 @@ describe("prepared provider reconfiguration handoff", () => {
};

it("accepts an exact handoff only for a locked authoritative rebuild resume (#6114)", () => {
expect(rebuildProviderFlowOptions(authorizedOptions, providerTarget)).toEqual({
expect(rebuildProviderFlowOptions(authorizedOptions, providerTarget)).toMatchObject({
authoritativeResumeConfig: true,
forceInferenceSetup: true,
});
expect(rebuildProviderFlowOptions({}, providerTarget)).toEqual({
expect(rebuildProviderFlowOptions({}, providerTarget)).toMatchObject({
authoritativeResumeConfig: false,
forceInferenceSetup: false,
});
});

it("authorizes incomplete-session recovery only for the locked rebuild context", () => {
const recoveryOptions = { ...authorizedOptions, rebuildProviderReconfigure: undefined };
expect(rebuildProviderFlowOptions(recoveryOptions, providerTarget)).toEqual({
expect(rebuildProviderFlowOptions(recoveryOptions, providerTarget)).toMatchObject({
authoritativeResumeConfig: true,
forceInferenceSetup: false,
});
Expand All @@ -134,6 +138,47 @@ describe("prepared provider reconfiguration handoff", () => {
}
});

it("activates a matching provider-recovery receipt and binds it to the session", () => {
const receiptTarget: ProviderRecoveryReceiptTarget = {
sandboxName: "alpha",
gatewayName: "nemoclaw-8081",
provider: "compatible-endpoint",
model: "nvidia/model",
route: {
provider: "compatible-endpoint",
model: "nvidia/model",
endpointUrl: "https://inference.example.test/v1",
preferredInferenceApi: "openai-completions",
source: "registry",
},
};
const receipt = mintProviderRecoveryReceipt(receiptTarget, {
nonce: "n-alpha",
expiresAtMs: Number.MAX_SAFE_INTEGER,
});
const flowContext = {
...providerTarget,
preferredInferenceApi: "openai-completions",
session: { sessionId: "sess-alpha" },
};

const activated = rebuildProviderFlowOptions(
{ ...authorizedOptions, providerRecoveryReceipt: receipt },
flowContext,
);
expect(activated.providerRecoveryReceipt?.sessionId).toBe("sess-alpha");

const wrongSandbox = rebuildProviderFlowOptions(
{
...authorizedOptions,
providerRecoveryReceipt: receipt,
rebuildProviderReconfigure: undefined,
},
{ ...flowContext, sandboxName: "beta" },
);
expect(wrongSandbox.providerRecoveryReceipt).toBeNull();
});

it("rejects an unauthorized or mismatched handoff (#6114)", () => {
expect(() =>
rebuildProviderFlowOptions(
Expand Down
45 changes: 42 additions & 3 deletions src/lib/onboard/authoritative-rebuild-target.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,11 @@ import { findDashboardForwardOwner } from "./dashboard-port";
import { resolveGatewayName } from "./gateway-binding";
import type { PortProbeResult } from "./preflight";
import { assertDashboardPortNotReserved } from "./preflight-ports";
import { validateRebuildProviderReconfigureHandoff } from "./rebuild-route-handoff";
import {
createProviderRecoveryReceiptLedger,
type ProviderRecoveryReceipt,
validateRebuildProviderReconfigureHandoff,
} from "./rebuild-route-handoff";
import type { OnboardOptions } from "./types";

export type AuthoritativeOnboardGatewayBinding = { name: string; port: number };
Expand Down Expand Up @@ -110,9 +114,19 @@ function validateRebuildHandoff(
/** Derive the provider-phase authority from one validated rebuild handoff. */
export function rebuildProviderFlowOptions(
opts: OnboardOptions,
target: Parameters<typeof validateRebuildHandoff>[1],
): { authoritativeResumeConfig: boolean; forceInferenceSetup: boolean } {
target: Parameters<typeof validateRebuildHandoff>[1] & {
session?: { sessionId: string } | null;
preferredInferenceApi?: string | null;
},
): {
authoritativeResumeConfig: boolean;
forceInferenceSetup: boolean;
providerRecoveryReceipt: ProviderRecoveryReceipt | null;
providerRecoveryReceiptLedger: ReturnType<typeof createProviderRecoveryReceiptLedger>;
} {
const authoritativeResumeConfig = opts.authoritativeResumeConfig === true;
const providerRecoveryReceiptLedger = createProviderRecoveryReceiptLedger();
let providerRecoveryReceipt: ProviderRecoveryReceipt | null = null;
if (authoritativeResumeConfig) {
const gateway = resolveAuthoritativeOnboardGatewayBinding(opts);
if (
Expand All @@ -128,10 +142,35 @@ export function rebuildProviderFlowOptions(
"Authoritative provider recovery requires a preflighted locked rebuild resume.",
);
}
const sessionId = target.session?.sessionId ?? null;
if (opts.providerRecoveryReceipt && sessionId) {
providerRecoveryReceipt = providerRecoveryReceiptLedger.activate(
opts.providerRecoveryReceipt,
{
target: {
sandboxName: target.sandboxName,
gatewayName: gateway.name,
provider: target.provider,
model: target.model,
route: {
provider: target.provider,
model: target.model,
endpointUrl: target.endpointUrl ?? null,
preferredInferenceApi: target.preferredInferenceApi ?? "",
source: "registry",
},
},
sessionId,
nowMs: Date.now(),
},
);
}
}
return {
authoritativeResumeConfig,
forceInferenceSetup: validateRebuildHandoff(opts, target),
providerRecoveryReceipt,
providerRecoveryReceiptLedger,
};
}

Expand Down
8 changes: 8 additions & 0 deletions src/lib/onboard/machine/core-flow-phases.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,10 @@

import type { WebSearchConfig } from "../../inference/web-search";
import type { DcodeAutoApprovalMode } from "../dcode-auto-approval";
import type {
createProviderRecoveryReceiptLedger,
ProviderRecoveryReceipt,
} from "../rebuild-route-handoff";
import {
mergeProviderModelSelectedContext,
mergeSandboxCreatedContext,
Expand Down Expand Up @@ -33,6 +37,8 @@ export interface CoreOnboardFlowPhaseOptions<
forceProviderSelection: boolean;
forceInferenceSetup?: boolean;
authoritativeResumeConfig?: boolean;
providerRecoveryReceipt?: ProviderRecoveryReceipt | null;
providerRecoveryReceiptLedger?: ReturnType<typeof createProviderRecoveryReceiptLedger>;
env: NodeJS.ProcessEnv;
constants: ProviderInferenceStateOptions<Context["gpu"], Context["agent"], Host>["constants"];
providerDeps: ProviderInferenceStateOptions<Context["gpu"], Context["agent"], Host>["deps"];
Expand Down Expand Up @@ -75,6 +81,8 @@ export function createCoreOnboardFlowPhases<
forceProviderSelection: options.forceProviderSelection,
forceInferenceSetup: options.forceInferenceSetup,
authoritativeResumeConfig: options.authoritativeResumeConfig,
providerRecoveryReceipt: options.providerRecoveryReceipt,
providerRecoveryReceiptLedger: options.providerRecoveryReceiptLedger,
initial: {
model: context.model,
provider: context.provider,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,12 @@ import { describe, expect, it, vi } from "vitest";

import { createSession } from "../../../state/onboard-session";
import { handleProviderInferenceState } from "./provider-inference";
import { baseOptions, baseSelection, createDeps } from "./provider-inference.test-support";
import {
activatedRecoveryReceipt,
baseOptions,
baseSelection,
createDeps,
} from "./provider-inference.test-support";

describe("authoritative provider inference recovery", () => {
it("stays enabled across messaging revalidation", async () => {
Expand Down Expand Up @@ -35,10 +40,17 @@ describe("authoritative provider inference recovery", () => {
isInferenceRouteReady: vi.fn(() => true),
});

const { receipt, ledger } = activatedRecoveryReceipt({
sandboxName: "my-assistant",
sessionId: session.sessionId,
});

const result = await handleProviderInferenceState({
...baseOptions(deps, session),
resume: true,
authoritativeResumeConfig: true,
providerRecoveryReceipt: receipt,
providerRecoveryReceiptLedger: ledger,
sandboxName: "my-assistant",
selectedMessagingChannels: ["telegram"],
});
Expand Down
Loading
Loading