Agentic Setup: Add observability

code/core/src/builder-manager/index.ts

@@ -2,6 +2,7 @@ import { cp, rm, writeFile } from 'node:fs/promises';
 
 import { stringifyProcessEnvs } from 'storybook/internal/common';
 import { logger } from 'storybook/internal/node-logger';
+import { getLastEvents, getSessionId } from 'storybook/internal/telemetry';
 
 import { globalExternals } from '@fal-works/esbuild-plugin-global-externals';
 import { resolveModulePath } from 'exsolve';
@@ -199,6 +200,9 @@ const starter: StarterFunction = async function* starterGeneratorFn({
 
   // Build additional global values
   const globals: Record<string, any> = await buildFrameworkGlobalsFromOptions(options);
+  // Pass along recent CLI events to customise user onboarding and telemetry after `sb ai` commands.
+  globals.STORYBOOK_LAST_EVENTS = await getLastEvents();
+  globals.STORYBOOK_SESSION_ID = await getSessionId();
 
   yield;
 
@@ -297,6 +301,9 @@ const builder: BuilderFunction = async function* builderGeneratorFn({ startTime,
 
   // Build additional global values
   const globals: Record<string, any> = await buildFrameworkGlobalsFromOptions(options);
+  // Builds are long-lasting and shouldn't account for CLI events prior to building.
+  globals.STORYBOOK_LAST_EVENTS = {};
+  globals.STORYBOOK_SESSION_ID = undefined;
 
   yield;
 

code/core/src/common/satellite-addons.ts

@@ -7,6 +7,7 @@ export default [
   '@storybook/addon-coverage',
   '@storybook/addon-webpack5-compiler-babel',
   '@storybook/addon-webpack5-compiler-swc',
+  '@storybook/addon-mcp',
   // Storybook for React Native related packages
   // TODO: For Storybook 10, we should check about possible automigrations
   '@storybook/addon-ondevice-actions',

code/core/src/core-events/index.ts

@@ -91,6 +91,9 @@ enum events {
   // Story discovery and testing flow
   GHOST_STORIES_REQUEST = 'ghostStoriesRequest',
   GHOST_STORIES_RESPONSE = 'ghostStoriesResponse',
+  // AI analytics - ai prepare command
+  AI_PREPARE_ANALYTICS_RESPONSE = 'aiPrepareAnalyticsResponse',
+  AI_PREPARE_ANALYTICS_REQUEST = 'aiPrepareAnalyticsRequest',
   // Open a file in the code editor
   OPEN_IN_EDITOR_REQUEST = 'openInEditorRequest',
   OPEN_IN_EDITOR_RESPONSE = 'openInEditorResponse',
@@ -168,6 +171,8 @@ export const {
   ARGTYPES_INFO_RESPONSE,
   GHOST_STORIES_REQUEST,
   GHOST_STORIES_RESPONSE,
+  AI_PREPARE_ANALYTICS_RESPONSE,
+  AI_PREPARE_ANALYTICS_REQUEST,
   OPEN_IN_EDITOR_REQUEST,
   OPEN_IN_EDITOR_RESPONSE,
   MANAGER_INERT_ATTRIBUTE_CHANGED,

code/core/src/core-server/presets/common-preset.ts

@@ -41,6 +41,7 @@ import { initializeSaveStory } from '../utils/save-story/save-story.ts';
 import { parseStaticDir } from '../utils/server-statics.ts';
 import { type OptionsWithRequiredCache, initializeWhatsNew } from '../utils/whats-new.ts';
 import { getWsToken } from './wsToken.ts';
+import { initAIAnalyticsChannel } from '../server-channel/ai-prepare-channel.ts';
 
 const interpolate = (string: string, data: Record<string, string> = {}) =>
   Object.entries(data).reduce((acc, [k, v]) => acc.replace(new RegExp(`%${k}%`, 'g'), v), string);
@@ -278,6 +279,7 @@ export const experimental_serverChannel = async (
   initializeWhatsNew(channel, options, coreOptions);
   initializeSaveStory(channel, options, coreOptions);
 
+  initAIAnalyticsChannel(channel, options, coreOptions, () => storyIndexGeneratorPromise);
   initFileSearchChannel(channel, options, coreOptions);
   initCreateNewStoryChannel(channel, options, coreOptions);
   initGhostStoriesChannel(channel, options, coreOptions);

code/core/src/core-server/server-channel/ai-prepare-channel.ts

@@ -0,0 +1,131 @@
+import type { Channel } from 'storybook/internal/channels';
+import {
+  AI_PREPARE_ANALYTICS_REQUEST,
+  AI_PREPARE_ANALYTICS_RESPONSE,
+} from 'storybook/internal/core-events';
+import {
+  getLastEvents,
+  getStorybookMetadata,
+  isStoryCreatedByAIPrepare,
+  telemetry,
+} from 'storybook/internal/telemetry';
+import type { CoreConfig, Options } from 'storybook/internal/types';
+import { logger } from 'storybook/internal/node-logger';
+
+import { runGhostStories } from '../utils/ghost-stories/run-story-tests.ts';
+import type { StoryIndexGenerator } from 'storybook/internal/core-server';
+import { waitForIdleVitest } from '../utils/wait-for-idle-vitest.ts';
+
+export function initAIAnalyticsChannel(
+  channel: Channel,
+  options: Options,
+  coreOptions: CoreConfig,
+  getStoryIndexGeneratorPromise?: () => Promise<StoryIndexGenerator> | undefined
+) {
+  if (coreOptions.disableTelemetry) {
+    return channel;
+  }
+
+  /** Send analytics about the ai prepare workflow when requested*/
+  channel.on(AI_PREPARE_ANALYTICS_REQUEST, async () => {
+    const stats: {
+      fileCount?: number;
+      storyCount?: number;
+      testRunDuration?: number;
+    } = {};
+
+    try {
+      const lastEvents = await getLastEvents();
+      const lastAIPrepare = lastEvents?.['ai-prepare'];
+      const lastPrepareStoryScoringRun = lastEvents?.['ai-prepare-story-scoring'];
+
+      // Only run if sb ai prepare has been called
+      if (!lastAIPrepare) {
+        return;
+      }
+
+      // Already ran once for this project — never run again
+      if (lastPrepareStoryScoringRun) {
+        return;
+      }
+
+      const metadata = await getStorybookMetadata(options.configDir);
+      const isReactStorybook = metadata?.renderer?.includes('@storybook/react');
+      const hasVitestAddon =
+        !!metadata?.addons &&
+        Object.keys(metadata.addons).some((addonKey) =>
+          addonKey.includes('@storybook/addon-vitest')
+        );
+
+      // For now this is gated by React + Vitest
+      if (!isReactStorybook || !hasVitestAddon) {
+        return;
+      }
+
+      // Wait for any running tests to finish before launching scoring, so we don't
+      // disturb end user activities.
+      const isIdle = await waitForIdleVitest();
+      if (!isIdle) {
+        logger.debug('AI_PREPARE_ANALYTICS_REQUEST timed out waiting for vitest to be available.');
+        return;
+      }
+
+      // Fetch AI-generated stories and score them with the ghost stories metrics, if any are found.
+      const generatorPromise = getStoryIndexGeneratorPromise?.();
+      if (!generatorPromise) {
+        logger.debug(
+          'AI_PREPARE_ANALYTICS_REQUEST could not proceed as the index generator is not ready.'
+        );
+        return;
+      }
+
+      const generator = await generatorPromise;
+      const indexAndStats = await generator.getIndexAndStats();
+      if (!indexAndStats) {
+        logger.debug('AI_PREPARE_ANALYTICS_REQUEST could not proceed as the index is not ready.');
+        return;
+      }
+
+      const aiStoryFiles = new Set<string>();
+      let aiStoryCount = 0;
+      for (const entry of Object.values(indexAndStats.storyIndex.entries)) {
+        if (isStoryCreatedByAIPrepare(entry)) {
+          aiStoryFiles.add(entry.importPath);
+          aiStoryCount++;
+        }
+      }
+
+      if (aiStoryFiles.size > 0) {
+        const aiTestRunResult = await runGhostStories([...aiStoryFiles]);
+        telemetry('ai-prepare-story-scoring', {
+          stats: {
+            fileCount: aiStoryFiles.size,
+            storyCount: aiStoryCount,
+            testRunDuration: aiTestRunResult.duration,
+          },
+          results: aiTestRunResult.summary,
+          ...(aiTestRunResult.runError ? { runError: aiTestRunResult.runError } : {}),
+        });
+      } else {
+        telemetry('ai-prepare-story-scoring', {
+          stats: {
+            fileCount: 0,
+            storyCount: 0,
+            testRunDuration: 0,
+          },
+          runError: 'No stories found that were generated by ai setup',
+        });
+      }
+    } catch {
+      telemetry('ai-prepare-story-scoring', {
+        stats,
+        runError: 'Unknown error during AI story scoring',
+      });
+    } finally {
+      // we don't currently do anything with this, but will be useful in the future
+      channel.emit(AI_PREPARE_ANALYTICS_RESPONSE);
+    }
+  });
+
+  return channel;
+}

code/core/src/core-server/server-channel/ghost-stories-channel.test.ts

@@ -343,7 +343,7 @@ describe('ghostStoriesChannel', () => {
         });
 
         expect(mockTelemetry.getLastEvents).toHaveBeenCalled();
-        expect(mockTelemetry.getSessionId).toHaveBeenCalled();
+        expect(mockTelemetry.getSessionId).not.toHaveBeenCalled();
         expect(mockTelemetry.getStorybookMetadata).not.toHaveBeenCalled();
         expect(mockStoryGeneration.getComponentCandidates).not.toHaveBeenCalled();
       });

code/core/src/core-server/server-channel/ghost-stories-channel.ts

@@ -6,10 +6,12 @@ import {
   getStorybookMetadata,
   telemetry,
 } from 'storybook/internal/telemetry';
+import { logger } from 'storybook/internal/node-logger';
 import type { CoreConfig, Options } from 'storybook/internal/types';
 
 import { getComponentCandidates } from '../utils/ghost-stories/get-candidates.ts';
 import { runGhostStories } from '../utils/ghost-stories/run-story-tests.ts';
+import { waitForIdleVitest } from '../utils/wait-for-idle-vitest.ts';
 
 export function initGhostStoriesChannel(
   channel: Channel,
@@ -36,16 +38,23 @@ export function initGhostStoriesChannel(
       const ghostRunStart = Date.now();
       const lastEvents = await getLastEvents();
       const lastInit = lastEvents?.init;
-      if (!lastEvents || !lastInit) {
+      const lastAIPrepare = lastEvents?.['ai-prepare'];
+      const lastGhostStoriesRun = lastEvents?.['ghost-stories'];
+
+      // We only want to run ghost stories immediately after init or ai prepare.
+      const lastRelevantEvent = lastAIPrepare ?? lastInit;
+      if (!lastRelevantEvent) {
+        return;
+      }
+
+      // Already ran once for this project — never run again
+      if (lastGhostStoriesRun) {
         return;
       }
 
       const sessionId = await getSessionId();
-      const lastGhostStoriesRun = lastEvents['ghost-stories'];
-      if (
-        lastGhostStoriesRun ||
-        (lastInit.body?.sessionId && lastInit.body.sessionId !== sessionId)
-      ) {
+      // We only capture ghost stories in the first ever session since init or ai prepare.
+      if (lastRelevantEvent.body?.sessionId && lastRelevantEvent.body.sessionId !== sessionId) {
         return;
       }
 
@@ -62,6 +71,14 @@ export function initGhostStoriesChannel(
         return;
       }
 
+      // Wait for any running tests to finish before launching scoring, so we don't
+      // disturb end user activities.
+      const isIdle = await waitForIdleVitest();
+      if (!isIdle) {
+        logger.debug('GHOST_STORIES_REQUEST timed out waiting for vitest to be available.');
+        return;
+      }
+
       // Phase 1: find candidates from components
       const candidateAnalysisStart = Date.now();
       const candidatesResult = await getComponentCandidates();

code/core/src/core-server/utils/doTelemetry.ts

@@ -1,4 +1,8 @@
-import { getPrecedingUpgrade, telemetry } from 'storybook/internal/telemetry';
+import {
+  collectAiPrepareEvidence,
+  getPrecedingUpgrade,
+  telemetry,
+} from 'storybook/internal/telemetry';
 import type { CoreConfig, Options } from 'storybook/internal/types';
 
 import type { Polka } from 'polka';
@@ -36,6 +40,14 @@ export async function doTelemetry(
       });
       return;
     }
+
+    // sb ai commands trigger side effects performed by agent harnesses, which can't be observed
+    // directly. This is the entry point for collecting evidence about those side effects and
+    // recording them in telemetry.
+    if (indexAndStats) {
+      collectAiPrepareEvidence('dev', options.configDir, indexAndStats.storyIndex);
+    }
+
     const { versionCheck, versionUpdates } = options;
     invariant(
       !versionUpdates || (versionUpdates && versionCheck),