muonroi-cli 1.2.0

package/src/usage/ledger.test.ts

@@ -0,0 +1,200 @@
+import * as fs from "node:fs/promises";
+import * as os from "node:os";
+import * as path from "node:path";
+import { afterEach, beforeEach, describe, expect, it } from "vitest";
+import { atomicReadJSON } from "../storage/atomic-io.js";
+import type { UsageState } from "../storage/usage-cap.js";
+import { commit, release, reserve } from "./ledger.js";
+import { subscribeThresholds } from "./thresholds.js";
+import type { ReservationToken, ThresholdEvent } from "./types.js";
+import { CapBreachError } from "./types.js";
+
+async function makeTmpHome(capUsd = 15): Promise<string> {
+  const home = await fs.mkdtemp(path.join(os.tmpdir(), "muonroi-ledger-"));
+  await fs.writeFile(path.join(home, "config.json"), JSON.stringify({ cap: { monthly_usd: capUsd } }));
+  return home;
+}
+
+async function readUsage(home: string): Promise<UsageState | null> {
+  return atomicReadJSON<UsageState>(path.join(home, "usage.json"));
+}
+
+describe("ledger", () => {
+  let home: string;
+
+  beforeEach(async () => {
+    home = await makeTmpHome(1.0); // $1.00 cap for easy math
+  });
+
+  describe("reserve()", () => {
+    it("returns ReservationToken when projection < cap", async () => {
+      const result = await reserve({
+        provider: "anthropic",
+        model: "claude-3-5-haiku-latest",
+        estInputTokens: 1000,
+        estOutputTokens: 500,
+        homeOverride: home,
+      });
+      expect(result).not.toBeInstanceOf(CapBreachError);
+      const tok = result as ReservationToken;
+      expect(tok.id).toBeTruthy();
+      expect(tok.model).toBe("claude-3-5-haiku-latest");
+      expect(tok.provider).toBe("anthropic");
+      expect(tok.projected_usd).toBeGreaterThan(0);
+
+      // usage.json should have a reservation
+      const state = await readUsage(home);
+      expect(state!.reservations).toHaveLength(1);
+      expect(state!.reservations[0].id).toBe(tok.id);
+    });
+
+    it("returns CapBreachError when projection would exceed cap", async () => {
+      // claude-3-5-sonnet-latest: $3/M in + $15/M out
+      // 1M input = $3 + 1M output = $15 = $18 total >> $1.00 cap
+      const result = await reserve({
+        provider: "anthropic",
+        model: "claude-3-5-sonnet-latest",
+        estInputTokens: 1_000_000,
+        estOutputTokens: 1_000_000,
+        homeOverride: home,
+      });
+      expect(result).toBeInstanceOf(CapBreachError);
+
+      // usage.json should be unchanged (no reservation added)
+      const state = await readUsage(home);
+      expect(state!.reservations).toHaveLength(0);
+    });
+  });
+
+  describe("commit()", () => {
+    it("removes reservation and increments current_usd by actual cost", async () => {
+      const tok = (await reserve({
+        provider: "anthropic",
+        model: "claude-3-5-haiku-latest",
+        estInputTokens: 10_000,
+        estOutputTokens: 5_000,
+        homeOverride: home,
+      })) as ReservationToken;
+
+      await commit(tok, 8_000, 4_000, home);
+
+      const state = await readUsage(home);
+      expect(state!.reservations).toHaveLength(0);
+      expect(state!.current_usd).toBeGreaterThan(0);
+    });
+
+    it("is idempotent -- second commit on same token is a no-op", async () => {
+      const tok = (await reserve({
+        provider: "anthropic",
+        model: "claude-3-5-haiku-latest",
+        estInputTokens: 10_000,
+        estOutputTokens: 5_000,
+        homeOverride: home,
+      })) as ReservationToken;
+
+      await commit(tok, 8_000, 4_000, home);
+      const stateAfterFirst = await readUsage(home);
+
+      await commit(tok, 8_000, 4_000, home);
+      const stateAfterSecond = await readUsage(home);
+
+      expect(stateAfterSecond!.current_usd).toBe(stateAfterFirst!.current_usd);
+    });
+  });
+
+  describe("release()", () => {
+    it("removes reservation without committing actual cost", async () => {
+      const tok = (await reserve({
+        provider: "anthropic",
+        model: "claude-3-5-haiku-latest",
+        estInputTokens: 10_000,
+        estOutputTokens: 5_000,
+        homeOverride: home,
+      })) as ReservationToken;
+
+      await release(tok, home);
+
+      const state = await readUsage(home);
+      expect(state!.reservations).toHaveLength(0);
+      expect(state!.current_usd).toBe(0);
+    });
+  });
+
+  describe("Pitfall 5: try/finally release pattern", () => {
+    it("reserve -> throw -> finally release cleans up reservation", async () => {
+      let tok: ReservationToken | undefined;
+      try {
+        tok = (await reserve({
+          provider: "anthropic",
+          model: "claude-3-5-haiku-latest",
+          estInputTokens: 10_000,
+          estOutputTokens: 5_000,
+          homeOverride: home,
+        })) as ReservationToken;
+
+        // Simulate an abort
+        throw new Error("stream aborted");
+      } catch {
+        // expected
+      } finally {
+        if (tok) await release(tok, home);
+      }
+
+      const state = await readUsage(home);
+      expect(state!.reservations).toHaveLength(0);
+      expect(state!.current_usd).toBe(0);
+    });
+  });
+
+  describe("subscribeThresholds integration", () => {
+    let unsub: (() => void) | undefined;
+
+    afterEach(() => {
+      if (unsub) unsub();
+    });
+
+    it("fires 80% threshold event when commit pushes past boundary", async () => {
+      // Use a very small cap so haiku costs are meaningful
+      home = await makeTmpHome(0.01); // $0.01 cap
+
+      const events: ThresholdEvent[] = [];
+      unsub = subscribeThresholds((e) => events.push(e));
+
+      // Reserve with haiku: 100k input + 25k output ~ $0.18 >> $0.01 cap
+      // Use smaller tokens to stay under cap for reservation
+      const tok = (await reserve({
+        provider: "anthropic",
+        model: "claude-3-5-haiku-latest",
+        estInputTokens: 1_000,
+        estOutputTokens: 500,
+        homeOverride: home,
+      })) as ReservationToken;
+
+      // Commit with actual tokens that push past 50% and 80%
+      // haiku: $0.80/M input, $4.00/M output
+      // 1000 in + 500 out = $0.0008 + $0.002 = $0.0028
+      // $0.0028 / $0.01 = 28% -- not enough for 50%
+
+      // Let's use a bigger cap scenario
+      await release(tok, home);
+
+      // Re-create home with better numbers
+      home = await makeTmpHome(0.005); // $0.005 cap
+
+      const tok2 = (await reserve({
+        provider: "anthropic",
+        model: "claude-3-5-haiku-latest",
+        estInputTokens: 1_000,
+        estOutputTokens: 500,
+        homeOverride: home,
+      })) as ReservationToken;
+
+      // Commit: 1000 in + 500 out = $0.0028
+      // $0.0028 / $0.005 = 56% -- should fire 50% threshold
+      await commit(tok2, 1_000, 500, home);
+
+      expect(events.length).toBeGreaterThanOrEqual(1);
+      expect(events.some((e) => e.level === 50)).toBe(true);
+    });
+  });
+});

