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 = Omit, '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 void>() private cache = new Map() private lastForFile = new Map() private inflight = new Set() private warmed = new Set() 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(kind: K, payload: Payload): Promise { 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(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 { 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(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(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|= 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)) }