@token-dashboard/codex-usage-uploader 0.1.0

package/src/state-db.js

@@ -0,0 +1,334 @@
+import fs from 'node:fs';
+import path from 'node:path';
+import { DatabaseSync } from 'node:sqlite';
+import {
+  MAX_BATCH_BYTES,
+  MAX_BUFFER_AGE_SECONDS,
+  MAX_EVENTS_PER_BATCH,
+} from './constants.js';
+import { nowTs, stableStringify } from './utils.js';
+
+export class StateDb {
+  constructor(dbPath) {
+    this.dbPath = dbPath;
+    fs.mkdirSync(path.dirname(dbPath), { recursive: true });
+    this.db = new DatabaseSync(dbPath, { readBigInts: true });
+    this.initSchema();
+  }
+
+  close() {
+    this.db.close();
+  }
+
+  initSchema() {
+    this.db.exec(`
+      CREATE TABLE IF NOT EXISTS identity_config (
+        key TEXT PRIMARY KEY,
+        value TEXT NOT NULL,
+        updated_at REAL NOT NULL
+      );
+
+      CREATE TABLE IF NOT EXISTS pending_batches (
+        id INTEGER PRIMARY KEY AUTOINCREMENT,
+        batch_key TEXT NOT NULL UNIQUE,
+        status TEXT NOT NULL,
+        payload_json TEXT NOT NULL,
+        payload_bytes INTEGER NOT NULL DEFAULT 0,
+        session_count INTEGER NOT NULL DEFAULT 0,
+        turn_count INTEGER NOT NULL DEFAULT 0,
+        event_count INTEGER NOT NULL DEFAULT 0,
+        attempt_count INTEGER NOT NULL DEFAULT 0,
+        next_retry_at REAL,
+        last_error TEXT,
+        created_at REAL NOT NULL,
+        updated_at REAL NOT NULL
+      );
+
+      CREATE TABLE IF NOT EXISTS upload_checkpoint (
+        key TEXT PRIMARY KEY,
+        value TEXT NOT NULL,
+        updated_at REAL NOT NULL
+      );
+    `);
+    this.ensureIngestionFilesTable();
+  }
+
+  ensureIngestionFilesTable() {
+    const existing = this.db
+      .prepare(
+        "SELECT name FROM sqlite_master WHERE type = 'table' AND name = 'ingestion_files'",
+      )
+      .get();
+    if (!existing) {
+      this.createIngestionFilesTable();
+      return;
+    }
+
+    const columns = this.db.prepare('PRAGMA table_info(ingestion_files)').all();
+    const hasSourceRoot = columns.some((column) => column.name === 'source_root');
+    if (hasSourceRoot) {
+      return;
+    }
+
+    this.db.exec('BEGIN');
+    try {
+      this.db.exec('ALTER TABLE ingestion_files RENAME TO ingestion_files_legacy');
+      this.createIngestionFilesTable();
+      this.db.prepare(`
+        INSERT INTO ingestion_files(
+          source_root, relpath, file_size, file_mtime_ns, last_line_no, state_json, updated_at
+        )
+        SELECT 'sessions', relpath, file_size, file_mtime_ns, last_line_no, state_json, updated_at
+        FROM ingestion_files_legacy
+      `).run();
+      this.db.exec('DROP TABLE ingestion_files_legacy');
+      this.db.exec('COMMIT');
+    } catch (error) {
+      this.db.exec('ROLLBACK');
+      throw error;
+    }
+  }
+
+  createIngestionFilesTable() {
+    this.db.exec(`
+      CREATE TABLE IF NOT EXISTS ingestion_files (
+        source_root TEXT NOT NULL,
+        relpath TEXT NOT NULL,
+        file_size INTEGER NOT NULL,
+        file_mtime_ns INTEGER NOT NULL,
+        last_line_no INTEGER NOT NULL,
+        state_json TEXT,
+        updated_at REAL NOT NULL,
+        PRIMARY KEY (source_root, relpath)
+      );
+    `);
+  }
+
+  getIdentity() {
+    const rows = this.db.prepare('SELECT key, value FROM identity_config').all();
+    return Object.fromEntries(rows.map((row) => [row.key, row.value]));
+  }
+
+  setIdentity(values) {
+    const ts = nowTs();
+    const upsert = this.db.prepare(`
+      INSERT INTO identity_config(key, value, updated_at)
+      VALUES (?, ?, ?)
+      ON CONFLICT(key) DO UPDATE SET value = excluded.value, updated_at = excluded.updated_at
+    `);
+    const del = this.db.prepare('DELETE FROM identity_config WHERE key = ?');
+    this.db.exec('BEGIN');
+    try {
+      for (const [key, value] of Object.entries(values)) {
+        if (value == null) {
+          del.run(key);
+        } else {
+          upsert.run(key, value, ts);
+        }
+      }
+      this.db.exec('COMMIT');
+    } catch (error) {
+      this.db.exec('ROLLBACK');
+      throw error;
+    }
+  }
+
+  getCheckpoint(key) {
+    const row = this.db.prepare('SELECT value FROM upload_checkpoint WHERE key = ?').get(key);
+    return row?.value ?? null;
+  }
+
+  setCheckpoint(key, value) {
+    const ts = nowTs();
+    this.db.prepare(`
+      INSERT INTO upload_checkpoint(key, value, updated_at)
+      VALUES (?, ?, ?)
+      ON CONFLICT(key) DO UPDATE SET value = excluded.value, updated_at = excluded.updated_at
+    `).run(key, value, ts);
+  }
+
+  getFileState(sourceRoot, relpath) {
+    const stmt = this.db.prepare(
+      'SELECT * FROM ingestion_files WHERE source_root = ? AND relpath = ?',
+    );
+    stmt.setReadBigInts(true);
+    return stmt.get(sourceRoot, relpath) ?? null;
+  }
+
+  upsertFileState(sourceRoot, relpath, fileSize, fileMtimeNs, lastLineNo, state) {
+    this.db.prepare(`
+      INSERT INTO ingestion_files(
+        source_root, relpath, file_size, file_mtime_ns, last_line_no, state_json, updated_at
+      )
+      VALUES (?, ?, ?, ?, ?, ?, ?)
+      ON CONFLICT(source_root, relpath) DO UPDATE SET
+        file_size = excluded.file_size,
+        file_mtime_ns = excluded.file_mtime_ns,
+        last_line_no = excluded.last_line_no,
+        state_json = excluded.state_json,
+        updated_at = excluded.updated_at
+    `).run(
+      sourceRoot,
+      relpath,
+      fileSize,
+      fileMtimeNs,
+      lastLineNo,
+      stableStringify(state),
+      nowTs(),
+    );
+  }
+
+  getBufferingBatch() {
+    const stmt = this.db.prepare(`
+      SELECT * FROM pending_batches
+      WHERE status = 'buffering'
+      ORDER BY id ASC
+      LIMIT 1
+    `);
+    stmt.setReadBigInts(true);
+    return stmt.get() ?? null;
+  }
+
+  saveBufferingPayload(payload) {
+    const existing = this.getBufferingBatch();
+    const payloadJson = stableStringify(payload);
+    const payloadBytes = Buffer.byteLength(payloadJson, 'utf8');
+    const sessionCount = payload.sessions?.length ?? 0;
+    const turnCount = payload.turns?.length ?? 0;
+    const eventCount = payload.events?.length ?? 0;
+    const ts = nowTs();
+    if (existing) {
+      this.db.prepare(`
+        UPDATE pending_batches
+        SET payload_json = ?, payload_bytes = ?, session_count = ?, turn_count = ?, event_count = ?, updated_at = ?
+        WHERE id = ?
+      `).run(payloadJson, payloadBytes, sessionCount, turnCount, eventCount, ts, existing.id);
+    } else {
+      this.db.prepare(`
+        INSERT INTO pending_batches(
+          batch_key, status, payload_json, payload_bytes, session_count, turn_count, event_count,
+          attempt_count, next_retry_at, created_at, updated_at
+        ) VALUES (?, 'buffering', ?, ?, ?, ?, ?, 0, NULL, ?, ?)
+      `).run(randomBatchKey(), payloadJson, payloadBytes, sessionCount, turnCount, eventCount, ts, ts);
+    }
+    return this.getBufferingBatch();
+  }
+
+  sealStaleBatches(force = false) {
+    const selectStmt = this.db.prepare(
+      "SELECT id, created_at FROM pending_batches WHERE status = 'buffering'",
+    );
+    selectStmt.setReadBigInts(true);
+    const rows = selectStmt.all();
+    let sealed = 0;
+    const ts = nowTs();
+    const updateStmt = this.db.prepare(`
+      UPDATE pending_batches
+      SET status = 'pending', next_retry_at = ?, updated_at = ?
+      WHERE id = ?
+    `);
+    for (const row of rows) {
+      if (force || ts - Number(row.created_at) >= MAX_BUFFER_AGE_SECONDS) {
+        updateStmt.run(ts, ts, row.id);
+        sealed += 1;
+      }
+    }
+    return sealed;
+  }
+
+  sealBufferIfThresholdHit() {
+    const row = this.getBufferingBatch();
+    if (!row) return false;
+    const shouldSeal =
+      Number(row.event_count) >= MAX_EVENTS_PER_BATCH ||
+      Number(row.payload_bytes) >= MAX_BATCH_BYTES ||
+      nowTs() - Number(row.created_at) >= MAX_BUFFER_AGE_SECONDS;
+    if (shouldSeal) {
+      const ts = nowTs();
+      this.db.prepare(`
+        UPDATE pending_batches
+        SET status = 'pending', next_retry_at = ?, updated_at = ?
+        WHERE id = ?
+      `).run(ts, ts, row.id);
+    }
+    return shouldSeal;
+  }
+
+  iterDuePendingBatches() {
+    const stmt = this.db.prepare(`
+      SELECT * FROM pending_batches
+      WHERE status = 'pending' AND (next_retry_at IS NULL OR next_retry_at <= ?)
+      ORDER BY created_at ASC, id ASC
+    `);
+    stmt.setReadBigInts(true);
+    return stmt.all(nowTs());
+  }
+
+  markAllPendingDue() {
+    this.db.prepare(`
+      UPDATE pending_batches
+      SET next_retry_at = 0, updated_at = ?
+      WHERE status = 'pending'
+    `).run(nowTs());
+  }
+
+  markBatchUploaded(batchId) {
+    this.db.prepare('DELETE FROM pending_batches WHERE id = ?').run(batchId);
+  }
+
+  markBatchFailed(batchId, attemptCount, errorMessage) {
+    const nextRetryAt = nowTs() + Math.min(300, 2 ** Math.min(attemptCount, 8));
+    this.db.prepare(`
+      UPDATE pending_batches
+      SET attempt_count = ?, next_retry_at = ?, last_error = ?, updated_at = ?
+      WHERE id = ?
+    `).run(attemptCount, nextRetryAt, String(errorMessage).slice(0, 1000), nowTs(), batchId);
+  }
+
+  resetBackfillState() {
+    this.db.exec('BEGIN');
+    try {
+      this.db.prepare('DELETE FROM ingestion_files').run();
+      this.db.prepare('DELETE FROM pending_batches').run();
+      this.db.prepare('DELETE FROM upload_checkpoint').run();
+      this.db.exec('COMMIT');
+    } catch (error) {
+      this.db.exec('ROLLBACK');
+      throw error;
+    }
+  }
+
+  getQueueStats() {
+    const stmt = this.db.prepare(`
+      SELECT
+        COALESCE(SUM(CASE WHEN status = 'buffering' THEN 1 ELSE 0 END), 0) AS buffering_batch_count,
+        COALESCE(SUM(CASE WHEN status = 'pending' AND attempt_count = 0 THEN 1 ELSE 0 END), 0) AS pending_batch_count,
+        COALESCE(SUM(CASE WHEN status = 'pending' AND attempt_count > 0 THEN 1 ELSE 0 END), 0) AS retrying_batch_count,
+        COALESCE(SUM(session_count), 0) AS queued_sessions,
+        COALESCE(SUM(turn_count), 0) AS queued_turns,
+        COALESCE(SUM(event_count), 0) AS queued_events,
+        MIN(created_at) AS oldest_created_at
+      FROM pending_batches
+    `);
+    stmt.setReadBigInts(true);
+    const row = stmt.get() ?? {};
+    const oldestCreatedAt =
+      row.oldest_created_at == null ? null : Number(row.oldest_created_at);
+    return {
+      bufferingBatchCount: Number(row.buffering_batch_count ?? 0),
+      pendingBatchCount: Number(row.pending_batch_count ?? 0),
+      retryingBatchCount: Number(row.retrying_batch_count ?? 0),
+      queuedSessions: Number(row.queued_sessions ?? 0),
+      queuedTurns: Number(row.queued_turns ?? 0),
+      queuedEvents: Number(row.queued_events ?? 0),
+      oldestPendingAgeSeconds:
+        oldestCreatedAt == null
+          ? null
+          : Math.max(0, Math.floor(nowTs() - oldestCreatedAt)),
+    };
+  }
+}
+
+function randomBatchKey() {
+  return `${Date.now().toString(16)}${Math.random().toString(16).slice(2, 14)}`;
+}

