offrouter-core 0.0.0 → 0.2.0

package/src/usage-file.ts

@@ -0,0 +1,435 @@
+/**
+ * File-backed usage store for OffRouter.
+ *
+ * Wraps an in-memory provider-health/quota engine and adds JSON persistence for
+ * per-account quota usage so near-limit detection survives across CLI
+ * invocations and MCP server restarts.
+ *
+ * On-disk layout (`$OFFROUTER_HOME/state/usage.json`):
+ *   { version: 1, accounts: Record<`${providerId}:${accountId}`, PersistedAccount> }
+ *
+ * Safety guarantees mirror FileSecretStore:
+ *   - atomic write via temp file + rename
+ *   - chmod 0600 on POSIX after every write
+ *   - missing or corrupt file => graceful fallback to empty state (never throws)
+ *
+ * Provider health/quota math (sliding-window error budget, cooldown) stays
+ * in-memory and process-local; only account quota usage is persisted, which is
+ * what near-limit detection and quota reporting depend on across sessions.
+ */
+import { readFileSync } from "node:fs";
+import {
+  chmod,
+  mkdir,
+  readFile,
+  rename,
+  unlink,
+  writeFile,
+} from "node:fs/promises";
+import { dirname } from "node:path";
+import { randomBytes } from "node:crypto";
+import {
+  computeAccountSnapshot,
+  InMemoryUsageStore,
+  type AccountUsageDoc,
+  type AccountUsageSnapshot,
+  type PaidFallbackExplanation,
+  type ProviderUsageSnapshot,
+  type RecordUsageResult,
+  type UsageEvent,
+  type UsageStore,
+} from "./usage.js";
+import type { ProviderHealth } from "./types.js";
+
+export { DEFAULT_NEAR_LIMIT_THRESHOLD } from "./usage.js";
+export { InMemoryUsageStore } from "./usage.js";
+export type { UsageStore } from "./usage.js";
+
+export interface FileUsageStoreFileSystem {
+  mkdir: typeof mkdir;
+  readFile: typeof readFile;
+  rename: typeof rename;
+  unlink: typeof unlink;
+  writeFile: typeof writeFile;
+  chmod: typeof chmod;
+}
+
+export interface FileUsageStoreOptions {
+  /** Absolute path to usage.json (typically `$OFFROUTER_HOME/state/usage.json`). */
+  filePath: string;
+  /** Injectable file operations for failure-path tests. Defaults to node fs. */
+  fileSystem?: Partial<FileUsageStoreFileSystem>;
+  /** Injectable clock for tests. */
+  now?: () => Date;
+  /**
+   * In-memory store that owns provider health/quota logic. Account usage is
+   * persisted separately and mirrored onto this delegate. Useful for tests.
+   */
+  delegate?: UsageStore;
+}
+
+/** On-disk schema for FileUsageStore. */
+interface PersistedAccount {
+  providerId: string;
+  accountId: string;
+  used: number;
+  limit?: number;
+  lastUpdated: string;
+  health?: ProviderHealth;
+  cooldownUntil?: number;
+  recentErrors?: number[];
+}
+
+interface PersistedUsageDocument {
+  version: 1;
+  accounts: Record<string, PersistedAccount>;
+}
+
+const EMPTY_DOC: PersistedUsageDocument = { version: 1, accounts: {} };
+
+function accountKey(providerId: string, accountId: string): string {
+  return `${providerId}:${accountId}`;
+}
+
+/**
+ * Parse a persisted usage document leniently. A missing or corrupt file is
+ * treated as empty state rather than a hard failure so a bad write can never
+ * wedge the CLI.
+ */
+function parseUsageDocument(raw: string): PersistedUsageDocument {
+  let parsed: unknown;
+  try {
+    parsed = JSON.parse(raw) as unknown;
+  } catch {
+    // Corrupt JSON: fall back to empty instead of throwing.
+    return { ...EMPTY_DOC, accounts: {} };
+  }
+
+  if (typeof parsed !== "object" || parsed === null || Array.isArray(parsed)) {
+    return { ...EMPTY_DOC, accounts: {} };
+  }
+
+  const record = parsed as Record<string, unknown>;
+  const accountsRaw = record.accounts;
+  if (
+    typeof accountsRaw !== "object" ||
+    accountsRaw === null ||
+    Array.isArray(accountsRaw)
+  ) {
+    return { ...EMPTY_DOC, accounts: {} };
+  }
+
+  const accounts: Record<string, PersistedAccount> = {};
+  for (const [key, value] of Object.entries(
+    accountsRaw as Record<string, unknown>,
+  )) {
+    if (typeof value !== "object" || value === null || Array.isArray(value)) {
+      continue;
+    }
+    const entry = value as Record<string, unknown>;
+    if (
+      typeof entry.providerId !== "string" ||
+      typeof entry.accountId !== "string" ||
+      typeof entry.used !== "number"
+    ) {
+      continue;
+    }
+    const doc: PersistedAccount = {
+      providerId: entry.providerId,
+      accountId: entry.accountId,
+      used: entry.used,
+      lastUpdated:
+        typeof entry.lastUpdated === "string" ? entry.lastUpdated : "",
+    };
+    if (typeof entry.limit === "number") {
+      doc.limit = entry.limit;
+    }
+    accounts[key] = doc;
+  }
+
+  return { version: 1, accounts };
+}
+
+/**
+ * File-backed UsageStore. Persists per-account quota usage atomically and
+ * delegates provider health/quota logic to an in-memory engine.
+ */
+export class FileUsageStore implements UsageStore {
+  readonly filePath: string;
+  readonly #fs: FileUsageStoreFileSystem;
+  readonly #now: () => Date;
+  readonly #delegate: UsageStore;
+  /** In-memory mirror of persisted account docs, hydrated lazily. */
+  #accounts: Map<string, AccountUsageDoc> | null = null;
+  /** Serializes mutations so concurrent writes never interleave. */
+  #chain: Promise<unknown> = Promise.resolve();
+
+  constructor(options: FileUsageStoreOptions) {
+    if (!options?.filePath || typeof options.filePath !== "string") {
+      throw new TypeError("FileUsageStore requires options.filePath");
+    }
+    this.filePath = options.filePath;
+    this.#now = options.now ?? (() => new Date());
+    this.#delegate =
+      options.delegate ?? new InMemoryUsageStore({ now: this.#now });
+    this.#fs = {
+      mkdir,
+      readFile,
+      rename,
+      unlink,
+      writeFile,
+      chmod,
+      ...options.fileSystem,
+    };
+  }
+
+  #exclusive<T>(fn: () => T | Promise<T>): Promise<T> {
+    const run = this.#chain.then(fn, fn);
+    this.#chain = run.then(
+      () => undefined,
+      () => undefined,
+    );
+    return run;
+  }
+
+  /** Load the file once into the in-memory account map. Never throws. */
+  async #load(): Promise<Map<string, AccountUsageDoc>> {
+    if (this.#accounts) {
+      return this.#accounts;
+    }
+    let doc: PersistedUsageDocument;
+    try {
+      const raw = await this.#fs.readFile(this.filePath, "utf8");
+      doc = parseUsageDocument(raw);
+    } catch (err) {
+      if (
+        err &&
+        typeof err === "object" &&
+        "code" in err &&
+        (err as { code?: string }).code === "ENOENT"
+      ) {
+        doc = { ...EMPTY_DOC, accounts: {} };
+      } else {
+        // Any read failure (including corruption surfaced as a throw) is
+        // treated as empty state so the CLI keeps working.
+        doc = { ...EMPTY_DOC, accounts: {} };
+      }
+    }
+    const map = new Map<string, AccountUsageDoc>();
+    for (const [key, entry] of Object.entries(doc.accounts)) {
+      map.set(key, {
+        providerId: entry.providerId,
+        accountId: entry.accountId,
+        used: entry.used,
+        limit: entry.limit,
+        updatedAt: entry.lastUpdated,
+      });
+    }
+    this.#accounts = map;
+    return map;
+  }
+
+  /** Synchronous load for toJSON diagnostics; never throws. */
+  #loadKnownSync(): Map<string, AccountUsageDoc> {
+    if (this.#accounts) {
+      return this.#accounts;
+    }
+    let map = new Map<string, AccountUsageDoc>();
+    try {
+      const raw = readFileSync(this.filePath, "utf8");
+      const doc = parseUsageDocument(raw);
+      map = new Map();
+      for (const [key, entry] of Object.entries(doc.accounts)) {
+        map.set(key, {
+          providerId: entry.providerId,
+          accountId: entry.accountId,
+          used: entry.used,
+          limit: entry.limit,
+          updatedAt: entry.lastUpdated,
+        });
+      }
+    } catch {
+      map = new Map();
+    }
+    return map;
+  }
+
+  async #persist(accounts: Map<string, AccountUsageDoc>): Promise<void> {
+    const persistedAccounts: Record<string, PersistedAccount> = {};
+    for (const [key, doc] of accounts) {
+      const snap = computeAccountSnapshot(doc);
+      const entry: PersistedAccount = {
+        providerId: doc.providerId,
+        accountId: doc.accountId,
+        used: doc.used,
+        lastUpdated: doc.updatedAt,
+        health: snap.health,
+        recentErrors: [],
+      };
+      if (doc.limit !== undefined) {
+        entry.limit = doc.limit;
+      }
+      persistedAccounts[key] = entry;
+    }
+    const payload = `${JSON.stringify(
+      { version: 1 as const, accounts: persistedAccounts },
+      null,
+      2,
+    )}\n`;
+
+    await this.#fs.mkdir(dirname(this.filePath), { recursive: true });
+    const tmp = `${this.filePath}.${randomBytes(8).toString("hex")}.tmp`;
+    try {
+      await this.#fs.writeFile(tmp, payload, {
+        encoding: "utf8",
+        mode: 0o600,
+      });
+      await this.#fs.chmod(tmp, 0o600);
+      await this.#fs.rename(tmp, this.filePath);
+      await this.#fs.chmod(this.filePath, 0o600);
+    } catch (err) {
+      await this.#fs.unlink(tmp).catch((unlinkErr: unknown) => {
+        if (
+          unlinkErr &&
+          typeof unlinkErr === "object" &&
+          "code" in unlinkErr &&
+          (unlinkErr as { code?: string }).code === "ENOENT"
+        ) {
+          return;
+        }
+        throw unlinkErr;
+      });
+      throw err;
+    }
+  }
+
+  // ---- Provider health/quota: delegated to the in-memory engine ----
+
+  record(event: UsageEvent): Promise<RecordUsageResult> {
+    return this.#delegate.record(event);
+  }
+
+  getProvider(providerId: string): Promise<ProviderUsageSnapshot | undefined> {
+    return this.#delegate.getProvider(providerId);
+  }
+
+  listProviders(): Promise<ProviderUsageSnapshot[]> {
+    return this.#delegate.listProviders();
+  }
+
+  setSubscriptionQuota(
+    providerId: string,
+    input: {
+      limit?: number;
+      used?: number;
+      status?: import("./types.js").SubscriptionStatus;
+      modelId?: string;
+      at?: Date;
+    },
+  ): Promise<ProviderUsageSnapshot> {
+    return this.#delegate.setSubscriptionQuota(providerId, input);
+  }
+
+  explainPaidFallback(
+    providerId: string,
+    options?: { hasApiKeyCandidate?: boolean },
+  ): Promise<PaidFallbackExplanation> {
+    return this.#delegate.explainPaidFallback(providerId, options);
+  }
+
+  // ---- Per-account quota usage: persisted ----
+
+  async recordUsage(
+    providerId: string,
+    accountId: string,
+    input: { used?: number; limit?: number; at?: Date },
+  ): Promise<AccountUsageSnapshot> {
+    return this.#exclusive(async () => {
+      const accounts = await this.#load();
+      const at = input.at ?? this.#now();
+      const key = accountKey(providerId, accountId);
+      let doc = accounts.get(key);
+      if (!doc) {
+        doc = {
+          providerId,
+          accountId,
+          used: 0,
+          updatedAt: at.toISOString(),
+        };
+        accounts.set(key, doc);
+      }
+      if (input.used !== undefined) doc.used = input.used;
+      if (input.limit !== undefined) doc.limit = input.limit;
+      doc.updatedAt = at.toISOString();
+      await this.#persist(accounts);
+      return computeAccountSnapshot(doc);
+    });
+  }
+
+  async isNearLimit(
+    providerId: string,
+    accountId: string,
+    threshold?: number,
+  ): Promise<boolean> {
+    return this.#exclusive(async () => {
+      const accounts = await this.#load();
+      const doc = accounts.get(accountKey(providerId, accountId));
+      if (!doc) return false;
+      return computeAccountSnapshot(doc, threshold).nearLimit;
+    });
+  }
+
+  async getNearLimitAccounts(
+    threshold?: number,
+  ): Promise<AccountUsageSnapshot[]> {
+    return this.#exclusive(async () => {
+      const accounts = await this.#load();
+      const out: AccountUsageSnapshot[] = [];
+      for (const doc of accounts.values()) {
+        const snap = computeAccountSnapshot(doc, threshold);
+        if (snap.nearLimit) out.push(snap);
+      }
+      return out.sort((a, b) => {
+        if (a.providerId !== b.providerId)
+          return a.providerId.localeCompare(b.providerId);
+        return a.accountId.localeCompare(b.accountId);
+      });
+    });
+  }
+
+  async getAccountUsage(
+    providerId: string,
+    accountId: string,
+  ): Promise<AccountUsageSnapshot | undefined> {
+    return this.#exclusive(async () => {
+      const accounts = await this.#load();
+      const doc = accounts.get(accountKey(providerId, accountId));
+      if (!doc) return undefined;
+      return computeAccountSnapshot(doc);
+    });
+  }
+
+  async listAccountUsage(): Promise<AccountUsageSnapshot[]> {
+    return this.#exclusive(async () => {
+      const accounts = await this.#load();
+      const out: AccountUsageSnapshot[] = [];
+      for (const doc of accounts.values()) {
+        out.push(computeAccountSnapshot(doc));
+      }
+      return out.sort((a, b) => {
+        if (a.providerId !== b.providerId)
+          return a.providerId.localeCompare(b.providerId);
+        return a.accountId.localeCompare(b.accountId);
+      });
+    });
+  }
+
+  toJSON(): Record<string, unknown> {
+    const accounts = this.#loadKnownSync();
+    return {
+      kind: "file-usage",
+      filePath: this.filePath,
+      accountCount: accounts.size,
+    };
+  }
+}

