@adecore/agents 0.41.1 → 0.43.0

package/src/usage/scanner.ts

@@ -1,18 +1,19 @@
 import { createHash, type Hash } from 'node:crypto';
-import { mkdir, open, readFile, readdir, stat } from 'node:fs/promises';
+import { open, readFile, readdir, stat } from 'node:fs/promises';
 import { join } from 'node:path';
 import { Readable } from 'node:stream';
 import { createZstdDecompress } from 'node:zlib';
 import type { UsageProvider, UsageRoot } from '@adecore/agent-contracts';
 import { errorText } from '../error-text.ts';
-import { isNotFound, writeAtomic } from '../fs.ts';
-import { decodeIndex, encodeIndex, type IndexedFile, type UsageIndex } from './index-file.ts';
+import { isNotFound } from '../fs.ts';
 import { foldByKey, type UsageRecord } from './record.ts';
 import { claudeMightCarryUsage, parseClaudeLine } from './readers/claude.ts';
 import { GROK_USAGE_FILE, parseGrokUsage } from './readers/grok.ts';
 import { hostLogMightCarryUsage, parseHostUsageLine } from './readers/host-log.ts';
 import { cloneCodexState, codexMightCarryUsage, createCodexState, parseCodexLine, type CodexParserState } from './readers/codex.ts';
 import { usageRoots, type UsageRootPath } from './roots.ts';
+import type { ScanState } from './scan-state.ts';
+import { UsageStore } from './usage-store.ts';
 
 const HASH_CHUNK_BYTES = 65536;
 // How much of a transcript is read at once. Codex writes sessions of more than a gigabyte, which read whole held several times their size.
@@ -115,6 +116,11 @@ function isUsageFile(provider: UsageProvider, name: string): boolean {
     return readsWhole(provider) ? name === GROK_USAGE_FILE : name.endsWith('.jsonl') || (provider === 'codex' && name.endsWith('.jsonl.zst'));
 }
 
