diff --git a/src/system-agent/setup-inference-detection.test.ts b/src/system-agent/setup-inference-detection.test.ts index 71e0b2853fb8..4e1858cb083e 100644 --- a/src/system-agent/setup-inference-detection.test.ts +++ b/src/system-agent/setup-inference-detection.test.ts @@ -1,5 +1,6 @@ import { createServer, get } from "node:http"; import type { AddressInfo } from "node:net"; +import { Worker } from "node:worker_threads"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { DEFAULT_AGENT_WORKSPACE_DIR } from "../agents/workspace-default.js"; import { listRecommendedToolInstalls } from "../plugins/recommended-tool-installs.js"; @@ -8,6 +9,9 @@ import type { SetupInferenceDetection } from "./setup-inference.js"; const blockingWorkerUrl = new URL( `data:text/javascript,${encodeURIComponent(` import { parentPort, workerData } from "node:worker_threads"; + if (workerData.started) { + Atomics.add(new Int32Array(workerData.started), 0, 1); + } parentPort.postMessage({ type: "partial", detection: workerData.partialDetection }); const deadline = Date.now() + workerData.blockMs; while (Date.now() < deadline) {} @@ -36,6 +40,22 @@ function emptyDetection(): SetupInferenceDetection { }; } +function detectedCodex(): SetupInferenceDetection { + return { + ...emptyDetection(), + candidates: [ + { + kind: "codex-cli", + modelRef: "openai/gpt-5.5", + label: "Codex", + detail: "logged in", + credentials: true, + recommended: false, + }, + ], + }; +} + const servers = new Set>(); beforeEach(() => { @@ -61,6 +81,7 @@ async function requestHealth(url: string): Promise<{ body: string; statusCode: n } afterEach(async () => { + vi.restoreAllMocks(); await Promise.all( [...servers].map( (server) => @@ -86,17 +107,19 @@ describe("isolated setup inference detection", () => { }); }); const address = server.address() as AddressInfo; - const fallback = emptyDetection(); + const fallback = detectedCodex(); const pendingStartedAt = performance.now(); const pending = detectSetupInferenceIsolated({ workerUrl: blockingWorkerUrl, workerData: { - blockMs: 10_000, + blockMs: 30_000, detection: emptyDetection(), partialDetection: fallback, }, - timeoutMs: 100, + // Generous timeout: the partial must arrive before the deadline even on a + // loaded CI runner, or the empty-timeout path rejects and this flakes. + timeoutMs: 3_000, fallbackEnv: {}, }); const startedAt = performance.now(); @@ -117,10 +140,47 @@ describe("isolated setup inference detection", () => { setupComplete: fallback.setupComplete, }); expect(detection.prepareOptions ?? []).toEqual([]); - expect(performance.now() - pendingStartedAt).toBeLessThan(1_000); + expect(performance.now() - pendingStartedAt).toBeLessThan(10_000); }); - it("preserves ambient API keys when detection times out", async () => { + it("rejects an empty timeout with an actionable typed error", async () => { + const { detectSetupInferenceIsolated } = await loadDetectionModule(); + + await expect( + detectSetupInferenceIsolated({ + workerUrl: silentBlockingWorkerUrl, + timeoutMs: 50, + fallbackEnv: {}, + }), + ).rejects.toMatchObject({ + name: "SetupInferenceDetectionTimeoutError", + message: + "Checking this Gateway for AI access timed out after 0.05s. " + + "The Gateway may be busy — try again.", + }); + }); + + it("returns partial candidates when detection times out", async () => { + const { detectSetupInferenceIsolated } = await loadDetectionModule(); + const partial = detectedCodex(); + + const detection = await detectSetupInferenceIsolated({ + workerUrl: blockingWorkerUrl, + workerData: { + blockMs: 30_000, + detection: emptyDetection(), + partialDetection: partial, + }, + // Generous timeout: the partial must beat the deadline on loaded CI + // runners, or the empty-timeout rejection makes this flake. + timeoutMs: 3_000, + fallbackEnv: {}, + }); + + expect(detection).toEqual(partial); + }); + + it("returns ambient API keys when detection times out", async () => { const { detectSetupInferenceIsolated } = await loadDetectionModule(); const detection = await detectSetupInferenceIsolated({ @@ -159,30 +219,85 @@ describe("isolated setup inference detection", () => { ]); }); - it("omits prepare choices when detection times out without a partial result", async () => { + it("waits for timed-out worker shutdown before running a fresh detection", async () => { const { detectSetupInferenceIsolated } = await loadDetectionModule(); + let releaseShutdown: (() => void) | undefined; + const shutdownGate = new Promise((resolve) => { + releaseShutdown = resolve; + }); + const terminateSpy = vi.spyOn(Worker.prototype, "terminate"); + terminateSpy.mockImplementationOnce(function (this: Worker) { + terminateSpy.mockRestore(); + return this.terminate().then(async (code) => { + await shutdownGate; + return code; + }); + }); + + await expect( + detectSetupInferenceIsolated({ + workerUrl: silentBlockingWorkerUrl, + timeoutMs: 50, + fallbackEnv: {}, + }), + ).rejects.toThrow("Checking this Gateway for AI access timed out"); + + let retrySettled = false; + const fresh = detectedCodex(); + const retry = detectSetupInferenceIsolated({ + workerUrl: blockingWorkerUrl, + workerData: { + blockMs: 0, + detection: fresh, + partialDetection: emptyDetection(), + }, + timeoutMs: 5_000, + fallbackEnv: {}, + }).then((detection) => { + retrySettled = true; + return detection; + }); + await new Promise((resolve) => { + setTimeout(resolve, 20); + }); + + expect(retrySettled).toBe(false); + releaseShutdown?.(); + await expect(retry).resolves.toEqual(fresh); + }); + + it("returns successful worker results unchanged", async () => { + const { detectSetupInferenceIsolated } = await loadDetectionModule(); + const detected = detectedCodex(); const detection = await detectSetupInferenceIsolated({ - workerUrl: silentBlockingWorkerUrl, - timeoutMs: 50, + workerUrl: blockingWorkerUrl, + workerData: { + blockMs: 0, + detection: detected, + partialDetection: emptyDetection(), + }, + timeoutMs: 5_000, fallbackEnv: {}, }); - expect(detection.prepareOptions).toBeUndefined(); + expect(detection).toEqual(detected); }); it("coalesces concurrent detections behind one bounded worker", async () => { const { detectSetupInferenceIsolated } = await loadDetectionModule(); - const fallback = vi.fn(async () => emptyDetection()); + const started = new SharedArrayBuffer(Int32Array.BYTES_PER_ELEMENT); + const detected = detectedCodex(); const options = { workerUrl: blockingWorkerUrl, workerData: { - blockMs: 10_000, - detection: emptyDetection(), + blockMs: 100, + detection: detected, partialDetection: emptyDetection(), + started, }, - timeoutMs: 50, - fallback, + timeoutMs: 5_000, + fallbackEnv: {}, }; const [first, second] = await Promise.all([ @@ -190,8 +305,8 @@ describe("isolated setup inference detection", () => { detectSetupInferenceIsolated(options), ]); - expect(first).toEqual(emptyDetection()); - expect(second).toEqual(emptyDetection()); - expect(fallback).toHaveBeenCalledOnce(); + expect(first).toEqual(detected); + expect(second).toEqual(detected); + expect(Atomics.load(new Int32Array(started), 0)).toBe(1); }); }); diff --git a/src/system-agent/setup-inference-detection.ts b/src/system-agent/setup-inference-detection.ts index 38b896eb9661..38e174874e89 100644 --- a/src/system-agent/setup-inference-detection.ts +++ b/src/system-agent/setup-inference-detection.ts @@ -12,6 +12,17 @@ const SETUP_INFERENCE_DETECTION_TIMEOUT_MS = 10_000; const log = createSubsystemLogger("system-agent/setup-inference-detection"); +class SetupInferenceDetectionTimeoutError extends Error { + override name = "SetupInferenceDetectionTimeoutError"; + + constructor(timeoutMs: number) { + super( + `Checking this Gateway for AI access timed out after ${timeoutMs / 1_000}s. ` + + "The Gateway may be busy — try again.", + ); + } +} + type DetectionWorkerMessage = | { type: "partial"; detection: SetupInferenceDetection } | { type: "result"; detection: SetupInferenceDetection } @@ -21,13 +32,11 @@ type DetectionWorkerOptions = { timeoutMs?: number; workerUrl?: URL; workerData?: WorkerOptions["workerData"]; - fallback?: () => Promise; fallbackEnv?: NodeJS.ProcessEnv; }; let inFlightDetection: Promise | undefined; let workerShutdown: Promise | undefined; -let workerShutdownResult: SetupInferenceDetection | undefined; function trackWorkerShutdown(worker: Worker): void { const current = worker.terminate().then( @@ -40,7 +49,6 @@ function trackWorkerShutdown(worker: Worker): void { void current.finally(() => { if (workerShutdown === current) { workerShutdown = undefined; - workerShutdownResult = undefined; } }); } @@ -97,21 +105,18 @@ function withAmbientCandidates( return { ...detection, candidates: [...detection.candidates, ...ambient] }; } -function createUndetectedFallback(env: NodeJS.ProcessEnv = process.env): SetupInferenceDetection { +function createUndetectedFallback(): SetupInferenceDetection { // This fallback must stay independent of the detection/plugin graph. The worker // supplies richer partial data when that graph loads before the deadline. - return withAmbientCandidates( - { - candidates: [], - unavailableCandidates: [], - manualProviders: [], - authOptions: [], - recommendedInstalls: listRecommendedToolInstalls(), - workspace: DEFAULT_AGENT_WORKSPACE_DIR, - setupComplete: false, - }, - env, - ); + return { + candidates: [], + unavailableCandidates: [], + manualProviders: [], + authOptions: [], + recommendedInstalls: listRecommendedToolInstalls(), + workspace: DEFAULT_AGENT_WORKSPACE_DIR, + setupComplete: false, + }; } async function runDetectionWorker( @@ -155,7 +160,6 @@ async function runDetectionWorker( reject(new Error(message.error)); return; } - workerShutdownResult = message.detection; resolve(message.detection); }); }); @@ -174,19 +178,18 @@ async function runDetectionWorker( const timer = setTimeout(() => { settle(() => { log.warn( - `Setup inference detection timed out after ${timeoutMs}ms; returning partial detection.`, + `Setup inference detection timed out after ${timeoutMs}ms; using partial signal if available.`, ); - if (options.fallback) { - void options.fallback().then(resolve, reject); - return; - } const env = options.fallbackEnv ?? process.env; const detection = withAmbientCandidates( - partialDetection ?? createUndetectedFallback(env), + partialDetection ?? createUndetectedFallback(), env, ); - workerShutdownResult = detection; - resolve(detection); + if (detection.candidates.length > 0 || detection.unavailableCandidates.length > 0) { + resolve(detection); + return; + } + reject(new SetupInferenceDetectionTimeoutError(timeoutMs)); }); }, timeoutMs); // Installing a message listener references the underlying MessagePort. @@ -202,10 +205,11 @@ export async function detectSetupInferenceIsolated( if (inFlightDetection) { return await inFlightDetection; } - // A native provider probe can delay Worker termination. Reuse the bounded - // result until exit instead of allowing repeat UI requests to stack threads. + // A native provider probe can delay Worker termination. Wait for exit before + // retrying so repeat UI requests neither stack threads nor reuse stale results. if (workerShutdown) { - return workerShutdownResult ?? createUndetectedFallback(); + await workerShutdown; + return await detectSetupInferenceIsolated(options); } const current = runDetectionWorker(options); inFlightDetection = current;