Fix: forward the Vertex AI express-mode API key into the Agent Engines client

core/src/memory/vertex_ai_memory_bank_service.ts

@@ -16,7 +16,10 @@ import {Content, createUserContent} from '@google/genai';
 import {Event} from '../events/event.js';
 import {Session} from '../sessions/session.js';
 import {logger} from '../utils/logger.js';
-import {getExpressModeApiKey} from '../utils/vertex_ai_utils.js';
+import {
+  createAgentEnginesClient,
+  getExpressModeApiKey,
+} from '../utils/vertex_ai_utils.js';
 import {
   BaseMemoryService,
   SearchMemoryRequest,
@@ -143,11 +146,11 @@ export class VertexAiMemoryBankService implements BaseMemoryService {
     if (options.client) {
       this.memories = options.client.agentEnginesInternal.memories;
     } else {
-      const client = new Client({
-        project: this.projectId,
+      this.memories = createAgentEnginesClient({
+        projectId: this.projectId,
         location: this.location,
-      });
-      this.memories = client.agentEnginesInternal.memories;
+        expressModeApiKey: this.expressModeApiKey,
+      }).memories;
     }
   }
 

core/src/sessions/vertex_ai_session_service.ts

@@ -4,7 +4,6 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import {Client} from '@google-cloud/vertexai/build/src/genai/client.js';
 import {Sessions} from '@google-cloud/vertexai/build/src/genai/sessions.js';
 import {
   AppendAgentEngineSessionEventConfig,
@@ -22,7 +21,10 @@ import {Event} from '../events/event.js';
 import {EventActions} from '../events/event_actions.js';
 import {ToolConfirmation} from '../tools/tool_confirmation.js';
 import {logger} from '../utils/logger.js';
-import {getExpressModeApiKey} from '../utils/vertex_ai_utils.js';
+import {
+  createAgentEnginesClient,
+  getExpressModeApiKey,
+} from '../utils/vertex_ai_utils.js';
 
 import {partialCopy} from '../utils/partial_copy.js';
 import {
@@ -102,11 +104,11 @@ export class VertexAiSessionService extends BaseSessionService {
     if (options.sessions) {
       this.sessions = options.sessions;
     } else {
-      const client = new Client({
-        project: this.projectId,
+      this.sessions = createAgentEnginesClient({
+        projectId: this.projectId,
         location: this.location,
-      });
-      this.sessions = client.agentEnginesInternal.sessions;
+        expressModeApiKey: this.expressModeApiKey,
+      }).sessions;
     }
   }
 

core/src/utils/vertex_ai_utils.ts

@@ -4,6 +4,17 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import {AgentEngines} from '@google-cloud/vertexai/build/src/genai/agentengines.js';
+import {
+  Client,
+  SDK_VERSION,
+} from '@google-cloud/vertexai/build/src/genai/client.js';
+import {
+  ApiClient,
+  NodeAuth,
+  NodeDownloader,
+  NodeUploader,
+} from '@google/genai/vertex_internal';
 import {getBooleanEnvVar} from './env_aware_utils.js';
 
 /**
@@ -32,3 +43,54 @@ export function getExpressModeApiKey(
 
   return undefined;
 }
+
+/**
+ * Builds the genai API client used for Vertex AI express mode.
+ *
+ * `@google-cloud/vertexai`'s `Client` has no API key option: it always
+ * authenticates with Application Default Credentials and requires a project
+ * and location, so express mode has to build the underlying client itself. The
+ * key is passed to `NodeAuth`, which emits the `x-goog-api-key` header, as well
+ * as to `ApiClient`, which uses it to skip the `projects/{p}/locations/{l}` URL
+ * prefix.
+ */
+export function createExpressModeApiClient(apiKey: string): ApiClient {
+  return new ApiClient({
+    auth: new NodeAuth({apiKey}),
+    uploader: new NodeUploader(),
+    downloader: new NodeDownloader(),
+    apiKey,
+    vertexai: true,
+    userAgentExtra: `vertex-genai-modules/${SDK_VERSION}`,
+  });
+}
+
+/** The `ApiClient` declaration that `AgentEngines` is compiled against. */
+type AgentEnginesApiClient = ConstructorParameters<typeof AgentEngines>[0];
+
+/**
+ * Creates the Agent Engines client for the given credentials.
+ *
+ * An express mode API key wins over a project/location pair, matching
+ * `_get_api_client` in the Python ADK. The two are mutually exclusive; see
+ * {@link getExpressModeApiKey}.
+ */
+export function createAgentEnginesClient(options: {
+  projectId?: string;
+  location?: string;
+  expressModeApiKey?: string;
+}): AgentEngines {
+  if (options.expressModeApiKey) {
+    const apiClient = createExpressModeApiClient(options.expressModeApiKey);
+    // `AgentEngines` is compiled against the `@google/genai` copy that
+    // `@google-cloud/vertexai@1.12.0` resolves (v1.52.0), while `core` resolves
+    // `@google/genai` v2.9.0. The two `ApiClient` classes are the same at
+    // runtime but nominally distinct to `tsc` because of the private
+    // `customBaseUrl` field.
+    return new AgentEngines(apiClient as unknown as AgentEnginesApiClient);
+  }
+  return new Client({
+    project: options.projectId,
+    location: options.location,
+  }).agentEnginesInternal;
+}

core/test/memory/vertex_ai_memory_bank_service_test.ts

@@ -15,7 +15,7 @@ import {
   VertexAiMemoryBankServiceOptions,
 } from '@google/adk';
 import {Content, Part} from '@google/genai';
-import {beforeEach, describe, expect, it, vi} from 'vitest';
+import {afterEach, beforeEach, describe, expect, it, vi} from 'vitest';
 
 describe('VertexAiMemoryBankService', () => {
   let service: VertexAiMemoryBankService;
@@ -613,3 +613,34 @@ describe('VertexAiMemoryBankService', () => {
     });
   });
 });
+
+describe('VertexAiMemoryBankService express mode', () => {
+  const FAKE_API_KEY = 'fake-express-key';
+  const originalEnv = process.env;
+
+  beforeEach(() => {
+    process.env = {...originalEnv, GOOGLE_GENAI_USE_VERTEXAI: 'true'};
+    delete process.env['GOOGLE_API_KEY'];
+  });
+
+  afterEach(() => {
+    process.env = originalEnv;
+  });
+
+  it('initializes from GOOGLE_API_KEY without a project or location', () => {
+    process.env['GOOGLE_API_KEY'] = FAKE_API_KEY;
+
+    expect(
+      new VertexAiMemoryBankService({agentEngineId: 'test-engine-id'}),
+    ).toBeDefined();
+  });
+
+  it('initializes from an explicit expressModeApiKey', () => {
+    expect(
+      new VertexAiMemoryBankService({
+        agentEngineId: 'test-engine-id',
+        expressModeApiKey: FAKE_API_KEY,
+      }),
+    ).toBeDefined();
+  });
+});

core/test/sessions/vertex_ai_session_service_test.ts

@@ -7,7 +7,7 @@
 import {Sessions} from '@google-cloud/vertexai/build/src/genai/sessions.js';
 import {createEvent, State, VertexAiSessionService} from '@google/adk';
 import {Session} from '@google/adk/sessions/session.js';
-import {beforeEach, describe, expect, it, vi} from 'vitest';
+import {afterEach, beforeEach, describe, expect, it, vi} from 'vitest';
 
 // Mock the unreleased nodejs-vertexai package so the import resolves
 vi.mock('nodejs-vertexai', () => ({
@@ -1093,3 +1093,63 @@ describe('VertexAiSessionService', () => {
     });
   });
 });
