ashlr-plugin 1.36.4

package/hooks/pretooluse-budget-guard.ts

@@ -0,0 +1,318 @@
+#!/usr/bin/env bun
+/**
+ * pretooluse-budget-guard.ts — Block tool calls when the session budget is exceeded.
+ *
+ * PreToolUse hook. Reads ASHLR_SESSION_BUDGET_USD or ASHLR_SESSION_BUDGET_TOKENS,
+ * estimates cumulative session spend from ~/.ashlr/session-log.jsonl, then:
+ *   - warns at 80%  (exits 0 with additionalContext nudge, event: budget_threshold_80)
+ *   - warns at 95%  (exits 0 with louder nudge, event: budget_threshold_95)
+ *   - blocks at 100% (exits 2 with stderr message, event: budget_exceeded)
+ *
+ * Cost heuristic: tokens ≈ (input_bytes + output_bytes) / 4
+ *   USD ≈ tokens / 1_000_000 * 12  (blended $4/Mtok input + $20/Mtok output, 60/40 mix)
+ *
+ * Design:
+ *   - Best-effort only — never throws. Falls through on any read error.
+ *   - Kill switch: ASHLR_SESSION_LOG=0 skips log reads (no usage data → never blocks).
+ *   - Only fires when ASHLR_SESSION_BUDGET_USD or ASHLR_SESSION_BUDGET_TOKENS is set.
+ *
+ * Hook contract (PreToolUse):
+ *   stdin  → { tool_name, tool_input? }
+ *   stdout → { hookSpecificOutput: { hookEventName: "PreToolUse", additionalContext? } }
+ *   exit 0 → allow
+ *   exit 2 → block (message on stderr shown to user)
+ */
+
+import { existsSync, readFileSync } from "fs";
+import { homedir } from "os";
+import { join } from "path";
+
+import { installHookTimeout } from "./pretooluse-common";
+
+if (import.meta.main) installHookTimeout("pretooluse-budget-guard");
+
+// ---------------------------------------------------------------------------
+// Constants
+// ---------------------------------------------------------------------------
+
+/** Blended price per million tokens (input $4 + output $20, 60/40 mix). */
+const BLENDED_USD_PER_MTOK = 4 * 0.6 + 20 * 0.4; // 9.2 $/Mtok
+
+/** Bytes per token heuristic. */
+const BYTES_PER_TOKEN = 4;
+
+const WARN_80_PCT = 0.80;
+const WARN_95_PCT = 0.95;
+const BLOCK_100_PCT = 1.00;
+
+// ---------------------------------------------------------------------------
+// Types
+// ---------------------------------------------------------------------------
+
+interface HookOutput {
+  hookSpecificOutput: {
+    hookEventName: "PreToolUse";
+    additionalContext?: string;
+  };
+}
+
+interface SessionLogEntry {
+  session?: string;
+  input_size?: number;
+  output_size?: number;
+}
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
+function passThrough(context?: string): HookOutput {
+  return {
+    hookSpecificOutput: {
+      hookEventName: "PreToolUse",
+      ...(context ? { additionalContext: context } : {}),
+    },
+  };
+}
+
+function home(): string {
+  return (
+    process.env.ASHLR_HOME_OVERRIDE?.trim() ||
+    process.env.HOME ||
+    homedir()
+  );
+}
+
+function sessionLogPath(homeDir: string): string {
+  return join(homeDir, ".ashlr", "session-log.jsonl");
+}
+
+/** Derive current session ID using the same strategy as session-log-append.ts */
+function currentSessionId(): string {
+  const explicit =
+    process.env.CLAUDE_SESSION_ID?.trim() ||
+    process.env.ASHLR_SESSION_ID?.trim();
+  if (explicit) return explicit;
+  const seed = `${process.cwd()}:${process.ppid ?? process.pid}`;
+  let h = 0;
+  for (let i = 0; i < seed.length; i++) {
+    h = ((h << 5) - h + seed.charCodeAt(i)) | 0;
+  }
+  return `h${(h >>> 0).toString(16)}`;
+}
+
+/** Read session-log.jsonl and sum input_size + output_size for the current session.
+ *
+ * Failures surface to stderr (not silent) so a corrupted session log can't
+ * underestimate spend without anyone noticing — direct $ / UX risk. Still
+ * returns 0 on any error: budget enforcement is best-effort. */
+export function readSessionBytes(homeDir: string, sessionId: string): number {
+  try {
+    const p = sessionLogPath(homeDir);
+    if (!existsSync(p)) return 0;
+    const raw = readFileSync(p, "utf-8");
+    let total = 0;
+    let malformed = 0;
+    for (const line of raw.split("\n")) {
+      const trimmed = line.trim();
+      if (!trimmed) continue;
+      try {
+        const entry = JSON.parse(trimmed) as SessionLogEntry;
+        if (entry.session && entry.session !== sessionId) continue;
+        total += (entry.input_size ?? 0) + (entry.output_size ?? 0);
+      } catch {
+        malformed++;
+      }
+    }
+    if (malformed > 0) {
+      process.stderr.write(
+        `[ashlr budget-guard] skipped ${malformed} malformed line${malformed === 1 ? "" : "s"} in session-log.jsonl — budget estimate may be low\n`,
+      );
+    }
+    return total;
+  } catch (e) {
+    const msg = e instanceof Error ? e.message : String(e);
+    process.stderr.write(`[ashlr budget-guard] session-log read failed: ${msg}\n`);
+    return 0;
+  }
+}
+
+/** Bytes → estimated tokens. */
+export function bytesToTokens(bytes: number): number {
+  return Math.ceil(bytes / BYTES_PER_TOKEN);
+}
+
+/** Tokens → estimated USD. */
+export function tokensToUsd(tokens: number): number {
+  return (tokens / 1_000_000) * BLENDED_USD_PER_MTOK;
+}
+
+// ---------------------------------------------------------------------------
+// Nudge event emission (fire-and-forget, no import from servers/ to avoid
+// circular deps in a hook context)
+// ---------------------------------------------------------------------------
+
+type BudgetEventKind =
+  | "budget_threshold_80"
+  | "budget_threshold_95"
+  | "budget_exceeded";
+
+function emitBudgetEvent(event: BudgetEventKind, pct: number): void {
+  try {
+    const p = join(home(), ".ashlr", "nudge-events.jsonl");
+    const line = JSON.stringify({
+      ts: new Date().toISOString(),
+      event,
+      pct: Math.round(pct * 1000) / 10,
+    });
+    try {
+      const { appendFileSync, mkdirSync } = require("fs");
+      const { dirname } = require("path");
+      mkdirSync(dirname(p), { recursive: true });
+      appendFileSync(p, line + "\n", "utf-8");
+    } catch { /* best-effort */ }
+  } catch { /* best-effort */ }
+}
+
+// ---------------------------------------------------------------------------
+// Core decision (exported for tests)
+// ---------------------------------------------------------------------------
+
+export interface DecideOpts {
+  home?: string;
+  sessionId?: string;
+  /** Override env reads for testing. */
+  budgetUsd?: string;
+  budgetTokens?: string;
+  sessionLog?: string; // 0 = disabled
+}
+
+export interface DecideResult {
+  action: "pass" | "warn-80" | "warn-95" | "block";
+  context?: string;
+  pct: number;
+  usedTokens: number;
+  budgetTokens: number;
+  usedUsd: number;
+  budgetUsd: number;
+}
+
+export function decide(opts: DecideOpts = {}): DecideResult {
+  const sessionLogDisabled =
+    (opts.sessionLog ?? process.env.ASHLR_SESSION_LOG) === "0";
+
+  const rawUsd = opts.budgetUsd ?? process.env.ASHLR_SESSION_BUDGET_USD;
+  const rawTokens = opts.budgetTokens ?? process.env.ASHLR_SESSION_BUDGET_TOKENS;
+
+  // No budget set → always pass.
+  if (!rawUsd && !rawTokens) {
+    return { action: "pass", pct: 0, usedTokens: 0, budgetTokens: 0, usedUsd: 0, budgetUsd: 0 };
+  }
+
+  // No session log data → pass (can't enforce without data).
+  if (sessionLogDisabled) {
+    return { action: "pass", pct: 0, usedTokens: 0, budgetTokens: 0, usedUsd: 0, budgetUsd: 0 };
+  }
+
+  const homeDir = opts.home ?? home();
+  const sessionId = opts.sessionId ?? currentSessionId();
+  const totalBytes = readSessionBytes(homeDir, sessionId);
+  const usedTokens = bytesToTokens(totalBytes);
+  const usedUsd = tokensToUsd(usedTokens);
+
+  let pct = 0;
+  let budgetTokensVal = 0;
+  let budgetUsdVal = 0;
+
+  if (rawUsd) {
+    budgetUsdVal = parseFloat(rawUsd);
+    if (!Number.isFinite(budgetUsdVal) || budgetUsdVal <= 0) {
+      return { action: "pass", pct: 0, usedTokens, budgetTokens: 0, usedUsd, budgetUsd: 0 };
+    }
+    pct = usedUsd / budgetUsdVal;
+    budgetTokensVal = Math.ceil((budgetUsdVal / BLENDED_USD_PER_MTOK) * 1_000_000);
+  } else if (rawTokens) {
+    budgetTokensVal = parseInt(rawTokens, 10);
+    if (!Number.isFinite(budgetTokensVal) || budgetTokensVal <= 0) {
+      return { action: "pass", pct: 0, usedTokens, budgetTokens: 0, usedUsd, budgetUsd: 0 };
+    }
+    pct = usedTokens / budgetTokensVal;
+    budgetUsdVal = tokensToUsd(budgetTokensVal);
+  }
+
+  const base = { usedTokens, budgetTokens: budgetTokensVal, usedUsd, budgetUsd: budgetUsdVal };
+
+  if (pct >= BLOCK_100_PCT) {
+    const msg = rawUsd
+      ? `[ashlr] Budget exceeded — $${usedUsd.toFixed(2)} of $${budgetUsdVal.toFixed(2)} used. Run /ashlr-budget off to clear.`
+      : `[ashlr] Budget exceeded — ${usedTokens.toLocaleString()} of ${budgetTokensVal.toLocaleString()} tokens used. Run /ashlr-budget off to clear.`;
+    emitBudgetEvent("budget_exceeded", pct);
+    return { action: "block", context: msg, pct, ...base };
+  }
+
+  if (pct >= WARN_95_PCT) {
+    const pctStr = Math.round(pct * 100);
+    const msg = rawUsd
+      ? `[ashlr] Budget at ${pctStr}% — $${usedUsd.toFixed(2)} of $${budgetUsdVal.toFixed(2)} used. Consider stopping soon.`
+      : `[ashlr] Budget at ${pctStr}% — ${usedTokens.toLocaleString()} of ${budgetTokensVal.toLocaleString()} tokens used. Consider stopping soon.`;
+    emitBudgetEvent("budget_threshold_95", pct);
+    return { action: "warn-95", context: msg, pct, ...base };
+  }
+
+  if (pct >= WARN_80_PCT) {
+    const pctStr = Math.round(pct * 100);
+    const msg = rawUsd
+      ? `[ashlr] Budget at ${pctStr}% — $${usedUsd.toFixed(2)} of $${budgetUsdVal.toFixed(2)} used.`
+      : `[ashlr] Budget at ${pctStr}% — ${usedTokens.toLocaleString()} of ${budgetTokensVal.toLocaleString()} tokens used.`;
+    emitBudgetEvent("budget_threshold_80", pct);
+    return { action: "warn-80", context: msg, pct, ...base };
+  }
+
+  return { action: "pass", pct, ...base };
+}
+
+// ---------------------------------------------------------------------------
+// Stdin reading
+// ---------------------------------------------------------------------------
+
+async function readStdin(): Promise<string> {
+  const chunks: Buffer[] = [];
+  for await (const chunk of process.stdin as AsyncIterable<Buffer | string>) {
+    chunks.push(typeof chunk === "string" ? Buffer.from(chunk) : chunk);
+  }
+  return Buffer.concat(chunks).toString("utf-8");
+}
+
+// ---------------------------------------------------------------------------
+// Entry point
+// ---------------------------------------------------------------------------
+
+async function main(): Promise<void> {
+  // Read stdin (required by hook contract) but we only need the tool name.
+  try {
+    await readStdin();
+  } catch {
+    // stdin read failure → pass through
+  }
+
+  let result: DecideResult;
+  try {
+    result = decide();
+  } catch {
+    process.stdout.write(JSON.stringify(passThrough()));
+    process.exit(0);
+  }
+
+  if (result!.action === "block") {
+    process.stderr.write(result!.context ?? "[ashlr] Budget exceeded.");
+    process.exit(2);
+  }
+
+  const context = result!.action !== "pass" ? result!.context : undefined;
+  process.stdout.write(JSON.stringify(passThrough(context)));
+  process.exit(0);
+}
+
+if (import.meta.main) {
+  await main();
+}