toknmtr-selfhost/agent/usage.ts

/**
 * toknmtr agent — `claude /usage` PTY-driven scraper.
 *
 * Ports `claude-usage-widget`'s `src-tauri/src/cli_usage.rs` (PTY-drive sequence + bar
 * parsing) to pure Node, without the `node-pty` native dependency. We don't have a true
 * PTY API in plain Node, so we shell out to a PTY-shim helper — `script` (util-linux,
 * preferred) or `unbuffer` (expect) — and drive `claude` through it. Both make `claude`
 * believe it's attached to a real terminal, which it needs in order to render the
 * interactive `/usage` TUI instead of falling back to a non-interactive mode.
 *
 * Sequence (mirrors the Rust reference): spawn → wait ~3.5s for the TUI to finish its
 * startup render → send "/usage\r" → drain output until it goes quiet for >1.2s (or a
 * 20s deadline trips) → send "/exit\r" as a best-effort clean shutdown → kill the child.
 * The captured bytes are then ANSI-stripped and the three rendered bars ("Current
 * session", "Current week (all models)", "Current week (Sonnet only)") are parsed out
 * with the same `NN% used` regex the Rust version uses.
 *
 * ASSUMPTIONS / things to watch if this breaks:
 *  - Verified against a live `script -qfc "claude" /dev/null` capture (Claude Code
 *    v2.1.197) on 2026-07-01: the rendered frame uses bare `\r` + a cursor-down escape
 *    (`\x1b[1B`) in place of `\n` for most line breaks, *not* `\r\n`. Once the CSI
 *    sequences are stripped, only the bare `\r` survives — so `stripAnsiCollapse` here
 *    splits on `\r\n`, lone `\r`, *and* `\n` (the Rust version relies on Rust's
 *    `str::lines()`, which only recognizes `\n`/`\r\n` — that's fine over there because
 *    `portable-pty` apparently surfaces real `\n`s; it would NOT be fine against the
 *    `script`-captured stream this file actually sees, so don't port that exact
 *    behavior back unmodified).
 *  - The "Current week (Sonnet only)" section is genuinely optional (absent entirely in
 *    the live capture above) — `week_sonnet_pct` legitimately being `null` is expected,
 *    not a parse failure.
 *  - `unbuffer` (expect) is supported as a fallback driver per the task spec, but was
 *    NOT available in the dev sandbox to test against; only `script` was exercised live.
 *  - This is brittle to Anthropic changing the rendered `/usage` output, same caveat as
 *    the widget. See that project's memory.md for prior gotchas.
 *
 * Standalone test: `node --experimental-strip-types agent/usage.ts`
 */
import { spawn, spawnSync } from 'node:child_process';
import { hostname } from 'node:os';

export interface UsageGauges {
	host: string;
	ts_utc: string;
	session_pct: number | null;
	week_all_pct: number | null;
	week_sonnet_pct: number | null;
}

const STARTUP_DELAY_MS = 3500; // let the TUI finish its startup render before typing
const POLL_INTERVAL_MS = 700; // how often we check for new output while draining
const QUIET_PERIOD_MS = 1200; // no new bytes for this long => render is done
const TOTAL_TIMEOUT_MS = 20000; // hard cap on the whole drain phase
const EXIT_DRAIN_MS = 500; // grace period after sending /exit before we kill

function commandExists(name: string): boolean {
	try {
		return spawnSync('which', [name], { stdio: 'ignore' }).status === 0;
	} catch {
		return false;
	}
}

/** Pick a PTY-shim driver for `claude`. Returns argv, or null if neither is installed. */
function pickDriver(): string[] | null {
	if (commandExists('script')) {
		// util-linux script: -q quiet (no "Script started/done" banner), -f flush output as
		// written, -c <command> run this under a pty. Typescript log target is /dev/null —
		// we only care about the copy `script` also streams to its own stdout.
		return ['script', '-qfc', 'claude', '/dev/null'];
	}
	if (commandExists('unbuffer')) {
		return ['unbuffer', 'claude'];
	}
	return null;
}

function sleep(ms: number): Promise<void> {
	return new Promise((resolve) => setTimeout(resolve, ms));
}

