From f1f9d2706171b47ea82578e06ce1c0bb27ad3b28 Mon Sep 17 00:00:00 2001 From: KunN-21 Date: Fri, 10 Jul 2026 16:00:15 +0700 Subject: [PATCH] feat(headroom): add extras detection and install UI (#2403) - add Headroom extras status + install endpoints - show Headroom version + code/ml extras in Token Saver UI - fix Windows interpreter selection to read from env with headroom-ai --- .../dashboard/token-saver/TokenSaverClient.js | 144 ++++++++++++++++++ src/app/api/headroom/extras/route.js | 30 ++++ src/lib/headroom/detect.js | 74 ++++++++- src/lib/headroom/process.js | 51 ++++++- tests/unit/headroom-detect.test.js | 74 ++++++++- 5 files changed, 368 insertions(+), 5 deletions(-) create mode 100644 src/app/api/headroom/extras/route.js diff --git a/src/app/(dashboard)/dashboard/token-saver/TokenSaverClient.js b/src/app/(dashboard)/dashboard/token-saver/TokenSaverClient.js index e77194d5..74b14d83 100644 --- a/src/app/(dashboard)/dashboard/token-saver/TokenSaverClient.js +++ b/src/app/(dashboard)/dashboard/token-saver/TokenSaverClient.js @@ -24,6 +24,15 @@ export default function TokenSaverClient() { useState(false); const [headroomActionLoading, setHeadroomActionLoading] = useState(false); const [headroomActionError, setHeadroomActionError] = useState(""); + const [headroomExtras, setHeadroomExtras] = useState({ + version: null, + extras: { code: false, ml: false }, + available: ["code", "ml"], + loading: false, + }); + const [pendingExtras, setPendingExtras] = useState([]); + const [extrasActionLoading, setExtrasActionLoading] = useState(false); + const [extrasActionError, setExtrasActionError] = useState(""); const [cavemanEnabled, setCavemanEnabled] = useState(false); const [cavemanLevel, setCavemanLevel] = useState("full"); const [ponytailEnabled, setPonytailEnabled] = useState(false); @@ -102,6 +111,39 @@ export default function TokenSaverClient() { }); const data = await res.json(); setHeadroomStatus({ ...data, loading: false }); + if (!data?.installed) { + setHeadroomExtras({ + version: null, + extras: { code: false, ml: false }, + available: ["code", "ml"], + loading: false, + }); + setPendingExtras([]); + return; + } + try { + const er = await fetch("/api/headroom/extras", { + headers: { "Cache-Control": "no-store" }, + }); + if (!er.ok) throw new Error("extras status failed"); + const ed = await er.json(); + setHeadroomExtras((s) => ({ + ...s, + version: ed.version ?? null, + extras: ed.extras || { code: false, ml: false }, + available: ed.available || ["code", "ml"], + loading: false, + })); + setPendingExtras([]); + } catch { + setHeadroomExtras({ + version: null, + extras: { code: false, ml: false }, + available: ["code", "ml"], + loading: false, + }); + setPendingExtras([]); + } } catch { setHeadroomStatus({ installed: false, @@ -109,6 +151,13 @@ export default function TokenSaverClient() { python: null, loading: false, }); + setHeadroomExtras({ + version: null, + extras: { code: false, ml: false }, + available: ["code", "ml"], + loading: false, + }); + setPendingExtras([]); } }, []); @@ -137,6 +186,37 @@ export default function TokenSaverClient() { } }, [refreshHeadroomStatus]); + const togglePendingExtra = (extra) => { + setPendingExtras((cur) => + cur.includes(extra) ? cur.filter((e) => e !== extra) : [...cur, extra] + ); + }; + + const handleInstallExtras = useCallback(async () => { + if (pendingExtras.length === 0) return; + setExtrasActionLoading(true); + setExtrasActionError(""); + try { + const res = await fetch("/api/headroom/extras", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ extras: pendingExtras }), + }); + const data = await res.json().catch(() => ({})); + if (!res.ok) throw new Error(data.error || "Install failed"); + setHeadroomExtras((s) => ({ + ...s, + version: data.version ?? s.version, + extras: data.extras || s.extras, + })); + setPendingExtras([]); + } catch (e) { + setExtrasActionError(e.message); + } finally { + setExtrasActionLoading(false); + } + }, [pendingExtras]); + const handleCavemanLevel = (level) => { setCavemanLevel(level); patchSetting({ cavemanLevel: level }); @@ -257,6 +337,70 @@ export default function TokenSaverClient() { onChange={() => handleHeadroomEnabled(!headroomEnabled)} /> + {headroomStatus.installed && ( +
+
+ + Compression extras + {headroomExtras.version ? ` · v${headroomExtras.version}` : ""}: + + {headroomExtras.available.map((extra) => { + const installed = !!headroomExtras.extras[extra]; + const pending = pendingExtras.includes(extra); + return ( + + ); + })} + {pendingExtras.length > 0 && ( + + )} +
+ {extrasActionError && ( +

{extrasActionError}

+ )} +

+ Default install is [proxy] only (SmartCrusher for + JSON). Adding [code] enables AST compression + (Python/JS/TS/Go/Rust/Java/C/C++/Perl). Adding [ml]{" "} + enables the Kompress-v2 HF model for prose/agentic traces but + adds ~1 GB (torch + huggingface-hub). +

+
+ )}

