@fragments-sdk/cli 0.2.2

package/src/commands/enhance.ts

@@ -0,0 +1,858 @@
+/**
+ * fragments enhance - AI-powered documentation generation
+ *
+ * Analyzes codebase for component usage patterns and generates
+ * when/whenNot documentation using AI.
+ *
+ * Supports:
+ * - Anthropic (Claude) API
+ * - OpenAI (GPT) API
+ * - Context-only mode for IDE AI (Cursor, Copilot, etc.)
+ */
+
+import pc from 'picocolors';
+import { BRAND } from '../core/index.js';
+import { loadConfig } from '../core/node.js';
+import { readFile, writeFile } from 'node:fs/promises';
+import { resolve, relative, join } from 'node:path';
+import {
+  scanCodebase,
+  parseAllStories,
+  generateComponentContext,
+  generateSystemPrompt,
+  generateUserPrompt,
+  generatePromptContext,
+  getScanStats,
+  extractPropsFromFile,
+  renderAllComponentVariants,
+  checkStorybookRunning,
+  shutdownSharedPool,
+  type ScanProgress,
+  type ComponentContext,
+  type UsageAnalysis,
+  type ParsedStoryFile,
+  type PropsExtractionResult,
+  type RenderedVariant,
+} from '../service/index.js';
+
+/**
+ * Supported AI providers
+ */
+export type AIProvider = 'anthropic' | 'openai' | 'none';
+
+/**
+ * Options for enhance command
+ */
+export interface EnhanceOptions {
+  /** Path to config file */
+  config?: string;
+  /** Component name to enhance (or 'all') */
+  component?: string;
+  /** Skip confirmation prompts */
+  yes?: boolean;
+  /** Only analyze, don't modify files */
+  dryRun?: boolean;
+  /** Output format: interactive, json, quiet, context */
+  format?: 'interactive' | 'json' | 'quiet' | 'context';
+  /** AI provider: anthropic, openai, or none */
+  provider?: AIProvider;
+  /** API key (or use ANTHROPIC_API_KEY/OPENAI_API_KEY env) */
+  apiKey?: string;
+  /** Model to use */
+  model?: string;
+  /** Root directory to scan */
+  root?: string;
+  /** Output context only (for use with IDE AI like Cursor) */
+  contextOnly?: boolean;
+  /** Render variants from running Storybook (captures complex stories that fail static extraction) */
+  renderVariants?: boolean;
+  /** Storybook URL for variant rendering (default: http://localhost:6006) */
+  storybookUrl?: string;
+}
+
+/**
+ * Enhanced segment data
+ */
+export interface EnhancedSegment {
+  componentName: string;
+  added: {
+    when: string[];
+    whenNot: string[];
+  };
+  confidence: number;
+  tokensUsed: number;
+  skipped: boolean;
+  reason?: string;
+}
+
+/**
+ * Result of enhance command
+ */
+export interface EnhanceResult {
+  success: boolean;
+  enhanced: EnhancedSegment[];
+  totalTokens: number;
+  estimatedCost: number;
+  /** Context output for IDE AI mode */
+  context?: string;
+}
+
+/**
+ * Default models for each provider
+ */
+const DEFAULT_MODELS: Record<AIProvider, string> = {
+  anthropic: 'claude-sonnet-4-20250514',
+  openai: 'gpt-4o',
+  none: '',
+};
+
+/**
+ * Run the enhance command
+ */
+export async function enhance(options: EnhanceOptions = {}): Promise<EnhanceResult> {
+  const {
+    config: configPath,
+    component,
+    yes = false,
+    dryRun = false,
+    format = 'interactive',
+    provider = detectProvider(options),
+    apiKey = getApiKey(provider, options.apiKey),
+    model = options.model || DEFAULT_MODELS[provider],
+    root = process.cwd(),
+    contextOnly = false,
+    renderVariants = false,
+    storybookUrl = 'http://localhost:6006',
+  } = options;
+
+  const isInteractive = format === 'interactive';
+  const isQuiet = format === 'quiet';
+  const isContextMode = contextOnly || format === 'context' || provider === 'none';
+
+  // Load config
+  const config = await loadConfig(configPath);
+  const rootDir = resolve(root);
+
+  if (isInteractive) {
+    console.log(pc.cyan(`\n${BRAND.name} AI Enhancement\n`));
+    if (isContextMode) {
+      console.log(pc.dim('Running in context-only mode (for use with IDE AI like Cursor)\n'));
+    } else {
+      console.log(pc.dim(`Using ${provider} API with model: ${model}\n`));
+    }
+  }
+
+  // Check for API key (unless context-only mode)
+  if (!isContextMode && !apiKey) {
+    const envVar = provider === 'openai' ? 'OPENAI_API_KEY' : 'ANTHROPIC_API_KEY';
+    const error = `API key required. Set ${envVar} or use --api-key, or use --context-only for IDE AI mode`;
+    if (format === 'json') {
+      console.log(JSON.stringify({ success: false, error }));
+    } else if (!isQuiet) {
+      console.error(pc.red('Error:'), error);
+      console.log(pc.dim('\nTip: Use --context-only to generate prompts for Cursor/Copilot/etc.'));
+    }
+    return { success: false, enhanced: [], totalTokens: 0, estimatedCost: 0 };
+  }
+
+  // Phase 1: Scan codebase for usage patterns
+  if (isInteractive) {
+    console.log(pc.dim('Phase 1: Scanning codebase for usage patterns...'));
+  }
+
+  let usageAnalysis: UsageAnalysis;
+  try {
+    usageAnalysis = await scanCodebase({
+      rootDir,
+      useCache: true,
+      onProgress: (progress: ScanProgress) => {
+        if (isInteractive && progress.phase === 'scanning') {
+          const pct = Math.round((progress.current / progress.total) * 100);
+          process.stdout.write(`\r  Scanning: ${pct}% (${progress.current}/${progress.total})`);
+        }
+      },
+    });
+    if (isInteractive) {
+      process.stdout.write('\r' + ' '.repeat(60) + '\r');
+      const stats = getScanStats(usageAnalysis);
+      console.log(pc.green(`  Found ${stats.totalUsages} usages across ${stats.totalFiles} files`));
+    }
+  } catch (error) {
+    const msg = error instanceof Error ? error.message : String(error);
+    if (format === 'json') {
+      console.log(JSON.stringify({ success: false, error: `Scan failed: ${msg}` }));
+    } else if (!isQuiet) {
+      console.error(pc.red('Scan failed:'), msg);
+    }
+    return { success: false, enhanced: [], totalTokens: 0, estimatedCost: 0 };
+  }
+
+  // Phase 2: Parse Storybook stories
+  if (isInteractive) {
+    console.log(pc.dim('Phase 2: Parsing Storybook stories...'));
+  }
+
+  let storyFiles: Map<string, ParsedStoryFile>;
+  try {
+    storyFiles = await parseAllStories(rootDir);
+    if (isInteractive) {
+      console.log(pc.green(`  Found ${storyFiles.size} story files`));
+    }
+  } catch {
+    storyFiles = new Map();
+    if (isInteractive) {
+      console.log(pc.yellow('  No Storybook stories found'));
+    }
+  }
+
+  // Phase 3: Load fragment files
+  if (isInteractive) {
+    console.log(pc.dim('Phase 3: Loading fragment files...'));
+  }
+
+  // Search entire root dir, not just src/
+  const segmentFiles = await findSegmentFiles(rootDir);
+
+  if (segmentFiles.length === 0) {
+    const msg = 'No fragment files found';
+    if (format === 'json') {
+      console.log(JSON.stringify({ success: false, error: msg }));
+    } else if (!isQuiet) {
+      console.log(pc.yellow(msg));
+    }
+    return { success: false, enhanced: [], totalTokens: 0, estimatedCost: 0 };
+  }
+
+  if (isInteractive) {
+    console.log(pc.green(`  Found ${segmentFiles.length} fragment files`));
+  }
+
+  // Filter components if specified
+  let componentsToEnhance: string[];
+  if (component && component !== 'all') {
+    componentsToEnhance = [component];
+  } else {
+    componentsToEnhance = segmentFiles.map(f => extractComponentName(f));
+  }
+
+  // Phase 4: Extract props from TypeScript source files
+  if (isInteractive) {
+    console.log(pc.dim('Phase 4: Extracting props from TypeScript interfaces...'));
+  }
+
+  const propsExtractions = new Map<string, PropsExtractionResult>();
+  for (const compName of componentsToEnhance) {
+    const segmentFile = segmentFiles.find(f => extractComponentName(f) === compName);
+    if (!segmentFile) continue;
+
+    // Try to find the component source file relative to the segment file
+    const segmentDir = segmentFile.replace(/\.segment\.(tsx?|jsx?)$/, '');
+    const possiblePaths = [
+      `${segmentDir}.tsx`,
+      `${segmentDir}.ts`,
+      `${segmentDir}/index.tsx`,
+      `${segmentDir}/index.ts`,
+      join(rootDir, 'src', 'components', `${compName}.tsx`),
+      join(rootDir, 'src', 'components', compName, `${compName}.tsx`),
+      join(rootDir, 'src', 'components', compName, 'index.tsx'),
+    ];
+
+    for (const srcPath of possiblePaths) {
+      try {
+        const fs = await import('node:fs');
+        if (fs.existsSync(srcPath)) {
+          const extraction = await extractPropsFromFile(srcPath, {
+            propsTypeName: `${compName}Props`,
+          });
+          if (extraction.success) {
+            propsExtractions.set(compName, extraction);
+            break;
+          }
+        }
+      } catch {
+        // Skip files that can't be parsed
+      }
+    }
+  }
+
+  if (isInteractive) {
+    const propsCount = Array.from(propsExtractions.values())
+      .reduce((sum, ext) => sum + ext.props.length, 0);
+    console.log(pc.green(`  Extracted ${propsCount} props from ${propsExtractions.size} component files`));
+  }
+
+  // Phase 5 (optional): Render variants from Storybook
+  const renderedVariants = new Map<string, RenderedVariant[]>();
+
+  if (renderVariants) {
+    if (isInteractive) {
+      console.log(pc.dim('Phase 5: Rendering variants from Storybook...'));
+    }
+
+    // Check if Storybook is running
+    const storybookRunning = await checkStorybookRunning(storybookUrl);
+
+    if (!storybookRunning) {
+      if (isInteractive) {
+        console.log(pc.yellow(`  Storybook not running at ${storybookUrl}. Skipping variant rendering.`));
+        console.log(pc.dim(`  Start Storybook with: npm run storybook (or yarn storybook)`));
+      }
+    } else {
+      // Render variants for each component
+      let totalVariants = 0;
+      let failedVariants = 0;
+
+      for (const compName of componentsToEnhance) {
+        try {
+          const result = await renderAllComponentVariants(storybookUrl, compName);
+          if (result.variants.length > 0) {
+            renderedVariants.set(compName, result.variants);
+            totalVariants += result.variants.length;
+          }
+          failedVariants += result.failed.length;
+        } catch (err) {
+          if (isInteractive) {
+            console.log(pc.dim(`  ${compName}: Failed to render (${(err as Error).message})`));
+          }
+        }
+      }
+
+      if (isInteractive) {
+        if (totalVariants > 0) {
+          console.log(pc.green(`  Rendered ${totalVariants} variants from ${renderedVariants.size} components`));
+        }
+        if (failedVariants > 0) {
+          console.log(pc.yellow(`  ${failedVariants} variant(s) failed to render`));
+        }
+      }
+
+      // Shutdown browser pool
+      await shutdownSharedPool();
+    }
+  }
+
+  // Build contexts for all components
+  const contexts: Array<{ name: string; context: ComponentContext; segmentFile: string }> = [];
+  for (const compName of componentsToEnhance) {
+    const analysis = usageAnalysis.components[compName];
+    const stories = storyFiles.get(compName);
+    const segmentFile = segmentFiles.find(f => extractComponentName(f) === compName);
+    const propsExtraction = propsExtractions.get(compName);
+
+    if (!segmentFile) continue;
+
+    const context = generateComponentContext(
+      compName,
+      analysis,
+      undefined,
+      stories,
+      propsExtraction
+    );
+    contexts.push({ name: compName, context, segmentFile });
+  }
+
+  // Context-only mode: output prompts for IDE AI
+  if (isContextMode) {
+    return handleContextOnlyMode(contexts, format, isInteractive);
+  }
+
+  // Phase 6: Generate AI enhancements
+  if (isInteractive) {
+    console.log(pc.dim(`\nPhase 6: Generating AI enhancements for ${componentsToEnhance.length} component(s)...\n`));
+  }
+
+  const enhanced: EnhancedSegment[] = [];
+  let totalTokens = 0;
+
+  // Initialize AI client
+  const aiClient = await createAIClient(provider, apiKey!);
+
+  for (const { name: compName, context, segmentFile } of contexts) {
+    // Check if we have enough data
+    if (!context.usageAnalysis || context.usageAnalysis.totalUsages < 2) {
+      enhanced.push({
+        componentName: compName,
+        added: { when: [], whenNot: [] },
+        confidence: 0,
+        tokensUsed: 0,
+        skipped: true,
+        reason: context.usageAnalysis?.totalUsages === 1 ? 'Only 1 usage found' : 'No usage data found',
+      });
+      if (isInteractive) {
+        console.log(pc.dim(`  ${compName}: Skipped (insufficient data)`));
+      }
+      continue;
+    }
+
+    if (isInteractive) {
+      process.stdout.write(`  ${compName}: Generating...`);
+    }
+
+    try {
+      const result = await generateEnhancement(aiClient, provider, model, context);
+
+      enhanced.push({
+        componentName: compName,
+        added: result.suggestions,
+        confidence: result.confidence,
+        tokensUsed: result.tokensUsed,
+        skipped: false,
+      });
+
+      totalTokens += result.tokensUsed;
+
+      if (isInteractive) {
+        process.stdout.write(`\r  ${compName}: ${pc.green('Done')} (${result.suggestions.when.length} when, ${result.suggestions.whenNot.length} whenNot)\n`);
+      }
+    } catch (error) {
+      const msg = error instanceof Error ? error.message : String(error);
+      enhanced.push({
+        componentName: compName,
+        added: { when: [], whenNot: [] },
+        confidence: 0,
+        tokensUsed: 0,
+        skipped: true,
+        reason: `AI error: ${msg}`,
+      });
+      if (isInteractive) {
+        process.stdout.write(`\r  ${compName}: ${pc.red('Failed')} - ${msg}\n`);
+      }
+    }
+  }
+
+  // Calculate cost
+  const estimatedCost = calculateCost(provider, totalTokens);
+
+  // Phase 7: Apply changes
+  if (!dryRun) {
+    if (isInteractive) {
+      console.log(pc.dim('\nPhase 7: Updating segment files...'));
+    }
+
+    for (const result of enhanced) {
+      if (result.skipped || (result.added.when.length === 0 && result.added.whenNot.length === 0)) {
+        continue;
+      }
+
+      const segmentFile = segmentFiles.find(f => extractComponentName(f) === result.componentName);
+      if (!segmentFile) continue;
+
+      try {
+        await updateSegmentFile(segmentFile, result.added);
+        if (isInteractive) {
+          console.log(pc.green(`  Updated: ${relative(rootDir, segmentFile)}`));
+        }
+      } catch {
+        if (isInteractive) {
+          console.log(pc.red(`  Failed to update: ${relative(rootDir, segmentFile)}`));
+        }
+      }
+    }
+  }
+
+  // Output results
+  const successCount = enhanced.filter(e => !e.skipped).length;
+
+  if (format === 'json') {
+    console.log(JSON.stringify({
+      success: true,
+      enhanced,
+      totalTokens,
+      estimatedCost,
+    }, null, 2));
+  } else if (isInteractive) {
+    console.log();
+    console.log(pc.bold('Summary:'));
+    console.log(`  Components processed: ${componentsToEnhance.length}`);
+    console.log(`  Successfully enhanced: ${successCount}`);
+    console.log(`  Skipped: ${enhanced.filter(e => e.skipped).length}`);
+    console.log(`  Total tokens used: ${totalTokens}`);
+    console.log(`  Estimated cost: $${estimatedCost.toFixed(4)}`);
+
+    if (dryRun) {
+      console.log();
+      console.log(pc.yellow('Dry run - no files were modified'));
+    }
+    console.log();
+  }
+
+  return {
+    success: true,
+    enhanced,
+    totalTokens,
+    estimatedCost,
+  };
+}
+
+/**
+ * Handle context-only mode for IDE AI
+ */
+function handleContextOnlyMode(
+  contexts: Array<{ name: string; context: ComponentContext; segmentFile: string }>,
+  format: string,
+  isInteractive: boolean
+): EnhanceResult {
+  const systemPrompt = generateSystemPrompt();
+  const componentContexts: string[] = [];
+
+  for (const { name, context } of contexts) {
+    if (!context.usageAnalysis || context.usageAnalysis.totalUsages < 2) {
+      continue;
+    }
+    componentContexts.push(generatePromptContext(context));
+  }
+
+  if (componentContexts.length === 0) {
+    if (format === 'json') {
+      console.log(JSON.stringify({ success: false, error: 'No components with sufficient usage data' }));
+    } else if (isInteractive) {
+      console.log(pc.yellow('\nNo components with sufficient usage data to analyze.'));
+    }
+    return { success: false, enhanced: [], totalTokens: 0, estimatedCost: 0 };
+  }
+
+  const fullContext = `${systemPrompt}
+
+---
+
+${componentContexts.join('\n\n---\n\n')}
+
+---
+
+Based on the usage analysis above, generate "when" and "whenNot" documentation for each component.
+
+For each component, provide your response in JSON format:
+\`\`\`json
+{
+  "componentName": "...",
+  "when": ["scenario 1", "scenario 2", ...],
+  "whenNot": ["anti-pattern 1", ...]
+}
+\`\`\``;
+
+  if (format === 'json') {
+    console.log(JSON.stringify({
+      success: true,
+      context: fullContext,
+      instructions: 'Copy this prompt into your IDE AI (Cursor, Copilot, etc.) to generate suggestions',
+    }, null, 2));
+  } else {
+    console.log(pc.bold('\n📋 AI Context Generated\n'));
+    console.log(pc.dim('Copy the following prompt into your IDE AI (Cursor, Copilot, etc.):\n'));
+    console.log(pc.dim('─'.repeat(60)));
+    console.log(fullContext);
+    console.log(pc.dim('─'.repeat(60)));
+    console.log();
+    console.log(pc.green('Tip: In Cursor, press Cmd+L to open chat and paste this prompt.'));
+    console.log(pc.dim('After getting suggestions, manually update your segment files.'));
+    console.log();
+  }
+
+  return {
+    success: true,
+    enhanced: [],
+    totalTokens: 0,
+    estimatedCost: 0,
+    context: fullContext,
+  };
+}
+
+/**
+ * Detect which AI provider to use based on available API keys
+ */
+function detectProvider(options: EnhanceOptions): AIProvider {
+  if (options.contextOnly) return 'none';
+  if (options.provider) return options.provider;
+  if (options.apiKey) {
+    // Try to guess from key format
+    if (options.apiKey.startsWith('sk-ant-')) return 'anthropic';
+    if (options.apiKey.startsWith('sk-')) return 'openai';
+  }
+  // Check environment variables
+  if (process.env.ANTHROPIC_API_KEY) return 'anthropic';
+  if (process.env.OPENAI_API_KEY) return 'openai';
+  return 'none';
+}
+
+/**
+ * Get API key for the provider
+ */
+function getApiKey(provider: AIProvider, explicitKey?: string): string | undefined {
+  if (explicitKey) return explicitKey;
+  if (provider === 'anthropic') return process.env.ANTHROPIC_API_KEY;
+  if (provider === 'openai') return process.env.OPENAI_API_KEY;
+  return undefined;
+}
+
+/**
+ * Create AI client for the provider
+ */
+async function createAIClient(provider: AIProvider, apiKey: string): Promise<unknown> {
+  if (provider === 'anthropic') {
+    const Anthropic = (await import('@anthropic-ai/sdk')).default;
+    return new Anthropic({ apiKey });
+  }
+  if (provider === 'openai') {
+    const OpenAI = (await import('openai')).default;
+    return new OpenAI({ apiKey });
+  }
+  throw new Error(`Unknown provider: ${provider}`);
+}
+
+/**
+ * Generate enhancement using the appropriate AI provider
+ */
+async function generateEnhancement(
+  client: unknown,
+  provider: AIProvider,
+  model: string,
+  context: ComponentContext
+): Promise<{
+  suggestions: { when: string[]; whenNot: string[] };
+  confidence: number;
+  tokensUsed: number;
+}> {
+  const systemPrompt = generateSystemPrompt();
+  const userPrompt = generateUserPrompt(context);
+
+  if (provider === 'anthropic') {
+    return generateWithAnthropic(client, model, systemPrompt, userPrompt, context);
+  }
+  if (provider === 'openai') {
+    return generateWithOpenAI(client, model, systemPrompt, userPrompt, context);
+  }
+  throw new Error(`Unknown provider: ${provider}`);
+}
+
+/**
+ * Generate with Anthropic API
+ */
+async function generateWithAnthropic(
+  client: unknown,
+  model: string,
+  systemPrompt: string,
+  userPrompt: string,
+  context: ComponentContext
+): Promise<{
+  suggestions: { when: string[]; whenNot: string[] };
+  confidence: number;
+  tokensUsed: number;
+}> {
+  const anthropic = client as import('@anthropic-ai/sdk').default;
+  const response = await anthropic.messages.create({
+    model,
+    max_tokens: 1024,
+    system: systemPrompt,
+    messages: [{ role: 'user', content: userPrompt }],
+  });
+
+  const content = response.content[0];
+  if (content.type !== 'text') {
+    throw new Error('Unexpected response type');
+  }
+
+  const suggestions = parseAIResponse(content.text);
+  const tokensUsed = (response.usage?.input_tokens || 0) + (response.usage?.output_tokens || 0);
+
+  return {
+    suggestions,
+    confidence: calculateConfidence(context, suggestions),
+    tokensUsed,
+  };
+}
+
+/**
+ * Generate with OpenAI API
+ */
+async function generateWithOpenAI(
+  client: unknown,
+  model: string,
+  systemPrompt: string,
+  userPrompt: string,
+  context: ComponentContext
+): Promise<{
+  suggestions: { when: string[]; whenNot: string[] };
+  confidence: number;
+  tokensUsed: number;
+}> {
+  const openai = client as import('openai').default;
+  const response = await openai.chat.completions.create({
+    model,
+    max_tokens: 1024,
+    messages: [
+      { role: 'system', content: systemPrompt },
+      { role: 'user', content: userPrompt },
+    ],
+  });
+
+  const content = response.choices[0]?.message?.content;
+  if (!content) {
+    throw new Error('No response from OpenAI');
+  }
+
+  const suggestions = parseAIResponse(content);
+  const tokensUsed = (response.usage?.prompt_tokens || 0) + (response.usage?.completion_tokens || 0);
+
+  return {
+    suggestions,
+    confidence: calculateConfidence(context, suggestions),
+    tokensUsed,
+  };
+}
+
+/**
+ * Parse AI response to extract suggestions
+ */
+function parseAIResponse(text: string): { when: string[]; whenNot: string[] } {
+  try {
+    // Extract JSON from response (may be wrapped in markdown code block)
+    const jsonMatch = text.match(/```json\n?([\s\S]*?)\n?```/) || text.match(/\{[\s\S]*\}/);
+    const jsonStr = jsonMatch ? (jsonMatch[1] || jsonMatch[0]) : text;
+    const parsed = JSON.parse(jsonStr);
+    return {
+      when: Array.isArray(parsed.when) ? parsed.when : [],
+      whenNot: Array.isArray(parsed.whenNot) ? parsed.whenNot : [],
+    };
+  } catch {
+    return extractSuggestionsFromText(text);
+  }
+}
+
+/**
+ * Extract suggestions from unstructured AI text response
+ */
+function extractSuggestionsFromText(text: string): { when: string[]; whenNot: string[] } {
+  const when: string[] = [];
+  const whenNot: string[] = [];
+
+  const lines = text.split('\n');
+  let currentSection: 'when' | 'whenNot' | null = null;
+
+  for (const line of lines) {
+    const trimmed = line.trim();
+    if (/^(when|use when|when to use)/i.test(trimmed)) {
+      currentSection = 'when';
+      continue;
+    }
+    if (/^(when not|do not use|avoid|whenNot)/i.test(trimmed)) {
+      currentSection = 'whenNot';
+      continue;
+    }
+
+    if (currentSection && /^[-*]\s+/.test(trimmed)) {
+      const item = trimmed.replace(/^[-*]\s+/, '').trim();
+      if (item && item.length > 10) {
+        if (currentSection === 'when') {
+          when.push(item);
+        } else {
+          whenNot.push(item);
+        }
+      }
+    }
+  }
+
+  return { when, whenNot };
+}
+
+/**
+ * Calculate confidence based on data quality and response
+ */
+function calculateConfidence(
+  context: ComponentContext,
+  suggestions: { when: string[]; whenNot: string[] }
+): number {
+  let confidence = 50;
+  if (context.usageAnalysis.totalUsages > 10) confidence += 20;
+  if (context.storybook && context.storybook.stories.length > 0) confidence += 15;
+  if (suggestions.when.length >= 2) confidence += 10;
+  if (suggestions.whenNot.length >= 1) confidence += 5;
+  return Math.min(confidence, 100);
+}
+
+/**
+ * Calculate cost based on provider and tokens
+ */
+function calculateCost(provider: AIProvider, tokens: number): number {
+  // Approximate costs per 1M tokens
+  const costsPer1M: Record<AIProvider, number> = {
+    anthropic: 3, // Claude Sonnet
+    openai: 5, // GPT-4o
+    none: 0,
+  };
+  return (tokens / 1_000_000) * costsPer1M[provider];
+}
+
+/**
+ * Find all segment files in a directory
+ */
+async function findSegmentFiles(dir: string): Promise<string[]> {
+  const fg = await import('fast-glob');
+  return fg.default(['**/*.segment.tsx', '**/*.segment.ts'], {
+    cwd: dir,
+    absolute: true,
+    ignore: ['**/node_modules/**', '**/dist/**'],
+  });
+}
+
+/**
+ * Extract component name from segment file path
+ */
+function extractComponentName(filePath: string): string {
+  const match = filePath.match(/([^/\\]+)\.segment\.(tsx?|jsx?)$/);
+  return match ? match[1] : '';
+}
+
+/**
+ * Update a segment file with new when/whenNot suggestions
+ */
+async function updateSegmentFile(
+  filePath: string,
+  suggestions: { when: string[]; whenNot: string[] }
+): Promise<void> {
+  const content = await readFile(filePath, 'utf-8');
+  let updated = content;
+
+  // Add 'when' items
+  if (suggestions.when.length > 0) {
+    const whenItems = suggestions.when.map(s => `    "${s}"`).join(',\n');
+    const whenMatch = updated.match(/when:\s*\[\s*([^\]]*)\]/);
+    if (whenMatch) {
+      const existingContent = whenMatch[1].trim();
+      if (existingContent) {
+        updated = updated.replace(
+          /when:\s*\[\s*([^\]]*)\]/,
+          `when: [\n${existingContent},\n${whenItems}\n  ]`
+        );
+      } else {
+        updated = updated.replace(
+          /when:\s*\[\s*\]/,
+          `when: [\n${whenItems}\n  ]`
+        );
+      }
+    }
+  }
+
+  // Add 'whenNot' items
+  if (suggestions.whenNot.length > 0) {
+    const whenNotItems = suggestions.whenNot.map(s => `    "${s}"`).join(',\n');
+    const whenNotMatch = updated.match(/whenNot:\s*\[\s*([^\]]*)\]/);
+    if (whenNotMatch) {
+      const existingContent = whenNotMatch[1].trim();
+      if (existingContent) {
+        updated = updated.replace(
+          /whenNot:\s*\[\s*([^\]]*)\]/,
+          `whenNot: [\n${existingContent},\n${whenNotItems}\n  ]`
+        );
+      } else {
+        updated = updated.replace(
+          /whenNot:\s*\[\s*\]/,
+          `whenNot: [\n${whenNotItems}\n  ]`
+        );
+      }
+    }
+  }
+
+  if (updated !== content) {
+    await writeFile(filePath, updated, 'utf-8');
+  }
+}