diff --git a/src/worker/worker.runtime.test.ts b/src/worker/worker.runtime.test.ts index 73a995b43ed7..ea5676473e0d 100644 --- a/src/worker/worker.runtime.test.ts +++ b/src/worker/worker.runtime.test.ts @@ -36,6 +36,7 @@ import { type WorkerInferenceTerminalFrame, type WorkerInferenceTerminalOutcome, } from "../../packages/gateway-protocol/src/schema/worker-inference.js"; +import { createDeferred, withTestTimeout } from "../../test/helpers/promise.js"; import { createOperationalRunInstanceRef } from "../agents/admitted-run-context.js"; import { listRunningSessions } from "../agents/bash-process-registry.js"; import { buildWorkerConnectParams, type WorkerLaunchDescriptor } from "./launch-descriptor.js"; @@ -94,6 +95,7 @@ const WORKER_LOOP_REPLAY = { }; const BUNDLE_HASH = Array.from({ length: 64 }, () => "a").join(""); const CREDENTIAL = ["worker", "fixture", "admission"].join("-"); +const WORKER_INFERENCE_START_TIMEOUT_MS = 90_000; type InferencePlan = | "text" @@ -165,6 +167,7 @@ class FakeWorkerGateway { private sentLiveResync = 0; private unavailable = false; private ignoredAdmission = false; + private readonly inferenceStarted = createDeferred(); socketPath = ""; connectionCount = 0; @@ -176,6 +179,14 @@ class FakeWorkerGateway { readonly sessionSpawnRequests: WorkerSessionsSpawnParams[] = []; readonly applicationOrder: string[] = []; + waitForInferenceStart(): Promise { + return withTestTimeout( + this.inferenceStarted.promise, + WORKER_INFERENCE_START_TIMEOUT_MS, + "worker inference start did not reach the fake Gateway", + ); + } + constructor(private readonly options: FakeGatewayOptions = {}) { this.httpServer = createServer(); this.webSocketServer = new WebSocketServer({ server: this.httpServer }); @@ -468,6 +479,7 @@ class FakeWorkerGateway { private handleInference(socket: WebSocket, frame: WorkerInferenceStartRequestFrame): void { this.methods.push(frame.method); this.inferenceRequests.push(structuredClone(frame.params)); + this.inferenceStarted.resolve(); if (this.options.silenceFirstInference && !this.droppedInference) { this.droppedInference = true; return; @@ -1168,7 +1180,7 @@ describe("worker runtime", () => { const { gateway, launch } = await setup({ inferencePlans: ["hold"] }); const controller = new AbortController(); const result = runWorkerDescriptor(launch, { signal: controller.signal }); - await waitForFast(() => expect(gateway.inferenceRequests).toHaveLength(1)); + await gateway.waitForInferenceStart(); controller.abort(new Error("operator stopped worker")); @@ -1187,7 +1199,7 @@ describe("worker runtime", () => { }); const controller = new AbortController(); const result = runWorkerDescriptor(launch, { signal: controller.signal }); - await waitForFast(() => expect(gateway.inferenceRequests).toHaveLength(1)); + await gateway.waitForInferenceStart(); vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout"] }); const timeoutSpy = vi.spyOn(globalThis, "setTimeout");