package/src/usage/ledger.ts

@@ -0,0 +1,213 @@
+/**
+ * src/usage/ledger.ts
+ *
+ * Reservation ledger for USAGE-03.
+ * Atomic reserve/commit/release primitives with file-lock via proper-lockfile.
+ *
+ * Invariant: current_usd + sum(reservations.usd) + projected <= cap.monthly_usd
+ * or reserve() returns CapBreachError.
+ *
+ * Pitfall 5 mitigation: callers MUST use try/finally with release() in finally
+ * to prevent reservation leaks on stream abort.
+ *
+ * Pitfall 7 mitigation: proper-lockfile exclusive lock serializes concurrent
+ * reserve() calls so 10-parallel tool-call bursts cannot collectively exceed cap.
+ */
+
+import { promises as fs } from "node:fs";
+import * as os from "node:os";
+import * as path from "node:path";
+import lockfile from "proper-lockfile";
+import { atomicReadJSON, atomicWriteJSON } from "../storage/atomic-io.js";
+import type { UsageState } from "../storage/usage-cap.js";
+import { projectCostUSD } from "./estimator.js";
+import { emit, evaluateThresholds } from "./thresholds.js";
+import { CapBreachError, type ReservationToken } from "./types.js";
+
+const DEFAULT_CAP_USD = 15;
+
+function muonroiHome(homeOverride?: string): string {
+  return homeOverride ?? process.env.MUONROI_CLI_HOME ?? path.join(os.homedir(), ".muonroi-cli");
+}
+
+async function loadCapUSD(homeOverride?: string): Promise<number> {
+  const cfg = await atomicReadJSON<{ cap?: { monthly_usd?: number } }>(
+    path.join(muonroiHome(homeOverride), "config.json"),
+  );
+  return cfg?.cap?.monthly_usd ?? DEFAULT_CAP_USD;
+}
+
+function emptyState(): UsageState {
+  const d = new Date();
+  const yyyy = d.getUTCFullYear();
+  const mm = String(d.getUTCMonth() + 1).padStart(2, "0");
+  return {
+    current_month_utc: `${yyyy}-${mm}`,
+    current_usd: 0,
+    reservations: [],
+    thresholds_fired_this_month: [],
+  };
+}
+
+async function ensureUsageFile(filePath: string): Promise<void> {
+  await fs.mkdir(path.dirname(filePath), { recursive: true });
+  try {
+    await fs.access(filePath);
+  } catch {
+    // Race-safe: if another process created the file between access() and writeFile(),
+    // the atomicWriteJSON may fail on rename. Swallow ENOENT/EBUSY and re-check.
+    try {
+      await atomicWriteJSON(filePath, emptyState());
+    } catch {
+      // If file now exists (race loser), that's fine -- proceed.
+      try {
+        await fs.access(filePath);
+      } catch {
+        throw new Error(`Failed to initialize usage file at ${filePath}`);
+      }
+    }
+  }
+}
+
+/**
+ * Execute a function under exclusive file lock on usage.json.
+ * The lock prevents racing readers/writers across CLI processes.
+ */
+async function withLock<T>(
+  filePath: string,
+  fn: (state: UsageState) => Promise<{ next: UsageState; result: T }>,
+): Promise<T> {
+  await ensureUsageFile(filePath);
+  const releaseLock = await lockfile.lock(filePath, {
+    retries: { retries: 10, minTimeout: 10, maxTimeout: 100 },
+    stale: 5_000,
+    realpath: false,
+  });
+  try {
+    const state = (await atomicReadJSON<UsageState>(filePath)) ?? emptyState();
+    const { next, result } = await fn(state);
+    await atomicWriteJSON(filePath, next);
+    return result;
+  } finally {
+    await releaseLock();
+  }
+}
+
+/**
+ * Reserve projected token spend against the monthly cap.
+ * Returns ReservationToken on success, CapBreachError if projection would exceed cap.
+ *
+ * Caller contract: MUST call commit() or release() on the token.
+ * Use try/finally: `const tok = await reserve(...); try { ... commit(tok) } finally { release(tok) }`
+ */
+export async function reserve(args: {
+  provider: string;
+  model: string;
+  estInputTokens: number;
+  estOutputTokens: number;
+  homeOverride?: string;
+}): Promise<ReservationToken | CapBreachError> {
+  const filePath = path.join(muonroiHome(args.homeOverride), "usage.json");
+  const cap = await loadCapUSD(args.homeOverride);
+
+  return withLock<ReservationToken | CapBreachError>(filePath, async (state) => {
+    const projected = projectCostUSD(args.provider, args.model, args.estInputTokens, args.estOutputTokens);
+    const reservedTotal = state.reservations.reduce((s, r) => s + r.usd, 0);
+
+    if (state.current_usd + reservedTotal + projected > cap) {
+      return {
+        next: state,
+        result: new CapBreachError(state.current_usd, reservedTotal, projected, cap),
+      };
+    }
+
+    const id = crypto.randomUUID();
+    const tok: ReservationToken = {
+      id,
+      model: args.model,
+      provider: args.provider,
+      projected_usd: projected,
+      est_input_tokens: args.estInputTokens,
+      est_output_tokens: args.estOutputTokens,
+      createdAtMs: Date.now(),
+    };
+
+    const next: UsageState = {
+      ...state,
+      reservations: [
+        ...state.reservations,
+        {
+          id,
+          usd: projected,
+          createdAtMs: tok.createdAtMs,
+          model: tok.model,
+          provider: tok.provider,
+          est_input_tokens: tok.est_input_tokens,
+          est_output_tokens: tok.est_output_tokens,
+        },
+      ],
+    };
+
+    return { next, result: tok };
+  });
+}
+
+/**
+ * Commit actual token spend. Removes the reservation and increments current_usd.
+ * Idempotent: if the reservation ID is not found, this is a no-op.
+ *
+ * Evaluates thresholds after state mutation and emits events post-lock-release.
+ */
+export async function commit(
+  token: ReservationToken,
+  actualInputTokens: number,
+  actualOutputTokens: number,
+  homeOverride?: string,
+): Promise<void> {
+  const filePath = path.join(muonroiHome(homeOverride), "usage.json");
+  const cap = await loadCapUSD(homeOverride);
+
+  const pendingEvents = await withLock(filePath, async (state) => {
+    const idx = state.reservations.findIndex((r) => r.id === token.id);
+    if (idx === -1) return { next: state, result: [] as import("./types.js").ThresholdEvent[] };
+
+    const actual = projectCostUSD(token.provider, token.model, actualInputTokens, actualOutputTokens);
+    const reservations = state.reservations.filter((r) => r.id !== token.id);
+    const prevUsd = state.current_usd;
+    const nextUsd = state.current_usd + actual;
+
+    // Evaluate thresholds
+    const thresholdResult = evaluateThresholds({
+      prevUsd,
+      nextUsd,
+      capUsd: cap,
+      firedThisMonth: state.thresholds_fired_this_month ?? [],
+    });
+
+    const next: UsageState = {
+      ...state,
+      current_usd: nextUsd,
+      reservations,
+      thresholds_fired_this_month: thresholdResult.nextFired,
+    };
+
+    return { next, result: thresholdResult.events };
+  });
+
+  // Emit events AFTER lock release to avoid holding the lock during listener callbacks
+  for (const ev of pendingEvents) {
+    emit(ev);
+  }
+}
+
+/**
+ * Release a reservation without committing actual cost.
+ * Used on stream abort (Pitfall 5 mitigation).
+ */
+export async function release(token: ReservationToken, homeOverride?: string): Promise<void> {
+  const filePath = path.join(muonroiHome(homeOverride), "usage.json");
+  await withLock(filePath, async (state) => {
+    const reservations = state.reservations.filter((r) => r.id !== token.id);
+    return { next: { ...state, reservations }, result: undefined };
+  });
+}

