feat(auth): Enhance non-interactive gcp auth (#4811)

2 files changed, 147 insertions, 0 deletions

diff --git a/packages/core/src/code_assist/oauth2.test.ts b/packages/core/src/code_assist/oauth2.test.ts
index ee63c02b..0dc77867 100644
--- a/packages/core/src/code_assist/oauth2.test.ts
+++ b/packages/core/src/code_assist/oauth2.test.ts
@@ -57,6 +57,8 @@ describe('oauth2', () => {
     fs.rmSync(tempHomeDir, { recursive: true, force: true });
     vi.clearAllMocks();
     delete process.env.CLOUD_SHELL;
+    delete process.env.GOOGLE_GENAI_USE_GCA;
+    delete process.env.GOOGLE_CLOUD_ACCESS_TOKEN;
   });
 
   it('should perform a web login', async () => {
@@ -332,4 +334,138 @@ describe('oauth2', () => {
       );
     });
   });
+
+  describe('with GCP environment variables', () => {
+    it('should use GOOGLE_CLOUD_ACCESS_TOKEN when GOOGLE_GENAI_USE_GCA is true', async () => {
+      process.env.GOOGLE_GENAI_USE_GCA = 'true';
+      process.env.GOOGLE_CLOUD_ACCESS_TOKEN = 'gcp-access-token';
+
+      const mockSetCredentials = vi.fn();
+      const mockGetAccessToken = vi
+        .fn()
+        .mockResolvedValue({ token: 'gcp-access-token' });
+      const mockOAuth2Client = {
+        setCredentials: mockSetCredentials,
+        getAccessToken: mockGetAccessToken,
+        on: vi.fn(),
+      } as unknown as OAuth2Client;
+      (OAuth2Client as unknown as Mock).mockImplementation(
+        () => mockOAuth2Client,
+      );
+
+      // Mock the UserInfo API response for fetchAndCacheUserInfo
+      (global.fetch as Mock).mockResolvedValue({
+        ok: true,
+        json: vi
+          .fn()
+          .mockResolvedValue({ email: '[email protected]' }),
+      } as unknown as Response);
+
+      const writeFileSpy = vi
+        .spyOn(fs.promises, 'writeFile')
+        .mockResolvedValue(undefined);
+
+      const client = await getOauthClient(
+        AuthType.LOGIN_WITH_GOOGLE,
+        mockConfig,
+      );
+
+      expect(client).toBe(mockOAuth2Client);
+      expect(mockSetCredentials).toHaveBeenCalledWith({
+        access_token: 'gcp-access-token',
+      });
+
+      // Verify fetchAndCacheUserInfo was effectively called
+      expect(mockGetAccessToken).toHaveBeenCalled();
+      expect(global.fetch).toHaveBeenCalledWith(
+        'https://www.googleapis.com/oauth2/v2/userinfo',
+        {
+          headers: {
+            Authorization: 'Bearer gcp-access-token',
+          },
+        },
+      );
+
+      // Verify Google Account was cached
+      const googleAccountPath = path.join(
+        tempHomeDir,
+        '.gemini',
+        'google_accounts.json',
+      );
+      expect(writeFileSpy).toHaveBeenCalledWith(
+        googleAccountPath,
+        JSON.stringify(
+          {
+            active: '[email protected]',
+            old: [],
+          },
+          null,
+          2,
+        ),
+        'utf-8',
+      );
+    });
+
+    it('should not use GCP token if GOOGLE_CLOUD_ACCESS_TOKEN is not set', async () => {
+      process.env.GOOGLE_GENAI_USE_GCA = 'true';
+
+      const mockSetCredentials = vi.fn();
+      const mockGetAccessToken = vi
+        .fn()
+        .mockResolvedValue({ token: 'cached-access-token' });
+      const mockGetTokenInfo = vi.fn().mockResolvedValue({});
+      const mockOAuth2Client = {
+        setCredentials: mockSetCredentials,
+        getAccessToken: mockGetAccessToken,
+        getTokenInfo: mockGetTokenInfo,
+        on: vi.fn(),
+      } as unknown as OAuth2Client;
+      (OAuth2Client as unknown as Mock).mockImplementation(
+        () => mockOAuth2Client,
+      );
+
+      // Make it fall through to cached credentials path
+      const cachedCreds = { refresh_token: 'cached-token' };
+      vi.spyOn(fs.promises, 'readFile').mockResolvedValue(
+        JSON.stringify(cachedCreds),
+      );
+
+      await getOauthClient(AuthType.LOGIN_WITH_GOOGLE, mockConfig);
+
+      // It should be called with the cached credentials, not the GCP access token.
+      expect(mockSetCredentials).toHaveBeenCalledTimes(1);
+      expect(mockSetCredentials).toHaveBeenCalledWith(cachedCreds);
+    });
+
+    it('should not use GCP token if GOOGLE_GENAI_USE_GCA is not set', async () => {
+      process.env.GOOGLE_CLOUD_ACCESS_TOKEN = 'gcp-access-token';
+
+      const mockSetCredentials = vi.fn();
+      const mockGetAccessToken = vi
+        .fn()
+        .mockResolvedValue({ token: 'cached-access-token' });
+      const mockGetTokenInfo = vi.fn().mockResolvedValue({});
+      const mockOAuth2Client = {
+        setCredentials: mockSetCredentials,
+        getAccessToken: mockGetAccessToken,
+        getTokenInfo: mockGetTokenInfo,
+        on: vi.fn(),
+      } as unknown as OAuth2Client;
+      (OAuth2Client as unknown as Mock).mockImplementation(
+        () => mockOAuth2Client,
+      );
+
+      // Make it fall through to cached credentials path
+      const cachedCreds = { refresh_token: 'cached-token' };
+      vi.spyOn(fs.promises, 'readFile').mockResolvedValue(
+        JSON.stringify(cachedCreds),
+      );
+
+      await getOauthClient(AuthType.LOGIN_WITH_GOOGLE, mockConfig);
+
+      // It should be called with the cached credentials, not the GCP access token.
+      expect(mockSetCredentials).toHaveBeenCalledTimes(1);
+      expect(mockSetCredentials).toHaveBeenCalledWith(cachedCreds);
+    });
+  });
 });
diff --git a/packages/core/src/code_assist/oauth2.ts b/packages/core/src/code_assist/oauth2.ts
index 5958625a..f1046416 100644
--- a/packages/core/src/code_assist/oauth2.ts
+++ b/packages/core/src/code_assist/oauth2.ts
@@ -78,6 +78,17 @@ export async function getOauthClient(
     },
   });
 
+  if (
+    process.env.GOOGLE_GENAI_USE_GCA &&
+    process.env.GOOGLE_CLOUD_ACCESS_TOKEN
+  ) {
+    client.setCredentials({
+      access_token: process.env.GOOGLE_CLOUD_ACCESS_TOKEN,
+    });
+    await fetchAndCacheUserInfo(client);
+    return client;
+  }
+
   client.on('tokens', async (tokens: Credentials) => {
     await cacheCredentials(tokens);
   });