comfyui-mcp/dist/services/generation-tracker.js

Local-first, agent-native control plane for ComfyUI — MCP server + autonomous sidebar agent that drives your live graph in natural language on ANY LLM: Claude/ChatGPT/Gemini on your subscription (no API key), free local models via Ollama (fully offline),

import { createHash } from "node:crypto";
import { join } from "node:path";
import { existsSync, mkdirSync, renameSync } from "node:fs";
import { homedir } from "node:os";
import BetterSqlite3 from "better-sqlite3";
import { config, getInstanceSlug } from "../config.js";
import { logger } from "../utils/logger.js";
import { FileHasher } from "./file-hasher.js";
// ---------------------------------------------------------------------------
// Schema
// ---------------------------------------------------------------------------
const INIT_SQL = `
CREATE TABLE IF NOT EXISTS generations (
  id              INTEGER PRIMARY KEY AUTOINCREMENT,
  created_at      TEXT    NOT NULL DEFAULT (datetime('now')),
  settings_hash   TEXT    NOT NULL,
  model_family    TEXT    NOT NULL,
  model_hash      TEXT    NOT NULL,
  model_name      TEXT,
  preset_name     TEXT,
  sampler         TEXT    NOT NULL,
  scheduler       TEXT    NOT NULL,
  steps           INTEGER NOT NULL,
  cfg             REAL    NOT NULL,
  denoise         REAL    DEFAULT 1.0,
  shift           REAL,
  width           INTEGER NOT NULL,
  height          INTEGER NOT NULL,
  lora_hash       TEXT,
  lora_name       TEXT,
  lora_strength   REAL,
  lora_civitai_id INTEGER,
  reuse_count     INTEGER NOT NULL DEFAULT 1,
  neg_prompt_hash TEXT
);

CREATE INDEX IF NOT EXISTS idx_gen_settings_hash ON generations(settings_hash);
CREATE INDEX IF NOT EXISTS idx_gen_model_family  ON generations(model_family);
CREATE INDEX IF NOT EXISTS idx_gen_model_hash    ON generations(model_hash);
CREATE INDEX IF NOT EXISTS idx_gen_lora_hash     ON generations(lora_hash);
CREATE INDEX IF NOT EXISTS idx_gen_lora_civitai  ON generations(lora_civitai_id);
CREATE INDEX IF NOT EXISTS idx_gen_created       ON generations(created_at);
`;
// ---------------------------------------------------------------------------
// Settings hash computation
// ---------------------------------------------------------------------------
/**
 * Compute the settings fingerprint. This determines what counts as "the same combo".
 * Resolution, display names, and preset names are intentionally excluded.
 */
