diff --git a/apps/cli/ai/mcp-client-tools.ts b/apps/cli/ai/mcp-client-tools.ts new file mode 100644 index 0000000000..60814f54e0 --- /dev/null +++ b/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; +} diff --git a/apps/cli/ai/mcp-config.ts b/apps/cli/ai/mcp-config.ts new file mode 100644 index 0000000000..2bd92786f8 --- /dev/null +++ b/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; +} diff --git a/apps/cli/ai/runtimes/pi/index.ts b/apps/cli/ai/runtimes/pi/index.ts index 64289648ae..b4a23d0159 100644 --- a/apps/cli/ai/runtimes/pi/index.ts +++ b/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 { diff --git a/apps/cli/ai/tests/mcp-config.test.ts b/apps/cli/ai/tests/mcp-config.test.ts new file mode 100644 index 0000000000..610b4317c3 --- /dev/null +++ b/apps/cli/ai/tests/mcp-config.test.ts @@ -0,0 +1,122 @@ +import fs from 'fs'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { + addStudioCodeMcpServer, + getStudioCodeMcpConfigPath, + readStudioCodeMcpConfig, + removeStudioCodeMcpServer, + saveStudioCodeMcpConfig, +} from 'cli/ai/mcp-config'; + +const { studioConfigRoot } = vi.hoisted( () => ( { + studioConfigRoot: `${ process.env.TMPDIR || '/tmp/' }studio-code-mcp-config-${ Math.random() + .toString( 36 ) + .slice( 2 ) }`, +} ) ); + +vi.mock( '@studio/common/lib/well-known-paths', async ( importOriginal ) => { + const actual = await importOriginal< typeof import('@studio/common/lib/well-known-paths') >(); + return { + ...actual, + getConfigDirectory: () => studioConfigRoot, + getStudioCodeMcpConfigPath: () => `${ studioConfigRoot }/mcp.json`, + }; +} ); + +describe( 'Studio Code MCP config', () => { + beforeEach( () => { + fs.rmSync( getStudioCodeMcpConfigPath(), { force: true } ); + } ); + + afterEach( () => { + fs.rmSync( studioConfigRoot, { recursive: true, force: true } ); + } ); + + it( 'returns an empty server list when no config file exists', async () => { + await expect( readStudioCodeMcpConfig() ).resolves.toEqual( { mcpServers: {} } ); + } ); + + it( 'reads stdio, http, and sse servers from config', async () => { + const configPath = getStudioCodeMcpConfigPath(); + fs.mkdirSync( studioConfigRoot, { recursive: true } ); + fs.writeFileSync( + configPath, + JSON.stringify( { + mcpServers: { + wporg: { + type: 'stdio', + command: 'node', + args: [ '/tmp/wporg.js' ], + env: { GITHUB_TOKEN: 'token' }, + }, + figma: { + type: 'http', + url: 'https://mcp.figma.com/mcp', + headers: { Authorization: 'Bearer token' }, + }, + legacy: { + type: 'sse', + url: 'https://example.com/sse', + }, + }, + } ) + ); + + await expect( readStudioCodeMcpConfig() ).resolves.toMatchObject( { + mcpServers: { + wporg: { type: 'stdio', command: 'node', args: [ '/tmp/wporg.js' ] }, + figma: { type: 'http', url: 'https://mcp.figma.com/mcp' }, + legacy: { type: 'sse', url: 'https://example.com/sse' }, + }, + } ); + } ); + + it( 'preserves unknown top-level MCP config fields', async () => { + await saveStudioCodeMcpConfig( { + mcpServers: { + wporg: { type: 'stdio', command: 'node' }, + }, + futureSdkField: true, + } ); + + await expect( readStudioCodeMcpConfig() ).resolves.toMatchObject( { + mcpServers: { + wporg: { type: 'stdio', command: 'node' }, + }, + futureSdkField: true, + } ); + } ); + + it( 'rejects the built-in studio server name', async () => { + await expect( + saveStudioCodeMcpConfig( { + mcpServers: { + studio: { type: 'stdio', command: 'node' }, + }, + } ) + ).rejects.toThrow( 'reserved' ); + } ); + + it( 'rejects unsafe server names', async () => { + await expect( + addStudioCodeMcpServer( 'bad name', { type: 'stdio', command: 'node' } ) + ).rejects.toThrow( 'Invalid MCP server name' ); + } ); + + it( 'adds and removes servers', async () => { + await addStudioCodeMcpServer( 'wporg', { + type: 'stdio', + command: 'node', + args: [ '/tmp/wporg.js' ], + } ); + await expect( readStudioCodeMcpConfig() ).resolves.toMatchObject( { + mcpServers: { + wporg: { command: 'node', args: [ '/tmp/wporg.js' ] }, + }, + } ); + + await expect( removeStudioCodeMcpServer( 'wporg' ) ).resolves.toBe( true ); + await expect( removeStudioCodeMcpServer( 'wporg' ) ).resolves.toBe( false ); + await expect( readStudioCodeMcpConfig() ).resolves.toEqual( { mcpServers: {} } ); + } ); +} ); diff --git a/apps/cli/ai/tests/pi-runtime.test.ts b/apps/cli/ai/tests/pi-runtime.test.ts index 7c5ed57253..8b53c69219 100644 --- a/apps/cli/ai/tests/pi-runtime.test.ts +++ b/apps/cli/ai/tests/pi-runtime.test.ts @@ -6,6 +6,7 @@ import type { AiModelId } from '@studio/common/ai/models'; const mocks = vi.hoisted( () => ( { createAgentSession: vi.fn(), + resolveCustomMcpTools: vi.fn(), createdSessions: [] as FakeSession[], nextEvents: null as AgentSessionEvent[] | null, studioRoot: '/tmp/studio-ai-pi-runtime', @@ -50,6 +51,10 @@ vi.mock( 'cli/lib/site-paths', () => ( { getDefaultSitePath: ( siteName: string ) => `${ mocks.studioRoot }/${ siteName }`, } ) ); +vi.mock( 'cli/ai/mcp-client-tools', () => ( { + resolveCustomMcpTools: mocks.resolveCustomMcpTools, +} ) ); + vi.mock( '@studio/common/lib/well-known-paths', async ( importOriginal ) => { const actual = await importOriginal< typeof import('@studio/common/lib/well-known-paths') >(); return { @@ -213,6 +218,8 @@ describe( 'pi runtime', () => { mocks.createdSessions.length = 0; mocks.nextEvents = null; mocks.createAgentSession.mockReset(); + mocks.resolveCustomMcpTools.mockReset(); + mocks.resolveCustomMcpTools.mockResolvedValue( [] ); mocks.createAgentSession.mockImplementation( async ( options: CreateAgentSessionOptions ) => { const session = new FakeSession( options ); mocks.createdSessions.push( session ); @@ -272,6 +279,36 @@ describe( 'pi runtime', () => { expect( mocks.createdSessions[ 0 ].options.model?.input ).toEqual( [ 'text', 'image' ] ); } ); + it( 'includes configured custom MCP tools in the pi tool registry', async () => { + mocks.resolveCustomMcpTools.mockResolvedValue( [ + { + name: 'mcp__figma__get_design_context', + label: 'mcp__figma__get_design_context', + description: 'Fetch Figma design context', + parameters: { type: 'object', properties: {} }, + execute: async () => ( { content: [ { type: 'text', text: 'figma result' } ] } ), + }, + ] ); + + await runRuntime( { + prompt: 'hello', + env: { + OPENAI_API_KEY: 'sk-test', + OPENAI_BASE_URL: 'https://proxy.example.com/v1', + }, + model: 'gpt-5.5', + session: newSession(), + } ); + + const options = mocks.createdSessions[ 0 ].options; + expect( options.tools ).toContain( 'mcp__figma__get_design_context' ); + expect( + ( options.customTools ?? [] ).some( + ( tool ) => tool.name === 'mcp__figma__get_design_context' + ) + ).toBe( true ); + } ); + it( 'rejects oversized direct Write, Edit, and Bash payloads', async () => { await runRuntime( { prompt: 'hello', diff --git a/apps/cli/commands/ai/mcp.ts b/apps/cli/commands/ai/mcp.ts new file mode 100644 index 0000000000..8a1f98db79 --- /dev/null +++ b/apps/cli/commands/ai/mcp.ts @@ -0,0 +1,179 @@ +import { __, sprintf } from '@wordpress/i18n'; +import { + addStudioCodeMcpServer, + getStudioCodeMcpConfigPath, + readStudioCodeMcpConfig, + removeStudioCodeMcpServer, + type StudioCodeMcpServerConfig, +} from 'cli/ai/mcp-config'; +import { Logger, LoggerError } from 'cli/logger'; +import { StudioArgv } from 'cli/types'; + +const logger = new Logger< string >(); + +function parseKeyValuePairs( + pairs: string[] | undefined, + optionName: string +): Record< string, string > | undefined { + if ( ! pairs?.length ) { + return undefined; + } + + const parsed: Record< string, string > = {}; + for ( const pair of pairs ) { + const separatorIndex = pair.indexOf( '=' ); + if ( separatorIndex <= 0 ) { + throw new Error( sprintf( __( '%s entries must use KEY=VALUE format.' ), optionName ) ); + } + parsed[ pair.slice( 0, separatorIndex ) ] = pair.slice( separatorIndex + 1 ); + } + return parsed; +} + +function buildServerConfig( argv: { + type?: string; + command?: string; + arg?: string[]; + url?: string; + env?: string[]; + header?: string[]; +} ): StudioCodeMcpServerConfig { + const type = argv.type ?? ( argv.url ? 'http' : 'stdio' ); + + if ( type === 'stdio' ) { + if ( ! argv.command ) { + throw new Error( __( '--command is required for stdio MCP servers.' ) ); + } + return { + type: 'stdio', + command: argv.command, + args: argv.arg, + env: parseKeyValuePairs( argv.env, '--env' ), + }; + } + + if ( type === 'http' || type === 'sse' ) { + if ( ! argv.url ) { + throw new Error( __( '--url is required for HTTP and SSE MCP servers.' ) ); + } + return { + type, + url: argv.url, + headers: parseKeyValuePairs( argv.header, '--header' ), + }; + } + + throw new Error( __( '--type must be one of: stdio, http, sse.' ) ); +} + +export const registerCommand = ( yargs: StudioArgv ) => { + yargs.command( { + command: 'list', + describe: __( 'List custom MCP servers' ), + handler: async () => { + const config = await readStudioCodeMcpConfig(); + const entries = Object.entries( config.mcpServers ); + if ( entries.length === 0 ) { + logger.reportSuccess( __( 'No custom MCP servers configured.' ) ); + return; + } + for ( const [ name, server ] of entries ) { + if ( 'command' in server ) { + console.log( `${ name }\t${ [ server.command, ...( server.args ?? [] ) ].join( ' ' ) }` ); + } else { + console.log( `${ name }\t${ server.type }: ${ server.url }` ); + } + } + }, + } ); + + yargs.command( { + command: 'add ', + describe: __( 'Add or replace a custom MCP server' ), + builder: ( addYargs ) => + addYargs + .positional( 'name', { + type: 'string', + description: __( 'MCP server name' ), + demandOption: true, + } ) + .option( 'type', { + type: 'string', + choices: [ 'stdio', 'http', 'sse' ] as const, + description: __( 'MCP server transport type' ), + } ) + .option( 'command', { + type: 'string', + description: __( 'Command for stdio MCP servers' ), + } ) + .option( 'arg', { + type: 'string', + array: true, + description: __( 'Argument for stdio MCP servers; repeat for multiple arguments' ), + } ) + .option( 'url', { + type: 'string', + description: __( 'URL for HTTP or SSE MCP servers' ), + } ) + .option( 'env', { + type: 'string', + array: true, + description: __( 'Environment variable for stdio servers in KEY=VALUE format' ), + } ) + .option( 'header', { + type: 'string', + array: true, + description: __( 'Header for HTTP or SSE servers in KEY=VALUE format' ), + } ), + handler: async ( argv ) => { + try { + const typedArgv = argv as { + name: string; + type?: string; + command?: string; + arg?: string[]; + url?: string; + env?: string[]; + header?: string[]; + }; + const server = buildServerConfig( typedArgv ); + await addStudioCodeMcpServer( typedArgv.name, server ); + logger.reportSuccess( sprintf( __( 'Added custom MCP server "%s".' ), typedArgv.name ) ); + } catch ( error ) { + logger.reportError( new LoggerError( __( 'Failed to add custom MCP server' ), error ) ); + } + }, + } ); + + yargs.command( { + command: 'remove ', + describe: __( 'Remove a custom MCP server' ), + builder: ( removeYargs ) => + removeYargs.positional( 'name', { + type: 'string', + description: __( 'MCP server name' ), + demandOption: true, + } ), + handler: async ( argv ) => { + const name = ( argv as { name: string } ).name; + const removed = await removeStudioCodeMcpServer( name ); + if ( removed ) { + logger.reportSuccess( sprintf( __( 'Removed custom MCP server "%s".' ), name ) ); + return; + } + logger.reportWarning( sprintf( __( 'No custom MCP server named "%s".' ), name ) ); + }, + } ); + + yargs.command( { + command: 'path', + describe: __( 'Print the custom MCP config path' ), + handler: () => { + console.log( getStudioCodeMcpConfigPath() ); + }, + } ); + + return yargs + .version( false ) + .demandCommand( 1, __( 'You must provide a valid code mcp command' ) ); +}; diff --git a/apps/cli/commands/ai/tests/mcp.test.ts b/apps/cli/commands/ai/tests/mcp.test.ts new file mode 100644 index 0000000000..f9016a5397 --- /dev/null +++ b/apps/cli/commands/ai/tests/mcp.test.ts @@ -0,0 +1,163 @@ +import fs from 'fs'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import yargs from 'yargs/yargs'; +import { getStudioCodeMcpConfigPath, readStudioCodeMcpConfig } from 'cli/ai/mcp-config'; +import { registerCommand } from 'cli/commands/ai/mcp'; +import { StudioArgv } from 'cli/types'; + +const { reportErrorMock, reportSuccessMock, reportWarningMock } = vi.hoisted( () => ( { + reportErrorMock: vi.fn(), + reportSuccessMock: vi.fn(), + reportWarningMock: vi.fn(), +} ) ); + +vi.mock( 'cli/logger', () => ( { + Logger: class { + reportError = reportErrorMock; + reportSuccess = reportSuccessMock; + reportWarning = reportWarningMock; + }, + LoggerError: class LoggerError extends Error { + constructor( message: string, previousError?: unknown ) { + super( + previousError instanceof Error ? `${ message }: ${ previousError.message }` : message + ); + } + }, +} ) ); + +const { studioSitesRoot } = vi.hoisted( () => ( { + studioSitesRoot: `${ process.env.TMPDIR || '/tmp/' }studio-code-mcp-command-root-${ Math.random() + .toString( 36 ) + .slice( 2 ) }`, +} ) ); + +vi.mock( 'cli/lib/site-paths', () => ( { + STUDIO_SITES_ROOT: studioSitesRoot, +} ) ); + +describe( 'CLI: studio code mcp', () => { + beforeEach( () => { + vi.clearAllMocks(); + fs.rmSync( getStudioCodeMcpConfigPath(), { force: true } ); + } ); + + afterEach( () => { + fs.rmSync( studioSitesRoot, { recursive: true, force: true } ); + } ); + + function buildParser(): StudioArgv { + const parser = yargs( [] ) + .scriptName( 'studio code mcp' ) + .strict() + .exitProcess( false ) as StudioArgv; + registerCommand( parser ); + return parser; + } + + it( 'adds a stdio server', async () => { + await buildParser().parseAsync( [ + 'add', + 'wporg', + '--command', + 'node', + '--arg', + '/tmp/wporg.js', + '--env', + 'GITHUB_TOKEN=test', + ] ); + + await expect( readStudioCodeMcpConfig() ).resolves.toEqual( { + mcpServers: { + wporg: { + type: 'stdio', + command: 'node', + args: [ '/tmp/wporg.js' ], + env: { GITHUB_TOKEN: 'test' }, + }, + }, + } ); + expect( reportSuccessMock ).toHaveBeenCalledWith( 'Added custom MCP server "wporg".' ); + } ); + + it( 'adds an http server', async () => { + await buildParser().parseAsync( [ + 'add', + 'figma', + '--type', + 'http', + '--url', + 'https://mcp.figma.com/mcp', + '--header', + 'Authorization=Bearer test', + ] ); + + await expect( readStudioCodeMcpConfig() ).resolves.toEqual( { + mcpServers: { + figma: { + type: 'http', + url: 'https://mcp.figma.com/mcp', + headers: { Authorization: 'Bearer test' }, + }, + }, + } ); + } ); + + it( 'reports invalid server definitions without throwing raw stacks', async () => { + await buildParser().parseAsync( [ 'add', 'wporg' ] ); + await buildParser().parseAsync( [ 'add', 'wporg', '--command', 'node', '--env', 'broken' ] ); + await buildParser().parseAsync( [ 'add', 'studio', '--command', 'node' ] ); + + expect( reportErrorMock ).toHaveBeenCalledTimes( 3 ); + expect( reportErrorMock.mock.calls[ 0 ][ 0 ].message ).toContain( '--command is required' ); + expect( reportErrorMock.mock.calls[ 1 ][ 0 ].message ).toContain( 'KEY=VALUE' ); + expect( reportErrorMock.mock.calls[ 2 ][ 0 ].message ).toContain( 'reserved' ); + } ); + + it( 'lists configured servers', async () => { + await buildParser().parseAsync( [ + 'add', + 'wporg', + '--command', + 'node', + '--arg', + '/tmp/wporg.js', + ] ); + vi.clearAllMocks(); + const output: string[] = []; + const consoleSpy = vi + .spyOn( console, 'log' ) + .mockImplementation( ( value ) => output.push( String( value ) ) ); + + await buildParser().parseAsync( [ 'list' ] ); + + expect( output ).toEqual( [ 'wporg\tnode /tmp/wporg.js' ] ); + consoleSpy.mockRestore(); + } ); + + it( 'removes a configured server', async () => { + await buildParser().parseAsync( [ 'add', 'wporg', '--command', 'node' ] ); + await buildParser().parseAsync( [ 'remove', 'wporg' ] ); + + await expect( readStudioCodeMcpConfig() ).resolves.toEqual( { mcpServers: {} } ); + expect( reportSuccessMock ).toHaveBeenCalledWith( 'Removed custom MCP server "wporg".' ); + } ); + + it( 'warns when removing an unknown server', async () => { + await buildParser().parseAsync( [ 'remove', 'wporg' ] ); + + expect( reportWarningMock ).toHaveBeenCalledWith( 'No custom MCP server named "wporg".' ); + } ); + + it( 'prints the config path', async () => { + const output: string[] = []; + const consoleSpy = vi + .spyOn( console, 'log' ) + .mockImplementation( ( value ) => output.push( String( value ) ) ); + + await buildParser().parseAsync( [ 'path' ] ); + + expect( output ).toEqual( [ getStudioCodeMcpConfigPath() ] ); + consoleSpy.mockRestore(); + } ); +} ); diff --git a/apps/cli/index.ts b/apps/cli/index.ts index 39f29c1214..df949c5e6f 100644 --- a/apps/cli/index.ts +++ b/apps/cli/index.ts @@ -160,6 +160,11 @@ async function main() { .demandCommand( 1, __( 'You must provide a valid code sessions command' ) ); } ); + aiYargs.command( 'mcp', __( 'Manage custom MCP servers' ), async ( mcpYargs ) => { + const { registerCommand: registerAiMcpCommand } = await import( 'cli/commands/ai/mcp' ); + registerAiMcpCommand( mcpYargs ); + mcpYargs.version( false ); + } ); aiYargs .example( [ [ 'studio code', __( 'Start an interactive chat with the AI agent' ) ], diff --git a/tools/common/lib/well-known-paths.ts b/tools/common/lib/well-known-paths.ts index d886e152eb..e5e6c5ea38 100644 --- a/tools/common/lib/well-known-paths.ts +++ b/tools/common/lib/well-known-paths.ts @@ -31,6 +31,10 @@ export function getCliConfigPath(): string { return path.join( getConfigDirectory(), 'cli.json' ); } +export function getStudioCodeMcpConfigPath(): string { + return path.join( getConfigDirectory(), 'mcp.json' ); +} + export function getCertificatesPath(): string { return path.join( getConfigDirectory(), 'certificates' ); }