tokenmaxxing 1.38.0 → 1.40.0

package/src/lib/sample.ts

@@ -1,18 +1,17 @@
-import { mkdirSync, rmSync, writeFileSync } from "node:fs";
+import { mkdirSync } from "node:fs";
 import { join } from "node:path";
 import { minBy } from "es-toolkit";
 import { z } from "zod";
-import { readItem, writeItem, deleteItem, liveTarget, parkedTarget, isolatedTarget, claudeAiOauthOnly, mergeIntoLive } from "./credstore.ts";
-import { claudePool, credItemFor, paths } from "./paths.ts";
-import { withClaudeRefreshLock } from "./claudelock.ts";
+import { writeFileAtomic } from "./atomic.ts";
+import { readStore } from "./credstore.ts";
 import { withLock } from "./lock.ts";
-import { readOAuthAccount } from "./claudejson.ts";
-import { claudeTierLabel, refreshCredential, isAccessTokenExpiring, isDeadCredential, fetchTokenIdentity, describeIdentity, IdentityUnavailableError, InvalidGrantError } from "./oauth.ts";
-import { mergeWindows, probeUsage, windowsOf } from "./usage.ts";
-import { POST_SWAP_COOLDOWN_MS, loadAccounts, loadLastSwapAt, saveAccounts } from "./state.ts";
-import { keepRotatedPair } from "./swap.ts";
 import { log } from "./log.ts";
