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

This commit is contained in:
Gal Zahavi
2025-07-25 10:19:38 -07:00
committed by GitHub
parent fb0db2dfd6
commit 6321442865
5 changed files with 200 additions and 8 deletions

View File

@@ -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: 'test-gcp-account@gmail.com' }),
} 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: 'test-gcp-account@gmail.com',
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);
});
});
});

View File

@@ -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);
});