summaryrefslogtreecommitdiff
path: root/packages/core/src/tools
diff options
context:
space:
mode:
Diffstat (limited to 'packages/core/src/tools')
-rw-r--r--packages/core/src/tools/mcp-client.test.ts59
-rw-r--r--packages/core/src/tools/mcp-client.ts26
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;