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/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/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++; } diff --git a/dev/test/cli/cli_deploy_agent_engine_test.ts b/dev/test/cli/cli_deploy_agent_engine_test.ts index d7121e9be..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); @@ -749,4 +755,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); });