@metabrain-labs/comfyui-mcp-server 1.0.0

package/src/services/storage/workflow-storage.ts

@@ -0,0 +1,192 @@
+import { ErrorCode, McpError } from "@modelcontextprotocol/sdk/types.js";
+import { mkdir, readFile, writeFile } from "fs/promises";
+import path from "path";
+import { COMMON } from "../../constants";
+import { CollectFormatTaskWorkflow, sourcePriority } from "../../types/common";
+
+export interface SaveWorkflowOptions {
+  dir?: string;
+  fileName?: string;
+  append?: boolean;
+}
+
+/**
+ * 保存工作流任务信息到本地
+ */
+export async function saveWorkflow(
+  data: CollectFormatTaskWorkflow[],
+  options: SaveWorkflowOptions = {},
+): Promise<{ filePath: string; itemsCollected: number }> {
+  try {
+    const {
+      dir = COMMON.WORKFLOW_DIR,
+      fileName = COMMON.WORKFLOW_FILE,
+      append = true,
+    } = options;
+
+    const itemsCollected = Object.keys(data).length;
+    await mkdir(dir, { recursive: true });
+
+    const filePath = COMMON.WORKFLOW_PATH;
+    let finalData: CollectFormatTaskWorkflow[] = [];
+
+    if (append) {
+      try {
+        const existingContent = await readFile(filePath, "utf-8");
+        const existingData = JSON.parse(existingContent);
+        if (Array.isArray(existingData)) {
+          finalData = existingData.flat(Infinity);
+        } else {
+          console.error("现有文件格式不是数组，将被覆盖");
+        }
+      } catch (error: any) {
+        if (error.code !== "ENOENT") {
+          throw new McpError(
+            ErrorCode.InternalError,
+            `读取现有文件失败: ${error.message}`,
+          );
+        }
+      }
+    }
+
+    const newData = data.flat(Infinity);
+    finalData = [...finalData, ...newData];
+    const uniqueData = deduplicateWorkflows(finalData);
+
+    const filteredData = uniqueData.filter((item) => {
+      if (typeof item === "object" && item !== null && !Array.isArray(item)) {
+        return Object.keys(item).length > 0;
+      }
+      return true;
+    });
+
+    await writeFile(filePath, JSON.stringify(filteredData, null, 2), "utf-8");
+    return { filePath, itemsCollected };
+  } catch (error) {
+    if (error instanceof McpError) {
+      throw error;
+    }
+    throw new McpError(
+      ErrorCode.InternalError,
+      `保存工作流失败: ${error instanceof Error ? error.message : String(error)}`,
+    );
+  }
+}
+
+/**
+ * 保存自定义工作流到本地
+ */
+export async function saveCustomWorkflow(
+  data: CollectFormatTaskWorkflow[],
+  options: SaveWorkflowOptions = {},
+): Promise<string> {
+  try {
+    const dir = COMMON.WORKFLOW_DIR;
+    const fileName = options.fileName!;
+
+    await mkdir(dir, { recursive: true });
+
+    const filePath = path.join(dir, fileName);
+
+    await writeFile(filePath, JSON.stringify(data, null, 2), "utf-8");
+    return filePath;
+  } catch (error) {
+    if (error instanceof McpError) {
+      throw error;
+    }
+    throw new McpError(
+      ErrorCode.InternalError,
+      `保存自定义工作流失败: ${error instanceof Error ? error.message : String(error)}`,
+    );
+  }
+}
+
+/**
+ * 根据 name 去重，保留最新的项
+ * 规则：
+ * 1. External 与 非External 重名时，优先保留 非External
+ * 2. 同类型时，保留 last_updated 最新的
+ * 3. InitialInspection 与 CompleteInspection 重名时，根据 userdata_modified 判断
+ */
+function deduplicateWorkflows(
+  workflows: CollectFormatTaskWorkflow[],
+): CollectFormatTaskWorkflow[] {
+  const map = new Map<string, CollectFormatTaskWorkflow>();
+
+  for (const workflow of workflows) {
+    const existing = map.get(workflow.name);
+    if (!existing) {
+      map.set(workflow.name, workflow);
+      continue;
+    }
+    if (shouldReplaceWorkflow(existing, workflow)) {
+      map.set(workflow.name, workflow);
+    }
+  }
+
+  return Array.from(map.values());
+}
+
+/**
+ * 判断是否应该替换工作流
+ * 优先级规则：
+ * 1. External 与 非External 重名时，优先保留 非External（External 被抛弃）
+ * 2. InitialInspection 与 CompleteInspection 重名时，根据 userdata_modified 判断
+ * 3. 同类型时，保留 last_updated 最新的
+ */
+function shouldReplaceWorkflow(
+  existing: CollectFormatTaskWorkflow,
+  candidate: CollectFormatTaskWorkflow,
+): boolean {
+  // 规则1: External 与 非External 重名时，优先保留 非External
+  const existingIsExternal = existing.inspection_status === "External";
+  const candidateIsExternal = candidate.inspection_status === "External";
+  
+  if (existingIsExternal && !candidateIsExternal) {
+    // existing 是 External，candidate 不是，用 candidate 替换 existing
+    return true;
+  }
+  
+  if (!existingIsExternal && candidateIsExternal) {
+    // existing 不是 External，candidate 是，保留 existing
+    return false;
+  }
+
+  // 规则2: 同类型时，保留 last_updated 最新的
+  if (existing.inspection_status === candidate.inspection_status) {
+    return candidate.last_updated > existing.last_updated;
+  }
+
+  // 规则3: InitialInspection 与 CompleteInspection 重名时，根据 userdata_modified 判断
+  const isInitialVsComplete =
+    (existing.inspection_status === "InitialInspection" &&
+      candidate.inspection_status === "CompleteInspection") ||
+    (existing.inspection_status === "CompleteInspection" &&
+      candidate.inspection_status === "InitialInspection");
+
+  if (isInitialVsComplete) {
+    const initial =
+      existing.inspection_status === "InitialInspection" ? existing : candidate;
+    const complete =
+      existing.inspection_status === "CompleteInspection"
+        ? existing
+        : candidate;
+
+    const initialModifiedTime = initial.userdata_modified ?? 0;
+    if (initialModifiedTime > complete.last_updated) {
+      return candidate === initial;
+    } else {
+      return candidate === complete;
+    }
+  }
+
+  // 其他情况（理论上不会走到这里，因为目前只有三种类型），按优先级处理
+  const existingPriority = sourcePriority[existing.inspection_status] ?? -1;
+  const candidatePriority = sourcePriority[candidate.inspection_status] ?? -1;
+
+  if (candidatePriority !== existingPriority) {
+    return candidatePriority > existingPriority;
+  }
+
+  return candidate.last_updated > existing.last_updated;
+}

