diff --git a/src/types/nodeSource.ts b/src/types/nodeSource.ts
index c20ecca00d..58c14c5447 100644
--- a/src/types/nodeSource.ts
+++ b/src/types/nodeSource.ts
@@ -37,7 +37,11 @@ export const getNodeSource = (python_module?: string): NodeSource => {
       badgeText: '🦊'
     }
   } else if (modules[0] === 'custom_nodes') {
-    const displayName = shortenNodeName(modules[1])
+    const moduleName = modules[1]
+    // Custom nodes installed via ComfyNodeRegistry will be in the format of
+    // custom_nodes.<custom node name>@<version>
+    const customNodeName = moduleName.split('@')[0]
+    const displayName = shortenNodeName(customNodeName)
     return {
       type: NodeSourceType.CustomNodes,
       className: 'comfy-custom-nodes',
diff --git a/tests-ui/tests/fast/nodeSource.test.ts b/tests-ui/tests/fast/nodeSource.test.ts
new file mode 100644
index 0000000000..a79aaeea3d
--- /dev/null
+++ b/tests-ui/tests/fast/nodeSource.test.ts
@@ -0,0 +1,63 @@
+import { getNodeSource, NodeSourceType } from '@/types/nodeSource'
+
+describe('getNodeSource', () => {
+  it('should return UNKNOWN_NODE_SOURCE when python_module is undefined', () => {
+    const result = getNodeSource(undefined)
+    expect(result).toEqual({
+      type: NodeSourceType.Unknown,
+      className: 'comfy-unknown',
+      displayText: 'Unknown',
+      badgeText: '?'
+    })
+  })
+
+  it('should identify core nodes from nodes module', () => {
+    const result = getNodeSource('nodes.some_module')
+    expect(result).toEqual({
+      type: NodeSourceType.Core,
+      className: 'comfy-core',
+      displayText: 'Comfy Core',
+      badgeText: '🦊'
+    })
+  })
+
+  it('should identify core nodes from comfy_extras module', () => {
+    const result = getNodeSource('comfy_extras.some_module')
+    expect(result).toEqual({
+      type: NodeSourceType.Core,
+      className: 'comfy-core',
+      displayText: 'Comfy Core',
+      badgeText: '🦊'
+    })
+  })
+
+  it('should identify custom nodes and format their names', () => {
+    const result = getNodeSource('custom_nodes.ComfyUI-Example')
+    expect(result).toEqual({
+      type: NodeSourceType.CustomNodes,
+      className: 'comfy-custom-nodes',
+      displayText: 'Example',
+      badgeText: 'Example'
+    })
+  })
+
+  it('should identify custom nodes with version and format their names', () => {
+    const result = getNodeSource('custom_nodes.ComfyUI-Example@1.0.0')
+    expect(result).toEqual({
+      type: NodeSourceType.CustomNodes,
+      className: 'comfy-custom-nodes',
+      displayText: 'Example',
+      badgeText: 'Example'
+    })
+  })
+
+  it('should return UNKNOWN_NODE_SOURCE for unrecognized modules', () => {
+    const result = getNodeSource('unknown_module.something')
+    expect(result).toEqual({
+      type: NodeSourceType.Unknown,
+      className: 'comfy-unknown',
+      displayText: 'Unknown',
+      badgeText: '?'
+    })
+  })
+})