Merge pull request #35235 from storybookjs/kasper/ai-agent-instance-selection

CLI: Prefer agent-matched Storybook instances

Author: kasperpeulen

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

@@ -131,6 +131,13 @@ describe('readRegistry', () => {
     await expect(readRegistry(REGISTRY_DIR)).resolves.toEqual([minimal]);
   });
 
+  it('accepts records with optional agent provenance', async () => {
+    const agentRecord = { ...aliveRecord, agent: 'claude-preview' };
+    vol.fromNestedJSON({ [REGISTRY_DIR]: { 'agent.json': JSON.stringify(agentRecord) } });
+
+    await expect(readRegistry(REGISTRY_DIR)).resolves.toEqual([agentRecord]);
+  });
+
   it('rejects out-of-range ports', async () => {
     vol.fromNestedJSON({
       [REGISTRY_DIR]: {

code/core/src/cli/ai/mcp/resolve-instance.test.ts

@@ -186,16 +186,201 @@ describe('resolveInstance', () => {
   });
 
   it('selects the instance matching BOTH cwd and port when a port is supplied', () => {
-    const a = record('/Users/x/projects/foo', 'ready', { pid: 100, port: 6006 });
-    const b = record('/Users/x/projects/foo', 'ready', { pid: 200, port: 6007 });
-    const result = resolveInstance([a, b], '/Users/x/projects/foo', 6007);
+    const a = record('/Users/x/projects/foo', 'ready', {
+      agent: 'claude-preview',
+      pid: 100,
+      port: 6006,
+    });
+    const b = record('/Users/x/projects/foo', 'ready', { agent: 'codex', pid: 200, port: 6007 });
+    const result = resolveInstance([a, b], '/Users/x/projects/foo', 6007, 'claude');
     expect(result.kind).toBe('instance');
     if (result.kind === 'instance') {
       expect(result.record).toBe(b);
       expect(result.matches).toEqual([b]);
     }
   });
 
+  it('prefers Claude preview records for Claude CLI invocations', () => {
+    const genericClaude = record('/Users/x/projects/foo', 'ready', {
+      agent: 'claude',
+      startedAt: '2026-06-09T11:00:00.000Z',
+    });
+    const claudePreview = record('/Users/x/projects/foo', 'ready', {
+      agent: 'claude-preview',
+      startedAt: '2026-06-09T10:00:00.000Z',
+    });
+    const newerCodex = record('/Users/x/projects/foo', 'ready', {
+      agent: 'codex',
+      startedAt: '2026-06-09T12:00:00.000Z',
+    });
+    const result = resolveInstance(
+      [genericClaude, claudePreview, newerCodex],
+      '/Users/x/projects/foo',
+      undefined,
+      'claude'
+    );
+
+    expect(result.kind).toBe('instance');
+    if (result.kind === 'instance') {
+      expect(result.record).toBe(claudePreview);
+      expect(result.matches).toEqual([claudePreview]);
+    }
+  });
+
+  it('falls back to generic Claude records when no Claude preview record matches', () => {
+    const genericClaude = record('/Users/x/projects/foo', 'ready', {
+      agent: 'claude',
+      startedAt: '2026-06-09T10:00:00.000Z',
+    });
+    const newerCodex = record('/Users/x/projects/foo', 'ready', {
+      agent: 'codex',
+      startedAt: '2026-06-09T11:00:00.000Z',
+    });
+    const result = resolveInstance(
+      [genericClaude, newerCodex],
+      '/Users/x/projects/foo',
+      undefined,
+      'claude'
+    );
+
+    expect(result.kind).toBe('instance');
+    if (result.kind === 'instance') {
+      expect(result.record).toBe(genericClaude);
+      expect(result.matches).toEqual([genericClaude]);
+    }
+  });
+
+  it('stays in the preferred agent bucket even when another bucket is ready', () => {
+    const startingPreview = record('/Users/x/projects/foo', 'starting', {
+      agent: 'claude-preview',
+      startedAt: '2026-06-09T11:00:00.000Z',
+    });
+    const readyCodex = record('/Users/x/projects/foo', 'ready', {
+      agent: 'codex',
+      startedAt: '2026-06-09T10:00:00.000Z',
+    });
+    const result = resolveInstance(
+      [startingPreview, readyCodex],
+      '/Users/x/projects/foo',
+      undefined,
+      'claude'
+    );
+
+    expect(result.kind).toBe('intercept');
+    if (result.kind === 'intercept') {
+      expect(result.reason).toBe('mcp-starting');
+      expect(result.matches).toEqual([startingPreview]);
+    }
+  });
+
+  it('reports errors from the preferred agent bucket instead of switching buckets', () => {
+    const erroredPreview = record('/Users/x/projects/foo', 'error', {
+      agent: 'claude-preview',
+      startedAt: '2026-06-09T11:00:00.000Z',
+    });
+    const readyClaude = record('/Users/x/projects/foo', 'ready', {
+      agent: 'claude',
+      startedAt: '2026-06-09T10:00:00.000Z',
+    });
+    const result = resolveInstance(
+      [erroredPreview, readyClaude],
+      '/Users/x/projects/foo',
+      undefined,
+      'claude'
+    );
+
+    expect(result.kind).toBe('intercept');
+    if (result.kind === 'intercept') {
+      expect(result.reason).toBe('mcp-error');
+      expect(result.matches).toEqual([erroredPreview]);
+    }
+  });
+
+  it('prefers records matching the current non-Claude agent', () => {
+    const codex = record('/Users/x/projects/foo', 'ready', {
+      agent: 'codex',
+      startedAt: '2026-06-09T10:00:00.000Z',
+    });
+    const newerCursor = record('/Users/x/projects/foo', 'ready', {
+      agent: 'cursor',
+      startedAt: '2026-06-09T11:00:00.000Z',
+    });
+    const result = resolveInstance(
+      [codex, newerCursor],
+      '/Users/x/projects/foo',
+      undefined,
+      'codex'
+    );
+
+    expect(result.kind).toBe('instance');
+    if (result.kind === 'instance') {
+      expect(result.record).toBe(codex);
+      expect(result.matches).toEqual([codex]);
+    }
+  });
+
+  it('chooses the latest-started ready instance inside the selected agent bucket', () => {
+    const olderPreview = record('/Users/x/projects/foo', 'ready', {
+      agent: 'claude-preview',
+      startedAt: '2026-06-09T10:00:00.000Z',
+    });
+    const newerPreview = record('/Users/x/projects/foo', 'ready', {
+      agent: 'claude-preview',
+      startedAt: '2026-06-09T11:00:00.000Z',
+    });
+    const newestCodex = record('/Users/x/projects/foo', 'ready', {
+      agent: 'codex',
+      startedAt: '2026-06-09T12:00:00.000Z',
+    });
+    const result = resolveInstance(
+      [olderPreview, newerPreview, newestCodex],
+      '/Users/x/projects/foo',
+      undefined,
+      'claude'
+    );
+
+    expect(result.kind).toBe('instance');
+    if (result.kind === 'instance') {
+      expect(result.record).toBe(newerPreview);
+      expect(result.matches).toEqual([newerPreview, olderPreview]);
+    }
+  });
+
+  it('falls back to latest-started behavior when no record matches the current agent', () => {
+    const older = record('/Users/x/projects/foo', 'ready', {
+      startedAt: '2026-06-09T10:00:00.000Z',
+    });
+    const newer = record('/Users/x/projects/foo', 'ready', {
+      agent: 'cursor',
+      startedAt: '2026-06-09T11:00:00.000Z',
+    });
+    const result = resolveInstance([older, newer], '/Users/x/projects/foo', undefined, 'codex');
+
+    expect(result.kind).toBe('instance');
+    if (result.kind === 'instance') {
+      expect(result.record).toBe(newer);
+      expect(result.matches).toEqual([newer, older]);
+    }
+  });
+
+  it('falls back to latest-started behavior when no current agent is detected', () => {
+    const olderPreview = record('/Users/x/projects/foo', 'ready', {
+      agent: 'claude-preview',
+      startedAt: '2026-06-09T10:00:00.000Z',
+    });
+    const newerCodex = record('/Users/x/projects/foo', 'ready', {
+      agent: 'codex',
+      startedAt: '2026-06-09T11:00:00.000Z',
+    });
+    const result = resolveInstance([olderPreview, newerCodex], '/Users/x/projects/foo');
+
+    expect(result.kind).toBe('instance');
+    if (result.kind === 'instance') {
+      expect(result.record).toBe(newerCodex);
+      expect(result.matches).toEqual([newerCodex, olderPreview]);
+    }
+  });
+
   it('ignores port when it is not supplied (routes by cwd alone)', () => {
     const a = record('/Users/x/projects/foo', 'ready', { pid: 100, port: 6006 });
     const b = record('/Users/x/projects/foo', 'ready', { pid: 200, port: 6007 });

code/core/src/cli/ai/mcp/resolve-instance.ts

@@ -1,5 +1,9 @@
 import { resolve } from 'node:path';
 
+import {
+  CLAUDE_AGENT_NAME,
+  CLAUDE_PREVIEW_AGENT_NAME,
+} from '../../../shared/constants/agent-provenance.ts';
 import type { InterceptReason, StorybookInstanceRecord } from './types.ts';
 
 export type ResolveResult =
@@ -34,15 +38,16 @@ export type ResolveResult =
  * - `error` → mcp-error intercept
  *
  * Zero matches → no-instance intercept (callers may surface running cwds). 2+ matches at the same
- * cwd → pick the most recently started instance (latest `startedAt` among `ready` records, else
- * latest overall), on the assumption that the freshest instance is the one the agent just started.
- * Records without a `startedAt` tie-break on lowest pid for determinism. All matches are returned
- * (most-recent first) as `matches` so callers can warn the agent without blocking the call.
+ * cwd → use the current agent to select the competing bucket, then pick the most recently started
+ * instance in that bucket (latest `startedAt` among `ready` records, else latest overall). Records
+ * without a `startedAt` tie-break on lowest pid for determinism. The selected bucket is returned
+ * (most-recent first) as `matches` so callers can warn only about instances that competed.
  */
 export function resolveInstance(
   records: StorybookInstanceRecord[],
   targetCwd: string,
-  targetPort?: number
+  targetPort?: number,
+  currentAgent?: string
 ): ResolveResult {
   const normalisedTarget = resolve(targetCwd);
   const cwdMatches = records.filter((r) => resolve(r.cwd) === normalisedTarget);
@@ -67,7 +72,7 @@ export function resolveInstance(
     };
   }
 
-  const sortedMatches = [...matches].sort(byMostRecentlyStarted);
+  const sortedMatches = selectCompetingBucket(matches, targetPort, currentAgent);
   const selected = sortedMatches.find((r) => r.mcp.status === 'ready') ?? sortedMatches[0];
 
   switch (selected.mcp.status) {
@@ -106,6 +111,28 @@ export function resolveInstance(
   }
 }
 
+function selectCompetingBucket(
+  matches: StorybookInstanceRecord[],
+  targetPort: number | undefined,
+  currentAgent: string | undefined
+) {
+  if (targetPort != null) {
+    return [...matches].sort(byMostRecentlyStarted);
+  }
+
+  // std-env reports Claude CLI as `claude`; preview-launched Storybooks record `claude-preview`.
+  const agentBuckets =
+    currentAgent === CLAUDE_AGENT_NAME
+      ? [CLAUDE_PREVIEW_AGENT_NAME, CLAUDE_AGENT_NAME]
+      : currentAgent
+        ? [currentAgent]
+        : [];
+  const selectedAgent = agentBuckets.find((agent) => matches.some((r) => r.agent === agent));
+  const bucket = selectedAgent ? matches.filter((r) => r.agent === selectedAgent) : matches;
+
+  return [...bucket].sort(byMostRecentlyStarted);
+}
+
 /**
  * `startedAt` as epoch millis, or `-Infinity` when absent/unparseable so such records sort as the
  * oldest (and fall through to the pid tie-break).

code/core/src/cli/ai/mcp/run-tool.test.ts

@@ -1,4 +1,4 @@
-import { beforeEach, describe, expect, it, vi } from 'vitest';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
 
 import { McpJsonRpcError, callMcpTool, listMcpTools } from './client.ts';
 import { loadStorybookAiMetadata, type StorybookAiMetadata } from './local-metadata.ts';
@@ -39,6 +39,10 @@ beforeEach(() => {
   vi.mocked(loadStorybookAiMetadata).mockReset().mockResolvedValue(defaultRuntimeMetadata);
 });
 
+afterEach(() => {
+  vi.unstubAllEnvs();
+});
+
 describe('runAiTool', () => {
   it('forwards the call to the matching instance and prints the markdown result', async () => {
     const result = await runAiTool('list-all-documentation', ['--withStoryIds', 'true'], {
@@ -459,12 +463,65 @@ describe('runAiTool', () => {
     vi.mocked(readRegistry).mockResolvedValue([record, sibling]);
     const result = await runAiTool('list-all-documentation', [], { cwd: '/projects/foo' });
     expect(result.exitCode).toBe(0);
-    expect(result.output).toContain('Multiple Storybook instances');
+    expect(result.output).toContain('Multiple matching Storybook instances');
+    expect(result.output).toContain('Matching instances at `/projects/foo`');
     expect(result.output).toContain('pid `1`');
     expect(result.output).toContain('pid `2`');
     expect(result.output).toContain('(used)');
     expect(result.output).toContain('upstream result');
   });
+
+  describe('when multiple instances exist in the selected agent bucket', () => {
+    const olderPreview = {
+      ...record,
+      agent: 'claude-preview',
+      instanceId: 'inst-2',
+      pid: 2,
+      port: 6007,
+      startedAt: '2026-06-09T10:00:00.000Z',
+      url: 'http://localhost:6007',
+    };
+    const selectedPreview = {
+      ...record,
+      agent: 'claude-preview',
+      instanceId: 'inst-3',
+      pid: 3,
+      port: 6008,
+      startedAt: '2026-06-09T11:00:00.000Z',
+      url: 'http://localhost:6008',
+    };
+    const newerCodex = {
+      ...record,
+      agent: 'codex',
+      instanceId: 'inst-4',
+      pid: 4,
+      port: 6009,
+      startedAt: '2026-06-09T12:00:00.000Z',
+      url: 'http://localhost:6009',
+    };
+
+    beforeEach(() => {
+      vi.stubEnv('AI_AGENT', 'claude');
+      vi.mocked(readRegistry).mockResolvedValue([olderPreview, selectedPreview, newerCodex]);
+    });
+
+    it('only warns about instances in the selected agent bucket', async () => {
+      const result = await runAiTool('list-all-documentation', [], { cwd: '/projects/foo' });
+
+      expect(callMcpTool).toHaveBeenCalledWith(
+        selectedPreview,
+        { name: 'list-all-documentation', arguments: {} },
+        undefined
+      );
+      expect(result.exitCode).toBe(0);
+      expect(result.output).toContain('Multiple matching Storybook instances');
+      expect(result.output).toContain('Matching instances at `/projects/foo`');
+      expect(result.output).toContain('pid `2`');
+      expect(result.output).toContain('pid `3`');
+      expect(result.output).not.toContain('pid `4`');
+      expect(result.output).not.toContain('http://localhost:6009');
+    });
+  });
 });
 
 describe('buildStorybookCommandsHelp', () => {