package/src/usage/midstream.test.ts

@@ -0,0 +1,55 @@
+import { beforeEach, describe, expect, it } from "vitest";
+import { midstreamPolicy } from "./midstream.js";
+import { emit } from "./thresholds.js";
+import type { ThresholdEvent } from "./types.js";
+
+describe("midstreamPolicy", () => {
+  beforeEach(() => {
+    midstreamPolicy.clear();
+  });
+
+  it("refuseNext() is false initially", () => {
+    expect(midstreamPolicy.refuseNext()).toBe(false);
+  });
+
+  it("refuseNext() flips to true after {level:100} threshold event", () => {
+    const ev: ThresholdEvent = {
+      level: 100,
+      current_pct: 100,
+      current_usd: 15,
+      cap_usd: 15,
+      atMs: Date.now(),
+    };
+    emit(ev);
+    expect(midstreamPolicy.refuseNext()).toBe(true);
+  });
+
+  it("does NOT flip on {level:50} or {level:80}", () => {
+    emit({ level: 50, current_pct: 50, current_usd: 7.5, cap_usd: 15, atMs: Date.now() });
+    expect(midstreamPolicy.refuseNext()).toBe(false);
+
+    emit({ level: 80, current_pct: 80, current_usd: 12, cap_usd: 15, atMs: Date.now() });
+    expect(midstreamPolicy.refuseNext()).toBe(false);
+  });
+
+  it("isStreamFinishAllowed() always returns true (in-flight finishes)", () => {
+    expect(midstreamPolicy.isStreamFinishAllowed()).toBe(true);
+    // Even after 100% breach
+    emit({ level: 100, current_pct: 101, current_usd: 15.15, cap_usd: 15, atMs: Date.now() });
+    expect(midstreamPolicy.isStreamFinishAllowed()).toBe(true);
+  });
+
+  it("currentPct() tracks the latest threshold event pct", () => {
+    expect(midstreamPolicy.currentPct()).toBe(0);
+    emit({ level: 50, current_pct: 52.3, current_usd: 7.8, cap_usd: 15, atMs: Date.now() });
+    expect(midstreamPolicy.currentPct()).toBeCloseTo(52.3);
+  });
+
+  it("clear() resets state (for tests + month rollover)", () => {
+    midstreamPolicy.forceRefuseNext();
+    expect(midstreamPolicy.refuseNext()).toBe(true);
+    midstreamPolicy.clear();
+    expect(midstreamPolicy.refuseNext()).toBe(false);
+    expect(midstreamPolicy.currentPct()).toBe(0);
+  });
+});

