From 9e10e55633eba17fda53897a3c1dfaaf5c93a476 Mon Sep 17 00:00:00 2001
From: Chenlei Hu <huchenlei@proton.me>
Date: Wed, 23 Apr 2025 10:47:08 -0400
Subject: [PATCH] [Schema] Add api_node flag to node def schema (#3586)

---
 src/schemas/nodeDef/nodeDefSchemaV2.ts |  3 ++-
 src/schemas/nodeDefSchema.ts           |  8 +++++++-
 src/stores/nodeDefStore.ts             |  2 ++
 tests-ui/tests/nodeDef.test.ts         | 19 +++++++++++++++++++
 4 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/src/schemas/nodeDef/nodeDefSchemaV2.ts b/src/schemas/nodeDef/nodeDefSchemaV2.ts
index 18a9f392e..077361d88 100644
--- a/src/schemas/nodeDef/nodeDefSchemaV2.ts
+++ b/src/schemas/nodeDef/nodeDefSchemaV2.ts
@@ -77,7 +77,8 @@ export const zComfyNodeDef = z.object({
   output_node: z.boolean(),
   python_module: z.string(),
   deprecated: z.boolean().optional(),
-  experimental: z.boolean().optional()
+  experimental: z.boolean().optional(),
+  api_node: z.boolean().optional()
 })
 
 // Export types
diff --git a/src/schemas/nodeDefSchema.ts b/src/schemas/nodeDefSchema.ts
index b9482172e..4be1765ef 100644
--- a/src/schemas/nodeDefSchema.ts
+++ b/src/schemas/nodeDefSchema.ts
@@ -222,7 +222,13 @@ export const zComfyNodeDef = z.object({
   output_node: z.boolean(),
   python_module: z.string(),
   deprecated: z.boolean().optional(),
-  experimental: z.boolean().optional()
+  experimental: z.boolean().optional(),
+  /**
+   * Whether the node is an API node. Running API nodes requires login to
+   * Comfy Org account.
+   * https://www.comfy.org/faq
+   */
+  api_node: z.boolean().optional()
 })
 
 // `/object_info`
diff --git a/src/stores/nodeDefStore.ts b/src/stores/nodeDefStore.ts
index 3541476a8..803159005 100644
--- a/src/stores/nodeDefStore.ts
+++ b/src/stores/nodeDefStore.ts
@@ -41,6 +41,7 @@ export class ComfyNodeDefImpl
   readonly deprecated: boolean
   readonly experimental: boolean
   readonly output_node: boolean
+  readonly api_node: boolean
   /**
    * @deprecated Use `inputs` instead
    */
@@ -121,6 +122,7 @@ export class ComfyNodeDefImpl
     this.experimental =
       obj.experimental ?? obj.category.startsWith('_for_testing')
     this.output_node = obj.output_node
+    this.api_node = !!obj.api_node
     this.input = obj.input ?? {}
     this.output = obj.output ?? []
     this.output_is_list = obj.output_is_list
diff --git a/tests-ui/tests/nodeDef.test.ts b/tests-ui/tests/nodeDef.test.ts
index 6a2fc1842..7250918f1 100644
--- a/tests-ui/tests/nodeDef.test.ts
+++ b/tests-ui/tests/nodeDef.test.ts
@@ -518,4 +518,23 @@ describe('ComfyNodeDefImpl', () => {
     expect(result.inputs['booleanInput']).toBeDefined()
     expect(result.inputs['floatInput']).toBeDefined()
   })
+
+  it.each([
+    { api_node: true, expected: true },
+    { api_node: false, expected: false },
+    { api_node: undefined, expected: false }
+  ] as { api_node: boolean | undefined; expected: boolean }[])(
+    'should handle api_node field: $api_node',
+    ({ api_node, expected }) => {
+      const result = new ComfyNodeDefImpl({
+        name: 'ApiNode',
+        display_name: 'API Node',
+        category: 'Test',
+        python_module: 'test_module',
+        description: 'A node with API',
+        api_node
+      } as ComfyNodeDefV1)
+      expect(result.api_node).toBe(expected)
+    }
+  )
 })