autocomplete-model/shared/symbols.ts

/**
 * "Symbols available at the cursor" context, from the TypeScript language
 * service, used both to build training data (pipeline/src/symbol-context.ts)
 * and at inference time (daemon). Keep the output format stable: Stage C was
 * trained on it.
 *
 * Context, in priority order:
 *   1. member completions when the cursor follows `expr.` / `expr.partial` (optional)
 *   2. hover signatures of imported symbols (ones used near the cursor first);
 *      repo-local interfaces/types/classes include their declaration source
 * Trimmed to a token budget.
 */
import fs from 'node:fs'
import path from 'node:path'
import ts from 'typescript'
import { parse as parseSfc } from '@vue/compiler-sfc'

import type { SymbolLang } from './symbol-lang.ts'
import type { Diagnostic } from './diagnostics.ts'
export { symbolLang, type SymbolLang } from './symbol-lang.ts'

export interface SymbolsRequest {
  /** Absolute path of the file being edited (it may not exist on disk yet). */
  filePath: string
  lang: SymbolLang
  /** The file as it is at completion time (prefix + suffix). */
  visible: string
  /** Cursor offset in `visible`. */
  cursor: number
  budget: number
  countTokens: (s: string) => number
  /** Include member completions when the cursor follows `expr.` (default true, as in training). */
  members?: boolean
  /**
   * Context format version. 1 = Stage C's training format (default, byte-identical to before).
   * 2 = also namespace imports (`import * as x`): x's members, the ones this file uses first, with signatures.
   */
  format?: 1 | 2
}

export interface SymbolsResult { context: string; items: number; hasMembers: boolean }

/** <script> regions of a Vue SFC, sorted by position (same as pipeline/src/parse.ts regionsOf). */
function vueScriptRegions(content: string): { start: number; end: number }[] {
  const { descriptor } = parseSfc(content, { sourceMap: false, ignoreEmpty: false })
  return [descriptor.script, descriptor.scriptSetup]
    .filter(b => b !== null)
    .map(b => ({ start: b!.loc.start.offset, end: b!.loc.end.offset }))
    .sort((a, b) => a.start - b.start)
}

function compilerOptionsFor(rootDir: string): ts.CompilerOptions {
  const fallback: ts.CompilerOptions = {
    target: ts.ScriptTarget.ESNext, module: ts.ModuleKind.ESNext, moduleResolution: ts.ModuleResolutionKind.Bundler,
    strict: true, allowJs: true, skipLibCheck: true, jsx: ts.JsxEmit.Preserve, allowImportingTsExtensions: true, noEmit: true,
  }
  const configPath = path.join(rootDir, 'tsconfig.json')
  let options = fallback
  if (fs.existsSync(configPath)) {
    const parsed = ts.getParsedCommandLineOfConfigFile(configPath, {}, { ...ts.sys, onUnRecoverableConfigFileDiagnostic: () => {} })
    if (parsed) options = { ...parsed.options, noEmit: true, skipLibCheck: true, allowJs: true }
  }
  // Common aliases in Quasar / Vite projects, if the tsconfig didn't resolve them.
  options.baseUrl ??= rootDir
  options.paths ??= { 'src/*': ['src/*'], '@/*': ['src/*'], 'components/*': ['src/components/*'], 'boot/*': ['src/boot/*'], 'stores/*': ['src/stores/*'] }
  return options
}

const displayText = (parts?: readonly ts.SymbolDisplayPart[]) => (parts ?? []).map(p => p.text).join('')
const clip = (s: string, max: number) => (s.length > max ? s.slice(0, max) + ' …' : s)

/** One TypeScript language service over a workspace root, with an overridable "current file". */
export class WorkspaceSymbols {
  readonly rootDir: string
  private service: ts.LanguageService
  private overrides = new Map<string, { text: string; version: number }>()
  private version = 0