package/src/usage/midstream.ts

@@ -0,0 +1,51 @@
+/**
+ * src/usage/midstream.ts
+ *
+ * Mid-stream policy for USAGE-05.
+ * In-flight stream finishes after threshold breach; next reservation refused.
+ * Acceptable single-stream overshoot ~101%.
+ *
+ * Subscribes to threshold events: {level:100} sets refuseNext=true.
+ */
+
+import { subscribeThresholds } from "./thresholds.js";
+
+let _refuseNext = false;
+let _capPct = 0;
+
+// Subscribe once at module load
+subscribeThresholds((ev) => {
+  _capPct = ev.current_pct;
+  if (ev.level === 100) _refuseNext = true;
+});
+
+export const midstreamPolicy = {
+  /** Whether the next stream reservation should be refused. */
+  refuseNext(): boolean {
+    return _refuseNext;
+  },
+
+  /** Force refuse (called when downgrade chain exhausted). */
+  forceRefuseNext(): void {
+    _refuseNext = true;
+  },
+
+  /** Reset state — for tests and month rollover. */
+  clear(): void {
+    _refuseNext = false;
+    _capPct = 0;
+  },
+
+  /** Current cap percentage from latest threshold event. */
+  currentPct(): number {
+    return _capPct;
+  },
+
+  /**
+   * In-flight stream may finish — caller checks BEFORE starting a new stream.
+   * Always returns true: once a stream is in-flight, it completes.
+   */
+  isStreamFinishAllowed(): boolean {
+    return true;
+  },
+};

