docfast/dist/middleware/usage.js

import { isProKey } from "../services/keys.js";
import logger from "../services/logger.js";
import { queryWithRetry, connectWithRetry } from "../services/db.js";
const FREE_TIER_LIMIT = 100;
const PRO_TIER_LIMIT = 5000;
// In-memory cache, periodically synced to PostgreSQL
let usage = new Map();
// Write-behind buffer for batching DB writes (Audit #10)
const dirtyKeys = new Set();
const retryCount = new Map();
const MAX_RETRIES = 3;
const FLUSH_INTERVAL_MS = 5000;
const FLUSH_THRESHOLD = 50;
function getMonthKey() {
    const d = new Date();
    return `${d.getFullYear()}-${String(d.getMonth() + 1).padStart(2, "0")}`;
}
export async function loadUsageData() {
    try {
        const result = await queryWithRetry("SELECT key, count, month_key FROM usage");
        usage = new Map();
        for (const row of result.rows) {
            usage.set(row.key, { count: row.count, monthKey: row.month_key });
        }
        logger.info(`Loaded usage data for ${usage.size} keys from PostgreSQL`);
    }
    catch (error) {
        logger.info("No existing usage data found, starting fresh");
        usage = new Map();
    }
}
// Batch flush dirty entries to DB (Audit #10 + #12)
export async function flushDirtyEntries() {
    if (dirtyKeys.size === 0)
        return;
    const keysToFlush = [...dirtyKeys];
    for (const key of keysToFlush) {
        const record = usage.get(key);
        if (!record)
            continue;
        const client = await connectWithRetry();
        try {
            await client.query(`INSERT INTO usage (key, count, month_key) VALUES ($1, $2, $3)
         ON CONFLICT (key) DO UPDATE SET count = $2, month_key = $3`, [key, record.count, record.monthKey]);
            dirtyKeys.delete(key);
            retryCount.delete(key);
        }
        catch (error) {
            // Audit #12: retry logic for failed writes
            const retries = (retryCount.get(key) || 0) + 1;
            if (retries >= MAX_RETRIES) {
                logger.error({ key: key.slice(0, 8) + "...", retries }, "CRITICAL: Usage write failed after max retries, data may diverge");
                dirtyKeys.delete(key);
                retryCount.delete(key);
            }
            else {
                retryCount.set(key, retries);
                logger.warn({ key: key.slice(0, 8) + "...", retries }, "Usage write failed, will retry");
            }
        }
        finally {
            client.release();
        }
    }
}
// Periodic flush
setInterval(flushDirtyEntries, FLUSH_INTERVAL_MS);
// Note: SIGTERM/SIGINT flush is handled by the shutdown orchestrator in index.ts
// to avoid race conditions with pool.end().
export function usageMiddleware(req, res, next) {
    const keyInfo = req.apiKeyInfo;
    const key = keyInfo?.key || "unknown";
    const monthKey = getMonthKey();
    if (isProKey(key)) {
        const record = usage.get(key);
        if (record && record.monthKey === monthKey && record.count >= PRO_TIER_LIMIT) {
            res.status(429).json({ error: "Pro tier limit reached (5,000/month). Contact support for higher limits." });
            return;
        }
        trackUsage(key, monthKey);
        next();
        return;
    }
    const record = usage.get(key);
    if (record && record.monthKey === monthKey && record.count >= FREE_TIER_LIMIT) {
        res.status(429).json({ error: "Free tier limit reached (100/month). Upgrade to Pro at https://docfast.dev/#pricing for 5,000 PDFs/month." });
        return;
    }
    trackUsage(key, monthKey);
    next();
}
function trackUsage(key, monthKey) {
    const record = usage.get(key);
    if (!record || record.monthKey !== monthKey) {
        usage.set(key, { count: 1, monthKey });
    }
    else {
        record.count++;
    }
    dirtyKeys.add(key);
    // Flush immediately if threshold reached
    if (dirtyKeys.size >= FLUSH_THRESHOLD) {
        flushDirtyEntries().catch((err) => logger.error({ err }, "Threshold flush failed"));
    }
}
export function getUsageForKey(key) {
    const monthKey = getMonthKey();
    const record = usage.get(key);
    if (record && record.monthKey === monthKey) {
        return { count: record.count, monthKey };
    }
    return { count: 0, monthKey };
}
export function getUsageStats(apiKey) {
    const stats = {};
    if (apiKey) {
        const record = usage.get(apiKey);
        if (record) {
            const masked = apiKey.slice(0, 8) + "...";
            stats[masked] = { count: record.count, month: record.monthKey };
        }
    }
    return stats;
}