Feat: Honor GOOGLE_GENAI_USE_ENTERPRISE in getExpressModeApiKey() (adk-python parity)

core/src/utils/vertex_ai_utils.ts

@@ -5,10 +5,42 @@
  */
 
 import {getBooleanEnvVar} from './env_aware_utils.js';
+import {logger} from './logger.js';
+
+const ENTERPRISE_MODE_ENV_VAR = 'GOOGLE_GENAI_USE_ENTERPRISE';
+const DEPRECATED_ENTERPRISE_MODE_ENV_VAR = 'GOOGLE_GENAI_USE_VERTEXAI';
+
+/**
+ * Returns whether Google GenAI enterprise mode is enabled.
+ *
+ * `GOOGLE_GENAI_USE_ENTERPRISE` takes precedence whenever it is set, even when
+ * it is set to a falsy value. `GOOGLE_GENAI_USE_VERTEXAI` is only consulted
+ * when `GOOGLE_GENAI_USE_ENTERPRISE` is absent, and using it logs a deprecation
+ * warning.
+ */
+function isEnterpriseModeEnabled(): boolean {
+  if (process.env?.[ENTERPRISE_MODE_ENV_VAR] !== undefined) {
+    return getBooleanEnvVar(ENTERPRISE_MODE_ENV_VAR);
+  }
+
+  if (process.env?.[DEPRECATED_ENTERPRISE_MODE_ENV_VAR] !== undefined) {
+    logger.warn(
+      `${DEPRECATED_ENTERPRISE_MODE_ENV_VAR} is deprecated, please use ` +
+        `${ENTERPRISE_MODE_ENV_VAR} instead`,
+    );
+    return getBooleanEnvVar(DEPRECATED_ENTERPRISE_MODE_ENV_VAR);
+  }
+
+  return false;
+}
 
 /**
  * Validates and returns the API key for Express Mode.
  *
+ * The key is only returned when enterprise mode is enabled via
+ * `GOOGLE_GENAI_USE_ENTERPRISE` (or the deprecated
+ * `GOOGLE_GENAI_USE_VERTEXAI`).
+ *
  * @param project The project id.
  * @param location The location.
  * @param expressModeApiKey The API key for Express Mode.
@@ -26,7 +58,7 @@ export function getExpressModeApiKey(
     );
   }
 
-  if (getBooleanEnvVar('GOOGLE_GENAI_USE_VERTEXAI')) {
+  if (isEnterpriseModeEnabled()) {
     return expressModeApiKey || process.env.GOOGLE_API_KEY;
   }
 

core/test/utils/vertex_ai_utils_test.ts

@@ -4,19 +4,34 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import {afterEach, beforeEach, describe, expect, it} from 'vitest';
+import {
+  afterEach,
+  beforeEach,
+  describe,
+  expect,
+  it,
+  vi,
+  type MockInstance,
+} from 'vitest';
+import {logger} from '../../src/utils/logger.js';
 import {getExpressModeApiKey} from '../../src/utils/vertex_ai_utils.js';
 
 describe('vertex_ai_utils', () => {
   describe('getExpressModeApiKey', () => {
     const originalEnv = process.env;
+    let warnSpy: MockInstance<typeof logger.warn>;
 
     beforeEach(() => {
       process.env = {...originalEnv};
+      delete process.env['GOOGLE_GENAI_USE_ENTERPRISE'];
+      delete process.env['GOOGLE_GENAI_USE_VERTEXAI'];
+      delete process.env['GOOGLE_API_KEY'];
+      warnSpy = vi.spyOn(logger, 'warn').mockImplementation(() => {});
     });
 
     afterEach(() => {
       process.env = originalEnv;
+      vi.restoreAllMocks();
     });
 
     it('should throw when both project and expressModeApiKey are provided', () => {
@@ -37,16 +52,19 @@ describe('vertex_ai_utils', () => {
       ).toThrow();
     });
 
-    it('should return undefined when GOOGLE_GENAI_USE_VERTEXAI is not set', () => {
-      delete process.env['GOOGLE_GENAI_USE_VERTEXAI'];
+    it('should return undefined and not warn when neither enterprise mode variable is set', () => {
+      process.env['GOOGLE_API_KEY'] = 'env-api-key';
       const result = getExpressModeApiKey();
       expect(result).toBeUndefined();
+      expect(warnSpy).not.toHaveBeenCalled();
     });
 
     it('should return undefined when GOOGLE_GENAI_USE_VERTEXAI is false', () => {
       process.env['GOOGLE_GENAI_USE_VERTEXAI'] = 'false';
+      process.env['GOOGLE_API_KEY'] = 'env-api-key';
       const result = getExpressModeApiKey();
       expect(result).toBeUndefined();
+      expect(warnSpy).toHaveBeenCalledOnce();
     });
 
     it('should return expressModeApiKey when GOOGLE_GENAI_USE_VERTEXAI is true', () => {
@@ -60,13 +78,43 @@ describe('vertex_ai_utils', () => {
       process.env['GOOGLE_API_KEY'] = 'env-api-key';
       const result = getExpressModeApiKey();
       expect(result).toBe('env-api-key');
+      expect(warnSpy).toHaveBeenCalledOnce();
+      expect(warnSpy).toHaveBeenCalledWith(
+        expect.stringContaining(
+          'GOOGLE_GENAI_USE_VERTEXAI is deprecated, please use ' +
+            'GOOGLE_GENAI_USE_ENTERPRISE',
+        ),
+      );
     });
 
     it('should return undefined when GOOGLE_GENAI_USE_VERTEXAI is true but no key available', () => {
       process.env['GOOGLE_GENAI_USE_VERTEXAI'] = 'true';
-      delete process.env['GOOGLE_API_KEY'];
       const result = getExpressModeApiKey();
       expect(result).toBeUndefined();
     });
+
+    it('should return expressModeApiKey when GOOGLE_GENAI_USE_ENTERPRISE is true', () => {
+      process.env['GOOGLE_GENAI_USE_ENTERPRISE'] = 'true';
+      const result = getExpressModeApiKey(undefined, undefined, 'my-api-key');
+      expect(result).toBe('my-api-key');
+    });
+
+    it('should prefer an enabled GOOGLE_GENAI_USE_ENTERPRISE over GOOGLE_GENAI_USE_VERTEXAI', () => {
+      process.env['GOOGLE_GENAI_USE_ENTERPRISE'] = 'true';
+      process.env['GOOGLE_GENAI_USE_VERTEXAI'] = 'false';
+      process.env['GOOGLE_API_KEY'] = 'env-api-key';
+      const result = getExpressModeApiKey();
+      expect(result).toBe('env-api-key');
+      expect(warnSpy).not.toHaveBeenCalled();
+    });
+
+    it('should not fall back to GOOGLE_GENAI_USE_VERTEXAI when GOOGLE_GENAI_USE_ENTERPRISE is set but disabled', () => {
+      process.env['GOOGLE_GENAI_USE_ENTERPRISE'] = '';
+      process.env['GOOGLE_GENAI_USE_VERTEXAI'] = 'true';
+      process.env['GOOGLE_API_KEY'] = 'env-api-key';
+      const result = getExpressModeApiKey();
+      expect(result).toBeUndefined();
+      expect(warnSpy).not.toHaveBeenCalled();
+    });
   });
 });