diff --git a/src/extensions/core/rerouteNode.ts b/src/extensions/core/rerouteNode.ts
index 42ab5e75d..1ace3e543 100644
--- a/src/extensions/core/rerouteNode.ts
+++ b/src/extensions/core/rerouteNode.ts
@@ -130,7 +130,6 @@ app.registerExtension({
                       : null
                   if (
                     inputType &&
-                    // @ts-expect-error Will self-resolve when LiteGraph types are generated
                     !LiteGraph.isValidConnection(inputType, nodeOutType)
                   ) {
                     // The output doesnt match our input so disconnect it
diff --git a/src/types/litegraph-core-augmentation.d.ts b/src/types/litegraph-core-augmentation.d.ts
index fafd1e901..d58f10ec8 100644
--- a/src/types/litegraph-core-augmentation.d.ts
+++ b/src/types/litegraph-core-augmentation.d.ts
@@ -11,6 +11,8 @@ declare module '@comfyorg/litegraph' {
     slot_types_out: string[]
     slot_types_default_out: Record<string, string[]>
     slot_types_default_in: Record<string, string[]>
+
+    isValidConnection(type_a: ISlotType, type_b: ISlotType): boolean
   }
 
   import type { LiteGraph as LG } from '@comfyorg/litegraph/dist/litegraph'
diff --git a/src/types/litegraphTypes.ts b/src/types/litegraphTypes.ts
index 9c93d4b83..c632302ec 100644
--- a/src/types/litegraphTypes.ts
+++ b/src/types/litegraphTypes.ts
@@ -1,4 +1,4 @@
-import {
+import type {
   ConnectingLink,
   LGraphNode,
   Vector2,
@@ -6,6 +6,7 @@ import {
   INodeOutputSlot,
   INodeSlot
 } from '@comfyorg/litegraph'
+import { LiteGraph } from '@comfyorg/litegraph'
 
 export class ConnectingLinkImpl implements ConnectingLink {
   node: LGraphNode
@@ -56,9 +57,8 @@ export class ConnectingLinkImpl implements ConnectingLink {
       this.releaseSlotType === 'output' ? newNode.outputs : newNode.inputs
     if (!newNodeSlots) return
 
-    const newNodeSlot = newNodeSlots.findIndex(
-      (slot: INodeSlot) =>
-        slot.type === this.type || slot.type === '*' || this.type === '*'
+    const newNodeSlot = newNodeSlots.findIndex((slot: INodeSlot) =>
+      LiteGraph.isValidConnection(slot.type, this.type)
     )
 
     if (newNodeSlot === -1) {