@metabrain-labs/comfyui-mcp-server 1.0.0

package/src/tools/handlers/queue-prompt.ts

@@ -0,0 +1,170 @@
+import { ErrorCode, McpError } from "@modelcontextprotocol/sdk/types.js";
+import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
+import { z } from "zod";
+import i18n from "../../i18n";
+import {
+  deleteDynamicTool,
+  executeDynamicWorkflowTool,
+  getAllDynamicTools,
+  getDynamicTool,
+} from "../../services/dynamic-tool";
+import {
+  executeWorkflowTaskByPrompts,
+  ExecutionProgress,
+} from "../../services/task/execution";
+import { waitForExecutionCompletion } from "../../services/task/wait";
+import { error, errorWithDetail, ok } from "../../types/result";
+import { buildComfyViewUrls } from "../../utils/mcp-helpers";
+import { ResultToMcpResponse, withMcpErrorHandling } from "../../utils/mcp-helpers";
+import { ComfyClient } from "../../utils/ws";
+
+export function registerQueuePrompt(server: McpServer, client: ComfyClient) {
+  server.registerTool(
+    "queue_prompt",
+    {
+      title: i18n.t("tool.queue_prompt.title"),
+      description: i18n.t("tool.queue_prompt.description"),
+      inputSchema: {
+        workflowName: z
+          .string()
+          .describe(i18n.t("tool.queue_prompt.inputSchema.workflowName")),
+        isAsync: z
+          .boolean()
+          .default(false)
+          .describe(i18n.t("tool.queue_prompt.inputSchema.isAsync")),
+        params: z
+          .record(z.string(), z.any())
+          .optional()
+          .describe(i18n.t("tool.queue_prompt.inputSchema.params")),
+      },
+    },
+    withMcpErrorHandling(
+      async ({ workflowName, isAsync, params = {} }, extra) => {
+        const startTime = Date.now();
+
+        const tool = getDynamicTool(workflowName);
+        if (!tool) {
+          return ResultToMcpResponse(
+            errorWithDetail(
+              i18n.t("error.toolNotAlreadyExistError", { workflowName }),
+              `availableTools: ${getAllDynamicTools().map((t) => t.name)}`,
+            ),
+          );
+        }
+
+        if (!client.isConnected()) {
+          await client.connect();
+        }
+        if (!client.isConnected()) {
+          throw new McpError(ErrorCode.InternalError, "WebSocket NOT CONNECTED");
+        }
+
+        const execResult = await executeDynamicWorkflowTool(
+          workflowName,
+          params,
+          client,
+        );
+
+        const clientId = client.getClientId();
+        const submitResult = await executeWorkflowTaskByPrompts({
+          prompts: execResult,
+          clientId,
+        });
+
+        console.error(`[工作流已提交] Prompt ID: ${submitResult.prompt_id}`);
+
+        if (!isAsync) {
+          const progressToken = extra?._meta?.progressToken as
+            | number
+            | string
+            | undefined;
+
+          if (progressToken) {
+            console.error(`[进度通知] 已启用，Token: ${progressToken}`);
+          } else {
+            console.error(`[进度通知] 未启用（AGENT 未传入 progressToken）`);
+          }
+
+          const progressInterval = 2000;
+          let lastProgressTime = 0;
+
+          const onProgress = async (progress: ExecutionProgress) => {
+            const now = Date.now();
+            if (progressToken && now - lastProgressTime >= progressInterval) {
+              lastProgressTime = now;
+              const progressValue =
+                progress.percent !== undefined
+                  ? progress.percent / 100
+                  : progress.stage === "completed"
+                    ? 1
+                    : 0.5;
+
+              try {
+                await extra?.sendNotification?.({
+                  method: "notifications/progress",
+                  params: {
+                    progressToken,
+                    progress: progressValue,
+                    total: 1,
+                    message: progress.message,
+                  },
+                });
+                console.error(
+                  `[进度通知] Token: ${progressToken}, Progress: ${(progressValue * 100).toFixed(1)}%`,
+                );
+              } catch (err) {
+                console.error(`[进度通知] 发送失败:`, err);
+              }
+            }
+          };
+
+          const executionResult = await waitForExecutionCompletion({
+            client,
+            promptId: submitResult.prompt_id,
+            timeout: 10 * 60 * 1000,
+            onProgress,
+          });
+
+          if (!executionResult.success) {
+            deleteDynamicTool(workflowName);
+            return ResultToMcpResponse(
+              errorWithDetail(
+                i18n.t("error.workflowExecuteFail", {
+                  error: executionResult.error,
+                }),
+                `Prompt ID: ${submitResult.prompt_id}`,
+              ),
+            );
+          }
+
+          const executionTime = Date.now() - startTime;
+          return ResultToMcpResponse(
+            ok(
+              i18n.t("tool.queue_prompt.success"),
+              {
+                promptId: submitResult.prompt_id,
+                img: buildComfyViewUrls(executionResult),
+                outputs: executionResult.outputs,
+              },
+              { action: "queue_prompt" },
+              executionTime,
+            ),
+          );
+        }
+
+        const executionTime = Date.now() - startTime;
+        return ResultToMcpResponse(
+          ok(
+            i18n.t("tool.queue_prompt.success"),
+            {
+              promptId: submitResult.prompt_id,
+              description: i18n.t("tool.queue_prompt.asyncSupplement"),
+            },
+            { action: "queue_prompt" },
+            executionTime,
+          ),
+        );
+      },
+    ),
+  );
+}

