offgrid-ai 0.4.1 → 0.4.3

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "offgrid-ai",
-  "version": "0.4.1",
+  "version": "0.4.3",
   "description": "Privacy-first CLI for running local LLMs — discover, configure, run, benchmark",
   "author": "Eeshan Srivastava (https://eeshans.com)",
   "type": "module",

package/src/autodetect.mjs

@@ -109,7 +109,7 @@ export function computeFlags(capabilities, modelPath, mmprojPath, draftModelPath
 
   // MTP flags
   if (mtp) {
-    argv.push("--spec-type", "draft-mtp", "--spec-draft-n-max", "2");
+    argv.push("--spec-type", "draft-mtp", "--spec-draft-n-max", "4");
   }
 
   return { flags, argv };

package/src/backends.mjs

@@ -27,7 +27,7 @@ export const BACKENDS = {
     defaultPort: LLAMA_CPP_PORT,
     defaultBaseUrl: baseUrlFor({ port: LLAMA_CPP_PORT }),
     needsCommandFile: true,
-    scanModels: () => scanGgufModels(),
+    scanModels: async () => (await scanGgufModels()).models,
   },
   "llama-cpp-mtp": {
     id: "llama-cpp-mtp",
@@ -38,7 +38,7 @@ export const BACKENDS = {
     defaultPort: LLAMA_CPP_MTP_PORT,
     defaultBaseUrl: baseUrlFor({ port: LLAMA_CPP_MTP_PORT }),
     needsCommandFile: true,
-    scanModels: () => scanGgufModels(),
+    scanModels: async () => (await scanGgufModels()).models,
   },
   "ollama": {
     id: "ollama",

package/src/cli.mjs

@@ -1,7 +1,8 @@
 import { totalmem } from "node:os";
 import { existsSync, statSync, rmSync } from "node:fs";
+import { basename } from "node:path";
 import { ensureDirs, findLlamaServer, hasHomebrew, DATA_DIR } from "./config.mjs";
-import { scanGgufModels } from "./scan.mjs";
+import { scanGgufModels, matchDrafter } from "./scan.mjs";
 import { createProfileFromModel, normalizeProfile, sanitizeProfileId } from "./profiles.mjs";
 import { readProfile, saveProfile, deleteProfile, loadProfiles, readCommandArgv } from "./profiles.mjs";
 import { backendFor, BACKENDS } from "./backends.mjs";
@@ -9,7 +10,7 @@ import { startServer, stopProfile, waitForReady, serverReady, serverMatchesProfi
 import { syncPiConfig, removeFromPiConfig, hasPiModel, launchPi, hasPi } from "./harness-pi.mjs";
 import { tailFriendly } from "./logs.mjs";
 import { estimateMemory } from "./estimate.mjs";
-import { pc, formatBytes, renderRows, renderSection, renderCard, humanCapabilitySummary, statusText, startInteractive, createPrompt, parseOptions } from "./ui.mjs";
+import { pc, formatBytes, renderRows, renderSection, renderCard, humanCapabilitySummary, startInteractive, createPrompt, parseOptions } from "./ui.mjs";
 import { checkForUpdate, currentPackageVersion, detectInvocation, updateCommand, runUpdateCommand } from "./updates.mjs";
 import { removeInstallerPathEntries } from "./shell-path.mjs";
 import { configureLocalProfile } from "./profile-setup.mjs";
@@ -78,7 +79,7 @@ export async function mainFlow() {
 
   // 1. Check what backends are available
   const llamaBinary = await findLlamaServer();
-  const ggufModels = await scanGgufModels();
+  const { models: ggufModels } = await scanGgufModels();
   const managedModels = await scanManagedModels();
   const profiles = await loadProfiles();
   const hasAnyBackend = llamaBinary || managedModels.some((m) => m.models.length > 0);
@@ -206,17 +207,17 @@ async function runCommand(argv) {
 }
 
 async function loadModelCatalog() {
-  const [profiles, ggufModels, managedModels] = await Promise.all([
+  const [profiles, { models: ggufModels, drafters }, managedModels] = await Promise.all([
     loadProfiles(),
     scanGgufModels(),
     scanManagedModels(),
   ]);
-  return normalizeCatalog({ profiles, ggufModels, managedModels });
+  return normalizeCatalog({ profiles, ggufModels, drafters, managedModels });
 }
 
 function normalizeCatalog(catalog) {
   if (catalog.newModels && catalog.managedItems) return catalog;
-  const { profiles, ggufModels, managedModels } = catalog;
+  const { profiles, ggufModels, drafters, managedModels } = catalog;
   const profiledPaths = new Set(profiles.map((p) => p.modelPath).filter(Boolean));
   const newModels = ggufModels.filter((m) => !profiledPaths.has(m.path));
   const managedItems = [];
@@ -228,10 +229,10 @@ function normalizeCatalog(catalog) {
       if (!profiledAliases.has(`${backendId}:${model.id}`)) managedItems.push({ model, backendId });
     }
   }
-  return { profiles, ggufModels, managedModels, newModels, managedItems };
+  return { profiles, ggufModels, drafters, managedModels, newModels, managedItems };
 }
 
-async function printModelCatalog({ profiles, newModels, managedItems }, items = modelCatalogItems({ profiles, newModels, managedItems })) {
+async function printModelCatalog({ profiles, newModels, managedItems, drafters }, items = modelCatalogItems({ profiles, newModels, managedItems })) {
   const itemNumber = (predicate) => {
     const index = items.findIndex(predicate);
     return index === -1 ? "  " : String(index + 1).padStart(2, " ");
@@ -243,13 +244,14 @@ async function printModelCatalog({ profiles, newModels, managedItems }, items =
 
   const fileMissingCount = profiles.filter((p) => isProfileFileMissing(p)).length;
 
+  const summaryBorder = fileMissingCount > 0 ? pc.red : newModels.length > 0 ? pc.yellow : pc.dim;
   console.log("\n" + renderCard("Your local AI workspace", renderRows([
-    ["Ready to chat", pc.green(`${profiles.length} saved setup${profiles.length === 1 ? "" : "s"}`)],
-    ["Need setup", newModels.length > 0 ? pc.yellow(`${newModels.length} downloaded model${newModels.length === 1 ? "" : "s"}`) : pc.green("none")],
-    ["Running now", runningProfilesNow.length > 0 ? pc.green(String(runningProfilesNow.length)) : pc.dim("none")],
-    ["File missing", fileMissingCount > 0 ? pc.red(`${fileMissingCount} setup${fileMissingCount === 1 ? "" : "s"}`) : pc.green("none")],
-    ["Next step", fileMissingCount > 0 ? pc.red("Remove or fix missing setups") : profiles.length > 0 ? "Start chatting" : newModels.length > 0 ? "Set up a downloaded model" : "Download a model"],
-  ]), { formatBorder: fileMissingCount > 0 ? pc.yellow : pc.cyan }));
+    ["Setups", `${profiles.length} saved`],
+    ["Need setup", newModels.length > 0 ? pc.yellow(`${newModels.length} model${newModels.length === 1 ? "" : "s"}`) : pc.dim("none")],
+    ["Running", runningProfilesNow.length > 0 ? pc.green(String(runningProfilesNow.length)) : pc.dim("none")],
+    ["File missing", fileMissingCount > 0 ? pc.red(`${fileMissingCount} setup${fileMissingCount === 1 ? "" : "s"}`) : pc.dim("none")],
+    ["Next step", fileMissingCount > 0 ? pc.red("Remove or fix missing setups") : profiles.length > 0 ? "Start chatting" : newModels.length > 0 ? pc.yellow("Set up a downloaded model") : "Download a model"],
+  ]), { formatBorder: summaryBorder }));
 
   console.log("\n" + pc.bold("Ready to chat"));
   if (profiles.length === 0) {
@@ -260,18 +262,19 @@ async function printModelCatalog({ profiles, newModels, managedItems }, items =
       const piConfigured = await hasPiModel(profile);
       const fileMissing = isProfileFileMissing(profile);
       const num = itemNumber((item) => item.type === "profile" && item.profile.id === profile.id);
-      console.log(profileCatalogCard(num, profile, { running, piConfigured, fileMissing }));
+      console.log(profileCatalogCard(num, profile, { running, piConfigured, fileMissing, drafters }));
     }
   }
 
-  console.log("\n" + pc.bold("Downloaded, needs one-time setup"));
+  console.log("\n" + pc.bold("Needs one-time setup"));
   if (newModels.length === 0) {
-    console.log(renderCard("All set", "Every downloaded local model already has a saved setup.", { formatBorder: pc.green }));
+    console.log(renderCard("All set", "Every downloaded local model already has a saved setup.", { formatBorder: pc.dim }));
   } else {
     for (const model of newModels.slice(0, 20)) {
       const caps = detectCapabilities(model.path, model.mmprojPath);
+      const drafter = matchDrafter(model.path, drafters);
       const num = itemNumber((item) => item.type === "new" && item.model.path === model.path);
-      console.log(downloadedModelCard(num, model, caps));
+      console.log(downloadedModelCard(num, model, caps, { mtpAvailable: caps.mtp || Boolean(drafter), drafter }));
     }
     if (newModels.length > 20) console.log(pc.dim(`  ... and ${newModels.length - 20} more`));
   }
@@ -296,46 +299,65 @@ function isProfileFileMissing(profile) {
   return !existsSync(profile.modelPath);
 }
 
-function profileCatalogCard(num, profile, { running, piConfigured, fileMissing }) {
+function profileCatalogCard(num, profile, { running, piConfigured, fileMissing, drafters }) {
   const backend = backendFor(profile.backend);
   const caps = profile.capabilities ?? {};
   let status;
   if (fileMissing) {
     status = pc.red("File missing");
   } else if (running) {
-    status = statusText("running", "Running now");
+    status = pc.green("Running now");
   } else {
-    status = statusText("ready", "Ready to chat");
+    status = "Ready";
   }
-  return renderCard(`${num}. ${profile.label}`, renderRows([
+  const border = fileMissing ? pc.red : running ? pc.green : pc.dim;
+  const mtpDrafter = profile.drafterPath
+    ? pc.green("enabled")
+    : (drafters ? matchDrafter(profile.modelPath, drafters) : null)
+      ? pc.yellow("available — re-setup to enable")
+      : (caps.architecture === "gemma4")
+        ? pc.yellow("drafter not found")
+        : null;
+  const rows = [
     ["Status", status],
     ["Good for", fileMissing ? pc.red("Model file not found") : humanCapabilitySummary(caps)],
-    ["Pi", piConfigured ? pc.green("Synced") : pc.yellow("Needs sync")],
-    ["Runs with", backend.label],
-  ]), { formatBorder: fileMissing ? pc.red : running ? pc.green : pc.cyan });
+  ];
+  if (mtpDrafter) rows.push(["MTP", mtpDrafter]);
+  rows.push(["Pi", piConfigured ? pc.dim("synced") : pc.yellow("Needs sync")]);
+  rows.push(["Runs with", backend.label]);
+  return renderCard(`${num}. ${profile.label}`, renderRows(rows), { formatBorder: border });
 }
 
-function downloadedModelCard(num, model, caps) {
+function downloadedModelCard(num, model, caps, { mtpAvailable } = {}) {
+  const mtpRow = mtpAvailable
+    ? [["MTP", pc.green("Drafter found — setup will enable 2× speedup")]]
+    : (caps.architecture === "gemma4" || caps.architecture === "gemma4")
+      ? [["MTP", pc.yellow("Drafter not found — download MTP model for 2× speedup")]]
+      : [];
   return renderCard(`${num}. ${model.label}`, renderRows([
-    ["Status", statusText("warning", "Needs one-time setup")],
+    ["Status", pc.yellow("Needs setup")],
     ["Good for", humanCapabilitySummary(caps)],
+    ...mtpRow,
     ["Size", formatBytes(model.sizeBytes)],
   ]), { formatBorder: pc.yellow });
 }
 
 function managedModelCard(num, model, backend) {
   return renderCard(`${num}. ${model.label}`, renderRows([
-    ["Status", statusText("info", `Local model via ${backend.label}`)],
+    ["Status", pc.dim(`Via ${backend.label}`)],
     ["Runs with", backend.label],
-    ["Model ID", pc.cyan(model.id)],
+    ["Model ID", model.id],
     ...(model.quant ? [["Size/type", model.quant]] : []),
-  ]), { formatBorder: pc.magenta });
+  ]), { formatBorder: pc.dim });
 }
 
-function modelCatalogItems({ profiles, newModels, managedItems }) {
+function modelCatalogItems({ profiles, newModels, managedItems, drafters }) {
   return [
     ...profiles.map((profile) => ({ type: "profile", profile, label: profile.label, hint: `${profile.modelAlias} · ${profile.baseUrl}`, fileMissing: isProfileFileMissing(profile) })),
-    ...newModels.map((model) => ({ type: "new", model, label: model.label, hint: `${model.quant ?? "GGUF"} · ${formatBytes(model.sizeBytes)}` })),
+    ...newModels.map((model) => {
+      const drafter = matchDrafter(model.path, drafters);
+      return { type: "new", model, label: model.label, hint: `${model.quant ?? "GGUF"} · ${formatBytes(model.sizeBytes)}`, drafter };
+    }),
     ...managedItems.map(({ model, backendId }) => ({ type: "managed", model, backendId, label: model.label, hint: BACKENDS[backendId].label })),
   ];
 }
@@ -366,7 +388,7 @@ async function handleCatalogAction(prompt, action, item) {
   if (action === "inspect") {
     if (item.type === "profile") return await printProfileDetails(await readProfile(item.profile.id));
     if (item.type === "managed") return printManagedModelDetails(item.model, BACKENDS[item.backendId]);
-    return printGgufModelDetails(item.model);
+    return printGgufModelDetails(item.model, item.drafter);
   }
 
   if (action === "setup") {
@@ -376,7 +398,7 @@ async function handleCatalogAction(prompt, action, item) {
       await saveProfile(profile);
       return await syncPiConfig(profile);
     }
-    const profile = await createProfileFromModel(item.model);
+    const profile = await createProfileFromModel(item.model, null, item.drafter?.path);
     const configured = await configureLocalProfile(prompt, profile);
     if (!configured) return;
     await saveProfile(configured);
@@ -391,7 +413,7 @@ async function handleCatalogAction(prompt, action, item) {
       await syncPiConfig(profile);
       return await runProfile(profile);
     }
-    const profile = await createProfileFromModel(item.model);
+    const profile = await createProfileFromModel(item.model, null, item.drafter?.path);
     const configured = await configureLocalProfile(prompt, profile);
     if (!configured) return;
     await saveProfile(configured);
@@ -410,26 +432,39 @@ async function printProfileDetails(profile) {
   const fileMissing = !isManaged && isProfileFileMissing(profile);
   const statusRow = fileMissing
     ? pc.red("File missing")
-    : running ? pc.green("Running now") : pc.green("Ready to chat");
-  console.log("\n" + renderSection("Model overview", renderRows([
+    : running ? pc.green("Running now") : "Ready";
+  const mtpStatus = profile.drafterPath
+    ? pc.green(`MTP enabled (drafter: ${basename(profile.drafterPath)})`)
+    : (profile.capabilities?.architecture === "gemma4")
+      ? pc.yellow("MTP available — download a drafter model to enable 2× speedup")
+      : null;
+  const overviewRows = [
     ["Name", pc.bold(profile.label)],
     ["Status", statusRow],
     ["Good for", fileMissing ? pc.red("Model file not found — remove or fix this setup") : humanCapabilitySummary(profile.capabilities ?? {})],
-    ["Pi", piConfigured ? pc.green("Synced") : pc.yellow("Needs sync")],
-    ["Server", fileMissing ? pc.red(profile.baseUrl) : pc.green(profile.baseUrl)],
-  ])));
+    ["Pi", piConfigured ? pc.dim("synced") : pc.yellow("Needs sync")],
+    ["Server", fileMissing ? pc.red(profile.baseUrl) : profile.baseUrl],
+  ];
+  if (mtpStatus) overviewRows.push(["MTP", mtpStatus]);
+  console.log("\n" + renderSection("Model overview", renderRows(overviewRows)));
 
-  console.log("\n" + renderSection("Model details", renderRows([
-    ["Setup ID", pc.cyan(profile.id)],
+  const detailRows = [
+    ["Setup ID", profile.id],
     ["Runs with", backend.label],
-    ["Model alias", pc.cyan(profile.modelAlias)],
+    ["Model alias", profile.modelAlias],
     ...(profile.capabilities ? [["Detected", capabilitySummary(profile.capabilities)]] : []),
-    ...(!isManaged ? [
+  ];
+  if (!isManaged) {
+    detailRows.push(
       ["Local file", fileMissing ? pc.red(`${profile.modelPath} (not found)`) : profile.modelPath ?? "unknown"],
       ["Vision file", profile.mmprojPath ? (existsSync(profile.mmprojPath) ? profile.mmprojPath : pc.red(`${profile.mmprojPath} (not found)`)) : "none"],
       ["Model size", profile.modelPath && existsSync(profile.modelPath) ? formatBytes(statSync(profile.modelPath).size) : "unknown"],
-    ] : []),
-  ]), { columns: 110 }));
+    );
+    if (profile.drafterPath) {
+      detailRows.push(["Drafter", existsSync(profile.drafterPath) ? profile.drafterPath : pc.red(`${profile.drafterPath} (not found)`)]);
+    }
+  }
+  console.log("\n" + renderSection("Model details", renderRows(detailRows), { columns: 110 }));
 
   if (fileMissing) {
     console.log("\n" + pc.red("⚠ This model's file is no longer on disk. Remove this setup or move the file back."));
@@ -441,20 +476,33 @@ async function printProfileDetails(profile) {
   }
 }
 
-function printGgufModelDetails(model) {
+function printGgufModelDetails(model, drafter) {
   const caps = detectCapabilities(model.path, model.mmprojPath);
-  console.log("\n" + renderSection("Downloaded model", renderRows([
+  const mtpAvailable = caps.mtp || Boolean(drafter);
+  const mtpRow = mtpAvailable
+    ? [["MTP", pc.green("Drafter found — setup will enable 2× speedup")]]
+    : (caps.architecture === "gemma4")
+      ? [["MTP", pc.yellow("Drafter not found — download MTP model for 2× speedup")]]
+      : [];
+  const overviewRows = [
     ["Name", pc.bold(model.label)],
     ["Status", pc.yellow("Needs one-time setup")],
     ["Good for", humanCapabilitySummary(caps)],
-    ["Size", formatBytes(model.sizeBytes)],
-  ])));
-  console.log("\n" + renderSection("Model details", renderRows([
+  ];
+  if (mtpRow.length) overviewRows.push(...mtpRow);
+  overviewRows.push(["Size", formatBytes(model.sizeBytes)]);
+  console.log("\n" + renderSection("Downloaded model", renderRows(overviewRows)));
+  const detailRows = [
     ["Local file", model.path],
     ["Vision file", model.mmprojPath ?? "none"],
     ["Detected", capabilitySummary(caps)],
     ["Quant", model.quant ?? "unknown"],
-  ]), { columns: 110 }));
+  ];
+  if (drafter) {
+    detailRows.push(["Drafter", drafter.path]);
+    detailRows.push(["Drafter size", formatBytes(drafter.sizeBytes)]);
+  }
+  console.log("\n" + renderSection("Model details", renderRows(detailRows), { columns: 110 }));
 }
 
 function printManagedModelDetails(model, backend) {
@@ -659,9 +707,9 @@ async function statusCommand() {
   if (running.length === 0) {
     console.log(renderCard("Status", renderRows([
       ["Running now", pc.dim("none")],
-      ["Ready setups", profiles.length > 0 ? pc.green(String(profiles.length)) : pc.dim("none")],
-      ["Next step", profiles.length > 0 ? "Run offgrid-ai to start chatting" : "Run offgrid-ai to set up a model"],
-    ]), { formatBorder: profiles.length > 0 ? pc.cyan : pc.yellow }));
+      ["Ready setups", profiles.length > 0 ? String(profiles.length) : pc.dim("none")],
+      ["Next step", profiles.length > 0 ? "Run offgrid-ai to start chatting" : pc.yellow("Run offgrid-ai to set up a model")],
+    ]), { formatBorder: pc.dim }));
     return;
   }
 
@@ -870,7 +918,7 @@ async function onboardFlow() {
     console.log(pc.green("✓ Pi found"));
 
     // 4. Model backends — at least one is mandatory
-    const ggufModels = await scanGgufModels();
+    const { models: ggufModels } = await scanGgufModels();
     const managedModels = await scanManagedModels();
     const totalManaged = managedModels.reduce((sum, m) => sum + m.models.length, 0);
     const hasModels = ggufModels.length > 0 || totalManaged > 0;

package/src/profile-setup.mjs

@@ -43,12 +43,14 @@ export async function configureLocalProfile(prompt, profile) {
   console.log(pc.dim("Sampling defaults are shown for transparency; you can edit command.json later if needed.\n"));
 
   if (caps.mtp) {
+    const drafterInfo = configured.drafterPath ? `\n  Drafter: ${configured.drafterPath}` : "";
     console.log(renderSection("MTP detected", renderRows([
       ["Backend", "llama.cpp MTP"],
       ["Port", String(LLAMA_CPP_MTP_PORT)],
-      ["Flags", "--spec-type draft-mtp --spec-draft-n-max 2"],
+      ["Flags", `--spec-type draft-mtp --spec-draft-n-max 4${configured.drafterPath ? " --spec-draft-model <drafter>" : ""}`],
     ])));
-    const useMtp = await prompt.yesNo("Use MTP speculative decoding flags?", true);
+    if (drafterInfo) console.log(pc.dim(drafterInfo));
+    const useMtp = await prompt.yesNo("Use MTP speculative decoding?", true);
     configured = useMtp ? applyMtpDefaults(configured) : removeMtpDefaults(configured);
   }
 
@@ -113,15 +115,19 @@ export function applyRuntimeFlagOverrides(profile, overrides) {
 
 function applyMtpDefaults(profile) {
   const flags = { ...profile.flags, port: LLAMA_CPP_MTP_PORT };
-  return applyProfileFlags({ ...profile, backend: "llama-cpp-mtp", providerId: "llama-cpp-mtp" }, flags, {
-    values: { "--spec-type": "draft-mtp", "--spec-draft-n-max": 2 },
-  });
+  const edits = {
+    values: { "--spec-type": "draft-mtp", "--spec-draft-n-max": 4 },
+  };
+  if (profile.drafterPath) {
+    edits.values["--spec-draft-model"] = profile.drafterPath;
+  }
+  return applyProfileFlags({ ...profile, backend: "llama-cpp-mtp", providerId: "llama-cpp-mtp" }, flags, edits);
 }
 
 function removeMtpDefaults(profile) {
   const flags = { ...profile.flags, port: LLAMA_CPP_PORT };
   return applyProfileFlags({ ...profile, backend: "llama-cpp", providerId: "llama-cpp" }, flags, {
-    remove: ["--spec-type", "--spec-draft-n-max"],
+    remove: ["--spec-type", "--spec-draft-n-max", "--spec-draft-model"],
   });
 }
 

package/src/profiles.mjs

@@ -131,22 +131,29 @@ export function normalizeProfile(profile) {
 
 // ── Auto-create profile from a discovered model ────────────────────────────
 
-export async function createProfileFromModel(model, backendId) {
+export async function createProfileFromModel(model, backendId, drafterPath) {
   const { detectCapabilities } = await import("./autodetect.mjs");
   const caps = detectCapabilities(model.path, model.mmprojPath);
-  const backend = backendId ?? (caps.mtp ? "llama-cpp-mtp" : "llama-cpp");
-  const id = slugFromLabel(model.label);
-  const { flags, argv } = computeFlags(caps, model.path, model.mmprojPath, null);
+  // If a drafter is provided, this model supports MTP regardless of filename
+  const hasMtp = caps.mtp || Boolean(drafterPath);
+  const backend = backendId ?? (hasMtp ? "llama-cpp-mtp" : "llama-cpp");
+  const { flags, argv } = computeFlags(
+    { ...caps, mtp: hasMtp },
+    model.path,
+    model.mmprojPath,
+    drafterPath ?? null,
+  );
 
   return normalizeProfile({
-    id,
+    id: slugFromLabel(model.label),
     label: model.label,
     backend,
     providerId: backend,
     modelAlias: model.aliasSuggestion,
     modelPath: model.path,
     mmprojPath: model.mmprojPath,
-    capabilities: summarizeCapabilities(caps),
+    drafterPath: drafterPath ?? null,
+    capabilities: summarizeCapabilities({ ...caps, mtp: hasMtp }),
     preset: null, // no presets — auto-detected
     flags,
     commandArgv: argv,

package/src/scan.mjs

@@ -2,47 +2,128 @@ import { statSync } from "node:fs";
 import { readdir } from "node:fs/promises";
 import { basename, dirname, join } from "node:path";
 import { getModelScanDirs } from "./config.mjs";
+import { readGgufMetadata } from "./gguf.mjs";
+
+// ── Scan for GGUF models and MTP drafters ────────────────────────────────
 
 export async function scanGgufModels(dirs) {
   const scanDirs = dirs ?? await getModelScanDirs();
   const allModels = [];
+  const allDrafters = [];
 
   for (const root of scanDirs) {
-    const models = await scanOneDir(root);
+    const { models, drafters } = await scanOneDir(root);
     allModels.push(...models);
+    allDrafters.push(...drafters);
   }
 
   // Deduplicate by path
   const seen = new Set();
-  return allModels.filter((m) => {
+  const models = allModels.filter((m) => {
     if (seen.has(m.path)) return false;
     seen.add(m.path);
     return true;
   }).sort((a, b) => a.label.localeCompare(b.label));
+
+  const drafterSeen = new Set();
+  const drafters = allDrafters.filter((d) => {
+    if (drafterSeen.has(d.path)) return false;
+    drafterSeen.add(d.path);
+    return true;
+  }).sort((a, b) => a.label.localeCompare(b.label));
+
+  return { models, drafters };
 }
 
 async function scanOneDir(root) {
   const files = await findFiles(root, (path) => path.toLowerCase().endsWith(".gguf"));
   const mmprojs = files.filter((path) => basename(path).toLowerCase().includes("mmproj"));
-  const models = files.filter((path) => !basename(path).toLowerCase().includes("mmproj"));
+  const candidates = files.filter((path) => !basename(path).toLowerCase().includes("mmproj"));
+
+  const models = [];
+  const drafters = [];
 
-  return models.map((path) => {
+  for (const path of candidates) {
     const dir = dirname(path);
     const mmprojPath = mmprojs.find((candidate) => dirname(candidate) === dir) ?? null;
     const name = basename(path).replace(/\.gguf$/i, "");
-    return {
-      path,
-      mmprojPath,
-      label: labelFromName(name),
-      aliasSuggestion: aliasFromName(name),
-      quant: quantFromName(name),
-      sizeBytes: statSync(path).size,
-      backend: "llama-cpp",
-      source: "local-gguf",
-    };
-  });
+    const sizeBytes = statSync(path).size;
+
+    // Read GGUF metadata to detect drafter architecture
+    const meta = safeReadGgufMetadata(path);
+    const architecture = typeof meta["general.architecture"] === "string" ? meta["general.architecture"] : null;
+
+    if (architecture === "gemma4_assistant") {
+      // This is an MTP drafter model, not a main model
+      drafters.push({
+        path,
+        label: labelFromName(name),
+        aliasSuggestion: aliasFromName(name),
+        quant: quantFromName(name),
+        sizeBytes,
+        architecture,
+        targetHint: drafterTargetHint(name),
+        backend: "llama-cpp",
+        source: "local-gguf",
+      });
+    } else {
+      models.push({
+        path,
+        mmprojPath,
+        label: labelFromName(name),
+        aliasSuggestion: aliasFromName(name),
+        quant: quantFromName(name),
+        sizeBytes,
+        backend: "llama-cpp",
+        source: "local-gguf",
+      });
+    }
+  }
+
+  return { models, drafters };
 }
 
+// ── Match drafters to target models ────────────────────────────────────
+
+// Map a drafter filename to a regex that matches its target model filenames.
+// e.g. "mtp-gemma-4-31B-it" matches "gemma-4-31B-it*"
+//      "gemma-4-31B-it-assistant-Q8_0" matches "gemma-4-31B-it*"
+//      "gemma-4-12b-it-MTP-Q8_0" matches "gemma-4-12b-it*"
+function drafterTargetHint(name) {
+  const lower = name.toLowerCase();
+  // Strip common prefixes/suffixes to get the base model identifier
+  let base = lower
+    .replace(/^mtp[-_]/i, "")
+    .replace(/[-_]mtp$/i, "")
+    .replace(/[-_]assistant([-_].*)?$/i, "")
+    .replace(/[-_]draft([-_].*)?$/i, "");
+  // Remove quant suffix for matching (Q4_K_M, Q8_0, UD-Q4_K_XL, etc.)
+  base = base.replace(/[-_]q\d_k_[a-z]+$/i, "").replace(/[-_]q\d_[01]$/i, "").replace(/[-_]ud-[a-z0-9_]+$/i, "");
+  return base;
+}
+
+// Find a drafter that matches a given target model path/name.
+// Returns the drafter object or null.
+export function matchDrafter(targetPath, drafters) {
+  if (!drafters || drafters.length === 0) return null;
+  const targetName = basename(targetPath).replace(/\.gguf$/i, "").toLowerCase();
+  // Remove quant suffix from target for matching
+  const targetBase = targetName
+    .replace(/[-_]q\d_k_[a-z]+$/i, "")
+    .replace(/[-_]q\d_[01]$/i, "")
+    .replace(/[-_]ud-[a-z0-9_]+$/i, "");
+
+  for (const drafter of drafters) {
+    const hint = drafter.targetHint;
+    if (targetBase.startsWith(hint) || targetBase === hint) {
+      return drafter;
+    }
+  }
+  return null;
+}
+
+// ── Helpers ─────────────────────────────────────────────────────────────
+
 async function findFiles(root, predicate) {
   const result = [];
   async function walk(dir) {
@@ -75,4 +156,12 @@ function aliasFromName(name) {
 
 function quantFromName(name) {
   return name.match(/(Q\d_K_[A-Z]+|Q\d_[01]|UD-[A-Z0-9_]+)/)?.[1];
+}
+
+function safeReadGgufMetadata(path) {
+  try {
+    return readGgufMetadata(path);
+  } catch {
+    return {};
+  }
 }