Extension: request the next edit right after accepting a suggestion (0.6.4)

Accepting a completion or a next-edit cancels the debounce timer and asks the daemon for
the follow-up immediately; next-edits chain. Setting localAutocomplete.nextEdit.afterAccept.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01J39imRDT8L7xYbWzsVWd8A

CHANGELOG.md

@@ -2,6 +2,36 @@
 
 Extension versions are released as `vscode-v<version>` on the [Releases](https://git.harvmaster.com/Harvmaster/autocomplete-model/releases) page. Models are versioned by their name in [`models/manifest.json`](models/manifest.json).
 
+## vscode 0.6.4
+
+- **Accepting a suggestion asks for the next edit immediately** (setting `nextEdit.afterAccept`, on by default). After
+  Tab on a completion or a next-edit, the follow-up is requested right away instead of after the typing pause, so
+  accepted next-edits chain (rename → next use → …).
+
+## vscode 0.6.3
+
+- **Fix: model choice didn't stick in managed mode.** The daemon's saved state (last model used) overrode both the
+  model picked in the menu and a new default from a release, so it kept restarting on the old model (e.g.
+  `stage-c`). A changed default now wins; switching through the daemon API still persists as before.
+
+## vscode 0.6.2
+
+- **Next-edit is ~3–5× faster:** the next-edit server uses n-gram speculative decoding (the rewritten region is mostly
+  a copy of the current one, so drafted tokens are verified in batches). Same output, since decoding is greedy.
+  Rewriting a 13-line region on a laptop CPU: ~2.1 s → ~0.7 s (~0.4 s when the prompt is cached).
+- Next-edit timings in "Show Last Prompt" now split TS errors / prompt / generation time.
+
+## vscode 0.6.1
+
+- **New default model `stage-d`**, with **`stage-e` for next-edit**. Both download automatically (~530 MB each); the
+  daemon runs `stage-e` in a second llama-server, only for next-edit. `stage-d` closes blocks (`}`) and predicts the
+  next line far better (held-out exact match: block close 0.27 → 0.79, next line 0.01 → 0.18), and sees TypeScript
+  errors near the cursor. Results in [models/README.md](models/README.md).
+- **Fix:** a "no edit" answer from a next-edit model could be applied as deleting the suggested region. llama-server
+  drops the model's no-edit token, so it arrived as empty text. Next-edit models are only enabled from this version on.
+- Manifest: `nextEditModel` (companion model) and `nextEditFormat` (replaces `nextEdit`, which older versions ignore
+  from now on).
+
 ## vscode 0.6.0
 
 - **Next-edit suggestions** (with models trained for it, e.g. the upcoming `stage-e`): after you pause typing, the model may suggest an edit around the cursor, such as removing a leftover `)` or updating a renamed variable's uses. Single-line changes show deleted text struck through and inserted text as ghost text; multi-line changes are highlighted with the diff on hover. **Tab** applies, **Esc** dismisses. Setting: `nextEdit.enabled`.

README.md

@@ -76,6 +76,7 @@ VS Code extension ──HTTP──▶ acd daemon ──▶ llama-server (llama.c
 |---|---|
 | [`models/`](models/) | GGUF models (Git LFS), `manifest.json`, model card, license |
 | [`extensions/vscode/`](extensions/vscode/) | the VS Code extension |
+| [`extensions/zed/`](extensions/zed/) | Zed setup (settings for the daemon's OpenAI-compatible endpoint) |
 | [`daemon/`](daemon/) | `acd`, the completion daemon (bundled into the extension; also runs standalone) |
 | [`shared/`](shared/) | prompt layout + symbol extraction, shared with the training pipeline so prompts match training exactly |
 | [`scripts/`](scripts/) | publish models and extension releases |
@@ -120,7 +121,7 @@ Running the daemon standalone (external mode) is covered in [daemon/README.md](d
 - **TypeScript errors as context:** the model sees nearby compiler errors and is trained on error → fix.
 - **Suggestions on cursor move:** predict what you came to a location to do, not just what you're typing.
 - Token healing for completions that start mid-identifier.
-- Zed support via the daemon's OpenAI-compatible endpoint.
+- Zed: the daemon's OpenAI-compatible endpoint is in place ([extensions/zed](extensions/zed/)); needs testing against real Zed traffic.
 
 ## License
 

daemon/src/config.ts

@@ -18,6 +18,8 @@ export interface ModelProfile {
   members?: boolean
   /** Trained for next-edit prediction (shared/nextedit.ts), enables /v1/next-edit. */
   nextEdit?: boolean
+  /** Another profile (trained with nextEdit) that serves /v1/next-edit for this model, in a second llama-server. */
+  nextEditModel?: string
 }
 
 export interface DaemonConfig {
@@ -30,6 +32,9 @@ export interface DaemonConfig {
     ctx: number
     /** Extra flags, e.g. ["-ngl", "99"] on Apple Silicon, ["-t", "16"] on the laptop CPU. */
     args: string[]
+    /** Extra flags for the next-edit companion server only. Default: n-gram speculative decoding (the output is
+     *  mostly a copy of the region, so drafts from the prompt are accepted; exact with greedy sampling). */
+    nextEditArgs?: string[]
   }
   models: Record<string, ModelProfile>
   defaultModel: string
@@ -58,6 +63,8 @@ export interface DaemonConfig {
   /** Where opted-in edit recordings from editors are appended (the opt-in is the editor setting). */
   recording: { dir: string }
   stateFile: string
+  /** Workspace roots for clients that can't send one (Zed via /v1/completions): files are looked up under these. */
+  workspaceRoots?: string[]
   /** Exit after this many minutes without any request (0 = never). Editors poll every 15 s while open. */
   idleShutdownMinutes?: number
   /** Opaque tag reported by /health (the VS Code extension uses it to replace daemons left over from older versions). */
@@ -69,6 +76,8 @@ export interface RuntimeState {
   enabled: boolean
   model: string
   mode: ContextMode
+  /** config.defaultModel when this state was saved: a changed default (editor model setting, new release) wins. */
+  defaultModel?: string
 }
 
 export function loadConfig(configPath: string): DaemonConfig {
@@ -84,17 +93,19 @@ export function loadConfig(configPath: string): DaemonConfig {
   if (!raw.models[raw.defaultModel]) throw new Error(`defaultModel ${raw.defaultModel} is not in models`)
   for (const [name, m] of Object.entries(raw.models)) {
     if (!m.modes.includes(m.defaultMode)) throw new Error(`model ${name}: defaultMode must be one of its modes`)
+    if (m.nextEditModel && !raw.models[m.nextEditModel]?.nextEdit) throw new Error(`model ${name}: nextEditModel ${m.nextEditModel} must be a profile with nextEdit`)
   }
   return raw
 }
 
 export function loadState(config: DaemonConfig): RuntimeState {
-  const fallback: RuntimeState = { enabled: true, model: config.defaultModel, mode: config.models[config.defaultModel].defaultMode }
+  const fallback: RuntimeState = { enabled: true, model: config.defaultModel, mode: config.models[config.defaultModel].defaultMode, defaultModel: config.defaultModel }
   try {
     const s = JSON.parse(fs.readFileSync(config.stateFile, 'utf8')) as Partial<RuntimeState>
-    const model = s.model && config.models[s.model] ? s.model : fallback.model
+    const keepModel = s.defaultModel === config.defaultModel && s.model && config.models[s.model]
+    const model = keepModel ? s.model! : fallback.model
     const mode = s.mode && config.models[model].modes.includes(s.mode) ? s.mode : config.models[model].defaultMode
-    return { enabled: s.enabled ?? true, model, mode }
+    return { enabled: s.enabled ?? true, model, mode, defaultModel: config.defaultModel }
   } catch {
     return fallback
   }

daemon/src/main.ts

@@ -9,6 +9,7 @@
  *   PUT  /v1/config     { enabled?, model?, mode? }   (changing model restarts llama-server)
  *   POST /v1/feedback   { requestId, event: shown|accepted|partial|rejected, acceptedText? }
  *   POST /v1/edits      { events: [...] }   opted-in edit recordings from editors -> recording.dir
+ *   POST /v1/completions OpenAI-compatible (Zed): Qwen FIM prompt in, completion text out (prompt rebuilt in our layout)
  *   POST /v1/prewarm    same body as /v1/fim: prefill the prompt in the background (first keystroke warm)
  *   POST /v1/next-edit  { repo, path, workspaceRoot?, languageId?, text, cursor } -> { noEdit, startLine, endLine, newText }
  *   GET  /v1/debug/last?n=1   last requests with decoded prompts and raw output
@@ -29,7 +30,8 @@ import { trimSuffixOverlap } from '../../shared/prompt.ts'
 import { symbolLang } from '../../shared/symbol-lang.ts'
 import { SymbolService, mightNeedThis } from './symbols.ts'
 import { EditTracker } from './edits.ts'
-import { NEXT_EDIT_TOKENS } from '../../shared/nextedit.ts'
+import { locateFile, modeFromModel, parseFimPrompt } from './openai-shim.ts'
+import { NEXT_EDIT_TOKENS, parseNextEditOutput } from '../../shared/nextedit.ts'
 
 const { values: args } = parseArgs({ options: { config: { type: 'string', default: 'acd.config.json' } } })
 const config = loadConfig(args.config!)
@@ -45,17 +47,35 @@ function loadModel(name: string) {
   loadError = null
   loading = llama.start(config.models[name].path)
     .catch(e => { loadError = (e as Error).message; console.error(loadError) })
-    .finally(() => { loading = null })
+    .finally(() => { loading = null; if (config.models[name].nextEditModel) void nextEditServer() })
   return loading
 }
 
+// Next-edit companion model (profile.nextEditModel): a second llama-server, started on first use.
+const NEXT_EDIT_SPECULATION = ['--spec-type', 'ngram-simple', '--spec-draft-n-max', '64', '--spec-ngram-simple-size-n', '3']
+const neLlama = new LlamaServer({ ...config.llamaServer, port: config.llamaServer.port + 1, args: [...config.llamaServer.args, ...(config.llamaServer.nextEditArgs ?? NEXT_EDIT_SPECULATION)] })
+let neLoading: Promise<void> | null = null
+/** The server that answers next-edit for the active model, or an error message. */
+async function nextEditServer(): Promise<LlamaServer | string> {
+  const profile = config.models[state.model]
+  if (profile.nextEdit) return loading || !llama.running ? (loadError ?? 'model loading') : llama
+  const companion = profile.nextEditModel
+  if (!companion) return `model ${state.model} is not trained for next-edit`
+  const wanted = config.models[companion].path
+  if (neLlama.model !== wanted && !neLoading) {
+    neLoading = neLlama.start(wanted).catch(e => console.error(`next-edit model ${companion}: ${(e as Error).message}`)).finally(() => { neLoading = null })
+  }
+  if (neLoading) await neLoading
+  return neLlama.running && neLlama.model === wanted ? neLlama : `next-edit model ${companion} failed to start`
+}
+
 let inflight: AbortController | null = null
 
 function describeConfig() {
   return {
     ...state,
     effectiveMode: state.mode,
-    models: Object.fromEntries(Object.entries(config.models).map(([n, m]) => [n, { order: m.order, modes: m.modes, defaultMode: m.defaultMode, nextEdit: m.nextEdit ?? false, errors: m.errors ?? false, contextFormat: m.contextFormat ?? 1 }])),
+    models: Object.fromEntries(Object.entries(config.models).map(([n, m]) => [n, { order: m.order, modes: m.modes, defaultMode: m.defaultMode, nextEdit: !!(m.nextEdit || m.nextEditModel), nextEditModel: m.nextEditModel ?? null, errors: m.errors ?? false, contextFormat: m.contextFormat ?? 1 }])),
     llama: { running: llama.running, loading: loading !== null, error: loadError },
   }
 }
@@ -185,6 +205,44 @@ async function handleFim(body: FimRequest, res: http.ServerResponse, clientGone:
   json(res, 200, response)
 }
 
+/**
+ * OpenAI-compatible /v1/completions for editors with built-in providers (Zed: edit_predictions.provider =
+ * "open_ai_compatible_api", prompt_format "qwen", model "ac:symbols"). The client's prompt is parsed back into
+ * prefix/suffix and rebuilt in our trained layout; the answer comes back in OpenAI text_completion format.
+ */
+async function handleOpenAiCompletion(body: Record<string, unknown>, res: http.ServerResponse, clientGone: AbortSignal) {
+  const prompt = String(body.prompt ?? '')
+  const parsed = parseFimPrompt(prompt) ?? { prefix: prompt, suffix: '' }
+  const located = locateFile(config.workspaceRoots ?? [], parsed.path, parsed.prefix)
+  const fimBody: FimRequest = {
+    requestId: `oai-${Date.now()}`,
+    repo: located ? path.basename(located.root) : parsed.repo ?? 'workspace',
+    path: located?.path ?? parsed.path ?? 'untitled.ts',
+    workspaceRoot: located?.root,
+    languageId: /\.vue$/.test(located?.path ?? parsed.path ?? '') ? 'vue' : 'typescript',
+    prefix: parsed.prefix,
+    suffix: parsed.suffix,
+    mode: modeFromModel(body.model),
+  }
+  // Reuse the normal pipeline by capturing its JSON response.
+  let captured: Record<string, unknown> = {}
+  const fake = { writeHead: () => fake, end: (s: string) => { captured = JSON.parse(s) }, writableEnded: false, destroyed: false } as unknown as http.ServerResponse
+  await handleFim(fimBody, fake, clientGone)
+  const text = String(captured.text ?? '')
+  const out = {
+    id: `cmpl-${Date.now()}`, object: 'text_completion', created: Math.floor(Date.now() / 1000), model: String(body.model ?? state.model),
+    choices: [{ text, index: 0, logprobs: null, finish_reason: 'stop' }],
+    usage: { prompt_tokens: (captured.timings as Record<string, number> | undefined)?.promptTokens ?? 0, completion_tokens: (captured.timings as Record<string, number> | undefined)?.genTokens ?? 0, total_tokens: 0 },
+  }
+  if (body.stream) {
+    res.writeHead(200, { 'content-type': 'text/event-stream', 'cache-control': 'no-cache' })
+    res.write(`data: ${JSON.stringify(out)}\n\n`)
+    res.end('data: [DONE]\n\n')
+    return
+  }
+  json(res, 200, out)
+}
+
 /** Prefill a file's prompt in the background (e.g. on editor focus) so the first keystroke is warm. */
 async function handlePrewarm(body: FimRequest, res: http.ServerResponse) {
   if (!state.enabled || loading || !llama.running) return json(res, 200, { warmed: false })
@@ -205,8 +263,8 @@ async function handleNextEdit(body: NextEditRequest, res: http.ServerResponse, c
   const t0 = performance.now()
   const profile = config.models[state.model]
   if (!state.enabled) return json(res, 200, { noEdit: true, disabled: true })
-  if (!profile.nextEdit) return json(res, 400, { error: `model ${state.model} is not trained for next-edit` })
-  if (loading || !llama.running) return json(res, 503, { error: loadError ?? 'model loading' })
+  const server = await nextEditServer()
+  if (typeof server === 'string') return json(res, server.includes('not trained') ? 400 : 503, { error: server })
   if (typeof body.text !== 'string' || typeof body.cursor !== 'number') return json(res, 400, { error: 'text and cursor are required' })
 
   const o = config.nextEditOptions ?? { regionBefore: 4, regionAfter: 8, maxTokens: 400, beforeTokens: 1024, afterTokens: 384 }
@@ -216,9 +274,11 @@ async function handleNextEdit(body: NextEditRequest, res: http.ServerResponse, c
   const endLine = Math.min(lines.length, cursorLine + o.regionAfter + 1)
   const recentEdits = edits.observe(`${body.workspaceRoot ?? body.repo}/${body.path}`, body.path, body.text, cursorLine)
   const lang = symbolLang(body.languageId ?? '')
+  const tErrors = performance.now()
   const errors = body.workspaceRoot && lang
     ? await symbols.errors({ root: body.workspaceRoot, filePath: path.join(body.workspaceRoot, body.path), lang, text: body.text, cursor: body.cursor, waitMs: 300 })
     : undefined
+  const errorsMs = performance.now() - tErrors
   const ids = builder.buildNextEdit({ repo: body.repo, path: body.path, file: body.text, regionStartLine: startLine, regionEndLine: endLine, cursor: body.cursor, recentEdits: recentEdits || undefined, errors },
     { beforeTokens: o.beforeTokens, afterTokens: o.afterTokens })
 
@@ -226,18 +286,18 @@ async function handleNextEdit(body: NextEditRequest, res: http.ServerResponse, c
   const controller = new AbortController()
   inflight = controller
   clientGone.addEventListener('abort', () => controller.abort())
-  const result = await llama.complete(ids, { nPredict: o.maxTokens, stop: [NEXT_EDIT_TOKENS.outputEnd, ...STOP_STRINGS] }, controller.signal)
+  const result = await server.complete(ids, { nPredict: o.maxTokens, stop: [NEXT_EDIT_TOKENS.outputEnd, ...STOP_STRINGS] }, controller.signal)
   if (inflight === controller) inflight = null
   if (result.stoppedBy === 'aborted') return json(res, 200, { noEdit: true, cancelled: true })
 
   const oldText = lines.slice(startLine, endLine).join('\n') + (endLine < lines.length ? '\n' : '')
   const raw = result.text
-  const newText = raw.startsWith(NEXT_EDIT_TOKENS.noEdit) ? null : raw
-  const noEdit = newText === null || newText === oldText || result.stoppedBy === 'limit'
+  const newText = parseNextEditOutput(raw, oldText)
+  const noEdit = newText === null || result.stoppedBy === 'limit'
   const response = {
     requestId: body.requestId, noEdit, startLine, endLine, oldText, newText: noEdit ? null : newText,
-    recentEdits, errors: errors ?? null, model: state.model,
-    timings: { promptTokens: ids.length, processedTokens: result.timings?.prompt_n ?? null, genTokens: result.timings?.predicted_n ?? null, totalMs: round(performance.now() - t0) },
+    recentEdits, errors: errors ?? null, model: profile.nextEditModel ?? state.model,
+    timings: { promptTokens: ids.length, processedTokens: result.timings?.prompt_n ?? null, genTokens: result.timings?.predicted_n ?? null, errorsMs: round(errorsMs), promptMs: round(result.timings?.prompt_ms ?? 0), genMs: round(result.timings?.predicted_ms ?? 0), totalMs: round(performance.now() - t0) },
   }
   remember({ kind: 'next-edit', requestId: body.requestId, path: body.path, model: state.model, prompt: builder.decode(ids), raw, noEdit, timings: response.timings })
   json(res, 200, response)
@@ -263,7 +323,7 @@ if (config.idleShutdownMinutes) {
   setInterval(() => {
     if (Date.now() - lastActivity > config.idleShutdownMinutes! * 60_000) {
       console.log(`idle for ${config.idleShutdownMinutes} min, shutting down`)
-      void llama.stop().then(() => process.exit(0))
+      void Promise.all([llama.stop(), neLlama.stop()]).then(() => process.exit(0))
     }
   }, 30_000).unref()
 }
@@ -278,7 +338,7 @@ const server = http.createServer(async (req, res) => {
     if (route === 'GET /health') return json(res, 200, { ok: llama.running, loading: loading !== null, model: state.model, error: loadError, instanceTag: config.instanceTag ?? null, pid: process.pid })
     if (route === 'POST /v1/shutdown') {
       json(res, 200, { ok: true })
-      void llama.stop().then(() => process.exit(0))
+      void Promise.all([llama.stop(), neLlama.stop()]).then(() => process.exit(0))
       return
     }
     if (route === 'GET /v1/config') return json(res, 200, describeConfig())
@@ -291,10 +351,14 @@ const server = http.createServer(async (req, res) => {
       const modelChanged = model !== state.model
       Object.assign(state, { model, mode, enabled: body.enabled ?? state.enabled })
       saveState(config, state)
-      if (modelChanged) void loadModel(model)
+      if (modelChanged) {
+        void loadModel(model)
+        if (neLlama.model && neLlama.model !== config.models[config.models[model].nextEditModel ?? '']?.path) void neLlama.stop()
+      }
       return json(res, 200, describeConfig())
     }
     if (route === 'POST /v1/fim') return await handleFim(await readBody(req) as FimRequest, res, gone.signal)
+    if (route === 'POST /v1/completions') return await handleOpenAiCompletion(await readBody(req) as Record<string, unknown>, res, gone.signal)
     if (route === 'POST /v1/prewarm') return await handlePrewarm(await readBody(req) as FimRequest, res)
     if (route === 'POST /v1/next-edit') return await handleNextEdit(await readBody(req) as NextEditRequest, res, gone.signal)
     if (route === 'GET /v1/debug/last') {
@@ -316,8 +380,8 @@ const server = http.createServer(async (req, res) => {
   }
 })
 
-process.on('SIGINT', () => void llama.stop().then(() => process.exit(0)))
-process.on('SIGTERM', () => void llama.stop().then(() => process.exit(0)))
+process.on('SIGINT', () => void Promise.all([llama.stop(), neLlama.stop()]).then(() => process.exit(0)))
+process.on('SIGTERM', () => void Promise.all([llama.stop(), neLlama.stop()]).then(() => process.exit(0)))
 
 server.listen(config.port, config.host, () => {
   console.log(`acd listening on http://${config.host}:${config.port} (model ${state.model}, mode ${state.mode})`)

daemon/src/openai-shim.ts

@@ -0,0 +1,69 @@
+import fs from 'node:fs'
+import path from 'node:path'
+
+/**
+ * Parses a FIM prompt built by another client (e.g. Zed's `open_ai_compatible_api` edit-prediction provider with
+ * prompt_format "qwen") back into prefix / suffix / file path, so the daemon can rebuild the prompt in the layout
+ * our models were trained on (suffix-first, repo header, symbols, healing).
+ *
+ * Accepts, in order of preference:
+ *   [<|repo_name|>repo\n] [<|file_sep|>path\n …] <|fim_prefix|>{prefix}<|fim_suffix|>{suffix}<|fim_middle|>
+ *   <|fim_prefix|><|fim_suffix|>{suffix}<|fim_middle|>{prefix}          (suffix-first)
+ * Returns null for anything else (e.g. a plain completion prompt): the caller then treats it as prefix-only.
+ */
+export function parseFimPrompt(prompt: string): { prefix: string; suffix: string; path?: string; repo?: string } | null {
+  const repo = /<\|repo_name\|>([^\n<]*)\n/.exec(prompt)?.[1]?.trim() || undefined
+  const seps = [...prompt.matchAll(/<\|file_sep\|>([^\n<]*)\n/g)]
+  const filePath = seps.length ? seps.at(-1)![1].trim() || undefined : undefined
+  const psm = /<\|fim_prefix\|>([\s\S]*?)<\|fim_suffix\|>([\s\S]*?)<\|fim_middle\|>\s*$/.exec(prompt)
+  if (psm && !psm[1].includes('<|fim_suffix|>')) {
+    if (psm[1] === '') {
+      // Possibly suffix-first: "<pre><suf>suffix<mid>prefix"
+      const spm = /<\|fim_prefix\|><\|fim_suffix\|>([\s\S]*?)<\|fim_middle\|>([\s\S]*)$/.exec(prompt)
+      if (spm) return { prefix: spm[2], suffix: spm[1], path: filePath, repo }
+    }
+    return { prefix: psm[1], suffix: psm[2], path: filePath, repo }
+  }
+  const spm = /<\|fim_prefix\|><\|fim_suffix\|>([\s\S]*?)<\|fim_middle\|>([\s\S]*)$/.exec(prompt)
+  if (spm) return { prefix: spm[2], suffix: spm[1], path: filePath, repo }
+  return null
+}
+
+/** "ac:symbols" / "ac:simple" / "ac:raw" select the context mode; anything else keeps the current toggle. */
+export function modeFromModel(model: unknown): 'symbols' | 'simple' | 'raw' | undefined {
+  const m = /^ac:(symbols|simple|raw)$/.exec(String(model ?? ''))
+  return m ? (m[1] as 'symbols' | 'simple' | 'raw') : undefined
+}
+
+/**
+ * Finds the workspace for a request: the root containing `relPath` (if the client sent a path), else the root
+ * containing a file whose content starts with the prompt's prefix (Zed may not send a path).
+ */
+export function locateFile(roots: string[], relPath: string | undefined, prefix: string): { root: string; path: string } | null {
+  if (relPath) {
+    for (const root of roots) {
+      if (path.isAbsolute(relPath) && relPath.startsWith(root + path.sep)) return { root, path: path.relative(root, relPath).split(path.sep).join('/') }
+      if (fs.existsSync(path.join(root, relPath))) return { root, path: relPath }
+    }
+  }
+  const probe = prefix.slice(0, 400)
+  if (probe.length < 40) return null
+  for (const root of roots) {
+    const stack = [root]
+    let visited = 0
+    while (stack.length && visited < 5000) {
+      const dir = stack.pop()!
+      for (const e of fs.readdirSync(dir, { withFileTypes: true })) {
+        if (e.name.startsWith('.') || e.name === 'node_modules' || e.name === 'dist') continue
+        const p = path.join(dir, e.name)
+        if (e.isDirectory()) stack.push(p)
+        else if (/\.(ts|tsx|js|jsx|vue|mts|cts)$/.test(e.name) && ++visited) {
+          try {
+            if (fs.readFileSync(p, 'utf8').startsWith(probe)) return { root, path: path.relative(root, p).split(path.sep).join('/') }
+          } catch { /* unreadable */ }
+        }
+      }
+    }
+  }
+  return null
+}

daemon/test/smoke.ts

@@ -124,6 +124,10 @@ if (cfg.models?.['nextedit-plumbing']) {
   for (let i = 0; i < 120 && !(await api('GET', '/health')).body.ok; i++) await new Promise(r => setTimeout(r, 500))
 }
 
+// ---------- OpenAI-compatible endpoint (Zed)
+const oai = (await api('POST', '/v1/completions', { model: 'ac:simple', prompt: "<|fim_prefix|>import { z } from 'zod'\n\nexport const configSchema = z.object({\n  port: z.coerce.number().default(3000),\n})\n\nexport type Config = <|fim_suffix|>\n<|fim_middle|>", max_tokens: 64 })).body
+check('OpenAI-compatible /v1/completions (Zed shim)', oai.object === 'text_completion' && String(oai.choices?.[0]?.text).includes('z.infer'), JSON.stringify(oai.choices?.[0]?.text))
+
 // ---------- edit recording endpoint
 const recorded = await api('POST', '/v1/edits', { events: [{ t: Date.now(), repo: 'smoke', path: 'src/a.ts', version: 2, changes: [{ offset: 0, length: 0, text: 'x' }], snapshot: 'x' }] })
 check('edit recordings are accepted', recorded.body.recorded === 1)
@@ -132,7 +136,7 @@ check('edit recordings are accepted', recorded.body.recorded === 1)
 // Optional typing replay over held-out FIM samples from the (private) training workspace.
 const dataDir = path.resolve(process.env.SMOKE_DATA_DIR ?? '../../data')
 const testFim = path.join(dataDir, 'fim/test-repo.jsonl')
-if (fs.existsSync(testFim)) {
+if (fs.existsSync(testFim) && Number(args['typing-samples']) > 0) {
   const files = new Map(fs.readFileSync(path.join(dataDir, 'corpus/files.jsonl'), 'utf8').trim().split('\n').map(l => JSON.parse(l)).map(r => [r.id, r]))
   const rank = (id: string) => crypto.createHash('sha256').update(`smoke:${id}`).digest('hex')
   const samples = fs.readFileSync(testFim, 'utf8').trim().split('\n').map(l => JSON.parse(l))

extensions/vscode/package.json

@@ -3,7 +3,7 @@
   "displayName": "Local Autocomplete",
   "description": "Inline code completion from a local fine-tuned model via the acd daemon",
   "publisher": "harvmaster",
-  "version": "0.6.0",
+  "version": "0.6.4",
   "license": "MIT",
   "engines": {
     "vscode": "^1.90.0"
@@ -218,6 +218,11 @@
           "default": 700,
           "description": "Pause after typing before asking for a next-edit suggestion."
         },
+        "localAutocomplete.nextEdit.afterAccept": {
+          "type": "boolean",
+          "default": true,
+          "description": "After you accept a completion or a next-edit suggestion, immediately ask for the next edit (without waiting for the typing pause). Accepted next-edits chain: a rename suggests the next use, and so on."
+        },
         "localAutocomplete.triggerOnCursorMove": {
           "type": "boolean",
           "default": false,

extensions/vscode/src/extension.ts

@@ -259,6 +259,7 @@ const modelSupportsNextEdit = () => !!(daemonConfig && daemonConfig.models[daemo
 // ---------- next-edit prediction (models trained for it, e.g. Stage E)
 const nextEdit = new NextEditController(s => {
   api('POST', '/v1/feedback', { event: 'next-edit-accepted', lines: s.newText.split('\n').length }).catch(() => {})
+  nextEditAfterAccept('accepted-next-edit')
 })
 let nextEditTimer: NodeJS.Timeout | null = null
 let nextEditAbort: AbortController | null = null
@@ -281,6 +282,15 @@ async function requestNextEdit(editor: vscode.TextEditor, reason: string) {
   }
 }
 
+/** Accepting a suggestion (next-edit or completion) asks for the follow-up right away instead of after the typing pause. */
+function nextEditAfterAccept(reason: string) {
+  const editor = vscode.window.activeTextEditor
+  if (!editor || !modelSupportsNextEdit() || !settings().get<boolean>('nextEdit.afterAccept', true)) return
+  if (nextEditTimer) clearTimeout(nextEditTimer)
+  // Next tick: the accepted text and the cursor move have been applied by then.
+  nextEditTimer = setTimeout(() => void requestNextEdit(editor, reason), 0)
+}
+
 function scheduleNextEdit(editor: vscode.TextEditor, reason: string, delay: number) {
   if (nextEditTimer) clearTimeout(nextEditTimer)
   nextEditTimer = setTimeout(() => void requestNextEdit(editor, reason), delay)
@@ -535,6 +545,7 @@ export async function activate(ctx: vscode.ExtensionContext) {
     vscode.commands.registerCommand('localAutocomplete.showLog', () => log.show()),
     vscode.commands.registerCommand('localAutocomplete.accepted', (requestId: string, text: string) => {
       api('POST', '/v1/feedback', { requestId, event: 'accepted', chars: text.length }).catch(() => {})
+      nextEditAfterAccept('accepted-completion')
     }),
     vscode.workspace.onDidChangeConfiguration(e => {
       if (!e.affectsConfiguration('localAutocomplete')) return

extensions/vscode/src/managed-core.ts

@@ -30,13 +30,19 @@ export interface ModelEntry {
   contextFormat?: 1 | 2
   errors?: boolean
   members?: boolean
+  /** Old (≤ 0.6.0) flag: ignored, because those daemons applied an empty "no edit" output as a deletion. */
   nextEdit?: boolean
+  /** Trained for next-edit (shared/nextedit.ts format version); read since 0.6.1. */
+  nextEditFormat?: 1
+  /** Companion model (another manifest entry) that serves next-edit for this one; downloaded alongside it. */
+  nextEditModel?: string
 }
 
-const profileOf = (m: ModelEntry) => ({
+const profileOf = (m: ModelEntry, available: Set<string>) => ({
   order: m.order, modes: m.modes, defaultMode: m.defaultMode,
   ...(m.contextFormat ? { contextFormat: m.contextFormat } : {}), ...(m.errors ? { errors: true } : {}),
-  ...(m.members !== undefined ? { members: m.members } : {}), ...(m.nextEdit ? { nextEdit: true } : {}),
+  ...(m.members !== undefined ? { members: m.members } : {}), ...(m.nextEditFormat ? { nextEdit: true } : {}),
+  ...(m.nextEditModel && available.has(m.nextEditModel) ? { nextEditModel: m.nextEditModel } : {}),
 })
 
 export interface Manifest {
@@ -229,12 +235,12 @@ export async function prepare(storage: string, daemonDir: string, extensionVersi
     if (!entry) throw new Error(`model "${s.model || manifest.default}" is not in the manifest (available: ${manifest.models.map(m => m.name).join(', ')})`)
     const file = await ensureFile(entry.url, path.join(storage, 'models', entry.file), entry.sha256, `model ${entry.name}`, progress, signal)
     defaultModel = entry.name
-    models[entry.name] = { path: file, ...profileOf(entry) }
-    // Other models already on disk stay switchable from the menu.
-    for (const m of manifest.models) {
-      const f = path.join(storage, 'models', m.file)
-      if (m.name !== entry.name && fs.existsSync(`${f}.verified`)) models[m.name] = { path: f, ...profileOf(m) }
-    }
+    const companion = entry.nextEditModel ? manifest.models.find(m => m.name === entry.nextEditModel && m.nextEditFormat) : undefined
+    if (companion) await ensureFile(companion.url, path.join(storage, 'models', companion.file), companion.sha256, `next-edit model ${companion.name}`, progress, signal)
+    // The selected model, its companion and other models already on disk (switchable from the menu).
+    const onDisk = manifest.models.filter(m => m.name === entry.name || fs.existsSync(`${path.join(storage, 'models', m.file)}.verified`))
+    const available = new Set(onDisk.filter(m => m.nextEditFormat).map(m => m.name))
+    for (const m of onDisk) models[m.name] = { path: path.join(storage, 'models', m.file), ...profileOf(m, available) }
   }
 
   const binary = s.llamaServerPath ? s.llamaServerPath.replace(/^~(?=$|\/)/, os.homedir()) : await ensureLlamaServer(storage, llamaCpp, progress, signal)

extensions/vscode/test/managed.e2e.mjs

@@ -36,20 +36,29 @@ const step = (name, fn) => fn().then(r => { console.log(`PASS  ${name}`); return
 
 // 1. Simulate an interrupted model download: first 100 MB already on disk.
 fs.mkdirSync(path.join(storage, 'models'), { recursive: true })
-const modelSrc = path.join(releaseDir, 'stage-c-q8_0.gguf')
+const modelSrc = path.join(releaseDir, 'stage-d-q8_0.gguf')
 const fd = fs.openSync(modelSrc, 'r'); const buf = Buffer.alloc(100e6); fs.readSync(fd, buf, 0, buf.length, 0); fs.closeSync(fd)
-fs.writeFileSync(path.join(storage, 'models', 'stage-c-q8_0.gguf.part'), buf)
+fs.writeFileSync(path.join(storage, 'models', 'stage-d-q8_0.gguf.part'), buf)
 
 let t0 = Date.now()
 const prepared = await step('first run: manifest, resumed + verified model download, llama.cpp download + unpack', () => core.prepare(storage, daemonDir, '0.4.0-test', settings, progress))
-console.log(`      ${((Date.now() - t0) / 1000).toFixed(1)} s; range requests: ${rangeRequests}; model verified: ${fs.existsSync(path.join(storage, 'models', 'stage-c-q8_0.gguf.verified'))}`)
+console.log(`      ${((Date.now() - t0) / 1000).toFixed(1)} s; range requests: ${rangeRequests}; model verified: ${fs.existsSync(path.join(storage, 'models', 'stage-d-q8_0.gguf.verified'))}`)
 assert.ok(rangeRequests >= 1, 'download should have resumed with a Range request')
 const cfg = JSON.parse(fs.readFileSync(prepared.configPath, 'utf8'))
 assert.ok(fs.existsSync(cfg.llamaServer.binary), 'llama-server unpacked')
+await step('next-edit companion downloaded and wired (stage-d → stage-e)', async () => {
+  assert.equal(cfg.defaultModel, 'stage-d')
+  assert.equal(cfg.models['stage-d'].nextEditModel, 'stage-e')
+  assert.equal(cfg.models['stage-e'].nextEdit, true)
+  assert.ok(fs.existsSync(path.join(storage, 'models', 'stage-e-q8_0.gguf.verified')))
+})
 
 await step('daemon starts on the editor runtime', async () => assert.equal(await core.ensureRunning(storage, daemonDir, prepared, settings.port, progress), 'started'))
 const fim = await (await fetch('http://127.0.0.1:7111/v1/fim', { method: 'POST', body: JSON.stringify({ repo: 'e2e', path: 'src/a.ts', prefix: "import { z } from 'zod'\nexport const s = z.object({ a: z.string() })\nexport type S = ", suffix: '\n' }) })).json()
 await step(`completion through the managed daemon (${JSON.stringify(fim.text)})`, async () => assert.match(fim.text, /z\.infer<typeof s>/))
+const neText = "export function f(s: string) {\n  const t = s.trim()\n  return parse(t))\n}\n"
+const ne = await (await fetch('http://127.0.0.1:7111/v1/next-edit', { method: 'POST', body: JSON.stringify({ repo: 'e2e', path: 'src/f.ts', text: neText, cursor: neText.indexOf('(t)') + 1 }) })).json()
+await step(`next-edit through the companion model (${JSON.stringify(ne.newText ?? ne.error ?? ne)})`, async () => assert.ok(ne.newText?.includes('return parse(t)\n')))
 
 t0 = Date.now()
 const again = await step('second window: no re-download, daemon reused', async () => {

extensions/zed/README.md

@@ -0,0 +1,32 @@
+# Zed
+
+Zed can't load third-party edit-prediction extensions, but its built-in `open_ai_compatible_api` provider can call
+the daemon's OpenAI-compatible endpoint. The daemon parses Zed's Qwen-format FIM prompt back into prefix/suffix and
+rebuilds it in the layout our models were trained on (suffix-first, repo header, symbols, token healing).
+
+1. Run the daemon standalone (see [daemon/README.md](../../daemon/README.md)). Add your project folders to
+   `"workspaceRoots"` in its config, so symbols mode and the `this.` repair can find the files. Zed may not send
+   the file path; the daemon then locates the file by content.
+2. In Zed's `settings.json`:
+
+```jsonc
+{
+  "edit_predictions": {
+    "provider": "open_ai_compatible_api",
+    "open_ai_compatible_api": {
+      "api_url": "http://127.0.0.1:7071/v1/completions",
+      "model": "ac:symbols",          // the context-mode toggle: "ac:symbols" | "ac:simple"
+      "prompt_format": "qwen"
+    }
+  }
+}
+```
+
+The mode can also be switched with `node daemon/src/acctl.ts mode simple`.
+
+**Status: untested against real Zed traffic.** The parser accepts prefix-first and suffix-first Qwen prompts, with
+or without `<|repo_name|>`/`<|file_sep|>` headers, and is covered by the daemon smoke test with a synthetic request.
+Still to check with a logging proxy:
+- whether Zed sends the file path
+- its prefix/suffix truncation
+- whether it streams and cancels superseded requests

models/NOTICE

@@ -2,6 +2,7 @@ These model files are modified versions of Qwen2.5-Coder-0.5B (https://huggingfa
 Copyright (c) Alibaba Cloud, licensed under the Apache License, Version 2.0 (see LICENSE).
 
 Modifications: full-parameter fine-tuning for suffix-first fill-in-the-middle completion
-(and, for stage-c, a repository/symbol-context header), followed by Q8_0 GGUF conversion.
+(and, from stage-c on, a repository/symbol-context header; stage-d adds TypeScript errors; stage-e adds
+next-edit prediction), followed by Q8_0 GGUF conversion.
 Fine-tuning data: permissively licensed (MIT) open-source TypeScript/Vue repositories and
 the publisher's own code. Secrets were scrubbed from the training data.

models/README.md

@@ -5,6 +5,8 @@ Q8_0 GGUF fill-in-the-middle models for the [Local Autocomplete](../extensions/v
 
 | model | description | prompt order | context modes | size |
 |---|---|---|---|---|
+| `stage-d` | Suffix-first FIM with symbols (incl. namespace imports) + TS errors; better next-line, block closing and multi-line completions | spm | symbols, simple | 531 MB |
+| `stage-e` | Next-edit prediction (rewrites a region around the cursor or answers "no edit"); also does FIM, slightly below stage-d | spm | symbols, simple | 531 MB |
 | `stage-c` | Suffix-first FIM with repo/symbol-context headers (default) | spm | symbols, simple | 531 MB |
 | `stage-b` | Suffix-first FIM without context headers (previous stage) | spm | raw | 531 MB |
 
@@ -21,6 +23,16 @@ runs it locally with llama.cpp. The files are stored with Git LFS; download link
 | Qwen2.5-Coder-1.5B (base) | 0.366 | 0.623 | 0.882 |
 | `stage-b` | 0.391 | 0.689 | 0.952 |
 | `stage-c` (with symbol context) | 0.397 | 0.681 | 0.954 |
+| `stage-c` + token healing (daemon default) | 0.405 | 0.689 | |
+| `stage-d` (symbols + TS errors, healed) | **0.426** | **0.700** | 0.954 |
+| `stage-e` (symbols + TS errors, healed) | 0.412 | 0.691 | 0.962 |
+
+New behaviours (886 further held-out samples), exact match, stage-c → stage-d: closing a block 0.27 → **0.79**,
+next line 0.01 → **0.18**, rest of block 0.08 → 0.11, member access 0.63 → 0.66.
+
+**Next-edit** (`stage-e`, 445 held-out samples): 80% exact region rewrites (edit similarity 0.99) where an edit
+was due, and "no edit" on 100% of samples where none was. `stage-d` uses `stage-e` as its next-edit companion
+(`nextEditModel` in the manifest): the extension downloads both and the daemon runs it in a second llama-server.
 
 The base models were evaluated in their native prefix-first order (they can't do suffix-first). Suffix-first prompts
 let the editor reuse the KV cache while typing: about 2 tokens re-processed per keystroke, ~90–100 ms end-to-end on a

models/manifest.json

@@ -1,7 +1,43 @@
 {
   "schema": 1,
-  "default": "stage-c",
+  "default": "stage-d",
   "models": [
+    {
+      "name": "stage-d",
+      "file": "stage-d-q8_0.gguf",
+      "url": "stage-d-q8_0.gguf",
+      "sha256": "8edaef02ae9e794929af0267f0f47f600824454b08df7d386e386eec1123c839",
+      "size": 531067552,
+      "order": "spm",
+      "modes": [
+        "symbols",
+        "simple"
+      ],
+      "defaultMode": "symbols",
+      "contextFormat": 2,
+      "errors": true,
+      "members": true,
+      "nextEditModel": "stage-e",
+      "description": "Suffix-first FIM with symbols (incl. namespace imports) + TS errors; better next-line, block closing and multi-line completions"
+    },
+    {
+      "name": "stage-e",
+      "file": "stage-e-q8_0.gguf",
+      "url": "stage-e-q8_0.gguf",
+      "sha256": "44d5fe87043f2fd3f8fba39c07a653ec7bdc228f83a2da4a5638be3ab77a962f",
+      "size": 531067552,
+      "order": "spm",
+      "modes": [
+        "symbols",
+        "simple"
+      ],
+      "defaultMode": "symbols",
+      "contextFormat": 2,
+      "errors": true,
+      "members": true,
+      "nextEditFormat": 1,
+      "description": "Next-edit prediction (rewrites a region around the cursor or answers \"no edit\"); also does FIM, slightly below stage-d"
+    },
     {
       "name": "stage-c",
       "file": "stage-c-q8_0.gguf",

scripts/make-model-release.mjs

@@ -23,9 +23,10 @@ const out = path.join(repo, 'models')
 // Prompt layout each model was trained with (the extension only offers matching context modes).
 const profiles = {
   // Stage E: Stage D + next-edit prediction (shared/nextedit.ts)
-  'stage-e': { order: 'spm', modes: ['symbols', 'simple'], defaultMode: 'symbols', contextFormat: 2, errors: true, members: true, nextEdit: true, description: 'Suffix-first FIM with symbols + TS errors, plus next-edit prediction (rewrites a region around the cursor)' },
+  // nextEditFormat (not the old nextEdit flag): extensions ≤ 0.6.0 must not use it (they applied "no edit" as a deletion).
+  'stage-e': { order: 'spm', modes: ['symbols', 'simple'], defaultMode: 'symbols', contextFormat: 2, errors: true, members: true, nextEditFormat: 1, description: 'Next-edit prediction (rewrites a region around the cursor or answers "no edit"); also does FIM, slightly below stage-d' },
   // Stage D: next-line / block-close / multi-line fixes, namespace-import symbols (format 2), TS errors in the header
-  'stage-d': { order: 'spm', modes: ['symbols', 'simple'], defaultMode: 'symbols', contextFormat: 2, errors: true, members: true, description: 'Suffix-first FIM with symbols (incl. namespace imports) + TS errors; better next-line, block closing and multi-line completions' },
+  'stage-d': { order: 'spm', modes: ['symbols', 'simple'], defaultMode: 'symbols', contextFormat: 2, errors: true, members: true, nextEditModel: 'stage-e', description: 'Suffix-first FIM with symbols (incl. namespace imports) + TS errors; better next-line, block closing and multi-line completions' },
   'stage-c': { order: 'spm', modes: ['symbols', 'simple'], defaultMode: 'simple', description: 'Suffix-first FIM with repo/symbol-context headers (default)' },
   'stage-b': { order: 'spm', modes: ['raw'], defaultMode: 'raw', description: 'Suffix-first FIM without context headers (previous stage)' },
 }
@@ -71,7 +72,8 @@ fs.writeFileSync(path.join(out, 'NOTICE'), `These model files are modified versi
 Copyright (c) Alibaba Cloud, licensed under the Apache License, Version 2.0 (see LICENSE).
 
 Modifications: full-parameter fine-tuning for suffix-first fill-in-the-middle completion
-(and, for stage-c, a repository/symbol-context header), followed by Q8_0 GGUF conversion.
+(and, from stage-c on, a repository/symbol-context header; stage-d adds TypeScript errors; stage-e adds
+next-edit prediction), followed by Q8_0 GGUF conversion.
 Fine-tuning data: permissively licensed (MIT) open-source TypeScript/Vue repositories and
 the publisher's own code. Secrets were scrubbed from the training data.
 `)
@@ -103,6 +105,16 @@ runs it locally with llama.cpp. The files are stored with Git LFS; download link
 | Qwen2.5-Coder-1.5B (base) | 0.366 | 0.623 | 0.882 |
 | \`stage-b\` | 0.391 | 0.689 | 0.952 |
 | \`stage-c\` (with symbol context) | 0.397 | 0.681 | 0.954 |
+| \`stage-c\` + token healing (daemon default) | 0.405 | 0.689 | |
+| \`stage-d\` (symbols + TS errors, healed) | **0.426** | **0.700** | 0.954 |
+| \`stage-e\` (symbols + TS errors, healed) | 0.412 | 0.691 | 0.962 |
+
+New behaviours (886 further held-out samples), exact match, stage-c → stage-d: closing a block 0.27 → **0.79**,
+next line 0.01 → **0.18**, rest of block 0.08 → 0.11, member access 0.63 → 0.66.
+
+**Next-edit** (\`stage-e\`, 445 held-out samples): 80% exact region rewrites (edit similarity 0.99) where an edit
+was due, and "no edit" on 100% of samples where none was. \`stage-d\` uses \`stage-e\` as its next-edit companion
+(\`nextEditModel\` in the manifest): the extension downloads both and the daemon runs it in a second llama-server.
 
 The base models were evaluated in their native prefix-first order (they can't do suffix-first). Suffix-first prompts
 let the editor reuse the KV cache while typing: about 2 tokens re-processed per keystroke, ~90–100 ms end-to-end on a

shared/nextedit.ts

@@ -124,11 +124,16 @@ export function buildNextEditPrompt(input: NextEditInput, encode: Encode, ids: N
   out.push(ids.fileSep, ...encode(`${input.path}\n`))
 
   // Before the region: last lines that fit; after: first lines that fit.
+  // (Window to ~6 chars/token first, then drop whole lines in proportion to the excess: a few encodes, not one per line.)
   let before = input.file.slice(0, rStart)
+  if (before.length > beforeTokens * 6) before = before.slice(before.indexOf('\n', before.length - beforeTokens * 6) + 1)
   let beforeIds = encode(before)
   while (beforeIds.length > beforeTokens && before.includes('\n')) {
-    before = before.slice(before.indexOf('\n') + 1)
-    beforeIds = beforeIds.length > beforeTokens * 2 ? encode(before.slice(-beforeTokens * 6)) : encode(before)
+    const excess = beforeIds.length - beforeTokens  // undershoot (×0.8), then single lines: same cut as line-by-line
+    const drop = excess > 32 ? Math.floor(0.8 * before.length * excess / beforeIds.length) : 1
+    const nl = before.indexOf('\n', drop - 1)
+    before = nl < 0 ? '' : before.slice(nl + 1)
+    beforeIds = encode(before)
   }
   let after = input.file.slice(rEnd)
   let afterIds = encode(after.slice(0, afterTokens * 8))
@@ -149,3 +154,12 @@ export function buildNextEditPrompt(input: NextEditInput, encode: Encode, ids: N
 export function nextEditTarget(newRegion: string | null, encode: Encode, ids: NextEditTokenIds): number[] {
   return newRegion === null ? [ids.noEdit, ids.outputEnd] : [...encode(newRegion), ids.outputEnd]
 }
+
+/**
+ * Model output → rewritten region, or null for "no edit". llama-server leaves special tokens out of `content`, so the
+ * NO_EDIT token usually arrives as an empty string; never treat empty output as "replace the region with nothing".
+ */
+export function parseNextEditOutput(raw: string, oldText: string): string | null {
+  if (raw.startsWith(NEXT_EDIT_TOKENS.noEdit) || raw.trim() === '' || raw === oldText) return null
+  return raw
+}