feat(mcp): Inject GoogleCredentialProvider headers in McpClient (#13783)

This commit is contained in:
sai-sunder-s
2025-11-26 12:08:19 -08:00
committed by GitHub
parent 3406dc5b2e
commit 0f12d6c426
7 changed files with 184 additions and 8 deletions
@@ -36,6 +36,8 @@ vi.mock('@google/genai');
vi.mock('../mcp/oauth-provider.js');
vi.mock('../mcp/oauth-token-storage.js');
vi.mock('../mcp/oauth-utils.js');
vi.mock('google-auth-library');
import { GoogleAuth } from 'google-auth-library';
vi.mock('../utils/events.js', () => ({
coreEvents: {
@@ -578,6 +580,16 @@ describe('mcp-client', () => {
});
describe('useGoogleCredentialProvider', () => {
beforeEach(() => {
// Mock GoogleAuth client
const mockClient = {
getAccessToken: vi.fn().mockResolvedValue({ token: 'test-token' }),
quotaProjectId: 'myproject',
};
GoogleAuth.prototype.getClient = vi.fn().mockResolvedValue(mockClient);
});
it('should use GoogleCredentialProvider when specified', async () => {
const transport = await createTransport(
'test-server',
@@ -605,6 +617,64 @@ describe('mcp-client', () => {
expect(googUserProject).toBe('myproject');
});
it('should use headers from GoogleCredentialProvider', async () => {
const mockGetRequestHeaders = vi.fn().mockResolvedValue({
'X-Goog-User-Project': 'provider-project',
});
vi.spyOn(
GoogleCredentialProvider.prototype,
'getRequestHeaders',
).mockImplementation(mockGetRequestHeaders);
const transport = await createTransport(
'test-server',
{
httpUrl: 'http://test.googleapis.com',
authProviderType: AuthProviderType.GOOGLE_CREDENTIALS,
oauth: {
scopes: ['scope1'],
},
},
false,
);
expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
expect(mockGetRequestHeaders).toHaveBeenCalled();
// eslint-disable-next-line @typescript-eslint/no-explicit-any
const headers = (transport as any)._requestInit?.headers;
expect(headers['X-Goog-User-Project']).toBe('provider-project');
});
it('should prioritize provider headers over config headers', async () => {
const mockGetRequestHeaders = vi.fn().mockResolvedValue({
'X-Goog-User-Project': 'provider-project',
});
vi.spyOn(
GoogleCredentialProvider.prototype,
'getRequestHeaders',
).mockImplementation(mockGetRequestHeaders);
const transport = await createTransport(
'test-server',
{
httpUrl: 'http://test.googleapis.com',
authProviderType: AuthProviderType.GOOGLE_CREDENTIALS,
oauth: {
scopes: ['scope1'],
},
headers: {
'X-Goog-User-Project': 'config-project',
},
},
false,
);
expect(transport).toBeInstanceOf(StreamableHTTPClientTransport);
// eslint-disable-next-line @typescript-eslint/no-explicit-any
const headers = (transport as any)._requestInit?.headers;
expect(headers['X-Goog-User-Project']).toBe('provider-project');
});
it('should use GoogleCredentialProvider with SSE transport', async () => {
const transport = await createTransport(
'test-server',
+6 -2
View File
@@ -35,6 +35,7 @@ import { DiscoveredMCPTool } from './mcp-tool.js';
import type { CallableTool, FunctionCall, Part, Tool } from '@google/genai';
import { basename } from 'node:path';
import { pathToFileURL } from 'node:url';
import type { McpAuthProvider } from '../mcp/auth-provider.js';
import { MCPOAuthProvider } from '../mcp/oauth-provider.js';
import { MCPOAuthTokenStorage } from '../mcp/oauth-token-storage.js';
import { OAuthUtils } from '../mcp/oauth-utils.js';
@@ -425,7 +426,9 @@ function createTransportRequestInit(
*
* @param mcpServerConfig The MCP server configuration
*/
function createAuthProvider(mcpServerConfig: MCPServerConfig) {
function createAuthProvider(
mcpServerConfig: MCPServerConfig,
): McpAuthProvider | undefined {
if (
mcpServerConfig.authProviderType ===
AuthProviderType.SERVICE_ACCOUNT_IMPERSONATION
@@ -1333,8 +1336,9 @@ export async function createTransport(
if (mcpServerConfig.httpUrl || mcpServerConfig.url) {
const authProvider = createAuthProvider(mcpServerConfig);
const headers: Record<string, string> =
(await authProvider?.getRequestHeaders?.()) ?? {};
const headers: Record<string, string> = {};
if (authProvider === undefined) {
// Check if we have OAuth configuration or stored tokens
let accessToken: string | null = null;