autocomplete-model/daemon/src/symbols.ts

import { Worker } from 'node:worker_threads'
import crypto from 'node:crypto'
import type { SymbolsResult as BaseSymbolsResult } from '../../shared/symbols.ts'
export type SymbolsResult = BaseSymbolsResult & { errors?: string }
import type { SymbolLang } from '../../shared/symbol-lang.ts'
import type { WorkerRequest, WorkerResponse } from './symbols-worker.ts'

type Payload<K extends WorkerRequest['kind']> = Omit<Extract<WorkerRequest, { kind: K }>, 'id' | 'kind'>

/**
 * Symbol context for `symbols` mode, with caching that keeps prompts stable.
 *
 * Context is cached per anchor: file + enclosing top-level declaration + the
 * file's import block. While you type inside one function the header is
 * identical, so llama-server's KV cache keeps hitting. When the anchor changes
 * we wait briefly for fresh context; if the language service is slower (e.g.
 * still building the program), the previous context for the file is used
 * (stale-while-revalidate) and the fresh one lands in the cache for the next keystroke.
 */
export class SymbolService {
  private worker: Worker
  private nextId = 1
  private pending = new Map<number, (r: WorkerResponse) => void>()
  private cache = new Map<string, SymbolsResult | null>()
  private lastForFile = new Map<string, SymbolsResult | null>()
  private inflight = new Set<string>()
  private warmed = new Set<string>()

  constructor(tokenizer: { json: string; config: string }) {
    // Source runs load the .ts worker (Node type stripping); the bundled extension build ships symbols-worker.mjs.
    const workerFile = import.meta.url.endsWith('.ts') ? './symbols-worker.ts' : './symbols-worker.mjs'
    this.worker = new Worker(new URL(workerFile, import.meta.url), { workerData: { tokenizerJson: tokenizer.json, tokenizerConfig: tokenizer.config } })
    this.worker.on('message', (r: WorkerResponse) => { this.pending.get(r.id)?.(r); this.pending.delete(r.id) })
    this.worker.on('error', e => console.error('symbols worker error:', e))
    this.worker.unref()
  }

  private call<K extends WorkerRequest['kind']>(kind: K, payload: Payload<K>): Promise<WorkerResponse> {
    const id = this.nextId++
    return new Promise(resolve => {
      this.pending.set(id, resolve)
      this.worker.postMessage({ id, kind, ...payload })
    })
  }

  /** Start building the language service for a workspace before it's needed. */
  warm(root: string) {
    if (this.warmed.has(root)) return
    this.warmed.add(root)
    void this.call('warm', { root })
  }

  async context(opts: { root: string; filePath: string; lang: SymbolLang; prefix: string; suffix: string; budget: number; members: boolean; waitMs: number; format?: 1 | 2; errors?: boolean }):
    Promise<{ result: SymbolsResult | null; status: 'hit' | 'fresh' | 'stale' | 'none' }> {
    // Errors go stale as you type; refresh them at most every 20 s (a header change costs a cold prefill).
    const errorsBucket = opts.errors ? `|e${Math.floor(Date.now() / 20_000)}` : ''
    const key = `${anchorKey(opts.filePath, opts.prefix, opts.members)}|f${opts.format ?? 1}${errorsBucket}`
    if (this.cache.has(key)) {
      const result = this.cache.get(key)!
      this.lastForFile.set(opts.filePath, result)
      return { result, status: 'hit' }
    }
    const request = this.inflight.has(key) ? null : this.call('context', {
      root: opts.root, filePath: opts.filePath, lang: opts.lang, visible: opts.prefix + opts.suffix, cursor: opts.prefix.length, budget: opts.budget, members: opts.members,
      format: opts.format ?? 1, errors: opts.errors ?? false,
    })
    if (request) {
      this.inflight.add(key)
      void request.then(r => {
        this.inflight.delete(key)
        if (r.error) console.error('symbols:', r.error)
        this.remember(key, (r.result as SymbolsResult | null) ?? null)
      })
      const fresh = await Promise.race([request, new Promise<null>(r => setTimeout(() => r(null), opts.waitMs))])
      if (fresh && !fresh.error) {
        const result = (fresh.result as SymbolsResult | null) ?? null
        this.lastForFile.set(opts.filePath, result)
        return { result, status: 'fresh' }
      }
    }
    const stale = this.lastForFile.get(opts.filePath)
    return stale !== undefined ? { result: stale, status: 'stale' } : { result: null, status: 'none' }
  }

  private remember(key: string, result: SymbolsResult | null) {
    this.cache.set(key, result)
    if (this.cache.size > 500) this.cache.delete(this.cache.keys().next().value!)
  }

  /** TS errors near the cursor (formatted), for next-edit prompts. */
  async errors(opts: { root: string; filePath: string; lang: SymbolLang; text: string; cursor: number; waitMs: number }): Promise<string | undefined> {
    const r = await Promise.race([
      this.call('errors', { root: opts.root, filePath: opts.filePath, lang: opts.lang, visible: opts.text, cursor: opts.cursor }),
      new Promise<null>(r => setTimeout(() => r(null), opts.waitMs)),
    ])
    return (r?.result as string | undefined) ?? undefined
  }

  async thisMember(opts: { root: string; filePath: string; lang: SymbolLang; prefix: string; suffix: string; waitMs: number }) {
    const r = await Promise.race([
      this.call('thisMember', { root: opts.root, filePath: opts.filePath, lang: opts.lang, visible: opts.prefix + opts.suffix, cursor: opts.prefix.length }),
      new Promise<null>(r => setTimeout(() => r(null), opts.waitMs)),
    ])
    return (r?.result as { name: string; identLength: number } | null) ?? null
  }
}

const TOP_LEVEL = /^(export\s|async\s|function[\s*]|class\s|const\s|let\s|var\s|interface\s|type\s|enum\s|abstract\s|declare\s|<script|<template|<style|@)/

/** file + enclosing top-level declaration (line index) + hash of the import lines (+ receiver when member context is on). */
export function anchorKey(filePath: string, prefix: string, members: boolean): string {
  const lines = prefix.split('\n')
  let anchor = 0
  for (let i = lines.length - 1; i >= 0; i--) if (TOP_LEVEL.test(lines[i])) { anchor = i; break }
  const imports = crypto.createHash('sha1').update(lines.filter(l => /^\s*import\s/.test(l)).join('\n')).digest('hex').slice(0, 12)
  const receiver = members ? (/([\w$.()[\]'"]+)\.[\w$]*$/.exec(prefix.slice(-200))?.[1] ?? '') : ''
  return `${filePath}|${anchor}|${imports}|${receiver}`
}

/** Cheap pre-check before asking the language service whether `this.` is missing. */
export function mightNeedThis(prefix: string): boolean {
  return /(^|[^\w$.])[A-Za-z_$][\w$]+$/.test(prefix.slice(-100)) && !/\b(function|const|let|var|class|interface|type|enum|import|export|as|new)\s+[A-Za-z_$][\w$]*$/.test(prefix.slice(-100))
}