Merge pull request #42 from storybookjs/mcp-telemetry

Add optional handlers to `@storybook/mcp` and telemetry to the new tools in `@storybook/addon-mcp`

Author: JReinhold

.changeset/chatty-bees-cut.md

@@ -0,0 +1,5 @@
+---
+'@storybook/mcp': patch
+---
+
+Add optional event handlers to the tool calls, so you can optionally run some functionality on all tool calls

.changeset/dark-chairs-cover.md

@@ -0,0 +1,5 @@
+---
+'@storybook/addon-mcp': patch
+---
+
+Log telemetry when the additional @storybook/mcp tools are called

.github/copilot-instructions.md

@@ -29,8 +29,12 @@ The `@storybook/mcp` package (in `packages/mcp`) is framework-agnostic:
 
 - Uses `tmcp` with HTTP transport and Valibot schema validation
 - Factory pattern: `createStorybookMcpHandler()` returns a request handler
-- Context-based: handlers accept `StorybookContext` to override source URLs
+- Context-based: handlers accept `StorybookContext` to override source URLs and provide optional callbacks
 - **Exports tools and types** for reuse by `addon-mcp` and other consumers
+- **Optional handlers**: `StorybookContext` supports optional handlers that are called at various points, allowing consumers to track usage or collect telemetry:
+  - `onSessionInitialize`: Called when an MCP session is initialized
+  - `onListAllComponents`: Called when the list-all-components tool is invoked
+  - `onGetComponentDocumentation`: Called when the get-component-documentation tool is invoked
 
 ## Development Environment
 
@@ -186,6 +190,11 @@ export { addMyTool, MY_TOOL_NAME } from './tools/my-tool.ts';
   - Checks for `experimental_componentManifestGenerator` preset
   - Only registers `addListAllComponentsTool` and `addGetComponentDocumentationTool` when enabled
 - Context includes `source` URL pointing to `/manifests/components.json` endpoint
+- **Optional handlers for tracking**:
+  - `onSessionInitialize`: Called when an MCP session is initialized, receives context
+  - `onListAllComponents`: Called when list tool is invoked, receives context and manifest
+  - `onGetComponentDocumentation`: Called when get tool is invoked, receives context, input, found components, and not found IDs
+  - Addon-mcp uses these handlers to collect telemetry on tool usage
 
 **Storybook internals used:**
 

.github/instructions/addon-mcp.instructions.md

@@ -350,6 +350,36 @@ The addon collects anonymous usage data:
 - Tool usage events (which tools are called and how often)
 - Client information (MCP client name)
 
+**For addon-specific tools**: Telemetry is collected directly in the tool implementation using `collectTelemetry()`.
+
+**For reused tools from `@storybook/mcp`**: The addon uses optional handlers (`onListAllComponents`, `onGetComponentDocumentation`) provided by the `StorybookContext` to track usage. These handlers are passed in the context when calling `transport.respond()` in `mcp-handler.ts`:
+
+```typescript
+const addonContext: AddonContext = {
+	// ... other context properties
+	onListAllComponents: async ({ manifest }) => {
+		if (!disableTelemetry && server) {
+			await collectTelemetry({
+				event: 'tool:listAllComponents',
+				server,
+				componentCount: Object.keys(manifest.components).length,
+			});
+		}
+	},
+	onGetComponentDocumentation: async ({ input, foundComponents, notFoundIds }) => {
+		if (!disableTelemetry && server) {
+			await collectTelemetry({
+				event: 'tool:getComponentDocumentation',
+				server,
+				inputComponentCount: input.componentIds.length,
+				foundCount: foundComponents.length,
+				notFoundCount: notFoundIds.length,
+			});
+		}
+	},
+};
+```
+
 Telemetry respects Storybook's `disableTelemetry` config:
 
 ```typescript

.github/instructions/mcp.instructions.md

@@ -13,10 +13,11 @@ This is a Model Context Protocol (MCP) server for Storybook that serves knowledg
 ### Key Components
 
 - **MCP Server**: Built using the `tmcp` library with HTTP transport
-- **Tools System**: Extensible tool registration system (includes `list_all_components` and `get_component_documentation`)
+- **Tools System**: Extensible tool registration system with optional handlers for tracking tool usage
 - **Component Manifest**: Parses and formats component documentation including React prop information from react-docgen
 - **Schema Validation**: Uses Valibot for JSON schema validation via `@tmcp/adapter-valibot`
 - **HTTP Transport**: Provides HTTP-based MCP communication via `@tmcp/transport-http`
+- **Context System**: `StorybookContext` allows passing optional handlers (`onSessionInitialize`, `onListAllComponents`, `onGetComponentDocumentation`) that are called at various points when provided
 
 ### File Structure
 
@@ -68,7 +69,8 @@ Component manifests can include a `reactDocgen` property containing prop informa
    ```
 
 **Type serialization examples:**