/** Spawn the driver, type "/usage", capture the raw (ANSI-laden) bytes, then kill it. */
async function driveClaudeUsage(argv: string[]): Promise<Buffer> {
	const [cmd, ...args] = argv;
	const child = spawn(cmd, args, {
		env: { ...process.env, TERM: 'xterm-256color' }
	});

	let out = Buffer.alloc(0);
	let lastGrowth = Date.now();
	const onData = (chunk: Buffer) => {
		out = Buffer.concat([out, chunk]);
		lastGrowth = Date.now();
	};
	child.stdout.on('data', onData);
	child.stderr.on('data', onData);
	// Swallow spawn-time errors here (e.g. ENOENT racing past commandExists) — the
	// caller's try/catch around fetchUsageGauges is the real backstop, but an
	// unhandled 'error' event on the child would otherwise crash the process.
	child.on('error', () => {});

	try {
		// 1. Let the TUI finish its startup render.
		await sleep(STARTUP_DELAY_MS);

		// 2. Send /usage.
		child.stdin.write('/usage\r');

		// 3. Drain until output goes quiet or we hit the deadline.
		const deadline = Date.now() + TOTAL_TIMEOUT_MS;
		for (;;) {
			await sleep(POLL_INTERVAL_MS);
			if (Date.now() - lastGrowth > QUIET_PERIOD_MS) break;
			if (Date.now() > deadline) break;
		}

		// 4. Best-effort clean exit, then a short grace drain.
		child.stdin.write('/exit\r');
		await sleep(EXIT_DRAIN_MS);
	} finally {
		child.stdout.off('data', onData);
		child.stderr.off('data', onData);
		child.kill('SIGKILL');
	}

	return out;
}

// CSI: ESC [ ... <terminator>; OSC: ESC ] ... BEL; DCS/SOS/PM/APC: ESC P|X|^|_ ... ESC \
// eslint-disable-next-line no-control-regex -- matching raw ANSI control bytes is the point
const CSI_RE = /\x1b\[[0-9;?]*[ -/]*[@-~]|\x1b\][^\x07]*\x07|\x1b[PX^_][\s\S]*?\x1b\\/g;

/**
 * Strip ANSI escapes and collapse repeated/blank lines (TUIs redraw the same content
 * many times). See the file-header note on why this splits on bare `\r` too, not just
 * `\r\n`/`\n`.
 */
export function stripAnsiCollapse(raw: Buffer): string {
	let text = raw.toString('utf8').replace(CSI_RE, '');
	// Drop stray BEL / lone ESC bytes left over from any sequence the regex didn't match.
	// eslint-disable-next-line no-control-regex -- matching raw ANSI control bytes is the point
	text = text.replace(/[\x07\x1b]/g, '');

	const deduped: string[] = [];
	for (const rawLine of text.split(/\r\n|\r|\n/)) {
		const trimmed = rawLine.replace(/\s+$/, '');
		if (deduped.length > 0 && deduped[deduped.length - 1] === trimmed) continue;
		deduped.push(trimmed);
	}

	const compressed: string[] = [];
	let prevBlank = false;
	for (const line of deduped) {
		const blank = line.length === 0;
		if (blank && prevBlank) continue;
		compressed.push(line);
		prevBlank = blank;
	}
	return compressed.join('\n');
}

const PCT_RE = /(\d{1,3})\s*%\s*used/;

/** Find `label`'s heading line, then look at the next few lines for "NN% used". */
function findSectionPct(lines: string[], label: string): number | null {
	const normLabel = label.replace(/\s+/g, '');
	const idx = lines.findIndex((l) => {
		const t = l.trim();
		return t === label || t.replace(/\s+/g, '') === normLabel;
	});
	if (idx === -1) return null;
	for (const line of lines.slice(idx + 1, idx + 7)) {
		const m = PCT_RE.exec(line);
		if (m) {
			const n = Number.parseInt(m[1], 10);
			return Number.isFinite(n) ? n : null;
		}
	}
	return null;
}

function parseUsageText(stripped: string): {
	session_pct: number | null;
	week_all_pct: number | null;
	week_sonnet_pct: number | null;
} {
	const lines = stripped.split('\n');
	return {
		session_pct: findSectionPct(lines, 'Current session'),
		week_all_pct: findSectionPct(lines, 'Current week (all models)'),
		week_sonnet_pct: findSectionPct(lines, 'Current week (Sonnet only)')
	};
}

/**
 * Fetch the three subscription-usage gauges by PTY-driving `claude /usage`.
 *
 * Best-effort: returns null (never throws) if no PTY-shim driver is installed, if
 * `claude` isn't reachable, or if the output couldn't be parsed at all. This must never
 * break the main agent run.
 */
export async function fetchUsageGauges(): Promise<UsageGauges | null> {
	try {
		const argv = pickDriver();
		if (!argv) return null;

		const raw = await driveClaudeUsage(argv);
		if (raw.length === 0) return null;

		const stripped = stripAnsiCollapse(raw);
		const { session_pct, week_all_pct, week_sonnet_pct } = parseUsageText(stripped);
		if (session_pct === null && week_all_pct === null && week_sonnet_pct === null) {
			return null;
		}

		return {
			host: hostname(),
			ts_utc: new Date().toISOString(),
			session_pct,
			week_all_pct,
			week_sonnet_pct
		};
	} catch {
		return null;
	}
}

async function main(): Promise<void> {
	const result = await fetchUsageGauges();
	console.log(JSON.stringify(result, null, 2));
}

const entryPoint = process.argv[1];
if (entryPoint && import.meta.url === `file://${entryPoint}`) {
	void main();
}