+/* A compressed Codex session is the same transcript as the plain one it was, so both count under one path. */
+function keyOf(path: string, provider: UsageProvider): string {
+    return provider === 'codex' && path.endsWith('.jsonl.zst') ? path.slice(0, -4) : path;
+}
+
 async function listUsageFiles(dir: string, provider: UsageProvider, out: string[] = []): Promise<string[]> {
     let entries;
 
@@ -266,47 +272,37 @@ function withAccount(records: UsageRecord[], accountOf: AccountOf): UsageRecord[
 
 /*
  * Parses only what a transcript appended since the last scan, and starts a file over once it was
- * rewritten or cut. A deleted transcript keeps its totals, since its usage still happened.
+ * rewritten or cut. A deleted transcript keeps its totals, since its usage still happened; only
+ * where the scan left it is forgotten.
  */
 export class UsageScanner {
-    private readonly home: string;
     private readonly roots: () => UsageRootPath[];
     private readonly sessionAccounts: () => ReadonlyMap<string, string>;
-    private index: UsageIndex = new Map();
+    private readonly store: UsageStore;
     private loaded = false;
-    private dirty = false;
 
     /* `roots` as a function is asked again on every scan, so an account added since is walked too. */
     constructor(home: string, roots: UsageRootPath[] | (() => UsageRootPath[]) = usageRoots(), options: UsageScannerOptions = {}) {
-        this.home = home;
+        this.store = new UsageStore(join(home, 'usage'));
         this.roots = typeof roots === 'function' ? roots : () => roots;
         this.sessionAccounts = options.sessionAccounts ?? (() => new Map());
     }
 
-    private get file(): string {
-        return join(this.home, 'usage', 'index.json');
-    }
-
+    /* A history that does not read fails the scan rather than being written over, which would lose what only it holds. */
     async load(): Promise<void> {
         if (this.loaded) {
             return;
         }
 
+        await this.store.load();
         this.loaded = true;
-
-        try {
-            this.index = decodeIndex(await readFile(this.file, 'utf8'));
-        } catch (e) {
-            if (!isNotFound(e)) {
-                console.warn('The usage index would not read; starting from a cold scan:', errorText(e));
-            }
-        }
     }
 
     async scan(): Promise<ScanReport> {
         const started = Date.now();
         await this.load();
         const roots: UsageRoot[] = [];
+        const seen = new Set<string>();
         let files = 0;
         let changedFiles = 0;
         let sessions: ReadonlyMap<string, string> | null = null;
@@ -350,6 +346,7 @@ export class UsageScanner {
                 }
 
                 files += 1;
+                seen.add(keyOf(path, root.provider));
 
                 try {
                     if (await this.readFileInto(path, root.provider, accountOf)) {
@@ -364,31 +361,28 @@ export class UsageScanner {
             roots.push({ provider: root.provider, path: root.path, status: failure === null ? 'ok' : 'failed', message: failure });
         }
 
-        if (this.dirty) {
-            this.dirty = false;
-            await mkdir(join(this.home, 'usage'), { recursive: true, mode: 0o700 });
-            await writeAtomic(this.file, encodeIndex(this.index));
+        // A root that is gone for now costs a read from the start once it is back, which replaces and never adds.
+        for (const path of [...this.store.state.keys()]) {
+            if (!seen.has(path)) {
+                this.store.forget(path);
+            }
         }
 
+        await this.store.commit();
+
         return { at: Date.now(), files, changedFiles, durationMs: Date.now() - started, roots };
     }
 
-    /* Every record the index holds, with the copies one message left in several transcripts folded. */
+    /* Every record the history holds, with the copies one message left in several transcripts folded. */
     records(): UsageRecord[] {
-        const all: UsageRecord[] = [];
-
-        for (const file of this.index.values()) {
-            all.push(...file.records, ...file.tail);
-        }
-
-        return foldByKey(all);
+        return this.store.records();
     }
 
     private async readFileInto(path: string, provider: UsageProvider, accountOf: AccountOf): Promise<boolean> {
         const info = await stat(path);
         const compressed = provider === 'codex' && path.endsWith('.jsonl.zst');
-        const key = compressed ? path.slice(0, -4) : path;
-        const known = this.index.get(key);
+        const key = keyOf(path, provider);
+        const known = this.store.state.get(key);
         const currentParser = provider !== 'codex' || known?.parserVersion === CODEX_PARSER_VERSION;
 
         if (
@@ -421,29 +415,30 @@ export class UsageScanner {
         const hash = grown === null || prefix === null ? createHash('sha256') : prefix;
         const from = grown?.offset ?? 0;
         const parsed = await parseLines(readBytes(path, from, compressed), hash, provider, from, grown?.codex ?? null);
-        const before = grown?.records ?? [];
+        const usage = { provider, records: withAccount(parsed.records, accountOf), tail: withAccount(parsed.tail, accountOf) };
 
-        const entry: IndexedFile = {
+        const state: ScanState = {
             provider,
             size: info.size,
             mtimeMs: info.mtimeMs,
             offset: parsed.offset,
-            // Folding here keeps the streaming copies of one message out of the index entirely.
-            records: foldByKey([...before, ...withAccount(parsed.records, accountOf)]),
-            tail: withAccount(parsed.tail, accountOf),
             codex: parsed.codex,
             head: headOf(hash),
             ...(provider === 'codex' ? { parserVersion: CODEX_PARSER_VERSION, compressed } : {})
         };
 
-        this.index.set(key, entry);
-        this.dirty = true;
+        if (grown === null) {
+            // Folding here keeps the streaming copies of one message out of the history entirely.
+            this.store.replace(key, { ...usage, records: foldByKey(usage.records) }, state);
+        } else {
+            this.store.append(key, usage, state);
+        }
 
         return true;
     }
 
     /*
-     * A file rewritten as a whole replaces what the index held of it, so a turn it still lists is never
+     * A file rewritten as a whole replaces what the history held of it, so a turn it still lists is never
      * counted twice. One that does not parse, such as one caught halfway through a write, keeps the
      * records it had and is read again on the next scan.
      */
@@ -454,18 +449,8 @@ export class UsageScanner {
             return false;
         }
 
-        this.index.set(path, {
-            provider,
-            size: info.size,
-            mtimeMs: info.mtimeMs,
-            offset: info.size,
-            records: foldByKey(withAccount(records, accountOf)),
-            tail: [],
-            codex: null,
-            head: null
-        });
-
-        this.dirty = true;
+        const state: ScanState = { provider, size: info.size, mtimeMs: info.mtimeMs, offset: info.size, codex: null, head: null };
+        this.store.replace(path, { provider, records: foldByKey(withAccount(records, accountOf)), tail: [] }, state);
 
         return true;
     }

package/src/usage/usage-store.ts

@@ -0,0 +1,237 @@
+import { randomUUID } from 'node:crypto';
+import { isDeepStrictEqual } from 'node:util';
+import { appendFile, mkdir, readFile, rename, truncate } from 'node:fs/promises';
+import { join } from 'node:path';
+import { errorText } from '../error-text.ts';
+import { isNotFound, removeFile, writeAtomic } from '../fs.ts';
+import { applyEntry, decodeHistory, encodeEntry, encodeHeader, type FileUsage, type HistoryEntry } from './history-file.ts';
+import { decodeIndex } from './index-file.ts';
+import { foldByKey, type UsageRecord } from './record.ts';
+import { decodeState, encodeState, type ScanState } from './scan-state.ts';
+
+const HISTORY_FILE = 'history.jsonl';
+const STATE_FILE = 'scan-state.json';
+const LEGACY_INDEX_FILE = 'index.json';
+// Below this, appended lines are not worth a rewrite of the whole history.
+const COMPACT_FLOOR_BYTES = 1 << 20;
+
+async function readOptional(path: string): Promise<Buffer | null> {
+    try {
+        return await readFile(path);
+    } catch (e) {
+        if (isNotFound(e)) {
+            return null;
+        }
+
+        throw e;
+    }
+}
+
+/*
+ * The usage of every transcript ever read, kept apart from where the scanner left each one. The state
+ * is saved after the lines it accounts for, against the history's size, so after a crash in between it
+ * does not match: it is dropped and every transcript is read from the start again, which replaces what
+ * a path counted instead of adding to it.
+ */
+export class UsageStore {
+    readonly usage = new Map<string, FileUsage>();
+    readonly state = new Map<string, ScanState>();
+    private readonly dir: string;
+    private pending: HistoryEntry[] = [];
+    private stateDirty = false;
+    private gen: string | null = null;
+    private base = 0;
+    private bytes = 0;
+    // Set when the history on disk may not end where this store thinks it does.
+    private rewriteNext = false;
+
+    constructor(dir: string) {
+        this.dir = dir;
+    }
+
+    async load(): Promise<void> {
+        this.usage.clear();
+        this.state.clear();
+        const history = await readOptional(join(this.dir, HISTORY_FILE));
+
+        if (history === null) {
+            await this.migrate();
+
+            return;
+        }
+
+        const decoded = decodeHistory(history);
+
+        for (const [path, usage] of decoded.usage) {
+            this.usage.set(path, usage);
+        }
+
+        if (decoded.complete < history.length) {
+            await truncate(join(this.dir, HISTORY_FILE), decoded.complete);
+        }
+
+        this.gen = decoded.header?.gen ?? null;
+        this.base = decoded.header?.base ?? 0;
+        this.bytes = decoded.complete;
+        this.rewriteNext = this.gen === null;
+
+        const saved = await readOptional(join(this.dir, STATE_FILE));
+        const state = saved === null ? null : decodeState(saved.toString('utf8'));
+
+        if (state !== null && state.history.gen === this.gen && state.history.bytes === this.bytes) {
+            for (const [path, file] of state.files) {
+                this.state.set(path, file);
+            }
+        }
+
+        // A migration that stopped after writing the history; the index has nothing the history lacks.
+        await removeFile(join(this.dir, LEGACY_INDEX_FILE));
+    }
+
+    records(): UsageRecord[] {
+        const all: UsageRecord[] = [];
+
+        for (const file of this.usage.values()) {
+            all.push(...file.records, ...file.tail);
+        }
+
+        return foldByKey(all);
+    }
+
+    /* A transcript read from the start: what it counts now replaces what it counted. */
+    replace(path: string, usage: FileUsage, state: ScanState): void {
+        this.setState(path, state);
+
+        if (isDeepStrictEqual(this.usage.get(path), usage)) {
+            return;
+        }
+
+        this.usage.set(path, usage);
+        this.pending.push({ path, op: 'set', usage });
+    }
+
+    /* A transcript read from where it was left: `added` counts on top of what it had. */
+    append(path: string, added: FileUsage, state: ScanState): void {
+        this.setState(path, state);
+        const known = this.usage.get(path);
+
+        if (added.records.length === 0 && known?.provider === added.provider && isDeepStrictEqual(known.tail, added.tail)) {
+            return;
+        }
+
+        const entry: HistoryEntry = { path, op: 'add', usage: { ...added, records: foldByKey(added.records) } };
+        applyEntry(this.usage, entry);
+        this.pending.push(entry);
+    }
+
+    /* Where the scanner left a transcript that is gone; what it counted stays. */
+    forget(path: string): void {
+        if (this.state.delete(path)) {
+            this.stateDirty = true;
+        }
+    }
+
+    async commit(): Promise<void> {
+        if (this.pending.length === 0 && !this.stateDirty) {
+            return;
+        }
+
+        await mkdir(this.dir, { recursive: true, mode: 0o700 });
+
+        try {
+            if (this.rewriteNext || (this.pending.length > 0 && this.gen === null)) {
+                await this.rewrite();
+            } else if (this.pending.length > 0) {
+                await this.appendPending();
+            }
+
+            await this.saveState();
+        } catch (e) {
+            // What failed to land is in memory, and a rewrite puts all of it on disk.
+            this.pending = [];
+            this.rewriteNext = true;
+
+            throw e;
+        }
+    }
+
+    private setState(path: string, state: ScanState): void {
+        this.state.set(path, state);
+        this.stateDirty = true;
+    }
+
+    private async appendPending(): Promise<void> {
+        const text = this.pending.map(encodeEntry).join('');
+        this.pending = [];
+        await appendFile(join(this.dir, HISTORY_FILE), text, { mode: 0o600 });
+        this.bytes += Buffer.byteLength(text);
+
+        if (this.bytes > 2 * this.base && this.bytes - this.base > COMPACT_FLOOR_BYTES) {
+            await this.rewrite();
+        }
+    }
+
+    /* The history as one `set` line per path, which drops every line a later one superseded. */
+    private async rewrite(): Promise<void> {
+        const lines = [...this.usage].map(([path, usage]) => encodeEntry({ path, op: 'set', usage }));
+        const gen = randomUUID();
+        const body = lines.join('');
+        const base = Buffer.byteLength(body);
+        const header = encodeHeader({ gen, base });
+
+        this.pending = [];
+        await writeAtomic(join(this.dir, HISTORY_FILE), header + body);
+        this.gen = gen;
+        this.base = base;
+        this.bytes = Buffer.byteLength(header + body);
+        this.rewriteNext = false;
+    }
+
+    private async saveState(): Promise<void> {
+        if (this.gen === null) {
+            return;
+        }
+
+        await writeAtomic(join(this.dir, STATE_FILE), encodeState(this.state, { gen: this.gen, bytes: this.bytes }));
+        this.stateDirty = false;
+    }
+
+    /*
+     * Moves an index of the single-file format into history and scan state. The history goes first:
+     * once it is there, a load treats the index as moved, so an interrupted migration finishes on the next load.
+     */
+    private async migrate(): Promise<void> {
+        const legacy = join(this.dir, LEGACY_INDEX_FILE);
+        const text = await readOptional(legacy);
+
+        if (text === null) {
+            return;
+        }
+
+        const index = decodeIndex(text.toString('utf8'));
+
+        if (index === null) {
+            console.warn('The usage index would not read; it is kept beside a fresh history.');
+            await rename(legacy, `${legacy}.unreadable`);
+
+            return;
+        }
+
+        for (const [path, file] of index) {
+            const { records, tail, provider, ...state } = file;
+            this.usage.set(path, { provider, records, tail });
+            this.state.set(path, { provider, ...state });
+        }
+
+        try {
+            await this.rewrite();
+            await this.saveState();
+            await removeFile(legacy);
+        } catch (e) {
+            // The index stays and is moved again on the next load; this run counts from memory.
+            this.rewriteNext = true;
+            this.stateDirty = true;
+            console.warn('The usage index could not be moved to the history yet:', errorText(e));
+        }
+    }
+}