feat: add mcp client (#1034)

* feat: add mcp client

Author: bywang56

app/aws-lsp-codewhisperer-runtimes/src/agent-standalone.ts

@@ -28,6 +28,7 @@ const props: RuntimeProps = {
         FsToolsServer,
         BashToolsServer,
         QLocalProjectContextServerTokenProxy,
+        // McpToolsServer,
         // LspToolsServer,
     ],
     name: 'AWS CodeWhisperer',

package-lock.json

@@ -51,7 +51,8 @@
         "picomatch": "^4.0.2",
         "shlex": "2.1.2",
         "uuid": "^11.0.5",
-        "vscode-uri": "^3.1.0"
+        "vscode-uri": "^3.1.0",
+        "@modelcontextprotocol/sdk": "^1.9.0"
     },
     "devDependencies": {
         "@types/adm-zip": "^0.5.5",

server/aws-lsp-codewhisperer/package.json

@@ -0,0 +1,67 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates.
+ * All Rights Reserved. SPDX-License-Identifier: Apache-2.0
+ */
+
+import { expect } from 'chai'
+import { McpManager } from './mcpManager'
+
+describe('McpManager (singleton & empty‑config)', () => {
+    const fakeLogging = {
+        log: (_: string) => {},
+        info: (_: string) => {},
+        warn: (_: string) => {},
+        error: (_: string) => {},
+        debug: (_: string) => {},
+    }
+
+    // Minimal stub for Features.workspace
+    const fakeWorkspace = {
+        fs: {
+            exists: (_: string) => Promise.resolve(false),
+            readFile: (_: string) => Promise.resolve(Buffer.from('')),
+        },
+        getUserHomeDir: () => '',
+    }
+
+    const fakeLsp = {}
+
+    const features = {
+        logging: fakeLogging,
+        workspace: fakeWorkspace,
+        lsp: fakeLsp,
+    } as unknown as Pick<
+        import('@aws/language-server-runtimes/server-interface/server').Features,
+        'logging' | 'workspace' | 'lsp'
+    >
+
+    afterEach(async () => {
+        try {
+            await McpManager.instance.close()
+        } catch {
+            /* ignore */
+        }
+    })
+
+    it('init() always returns the same instance', async () => {
+        const m1 = await McpManager.init([], features)
+        const m2 = await McpManager.init([], features)
+        expect(m1).to.equal(m2)
+    })
+
+    it('with no config paths it discovers zero tools', async () => {
+        const mgr = await McpManager.init([], features)
+        expect(mgr.getAllTools()).to.be.an('array').that.is.empty
+    })
+
+    it('callTool() on non‑existent server throws', async () => {
+        const mgr = await McpManager.init([], features)
+        try {
+            await mgr.callTool('nope', 'foo', {})
+            throw new Error('Expected callTool to throw, but it did not')
+        } catch (err: any) {
+            expect(err).to.be.instanceOf(Error)
+            expect(err.message).to.equal(`MCP: server 'nope' not connected`)
+        }
+    })
+})

server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpManager.test.ts

@@ -0,0 +1,113 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates.
+ * All Rights Reserved. SPDX-License-Identifier: Apache-2.0
+ */
+
+import type { Features } from '@aws/language-server-runtimes/server-interface/server'
+import { Client } from '@modelcontextprotocol/sdk/client/index.js'
+import { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js'
+import type { MCPServerConfig, McpToolDefinition, ListToolsResponse } from './mcpTypes'
+import { loadMcpServerConfigs } from './mcpUtils'
+
+export class McpManager {
+    static #instance?: McpManager
+    private clients = new Map<string, Client>()
+    private mcpTools: McpToolDefinition[] = []
+
+    private constructor(
+        private configPaths: string[],
+        private features: Pick<Features, 'logging' | 'workspace' | 'lsp'>
+    ) {}
+
+    public static async init(
+        configPaths: string[],
+        features: Pick<Features, 'logging' | 'workspace' | 'lsp'>
+    ): Promise<McpManager> {
+        if (!McpManager.#instance) {
+            const mgr = new McpManager(configPaths, features)
+            McpManager.#instance = mgr
+            await mgr.discoverAllServers()
+            features.logging.info(`MCP: discovered ${mgr.mcpTools.length} tools across all servers`)
+        }
+        return McpManager.#instance
+    }
+
+    public static get instance(): McpManager {
+        if (!McpManager.#instance) {
+            throw new Error('McpManager not initialized—call McpManager.init(...) first')
+        }
+        return McpManager.#instance
+    }
+
+    public async close(): Promise<void> {
+        this.features.logging.info('MCP: closing all clients')
+        for (const [name, client] of this.clients.entries()) {
+            try {
+                await client.close()
+                this.features.logging.info(`MCP: closed client for ${name}`)
+            } catch (e: any) {
+                this.features.logging.error(`MCP: error closing client ${name}: ${e.message}`)
+            }
+        }
+        this.clients.clear()
+        this.mcpTools = []
+        McpManager.#instance = undefined
+    }
+
+    private async discoverAllServers(): Promise<void> {
+        const configs = await loadMcpServerConfigs(this.features.workspace, this.features.logging, this.configPaths)
+
+        for (const [name, cfg] of configs.entries()) {
+            if (cfg.disabled) {
+                this.features.logging.info(`MCP: server '${name}' is disabled, skipping`)
+                continue
+            }
+            await this.initOneServer(name, cfg)
+        }
+    }
+
+    private async initOneServer(serverName: string, cfg: MCPServerConfig): Promise<void> {
+        try {
+            this.features.logging.debug(`MCP: initializing server [${serverName}]`)
+            const mergedEnv = {
+                ...(process.env as Record<string, string>),
+                ...(cfg.env ?? {}),
+            }
+            const transport = new StdioClientTransport({
+                command: cfg.command,
+                args: cfg.args ?? [],
+                env: mergedEnv,
+            })
+            const client = new Client({
+                name: `mcp-client-${serverName}`,
+                version: '1.0.0',
+            })
+            await client.connect(transport)
+            this.clients.set(serverName, client)
+
+            const resp = (await client.listTools()) as ListToolsResponse
+            for (const t of resp.tools) {
+                const toolName = t.name!
+                this.features.logging.info(`MCP: discovered tool ${serverName}::${toolName}`)
+                this.mcpTools.push({
+                    serverName,
+                    toolName,
+                    description: t.description ?? '',
+                    inputSchema: t.inputSchema ?? {},
+                })
+            }
+        } catch (e: any) {
+            this.features.logging.warn(`MCP: server [${serverName}] init failed: ${e.message}`)
+        }
+    }
+
+    public getAllTools(): McpToolDefinition[] {
+        return [...this.mcpTools]
+    }
+
+    public async callTool(server: string, tool: string, args: any): Promise<any> {
+        const c = this.clients.get(server)
+        if (!c) throw new Error(`MCP: server '${server}' not connected`)
+        return c.callTool({ name: tool, arguments: args })
+    }
+}

server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpManager.ts

@@ -0,0 +1,53 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates.
+ * All Rights Reserved. SPDX-License-Identifier: Apache-2.0
+ */
+
+import { expect } from 'chai'
+import { McpTool } from './mcpTool'
+import { McpManager } from './mcpManager'
+import type { McpToolDefinition } from './mcpTypes'
+
+describe('McpTool', () => {
+    const fakeFeatures = {
+        logging: { info: () => {}, warn: () => {}, error: () => {}, debug: () => {}, log: () => {} },
+        workspace: {
+            fs: { exists: () => Promise.resolve(false), readFile: () => Promise.resolve(Buffer.from('')) },
+            getUserHomeDir: () => '',
+        },
+        lsp: {},
+    } as unknown as Pick<
+        import('@aws/language-server-runtimes/server-interface/server').Features,
+        'logging' | 'workspace' | 'lsp'
+    >
+
+    const definition: McpToolDefinition = {
+        serverName: 'nope',
+        toolName: 'doesNotExist',
+        description: 'desc',
+        inputSchema: {},
+    }
+
+    afterEach(async () => {
+        try {
+            await McpManager.instance.close()
+        } catch {
+            // ignore
+        }
+    })
+
+    it('invoke() throws when server is not connected', async () => {
+        // Initialize McpManager with no config paths → zero servers
+        await McpManager.init([], fakeFeatures)
+
+        const tool = new McpTool(fakeFeatures, definition)
+
+        try {
+            await tool.invoke({}) // should throw
+            throw new Error('Expected invoke() to throw')
+        } catch (err: any) {
+            expect(err).to.be.instanceOf(Error)
+            expect(err.message).to.equal(`Failed to invoke MCP tool: MCP: server 'nope' not connected`)
+        }
+    })
+})

server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpTool.test.ts

@@ -0,0 +1,55 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates.
+ * All Rights Reserved. SPDX-License-Identifier: Apache-2.0
+ */
+
+import { CommandValidation, InvokeOutput, OutputKind } from '../toolShared'
+import type { McpToolDefinition } from './mcpTypes'
+import type { Features } from '@aws/language-server-runtimes/server-interface/server'
+import { McpManager } from './mcpManager'
+
+export class McpTool {
+    constructor(
+        private readonly features: Pick<Features, 'logging' | 'workspace' | 'lsp'>,
+        private readonly def: McpToolDefinition
+    ) {}
+
+    public getSpec() {
+        return {
+            name: this.def.toolName,
+            description: this.def.description,
+            inputSchema: this.def.inputSchema,
+        } as const
+    }
+
+    public validate(_input: any): Promise<void> {
+        return Promise.resolve()
+    }
+
+    public async queueDescription(command: string, updates: WritableStream) {
+        const writer = updates.getWriter()
+        await writer.write(`Invoking MCP tool: ${this.def.toolName} on server ${this.def.serverName}`)
+        await writer.close()
+        writer.releaseLock()
+    }
+
+    public requiresAcceptance(_input: any): CommandValidation {
+        return { requiresAcceptance: false }
+    }
+
+    public async invoke(input: any): Promise<InvokeOutput> {
+        try {
+            const result = await McpManager.instance.callTool(this.def.serverName, this.def.toolName, input)
+            const content = typeof result === 'object' ? JSON.stringify(result) : String(result)
+            return {
+                output: {
+                    kind: OutputKind.Text,
+                    content: content,
+                },
+            }
+        } catch (err: any) {
+            this.features.logging.error(`MCP tool ${this.def.toolName} failed: ${err.message}`)
+            throw new Error(`Failed to invoke MCP tool: ${err.message}`)
+        }
+    }
+}

server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpTool.ts

@@ -0,0 +1,32 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+export interface McpToolDefinition {
+    serverName: string
+    toolName: string
+    description: string
+    inputSchema: any // schema from the server
+}
+
+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
+    }[]
+}