import * as vscode from "vscode"; import type { McpManagedServerInfo, McpServerConfig as SdkMcpServerConfig, McpTestResult, } from "@moonshot-ai/kimi-code-sdk"; import { Events, Methods } from "../../shared/bridge"; import { MCP_SECRET_MASK, type MCPServerConfig, type MCPTestResult, type UpdateMCPServerRequest, } from "../../shared/legacy-sdk"; import type { Handler } from "./types"; const SENSITIVE_MCP_KEY_WORDS = new Set([ "authorization", "cookie", "credential", "credentials", "passwd", "password", "secret", "token", ]); interface NameParams { name: string } export const mcpHandlers: Record> = { [Methods.GetMCPServers]: async (_, ctx): Promise => { return listWorkspaceServers(ctx); }, [Methods.AddMCPServer]: async (params: MCPServerConfig, ctx): Promise => { const server = restoreMaskedSecrets(undefined, params); await ctx.harness.addMcpServer(toSdkServer(server)); const servers = await listWorkspaceServers(ctx); ctx.broadcast(Events.MCPServersChanged, servers); return servers; }, [Methods.UpdateMCPServer]: async ( params: UpdateMCPServerRequest | MCPServerConfig, ctx, ): Promise => { const request = normalizeUpdateRequest(params); const current = ( await ctx.harness.listMcpServers({ cwd: ctx.workDir ?? undefined }) ).find((server) => server.name === request.originalName); const edited = restoreMaskedSecrets(current, request.server); const next = mergeEditableServer(current, edited, request.replaceEditableFields); await updateOrRenameServer(ctx.harness, request.originalName, current, next); const servers = await listWorkspaceServers(ctx); ctx.broadcast(Events.MCPServersChanged, servers); return servers; }, [Methods.RemoveMCPServer]: async ({ name }: NameParams, ctx): Promise => { await ctx.harness.removeMcpServer(name); const servers = await listWorkspaceServers(ctx); ctx.broadcast(Events.MCPServersChanged, servers); return servers; }, [Methods.AuthMCP]: async ({ name }: NameParams, ctx) => { await vscode.window.withProgress( { location: vscode.ProgressLocation.Notification, title: `Kimi: Authenticating "${name}"...`, cancellable: false, }, async () => { try { await ctx.harness.authenticateMcpServer(name, { onAuthorizationUrl: async (url) => vscode.env.openExternal(vscode.Uri.parse(url)), }); await vscode.window.showInformationMessage(`Kimi: OAuth completed for "${name}"`); } catch (error) { const message = error instanceof Error ? error.message : String(error); await vscode.window.showErrorMessage(`Kimi: OAuth failed for "${name}": ${message}`); throw error; } }, ); return { ok: true }; }, [Methods.ResetAuthMCP]: async ({ name }: NameParams, ctx) => { await vscode.window.withProgress( { location: vscode.ProgressLocation.Notification, title: `Kimi: Resetting auth for "${name}"...`, cancellable: false, }, async () => { try { await ctx.harness.resetMcpServerAuth(name); await vscode.window.showInformationMessage(`Kimi: Auth reset for "${name}"`); } catch (error) { const message = error instanceof Error ? error.message : String(error); await vscode.window.showErrorMessage(`Kimi: Reset auth failed for "${name}": ${message}`); throw error; } }, ); return { ok: true }; }, [Methods.TestMCP]: async ({ name }: NameParams, ctx): Promise => { void vscode.window.showInformationMessage(`Kimi: Testing MCP server "${name}"...`); const result = toWebviewTestResult(await ctx.harness.testMcpServer(name, { cwd: ctx.workDir ?? undefined, })); if (!result.success) { ctx.logError(`MCP server test failed for "${name}"`, new Error(result.output)); } return result; }, }; /** * The workspace-aware server list shown in the modal. The mutation RPCs * (add/update/remove) return a list resolved without a cwd, so the webview * refresh after every mutation must re-list with the workspace cwd — * otherwise project-layer entries drop out of the modal until the next full * load. */ async function listWorkspaceServers(ctx: Parameters[1]): Promise { return toWebviewServers(await ctx.harness.listMcpServers({ cwd: ctx.workDir ?? undefined })); } function toWebviewServers(servers: readonly McpManagedServerInfo[]): MCPServerConfig[] { return servers .filter((server) => server.transport === "stdio" || server.transport === "http") .map((server) => { // The management view's source/origin/mutable tags stay in the webview // payload so the panel can hide mutating controls on read-only entries; // only the nested plugin origin detail is dropped. const { plugin: _plugin, ...config } = server; if (config.transport === "stdio") { return { ...config, env: maskSecretValues(config.env) } as MCPServerConfig; } return { ...config, headers: maskSecretValues(config.headers) } as MCPServerConfig; }); } function maskSecretValues(values: Record | undefined): Record | undefined { if (values === undefined) return undefined; return Object.fromEntries(Object.entries(values).map(([key, value]) => [ key, isSensitiveMcpKey(key) ? MCP_SECRET_MASK : value, ])); } function isSensitiveMcpKey(key: string): boolean { const words = key .replaceAll(/([a-z0-9])([A-Z])/g, "$1 $2") .toLowerCase() .split(/[^a-z0-9]+/) .filter(Boolean); const compact = words.join(""); return words.some((word) => SENSITIVE_MCP_KEY_WORDS.has(word)) || words.includes("key") || compact === "proxyauthorization" || compact === "setcookie" || compact.endsWith("apikey") || compact.endsWith("accesskey") || compact.endsWith("privatekey"); } function restoreMaskedSecrets( current: SdkMcpServerConfig | undefined, edited: MCPServerConfig, ): MCPServerConfig { if (edited.transport === "stdio") { const currentEnv = current?.transport === "stdio" ? current.env : undefined; return { ...edited, env: restoreMaskedRecord("environment variable", currentEnv, edited.env, false), }; } const currentHeaders = current?.transport === "http" ? current.headers : undefined; return { ...edited, headers: restoreMaskedRecord("header", currentHeaders, edited.headers, true), }; } function restoreMaskedRecord( label: string, current: Record | undefined, edited: Record | undefined, caseInsensitiveKeys: boolean, ): Record | undefined { if (edited === undefined) return undefined; return Object.fromEntries(Object.entries(edited).map(([key, value]) => { if (value !== MCP_SECRET_MASK) return [key, value]; const stored = findStoredSecret(current, key, caseInsensitiveKeys); if (stored === undefined) { throw new Error(`Cannot preserve masked MCP ${label} "${key}" because no stored value exists`); } return [key, stored]; })); } function findStoredSecret( current: Record | undefined, key: string, caseInsensitiveKeys: boolean, ): string | undefined { if (current === undefined) return undefined; if (Object.hasOwn(current, key)) return current[key]; if (!caseInsensitiveKeys) return undefined; const normalized = key.toLowerCase(); const match = Object.entries(current).find(([storedKey]) => storedKey.toLowerCase() === normalized); return match?.[1]; } function toSdkServer(server: MCPServerConfig): SdkMcpServerConfig { const name = server.name.trim(); if (server.transport === "stdio") { return { name, transport: "stdio", command: server.command?.trim() ?? "", args: server.args, env: server.env, }; } return { name, transport: "http", url: server.url?.trim() ?? "", headers: server.headers, auth: server.auth, bearerTokenEnvVar: server.bearerTokenEnvVar, }; } function mergeEditableServer( current: SdkMcpServerConfig | undefined, edited: MCPServerConfig, replaceEditableFields: boolean, ): SdkMcpServerConfig { const next = toSdkServer(edited); if (current === undefined || current.transport !== next.transport) return next; if (!replaceEditableFields) { return mergeReleasedFormUpdate(current, next); } return { ...current, ...next } as SdkMcpServerConfig; } function mergeReleasedFormUpdate( current: SdkMcpServerConfig, next: SdkMcpServerConfig, ): SdkMcpServerConfig { const defined = Object.fromEntries( Object.entries(next).filter(([, value]) => value !== undefined), ); return { ...current, ...defined } as SdkMcpServerConfig; } function normalizeUpdateRequest( params: UpdateMCPServerRequest | MCPServerConfig, ): { readonly originalName: string; readonly server: MCPServerConfig; readonly replaceEditableFields: boolean; } { if ("server" in params) { return { originalName: params.originalName.trim(), server: params.server, replaceEditableFields: true, }; } return { originalName: params.name.trim(), server: params, replaceEditableFields: false, }; } async function updateOrRenameServer( harness: Pick< Parameters[1]["harness"], "addMcpServer" | "updateMcpServer" | "removeMcpServer" >, originalName: string, current: SdkMcpServerConfig | undefined, next: SdkMcpServerConfig, ): Promise { if (next.name === originalName) { await harness.updateMcpServer(next); return; } if (current === undefined) { throw new Error(`MCP server "${originalName}" was not found`); } await harness.addMcpServer(next); try { await harness.removeMcpServer(originalName); } catch (error) { await harness.removeMcpServer(next.name).catch(() => undefined); throw error; } } function toWebviewTestResult(result: McpTestResult): MCPTestResult { return { success: result.success, output: sanitizeMcpDiagnostic(result.output) }; } function sanitizeMcpDiagnostic(output: string): string { return output .replaceAll(/\bBearer\s+[^\s,;]+/gi, "Bearer [redacted]") .replaceAll( /(["']?(?:authorization|cookie|credentials?|password|passwd|secret|token|api[-_ ]?key)["']?\s*[:=]\s*)("[^"]*"|'[^']*'|[^\s,;]+)/gi, "$1[redacted]", ); }