  constructor(rootDir: string) {
    this.rootDir = rootDir
    const options = compilerOptionsFor(rootDir)
    const host: ts.LanguageServiceHost = {
      getCompilationSettings: () => options,
      getScriptFileNames: () => [...this.overrides.keys()],
      getScriptVersion: f => String(this.overrides.get(f)?.version ?? 0),
      getScriptSnapshot: f => {
        const o = this.overrides.get(f)
        if (o) return ts.ScriptSnapshot.fromString(o.text)
        return fs.existsSync(f) ? ts.ScriptSnapshot.fromString(fs.readFileSync(f, 'utf8')) : undefined
      },
      getCurrentDirectory: () => rootDir,
      getDefaultLibFileName: o => ts.getDefaultLibFilePath(o),
      fileExists: f => this.overrides.has(f) || ts.sys.fileExists(f),
      readFile: f => this.overrides.get(f)?.text ?? ts.sys.readFile(f),
      readDirectory: ts.sys.readDirectory,
      directoryExists: ts.sys.directoryExists,
      getDirectories: ts.sys.getDirectories,
    }
    this.service = ts.createLanguageService(host, ts.createDocumentRegistry())
  }

  /** Makes `text` the content of the file being edited; returns the TS-visible file name, text and cursor. */
  private load(filePath: string, lang: SymbolLang, visible: string, cursor: number) {
    let tsText = visible
    let tsCursor: number | null = cursor
    let fileName = filePath
    if (lang === 'vue') {
      const scripts = vueScriptRegions(visible)
      if (!scripts.length) return null
      tsText = scripts.map(r => visible.slice(r.start, r.end)).join('\n')
      const inScript = scripts.find(r => r.start <= cursor && cursor <= r.end)
      tsCursor = inScript ? cursor - inScript.start + scripts.slice(0, scripts.indexOf(inScript)).reduce((n, r) => n + r.end - r.start + 1, 0) : null
      fileName += '.ts'
    }
    this.overrides.clear()
    this.overrides.set(fileName, { text: tsText, version: ++this.version })
    return { fileName, tsText, tsCursor }
  }

