Skills M4 (5/9): the Vitest addon hosts the test toolset and serves the story-test tool

code/addons/mcp/src/mcp-handler.ts

@@ -152,6 +152,7 @@ export const mcpServerHandler = async ({
     endpoint,
     toolsets: getToolsets(webRequest, addonOptions),
     reviewEnabled: isReviewEnabledForRequest(webRequest, reviewGates!),
+    cliClient: webRequest.headers.get(STORYBOOK_MCP_PROXY_HEADER) === 'true',
     origin: origin!,
     disableTelemetry: disableTelemetry!,
     a11yEnabled,

code/addons/mcp/src/preset.test.ts

@@ -2,7 +2,7 @@ import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
 import type { Options } from 'storybook/internal/types';
 import { experimental_devServer } from './preset.ts';
 import * as mcpHandlerModule from './mcp-handler.ts';
-import * as runStoryTests from './tools/run-story-tests.ts';
+import * as runStoryTests from './utils/addon-vitest.ts';
 import * as moduleGraph from './utils/module-graph.ts';
 import { registerCoreToolsetsForTest } from './test-support/register-core-toolsets.ts';
 
@@ -419,7 +419,7 @@ describe('experimental_devServer', () => {
   });
 
   it('should show Storybook version requirement for addon-vitest and a manual manifest link', async () => {
-    vi.spyOn(runStoryTests, 'getAddonVitestConstants').mockResolvedValue(undefined);
+    vi.spyOn(runStoryTests, 'isAddonVitestEnabled').mockResolvedValue(false);
     const manifestEnabledOptions = {
       presets: {
         apply: vi.fn((key: string) => {

code/addons/mcp/src/storybook-ai-metadata.test.ts

@@ -2,15 +2,15 @@ import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
 import { registerCoreToolsetsForTest } from './test-support/register-core-toolsets.ts';
 import { McpServer } from 'tmcp';
 import { ValibotJsonSchemaAdapter } from '@tmcp/adapter-valibot';
-import { GET_TOOL_NAME } from '@storybook/mcp';
 import { logger } from 'storybook/internal/node-logger';
 import { buildStorybookAiMetadata } from './storybook-ai-metadata.ts';
-import { getAddonVitestConstants } from './tools/run-story-tests.ts';
+import { isAddonVitestEnabled } from './utils/addon-vitest.ts';
 import type { AddonContext } from './types.ts';
 import { getManifestStatus } from './tools/is-manifest-available.ts';
 import { isAddonA11yEnabled } from './utils/is-addon-a11y-enabled.ts';
 import { getReviewStatus } from './utils/is-review-available.ts';
 import { isModuleGraphSupported, isModuleGraphSupportedByBuilder } from './utils/module-graph.ts';
+import { MCP_TOOL_NAMES } from 'storybook/open-service';
 import {
   DISPLAY_REVIEW_TOOL_NAME,
   GET_STORIES_BY_COMPONENT_TOOL_NAME,
@@ -41,10 +41,7 @@ vi.mock('./utils/is-addon-a11y-enabled.ts', () => ({
   isAddonA11yEnabled: vi.fn(),
 }));
 
-vi.mock('./tools/run-story-tests.ts', async (importActual) => ({
-  ...(await importActual<typeof import('./tools/run-story-tests.ts')>()),
-  getAddonVitestConstants: vi.fn(),
-}));
+vi.mock('./utils/addon-vitest.ts', { spy: true });
 
 describe('buildStorybookAiMetadata', () => {
   beforeEach(() => {
@@ -65,10 +62,7 @@ describe('buildStorybookAiMetadata', () => {
       hasFeatureFlag: true,
       docgenServer: false,
     });
-    vi.mocked(getAddonVitestConstants).mockResolvedValue({
-      TRIGGER_TEST_RUN_REQUEST: 'TRIGGER_TEST_RUN_REQUEST',
-      TRIGGER_TEST_RUN_RESPONSE: 'TRIGGER_TEST_RUN_RESPONSE',
-    });
+    vi.mocked(isAddonVitestEnabled).mockResolvedValue(true);
     vi.mocked(isAddonA11yEnabled).mockResolvedValue(true);
   });
 
@@ -158,7 +152,7 @@ describe('buildStorybookAiMetadata', () => {
   });
 
   it('keeps addon-vitest availability aligned between metadata and live tools/list', async () => {
-    vi.mocked(getAddonVitestConstants).mockResolvedValue(undefined);
+    vi.mocked(isAddonVitestEnabled).mockResolvedValue(false);
     const options = createOptions();
 
     const metadata = await buildStorybookAiMetadata(options);
@@ -493,6 +487,50 @@ function createOptions({
   } as any;
 }
 
+/**
+ * The published tool surface for each availability combination.
+ *
+ * Availability decides which tools an agent is offered at all, so it is part of the wire contract
+ * and not just internal wiring. These pin the surface per combination; the descriptions and schemas
+ * behind each name are snapshotted by the e2e suite against a live server.
+ */
+const LIST_TOOL_NAME = MCP_TOOL_NAMES['docs.list'];
+const GET_TOOL_NAME = MCP_TOOL_NAMES['docs.show'];
+
+describe('tool availability variants', () => {
+  const names = async (overrides: Partial<ToolAvailability>) =>
+    (await listRegisteredTools(createOptions(), { availability: createAvailability(overrides) }))
+      .map((tool: { name: string }) => tool.name)
+      .sort();
+
+  it('offers no review tool when review is off', async () => {
+    const withReview = await names({});
+    const withoutReview = await names({ reviewEnabled: false, changeDetectionEnabled: false });
+
+    expect(withReview).toContain(DISPLAY_REVIEW_TOOL_NAME);
+    expect(withoutReview).not.toContain(DISPLAY_REVIEW_TOOL_NAME);
+  });
+
+  it('offers no test tool when addon-vitest is absent', async () => {
+    expect(await names({ testSupported: false })).not.toContain(RUN_STORY_TESTS_TOOL_NAME);
+  });
+
+  it('keeps the docs tools when a11y is off, which only varies their prose', async () => {
+    expect(await names({ a11yEnabled: false })).toEqual(await names({ a11yEnabled: true }));
+  });
+
+  it('offers no docs tools when docs are unavailable', async () => {
+    const withoutDocs = await names({
+      docsEnabled: false,
+      docsHasManifests: false,
+      docsFeatureEnabled: false,
+    });
+
+    expect(withoutDocs).not.toContain(LIST_TOOL_NAME);
+    expect(withoutDocs).not.toContain(GET_TOOL_NAME);
+  });
+});
+
 function createAvailability(overrides: Partial<ToolAvailability> = {}): ToolAvailability {
   return {
     moduleGraphSupported: true,

code/addons/mcp/src/telemetry.test.ts

@@ -47,9 +47,25 @@ describe('collectTelemetry', () => {
         roots: { listChanged: true },
       },
       customField: 'custom-value',
+      consumer: 'mcp',
     });
   });
 
+  it('reports the storybook ai CLI channel as the cli consumer', async () => {
+    vi.mocked(telemetry).mockResolvedValue(undefined);
+
+    const cliServer = Object.assign(Object.create(mockServer), {
+      ctx: { ...mockServer.ctx, custom: { cliClient: true } },
+    }) as any;
+
+    await collectTelemetry({ event: 'test-event', server: cliServer });
+
+    expect(telemetry).toHaveBeenCalledWith(
+      'addon-mcp',
+      expect.objectContaining({ consumer: 'cli' })
+    );
+  });
+
   it('should pass through additional payload fields', async () => {
     vi.mocked(telemetry).mockResolvedValue(undefined);
 
@@ -69,6 +85,7 @@ describe('collectTelemetry', () => {
       toolName: 'list-all-documentation',
       duration: 123,
       success: true,
+      consumer: 'mcp',
     });
   });
 
@@ -103,6 +120,7 @@ describe('collectTelemetry', () => {
       mcpSessionId: undefined,
       clientInfo: undefined,
       clientCapabilities: undefined,
+      consumer: 'mcp',
     });
   });
 });

code/addons/mcp/src/telemetry.ts

@@ -18,6 +18,9 @@ export async function collectTelemetry({
       mcpSessionId: server.ctx.sessionId,
       clientInfo: server.ctx.sessionInfo?.clientInfo,
       clientCapabilities: server.ctx.sessionInfo?.clientCapabilities,
+      // Which channel drove the call: the `storybook ai` CLI marks its requests with the proxy
+      // header; everything else is a direct MCP client.
+      consumer: server.ctx.custom?.cliClient ? 'cli' : 'mcp',
       ...payload,
     });
   } catch (error) {

code/addons/mcp/src/test-support/register-core-toolsets.ts

@@ -11,6 +11,7 @@ import type { StoryIndex } from 'storybook/internal/types';
 import { clearToolsetRegistry } from 'storybook/open-service';
 import {
   createStoriesToolset,
+  createTestToolset,
   registerToolset,
   reviewToolset,
 } from 'storybook/internal/core-server';
@@ -20,7 +21,8 @@ const EMPTY_INDEX: StoryIndex = { v: 5, entries: {} };
 export function registerCoreToolsetsForTest({
   index = EMPTY_INDEX,
   reviewEnabled = true,
-}: { index?: StoryIndex; reviewEnabled?: boolean } = {}) {
+  testToolset = true,
+}: { index?: StoryIndex; reviewEnabled?: boolean; testToolset?: boolean } = {}) {
   clearToolsetRegistry();
 
   const storyIndex = { getIndex: async () => index };
@@ -37,4 +39,15 @@ export function registerCoreToolsetsForTest({
     })
   );
   registerToolset(reviewToolset);
+  // `testToolset: false` mirrors addon-vitest being absent or not enabled, where its `services`
+  // hook never runs and the `test` toolset stays unregistered.
+  if (testToolset) {
+    registerToolset(
+      createTestToolset({
+        channel: { on: () => {}, off: () => {}, emit: () => {} } as never,
+        storyIndex,
+        a11yEnabled: false,
+      })
+    );
+  }
 }

code/addons/mcp/src/tools/get-storybook-story-instructions.test.ts

@@ -1,7 +1,7 @@
 import { describe, it, expect, vi, beforeEach } from 'vitest';
 import { McpServer } from 'tmcp';
 import { ValibotJsonSchemaAdapter } from '@tmcp/adapter-valibot';
-import { getAddonVitestConstants } from './run-story-tests.ts';
+import { isAddonVitestEnabled } from '../utils/addon-vitest.ts';
 import {
   addGetUIBuildingInstructionsTool,
   buildStorybookStoryInstructions,
@@ -14,8 +14,8 @@ import {
   GET_UI_BUILDING_INSTRUCTIONS_TOOL_NAME,
 } from './tool-names.ts';
 
-vi.mock('./run-story-tests.ts', () => ({
-  getAddonVitestConstants: vi.fn(),
+vi.mock('../utils/addon-vitest.ts', () => ({
+  isAddonVitestEnabled: vi.fn(),
 }));
 
 vi.mock('../utils/is-review-available.ts', () => ({
@@ -26,10 +26,7 @@ describe('getUIBuildingInstructionsTool', () => {
   let server: McpServer<any, AddonContext>;
 
   beforeEach(async () => {
-    vi.mocked(getAddonVitestConstants).mockResolvedValue({
-      TRIGGER_TEST_RUN_REQUEST: 'TRIGGER_TEST_RUN_REQUEST',
-      TRIGGER_TEST_RUN_RESPONSE: 'TRIGGER_TEST_RUN_RESPONSE',
-    });
+    vi.mocked(isAddonVitestEnabled).mockResolvedValue(true);
 
     vi.mocked(getReviewStatus).mockResolvedValue({
       available: false,
@@ -69,7 +66,7 @@ describe('getUIBuildingInstructionsTool', () => {
       }
     );
 
-    await addGetUIBuildingInstructionsTool(server);
+    await addGetUIBuildingInstructionsTool(server, undefined, { addonVitestAvailable: true });
   });
 
   async function getToolDescription(context: AddonContext) {

code/addons/mcp/src/tools/get-storybook-story-instructions.ts

@@ -1,6 +1,6 @@
 import type { McpServer } from 'tmcp';
 import type { Options } from 'storybook/internal/types';
-import { getAddonVitestConstants } from './run-story-tests.ts';
+import { isAddonVitestEnabled } from '../utils/addon-vitest.ts';
 import { collectTelemetry } from '../telemetry.ts';
 import { getFinalLinksGuidance } from '../instructions/build-server-instructions.ts';
 import { getReviewStatus } from '../utils/is-review-available.ts';
@@ -47,10 +47,11 @@ export async function addGetUIBuildingInstructionsTool(
   server: McpServer<any, AddonContext>,
   enabled: Parameters<McpServer<any, AddonContext>['tool']>[0]['enabled'] = () =>
     server.ctx.custom?.toolsets?.dev ?? true,
-  { docsAvailable = false }: { docsAvailable?: boolean } = {}
+  {
+    docsAvailable = false,
+    addonVitestAvailable = false,
+  }: { docsAvailable?: boolean; addonVitestAvailable?: boolean } = {}
 ) {
-  const addonVitestAvailable = !!(await getAddonVitestConstants());
-
   server.tool(
     {
       name: GET_UI_BUILDING_INSTRUCTIONS_TOOL_NAME,
@@ -198,7 +199,8 @@ export async function buildStorybookStoryInstructions(
     .replace('{{FINAL_LINKS_GUIDANCE}}', getFinalLinksGuidance(reviewEnabled))
     .replace('{{CHANGED_STORY_FALLBACK_LINK_GUIDANCE}}', changedStoryFallbackLinkGuidance);
 
-  const resolvedAddonVitestAvailable = addonVitestAvailable ?? !!(await getAddonVitestConstants());
+  const resolvedAddonVitestAvailable =
+    addonVitestAvailable ?? (await isAddonVitestEnabled(options));
   const testToolsetAvailable = (toolsets?.test ?? true) && resolvedAddonVitestAvailable;
 
   if (testToolsetAvailable) {