diff options
Diffstat (limited to 'packages/core/src/tools')
| -rw-r--r-- | packages/core/src/tools/mcp-client.test.ts | 59 | ||||
| -rw-r--r-- | packages/core/src/tools/mcp-client.ts | 26 |
2 files changed, 84 insertions, 1 deletions
diff --git a/packages/core/src/tools/mcp-client.test.ts b/packages/core/src/tools/mcp-client.test.ts index fbd2a2d4..428c9d2d 100644 --- a/packages/core/src/tools/mcp-client.test.ts +++ b/packages/core/src/tools/mcp-client.test.ts @@ -16,6 +16,8 @@ import { SSEClientTransport } from '@modelcontextprotocol/sdk/client/sse.js'; import * as SdkClientStdioLib from '@modelcontextprotocol/sdk/client/stdio.js'; import * as ClientLib from '@modelcontextprotocol/sdk/client/index.js'; import * as GenAiLib from '@google/genai'; +import { GoogleCredentialProvider } from '../mcp/google-auth-provider.js'; +import { AuthProviderType } from '../config/config.js'; vi.mock('@modelcontextprotocol/sdk/client/stdio.js'); vi.mock('@modelcontextprotocol/sdk/client/index.js'); @@ -173,6 +175,63 @@ describe('mcp-client', () => { stderr: 'pipe', }); }); + + describe('useGoogleCredentialProvider', () => { + it('should use GoogleCredentialProvider when specified', async () => { + const transport = await createTransport( + 'test-server', + { + httpUrl: 'http://test-server', + authProviderType: AuthProviderType.GOOGLE_CREDENTIALS, + oauth: { + scopes: ['scope1'], + }, + }, + false, + ); + + expect(transport).toBeInstanceOf(StreamableHTTPClientTransport); + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const authProvider = (transport as any)._authProvider; + expect(authProvider).toBeInstanceOf(GoogleCredentialProvider); + }); + + it('should use GoogleCredentialProvider with SSE transport', async () => { + const transport = await createTransport( + 'test-server', + { + url: 'http://test-server', + authProviderType: AuthProviderType.GOOGLE_CREDENTIALS, + oauth: { + scopes: ['scope1'], + }, + }, + false, + ); + + expect(transport).toBeInstanceOf(SSEClientTransport); + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const authProvider = (transport as any)._authProvider; + expect(authProvider).toBeInstanceOf(GoogleCredentialProvider); + }); + + it('should throw an error if no URL is provided with GoogleCredentialProvider', async () => { + await expect( + createTransport( + 'test-server', + { + authProviderType: AuthProviderType.GOOGLE_CREDENTIALS, + oauth: { + scopes: ['scope1'], + }, + }, + false, + ), + ).rejects.toThrow( + 'No URL configured for Google Credentials MCP server', + ); + }); + }); }); describe('isEnabled', () => { const funcDecl = { name: 'myTool' }; diff --git a/packages/core/src/tools/mcp-client.ts b/packages/core/src/tools/mcp-client.ts index 457259e5..3c482100 100644 --- a/packages/core/src/tools/mcp-client.ts +++ b/packages/core/src/tools/mcp-client.ts @@ -16,7 +16,8 @@ import { StreamableHTTPClientTransportOptions, } from '@modelcontextprotocol/sdk/client/streamableHttp.js'; import { parse } from 'shell-quote'; -import { MCPServerConfig } from '../config/config.js'; +import { AuthProviderType, MCPServerConfig } from '../config/config.js'; +import { GoogleCredentialProvider } from '../mcp/google-auth-provider.js'; import { DiscoveredMCPTool } from './mcp-tool.js'; import { FunctionDeclaration, mcpToTool } from '@google/genai'; import { ToolRegistry } from './tool-registry.js'; @@ -852,6 +853,29 @@ export async function createTransport( mcpServerConfig: MCPServerConfig, debugMode: boolean, ): Promise<Transport> { + if ( + mcpServerConfig.authProviderType === AuthProviderType.GOOGLE_CREDENTIALS + ) { + const provider = new GoogleCredentialProvider(mcpServerConfig); + const transportOptions: + | StreamableHTTPClientTransportOptions + | SSEClientTransportOptions = { + authProvider: provider, + }; + if (mcpServerConfig.httpUrl) { + return new StreamableHTTPClientTransport( + new URL(mcpServerConfig.httpUrl), + transportOptions, + ); + } else if (mcpServerConfig.url) { + return new SSEClientTransport( + new URL(mcpServerConfig.url), + transportOptions, + ); + } + throw new Error('No URL configured for Google Credentials MCP server'); + } + // Check if we have OAuth configuration or stored tokens let accessToken: string | null = null; let hasOAuthConfig = mcpServerConfig.oauth?.enabled; |
