diff --git a/src/app/(dashboard)/dashboard/token-saver/TokenSaverClient.js b/src/app/(dashboard)/dashboard/token-saver/TokenSaverClient.js index fc1f9152..4e641f3d 100644 --- a/src/app/(dashboard)/dashboard/token-saver/TokenSaverClient.js +++ b/src/app/(dashboard)/dashboard/token-saver/TokenSaverClient.js @@ -1,7 +1,7 @@ "use client"; -import { useState, useEffect, useCallback } from "react"; -import { Card, Button, Input, Modal, Toggle } from "@/shared/components"; +import { useState, useEffect, useCallback, useRef } from "react"; +import { Card, Button, Input, Modal, Toggle, ConfirmModal } from "@/shared/components"; import { useCopyToClipboard } from "@/shared/hooks/useCopyToClipboard"; import { getCurrentLocale, onLocaleChange } from "@/i18n/runtime"; import { @@ -33,6 +33,13 @@ export default function TokenSaverClient() { const [pendingExtras, setPendingExtras] = useState([]); const [extrasActionLoading, setExtrasActionLoading] = useState(false); const [extrasActionError, setExtrasActionError] = useState(""); + const [removingExtra, setRemovingExtra] = useState(null); + const [installLog, setInstallLog] = useState(""); + const [extrasConfirm, setExtrasConfirm] = useState(null); + const [codeAware, setCodeAware] = useState(false); + const [kompress, setKompress] = useState(true); + const [restartingProxy, setRestartingProxy] = useState(false); + const logPollRef = useRef(null); const [cavemanEnabled, setCavemanEnabled] = useState(false); const [cavemanLevel, setCavemanLevel] = useState("full"); const [ponytailEnabled, setPonytailEnabled] = useState(false); @@ -205,10 +212,37 @@ export default function TokenSaverClient() { ); }; - const handleInstallExtras = useCallback(async () => { + // Poll the install log tail while a pip install/uninstall is running. + const startLogPolling = useCallback(() => { + setInstallLog(""); + if (logPollRef.current) clearInterval(logPollRef.current); + const tick = async () => { + try { + const r = await fetch("/api/headroom/extras?log=1", { + headers: { "Cache-Control": "no-store" }, + }); + const d = await r.json().catch(() => ({})); + if (typeof d.log === "string") setInstallLog(d.log); + } catch { /* ignore transient poll errors */ } + }; + tick(); + logPollRef.current = setInterval(tick, 1500); + }, []); + + const stopLogPolling = useCallback(() => { + if (logPollRef.current) { + clearInterval(logPollRef.current); + logPollRef.current = null; + } + }, []); + + useEffect(() => () => stopLogPolling(), [stopLogPolling]); + + const installExtrasConfirmed = useCallback(async () => { if (pendingExtras.length === 0) return; setExtrasActionLoading(true); setExtrasActionError(""); + startLogPolling(); try { const res = await fetch("/api/headroom/extras", { method: "POST", @@ -226,9 +260,83 @@ export default function TokenSaverClient() { } catch (e) { setExtrasActionError(e.message); } finally { + stopLogPolling(); setExtrasActionLoading(false); } - }, [pendingExtras]); + }, [pendingExtras, startLogPolling, stopLogPolling]); + + const removeExtraConfirmed = useCallback(async (extra) => { + setRemovingExtra(extra); + setExtrasActionError(""); + startLogPolling(); + try { + const res = await fetch("/api/headroom/extras", { + method: "DELETE", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ extras: [extra] }), + }); + const data = await res.json().catch(() => ({})); + if (!res.ok) throw new Error(data.error || "Remove failed"); + setHeadroomExtras((s) => ({ + ...s, + version: data.version ?? s.version, + extras: data.extras || s.extras, + })); + } catch (e) { + setExtrasActionError(e.message); + } finally { + stopLogPolling(); + setRemovingExtra(null); + } + }, [startLogPolling, stopLogPolling]); + + const handleInstallExtras = useCallback(() => { + if (pendingExtras.length === 0) return; + // Warn about the heavy ~1GB torch download before installing [ml]. + if (pendingExtras.includes("ml")) { + setExtrasConfirm({ + title: "Install [ml]", + message: "[ml] downloads ~1 GB (torch + huggingface-hub). Continue?", + confirmText: "Install", + variant: "primary", + onConfirm: installExtrasConfirmed, + }); + return; + } + installExtrasConfirmed(); + }, [pendingExtras, installExtrasConfirmed]); + + const handleRemoveExtra = useCallback((extra) => { + setExtrasConfirm({ + title: `Remove [${extra}]`, + message: `Remove [${extra}] and its packages?`, + confirmText: "Remove", + variant: "danger", + onConfirm: () => removeExtraConfirmed(extra), + }); + }, [removeExtraConfirmed]); + + // Toggle an extra's active state (persist setting), then restart the proxy so + // the new --code-aware / --disable-kompress flags take effect. + const toggleExtraActive = useCallback(async (extra, value) => { + setExtrasActionError(""); + if (extra === "code") setCodeAware(value); + if (extra === "ml") setKompress(value); + const key = extra === "code" ? "headroomCodeAware" : "headroomKompress"; + await patchSetting({ [key]: value }); + if (!headroomStatus.running) return; + setRestartingProxy(true); + try { + const res = await fetch("/api/headroom/restart", { method: "POST" }); + const data = await res.json().catch(() => ({})); + if (!res.ok) throw new Error(data.error || "Restart failed"); + await refreshHeadroomStatus(); + } catch (e) { + setExtrasActionError(e.message); + } finally { + setRestartingProxy(false); + } + }, [headroomStatus.running, refreshHeadroomStatus]); const handleCavemanLevel = (level) => { setCavemanLevel(level); @@ -307,6 +415,8 @@ export default function TokenSaverClient() { setRtkEnabledState(data.rtkEnabled !== false); setHeadroomEnabled(!!data.headroomEnabled); setHeadroomUrl(data.headroomUrl || "http://localhost:8787"); + setCodeAware(data.headroomCodeAware === true); + setKompress(data.headroomKompress !== false); setCavemanEnabled(!!data.cavemanEnabled); setCavemanLevel(data.cavemanLevel || "full"); setPonytailEnabled(!!data.ponytailEnabled); @@ -387,7 +497,7 @@ export default function TokenSaverClient() { onChange={() => handleRtkEnabled(!rtkEnabled)} /> -
+

@@ -425,7 +535,7 @@ export default function TokenSaverClient() { />

{headroomStatus.installed && ( -
+
Compression extras @@ -434,33 +544,57 @@ export default function TokenSaverClient() { {headroomExtras.available.map((extra) => { const installed = !!headroomExtras.extras[extra]; const pending = pendingExtras.includes(extra); + const extraTitle = + extra === "code" + ? "tree-sitter AST compression for code responses" + : "Kompress-v2 HF model for prose/agentic traces (~+1GB)"; + + if (installed) { + const active = extra === "code" ? codeAware : kompress; + return ( +
+ toggleExtraActive(extra, !active)} + /> + [{extra}] + +
+ ); + } + return ( ); })} @@ -479,16 +613,26 @@ export default function TokenSaverClient() { {extrasActionError && (

{extrasActionError}

)} + {restartingProxy && ( +

Restarting proxy…

+ )} + {(extrasActionLoading || removingExtra) && installLog && ( + 
+                {installLog}
+
+ )}

- Default install is [proxy] only (SmartCrusher for - JSON). Adding [code] enables AST compression + Installing adds the package; use on/off{" "} + to activate it (restarts the proxy). 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).

)} -
+

Compress LLM output{" "} @@ -589,6 +733,8 @@ export default function TokenSaverClient() { />

+ {/* PXPIPE hidden from UI — experimental, not exposed to users yet */} + {false && (
@@ -632,6 +778,7 @@ export default function TokenSaverClient() { onChange={() => handlePxpipeEnabled(!pxpipeEnabled)} />
+ )} setShowPxpipeModal(false)} > @@ -849,6 +996,20 @@ export default function TokenSaverClient() {
+ + setExtrasConfirm(null)} + onConfirm={() => { + const fn = extrasConfirm?.onConfirm; + setExtrasConfirm(null); + fn?.(); + }} + title={extrasConfirm?.title} + message={extrasConfirm?.message} + confirmText={extrasConfirm?.confirmText} + variant={extrasConfirm?.variant} + />
); } diff --git a/src/app/api/headroom/extras/route.js b/src/app/api/headroom/extras/route.js index 4c1d7230..3dffaa5b 100644 --- a/src/app/api/headroom/extras/route.js +++ b/src/app/api/headroom/extras/route.js @@ -1,11 +1,15 @@ import { NextResponse } from "next/server"; import { findPython310, getInstalledHeadroomExtras, HEADROOM_COMPRESSION_EXTRAS } from "@/lib/headroom/detect"; -import { installHeadroomExtras } from "@/lib/headroom/process"; +import { installHeadroomExtras, uninstallHeadroomExtras, getInstallLogTail } from "@/lib/headroom/process"; export const dynamic = "force-dynamic"; -export async function GET() { +export async function GET(req) { try { + // `?log=1` returns the live install/uninstall log tail for progress polling. + if (new URL(req.url).searchParams.get("log") === "1") { + return NextResponse.json({ log: getInstallLogTail() }); + } const python = findPython310(); const status = getInstalledHeadroomExtras(python); return NextResponse.json({ @@ -28,3 +32,15 @@ export async function POST(req) { return NextResponse.json({ error: error.message, code: error.code || null }, { status }); } } + +export async function DELETE(req) { + try { + const body = await req.json().catch(() => ({})); + const requested = Array.isArray(body?.extras) ? body.extras : []; + const result = await uninstallHeadroomExtras(requested); + return NextResponse.json(result); + } catch (error) { + const status = error.code === "NO_PYTHON" || error.code === "INVALID_EXTRAS" ? 400 : 500; + return NextResponse.json({ error: error.message, code: error.code || null }, { status }); + } +} diff --git a/src/app/api/headroom/restart/route.js b/src/app/api/headroom/restart/route.js new file mode 100644 index 00000000..7d98e80f --- /dev/null +++ b/src/app/api/headroom/restart/route.js @@ -0,0 +1,35 @@ +import { NextResponse } from "next/server"; +import { getSettings } from "@/lib/localDb"; +import { restartHeadroomProxy } from "@/lib/headroom/process"; +import { DEFAULT_HEADROOM_URL, isLoopbackHeadroomUrl } from "@/lib/headroom/detect"; + +export const dynamic = "force-dynamic"; + +function parsePortFromUrl(url) { + try { + const u = new URL(url); + const p = parseInt(u.port, 10); + if (p > 0 && p < 65536) return p; + } catch { /* ignore, fall through to default */ } + return null; +} + +export async function POST() { + try { + const settings = await getSettings(); + const url = settings.headroomUrl || DEFAULT_HEADROOM_URL; + if (!isLoopbackHeadroomUrl(url)) { + return NextResponse.json({ error: "External Headroom proxies must be started outside 9Router", code: "EXTERNAL_PROXY" }, { status: 400 }); + } + const port = parsePortFromUrl(url) || 8787; + const result = await restartHeadroomProxy({ + port, + codeAware: settings.headroomCodeAware === true, + kompress: settings.headroomKompress !== false, + }); + return NextResponse.json({ success: true, ...result }); + } catch (error) { + const status = error.code === "NOT_INSTALLED" ? 400 : 500; + return NextResponse.json({ error: error.message, code: error.code || null }, { status }); + } +} diff --git a/src/app/api/headroom/start/route.js b/src/app/api/headroom/start/route.js index 56af456c..9be6d88d 100644 --- a/src/app/api/headroom/start/route.js +++ b/src/app/api/headroom/start/route.js @@ -22,7 +22,11 @@ export async function POST() { return NextResponse.json({ error: "External Headroom proxies must be started outside 9Router", code: "EXTERNAL_PROXY" }, { status: 400 }); } const port = parsePortFromUrl(url) || 8787; - const result = await startHeadroomProxy({ port }); + const result = await startHeadroomProxy({ + port, + codeAware: settings.headroomCodeAware === true, + kompress: settings.headroomKompress !== false, + }); return NextResponse.json({ success: true, ...result }); } catch (error) { const status = error.code === "NOT_INSTALLED" ? 400 : 500; diff --git a/src/lib/headroom/detect.js b/src/lib/headroom/detect.js index 91b245f7..0ba9e2b4 100644 --- a/src/lib/headroom/detect.js +++ b/src/lib/headroom/detect.js @@ -9,7 +9,7 @@ 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 = { +export const EXTRA_MARKERS = { code: ["tree-sitter", "tree-sitter-language-pack"], ml: ["torch", "huggingface-hub"], }; @@ -64,11 +64,33 @@ 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 +// `python3`, `python3.13`, `python` can point at different envs on any OS. Prefer +// the interpreter that can also see the installed `headroom-ai` package so the // dashboard probes and install action operate on the same interpreter as the CLI. +// Falls back to the first version-eligible candidate when headroom-ai is not yet +// installed anywhere (needed for the initial install). +// Interpreters to probe, most specific first: the python next to the headroom +// binary (guaranteed to have headroom-ai), then full paths from EXTRA_BINS, then +// bare names resolved via PATH. +function pythonCandidates() { + const list = []; + const bin = findHeadroomBinary(); + if (bin) { + const dir = path.dirname(bin); + const names = IS_WIN ? ["python.exe", "python3.exe"] : ["python3", "python3.13", "python"]; + for (const n of names) list.push(path.join(dir, n)); + } + for (const dir of EXTRA_BINS) { + if (!dir) continue; + for (const n of PYTHON_CANDIDATES) list.push(path.join(dir, IS_WIN ? `${n}.exe` : n)); + } + list.push(...PYTHON_CANDIDATES); + return list; +} + export function findPython310() { - for (const candidate of PYTHON_CANDIDATES) { + let fallback = null; + for (const candidate of pythonCandidates()) { try { const ver = execSync(`${candidate} --version`, { stdio: ["ignore", "pipe", "ignore"], @@ -79,7 +101,7 @@ export function findPython310() { 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]))) continue; - if (!IS_WIN) return candidate; + if (!fallback) fallback = candidate; try { execFileSync(candidate, ["-m", "pip", "show", "headroom-ai"], { stdio: ["ignore", "pipe", "ignore"], @@ -89,13 +111,13 @@ export function findPython310() { }); return candidate; } catch { - // Keep scanning on Windows until an interpreter sees headroom-ai. + // Keep scanning until an interpreter that sees headroom-ai is found. } } catch { // candidate not present, try next } } - return null; + return fallback; } // Probe whether a Headroom proxy is reachable at the given URL by hitting /health. diff --git a/src/lib/headroom/process.js b/src/lib/headroom/process.js index a404ad4f..41720875 100644 --- a/src/lib/headroom/process.js +++ b/src/lib/headroom/process.js @@ -2,11 +2,12 @@ import fs from "fs"; import path from "path"; import { spawn } from "child_process"; import { DATA_DIR } from "@/lib/dataDir.js"; -import { findHeadroomBinary, findPython310, HEADROOM_COMPRESSION_EXTRAS, getInstalledHeadroomExtras } from "./detect.js"; +import { findHeadroomBinary, findPython310, HEADROOM_COMPRESSION_EXTRAS, EXTRA_MARKERS, getInstalledHeadroomExtras } from "./detect.js"; const HEADROOM_DIR = path.join(DATA_DIR, "headroom"); const PID_FILE = path.join(HEADROOM_DIR, "proxy.pid"); const LOG_FILE = path.join(HEADROOM_DIR, "proxy.log"); +const INSTALL_LOG_FILE = path.join(HEADROOM_DIR, "install.log"); const DEFAULT_PORT = 8787; const STARTUP_TIMEOUT_MS = 8000; @@ -41,7 +42,17 @@ export function getManagedPid() { return pid && isPidAlive(pid) ? pid : null; } -export async function startHeadroomProxy({ port = DEFAULT_PORT } = {}) { +// Build proxy CLI flags for the active compression extras. `[code]` (AST +// compression) is off by default in headroom → pass --code-aware to turn it on; +// `[ml]` (Kompress) is on by default → pass --disable-kompress to turn it off. +function extrasProxyArgs({ codeAware, kompress } = {}) { + const args = []; + if (codeAware) args.push("--code-aware"); + if (kompress === false) args.push("--disable-kompress"); + return args; +} + +export async function startHeadroomProxy({ port = DEFAULT_PORT, codeAware = false, kompress = true } = {}) { const safePort = Number(port) > 0 && Number(port) < 65536 ? Number(port) : DEFAULT_PORT; const binary = findHeadroomBinary(); if (!binary) { @@ -57,7 +68,8 @@ export async function startHeadroomProxy({ port = DEFAULT_PORT } = {}) { // spawn stdio requires fd numbers, not WriteStream objects. const outFd = fs.openSync(LOG_FILE, "a"); - const child = spawn(binary, ["proxy", "--port", String(safePort)], { + const args = ["proxy", "--port", String(safePort), ...extrasProxyArgs({ codeAware, kompress })]; + const child = spawn(binary, args, { stdio: ["ignore", outFd, outFd], detached: true, windowsHide: true, @@ -118,6 +130,25 @@ export function stopHeadroomProxy() { } } +// Stop the managed proxy (if any), wait for the pid to die, then start again +// with the given flags. Used when toggling active extras that require a restart. +export async function restartHeadroomProxy(opts = {}) { + const pid = getManagedPid(); + if (pid) { + try { process.kill(pid, "SIGTERM"); } catch { /* already gone */ } + // Wait up to ~3s for graceful exit, force-kill if still alive. + for (let i = 0; i < 30 && isPidAlive(pid); i++) { + await new Promise((r) => setTimeout(r, 100)); + } + if (isPidAlive(pid)) { + try { process.kill(pid, "SIGKILL"); } catch { /* already gone */ } + await new Promise((r) => setTimeout(r, 300)); + } + clearPid(); + } + return startHeadroomProxy(opts); +} + export function getHeadroomLogTail(maxLines = 200) { try { if (!fs.existsSync(LOG_FILE)) return ""; @@ -153,7 +184,8 @@ export async function installHeadroomExtras(extras = []) { const args = ["-m", "pip", "install", "--upgrade", spec]; ensureDir(); - const outFd = fs.openSync(path.join(HEADROOM_DIR, "install.log"), "a"); + // Truncate ("w") so the log reflects only the current install for live progress. + const outFd = fs.openSync(INSTALL_LOG_FILE, "w"); const child = spawn(py, args, { stdio: ["ignore", outFd, outFd], windowsHide: true, @@ -175,3 +207,54 @@ export async function installHeadroomExtras(extras = []) { }); }); } + +// Uninstall the marker packages that back a single extra (e.g. `ml` → torch, +// huggingface-hub). `headroom-ai` base and the `proxy` extra are never removed. +export async function uninstallHeadroomExtras(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; + } + const pkgs = [...new Set(requested.flatMap((e) => EXTRA_MARKERS[e] || []))]; + if (pkgs.length === 0) { + const err = new Error("No valid extras to remove"); + err.code = "INVALID_EXTRAS"; + throw err; + } + const args = ["-m", "pip", "uninstall", "-y", ...pkgs]; + + ensureDir(); + const outFd = fs.openSync(INSTALL_LOG_FILE, "w"); + 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, removed: pkgs, extras: requested, ...status }); + } else { + const err = new Error(`pip uninstall exited with code=${code} — see headroom/install.log`); + err.code = "UNINSTALL_FAILED"; + reject(err); + } + }); + }); +} + +// Read the tail of the install/uninstall log for live progress in the UI. +export function getInstallLogTail(maxLines = 15) { + try { + if (!fs.existsSync(INSTALL_LOG_FILE)) return ""; + const lines = fs.readFileSync(INSTALL_LOG_FILE, "utf8").split(/\r?\n/).filter(Boolean); + return lines.slice(-maxLines).join("\n"); + } catch { return ""; } +} diff --git a/tests/unit/headroom-detect.test.js b/tests/unit/headroom-detect.test.js index 1f0d3c30..d1239a6b 100644 --- a/tests/unit/headroom-detect.test.js +++ b/tests/unit/headroom-detect.test.js @@ -38,23 +38,22 @@ describe("headroom detect", () => { }); it("prefers the interpreter that actually has headroom-ai installed", () => { + // headroom binary lives in a bin dir; the python next to it has headroom-ai. + const binPython = "/opt/hr/bin/python3"; 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"); + if (String(cmd).includes("where") || String(cmd).includes("which")) return Buffer.from("/opt/hr/bin/headroom\n"); + if (String(cmd).includes("--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" }, - ])); + if (args.join(" ") === "-m pip show headroom-ai") { + if (py === binPython) return Buffer.from("Name: headroom-ai\nVersion: 0.26.0\n"); + throw new Error(`not installed in ${py}`); + } throw new Error(`unexpected execFileSync: ${py} ${args.join(" ")}`); }); - expect(findPython310()).toBe("python"); + expect(findPython310()).toBe(binPython); }); it("keeps top-level installed flag true when extras are readable", async () => {