Merge pull request #9 from storybookjs/telemetry

Add telemetry for session starting, tool calling

Author: JReinhold

.changeset/modern-donkeys-cut.md

@@ -0,0 +1,5 @@
+---
+"@storybook/addon-mcp": patch
+---
+
+Add basic telemetry for sessions and tool calls

src/mcp-handler.ts

@@ -6,22 +6,43 @@ import { isInitializeRequest } from "@modelcontextprotocol/sdk/types.js";
 import pkgJson from "../package.json" with { type: "json" };
 import { registerStoryUrlsTool } from "./tools/get-story-urls";
 import { registerUIBuildingTool } from "./tools/get-ui-building-instructions";
-import type { Options } from "storybook/internal/types";
+import type { Options, CoreConfig } from "storybook/internal/types";
 import type { IncomingMessage, ServerResponse } from "node:http";
+import {
+  collectTelemetry,
+  mcpSessionIdToClientMap,
+  setDisableTelemetry,
+} from "./telemetry";
 
-function createMcpServer(options: Options) {
+async function createMcpServer(options: Options, client: string) {
   // New initialization request
   const transport = new StreamableHTTPServerTransport({
     sessionIdGenerator: () => randomUUID(),
-    onsessioninitialized: (sessionId) => {
-      // Store the transport by session ID
+    onsessioninitialized: async (sessionId) => {
       transports[sessionId] = transport;
+
+      const { disableTelemetry } = await options.presets.apply<CoreConfig>(
+        "core",
+        {},
+      );
+      setDisableTelemetry(disableTelemetry);
+      mcpSessionIdToClientMap[sessionId] = client;
+
+      await collectTelemetry({
+        event: "session:initialized",
+        mcpSessionId: sessionId,
+      });
     },
   });
 
   transport.onclose = () => {
-    if (transport.sessionId) {
-      delete transports[transport.sessionId];
+    if (!transport.sessionId) {
+      return;
+    }
+
+    delete transports[transport.sessionId];
+    if (mcpSessionIdToClientMap[transport.sessionId]) {
+      delete mcpSessionIdToClientMap[transport.sessionId];
     }
   };
 
@@ -61,7 +82,7 @@ const handlePostRequest = async (
     // Reuse existing transport
     transport = transports[sessionId];
   } else if (!sessionId && isInitializeRequest(body)) {
-    transport = await createMcpServer(options);
+    transport = await createMcpServer(options, body.params.clientInfo.name);
   } else {
     // Invalid request
     res.statusCode = 400;

src/telemetry.ts

@@ -0,0 +1,34 @@
+import { logger } from "storybook/internal/node-logger";
+import { telemetry } from "storybook/internal/telemetry";
+
+export async function collectTelemetry({
+  event,
+  mcpSessionId,
+  ...payload
+}: {
+  event: string;
+  mcpSessionId: string;
+  [key: string]: any;
+}) {
+  if (disableTelemetry) {
+    return;
+  }
+
+  try {
+    return await telemetry("addon-mcp" as any, {
+      event,
+      mcpSessionId,
+      client: mcpSessionIdToClientMap[mcpSessionId!],
+      ...payload,
+    });
+  } catch (error) {
+    logger.debug("Error collecting telemetry:", error);
+  }
+}
+
+export const mcpSessionIdToClientMap: Record<string, string> = {};
+
+let disableTelemetry = false;
+export const setDisableTelemetry = (value = false) => {
+  disableTelemetry = value;
+};

src/tools/get-story-urls.ts

@@ -4,6 +4,7 @@ import { storyNameFromExport } from "storybook/internal/csf";
 import type { Options, StoryIndex } from "storybook/internal/types";
 import { logger } from "storybook/internal/node-logger";
 import z from "zod";
+import { collectTelemetry } from "../telemetry";
 
 const inputStoriesSchema = z.array(
   z.object({
@@ -39,7 +40,7 @@ export function registerStoryUrlsTool({
         urls: outputUrlsSchema,
       },
     },
-    async ({ stories }) => {
+    async ({ stories }, { sessionId }) => {
       const index: StoryIndex = await (
         await fetch(`${origin}/index.json`)
       ).json();
@@ -48,6 +49,7 @@ export function registerStoryUrlsTool({
       logger.debug("index entries found:", entriesList.length);
 
       const result: z.infer<typeof outputUrlsSchema> = [];
+      let foundStoryCount = 0;
 
       for (const {
         exportName,
@@ -75,6 +77,7 @@ export function registerStoryUrlsTool({
         if (foundStoryId) {
           logger.debug("Found story ID:", foundStoryId);
           result.push(`${origin}/?path=/story/${foundStoryId}`);
+          foundStoryCount++;
         } else {
           logger.debug("No story found");
           let errorMessage = `No story found for export name "${exportName}" with absolute file path "${absoluteStoryPath}"`;
@@ -85,6 +88,13 @@ export function registerStoryUrlsTool({
         }
       }
 
+      await collectTelemetry({
+        event: "tool:getStoryUrls",
+        mcpSessionId: sessionId!,
+        inputStoryCount: stories.length,
+        outputStoryCount: foundStoryCount,
+      });
+
       return {
         content: [
           {

src/tools/get-ui-building-instructions.ts

@@ -1,5 +1,6 @@
 import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
 import { GET_STORY_URLS_TOOL_NAME } from "./get-story-urls";
+import { collectTelemetry } from "../telemetry";
 
 const INSTRUCTIONS = `
   # Writing User Interfaces Components
@@ -51,8 +52,14 @@ export function registerUIBuildingTool(server: McpServer) {
       limited to adding or updating new components, pages, screens or layouts.`,
       inputSchema: {},
     },
-    async () => ({
-      content: [{ type: "text", text: INSTRUCTIONS }],
-    }),
+    async (_, { sessionId }) => {
+      await collectTelemetry({
+        event: "tool:getUIBuildingInstructions",
+        mcpSessionId: sessionId!,
+      });
+      return {
+        content: [{ type: "text", text: INSTRUCTIONS }],
+      };
+    },
   );
 }