package/src/usage/thresholds.test.ts

@@ -0,0 +1,83 @@
+import { describe, expect, it } from "vitest";
+import { evaluateThresholds } from "./thresholds.js";
+
+describe("thresholds", () => {
+  describe("evaluateThresholds", () => {
+    it("fires ThresholdEvent at 50% crossing", () => {
+      const result = evaluateThresholds({
+        prevUsd: 0.4,
+        nextUsd: 0.55,
+        capUsd: 1.0,
+        firedThisMonth: [],
+      });
+      expect(result.events).toHaveLength(1);
+      expect(result.events[0].level).toBe(50);
+      expect(result.nextFired).toContain(50);
+    });
+
+    it("does not re-fire 50% if already in firedThisMonth", () => {
+      const result = evaluateThresholds({
+        prevUsd: 0.55,
+        nextUsd: 0.6,
+        capUsd: 1.0,
+        firedThisMonth: [50],
+      });
+      expect(result.events).toHaveLength(0);
+    });
+
+    it("fires 80% and 100% in sequence", () => {
+      const r80 = evaluateThresholds({
+        prevUsd: 0.7,
+        nextUsd: 0.85,
+        capUsd: 1.0,
+        firedThisMonth: [50],
+      });
+      expect(r80.events).toHaveLength(1);
+      expect(r80.events[0].level).toBe(80);
+
+      const r100 = evaluateThresholds({
+        prevUsd: 0.85,
+        nextUsd: 1.05,
+        capUsd: 1.0,
+        firedThisMonth: r80.nextFired,
+      });
+      expect(r100.events).toHaveLength(1);
+      expect(r100.events[0].level).toBe(100);
+    });
+
+    it("fires multiple thresholds in one jump", () => {
+      const result = evaluateThresholds({
+        prevUsd: 0.0,
+        nextUsd: 1.05,
+        capUsd: 1.0,
+        firedThisMonth: [],
+      });
+      expect(result.events).toHaveLength(3);
+      expect(result.events.map((e) => e.level)).toEqual([50, 80, 100]);
+    });
+
+    it("resets on month rollover (empty firedThisMonth)", () => {
+      // After month rollover, firedThisMonth is empty again
+      const result = evaluateThresholds({
+        prevUsd: 0.0,
+        nextUsd: 0.55,
+        capUsd: 1.0,
+        firedThisMonth: [], // reset after rollover
+      });
+      expect(result.events).toHaveLength(1);
+      expect(result.events[0].level).toBe(50);
+    });
+
+    it("includes correct current_pct and cap_usd in event", () => {
+      const result = evaluateThresholds({
+        prevUsd: 0.0,
+        nextUsd: 0.55,
+        capUsd: 1.0,
+        firedThisMonth: [],
+      });
+      expect(result.events[0].current_pct).toBeCloseTo(55, 0);
+      expect(result.events[0].cap_usd).toBe(1.0);
+      expect(result.events[0].current_usd).toBe(0.55);
+    });
+  });
+});