package/src/usage.test.ts

@@ -281,3 +281,55 @@ describe("InMemoryStateStore", () => {
     expect(JSON.stringify(json)).not.toMatch(/sk-/);
   });
 });
+
+describe("near-limit / per-account usage", () => {
+  it("recordUsage returns near-limit snapshot when quota is nearly consumed", async () => {
+    const store = new InMemoryUsageStore();
+    const snap = await store.recordUsage("anthropic", "anthropic-1", { used: 85, limit: 100 });
+    expect(snap.remaining).toBe(15);
+    expect(snap.percentRemaining).toBe(15);
+    expect(snap.nearLimit).toBe(true);
+    expect(snap.subscriptionStatus).toBe("near-limit");
+  });
+
+  it("isNearLimit respects custom threshold", async () => {
+    const store = new InMemoryUsageStore();
+    await store.recordUsage("anthropic", "anthropic-1", { used: 85, limit: 100 });
+    expect(await store.isNearLimit("anthropic", "anthropic-1")).toBe(true);
+    expect(await store.isNearLimit("anthropic", "anthropic-1", 0.1)).toBe(false);
+  });
+
+  it("recordUsage detects exhaustion", async () => {
+    const store = new InMemoryUsageStore();
+    const snap = await store.recordUsage("anthropic", "anthropic-2", { used: 100, limit: 100 });
+    expect(snap.nearLimit).toBe(false);
+    expect(snap.subscriptionStatus).toBe("exhausted");
+  });
+
+  it("getNearLimitAccounts returns only accounts near their limit", async () => {
+    const store = new InMemoryUsageStore();
+    await store.recordUsage("anthropic", "anthropic-1", { used: 85, limit: 100 });
+    await store.recordUsage("anthropic", "anthropic-2", { used: 100, limit: 100 });
+    const nearLimit = await store.getNearLimitAccounts();
+    expect(nearLimit).toHaveLength(1);
+    expect(nearLimit[0]!.accountId).toBe("anthropic-1");
+  });
+
+  it("getAccountUsage returns correct snapshots for different accounts", async () => {
+    const store = new InMemoryUsageStore();
+    await store.recordUsage("anthropic", "anthropic-1", { used: 85, limit: 100 });
+    await store.recordUsage("anthropic", "anthropic-2", { used: 10, limit: 100 });
+    const snap1 = await store.getAccountUsage("anthropic", "anthropic-1");
+    const snap2 = await store.getAccountUsage("anthropic", "anthropic-2");
+    expect(snap1?.used).toBe(85);
+    expect(snap2?.used).toBe(10);
+  });
+
+  it("existing setSubscriptionQuota and record still behave as before", async () => {
+    const store = new InMemoryUsageStore();
+    await store.setSubscriptionQuota("test", { limit: 50, used: 10, status: "active" });
+    const snap = await store.getProvider("test");
+    expect(snap?.subscriptionQuotaLimit).toBe(50);
+    expect(snap?.subscriptionQuotaUsed).toBe(10);
+  });
+});