+
+describe('VertexAiSessionService express mode', () => {
+  const FAKE_API_KEY = 'fake-express-key';
+  const originalEnv = process.env;
+
+  beforeEach(() => {
+    process.env = {...originalEnv, GOOGLE_GENAI_USE_VERTEXAI: 'true'};
+    delete process.env['GOOGLE_API_KEY'];
+  });
+
+  afterEach(() => {
+    process.env = originalEnv;
+    vi.restoreAllMocks();
+  });
+
+  it('initializes from GOOGLE_API_KEY without a project or location', () => {
+    process.env['GOOGLE_API_KEY'] = FAKE_API_KEY;
+
+    expect(new VertexAiSessionService({})).toBeDefined();
+  });
+
+  it('initializes from an explicit expressModeApiKey', () => {
+    expect(
+      new VertexAiSessionService({expressModeApiKey: FAKE_API_KEY}),
+    ).toBeDefined();
+  });
+
+  it('authenticates with the key and omits the project path prefix', async () => {
+    const fetchSpy = vi.spyOn(globalThis, 'fetch').mockResolvedValue(
+      new Response(
+        JSON.stringify({
+          name: 'operations/test-operation-id',
+          done: true,
+          response: {
+            name: 'reasoningEngines/12345/sessions/test-session-id',
+            sessionState: {},
+            updateTime: '2026-04-21T12:00:00Z',
+          },
+        }),
+        {status: 200, headers: {'content-type': 'application/json'}},
+      ),
+    );
+    const service = new VertexAiSessionService({
+      expressModeApiKey: FAKE_API_KEY,
+    });
+
+    const session = await service.createSession({
+      appName: '12345',
+      userId: 'testUser',
+    });
+
+    expect(session.id).toBe('test-session-id');
+    expect(fetchSpy).toHaveBeenCalledTimes(1);
+    const [url, init] = fetchSpy.mock.calls[0];
+    expect(String(url)).toBe(
+      'https://aiplatform.googleapis.com/v1beta1/reasoningEngines/12345/sessions',
+    );
+    expect(new Headers(init?.headers).get('x-goog-api-key')).toBe(FAKE_API_KEY);
+  });
+});

core/test/utils/vertex_ai_utils_test.ts

@@ -4,8 +4,15 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import {SDK_VERSION} from '@google-cloud/vertexai/build/src/genai/client.js';
 import {afterEach, beforeEach, describe, expect, it} from 'vitest';