package/src/usage/thresholds.ts

@@ -0,0 +1,74 @@
+/**
+ * src/usage/thresholds.ts
+ *
+ * Threshold event system for USAGE-02.
+ * Fires at 50/80/100% of monthly cap crossing.
+ * Events are deduplicated per month via thresholds_fired_this_month in UsageState.
+ *
+ * Plan 06 (status bar) and Plan 05 (downgrade chain) consume these events.
+ */
+
+import type { ThresholdEvent, ThresholdLevel } from "./types.js";
+
+type Listener = (e: ThresholdEvent) => void;
+const listeners = new Set<Listener>();
+
+/**
+ * Subscribe to threshold crossing events.
+ * Returns an unsubscribe function.
+ */
+export function subscribeThresholds(fn: Listener): () => void {
+  listeners.add(fn);
+  return () => {
+    listeners.delete(fn);
+  };
+}
+
+/**
+ * Emit a threshold event to all subscribers.
+ * Called by ledger.ts commit() after lock release.
+ */
+export function emit(e: ThresholdEvent): void {
+  for (const l of listeners) l(e);
+}
+
+const LEVELS: ThresholdLevel[] = [50, 80, 100];
+
+export interface EvalArgs {
+  prevUsd: number;
+  nextUsd: number;
+  capUsd: number;
+  firedThisMonth: number[];
+}
+
+export interface EvalResult {
+  events: ThresholdEvent[];
+  nextFired: number[];
+}
+
+/**
+ * Evaluate which threshold boundaries were crossed.
+ * Returns events to emit and the updated fired-this-month array.
+ * Pure function -- no side effects.
+ */
+export function evaluateThresholds(args: EvalArgs): EvalResult {
+  const out: ThresholdEvent[] = [];
+  const nextFired = [...args.firedThisMonth];
+
+  for (const lv of LEVELS) {
+    const boundary = (args.capUsd * lv) / 100;
+    if (args.prevUsd < boundary && args.nextUsd >= boundary && !nextFired.includes(lv)) {
+      const ev: ThresholdEvent = {
+        level: lv,
+        current_pct: (args.nextUsd / args.capUsd) * 100,
+        current_usd: args.nextUsd,
+        cap_usd: args.capUsd,
+        atMs: Date.now(),
+      };
+      out.push(ev);
+      nextFired.push(lv);
+    }
+  }
+
+  return { events: out, nextFired };
+}