package/src/utils.js

@@ -0,0 +1,53 @@
+import fs from 'node:fs';
+import os from 'node:os';
+import path from 'node:path';
+import { Buffer } from 'node:buffer';
+import { createHash } from 'node:crypto';
+
+export function nowTs() {
+  return Date.now() / 1000;
+}
+
+export function sleep(ms) {
+  return new Promise((resolve) => setTimeout(resolve, ms));
+}
+
+export function stableStringify(value) {
+  return JSON.stringify(sortForJson(value));
+}
+
+function sortForJson(value) {
+  if (Array.isArray(value)) {
+    return value.map(sortForJson);
+  }
+  if (value && typeof value === 'object') {
+    return Object.fromEntries(
+      Object.keys(value)
+        .sort()
+        .map((key) => [key, sortForJson(value[key])]),
+    );
+  }
+  return value;
+}
+
+export function isoToEpochMs(value) {
+  if (!value) return null;
+  const normalized = String(value).replace('Z', '+00:00');
+  const timestamp = Date.parse(normalized);
+  return Number.isNaN(timestamp) ? null : timestamp;
+}
+
+export function deriveRepoName(repositoryUrl, cwd) {
+  if (repositoryUrl) {
+    let last = repositoryUrl.replace(/\/+$/, '').split('/').pop() ?? '';
+    if (last.endsWith('.git')) last = last.slice(0, -4);
+    if (last) return last;
+  }
+  if (cwd) return path.basename(cwd) || null;
+  return null;
+}
+
+export function computeEventUid(collectorId, relpath, lineNo) {
+  const raw = Buffer.from(`${collectorId}|${relpath}|${lineNo}`, 'utf8');
+  return createHash('sha1').update(raw).digest('hex');
+}