package/src/tools/handlers/save-custom-workflow.ts

@@ -0,0 +1,67 @@
+import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
+import { z } from "zod";
+import i18n from "../../i18n";
+import { collectAndSaveFormatTaskFromExternal } from "../../services";
+import { error, ok } from "../../types/result";
+import {
+  ResultToMcpResponse,
+  withMcpErrorHandling,
+} from "../../utils/mcp-helpers";
+import { ComfyClient } from "../../utils/ws";
+
+export function registerSaveCustomWorkflow(
+  server: McpServer,
+  client: ComfyClient,
+) {
+  server.registerTool(
+    "save_custom_workflow",
+    {
+      title: i18n.t("tool.save_custom_workflow.title"),
+      description: i18n.t("tool.save_custom_workflow.description"),
+      inputSchema: {
+        filename: z
+          .string()
+          .describe(i18n.t("tool.save_custom_workflow.inputSchema.filename")),
+        apiJson: z
+          .record(z.string(), z.any())
+          .describe(i18n.t("tool.save_custom_workflow.inputSchema.apiJson")),
+      },
+    },
+    withMcpErrorHandling(async ({ filename, apiJson }) => {
+      const startTime = Date.now();
+
+      const fileInfo = filename.split(".");
+      if (fileInfo.length > 1 && !filename.endsWith(".json")) {
+        return ResultToMcpResponse(
+          error(
+            i18n.t("error.hasExtendName", {
+              current: "." + fileInfo[1],
+              need: ".json",
+            }),
+          ),
+        );
+      }
+
+      if (!filename.endsWith(".json")) {
+        filename += ".json";
+      }
+
+      const { filePath, promptId } = await collectAndSaveFormatTaskFromExternal(
+        filename,
+        apiJson,
+        client,
+      );
+
+      const executionTime = Date.now() - startTime;
+
+      return ResultToMcpResponse(
+        ok(
+          i18n.t("tool.save_custom_workflow.success", { filePath }),
+          { filePath, promptId },
+          { action: "save_custom_workflow" },
+          executionTime,
+        ),
+      );
+    }),
+  );
+}

package/src/tools/handlers/save-task-assets.ts

