From f1647d9e19bf6930bc50bd2e66d2929f8f771503 Mon Sep 17 00:00:00 2001 From: Marat Boshernitsan Date: Tue, 8 Jul 2025 09:37:10 -0700 Subject: [PATCH] Improve auth env var validation logic and messaging to detect settings that confuse GenAI SDK (#1381) Co-authored-by: Scott Densmore Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- docs/cli/authentication.md | 6 +- docs/cli/configuration.md | 7 +- packages/cli/src/config/auth.test.ts | 6 +- packages/cli/src/config/auth.ts | 6 +- packages/core/src/code_assist/setup.test.ts | 82 +++++++++++++++++++ packages/core/src/code_assist/setup.ts | 2 +- .../core/src/core/contentGenerator.test.ts | 79 +++++++++++++++++- packages/core/src/core/contentGenerator.ts | 17 ++-- 8 files changed, 175 insertions(+), 30 deletions(-) create mode 100644 packages/core/src/code_assist/setup.test.ts diff --git a/docs/cli/authentication.md b/docs/cli/authentication.md index 3dddb897..f797e454 100644 --- a/docs/cli/authentication.md +++ b/docs/cli/authentication.md @@ -47,18 +47,16 @@ The Gemini CLI requires you to authenticate with Google's AI services. On initia gcloud auth application-default login ``` For more information, see [Set up Application Default Credentials for Google Cloud](https://cloud.google.com/docs/authentication/provide-credentials-adc). - - Set the `GOOGLE_CLOUD_PROJECT`, `GOOGLE_CLOUD_LOCATION`, and `GOOGLE_GENAI_USE_VERTEXAI` environment variables. In the following methods, replace `YOUR_PROJECT_ID` and `YOUR_PROJECT_LOCATION` with the relevant values for your project: + - Set the `GOOGLE_CLOUD_PROJECT` and `GOOGLE_CLOUD_LOCATION` environment variables. In the following methods, replace `YOUR_PROJECT_ID` and `YOUR_PROJECT_LOCATION` with the relevant values for your project: - You can temporarily set these environment variables in your current shell session using the following commands: ```bash export GOOGLE_CLOUD_PROJECT="YOUR_PROJECT_ID" export GOOGLE_CLOUD_LOCATION="YOUR_PROJECT_LOCATION" # e.g., us-central1 - export GOOGLE_GENAI_USE_VERTEXAI=true ``` - For repeated use, you can add the environment variables to your [.env file](#persisting-environment-variables-with-env-files) or your shell's configuration file (like `~/.bashrc`, `~/.zshrc`, or `~/.profile`). For example, the following commands add the environment variables to a `~/.bashrc` file: ```bash echo 'export GOOGLE_CLOUD_PROJECT="YOUR_PROJECT_ID"' >> ~/.bashrc echo 'export GOOGLE_CLOUD_LOCATION="YOUR_PROJECT_LOCATION"' >> ~/.bashrc - echo 'export GOOGLE_GENAI_USE_VERTEXAI=true' >> ~/.bashrc source ~/.bashrc ``` - If using express mode: @@ -66,12 +64,10 @@ The Gemini CLI requires you to authenticate with Google's AI services. On initia - You can temporarily set these environment variables in your current shell session using the following commands: ```bash export GOOGLE_API_KEY="YOUR_GOOGLE_API_KEY" - export GOOGLE_GENAI_USE_VERTEXAI=true ``` - For repeated use, you can add the environment variables to your [.env file](#persisting-environment-variables-with-env-files) or your shell's configuration file (like `~/.bashrc`, `~/.zshrc`, or `~/.profile`). For example, the following commands add the environment variables to a `~/.bashrc` file: ```bash echo 'export GOOGLE_API_KEY="YOUR_GOOGLE_API_KEY"' >> ~/.bashrc - echo 'export GOOGLE_GENAI_USE_VERTEXAI=true' >> ~/.bashrc source ~/.bashrc ``` 4. **Cloud Shell:** diff --git a/docs/cli/configuration.md b/docs/cli/configuration.md index 31931391..5f0514b0 100644 --- a/docs/cli/configuration.md +++ b/docs/cli/configuration.md @@ -242,12 +242,12 @@ The CLI automatically loads environment variables from an `.env` file. The loadi - **`GOOGLE_API_KEY`**: - Your Google Cloud API key. - Required for using Vertex AI in express mode. - - Ensure you have the necessary permissions and set the `GOOGLE_GENAI_USE_VERTEXAI=true` environment variable. + - Ensure you have the necessary permissions. - Example: `export GOOGLE_API_KEY="YOUR_GOOGLE_API_KEY"`. - **`GOOGLE_CLOUD_PROJECT`**: - Your Google Cloud Project ID. - Required for using Code Assist or Vertex AI. - - If using Vertex AI, ensure you have the necessary permissions and set the `GOOGLE_GENAI_USE_VERTEXAI=true` environment variable. + - If using Vertex AI, ensure you have the necessary permissions in this project. - **Cloud Shell Note:** When running in a Cloud Shell environment, this variable defaults to a special project allocated for Cloud Shell users. If you have `GOOGLE_CLOUD_PROJECT` set in your global environment in Cloud Shell, it will be overridden by this default. To use a different project in Cloud Shell, you must define `GOOGLE_CLOUD_PROJECT` in a `.env` file. - Example: `export GOOGLE_CLOUD_PROJECT="YOUR_PROJECT_ID"`. - **`GOOGLE_APPLICATION_CREDENTIALS`** (string): @@ -258,8 +258,7 @@ The CLI automatically loads environment variables from an `.env` file. The loadi - Example: `export OTLP_GOOGLE_CLOUD_PROJECT="YOUR_PROJECT_ID"`. - **`GOOGLE_CLOUD_LOCATION`**: - Your Google Cloud Project Location (e.g., us-central1). - - Required for using Vertex AI in non-express mode. - - If using Vertex AI, ensure you have the necessary permissions and set the `GOOGLE_GENAI_USE_VERTEXAI=true` environment variable. + - Required for using Vertex AI in non express mode. - Example: `export GOOGLE_CLOUD_LOCATION="YOUR_PROJECT_LOCATION"`. - **`GEMINI_SANDBOX`**: - Alternative to the `sandbox` setting in `settings.json`. diff --git a/packages/cli/src/config/auth.test.ts b/packages/cli/src/config/auth.test.ts index a4c0cf25..96defb1e 100644 --- a/packages/cli/src/config/auth.test.ts +++ b/packages/cli/src/config/auth.test.ts @@ -40,7 +40,7 @@ describe('validateAuthMethod', () => { it('should return an error message if GEMINI_API_KEY is not set', () => { expect(validateAuthMethod(AuthType.USE_GEMINI)).toBe( - 'GEMINI_API_KEY environment variable not found. Add that to your .env and try again, no reload needed!', + 'GEMINI_API_KEY environment variable not found. Add that to your environment and try again (no reload needed if using .env)!', ); }); }); @@ -59,10 +59,10 @@ describe('validateAuthMethod', () => { it('should return an error message if no required environment variables are set', () => { expect(validateAuthMethod(AuthType.USE_VERTEX_AI)).toBe( - 'Must specify GOOGLE_GENAI_USE_VERTEXAI=true and either:\n' + + 'When using Vertex AI, you must specify either:\n' + '• GOOGLE_CLOUD_PROJECT and GOOGLE_CLOUD_LOCATION environment variables.\n' + '• GOOGLE_API_KEY environment variable (if using express mode).\n' + - 'Update your .env and try again, no reload needed!', + 'Update your environment and try again (no reload needed if using .env)!', ); }); }); diff --git a/packages/cli/src/config/auth.ts b/packages/cli/src/config/auth.ts index 0f97b4d5..91d4eee3 100644 --- a/packages/cli/src/config/auth.ts +++ b/packages/cli/src/config/auth.ts @@ -18,7 +18,7 @@ export const validateAuthMethod = (authMethod: string): string | null => { if (authMethod === AuthType.USE_GEMINI) { if (!process.env.GEMINI_API_KEY) { - return 'GEMINI_API_KEY environment variable not found. Add that to your .env and try again, no reload needed!'; + return 'GEMINI_API_KEY environment variable not found. Add that to your environment and try again (no reload needed if using .env)!'; } return null; } @@ -29,10 +29,10 @@ export const validateAuthMethod = (authMethod: string): string | null => { const hasGoogleApiKey = !!process.env.GOOGLE_API_KEY; if (!hasVertexProjectLocationConfig && !hasGoogleApiKey) { return ( - 'Must specify GOOGLE_GENAI_USE_VERTEXAI=true and either:\n' + + 'When using Vertex AI, you must specify either:\n' + '• GOOGLE_CLOUD_PROJECT and GOOGLE_CLOUD_LOCATION environment variables.\n' + '• GOOGLE_API_KEY environment variable (if using express mode).\n' + - 'Update your .env and try again, no reload needed!' + 'Update your environment and try again (no reload needed if using .env)!' ); } return null; diff --git a/packages/core/src/code_assist/setup.test.ts b/packages/core/src/code_assist/setup.test.ts new file mode 100644 index 00000000..479abae0 --- /dev/null +++ b/packages/core/src/code_assist/setup.test.ts @@ -0,0 +1,82 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { setupUser, ProjectIdRequiredError } from './setup.js'; +import { CodeAssistServer } from '../code_assist/server.js'; +import { OAuth2Client } from 'google-auth-library'; +import { GeminiUserTier, UserTierId } from './types.js'; + +vi.mock('../code_assist/server.js'); + +const mockPaidTier: GeminiUserTier = { + id: UserTierId.STANDARD, + name: 'paid', + description: 'Paid tier', +}; + +describe('setupUser', () => { + let mockLoad: ReturnType; + let mockOnboardUser: ReturnType; + + beforeEach(() => { + vi.resetAllMocks(); + mockLoad = vi.fn(); + mockOnboardUser = vi.fn().mockResolvedValue({ + done: true, + response: { + cloudaicompanionProject: { + id: 'server-project', + }, + }, + }); + vi.mocked(CodeAssistServer).mockImplementation( + () => + ({ + loadCodeAssist: mockLoad, + onboardUser: mockOnboardUser, + }) as unknown as CodeAssistServer, + ); + }); + + it('should use GOOGLE_CLOUD_PROJECT when set', async () => { + process.env.GOOGLE_CLOUD_PROJECT = 'test-project'; + mockLoad.mockResolvedValue({ + currentTier: mockPaidTier, + }); + await setupUser({} as OAuth2Client); + expect(CodeAssistServer).toHaveBeenCalledWith( + expect.any(Object), + 'test-project', + ); + }); + + it('should treat empty GOOGLE_CLOUD_PROJECT as undefined and use project from server', async () => { + process.env.GOOGLE_CLOUD_PROJECT = ''; + mockLoad.mockResolvedValue({ + cloudaicompanionProject: 'server-project', + currentTier: mockPaidTier, + }); + const projectId = await setupUser({} as OAuth2Client); + expect(CodeAssistServer).toHaveBeenCalledWith( + expect.any(Object), + undefined, + ); + expect(projectId).toBe('server-project'); + }); + + it('should throw ProjectIdRequiredError when no project ID is available', async () => { + delete process.env.GOOGLE_CLOUD_PROJECT; + // And the server itself requires a project ID internally + vi.mocked(CodeAssistServer).mockImplementation(() => { + throw new ProjectIdRequiredError(); + }); + + await expect(setupUser({} as OAuth2Client)).rejects.toThrow( + ProjectIdRequiredError, + ); + }); +}); diff --git a/packages/core/src/code_assist/setup.ts b/packages/core/src/code_assist/setup.ts index 7db6bdcd..3c7b81b0 100644 --- a/packages/core/src/code_assist/setup.ts +++ b/packages/core/src/code_assist/setup.ts @@ -28,7 +28,7 @@ export class ProjectIdRequiredError extends Error { * @returns the user's actual project id */ export async function setupUser(client: OAuth2Client): Promise { - let projectId = process.env.GOOGLE_CLOUD_PROJECT; + let projectId = process.env.GOOGLE_CLOUD_PROJECT || undefined; const caServer = new CodeAssistServer(client, projectId); const clientMetadata: ClientMetadata = { diff --git a/packages/core/src/core/contentGenerator.test.ts b/packages/core/src/core/contentGenerator.test.ts index 4c6134f2..eb480710 100644 --- a/packages/core/src/core/contentGenerator.test.ts +++ b/packages/core/src/core/contentGenerator.test.ts @@ -4,15 +4,19 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { describe, it, expect, vi } from 'vitest'; -import { createContentGenerator, AuthType } from './contentGenerator.js'; +import { describe, it, expect, vi, beforeEach, afterAll } from 'vitest'; +import { + createContentGenerator, + AuthType, + createContentGeneratorConfig, +} from './contentGenerator.js'; import { createCodeAssistContentGenerator } from '../code_assist/codeAssist.js'; import { GoogleGenAI } from '@google/genai'; vi.mock('../code_assist/codeAssist.js'); vi.mock('@google/genai'); -describe('contentGenerator', () => { +describe('createContentGenerator', () => { it('should create a CodeAssistContentGenerator', async () => { const mockGenerator = {} as unknown; vi.mocked(createCodeAssistContentGenerator).mockResolvedValue( @@ -48,3 +52,72 @@ describe('contentGenerator', () => { expect(generator).toBe((mockGenerator as GoogleGenAI).models); }); }); + +describe('createContentGeneratorConfig', () => { + const originalEnv = process.env; + + beforeEach(() => { + // Reset modules to re-evaluate imports and environment variables + vi.resetModules(); + // Restore process.env before each test + process.env = { ...originalEnv }; + }); + + afterAll(() => { + // Restore original process.env after all tests + process.env = originalEnv; + }); + + it('should configure for Gemini using GEMINI_API_KEY when set', async () => { + process.env.GEMINI_API_KEY = 'env-gemini-key'; + const config = await createContentGeneratorConfig( + undefined, + AuthType.USE_GEMINI, + ); + expect(config.apiKey).toBe('env-gemini-key'); + expect(config.vertexai).toBe(false); + }); + + it('should not configure for Gemini if GEMINI_API_KEY is empty', async () => { + process.env.GEMINI_API_KEY = ''; + const config = await createContentGeneratorConfig( + undefined, + AuthType.USE_GEMINI, + ); + expect(config.apiKey).toBeUndefined(); + expect(config.vertexai).toBeUndefined(); + }); + + it('should configure for Vertex AI using GOOGLE_API_KEY when set', async () => { + process.env.GOOGLE_API_KEY = 'env-google-key'; + const config = await createContentGeneratorConfig( + undefined, + AuthType.USE_VERTEX_AI, + ); + expect(config.apiKey).toBe('env-google-key'); + expect(config.vertexai).toBe(true); + }); + + it('should configure for Vertex AI using GCP project and location when set', async () => { + process.env.GOOGLE_CLOUD_PROJECT = 'env-gcp-project'; + process.env.GOOGLE_CLOUD_LOCATION = 'env-gcp-location'; + const config = await createContentGeneratorConfig( + undefined, + AuthType.USE_VERTEX_AI, + ); + expect(config.vertexai).toBe(true); + expect(config.apiKey).toBeUndefined(); + }); + + it('should not configure for Vertex AI if required env vars are empty', async () => { + process.env.GOOGLE_API_KEY = ''; + process.env.GOOGLE_CLOUD_PROJECT = ''; + process.env.GOOGLE_CLOUD_LOCATION = ''; + const config = await createContentGeneratorConfig( + undefined, + AuthType.USE_VERTEX_AI, + ); + expect(config.apiKey).toBeUndefined(); + expect(config.vertexai).toBeUndefined(); + }); +}); diff --git a/packages/core/src/core/contentGenerator.ts b/packages/core/src/core/contentGenerator.ts index ce3c11a9..e9e1138f 100644 --- a/packages/core/src/core/contentGenerator.ts +++ b/packages/core/src/core/contentGenerator.ts @@ -52,10 +52,10 @@ export async function createContentGeneratorConfig( model: string | undefined, authType: AuthType | undefined, ): Promise { - const geminiApiKey = process.env.GEMINI_API_KEY; - const googleApiKey = process.env.GOOGLE_API_KEY; - const googleCloudProject = process.env.GOOGLE_CLOUD_PROJECT; - const googleCloudLocation = process.env.GOOGLE_CLOUD_LOCATION; + const geminiApiKey = process.env.GEMINI_API_KEY || undefined; + const googleApiKey = process.env.GOOGLE_API_KEY || undefined; + const googleCloudProject = process.env.GOOGLE_CLOUD_PROJECT || undefined; + const googleCloudLocation = process.env.GOOGLE_CLOUD_LOCATION || undefined; // Use runtime model from config if available, otherwise fallback to parameter or default const effectiveModel = model || DEFAULT_GEMINI_MODEL; @@ -75,6 +75,7 @@ export async function createContentGeneratorConfig( if (authType === AuthType.USE_GEMINI && geminiApiKey) { contentGeneratorConfig.apiKey = geminiApiKey; + contentGeneratorConfig.vertexai = false; contentGeneratorConfig.model = await getEffectiveModel( contentGeneratorConfig.apiKey, contentGeneratorConfig.model, @@ -85,16 +86,10 @@ export async function createContentGeneratorConfig( if ( authType === AuthType.USE_VERTEX_AI && - !!googleApiKey && - googleCloudProject && - googleCloudLocation + (googleApiKey || (googleCloudProject && googleCloudLocation)) ) { contentGeneratorConfig.apiKey = googleApiKey; contentGeneratorConfig.vertexai = true; - contentGeneratorConfig.model = await getEffectiveModel( - contentGeneratorConfig.apiKey, - contentGeneratorConfig.model, - ); return contentGeneratorConfig; }