package/src/services/task/execution.ts

@@ -0,0 +1,69 @@
+import { ErrorCode, McpError } from "@modelcontextprotocol/sdk/types.js";
+import { api } from "../../api/api";
+import {
+  ExecutePromptRequest,
+  ExecutePromptResult,
+} from "../../types/execute";
+import { ComfyPromptConfig } from "../../types/task";
+
+export interface ExecuteTaskOptions {
+  prompts: ComfyPromptConfig;
+  clientId: string;
+}
+
+/** 执行进度信息 */
+export interface ExecutionProgress {
+  stage: "starting" | "executing" | "progress" | "completed" | "error";
+  message: string;
+  percent?: number;
+  nodeId?: string;
+  current?: number;
+  max?: number;
+  elapsedTime?: number;
+}
+
+/**
+ * 执行工作流任务
+ */
+export async function executeWorkflowTaskByPrompts(
+  options: ExecuteTaskOptions,
+): Promise<ExecutePromptResult> {
+  const { prompts, clientId } = options;
+
+  if (!clientId) {
+    throw new McpError(
+      ErrorCode.InternalError,
+      `不存在WS客户端ID，请检查是否正确连接WS服务器`,
+    );
+  }
+
+  const data: ExecutePromptRequest = {
+    client_id: clientId,
+    prompt: prompts,
+  };
+
+  return await api.prompt(data);
+}
+
+/**
+ * 执行自定义工作流任务
+ */
+export async function executeCustomWorkflowTaskByPrompts(
+  options: ExecuteTaskOptions,
+): Promise<ExecutePromptResult> {
+  const { prompts, clientId } = options;
+
+  if (!clientId) {
+    throw new McpError(
+      ErrorCode.InternalError,
+      `不存在WS客户端ID，请检查是否正确连接WS服务器`,
+    );
+  }
+
+  const data: ExecutePromptRequest = {
+    client_id: clientId,
+    prompt: prompts,
+  };
+
+  return await api.prompt(data);
+}

package/src/services/task/fetch.ts

