Fix: poll agent engine operations without a trailing sleep

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++;
     }
 

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',

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++;
     }
 

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<void>((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);
 });