-- Unions: `"primary" | "secondary"` 
+
+- Unions: `"primary" | "secondary"`
 - Functions: `(event: MouseEvent) => void`
 - Objects: `{ name: string; age?: number }`
 - Arrays: `string[]`
@@ -192,20 +194,37 @@ To add a new MCP tool:
 1. Create a new file in `src/tools/` (e.g., `src/tools/my-tool.ts`)
 2. Export a constant for the tool name
 3. Export an async function that adds the tool to the server:
+
    ```typescript
-   export async function addMyTool(server: McpServer) {
+   export async function addMyTool(server: McpServer<any, StorybookContext>) {
    	server.tool(
    		{
    			name: 'my_tool_name',
    			description: 'Tool description',
    		},
-   		() => ({
-   			content: [{ type: 'text', text: 'result' }],
-   		}),
+   		async () => {
+   			// Tool implementation
+   			const result = 'result';
+
+   			// Call optional handler if provided
+   			await server.ctx.custom?.onMyTool?.({
+   				context: server.ctx.custom,
+   				// ... any relevant data
+   			});
+
+   			return {
+   				content: [{ type: 'text', text: result }],
+   			};
+   		},
    	);
    }
    ```
-4. Import and call the function in `src/index.ts` after `addListTool(server)`
+
+4. Import and call the function in `src/index.ts` in the `createStorybookMcpHandler` function
+5. If adding an optional handler:
+   - Add the handler type to `StorybookContext` in `src/types.ts`
+   - Document what parameters the handler receives
+   - Export the handler type from `src/index.ts` if it should be usable by consumers
 
 ## MCP Protocol
 

packages/addon-mcp/src/mcp-handler.ts

@@ -19,7 +19,7 @@ import { logger } from 'storybook/internal/node-logger';
 let transport: HttpTransport<AddonContext> | undefined;
 let origin: string | undefined;
 // Promise that ensures single initialization, even with concurrent requests