-import { CredentialBlobSchema, TokenIdentitySchema, UsageWindowsSchema, type Account, type Config, type OAuthCreds, type TokenIdentity } from "./types.ts";
+import { claudeTierLabel, isDeadCredential } from "./oauth.ts";
+import { claudePool, sampleDirFor, storeDirFor } from "./paths.ts";
+import type { Observation } from "./provider.ts";
+import { loadAccounts, loadUsageSnapshot, saveAccounts } from "./state.ts";
+import { mergeWindows, probeUsage, windowsOf } from "./usage.ts";
+import { UsageWindowsSchema, type Account, type Config } from "./types.ts";
 
 const SampleOutcomeSchema = z.discriminatedUnion("ok", [
   z.object({ ok: z.literal(true), usage: UsageWindowsSchema }),
@@ -20,288 +19,96 @@ const SampleOutcomeSchema = z.discriminatedUnion("ok", [
 ]);
 export type SampleOutcome = z.infer<typeof SampleOutcomeSchema>;
 
-const IdentityCheckSchema = z.discriminatedUnion("status", [
-  z.object({ status: z.literal("match") }),
-  z.object({ status: z.literal("mismatch"), reason: z.string(), owner: TokenIdentitySchema }),
-  z.object({ status: z.literal("unavailable"), reason: z.string(), stale: z.boolean() }),
+const PreparedProbeSchema = z.discriminatedUnion("ok", [
+  z.object({ ok: z.literal(true), dir: z.string() }),
+  z.object({ ok: z.literal(false), reason: z.string() }),
 ]);
-type IdentityCheck = z.infer<typeof IdentityCheckSchema>;
+type PreparedProbe = z.infer<typeof PreparedProbeSchema>;
 
-async function checkIdentity(creds: OAuthCreds, account: Account, signal?: AbortSignal): Promise<IdentityCheck> {
-  let identity: TokenIdentity;
-  try {
-    identity = await fetchTokenIdentity(creds.accessToken, signal);
-  } catch (e) {
-    return {
-      status: "unavailable",
-      reason: `credential identity check failed: ${e instanceof Error ? e.message : String(e)}`,
-      stale: e instanceof IdentityUnavailableError && e.status === 401,
-    };
-  }
-  if (identity.accountUuid === account.id) return { status: "match" };
-  return { status: "mismatch", reason: `credential actually belongs to ${describeIdentity(identity)}`, owner: identity };
+export function teeObservation(account: Account): Observation | null {
+  const stored = account.lastUsageAt != null ? { windows: account.windows, at: account.lastUsageAt } : null;
+  const snap = loadUsageSnapshot(account.id);
+  if (!snap) return stored;
+  const at = snap.state.sampledAt ?? snap.state.ts;
+  if (stored != null && at <= stored.at) return stored;
+  const aggregate = windowsOf({ fiveHour: snap.state.fiveHour, sevenDay: snap.state.sevenDay, perModel: {} }, at);
+  return { windows: mergeWindows(aggregate, account.windows), at };
 }
 
-function refreshTier(account: Account, creds: OAuthCreds): void {
-  account.tier = claudeTierLabel(creds) ?? account.tier;
+export function foldTee(account: Account): boolean {
+  const observed = teeObservation(account);
+  if (!observed || (account.lastUsageAt != null && observed.at <= account.lastUsageAt)) return false;
+  account.windows = observed.windows;
+  account.lastUsageAt = observed.at;
+  return true;
 }
 
-const ProbeFailureSchema = z.object({ ok: z.literal(false), reason: z.string() });
-type ProbeFailure = z.infer<typeof ProbeFailureSchema>;
-const PreparedProbeSchema = z.discriminatedUnion("ok", [z.object({ ok: z.literal(true), dir: z.string(), installed: z.string() }), ProbeFailureSchema]);
-type PreparedProbe = z.infer<typeof PreparedProbeSchema>;
-
-async function prepareParkedProbe(account: Account, dir: string, signal?: AbortSignal): Promise<PreparedProbe> {
-  const backup = parkedTarget(credItemFor(account.id));
-  const parkedRaw = await readItem(backup);
-  if (!parkedRaw) return { ok: false, reason: "no parked credential - run `tokenmaxxing auth`" };
-  if (!account.oauthAccount) return { ok: false, reason: "account record has no oauthAccount - run `tokenmaxxing auth`" };
-
-  let creds: OAuthCreds;
+export async function prepareProbe(account: Account, suffix = ""): Promise<PreparedProbe> {
+  let creds;
   try {
-    creds = CredentialBlobSchema.parse(JSON.parse(parkedRaw)).claudeAiOauth;
+    creds = await readStore(account.id);
   } catch (e) {
-    return { ok: false, reason: `parked credential unreadable (${(e instanceof Error ? e.message : String(e)).slice(0, 80)}) - run \`tokenmaxxing auth\`` };
+    return { ok: false, reason: `store credential unreadable (${(e instanceof Error ? e.message : String(e)).slice(0, 80)}) - run \`tokenmaxxing auth\`` };
   }
+  if (!creds) return { ok: false, reason: "no credential in this account's store - run `tokenmaxxing auth`" };
   if (isDeadCredential(creds)) {
     account.needsReauth = true;
-    return { ok: false, reason: "parked credential was cleared after a failed refresh - re-auth with `tokenmaxxing auth`" };
-  }
-
-  const liveRaw = await readItem(liveTarget());
-  let liveAccount: string | null = null;
-  let liveToken: string | null = null;
-  let liveRefresh: string | null = null;
-  if (liveRaw != null) {
-    let liveCreds: OAuthCreds;
-    try {
-      liveCreds = CredentialBlobSchema.parse(JSON.parse(liveRaw)).claudeAiOauth;
-    } catch (e) {
-      return { ok: false, reason: `cannot read the live credential (${(e instanceof Error ? e.message : String(e)).slice(0, 80)}) - refusing to sample a possibly-live account` };
-    }
-    if (!isDeadCredential(liveCreds)) {
-      liveToken = liveCreds.accessToken;
-      liveRefresh = liveCreds.refreshToken;
-      try {
-        liveAccount = (await fetchTokenIdentity(liveCreds.accessToken, signal)).accountUuid;
-      } catch (e) {
-        return { ok: false, reason: `cannot verify the live credential's owner (${(e instanceof Error ? e.message : String(e)).slice(0, 80)}) - refusing to sample a possibly-live account` };
-      }
-      if (liveAccount === account.id) {
-        return { ok: false, reason: "this account holds the LIVE login (active label drifted) - run `tokenmaxxing switch` to reconcile" };
-      }
-    }
+    return { ok: false, reason: "the store's credential was cleared after a failed refresh - re-auth with `tokenmaxxing auth`" };
   }
-
-  const relocate = async (owner: TokenIdentity): Promise<string> => {
-    const pooled = loadAccounts(claudePool).accounts.find((a) => a.id === owner.accountUuid) ?? null;
-    if (pooled == null || pooled.id === liveAccount) return "was left in place";
-    await writeItem(parkedTarget(credItemFor(pooled.id)), JSON.stringify({ claudeAiOauth: creds }));
-    log("swap.parked_relocated", { account: pooled.id.slice(0, 8) });
-    return "was copied to that account's slot";
-  };
-
-  if (isAccessTokenExpiring(creds, 300_000)) {
-    const owner = await checkIdentity(creds, account, signal);
-    if (owner.status === "mismatch") {
-      const kept = await relocate(owner.owner);
-      account.needsReauth = true;
-      return { ok: false, reason: `${owner.reason} - refusing to spend another account's grant; the pair ${kept}; re-auth with \`tokenmaxxing auth\`` };
-    }
-    if (owner.status === "unavailable" && !owner.stale) {
-      return { ok: false, reason: `${owner.reason} - refusing to refresh a parked credential whose owner cannot be verified` };
-    }
-    const rotate = async (): Promise<{ fresh: OAuthCreds } | { failed: ProbeFailure }> => {
-      try {
-        return { fresh: await refreshCredential(creds, Date.now(), signal) };
-      } catch (e) {
-        if (e instanceof InvalidGrantError) {
-          account.needsReauth = true;
-          return { failed: { ok: false, reason: "refresh token dead - re-auth with `tokenmaxxing auth`" } };
-        }
-        return { failed: { ok: false, reason: `token refresh failed: ${e instanceof Error ? e.message : String(e)}` } };
-      }
-    };
-    let result: { fresh: OAuthCreds } | { failed: ProbeFailure };
-    if (owner.status === "match") {
-      result = await rotate();
-    } else {
-      if (liveRefresh != null && creds.refreshToken === liveRefresh) {
-        return { ok: false, reason: "this slot holds a copy of the LIVE login's grant - refusing to rotate it from a parked slot; run `tokenmaxxing switch` to reconcile" };
-      }
-      try {
-        signal?.throwIfAborted();
-        result = await withClaudeRefreshLock(async (lock) => {
-          const rotated = await rotate();
-          if ("failed" in rotated) return rotated;
-          const verified = await checkIdentity(rotated.fresh, account, signal);
-          if (verified.status === "match") return rotated;
-          if (verified.status === "unavailable") {
-            await writeItem(backup, JSON.stringify({ claudeAiOauth: rotated.fresh }));
-            return { failed: { ok: false, reason: `${verified.reason} - the rotated pair stays in this slot unverified until the next pass` } };
-          }
-          const trueOwner = loadAccounts(claudePool).accounts.find((a) => a.id === verified.owner.accountUuid) ?? null;
-          const kept = await keepRotatedPair({ fresh: rotated.fresh, owner: trueOwner, fallback: account, liveOwnerId: liveAccount, expectedLiveToken: liveToken, lock });
-          account.needsReauth = true;
-          return { failed: { ok: false, reason: `${verified.reason}, whose grant this refresh rotated - the rotated token ${kept}; re-auth with \`tokenmaxxing auth\`` } };
-        });
-      } catch (e) {
-        return { ok: false, reason: `cannot arbitrate a stale parked credential: ${e instanceof Error ? e.message : String(e)}` };
-      }
-    }
-    if ("failed" in result) return result.failed;
-    creds = result.fresh;
-    await writeItem(backup, JSON.stringify({ claudeAiOauth: creds }));
-  }
-
-  const identity = await checkIdentity(creds, account, signal);
-  if (identity.status === "mismatch") {
-    const kept = await relocate(identity.owner);
-    account.needsReauth = true;
-    return { ok: false, reason: `${identity.reason} - this account's own credential is gone; the pair ${kept}; re-auth with \`tokenmaxxing auth\`` };
-  }
-  if (identity.status === "unavailable") {
-    return { ok: false, reason: identity.reason };
-  }
-  refreshTier(account, creds);
-
-  rmSync(dir, { recursive: true, force: true });
+  if (!account.oauthAccount) return { ok: false, reason: "account record has no oauthAccount - run `tokenmaxxing auth`" };
+  account.tier = claudeTierLabel(creds) ?? account.tier;
+  const dir = sampleDirFor(account.id, suffix);
   mkdirSync(dir, { recursive: true });
-  const installed = JSON.stringify({ claudeAiOauth: creds });
-  try {
-    await writeItem(isolatedTarget(dir), installed);
-    writeFileSync(join(dir, ".claude.json"), JSON.stringify({ oauthAccount: account.oauthAccount, hasCompletedOnboarding: true }));
-  } catch (e) {
-    await finishParkedProbe(account, { dir, installed });
-    throw e;
-  }
-  return { ok: true, dir, installed };
+  writeFileAtomic(join(dir, ".claude.json"), JSON.stringify({ oauthAccount: account.oauthAccount, hasCompletedOnboarding: true }));
+  return { ok: true, dir };
 }
 
-async function runParkedProbe(dir: string, opts: { retries?: number }): Promise<SampleOutcome> {
-  const usage = await probeUsage(dir, Date.now(), { retries: opts.retries });
+export async function runProbe(account: Account, dir: string, opts: { retries?: number }): Promise<SampleOutcome> {
+  const usage = await probeUsage({ configDir: dir, store: storeDirFor(account.id) }, Date.now(), { retries: opts.retries });
   return usage ? { ok: true, usage } : { ok: false, reason: "`/usage` returned no limit data (see log)" };
 }
 
-async function finishParkedProbe(account: Account, prepared: { dir: string; installed: string }): Promise<void> {
-  const isoTarget = isolatedTarget(prepared.dir);
-  const backup = parkedTarget(credItemFor(account.id));
-  try {
-    const afterIso = await readItem(isoTarget);
-    if (afterIso && afterIso !== prepared.installed) {
-      const parkedNow = await readItem(backup);
-      const installedToken = CredentialBlobSchema.parse(JSON.parse(prepared.installed)).claudeAiOauth.accessToken;
-      const parkedToken = parkedNow == null ? null : CredentialBlobSchema.parse(JSON.parse(parkedNow)).claudeAiOauth.accessToken;
-      if (parkedToken === installedToken) await writeItem(backup, claudeAiOauthOnly(afterIso));
-      else log("sample.harvest_skipped", { account: account.id.slice(0, 8) });
-    }
-  } finally {
-    await deleteItem(isoTarget);
-    rmSync(prepared.dir, { recursive: true, force: true });
-  }
-}
-
-export async function probeParkedUsage(account: Account): Promise<SampleOutcome> {
-  const prepared = await prepareParkedProbe(account, join(paths.sampleDir, credItemFor(account.id)));
-  if (!prepared.ok) return prepared;
-  try {
-    return await runParkedProbe(prepared.dir, {});
-  } finally {
-    await finishParkedProbe(account, prepared);
-  }
+export async function probeAccountUsage(account: Account): Promise<SampleOutcome> {
+  const prepared = await prepareProbe(account);
+  return prepared.ok ? runProbe(account, prepared.dir, {}) : prepared;
 }
 
-const PREPARE_DEADLINE_MS = 20_000;
-
-export async function sampleOldestParked(cfg: Config): Promise<void> {
+export async function sampleOldest(cfg: Config): Promise<void> {
   const reserved = await withLock(claudePool.lockFile, async () => {
     const now = Date.now();
-    const lastSwapAt = loadLastSwapAt(claudePool);
-    if (lastSwapAt != null && now - lastSwapAt < POST_SWAP_COOLDOWN_MS) return null;
     const idx = loadAccounts(claudePool);
-    const live = readOAuthAccount()?.accountUuid ?? null;
+    let dirty = false;
+    for (const a of idx.accounts) dirty = foldTee(a) || dirty;
     const sampledAt = (a: Account) => Math.max(a.lastUsageAt ?? 0, a.lastProbeAt ?? 0);
-    const stale = idx.accounts.filter((a) => a.id !== live && a.needsReauth !== true && now - sampledAt(a) > cfg.policy.usagePollTtlMs);
+    const stale = idx.accounts.filter((a) => a.needsReauth !== true && now - sampledAt(a) > cfg.policy.usagePollTtlMs);
     const target = minBy(stale, sampledAt);
-    if (!target) return null;
+    if (!target) {
+      if (dirty) saveAccounts(claudePool, idx);
+      return null;
+    }
     target.lastProbeAt = now;
-    const prepared = await prepareParkedProbe(target, join(paths.sampleDir, `${credItemFor(target.id)}-tick`), AbortSignal.timeout(PREPARE_DEADLINE_MS));
+    const prepared = await prepareProbe(target, "-tick");
     saveAccounts(claudePool, idx);
     if (!prepared.ok) {
-      log("sample.parked_failed", { account: target.id.slice(0, 8), reason: prepared.reason.slice(0, 200) });
+      log("sample.failed", { account: target.id.slice(0, 8), reason: prepared.reason.slice(0, 200) });
       return null;
     }
-    return { account: target, prepared };
+    return { account: target, dir: prepared.dir };
   });
   if (!reserved) return;
-  const { account, prepared } = reserved;
   const startedAt = Date.now();
-  let outcome: SampleOutcome | null = null;
-  try {
-    outcome = await runParkedProbe(prepared.dir, { retries: 0 });
-  } finally {
-    await withLock(claudePool.lockFile, async () => {
-      await finishParkedProbe(account, prepared);
-      if (outcome == null) return;
-      const idx = loadAccounts(claudePool);
-      const stored = idx.accounts.find((a) => a.id === account.id);
-      if (stored && outcome.ok && (stored.lastUsageAt == null || startedAt > stored.lastUsageAt)) {
-        stored.windows = mergeWindows(windowsOf(outcome.usage, startedAt), stored.windows);
-        stored.lastUsageAt = startedAt;
-        saveAccounts(claudePool, idx);
-      }
-      log(outcome.ok ? "sample.parked_ok" : "sample.parked_failed", {
-        account: account.id.slice(0, 8),
-        ...(outcome.ok ? {} : { reason: outcome.reason.slice(0, 200) }),
-      });
+  const outcome = await runProbe(reserved.account, reserved.dir, { retries: 0 });
+  await withLock(claudePool.lockFile, () => {
+    const idx = loadAccounts(claudePool);
+    const stored = idx.accounts.find((a) => a.id === reserved.account.id);
+    if (stored && outcome.ok && (stored.lastUsageAt == null || startedAt > stored.lastUsageAt)) {
+      stored.windows = mergeWindows(windowsOf(outcome.usage, startedAt), stored.windows);
+      stored.lastUsageAt = startedAt;
+      saveAccounts(claudePool, idx);
+    }
+    log(outcome.ok ? "sample.ok" : "sample.failed", {
+      account: reserved.account.id.slice(0, 8),
+      ...(outcome.ok ? {} : { reason: outcome.reason.slice(0, 200) }),
     });
-  }
-}
-
-export async function ensureLiveTokenFresh(): Promise<void> {
-  const liveRaw = await readItem(liveTarget());
-  if (!liveRaw) return;
-  let creds: OAuthCreds;
-  try {
-    creds = CredentialBlobSchema.parse(JSON.parse(liveRaw)).claudeAiOauth;
-  } catch {
-    return;
-  }
-  if (isDeadCredential(creds)) throw new InvalidGrantError("live credential was cleared after a failed refresh");
-  if (!isAccessTokenExpiring(creds, 300_000)) return;
-  await withClaudeRefreshLock(async (lock) => {
-    const raw2 = await readItem(liveTarget());
-    if (raw2 == null) throw new Error("live credential vanished while waiting for the refresh lock");
-    const current = CredentialBlobSchema.parse(JSON.parse(raw2)).claudeAiOauth;
-    if (isDeadCredential(current)) throw new InvalidGrantError("live credential was cleared after a failed refresh");
-    const next = isAccessTokenExpiring(current, 300_000) ? await refreshCredential(current) : current;
-    if (next === current) return;
-    if (lock.compromised()) throw new Error("refresh lock compromised mid-refresh - discarding the live rewrite");
-    await writeItem(liveTarget(), mergeIntoLive(raw2, next));
   });
 }
-
-export async function probeActiveUsage(account: Account): Promise<SampleOutcome> {
-  try {
-    await ensureLiveTokenFresh();
-  } catch (e) {
-    if (e instanceof InvalidGrantError) return { ok: false, reason: "live refresh token dead - run `claude` and `/login`" };
-    return { ok: false, reason: `token refresh failed: ${e instanceof Error ? e.message : String(e)}` };
-  }
-  const liveRaw = await readItem(liveTarget());
-  if (!liveRaw) return { ok: false, reason: "no live credential - run `claude` and `/login`" };
-  let creds: OAuthCreds;
-  try {
-    creds = CredentialBlobSchema.parse(JSON.parse(liveRaw)).claudeAiOauth;
-  } catch (e) {
-    return { ok: false, reason: `live credential blob unreadable (${(e instanceof Error ? e.message : String(e)).slice(0, 80)})` };
-  }
-
-  const identity = await checkIdentity(creds, account);
-  if (identity.status === "mismatch") return { ok: false, reason: `live ${identity.reason} - active label drifted; run \`tokenmaxxing switch\`` };
-  if (identity.status === "unavailable") return { ok: false, reason: identity.reason };
-  refreshTier(account, creds);
-
-  const usage = await probeUsage();
-  return usage ? { ok: true, usage } : { ok: false, reason: "`/usage` returned no limit data (see log)" };
-}

package/src/lib/sessions.ts

@@ -3,6 +3,7 @@ import { join } from "node:path";
 import { z } from "zod";
 import { paths } from "./paths.ts";
 import { writeFileAtomic } from "./atomic.ts";
+import { RespawnMarkerSchema } from "./types.ts";
 
 const SessionSchema = z.object({ flags: z.array(z.string()), cwd: z.string() });
 
@@ -34,3 +35,27 @@ export function pruneStaleSessions(now: number): void {
     }
   }
 }
+
+const SupervisedSessionSchema = z.object({ sid: z.string().min(1), launchedAt: z.number().finite().nullable() });
+export type SupervisedSession = z.infer<typeof SupervisedSessionSchema>;
+
+const LaunchedAtSchema = z.coerce.number().finite().optional().catch(undefined);
+
+export function supervisedSession(env: Record<string, string | undefined> = process.env): SupervisedSession | null {
+  if (env.TOKENMAXXING_SUPERVISED !== "1") return null;
+  const sid = env.TOKENMAXXING_SESSION_ID;
+  if (sid == null || sid === "") return null;
+  return SupervisedSessionSchema.parse({ sid, launchedAt: LaunchedAtSchema.parse(env.TOKENMAXXING_LAUNCHED_AT) ?? null });
+}
+
+export function writeRespawnMarker(input: { session: SupervisedSession; sessionId: string; accountId: string; waitUntil: number }): void {
+  mkdirSync(paths.respawnDir, { recursive: true });
+  const payload = RespawnMarkerSchema.parse({
+    accountId: input.accountId,
+    ts: Date.now(),
+    waitUntil: input.waitUntil,
+    sessionId: input.sessionId,
+    ...(input.session.launchedAt != null ? { launchedAt: input.session.launchedAt } : {}),
+  });
+  writeFileAtomic(join(paths.respawnDir, input.session.sid), JSON.stringify(payload));
+}

package/src/lib/state.ts

@@ -1,7 +1,7 @@
 import { closeSync, existsSync, fstatSync, openSync, readFileSync, rmSync, utimesSync } from "node:fs";
 import { isEqual } from "es-toolkit";
 import { z } from "zod";
-import { paths, realClaudeBinFromEnv, realCodexBinFromEnv, type PoolPaths } from "./paths.ts";
+import { paths, realClaudeBinFromEnv, realCodexBinFromEnv, usageJsonFor, type PoolPaths } from "./paths.ts";
 import { writeFileAtomic } from "./atomic.ts";
 import {
   AccountsIndexSchema,
@@ -99,17 +99,19 @@ export function upsertAccount(
   return existing ?? fresh;
 }
 
-export function loadUsageSnapshot(): { state: UsageState; at: number } | null {
+export type UsageSnapshot = { state: UsageState; at: number };
+
+export function loadUsageSnapshot(accountId: string): UsageSnapshot | null {
   let fd: number;
   try {
-    fd = openSync(paths.usageJson, "r");
+    fd = openSync(usageJsonFor(accountId), "r");
   } catch {
     return null;
   }
   try {
     const at = fstatSync(fd).mtimeMs;
     const parsed = UsageStateSchema.safeParse(JSON.parse(readFileSync(fd, "utf8")));
-    return parsed.success ? { state: parsed.data, at } : null;
+    return parsed.success && parsed.data.account === accountId ? { state: parsed.data, at } : null;
   } catch {
     return null;
   } finally {
@@ -117,16 +119,12 @@ export function loadUsageSnapshot(): { state: UsageState; at: number } | null {
   }
 }
 
-export function loadUsage(): UsageState | null {
-  return loadUsageSnapshot()?.state ?? null;
-}
-
-export function clearUsageSnapshots(): void {
-  rmSync(paths.usageJson, { force: true });
+export function clearUsageSnapshot(accountId: string): void {
+  rmSync(usageJsonFor(accountId), { force: true });
 }
 
 export function loadLastSwapAt(pool: PoolPaths): number | null {
-  if (!existsSync(pool.lastSwapJson)) return null;
+  if (pool.lastSwapJson == null || !existsSync(pool.lastSwapJson)) return null;
   let json: unknown;
   try {
     json = JSON.parse(readFileSync(pool.lastSwapJson, "utf8"));
@@ -137,55 +135,32 @@ export function loadLastSwapAt(pool: PoolPaths): number | null {
 }
 
 export function saveLastSwapAt(pool: PoolPaths, ts: number): void {
+  if (pool.lastSwapJson == null) throw new Error("this pool keeps no swap clock");
   writeFileAtomic(pool.lastSwapJson, JSON.stringify(LastSwapSchema.parse({ ts })));
 }
 
-const DepletedWaitSchema = z.object({ waitUntil: z.number(), id: z.string(), ts: z.number() });
-export type DepletedWait = z.infer<typeof DepletedWaitSchema>;
-
-export function loadDepletedWait(): DepletedWait | null {
-  if (!existsSync(paths.depletedJson)) return null;
-  let json: unknown;
-  try {
-    json = JSON.parse(readFileSync(paths.depletedJson, "utf8"));
-  } catch {
-    throw new Error(`${paths.depletedJson} is corrupt (unparsable JSON) - repair or remove the file`);
-  }
-  return DepletedWaitSchema.parse(json);
-}
-
-export function saveDepletedWait(rec: DepletedWait): void {
-  writeFileAtomic(paths.depletedJson, JSON.stringify(DepletedWaitSchema.parse(rec)));
-}
-
-export function clearDepletedWait(): void {
-  rmSync(paths.depletedJson, { force: true });
-}
-
 export const POST_SWAP_COOLDOWN_MS = 45_000;
 
 const USAGE_TS_REFRESH_MS = 10 * 60_000;
 const SAMPLED_AT_REFRESH_MS = 30_000;
 
-export function writeUsage(input: UsageState, opts: { stamp?: boolean } = {}): boolean {
-  const prev = loadUsage();
-  const stamp = opts.stamp === true;
-  const carriedSampleAt = prev != null && prev.account === input.account ? (prev.sampledAt ?? prev.ts) : undefined;
-  const sampledAt = stamp ? (carriedSampleAt ?? input.sampledAt) : input.ts;
-  const next: UsageState = { ...input, ...(sampledAt != null ? { sampledAt } : {}) };
+export function writeUsage(input: UsageState): boolean {
+  const file = usageJsonFor(input.account);
+  const prev = loadUsageSnapshot(input.account)?.state ?? null;
+  const next: UsageState = { ...input, sampledAt: input.ts };
   if (
     prev &&
     isEqual({ ...prev, ts: 0, sampledAt: 0 }, { ...next, ts: 0, sampledAt: 0 }) &&
     next.ts - prev.ts < USAGE_TS_REFRESH_MS &&
-    (stamp || next.ts - (prev.sampledAt ?? prev.ts) < SAMPLED_AT_REFRESH_MS)
+    next.ts - (prev.sampledAt ?? prev.ts) < SAMPLED_AT_REFRESH_MS
   ) {
     try {
-      utimesSync(paths.usageJson, new Date(next.ts), new Date(next.ts));
+      utimesSync(file, new Date(next.ts), new Date(next.ts));
     } catch (e) {
       if (ErrnoSchema.safeParse(e).data?.code !== "ENOENT") throw e;
     }
     return false;
   }
-  writeFileAtomic(paths.usageJson, JSON.stringify(next));
+  writeFileAtomic(file, JSON.stringify(next));
   return true;
 }

package/src/lib/types.ts

@@ -54,7 +54,7 @@ const ModelInfoSchema = z.object({ id: z.string(), display: z.string() });
 export type ModelInfo = z.infer<typeof ModelInfoSchema>;
 
 export const UsageStateSchema = AggregateWindowsSchema.extend({
-  account: z.string().nullable(),
+  account: z.string(),
   ts: z.number(),
   model: ModelInfoSchema.nullable().default(null),
   sampledAt: z.number().optional(),
@@ -130,7 +130,7 @@ export const ConfigSchema = z
 export type Config = z.infer<typeof ConfigSchema>;
 
 export const RespawnMarkerSchema = z.object({
-  account: z.string(),
+  accountId: z.string(),
   ts: z.number(),
   waitUntil: z.number(),
   sessionId: z.string(),
@@ -186,15 +186,6 @@ export const SubagentStatusLineStdinSchema = z.looseObject({
     .catch(undefined),
 });
 
-export const RefreshResponseSchema = z.looseObject({
-  access_token: z.string(),
-  refresh_token: z.string().optional(),
-  expires_in: z.number().optional(),
-  refresh_token_expires_in: z.number().optional(),
-  scope: z.string().optional(),
-  token_type: z.string().optional(),
-});
-
 export const ProfileResponseSchema = z.looseObject({
   account: z.looseObject({ uuid: z.string(), email: z.string().nullish() }),
   organization: z.looseObject({ uuid: z.string(), name: z.string().nullish() }),