diff --git a/src/lib/onboard/model-router-process.ts b/src/lib/onboard/model-router-process.ts index ed3a09aa73..e642474adb 100644 --- a/src/lib/onboard/model-router-process.ts +++ b/src/lib/onboard/model-router-process.ts @@ -41,8 +41,8 @@ const ROUTER_HEALTH_BODY_MAX_BYTES = 64 * 1024; /** * Fetch /health and keep the response body for diagnosis (#8962). Unlike - * `isRouterHealthy`, this waits for the body, so callers pass a longer - * timeout; `startModelRouter` uses 30 seconds. LiteLLM's /health probes + * `isRouterHealthy`, this waits for the body, so a caller that must read it + * budgets for it; the final startup snapshot uses 30 seconds. LiteLLM's /health probes * every upstream endpoint per request and can answer well after the * 3-second liveness budget. The timeout is a wall-clock deadline, not a * socket idle timeout, so a responder that trickles bytes cannot hold the diff --git a/src/lib/onboard/model-router-reconcile.test.ts b/src/lib/onboard/model-router-reconcile.test.ts new file mode 100644 index 0000000000..072b06bc96 --- /dev/null +++ b/src/lib/onboard/model-router-reconcile.test.ts @@ -0,0 +1,89 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { beforeEach, describe, expect, it, vi } from "vitest"; + +import { reconcileModelRouter } from "./model-router"; + +const RECORDED_ROUTER_PID = 4321; + +const holder = vi.hoisted(() => ({ + snapshotBody: null as string | null, + stopped: [] as Array<[number, number]>, + reachabilityProbes: 0, +})); + +// `stopModelRouterProcess` throws a sentinel so each case ends at the +// reuse decision. Restarting the router is `startModelRouter`'s contract and +// is covered by `test/onboard-model-router.test.ts`. +vi.mock("./model-router-process", () => ({ + ROUTER_HEALTH_TIMEOUT_MS: 3_000, + getRouterHealthSnapshot: vi.fn(async () => ({ healthy: true, body: holder.snapshotBody })), + isRouterHealthy: vi.fn(async () => true), + doesModelRouterProcessOwnPort: vi.fn(() => true), + inspectModelRouterProcessForPort: vi.fn(() => ({ status: "missing" as const })), + stopModelRouterProcess: vi.fn(async (pid: number, port: number) => { + holder.stopped.push([pid, port]); + throw new Error("router restart reached"); + }), +})); + +vi.mock("../credentials/store", () => ({ + normalizeCredentialValue: (value: string) => value, + resolveProviderCredential: () => "", + saveCredential: vi.fn(), +})); + +vi.mock("./credential-env", () => ({ + hydrateCredentialEnv: () => "nvapi-TEST-NOT-A-REAL-ROUTER-KEY", +})); + +vi.mock("../state/onboard-session", () => ({ + loadSession: () => ({ + routerPid: RECORDED_ROUTER_PID, + routerCredentialHash: "MATCHING-HASH", + }), + updateSession: vi.fn(), +})); + +vi.mock("../security/credential-hash", () => ({ hashCredential: () => "MATCHING-HASH" })); + +vi.mock("./host-service-reachability", () => ({ + probeHostServiceSandboxReachability: vi.fn(async () => { + holder.reachabilityProbes += 1; + return { ok: true }; + }), + formatHostServiceUnreachableMessage: () => "", +})); + +describe("model router reconciliation", () => { + beforeEach(() => { + holder.snapshotBody = null; + holder.stopped = []; + holder.reachabilityProbes = 0; + }); + + it("reuses a recorded router whose health snapshot names a healthy endpoint", async () => { + holder.snapshotBody = JSON.stringify({ + healthy_endpoints: [{ api_base: "https://integrate.api.nvidia.com/v1" }], + unhealthy_endpoints: [], + }); + + await reconcileModelRouter(); + + expect(holder.stopped).toEqual([]); + expect(holder.reachabilityProbes).toBe(1); + }); + + it("restarts a recorded router that answers 2xx with no healthy endpoint (#9437)", async () => { + holder.snapshotBody = JSON.stringify({ + healthy_endpoints: [], + unhealthy_endpoints: [{ api_base: "https://integrate.api.nvidia.com/v1" }], + }); + + await expect(reconcileModelRouter()).rejects.toThrow("router restart reached"); + + expect(holder.stopped).toEqual([[RECORDED_ROUTER_PID, expect.any(Number)]]); + expect(holder.reachabilityProbes).toBe(0); + }); +}); diff --git a/src/lib/onboard/model-router.ts b/src/lib/onboard/model-router.ts index ac6249fed9..d8c5985d7c 100644 --- a/src/lib/onboard/model-router.ts +++ b/src/lib/onboard/model-router.ts @@ -61,9 +61,9 @@ const ROUTER_HEALTH_INTERVAL_MS = 2000; const ROUTER_STARTUP_TIMEOUT_MS = 10 * 60_000; // LiteLLM's /health live-probes every upstream endpoint per request, so it // can need far longer than the 3-second liveness budget to answer (#8962). -// The startup poll keeps the 3-second budget: the status-only liveness -// probe must not accept a fast 200 that names zero healthy endpoints, so -// recovery for a slow-but-healthy router runs through the body-checked +// The startup poll keeps the 3-second budget and reads the body, so it +// never accepts a fast 200 that names zero healthy endpoints; recovery for +// a router whose /health outruns that budget runs through the body-checked // final snapshot after the poll exhausts its retries. const ROUTER_FINAL_HEALTH_SNAPSHOT_TIMEOUT_MS = 30_000; const ROUTER_LOG_TAIL_LINES = 20; @@ -468,7 +468,8 @@ export async function startModelRouter( Math.min(ROUTER_HEALTH_REQUEST_TIMEOUT_MS, Math.ceil(remainingMs)), ); healthAttempts += 1; - const healthy = await deps.isRouterHealthy(port, healthTimeoutMs); + const pollSnapshot = await deps.getRouterHealthSnapshot(port, healthTimeoutMs); + const healthy = isRouterSnapshotReady(pollSnapshot); const processAlive = deps.isProcessAlive(pid); if (healthy && processAlive) return pid; if (!processAlive) { @@ -484,7 +485,7 @@ export async function startModelRouter( const finalSnapshot: RouterHealthSnapshot = childExited ? { healthy: false, body: null } : await deps.getRouterHealthSnapshot(port, ROUTER_FINAL_HEALTH_SNAPSHOT_TIMEOUT_MS); - if (finalSnapshot.healthy && hasHealthyEndpoint(finalSnapshot.body) && deps.isProcessAlive(pid)) { + if (isRouterSnapshotReady(finalSnapshot) && deps.isProcessAlive(pid)) { return pid; } try { @@ -503,11 +504,11 @@ export async function startModelRouter( ); } -/** True when the parsed /health body names at least one healthy endpoint. */ -function hasHealthyEndpoint(body: string | null): boolean { - if (!body) return false; +/** Router readiness: /health answered 2xx and names at least one healthy endpoint. */ +function isRouterSnapshotReady(snapshot: RouterHealthSnapshot): boolean { + if (!snapshot.healthy || !snapshot.body) return false; try { - const parsed = JSON.parse(body) as { healthy_endpoints?: readonly unknown[] }; + const parsed = JSON.parse(snapshot.body) as { healthy_endpoints?: readonly unknown[] }; return Array.isArray(parsed?.healthy_endpoints) && parsed.healthy_endpoints.length > 0; } catch { return false; @@ -595,7 +596,7 @@ const MODEL_ROUTER_SERVICE_LABEL = "Model Router"; /** * Verify the host Model Router is reachable from the OpenShell Docker network. * - * `isRouterHealthy()` only proves the router answers on the host loopback. On + * A healthy /health answer only proves the router responds on the host loopback. On * Linux Docker-driver hosts with UFW default-deny, a sandbox container can * still fail to reach `host.openshell.internal:` even though the * host curl succeeds (#4564). This mirrors the Ollama auth-proxy probe: on a @@ -636,19 +637,29 @@ export async function reconcileModelRouter(): Promise { const recordedPid = session?.routerPid ?? null; const recordedCredentialHash = session?.routerCredentialHash ?? null; - if (await isRouterHealthy(routerPort)) { + // One snapshot answers both questions: `healthy` is the occupied-port check, + // and `isRouterSnapshotReady` is the single authority for declaring the + // router usable, exactly as it is for the startup poll. Budget the body read, + // because /health probes every upstream endpoint and can answer well after + // the 3-second liveness budget. + const snapshot = await getRouterHealthSnapshot( + routerPort, + ROUTER_FINAL_HEALTH_SNAPSHOT_TIMEOUT_MS, + ); + if (snapshot.healthy) { const recordedProcessOwnsRouter = doesModelRouterProcessOwnPort(recordedPid, routerPort); if ( routerCredentialHash && recordedCredentialHash === routerCredentialHash && - recordedProcessOwnsRouter + recordedProcessOwnsRouter && + isRouterSnapshotReady(snapshot) ) { console.log(` ✓ Model router is already healthy on port ${routerPort}`); await verifyModelRouterSandboxReachability(routerPort); return; } if (recordedProcessOwnsRouter) { - console.log(" Restarting model router with updated credentials..."); + console.log(" Restarting model router..."); await stopModelRouterProcess( requireValue(recordedPid, "Expected recorded router PID"), routerPort, diff --git a/test/onboard-model-router.test.ts b/test/onboard-model-router.test.ts index 8c03f0b996..72c4f7e73f 100644 --- a/test/onboard-model-router.test.ts +++ b/test/onboard-model-router.test.ts @@ -43,6 +43,10 @@ const MODEL_ROUTER_FINGERPRINT_FILE = ".nemoclaw-source-fingerprint"; const MODEL_ROUTER_TEST_SOURCE_SHA = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"; const MODEL_ROUTER_TEST_VERSION = "0.1.0"; const NVIDIA_TEST_CREDENTIAL = "nvapi-TEST-NOT-A-REAL-ROUTER-KEY"; +const ROUTER_HEALTHY_BODY = JSON.stringify({ + healthy_endpoints: [{ api_base: "https://integrate.api.nvidia.com/v1" }], + unhealthy_endpoints: [], +}); type PrepareCall = { venvDir: string; @@ -295,7 +299,6 @@ describe("onboard Model Router setup", () => { const port = 45_678; const healthChecks: number[] = []; const sleepCalls: number[] = []; - let healthProbe = 0; let pid: number | null = null; const blueprintDir = path.join(rootDir, "nemoclaw-blueprint"); @@ -350,8 +353,11 @@ describe("onboard Model Router setup", () => { name === "ROUTER_API_KEY" ? "router-secret" : null, isRouterHealthy: async (routerPort) => { healthChecks.push(routerPort); - healthProbe += 1; - return healthProbe > 1; + return false; + }, + getRouterHealthSnapshot: async (routerPort) => { + healthChecks.push(routerPort); + return { healthy: true, body: ROUTER_HEALTHY_BODY }; }, sleep: async (milliseconds) => { sleepCalls.push(milliseconds); @@ -406,7 +412,6 @@ describe("onboard Model Router setup", () => { const homeDir = path.join(tmpDir, "home"); const routerCommand = path.join(tmpDir, "managed", "model-router"); const port = 45_692; - let healthProbe = 0; let pid: number | null = null; fs.mkdirSync(path.join(rootDir, "nemoclaw-blueprint", "router"), { recursive: true }); @@ -433,10 +438,8 @@ describe("onboard Model Router setup", () => { homeDir, ensureModelRouterCommand: () => routerCommand, resolveProviderCredential: () => null, - isRouterHealthy: async () => { - healthProbe += 1; - return healthProbe > 1; - }, + isRouterHealthy: async () => false, + getRouterHealthSnapshot: async () => ({ healthy: true, body: ROUTER_HEALTHY_BODY }), sleep: async () => undefined, }, ); @@ -475,9 +478,10 @@ describe("onboard Model Router setup", () => { }), resolveProviderCredential: () => null, buildSubprocessEnv: () => ({}), - isRouterHealthy: async () => { + isRouterHealthy: async () => false, + getRouterHealthSnapshot: async () => { healthProbe += 1; - return healthProbe > 61; + return { healthy: healthProbe > 60, body: ROUTER_HEALTHY_BODY }; }, sleep, isProcessAlive: () => true, @@ -487,7 +491,7 @@ describe("onboard Model Router setup", () => { ); assert.equal(startedPid, pid); - assert.equal(healthProbe, 62); + assert.equal(healthProbe, 61); assert.equal(sleep.mock.calls.length, 61); assert.equal(terminateProcess.mock.calls.length, 0); }); @@ -496,7 +500,7 @@ describe("onboard Model Router setup", () => { const pid = 12_345; const sleep = vi.fn(async () => undefined); const terminateProcess = vi.fn(); - const isRouterHealthy = vi.fn(async () => false); + const getRouterHealthSnapshot = vi.fn(async () => ({ healthy: false, body: null })); await assert.rejects( startModelRouter( @@ -515,8 +519,8 @@ describe("onboard Model Router setup", () => { }), resolveProviderCredential: () => null, buildSubprocessEnv: () => ({}), - isRouterHealthy, - getRouterHealthSnapshot: async () => ({ healthy: false, body: null }), + isRouterHealthy: async () => false, + getRouterHealthSnapshot, sleep, isProcessAlive: () => true, terminateProcess, @@ -526,7 +530,7 @@ describe("onboard Model Router setup", () => { /failed to become healthy on port 45680 within 600 seconds \(completed health checks: 300\)/, ); - assert.equal(isRouterHealthy.mock.calls.length, 301); + assert.equal(getRouterHealthSnapshot.mock.calls.length, 301); assert.equal(sleep.mock.calls.length, 300); assert.deepEqual(terminateProcess.mock.calls, [[pid]]); }); @@ -534,7 +538,6 @@ describe("onboard Model Router setup", () => { it("sets OPENAI_API_KEY to the routed credential when an ambient OPENAI_API_KEY exists (#8962)", async () => { const pid = 12_345; let spawnedEnv: Record | null = null; - let healthProbe = 0; await startModelRouter( { port: 45_690, pool_config_path: "router/test-pool.yaml", credential_env: "ROUTER_API_KEY" }, @@ -558,10 +561,8 @@ describe("onboard Model Router setup", () => { resolveProviderCredential: (name) => ({ ROUTER_API_KEY: "router-secret", OPENAI_API_KEY: "stale-openai" })[name] ?? null, buildSubprocessEnv: (extra) => ({ ...extra }), - isRouterHealthy: async () => { - healthProbe += 1; - return healthProbe > 1; - }, + isRouterHealthy: async () => false, + getRouterHealthSnapshot: async () => ({ healthy: true, body: ROUTER_HEALTHY_BODY }), sleep: async () => undefined, isProcessAlive: () => true, terminateProcess: () => undefined, @@ -643,10 +644,6 @@ describe("onboard Model Router setup", () => { it("returns the router PID when the final health snapshot proves recovery (#8962)", async () => { const pid = 12_345; const terminateProcess = vi.fn(); - const healthyBody = JSON.stringify({ - healthy_endpoints: [{ api_base: "https://integrate.api.nvidia.com/v1" }], - unhealthy_endpoints: [], - }); const startedPid = await startModelRouter( { port: 45_693, pool_config_path: "router/test-pool.yaml" }, @@ -665,7 +662,12 @@ describe("onboard Model Router setup", () => { resolveProviderCredential: () => null, buildSubprocessEnv: () => ({}), isRouterHealthy: async () => false, - getRouterHealthSnapshot: async () => ({ healthy: true, body: healthyBody }), + // /health outruns the poll's 3-second budget and answers only within + // the 30-second final-snapshot budget. + getRouterHealthSnapshot: async (_port: number, timeoutMs = 0) => ({ + healthy: timeoutMs >= 30_000, + body: timeoutMs >= 30_000 ? ROUTER_HEALTHY_BODY : null, + }), sleep: async () => undefined, isProcessAlive: () => true, terminateProcess, @@ -722,7 +724,7 @@ describe("onboard Model Router setup", () => { ); }); - it("still fails when the final snapshot is 2xx with zero healthy endpoints (#8962)", async () => { + it("still fails when the poll and final snapshot are 2xx with zero healthy endpoints (#8962)", async () => { const pid = 12_345; const terminateProcess = vi.fn(); const allUnhealthyBody = JSON.stringify({ @@ -747,7 +749,10 @@ describe("onboard Model Router setup", () => { }), resolveProviderCredential: () => null, buildSubprocessEnv: () => ({}), - isRouterHealthy: async () => false, + // The pre-spawn port guard calls isRouterHealthy without a timeout. + // Return true for timeout-bearing calls so a regression to the old + // boolean startup poll cannot accept zero healthy endpoints. + isRouterHealthy: async (_port: number, timeoutMs) => timeoutMs !== undefined, getRouterHealthSnapshot: async () => ({ healthy: true, body: allUnhealthyBody }), sleep: async () => undefined, isProcessAlive: () => true, @@ -827,7 +832,6 @@ describe("onboard Model Router setup", () => { const pid = 12_345; let spawnedEnv: Record | null = null; - let healthProbe = 0; await startModelRouter( { port: 45_695, @@ -853,10 +857,8 @@ describe("onboard Model Router setup", () => { resolveProviderCredential: (name) => ({ ROUTER_API_KEY: "router-secret", OPENAI_API_KEY: "operator-openai" })[name] ?? null, buildSubprocessEnv: (extra) => ({ ...extra }), - isRouterHealthy: async () => { - healthProbe += 1; - return healthProbe > 1; - }, + isRouterHealthy: async () => false, + getRouterHealthSnapshot: async () => ({ healthy: true, body: ROUTER_HEALTHY_BODY }), sleep: async () => undefined, isProcessAlive: () => true, terminateProcess: () => undefined, @@ -874,7 +876,6 @@ describe("onboard Model Router setup", () => { it("preserves routed credential fallback for an unproven pool (#8962)", async () => { const pid = 12_345; let spawnedEnv: Record | null = null; - let healthProbe = 0; await startModelRouter( { @@ -900,10 +901,8 @@ describe("onboard Model Router setup", () => { readPoolConfig: () => 'models:\n - litellm_model: "openai/gpt-test"\n', resolveProviderCredential: (name) => (name === "ROUTER_API_KEY" ? "router-secret" : null), buildSubprocessEnv: (extra) => ({ ...extra }), - isRouterHealthy: async () => { - healthProbe += 1; - return healthProbe > 1; - }, + isRouterHealthy: async () => false, + getRouterHealthSnapshot: async () => ({ healthy: true, body: ROUTER_HEALTHY_BODY }), sleep: async () => undefined, isProcessAlive: () => true, terminateProcess: () => undefined, @@ -924,9 +923,9 @@ describe("onboard Model Router setup", () => { const sleep = vi.fn(async (milliseconds: number) => { nowMs += milliseconds; }); - const isRouterHealthy = vi.fn(async (_port: number, timeoutMs = 0) => { + const getRouterHealthSnapshot = vi.fn(async (_port: number, timeoutMs = 0) => { nowMs += timeoutMs; - return false; + return { healthy: false, body: null }; }); await assert.rejects( @@ -946,11 +945,8 @@ describe("onboard Model Router setup", () => { }), resolveProviderCredential: () => null, buildSubprocessEnv: () => ({}), - isRouterHealthy, - getRouterHealthSnapshot: async (_port: number, timeoutMs = 0) => { - nowMs += timeoutMs; - return { healthy: false, body: null }; - }, + isRouterHealthy: async () => false, + getRouterHealthSnapshot, sleep, now: () => nowMs, isProcessAlive: () => true, @@ -965,7 +961,7 @@ describe("onboard Model Router setup", () => { ); assert.equal(nowMs, 600_000); - assert.equal(isRouterHealthy.mock.calls.length, 115); + assert.equal(getRouterHealthSnapshot.mock.calls.length, 115); assert.equal(sleep.mock.calls.length, 114); assert.deepEqual(terminateProcess.mock.calls, [[pid]]); }); @@ -980,7 +976,6 @@ describe("onboard Model Router setup", () => { const mkdirSync = vi.fn(); const proxyConfigArgs: string[][] = []; const proxyArgs: string[][] = []; - let healthProbe = 0; vi.stubEnv("HOME", homeDir); vi.stubEnv("NEMOCLAW_GATEWAY_PORT", "9123"); vi.resetModules(); @@ -1008,10 +1003,8 @@ describe("onboard Model Router setup", () => { }, resolveProviderCredential: () => null, buildSubprocessEnv: () => ({}), - isRouterHealthy: async () => { - healthProbe += 1; - return healthProbe > 1; - }, + isRouterHealthy: async () => false, + getRouterHealthSnapshot: async () => ({ healthy: true, body: ROUTER_HEALTHY_BODY }), sleep: async () => undefined, isProcessAlive: () => true, terminateProcess: () => undefined,