diff --git a/integration-tests/mcp-sampling.test.ts b/integration-tests/mcp-sampling.test.ts new file mode 100644 index 00000000000..4edf852401a --- /dev/null +++ b/integration-tests/mcp-sampling.test.ts @@ -0,0 +1,252 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * This test verifies MCP (Model Context Protocol) server integration. + * It uses a minimal MCP server implementation that doesn't require + * external dependencies, making it compatible with Docker sandbox mode. + */ + +import { describe, it, beforeAll, expect } from 'vitest'; +import { TestRig, poll, validateModelOutput } from './test-helper.js'; +import { join } from 'node:path'; +import { writeFileSync } from 'node:fs'; + +const serverScript = `#!/usr/bin/env node +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +const readline = require('readline'); +const fs = require('fs'); + +const debugEnabled = process.env['MCP_DEBUG'] === 'true' || process.env['VERBOSE'] === 'true'; +function debug(msg) { + if (debugEnabled) { + fs.writeSync(2, '[MCP-DEBUG] ' + msg + '\n'); + } +} + +debug('MCP sampling server starting...'); + +class SimpleJSONRPC { + constructor() { + this.handlers = new Map(); + this.rl = readline.createInterface({ + input: process.stdin, + output: process.stdout, + terminal: false + }); + + this.rl.on('line', (line) => { + debug('Received line: ' + line); + try { + const message = JSON.parse(line); + debug('Parsed message: ' + JSON.stringify(message)); + this.handleMessage(message); + } catch (e) { + debug('Parse error: ' + e.message); + } + }); + } + + send(message) { + const msgStr = JSON.stringify(message); + debug('Sending message: ' + msgStr); + process.stdout.write(msgStr + '\n'); + } + + request(method, params) { + return new Promise((resolve, reject) => { + const id = Math.random().toString(36).substring(2); + this.handlers.set(id, (response) => { + if (response.error) { + reject(new Error(response.error.message)); + } + else { + resolve(response.result); + } + }); + this.send({ + jsonrpc: '2.0', + id, + method, + params + }); + }); + } + + async handleMessage(message) { + if (message.id && this.handlers.has(message.id)) { + this.handlers.get(message.id)(message); + this.handlers.delete(message.id); + return; + } + + if (message.method && this.handlers.has(message.method)) { + try { + const result = await this.handlers.get(message.method)(message.params || {}); + if (message.id !== undefined) { + this.send({ + jsonrpc: '2.0', + id: message.id, + result + }); + } + } + catch (error) { + if (message.id !== undefined) { + this.send({ + jsonrpc: '2.0', + id: message.id, + error: { + code: -32603, + message: error.message + } + }); + } + } + } + else if (message.id !== undefined) { + this.send({ + jsonrpc: '2.0', + id: message.id, + error: { + code: -32601, + message: 'Method not found' + } + }); + } + } + + on(method, handler) { + this.handlers.set(method, handler); + } +} + +const rpc = new SimpleJSONRPC(); + +rpc.on('initialize', async (params) => { + debug('Handling initialize request'); + return { + protocolVersion: '2024-11-05', + capabilities: { + sampling: true, + }, + serverInfo: { + name: 'sampler-server', + version: '1.0.0' + } + }; +}); + +rpc.on('tools/list', async () => { + debug('Handling tools/list request'); + return { + tools: [{ + name: 'sample', + description: 'Uses the LLM to sample a response to a prompt.', + inputSchema: { + type: 'object', + properties: { + prompt: { type: 'string', description: 'The prompt to sample' }, + }, + required: ['prompt'] + } + }] + }; +}); + +rpc.on('tools/call', async (params) => { + debug('Handling tools/call request for tool: ' + params.name); + if (params.name === 'sample') { + const { prompt } = params.arguments; + + debug('Requesting sampling from client...'); + // MCP spec: messages at top level, content is a single object (not array) + const result = await rpc.request('sampling/createMessage', { + messages: [ + { + role: 'user', + content: { type: 'text', text: prompt } + } + ] + }); + debug('Received sampling result: ' + JSON.stringify(result)); + + return result; + } + throw new Error('Unknown tool: ' + params.name); +}); + +rpc.send({ + jsonrpc: '2.0', + method: 'initialized' +}); +`; + +describe('mcp-sampling', () => { + const rig = new TestRig(); + + beforeAll(async () => { + await rig.setup('mcp-sampling', { + settings: { + mcpServers: { + 'sampler-server': { + command: 'node', + args: ['mcp-server.cjs'], + }, + }, + }, + }); + + const testServerPath = join(rig.testDir!, 'mcp-server.cjs'); + writeFileSync(testServerPath, serverScript); + + if (process.platform !== 'win32') { + const { chmodSync } = await import('node:fs'); + chmodSync(testServerPath, 0o755); + } + + const { accessSync, constants } = await import('node:fs'); + const isReady = await poll( + () => { + try { + accessSync(testServerPath, constants.F_OK); + return true; + } catch { + return false; + } + }, + 5000, + 100, + ); + + if (!isReady) { + throw new Error('MCP server script was not ready in time.'); + } + }); + + it('should use the sample tool and get a response', async () => { + const child = rig.run( + "Use the sample tool to ask 'what is the capital of France?'", + ['--auto-confirm-mcp-sampling'], + ); + + const foundToolCall = await rig.waitForToolCall('sample'); + expect(foundToolCall, 'Expected to find a sample tool call').toBeTruthy(); + + const output = await child; + + validateModelOutput(output, 'Paris', 'MCP sampling test'); + expect( + output.includes('Paris'), + 'Expected output to contain the capital of France (Paris)', + ).toBeTruthy(); + }); +}); diff --git a/package-lock.json b/package-lock.json index 8ef2b61688c..c4a5085a999 100644 --- a/package-lock.json +++ b/package-lock.json @@ -2474,7 +2474,6 @@ "integrity": "sha512-t54CUOsFMappY1Jbzb7fetWeO0n6K0k/4+/ZpkS+3Joz8I4VcvY9OiEBFRYISqaI2fq5sCiPtAjRDOzVYG8m+Q==", "dev": true, "license": "MIT", - "peer": true, "dependencies": { "@octokit/auth-token": "^6.0.0", "@octokit/graphql": "^9.0.2", @@ -2655,7 +2654,6 @@ "resolved": "https://registry.npmjs.org/@opentelemetry/api/-/api-1.9.0.tgz", "integrity": "sha512-3giAOQvZiH5F9bMlMiv8+GSPMeqg0dbaeo58/0SlA9sxSqZhnUtxzX9/2FzyhS9sWQf5S0GJE0AKBrFqjpeYcg==", "license": "Apache-2.0", - "peer": true, "engines": { "node": ">=8.0.0" } @@ -2689,7 +2687,6 @@ "resolved": "https://registry.npmjs.org/@opentelemetry/core/-/core-2.0.1.tgz", "integrity": "sha512-MaZk9SJIDgo1peKevlbhP6+IwIiNPNmswNL4AF0WaQJLbHXjr9SrZMgS12+iqr9ToV4ZVosCcc0f8Rg67LXjxw==", "license": "Apache-2.0", - "peer": true, "dependencies": { "@opentelemetry/semantic-conventions": "^1.29.0" }, @@ -3058,7 +3055,6 @@ "resolved": "https://registry.npmjs.org/@opentelemetry/resources/-/resources-2.0.1.tgz", "integrity": "sha512-dZOB3R6zvBwDKnHDTB4X1xtMArB/d324VsbiPkX/Yu0Q8T2xceRthoIVFhJdvgVM2QhGVUyX9tzwiNxGtoBJUw==", "license": "Apache-2.0", - "peer": true, "dependencies": { "@opentelemetry/core": "2.0.1", "@opentelemetry/semantic-conventions": "^1.29.0" @@ -3092,7 +3088,6 @@ "resolved": "https://registry.npmjs.org/@opentelemetry/sdk-metrics/-/sdk-metrics-2.0.1.tgz", "integrity": "sha512-wf8OaJoSnujMAHWR3g+/hGvNcsC16rf9s1So4JlMiFaFHiE4HpIA3oUh+uWZQ7CNuK8gVW/pQSkgoa5HkkOl0g==", "license": "Apache-2.0", - "peer": true, "dependencies": { "@opentelemetry/core": "2.0.1", "@opentelemetry/resources": "2.0.1" @@ -3145,7 +3140,6 @@ "resolved": "https://registry.npmjs.org/@opentelemetry/sdk-trace-base/-/sdk-trace-base-2.0.1.tgz", "integrity": "sha512-xYLlvk/xdScGx1aEqvxLwf6sXQLXCjk3/1SQT9X9AoN5rXRhkdvIFShuNNmtTEPRBqcsMbS4p/gJLNI2wXaDuQ==", "license": "Apache-2.0", - "peer": true, "dependencies": { "@opentelemetry/core": "2.0.1", "@opentelemetry/resources": "2.0.1", @@ -4358,7 +4352,6 @@ "integrity": "sha512-6mDvHUFSjyT2B2yeNx2nUgMxh9LtOWvkhIU3uePn2I2oyNymUAX1NIsdgviM4CH+JSrp2D2hsMvJOkxY+0wNRA==", "devOptional": true, "license": "MIT", - "peer": true, "dependencies": { "csstype": "^3.0.2" } @@ -4636,7 +4629,6 @@ "integrity": "sha512-6sMvZePQrnZH2/cJkwRpkT7DxoAWh+g6+GFRK6bV3YQo7ogi3SX5rgF6099r5Q53Ma5qeT7LGmOmuIutF4t3lA==", "dev": true, "license": "MIT", - "peer": true, "dependencies": { "@typescript-eslint/scope-manager": "8.35.0", "@typescript-eslint/types": "8.35.0", @@ -5641,7 +5633,6 @@ "resolved": "https://registry.npmjs.org/acorn/-/acorn-8.15.0.tgz", "integrity": "sha512-NZyJarBfL7nWwIq+FDL6Zp/yHEhePMNnnJ0y3qfieCrmNvYct8uvtiV41UvlSe6apAfk0fY1FbWx+NwfmpvtTg==", "license": "MIT", - "peer": true, "bin": { "acorn": "bin/acorn" }, @@ -6086,7 +6077,8 @@ "version": "1.1.1", "resolved": "https://registry.npmjs.org/array-flatten/-/array-flatten-1.1.1.tgz", "integrity": "sha512-PCVAQswWemu6UdxsDFFX/+gVeYqKAod3D3UVm91jHwynguOwAvYPhx8nNlM++NqRcK6CxxpUafjmhIdKiHibqg==", - "license": "MIT" + "license": "MIT", + "peer": true }, "node_modules/array-includes": { "version": "3.1.9", @@ -7370,6 +7362,7 @@ "resolved": "https://registry.npmjs.org/content-disposition/-/content-disposition-0.5.4.tgz", "integrity": "sha512-FveZTNuGw04cxlAiWbzi6zTAL/lhehaWbTtgluJh4/E95DqMwTmha3KZN1aAWA8cFIhHzMZUvLevkw5Rqk+tSQ==", "license": "MIT", + "peer": true, "dependencies": { "safe-buffer": "5.2.1" }, @@ -8689,7 +8682,6 @@ "integrity": "sha512-GsGizj2Y1rCWDu6XoEekL3RLilp0voSePurjZIkxL3wlm5o5EC9VpgaP7lrCvjnkuLvzFBQWB3vWB3K5KQTveQ==", "dev": true, "license": "MIT", - "peer": true, "dependencies": { "@eslint-community/eslint-utils": "^4.2.0", "@eslint-community/regexpp": "^4.12.1", @@ -9292,6 +9284,7 @@ "resolved": "https://registry.npmjs.org/cookie/-/cookie-0.7.1.tgz", "integrity": "sha512-6DnInpx7SJ2AK3+CTUE/ZM0vWTUboZCegxhC2xiIydHR9jNuTAASBrfEpHhiGOZw/nX51bHt6YQl8jsGo4y/0w==", "license": "MIT", + "peer": true, "engines": { "node": ">= 0.6" } @@ -9301,6 +9294,7 @@ "resolved": "https://registry.npmjs.org/debug/-/debug-2.6.9.tgz", "integrity": "sha512-bC7ElrdJaJnPbAP+1EotYvqZsb3ecl5wi6Bfi6BJTUcNowp6cvspg0jXznRTKDjm/E7AdgFBVeAPVMNcKGsHMA==", "license": "MIT", + "peer": true, "dependencies": { "ms": "2.0.0" } @@ -9310,6 +9304,7 @@ "resolved": "https://registry.npmjs.org/statuses/-/statuses-2.0.1.tgz", "integrity": "sha512-RwNA9Z/7PrK06rYLIzFMlaF+l73iwpzsqRIFgbMLbTcLD6cOao82TaWefPXQvB2fOC4AjuYSEndS7N/mTCbkdQ==", "license": "MIT", + "peer": true, "engines": { "node": ">= 0.8" } @@ -9563,6 +9558,7 @@ "resolved": "https://registry.npmjs.org/finalhandler/-/finalhandler-1.3.1.tgz", "integrity": "sha512-6BN9trH7bp3qvnrRyzsBz+g3lZxTNZTbVO2EV1CS0WIcDbawYVdYvGflME/9QP0h0pYlCDBCTjYa9nZzMDpyxQ==", "license": "MIT", + "peer": true, "dependencies": { "debug": "2.6.9", "encodeurl": "~2.0.0", @@ -9581,6 +9577,7 @@ "resolved": "https://registry.npmjs.org/debug/-/debug-2.6.9.tgz", "integrity": "sha512-bC7ElrdJaJnPbAP+1EotYvqZsb3ecl5wi6Bfi6BJTUcNowp6cvspg0jXznRTKDjm/E7AdgFBVeAPVMNcKGsHMA==", "license": "MIT", + "peer": true, "dependencies": { "ms": "2.0.0" } @@ -9589,13 +9586,15 @@ "version": "2.0.0", "resolved": "https://registry.npmjs.org/ms/-/ms-2.0.0.tgz", "integrity": "sha512-Tpp60P6IUJDTuOq/5Z8cdskzJujfwqfOTkrwIwj7IRISpnkJnT6SyJ4PCPnGMoFjC9ddhal5KVIYtAt97ix05A==", - "license": "MIT" + "license": "MIT", + "peer": true }, "node_modules/finalhandler/node_modules/statuses": { "version": "2.0.1", "resolved": "https://registry.npmjs.org/statuses/-/statuses-2.0.1.tgz", "integrity": "sha512-RwNA9Z/7PrK06rYLIzFMlaF+l73iwpzsqRIFgbMLbTcLD6cOao82TaWefPXQvB2fOC4AjuYSEndS7N/mTCbkdQ==", "license": "MIT", + "peer": true, "engines": { "node": ">= 0.8" } @@ -10878,7 +10877,6 @@ "resolved": "https://registry.npmjs.org/@jrichman/ink/-/ink-6.4.7.tgz", "integrity": "sha512-QHyxhNF5VonF5cRmdAJD/UPucB9nRx3FozWMjQrDGfBxfAL9lpyu72/MlFPgloS1TMTGsOt7YN6dTPPA6mh0Aw==", "license": "MIT", - "peer": true, "dependencies": { "@alcalzone/ansi-tokenize": "^0.2.1", "ansi-escapes": "^7.0.0", @@ -14063,7 +14061,8 @@ "version": "0.1.12", "resolved": "https://registry.npmjs.org/path-to-regexp/-/path-to-regexp-0.1.12.tgz", "integrity": "sha512-RA1GjUVMnvYFxuqovrEqZoxxW5NUZqbwKtYz/Tt7nXerk0LbLblQmrsgdeOxV5SFHf0UDggjS/bSeOZwt1pmEQ==", - "license": "MIT" + "license": "MIT", + "peer": true }, "node_modules/path-type": { "version": "3.0.0", @@ -14640,7 +14639,6 @@ "resolved": "https://registry.npmjs.org/react/-/react-19.2.0.tgz", "integrity": "sha512-tmbWg6W31tQLeB5cdIBOicJDJRR2KzXsV7uSK9iNfLWQ5bIZfxuPEHp7M8wiHyHnn0DD1i7w3Zmin0FtkrwoCQ==", "license": "MIT", - "peer": true, "engines": { "node": ">=0.10.0" } @@ -14651,7 +14649,6 @@ "integrity": "sha512-ePrwPfxAnB+7hgnEr8vpKxL9cmnp7F322t8oqcPshbIQQhDKgFDW4tjhF2wjVbdXF9O/nyuy3sQWd9JGpiLPvA==", "devOptional": true, "license": "MIT", - "peer": true, "dependencies": { "shell-quote": "^1.6.1", "ws": "^7" @@ -16911,7 +16908,6 @@ "resolved": "https://registry.npmjs.org/picomatch/-/picomatch-4.0.3.tgz", "integrity": "sha512-5gTmgEY/sqK6gFXLIsQNH19lWb4ebPDLA4SdLP7dsWkIXHWlG66oPuVvXSGFPppYZz8ZDZq0dYYrbHfBCVUb1Q==", "license": "MIT", - "peer": true, "engines": { "node": ">=12" }, @@ -17135,8 +17131,7 @@ "resolved": "https://registry.npmjs.org/tslib/-/tslib-2.8.1.tgz", "integrity": "sha512-oJFu94HQb+KVduSUQL7wnpmqnfmLsOA/nAh6b6EH0wCEoK0/mPeXU6c3wKDV83MkOuHPRHtSXKKU99IBazS/2w==", "dev": true, - "license": "0BSD", - "peer": true + "license": "0BSD" }, "node_modules/tsx": { "version": "4.20.3", @@ -17144,7 +17139,6 @@ "integrity": "sha512-qjbnuR9Tr+FJOMBqJCW5ehvIo/buZq7vH7qD7JziU98h6l3qGy0a/yPFjwO+y0/T7GFpNgNAvEcPPVfyT8rrPQ==", "devOptional": true, "license": "MIT", - "peer": true, "dependencies": { "esbuild": "~0.25.0", "get-tsconfig": "^4.7.5" @@ -17328,7 +17322,6 @@ "integrity": "sha512-p1diW6TqL9L07nNxvRMM7hMMw4c5XOo/1ibL4aAIGmSAt9slTE1Xgw5KWuof2uTOvCg9BY7ZRi+GaF+7sfgPeQ==", "devOptional": true, "license": "Apache-2.0", - "peer": true, "bin": { "tsc": "bin/tsc", "tsserver": "bin/tsserver" @@ -17491,6 +17484,7 @@ "resolved": "https://registry.npmjs.org/utils-merge/-/utils-merge-1.0.1.tgz", "integrity": "sha512-pMZTvIkT1d+TFGvDOqodOclx0QWkkgi6Tdoa8gC8ffGAAqz9pzPTZWAybbsHHoED/ztMtkv/VoYTYyShUn81hA==", "license": "MIT", + "peer": true, "engines": { "node": ">= 0.4.0" } @@ -17545,7 +17539,6 @@ "resolved": "https://registry.npmjs.org/vite/-/vite-7.2.2.tgz", "integrity": "sha512-BxAKBWmIbrDgrokdGZH1IgkIk/5mMHDreLDmCJ0qpyJaAteP8NvMhkwr/ZCQNqNH97bw/dANTE9PDzqwJghfMQ==", "license": "MIT", - "peer": true, "dependencies": { "esbuild": "^0.25.0", "fdir": "^6.5.0", @@ -17659,7 +17652,6 @@ "resolved": "https://registry.npmjs.org/picomatch/-/picomatch-4.0.3.tgz", "integrity": "sha512-5gTmgEY/sqK6gFXLIsQNH19lWb4ebPDLA4SdLP7dsWkIXHWlG66oPuVvXSGFPppYZz8ZDZq0dYYrbHfBCVUb1Q==", "license": "MIT", - "peer": true, "engines": { "node": ">=12" }, @@ -17672,7 +17664,6 @@ "resolved": "https://registry.npmjs.org/vitest/-/vitest-3.2.4.tgz", "integrity": "sha512-LUCP5ev3GURDysTWiP47wRRUpLKMOfPh+yKTx3kVIEiu5KOMeqzpnYNsKyOoVrULivR8tLcks4+lga33Whn90A==", "license": "MIT", - "peer": true, "dependencies": { "@types/chai": "^5.2.2", "@vitest/expect": "3.2.4", @@ -18377,7 +18368,6 @@ "resolved": "https://registry.npmjs.org/zod/-/zod-3.25.76.tgz", "integrity": "sha512-gzUt/qt81nXsFGKIFcC3YnfEAx5NkunCfnDlvuBSSFS02bcXu4Lmea0AFIUwbLWxWPx3d9p8S5QoaujKcNQxcQ==", "license": "MIT", - "peer": true, "funding": { "url": "https://github.com/sponsors/colinhacks" } @@ -18944,7 +18934,6 @@ "resolved": "https://registry.npmjs.org/picomatch/-/picomatch-4.0.3.tgz", "integrity": "sha512-5gTmgEY/sqK6gFXLIsQNH19lWb4ebPDLA4SdLP7dsWkIXHWlG66oPuVvXSGFPppYZz8ZDZq0dYYrbHfBCVUb1Q==", "license": "MIT", - "peer": true, "engines": { "node": ">=12" }, diff --git a/packages/cli/src/config/config.ts b/packages/cli/src/config/config.ts index c5c61ce748d..7f2c22beffe 100755 --- a/packages/cli/src/config/config.ts +++ b/packages/cli/src/config/config.ts @@ -83,6 +83,7 @@ export interface CliArgs { outputFormat: string | undefined; fakeResponses: string | undefined; recordResponses: string | undefined; + autoConfirmMcpSampling: boolean | undefined; } export async function parseArguments( @@ -248,6 +249,11 @@ export async function parseArguments( type: 'string', description: 'Path to a file to record model responses for testing.', hidden: true, + }) + .option('auto-confirm-mcp-sampling', { + type: 'boolean', + description: 'Automatically confirm all MCP sampling requests.', + hidden: true, }), ) // Register MCP subcommands @@ -756,6 +762,7 @@ export async function loadCliConfig( ptyInfo: ptyInfo?.name, disableLLMCorrection: settings.tools?.disableLLMCorrection, modelConfigServiceConfig: settings.modelConfigs, + autoConfirmMcpSampling: argv.autoConfirmMcpSampling, // TODO: loading of hooks based on workspace trust enableHooks: (settings.tools?.enableHooks ?? true) && diff --git a/packages/cli/src/gemini.test.tsx b/packages/cli/src/gemini.test.tsx index 896f89e3c82..7b0a93d8eb5 100644 --- a/packages/cli/src/gemini.test.tsx +++ b/packages/cli/src/gemini.test.tsx @@ -490,6 +490,7 @@ describe('gemini.tsx main function kitty protocol', () => { outputFormat: undefined, fakeResponses: undefined, recordResponses: undefined, + autoConfirmMcpSampling: undefined, }); await act(async () => { diff --git a/packages/cli/src/gemini.tsx b/packages/cli/src/gemini.tsx index da6643349b5..33c9795a317 100644 --- a/packages/cli/src/gemini.tsx +++ b/packages/cli/src/gemini.tsx @@ -5,7 +5,7 @@ */ import React from 'react'; -import { render } from 'ink'; +import { render, type RenderOptions } from 'ink'; import { AppContainer } from './ui/AppContainer.js'; import { loadCliConfig, parseArguments } from './config/config.js'; import * as cliConfig from './config/config.js'; @@ -265,7 +265,7 @@ export async function startInteractiveUI( alternateBuffer: useAlternateBuffer, incrementalRendering: settings.merged.ui.incrementalRendering !== false && useAlternateBuffer, - }, + } as RenderOptions & { incrementalRendering?: boolean }, ); checkForUpdates(settings) diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index 72f74a76f0a..1ceedefd7f5 100644 --- a/packages/cli/src/ui/AppContainer.tsx +++ b/packages/cli/src/ui/AppContainer.tsx @@ -252,6 +252,13 @@ export const AppContainer = (props: AppContainerProps) => { setPermissionsDialogProps(null); }, []); + const [mcpSamplingRequest, setMcpSamplingRequest] = useState<{ + serverName: string; + prompt: unknown; + resolve: () => void; + reject: (reason?: unknown) => void; + } | null>(null); + const toggleDebugProfiler = useCallback( () => setShowDebugProfiler((prev) => !prev), [], @@ -354,9 +361,33 @@ export const AppContainer = (props: AppContainerProps) => { setCurrentModel(config.getModel()); }; + const handleMcpSamplingRequest = (payload: { + serverName: string; + prompt: unknown; + resolve: () => void; + reject: (reason?: unknown) => void; + }) => { + // Wrap the resolve and reject to clear the state after calling them + const wrappedResolve = () => { + payload.resolve(); + setMcpSamplingRequest(null); + }; + const wrappedReject = (reason?: unknown) => { + payload.reject(reason); + setMcpSamplingRequest(null); + }; + setMcpSamplingRequest({ + ...payload, + resolve: wrappedResolve, + reject: wrappedReject, + }); + }; + coreEvents.on(CoreEvent.ModelChanged, handleModelChanged); + coreEvents.on(CoreEvent.McpSamplingRequest, handleMcpSamplingRequest); return () => { coreEvents.off(CoreEvent.ModelChanged, handleModelChanged); + coreEvents.off(CoreEvent.McpSamplingRequest, handleMcpSamplingRequest); }; }, [config]); @@ -1473,6 +1504,7 @@ Logging in with Google... Restarting Gemini CLI to continue. !!proQuotaRequest || isSessionBrowserOpen || isAuthDialogOpen || + !!mcpSamplingRequest || authState === AuthState.AwaitingApiKeyInput; const pendingHistoryItems = useMemo( @@ -1612,8 +1644,10 @@ Logging in with Google... Restarting Gemini CLI to continue. activePtyId, embeddedShellFocused, showDebugProfiler, + authState, customDialog, copyModeEnabled, + mcpSamplingRequest, warningMessage, bannerData, bannerVisible, @@ -1708,6 +1742,7 @@ Logging in with Google... Restarting Gemini CLI to continue. apiKeyDefaultValue, authState, copyModeEnabled, + mcpSamplingRequest, warningMessage, bannerData, bannerVisible, diff --git a/packages/cli/src/ui/components/ConsentPrompt.tsx b/packages/cli/src/ui/components/ConsentPrompt.tsx index efa6b136a3e..6eb806d2f8e 100644 --- a/packages/cli/src/ui/components/ConsentPrompt.tsx +++ b/packages/cli/src/ui/components/ConsentPrompt.tsx @@ -9,30 +9,32 @@ import { type ReactNode } from 'react'; import { theme } from '../semantic-colors.js'; import { MarkdownDisplay } from '../utils/MarkdownDisplay.js'; import { RadioButtonSelect } from './shared/RadioButtonSelect.js'; +import { Scrollable } from './shared/Scrollable.js'; type ConsentPromptProps = { // If a simple string is given, it will render using markdown by default. prompt: ReactNode; onConfirm: (value: boolean) => void; terminalWidth: number; + availableTerminalHeight?: number; }; export const ConsentPrompt = (props: ConsentPromptProps) => { - const { prompt, onConfirm, terminalWidth } = props; + const { prompt, onConfirm, terminalWidth, availableTerminalHeight } = props; - return ( - + // Account for border (2) + paddingY (2) + marginTop (1) + RadioButtonSelect height (~3) + const scrollableHeight = availableTerminalHeight + ? availableTerminalHeight - 8 + : undefined; + + const content = ( + <> {typeof prompt === 'string' ? ( ) : ( prompt @@ -46,6 +48,24 @@ export const ConsentPrompt = (props: ConsentPromptProps) => { onSelect={onConfirm} /> + + ); + + return ( + + {scrollableHeight ? ( + + {content} + + ) : ( + content + )} ); }; diff --git a/packages/cli/src/ui/components/DialogManager.tsx b/packages/cli/src/ui/components/DialogManager.tsx index f915bc78525..995cc4d248f 100644 --- a/packages/cli/src/ui/components/DialogManager.tsx +++ b/packages/cli/src/ui/components/DialogManager.tsx @@ -22,6 +22,7 @@ import { RELAUNCH_EXIT_CODE } from '../../utils/processUtils.js'; import { SessionBrowser } from './SessionBrowser.js'; import { PermissionsModifyTrustDialog } from './PermissionsModifyTrustDialog.js'; import { ModelDialog } from './ModelDialog.js'; +import { McpSamplingDialog } from './McpSamplingDialog.js'; import { theme } from '../semantic-colors.js'; import { useUIState } from '../contexts/UIStateContext.js'; import { useUIActions } from '../contexts/UIActionsContext.js'; @@ -31,6 +32,7 @@ import process from 'node:process'; import { type UseHistoryManagerReturn } from '../hooks/useHistoryManager.js'; import { AdminSettingsChangedDialog } from './AdminSettingsChangedDialog.js'; import { IdeTrustChangeDialog } from './IdeTrustChangeDialog.js'; +import type { SamplingMessage } from '@modelcontextprotocol/sdk/types.js'; interface DialogManagerProps { addItem: UseHistoryManagerReturn['addItem']; @@ -97,6 +99,7 @@ export const DialogManager = ({ prompt={uiState.confirmationRequest.prompt} onConfirm={uiState.confirmationRequest.onConfirm} terminalWidth={terminalWidth} + availableTerminalHeight={terminalHeight - staticExtraHeight} /> ); } @@ -107,6 +110,7 @@ export const DialogManager = ({ prompt={request.prompt} onConfirm={request.onConfirm} terminalWidth={terminalWidth} + availableTerminalHeight={terminalHeight - staticExtraHeight} /> ); } @@ -231,5 +235,17 @@ export const DialogManager = ({ ); } + if (uiState.mcpSamplingRequest) { + return ( + + ); + } + return null; }; diff --git a/packages/cli/src/ui/components/McpSamplingDialog.tsx b/packages/cli/src/ui/components/McpSamplingDialog.tsx new file mode 100644 index 00000000000..791d6d8e51c --- /dev/null +++ b/packages/cli/src/ui/components/McpSamplingDialog.tsx @@ -0,0 +1,96 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { Box, Text } from 'ink'; +import type { SamplingMessage } from '@modelcontextprotocol/sdk/types.js'; +import { RadioButtonSelect } from './shared/RadioButtonSelect.js'; +import { Scrollable } from './shared/Scrollable.js'; +import { theme } from '../semantic-colors.js'; + +/** + * Extract text content from an MCP message content object. + */ +function getMessageText(content: SamplingMessage['content']): string { + if (typeof content === 'object' && content !== null) { + if ('text' in content && typeof content.text === 'string') { + return content.text; + } + if ('type' in content) { + // For non-text content types (image, audio), show the type + return `[${content.type} content]`; + } + } + return JSON.stringify(content); +} + +export function McpSamplingDialog({ + serverName, + prompt, + onConfirm, + onReject, + availableTerminalHeight, +}: { + serverName: string; + prompt: SamplingMessage[]; + onConfirm: () => void; + onReject: () => void; + availableTerminalHeight?: number; +}) { + // Account for border (2) + padding (2) + title (1) + margin (1) + + // prompt border (2) + prompt padding (2) + margin (1) + buttons (2) + const scrollableHeight = availableTerminalHeight + ? availableTerminalHeight - 13 + : undefined; + + const content = ( + + {prompt.map((message, index) => ( + + + {message.role === 'user' ? 'User' : 'Assistant'}: + + + {getMessageText(message.content)} + + + ))} + + ); + + return ( + + + + MCP server {serverName} wants to run a prompt: + + + + {scrollableHeight ? ( + + {content} + + ) : ( + content + )} + + + { + if (value) { + onConfirm(); + } else { + onReject(); + } + }} + /> + + + ); +} diff --git a/packages/cli/src/ui/components/shared/Scrollable.test.tsx b/packages/cli/src/ui/components/shared/Scrollable.test.tsx index 22c2055f49b..9c19c7303e0 100644 --- a/packages/cli/src/ui/components/shared/Scrollable.test.tsx +++ b/packages/cli/src/ui/components/shared/Scrollable.test.tsx @@ -101,20 +101,20 @@ describe('', () => { throw new Error('capturedEntry is undefined'); } - // Initial state (starts at bottom because of auto-scroll logic) - expect(capturedEntry.getScrollState().scrollTop).toBe(5); + // Initial state (starts at top on initial render) + expect(capturedEntry.getScrollState().scrollTop).toBe(0); - // Call scrollBy multiple times (upwards) in the same tick + // Call scrollBy multiple times (downwards) in the same tick act(() => { - capturedEntry!.scrollBy(-1); - capturedEntry!.scrollBy(-1); + capturedEntry!.scrollBy(1); + capturedEntry!.scrollBy(1); }); - // Should have moved up by 2 - expect(capturedEntry.getScrollState().scrollTop).toBe(3); + // Should have moved down by 2 + expect(capturedEntry.getScrollState().scrollTop).toBe(2); act(() => { - capturedEntry!.scrollBy(-2); + capturedEntry!.scrollBy(2); }); - expect(capturedEntry.getScrollState().scrollTop).toBe(1); + expect(capturedEntry.getScrollState().scrollTop).toBe(4); }); }); diff --git a/packages/cli/src/ui/components/shared/Scrollable.tsx b/packages/cli/src/ui/components/shared/Scrollable.tsx index 16436be7c6f..f77346fa030 100644 --- a/packages/cli/src/ui/components/shared/Scrollable.tsx +++ b/packages/cli/src/ui/components/shared/Scrollable.tsx @@ -51,6 +51,7 @@ export const Scrollable: React.FC = ({ }, [size]); const childrenCountRef = useRef(0); + const hasInitializedRef = useRef(false); // This effect needs to run on every render to correctly measure the container // and scroll to the bottom if new children are added. The if conditions @@ -70,7 +71,9 @@ export const Scrollable: React.FC = ({ size.scrollHeight !== scrollHeight ) { setSize({ innerHeight, scrollHeight }); - if (isAtBottom) { + // Only maintain scroll position at bottom after initial render. + // On initial render, start at top unless scrollToBottom is explicitly set. + if (isAtBottom && hasInitializedRef.current) { setScrollTop(Math.max(0, scrollHeight - innerHeight)); } } @@ -80,6 +83,7 @@ export const Scrollable: React.FC = ({ setScrollTop(Math.max(0, scrollHeight - innerHeight)); } childrenCountRef.current = childCountCurrent; + hasInitializedRef.current = true; }); const { getScrollTop, setPendingScrollTop } = useBatchedScroll(scrollTop); diff --git a/packages/cli/src/ui/contexts/ScrollProvider.test.tsx b/packages/cli/src/ui/contexts/ScrollProvider.test.tsx index 021e10e280a..41b96822da9 100644 --- a/packages/cli/src/ui/contexts/ScrollProvider.test.tsx +++ b/packages/cli/src/ui/contexts/ScrollProvider.test.tsx @@ -276,9 +276,9 @@ describe('ScrollProvider', () => { // Advance timers to trigger the batched update await vi.runAllTimersAsync(); - // Should have called scrollBy once with accumulated delta (3) + // Should have called scrollBy once with accumulated delta (3 events * 3 lines = 9) expect(scrollBy).toHaveBeenCalledTimes(1); - expect(scrollBy).toHaveBeenCalledWith(3); + expect(scrollBy).toHaveBeenCalledWith(9); }); it('handles mixed direction scroll events in batch', async () => { @@ -335,7 +335,7 @@ describe('ScrollProvider', () => { await vi.runAllTimersAsync(); expect(scrollBy).toHaveBeenCalledTimes(1); - expect(scrollBy).toHaveBeenCalledWith(1); // 1 + 1 - 1 = 1 + expect(scrollBy).toHaveBeenCalledWith(3); // 3 + 3 - 3 = 3 (3 lines per scroll event) }); it('respects scroll limits during batching', async () => { @@ -390,13 +390,14 @@ describe('ScrollProvider', () => { await vi.runAllTimersAsync(); - // Should have accumulated only 1, because subsequent scrolls would be blocked - // Actually, the logic in ScrollProvider uses effectiveScrollTop to check bounds. - // scrollTop=89, max=90. - // 1st scroll: pending=1, effective=90. Allowed. - // 2nd scroll: pending=1, effective=90. canScrollDown checks effective < 90. 90 < 90 is false. Blocked. + // Should have accumulated only 3 (one scroll event), because subsequent scrolls would be blocked + // The logic in ScrollProvider uses effectiveScrollTop to check bounds. + // scrollTop=89, max=90, scroll step=3. + // 1st scroll: pending=3, effective=92. canScrollDown checks 89 < 89.999. True, allowed. + // 2nd scroll: pending=6, effective=95. canScrollDown checks 92 < 89.999. False. Blocked. + // Note: The actual clamping to max happens in the Scrollable component's scrollBy. expect(scrollBy).toHaveBeenCalledTimes(1); - expect(scrollBy).toHaveBeenCalledWith(1); + expect(scrollBy).toHaveBeenCalledWith(3); }); it('calls scrollTo when dragging scrollbar thumb if available', async () => { diff --git a/packages/cli/src/ui/contexts/ScrollProvider.tsx b/packages/cli/src/ui/contexts/ScrollProvider.tsx index b461622fb2f..041b5c5eeb4 100644 --- a/packages/cli/src/ui/contexts/ScrollProvider.tsx +++ b/packages/cli/src/ui/contexts/ScrollProvider.tsx @@ -80,7 +80,17 @@ export const ScrollProvider: React.FC<{ children: React.ReactNode }> = ({ ); const register = useCallback((entry: ScrollableEntry) => { - setScrollables((prev) => new Map(prev).set(entry.id, entry)); + setScrollables((prev) => { + // Check if entry already exists with same id - if so, just update the reference + // without creating a new Map to avoid unnecessary re-renders + const existing = prev.get(entry.id); + if (existing) { + // Update the existing entry in place to avoid Map recreation + prev.set(entry.id, entry); + return prev; + } + return new Map(prev).set(entry.id, entry); + }); }, []); const unregister = useCallback((id: string) => { @@ -125,40 +135,47 @@ export const ScrollProvider: React.FC<{ children: React.ReactNode }> = ({ } }, []); - const handleScroll = (direction: 'up' | 'down', mouseEvent: MouseEvent) => { - const delta = direction === 'up' ? -1 : 1; - const candidates = findScrollableCandidates( - mouseEvent, - scrollablesRef.current, - ); - - for (const candidate of candidates) { - const { scrollTop, scrollHeight, innerHeight } = - candidate.getScrollState(); - const pendingDelta = pendingScrollsRef.current.get(candidate.id) || 0; - const effectiveScrollTop = scrollTop + pendingDelta; - - // Epsilon to handle floating point inaccuracies. - const canScrollUp = effectiveScrollTop > 0.001; - const canScrollDown = - effectiveScrollTop < scrollHeight - innerHeight - 0.001; - - if (direction === 'up' && canScrollUp) { - pendingScrollsRef.current.set(candidate.id, pendingDelta + delta); - scheduleFlush(); - return true; - } + // Scroll 3 lines per scroll event for smoother trackpad/mouse wheel scrolling + const SCROLL_LINES_PER_EVENT = 3; + + const handleScroll = useCallback( + (direction: 'up' | 'down', mouseEvent: MouseEvent) => { + const delta = + direction === 'up' ? -SCROLL_LINES_PER_EVENT : SCROLL_LINES_PER_EVENT; + const candidates = findScrollableCandidates( + mouseEvent, + scrollablesRef.current, + ); + + for (const candidate of candidates) { + const { scrollTop, scrollHeight, innerHeight } = + candidate.getScrollState(); + const pendingDelta = pendingScrollsRef.current.get(candidate.id) || 0; + const effectiveScrollTop = scrollTop + pendingDelta; + + // Epsilon to handle floating point inaccuracies. + const canScrollUp = effectiveScrollTop > 0.001; + const canScrollDown = + effectiveScrollTop < scrollHeight - innerHeight - 0.001; + + if (direction === 'up' && canScrollUp) { + pendingScrollsRef.current.set(candidate.id, pendingDelta + delta); + scheduleFlush(); + return true; + } - if (direction === 'down' && canScrollDown) { - pendingScrollsRef.current.set(candidate.id, pendingDelta + delta); - scheduleFlush(); - return true; + if (direction === 'down' && canScrollDown) { + pendingScrollsRef.current.set(candidate.id, pendingDelta + delta); + scheduleFlush(); + return true; + } } - } - return false; - }; + return false; + }, + [scheduleFlush], + ); - const handleLeftPress = (mouseEvent: MouseEvent) => { + const handleLeftPress = useCallback((mouseEvent: MouseEvent) => { // Check for scrollbar interaction first for (const entry of scrollablesRef.current.values()) { if (!entry.ref.current || !entry.hasFocus()) { @@ -258,9 +275,9 @@ export const ScrollProvider: React.FC<{ children: React.ReactNode }> = ({ return false; } return false; - }; + }, []); - const handleMove = (mouseEvent: MouseEvent) => { + const handleMove = useCallback((mouseEvent: MouseEvent) => { const state = dragStateRef.current; if (!state.active || !state.id) return false; @@ -303,9 +320,9 @@ export const ScrollProvider: React.FC<{ children: React.ReactNode }> = ({ entry.scrollBy(targetScrollTop - scrollTop); } return true; - }; + }, []); - const handleLeftRelease = () => { + const handleLeftRelease = useCallback(() => { if (dragStateRef.current.active) { dragStateRef.current = { active: false, @@ -315,9 +332,9 @@ export const ScrollProvider: React.FC<{ children: React.ReactNode }> = ({ return true; } return false; - }; + }, []); - useMouse( + const mouseHandler = useCallback( (event: MouseEvent) => { if (event.name === 'scroll-up') { return handleScroll('up', event); @@ -332,9 +349,11 @@ export const ScrollProvider: React.FC<{ children: React.ReactNode }> = ({ } return false; }, - { isActive: true }, + [handleScroll, handleLeftPress, handleMove, handleLeftRelease], ); + useMouse(mouseHandler, { isActive: true }); + const contextValue = useMemo( () => ({ register, unregister }), [register, unregister], diff --git a/packages/cli/src/ui/contexts/UIStateContext.tsx b/packages/cli/src/ui/contexts/UIStateContext.tsx index 80db5782ffe..9b9632af68d 100644 --- a/packages/cli/src/ui/contexts/UIStateContext.tsx +++ b/packages/cli/src/ui/contexts/UIStateContext.tsx @@ -13,6 +13,7 @@ import type { LoopDetectionConfirmationRequest, HistoryItemWithoutId, StreamingState, + AuthState, ActiveHook, } from '../types.js'; import type { CommandContext, SlashCommand } from '../commands/types.js'; @@ -38,6 +39,13 @@ export interface ProQuotaDialogRequest { resolve: (intent: FallbackIntent) => void; } +export interface McpSamplingRequest { + serverName: string; + prompt: unknown; + resolve: () => void; + reject: (reason?: unknown) => void; +} + import { type UseHistoryManagerReturn } from '../hooks/useHistoryManager.js'; import { type RestartReason } from '../hooks/useIdeTrustListener.js'; import type { TerminalBackgroundColor } from '../utils/terminalCapabilityManager.js'; @@ -130,6 +138,7 @@ export interface UIState { showDebugProfiler: boolean; showFullTodos: boolean; copyModeEnabled: boolean; + mcpSamplingRequest: McpSamplingRequest | null; warningMessage: string | null; bannerData: { defaultText: string; @@ -137,6 +146,7 @@ export interface UIState { }; bannerVisible: boolean; customDialog: React.ReactNode | null; + authState: AuthState; terminalBackgroundColor: TerminalBackgroundColor; settingsNonce: number; adminSettingsChanged: boolean; diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index c7d5f66a151..c3e7c910d8a 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -106,6 +106,7 @@ vi.mock('../core/contentGenerator.js'); vi.mock('../core/client.js', () => ({ GeminiClient: vi.fn().mockImplementation(() => ({ initialize: vi.fn().mockResolvedValue(undefined), + setTools: vi.fn().mockResolvedValue(undefined), stripThoughtsFromHistory: vi.fn(), isInitialized: vi.fn().mockReturnValue(false), })), diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 80b1dece501..1db1cf99d83 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -370,6 +370,7 @@ export interface ConfigParameters { ptyInfo?: string; disableYoloMode?: boolean; modelConfigServiceConfig?: ModelConfigServiceConfig; + autoConfirmMcpSampling?: boolean; enableHooks?: boolean; enableHooksUI?: boolean; experiments?: Experiments; @@ -509,6 +510,7 @@ export class Config { readonly fakeResponses?: string; readonly recordResponses?: string; private readonly disableYoloMode: boolean; + private readonly autoConfirmMcpSampling: boolean; private pendingIncludeDirectories: string[]; private readonly enableHooks: boolean; private readonly enableHooksUI: boolean; @@ -709,6 +711,7 @@ export class Config { }; this.retryFetchErrors = params.retryFetchErrors ?? false; this.disableYoloMode = params.disableYoloMode ?? false; + this.autoConfirmMcpSampling = params.autoConfirmMcpSampling ?? false; if (params.hooks) { const { disabled: _, ...restOfHooks } = params.hooks; @@ -717,7 +720,6 @@ export class Config { if (params.projectHooks) { this.projectHooks = params.projectHooks; } - this.experiments = params.experiments; this.onModelChange = params.onModelChange; this.onReload = params.onReload; @@ -844,6 +846,11 @@ export class Config { } await this.geminiClient.initialize(); + // Ensure any MCP tools discovered during async initialization are loaded. + // MCP servers start asynchronously above (not awaited), so tools may have + // been discovered before geminiClient.initialize() completed. This call + // ensures the LLM has access to all discovered tools. + await this.geminiClient.setTools(); } getContentGenerator(): ContentGenerator { @@ -1081,6 +1088,10 @@ export class Config { return this.sandbox; } + getAutoConfirmMcpSampling(): boolean { + return this.autoConfirmMcpSampling; + } + isRestrictiveSandbox(): boolean { const sandboxConfig = this.getSandbox(); const seatbeltProfile = process.env['SEATBELT_PROFILE']; diff --git a/packages/core/src/mcp/oauth-provider.ts b/packages/core/src/mcp/oauth-provider.ts index b79ec693a30..8e89e11cad3 100644 --- a/packages/core/src/mcp/oauth-provider.ts +++ b/packages/core/src/mcp/oauth-provider.ts @@ -463,17 +463,25 @@ export class MCPOAuthProvider { // Add resource parameter for MCP OAuth spec compliance // Only add if we have an MCP server URL (indicates MCP OAuth flow, not standard OAuth) + debugLogger.debug( + `buildAuthorizationUrl: mcpServerUrl=${mcpServerUrl}, adding resource parameter: ${!!mcpServerUrl}`, + ); if (mcpServerUrl) { try { - params.append( - 'resource', - OAuthUtils.buildResourceParameter(mcpServerUrl), + const resourceParam = OAuthUtils.buildResourceParameter(mcpServerUrl); + params.append('resource', resourceParam); + debugLogger.debug( + `Added resource parameter to authorization URL: ${resourceParam}`, ); } catch (error) { debugLogger.warn( `Could not add resource parameter: ${getErrorMessage(error)}`, ); } + } else { + debugLogger.debug( + 'No mcpServerUrl provided, skipping resource parameter', + ); } const url = new URL(config.authorizationUrl!); @@ -521,18 +529,26 @@ export class MCPOAuthProvider { // Add resource parameter for MCP OAuth spec compliance // Only add if we have an MCP server URL (indicates MCP OAuth flow, not standard OAuth) + debugLogger.debug( + `exchangeCodeForToken: mcpServerUrl=${mcpServerUrl}, adding resource parameter: ${!!mcpServerUrl}`, + ); if (mcpServerUrl) { const resourceUrl = mcpServerUrl; try { - params.append( - 'resource', - OAuthUtils.buildResourceParameter(resourceUrl), + const resourceParam = OAuthUtils.buildResourceParameter(resourceUrl); + params.append('resource', resourceParam); + debugLogger.debug( + `Added resource parameter to token exchange: ${resourceParam}`, ); } catch (error) { debugLogger.warn( `Could not add resource parameter: ${getErrorMessage(error)}`, ); } + } else { + debugLogger.debug( + 'No mcpServerUrl provided to token exchange, skipping resource parameter', + ); } const response = await fetch(config.tokenUrl!, { @@ -644,17 +660,25 @@ export class MCPOAuthProvider { // Add resource parameter for MCP OAuth spec compliance // Only add if we have an MCP server URL (indicates MCP OAuth flow, not standard OAuth) + debugLogger.debug( + `refreshAccessToken: mcpServerUrl=${mcpServerUrl}, adding resource parameter: ${!!mcpServerUrl}`, + ); if (mcpServerUrl) { try { - params.append( - 'resource', - OAuthUtils.buildResourceParameter(mcpServerUrl), + const resourceParam = OAuthUtils.buildResourceParameter(mcpServerUrl); + params.append('resource', resourceParam); + debugLogger.debug( + `Added resource parameter to refresh request: ${resourceParam}`, ); } catch (error) { debugLogger.warn( `Could not add resource parameter: ${getErrorMessage(error)}`, ); } + } else { + debugLogger.debug( + 'No mcpServerUrl provided to refresh, skipping resource parameter', + ); } const response = await fetch(tokenUrl, { @@ -746,6 +770,10 @@ export class MCPOAuthProvider { mcpServerUrl?: string, events?: EventEmitter, ): Promise { + debugLogger.debug( + `authenticate called for server '${serverName}' with mcpServerUrl: ${mcpServerUrl}`, + ); + // Helper function to display messages through handler or fallback to console.log const displayMessage = (message: string) => { if (events) { diff --git a/packages/core/src/tools/mcp-client.test.ts b/packages/core/src/tools/mcp-client.test.ts index a448fd288b8..a407eab28d7 100644 --- a/packages/core/src/tools/mcp-client.test.ts +++ b/packages/core/src/tools/mcp-client.test.ts @@ -1608,6 +1608,14 @@ describe('connectToMcpServer with OAuth', () => { let testWorkspace: string; let mockAuthProvider: MCPOAuthProvider; let mockTokenStorage: MCPOAuthTokenStorage; + const mockConfig = { + getAutoConfirmMcpSampling: () => false, + getGeminiClient: () => ({ + generateContent: vi.fn(), + }), + getModel: () => 'gemini-2.0-flash', + sanitizationConfig: EMPTY_CONFIG, + } as unknown as Config; beforeEach(() => { mockedClient = { @@ -1679,7 +1687,7 @@ describe('connectToMcpServer with OAuth', () => { { httpUrl: serverUrl, oauth: { enabled: true } }, false, workspaceContext, - EMPTY_CONFIG, + mockConfig, ); expect(client).toBe(mockedClient); @@ -1724,7 +1732,7 @@ describe('connectToMcpServer with OAuth', () => { { httpUrl: serverUrl, oauth: { enabled: true } }, false, workspaceContext, - EMPTY_CONFIG, + mockConfig, ); expect(client).toBe(mockedClient); @@ -1742,6 +1750,14 @@ describe('connectToMcpServer - HTTP→SSE fallback', () => { let mockedClient: ClientLib.Client; let workspaceContext: WorkspaceContext; let testWorkspace: string; + const mockConfig = { + getAutoConfirmMcpSampling: () => false, + getGeminiClient: () => ({ + generateContent: vi.fn(), + }), + getModel: () => 'gemini-2.0-flash', + sanitizationConfig: EMPTY_CONFIG, + } as unknown as Config; beforeEach(() => { mockedClient = { @@ -1779,7 +1795,7 @@ describe('connectToMcpServer - HTTP→SSE fallback', () => { { url: 'http://test-server', type: 'http' }, false, workspaceContext, - EMPTY_CONFIG, + mockConfig, ), ).rejects.toThrow('Connection failed'); @@ -1798,7 +1814,7 @@ describe('connectToMcpServer - HTTP→SSE fallback', () => { { url: 'http://test-server', type: 'sse' }, false, workspaceContext, - EMPTY_CONFIG, + mockConfig, ), ).rejects.toThrow('Connection failed'); @@ -1816,7 +1832,7 @@ describe('connectToMcpServer - HTTP→SSE fallback', () => { { url: 'http://test-server' }, false, workspaceContext, - EMPTY_CONFIG, + mockConfig, ); expect(client).toBe(mockedClient); @@ -1838,7 +1854,7 @@ describe('connectToMcpServer - HTTP→SSE fallback', () => { { url: 'http://test-server' }, false, workspaceContext, - EMPTY_CONFIG, + mockConfig, ), ).rejects.toThrow('Server error'); @@ -1855,7 +1871,7 @@ describe('connectToMcpServer - HTTP→SSE fallback', () => { { url: 'http://test-server' }, false, workspaceContext, - EMPTY_CONFIG, + mockConfig, ); expect(client).toBe(mockedClient); @@ -1869,6 +1885,14 @@ describe('connectToMcpServer - OAuth with transport fallback', () => { let testWorkspace: string; let mockAuthProvider: MCPOAuthProvider; let mockTokenStorage: MCPOAuthTokenStorage; + const mockConfig = { + getAutoConfirmMcpSampling: () => false, + getGeminiClient: () => ({ + generateContent: vi.fn(), + }), + getModel: () => 'gemini-2.0-flash', + sanitizationConfig: EMPTY_CONFIG, + } as unknown as Config; beforeEach(() => { mockedClient = { @@ -1925,7 +1949,7 @@ describe('connectToMcpServer - OAuth with transport fallback', () => { { url: 'http://test-server', oauth: { enabled: true } }, false, workspaceContext, - EMPTY_CONFIG, + mockConfig, ); expect(client).toBe(mockedClient); diff --git a/packages/core/src/tools/mcp-client.ts b/packages/core/src/tools/mcp-client.ts index 1f96d34169a..250bbebc385 100644 --- a/packages/core/src/tools/mcp-client.ts +++ b/packages/core/src/tools/mcp-client.ts @@ -24,6 +24,7 @@ import type { Resource, } from '@modelcontextprotocol/sdk/types.js'; import { + CreateMessageRequestSchema, ListResourcesResultSchema, ListRootsRequestSchema, ReadResourceResultSchema, @@ -34,6 +35,10 @@ import { import { parse } from 'shell-quote'; import type { Config, MCPServerConfig } from '../config/config.js'; import { AuthProviderType } from '../config/config.js'; +import { + DEFAULT_GEMINI_FLASH_MODEL, + DEFAULT_GEMINI_MODEL_AUTO, +} from '../config/models.js'; import { GoogleCredentialProvider } from '../mcp/google-auth-provider.js'; import { ServiceAccountImpersonationProvider } from '../mcp/sa-impersonation-provider.js'; import { DiscoveredMCPTool } from './mcp-tool.js'; @@ -58,7 +63,7 @@ import type { import type { ToolRegistry } from './tool-registry.js'; import { debugLogger } from '../utils/debugLogger.js'; import { type MessageBus } from '../confirmation-bus/message-bus.js'; -import { coreEvents } from '../utils/events.js'; +import { coreEvents, CoreEvent } from '../utils/events.js'; import type { ResourceRegistry } from '../resources/resource-registry.js'; import { sanitizeEnvironment, @@ -141,7 +146,7 @@ export class McpClient { this.serverConfig, this.debugMode, this.workspaceContext, - this.cliConfig.sanitizationConfig, + this.cliConfig, ); this.registerNotificationHandlers(); @@ -825,7 +830,7 @@ export async function connectAndDiscover( mcpServerConfig, debugMode, workspaceContext, - cliConfig.sanitizationConfig, + cliConfig, ); mcpClient.onerror = (error) => { @@ -1327,7 +1332,7 @@ async function retryWithOAuth( * * @param mcpServerName The name of the MCP server, used for logging and identification. * @param mcpServerConfig The configuration specifying how to connect to the server. - * @returns A promise that resolves to a connected MCP `Client` instance. + * @returns A promise that resolves to a connected MCP `Client` instance and its `Transport`. * @throws An error if the connection fails or the configuration is invalid. */ export async function connectToMcpServer( @@ -1335,7 +1340,7 @@ export async function connectToMcpServer( mcpServerConfig: MCPServerConfig, debugMode: boolean, workspaceContext: WorkspaceContext, - sanitizationConfig: EnvironmentSanitizationConfig, + cliConfig: Config, ): Promise { const mcpClient = new Client( { @@ -1352,6 +1357,7 @@ export async function connectToMcpServer( roots: { listChanged: true, }, + sampling: {}, }); mcpClient.setRequestHandler(ListRootsRequestSchema, async () => { @@ -1367,6 +1373,170 @@ export async function connectToMcpServer( }; }); + // Timeout for sampling consent dialog (5 minutes) + const SAMPLING_CONSENT_TIMEOUT_MS = 5 * 60 * 1000; + + mcpClient.setRequestHandler(CreateMessageRequestSchema, async (req) => { + const autoConfirm = cliConfig.getAutoConfirmMcpSampling(); + + let rejectSampling: ((reason?: unknown) => void) | undefined; + + if (!autoConfirm) { + const consentPromise = new Promise((resolve, reject) => { + const timeoutId = setTimeout(() => { + reject(new Error('Sampling consent request timed out')); + }, SAMPLING_CONSENT_TIMEOUT_MS); + + rejectSampling = (reason) => { + clearTimeout(timeoutId); + reject(reason); + }; + + coreEvents.emit(CoreEvent.McpSamplingRequest, { + serverName: mcpServerName, + prompt: req.params.messages, + resolve: () => { + clearTimeout(timeoutId); + resolve(); + }, + reject: (_reason?: unknown) => { + clearTimeout(timeoutId); + reject(new Error('User rejected sampling request')); + }, + }); + }); + + await consentPromise; + } + + try { + const geminiClient = cliConfig.getGeminiClient(); + const contents = req.params.messages.map((message) => { + // MCP spec: message.content is a single object (not an array) + const content = message.content as { + type: string; + text?: string; + data?: string; + mimeType?: string; + }; + + let parts; + if (content.type === 'text' && content.text) { + parts = [{ text: content.text }]; + } else if ( + content.type === 'image' && + content.data && + content.mimeType + ) { + // For image content, convert to Gemini format + parts = [ + { + inlineData: { + mimeType: content.mimeType, + data: content.data, + }, + }, + ]; + } else if ( + content.type === 'audio' && + content.data && + content.mimeType + ) { + // For audio content, convert to Gemini format + parts = [ + { + inlineData: { + mimeType: content.mimeType, + data: content.data, + }, + }, + ]; + } else { + throw new Error( + `Unsupported or invalid content type: ${content.type}`, + ); + } + + // Map MCP roles to Gemini roles + // MCP: 'user' | 'assistant' + // Gemini: 'user' | 'model' + const geminiRole = + message.role === 'assistant' ? 'model' : message.role; + + return { + role: geminiRole, + parts, + }; + }); + + // Resolve the model to use for sampling + // If the config model is "auto", use flash as it's faster and cheaper for sampling + let modelToUse = cliConfig.getModel(); + if (modelToUse === DEFAULT_GEMINI_MODEL_AUTO) { + modelToUse = DEFAULT_GEMINI_FLASH_MODEL; + } + + // TODO: Consider req.params.modelPreferences to select model based on server hints + // For now, we just use the resolved model from config + + const result = await geminiClient.generateContent( + { model: modelToUse }, + contents, + new AbortController().signal, + ); + + const firstCandidate = result.candidates?.[0]; + if ( + !firstCandidate || + !firstCandidate.content || + !firstCandidate.content.parts + ) { + throw new Error('No response from Gemini'); + } + + // MCP spec: response content should be a single object (not an array) + // Gemini can return multiple parts, so we'll concatenate text parts + const textParts: string[] = []; + for (const part of firstCandidate.content.parts) { + if ('text' in part && part.text) { + textParts.push(part.text); + } else { + throw new Error( + 'Unsupported response part type - only text parts are supported', + ); + } + } + + const responseContent = { + type: 'text' as const, + text: textParts.join(''), + }; + + // Map Gemini's finish reason to MCP's stopReason + // Gemini: STOP, MAX_TOKENS, SAFETY, RECITATION, OTHER, BLOCKLIST, PROHIBITED_CONTENT, SPII + // MCP: endTurn, maxTokens, stopSequence (we'll map to endTurn or maxTokens) + let stopReason: 'endTurn' | 'maxTokens' | 'stopSequence' = 'endTurn'; + if (firstCandidate.finishReason === 'MAX_TOKENS') { + stopReason = 'maxTokens'; + } + + // MCP spec: return role, content, model, and stopReason at top level (not wrapped in "message") + return { + role: 'assistant' as const, + content: responseContent, + model: modelToUse, + stopReason, + }; + } catch (error) { + // Close the dialog with an error message if consent was required + if (rejectSampling) { + rejectSampling(error); + } + // Re-throw to propagate to MCP server + throw error; + } + }); + let unlistenDirectories: Unsubscribe | undefined = workspaceContext.onDirectoriesChanged(async () => { try { @@ -1402,7 +1572,7 @@ export async function connectToMcpServer( mcpServerName, mcpServerConfig, debugMode, - sanitizationConfig, + cliConfig.sanitizationConfig, ); try { await mcpClient.connect(transport, { @@ -1772,7 +1942,20 @@ export async function createTransport( } } if (accessToken) { - headers['Authorization'] = `Bearer ${accessToken}`; + // Only use the stored token if no Authorization header is manually configured + // or if OAuth is explicitly enabled (which implies we should use the token) + const hasConfiguredAuth = Object.keys( + mcpServerConfig.headers || {}, + ).some((k) => k.toLowerCase() === 'authorization'); + const isOAuthEnabled = !!mcpServerConfig.oauth?.enabled; + + if (!hasConfiguredAuth || isOAuthEnabled) { + headers['Authorization'] = `Bearer ${accessToken}`; + } else { + debugLogger.log( + `Ignoring stored OAuth token for server '${mcpServerName}' because an Authorization header is explicitly configured.`, + ); + } } } diff --git a/packages/core/src/utils/events.ts b/packages/core/src/utils/events.ts index 79e440e9ad7..086124d1b86 100644 --- a/packages/core/src/utils/events.ts +++ b/packages/core/src/utils/events.ts @@ -108,12 +108,23 @@ export interface RetryAttemptPayload { model: string; } +/** + * Payload for the 'mcp-sampling-request' event. + */ +export interface McpSamplingRequestPayload { + serverName: string; + prompt: unknown; + resolve: () => void; + reject: (reason?: unknown) => void; +} + export enum CoreEvent { UserFeedback = 'user-feedback', ModelChanged = 'model-changed', ConsoleLog = 'console-log', Output = 'output', MemoryChanged = 'memory-changed', + McpSamplingRequest = 'mcp-sampling-request', ExternalEditorClosed = 'external-editor-closed', SettingsChanged = 'settings-changed', HookStart = 'hook-start', @@ -129,6 +140,7 @@ export interface CoreEvents { [CoreEvent.ConsoleLog]: [ConsoleLogPayload]; [CoreEvent.Output]: [OutputPayload]; [CoreEvent.MemoryChanged]: [MemoryChangedPayload]; + [CoreEvent.McpSamplingRequest]: [McpSamplingRequestPayload]; [CoreEvent.ExternalEditorClosed]: never[]; [CoreEvent.SettingsChanged]: never[]; [CoreEvent.HookStart]: [HookStartPayload];