diff --git a/code/core/src/cli/ai/mcp/run-tool.test.ts b/code/core/src/cli/ai/mcp/run-tool.test.ts index 01c4f4dc6b77..d911492628ff 100644 --- a/code/core/src/cli/ai/mcp/run-tool.test.ts +++ b/code/core/src/cli/ai/mcp/run-tool.test.ts @@ -742,6 +742,154 @@ describe('runAiToolHelp', () => { expect(result.outcome).toEqual({ kind: 'help' }); }); + it('recurses into array item and object property schemas so nested fields self-document', async () => { + vi.mocked(loadStorybookAiMetadata).mockResolvedValue({ + instructions: 'Follow the story workflow.', + tools: [ + { + name: 'display-review', + description: 'Publish a review.', + inputSchema: { + properties: { + collections: { + type: 'array', + items: { + type: 'object', + properties: { + title: { type: 'string', description: 'What the group is' }, + storyIds: { + type: 'array', + items: { type: 'string' }, + description: 'IDs the group renders', + }, + }, + required: ['title', 'storyIds'], + }, + }, + changedFiles: { + type: 'array', + items: { type: 'string' }, + description: 'Paths you changed', + }, + }, + required: ['collections', 'changedFiles'], + }, + }, + ], + }); + + const result = await runAiToolHelp('display-review', { cwd: '/projects/foo' }); + + expect(result.exitCode).toBe(0); + // Nested array-item fields self-document; the primitive `changedFiles` array + // stays flat (no "each item:" expansion). + expect(result.output).toMatchInlineSnapshot(` + "Usage: storybook ai display-review [--key value ...] + + Publish a review. + + Execution: requires a running Storybook. + + Arguments: + - \`--collections\` (array of object, required) + each item: + - \`title\` (string, required): What the group is + - \`storyIds\` (array of string, required): IDs the group renders + - \`--changedFiles\` (array of string, required): Paths you changed" + `); + }); + + it('describes anyOf/oneOf union variants in help', async () => { + vi.mocked(loadStorybookAiMetadata).mockResolvedValue({ + instructions: 'Follow the story workflow.', + tools: [ + { + name: 'preview-stories', + description: 'Preview stories.', + inputSchema: { + properties: { + stories: { + type: 'array', + items: { + anyOf: [ + { + type: 'object', + properties: { storyId: { type: 'string', description: 'A story id' } }, + required: ['storyId'], + }, + { + type: 'object', + properties: { + exportName: { type: 'string', description: 'An export name' }, + }, + required: ['exportName'], + }, + ], + }, + description: 'Stories to preview', + }, + }, + required: ['stories'], + }, + }, + ], + }); + + const result = await runAiToolHelp('preview-stories', { cwd: '/projects/foo' }); + + expect(result.exitCode).toBe(0); + expect(result.output).toMatchInlineSnapshot(` + "Usage: storybook ai preview-stories [--key value ...] + + Preview stories. + + Execution: requires a running Storybook. + + Arguments: + - \`--stories\` (array, required): Stories to preview + each item: + option 1 + - \`storyId\` (string, required): A story id + option 2 + - \`exportName\` (string, required): An export name" + `); + }); + + it('falls back to the top-level type and description when an argument schema cannot be modeled', async () => { + vi.mocked(loadStorybookAiMetadata).mockResolvedValue({ + instructions: 'Follow the story workflow.', + tools: [ + { + name: 'exotic-tool', + description: 'Has an unmodeled schema.', + inputSchema: { + properties: { + // `items: false` is valid JSON Schema but outside the node schema, so + // validation fails and the help must degrade rather than throw or drop it. + weird: { type: 'string', description: 'An odd one', items: false }, + }, + required: ['weird'], + }, + }, + ], + }); + + const result = await runAiToolHelp('exotic-tool', { cwd: '/projects/foo' }); + + expect(result.exitCode).toBe(0); + // Degrades to the top-level type/description; no crash, no dropped argument. + expect(result.output).toMatchInlineSnapshot(` + "Usage: storybook ai exotic-tool [--key value ...] + + Has an unmodeled schema. + + Execution: requires a running Storybook. + + Arguments: + - \`--weird\` (string, required): An odd one" + `); + }); + it('marks local commands in single-command help', async () => { vi.mocked(loadStorybookAiMetadata).mockResolvedValue({ instructions: 'Follow the story workflow.', diff --git a/code/core/src/cli/ai/mcp/run-tool.ts b/code/core/src/cli/ai/mcp/run-tool.ts index 7741e8955044..9189d6591a4e 100644 --- a/code/core/src/cli/ai/mcp/run-tool.ts +++ b/code/core/src/cli/ai/mcp/run-tool.ts @@ -2,6 +2,7 @@ import { resolve } from 'node:path'; import * as v from 'valibot'; +import { detectAgent } from '../../../telemetry/detect-agent.ts'; import { McpJsonRpcError, callMcpTool, listMcpTools } from './client.ts'; import { getInterceptMarkdown } from './intercepts.ts'; import { @@ -9,17 +10,17 @@ import { resolveStorybookConfigDir, type StorybookAiMetadata, } from './local-metadata.ts'; -import { detectAgent } from '../../../telemetry/detect-agent.ts'; import { readRegistry } from './registry.ts'; import { resolveInstance } from './resolve-instance.ts'; import { parsePort, parseToolArgs } from './tool-args.ts'; -import { ToolCallResultSchema } from './types.ts'; import type { InterceptReason, + JsonSchemaNode, McpToolDescriptor, StorybookInstanceRecord, ToolCallResult, } from './types.ts'; +import { JsonSchemaNodeSchema, ToolCallResultSchema } from './types.ts'; /** * Why an invocation failed before any command executed, for the `ai-command` telemetry event @@ -543,6 +544,69 @@ function getLocalToolNames(metadata: StorybookAiMetadata): Set { return new Set(Object.keys(metadata.localTools ?? {})); } +const MAX_SCHEMA_DEPTH = 4; + +function schemaItem(schema: JsonSchemaNode): JsonSchemaNode | undefined { + return Array.isArray(schema.items) ? schema.items[0] : schema.items; +} + +function schemaTypeLabel(schema: JsonSchemaNode): string | undefined { + if (schema.type === 'array') { + const item = schemaItem(schema); + return item?.type ? `array of ${item.type}` : 'array'; + } + if (schema.anyOf || schema.oneOf) { + return 'one of'; + } + return schema.type; +} + +function schemaChildLines(schema: JsonSchemaNode, indent: string, depth: number): string[] { + if (depth <= 0) { + return []; + } + const lines: string[] = []; + + if (schema.type === 'object' && schema.properties) { + const required = new Set(schema.required ?? []); + for (const [name, child] of Object.entries(schema.properties)) { + lines.push(...schemaLines(`\`${name}\``, child, required.has(name), indent, depth)); + } + } + + const item = schemaItem(schema); + if (item && ((item.type === 'object' && item.properties) || item.anyOf || item.oneOf)) { + lines.push(`${indent}each item:`); + lines.push(...schemaChildLines(item, `${indent} `, depth - 1)); + } + + const variants = schema.anyOf ?? schema.oneOf; + if (variants) { + variants.forEach((variant, index) => { + const suffix = variant.description ? `: ${variant.description}` : ''; + lines.push(`${indent}option ${index + 1}${suffix}`); + lines.push(...schemaChildLines(variant, `${indent} `, depth - 1)); + }); + } + + return lines; +} + +function schemaLines( + label: string, + schema: JsonSchemaNode, + isRequired: boolean, + indent: string, + depth: number +): string[] { + const meta = [schemaTypeLabel(schema), isRequired ? 'required' : undefined] + .filter(Boolean) + .join(', '); + const description = schema.description ? `: ${schema.description}` : ''; + const head = `${indent}- ${label}${meta ? ` (${meta})` : ''}${description}`; + return [head, ...schemaChildLines(schema, `${indent} `, depth - 1)]; +} + function formatToolHelp(tool: McpToolDescriptor, { local }: { local: boolean }): string { const lines = [`Usage: storybook ai ${tool.name} [--key value ...]`]; if (tool.description) { @@ -557,17 +621,17 @@ function formatToolHelp(tool: McpToolDescriptor, { local }: { local: boolean }): 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}`; - }) - ); + lines.push('', 'Arguments:'); + for (const [name, schema] of properties) { + // Validate at this boundary rather than tightening the lenient tools/list + // parse: an unmodeled schema shape falls back to its top-level type and + // description instead of dropping the argument or failing the command. + const parsed = v.safeParse(JsonSchemaNodeSchema, schema); + const node: JsonSchemaNode = parsed.success + ? parsed.output + : { type: schema.type, description: schema.description }; + lines.push(...schemaLines(`\`--${name}\``, node, required.has(name), '', MAX_SCHEMA_DEPTH)); + } } return lines.join('\n'); } diff --git a/code/core/src/cli/ai/mcp/types.ts b/code/core/src/cli/ai/mcp/types.ts index 4853e955ae6a..9a9f3dbff25f 100644 --- a/code/core/src/cli/ai/mcp/types.ts +++ b/code/core/src/cli/ai/mcp/types.ts @@ -67,6 +67,34 @@ export const ToolCallResultSchema = v.looseObject({ }); export type ToolCallResult = v.InferOutput; +/** + * A JSON Schema node, as far as the `storybook ai --help` renderer walks it: + * object `properties`, array `items`, and `anyOf`/`oneOf` variants. Recursive, so the + * schema is built with `v.lazy` and annotated with the interface (valibot can't infer + * a recursive type). Kept a `looseObject` so unknown JSON Schema keywords pass through. + */ +export interface JsonSchemaNode { + type?: string; + description?: string; + properties?: Record; + required?: string[]; + items?: JsonSchemaNode | JsonSchemaNode[]; + anyOf?: JsonSchemaNode[]; + oneOf?: JsonSchemaNode[]; +} + +export const JsonSchemaNodeSchema: v.GenericSchema = v.lazy(() => + v.looseObject({ + type: v.optional(v.string()), + description: v.optional(v.string()), + properties: v.optional(v.record(v.string(), JsonSchemaNodeSchema)), + required: v.optional(v.array(v.string())), + items: v.optional(v.union([JsonSchemaNodeSchema, v.array(JsonSchemaNodeSchema)])), + anyOf: v.optional(v.array(JsonSchemaNodeSchema)), + oneOf: v.optional(v.array(JsonSchemaNodeSchema)), + }) +); + /** A tool descriptor from an MCP `tools/list` response. */ export const McpToolDescriptorSchema = v.looseObject({ name: v.string(),