@@ -0,0 +1,82 @@
+import fs from "fs";
+import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
+import { z } from "zod";
+import { COMMON } from "../../constants";
+import i18n from "../../i18n";
+import { saveAssetsByPromptId } from "../../services";
+import { error, ok } from "../../types/result";
+import {
+  ResultToMcpResponse,
+  withMcpErrorHandling,
+} from "../../utils/mcp-helpers";
+
+export function registerSaveTaskAssets(server: McpServer) {
+  server.registerTool(
+    "save_task_assets",
+    {
+      title: i18n.t("tool.save_task_assets.title"),
+      description: i18n.t("tool.save_task_assets.description"),
+      inputSchema: {
+        promptId: z
+          .string()
+          .describe(i18n.t("tool.save_task_assets.inputSchema.promptId")),
+        destinationDir: z
+          .string()
+          .optional()
+          .describe(i18n.t("tool.save_task_assets.inputSchema.destinationDir")),
+        overwrite: z
+          .boolean()
+          .optional()
+          .default(false)
+          .describe(i18n.t("tool.save_task_assets.inputSchema.overwrite")),
+      },
+    },
+    withMcpErrorHandling(async ({ promptId, destinationDir, overwrite }) => {
+      const startTime = Date.now();
+
+      if (destinationDir) {
+        if (!fs.existsSync(destinationDir)) {
+          return ResultToMcpResponse(
+            error(i18n.t("error.dirNotExist", { destinationDir })),
+          );
+        }
+
+        const stat = fs.statSync(destinationDir);
+        if (!stat.isDirectory()) {
+          return ResultToMcpResponse(
+            error(i18n.t("error.notDir", { destinationDir })),
+          );
+        }
+
+        try {
+          fs.accessSync(destinationDir, fs.constants.W_OK);
+        } catch {
+          return ResultToMcpResponse(
+            error(i18n.t("error.notWritableDir", { destinationDir })),
+          );
+        }
+      } else {
+        if (!fs.existsSync(COMMON.ASSETS_DIR)) {
+          fs.mkdirSync(COMMON.ASSETS_DIR, { recursive: true });
+        }
+      }
+
+      const { assetsNames, filePath } = await saveAssetsByPromptId(
+        promptId,
+        overwrite,
+        destinationDir,
+      );
+
+      const executionTime = Date.now() - startTime;
+
+      return ResultToMcpResponse(
+        ok(
+          i18n.t("tool.save_task_assets.success", { assetsNames, filePath }),
+          {},
+          { action: "save_task_assets" },
+          executionTime,
+        ),
+      );
+    }),
+  );
+}

package/src/tools/handlers/system-status.ts

@@ -0,0 +1,30 @@
+import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
+import { api } from "../../api/api";
+import i18n from "../../i18n";
+import { ok } from "../../types/result";
+import { ResultToMcpResponse, withMcpErrorHandling } from "../../utils/mcp-helpers";
+
+export function registerGetSystemStatus(server: McpServer) {
+  server.registerTool(
+    "get_system_status",
+    {
+      title: i18n.t("tool.get_system_status.title"),
+      description: i18n.t("tool.get_system_status.description"),
+      inputSchema: {},
+    },
+    withMcpErrorHandling(async () => {
+      const startTime = Date.now();
+      const result = await api.getSystemStatus();
+      const executionTime = Date.now() - startTime;
+
+      return ResultToMcpResponse(
+        ok(
+          i18n.t("tool.get_system_status.success"),
+          result,
+          { action: "get_system_status" },
+          executionTime,
+        ),
+      );
+    }),
+  );
+}

package/src/tools/handlers/task-detail.ts

@@ -0,0 +1,35 @@
+import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
+import { z } from "zod";
+import i18n from "../../i18n";
+import { getTaskDetailByPromptId } from "../../services";
+import { ok } from "../../types/result";
+import { ResultToMcpResponse, withMcpErrorHandling } from "../../utils/mcp-helpers";
+
+export function registerGetTaskDetail(server: McpServer) {
+  server.registerTool(
+    "cui_get_task_detail",
+    {
+      title: i18n.t("tool.cui_get_task_detail.title"),
+      description: i18n.t("tool.cui_get_task_detail.description"),
+      inputSchema: {
+        promptId: z
+          .string()
+          .describe(i18n.t("tool.cui_get_task_detail.inputSchema.promptId")),
+      },
+    },
+    withMcpErrorHandling(async ({ promptId }) => {
+      const startTime = Date.now();
+      const result = await getTaskDetailByPromptId({ promptId });
+      const executionTime = Date.now() - startTime;
+
+      return ResultToMcpResponse(
+        ok(
+          i18n.t("tool.cui_get_task_detail.success"),
+          result,
+          { action: "cui_get_task_detail" },
+          executionTime,
+        ),
+      );
+    }),
+  );
+}

package/src/tools/handlers/tool-status-map.ts

