dify/cli/scripts/generate-command-tree.ts

import { readdir, readFile, rename, writeFile } from 'node:fs/promises'
import { join, relative, sep } from 'node:path'
import { fileURLToPath } from 'node:url'
import { isExcludedCommandPath } from '../src/framework/command-fs.js'

const RESERVED_JS_KEYWORDS = new Set([
  'break',
  'case',
  'catch',
  'class',
  'const',
  'continue',
  'debugger',
  'default',
  'delete',
  'do',
  'else',
  'enum',
  'export',
  'extends',
  'false',
  'finally',
  'for',
  'function',
  'if',
  'import',
  'in',
  'instanceof',
  'new',
  'null',
  'return',
  'super',
  'switch',
  'this',
  'throw',
  'true',
  'try',
  'typeof',
  'var',
  'void',
  'while',
  'with',
  'yield',
])

export type CommandEntry = {
  readonly tokens: readonly string[]
  readonly identifier: string
  readonly importPath: string
}

export type TreeNode = {
  command?: string
  subcommands: Map<string, TreeNode>
}

export function pathToTokens(filePath: string, commandsRoot: string): string[] {
  const normalized = filePath.replace(/\\/g, '/')
  const root = commandsRoot.replace(/\\/g, '/').replace(/\/$/, '')
  const trimmed = normalized.startsWith(`${root}/`)
    ? normalized.slice(root.length + 1)
    : normalized
  const withoutIndex = trimmed.replace(/\/index\.ts$/, '')
  return withoutIndex.split('/').filter(s => s.length > 0)
}

function capitalize(part: string): string {
  if (part.length === 0)
    return ''
  return part[0]!.toUpperCase() + part.slice(1)
}

export function tokensToIdentifier(tokens: readonly string[]): string {
  const id = tokens
    .flatMap(t => t.split(/[-_]/))
    .map(capitalize)
    .join('')
  if (RESERVED_JS_KEYWORDS.has(id.toLowerCase()))
    return `_${id}`
  return id
}

export function buildTree(entries: readonly CommandEntry[]): TreeNode {
  const root: TreeNode = { subcommands: new Map() }
  for (const entry of entries) {
    let node = root
    for (let i = 0; i < entry.tokens.length; i++) {
      const tok = entry.tokens[i]!
      let next = node.subcommands.get(tok)
      if (!next) {
        next = { subcommands: new Map() }
        node.subcommands.set(tok, next)
      }
      node = next
    }
    node.command = entry.identifier
  }
  return root
}

const HEADER = `// @generated by scripts/generate-command-tree.ts — DO NOT EDIT.
// Regenerate via \`pnpm tree:gen\`. Drift gated by \`pnpm tree:check\` in CI.
`

function compareStrings(a: string, b: string): number {
  if (a < b)
    return -1
  if (a > b)
    return 1
  return 0
}

function emitImports(entries: readonly CommandEntry[]): string {
  const sorted = [...entries].sort((a, b) => compareStrings(a.importPath, b.importPath))
  const lines = [`import type { CommandTree } from '../framework/registry.js'`]
  for (const e of sorted)
    lines.push(`import ${e.identifier} from '${e.importPath}'`)
  return lines.join('\n')
}

function emitNode(node: TreeNode, indent: string): string {
  const inner = `${indent}  `
  const keys = [...node.subcommands.keys()].sort()
  const parts: string[] = []

  if (node.command !== undefined)
    parts.push(`${inner}command: ${node.command},`)

  if (keys.length === 0) {
    parts.push(`${inner}subcommands: {},`)
  }
  else {
    parts.push(`${inner}subcommands: {`)
    for (const key of keys) {
      const child = node.subcommands.get(key)!
      parts.push(emitEntry(key, child, `${inner}  `))
    }
    parts.push(`${inner}},`)
  }

  return parts.join('\n')
}

function emitEntry(key: string, node: TreeNode, indent: string): string {
  const isLeaf = node.subcommands.size === 0 && node.command !== undefined
  if (isLeaf)
    return `${indent}${key}: { command: ${node.command}, subcommands: {} },`

  return [
    `${indent}${key}: {`,
    emitNode(node, indent),
    `${indent}},`,
  ].join('\n')
}

export function formatModule(entries: readonly CommandEntry[], tree: TreeNode): string {
  const importsBlock = emitImports(entries)
  const topKeys = [...tree.subcommands.keys()].sort()
  const literalParts = ['export const commandTree: CommandTree = {']
  for (const key of topKeys) {
    const child = tree.subcommands.get(key)!
    literalParts.push(emitEntry(key, child, '  '))
  }
  literalParts.push('}')
  return `${HEADER}\n${importsBlock}\n\n${literalParts.join('\n')}\n`
}

async function walk(dir: string): Promise<string[]> {
  const out: string[] = []
  const entries = await readdir(dir, { withFileTypes: true })
  for (const e of entries) {
    const full = join(dir, e.name)
    if (e.isDirectory())
      out.push(...await walk(full))
    else if (e.isFile())
      out.push(full)
  }
  return out
}

