Add custom MCP server management for Studio Code

apps/cli/ai/mcp-client-tools.ts

@@ -0,0 +1,148 @@
+// eslint-disable-next-line import-x/no-unresolved -- subpath resolved via package exports, which the lint resolver doesn't follow
+import { Client } from '@modelcontextprotocol/sdk/client/index.js';
+// eslint-disable-next-line import-x/no-unresolved -- subpath resolved via package exports, which the lint resolver doesn't follow
+import { SSEClientTransport } from '@modelcontextprotocol/sdk/client/sse.js';
+// eslint-disable-next-line import-x/no-unresolved -- subpath resolved via package exports, which the lint resolver doesn't follow
+import { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js';
+// eslint-disable-next-line import-x/no-unresolved -- subpath resolved via package exports, which the lint resolver doesn't follow
+import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js';
+import { readStudioCodeMcpConfig, type StudioCodeMcpServerConfig } from 'cli/ai/mcp-config';
+import type { Transport } from '@modelcontextprotocol/sdk/shared/transport.js';
+import type { ToolContent } from 'cli/ai/tools/define-tool';
+
+type McpClientTool = {
+	name: string;
+	label: string;
+	description: string;
+	parameters: unknown;
+	execute: (
+		toolCallId: string,
+		params: Record< string, unknown >
+	) => Promise< { content: ToolContent[]; details?: unknown } >;
+};
+
+type ConnectedMcpServer = {
+	client: Client;
+	transport: Transport;
+};
+
+function normalizeMcpToolName( serverName: string, toolName: string ): string {
+	return `mcp__${ serverName }__${ toolName.replace( /[^a-zA-Z0-9_-]/g, '_' ) }`;
+}
+
+function buildRequestInit(
+	headers: Record< string, string > | undefined
+): RequestInit | undefined {
+	return headers ? { headers } : undefined;
+}
+
+function buildStdioEnv( env: Record< string, string > | undefined ): Record< string, string > {
+	const inherited = Object.fromEntries(
+		Object.entries( process.env ).filter(
+			( entry ): entry is [ string, string ] => typeof entry[ 1 ] === 'string'
+		)
+	);
+	return { ...inherited, ...( env ?? {} ) };
+}
+
+function createTransport( server: StudioCodeMcpServerConfig ): Transport {
+	if ( 'command' in server ) {
+		return new StdioClientTransport( {
+			command: server.command,
+			args: server.args,
+			env: buildStdioEnv( server.env ),
+			cwd: process.cwd(),
+			stderr: 'pipe',
+		} );
+	}
+
+	if ( server.type === 'sse' ) {
+		return new SSEClientTransport( new URL( server.url ), {
+			requestInit: buildRequestInit( server.headers ),
+			eventSourceInit: buildRequestInit( server.headers ) as never,
+		} );
+	}
+
+	return new StreamableHTTPClientTransport( new URL( server.url ), {
+		requestInit: buildRequestInit( server.headers ),
+	} );
+}
+
+async function connectMcpServer(
+	serverName: string,
+	server: StudioCodeMcpServerConfig
+): Promise< ConnectedMcpServer > {
+	const client = new Client( { name: `studio-code-${ serverName }`, version: '1.0.0' }, {} );
+	const transport = createTransport( server );
+	try {
+		await client.connect( transport );
+		return { client, transport };
+	} catch ( error ) {
+		await transport.close().catch( () => {} );
+		throw error;
+	}
+}
+
+function toToolContent( content: unknown ): ToolContent[] {
+	if ( ! Array.isArray( content ) ) {
+		return [ { type: 'text', text: JSON.stringify( content ) } ];
+	}
+
+	return content.map( ( item ) => {
+		if ( item && typeof item === 'object' ) {
+			const block = item as Record< string, unknown >;
+			if ( block.type === 'text' && typeof block.text === 'string' ) {
+				return { type: 'text', text: block.text };
+			}
+			if (
+				block.type === 'image' &&
+				typeof block.data === 'string' &&
+				typeof block.mimeType === 'string'
+			) {
+				return { type: 'image', data: block.data, mimeType: block.mimeType };
+			}
+		}
+
+		return { type: 'text', text: JSON.stringify( item ) };
+	} );
+}
+
+export async function resolveCustomMcpTools(): Promise< McpClientTool[] > {
+	const config = await readStudioCodeMcpConfig();
+	const tools: McpClientTool[] = [];
+
+	for ( const [ serverName, server ] of Object.entries( config.mcpServers ) ) {
+		const { client } = await connectMcpServer( serverName, server );
+		const { tools: serverTools } = await client.listTools();
+
+		for ( const tool of serverTools ) {
+			const name = normalizeMcpToolName( serverName, tool.name );
+			tools.push( {
+				name,
+				label: name,
+				description:
+					tool.description ?? `Tool ${ tool.name } from custom MCP server ${ serverName }`,
+				parameters: tool.inputSchema,
+				execute: async ( _toolCallId, params ) => {
+					const result = await client.callTool( {
+						name: tool.name,
+						arguments: params,
+					} );
+
+					if ( 'toolResult' in result ) {
+						return { content: [ { type: 'text', text: JSON.stringify( result.toolResult ) } ] };
+					}
+
+					return {
+						content: toToolContent( result.content ),
+						details: result.structuredContent
+							? { structuredContent: result.structuredContent }
+							: undefined,
+					};
+				},
+			} );
+		}
+	}
+
+	return tools;
+}

apps/cli/ai/mcp-config.ts

@@ -0,0 +1,126 @@
+import fs from 'fs';
+import path from 'path';
+import { getStudioCodeMcpConfigPath } from '@studio/common/lib/well-known-paths';
+import { readFile, writeFile } from 'atomically';
+import { z } from 'zod';
+
+const RESERVED_SERVER_NAMES = new Set( [ 'studio' ] );
+const SERVER_NAME_PATTERN = /^[a-zA-Z0-9_-]+$/;
+
+const stringRecordSchema = z.record( z.string(), z.string() );
+
+const stdioServerSchema = z
+	.object( {
+		type: z.literal( 'stdio' ).optional(),
+		command: z.string().trim().min( 1 ),
+		args: z.array( z.string() ).optional(),
+		env: stringRecordSchema.optional(),
+	} )
+	.strict();
+
+const httpServerSchema = z
+	.object( {
+		type: z.literal( 'http' ),
+		url: z.string().trim().url(),
+		headers: stringRecordSchema.optional(),
+	} )
+	.strict();
+
+const sseServerSchema = z
+	.object( {
+		type: z.literal( 'sse' ),
+		url: z.string().trim().url(),
+		headers: stringRecordSchema.optional(),
+	} )
+	.strict();
+
+const serverSchema = z.union( [ stdioServerSchema, httpServerSchema, sseServerSchema ] );
+
+const configSchema = z
+	.object( {
+		mcpServers: z.record( z.string(), serverSchema ).default( {} ),
+	} )
+	.passthrough();
+
+export type StudioCodeMcpServerConfig = z.infer< typeof serverSchema >;
+export type StudioCodeMcpConfig = z.infer< typeof configSchema >;
+
+export { getStudioCodeMcpConfigPath };
+
+function validateServerName( name: string ): void {
+	if ( ! SERVER_NAME_PATTERN.test( name ) ) {
+		throw new Error(
+			`Invalid MCP server name "${ name }". Use only letters, numbers, underscores, and hyphens.`
+		);
+	}
+
+	if ( RESERVED_SERVER_NAMES.has( name ) ) {
+		throw new Error(
+			`Invalid MCP server name "${ name }". This name is reserved for Studio's built-in MCP server.`
+		);
+	}
+}
+
+function parseStudioCodeMcpConfig( data: unknown ): StudioCodeMcpConfig {
+	const config = configSchema.parse( data );
+	for ( const name of Object.keys( config.mcpServers ) ) {
+		validateServerName( name );
+	}
+	return config;
+}
+
+export async function readStudioCodeMcpConfig(): Promise< StudioCodeMcpConfig > {
+	const configPath = getStudioCodeMcpConfigPath();
+
+	if ( ! fs.existsSync( configPath ) ) {
+		return { mcpServers: {} };
+	}
+
+	try {
+		const fileContent = await readFile( configPath, { encoding: 'utf8' } );
+		return parseStudioCodeMcpConfig( JSON.parse( fileContent ) );
+	} catch ( error ) {
+		if ( error instanceof z.ZodError ) {
+			throw new Error( `Invalid Studio Code MCP config at ${ configPath }: ${ error.message }` );
+		}
+		if ( error instanceof SyntaxError ) {
+			throw new Error( `Studio Code MCP config at ${ configPath } is not valid JSON.` );
+		}
+		throw error;
+	}
+}
+
+export async function saveStudioCodeMcpConfig( config: StudioCodeMcpConfig ): Promise< void > {
+	const validatedConfig = parseStudioCodeMcpConfig( config );
+	const configDirectory = path.dirname( getStudioCodeMcpConfigPath() );
+	if ( ! fs.existsSync( configDirectory ) ) {
+		fs.mkdirSync( configDirectory, { recursive: true } );
+	}
+
+	await writeFile(
+		getStudioCodeMcpConfigPath(),
+		JSON.stringify( validatedConfig, null, 2 ) + '\n',
+		{ encoding: 'utf8' }
+	);
+}
+
+export async function addStudioCodeMcpServer(
+	name: string,
+	server: StudioCodeMcpServerConfig
+): Promise< void > {
+	validateServerName( name );
+	const config = await readStudioCodeMcpConfig();
+	config.mcpServers[ name ] = server;
+	await saveStudioCodeMcpConfig( config );
+}
+
+export async function removeStudioCodeMcpServer( name: string ): Promise< boolean > {
+	validateServerName( name );
+	const config = await readStudioCodeMcpConfig();
+	if ( ! config.mcpServers[ name ] ) {
+		return false;
+	}
+	delete config.mcpServers[ name ];
+	await saveStudioCodeMcpConfig( config );
+	return true;
+}

apps/cli/ai/runtimes/pi/index.ts

@@ -28,6 +28,7 @@ import {
 	type AiModelId,
 } from '@studio/common/ai/models';
 import { getAiPayloadsPath, getConfigDirectory } from '@studio/common/lib/well-known-paths';
+import { resolveCustomMcpTools } from 'cli/ai/mcp-client-tools';
 import { buildSystemPrompt } from 'cli/ai/system-prompt';
 import { resolveStudioToolDefinitions } from 'cli/ai/tools';
 import { createAskUserQuestionTool } from 'cli/ai/tools/ask-user-question';
@@ -279,7 +280,7 @@ async function createStudioAgentSession(
 			: { chatArtifactsEnabled, remoteSession }
 	);
 
-	const tools = buildAgentTools( config, chatArtifactsEnabled, remoteSession );
+	const tools = await buildAgentTools( config, chatArtifactsEnabled, remoteSession );
 	const toolDefinitions = tools.map( ( tool ) => toToolDefinition( tool, payloadGuardState ) );
 	const { authStorage, modelRegistry } = createModelRegistry( model, family, creds );
 	const settingsManager = createSettingsManager( config.env );
@@ -461,11 +462,11 @@ function toToolDefinition(
 	};
 }
 
