diff --git a/src/agents/run-wait.test.ts b/src/agents/run-wait.test.ts index 128be707b04..adf4f36f81a 100644 --- a/src/agents/run-wait.test.ts +++ b/src/agents/run-wait.test.ts @@ -15,6 +15,55 @@ import { waitForAgentRunAndReadUpdatedAssistantReply, } from "./run-wait.js"; +type AgentWaitGatewayRequest = { + method?: string; + params?: { + runId?: string; + timeoutMs?: unknown; + }; + timeoutMs?: unknown; +}; + +function expectNumber(value: unknown, label: string): number { + expect(typeof value).toBe("number"); + if (typeof value !== "number") { + throw new Error(`expected ${label} to be a number`); + } + return value; +} + +function gatewayWaitRequests(): AgentWaitGatewayRequest[] { + return callGatewayMock.mock.calls.map(([request]) => request as AgentWaitGatewayRequest); +} + +function requireRequestAt( + requests: readonly AgentWaitGatewayRequest[], + index: number, +): AgentWaitGatewayRequest { + const request = requests.at(index); + if (!request) { + throw new Error(`expected gateway request at index ${index}`); + } + return request; +} + +function expectAgentWaitRequest( + request: AgentWaitGatewayRequest, + runId: string, + maxParamTimeoutMs: number, +): void { + expect(request.method).toBe("agent.wait"); + expect(request.params?.runId).toBe(runId); + + const paramTimeoutMs = expectNumber(request.params?.timeoutMs, `${runId} param timeoutMs`); + const requestTimeoutMs = expectNumber(request.timeoutMs, `${runId} request timeoutMs`); + expect(requestTimeoutMs).toBeGreaterThan(0); + expect(requestTimeoutMs).toBeLessThanOrEqual(maxParamTimeoutMs + 2_000); + expect(paramTimeoutMs).toBe(requestTimeoutMs - 2_000); + expect(paramTimeoutMs).toBeGreaterThan(0); + expect(paramTimeoutMs).toBeLessThanOrEqual(maxParamTimeoutMs); +} + describe("readLatestAssistantReply", () => { beforeEach(() => { callGatewayMock.mockClear(); @@ -287,29 +336,14 @@ describe("waitForAgentRunsToDrain", () => { getPendingRunIds: () => activeRunIds, }); - expect(result).toEqual({ - timedOut: false, - pendingRunIds: [], - deadlineAtMs: expect.any(Number), - }); - expect(callGatewayMock.mock.calls.map((call) => call[0])).toEqual([ - { - method: "agent.wait", - params: { - runId: "run-1", - timeoutMs: expect.any(Number), - }, - timeoutMs: expect.any(Number), - }, - { - method: "agent.wait", - params: { - runId: "run-2", - timeoutMs: expect.any(Number), - }, - timeoutMs: expect.any(Number), - }, - ]); + expect(result.timedOut).toBe(false); + expect(result.pendingRunIds).toStrictEqual([]); + expectNumber(result.deadlineAtMs, "deadlineAtMs"); + + const requests = gatewayWaitRequests(); + expect(requests).toHaveLength(2); + expectAgentWaitRequest(requireRequestAt(requests, 0), "run-1", 1_000); + expectAgentWaitRequest(requireRequestAt(requests, 1), "run-2", 1_000); }); it("deduplicates and trims pending run ids", async () => { @@ -344,15 +378,9 @@ describe("waitForAgentRunsToDrain", () => { }); expect(result.timedOut).toBe(false); - expect(callGatewayMock.mock.calls.map((call) => call[0])).toEqual([ - expect.objectContaining({ - method: "agent.wait", - params: expect.objectContaining({ runId: "run-1" }), - }), - expect.objectContaining({ - method: "agent.wait", - params: expect.objectContaining({ runId: "run-2" }), - }), - ]); + const requests = gatewayWaitRequests(); + expect(requests).toHaveLength(2); + expectAgentWaitRequest(requireRequestAt(requests, 0), "run-1", 1_000); + expectAgentWaitRequest(requireRequestAt(requests, 1), "run-2", 1_000); }); });