@@ -0,0 +1,33 @@
+/**
+ * 工具与 inspection_status 的映射关系
+ * key: ToolName
+ * value: 该Tool归属的inspection_status列表，一个Tool可归属多个inspection_status
+ */
+
+import { SourceType } from "../../types/common";
+
+export const toolInspectionStatusMap: Record<string, SourceType[]> = {
+  // 基础核心工具
+  get_core_manual: ["External", "InitialInspection", "CompleteInspection"],
+
+  // 工作流目录相关
+  get_workflows_catalog: ["InitialInspection", "CompleteInspection"],
+  mount_workflow: ["External", "InitialInspection", "CompleteInspection"],
+  get_workflow_api: ["External", "InitialInspection", "CompleteInspection"],
+
+  // 任务执行相关
+  queue_prompt: ["External", "InitialInspection", "CompleteInspection"],
+  queue_custom_prompt: [], // 不可被任何inspection_status使用
+  interrupt_prompt: ["External", "InitialInspection", "CompleteInspection"],
+  get_prompt_result: ["External", "InitialInspection", "CompleteInspection"],
+  get_task_detail: ["External", "InitialInspection", "CompleteInspection"],
+
+  // 资产与存储相关
+  save_custom_workflow: ["External"], // 只有External可用
+  save_task_assets: ["External", "InitialInspection", "CompleteInspection"],
+  upload_assets: ["External", "InitialInspection", "CompleteInspection"],
+
+  // 系统工具
+  get_system_status: ["External", "InitialInspection", "CompleteInspection"],
+  list_models: ["External", "InitialInspection", "CompleteInspection"],
+};

package/src/tools/handlers/upload-assets.ts

@@ -0,0 +1,82 @@
+import fs from "fs";
+import path from "path";
+import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
+import { z } from "zod";
+import { api } from "../../api/api";
+import i18n from "../../i18n";
+import { error, ok } from "../../types/result";
+import { ResultToMcpResponse, withMcpErrorHandling } from "../../utils/mcp-helpers";
+
+export function registerUploadAssets(server: McpServer) {
+  server.registerTool(
+    "upload_assets",
+    {
+      title: i18n.t("tool.upload_assets.title"),
+      description: i18n.t("tool.upload_assets.description"),
+      inputSchema: {
+        fileSource: z
+          .string()
+          .describe(i18n.t("tool.upload_assets.inputSchema.fileSource")),
+        mimeType: z
+          .string()
+          .optional()
+          .describe(i18n.t("tool.upload_assets.inputSchema.mimeType")),
+      },
+    },
+    withMcpErrorHandling(async ({ fileSource, mimeType }) => {
+      try {
+        const startTime = Date.now();
+        let buffer;
+        let finalFileName = "uploaded_image.png";
+        const isUrl = /^https?:\/\//i.test(fileSource);
+
+        if (isUrl) {
+          const res = await fetch(fileSource);
+          if (!res.ok) {
+            return ResultToMcpResponse(
+              error(
+                i18n.t("error.downloadAssetsFail", { status: res.status }),
+              ),
+            );
+          }
+          const arrayBuffer = await res.arrayBuffer();
+          buffer = Buffer.from(arrayBuffer);
+          const urlWithoutQuery = fileSource.split("?")[0];
+          finalFileName = urlWithoutQuery.split("/").pop() || "downloaded.png";
+        } else {
+          if (!fs.existsSync(fileSource)) {
+            return ResultToMcpResponse(
+              error(i18n.t("error.fileNotExistError", { fileSource })),
+            );
+          }
+          buffer = fs.readFileSync(fileSource);
+          finalFileName = path.basename(fileSource);
+        }
+
+        const form = new FormData();
+        const blob = new Blob([buffer], {
+          type: mimeType || "application/octet-stream",
+        });
+        form.append("image", blob, finalFileName);
+
+        const response = await api.uploadImg(form);
+        const executionTime = Date.now() - startTime;
+
+        return ResultToMcpResponse(
+          ok(
+            i18n.t("tool.upload_assets.success", { finalFileName }),
+            response,
+            { action: "cui_upload_file" },
+            executionTime,
+          ),
+        );
+      } catch (err: any) {
+        return ResultToMcpResponse(
+          error(
+            i18n.t("error.uploadFail", { message: err.message }),
+          ),
+        );
+      }
+    }),
+  );
+}

package/src/tools/handlers/workflow-api.ts

