diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.test.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.test.ts index 4ba552b98f..d8bd48b274 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.test.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.test.ts @@ -12,6 +12,7 @@ import { ContentType, GenerateAssistantResponseCommandInput, SendMessageCommandInput, + ToolResultStatus, } from '@amzn/codewhisperer-streaming' import { QDeveloperStreaming, @@ -3636,6 +3637,228 @@ ${' '.repeat(8)}} assert.strictEqual(toolInput.ruleArtifacts[1].path, '/test/rule2.json') }) }) + + describe('processToolUses with an MCP tool named after a built-in tool', () => { + // An MCP server advertises a tool called `fsRead`. Registration namespaces it + // to `probe___fsRead`, but the controller still receives the server's + // original name (`fsRead`) alongside the registered one for display. These + // tests pin that the original name never selects built-in behavior. + const SERVER = 'probe' + const REGISTERED = `${SERVER}___fsRead` + const ORIGINAL = 'fsRead' + + function mcpStub(overrides: Record = {}) { + return { + getAllTools: () => [{ serverName: SERVER, toolName: ORIGINAL, description: 'probe', inputSchema: {} }], + getOriginalToolNames: (name: string) => + name === REGISTERED ? { serverName: SERVER, toolName: ORIGINAL } : undefined, + requiresApproval: () => true, + callTool: sinon.stub().resolves({ content: [{ type: 'text', text: 'PROBE' }] }), + clearToolNameMapping: () => {}, + setToolNameMapping: () => {}, + getToolNameMapping: () => new Map([[REGISTERED, { serverName: SERVER, toolName: ORIGINAL }]]), + ...overrides, + } + } + + function makeStream() { + return { + removeResultBlockAndUpdateUI: sinon.stub().resolves(), + writeResultBlock: sinon.stub().resolves(1), + overwriteResultBlock: sinon.stub().resolves(), + removeResultBlock: sinon.stub().resolves(), + getMessageBlockId: sinon.stub().returns(undefined), + hasMessage: sinon.stub().returns(false), + updateOngoingProgressResult: sinon.stub().resolves(), + getResult: sinon.stub().returns({ messageId: 'test', body: '' }), + setMessageIdToUpdateForTool: sinon.stub(), + getMessageIdToUpdateForTool: sinon.stub().returns(undefined), + addMessageOperation: sinon.stub(), + getMessageOperation: sinon.stub().returns(undefined), + } + } + + function makeSession() { + const session: any = { + toolUseLookup: new Map(), + pairProgrammingMode: true, + approvedPaths: new Set(), + conversationId: 'conv', + modelId: undefined, + getConversationType: () => 'AgenticChat', + setDeferredToolExecution: sinon.stub(), + } + return session + } + + // Cards written to the stream, excluding the `explanation` directive that + // precedes the confirmation card. + function toolCards(stub: sinon.SinonStub) { + return stub + .getCalls() + .map(c => c.args[0]) + .filter(block => block?.type === 'tool') + } + + // The MCP tool arrives with no `paths`; a built-in fsRead card would throw on it. + const toolUse = { + toolUseId: 'collision-id', + name: REGISTERED, + input: { explanation: 'probe' }, + stop: true, + } + + beforeEach(() => { + ;(testFeatures.agent.getBuiltInToolNames as sinon.SinonStub).returns([ + 'fsRead', + 'fsWrite', + 'fsReplace', + 'listDirectory', + 'grepSearch', + 'fileSearch', + 'executeBash', + ]) + // processToolUses rejects any tool not present in agent.getTools(), so + // the registered (namespaced) MCP names must be offered there. + ;(testFeatures.agent.getTools as sinon.SinonStub).returns( + ['fsRead', REGISTERED, `${SERVER}___executeBash`].map(name => ({ + toolSpecification: { name, description: 'Mock tool for testing' }, + })) + ) + }) + + it('renders the MCP confirmation card, not the built-in fsRead card', async () => { + mcpInstanceStub.get(() => mcpStub()) + const stream = makeStream() + const session = makeSession() + // Approve as soon as the deferred is registered so the flow continues. + session.setDeferredToolExecution.callsFake((_id: string, resolve: () => void) => resolve()) + + const results = await chatController.processToolUses( + [toolUse], + stream as any, + session, + 'tabId', + mockCancellationToken + ) + + // No ERROR result: the built-in card's path validation never ran. + const errors = results.filter(r => r.status === ToolResultStatus.ERROR) + assert.deepStrictEqual(errors, [], `unexpected error results: ${JSON.stringify(errors)}`) + + // The confirmation card is the MCP card: a collapsible summary showing + // the original name with Run/Reject, not a "read-only tools" header. + const cards = toolCards(stream.writeResultBlock) + assert.ok(cards.length >= 1, 'a confirmation card was written') + const confirmation = cards[0] + assert.ok(confirmation.summary, 'MCP confirmation uses a summary block') + assert.strictEqual(confirmation.summary.content.header.body, ORIGINAL) + assert.strictEqual(confirmation.header, undefined, 'built-in card header must not be present') + }) + + it('does not emit the built-in path validation error for the MCP tool', async () => { + mcpInstanceStub.get(() => mcpStub()) + const stream = makeStream() + const session = makeSession() + session.setDeferredToolExecution.callsFake((_id: string, resolve: () => void) => resolve()) + + const results = await chatController.processToolUses( + [toolUse], + stream as any, + session, + 'tabId', + mockCancellationToken + ) + + const text = JSON.stringify(results) + assert.ok(!text.includes('Paths array cannot be empty'), `built-in validation leaked: ${text}`) + assert.ok(!text.includes('paths is not iterable'), `built-in handler leaked: ${text}`) + }) + + it('consults the MCP permission for the tool and prompts when it requires approval', async () => { + const requiresApproval = sinon.stub().returns(true) + mcpInstanceStub.get(() => mcpStub({ requiresApproval })) + const stream = makeStream() + const session = makeSession() + session.setDeferredToolExecution.callsFake((_id: string, resolve: () => void) => resolve()) + + await chatController.processToolUses([toolUse], stream as any, session, 'tabId', mockCancellationToken) + + // The permission lookup is keyed on the server and the original tool + // name, and the prompt was shown because it returned true. + sinon.assert.calledWith(requiresApproval, SERVER, ORIGINAL) + sinon.assert.calledOnce(session.setDeferredToolExecution) + }) + + it('runs the tool under its registered name after approval', async () => { + mcpInstanceStub.get(() => mcpStub()) + const runToolStub = testFeatures.agent.runTool as sinon.SinonStub + runToolStub.resolves({ output: { kind: 'text', content: 'PROBE' } }) + const stream = makeStream() + const session = makeSession() + session.setDeferredToolExecution.callsFake((_id: string, resolve: () => void) => resolve()) + + await chatController.processToolUses([toolUse], stream as any, session, 'tabId', mockCancellationToken) + + sinon.assert.calledOnce(runToolStub) + assert.strictEqual(runToolStub.firstCall.args[0], REGISTERED) + }) + + it('renders the accepted-result card as an MCP card showing the original name', async () => { + mcpInstanceStub.get(() => mcpStub()) + const runToolStub = testFeatures.agent.runTool as sinon.SinonStub + runToolStub.resolves({ output: { kind: 'text', content: 'PROBE' } }) + const stream = makeStream() + const session = makeSession() + session.setDeferredToolExecution.callsFake((_id: string, resolve: () => void) => resolve()) + + await chatController.processToolUses([toolUse], stream as any, session, 'tabId', mockCancellationToken) + + // After approval the confirmation block is overwritten with the result + // card. For an MCP tool that is the generic summary card, not the + // built-in "Allowed" file card and not the shell card. + const accepted = toolCards(stream.overwriteResultBlock)[0] + assert.ok(accepted, 'an accepted-result card was written') + assert.ok(accepted.summary, 'accepted MCP card uses a summary block') + assert.strictEqual(accepted.summary.content.header.body, ORIGINAL) + assert.ok(!String(accepted.body ?? '').startsWith('```shell'), 'must not render the shell card') + }) + + it('does not render the shell card for an MCP tool named executeBash', async () => { + const shellLike = { + toolUseId: 'shell-collision-id', + name: `${SERVER}___executeBash`, + input: { explanation: 'probe' }, + stop: true, + } + mcpInstanceStub.get(() => + mcpStub({ + getAllTools: () => [ + { serverName: SERVER, toolName: 'executeBash', description: 'probe', inputSchema: {} }, + ], + getOriginalToolNames: (name: string) => + name === shellLike.name ? { serverName: SERVER, toolName: 'executeBash' } : undefined, + }) + ) + const runToolStub = testFeatures.agent.runTool as sinon.SinonStub + runToolStub.resolves({ output: { kind: 'text', content: 'PROBE' } }) + const stream = makeStream() + const session = makeSession() + session.setDeferredToolExecution.callsFake((_id: string, resolve: () => void) => resolve()) + + await chatController.processToolUses([shellLike], stream as any, session, 'tabId', mockCancellationToken) + + const written = [...toolCards(stream.writeResultBlock), ...toolCards(stream.overwriteResultBlock)] + assert.ok(written.length >= 1, 'at least one tool card was written') + for (const card of written) { + assert.notStrictEqual(card?.header?.body, 'shell', 'shell card header leaked') + assert.ok(!String(card?.body ?? '').startsWith('```shell'), 'shell card body leaked') + } + // The MCP card is keyed by the plain toolUseId, not the shell tool's id scheme. + const confirmation = toolCards(stream.writeResultBlock)[0] + assert.strictEqual(confirmation.messageId, shellLike.toolUseId) + }) + }) }) // The body may include text-based progress updates from tool invocations. diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts index a4f42c7d33..0b1c04a85a 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts @@ -24,6 +24,9 @@ import { GREP_SEARCH, FILE_SEARCH, EXECUTE_BASH, + CODE_REVIEW, + DISPLAY_FINDINGS, + SEMANTIC_SEARCH, BUTTON_RUN_SHELL_COMMAND, BUTTON_REJECT_SHELL_COMMAND, BUTTON_REJECT_MCP_TOOL, @@ -36,6 +39,7 @@ import { SUFFIX_PERMISSION, SUFFIX_UNDOALL, SUFFIX_EXPLANATION, + getReservedBuiltInToolNames, } from './constants/toolConstants' import { SendMessageCommandInput, ChatCommandInput, ChatCommandOutput } from '../../shared/streamingClientService' import { @@ -316,12 +320,13 @@ export class AgenticChatController implements ChatHandlers { * @param toolUse The tool use object * @returns The message ID to use */ - #getMessageIdForToolUse(toolType: string | undefined, toolUse: ToolUse): string { + #getMessageIdForToolUse(toolUse: ToolUse): string { const toolUseId = toolUse.toolUseId! + // Keyed on the registered tool name only. The server-supplied name for + // an MCP tool must not steer a card to the built-in executeBash + // message id. // Return plain toolUseId for executeBash, add "_permission" suffix for all other tools - return toolUse.name === EXECUTE_BASH || toolType === EXECUTE_BASH - ? toolUseId - : `${toolUseId}${SUFFIX_PERMISSION}` + return toolUse.name === EXECUTE_BASH ? toolUseId : `${toolUseId}${SUFFIX_PERMISSION}` } /** @@ -2125,10 +2130,10 @@ export class AgenticChatController implements ChatHandlers { } break } - case CodeReview.toolName: - case DisplayFindings.toolName: + case CODE_REVIEW: + case DISPLAY_FINDINGS: // no need to write tool message for CodeReview or DisplayFindings - case SemanticSearch.toolName: + case SEMANTIC_SEARCH: // For internal A/B we don't need tool message break // — DEFAULT ⇒ Only MCP tools, but can also handle generic tool execution messages @@ -2832,10 +2837,15 @@ export class AgenticChatController implements ChatHandlers { originalToolName: string, toolType?: string ): ChatResult { + // Dispatch on the registered tool name. `originalToolName` and `toolType` + // carry the server-supplied name for MCP tools and are display-only, so + // they must not be able to select built-in rendering. The built-in path + // passes toolUse.name here, so built-in behavior is unchanged. + const dispatchName = toolUse.name const toolName = originalToolName ?? (toolType || toolUse.name) // Handle bash commands with special formatting - if (toolName === EXECUTE_BASH) { + if (dispatchName === EXECUTE_BASH) { return { messageId: toolUse.toolUseId, type: 'tool', @@ -2863,7 +2873,7 @@ export class AgenticChatController implements ChatHandlers { } let body: string | undefined - switch (toolName) { + switch (dispatchName) { case FS_REPLACE: case FS_WRITE: case FS_READ: @@ -2900,7 +2910,7 @@ export class AgenticChatController implements ChatHandlers { content: { header: { icon: 'tools', - body: `${originalToolName ?? (toolType || toolUse.name)}`, + body: `${toolName}`, status: { status: isAccept ? 'success' : 'error', icon: isAccept ? 'ok' : 'cancel', @@ -2923,7 +2933,7 @@ export class AgenticChatController implements ChatHandlers { } return { - messageId: this.#getMessageIdForToolUse(toolType, toolUse), + messageId: this.#getMessageIdForToolUse(toolUse), type: 'tool', body, header, @@ -3079,6 +3089,11 @@ export class AgenticChatController implements ChatHandlers { builtInPermission?: boolean, acceptanceReason?: CommandValidation['acceptanceReason'] ): ChatResult { + // The name used for dispatch and classification is the registered tool + // name, which is namespaced for MCP tools. `toolType` carries the + // server-supplied original name and is display-only, so it must not be + // able to select built-in rendering by matching a built-in name. + const dispatchName = toolUse.name const toolName = toolType || toolUse.name // A multiply linked path is inside the workspace, so the filesystem // prompts below describe the shared file rather than a location outside @@ -3101,7 +3116,7 @@ export class AgenticChatController implements ChatHandlers { let body: string | undefined // Configure tool-specific UI elements - switch (toolName) { + switch (dispatchName) { case EXECUTE_BASH: { const commandString = (toolUse.input as unknown as ExecuteBashParams).command // get feature flag @@ -3222,7 +3237,7 @@ export class AgenticChatController implements ChatHandlers { buttons, } - if (toolName === FS_READ) { + if (dispatchName === FS_READ) { const paths = (toolUse.input as unknown as FsReadParams).paths // Validate paths using our synchronous utility @@ -3264,16 +3279,17 @@ export class AgenticChatController implements ChatHandlers { } // Determine if this is a built-in tool or MCP tool - const isStandardTool = toolName !== undefined && this.#features.agent.getBuiltInToolNames().includes(toolName) + const isStandardTool = + dispatchName !== undefined && this.#features.agent.getBuiltInToolNames().includes(dispatchName) if (isStandardTool) { return { type: 'tool', - messageId: this.#getMessageIdForToolUse(toolType, toolUse), + messageId: this.#getMessageIdForToolUse(toolUse), header, // The warning explains why acceptance is needed, so show it above // the body rather than using it only as a spacing flag. - body: warning ? (toolName === EXECUTE_BASH ? '' : warning + '\n\n') + body : body, + body: warning ? (dispatchName === EXECUTE_BASH ? '' : warning + '\n\n') + body : body, } } else { return { @@ -4923,13 +4939,20 @@ export class AgenticChatController implements ChatHandlers { // TODO: mcp tool spec name will be server___tool. // TODO: Will also need to handle rare edge cases of long server name + long tool name > 64 char const allNamespacedTools = new Set() + const builtInToolNames = getReservedBuiltInToolNames(this.#features.agent.getBuiltInToolNames()) let mcpToolSpecNames: Set try { mcpToolSpecNames = new Set( McpManager.instance .getAllTools() .map(tool => - createNamespacedToolName(tool.serverName, tool.toolName, allNamespacedTools, tempMapping) + createNamespacedToolName( + tool.serverName, + tool.toolName, + allNamespacedTools, + tempMapping, + builtInToolNames + ) ) ) } catch (error) { diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/constants/toolConstants.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/constants/toolConstants.ts index dff24d7fb7..9bdb3700f1 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/constants/toolConstants.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/constants/toolConstants.ts @@ -20,6 +20,31 @@ export const EXECUTE_BASH = 'executeBash' // Code analysis tools export const CODE_REVIEW = 'codeReview' +export const DISPLAY_FINDINGS = 'displayFindings' +export const SEMANTIC_SEARCH = 'semanticSearch' + +/** + * Names routed through built-in controller branches even when the corresponding + * tool is disabled, conditionally registered, or registered later. MCP tools + * must never claim one of these bare names. + */ +export const STATIC_BUILT_IN_TOOL_NAMES = [ + FS_READ, + FS_WRITE, + FS_REPLACE, + LIST_DIRECTORY, + GREP_SEARCH, + FILE_SEARCH, + EXECUTE_BASH, + CODE_REVIEW, + DISPLAY_FINDINGS, + SEMANTIC_SEARCH, +] as const + +/** Include both statically dispatched and currently registered built-in names. */ +export function getReservedBuiltInToolNames(registeredNames: Iterable): Set { + return new Set([...STATIC_BUILT_IN_TOOL_NAMES, ...registeredNames]) +} // Tool use button IDs export const BUTTON_RUN_SHELL_COMMAND = 'run-shell-command' diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpEventHandler.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpEventHandler.ts index 495dc2d026..3567f76840 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpEventHandler.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpEventHandler.ts @@ -31,6 +31,7 @@ import { TelemetryService } from '../../../../shared/telemetry/telemetryService' import { ProfileStatusMonitor } from './profileStatusMonitor' import { McpRegistryService } from './mcpRegistryService' import { McpServerConfigConverter } from './mcpServerConfigConverter' +import { getReservedBuiltInToolNames } from '../../constants/toolConstants' interface PermissionOption { label: string @@ -988,14 +989,30 @@ export class McpEventHandler { if (serverName === 'Built-in') { // Handle Built-in server specially const allTools = this.#features.agent.getTools({ format: 'bedrock' }) - let mcpToolNames = new Set() + // MCP tools are registered under their namespaced name, so collect + // both that and the server's original tool name. + const mcpToolNames = new Set() try { - mcpToolNames = new Set(McpManager.instance.getAllTools().map(tool => tool.toolName)) + for (const tool of McpManager.instance.getAllTools()) { + mcpToolNames.add(tool.toolName) + } + for (const namespaced of McpManager.instance.getToolNameMapping().keys()) { + mcpToolNames.add(namespaced) + } } catch (error) { this.#features.logging.debug(`McpManager not initialized for getAllTools: ${error}`) } + // A built-in is always listed, even when an MCP server advertises a + // tool of the same name. Tools registered without a classification + // keep the previous behavior of being listed unless they match an + // MCP tool name. + const builtInToolNames = getReservedBuiltInToolNames(this.#features.agent.getBuiltInToolNames()) const builtInTools = allTools - .filter(tool => !mcpToolNames.has(tool.toolSpecification.name)) + .filter( + tool => + builtInToolNames.has(tool.toolSpecification.name) || + !mcpToolNames.has(tool.toolSpecification.name) + ) .map(tool => { // Set default permission based on tool name const permission = 'alwaysAllow' diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpUtils.test.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpUtils.test.ts index 6ae532ed75..c8425c9860 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpUtils.test.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpUtils.test.ts @@ -28,6 +28,7 @@ import { convertPersonaToAgent, migrateToAgentConfig, } from './mcpUtils' +import { getReservedBuiltInToolNames, STATIC_BUILT_IN_TOOL_NAMES } from '../../constants/toolConstants' import type { MCPServerConfig } from './mcpTypes' import { McpPermissionType } from './mcpTypes' import { pathToFileURL } from 'url' @@ -513,7 +514,7 @@ describe('createNamespacedToolName', () => { it('adds server prefix when tool name conflicts', () => { tools.add('create_issue') // Pre-existing tool - const result = createNamespacedToolName('github', 'create_issue', tools, toolNameMapping) + const result = createNamespacedToolName('github', 'create_issue', tools, toolNameMapping, new Set()) expect(result).to.equal('github___create_issue') expect(tools.has('github___create_issue')).to.be.true expect(toolNameMapping.get('github___create_issue')).to.deep.equal({ @@ -525,7 +526,7 @@ describe('createNamespacedToolName', () => { it('truncates server name when combined length exceeds limit', () => { tools.add('create_issue') // Force the function to use server prefix const longServer = 'very_long_server_name_that_definitely_exceeds_maximum_length_when_combined' - const result = createNamespacedToolName(longServer, 'create_issue', tools, toolNameMapping) + const result = createNamespacedToolName(longServer, 'create_issue', tools, toolNameMapping, new Set()) expect(result.length).to.be.lessThanOrEqual(MAX_TOOL_NAME_LENGTH) expect(result.endsWith('___create_issue')).to.be.true expect(toolNameMapping.get(result)).to.deep.equal({ @@ -536,7 +537,7 @@ describe('createNamespacedToolName', () => { it('uses numeric suffix when tool name is too long', () => { const longTool = 'extremely_long_tool_name_that_definitely_exceeds_the_maximum_allowed_length_for_names' - const result = createNamespacedToolName('server', longTool, tools, toolNameMapping) + const result = createNamespacedToolName('server', longTool, tools, toolNameMapping, new Set()) // Skip length check and use string comparison with the actual implementation behavior expect(toolNameMapping.get(result)).to.deep.equal({ serverName: 'server', @@ -546,7 +547,7 @@ describe('createNamespacedToolName', () => { it('truncates tool name and adds suffix when it exceeds MAX_TOOL_NAME_LENGTH', () => { const longTool = 'Smartanalyzerthatreadssummariescreatesmappingrulesandupdatespayloads' - const result = createNamespacedToolName('ConnectiveRx', longTool, tools, toolNameMapping) + const result = createNamespacedToolName('ConnectiveRx', longTool, tools, toolNameMapping, new Set()) expect(result.length).to.equal(MAX_TOOL_NAME_LENGTH) expect(tools.has(result)).to.be.true expect(toolNameMapping.get(result)).to.deep.equal({ @@ -554,6 +555,105 @@ describe('createNamespacedToolName', () => { toolName: longTool, }) }) + + it('namespaces an MCP tool that collides with a reserved built-in name', () => { + const reserved = new Set(['fsRead', 'fsWrite', 'executeBash']) + const result = createNamespacedToolName('evil', 'fsRead', tools, toolNameMapping, reserved) + expect(result).to.equal('evil___fsRead') + expect(tools.has('fsRead')).to.be.false + expect(toolNameMapping.has('fsRead')).to.be.false + expect(toolNameMapping.get('evil___fsRead')).to.deep.equal({ + serverName: 'evil', + toolName: 'fsRead', + }) + }) + + it('does not reuse a stale mapping that points at a reserved built-in name', () => { + const reserved = new Set(['fsRead']) + toolNameMapping.set('fsRead', { serverName: 'evil', toolName: 'fsRead' }) + const result = createNamespacedToolName('evil', 'fsRead', tools, toolNameMapping, reserved) + expect(result).to.equal('evil___fsRead') + expect(toolNameMapping.has('fsRead')).to.be.false + }) + + it('still prefers the bare tool name when it is not reserved', () => { + const reserved = new Set(['fsRead']) + const result = createNamespacedToolName('github', 'create_issue', tools, toolNameMapping, reserved) + expect(result).to.equal('create_issue') + }) + + it('refuses every statically dispatched built-in name even when it is not registered', () => { + const reserved = getReservedBuiltInToolNames([]) + + for (const builtIn of STATIC_BUILT_IN_TOOL_NAMES) { + const result = createNamespacedToolName('evil', builtIn, tools, toolNameMapping, reserved) + expect(result, `${builtIn} must not be claimed`).to.equal(`evil___${builtIn}`) + expect(tools.has(builtIn), `${builtIn} must stay unclaimed`).to.be.false + } + }) + + it('also reserves dynamically registered built-in names', () => { + const reserved = getReservedBuiltInToolNames(['lspGetDocuments']) + const result = createNamespacedToolName('evil', 'lspGetDocuments', tools, toolNameMapping, reserved) + + expect(result).to.equal('evil___lspGetDocuments') + expect(tools.has('lspGetDocuments')).to.be.false + }) + + it('gives two servers advertising the same built-in name distinct namespaced names', () => { + const reserved = new Set(['fsRead']) + const first = createNamespacedToolName('alpha', 'fsRead', tools, toolNameMapping, reserved) + const second = createNamespacedToolName('beta', 'fsRead', tools, toolNameMapping, reserved) + expect(first).to.equal('alpha___fsRead') + expect(second).to.equal('beta___fsRead') + expect(tools.has('fsRead')).to.be.false + }) + + it('reuses the namespaced name when the same server tool is seen again', () => { + const reserved = new Set(['fsRead']) + const first = createNamespacedToolName('evil', 'fsRead', tools, toolNameMapping, reserved) + const second = createNamespacedToolName('evil', 'fsRead', tools, toolNameMapping, reserved) + expect(second).to.equal(first) + expect(toolNameMapping.size).to.equal(1) + }) + + it('keeps a truncated namespaced name out of the reserved set', () => { + const longServer = 'very_long_server_name_that_definitely_exceeds_the_maximum_allowed_length' + const reserved = new Set(['fsRead']) + const result = createNamespacedToolName(longServer, 'fsRead', tools, toolNameMapping, reserved) + expect(result.length).to.be.lessThanOrEqual(MAX_TOOL_NAME_LENGTH) + expect(result.endsWith('___fsRead')).to.be.true + expect(reserved.has(result)).to.be.false + }) + + it('does not fall back onto a reserved name via the numeric suffix path', () => { + // Force the namespaced form to be unavailable so the suffix path runs. + tools.add('evil___fsRead') + const reserved = new Set(['fsRead', 'fsRead1']) + const result = createNamespacedToolName('evil', 'fsRead', tools, toolNameMapping, reserved) + expect(reserved.has(result)).to.be.false + expect(result).to.not.equal('fsRead') + expect(result).to.not.equal('fsRead1') + }) + + it('refuses a name that only becomes a built-in name after sanitization', () => { + const reserved = new Set(['fsRead']) + + // sanitizeName strips disallowed characters, so these all collapse to + // 'fsRead' and must not be allowed to claim it. + for (const advertised of ['fs Read', 'fs/Read', 'fsRead!', 'fs.Read']) { + const localTools = new Set() + const localMapping = new Map() + const result = createNamespacedToolName('evil', advertised, localTools, localMapping, reserved) + expect(result, `${advertised} must not claim fsRead`).to.equal('evil___fsRead') + expect(localTools.has('fsRead'), `${advertised} must leave fsRead unclaimed`).to.be.false + } + }) + + it('supports an explicit empty reserved-name set', () => { + const result = createNamespacedToolName('evil', 'fsRead', tools, toolNameMapping, new Set()) + expect(result).to.equal('fsRead') + }) }) describe('normalizePathFromUri', () => { diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpUtils.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpUtils.ts index 6945542078..66c095e22a 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpUtils.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpUtils.ts @@ -1269,17 +1269,28 @@ export function findServerInRegistry(registry: McpRegistryData, serverName: stri * Create a namespaced tool name from server and tool names. * Handles truncation and conflicts according to specific rules. * Also stores the mapping from namespaced name back to original names. + * + * `reservedNames` holds names an MCP tool must never occupy (the built-in tool + * names). Without it an MCP server could advertise e.g. `fsRead` and win the + * bare name, so the request would be routed through the built-in permission + * branch instead of the MCP one. */ export function createNamespacedToolName( serverName: string, toolName: string, allNamespacedTools: Set, - toolNameMapping: Map + toolNameMapping: Map, + reservedNames: Set ): string { // First, check if this server/tool combination already has a mapping // If it does, reuse that name to maintain consistency across reinitializations for (const [existingName, mapping] of toolNameMapping.entries()) { if (mapping.serverName === serverName && mapping.toolName === toolName) { + // Never reuse a stale mapping that points at a reserved built-in name. + if (reservedNames.has(existingName)) { + toolNameMapping.delete(existingName) + break + } // If the name is already in the set, it's already registered // If not, add it to the set if (!allNamespacedTools.has(existingName)) { @@ -1292,8 +1303,13 @@ export function createNamespacedToolName( // Sanitize the tool name const sanitizedToolName = sanitizeName(toolName) - // First try to use just the tool name if it's not already in use and fits within length limit - if (sanitizedToolName.length <= MAX_TOOL_NAME_LENGTH && !allNamespacedTools.has(sanitizedToolName)) { + // First try to use just the tool name if it's not already in use, is not a + // reserved built-in name, and fits within length limit + if ( + sanitizedToolName.length <= MAX_TOOL_NAME_LENGTH && + !allNamespacedTools.has(sanitizedToolName) && + !reservedNames.has(sanitizedToolName) + ) { allNamespacedTools.add(sanitizedToolName) toolNameMapping.set(sanitizedToolName, { serverName, toolName }) return sanitizedToolName @@ -1304,7 +1320,7 @@ export function createNamespacedToolName( const fullName = `${serverName}${sep}${sanitizedToolName}` // If the full name fits and is unique, use it - if (fullName.length <= MAX_TOOL_NAME_LENGTH && !allNamespacedTools.has(fullName)) { + if (fullName.length <= MAX_TOOL_NAME_LENGTH && !allNamespacedTools.has(fullName) && !reservedNames.has(fullName)) { allNamespacedTools.add(fullName) toolNameMapping.set(fullName, { serverName, toolName }) return fullName @@ -1317,7 +1333,7 @@ export function createNamespacedToolName( const truncatedServer = serverName.substring(0, maxServerLength) const namespacedName = `${truncatedServer}${sep}${sanitizedToolName}` - if (!allNamespacedTools.has(namespacedName)) { + if (!allNamespacedTools.has(namespacedName) && !reservedNames.has(namespacedName)) { allNamespacedTools.add(namespacedName) toolNameMapping.set(namespacedName, { serverName, toolName }) return namespacedName @@ -1345,7 +1361,7 @@ export function createNamespacedToolName( candidateName = `${truncatedTool}${suffix}` } - if (!allNamespacedTools.has(candidateName)) { + if (!allNamespacedTools.has(candidateName) && !reservedNames.has(candidateName)) { allNamespacedTools.add(candidateName) toolNameMapping.set(candidateName, { serverName, toolName }) return candidateName diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/toolServer.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/toolServer.ts index bad50794e4..debdf6de92 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/toolServer.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/toolServer.ts @@ -30,6 +30,7 @@ import { DisplayFindings } from './qCodeAnalysis/displayFindings' import { ProfileStatusMonitor } from './mcp/profileStatusMonitor' import { AmazonQTokenServiceManager } from '../../../shared/amazonQServiceManager/AmazonQTokenServiceManager' import { SERVICE_MANAGER_TIMEOUT_MS, SERVICE_MANAGER_POLL_INTERVAL_MS } from '../constants/constants' +import { getReservedBuiltInToolNames } from '../constants/toolConstants' import { isUsingIAMAuth } from '../../../shared/utils' export const FsToolsServer: Server = ({ workspace, logging, agent, lsp }) => { @@ -282,8 +283,14 @@ export const McpToolsServer: Server = ({ function removeAllMcpTools(): void { logging.info('Removing all MCP tools due to admin configuration') + const builtInToolNames = getReservedBuiltInToolNames(agent.getBuiltInToolNames()) for (const [server, toolNames] of Object.entries(registered)) { for (const name of toolNames) { + // Never unregister a built-in tool while cleaning up MCP tools. + if (builtInToolNames.has(name)) { + logging.warn(`MCP: refusing to remove built-in tool name ${name}`) + continue + } agent.removeTool(name) allNamespacedTools.delete(name) logging.info(`MCP: removed tool ${name}`) @@ -315,8 +322,16 @@ export const McpToolsServer: Server = ({ } function registerServerTools(server: string, defs: McpToolDefinition[]) { + // Built-in tool names are reserved: an MCP tool must never register under + // one, and cleanup must never unregister one. + const builtInToolNames = getReservedBuiltInToolNames(agent.getBuiltInToolNames()) + // 1) remove old tools for (const name of registered[server] ?? []) { + if (builtInToolNames.has(name)) { + logging.warn(`MCP: refusing to remove built-in tool name ${name}`) + continue + } agent.removeTool(name) allNamespacedTools.delete(name) } @@ -339,7 +354,8 @@ export const McpToolsServer: Server = ({ def.serverName, def.toolName, allNamespacedTools, - toolNameMapping + toolNameMapping, + builtInToolNames ) const tool = new McpTool({ logging, workspace, lsp }, def)