-function buildAgentTools(
+async function buildAgentTools(
 	config: ResolvedStudioAgentTurnConfig,
 	chatArtifactsEnabled: boolean,
 	remoteSession: boolean
-): AgentToolAny[] {
+): Promise< AgentToolAny[] > {
 	const isRemoteSite = Boolean(
 		config.activeSite?.remote && config.activeSite?.wpcomSiteId && config.wpcomAccessToken
 	);
@@ -476,6 +477,7 @@ function buildAgentTools(
 
 	const skillToolDef = createSkillTool();
 	const skillTool: AgentToolAny[] = skillToolDef ? [ skillToolDef ] : [];
+	const customMcpTools = ( await resolveCustomMcpTools() ) as unknown as AgentToolAny[];
 
 	const renameTool = ( tool: AgentToolAny, name: string ): AgentToolAny => ( {
 		...tool,
@@ -496,6 +498,7 @@ function buildAgentTools(
 			takeScreenshotTool,
 			createSiteTool,
 			pullSiteTool,
+			...customMcpTools,
 			...remoteScratchTools,
 			...askUserTool,
 			...skillTool,
@@ -515,7 +518,7 @@ function buildAgentTools(
 		emitChatArtifacts: chatArtifactsEnabled,
 		remoteSession,
 	} ) as unknown as AgentToolAny[];
-	return [ ...studioTools, ...askUserTool, ...skillTool, ...piTools ];
+	return [ ...studioTools, ...customMcpTools, ...askUserTool, ...skillTool, ...piTools ];
 }
 
 function parseJsonHeaderEnv( value: string | undefined ): Record< string, string > | undefined {