Merge pull request #34706 from storybookjs/yann/ai-feature-fixes

UI: Fix showing and hiding copy prompt in the correct scenarios

Author: yannbf

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

@@ -273,7 +273,7 @@ export const experimental_serverChannel = async (
   options: OptionsWithRequiredCache
 ) => {
   initAIAnalyticsChannel(channel, options, () => storyIndexGeneratorPromise);
-  initializeChecklist(channel);
+  initializeChecklist(channel, () => storyIndexGeneratorPromise, options.configDir);
   initializeWhatsNew(channel, options);
   initializeSaveStory(channel, options);
   initFileSearchChannel(channel, options);

code/core/src/core-server/utils/ai-checklist-flags.test.ts

@@ -0,0 +1,105 @@
+import { resolve } from 'node:path';
+
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+const { mockCacheStore, mockCache } = vi.hoisted(() => {
+  const store = new Map<string, unknown>();
+  return {
+    mockCacheStore: store,
+    mockCache: {
+      get: async (key: string) => store.get(key),
+      set: async (key: string, value: unknown) => {
+        store.set(key, value);
+      },
+    },
+  };
+});
+
+vi.mock('storybook/internal/common', () => ({
+  cache: mockCache,
+}));
+
+describe('ai-checklist-flags', () => {
+  beforeEach(() => {
+    mockCacheStore.clear();
+  });
+
+  afterEach(() => {
+    vi.resetModules();
+  });
+
+  describe('hasAiInitOptIn', () => {
+    it('returns false when nothing is cached', async () => {
+      const { hasAiInitOptIn } = await import('./ai-checklist-flags.ts');
+      expect(await hasAiInitOptIn('/some/project/.storybook')).toBe(false);
+    });
+
+    it('returns true when the cached configDir matches the resolved input', async () => {
+      mockCacheStore.set('ai-init-opt-in', {
+        timestamp: Date.now(),
+        configDir: resolve('/repo/apps/web/.storybook'),
+      });
+      const { hasAiInitOptIn } = await import('./ai-checklist-flags.ts');
+      expect(await hasAiInitOptIn('/repo/apps/web/.storybook')).toBe(true);
+    });
+
+    it('returns false when the cached configDir is for a different project', async () => {
+      mockCacheStore.set('ai-init-opt-in', {
+        timestamp: Date.now(),
+        configDir: resolve('/repo/apps/web/.storybook'),
+      });
+      const { hasAiInitOptIn } = await import('./ai-checklist-flags.ts');
+      expect(await hasAiInitOptIn('/repo/packages/ui/.storybook')).toBe(false);
+    });
+
+    it('returns false when the cached entry lacks a configDir field', async () => {
+      // Defensive — should never happen in practice because the CLI always
+      // writes configDir, but a corrupted cache shouldn't unlock this flag.
+      mockCacheStore.set('ai-init-opt-in', { timestamp: Date.now() });
+      const { hasAiInitOptIn } = await import('./ai-checklist-flags.ts');
+      expect(await hasAiInitOptIn('/any/project/.storybook')).toBe(false);
+    });
+  });
+
+  describe('hasAiSetupRun', () => {
+    it('returns false when nothing is cached', async () => {
+      const { hasAiSetupRun } = await import('./ai-checklist-flags.ts');
+      expect(await hasAiSetupRun('/some/project/.storybook')).toBe(false);
+    });
+
+    it('returns true when the cached configDir matches', async () => {
+      mockCacheStore.set('ai-setup-ran', {
+        timestamp: Date.now(),
+        configDir: resolve('/repo/apps/web/.storybook'),
+      });
+      const { hasAiSetupRun } = await import('./ai-checklist-flags.ts');
+      expect(await hasAiSetupRun('/repo/apps/web/.storybook')).toBe(true);
+    });
+
+    it('returns false when the cached configDir is for a sibling monorepo project', async () => {
+      // Regression: running `storybook ai setup` in one repo must not flip
+      // another repo's checklist to "done".
+      mockCacheStore.set('ai-setup-ran', {
+        timestamp: Date.now(),
+        configDir: resolve('/repo/apps/web/.storybook'),
+      });
+      const { hasAiSetupRun } = await import('./ai-checklist-flags.ts');
+      expect(await hasAiSetupRun('/repo/packages/ui/.storybook')).toBe(false);
+    });
+
+    it('treats relative input as resolved against cwd', async () => {
+      mockCacheStore.set('ai-setup-ran', {
+        timestamp: Date.now(),
+        configDir: resolve('.storybook'),
+      });
+      const { hasAiSetupRun } = await import('./ai-checklist-flags.ts');
+      expect(await hasAiSetupRun('.storybook')).toBe(true);
+    });
+
+    it('returns false when the cached entry lacks a configDir field', async () => {
+      mockCacheStore.set('ai-setup-ran', { timestamp: Date.now() });
+      const { hasAiSetupRun } = await import('./ai-checklist-flags.ts');
+      expect(await hasAiSetupRun('/any/project/.storybook')).toBe(false);
+    });
+  });
+});

code/core/src/core-server/utils/ai-checklist-flags.ts

@@ -0,0 +1,56 @@
+import { resolve } from 'node:path';
+
+import { cache } from 'storybook/internal/common';
+
+/**
+ * Flags persisted to the regular fs cache by the CLI to drive AI-related UI in
+ * the dev server. They live OUTSIDE the telemetry event cache on purpose:
+ * Storybook's UI behavior must not depend on whether telemetry happens to be
+ * enabled. Both flags are tiny local files containing no PII.
+ *
+ * Both flags are scoped to a Storybook project via `configDir`. In monorepos
+ * with hoisted `node_modules`, multiple Storybook projects share the same
+ * `node_modules/.cache/storybook/...` directory — without scoping, running
+ * `storybook ai setup` (or `storybook init` with AI accepted) in package A
+ * would falsely flip package B's checklist or copy-prompt UI.
+ *
+ * The CLI writes `{ timestamp, configDir }` (absolute, resolved). The dev
+ * server compares the cached `configDir` against its own resolved
+ * `options.configDir` and only honors the flag on a match.
+ */
+
+interface ProjectScopedFlag {
+  timestamp: number;
+  configDir: string;
+}
+
+function isProjectScopedFlag(value: unknown): value is ProjectScopedFlag {
+  return (
+    typeof value === 'object' &&
+    value !== null &&
+    'configDir' in value &&
+    typeof (value as ProjectScopedFlag).configDir === 'string'
+  );
+}
+
+async function readProjectScopedFlag(key: string, configDir: string): Promise<boolean> {
+  try {
+    const value = await cache.get(key);
+    if (!isProjectScopedFlag(value)) {
+      return false;
+    }
+    return value.configDir === resolve(configDir);
+  } catch {
+    return false;
+  }
+}
+
+/** Written by `storybook init` when the user accepted the AI feature. */
+export async function hasAiInitOptIn(configDir: string): Promise<boolean> {
+  return readProjectScopedFlag('ai-init-opt-in', configDir);
+}
+
+/** Written by `storybook ai setup` when the prompt CLI ran in this project. */
+export async function hasAiSetupRun(configDir: string): Promise<boolean> {
+  return readProjectScopedFlag('ai-setup-ran', configDir);
+}