platform/src/servers/api/chat/usage-store.ts

import { listProviderUsage, recordProviderUsage } from 'officerdb';
import { readClaudeUsage, type UsageDay, type UsageLimit, type UsageModel } from './claude-usage';
import { logger } from './logger';

/**
 * Usage readings across every AI account this machine has seen — not just the one signed in.
 *
 * ── The problem this exists for ──
 *
 * Half a usage reading cannot be taken at all for an account you are not signed into. The limits come from
 * Anthropic's OAuth endpoint using the credential the CLI holds for the CURRENT login, and the token counts
 * come from that login's transcripts. Switch accounts and the previous one is not slow to read — it is
 * unreadable.
 *
 * The owner switches between a personal account and a company one, and the question worth answering is
 * "which of my accounts has room right now". Only a stored reading per account can answer it, so a snapshot
 * is taken at the end of every turn and kept in `ai_provider_usage`.
 *
 * ── Why the countdowns still work on a stale row ──
 *
 * Because `resetsAt` is an ABSOLUTE instant, not a duration. A row measured six hours ago has a percentage
 * that is as-of, and a "resets in 14h 20m" that is exactly right — the client ticks it off the stored
 * timestamp with no network and no signed-in session. That asymmetry is the whole reason a cached row is
 * worth showing rather than apologising for.
 */

export const CLAUDE_PROVIDER = 'claude-code';

/** The reading as stored — the three display sections, without the account identity wrapped around them. */
export type UsageSnapshot = {
  limits: UsageLimit[];
  days: UsageDay[];
  models: UsageModel[];
  limitsProblem?: string;
};

export type UsageAccount = {
  provider: string;
  email: string;
  plan: string | null;
  /** Whether this is the account the CLI is signed in as right now. Exactly one row is, or none. */
  current: boolean;
  /** When the reading was MEASURED. Drives "as of N minutes ago" and nothing else. */
  measuredAt: string;
} & UsageSnapshot;

/**
 * Take a reading for whoever is signed in, and store it.
 *
 * Best-effort by design and it never throws: the caller is a turn ending, and a chat turn must not fail
 * because a display row would not write. A missed snapshot costs nothing — the next turn takes another.
 *
 * Returns the account it stored, or null when there was nothing to store (no login, so no account to key a
 * row on).
 */
export async function snapshotClaudeUsage(userId: number): Promise<UsageAccount | null> {
  let usage: Awaited<ReturnType<typeof readClaudeUsage>>;
  try {
    usage = await readClaudeUsage();
  } catch (err) {
    logger.warn('Could not read Claude usage', { userId, error: String(err) });
    return null;
  }

  const email = usage.account.email;
  // No account means no key. A reading that cannot be attributed is worse than no reading: it would land
  // under whichever account happened to be signed in last.
  if (!email) return null;

  const snapshot: UsageSnapshot = {
    limits: usage.limits,
    days: usage.days,
    models: usage.models,
    ...(usage.limitsProblem ? { limitsProblem: usage.limitsProblem } : {}),
  };

  /*
    The write is separate from the read, and its failure is not the read's failure.

    These were one `try` until 2026-08-18, and the consequence was immediate and total: a machine that had
    not run `db:push` had no `ai_provider_usage` table, so the upsert threw, so a perfectly good reading was
    discarded and the panel rendered "Not signed in" over an account that was signed in fine. The stored
    half is an ENHANCEMENT — it buys tabs for accounts that cannot be measured — and an enhancement that
    fails must cost only itself.

    Observed rather than reasoned about: the table was missing on this machine and the panel went blank.
  */
  try {
    await recordProviderUsage({
      userId,
      aiProvider: CLAUDE_PROVIDER,
      emailAccount: email,
      plan: usage.account.plan,
      payload: snapshot,
    });
  } catch (err) {
    logger.warn('Could not store the Claude usage snapshot', { userId, error: String(err) });
  }

  return {
    provider: CLAUDE_PROVIDER,
    email,
    plan: usage.account.plan,
    current: true,
    measuredAt: usage.fetchedAt,
    ...snapshot,
  };
}

/**
 * Every account with a stored reading, current one first.
 *
 * Reads only Postgres — no transcript scan, no network — so this is what the panel renders the instant it
 * opens. The fresh reading arrives separately and replaces one tab's worth.
 *
 * A missing table degrades to an empty list rather than a 500. `ai_provider_usage` arrived after this route
 * did, and a machine that has not run `db:push` should lose the account tabs, not the panel.
 */
export async function listStoredUsage(userId: number, currentEmail: string | null): Promise<UsageAccount[]> {
  let rows: Awaited<ReturnType<typeof listProviderUsage>>;
  try {
    rows = await listProviderUsage(userId, CLAUDE_PROVIDER);
  } catch (err) {
    logger.warn('Could not read stored usage', { userId, error: String(err) });
    return [];
  }

  const accounts = rows.map((row) => {
    // The payload is a document written by an earlier version of this code, so its shape is whatever was
    // current the day the row was written. Defaulted rather than trusted: the three sections are mapped over
    // directly in the panel, and one missing array is a crash in the composer rather than a blank section.
    const payload = (row.payload ?? {}) as Partial<UsageSnapshot>;
    return {
      provider: row.aiProvider,
      email: row.emailAccount,
      plan: row.plan,
      current: row.emailAccount === currentEmail,
      measuredAt: row.updatedAt.toISOString(),
      limits: payload.limits ?? [],
      days: payload.days ?? [],
      models: payload.models ?? [],
      ...(payload.limitsProblem ? { limitsProblem: payload.limitsProblem } : {}),
    };
  });

  // Current account first — it is the one the panel opens on, and the only one whose numbers are live.
  return accounts.sort((a, b) => Number(b.current) - Number(a.current));
}

/**
 * Fold a fresh reading into the stored list.
 *
 * The fresh one wins for its own account and the others are left exactly as they were, which is the whole
 * contract of the panel: one tab is live, the rest are as-of.
 */
export function mergeUsage(stored: UsageAccount[], fresh: UsageAccount | null): UsageAccount[] {
  if (!fresh) return stored;
  const rest = stored.filter((a) => !(a.provider === fresh.provider && a.email === fresh.email));
  return [fresh, ...rest.map((a) => ({ ...a, current: false }))];
}