@@ -0,0 +1,131 @@
+import { ErrorCode, McpError } from "@modelcontextprotocol/sdk/types.js";
+import { api } from "../../api/api";
+import {
+  ComfyNodeOutput,
+  ComfyPromptConfig,
+  ComfyTaskItem,
+  ComfyTaskResponse,
+} from "../../types/task";
+
+export interface FetchTasksOptions {
+  maxItems?: number;
+  offset?: number;
+}
+
+export interface FetchTaskOptions {
+  promptId: string;
+}
+
+/**
+ * 获取历史任务（支持分页）
+ */
+export async function fetchHistoryTasks(
+  options: FetchTasksOptions,
+): Promise<{ successTasks: ComfyTaskResponse; total: number; fail: number }> {
+  const { maxItems, offset } = options;
+  const res = await api.pageHistoryTasks(maxItems, offset);
+  const total = Object.entries(res).length;
+
+  const successTasks = Object.fromEntries(
+    Object.entries(res).filter(
+      ([uuid, item]) => item.status.status_str === "success",
+    ),
+  ) as ComfyTaskResponse;
+
+  const fail = total - Object.entries(successTasks).length;
+  return { successTasks, total, fail };
+}
+
+/**
+ * 获取指定promptIds对应的历史任务信息
+ */
+export async function fetchUserWorkflow(
+  availableWorkflow: string[],
+): Promise<ComfyTaskResponse> {
+  let historyTasks: [string, ComfyTaskItem][] = [];
+
+  for (const promptId of availableWorkflow) {
+    const res = await api.getDetailHistoryTasks(promptId);
+    historyTasks = historyTasks.concat(Object.entries(res));
+  }
+
+  const successTasks = Object.fromEntries(historyTasks) as ComfyTaskResponse;
+
+  if (successTasks === null) {
+    throw new McpError(
+      ErrorCode.InternalError,
+      "Invalid detail tasks response",
+    );
+  }
+
+  return successTasks;
+}
+
+/**
+ * 根据 promptId 获取任务详情
+ */
+export async function fetchTaskByPromptId(
+  options: FetchTaskOptions,
+): Promise<ComfyPromptConfig> {
+  const { promptId } = options;
+  const res = await api.getDetailHistoryTasks(promptId);
+
+  const successTasks = Object.fromEntries(
+    Object.entries(res).filter(([uuid, item]) => {
+      if (item.status.status_str === "success") {
+        return true;
+      } else {
+        const isExecutionInterrupted = item.status.messages.filter((item) => {
+          return item[0] === "execution_interrupted";
+        });
+        return isExecutionInterrupted.length > 0;
+      }
+    }),
+  ) as ComfyTaskResponse;
+
+  if (successTasks === null) {
+    throw new McpError(
+      ErrorCode.InternalError,
+      "Invalid detail tasks response",
+    );
+  }
+
+  return Object.values(successTasks)[0].prompt[2];
+}
+
+/**
+ * 获取成功历史任务的资产信息
+ */
+export async function fetchAssetsByPromptId(
+  promptId: string,
+): Promise<ComfyNodeOutput[]> {
+  const res = await api.getDetailHistoryTasks(promptId);
+
+  const successTasks = Object.fromEntries(
+    Object.entries(res).filter(([uuid, item]) => {
+      if (item.status.status_str === "success") {
+        return true;
+      } else {
+        const isExecutionInterrupted = item.status.messages.filter((item) => {
+          return item[0] === "execution_interrupted";
+        });
+        return isExecutionInterrupted.length > 0;
+      }
+    }),
+  ) as ComfyTaskResponse;
+
+  if (successTasks === null) {
+    throw new McpError(
+      ErrorCode.InternalError,
+      "Invalid detail tasks response",
+    );
+  }
+
+  const outputs = Object.values(successTasks)[0].outputs;
+
+  if (Object.keys(outputs).length === 0) {
+    throw new McpError(ErrorCode.InternalError, "Invalid outputs response");
+  }
+
+  return Object.values(outputs);
+}

package/src/services/task/index.ts

@@ -0,0 +1,3 @@
+export * from "./fetch";
+export * from "./execution";
+export * from "./wait";

package/src/services/task/wait.ts