@@ -0,0 +1,46 @@
+import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
+import { z } from "zod";
+import { api } from "../../api/api";
+import i18n from "../../i18n";
+import { error, ok } from "../../types/result";
+import {
+  ResultToMcpResponse,
+  withMcpErrorHandling,
+} from "../../utils/mcp-helpers";
+
+export function registerGetWorkflowAPI(server: McpServer) {
+  server.registerTool(
+    "get_workflow_API",
+    {
+      title: i18n.t("tool.get_workflow_API.title"),
+      description: i18n.t("tool.get_workflow_API.description"),
+      inputSchema: {
+        workflowPath: z
+          .string()
+          .describe(i18n.t("tool.get_workflow_API.inputSchema.workflowPath")),
+      },
+    },
+    withMcpErrorHandling(async ({ workflowPath }) => {
+      if (!workflowPath.endsWith(".json")) {
+        return ResultToMcpResponse(
+          error(i18n.t("error.getWorkflowFormatError", { workflowPath })),
+        );
+      }
+
+      const startTime = Date.now();
+      const result = await api.getDetailUserData(
+        encodeURIComponent(workflowPath),
+      );
+      const executionTime = Date.now() - startTime;
+
+      return ResultToMcpResponse(
+        ok(
+          i18n.t("tool.get_workflow_API.success"),
+          result,
+          { action: "get_workflow_API" },
+          executionTime,
+        ),
+      );
+    }),
+  );
+}

package/src/tools/handlers/workflows-catalog.ts

@@ -0,0 +1,79 @@
+import { readFile } from "fs/promises";
+import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
+import { z } from "zod";
+import { COMMON } from "../../constants";
+import i18n from "../../i18n";
+import {
+  collectAndSaveFormatTask,
+  collectAndSaveFormatTaskFromWorkflows,
+  collectExternalWorkflowsFromDirectory,
+  saveWorkflow,
+} from "../../services";
+import { ComfyClient } from "../../utils/ws";
+import { WorkflowConverter } from "../../utils/workflow-converter";
+import {
+  ResultToMcpStringResponse,
+  withMcpErrorHandling,
+} from "../../utils/mcp-helpers";
+
+export function registerGetWorkflowsCatalog(
+  server: McpServer,
+  client: ComfyClient,
+  converter: WorkflowConverter,
+) {
+  server.registerTool(
+    "get_workflows_catalog",
+    {
+      title: i18n.t("tool.get_workflows_catalog.title"),
+      description: i18n.t("tool.get_workflows_catalog.description"),
+      inputSchema: {
+        maxItems: z
+          .number()
+          .min(1)
+          .max(10)
+          .optional()
+          .default(10)
+          .describe(i18n.t("tool.get_workflows_catalog.inputSchema.maxItems")),
+        useExistingCatalog: z
+          .boolean()
+          .optional()
+          .default(false)
+          .describe(i18n.t("tool.get_workflows_catalog.inputSchema.useExistingCatalog")),
+      },
+    },
+    withMcpErrorHandling(async ({ maxItems, useExistingCatalog }) => {
+      // 如果 useExistingCatalog 为 true，直接返回已有的 workflow.json
+      if (useExistingCatalog) {
+        const content = await readFile(COMMON.WORKFLOW_PATH, "utf-8");
+        return ResultToMcpStringResponse(content);
+      }
+
+      let hasMore: boolean = true;
+
+      // 1. 收集历史任务（CompleteInspection）
+      for (let i = 0; hasMore; i++) {
+        hasMore = await collectAndSaveFormatTask({
+          maxItems,
+          offset: i * maxItems,
+          append: i > 0,
+        });
+      }
+
+      // 2. 收集用户工作流（InitialInspection）
+      await collectAndSaveFormatTaskFromWorkflows(client, converter);
+
+      // 3. 收集本地 workflow 目录的 External 类型工作流
+      // 由于去重逻辑优先保留非 External 类型，External 只作为补充
+      const externalResult = await collectExternalWorkflowsFromDirectory();
+      const externalWorkflows = externalResult || [];
+
+      // 只有当有 External 工作流时才保存
+      if (externalWorkflows.length > 0) {
+        await saveWorkflow(externalWorkflows, { append: true });
+      }
+
+      const content = await readFile(COMMON.WORKFLOW_PATH, "utf-8");
+      return ResultToMcpStringResponse(content);
+    }),
+  );
+}