comfyui-mcp
Version:
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),
479 lines • 20.9 kB
JavaScript
import { mkdir, writeFile, readdir, unlink, stat } from "node:fs/promises";
import { join } from "node:path";
import { tmpdir } from "node:os";
import { getClient, ensureConnected, getHistory, } from "../comfyui/client.js";
import { getCloudUrl, getComfyUIBaseUrl, isCloudMode, } from "../config.js";
import { attachExecutionListeners } from "../comfyui/events.js";
import { logger } from "../utils/logger.js";
import { AssetRegistry, normalizeAssetType } from "./asset-registry.js";
import { analyzeHistoryEntry, hasAffirmativeSuccessStatus, historyCompletionTimeMs, historyTerminalTimeMs, normalizeHistoryMessages, } from "./job-history.js";
// ── Constants ──────────────────────────────────────────────────────────
const COMPLETIONS_DIR = join(tmpdir(), "comfyui-mcp-completions");
const HISTORY_FLUSH_DELAY_MS = 500;
const REPORTED_MAX_AGE_MS = 60 * 60 * 1000; // 1 hour
// Job-watch timeout + poll cadence are env-tunable: long video renders
// (LTX / WAN can take 15+ min) were getting killed by the previous
// hardcoded 10-minute ceiling. Defaults: 30-minute timeout, 2 s poll.
// Read at call time (not module load) so tests and long-lived servers
// can adjust without a restart. Gap flagged by josephoibrahim/comfy-cozy.
export function watcherTimeoutMs() {
const raw = Number(process.env.COMFYUI_JOB_TIMEOUT_S ?? "1800");
return Number.isFinite(raw) && raw > 0 ? raw * 1000 : 1800 * 1000;
}
export function pollIntervalMs() {
const raw = Number(process.env.COMFYUI_JOB_POLL_INTERVAL_S ?? "2");
return Number.isFinite(raw) && raw > 0 ? raw * 1000 : 2000;
}
// ── Module-level state ─────────────────────────────────────────────────
const activeWatchers = new Map();
// ── Private helpers ────────────────────────────────────────────────────
function sleep(ms) {
return new Promise((resolve) => setTimeout(resolve, ms));
}
async function ensureCompletionsDir() {
await mkdir(COMPLETIONS_DIR, { recursive: true });
}
function buildImageUrl(filename, subfolder, type) {
const params = new URLSearchParams({ filename, subfolder, type });
if (isCloudMode()) {
const base = getCloudUrl().replace(/\/+$/, "");
return `${base}/api/view?${params.toString()}`;
}
return `${getComfyUIBaseUrl()}/view?${params.toString()}`;
}
const KNOWN_IMAGE_KEYS = new Set(["images"]);
const KNOWN_VIDEO_KEYS = new Set(["videos", "video", "gifs"]);
const IMAGE_EXTS = new Set([".png", ".jpg", ".jpeg", ".bmp"]);
const VIDEO_EXTS = new Set([".mp4", ".webm", ".mov", ".mkv", ".m4v", ".avi", ".gif", ".webp"]);
function isHistoryMediaRef(value) {
if (value === null || typeof value !== "object" || Array.isArray(value))
return false;
const rec = value;
if (typeof rec.filename !== "string" || rec.filename.length === 0)
return false;
if (rec.subfolder != null && typeof rec.subfolder !== "string")
return false;
if (rec.type != null && typeof rec.type !== "string")
return false;
return true;
}
/** One-level bags only: a media ref, or an array of them. Nested objects are not walked. */
function historyMediaRefs(value) {
if (Array.isArray(value))
return value.filter(isHistoryMediaRef);
return isHistoryMediaRef(value) ? [value] : [];
}
function filenameExt(filename) {
const base = filename.slice(Math.max(filename.lastIndexOf("/"), filename.lastIndexOf("\\")) + 1);
const dot = base.lastIndexOf(".");
return dot >= 0 ? base.slice(dot).toLowerCase() : "";
}
function unknownKeyMediaKind(filename) {
const ext = filenameExt(filename);
if (VIDEO_EXTS.has(ext))
return "video";
if (IMAGE_EXTS.has(ext))
return "image";
return null;
}
function toMediaOutput(ref) {
const type = normalizeAssetType(ref.type);
return {
filename: ref.filename,
subfolder: ref.subfolder ?? "",
type,
url: buildImageUrl(ref.filename, ref.subfolder ?? "", type),
};
}
export function buildCompletionNotification(promptId, entry, startTime) {
const messages = normalizeHistoryMessages(entry);
const analysis = analyzeHistoryEntry(entry);
// Timing — interrupt is a terminal end event, same as success/error (#2512).
// Duration and finished-at come from ComfyUI's execution record, never from
// the moment this notification is later delivered.
const observedAt = Date.now();
const terminalAt = historyTerminalTimeMs(entry, observedAt);
const durationMs = analysis.execution_stats?.total_duration_ms ??
(terminalAt !== undefined ? Math.max(0, terminalAt - startTime) : observedAt - startTime);
// Status
const errorMsg = messages.find((m) => m[0] === "execution_error");
const interruptMsg = messages.find((m) => m[0] === "execution_interrupted");
let status = "success";
if (errorMsg)
status = "error";
else if (interruptMsg)
status = "interrupted";
// Error details
const error = errorMsg
? analysis.error
: undefined;
// Cached nodes
const cachedMsg = messages.find((m) => m[0] === "execution_cached");
const cachedNodesRaw = cachedMsg?.[1].nodes;
const cachedNodes = Array.isArray(cachedNodesRaw)
? cachedNodesRaw.map((node) => String(node))
: [];
// Output images and videos
const outputs = [];
const video_outputs = [];
for (const [nodeId, nodeOutput] of Object.entries(entry.outputs || {})) {
if (!nodeOutput || typeof nodeOutput !== "object")
continue;
const out = nodeOutput;
// Extract image outputs (SaveImage, PreviewImage) — malformed refs are
// filtered BEFORE any use, so a bad /history entry can neither crash the
// parse nor surface as an output with filename=undefined.
const images = [];
if (Array.isArray(out.images)) {
for (const img of out.images.filter(isHistoryMediaRef)) {
images.push(toMediaOutput(img));
}
}
// Extract video outputs (SaveVideo, CreateVideo, VHS_VideoCombine).
// These nodes emit under 'videos', 'video', or 'gifs' keys. Accumulate all
// keys into a single entry per node, mirroring how images group once.
// Adapted from jcd315's fork (jcd315/comfyui-mcp-muse, commit e13342ec).
const videos = [];
for (const videoKey of KNOWN_VIDEO_KEYS) {
const videoData = out[videoKey];
if (!Array.isArray(videoData))
continue;
for (const vid of videoData.filter(isHistoryMediaRef)) {
videos.push(toMediaOutput(vid));
}
}
// #2845 — SaveVideoDexter (`dexter_video`) and NKDVideoViewer (`nkd_video`)
// emit a `{filename, subfolder?, type?}` ref under a custom key, as a single
// object or a one-item array. Harvest that one level; nested bags stay closed.
for (const [key, raw] of Object.entries(out)) {
if (KNOWN_IMAGE_KEYS.has(key) || KNOWN_VIDEO_KEYS.has(key))
continue;
for (const ref of historyMediaRefs(raw)) {
const kind = unknownKeyMediaKind(ref.filename);
if (kind === "video")
videos.push(toMediaOutput(ref));
else if (kind === "image")
images.push(toMediaOutput(ref));
}
}
if (images.length > 0) {
outputs.push({ node_id: nodeId, images });
}
if (videos.length > 0) {
video_outputs.push({ node_id: nodeId, videos });
}
}
return {
prompt_id: promptId,
status,
duration_ms: Math.round(durationMs),
timestamp: new Date(terminalAt ?? observedAt).toISOString(),
error,
outputs,
video_outputs,
...(analysis.text_outputs ? { text_outputs: analysis.text_outputs } : {}),
cached_nodes: cachedNodes,
execution_stats: analysis.execution_stats,
};
}
async function handleCompletion(promptId, state, detectedBy) {
// Race guard: only first detector proceeds
if (state.completed)
return;
state.completed = true;
// The watch's OWN observed finish time: the moment this watcher detected
// the completion (WS event or poll tick). For a live-watched job this is a
// real, truthful observation — within one poll interval / WS latency of the
// actual finish — and serves as the createdAt fallback when the history
// entry carries no usable execution_success timestamp (never the registry's
// silent default-now; provenance is recorded on the record, #751 r4 gate).
const observedFinishAt = Date.now();
logger.info(`Completion detected via ${detectedBy}`, { prompt_id: promptId });
// Stop both monitoring tracks
cleanup(state);
// Wait for history to flush (WS fires before outputs are written)
if (detectedBy === "ws") {
await sleep(HISTORY_FLUSH_DELAY_MS);
}
// Fetch history — retry once if not found
let entry;
for (let attempt = 0; attempt < 2; attempt++) {
try {
const history = await getHistory(promptId);
entry = history[promptId];
if (entry?.status.completed)
break;
}
catch (err) {
logger.warn(`History fetch attempt ${attempt + 1} failed`, {
prompt_id: promptId,
error: err instanceof Error ? err.message : err,
});
}
if (attempt === 0)
await sleep(500);
}
if (!entry) {
logger.error("Could not fetch history for completed job", {
prompt_id: promptId,
});
activeWatchers.delete(promptId);
return;
}
// Build and write notification
try {
const notification = buildCompletionNotification(promptId, entry, state.startTime);
// Register outputs with the AssetRegistry so they can be referenced by
// asset_id for get_image (action:"view") and for
// generate_image (action:"regenerate"). Registration requires AFFIRMATIVE
// success evidence — the shared predicate (job-history) both registration
// paths use. The live watch only detects "finished" (WS and poll both
// route success AND error to this handler), and notification.status is a
// derivation that defaults to success when messages are absent/malformed;
// the history entry's own status_str + messages are therefore the
// AUTHORITATIVE status source for registration on this path too. (The
// derived status still drives the completion FILE, where presenting a
// best-effort outcome is a display matter, not an asset guarantee.)
if (hasAffirmativeSuccessStatus(entry) && state.workflow) {
try {
// createdAt is a REAL time, never the registry's silent default-now:
// prefer the entry's recorded execution_success timestamp; fall back
// to this watcher's own observed finish time — distinguishable via
// createdAtSource ("history" vs "observed").
const recordedAt = historyCompletionTimeMs(entry, observedFinishAt);
AssetRegistry.register({
promptId,
workflow: state.workflow,
outputs: notification.outputs.map((o) => ({
node_id: o.node_id,
images: o.images.map((img) => ({
filename: img.filename,
subfolder: img.subfolder,
type: img.type,
url: img.url,
})),
})),
createdAt: recordedAt ?? observedFinishAt,
createdAtSource: recordedAt !== undefined ? "history" : "observed",
});
for (const output of notification.outputs) {
for (const img of output.images) {
const record = AssetRegistry.find(promptId, img);
if (record)
img.asset_id = record.assetId;
}
}
}
catch (regErr) {
logger.warn("AssetRegistry.register failed", {
prompt_id: promptId,
error: regErr instanceof Error ? regErr.message : regErr,
});
}
}
await ensureCompletionsDir();
const filePath = join(COMPLETIONS_DIR, `${promptId}.json`);
await writeFile(filePath, JSON.stringify(notification, null, 2), "utf-8");
logger.info("Completion file written", {
prompt_id: promptId,
status: notification.status,
duration_ms: notification.duration_ms,
images: notification.outputs.reduce((n, o) => n + o.images.length, 0),
videos: notification.video_outputs.reduce((n, o) => n + o.videos.length, 0),
text_nodes: notification.text_outputs?.length ?? 0,
});
}
catch (err) {
logger.error("Failed to write completion file", {
prompt_id: promptId,
error: err instanceof Error ? err.message : err,
});
}
activeWatchers.delete(promptId);
}
function cleanup(state) {
if (state.wsCleanup) {
try {
state.wsCleanup();
}
catch {
// Ignore cleanup errors
}
state.wsCleanup = undefined;
}
if (state.pollTimer) {
clearInterval(state.pollTimer);
state.pollTimer = undefined;
}
if (state.timeoutTimer) {
clearTimeout(state.timeoutTimer);
state.timeoutTimer = undefined;
}
}
// ── Public API ─────────────────────────────────────────────────────────
export const JobWatcher = {
/**
* Start monitoring a prompt_id for completion via WS + polling dual-track.
* Optionally pass the submitted workflow so completed outputs can be
* registered with the AssetRegistry for get_image (action:"view") and for
* generate_image (action:"regenerate").
*/
watch(promptId, workflow) {
// Don't double-watch
if (activeWatchers.has(promptId)) {
logger.warn("Already watching prompt", { prompt_id: promptId });
return;
}
const state = {
promptId,
startTime: Date.now(),
completed: false,
workflow,
};
activeWatchers.set(promptId, state);
// ── WebSocket track (best-effort, local/remote only) ──
// Comfy Cloud has no WebSocket; rely on the HTTP polling track instead.
const skipWs = isCloudMode();
if (skipWs) {
logger.info("Cloud mode — skipping WS attach, polling-only", {
prompt_id: promptId,
});
}
if (!skipWs)
try {
const client = getClient();
if (!client.closed) {
state.wsCleanup = attachExecutionListeners(client, promptId, {
onComplete: () => {
handleCompletion(promptId, state, "ws").catch((err) => logger.error("WS completion handler error", {
prompt_id: promptId,
error: err instanceof Error ? err.message : err,
}));
},
onError: () => {
handleCompletion(promptId, state, "ws").catch((err) => logger.error("WS error handler error", {
prompt_id: promptId,
error: err instanceof Error ? err.message : err,
}));
},
});
logger.info("WS listeners attached", { prompt_id: promptId });
}
else {
logger.info("WS not connected, polling-only mode", {
prompt_id: promptId,
});
// Try to connect in background for future use
ensureConnected().then((client) => {
if (!state.completed) {
state.wsCleanup = attachExecutionListeners(client, promptId, {
onComplete: () => {
handleCompletion(promptId, state, "ws").catch(() => { });
},
onError: () => {
handleCompletion(promptId, state, "ws").catch(() => { });
},
});
logger.info("Late WS listeners attached", {
prompt_id: promptId,
});
}
}, () => {
// WS connection failed — polling will handle it
});
}
}
catch {
logger.info("WS setup failed, polling-only mode", {
prompt_id: promptId,
});
}
// ── Polling track ──
state.pollTimer = setInterval(() => {
if (state.completed)
return;
getHistory(promptId)
.then((history) => {
const entry = history[promptId];
if (entry?.status.completed) {
handleCompletion(promptId, state, "poll").catch((err) => logger.error("Poll completion handler error", {
prompt_id: promptId,
error: err instanceof Error ? err.message : err,
}));
}
})
.catch((err) => {
logger.warn("Poll getHistory error", {
prompt_id: promptId,
error: err instanceof Error ? err.message : err,
});
});
}, pollIntervalMs());
// ── Timeout ──
state.timeoutTimer = setTimeout(() => {
if (!state.completed) {
logger.warn("Watcher timeout, cleaning up", {
prompt_id: promptId,
elapsed_ms: Date.now() - state.startTime,
timeout_ms: watcherTimeoutMs(),
});
cleanup(state);
activeWatchers.delete(promptId);
}
}, watcherTimeoutMs());
logger.info("Job watcher started", { prompt_id: promptId });
},
/**
* Stop monitoring a prompt_id.
*/
unwatch(promptId) {
const state = activeWatchers.get(promptId);
if (state) {
cleanup(state);
activeWatchers.delete(promptId);
logger.info("Job watcher stopped", { prompt_id: promptId });
}
},
/**
* List actively watched prompt_ids.
*/
listActive() {
return [...activeWatchers.keys()];
},
/**
* Clean up old .reported completion files (called on startup).
*/
async cleanupOldFiles() {
try {
const files = await readdir(COMPLETIONS_DIR);
const now = Date.now();
let cleaned = 0;
for (const file of files) {
if (!file.endsWith(".reported"))
continue;
const filePath = join(COMPLETIONS_DIR, file);
try {
const info = await stat(filePath);
if (now - info.mtimeMs > REPORTED_MAX_AGE_MS) {
await unlink(filePath);
cleaned++;
}
}
catch {
// File may have been deleted concurrently
}
}
if (cleaned > 0) {
logger.info(`Cleaned up ${cleaned} old completion files`);
}
}
catch {
// Directory doesn't exist yet — nothing to clean
}
},
/** Exposed for testing/debugging */
get COMPLETIONS_DIR() {
return COMPLETIONS_DIR;
},
};
//# sourceMappingURL=job-watcher.js.map