apunkt/ktx
1
0
Fork 0
mirror of https://github.com/Kaelio/ktx.git synced 2026-06-07 07:55:13 +02:00
Code Issues Projects Releases Packages Wiki Activity

fix: restrict claude-code mcp servers

Browse source
This commit is contained in:
Andrey Avtomonov 2026-05-18 09:40:24 +02:00
parent 47c4377816
commit fdd66ebf59
2 changed files with 54 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

33
packages/context/src/llm/claude-code-runtime.test.ts
View file

@ -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' }),

23
packages/context/src/llm/claude-code-runtime.ts
View file

@ -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<string> {
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<Options['managedSettings']> {
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),
}),
},
}
: {}),
};
}