From 5b46c0c4f1f11b507d6ec883e0d65fc8d735545b Mon Sep 17 00:00:00 2001 From: Amaad Martin Date: Fri, 7 Aug 2026 07:56:38 -0700 Subject: [PATCH 1/3] fix(sessions): poll agent engine operations without a trailing sleep Both LRO poll loops raced the poll against a fixed sleep with Promise.all, so each iteration cost max(poll, interval). The iteration that observed done still paid its own sleep, adding 1000 ms to every createSession and 5000 ms to every agent engine deploy. Sleep between polls instead. The first poll still fires immediately, the attempt budget is still 30 polls, and the timeout error strings are unchanged. --- core/src/sessions/vertex_ai_session_service.ts | 16 +++++++++------- dev/src/cli/deploy/cli_deploy_agent_engine.ts | 16 +++++++++------- 2 files changed, 18 insertions(+), 14 deletions(-) diff --git a/core/src/sessions/vertex_ai_session_service.ts b/core/src/sessions/vertex_ai_session_service.ts index 6398b73f8..77f08b688 100644 --- a/core/src/sessions/vertex_ai_session_service.ts +++ b/core/src/sessions/vertex_ai_session_service.ts @@ -41,6 +41,7 @@ import { import {createSession, Session} from './session.js'; const DEFAULT_MAX_ATTEMPTS = 30; +const POLL_INTERVAL_MS = 1_000; const GRPC_NOT_FOUND = 5; const HTTP_NOT_FOUND = 404; @@ -181,13 +182,14 @@ export class VertexAiSessionService extends BaseSessionService { let attempts = 0; while (!apiResponse.done && attempts < DEFAULT_MAX_ATTEMPTS) { - const [nextResponse] = await Promise.all([ - this.sessions.getSessionOperationInternal({ - operationName: operationName, - }), - new Promise((resolve) => setTimeout(resolve, 1000)), - ]); - apiResponse = nextResponse; + // Delay between polls only, so the poll that observes `done` returns + // immediately instead of waiting out an interval nothing depends on. + if (attempts > 0) { + await new Promise((resolve) => setTimeout(resolve, POLL_INTERVAL_MS)); + } + apiResponse = await this.sessions.getSessionOperationInternal({ + operationName, + }); attempts++; } diff --git a/dev/src/cli/deploy/cli_deploy_agent_engine.ts b/dev/src/cli/deploy/cli_deploy_agent_engine.ts index 9d3f18031..be01b4af4 100644 --- a/dev/src/cli/deploy/cli_deploy_agent_engine.ts +++ b/dev/src/cli/deploy/cli_deploy_agent_engine.ts @@ -21,6 +21,7 @@ import { } from './deploy_utils.js'; const DEFAULT_MAX_ATTEMPTS = 30; +const POLL_INTERVAL_MS = 5_000; export interface DeployToAgentEngineOptions extends BaseDeployOptions { displayName?: string; @@ -180,13 +181,14 @@ export async function deployToAgentEngine(options: DeployToAgentEngineOptions) { let attempts = 0; while (!apiResponse.done && attempts < DEFAULT_MAX_ATTEMPTS) { - const [nextResponse] = await Promise.all([ - client.agentEnginesInternal.getAgentOperationInternal({ - operationName, - }), - new Promise((resolve) => setTimeout(resolve, 5000)), - ]); - apiResponse = nextResponse; + // Delay between polls only, so the poll that observes `done` returns + // immediately instead of waiting out an interval nothing depends on. + if (attempts > 0) { + await new Promise((resolve) => setTimeout(resolve, POLL_INTERVAL_MS)); + } + apiResponse = await client.agentEnginesInternal.getAgentOperationInternal( + {operationName}, + ); attempts++; } From 0e12893ab39515607e7e6d822ecf5b481bcc3bcd Mon Sep 17 00:00:00 2001 From: Amaad Martin Date: Fri, 7 Aug 2026 08:01:30 -0700 Subject: [PATCH 2/3] test: pin the LRO poll schedule in both agent engine poll loops Four tests: the first poll must settle the promise with the clock still at its start value and no pending timer, and the not-done path must still issue 30 polls separated by 29 interval sleeps before it times out. --- .../vertex_ai_session_service_test.ts | 82 +++++++++++++++++++ dev/test/cli/cli_deploy_agent_engine_test.ts | 68 +++++++++++++++ 2 files changed, 150 insertions(+) diff --git a/core/test/sessions/vertex_ai_session_service_test.ts b/core/test/sessions/vertex_ai_session_service_test.ts index 97396d11a..4fea570db 100644 --- a/core/test/sessions/vertex_ai_session_service_test.ts +++ b/core/test/sessions/vertex_ai_session_service_test.ts @@ -321,6 +321,88 @@ describe('VertexAiSessionService', () => { vi.useRealTimers(); }); + it('resolves as soon as the first poll reports done, without waiting an interval', async () => { + mockClient.createInternal.mockResolvedValue({ + name: 'operations/op-1', + done: false, + }); + mockClient.getSessionOperationInternal.mockResolvedValue({ + done: true, + response: { + name: 'projects/p/locations/l/sessions/test-id', + sessionState: {}, + }, + }); + + vi.useFakeTimers(); + try { + const start = Date.now(); + let resolved = false; + const createPromise = service + .createSession({appName: '12345', userId: 'testUser'}) + .then((session) => { + resolved = true; + return session; + }); + + // Crosses a macrotask boundary, draining the microtask chain while + // leaving the fake clock at its start value. + await vi.advanceTimersByTimeAsync(0); + + expect(resolved).toBe(true); + expect(Date.now() - start).toBe(0); + expect(vi.getTimerCount()).toBe(0); + expect(mockClient.getSessionOperationInternal).toHaveBeenCalledTimes(1); + await expect(createPromise).resolves.toMatchObject({id: 'test-id'}); + } finally { + vi.useRealTimers(); + } + }); + + it('spaces polls one interval apart and throws after the maximum attempts', async () => { + mockClient.createInternal.mockResolvedValue({ + name: 'operation-456', + done: false, + }); + mockClient.getSessionOperationInternal.mockResolvedValue({ + name: 'operation-456', + done: false, + }); + + vi.useFakeTimers(); + try { + const start = Date.now(); + const createPromise = service.createSession({ + appName: '12345', + userId: 'testUser', + }); + const rejection = expect(createPromise).rejects.toThrow( + 'Session creation operation operation-456 did not complete in time.', + ); + + await vi.advanceTimersByTimeAsync(0); + expect(mockClient.getSessionOperationInternal).toHaveBeenCalledTimes(1); + + await vi.advanceTimersByTimeAsync(999); + expect(mockClient.getSessionOperationInternal).toHaveBeenCalledTimes(1); + + await vi.advanceTimersByTimeAsync(1); + expect(mockClient.getSessionOperationInternal).toHaveBeenCalledTimes(2); + + await vi.advanceTimersByTimeAsync(28 * 1000); + expect(mockClient.getSessionOperationInternal).toHaveBeenCalledTimes( + 30, + ); + // 30 polls separated by 29 sleeps, and nothing pending afterwards. + expect(Date.now() - start).toBe(29 * 1000); + expect(vi.getTimerCount()).toBe(0); + + await rejection; + } finally { + vi.useRealTimers(); + } + }); + it('falls back to Date.now if update_time is missing in createSession', async () => { mockClient.createInternal.mockResolvedValue({ name: 'projects/p/locations/l/operations/o', diff --git a/dev/test/cli/cli_deploy_agent_engine_test.ts b/dev/test/cli/cli_deploy_agent_engine_test.ts index d7121e9be..98b1420ce 100644 --- a/dev/test/cli/cli_deploy_agent_engine_test.ts +++ b/dev/test/cli/cli_deploy_agent_engine_test.ts @@ -749,4 +749,72 @@ describe('deployToAgentEngine', () => { 'Reasoning Engine update failed: [Code 404] Resource not found', ); }); + + it('should resolve as soon as the first poll reports done, without waiting an interval', async () => { + // The beforeEach stub makes every sleep instantaneous, which would hide a + // trailing sleep. Use a fake clock so the elapsed time is measurable. + vi.unstubAllGlobals(); + vi.useFakeTimers(); + try { + const start = Date.now(); + + await deployToAgentEngine(defaultOptions); + + expect(Date.now() - start).toBe(0); + expect(vi.getTimerCount()).toBe(0); + expect(mockGetAgentOperationInternal).toHaveBeenCalledTimes(1); + } finally { + vi.useRealTimers(); + } + }, 5000); + + it('should space polls one interval apart and throw after the maximum attempts', async () => { + let resolveReachedLoop: () => void = () => {}; + const reachedLoopPromise = new Promise((r) => { + resolveReachedLoop = r; + }); + + mockCreateInternal.mockImplementation(() => { + resolveReachedLoop(); + return Promise.resolve({ + name: 'operations/test-operation', + done: false, + }); + }); + + mockGetAgentOperationInternal.mockResolvedValue({ + name: 'operations/test-operation', + done: false, + }); + + vi.unstubAllGlobals(); + vi.useFakeTimers(); + try { + const start = Date.now(); + const deployPromise = deployToAgentEngine(defaultOptions); + const rejection = expect(deployPromise).rejects.toThrow( + 'Reasoning Engine creation operation operations/test-operation did not complete in time.', + ); + + await reachedLoopPromise; + await vi.advanceTimersByTimeAsync(0); + expect(mockGetAgentOperationInternal).toHaveBeenCalledTimes(1); + + await vi.advanceTimersByTimeAsync(4999); + expect(mockGetAgentOperationInternal).toHaveBeenCalledTimes(1); + + await vi.advanceTimersByTimeAsync(1); + expect(mockGetAgentOperationInternal).toHaveBeenCalledTimes(2); + + await vi.advanceTimersByTimeAsync(28 * 5000); + expect(mockGetAgentOperationInternal).toHaveBeenCalledTimes(30); + // 30 polls separated by 29 sleeps, and nothing pending afterwards. + expect(Date.now() - start).toBe(29 * 5000); + expect(vi.getTimerCount()).toBe(0); + + await rejection; + } finally { + vi.useRealTimers(); + } + }, 30000); }); From 954fb5e43bfcd141753d57fba2179aa8878e4f0f Mon Sep 17 00:00:00 2001 From: Amaad Martin Date: Fri, 7 Aug 2026 08:13:05 -0700 Subject: [PATCH 3/3] test: attach the deploy timeout rejection handler before advancing timers Both agent engine deploy timeout tests attached their rejects assertion after the timer advance loop. The loop now drops its trailing sleep, so the deploy rejects one advance earlier and Node reported the rejection as unhandled, which failed the run under load. The assertion and its message are unchanged. --- dev/test/cli/cli_deploy_agent_engine_test.ts | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/dev/test/cli/cli_deploy_agent_engine_test.ts b/dev/test/cli/cli_deploy_agent_engine_test.ts index 98b1420ce..e160b4aa2 100644 --- a/dev/test/cli/cli_deploy_agent_engine_test.ts +++ b/dev/test/cli/cli_deploy_agent_engine_test.ts @@ -586,6 +586,11 @@ describe('deployToAgentEngine', () => { vi.useFakeTimers(); const deployPromise = deployToAgentEngine(defaultOptions); + // Attach the handler before advancing: the deploy now rejects during the + // advance loop, and an unattached rejection is reported as unhandled. + const rejection = expect(deployPromise).rejects.toThrow( + 'Reasoning Engine creation operation operations/test-operation did not complete in time.', + ); await reachedLoopPromise; await Promise.resolve(); // yield @@ -594,9 +599,7 @@ describe('deployToAgentEngine', () => { await vi.advanceTimersByTimeAsync(5000); } - await expect(deployPromise).rejects.toThrow( - 'Reasoning Engine creation operation operations/test-operation did not complete in time.', - ); + await rejection; vi.useRealTimers(); }, 30000); @@ -710,6 +713,11 @@ describe('deployToAgentEngine', () => { vi.useFakeTimers(); const deployPromise = deployToAgentEngine(options); + // Attach the handler before advancing: the deploy now rejects during the + // advance loop, and an unattached rejection is reported as unhandled. + const rejection = expect(deployPromise).rejects.toThrow( + 'Reasoning Engine update operation operations/test-update-op did not complete in time.', + ); await reachedLoopPromise; await Promise.resolve(); // yield @@ -718,9 +726,7 @@ describe('deployToAgentEngine', () => { await vi.advanceTimersByTimeAsync(5000); } - await expect(deployPromise).rejects.toThrow( - 'Reasoning Engine update operation operations/test-update-op did not complete in time.', - ); + await rejection; vi.useRealTimers(); }, 30000);