diff --git a/src/app/api/headroom/extras/route.js b/src/app/api/headroom/extras/route.js new file mode 100644 index 00000000..4c1d7230 --- /dev/null +++ b/src/app/api/headroom/extras/route.js @@ -0,0 +1,30 @@ +import { NextResponse } from "next/server"; +import { findPython310, getInstalledHeadroomExtras, HEADROOM_COMPRESSION_EXTRAS } from "@/lib/headroom/detect"; +import { installHeadroomExtras } from "@/lib/headroom/process"; + +export const dynamic = "force-dynamic"; + +export async function GET() { + try { + const python = findPython310(); + const status = getInstalledHeadroomExtras(python); + return NextResponse.json({ + available: HEADROOM_COMPRESSION_EXTRAS, + ...status, + }); + } catch (error) { + return NextResponse.json({ error: error.message }, { status: 500 }); + } +} + +export async function POST(req) { + try { + const body = await req.json().catch(() => ({})); + const requested = Array.isArray(body?.extras) ? body.extras : []; + const result = await installHeadroomExtras(requested); + return NextResponse.json(result); + } catch (error) { + const status = error.code === "NOT_INSTALLED" || error.code === "NO_PYTHON" ? 400 : 500; + return NextResponse.json({ error: error.message, code: error.code || null }, { status }); + } +} diff --git a/src/lib/headroom/detect.js b/src/lib/headroom/detect.js index ac64dc87..91b245f7 100644 --- a/src/lib/headroom/detect.js +++ b/src/lib/headroom/detect.js @@ -1,6 +1,21 @@ -import { execSync } from "child_process"; +import { execFileSync, execSync } from "child_process"; import path from "path"; +// Extras that improve headroom compression quality. `proxy` is the base; +// `code` adds tree-sitter AST compression; `ml` adds Kompress-v2 HF model. +// Other `[all]` extras (image, voice, otel, reports, evals, ...) are not +// useful for the 9router proxy use case, so we don't track them here. +export const HEADROOM_COMPRESSION_EXTRAS = ["code", "ml"]; + +// Marker packages that each extra pulls in. Detected from `pip list --format=json` +// so one call can answer both the installed version and active extras. +const EXTRA_MARKERS = { + code: ["tree-sitter", "tree-sitter-language-pack"], + ml: ["torch", "huggingface-hub"], +}; + +const HEADROOM_PIP_TIMEOUT_MS = 8000; + const IS_WIN = process.platform === "win32"; const WHICH_CMD = IS_WIN ? "where" : "which"; @@ -49,6 +64,9 @@ export function findHeadroomBinary() { } // Find a Python interpreter >= 3.10 (headroom-ai requires it). Returns null if none. +// On Windows, `python3` and `python` can point at different envs. Prefer the +// first candidate that can also see the installed `headroom-ai` package so the +// dashboard probes and install action operate on the same interpreter as the CLI. export function findPython310() { for (const candidate of PYTHON_CANDIDATES) { try { @@ -60,8 +78,18 @@ export function findPython310() { const match = ver.match(/(\d+)\.(\d+)/); if (!match) continue; const [major, minor] = [parseInt(match[1], 10), parseInt(match[2], 10)]; - if (major > MIN_VERSION[0] || (major === MIN_VERSION[0] && minor >= MIN_VERSION[1])) { + if (!(major > MIN_VERSION[0] || (major === MIN_VERSION[0] && minor >= MIN_VERSION[1]))) continue; + if (!IS_WIN) return candidate; + try { + execFileSync(candidate, ["-m", "pip", "show", "headroom-ai"], { + stdio: ["ignore", "pipe", "ignore"], + windowsHide: true, + timeout: HEADROOM_PIP_TIMEOUT_MS, + env: { ...process.env, PATH: EXTENDED_PATH }, + }); return candidate; + } catch { + // Keep scanning on Windows until an interpreter sees headroom-ai. } } catch { // candidate not present, try next @@ -98,5 +126,45 @@ export async function getHeadroomStatus(url) { const installed = Boolean(path); const running = await probeProxyRunning(url); const localUrl = isLoopbackHeadroomUrl(url); - return { installed, path, running, python, localUrl, canStart: installed && localUrl }; + const extrasStatus = installed ? getInstalledHeadroomExtras(python) : { installed: false, version: null, extras: { code: false, ml: false } }; + return { + installed, + path, + running, + python, + localUrl, + canStart: installed && localUrl, + version: extrasStatus.version, + extras: extrasStatus.extras, + }; +} + +// Parse installed headroom-ai version + which compression extras are +// actually installed (detected via marker package presence). One `pip list` +// call is enough to answer both questions. +// +// Returns: { installed: bool, version: string|null, extras: { code, ml } } +export function getInstalledHeadroomExtras(python) { + const py = python || findPython310(); + if (!py) return { installed: false, version: null, extras: { code: false, ml: false } }; + try { + const out = execFileSync(py, ["-m", "pip", "list", "--format=json", "--disable-pip-version-check"], { + stdio: ["ignore", "pipe", "ignore"], + windowsHide: true, + timeout: HEADROOM_PIP_TIMEOUT_MS, + env: { ...process.env, PATH: EXTENDED_PATH }, + }).toString(); + const packages = JSON.parse(out); + const names = new Set(packages.map((p) => String(p.name || "").toLowerCase())); + const installed = names.has("headroom-ai"); + if (!installed) return { installed: false, version: null, extras: { code: false, ml: false } }; + const version = packages.find((p) => p.name?.toLowerCase() === "headroom-ai")?.version || null; + const extras = {}; + for (const extra of HEADROOM_COMPRESSION_EXTRAS) { + extras[extra] = EXTRA_MARKERS[extra].some((m) => names.has(m)); + } + return { installed: true, version, extras }; + } catch { + return { installed: false, version: null, extras: { code: false, ml: false } }; + } } diff --git a/src/lib/headroom/process.js b/src/lib/headroom/process.js index d50bc7ec..a404ad4f 100644 --- a/src/lib/headroom/process.js +++ b/src/lib/headroom/process.js @@ -2,7 +2,7 @@ import fs from "fs"; import path from "path"; import { spawn } from "child_process"; import { DATA_DIR } from "@/lib/dataDir.js"; -import { findHeadroomBinary } from "./detect.js"; +import { findHeadroomBinary, findPython310, HEADROOM_COMPRESSION_EXTRAS, getInstalledHeadroomExtras } from "./detect.js"; const HEADROOM_DIR = path.join(DATA_DIR, "headroom"); const PID_FILE = path.join(HEADROOM_DIR, "proxy.pid"); @@ -126,3 +126,52 @@ export function getHeadroomLogTail(maxLines = 200) { return lines.slice(-maxLines).join("\n"); } catch { return ""; } } + +// Install (or upgrade) headroom-ai with the requested compression extras. +// `extras` is a whitelist from HEADROOM_COMPRESSION_EXTRAS — anything else +// is rejected to keep the install surface predictable. Always installs the +// `proxy` base + whatever extras the user picked, regardless of what is +// already present. +export async function installHeadroomExtras(extras = []) { + const requested = Array.isArray(extras) ? extras.filter((e) => HEADROOM_COMPRESSION_EXTRAS.includes(e)) : []; + const py = findPython310(); + if (!py) { + const err = new Error("Python >= 3.10 not found"); + err.code = "NO_PYTHON"; + throw err; + } + if (!findHeadroomBinary()) { + const err = new Error("headroom-ai not installed (run `pip install headroom-ai[proxy]` first)"); + err.code = "NOT_INSTALLED"; + throw err; + } + // pip install string is built from a closed set (HEADROOM_COMPRESSION_EXTRAS), + // so it cannot be poisoned by caller input — the comma-list is a fixed + // ['proxy', ...requested]. No shell interpolation. + const extrasList = ["proxy", ...requested].join(","); + const spec = `headroom-ai[${extrasList}]`; + const args = ["-m", "pip", "install", "--upgrade", spec]; + + ensureDir(); + const outFd = fs.openSync(path.join(HEADROOM_DIR, "install.log"), "a"); + const child = spawn(py, args, { + stdio: ["ignore", outFd, outFd], + windowsHide: true, + env: { ...process.env }, + }); + + return new Promise((resolve, reject) => { + child.once("error", (e) => { fs.closeSync(outFd); reject(e); }); + child.once("exit", (code) => { + fs.closeSync(outFd); + if (code === 0) { + const status = getInstalledHeadroomExtras(py); + resolve({ success: true, code, spec, extras: requested, ...status }); + } else { + const err = new Error(`pip install exited with code=${code} — see headroom/install.log`); + err.code = "INSTALL_FAILED"; + reject(err); + } + }); + }); +} diff --git a/tests/unit/headroom-detect.test.js b/tests/unit/headroom-detect.test.js index cbe16b25..1f0d3c30 100644 --- a/tests/unit/headroom-detect.test.js +++ b/tests/unit/headroom-detect.test.js @@ -2,21 +2,93 @@ import { describe, it, expect, vi, afterEach } from "vitest"; const mocks = vi.hoisted(() => ({ execSync: vi.fn(() => { throw new Error("not found"); }), + execFile: vi.fn(() => ({ toString: () => "[object Object]" })), + execFileSync: vi.fn(() => Buffer.from(JSON.stringify([ + { name: "headroom-ai", version: "0.26.0" }, + { name: "tree-sitter", version: "0.25.0" }, + ]))), })); vi.mock("child_process", () => ({ execSync: mocks.execSync, + execFile: mocks.execFile, + execFileSync: mocks.execFileSync, })); -import { getHeadroomStatus, isLoopbackHeadroomUrl } from "../../src/lib/headroom/detect.js"; +import { findPython310, getHeadroomStatus, getInstalledHeadroomExtras, isLoopbackHeadroomUrl } from "../../src/lib/headroom/detect.js"; afterEach(() => { vi.clearAllMocks(); }); describe("headroom detect", () => { + it("detects installed headroom version and extras from pip list", () => { + const result = getInstalledHeadroomExtras("python3"); + + expect(mocks.execFileSync).toHaveBeenCalledWith( + "python3", + ["-m", "pip", "list", "--format=json", "--disable-pip-version-check"], + expect.objectContaining({ windowsHide: true, timeout: 8000 }), + ); + expect(result).toEqual({ + installed: true, + version: "0.26.0", + extras: { code: true, ml: false }, + }); + }); + + it("prefers the interpreter that actually has headroom-ai installed", () => { + mocks.execSync.mockImplementation((cmd) => { + if (String(cmd).includes("where") || String(cmd).includes("which")) return Buffer.from("C:/Python/Scripts/headroom.exe\n"); + if (String(cmd).includes("python3 --version")) return Buffer.from("Python 3.13.0\n"); + if (String(cmd).includes("python --version")) return Buffer.from("Python 3.13.0\n"); + throw new Error("unexpected execSync"); + }); + mocks.execFileSync.mockImplementation((py, args) => { + if (py === "python3" && args.join(" ") === "-m pip show headroom-ai") throw new Error("not installed in python3"); + if (py === "python" && args.join(" ") === "-m pip show headroom-ai") return Buffer.from("Name: headroom-ai\nVersion: 0.26.0\n"); + if (py === "python" && args.join(" ").startsWith("-m pip list ")) return Buffer.from(JSON.stringify([ + { name: "headroom-ai", version: "0.26.0" }, + { name: "tree-sitter", version: "0.25.0" }, + ])); + throw new Error(`unexpected execFileSync: ${py} ${args.join(" ")}`); + }); + + expect(findPython310()).toBe("python"); + }); + + it("keeps top-level installed flag true when extras are readable", async () => { + global.fetch = vi.fn(async () => new Response("ok", { status: 200 })); + mocks.execSync.mockImplementation((cmd) => { + if (String(cmd).includes("where") || String(cmd).includes("which")) return Buffer.from("C:/Python/Scripts/headroom.exe\n"); + if (String(cmd).includes("python3 --version")) return Buffer.from("Python 3.13.0\n"); + if (String(cmd).includes("python --version")) return Buffer.from("Python 3.13.0\n"); + throw new Error("unexpected execSync"); + }); + mocks.execFileSync.mockImplementation((py, args) => { + if (py === "python3" && args.join(" ") === "-m pip show headroom-ai") throw new Error("not installed in python3"); + if (py === "python" && args.join(" ") === "-m pip show headroom-ai") return Buffer.from("Name: headroom-ai\nVersion: 0.26.0\n"); + if (py === "python" && args.join(" ").startsWith("-m pip list ")) return Buffer.from(JSON.stringify([ + { name: "headroom-ai", version: "0.26.0" }, + { name: "tree-sitter", version: "0.25.0" }, + ])); + throw new Error(`unexpected execFileSync: ${py} ${args.join(" ")}`); + }); + + const status = await getHeadroomStatus("http://localhost:8787"); + + expect(status.installed).toBe(true); + expect(status.version).toBe("0.26.0"); + expect(status.extras).toEqual({ code: true, ml: false }); + }); + it("treats a reachable external proxy as running without local CLI", async () => { global.fetch = vi.fn(async () => new Response("ok", { status: 200 })); + mocks.execSync.mockImplementation((cmd) => { + if (String(cmd).includes("where") || String(cmd).includes("which")) throw new Error("not found"); + throw new Error("unexpected execSync"); + }); + mocks.execFileSync.mockImplementation(() => { throw new Error("pip unavailable"); }); const status = await getHeadroomStatus("http://headroom:8787");