-import {getExpressModeApiKey} from '../../src/utils/vertex_ai_utils.js';
+import {
+  createAgentEnginesClient,
+  createExpressModeApiClient,
+  getExpressModeApiKey,
+} from '../../src/utils/vertex_ai_utils.js';
+
+const FAKE_API_KEY = 'fake-express-key';
 
 describe('vertex_ai_utils', () => {
   describe('getExpressModeApiKey', () => {
@@ -69,4 +76,57 @@ describe('vertex_ai_utils', () => {
       expect(result).toBeUndefined();
     });
   });
+
+  describe('createExpressModeApiClient', () => {
+    it('should authenticate with the key instead of a project and location', () => {
+      const client = createExpressModeApiClient(FAKE_API_KEY);
+
+      expect(client.getApiKey()).toBe(FAKE_API_KEY);
+      expect(client.getProject()).toBeUndefined();
+      expect(client.getLocation()).toBeUndefined();
+      expect(client.isVertexAI()).toBe(true);
+    });
+
+    it('should send the key as the x-goog-api-key header', async () => {
+      const headers =
+        await createExpressModeApiClient(FAKE_API_KEY).getAuthHeaders();
+
+      expect(headers.get('x-goog-api-key')).toBe(FAKE_API_KEY);
+    });
+
+    it('should report the same user agent as the vendor client', () => {
+      const client = createExpressModeApiClient(FAKE_API_KEY);
+
+      expect(client.clientOptions.userAgentExtra).toBe(
+        `vertex-genai-modules/${SDK_VERSION}`,
+      );
+    });
+  });
+
+  describe('createAgentEnginesClient', () => {
+    it('should build a client from an express mode key alone', () => {
+      const client = createAgentEnginesClient({
+        expressModeApiKey: FAKE_API_KEY,
+      });
+
+      expect(client.sessions).toBeDefined();
+      expect(client.memories).toBeDefined();
+    });
+
+    it('should build a client from a project and location', () => {
+      const client = createAgentEnginesClient({
+        projectId: 'test-project',
+        location: 'us-central1',
+      });
+
+      expect(client.sessions).toBeDefined();
+      expect(client.memories).toBeDefined();
+    });
+
+    it('should throw when given neither a key nor a project and location', () => {
+      expect(() => createAgentEnginesClient({})).toThrow(
+        'Authentication is not set up.',
+      );
+    });
+  });
 });

tests/integration/memory/vertex_ai_memory_bank_service_test.ts

@@ -14,7 +14,7 @@ import {
   VertexAiMemoryBankService,
 } from '@google/adk';
 import {createUserContent} from '@google/genai';
-import {beforeEach, describe, expect, it, vi} from 'vitest';
+import {afterEach, beforeEach, describe, expect, it, vi} from 'vitest';
 import {GeminiWithMockResponses} from '../test_case_utils.js';
 
 describe('VertexAiMemoryBankService Integration', () => {
@@ -161,3 +161,57 @@ describe('VertexAiMemoryBankService Integration', () => {
     );
   });
 });
+
+describe('VertexAiMemoryBankService Express Mode Integration', () => {
+  const FAKE_API_KEY = 'fake-express-key';
+  const originalEnv = process.env;
+
+  beforeEach(() => {
+    process.env = {...originalEnv, GOOGLE_GENAI_USE_VERTEXAI: 'true'};
+    delete process.env['GOOGLE_API_KEY'];
+  });
+
+  afterEach(() => {
+    process.env = originalEnv;
+    vi.restoreAllMocks();
+  });
+
+  it('should search memory over a real express mode client', async () => {
+    const fetchSpy = vi.spyOn(globalThis, 'fetch').mockResolvedValue(
+      new Response(
+        JSON.stringify({
+          retrievedMemories: [
+            {
+              memory: {
+                fact: 'Your favorite color is green.',
+                updateTime: '2026-04-21T12:00:00Z',
+              },
+              distance: 0.1,
+            },
+          ],
+        }),
+        {status: 200, headers: {'content-type': 'application/json'}},
+      ),
+    );
+    const service = new VertexAiMemoryBankService({
+      agentEngineId: 'test-engine-id',
+      expressModeApiKey: FAKE_API_KEY,
+    });
+
+    const response = await service.searchMemory({
+      appName: 'test_memory_app',
+      userId: 'test_user',
+      query: 'favorite color',
+    });
+
+    expect(response.memories[0].content?.parts?.[0]?.text).toBe(
+      'Your favorite color is green.',
+    );
+    expect(fetchSpy).toHaveBeenCalledTimes(1);
+    const [url, init] = fetchSpy.mock.calls[0];
+    expect(String(url)).toBe(
+      'https://aiplatform.googleapis.com/v1beta1/reasoningEngines/test-engine-id/memories:retrieve',
+    );
+    expect(new Headers(init?.headers).get('x-goog-api-key')).toBe(FAKE_API_KEY);
+  });
+});