diff --git a/code/lib/cli-storybook/package.json b/code/lib/cli-storybook/package.json index 2e91d0b1b5e5..85b00077aece 100644 --- a/code/lib/cli-storybook/package.json +++ b/code/lib/cli-storybook/package.json @@ -54,13 +54,15 @@ "envinfo": "^7.14.0", "globby": "^14.1.0", "leven": "^4.0.0", + "memfs": "^4.11.1", "p-limit": "^7.2.0", "picocolors": "^1.1.0", "semver": "^7.7.3", "slash": "^5.0.0", "tiny-invariant": "^1.3.3", "tinyclip": "^0.1.12", - "typescript": "^5.8.3" + "typescript": "^5.8.3", + "valibot": "^1.4.0" }, "publishConfig": { "access": "public" diff --git a/code/lib/cli-storybook/src/ai/mcp/client.test.ts b/code/lib/cli-storybook/src/ai/mcp/client.test.ts new file mode 100644 index 000000000000..6dc23864a145 --- /dev/null +++ b/code/lib/cli-storybook/src/ai/mcp/client.test.ts @@ -0,0 +1,214 @@ +import { describe, expect, it, vi } from 'vitest'; + +import { McpJsonRpcError, callMcpTool, listMcpTools } from './client.ts'; +import type { StorybookInstanceRecord } from './types.ts'; + +const record: StorybookInstanceRecord = { + schemaVersion: 1, + instanceId: 'i-1', + pid: 1, + cwd: '/projects/foo', + url: 'http://localhost:6006', + port: 6006, + mcp: { status: 'ready', endpoint: '/mcp' }, +}; + +const jsonResponse = (body: unknown, status = 200) => + new Response(JSON.stringify(body), { + status, + headers: { 'Content-Type': 'application/json' }, + }); + +const sseResponse = (body: string, status = 200) => + new Response(body, { + status, + headers: { 'Content-Type': 'text/event-stream' }, + }); + +describe('callMcpTool', () => { + it('POSTs a JSON-RPC tools/call request to the endpoint (application/json)', async () => { + const fetchImpl = vi.fn(async () => + jsonResponse({ + jsonrpc: '2.0', + id: 'whatever', + result: { content: [{ type: 'text', text: 'hello' }] }, + }) + ) as unknown as typeof fetch; + + const result = await callMcpTool( + record, + { name: 'list-all-documentation', arguments: { withStoryIds: true } }, + fetchImpl + ); + + expect(result.content).toEqual([{ type: 'text', text: 'hello' }]); + + const call = vi.mocked(fetchImpl).mock.calls[0]; + expect(call[0]).toBe('http://localhost:6006/mcp'); + const init = call[1] as RequestInit; + const headers = init.headers as Record; + expect(headers.Accept).toBe('application/json, text/event-stream'); + expect(headers['X-Storybook-MCP-Proxy']).toBe('true'); + expect(init.signal).toBeInstanceOf(AbortSignal); + const body = JSON.parse(init.body as string); + expect(body).toMatchObject({ + jsonrpc: '2.0', + method: 'tools/call', + params: { + name: 'list-all-documentation', + arguments: { withStoryIds: true }, + }, + }); + expect(typeof body.id).toBe('string'); + }); + + it('resolves the endpoint path against the instance url without mangling the scheme', async () => { + const fetchImpl = vi.fn(async () => + jsonResponse({ jsonrpc: '2.0', id: 'whatever', result: { content: [] } }) + ) as unknown as typeof fetch; + + await callMcpTool( + { ...record, url: 'http://127.0.0.1:6007', mcp: { status: 'ready', endpoint: '/mcp' } }, + { name: 'list-all-documentation' }, + fetchImpl + ); + + expect(vi.mocked(fetchImpl).mock.calls[0][0]).toBe('http://127.0.0.1:6007/mcp'); + }); + + it('parses a single-event SSE response (text/event-stream)', async () => { + const sseBody = + 'event: message\n' + + 'data: {"jsonrpc":"2.0","id":1,"result":{"content":[{"type":"text","text":"hi"}]}}\n' + + '\n'; + const fetchImpl = (async () => sseResponse(sseBody)) as typeof fetch; + + const result = await callMcpTool(record, { name: 'list-all-documentation' }, fetchImpl); + expect(result.content).toEqual([{ type: 'text', text: 'hi' }]); + }); + + it('joins multi-line SSE data correctly', async () => { + const envelope = { + jsonrpc: '2.0', + id: 1, + result: { content: [{ type: 'text', text: 'line\nwith newline' }] }, + }; + const dataLines = JSON.stringify(envelope, null, 2) + .split('\n') + .map((l) => `data: ${l}`) + .join('\n'); + const sseBody = `event: message\n${dataLines}\n\n`; + const fetchImpl = (async () => sseResponse(sseBody)) as typeof fetch; + + const result = await callMcpTool(record, { name: 'list-all-documentation' }, fetchImpl); + expect(result.content?.[0]).toEqual({ type: 'text', text: 'line\nwith newline' }); + }); + + it('throws on SSE responses that contain no data event', async () => { + const fetchImpl = (async () => sseResponse('event: ping\n\n')) as typeof fetch; + await expect( + callMcpTool(record, { name: 'list-all-documentation' }, fetchImpl) + ).rejects.toThrow(/SSE response with no data event/); + }); + + it('throws when the record has no mcp.endpoint', async () => { + const noEndpoint: StorybookInstanceRecord = { ...record, mcp: { status: 'ready' } }; + const fetchImpl = vi.fn() as unknown as typeof fetch; + await expect( + callMcpTool(noEndpoint, { name: 'list-all-documentation' }, fetchImpl) + ).rejects.toThrow(/has no server endpoint registered/); + }); + + it('throws when the response is not ok', async () => { + const fetchImpl = (async () => + new Response('boom', { status: 500, statusText: 'Server Error' })) as typeof fetch; + await expect( + callMcpTool(record, { name: 'list-all-documentation' }, fetchImpl) + ).rejects.toThrow(/responded with 500/); + }); + + it('throws when the response content-type is neither JSON nor SSE', async () => { + const fetchImpl = (async () => + new Response('', { + status: 200, + headers: { 'Content-Type': 'text/html' }, + })) as typeof fetch; + await expect( + callMcpTool(record, { name: 'list-all-documentation' }, fetchImpl) + ).rejects.toThrow(/unsupported content-type "text\/html"/); + }); + + it('throws an McpJsonRpcError when the JSON-RPC payload carries an error', async () => { + const fetchImpl = (async () => + jsonResponse({ + jsonrpc: '2.0', + id: 'whatever', + error: { code: -32601, message: 'unknown tool' }, + })) as typeof fetch; + const promise = callMcpTool(record, { name: 'nope' }, fetchImpl); + await expect(promise).rejects.toThrow(/Storybook server error -32601: unknown tool/); + await expect(promise).rejects.toBeInstanceOf(McpJsonRpcError); + }); + + it.each([ + ['a primitive result', { jsonrpc: '2.0', id: 1, result: 'hello' }], + ['a null result', { jsonrpc: '2.0', id: 1, result: null }], + [ + 'a content item without a type', + { jsonrpc: '2.0', id: 1, result: { content: [{ text: 'x' }] } }, + ], + ['a malformed error object', { jsonrpc: '2.0', id: 1, error: { code: 'x' } }], + ])('rejects %s as an unexpected response shape', async (_label, body) => { + const fetchImpl = (async () => jsonResponse(body)) as typeof fetch; + await expect( + callMcpTool(record, { name: 'list-all-documentation' }, fetchImpl) + ).rejects.toThrow(/unexpected response shape/); + }); + + it('passes through extra content fields and result keys (loose validation)', async () => { + const fetchImpl = (async () => + jsonResponse({ + jsonrpc: '2.0', + id: 1, + result: { + content: [{ type: 'resource_link', uri: 'http://x' }], + _meta: { 'storybook.dev/foo': 1 }, + }, + })) as typeof fetch; + const result = await callMcpTool(record, { name: 'x' }, fetchImpl); + expect(result.content?.[0]).toMatchObject({ type: 'resource_link', uri: 'http://x' }); + }); +}); + +describe('listMcpTools', () => { + it('POSTs a JSON-RPC tools/list request and returns the tool descriptors', async () => { + const tools = [ + { name: 'get-documentation', description: 'Docs', inputSchema: { properties: {} } }, + { name: 'list-all-documentation' }, + ]; + const fetchImpl = vi.fn(async () => + jsonResponse({ jsonrpc: '2.0', id: 'x', result: { tools } }) + ) as unknown as typeof fetch; + + await expect(listMcpTools(record, fetchImpl)).resolves.toEqual(tools); + + const body = JSON.parse(vi.mocked(fetchImpl).mock.calls[0][1]?.body as string); + expect(body).toMatchObject({ method: 'tools/list', params: {} }); + }); + + it('returns [] when the result has no tools array', async () => { + const fetchImpl = (async () => + jsonResponse({ jsonrpc: '2.0', id: 'x', result: {} })) as typeof fetch; + await expect(listMcpTools(record, fetchImpl)).resolves.toEqual([]); + }); + + it('rejects tool descriptors without a name as an unexpected response shape', async () => { + const fetchImpl = (async () => + jsonResponse({ + jsonrpc: '2.0', + id: 'x', + result: { tools: [{ description: 'nameless' }] }, + })) as typeof fetch; + await expect(listMcpTools(record, fetchImpl)).rejects.toThrow(/unexpected response shape/); + }); +}); diff --git a/code/lib/cli-storybook/src/ai/mcp/client.ts b/code/lib/cli-storybook/src/ai/mcp/client.ts new file mode 100644 index 000000000000..d7800642d6de --- /dev/null +++ b/code/lib/cli-storybook/src/ai/mcp/client.ts @@ -0,0 +1,194 @@ +import * as v from 'valibot'; + +import { + type McpToolDescriptor, + McpToolDescriptorSchema, + type StorybookInstanceRecord, + type ToolCallResult, + ToolCallResultSchema, +} from './types.ts'; + +/** + * Marks the request as coming from a trusted local Storybook client. `@storybook/addon-mcp` uses + * this header to skip auth flows meant for remote (composed) Storybooks. + */ +const STORYBOOK_MCP_PROXY_HEADER = 'X-Storybook-MCP-Proxy'; +const STORYBOOK_MCP_PROXY_HEADER_VALUE = 'true'; + +/** + * Upper bound on a single request so a hung server cannot stall the CLI forever. Generous because + * `run-story-tests` on a full suite legitimately runs for minutes. + */ +const REQUEST_TIMEOUT_MS = 10 * 60 * 1000; + +export type ToolCallParams = { + name: string; + arguments?: Record; +}; + +/** A JSON-RPC level error returned by the Storybook MCP server (e.g. unknown tool). */ +export class McpJsonRpcError extends Error { + constructor( + public readonly code: number, + message: string + ) { + super(`Storybook server error ${code}: ${message}`); + this.name = 'McpJsonRpcError'; + } +} + +const JsonRpcEnvelopeSchema = v.looseObject({ + result: v.optional(v.unknown()), + error: v.optional(v.looseObject({ code: v.number(), message: v.string() })), +}); + +const ToolListResultSchema = v.looseObject({ + tools: v.optional(v.array(McpToolDescriptorSchema)), +}); + +/** Forward an MCP `tools/call` JSON-RPC request to a local Storybook MCP server. */ +export async function callMcpTool( + record: StorybookInstanceRecord, + params: ToolCallParams, + fetchImpl: typeof fetch = fetch +): Promise { + return sendJsonRpcRequest(record, 'tools/call', params, ToolCallResultSchema, fetchImpl); +} + +/** List the tools exposed by a local Storybook MCP server via `tools/list`. */ +export async function listMcpTools( + record: StorybookInstanceRecord, + fetchImpl: typeof fetch = fetch +): Promise { + const result = await sendJsonRpcRequest( + record, + 'tools/list', + {}, + ToolListResultSchema, + fetchImpl + ); + return result.tools ?? []; +} + +/** + * Send a single JSON-RPC request to the instance's MCP endpoint over HTTP. + * + * This is deliberately NOT a full MCP client: there is no `initialize` handshake, session, or + * protocol-version negotiation. The downstream is always `@storybook/addon-mcp`, whose tmcp + * HttpTransport is stateless and serves `tools/*` per-request — the same local shortcut + * `@storybook/mcp-proxy` takes in its proxy-client. If the CLI ever needs to talk to arbitrary + * MCP servers, replace this with a real client instead of extending it. + * + * tmcp hardcodes `text/event-stream` for any request with an id, so we accept both content-types + * and parse the SSE envelope when needed. + */ +async function sendJsonRpcRequest( + record: StorybookInstanceRecord, + method: 'tools/call' | 'tools/list', + params: unknown, + resultSchema: v.GenericSchema, + fetchImpl: typeof fetch +): Promise { + const endpoint = record.mcp.endpoint; + if (!endpoint) { + throw new Error(`The Storybook instance at ${record.cwd} has no server endpoint registered`); + } + + const target = new URL(endpoint, record.url).href; + + const response = await fetchImpl(target, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + Accept: 'application/json, text/event-stream', + [STORYBOOK_MCP_PROXY_HEADER]: STORYBOOK_MCP_PROXY_HEADER_VALUE, + }, + body: JSON.stringify({ + jsonrpc: '2.0', + id: crypto.randomUUID(), + method, + params, + }), + signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS), + }); + + if (!response.ok) { + throw new Error( + `The Storybook server at ${target} responded with ${response.status} ${response.statusText}` + ); + } + + const payload = await readJsonRpcResponse(response, target); + + const envelope = v.safeParse(JsonRpcEnvelopeSchema, payload); + if (!envelope.success) { + throw unexpectedShapeError(target); + } + if (envelope.output.error) { + throw new McpJsonRpcError(envelope.output.error.code, envelope.output.error.message); + } + if (envelope.output.result === undefined) { + throw new Error('The Storybook server returned no result'); + } + + const result = v.safeParse(resultSchema, envelope.output.result); + if (!result.success) { + throw unexpectedShapeError(target); + } + return result.output; +} + +function unexpectedShapeError(target: string): Error { + return new Error(`The Storybook server at ${target} returned an unexpected response shape`); +} + +async function readJsonRpcResponse(response: Response, endpoint: string): Promise { + const contentType = (response.headers.get('content-type') ?? '').toLowerCase(); + + if (contentType.includes('application/json')) { + return await response.json(); + } + + if (contentType.includes('text/event-stream')) { + return parseSseEnvelope(await response.text(), endpoint); + } + + throw new Error( + `The Storybook server at ${endpoint} returned unsupported content-type "${contentType}". Expected application/json or text/event-stream.` + ); +} + +/** + * Parse an MCP Streamable HTTP SSE response containing a single JSON-RPC envelope. Format per the + * SSE spec: lines starting with `data:` hold payload bytes; multiple `data:` lines in one event + * are joined with `\n`; the event terminates at the first blank line. We only care about the first + * event because a tools/call or tools/list response is always a single message. + */ +function parseSseEnvelope(body: string, endpoint: string): unknown { + const dataLines: string[] = []; + for (const rawLine of body.split('\n')) { + const line = rawLine.replace(/\r$/, ''); + if (line.startsWith('data:')) { + const value = line.slice(5); + dataLines.push(value.startsWith(' ') ? value.slice(1) : value); + continue; + } + if (line === '' && dataLines.length > 0) { + break; + } + } + if (dataLines.length === 0) { + throw new Error( + `The Storybook server at ${endpoint} returned an SSE response with no data event` + ); + } + try { + return JSON.parse(dataLines.join('\n')); + } catch (error) { + throw new Error( + `The Storybook server at ${endpoint} returned an SSE event whose data could not be parsed as JSON: ${ + error instanceof Error ? error.message : String(error) + }` + ); + } +} diff --git a/code/lib/cli-storybook/src/ai/mcp/intercepts.test.ts b/code/lib/cli-storybook/src/ai/mcp/intercepts.test.ts new file mode 100644 index 000000000000..df444e6efcd6 --- /dev/null +++ b/code/lib/cli-storybook/src/ai/mcp/intercepts.test.ts @@ -0,0 +1,55 @@ +import { describe, expect, it } from 'vitest'; + +import { getInterceptMarkdown } from './intercepts.ts'; +import type { StorybookInstanceRecord } from './types.ts'; + +const record = (cwd: string, url: string): StorybookInstanceRecord => ({ + schemaVersion: 1, + instanceId: 'i-1', + pid: 1, + cwd, + url, + port: 6006, + mcp: { status: 'ready', endpoint: '/mcp' }, +}); + +describe('getInterceptMarkdown', () => { + it('no-instance without candidates tells the agent to start storybook dev', () => { + const markdown = getInterceptMarkdown('no-instance'); + expect(markdown).toContain('Storybook is not running at this cwd'); + expect(markdown).toContain('storybook dev'); + }); + + it('no-instance with candidates lists the running cwds', () => { + const markdown = getInterceptMarkdown('no-instance', { + records: [record('/projects/foo', 'http://localhost:6006')], + }); + expect(markdown).toContain('Running Storybooks:'); + expect(markdown).toContain('- `/projects/foo` (http://localhost:6006)'); + expect(markdown).toContain('--cwd'); + }); + + it('port-mismatch lists the ports running at the cwd', () => { + const markdown = getInterceptMarkdown('port-mismatch', { + port: 9999, + records: [record('/projects/foo', 'http://localhost:6006')], + }); + expect(markdown).toContain('not on port `9999`'); + expect(markdown).toContain('- port `6006`'); + expect(markdown).toContain('omit `--port`'); + }); + + it('addon-missing instructs installing the MCP addon', () => { + const markdown = getInterceptMarkdown('addon-missing'); + expect(markdown).toContain('`@storybook/addon-mcp` addon is missing'); + expect(markdown).toContain('npx storybook add @storybook/addon-mcp'); + }); + + it('mcp-starting asks to wait and retry', () => { + expect(getInterceptMarkdown('mcp-starting')).toContain('still starting up'); + }); + + it('mcp-error points at the Storybook terminal output', () => { + expect(getInterceptMarkdown('mcp-error')).toContain('Inspect the Storybook terminal output'); + }); +}); diff --git a/code/lib/cli-storybook/src/ai/mcp/intercepts.ts b/code/lib/cli-storybook/src/ai/mcp/intercepts.ts new file mode 100644 index 000000000000..3390445461c0 --- /dev/null +++ b/code/lib/cli-storybook/src/ai/mcp/intercepts.ts @@ -0,0 +1,62 @@ +import type { InterceptReason, StorybookInstanceRecord } from './types.ts'; + +/** + * Repair-instruction markdown for agents, mirroring `@storybook/mcp-proxy` (storybookjs/mcp) so + * the CLI and the proxy give the same guidance — keep the two in sync when updating either. + */ +const NO_INSTANCE_EMPTY = `Storybook is not running at this cwd. Start \`storybook dev\` from the project's cwd and retry the command.`; + +const buildNoInstanceWithCandidates = (records: StorybookInstanceRecord[]) => + `No Storybook is running at this cwd. Either start Storybook from the project's cwd, or retry with \`--cwd\` set to one of the running cwds below. + +Running Storybooks: +${records.map((r) => `- \`${r.cwd}\` (${r.url})`).join('\n')}`; + +const buildPortMismatch = (port: number | undefined, records: StorybookInstanceRecord[]) => + `Storybook is running at this cwd, but not on port \`${port ?? 'unknown'}\`. Retry with one of the running ports below, or omit \`--port\` to route by cwd alone. + +Running Storybooks at this cwd: +${records.map((r) => `- port \`${r.port}\` (${r.url}, status: \`${r.mcp.status}\`)`).join('\n')}`; + +const ADDON_MISSING = `Storybook is running but does not provide these commands. The \`@storybook/addon-mcp\` addon is missing. + +Install it: +\`\`\` +npx storybook add @storybook/addon-mcp +\`\`\` + +Restart Storybook, then retry the command.`; + +const MCP_STARTING = `Storybook is running but its command server is still starting up. Wait a moment and retry the command.`; + +const MCP_ERROR = `Storybook is running but its command server reported an error. Inspect the Storybook terminal output, fix the underlying issue, then retry the command.`; + +export type InterceptExtras = { + records?: StorybookInstanceRecord[]; + port?: number; +}; + +export function getInterceptMarkdown( + reason: InterceptReason, + extras: InterceptExtras = {} +): string { + const { records, port } = extras; + switch (reason) { + case 'no-instance': + return records && records.length > 0 + ? buildNoInstanceWithCandidates(records) + : NO_INSTANCE_EMPTY; + case 'port-mismatch': + return buildPortMismatch(port, records ?? []); + case 'addon-missing': + return ADDON_MISSING; + case 'mcp-starting': + return MCP_STARTING; + case 'mcp-error': + return MCP_ERROR; + default: { + const unhandled: never = reason; + throw new Error(`Unhandled intercept reason: ${unhandled as string}`); + } + } +} diff --git a/code/lib/cli-storybook/src/ai/mcp/register.test.ts b/code/lib/cli-storybook/src/ai/mcp/register.test.ts new file mode 100644 index 000000000000..6d344a014c83 --- /dev/null +++ b/code/lib/cli-storybook/src/ai/mcp/register.test.ts @@ -0,0 +1,212 @@ +import { writeFile } from 'node:fs/promises'; + +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; + +import { Command } from 'commander'; + +import { isAiCliFeatureEnabled, registerAiMcpPassthrough } from './register.ts'; +import { buildStorybookCommandsHelp, runAiTool, runAiToolHelp } from './run-tool.ts'; + +vi.mock('./run-tool.ts', { spy: true }); + +vi.mock('node:fs/promises', { spy: true }); + +describe('isAiCliFeatureEnabled', () => { + it.each([ + ['1', true], + ['true', true], + ['0', false], + ['false', false], + ['', false], + [undefined, false], + ])('STORYBOOK_FEATURE_AI_CLI=%j → %j', (value, expected) => { + expect(isAiCliFeatureEnabled({ STORYBOOK_FEATURE_AI_CLI: value })).toBe(expected); + }); +}); + +/** Replicate the `ai` command tree from `bin/run.ts`: a `setup` subcommand plus a help action. */ +function buildProgram({ withPassthrough }: { withPassthrough: boolean }) { + const program = new Command(); + program.exitOverride(); + const setupAction = vi.fn(); + const helpAction = vi.fn(); + + const aiCommand = program + .command('ai') + .description('AI agent helpers for Storybook') + .option('-o, --output ', 'Write the prompt output to a file') + .exitOverride(); + aiCommand.configureOutput({ writeOut: () => {}, writeErr: () => {} }); + aiCommand.command('setup').action(setupAction); + aiCommand.action(helpAction); + + if (withPassthrough) { + registerAiMcpPassthrough(program, aiCommand); + } + + return { program, aiCommand, setupAction, helpAction }; +} + +function parse(program: Command, argv: string[]) { + return program.parseAsync(['node', 'storybook', ...argv]); +} + +function stdoutText(): string { + return vi + .mocked(process.stdout.write) + .mock.calls.map(([chunk]) => String(chunk)) + .join(''); +} + +beforeEach(() => { + vi.mocked(runAiTool).mockResolvedValue({ exitCode: 0, output: 'ok' }); + vi.mocked(runAiToolHelp).mockResolvedValue({ exitCode: 0, output: 'tool help' }); + vi.mocked(buildStorybookCommandsHelp).mockResolvedValue( + 'Storybook commands (from the running Storybook):' + ); + vi.mocked(writeFile).mockResolvedValue(undefined); + vi.spyOn(process.stdout, 'write').mockImplementation(() => true); +}); + +afterEach(() => { + vi.restoreAllMocks(); + vi.mocked(runAiTool).mockReset(); + vi.mocked(runAiToolHelp).mockReset(); + vi.mocked(buildStorybookCommandsHelp).mockReset(); + process.exitCode = undefined; +}); + +describe('without the feature flag (no registration)', () => { + it('keeps `setup` as the only subcommand', () => { + const { aiCommand } = buildProgram({ withPassthrough: false }); + expect(aiCommand.commands.map((c) => c.name())).toEqual(['setup']); + }); + + it('rejects tool names like today (excess arguments)', async () => { + const { program } = buildProgram({ withPassthrough: false }); + await expect(parse(program, ['ai', 'list-all-documentation'])).rejects.toMatchObject({ + code: 'commander.excessArguments', + }); + expect(runAiTool).not.toHaveBeenCalled(); + }); + + it('keeps the bare `ai` help action', async () => { + const { program, helpAction } = buildProgram({ withPassthrough: false }); + await parse(program, ['ai']); + expect(helpAction).toHaveBeenCalled(); + expect(buildStorybookCommandsHelp).not.toHaveBeenCalled(); + }); +}); + +describe('with the feature flag (passthrough registered)', () => { + it('forwards `ai ` with pass-through tokens to runAiTool', async () => { + const { program } = buildProgram({ withPassthrough: true }); + await parse(program, ['ai', 'get-documentation', '--id', 'button-docs']); + expect(runAiTool).toHaveBeenCalledWith('get-documentation', ['--id', 'button-docs'], { + cwd: undefined, + port: undefined, + json: undefined, + }); + }); + + it('parses --cwd, --port and --json before the tool name as CLI options', async () => { + const { program } = buildProgram({ withPassthrough: true }); + await parse(program, [ + 'ai', + '--cwd', + '/x', + '--port', + '6006', + '--json', + '{"a":1}', + 'get-documentation', + ]); + expect(runAiTool).toHaveBeenCalledWith('get-documentation', [], { + cwd: '/x', + port: '6006', + json: '{"a":1}', + }); + }); + + it('passes tokens after the tool name through verbatim, even option-like ones', async () => { + const { program } = buildProgram({ withPassthrough: true }); + await parse(program, ['ai', 'tool-x', '--cwd', '/y', '--output', 'z']); + expect(runAiTool).toHaveBeenCalledWith('tool-x', ['--cwd', '/y', '--output', 'z'], { + cwd: undefined, + port: undefined, + json: undefined, + }); + }); + + it('writes the result to the file given via --output instead of stdout', async () => { + const { program } = buildProgram({ withPassthrough: true }); + vi.mocked(runAiTool).mockResolvedValue({ exitCode: 0, output: 'markdown result' }); + await parse(program, ['ai', '-o', '/out/result.md', 'tool-x']); + expect(writeFile).toHaveBeenCalledWith('/out/result.md', 'markdown result\n', 'utf-8'); + expect(process.stdout.write).not.toHaveBeenCalledWith('markdown result\n'); + }); + + it('writes the result to stdout', async () => { + const { program } = buildProgram({ withPassthrough: true }); + vi.mocked(runAiTool).mockResolvedValue({ exitCode: 0, output: 'markdown result' }); + await parse(program, ['ai', 'tool-x']); + expect(process.stdout.write).toHaveBeenCalledWith('markdown result\n'); + expect(process.exitCode).toBeUndefined(); + }); + + it('sets a non-zero exit code on failure', async () => { + const { program } = buildProgram({ withPassthrough: true }); + vi.mocked(runAiTool).mockResolvedValue({ exitCode: 1, output: 'repair instructions' }); + await parse(program, ['ai', 'tool-x']); + expect(process.stdout.write).toHaveBeenCalledWith('repair instructions\n'); + expect(process.exitCode).toBe(1); + }); + + it('still dispatches `ai setup` to the setup subcommand', async () => { + const { program, setupAction } = buildProgram({ withPassthrough: true }); + await parse(program, ['ai', 'setup']); + expect(setupAction).toHaveBeenCalled(); + expect(runAiTool).not.toHaveBeenCalled(); + }); + + it.each([[['ai']], [['ai', '--help']], [['ai', '-h']]])( + 'shows commander help plus the tool commands section for %j', + async (argv) => { + const { program } = buildProgram({ withPassthrough: true }); + await parse(program, argv); + expect(buildStorybookCommandsHelp).toHaveBeenCalledWith({ cwd: undefined, port: undefined }); + const output = stdoutText(); + expect(output).toContain('Usage:'); + expect(output).toContain('setup'); + expect(output).toContain('Storybook commands (from the running Storybook):'); + expect(runAiTool).not.toHaveBeenCalled(); + } + ); + + it('passes --cwd and --port through to the tool commands section', async () => { + const { program } = buildProgram({ withPassthrough: true }); + await parse(program, ['ai', '--cwd', '/x', '--port', '6006', '--help']); + expect(buildStorybookCommandsHelp).toHaveBeenCalledWith({ cwd: '/x', port: '6006' }); + }); + + it('shows single-tool help for `ai --help `', async () => { + const { program } = buildProgram({ withPassthrough: true }); + await parse(program, ['ai', '--help', 'get-documentation']); + expect(runAiToolHelp).toHaveBeenCalledWith('get-documentation', { + cwd: undefined, + port: undefined, + }); + expect(process.stdout.write).toHaveBeenCalledWith('tool help\n'); + expect(runAiTool).not.toHaveBeenCalled(); + }); + + it('passes a --help token after the tool name through to runAiTool', async () => { + const { program } = buildProgram({ withPassthrough: true }); + await parse(program, ['ai', 'get-documentation', '--help']); + expect(runAiTool).toHaveBeenCalledWith('get-documentation', ['--help'], { + cwd: undefined, + port: undefined, + json: undefined, + }); + }); +}); diff --git a/code/lib/cli-storybook/src/ai/mcp/register.ts b/code/lib/cli-storybook/src/ai/mcp/register.ts new file mode 100644 index 000000000000..411d04c73224 --- /dev/null +++ b/code/lib/cli-storybook/src/ai/mcp/register.ts @@ -0,0 +1,95 @@ +import { writeFile } from 'node:fs/promises'; +import { resolve } from 'node:path'; + +import { optionalEnvToBoolean } from 'storybook/internal/common'; +import { logger } from 'storybook/internal/node-logger'; + +import type { Command } from 'commander'; + +import { + type AiToolRunResult, + buildStorybookCommandsHelp, + runAiTool, + runAiToolHelp, +} from './run-tool.ts'; + +/** + * The `storybook ai ` MCP passthrough is experimental (storybookjs/storybook#35124) and only + * registered when this feature flag is set; without it, `storybook ai` exposes `setup` only. + */ +export function isAiCliFeatureEnabled(env: NodeJS.ProcessEnv = process.env): boolean { + return optionalEnvToBoolean(env.STORYBOOK_FEATURE_AI_CLI) === true; +} + +const CWD_DESCRIPTION = + 'Project directory of the target Storybook (defaults to the current working directory)'; +const PORT_DESCRIPTION = + 'Port of the target Storybook, to address one specific instance when several run at the same cwd'; + +/** + * Register the passthrough on the `ai` command: a generic `[command] [args...]` argument pair that + * forwards any command to the running Storybook's server (MCP under the hood, but that is an + * implementation detail — user-facing copy says "commands"). `passThroughOptions` hands every + * token after the command name to the command untouched, which requires positional options on the + * program. + * + * Commander's built-in (synchronous) help is replaced with our own `-h, --help` option so the help + * output can include the commands fetched from the running Storybook. + */ +export function registerAiMcpPassthrough(program: Command, aiCommand: Command): void { + program.enablePositionalOptions(); + + aiCommand + .helpOption(false) + .usage('[options] [command] [args...]') + .argument('[command]', 'A command provided by the running Storybook') + .argument( + '[args...]', + 'Command arguments as `--key value` flags; values are JSON-parsed when possible' + ) + .option('--cwd ', CWD_DESCRIPTION) + .option('--port ', PORT_DESCRIPTION) + .option( + '--json ', + 'Raw JSON object with the command arguments (escape hatch for complex values)' + ) + .option('-h, --help', 'Show help, including the commands provided by the running Storybook') + .passThroughOptions() + .action( + async ( + command: string | undefined, + commandArgs: string[], + options: { cwd?: string; port?: string; json?: string; output?: string; help?: boolean } + ) => { + const target = { cwd: options.cwd, port: options.port }; + if (options.help && command) { + await printResult(await runAiToolHelp(command, target), options.output); + return; + } + if (options.help || !command) { + const commandsSection = await buildStorybookCommandsHelp(target); + process.stdout.write(`${aiCommand.helpInformation()}\n${commandsSection}\n`); + return; + } + const result = await runAiTool(command, commandArgs, { ...target, json: options.json }); + await printResult(result, options.output); + } + ); +} + +/** Print to stdout, or to the file given via the `ai` command's `-o, --output` option. */ +async function printResult( + { output, exitCode }: AiToolRunResult, + outputPath: string | undefined +): Promise { + if (outputPath) { + const resolvedPath = resolve(outputPath); + await writeFile(resolvedPath, `${output}\n`, 'utf-8'); + logger.log(`Output written to ${resolvedPath}`); + } else { + process.stdout.write(`${output}\n`); + } + if (exitCode !== 0) { + process.exitCode = exitCode; + } +} diff --git a/code/lib/cli-storybook/src/ai/mcp/registry.test.ts b/code/lib/cli-storybook/src/ai/mcp/registry.test.ts new file mode 100644 index 000000000000..0e0515ceba52 --- /dev/null +++ b/code/lib/cli-storybook/src/ai/mcp/registry.test.ts @@ -0,0 +1,143 @@ +import { readFile, readdir, rm } from 'node:fs/promises'; + +import { vol } from 'memfs'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; + +import { readRegistry } from './registry.ts'; + +// Spy-only mock: keep the real `node:fs/promises` module shape, then redirect the calls used by +// the registry reader to `memfs` so disk state stays scoped to `vol`. +vi.mock('node:fs/promises', { spy: true }); + +const REGISTRY_DIR = '/registry'; + +beforeEach(async () => { + const memfs = await vi.importActual('memfs'); + + vi.mocked(readdir).mockImplementation( + memfs.fs.promises.readdir as unknown as typeof import('node:fs/promises').readdir + ); + vi.mocked(readFile).mockImplementation( + memfs.fs.promises.readFile as unknown as typeof import('node:fs/promises').readFile + ); + vi.mocked(rm).mockImplementation( + memfs.fs.promises.rm as unknown as typeof import('node:fs/promises').rm + ); + + // Deterministic PID liveness: in these tests only the current process counts as alive. + vi.spyOn(process, 'kill').mockImplementation((pid) => { + if (pid !== process.pid) { + const error = new Error('ESRCH') as NodeJS.ErrnoException; + error.code = 'ESRCH'; + throw error; + } + return true; + }); +}); + +afterEach(() => { + vol.reset(); + vi.restoreAllMocks(); +}); + +const aliveRecord = { + schemaVersion: 1, + instanceId: 'alive-uuid', + pid: process.pid, + cwd: '/projects/alive', + url: 'http://localhost:6006', + port: 6006, + storybookVersion: '10.5.0', + startedAt: '2026-05-18T12:00:00.000Z', + updatedAt: '2026-05-18T12:00:03.000Z', + mcp: { status: 'ready', endpoint: '/mcp' }, +}; + +describe('readRegistry', () => { + it('returns [] when the registry dir does not exist', async () => { + vol.fromNestedJSON({ '/elsewhere': {} }); + await expect(readRegistry('/registry-does-not-exist')).resolves.toEqual([]); + }); + + it('returns [] when the registry dir is empty', async () => { + vol.fromNestedJSON({ [REGISTRY_DIR]: {} }); + await expect(readRegistry(REGISTRY_DIR)).resolves.toEqual([]); + }); + + it('parses valid records and skips dead PIDs, bad schemas, malformed JSON and non-JSON files', async () => { + const dead = { ...aliveRecord, instanceId: 'dead-uuid', pid: 2147483646 }; + const unknownStatus = { + ...aliveRecord, + instanceId: 'bad-uuid', + mcp: { status: 'unrecognised', endpoint: '/mcp' }, + }; + + vol.fromNestedJSON({ + [REGISTRY_DIR]: { + 'alive.json': JSON.stringify(aliveRecord), + 'dead.json': JSON.stringify(dead), + 'bad-status.json': JSON.stringify(unknownStatus), + 'malformed.json': '{ not json', + 'wrong-shape.json': JSON.stringify({ foo: 'bar' }), + 'ignored.txt': 'should be ignored', + }, + }); + + await expect(readRegistry(REGISTRY_DIR)).resolves.toEqual([aliveRecord]); + }); + + it('removes the registry file of a dead PID', async () => { + const dead = { ...aliveRecord, instanceId: 'dead-uuid', pid: 2147483646 }; + vol.fromNestedJSON({ [REGISTRY_DIR]: { 'dead.json': JSON.stringify(dead) } }); + + await expect(readRegistry(REGISTRY_DIR)).resolves.toEqual([]); + expect(vol.toJSON()[`${REGISTRY_DIR}/dead.json`]).toBeUndefined(); + }); + + it('filters records with non-positive PIDs (process-group sentinels)', async () => { + vol.fromNestedJSON({ + [REGISTRY_DIR]: { + 'zero.json': JSON.stringify({ ...aliveRecord, instanceId: 'zero', pid: 0 }), + 'negative.json': JSON.stringify({ ...aliveRecord, instanceId: 'neg', pid: -1 }), + }, + }); + + await expect(readRegistry(REGISTRY_DIR)).resolves.toEqual([]); + }); + + it('treats EPERM on the liveness signal as alive (foreign-user process)', async () => { + vi.mocked(process.kill).mockImplementation(() => { + const error = new Error('EPERM') as NodeJS.ErrnoException; + error.code = 'EPERM'; + throw error; + }); + vol.fromNestedJSON({ [REGISTRY_DIR]: { 'alive.json': JSON.stringify(aliveRecord) } }); + + await expect(readRegistry(REGISTRY_DIR)).resolves.toEqual([aliveRecord]); + }); + + it('accepts records without the optional version and timestamp fields', async () => { + const minimal = { + schemaVersion: 1, + instanceId: 'minimal', + pid: process.pid, + cwd: '/projects/minimal', + url: 'http://localhost:6007', + port: 6007, + mcp: { status: 'starting' }, + }; + vol.fromNestedJSON({ [REGISTRY_DIR]: { 'minimal.json': JSON.stringify(minimal) } }); + + await expect(readRegistry(REGISTRY_DIR)).resolves.toEqual([minimal]); + }); + + it('rejects out-of-range ports', async () => { + vol.fromNestedJSON({ + [REGISTRY_DIR]: { + 'bad-port.json': JSON.stringify({ ...aliveRecord, port: 65536 }), + }, + }); + + await expect(readRegistry(REGISTRY_DIR)).resolves.toEqual([]); + }); +}); diff --git a/code/lib/cli-storybook/src/ai/mcp/registry.ts b/code/lib/cli-storybook/src/ai/mcp/registry.ts new file mode 100644 index 000000000000..9cc702d43ef1 --- /dev/null +++ b/code/lib/cli-storybook/src/ai/mcp/registry.ts @@ -0,0 +1,82 @@ +import * as fs from 'node:fs/promises'; +import { homedir } from 'node:os'; +import { join } from 'node:path'; + +import * as v from 'valibot'; + +import { type StorybookInstanceRecord, StorybookInstanceRecordSchema } from './types.ts'; + +/** + * Must stay in sync with `getDefaultRuntimeInstanceRegistryDir` in + * `code/core/src/core-server/utils/runtime-instance-registry.ts` (the writer side). Duplicated + * here so this reader does not pull the core-server module graph into the CLI's unit tests; the + * path is specified in storybookjs/storybook#34826. + */ +export const DEFAULT_REGISTRY_DIR = join(homedir(), '.storybook', 'instances'); + +/** + * Errno codes for which we degrade to "no instance" rather than throwing. The command is meant to + * fail-soft for environmental issues; a noisy stack trace would be worse UX than the + * missing-instance repair instructions. + */ +const SOFT_REGISTRY_ERRORS = new Set(['ENOENT', 'EACCES', 'EPERM', 'ENOTDIR']); + +/** + * Read all Storybook instance records from `registryDir`. + * + * Each file is expected to be a single JSON object matching {@link StorybookInstanceRecord}. + * Records whose PID is no longer alive are filtered out (and their files removed). Malformed files + * are skipped silently — the command should degrade to "no instance" rather than fail loudly. + */ +export async function readRegistry( + registryDir: string = DEFAULT_REGISTRY_DIR +): Promise { + let entries: string[]; + try { + entries = await fs.readdir(registryDir); + } catch (error) { + if (SOFT_REGISTRY_ERRORS.has((error as NodeJS.ErrnoException).code ?? '')) { + return []; + } + throw error; + } + + const records = await Promise.all( + entries + .filter((name) => name.endsWith('.json')) + .map(async (name) => { + try { + const raw = await fs.readFile(join(registryDir, name), 'utf-8'); + const parsed = v.safeParse(StorybookInstanceRecordSchema, JSON.parse(raw)); + if (!parsed.success) { + return null; + } + if (!isProcessAlive(parsed.output.pid)) { + await fs.rm(join(registryDir, name), { force: true }).catch(() => {}); + return null; + } + return parsed.output; + } catch { + return null; + } + }) + ); + + return records.filter((r): r is StorybookInstanceRecord => r !== null); +} + +/** + * Liveness check by sending signal 0. `EPERM` means the PID exists but we lack permission to + * signal it (foreign user), which still counts as alive. + */ +function isProcessAlive(pid: number): boolean { + if (!Number.isInteger(pid) || pid <= 0) { + return false; + } + try { + process.kill(pid, 0); + return true; + } catch (error) { + return (error as NodeJS.ErrnoException).code === 'EPERM'; + } +} diff --git a/code/lib/cli-storybook/src/ai/mcp/resolve-instance.test.ts b/code/lib/cli-storybook/src/ai/mcp/resolve-instance.test.ts new file mode 100644 index 000000000000..79d454a937ec --- /dev/null +++ b/code/lib/cli-storybook/src/ai/mcp/resolve-instance.test.ts @@ -0,0 +1,230 @@ +import { describe, expect, it } from 'vitest'; + +import { resolveInstance } from './resolve-instance.ts'; +import type { McpStatus, StorybookInstanceRecord } from './types.ts'; + +let nextInstance = 0; + +function record( + cwd: string, + status: McpStatus = 'ready', + overrides: Partial = {} +): StorybookInstanceRecord { + nextInstance += 1; + return { + schemaVersion: 1, + instanceId: `inst-${nextInstance}`, + pid: 1000 + nextInstance, + cwd, + url: `http://localhost:${6000 + nextInstance}`, + port: 6000 + nextInstance, + mcp: { + status, + endpoint: + status === 'ready' || status === 'error' + ? `http://localhost:${6000 + nextInstance}/mcp` + : undefined, + }, + ...overrides, + }; +} + +describe('resolveInstance', () => { + it('returns no-instance with empty candidates when registry is empty', () => { + const result = resolveInstance([], '/Users/x/projects/foo'); + expect(result).toEqual({ kind: 'intercept', reason: 'no-instance', records: [], matches: [] }); + }); + + it('returns no-instance with candidates when no record cwd matches', () => { + const a = record('/Users/x/projects/foo'); + const b = record('/Users/x/projects/bar'); + const result = resolveInstance([a, b], '/Users/x/projects/baz'); + expect(result.kind).toBe('intercept'); + if (result.kind === 'intercept') { + expect(result.reason).toBe('no-instance'); + expect(result.records).toEqual([a, b]); + } + }); + + it('matches a record by exact normalized cwd', () => { + const r = record('/Users/x/projects/foo'); + const result = resolveInstance([r], '/Users/x/projects/foo'); + expect(result).toEqual({ kind: 'instance', record: r, matches: [r] }); + }); + + it('normalizes trailing slashes and dot segments before matching', () => { + const r = record('/Users/x/projects/foo'); + const result = resolveInstance([r], '/Users/x/projects/foo/./'); + expect(result).toEqual({ kind: 'instance', record: r, matches: [r] }); + }); + + it('does NOT match a child path of a record cwd (exact only)', () => { + const r = record('/Users/x/projects/foo'); + const result = resolveInstance([r], '/Users/x/projects/foo/src/Button.tsx'); + expect(result.kind).toBe('intercept'); + if (result.kind === 'intercept') { + expect(result.reason).toBe('no-instance'); + } + }); + + it('does NOT match a sibling string prefix', () => { + const r = record('/Users/x/projects/foo'); + const result = resolveInstance([r], '/Users/x/projects/foobar'); + expect(result.kind).toBe('intercept'); + if (result.kind === 'intercept') { + expect(result.reason).toBe('no-instance'); + } + }); + + it('tie-breaks on lowest pid when 2+ records share the cwd and none carry a startedAt', () => { + const a = record('/Users/x/projects/foo', 'ready', { pid: 200 }); + const b = record('/Users/x/projects/foo', 'ready', { pid: 100 }); + const result = resolveInstance([a, b], '/Users/x/projects/foo'); + expect(result.kind).toBe('instance'); + if (result.kind === 'instance') { + expect(result.record).toBe(b); + expect(result.matches).toEqual([b, a]); + } + }); + + it('picks the most recently started ready instance when 2+ records share the cwd', () => { + const older = record('/Users/x/projects/foo', 'ready', { + pid: 100, + startedAt: '2026-06-09T10:00:00.000Z', + }); + const newer = record('/Users/x/projects/foo', 'ready', { + pid: 200, + startedAt: '2026-06-09T11:00:00.000Z', + }); + const result = resolveInstance([older, newer], '/Users/x/projects/foo'); + expect(result.kind).toBe('instance'); + if (result.kind === 'instance') { + expect(result.record).toBe(newer); + expect(result.matches).toEqual([newer, older]); + } + }); + + it('treats a record without startedAt as older than one with a startedAt', () => { + const noStamp = record('/Users/x/projects/foo', 'ready', { pid: 100 }); + const stamped = record('/Users/x/projects/foo', 'ready', { + pid: 200, + startedAt: '2026-06-09T11:00:00.000Z', + }); + const result = resolveInstance([noStamp, stamped], '/Users/x/projects/foo'); + expect(result.kind).toBe('instance'); + if (result.kind === 'instance') { + expect(result.record).toBe(stamped); + } + }); + + it('prefers a ready record over a more recently started non-ready one', () => { + const ready = record('/Users/x/projects/foo', 'ready', { + pid: 100, + startedAt: '2026-06-09T10:00:00.000Z', + }); + const newerStarting = record('/Users/x/projects/foo', 'starting', { + pid: 200, + startedAt: '2026-06-09T11:00:00.000Z', + }); + const result = resolveInstance([ready, newerStarting], '/Users/x/projects/foo'); + expect(result.kind).toBe('instance'); + if (result.kind === 'instance') { + expect(result.record).toBe(ready); + } + }); + + it('dispatches the most recently started instance status when none are ready', () => { + const olderError = record('/Users/x/projects/foo', 'error', { + pid: 100, + startedAt: '2026-06-09T10:00:00.000Z', + }); + const newerStarting = record('/Users/x/projects/foo', 'starting', { + pid: 200, + startedAt: '2026-06-09T11:00:00.000Z', + }); + const result = resolveInstance([olderError, newerStarting], '/Users/x/projects/foo'); + expect(result.kind).toBe('intercept'); + if (result.kind === 'intercept') { + expect(result.reason).toBe('mcp-starting'); + } + }); + + it('prefers a ready record over non-ready ones when multiple records share the cwd', () => { + const starting = record('/Users/x/projects/foo', 'starting', { pid: 100 }); + const ready = record('/Users/x/projects/foo', 'ready', { pid: 200 }); + const result = resolveInstance([starting, ready], '/Users/x/projects/foo'); + expect(result.kind).toBe('instance'); + if (result.kind === 'instance') { + expect(result.record).toBe(ready); + expect(result.matches).toEqual([starting, ready]); + } + }); + + it('falls back to dispatching the lowest-pid status when no record at the cwd is ready', () => { + const a = record('/Users/x/projects/foo', 'starting', { pid: 200 }); + const b = record('/Users/x/projects/foo', 'error', { pid: 100 }); + const result = resolveInstance([a, b], '/Users/x/projects/foo'); + expect(result).toEqual({ kind: 'intercept', reason: 'mcp-error', matches: [b, a] }); + }); + + it('dispatches mcp.status=starting as mcp-starting intercept', () => { + const r = record('/p', 'starting'); + const result = resolveInstance([r], '/p'); + expect(result).toEqual({ kind: 'intercept', reason: 'mcp-starting', matches: [r] }); + }); + + it('dispatches mcp.status=not-installed as addon-missing intercept', () => { + const r = record('/p', 'not-installed'); + const result = resolveInstance([r], '/p'); + expect(result).toEqual({ kind: 'intercept', reason: 'addon-missing', matches: [r] }); + }); + + it('dispatches mcp.status=error as mcp-error intercept', () => { + const r = record('/p', 'error'); + const result = resolveInstance([r], '/p'); + expect(result).toEqual({ kind: 'intercept', reason: 'mcp-error', matches: [r] }); + }); + + it('selects the instance matching BOTH cwd and port when a port is supplied', () => { + const a = record('/Users/x/projects/foo', 'ready', { pid: 100, port: 6006 }); + const b = record('/Users/x/projects/foo', 'ready', { pid: 200, port: 6007 }); + const result = resolveInstance([a, b], '/Users/x/projects/foo', 6007); + expect(result.kind).toBe('instance'); + if (result.kind === 'instance') { + expect(result.record).toBe(b); + expect(result.matches).toEqual([b]); + } + }); + + it('ignores port when it is not supplied (routes by cwd alone)', () => { + const a = record('/Users/x/projects/foo', 'ready', { pid: 100, port: 6006 }); + const b = record('/Users/x/projects/foo', 'ready', { pid: 200, port: 6007 }); + const result = resolveInstance([a, b], '/Users/x/projects/foo'); + expect(result.kind).toBe('instance'); + if (result.kind === 'instance') { + expect(result.record).toBe(a); + expect(result.matches).toEqual([a, b]); + } + }); + + it('returns port-mismatch with the cwd instances as candidates when cwd matches but no instance is on the port', () => { + const a = record('/Users/x/projects/foo', 'ready', { pid: 100, port: 6006 }); + const b = record('/Users/x/projects/foo', 'ready', { pid: 200, port: 6007 }); + const result = resolveInstance([a, b], '/Users/x/projects/foo', 9999); + expect(result.kind).toBe('intercept'); + if (result.kind === 'intercept') { + expect(result.reason).toBe('port-mismatch'); + expect(result.records).toEqual([a, b]); + expect(result.matches).toEqual([]); + } + }); + + it('returns no-instance (not port-mismatch) when the cwd itself does not match', () => { + const a = record('/Users/x/projects/foo', 'ready', { port: 6006 }); + const result = resolveInstance([a], '/Users/x/projects/bar', 6006); + expect(result.kind).toBe('intercept'); + if (result.kind === 'intercept') { + expect(result.reason).toBe('no-instance'); + } + }); +}); diff --git a/code/lib/cli-storybook/src/ai/mcp/resolve-instance.ts b/code/lib/cli-storybook/src/ai/mcp/resolve-instance.ts new file mode 100644 index 000000000000..4af476eab3ae --- /dev/null +++ b/code/lib/cli-storybook/src/ai/mcp/resolve-instance.ts @@ -0,0 +1,132 @@ +import { resolve } from 'node:path'; + +import type { InterceptReason, StorybookInstanceRecord } from './types.ts'; + +export type ResolveResult = + | { + kind: 'instance'; + record: StorybookInstanceRecord; + matches: StorybookInstanceRecord[]; + } + | { + kind: 'intercept'; + reason: InterceptReason; + records?: StorybookInstanceRecord[]; + matches: StorybookInstanceRecord[]; + }; + +/** + * Pick the Storybook instance whose cwd exactly matches `targetCwd` after normalisation. Per + * milestone 2 of storybookjs/storybook#34826: matching is exact-normalized, with no longest-prefix + * or fallback behaviour. + * + * When `targetPort` is supplied (e.g. an agent that launched Storybook on a known port and wants + * to address that exact instance), it further constrains the cwd matches: an instance must match + * BOTH cwd and port. If the cwd matches but no instance there is on `targetPort`, a + * `port-mismatch` intercept is returned with the cwd's instances as candidates so callers can + * surface the running ports. + * + * If at least one record matches, dispatch based on the selected instance's `mcp.status`: + * + * - `ready` → forward the call + * - `starting` → mcp-starting intercept + * - `not-installed` → addon-missing intercept + * - `error` → mcp-error intercept + * + * Zero matches → no-instance intercept (callers may surface running cwds). 2+ matches at the same + * cwd → pick the most recently started instance (latest `startedAt` among `ready` records, else + * latest overall), on the assumption that the freshest instance is the one the agent just started. + * Records without a `startedAt` tie-break on lowest pid for determinism. All matches are returned + * (most-recent first) as `matches` so callers can warn the agent without blocking the call. + */ +export function resolveInstance( + records: StorybookInstanceRecord[], + targetCwd: string, + targetPort?: number +): ResolveResult { + const normalisedTarget = resolve(targetCwd); + const cwdMatches = records.filter((r) => resolve(r.cwd) === normalisedTarget); + const matches = targetPort == null ? cwdMatches : cwdMatches.filter((r) => r.port === targetPort); + + if (matches.length === 0) { + // cwd matched, but no instance there is on the requested port: a distinct, + // more actionable failure than "nothing is running here". + if (targetPort != null && cwdMatches.length > 0) { + return { + kind: 'intercept', + reason: 'port-mismatch', + records: cwdMatches, + matches: [], + }; + } + return { + kind: 'intercept', + reason: 'no-instance', + records, + matches: [], + }; + } + + const sortedMatches = [...matches].sort(byMostRecentlyStarted); + const selected = sortedMatches.find((r) => r.mcp.status === 'ready') ?? sortedMatches[0]; + + switch (selected.mcp.status) { + case 'ready': + return { + kind: 'instance', + record: selected, + matches: sortedMatches, + }; + + case 'starting': + return { + kind: 'intercept', + reason: 'mcp-starting', + matches: sortedMatches, + }; + + case 'not-installed': + return { + kind: 'intercept', + reason: 'addon-missing', + matches: sortedMatches, + }; + + case 'error': + return { + kind: 'intercept', + reason: 'mcp-error', + matches: sortedMatches, + }; + + default: { + const unhandled: never = selected.mcp.status; + throw new Error(`Unhandled MCP status: ${unhandled as string}`); + } + } +} + +/** + * `startedAt` as epoch millis, or `-Infinity` when absent/unparseable so such records sort as the + * oldest (and fall through to the pid tie-break). + */ +function startedAtMs(r: StorybookInstanceRecord): number { + if (!r.startedAt) { + return Number.NEGATIVE_INFINITY; + } + const t = Date.parse(r.startedAt); + return Number.isNaN(t) ? Number.NEGATIVE_INFINITY : t; +} + +/** + * Sort comparator: most recently started first, tie-breaking on lowest pid so ordering stays + * deterministic when timestamps are equal or missing. + */ +function byMostRecentlyStarted(a: StorybookInstanceRecord, b: StorybookInstanceRecord): number { + const ta = startedAtMs(a); + const tb = startedAtMs(b); + if (ta !== tb) { + return tb > ta ? 1 : -1; + } + return a.pid - b.pid; +} diff --git a/code/lib/cli-storybook/src/ai/mcp/run-tool.test.ts b/code/lib/cli-storybook/src/ai/mcp/run-tool.test.ts new file mode 100644 index 000000000000..6fce644ccd97 --- /dev/null +++ b/code/lib/cli-storybook/src/ai/mcp/run-tool.test.ts @@ -0,0 +1,359 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +import { McpJsonRpcError, callMcpTool, listMcpTools } from './client.ts'; +import { readRegistry } from './registry.ts'; +import { buildStorybookCommandsHelp, runAiTool, runAiToolHelp } from './run-tool.ts'; +import type { StorybookInstanceRecord } from './types.ts'; + +vi.mock('./registry.ts', { spy: true }); +vi.mock('./client.ts', { spy: true }); + +const record: StorybookInstanceRecord = { + schemaVersion: 1, + instanceId: 'inst-1', + pid: 1, + cwd: '/projects/foo', + url: 'http://localhost:6006', + port: 6006, + mcp: { status: 'ready', endpoint: '/mcp' }, +}; + +beforeEach(() => { + vi.mocked(readRegistry).mockReset().mockResolvedValue([record]); + vi.mocked(callMcpTool) + .mockReset() + .mockResolvedValue({ content: [{ type: 'text', text: 'upstream result' }] }); + vi.mocked(listMcpTools) + .mockReset() + .mockResolvedValue([{ name: 'list-all-documentation', description: 'List docs' }]); +}); + +describe('runAiTool', () => { + it('forwards the call to the matching instance and prints the markdown result', async () => { + const result = await runAiTool('list-all-documentation', ['--withStoryIds', 'true'], { + cwd: '/projects/foo', + }); + + expect(callMcpTool).toHaveBeenCalledWith( + record, + { name: 'list-all-documentation', arguments: { withStoryIds: true } }, + undefined + ); + expect(result).toEqual({ exitCode: 0, output: 'upstream result' }); + }); + + it('defaults the cwd to process.cwd()', async () => { + vi.mocked(readRegistry).mockResolvedValue([{ ...record, cwd: process.cwd() }]); + const result = await runAiTool('list-all-documentation', []); + expect(result.exitCode).toBe(0); + }); + + it('merges --json arguments with --key overrides', async () => { + await runAiTool('get-documentation', ['--id', 'override'], { + cwd: '/projects/foo', + json: '{"id":"base","verbose":true}', + }); + + expect(callMcpTool).toHaveBeenCalledWith( + record, + { name: 'get-documentation', arguments: { id: 'override', verbose: true } }, + undefined + ); + }); + + it('returns the arg-parsing error without contacting the registry', async () => { + const result = await runAiTool('get-documentation', ['positional'], { cwd: '/projects/foo' }); + expect(result.exitCode).toBe(1); + expect(result.output).toContain('Unexpected argument'); + expect(readRegistry).not.toHaveBeenCalled(); + }); + + it('prints the no-instance repair markdown and exits non-zero when nothing runs at the cwd', async () => { + const result = await runAiTool('get-documentation', [], { cwd: '/projects/other' }); + expect(result.exitCode).toBe(1); + expect(result.output).toContain('No Storybook is running at this cwd'); + expect(result.output).toContain('/projects/foo'); + }); + + it('routes to the instance on the requested --port when several share the cwd', async () => { + const onOtherPort = { ...record, instanceId: 'inst-2', pid: 2, port: 6007 }; + vi.mocked(readRegistry).mockResolvedValue([record, onOtherPort]); + const result = await runAiTool('list-all-documentation', ['--port', '6007'], { + cwd: '/projects/foo', + }); + expect(callMcpTool).toHaveBeenCalledWith(onOtherPort, expect.anything(), undefined); + expect(result.exitCode).toBe(0); + }); + + it('prints the port-mismatch repair markdown when no instance at the cwd is on the port', async () => { + const result = await runAiTool('list-all-documentation', [], { + cwd: '/projects/foo', + port: '9999', + }); + expect(result.exitCode).toBe(1); + expect(result.output).toContain('not on port `9999`'); + expect(result.output).toContain('- port `6006`'); + }); + + it('rejects an invalid --port without contacting the registry', async () => { + const result = await runAiTool('list-all-documentation', ['--port', 'abc'], { + cwd: '/projects/foo', + }); + expect(result.exitCode).toBe(1); + expect(result.output).toContain('`--port` must be a port number'); + expect(readRegistry).not.toHaveBeenCalled(); + }); + + it.each([ + ['starting', 'still starting up'], + ['not-installed', '`@storybook/addon-mcp` addon is missing'], + ['error', 'Inspect the Storybook terminal output'], + ] as const)('prints the repair markdown for mcp.status=%s', async (status, expected) => { + vi.mocked(readRegistry).mockResolvedValue([{ ...record, mcp: { status } }]); + const result = await runAiTool('get-documentation', [], { cwd: '/projects/foo' }); + expect(result.exitCode).toBe(1); + expect(result.output).toContain(expected); + }); + + it('prints a placeholder when the tool returns no content', async () => { + vi.mocked(callMcpTool).mockResolvedValue({ content: [] }); + const result = await runAiTool('list-all-documentation', [], { cwd: '/projects/foo' }); + expect(result).toEqual({ exitCode: 0, output: '(the command returned no content)' }); + }); + + it('surfaces a clean error when a ready record is missing its endpoint', async () => { + vi.mocked(callMcpTool).mockRejectedValue( + new Error('The Storybook instance at /projects/foo has no server endpoint registered') + ); + vi.mocked(readRegistry).mockResolvedValue([{ ...record, mcp: { status: 'ready' } }]); + const result = await runAiTool('list-all-documentation', [], { cwd: '/projects/foo' }); + expect(result.exitCode).toBe(1); + expect(result.output).toContain('Failed to reach the Storybook server at (no endpoint)'); + }); + + it('renders non-text content items as JSON blocks', async () => { + vi.mocked(callMcpTool).mockResolvedValue({ + content: [ + { type: 'text', text: 'intro' }, + { type: 'resource_link', uri: 'http://x' }, + ], + }); + const result = await runAiTool('get-documentation', [], { cwd: '/projects/foo' }); + expect(result.output).toContain('intro'); + expect(result.output).toContain('```json'); + expect(result.output).toContain('"resource_link"'); + }); + + it('lists the available tools when the call fails because the tool is unknown', async () => { + vi.mocked(callMcpTool).mockRejectedValue(new McpJsonRpcError(-32601, 'unknown tool')); + const result = await runAiTool('no-such-tool', [], { cwd: '/projects/foo' }); + expect(result.exitCode).toBe(1); + expect(result.output).toContain('Unknown command `no-such-tool`'); + expect(result.output).toContain('- `list-all-documentation`'); + }); + + it('lists the available tools when the server reports the unknown tool as an error result', async () => { + // addon-mcp (tmcp) reports unknown tools as an isError result, not a JSON-RPC error. + vi.mocked(callMcpTool).mockResolvedValue({ + content: [{ type: 'text', text: 'Tool no-such-tool not found' }], + isError: true, + }); + const result = await runAiTool('no-such-tool', [], { cwd: '/projects/foo' }); + expect(result.exitCode).toBe(1); + expect(result.output).toContain('Unknown command `no-such-tool`'); + expect(result.output).toContain('- `list-all-documentation`'); + }); + + it('keeps the original error result when the failing tool does exist', async () => { + vi.mocked(callMcpTool).mockResolvedValue({ + content: [{ type: 'text', text: 'tests failed' }], + isError: true, + }); + const result = await runAiTool('list-all-documentation', [], { cwd: '/projects/foo' }); + expect(result).toEqual({ exitCode: 1, output: 'tests failed' }); + }); + + it('prints the original JSON-RPC error when the tool exists', async () => { + vi.mocked(callMcpTool).mockRejectedValue(new McpJsonRpcError(-32602, 'invalid arguments')); + const result = await runAiTool('list-all-documentation', [], { cwd: '/projects/foo' }); + expect(result.exitCode).toBe(1); + expect(result.output).toContain('Storybook server error -32602: invalid arguments'); + }); + + it('prints the original JSON-RPC error when the tool list cannot be fetched', async () => { + vi.mocked(callMcpTool).mockRejectedValue(new McpJsonRpcError(-32601, 'unknown tool')); + vi.mocked(listMcpTools).mockRejectedValue(new Error('boom')); + const result = await runAiTool('no-such-tool', [], { cwd: '/projects/foo' }); + expect(result.exitCode).toBe(1); + expect(result.output).toContain('Storybook server error -32601: unknown tool'); + }); + + it('surfaces a friendly error when the MCP server is unreachable', async () => { + vi.mocked(callMcpTool).mockRejectedValue(new Error('connection refused')); + const result = await runAiTool('get-documentation', [], { cwd: '/projects/foo' }); + expect(result.exitCode).toBe(1); + expect(result.output).toContain('Failed to reach the Storybook server at /mcp'); + expect(result.output).toContain('connection refused'); + }); + + it('prepends a warning when multiple instances run at the same cwd', async () => { + const sibling = { ...record, instanceId: 'inst-2', pid: 2, url: 'http://localhost:6007' }; + vi.mocked(readRegistry).mockResolvedValue([record, sibling]); + const result = await runAiTool('list-all-documentation', [], { cwd: '/projects/foo' }); + expect(result.exitCode).toBe(0); + expect(result.output).toContain('Multiple Storybook instances'); + expect(result.output).toContain('pid `1`'); + expect(result.output).toContain('pid `2`'); + expect(result.output).toContain('(used)'); + expect(result.output).toContain('upstream result'); + }); +}); + +describe('buildStorybookCommandsHelp', () => { + it('lists each tool with the first line of its description', async () => { + vi.mocked(listMcpTools).mockResolvedValue([ + { + name: 'get-documentation', + description: 'Get docs for a component.\n\nLong details that should not appear.', + }, + { name: 'list-all-documentation' }, + ]); + + const section = await buildStorybookCommandsHelp({ cwd: '/projects/foo' }); + expect(section).toContain( + 'Storybook commands (from the Storybook running at http://localhost:6006):' + ); + expect(section).toContain('get-documentation'); + expect(section).toContain('Get docs for a component.'); + expect(section).not.toContain('Long details'); + expect(section).toContain("Run 'storybook ai --help'"); + }); + + it('degrades to a note when no Storybook is running (help must not fail)', async () => { + vi.mocked(readRegistry).mockResolvedValue([]); + const section = await buildStorybookCommandsHelp({ cwd: '/projects/foo' }); + expect(section).toContain('Storybook commands: (unavailable'); + expect(section).toContain('storybook dev'); + }); + + it('lists the sibling ports when several instances run at the cwd', async () => { + const older = { ...record, instanceId: 'inst-2', pid: 2, port: 6007 }; + const newest = { ...record, startedAt: '2026-06-10T12:00:00.000Z' }; + vi.mocked(readRegistry).mockResolvedValue([newest, older]); + const section = await buildStorybookCommandsHelp({ cwd: '/projects/foo' }); + expect(section).toContain('2 instances are running at this cwd'); + expect(section).toContain('port 6006'); + expect(section).toContain('other ports: 6007'); + expect(section).toContain('`--port`'); + }); + + it('names the port mismatch instead of claiming nothing is running', async () => { + const section = await buildStorybookCommandsHelp({ cwd: '/projects/foo', port: '9999' }); + expect(section).toContain('Storybook commands: (unavailable'); + expect(section).toContain('no instance on port `9999`'); + expect(section).toContain('running ports: 6006'); + expect(section).not.toContain('no running Storybook detected'); + }); + + it('says the Storybook is starting up instead of claiming nothing is running', async () => { + vi.mocked(readRegistry).mockResolvedValue([{ ...record, mcp: { status: 'starting' } }]); + const section = await buildStorybookCommandsHelp({ cwd: '/projects/foo' }); + expect(section).toContain('still starting up'); + }); + + it('points at the missing addon instead of claiming nothing is running', async () => { + vi.mocked(readRegistry).mockResolvedValue([{ ...record, mcp: { status: 'not-installed' } }]); + const section = await buildStorybookCommandsHelp({ cwd: '/projects/foo' }); + expect(section).toContain('install `@storybook/addon-mcp`'); + }); + + it('shows the Storybook version reported by the running instance', async () => { + vi.mocked(readRegistry).mockResolvedValue([{ ...record, storybookVersion: '10.5.0' }]); + const section = await buildStorybookCommandsHelp({ cwd: '/projects/foo' }); + expect(section).toContain( + 'Storybook commands (from the Storybook running at http://localhost:6006, Storybook 10.5.0):' + ); + }); + + it('degrades to a note when the MCP server is unreachable', async () => { + vi.mocked(listMcpTools).mockRejectedValue(new Error('connection refused')); + const section = await buildStorybookCommandsHelp({ cwd: '/projects/foo' }); + expect(section).toContain('Storybook commands: (unavailable'); + expect(section).toContain('could not be reached'); + }); + + it('degrades to a note when no tools are exposed', async () => { + vi.mocked(listMcpTools).mockResolvedValue([]); + const section = await buildStorybookCommandsHelp({ cwd: '/projects/foo' }); + expect(section).toContain('provides no commands'); + }); +}); + +describe('runAiToolHelp', () => { + it('prints the description and arguments of a single tool', async () => { + vi.mocked(listMcpTools).mockResolvedValue([ + { + name: 'get-documentation', + description: 'Get docs for a component.', + inputSchema: { + properties: { id: { type: 'string', description: 'Documentation id' } }, + required: ['id'], + }, + }, + ]); + + const result = await runAiToolHelp('get-documentation', { cwd: '/projects/foo' }); + expect(result.exitCode).toBe(0); + expect(result.output).toContain('Usage: storybook ai get-documentation'); + expect(result.output).toContain('Get docs for a component.'); + expect(result.output).toContain('- `--id` (string, required): Documentation id'); + }); + + it('is reachable through runAiTool via a --help token after the tool name', async () => { + vi.mocked(listMcpTools).mockResolvedValue([ + { name: 'get-documentation', description: 'Get docs.' }, + ]); + const result = await runAiTool('get-documentation', ['--help'], { cwd: '/projects/foo' }); + expect(result.exitCode).toBe(0); + expect(result.output).toContain('Usage: storybook ai get-documentation'); + expect(callMcpTool).not.toHaveBeenCalled(); + }); + + it('honors a --port token given after the command name on the help path', async () => { + const result = await runAiTool('get-documentation', ['--port', '9999', '--help'], { + cwd: '/projects/foo', + }); + expect(result.exitCode).toBe(1); + expect(result.output).toContain('not on port `9999`'); + }); + + it('lists the available tools for an unknown tool name', async () => { + const result = await runAiToolHelp('no-such-tool', { cwd: '/projects/foo' }); + expect(result.exitCode).toBe(1); + expect(result.output).toContain('Unknown command `no-such-tool`'); + expect(result.output).toContain('- `list-all-documentation`'); + }); + + it('prints repair markdown and exits non-zero on intercepts', async () => { + vi.mocked(readRegistry).mockResolvedValue([]); + const result = await runAiToolHelp('get-documentation', { cwd: '/projects/foo' }); + expect(result.exitCode).toBe(1); + expect(result.output).toContain('Storybook is not running at this cwd'); + }); + + it('rejects an invalid --port', async () => { + const result = await runAiToolHelp('get-documentation', { + cwd: '/projects/foo', + port: 'abc', + }); + expect(result.exitCode).toBe(1); + expect(result.output).toContain('`--port` must be a port number'); + }); + + it('surfaces a friendly error when the MCP server is unreachable', async () => { + vi.mocked(listMcpTools).mockRejectedValue(new Error('connection refused')); + const result = await runAiToolHelp('get-documentation', { cwd: '/projects/foo' }); + expect(result.exitCode).toBe(1); + expect(result.output).toContain('Failed to reach the Storybook server at /mcp'); + }); +}); diff --git a/code/lib/cli-storybook/src/ai/mcp/run-tool.ts b/code/lib/cli-storybook/src/ai/mcp/run-tool.ts new file mode 100644 index 000000000000..596834ece95d --- /dev/null +++ b/code/lib/cli-storybook/src/ai/mcp/run-tool.ts @@ -0,0 +1,336 @@ +import { resolve } from 'node:path'; + +import { McpJsonRpcError, callMcpTool, listMcpTools } from './client.ts'; +import { getInterceptMarkdown } from './intercepts.ts'; +import { readRegistry } from './registry.ts'; +import { resolveInstance } from './resolve-instance.ts'; +import { parseToolArgs } from './tool-args.ts'; +import type { + InterceptReason, + McpToolDescriptor, + StorybookInstanceRecord, + ToolCallResult, +} from './types.ts'; + +export type AiToolRunResult = { exitCode: 0 | 1; output: string }; + +/** Injectable dependencies for tests. */ +export type AiToolRunDeps = { + registryDir?: string; + fetchImpl?: typeof fetch; +}; + +export type AiToolOptions = { + /** Project directory of the target Storybook; defaults to `process.cwd()`. */ + cwd?: string; + /** Port of the target Storybook, to address one specific instance when several share the cwd. */ + port?: string; + /** Raw JSON object with tool arguments (escape hatch for complex values). */ + json?: string; +}; + +/** + * Run a single MCP tool against the Storybook running at the target cwd and return its result as + * markdown. Intercept conditions (no running instance, addon missing, ...) return the same + * repair-instruction markdown as `@storybook/mcp-proxy`, with exit code 1. + */ +export async function runAiTool( + toolName: string, + toolArgTokens: string[], + options: AiToolOptions = {}, + deps: AiToolRunDeps = {} +): Promise { + const parsed = parseToolArgs(toolArgTokens, { + cwd: options.cwd, + port: options.port, + json: options.json, + }); + if (!parsed.ok) { + return { exitCode: 1, output: parsed.error }; + } + if (parsed.help) { + return toolHelp(toolName, parsed.cwd, parsed.port, deps); + } + + const resolution = await resolveReadyInstance(parsed.cwd, parsed.port, deps); + if (resolution.kind === 'error') { + return { exitCode: 1, output: resolution.output }; + } + const { record, matches } = resolution; + + try { + const result = await callMcpTool( + record, + { name: toolName, arguments: parsed.args }, + deps.fetchImpl + ); + if (result.isError) { + // addon-mcp reports unknown tools as an error *result* rather than a JSON-RPC error. + const unknownTool = await describeUnknownTool(record, toolName, deps.fetchImpl); + if (unknownTool) { + return { exitCode: 1, output: unknownTool }; + } + } + const siblings = matches.filter((r) => r !== record); + const sections = [ + ...(siblings.length > 0 ? [formatMultiInstanceWarning(record, siblings)] : []), + formatToolResult(result), + ]; + return { exitCode: result.isError ? 1 : 0, output: sections.join('\n\n') }; + } catch (error) { + if (error instanceof McpJsonRpcError) { + const unknownTool = await describeUnknownTool(record, toolName, deps.fetchImpl); + return { exitCode: 1, output: unknownTool ?? error.message }; + } + return { exitCode: 1, output: formatServerUnreachable(record, error) }; + } +} + +/** + * Build the "Storybook commands" help section listing the commands provided by the running + * Storybook, appended to `storybook ai --help`. Help must never fail, so any error degrades to a + * short note explaining why no commands are listed. + */ +export async function buildStorybookCommandsHelp( + options: AiToolOptions = {}, + deps: AiToolRunDeps = {} +): Promise { + const unavailable = (note: string) => `Storybook commands: (unavailable — ${note})`; + + const parsed = parseToolArgs([], { cwd: options.cwd, port: options.port }); + if (!parsed.ok) { + return unavailable(parsed.error); + } + + const resolution = await resolveReadyInstance(parsed.cwd, parsed.port, deps); + if (resolution.kind === 'error') { + return unavailable(helpUnavailableNote(resolution, parsed.port)); + } + const { record, matches } = resolution; + + let tools: McpToolDescriptor[]; + try { + tools = await listMcpTools(record, deps.fetchImpl); + } catch { + return unavailable(`the Storybook at ${record.url} could not be reached`); + } + if (tools.length === 0) { + return unavailable(`the Storybook at ${record.url} provides no commands`); + } + + const siblingPorts = matches.filter((r) => r !== record).map((r) => r.port); + const siblingNote = + siblingPorts.length > 0 + ? [ + `(${matches.length} instances are running at this cwd — using the most recently started, port ${record.port}; other ports: ${siblingPorts.join(', ')}. Pass \`--port\` to target a specific one.)`, + ] + : []; + + const width = Math.max(...tools.map((tool) => tool.name.length)) + 2; + const lines = tools.map((tool) => { + const summary = tool.description?.trim().split('\n')[0] ?? ''; + return ` ${tool.name.padEnd(width)}${summary}`; + }); + const version = record.storybookVersion ? `, Storybook ${record.storybookVersion}` : ''; + return [ + `Storybook commands (from the Storybook running at ${record.url}${version}):`, + ...siblingNote, + ...lines, + '', + `Run 'storybook ai --help' for a command's description and arguments.`, + ].join('\n'); +} + +/** One-line reason why the help section cannot list commands, accurate per intercept. */ +function helpUnavailableNote( + error: Extract, + port: number | undefined +): string { + switch (error.reason) { + case 'no-instance': + return 'no running Storybook detected at this cwd; start `storybook dev` to list its commands'; + case 'port-mismatch': + return `no instance on port \`${port}\` at this cwd — running ports: ${error.records + .map((r) => r.port) + .join(', ')}`; + case 'mcp-starting': + return 'the Storybook at this cwd is still starting up; retry in a moment'; + case 'addon-missing': + return 'the running Storybook does not provide commands — install `@storybook/addon-mcp`'; + case 'mcp-error': + return "the running Storybook's command server reported an error"; + default: { + const unhandled: never = error.reason; + throw new Error(`Unhandled intercept reason: ${unhandled as string}`); + } + } +} + +/** Show the description and arguments of a single command (`storybook ai --help`). */ +export async function runAiToolHelp( + toolName: string, + options: AiToolOptions = {}, + deps: AiToolRunDeps = {} +): Promise { + const parsed = parseToolArgs([], { cwd: options.cwd, port: options.port }); + if (!parsed.ok) { + return { exitCode: 1, output: parsed.error }; + } + return toolHelp(toolName, parsed.cwd, parsed.port, deps); +} + +async function toolHelp( + toolName: string, + cwd: string | undefined, + port: number | undefined, + deps: AiToolRunDeps +): Promise { + const resolution = await resolveReadyInstance(cwd, port, deps); + if (resolution.kind === 'error') { + return { exitCode: 1, output: resolution.output }; + } + const { record } = resolution; + + let tools: McpToolDescriptor[]; + try { + tools = await listMcpTools(record, deps.fetchImpl); + } catch (error) { + return { exitCode: 1, output: formatServerUnreachable(record, error) }; + } + + const tool = tools.find((candidate) => candidate.name === toolName); + if (!tool) { + return { exitCode: 1, output: formatUnknownTool(toolName, tools, record) }; + } + return { exitCode: 0, output: formatToolHelp(tool) }; +} + +function formatServerUnreachable(record: StorybookInstanceRecord, error: unknown): string { + return `Failed to reach the Storybook server at ${record.mcp.endpoint ?? '(no endpoint)'}: ${ + error instanceof Error ? error.message : String(error) + }`; +} + +type InstanceResolution = + | { kind: 'ok'; record: StorybookInstanceRecord; matches: StorybookInstanceRecord[] } + | { + kind: 'error'; + output: string; + reason: InterceptReason; + records: StorybookInstanceRecord[]; + }; + +/** + * Resolve the running Storybook instance for `cwdInput` via the registry. No version or + * installed checks: the CLI is invoked as `npx storybook`, so the fact that it is executing + * already proves the project has a compatible Storybook. + */ +async function resolveReadyInstance( + cwdInput: string | undefined, + port: number | undefined, + deps: AiToolRunDeps +): Promise { + const cwd = resolve(cwdInput ?? process.cwd()); + + const records = await readRegistry(deps.registryDir); + const resolution = resolveInstance(records, cwd, port); + + if (resolution.kind === 'intercept') { + return { + kind: 'error', + output: getInterceptMarkdown(resolution.reason, { records: resolution.records, port }), + reason: resolution.reason, + records: resolution.records ?? [], + }; + } + + return { kind: 'ok', record: resolution.record, matches: resolution.matches }; +} + +/** + * Build the "unknown tool" error listing the available tools, or null when the tool does exist + * (the JSON-RPC error had another cause) or the tool list cannot be fetched. + */ +async function describeUnknownTool( + record: StorybookInstanceRecord, + toolName: string, + fetchImpl?: typeof fetch +): Promise { + let tools: McpToolDescriptor[]; + try { + tools = await listMcpTools(record, fetchImpl); + } catch { + return null; + } + if (tools.some((tool) => tool.name === toolName)) { + return null; + } + return formatUnknownTool(toolName, tools, record); +} + +function formatUnknownTool( + toolName: string, + tools: McpToolDescriptor[], + record: StorybookInstanceRecord +): string { + return `Unknown command \`${toolName}\`. The Storybook running at ${record.url} provides: + +${tools.map((tool) => `- \`${tool.name}\``).join('\n')} + +Run \`storybook ai --help\` for all commands, or \`storybook ai --help\` for a command's arguments.`; +} + +/** Render a tools/call result as markdown: text content verbatim, other content as JSON blocks. */ +function formatToolResult(result: ToolCallResult): string { + const content = result.content ?? []; + if (content.length === 0) { + return '(the command returned no content)'; + } + return content + .map((item) => + item.type === 'text' && typeof item.text === 'string' + ? item.text + : `\`\`\`json\n${JSON.stringify(item, null, 2)}\n\`\`\`` + ) + .join('\n\n'); +} + +function formatToolHelp(tool: McpToolDescriptor): string { + const lines = [`Usage: storybook ai ${tool.name} [--key value ...]`]; + if (tool.description) { + lines.push('', tool.description.trim()); + } + const properties = Object.entries(tool.inputSchema?.properties ?? {}); + if (properties.length > 0) { + const required = new Set(tool.inputSchema?.required ?? []); + lines.push( + '', + 'Arguments:', + ...properties.map(([name, schema]) => { + const meta = [schema.type, required.has(name) ? 'required' : undefined] + .filter(Boolean) + .join(', '); + const description = schema.description ? `: ${schema.description}` : ''; + return `- \`--${name}\`${meta ? ` (${meta})` : ''}${description}`; + }) + ); + } + return lines.join('\n'); +} + +function formatMultiInstanceWarning( + chosen: StorybookInstanceRecord, + siblings: StorybookInstanceRecord[] +): string { + const all = [chosen, ...siblings]; + const lines = all.map((r) => { + const marker = r === chosen ? ' (used)' : ''; + return `> - pid \`${r.pid}\` at ${r.url} (status: \`${r.mcp.status}\`)${marker}`; + }); + return `> Warning: Multiple Storybook instances are running at this cwd. This call was sent to pid \`${chosen.pid}\`. +> +> Instances at \`${chosen.cwd}\`: +${lines.join('\n')} +> +> If results look unexpected, ask the user whether they want to stop the other instance(s).`; +} diff --git a/code/lib/cli-storybook/src/ai/mcp/tool-args.test.ts b/code/lib/cli-storybook/src/ai/mcp/tool-args.test.ts new file mode 100644 index 000000000000..c41d3e138efb --- /dev/null +++ b/code/lib/cli-storybook/src/ai/mcp/tool-args.test.ts @@ -0,0 +1,172 @@ +import { describe, expect, it } from 'vitest'; + +import { parseToolArgs } from './tool-args.ts'; + +function args(tokens: string[], defaults?: { cwd?: string; port?: string; json?: string }) { + const result = parseToolArgs(tokens, defaults); + if (!result.ok) { + throw new Error(`expected ok, got error: ${result.error}`); + } + return result; +} + +function error(tokens: string[], defaults?: { cwd?: string; port?: string; json?: string }) { + const result = parseToolArgs(tokens, defaults); + if (result.ok) { + throw new Error(`expected error, got ok: ${JSON.stringify(result.args)}`); + } + return result.error; +} + +describe('parseToolArgs', () => { + it('returns empty args for no tokens', () => { + expect(args([])).toEqual({ ok: true, cwd: undefined, port: undefined, help: false, args: {} }); + }); + + it('consumes --help and -h as a help request instead of forwarding them', () => { + expect(args(['--help'])).toMatchObject({ help: true, args: {} }); + expect(args(['-h'])).toMatchObject({ help: true, args: {} }); + expect(args(['--id', 'x', '--help'])).toMatchObject({ help: true, args: { id: 'x' } }); + }); + + it('maps `--key value` pairs to tool arguments', () => { + expect(args(['--id', 'button-docs']).args).toEqual({ id: 'button-docs' }); + }); + + it('supports `--key=value`', () => { + expect(args(['--id=button-docs']).args).toEqual({ id: 'button-docs' }); + }); + + describe('JSON-parse coercion', () => { + it('coerces booleans, numbers and null', () => { + expect(args(['--a', 'true', '--b', '42', '--c', 'null']).args).toEqual({ + a: true, + b: 42, + c: null, + }); + }); + + it('coerces JSON arrays and objects', () => { + expect(args(['--ids', '["a","b"]', '--filter', '{"tag":"x"}']).args).toEqual({ + ids: ['a', 'b'], + filter: { tag: 'x' }, + }); + }); + + it('falls back to the raw string when the value is not valid JSON', () => { + expect(args(['--id', 'button-docs', '--path', 'src/Button.tsx']).args).toEqual({ + id: 'button-docs', + path: 'src/Button.tsx', + }); + }); + + it('unquotes explicitly quoted JSON strings', () => { + expect(args(['--id', '"true"']).args).toEqual({ id: 'true' }); + }); + + it('accepts negative numbers as values', () => { + expect(args(['--offset', '-1']).args).toEqual({ offset: -1 }); + }); + }); + + it('treats a bare `--flag` as true', () => { + expect(args(['--withStoryIds']).args).toEqual({ withStoryIds: true }); + expect(args(['--withStoryIds', '--id', 'x']).args).toEqual({ withStoryIds: true, id: 'x' }); + }); + + it('lets the last occurrence of a repeated key win', () => { + expect(args(['--id', 'a', '--id', 'b']).args).toEqual({ id: 'b' }); + }); + + describe('--cwd', () => { + it('consumes --cwd instead of forwarding it', () => { + expect(args(['--cwd', '/projects/foo', '--id', 'x'])).toEqual({ + ok: true, + cwd: '/projects/foo', + port: undefined, + help: false, + args: { id: 'x' }, + }); + }); + + it('uses the commander-parsed default when not repeated in the tokens', () => { + expect(args(['--id', 'x'], { cwd: '/projects/foo' }).cwd).toBe('/projects/foo'); + }); + + it('prefers a --cwd token over the commander-parsed default', () => { + expect(args(['--cwd', '/b'], { cwd: '/a' }).cwd).toBe('/b'); + }); + + it('errors when --cwd has no value', () => { + expect(error(['--cwd'])).toContain('`--cwd` requires a value'); + }); + }); + + describe('--port', () => { + it('consumes --port as a number instead of forwarding it', () => { + expect(args(['--port', '6006', '--id', 'x'])).toEqual({ + ok: true, + cwd: undefined, + port: 6006, + help: false, + args: { id: 'x' }, + }); + }); + + it('uses the commander-parsed default and lets a token override it', () => { + expect(args(['--id', 'x'], { port: '6006' }).port).toBe(6006); + expect(args(['--port', '6007'], { port: '6006' }).port).toBe(6007); + }); + + it('errors on non-numeric or out-of-range ports', () => { + expect(error(['--port', 'abc'])).toContain('`--port` must be a port number'); + expect(error(['--port', '0'])).toContain('`--port` must be a port number'); + expect(error(['--port', '65536'])).toContain('`--port` must be a port number'); + expect(error(['--port', '6006.5'])).toContain('`--port` must be a port number'); + }); + + it('errors when --port has no value', () => { + expect(error(['--port'])).toContain('`--port` requires a value'); + }); + }); + + describe('--json escape hatch', () => { + it('uses the JSON object as the tool arguments', () => { + expect(args(['--json', '{"id":"x","n":1}']).args).toEqual({ id: 'x', n: 1 }); + }); + + it('lets explicit --key flags override --json entries', () => { + expect(args(['--json', '{"id":"x","n":1}', '--id', 'y']).args).toEqual({ id: 'y', n: 1 }); + }); + + it('accepts --json parsed by commander before the tool name', () => { + expect(args(['--id', 'y'], { json: '{"id":"x","n":1}' }).args).toEqual({ id: 'y', n: 1 }); + }); + + it('errors on invalid JSON', () => { + expect(error(['--json', '{nope'])).toContain('`--json` must be valid JSON'); + }); + + it('errors when the JSON is not an object', () => { + expect(error(['--json', '[1,2]'])).toContain('must be a JSON object'); + expect(error(['--json', '"text"'])).toContain('must be a JSON object'); + expect(error(['--json', 'null'])).toContain('must be a JSON object'); + }); + + it('errors when --json has no value', () => { + expect(error(['--json'])).toContain('`--json` requires a value'); + }); + }); + + it('errors on positional tokens', () => { + expect(error(['positional'])).toContain('Unexpected argument `positional`'); + }); + + it('errors on a bare `--` separator', () => { + expect(error(['--'])).toContain('Unexpected argument `--`'); + }); + + it('errors on `--=value`', () => { + expect(error(['--=x'])).toContain('Invalid flag'); + }); +}); diff --git a/code/lib/cli-storybook/src/ai/mcp/tool-args.ts b/code/lib/cli-storybook/src/ai/mcp/tool-args.ts new file mode 100644 index 000000000000..ee1357fba742 --- /dev/null +++ b/code/lib/cli-storybook/src/ai/mcp/tool-args.ts @@ -0,0 +1,134 @@ +export type ParsedToolArgs = + | { + ok: true; + cwd: string | undefined; + port: number | undefined; + help: boolean; + args: Record; + } + | { ok: false; error: string }; + +/** + * Parse the pass-through tokens after `storybook ai ` into MCP tool arguments. + * + * - `--key value` and `--key=value` become tool arguments; values are coerced by attempting + * `JSON.parse`, falling back to the raw string. + * - A bare `--key` (no value) becomes `true`. + * - `--json ''` is an escape hatch providing the raw argument object; explicit `--key` + * flags override its entries. + * - `--cwd `, `--port ` and `--help`/`-h` are consumed by the CLI itself and never + * forwarded to the tool. + * + * `defaults` carries `--cwd`/`--port`/`--json` values that commander already parsed before the + * tool name; the same flags appearing after the tool name take precedence. + */ +export function parseToolArgs( + tokens: string[], + defaults: { cwd?: string; port?: string; json?: string } = {} +): ParsedToolArgs { + let cwd = defaults.cwd; + let rawPort = defaults.port; + let rawJson = defaults.json; + let help = false; + const flagArgs: Record = {}; + + let i = 0; + while (i < tokens.length) { + const token = tokens[i]; + i += 1; + + if (token === '--help' || token === '-h') { + help = true; + continue; + } + + if (!token.startsWith('--') || token === '--') { + return { + ok: false, + error: `Unexpected argument \`${token}\`. Command arguments must be passed as \`--key value\` flags (or via \`--json ''\`).`, + }; + } + + let key = token.slice(2); + let value: string | undefined; + const equalsIndex = key.indexOf('='); + if (equalsIndex !== -1) { + value = key.slice(equalsIndex + 1); + key = key.slice(0, equalsIndex); + } else if (i < tokens.length && !tokens[i].startsWith('--')) { + value = tokens[i]; + i += 1; + } + + if (key === '') { + return { ok: false, error: `Invalid flag \`${token}\`.` }; + } + + if (key === 'cwd') { + if (value === undefined) { + return { ok: false, error: '`--cwd` requires a value.' }; + } + cwd = value; + continue; + } + + if (key === 'port') { + if (value === undefined) { + return { ok: false, error: '`--port` requires a value.' }; + } + rawPort = value; + continue; + } + + if (key === 'json') { + if (value === undefined) { + return { ok: false, error: '`--json` requires a value.' }; + } + rawJson = value; + continue; + } + + flagArgs[key] = value === undefined ? true : coerceValue(value); + } + + let port: number | undefined; + if (rawPort !== undefined) { + port = Number(rawPort); + if (!Number.isInteger(port) || port < 1 || port > 65535) { + return { + ok: false, + error: `\`--port\` must be a port number (1-65535), got \`${rawPort}\`.`, + }; + } + } + + let jsonArgs: Record = {}; + if (rawJson !== undefined) { + let parsed: unknown; + try { + parsed = JSON.parse(rawJson); + } catch (error) { + return { + ok: false, + error: `\`--json\` must be valid JSON: ${error instanceof Error ? error.message : String(error)}`, + }; + } + if (typeof parsed !== 'object' || parsed === null || Array.isArray(parsed)) { + return { + ok: false, + error: '`--json` must be a JSON object, e.g. \'{"id": "button-docs"}\'.', + }; + } + jsonArgs = parsed as Record; + } + + return { ok: true, cwd, port, help, args: { ...jsonArgs, ...flagArgs } }; +} + +function coerceValue(raw: string): unknown { + try { + return JSON.parse(raw); + } catch { + return raw; + } +} diff --git a/code/lib/cli-storybook/src/ai/mcp/types.ts b/code/lib/cli-storybook/src/ai/mcp/types.ts new file mode 100644 index 000000000000..fc6329479128 --- /dev/null +++ b/code/lib/cli-storybook/src/ai/mcp/types.ts @@ -0,0 +1,82 @@ +/** + * Reader-side types for the `storybook ai ` MCP passthrough, copied from + * `@storybook/mcp-proxy` (storybookjs/mcp) per storybookjs/storybook#35124. The + * writer side lives in `code/core/src/core-server/utils/runtime-instance-registry.ts`; + * this reader is intentionally more lenient (extra statuses, optional fields) so it + * also accepts records written by other Storybook versions and wrappers. + */ +import * as v from 'valibot'; + +/** + * The in-repo writer only emits `not-installed` and `ready`; `starting` and `error` are written by + * external wrappers (e.g. the storybookjs/mcp launch script) and must keep being dispatched here. + */ +export const McpStatusSchema = v.picklist(['not-installed', 'starting', 'ready', 'error']); +export type McpStatus = v.InferOutput; + +/** + * A single Storybook runtime record written under the registry dir (default + * `~/.storybook/instances`). One file per running `storybook dev` instance. + * Spec: storybookjs/storybook#34826. + */ +export const StorybookInstanceRecordSchema = v.object({ + schemaVersion: v.literal(1), + instanceId: v.string(), + pid: v.pipe(v.number(), v.minValue(1), v.integer()), + cwd: v.string(), + url: v.string(), + port: v.pipe(v.number(), v.minValue(1), v.maxValue(65535), v.integer()), + storybookVersion: v.optional(v.string()), + startedAt: v.optional(v.string()), + updatedAt: v.optional(v.string()), + mcp: v.object({ + status: McpStatusSchema, + endpoint: v.optional(v.string()), + }), +}); +export type StorybookInstanceRecord = v.InferOutput; + +export type InterceptReason = + | 'no-instance' + | 'port-mismatch' + | 'addon-missing' + | 'mcp-starting' + | 'mcp-error'; + +/** + * Result of an MCP `tools/call` request, as returned by `@storybook/addon-mcp`. Loose: servers may + * legally attach extra fields (`_meta`, `structuredContent`, image/audio content properties); we + * validate only what the CLI renders and pass the rest through. + */ +export const ToolResultContentItemSchema = v.looseObject({ + type: v.string(), + text: v.optional(v.string()), +}); +export type ToolResultContentItem = v.InferOutput; + +export const ToolCallResultSchema = v.looseObject({ + content: v.optional(v.array(ToolResultContentItemSchema)), + isError: v.optional(v.boolean()), +}); +export type ToolCallResult = v.InferOutput; + +/** A tool descriptor from an MCP `tools/list` response. */ +export const McpToolDescriptorSchema = v.looseObject({ + name: v.string(), + description: v.optional(v.string()), + inputSchema: v.optional( + v.looseObject({ + properties: v.optional( + v.record( + v.string(), + v.looseObject({ + type: v.optional(v.string()), + description: v.optional(v.string()), + }) + ) + ), + required: v.optional(v.array(v.string())), + }) + ), +}); +export type McpToolDescriptor = v.InferOutput; diff --git a/code/lib/cli-storybook/src/bin/run.ts b/code/lib/cli-storybook/src/bin/run.ts index b784848b26aa..5f0ac08b0de6 100644 --- a/code/lib/cli-storybook/src/bin/run.ts +++ b/code/lib/cli-storybook/src/bin/run.ts @@ -24,6 +24,7 @@ import { link } from '../link.ts'; import { migrate } from '../migrate.ts'; import { sandbox } from '../sandbox.ts'; import { aiSetup } from '../ai/index.ts'; +import { isAiCliFeatureEnabled, registerAiMcpPassthrough } from '../ai/mcp/register.ts'; import { type UpgradeOptions, upgrade } from '../upgrade.ts'; addToGlobalContext('cliVersion', versions.storybook); @@ -329,6 +330,12 @@ aiCommand.action(() => { aiCommand.outputHelp(); }); +// Experimental `storybook ai ` passthrough to the local Storybook MCP server +// (storybookjs/storybook#35124). Overrides the help-only action above when enabled. +if (isAiCliFeatureEnabled()) { + registerAiMcpPassthrough(program, aiCommand); +} + program.on('command:*', ([invalidCmd]) => { let errorMessage = ` Invalid command: ${picocolors.bold(invalidCmd)}.\n See --help for a list of available commands.`; const availableCommands = program.commands.map((cmd) => cmd.name()); diff --git a/yarn.lock b/yarn.lock index 6635182c9967..9c5289d30a31 100644 --- a/yarn.lock +++ b/yarn.lock @@ -8507,6 +8507,7 @@ __metadata: globby: "npm:^14.1.0" jscodeshift: "npm:^0.15.1" leven: "npm:^4.0.0" + memfs: "npm:^4.11.1" p-limit: "npm:^7.2.0" picocolors: "npm:^1.1.0" semver: "npm:^7.7.3" @@ -8516,6 +8517,7 @@ __metadata: tinyclip: "npm:^0.1.12" ts-dedent: "npm:^2.0.0" typescript: "npm:^5.8.3" + valibot: "npm:^1.4.0" bin: cli: ./dist/bin/index.js languageName: unknown