diff --git a/vite.config.mts b/vite.config.mts
index 731916b4b..6df89c707 100644
--- a/vite.config.mts
+++ b/vite.config.mts
@@ -11,6 +11,13 @@ interface ShimResult {
   exports: string[];
 }
 
+function isLegacyFile(id: string): boolean {
+  return id.endsWith('.ts') && (
+    id.includes("src/extensions/core") ||
+    id.includes("src/scripts")
+  );
+}
+
 function comfyAPIPlugin(): Plugin {
   return {
     name: 'comfy-api-plugin',
@@ -18,8 +25,7 @@ function comfyAPIPlugin(): Plugin {
       if (IS_DEV)
         return null;
 
-      // TODO: Remove second condition after all js files are converted to ts
-      if (id.endsWith('.ts') || (id.endsWith('.js') && id.includes("extensions/core"))) {
+      if (isLegacyFile(id)) {
         const result = transformExports(code, id);
 
         if (result.exports.length > 0) {