From fdd66ebf59c878b01feff5a00415e577f85153c1 Mon Sep 17 00:00:00 2001 From: Andrey Avtomonov Date: Mon, 18 May 2026 09:40:24 +0200 Subject: [PATCH] fix: restrict claude-code mcp servers --- .../src/llm/claude-code-runtime.test.ts | 33 +++++++++++++++++++ .../context/src/llm/claude-code-runtime.ts | 23 +++++++++++-- 2 files changed, 54 insertions(+), 2 deletions(-) diff --git a/packages/context/src/llm/claude-code-runtime.test.ts b/packages/context/src/llm/claude-code-runtime.test.ts index f69c5d75..38959140 100644 --- a/packages/context/src/llm/claude-code-runtime.test.ts +++ b/packages/context/src/llm/claude-code-runtime.test.ts @@ -78,6 +78,11 @@ describe('ClaudeCodeKtxLlmRuntime', () => { skills: [], plugins: [], tools: [], + managedSettings: { + allowManagedMcpServersOnly: true, + allowedMcpServers: [], + }, + strictMcpConfig: true, allowedTools: [], permissionMode: 'dontAsk', persistSession: false, @@ -144,6 +149,11 @@ describe('ClaudeCodeKtxLlmRuntime', () => { const options = query.mock.calls[0][0].options; expect(options.allowedTools).toEqual(['mcp__ktx__load_skill']); + expect(options.managedSettings).toEqual({ + allowManagedMcpServersOnly: true, + allowedMcpServers: [{ serverName: 'ktx' }], + }); + expect(options.strictMcpConfig).toBe(true); expect(await options.canUseTool('mcp__ktx__load_skill', {}, { signal: new AbortController().signal, toolUseID: '1' })).toEqual({ behavior: 'allow', toolUseID: '1', @@ -176,6 +186,11 @@ describe('ClaudeCodeKtxLlmRuntime', () => { skills: [], plugins: [], tools: [], + managedSettings: { + allowManagedMcpServersOnly: true, + allowedMcpServers: [], + }, + strictMcpConfig: true, allowedTools: [], permissionMode: 'dontAsk', persistSession: false, @@ -268,6 +283,11 @@ describe('ClaudeCodeKtxLlmRuntime', () => { const options = query.mock.calls[0][0].options; expect(options.allowedTools).toEqual(['mcp__ktx__load_skill']); + expect(options.managedSettings).toEqual({ + allowManagedMcpServersOnly: true, + allowedMcpServers: [{ serverName: 'ktx' }], + }); + expect(options.strictMcpConfig).toBe(true); expect(await options.canUseTool('mcp__ktx__load_skill', {}, { signal: new AbortController().signal, toolUseID: '1' })).toEqual({ behavior: 'allow', toolUseID: '1', @@ -334,6 +354,10 @@ describe('ClaudeCodeKtxLlmRuntime', () => { answer: 'yes', }); expect(objectQuery.mock.calls[0][0].options.env).toEqual(expect.objectContaining({ PATH: '/usr/bin' })); + expect(objectQuery.mock.calls[0][0].options.managedSettings).toEqual({ + allowManagedMcpServersOnly: true, + allowedMcpServers: [], + }); expect(objectQuery.mock.calls[0][0].options.env).not.toEqual( expect.objectContaining({ ANTHROPIC_API_KEY: 'sk-ant-test', AWS_PROFILE: 'prod' }), // pragma: allowlist secret ); @@ -374,6 +398,10 @@ describe('ClaudeCodeKtxLlmRuntime', () => { telemetryTags: { operationName: 'test' }, }); expect(agentQuery.mock.calls[0][0].options.env).toEqual(expect.objectContaining({ HOME: '/Users/test' })); + expect(agentQuery.mock.calls[0][0].options.managedSettings).toEqual({ + allowManagedMcpServersOnly: true, + allowedMcpServers: [{ serverName: 'ktx' }], + }); expect(agentQuery.mock.calls[0][0].options.env).not.toEqual( expect.objectContaining({ ANTHROPIC_AUTH_TOKEN: 'token', CLAUDE_CODE_USE_VERTEX: '1' }), ); @@ -442,6 +470,11 @@ describe('ClaudeCodeKtxLlmRuntime', () => { skills: [], plugins: [], tools: [], + managedSettings: { + allowManagedMcpServersOnly: true, + allowedMcpServers: [], + }, + strictMcpConfig: true, allowedTools: [], persistSession: false, env: expect.not.objectContaining({ ANTHROPIC_API_KEY: 'sk-ant-test' }), diff --git a/packages/context/src/llm/claude-code-runtime.ts b/packages/context/src/llm/claude-code-runtime.ts index 5d8edf26..bf815445 100644 --- a/packages/context/src/llm/claude-code-runtime.ts +++ b/packages/context/src/llm/claude-code-runtime.ts @@ -45,6 +45,8 @@ const BUILTIN_TOOLS = [ 'TodoWrite', ]; +const KTX_MCP_SERVER_NAME = 'ktx'; + function isResult(message: SDKMessage): message is SDKResultMessage { return message.type === 'result'; } @@ -113,7 +115,14 @@ function assertInitIsolation( } function expectedMcpServerNames(tools: KtxRuntimeToolSet | undefined): Set { - return tools && Object.keys(tools).length > 0 ? new Set(['ktx']) : new Set(); + return tools && Object.keys(tools).length > 0 ? new Set([KTX_MCP_SERVER_NAME]) : new Set(); +} + +function managedMcpSettings(serverNames: string[]): NonNullable { + return { + allowManagedMcpServersOnly: true, + allowedMcpServers: serverNames.map((serverName) => ({ serverName })), + }; } function baseOptions(input: { @@ -125,6 +134,7 @@ function baseOptions(input: { }): Options { const toolIds = mcpToolIds(input.tools ?? {}); const allowedToolIds = new Set(toolIds); + const expectedServerNames = [...expectedMcpServerNames(input.tools)]; return { cwd: input.projectDir, model: input.model, @@ -133,6 +143,8 @@ function baseOptions(input: { skills: [], plugins: [], tools: [], + managedSettings: managedMcpSettings(expectedServerNames), + strictMcpConfig: true, allowedTools: toolIds, disallowedTools: BUILTIN_TOOLS, canUseTool: async (toolName, _toolInput, options) => @@ -147,7 +159,14 @@ function baseOptions(input: { persistSession: false, env: createKtxClaudeCodeEnv(input.env), ...(input.tools && Object.keys(input.tools).length > 0 - ? { mcpServers: { ktx: createSdkMcpServer({ name: 'ktx', tools: createClaudeSdkTools(input.tools) }) } } + ? { + mcpServers: { + [KTX_MCP_SERVER_NAME]: createSdkMcpServer({ + name: KTX_MCP_SERVER_NAME, + tools: createClaudeSdkTools(input.tools), + }), + }, + } : {}), }; }