fix: Add .js extensions to relative imports in ESM TypeScript projects

codeflash/languages/javascript/module_system.py

@@ -513,3 +513,54 @@ def ensure_vitest_imports(code: str, test_framework: str) -> str:
 
     logger.debug("Added vitest imports: %s", used_globals)
     return "\n".join(lines)
+
+
+def add_js_extensions_to_relative_imports(code: str) -> str:
+    """Add .js extensions to relative imports in ESM code.
+
+    In ESM mode with TypeScript, Node.js requires explicit .js extensions
+    for relative imports, even though the source files are .ts files.
+
+    This function adds .js extensions to relative imports that don't already
+    have a file extension.
+
+    Args:
+        code: JavaScript/TypeScript code with import statements.
+
+    Returns:
+        Code with .js extensions added to relative imports.
+
+    Examples:
+        >>> add_js_extensions_to_relative_imports("import X from './module';")
+        "import X from './module.js';"
+
+        >>> add_js_extensions_to_relative_imports("import X from './module.js';")
+        "import X from './module.js';"
+
+        >>> add_js_extensions_to_relative_imports("import X from 'node:assert';")
+        "import X from 'node:assert';"
+
+    """
+    # Pattern to match ES module import statements with relative paths
+    # Matches: import ... from './path' or import ... from "../path"
+    # Groups: (import statement)(quote char)(relative path)(quote char)
+    import_pattern = re.compile(
+        r"(import\s+(?:(?:\{[^}]*\})|(?:\*\s+as\s+\w+)|(?:\w+))\s+from\s+)(['\"])(\.\.?[^'\"]+)(['\"])"
+    )
+
+    def add_extension(match):
+        """Add .js extension if the import path doesn't have one."""
+        prefix = match.group(1)  # "import ... from "
+        quote_open = match.group(2)  # ' or "
+        path = match.group(3)  # The relative path (e.g., "./module" or "../foo/bar")
+        quote_close = match.group(4)  # ' or "
+
+        # Check if path already has an extension
+        # Common extensions: .js, .ts, .jsx, .tsx, .mjs, .mts, .json
+        if re.search(r"\.(js|ts|jsx|tsx|mjs|mts|json)$", path):
+            return match.group(0)
+
+        # Add .js extension
+        return f"{prefix}{quote_open}{path}.js{quote_close}"
+
+    return import_pattern.sub(add_extension, code)

codeflash/languages/javascript/support.py

@@ -2012,6 +2012,7 @@ def process_generated_test_strings(
             validate_and_fix_import_style,
         )
         from codeflash.languages.javascript.module_system import (
+            ModuleSystem,
             ensure_module_system_compatibility,
             ensure_vitest_imports,
         )
@@ -2036,6 +2037,13 @@ def process_generated_test_strings(
             generated_test_source, project_module_system, test_cfg.tests_project_rootdir
         )
 
+        # Add .js extensions to relative imports for ESM projects
+        # TypeScript + ESM requires explicit .js extensions even for .ts source files
+        if project_module_system == ModuleSystem.ES_MODULE:
+            from codeflash.languages.javascript.module_system import add_js_extensions_to_relative_imports
+
+            generated_test_source = add_js_extensions_to_relative_imports(generated_test_source)
+
         # Ensure vitest imports are present when using vitest framework
         generated_test_source = ensure_vitest_imports(generated_test_source, test_cfg.test_framework)
 

tests/test_languages/test_javascript_module_system.py

@@ -284,3 +284,80 @@ def test_real_world_budibase_import(self):
         result = convert_commonjs_to_esm(code)
         expected = "import { queue, context, db as dbCore, cache, events } from '@budibase/backend-core';"
         assert result == expected
+
+
+class TestAddJsExtensionsToRelativeImports:
+    """Tests for adding .js extensions to relative imports in ESM mode."""
+
+    def test_add_js_extension_to_relative_import(self):
+        """Test adding .js extension to relative import without extension."""
+        from codeflash.languages.javascript.module_system import add_js_extensions_to_relative_imports
+
+        code = "import TreeNode from '../../injector/topology-tree/tree-node';"
+        result = add_js_extensions_to_relative_imports(code)
+        expected = "import TreeNode from '../../injector/topology-tree/tree-node.js';"
+        assert result == expected
+
+    def test_add_js_extension_to_single_dot_import(self):
+        """Test adding .js extension to same-directory import."""
+        from codeflash.languages.javascript.module_system import add_js_extensions_to_relative_imports
+
+        code = "import { foo } from './module';"
+        result = add_js_extensions_to_relative_imports(code)
+        expected = "import { foo } from './module.js';"
+        assert result == expected
+
+    def test_skip_imports_with_existing_extensions(self):
+        """Test that imports with extensions are left unchanged."""
+        from codeflash.languages.javascript.module_system import add_js_extensions_to_relative_imports
+
+        code = "import TreeNode from '../../tree-node.js';"
+        result = add_js_extensions_to_relative_imports(code)
+        assert result == code
+
+        code2 = "import TreeNode from '../../tree-node.ts';"
+        result2 = add_js_extensions_to_relative_imports(code2)
+        assert result2 == code2
+
+    def test_skip_node_modules_imports(self):
+        """Test that node_modules imports are left unchanged."""
+        from codeflash.languages.javascript.module_system import add_js_extensions_to_relative_imports
+
+        code = "import assert from 'node:assert/strict';"
+        result = add_js_extensions_to_relative_imports(code)
+        assert result == code
+
+        code2 = "import { describe } from 'mocha';"
+        result2 = add_js_extensions_to_relative_imports(code2)
+        assert result2 == code2
+
+    def test_multiple_imports(self):
+        """Test handling multiple imports in one code block."""
+        from codeflash.languages.javascript.module_system import add_js_extensions_to_relative_imports
+
+        code = """import assert from 'node:assert/strict';
+import TreeNode from '../../injector/topology-tree/tree-node';
+import { helper } from './helper';"""
+        result = add_js_extensions_to_relative_imports(code)
+        expected = """import assert from 'node:assert/strict';
+import TreeNode from '../../injector/topology-tree/tree-node.js';
+import { helper } from './helper.js';"""
+        assert result == expected
+
+    def test_named_imports(self):
+        """Test adding extensions to named imports."""
+        from codeflash.languages.javascript.module_system import add_js_extensions_to_relative_imports
+
+        code = "import { foo, bar } from '../utils/helpers';"
+        result = add_js_extensions_to_relative_imports(code)
+        expected = "import { foo, bar } from '../utils/helpers.js';"
+        assert result == expected
+
+    def test_namespace_imports(self):
+        """Test adding extensions to namespace imports."""
+        from codeflash.languages.javascript.module_system import add_js_extensions_to_relative_imports
+
+        code = "import * as helpers from '../utils';"
+        result = add_js_extensions_to_relative_imports(code)
+        expected = "import * as helpers from '../utils.js';"
+        assert result == expected