function toPosix(p: string): string {
  return p.split(sep).join('/')
}

export async function discoverCommands(commandsDir: string): Promise<CommandEntry[]> {
  const all = await walk(commandsDir)
  const tsFiles = all.filter(f => f.endsWith('.ts') && !f.endsWith('.test.ts') && !f.endsWith('.d.ts'))

  const loose: string[] = []
  for (const abs of tsFiles) {
    const rel = toPosix(relative(commandsDir, abs))
    if (isExcludedCommandPath(rel))
      continue
    if (rel === 'tree.ts' || rel === 'tree.generated.ts')
      continue
    // Only flag files directly under commands/ (no path separator — no parent folder)
    if (!rel.includes('/'))
      loose.push(rel)
  }
  if (loose.length > 0) {
    const list = loose.map(p => `  - src/commands/${p}`).join('\n')
    throw new Error(
      `commands must live under their own folder (see CLAUDE memory: feedback_cli_command_structure). Found:\n${list}`,
    )
  }

  const entries: CommandEntry[] = []
  for (const abs of tsFiles) {
    const rel = toPosix(relative(commandsDir, abs))
    if (isExcludedCommandPath(rel))
      continue
    if (!rel.endsWith('/index.ts'))
      continue
    const tokens = pathToTokens(rel, '')
    if (tokens.length === 0)
      continue
    if (tokens[0]!.startsWith('-'))
      throw new Error(`command token cannot start with '-': ${rel}`)
    entries.push({
      tokens,
      identifier: tokensToIdentifier(tokens),
      importPath: `./${tokens.join('/')}/index.js`,
    })
  }

  entries.sort((a, b) => compareStrings(a.importPath, b.importPath))

  if (entries.length === 0)
    throw new Error(`no commands found under ${commandsDir}`)

  assertUniqueIdentifiers(entries)
  return entries
}

function assertUniqueIdentifiers(entries: readonly CommandEntry[]): void {
  const seen = new Map<string, string>()
  for (const e of entries) {
    const prev = seen.get(e.identifier)
    if (prev !== undefined)
      throw new Error(`identifier collision: ${e.identifier} from ${prev} and ${e.importPath}`)
    seen.set(e.identifier, e.importPath)
  }
}

export type GenerateOptions = {
  readonly commandsDir: string
  readonly mode: 'write' | 'check'
}

export type GenerateResult
  = | { mode: 'write', wrote: boolean, path: string }
    | { mode: 'check', ok: true, path: string }
    | { mode: 'check', ok: false, path: string, diff: string }

export async function generate(opts: GenerateOptions): Promise<GenerateResult> {
  const entries = await discoverCommands(opts.commandsDir)
  const tree = buildTree(entries)
  const content = formatModule(entries, tree)
  const target = join(opts.commandsDir, 'tree.generated.ts')

  if (opts.mode === 'check') {
    let onDisk = ''
    try {
      onDisk = await readFile(target, 'utf8')
    }
    catch {
      onDisk = ''
    }
    if (onDisk === content)
      return { mode: 'check', ok: true, path: target }
    return { mode: 'check', ok: false, path: target, diff: shortDiff(onDisk, content) }
  }

  const tmp = `${target}.tmp-${process.pid}-${Date.now()}`
  await writeFile(tmp, content, 'utf8')
  await rename(tmp, target)
  return { mode: 'write', wrote: true, path: target }
}

function shortDiff(a: string, b: string): string {
  const aLines = a.split('\n')
  const bLines = b.split('\n')
  const lines: string[] = []
  const max = Math.max(aLines.length, bLines.length)
  for (let i = 0; i < max; i++) {
    if (aLines[i] !== bLines[i]) {
      if (aLines[i] !== undefined)
        lines.push(`- ${aLines[i]}`)
      if (bLines[i] !== undefined)
        lines.push(`+ ${bLines[i]}`)
    }
  }
  return lines.slice(0, 40).join('\n')
}

async function main(): Promise<void> {
  const here = fileURLToPath(import.meta.url)
  const commandsDir = join(here, '..', '..', 'src', 'commands')
  const checkMode = process.argv.includes('--check')
  const result = await generate({ commandsDir, mode: checkMode ? 'check' : 'write' })

  if (result.mode === 'write') {
    process.stdout.write(`tree:gen wrote ${result.path}\n`)
    return
  }
  if (result.ok) {
    process.stdout.write(`tree:check ok\n`)
    return
  }
  process.stderr.write(`tree:check FAILED — tree.generated.ts is stale.\nDiff (first 40 lines):\n${result.diff}\n\nRun \`pnpm tree:gen\` and commit.\n`)
  process.exit(1)
}

const invokedDirectly = process.argv[1] !== undefined
  && fileURLToPath(import.meta.url) === process.argv[1]

if (invokedDirectly)
  await main()