From f9157ee05f3f4a4cb0097315bd74d7533e1345ab Mon Sep 17 00:00:00 2001
From: bymyself <cbyrne@comfy.org>
Date: Thu, 27 Feb 2025 09:15:31 -0700
Subject: [PATCH] Update workflow schema to include node pack ID and version
 (#2751)

---
 src/types/comfyWorkflow.ts           | 52 ++++++++++++++++++++-
 tests-ui/tests/comfyWorkflow.test.ts | 70 ++++++++++++++++++++++++++++
 2 files changed, 121 insertions(+), 1 deletion(-)

diff --git a/src/types/comfyWorkflow.ts b/src/types/comfyWorkflow.ts
index 492ae6891..c091800b4 100644
--- a/src/types/comfyWorkflow.ts
+++ b/src/types/comfyWorkflow.ts
@@ -114,9 +114,59 @@ const zFlags = z
   })
   .passthrough()
 
+const repoLikeIdPattern = /^[a-zA-Z0-9](?:[a-zA-Z0-9._-]*[a-zA-Z0-9])?$/
+const githubUsernamePattern = /^(?!-)(?!.*--)[a-zA-Z0-9-]+(?<!-)$/
+const gitHashPattern = /^[0-9a-f]{4,40}$/i
+const semverPattern =
+  /^(0|[1-9]\d*)\.(0|[1-9]\d*)\.(0|[1-9]\d*)(?:-([\da-z-]+(?:\.[\da-z-]+)*))?(?:\+([\da-z-]+(?:\.[\da-z-]+)*))?$/
+
+// Shared schema for Comfy Node Registry IDs and GitHub repo names
+const zRepoLikeId = z
+  .string()
+  .min(1)
+  .max(100)
+  .regex(repoLikeIdPattern, {
+    message: "ID can only contain ASCII letters, digits, '_', '-', and '.'"
+  })
+  .refine((id) => !/^[_\-.]|[_\-.]$/.test(id), {
+    message: "ID must not start or end with '_', '-', or '.'"
+  })
+
+const zCnrId = zRepoLikeId
+const zGithubRepoName = zRepoLikeId
+
+// GitHub username/organization schema
+const zGithubUsername = z
+  .string()
+  .min(1)
+  .max(39)
+  .regex(githubUsernamePattern, 'Invalid GitHub username/org')
+
+// Auxiliary ID identifies node packs not installed via the Comfy Node Registry
+const zAuxId = z
+  .string()
+  .regex(/^[^/]+\/[^/]+$/, "Invalid format. Must be 'github-user/repo-name'")
+  .transform((id) => id.split('/'))
+  .refine(
+    ([username, repo]) =>
+      zGithubUsername.safeParse(username).success &&
+      zGithubRepoName.safeParse(repo).success,
+    "Invalid aux_id: Must be valid 'github-username/github-repo-name'"
+  )
+  .transform(([username, repo]) => `${username}/${repo}`)
+
+const zSemVer = z
+  .string()
+  .regex(semverPattern, 'Invalid semantic version (x.y.z)')
+const zGitHash = z.string().regex(gitHashPattern, 'Invalid Git commit hash')
+const zVersion = z.union([zSemVer, zGitHash])
+
 const zProperties = z
   .object({
-    ['Node name for S&R']: z.string().optional()
+    ['Node name for S&R']: z.string().optional(),
+    cnr_id: zCnrId.optional(),
+    aux_id: zAuxId.optional(),
+    ver: zVersion.optional()
   })
   .passthrough()
 
diff --git a/tests-ui/tests/comfyWorkflow.test.ts b/tests-ui/tests/comfyWorkflow.test.ts
index 32db20d36..86176b60a 100644
--- a/tests-ui/tests/comfyWorkflow.test.ts
+++ b/tests-ui/tests/comfyWorkflow.test.ts
@@ -124,4 +124,74 @@ describe('parseComfyWorkflow', () => {
     ]
     expect(await validateComfyWorkflow(workflow)).not.toBeNull()
   })
+
+  describe('workflow.nodes.properties.aux_id', () => {
+    const validAuxIds = [
+      'valid/valid',
+      'valid-username-with-dash/valid_github-repo-name-with-underscore'
+    ]
+    it.each(validAuxIds)('valid aux_id: %s', async (aux_id) => {
+      const workflow = JSON.parse(JSON.stringify(defaultGraph))
+      workflow.nodes[0].properties.aux_id = aux_id
+      expect(await validateComfyWorkflow(workflow)).not.toBeNull()
+    })
+    const invalidAuxIds = [
+      'invalid spaces in username/repo',
+      'invalid-chars-name-$/repo',
+      'github-name/invalid spaces in repo',
+      'not-both-names-with-slash'
+    ]
+    it.each(invalidAuxIds)('invalid aux_id: %s', async (aux_id) => {
+      const workflow = JSON.parse(JSON.stringify(defaultGraph))
+      workflow.nodes[0].properties.aux_id = aux_id
+      expect(await validateComfyWorkflow(workflow)).toBeNull()
+    })
+  })
+
+  describe('workflow.nodes.properties.cnr_id', () => {
+    const validCnrIds = ['valid', 'valid-with-dash', 'valid_with_underscores']
+    it.each(validCnrIds)('valid cnr_id: %s', async (cnr_id) => {
+      const workflow = JSON.parse(JSON.stringify(defaultGraph))
+      workflow.nodes[0].properties.cnr_id = cnr_id
+      expect(await validateComfyWorkflow(workflow)).not.toBeNull()
+    })
+
+    const invalidCnrIds = ['invalid cnr-id', 'invalid^cnr-id', 'invalid cnr id']
+    it.each(invalidCnrIds)('invalid cnr_id: %s', async (cnr_id) => {
+      const workflow = JSON.parse(JSON.stringify(defaultGraph))
+      workflow.nodes[0].properties.cnr_id = cnr_id
+      expect(await validateComfyWorkflow(workflow)).toBeNull()
+    })
+  })
+
+  describe('workflow.nodes.properties.ver', () => {
+    const validVersionStrings = [
+      // Semver
+      '0.1.0',
+      '0.1.0-alpha',
+      '0.1.0-alpha.1',
+      '1.3.321',
+      // Git hash
+      '080e6d4af809a46852d1c4b7ed85f06e8a3a72be'
+    ]
+    it.each(validVersionStrings)('valid version: %s', async (ver) => {
+      const workflow = JSON.parse(JSON.stringify(defaultGraph))
+      workflow.nodes[0].properties.ver = ver
+      expect(await validateComfyWorkflow(workflow)).not.toBeNull()
+    })
+
+    const invalidVersionStrings = [
+      // Semver
+      '0.1-alpha',
+      '0. 1.0',
+      '0.0.0.0',
+      // Git hash
+      '080e6d4af809a46852d1c4b7ed85f06e8a3a72be-invalid'
+    ]
+    it.each(invalidVersionStrings)('invalid version: %s', async (ver) => {
+      const workflow = JSON.parse(JSON.stringify(defaultGraph))
+      workflow.nodes[0].properties.ver = ver
+      expect(await validateComfyWorkflow(workflow)).toBeNull()
+    })
+  })
 })