Feat(security) - Make the OAuthTokenStorage non static (#7716)

Co-authored-by: Shi Shu <shii@google.com>
This commit is contained in:
shishu314
2025-09-04 16:42:47 -04:00
committed by GitHub
parent e088c06a9a
commit 35a841f71a
7 changed files with 188 additions and 149 deletions

View File

@@ -22,17 +22,18 @@ import { Type } from '@google/genai';
vi.mock('@google/gemini-cli-core', async (importOriginal) => {
const actual =
await importOriginal<typeof import('@google/gemini-cli-core')>();
const mockAuthenticate = vi.fn();
return {
...actual,
getMCPServerStatus: vi.fn(),
getMCPDiscoveryState: vi.fn(),
MCPOAuthProvider: {
authenticate: vi.fn(),
},
MCPOAuthTokenStorage: {
MCPOAuthProvider: vi.fn(() => ({
authenticate: mockAuthenticate,
})),
MCPOAuthTokenStorage: vi.fn(() => ({
getToken: vi.fn(),
isTokenExpired: vi.fn(),
},
})),
};
});
@@ -892,13 +893,14 @@ describe('mcpCommand', () => {
context.ui.reloadCommands = vi.fn();
const { MCPOAuthProvider } = await import('@google/gemini-cli-core');
const mockAuthProvider = new MCPOAuthProvider();
const authCommand = mcpCommand.subCommands?.find(
(cmd) => cmd.name === 'auth',
);
const result = await authCommand!.action!(context, 'test-server');
expect(MCPOAuthProvider.authenticate).toHaveBeenCalledWith(
expect(mockAuthProvider.authenticate).toHaveBeenCalledWith(
'test-server',
{ enabled: true },
'http://localhost:3000',
@@ -928,9 +930,10 @@ describe('mcpCommand', () => {
});
const { MCPOAuthProvider } = await import('@google/gemini-cli-core');
(
MCPOAuthProvider.authenticate as ReturnType<typeof vi.fn>
).mockRejectedValue(new Error('Auth failed'));
const mockAuthProvider = new MCPOAuthProvider();
vi.mocked(mockAuthProvider.authenticate).mockRejectedValue(
new Error('Auth failed'),
);
const authCommand = mcpCommand.subCommands?.find(
(cmd) => cmd.name === 'auth',

View File

@@ -20,6 +20,7 @@ import {
MCPServerStatus,
mcpServerRequiresOAuth,
getErrorMessage,
MCPOAuthTokenStorage,
} from '@google/gemini-cli-core';
const COLOR_GREEN = '\u001b[32m';
@@ -141,9 +142,10 @@ const getMcpStatus = async (
const { MCPOAuthTokenStorage } = await import(
'@google/gemini-cli-core'
);
const hasToken = await MCPOAuthTokenStorage.getToken(serverName);
const tokenStorage = new MCPOAuthTokenStorage();
const hasToken = await tokenStorage.getToken(serverName);
if (hasToken) {
const isExpired = MCPOAuthTokenStorage.isTokenExpired(hasToken.token);
const isExpired = tokenStorage.isTokenExpired(hasToken.token);
if (isExpired) {
message += ` ${COLOR_YELLOW}(OAuth token expired)${RESET_COLOR}`;
} else {
@@ -385,11 +387,8 @@ const authCommand: SlashCommand = {
// Pass the MCP server URL for OAuth discovery
const mcpServerUrl = server.httpUrl || server.url;
await MCPOAuthProvider.authenticate(
serverName,
oauthConfig,
mcpServerUrl,
);
const authProvider = new MCPOAuthProvider(new MCPOAuthTokenStorage());
await authProvider.authenticate(serverName, oauthConfig, mcpServerUrl);
context.ui.addItem(
{