MCP client POC

Author: bywang56

package-lock.json

@@ -13,7 +13,8 @@
             "args": ["--extensionDevelopmentPath=${workspaceFolder}"],
             "env": {
                 "SSMDOCUMENT_LANGUAGESERVER_PORT": "6010",
-                "WEBPACK_DEVELOPER_SERVER": "http://localhost:8080"
+                "WEBPACK_DEVELOPER_SERVER": "http://localhost:8080",
+                "PATH": "/Users/bywang/.toolbox/bin:/usr/local/bin:/usr/bin:/bin:/usr/sbin:/sbin"
                 // "__AMAZONQLSP_PATH": "${workspaceFolder}/../../../language-servers/app/aws-lsp-codewhisperer-runtimes/out/token-standalone.js",
                 // "__AMAZONQLSP_UI": "${workspaceFolder}/../../../language-servers/chat-client/build/amazonq-ui.js"
             },

packages/amazonq/.vscode/launch.json

@@ -576,7 +576,8 @@
         "winston": "^3.11.0",
         "winston-transport": "^4.6.0",
         "xml2js": "^0.6.1",
-        "yaml-cfn": "^0.3.2"
+        "yaml-cfn": "^0.3.2",
+        "@modelcontextprotocol/sdk": "^1.9.0"
     },
     "overrides": {
         "webfont": {

packages/core/package.json

@@ -92,6 +92,7 @@ import {
     additionalContentInnerContextLimit,
     workspaceChunkMaxSize,
     defaultContextLengths,
+    tools,
 } from '../../constants'
 import { ChatSession } from '../../clients/chat/v0/chat'
 import { amazonQTabSuffix } from '../../../shared/constants'
@@ -102,6 +103,7 @@ import { tempDirPath } from '../../../shared/filesystemUtilities'
 import { Database } from '../../../shared/db/chatDb/chatDb'
 import { TabBarController } from './tabBarController'
 import { messageToChatMessage } from '../../../shared/db/chatDb/util'
+import { McpManager } from '../../tools/mcp/mcpManager'
 
 export interface ChatControllerMessagePublishers {
     readonly processPromptChatMessage: MessagePublisher<PromptMessage>
@@ -178,6 +180,7 @@ export class ChatController {
     private userPromptsWatcher: vscode.FileSystemWatcher | undefined
     private chatHistoryDb = Database.getInstance()
     private cancelTokenSource: vscode.CancellationTokenSource = new vscode.CancellationTokenSource()
+    private mcpManager: McpManager | undefined
 
     public constructor(
         private readonly chatControllerMessageListeners: ChatControllerMessageListeners,
@@ -197,6 +200,28 @@ export class ChatController {
         this.userIntentRecognizer = new UserIntentRecognizer()
         this.tabBarController = new TabBarController(this.messenger)
 
+        // MCP intitialzation
+        const mcpConfigPath = path.join(process.env.HOME ?? '', '.aws', 'amazonq', 'mcp.json')
+        McpManager.create(mcpConfigPath)
+            .then((manager) => {
+                this.mcpManager = manager
+                ToolUtils.mcpManager = manager
+                const discovered = manager.getAllMcpTools()
+                for (const def of discovered) {
+                    tools.push({
+                        toolSpecification: {
+                            name: def.toolName,
+                            description: def.description,
+                            inputSchema: { json: def.inputSchema },
+                        },
+                    })
+                }
+                getLogger().info(`MCP: successfully discovered ${discovered.length} new tools.`)
+            })
+            .catch((err) => {
+                getLogger().error(`Failed to init MCP manager: ${err}`)
+            })
+
         onDidChangeAmazonQVisibility((visible) => {
             if (visible) {
                 this.telemetryHelper.recordOpenChat()

packages/core/src/codewhispererChat/controllers/chat/controller.ts

@@ -0,0 +1,104 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { ListToolsResponse, MCPConfig, MCPServerConfig } from './mcpTypes'
+import { Client } from '@modelcontextprotocol/sdk/client/index.js'
+import { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js'
+import fs from '../../../shared/fs/fs'
+import { getLogger } from '../../../shared/logger/logger'
+
+export interface McpToolDefinition {
+    serverName: string
+    toolName: string
+    description: string
+    inputSchema: any // schema from the server
+}
+
+export class McpManager {
+    private mcpServers: Record<string, MCPServerConfig> = {}
+    private clients: Map<string, Client> = new Map() // key: serverName, val: MCP client
+    private mcpTools: McpToolDefinition[] = []
+
+    private constructor(private configPath: string) {}
+
+    public static async create(configPath: string): Promise<McpManager> {
+        const instance = new McpManager(configPath)
+        await instance.loadConfig()
+        await instance.initAllServers()
+        return instance
+    }
+
+    public async loadConfig(): Promise<void> {
+        if (!(await fs.exists(this.configPath))) {
+            throw new Error(`Could not load the MCP config at ${this.configPath}`)
+        }
+        const raw = await fs.readFileText(this.configPath)
+        const json = JSON.parse(raw) as MCPConfig
+        if (!json.mcpServers) {
+            throw new Error(`No "mcpServers" field found in config: ${this.configPath}`)
+        }
+        this.mcpServers = json.mcpServers
+    }
+
+    public async initAllServers(): Promise<void> {
+        for (const [serverName, serverConfig] of Object.entries(this.mcpServers)) {
+            if (serverConfig.disabled) {
+                getLogger().info(`MCP server [${serverName}] is disabled, skipping.`)
+                continue
+            }
+            await this.initOneServer(serverName, serverConfig)
+        }
+    }
+
+    private async initOneServer(serverName: string, serverConfig: MCPServerConfig): Promise<void> {
+        try {
+            getLogger().debug(`Initializing MCP server [${serverName}] with command: ${serverConfig.command}`)
+            const transport = new StdioClientTransport({
+                command: serverConfig.command,
+                args: serverConfig.args ?? [],
+                env: process.env as Record<string, string>,
+            })
+            const client = new Client({
+                name: `q-agentic-chat-mcp-client-${serverName}`,
+                version: '1.0.0',
+            })
+            await client.connect(transport)
+            this.clients.set(serverName, client)
+
+            const toolsResult = (await client.listTools()) as ListToolsResponse
+            for (const toolInfo of toolsResult.tools) {
+                const toolDef: McpToolDefinition = {
+                    serverName,
+                    toolName: toolInfo.name ?? 'unknown',
+                    description: toolInfo.description ?? '',
+                    inputSchema: toolInfo.inputSchema ?? {},
+                }
+                this.mcpTools.push(toolDef)
+                getLogger().info(`Found MCP tool [${toolDef.toolName}] from server [${serverName}]`)
+            }
+        } catch (err) {
+            getLogger().error(`Failed to init server [${serverName}]: ${(err as Error).message}`)
+        }
+    }
+
+    public getAllMcpTools(): McpToolDefinition[] {
+        return [...this.mcpTools]
+    }
+
+    public async callTool(serverName: string, toolName: string, args: any): Promise<any> {
+        const client = this.clients.get(serverName)
+        if (!client) {
+            throw new Error(`MCP server [${serverName}] not connected or not found in clients.`)
+        }
+        return await client.callTool({
+            name: toolName,
+            arguments: args,
+        })
+    }
+
+    public findTool(toolName: string): McpToolDefinition | undefined {
+        return this.mcpTools.find((t) => t.toolName === toolName)
+    }
+}

packages/core/src/codewhispererChat/tools/mcp/mcpManager.ts

@@ -0,0 +1,55 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+import { Writable } from 'stream'
+import { getLogger } from '../../../shared/logger/logger'
+import { ToolUtils } from '../toolUtils'
+import { CommandValidation, InvokeOutput, OutputKind } from '../toolShared'
+
+export interface McpToolParams {
+    serverName: string
+    toolName: string
+    input?: any
+}
+
+export class McpTool {
+    private readonly logger = getLogger('mcp')
+    private serverName: string
+    private toolName: string
+    private input: any
+
+    public constructor(params: McpToolParams) {
+        this.serverName = params.serverName
+        this.toolName = params.toolName
+        this.input = params.input
+    }
+
+    public async validate(): Promise<void> {}
+
+    public queueDescription(updates: Writable): void {
+        updates.write(`Invoking remote MCP tool: ${this.toolName} on server ${this.serverName}`)
+        updates.end()
+    }
+
+    public requiresAcceptance(): CommandValidation {
+        return { requiresAcceptance: true }
+    }
+
+    public async invoke(updates?: Writable): Promise<InvokeOutput> {
+        try {
+            const result = await ToolUtils.mcpManager!.callTool(this.serverName, this.toolName, this.input)
+            const content = typeof result === 'object' ? JSON.stringify(result) : String(result)
+
+            return {
+                output: {
+                    kind: OutputKind.Text,
+                    content,
+                },
+            }
+        } catch (error: any) {
+            this.logger.error(`Failed to invoke MCP tool: ${error.message ?? error}`)
+            throw new Error(`Failed to invoke MCP tool: ${error.message ?? error}`)
+        }
+    }
+}

packages/core/src/codewhispererChat/tools/mcp/mcpTool.ts

@@ -0,0 +1,25 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+export interface MCPServerConfig {
+    command: string
+    args?: string[]
+    env?: Record<string, string>
+    disabled?: boolean
+    autoApprove?: string[]
+}
+
+export interface MCPConfig {
+    mcpServers: Record<string, MCPServerConfig>
+}
+
+export interface ListToolsResponse {
+    tools: {
+        name?: string
+        description?: string
+        inputSchema?: object
+        [key: string]: any
+    }[]
+}

packages/core/src/codewhispererChat/tools/mcp/mcpTypes.ts

@@ -15,21 +15,27 @@ import {
     fsReadToolResponseSize,
 } from './toolShared'
 import { ListDirectory, ListDirectoryParams } from './listDirectory'
+import { McpTool } from './mcp/mcpTool'
+import { McpManager } from './mcp/mcpManager'
 
 export enum ToolType {
     FsRead = 'fsRead',
     FsWrite = 'fsWrite',
     ExecuteBash = 'executeBash',
     ListDirectory = 'listDirectory',
+    Mcp = 'mcp',
 }
 
 export type Tool =
     | { type: ToolType.FsRead; tool: FsRead }
     | { type: ToolType.FsWrite; tool: FsWrite }
     | { type: ToolType.ExecuteBash; tool: ExecuteBash }
     | { type: ToolType.ListDirectory; tool: ListDirectory }
+    | { type: ToolType.Mcp; tool: McpTool }
 
 export class ToolUtils {
+    static mcpManager?: McpManager
+
     static displayName(tool: Tool): string {
         switch (tool.type) {
             case ToolType.FsRead:
@@ -40,6 +46,8 @@ export class ToolUtils {
                 return 'Execute shell command'
             case ToolType.ListDirectory:
                 return 'List directory from filesystem'
+            case ToolType.Mcp:
+                return 'Execute MCP tool'
         }
     }
 
@@ -53,6 +61,8 @@ export class ToolUtils {
                 return tool.tool.requiresAcceptance()
             case ToolType.ListDirectory:
                 return tool.tool.requiresAcceptance()
+            case ToolType.Mcp:
+                return { requiresAcceptance: false }
         }
     }
 
@@ -75,6 +85,8 @@ export class ToolUtils {
                 return tool.tool.invoke(updates ?? undefined, cancellationToken)
             case ToolType.ListDirectory:
                 return tool.tool.invoke(updates)
+            case ToolType.Mcp:
+                return tool.tool.invoke(updates)
         }
     }
 
@@ -159,7 +171,18 @@ export class ToolUtils {
                         type: ToolType.ListDirectory,
                         tool: new ListDirectory(value.input as unknown as ListDirectoryParams),
                     }
-                default:
+                default: {
+                    const mcpToolDef = ToolUtils.mcpManager?.findTool(value.name as string)
+                    if (mcpToolDef) {
+                        return {
+                            type: ToolType.Mcp,
+                            tool: new McpTool({
+                                serverName: mcpToolDef.serverName,
+                                toolName: mcpToolDef.toolName,
+                                input: value.input, // pass LLM's JSON input
+                            }),
+                        }
+                    }
                     return {
                         toolUseId: value.toolUseId,
                         content: [
@@ -169,6 +192,7 @@ export class ToolUtils {
                             } as ToolResultContentBlock,
                         ],
                     }
+                }
             }
         } catch (error) {
             return mapErr(error)

packages/core/src/codewhispererChat/tools/toolUtils.ts

@@ -22,6 +22,7 @@ export type LogTopic =
     | 'listDirectory'
     | 'chatStream'
     | 'chatHistoryDb'
+    | 'mcp'
     | 'unknown'
 
 class ErrorLog {