  /** Source text of the declaration at a definition location, for repo-local types (not node_modules). */
  private declarationSource(fileName: string, pos: number): string | null {
    const def = this.service.getDefinitionAtPosition(fileName, pos)?.[0]
    if (!def || def.fileName.includes('node_modules') || def.fileName.endsWith('.d.ts')) return null
    if (!['interface', 'type', 'class', 'enum'].includes(def.kind)) return null
    const sf = this.service.getProgram()?.getSourceFile(def.fileName)
    if (!sf) return null
    let node: ts.Node | undefined
    const find = (n: ts.Node) => {
      if (n.getStart(sf) <= def.textSpan.start && def.textSpan.start + def.textSpan.length <= n.end) {
        if (ts.isInterfaceDeclaration(n) || ts.isTypeAliasDeclaration(n) || ts.isClassDeclaration(n) || ts.isEnumDeclaration(n)) node = n
        ts.forEachChild(n, find)
      }
    }
    find(sf)
    if (!node) return null
    let text = (node as ts.Node).getText(sf)
    if (ts.isClassDeclaration(node)) {
      // Signatures only: header + each member up to (not including) its body.
      const cls = node
      const header = sf.text.slice(cls.getStart(sf), cls.members.pos).trimEnd()
      const members = cls.members.map(m => {
        const body = (m as { body?: ts.Node }).body
        const sig = sf.text.slice(m.getStart(sf), body ? body.getStart(sf) : m.end).trim().replace(/\s*=\s*[\s\S]{60,}$/, '')
        return '  ' + sig.replace(/\s+/g, ' ').replace(/[;{]?$/, ';')
      })
      text = [header, ...members, '}'].join('\n')
    }
    const lines = text.split('\n')
    return lines.length > 25 ? [...lines.slice(0, 24), '  // …', lines.at(-1)!].join('\n') : text
  }

  context(req: SymbolsRequest): SymbolsResult | null {
    const loaded = this.load(req.filePath, req.lang, req.visible, req.cursor)
    if (!loaded) return null
    const { fileName, tsText, tsCursor } = loaded
    const items: { priority: number; text: string }[] = []

    // 1. Members at the cursor
    if (tsCursor !== null && req.members !== false) {
      const m = /\.([A-Za-z_$][\w$]*)?$/.exec(tsText.slice(Math.max(0, tsCursor - 200), tsCursor))
      if (m) {
        const pos = tsCursor - (m[1]?.length ?? 0)
        const completions = this.service.getCompletionsAtPosition(fileName, pos, { includeCompletionsWithInsertText: true })
        if (completions?.isMemberCompletion) {
          const entries = [...completions.entries].sort((a, b) => a.sortText.localeCompare(b.sortText) || a.name.localeCompare(b.name)).slice(0, 40)
          const detailed = entries.slice(0, 12).map(e => {
            const d = this.service.getCompletionEntryDetails(fileName, pos, e.name, undefined, undefined, undefined, undefined)
            return clip(displayText(d?.displayParts).replace(/^\((method|property|getter|setter)\) /, ''), 200)
          })
          const receiver = /([\w$.()[\]'"]+)\.[\w$]*$/.exec(tsText.slice(Math.max(0, tsCursor - 200), tsCursor))?.[1] ?? 'expression'
          const rest = entries.slice(12).map(e => e.name)
          items.push({ priority: 0, text: `// members of ${receiver}\n${detailed.join('\n')}${rest.length ? `\n// also: ${rest.join(', ')}` : ''}` })
        }
      }
    }

    // 2. Imported symbols
    const sf = ts.createSourceFile(fileName, tsText, ts.ScriptTarget.Latest, true)
    const near = tsCursor !== null
      ? tsText.slice(Math.max(0, tsCursor - 2000), tsCursor + 500)
      : req.visible.slice(Math.max(0, req.cursor - 2000), req.cursor + 500)
    for (const stmt of sf.statements) {
      if (!ts.isImportDeclaration(stmt) || !stmt.importClause) continue
      const module = (stmt.moduleSpecifier as ts.StringLiteral).text
      // Frameworks the model already knows well: their signatures would only spend budget.
      if (/^(vue|vue-router|pinia|node:.*|fs|path|events|crypto)$/.test(module)) continue
      const names: ts.Identifier[] = []
      if (stmt.importClause.name) names.push(stmt.importClause.name)
      const bindings = stmt.importClause.namedBindings
      if ((req.format ?? 1) >= 2 && bindings && ts.isNamespaceImport(bindings)) {
        const ns = this.namespaceMembers(fileName, tsText, bindings.name.text)
        if (ns) {
          const used = new RegExp(`\\b${bindings.name.text.replace(/\$/g, '\\$')}\\.`).test(near)
          items.push({ priority: used ? 1 : 2, text: `// from '${module}' (namespace ${bindings.name.text})\n${ns}` })
        }
        continue
      }
      if (bindings && ts.isNamedImports(bindings)) for (const el of bindings.elements) names.push(el.name)
      for (const id of names) {
        const pos = id.getStart(sf)
        const info = this.service.getQuickInfoAtPosition(fileName, pos)
        let text = displayText(info?.displayParts).replace(/^\(alias\) /, '').split(/\nimport /)[0]
        if (!text || / any$/.test(text) || text.endsWith(': any')) continue
        const decl = this.declarationSource(fileName, pos)
        if (decl) text = decl
        // Skip hovers that carry no signature (e.g. "import z", "namespace debug", "class Mutex").
        else if (!/[:(]/.test(text.split('\n')[0]) || /^(import|namespace|module) /.test(text)) continue
        const used = new RegExp(`\\b${id.text.replace(/\$/g, '\\$')}\\b`).test(near)
        items.push({ priority: used ? 1 : 2, text: `// from '${module}'\n${clip(text, 900)}` })
      }
    }

    if (!items.length) return null
    items.sort((a, b) => a.priority - b.priority)
    const chosen: string[] = []
    let used = 0
    for (const item of items) {
      const t = req.countTokens(item.text) + 1
      if (used + t > req.budget) continue
      chosen.push(item.text)
      used += t
    }
    return chosen.length ? { context: chosen.join('\n'), items: chosen.length, hasMembers: chosen[0].startsWith('// members of') } : null
  }

  /**
   * Members of a namespace import, as signature lines. Ranked by how often this file already uses each
   * member (`ns.member`), then by TypeScript's own completion order; details for the top 12.
   */
  private namespaceMembers(fileName: string, tsText: string, ns: string): string | null {
    // Probe: ask for completions at `ns.` appended to the file, then restore.
    const probe = `${tsText}\n;${ns}.`
    this.overrides.set(fileName, { text: probe, version: ++this.version })
    try {
      const pos = probe.length
      const completions = this.service.getCompletionsAtPosition(fileName, pos, {})
      if (!completions?.entries.length) return null
      const usage = new Map<string, number>()
      for (const m of tsText.matchAll(new RegExp(`\\b${ns.replace(/\$/g, '\\$')}\\.([A-Za-z_$][\\w$]*)`, 'g'))) usage.set(m[1], (usage.get(m[1]) ?? 0) + 1)
      const entries = [...completions.entries].sort((a, b) => (usage.get(b.name) ?? 0) - (usage.get(a.name) ?? 0) || a.sortText.localeCompare(b.sortText) || a.name.localeCompare(b.name)).slice(0, 40)
      const detailed = entries.slice(0, 12).map(e => {
        const d = this.service.getCompletionEntryDetails(fileName, pos, e.name, undefined, undefined, undefined, undefined)
        return clip(displayText(d?.displayParts).replace(/^\((method|property|function|const|alias|let|var)\) /, ''), 200)
      })
      const rest = entries.slice(12).map(e => e.name)
      return `${detailed.join('\n')}${rest.length ? `\n// also: ${rest.join(', ')}` : ''}`
    } finally {
      this.overrides.set(fileName, { text: tsText, version: ++this.version })
    }
  }

  /**
   * Syntactic + semantic diagnostics of the visible file (Vue: its script blocks), with lines mapped back to the
   * visible file. Used for "TS errors as context" and to verify generated next-edit data.
   */
  diagnostics(filePath: string, lang: SymbolLang, visible: string): Diagnostic[] {
    const loaded = this.load(filePath, lang, visible, 0)
    if (!loaded) return []
    // Map TS-text offsets back to visible-file lines (Vue script blocks are concatenated with '\n').
    const segments: { tsStart: number; visStart: number; length: number }[] = []
    if (lang === 'vue') {
      let tsPos = 0
      for (const r of vueScriptRegions(visible)) {
        segments.push({ tsStart: tsPos, visStart: r.start, length: r.end - r.start })
        tsPos += r.end - r.start + 1
      }
    } else {
      segments.push({ tsStart: 0, visStart: 0, length: visible.length })
    }
    const toVisibleLine = (tsOffset: number) => {
      const seg = segments.find(g => tsOffset >= g.tsStart && tsOffset <= g.tsStart + g.length) ?? segments[0]
      const vis = seg.visStart + (tsOffset - seg.tsStart)
      return visible.slice(0, vis).split('\n').length - 1
    }
    const all = [...this.service.getSyntacticDiagnostics(loaded.fileName), ...this.service.getSemanticDiagnostics(loaded.fileName)]
    return all
      .filter(d => d.start !== undefined && (d.category === ts.DiagnosticCategory.Error || d.category === ts.DiagnosticCategory.Warning))
      .map(d => ({
        line: toVisibleLine(d.start!),
        code: d.code,
        message: ts.flattenDiagnosticMessageText(d.messageText, '\n'),
        category: d.category === ts.DiagnosticCategory.Error ? 'error' as const : 'warning' as const,
      }))
  }

  /**
   * If the identifier being typed at the cursor only resolves as a member of
   * the enclosing class (the user forgot `this.`), returns the member name.
   * Uses the language service's own completions: TS offers class members with
   * insertText `this.x` at a lower priority than locals and parameters.
   */
  thisMember(filePath: string, lang: SymbolLang, visible: string, cursor: number): { name: string; identLength: number } | null {
    const before = visible.slice(Math.max(0, cursor - 200), cursor)
    const m = /(^|[^\w$.])([A-Za-z_$][\w$]*)$/.exec(before)
    if (!m || m[2].length < 2) return null
    // Not when declaring something with that name.
    if (/\b(function|const|let|var|class|interface|type|enum|import|export|as|new)\s+[A-Za-z_$][\w$]*$/.test(before)) return null
    const ident = m[2]
    const loaded = this.load(filePath, lang, visible, cursor)
    if (!loaded || loaded.tsCursor === null) return null
    const completions = this.service.getCompletionsAtPosition(loaded.fileName, loaded.tsCursor, { includeCompletionsWithInsertText: true })
    const matching = (completions?.entries ?? []).filter(e => e.name.startsWith(ident) && e.kind !== 'keyword')
    if (!matching.length) return null
    matching.sort((a, b) => a.sortText.localeCompare(b.sortText))
    const best = matching[0]
    return best.insertText?.startsWith('this.') ? { name: best.name, identLength: ident.length } : null
  }
}