Merge pull request #35212 from storybookjs/kasper/ai-serverless-help

CLI: Load Storybook AI help from preset metadata

Author: kasperpeulen

code/core/src/cli/ai/mcp/client.test.ts

@@ -2,13 +2,7 @@ import { versions } from 'storybook/internal/common';
 
 import { describe, expect, it, vi } from 'vitest';
 
-import {
-  MCP_CLIENT_INFO,
-  McpJsonRpcError,
-  callMcpTool,
-  listMcpTools,
-  listMcpToolsWithServerMetadata,
-} from './client.ts';
+import { MCP_CLIENT_INFO, McpJsonRpcError, callMcpTool, listMcpTools } from './client.ts';
 import type { StorybookInstanceRecord } from './types.ts';
 
 const record: StorybookInstanceRecord = {
@@ -228,38 +222,20 @@ describe('listMcpTools', () => {
 });
 
 describe('initialize handshake (clientInfo for telemetry segmentation)', () => {
-  const initializeResponse = (sessionId?: string, instructions?: unknown) =>
+  const initializeResponse = (sessionId?: string) =>
     jsonResponse(
       {
         jsonrpc: '2.0',
         id: 'init',
         result: {
           protocolVersion: '2025-06-18',
           serverInfo: {},
-          ...(instructions !== undefined ? { instructions } : {}),
         },
       },
       200,
       sessionId ? { 'mcp-session-id': sessionId } : {}
     );
 
-  const initializeSseResponse = (sessionId: string | undefined, instructions: unknown) => {
-    const envelope = {
-      jsonrpc: '2.0',
-      id: 'init',
-      result: {
-        protocolVersion: '2025-06-18',
-        serverInfo: {},
-        instructions,
-      },
-    };
-    return sseResponse(
-      `event: message\ndata: ${JSON.stringify(envelope)}\n\n`,
-      200,
-      sessionId ? { 'mcp-session-id': sessionId } : {}
-    );
-  };
-
   const toolResult = () =>
     jsonResponse({
       jsonrpc: '2.0',
@@ -406,51 +382,6 @@ describe('initialize handshake (clientInfo for telemetry segmentation)', () => {
     expect(headers['Mcp-Session-Id']).toBe('session-7');
   });
 
-  it('returns server instructions from the initialize response alongside tools/list results', async () => {
-    const tools = [{ name: 'get-documentation', description: 'Get docs' }];
-    const fetchImpl = vi
-      .fn()
-      .mockResolvedValueOnce(initializeResponse('session-1', '  Follow the story workflow.  '))
-      .mockResolvedValueOnce(
-        jsonResponse({ jsonrpc: '2.0', id: 'x', result: { tools } })
-      ) as unknown as typeof fetch;
-
-    await expect(listMcpToolsWithServerMetadata(record, fetchImpl)).resolves.toEqual({
-      tools,
-      serverMetadata: { instructions: 'Follow the story workflow.' },
-    });
-  });
-
-  it('returns server instructions from an SSE initialize response', async () => {
-    const tools = [{ name: 'get-documentation', description: 'Get docs' }];
-    const fetchImpl = vi
-      .fn()
-      .mockResolvedValueOnce(initializeSseResponse('session-1', '  Follow the story workflow.  '))
-      .mockResolvedValueOnce(toolListResult(tools)) as unknown as typeof fetch;
-
-    await expect(listMcpToolsWithServerMetadata(record, fetchImpl)).resolves.toEqual({
-      tools,
-      serverMetadata: { instructions: 'Follow the story workflow.' },
-    });
-  });
-
-  it.each([undefined, '', '   ', 123])(
-    'omits server instructions when initialize returns %j',
-    async (instructions) => {
-      const fetchImpl = vi
-        .fn()
-        .mockResolvedValueOnce(initializeResponse('session-1', instructions))
-        .mockResolvedValueOnce(
-          jsonResponse({ jsonrpc: '2.0', id: 'x', result: { tools: [] } })
-        ) as unknown as typeof fetch;
-
-      await expect(listMcpToolsWithServerMetadata(record, fetchImpl)).resolves.toEqual({
-        tools: [],
-        serverMetadata: {},
-      });
-    }
-  );
-
   it.each([
     [
       'malformed JSON',
@@ -472,17 +403,14 @@ describe('initialize handshake (clientInfo for telemetry segmentation)', () => {
       .mockResolvedValueOnce(initResponse())
       .mockResolvedValueOnce(toolListResult(tools)) as unknown as typeof fetch;
 
-    await expect(listMcpToolsWithServerMetadata(record, fetchImpl)).resolves.toEqual({
-      tools,
-      serverMetadata: {},
-    });
+    await expect(listMcpTools(record, fetchImpl)).resolves.toEqual(tools);
   });
 
   it('keeps listMcpTools returning only the tool descriptors', async () => {
     const tools = [{ name: 'get-documentation', description: 'Get docs' }];
     const fetchImpl = vi
       .fn()
-      .mockResolvedValueOnce(initializeResponse('session-1', 'Follow the story workflow.'))
+      .mockResolvedValueOnce(initializeResponse('session-1'))
       .mockResolvedValueOnce(
         jsonResponse({ jsonrpc: '2.0', id: 'x', result: { tools } })
       ) as unknown as typeof fetch;

code/core/src/cli/ai/mcp/client.ts

@@ -54,23 +54,10 @@ const JsonRpcEnvelopeSchema = v.looseObject({
   error: v.optional(v.looseObject({ code: v.number(), message: v.string() })),
 });
 
-const InitializeResultSchema = v.looseObject({
-  instructions: v.optional(v.string()),
-});
-
 const ToolListResultSchema = v.looseObject({
   tools: v.optional(v.array(McpToolDescriptorSchema)),
 });
 
-export type McpServerMetadata = {
-  instructions?: string;
-};
-
-export type McpToolList = {
-  tools: McpToolDescriptor[];
-  serverMetadata: McpServerMetadata;
-};
-
 /** Forward an MCP `tools/call` JSON-RPC request to a local Storybook MCP server. */
 export async function callMcpTool(
   record: StorybookInstanceRecord,
@@ -92,23 +79,14 @@ export async function listMcpTools(
   record: StorybookInstanceRecord,
   fetchImpl: typeof fetch = fetch
 ): Promise<McpToolDescriptor[]> {
-  const { tools } = await listMcpToolsWithServerMetadata(record, fetchImpl);
-  return tools;
-}
-
-/** List tools and include server metadata from the preceding MCP `initialize` response. */
-export async function listMcpToolsWithServerMetadata(
-  record: StorybookInstanceRecord,
-  fetchImpl: typeof fetch = fetch
-): Promise<McpToolList> {
-  const { result, serverMetadata } = await sendJsonRpcRequest(
+  const { result } = await sendJsonRpcRequest(
     record,
     'tools/list',
     {},
     ToolListResultSchema,
     fetchImpl
   );
-  return { tools: result.tools ?? [], serverMetadata };
+  return result.tools ?? [];
 }
 
 const REQUEST_HEADERS = {
@@ -117,29 +95,21 @@ const REQUEST_HEADERS = {
   [STORYBOOK_MCP_PROXY_HEADER]: STORYBOOK_MCP_PROXY_HEADER_VALUE,
 };
 
-type InitializeMcpSessionResult = {
-  sessionId: string | null;
-  serverMetadata: McpServerMetadata;
-};
-
 /**
  * Send a minimal MCP `initialize` request carrying {@link MCP_CLIENT_INFO} and return the session
- * id plus any best-effort server metadata from the initialize response.
+ * id when the server assigns one.
  *
  * The session id is MCP Streamable HTTP spec behavior, not a tmcp implementation detail: the
  * server assigns it during initialization, returns it in the `Mcp-Session-Id` response header,
  * and associates the session's clientInfo with later requests echoing that header. The same
- * initialize result can also carry server instructions, so metadata is parsed from the response
- * body even when no session id is assigned.
+ * initialize response body is still consumed so the follow-up request cannot race ahead of the
+ * server storing clientInfo for telemetry segmentation.
  *
  * The handshake is best-effort — when it fails (or a future server ignores sessions), the actual
- * request proceeds without a session and keeps working; only the telemetry segmentation and
- * initialize metadata are lost, and error reporting stays anchored on the real call. It shares the
- * full {@link REQUEST_TIMEOUT_MS} budget rather than a tighter one: `storybook ai --help` renders the
- * server instructions carried here, and the only thing that slows the handshake is the dev server
- * still starting up — which the command must wait through anyway. A tighter budget would just drop
- * the instructions on that first slow request while the command list (sent right after) comes back
- * fine.
+ * request proceeds without a session and keeps working; only telemetry segmentation is lost, and
+ * error reporting stays anchored on the real call. It shares the full {@link REQUEST_TIMEOUT_MS}
+ * budget rather than a tighter one because the only thing that slows the handshake is the dev
+ * server still starting up, which the command must wait through anyway.
  *
  * Sessions are deliberately one-shot: each JSON-RPC request gets its own handshake and the session
  * is never reused or closed. A CLI invocation makes one request on the happy path (two on error
@@ -153,7 +123,7 @@ type InitializeMcpSessionResult = {
 async function initializeMcpSession(
   target: string,
   fetchImpl: typeof fetch
-): Promise<InitializeMcpSessionResult> {
+): Promise<string | null> {
   try {
     const response = await fetchImpl(target, {
       method: 'POST',
@@ -173,22 +143,18 @@ async function initializeMcpSession(
     const sessionId = response.headers.get('mcp-session-id');
     if (!response.ok) {
       await response.body?.cancel();
-      return { sessionId: null, serverMetadata: {} };
+      return null;
     }
 
-    let serverMetadata: McpServerMetadata = {};
     try {
-      serverMetadata = parseInitializeServerMetadata(
-        await readJsonRpcResponse(response, target),
-        target
-      );
+      await response.text();
     } catch {
-      // The initialize request is best-effort metadata; a malformed response must not block the
+      // The initialize request is best-effort telemetry setup; a malformed body must not block the
       // actual tools/list or tools/call request from preserving its existing behavior.
     }
-    return { sessionId, serverMetadata };
+    return sessionId;
   } catch {
-    return { sessionId: null, serverMetadata: {} };
+    return null;
   }
 }
 
@@ -212,15 +178,15 @@ async function sendJsonRpcRequest<TResult>(
   params: unknown,
   resultSchema: v.GenericSchema<unknown, TResult>,
   fetchImpl: typeof fetch
-): Promise<{ result: TResult; serverMetadata: McpServerMetadata }> {
+): Promise<{ result: TResult }> {
   const endpoint = record.mcp.endpoint;
   if (!endpoint) {
     throw new Error(`The Storybook instance at ${record.cwd} has no server endpoint registered`);
   }
 
   const target = new URL(endpoint, record.url).href;
 
-  const { sessionId, serverMetadata } = await initializeMcpSession(target, fetchImpl);
+  const sessionId = await initializeMcpSession(target, fetchImpl);
 
   const response = await fetchImpl(target, {
     method: 'POST',
@@ -254,7 +220,7 @@ async function sendJsonRpcRequest<TResult>(
   if (!result.success) {
     throw unexpectedShapeError(target);
   }
-  return { result: result.output, serverMetadata };
+  return { result: result.output };
 }
 
 function unexpectedShapeError(target: string): Error {
@@ -263,8 +229,7 @@ function unexpectedShapeError(target: string): Error {
 
 /**
  * Unwrap a parsed JSON-RPC payload into its `result`, or report why it isn't usable. The command
- * path throws on the reported error; the best-effort initialize-metadata parse falls back to empty
- * — sharing this keeps the envelope handling in one place.
+ * path throws on the reported error.
  */
 function unwrapJsonRpcResult(
   payload: unknown,
@@ -286,21 +251,6 @@ function unwrapJsonRpcResult(
   return { ok: true, result: envelope.output.result };
 }
 
-function parseInitializeServerMetadata(payload: unknown, target: string): McpServerMetadata {
-  const unwrapped = unwrapJsonRpcResult(payload, target);
-  if (!unwrapped.ok) {
-    return {};
-  }
-
-  const result = v.safeParse(InitializeResultSchema, unwrapped.result);
-  if (!result.success) {
-    return {};
-  }
-
-  const instructions = result.output.instructions?.trim();
-  return instructions ? { instructions } : {};
-}
-
 async function readJsonRpcResponse(response: Response, endpoint: string): Promise<unknown> {
   const contentType = (response.headers.get('content-type') ?? '').toLowerCase();
   const body = await response.text();