@@ -0,0 +1,294 @@
+import { ExecutionResult } from "../../types/execute";
+import { ComfyClient } from "../../utils/ws";
+import { ExecutionProgress } from "./execution";
+
+export interface WaitForExecutionOptions {
+  client: ComfyClient;
+  promptId: string;
+  timeout?: number;
+  onProgress?: (progress: ExecutionProgress) => void;
+}
+
+/**
+ * 等待工作流执行完成
+ */
+export async function waitForExecutionCompletion(
+  options: WaitForExecutionOptions,
+): Promise<ExecutionResult> {
+  const { client, promptId, timeout = 10 * 60 * 1000, onProgress } = options;
+  const startTime = Date.now();
+
+  console.error(`[等待执行] 开始监听 Prompt ID: ${promptId}`);
+
+  const reportProgress = (progress: ExecutionProgress) => {
+    if (onProgress) {
+      onProgress({
+        ...progress,
+        elapsedTime: Date.now() - startTime,
+      });
+    }
+  };
+
+  return new Promise((resolve, reject) => {
+    let isCompleted = false;
+    const outputs: Record<string, any> = {};
+
+    const timeoutTimer = setTimeout(() => {
+      if (!isCompleted) {
+        isCompleted = true;
+        cleanup();
+        resolve({
+          success: false,
+          promptId,
+          error: `工作流执行超时（${timeout / 1000}秒）`,
+          executionTime: Date.now() - startTime,
+        });
+      }
+    }, timeout);
+
+    const originalHooks = {
+      onExecutionStart: client.hook.onExecutionStart.bind(client.hook),
+      onNodeExecuting: client.hook.onNodeExecuting.bind(client.hook),
+      onProgress: client.hook.onProgress.bind(client.hook),
+      onExecuted: client.hook.onExecuted.bind(client.hook),
+      onExecutionSuccess: client.hook.onExecutionSuccess.bind(client.hook),
+      onExecutionError: client.hook.onExecutionError.bind(client.hook),
+      onExecutionInterrupted: client.hook.onExecutionInterrupted.bind(
+        client.hook,
+      ),
+    };
+
+    const cleanup = () => {
+      clearTimeout(timeoutTimer);
+      client.hook.onExecutionStart = originalHooks.onExecutionStart;
+      client.hook.onNodeExecuting = originalHooks.onNodeExecuting;
+      client.hook.onProgress = originalHooks.onProgress;
+      client.hook.onExecuted = originalHooks.onExecuted;
+      client.hook.onExecutionSuccess = originalHooks.onExecutionSuccess;
+      client.hook.onExecutionError = originalHooks.onExecutionError;
+      client.hook.onExecutionInterrupted = originalHooks.onExecutionInterrupted;
+    };
+
+    client.hook.onExecutionStart = (data) => {
+      console.error(`[onExecutionStart] 收到消息:`, data);
+      originalHooks.onExecutionStart(data);
+      if (data.prompt_id === promptId) {
+        console.error(`[执行开始] Prompt ID: ${promptId}`);
+        reportProgress({
+          stage: "starting",
+          message: "工作流开始执行",
+          percent: 0,
+        });
+      }
+    };
+
+    client.hook.onNodeExecuting = (data) => {
+      console.error(`[onNodeExecuting] 收到消息:`, data);
+      originalHooks.onNodeExecuting(data);
+      if (data.prompt_id === promptId && data.node) {
+        console.error(`[执行节点] Node: ${data.node}`);
+        reportProgress({
+          stage: "executing",
+          message: `正在执行节点: ${data.node}`,
+          nodeId: data.node,
+        });
+      }
+    };
+
+    client.hook.onProgress = (data) => {
+      originalHooks.onProgress(data);
+      if (data.prompt_id === promptId) {
+        const percent = Math.round((data.value / data.max) * 100);
+        console.error(`[进度更新] ${data.value}/${data.max} (${percent}%)`);
+        reportProgress({
+          stage: "progress",
+          message: `节点 ${data.node} 执行中: ${data.value}/${data.max}`,
+          percent,
+          nodeId: data.node,
+          current: data.value,
+          max: data.max,
+        });
+      }
+    };
+
+    client.hook.onExecuted = (data) => {
+      console.error(`[onExecuted] 收到消息:`, data);
+      originalHooks.onExecuted(data);
+      if (data.prompt_id === promptId) {
+        console.error(`[节点完成] Node: ${data.node}`);
+        if (data.output) {
+          outputs[data.node] = data.output;
+        }
+      }
+    };
+
+    client.hook.onExecutionSuccess = (data) => {
+      console.error(`[onExecutionSuccess] 收到消息:`, data);
+      originalHooks.onExecutionSuccess(data);
+      if (data.prompt_id === promptId && !isCompleted) {
+        isCompleted = true;
+        cleanup();
+        console.error(`[执行成功] Prompt ID: ${promptId}`);
+        reportProgress({
+          stage: "completed",
+          message: "工作流执行完成",
+          percent: 100,
+        });
+        resolve({
+          success: true,
+          promptId,
+          outputs,
+          executionTime: Date.now() - startTime,
+        });
+      }
+    };
+
+    client.hook.onExecutionError = (data) => {
+      console.error(`[onExecutionError] 收到消息:`, data);
+      originalHooks.onExecutionError(data);
+      if (data.prompt_id === promptId && !isCompleted) {
+        isCompleted = true;
+        cleanup();
+        console.error(`[执行错误] Prompt ID: ${promptId}`, data);
+        const errorMsg =
+          data.exception_message || data.error || JSON.stringify(data);
+        reportProgress({
+          stage: "error",
+          message: `执行出错: ${errorMsg}`,
+        });
+        resolve({
+          success: false,
+          promptId,
+          error: errorMsg,
+          executionTime: Date.now() - startTime,
+        });
+      }
+    };
+
+    client.hook.onExecutionInterrupted = (data) => {
+      console.error(`[onExecutionInterrupted] 收到消息:`, data);
+      originalHooks.onExecutionInterrupted(data);
+      if (data.prompt_id === promptId && !isCompleted) {
+        isCompleted = true;
+        cleanup();
+        console.error(`[执行中断] Prompt ID: ${promptId}`);
+        resolve({
+          success: false,
+          promptId,
+          error: "工作流执行被中断",
+          executionTime: Date.now() - startTime,
+        });
+      }
+    };
+  });
+}
+
+/**
+ * 等待工作流执行中断
+ */
+export async function waitForExecutionInterrupt(
+  options: WaitForExecutionOptions,
+): Promise<ExecutionResult> {
+  const { client, promptId, timeout = 20 * 1000 } = options;
+  const startTime = Date.now();
+
+  console.error(`[等待中断信号] 开始监听 Prompt ID: ${promptId}`);
+
+  return new Promise((resolve, reject) => {
+    let isCompleted = false;
+
+    const timeoutTimer = setTimeout(() => {
+      if (!isCompleted) {
+        isCompleted = true;
+        cleanup();
+        resolve({
+          success: false,
+          promptId,
+          error: `工作流等待中断信号超时（${timeout / 1000}秒）`,
+          executionTime: Date.now() - startTime,
+        });
+      }
+    }, timeout);
+
+    const originalHooks = {
+      onExecutionInterrupted: client.hook.onExecutionInterrupted.bind(
+        client.hook,
+      ),
+    };
+
+    const cleanup = () => {
+      clearTimeout(timeoutTimer);
+      client.hook.onExecutionInterrupted = originalHooks.onExecutionInterrupted;
+    };
+
+    client.hook.onExecutionInterrupted = (data) => {
+      console.error(`[onExecutionInterrupted] 收到消息:`, data);
+      originalHooks.onExecutionInterrupted(data);
+      if (data.prompt_id === promptId && !isCompleted) {
+        isCompleted = true;
+        cleanup();
+        console.error(`[执行中断] Prompt ID: ${promptId}`);
+        resolve({
+          success: true,
+          promptId,
+          error: "工作流已收到中断信号",
+          executionTime: Date.now() - startTime,
+        });
+      }
+    };
+  });
+}
+
+/**
+ * 等待工作流开始执行
+ */
+export async function waitForExecutionStart(
+  options: WaitForExecutionOptions,
+): Promise<ExecutionResult> {
+  const { client, promptId, timeout = 20 * 1000 } = options;
+  const startTime = Date.now();
+
+  console.error(`[等待开始信号] 开始监听 Prompt ID: ${promptId}`);
+
+  return new Promise((resolve, reject) => {
+    let isCompleted = false;
+
+    const timeoutTimer = setTimeout(() => {
+      if (!isCompleted) {
+        isCompleted = true;
+        cleanup();
+        resolve({
+          success: false,
+          promptId,
+          error: `工作流等待开始信号超时（${timeout / 1000}秒）`,
+          executionTime: Date.now() - startTime,
+        });
+      }
+    }, timeout);
+
+    const originalHooks = {
+      onExecutionStart: client.hook.onExecutionStart.bind(client.hook),
+    };
+
+    const cleanup = () => {
+      clearTimeout(timeoutTimer);
+      client.hook.onExecutionStart = originalHooks.onExecutionStart;
+    };
+
+    client.hook.onExecutionStart = (data) => {
+      console.error(`[onExecutionStart] 收到消息:`, data);
+      originalHooks.onExecutionStart(data);
+      if (data.prompt_id === promptId && !isCompleted) {
+        isCompleted = true;
+        cleanup();
+        console.error(`[执行开始] Prompt ID: ${promptId}`);
+        resolve({
+          success: true,
+          promptId,
+          error: "工作流已收到开始信号",
+          executionTime: Date.now() - startTime,
+        });
+      }
+    };
+  });
+}