-let initialize: Promise<void> | undefined;
+let initialize: Promise<McpServer<any, AddonContext>> | undefined;
 
 const initializeMCPServer = async (options: Options) => {
 	const server = new McpServer(
@@ -67,6 +67,7 @@ const initializeMCPServer = async (options: Options) => {
 
 	origin = `http://localhost:${options.port}`;
 	logger.debug('MCP server origin:', origin);
+	return server;
 };
 
 /**
@@ -79,16 +80,13 @@ export const mcpServerHandler = async (
 	next: Connect.NextFunction,
 	options: Options,
 ) => {
-	const { disableTelemetry = false } = await options.presets.apply<CoreConfig>(
-		'core',
-		{},
-	);
+	const disableTelemetry = options.disableTelemetry ?? false;
 
 	// Initialize MCP server and transport on first request, with concurrency safety
 	if (!initialize) {
 		initialize = initializeMCPServer(options);
 	}
-	await initialize;
+	const server = await initialize;
 
 	// Convert Node.js request to Web API Request
 	const webRequest = await incomingMessageToWebRequest(req);
@@ -100,6 +98,29 @@ export const mcpServerHandler = async (
 		disableTelemetry,
 		// Source URL for component manifest tools - points to the manifest endpoint
 		source: `${origin}/manifests/components.json`,
+		// Telemetry handlers for component manifest tools
+		...(!disableTelemetry && {
+			onListAllComponents: async ({ manifest }) => {
+				await collectTelemetry({
+					event: 'tool:listAllComponents',
+					server,
+					componentCount: Object.keys(manifest.components).length,
+				});
+			},
+			onGetComponentDocumentation: async ({
+				input,
+				foundComponents,
+				notFoundIds,
+			}) => {
+				await collectTelemetry({
+					event: 'tool:getComponentDocumentation',
+					server,
+					inputComponentCount: input.componentIds.length,
+					foundCount: foundComponents.length,
+					notFoundCount: notFoundIds.length,
+				});
+			},
+		}),
 	};
 
 	const response = await transport!.respond(webRequest, addonContext);

packages/mcp/src/index.ts

@@ -19,10 +19,54 @@ export {
 // Export types for reuse
 export type { StorybookContext } from './types.ts';
 
+// copied from tmcp internals as it's not exposed
+type InitializeRequestParams = {
+	protocolVersion: string;
+	capabilities: {
+		experimental?: {} | undefined;
+		sampling?: {} | undefined;
+		elicitation?: {} | undefined;
+		roots?:
+			| {
+					listChanged?: boolean | undefined;
+			  }
+			| undefined;
+	};
+	clientInfo: {
+		icons?:
+			| {
+					src: string;
+					mimeType?: string | undefined;
+					sizes?: string[] | undefined;
+			  }[]
+			| undefined;
+		version: string;
+		websiteUrl?: string | undefined;
+		name: string;
+		title?: string | undefined;
+	};
+};
+
+/**
+ * Options for creating a Storybook MCP handler.
+ * Extends StorybookContext with server-level configuration.
+ */
+export interface StorybookMcpHandlerOptions extends StorybookContext {
+	/**
+	 * Optional handler called when an MCP session is initialized.
+	 * This is only valid at the handler creation level, not per-request.
+	 * Receives the initialize request parameters from the MCP protocol.
+	 */
+	onSessionInitialize?: (
+		initializeRequestParams: InitializeRequestParams,
+	) => void | Promise<void>;
+}
+export type { ComponentManifest, ComponentManifestMap } from './types.ts';
+
 type Handler = (req: Request, context?: StorybookContext) => Promise<Response>;
 
 export const createStorybookMcpHandler = async (
-	options: StorybookContext = {},
+	options: StorybookMcpHandlerOptions = {},
 ): Promise<Handler> => {
 	const adapter = new ValibotJsonSchemaAdapter();
 	const server = new McpServer(
@@ -41,16 +85,24 @@ export const createStorybookMcpHandler = async (
 		},
 	).withContext<StorybookContext>();
 
+	if (options.onSessionInitialize) {
+		server.on('initialize', options.onSessionInitialize);
+	}
+
 	await addListAllComponentsTool(server);
 	await addGetComponentDocumentationTool(server);
 
 	const transport = new HttpTransport(server, { path: null });
 
 	return (async (req, context) => {
-		const source = context?.source ?? options.source;
-		const manifestProvider =
-			context?.manifestProvider ?? options.manifestProvider;
-
-		return await transport.respond(req, { source, manifestProvider });
+		return await transport.respond(req, {
+			source: context?.source ?? options.source,
+			manifestProvider: context?.manifestProvider ?? options.manifestProvider,
+			onListAllComponents:
+				context?.onListAllComponents ?? options.onListAllComponents,
+			onGetComponentDocumentation:
+				context?.onGetComponentDocumentation ??
+				options.onGetComponentDocumentation,
+		});
 	}) as Handler;
 };

packages/mcp/src/tools/get-component-documentation.test.ts

@@ -300,6 +300,40 @@ describe('getComponentDocumentationTool', () => {
 		`);
 	});
 
+	it('should call onGetComponentDocumentation handler when provided', async () => {
+		const handler = vi.fn();
+
+		const request = {
+			jsonrpc: '2.0' as const,
+			id: 2,
+			method: 'tools/call',
+			params: {
+				name: GET_TOOL_NAME,
+				arguments: {
+					componentIds: ['button', 'card', 'non-existent'],
+				},
+			},
+		};
+
+		// Pass the handler in the context for this specific request
+		await server.receive(request, {
+			custom: { onGetComponentDocumentation: handler },
+		});
+
+		expect(handler).toHaveBeenCalledTimes(1);
+		expect(handler).toHaveBeenCalledWith({
+			context: expect.objectContaining({
+				onGetComponentDocumentation: handler,
+			}),
+			input: { componentIds: ['button', 'card', 'non-existent'] },
+			foundComponents: [
+				expect.objectContaining({ id: 'button', name: 'Button' }),
+				expect.objectContaining({ id: 'card', name: 'Card' }),
+			],
+			notFoundIds: ['non-existent'],
+		});
+	});
+
 	it('should include props section when reactDocgen is present', async () => {
 		const manifestWithReactDocgen = {
 			v: 1,

packages/mcp/src/tools/get-component-documentation.ts

@@ -1,6 +1,6 @@
 import * as v from 'valibot';
 import type { McpServer } from 'tmcp';
-import type { StorybookContext } from '../types.ts';
+import type { StorybookContext, ComponentManifest } from '../types.ts';
 import { getManifest, errorToMCPContent } from '../utils/get-manifest.ts';
 import { formatComponentManifest } from '../utils/format-manifest.ts';
 
@@ -36,6 +36,7 @@ export async function addGetComponentDocumentationTool(
 
 				const content = [];
 				const notFoundIds: string[] = [];
+				const foundComponents: ComponentManifest[] = [];
 
 				for (const componentId of input.componentIds) {
 					const component = manifest.components[componentId];
@@ -45,6 +46,7 @@ export async function addGetComponentDocumentationTool(
 						continue;
 					}
 
+					foundComponents.push(component);
 					content.push({
 						type: 'text' as const,
 						text: formatComponentManifest(component),
@@ -62,6 +64,13 @@ export async function addGetComponentDocumentationTool(
 					});
 				}
 
+				await server.ctx.custom?.onGetComponentDocumentation?.({
+					context: server.ctx.custom,
+					input: { componentIds: input.componentIds },
+					foundComponents,
+					notFoundIds,
+				});
+
 				return {
 					content,
 					...(allNotFound && { isError: true }),