function computeSettingsHash(entry) {
    const canonical = {
        cfg: entry.cfg,
        denoise: entry.denoise,
        model_family: entry.modelFamily,
        model_hash: entry.modelHash,
        sampler: entry.sampler,
        scheduler: entry.scheduler,
        steps: entry.steps,
    };
    // Only include optional fields if they have values
    if (entry.shift != null)
        canonical.shift = entry.shift;
    if (entry.loraHash)
        canonical.lora_hash = entry.loraHash;
    if (entry.loraStrength != null && entry.loraHash)
        canonical.lora_strength = entry.loraStrength;
    // Sort keys for deterministic output
    const sorted = Object.keys(canonical)
        .sort()
        .reduce((acc, key) => {
        acc[key] = canonical[key];
        return acc;
    }, {});
    return createHash("sha256").update(JSON.stringify(sorted)).digest("hex");
}
// ---------------------------------------------------------------------------
// GenerationTracker class
// ---------------------------------------------------------------------------
export class GenerationTracker {
    db;
    fileHasher;
    constructor(dbPath) {
        const resolvedPath = dbPath ?? this.defaultDbPath();
        // Ensure parent directory exists
        const dir = join(resolvedPath, "..");
        mkdirSync(dir, { recursive: true });
        // Existing installs have real generation history at the OLD in-tree path.
        // Moving where we look without bringing it along would silently orphan it —
        // the tracker would come up empty and every past generation would look like
        // it never happened. Only when the caller did not name a path: an explicit
        // `dbPath` is the caller's business.
        if (!dbPath)
            this.migrateLegacyInTreeDb(resolvedPath);
        this.db = new BetterSqlite3(resolvedPath);
        this.db.pragma("journal_mode = WAL");
        this.db.pragma("foreign_keys = ON");
        this.db.exec(INIT_SQL);
        this.fileHasher = new FileHasher(this.db);
        logger.info(`Generation tracker DB: ${resolvedPath}`);
    }
    /**
     * One-time move of the pre-#872 database out of the ComfyUI git tree.
     *
     * WAL means THREE files (`.db`, `-wal`, `-shm`) and all of them have to travel
     * together: a `.db` carried across without its `-wal` loses every transaction
     * that had not been checkpointed, which is silent data loss dressed as a
     * successful migration.
     *
     * Best-effort by design. This runs at construction, and failing to move an old
     * database is not a reason to refuse to track new generations — but it IS a
     * reason to say so loudly, because the alternative is a user quietly wondering
     * where their history went.
     */
    migrateLegacyInTreeDb(target) {
        if (!config.comfyuiPath)
            return;
        const legacy = join(config.comfyuiPath, "comfyui-mcp", "generations.db");
        if (legacy === target || !existsSync(legacy))
            return;
        if (existsSync(target)) {
            // Both exist. The new one is authoritative — this process may already have
            // written to it — and overwriting either would destroy history. Say where
            // the old one is and leave it alone.
            logger.warn(`[generation-tracker] a legacy generation database is still at ${legacy}, inside the ` +
                `ComfyUI install, but ${target} already exists so it was NOT moved. The legacy file ` +
                `is not being read. Note that its presence can break ComfyUI's own git updater ` +
                `(#872) — move or delete it once you have salvaged anything you need.`);
            return;
        }
        const suffixes = ["", "-wal", "-shm"];
        try {
            for (const suffix of suffixes) {
                const from = `${legacy}${suffix}`;
                if (existsSync(from))
                    renameSync(from, `${target}${suffix}`);
            }
            logger.info(`[generation-tracker] moved the generation database out of the ComfyUI install ` +
                `(${legacy} → ${target}). It lived inside ComfyUI's git working tree, where it could ` +
                `break ComfyUI's updater (#872).`);
        }
        catch (err) {
            // A partial move is the dangerous outcome: the `.db` here and its `-wal`
            // there is a corrupt pair. Report the whole state rather than a bare error.
            logger.warn(`[generation-tracker] could not move the legacy generation database out of the ComfyUI ` +
                `install (${legacy} → ${target}): ${err instanceof Error ? err.message : String(err)}. ` +
                `Some of the .db/-wal/-shm set may have moved and some may not — check BOTH locations ` +
                `before deleting anything, and keep the .db together with its -wal.`);
        }
    }
    defaultDbPath() {
        // NEVER inside the ComfyUI install (#872). This used to return
        // `<comfyuiPath>/comfyui-mcp/generations.db` for every local install, which
        // put a live WAL SQLite database — three open files — inside ComfyUI's own
        // GIT WORKING TREE, in a directory its `.gitignore` does not cover.
        //
        // ComfyUI's updater saves uncommitted work with `git add -A`, so it commits
        // our database, then `git reset --hard` tries to remove it — and Windows
        // refuses to delete a file whose handle we hold for the orchestrator's entire
        // lifetime. The user's ComfyUI update ABORTS, and its rollback fails too,
        // leaving them on a backup branch. Writing our state into someone else's
        // version-controlled tree was the mistake; the file being locked is just what
        // made it visible.
        //
        // The per-user data dir was already the path for remote/cloud, and its
        // instance slug (host+port) keeps separate installs separate — which is the
        // only thing the in-tree location was really achieving.
        const baseDir = process.env.COMFYUI_MCP_DATA_DIR?.trim() || join(homedir(), ".comfyui-mcp");
        return join(baseDir, "instances", getInstanceSlug(), "generations.db");
    }
    /**
     * Log a generation. If the same settings_hash already exists,
     * increments reuse_count and updates timestamp instead of inserting.
     */
    logGeneration(entry) {
        const settingsHash = computeSettingsHash(entry);
        const existing = this.db
            .prepare(`SELECT id, reuse_count FROM generations WHERE settings_hash = ?`)
            .get(settingsHash);
        if (existing) {
            this.db
                .prepare(`UPDATE generations
           SET reuse_count = reuse_count + 1,
               created_at = datetime('now'),
               width = ?,
               height = ?
           WHERE id = ?`)
                .run(entry.width, entry.height, existing.id);
            const newCount = existing.reuse_count + 1;
            logger.info(`Generation logged (reuse #${newCount}): ${entry.modelFamily} / ${entry.sampler}/${entry.scheduler} / ${entry.steps} steps`);
            return { settingsHash, reuseCount: newCount };
        }
        this.db
            .prepare(`INSERT INTO generations (
           settings_hash, model_family, model_hash, model_name, preset_name,
           sampler, scheduler, steps, cfg, denoise, shift,
           width, height,
           lora_hash, lora_name, lora_strength, lora_civitai_id,
           neg_prompt_hash
         ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`)
            .run(settingsHash, entry.modelFamily, entry.modelHash, entry.modelName, entry.presetName, entry.sampler, entry.scheduler, entry.steps, entry.cfg, entry.denoise, entry.shift, entry.width, entry.height, entry.loraHash, entry.loraName, entry.loraStrength, entry.loraCivitaiId, entry.negPromptHash);
        logger.info(`Generation logged (new): ${entry.modelFamily} / ${entry.sampler}/${entry.scheduler} / ${entry.steps} steps`);
        return { settingsHash, reuseCount: 1 };
    }
    /**
     * Get top settings for a model family, ordered by reuse count.
     */
    suggestSettings(modelFamily, limit = 10) {
        const rows = this.db
            .prepare(`SELECT settings_hash, model_family, model_hash, model_name, preset_name,
                sampler, scheduler, steps, cfg, denoise, shift,
                lora_hash, lora_name, lora_strength, reuse_count, created_at
         FROM generations
         WHERE model_family = ?
         ORDER BY reuse_count DESC, created_at DESC
         LIMIT ?`)
            .all(modelFamily, limit);
        return rows.map(mapRowToTopSetting);
    }
    /**
     * Get top settings for a specific LoRA hash.
     */
    suggestSettingsForLora(loraHash, limit = 10) {
        const rows = this.db
            .prepare(`SELECT settings_hash, model_family, model_hash, model_name, preset_name,
                sampler, scheduler, steps, cfg, denoise, shift,
                lora_hash, lora_name, lora_strength, reuse_count, created_at
         FROM generations
         WHERE lora_hash = ?
         ORDER BY reuse_count DESC, created_at DESC
         LIMIT ?`)
            .all(loraHash, limit);
        return rows.map(mapRowToTopSetting);
    }
    /**
     * Search generations by model or LoRA name (substring match).
     */
    searchByName(query, limit = 20) {
        const pattern = `%${query}%`;
        const rows = this.db
            .prepare(`SELECT settings_hash, model_family, model_hash, model_name, preset_name,
                sampler, scheduler, steps, cfg, denoise, shift,
                lora_hash, lora_name, lora_strength, reuse_count, created_at
         FROM generations
         WHERE model_name LIKE ? OR lora_name LIKE ?
         ORDER BY reuse_count DESC, created_at DESC
         LIMIT ?`)
            .all(pattern, pattern, limit);
        return rows.map(mapRowToTopSetting);
    }
    /**
     * Get overall generation statistics.
     */
    getStats(modelFamily) {
        const whereClause = modelFamily ? `WHERE model_family = ?` : "";
        const params = modelFamily ? [modelFamily] : [];
        const totalRow = this.db
            .prepare(`SELECT COALESCE(SUM(reuse_count), 0) AS total FROM generations ${whereClause}`)
            .get(...params);
        const uniqueRow = this.db
            .prepare(`SELECT COUNT(*) AS cnt FROM generations ${whereClause}`)
            .get(...params);
        const breakdownRows = this.db
            .prepare(`SELECT model_family, COALESCE(SUM(reuse_count), 0) AS total
         FROM generations
         GROUP BY model_family
         ORDER BY total DESC`)
            .all();
        const topRows = this.db
            .prepare(`SELECT settings_hash, model_family, model_hash, model_name, preset_name,
                sampler, scheduler, steps, cfg, denoise, shift,
                lora_hash, lora_name, lora_strength, reuse_count, created_at
         FROM generations ${whereClause}
         ORDER BY reuse_count DESC, created_at DESC
         LIMIT 5`)
            .all(...params);
        return {
            totalGenerations: totalRow.total,
            uniqueCombos: uniqueRow.cnt,
            modelBreakdown: breakdownRows.map((r) => ({
                modelFamily: r.model_family,
                count: r.total,
            })),
            topSettings: topRows.map(mapRowToTopSetting),
        };
    }
    /**
     * Get shareable entries (only those with public LoRAs or no LoRA).
     */
    getShareableEntries() {
        const rows = this.db
            .prepare(`SELECT settings_hash, model_family, model_hash, model_name, preset_name,
                sampler, scheduler, steps, cfg, denoise, shift,
                lora_hash, lora_name, lora_strength, reuse_count, created_at
         FROM generations
         WHERE lora_hash IS NULL OR lora_civitai_id IS NOT NULL
         ORDER BY reuse_count DESC`)
            .all();
        return rows.map(mapRowToTopSetting);
    }
    close() {
        this.db.close();
    }
}
// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------
function mapRowToTopSetting(row) {
    return {
        settingsHash: row.settings_hash,
        modelFamily: row.model_family,
        modelHash: row.model_hash,
        modelName: row.model_name,
        presetName: row.preset_name,
        sampler: row.sampler,
        scheduler: row.scheduler,
        steps: row.steps,
        cfg: row.cfg,
        denoise: row.denoise,
        shift: row.shift,
        loraHash: row.lora_hash,
        loraName: row.lora_name,
        loraStrength: row.lora_strength,
        reuseCount: row.reuse_count,
        createdAt: row.created_at,
    };
}
// ---------------------------------------------------------------------------
// Singleton
// ---------------------------------------------------------------------------
let _tracker = null;
export function getTracker() {
    if (!_tracker) {
        _tracker = new GenerationTracker();
    }
    return _tracker;
}
//# sourceMappingURL=generation-tracker.js.map