fix: strip file extensions from JS/TS import paths in generated tests

LLMs often add .js extensions to TypeScript import paths (e.g.,
`import { func } from '../module.js'`), but TypeScript's module resolution
doesn't require explicit extensions. This causes "Cannot find module" errors
when Jest tries to resolve these imports.

This change adds `strip_js_extensions()` to remove .js/.ts/.tsx/.jsx/.mjs/.mts
extensions from relative import paths in generated tests. The function handles:
- ES module imports: import { x } from '../path/file.js'
- CommonJS requires: require('../path/file.js')
- Jest mocks: jest.mock('../path/file.js'), jest.doMock(), etc.

External package imports (lodash, react, etc.) and alias imports (@/components)
are preserved.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Author: misrasaurabh1

codeflash/languages/javascript/instrument.py

@@ -18,6 +18,43 @@
     from codeflash.discovery.functions_to_optimize import FunctionToOptimize
 
 
+# Patterns to strip file extensions from import paths
+# LLMs sometimes add .js extensions to TypeScript imports, which breaks module resolution
+_JS_EXTENSION_PATTERN = re.compile(
+    r"""(from\s+['"])(\.{0,2}/[^'"]+?)(\.(?:js|ts|tsx|jsx|mjs|mts))(['"])"""
+)
+_REQUIRE_EXTENSION_PATTERN = re.compile(
+    r"""(require\s*\(\s*['"])(\.{0,2}/[^'"]+?)(\.(?:js|ts|tsx|jsx|mjs|mts))(['"]\s*\))"""
+)
+_JEST_MOCK_EXTENSION_PATTERN = re.compile(
+    r"""(jest\.(?:mock|doMock|unmock|requireActual|requireMock)\s*\(\s*['"])(\.{0,2}/[^'"]+?)(\.(?:js|ts|tsx|jsx|mjs|mts))(['"])"""
+)
+
+
+def strip_js_extensions(source: str) -> str:
+    """Strip .js/.ts/.tsx/.jsx extensions from relative import paths.
+
+    TypeScript and Jest's module resolution automatically resolve file extensions,
+    so adding them explicitly can cause "Cannot find module" errors when the LLM
+    adds incorrect extensions (e.g., .js to a .ts file).
+
+    This function removes extensions from:
+    - ES module imports: import { x } from '../path/file.js'
+    - CommonJS requires: require('../path/file.js')
+    - Jest mocks: jest.mock('../path/file.js')
+
+    Args:
+        source: The test source code.
+
+    Returns:
+        Source code with extensions stripped from relative import paths.
+
+    """
+    source = _JS_EXTENSION_PATTERN.sub(r"\1\2\4", source)
+    source = _REQUIRE_EXTENSION_PATTERN.sub(r"\1\2\4", source)
+    return _JEST_MOCK_EXTENSION_PATTERN.sub(r"\1\2\4", source)
+
+
 class TestingMode:
     """Testing mode constants."""
 

codeflash/languages/javascript/module_system.py

@@ -184,10 +184,18 @@ def _get_relative_import_path(target_path: Path, source_path: Path) -> str:
 
 
 def add_js_extension(module_path: str) -> str:
-    """Add .js extension to relative module paths for ESM compatibility."""
-    if module_path.startswith(("./", "../")):
-        if not module_path.endswith(".js") and not module_path.endswith(".mjs"):
-            return module_path + ".js"
+    """Process module path for ESM compatibility.
+
+    NOTE: This function intentionally does NOT add extensions because:
+    1. TypeScript projects resolve modules without explicit extensions
+    2. Adding .js to .ts imports causes "Cannot find module" errors
+    3. Modern bundlers (webpack, vite, etc.) handle extension resolution automatically
+
+    The function name is preserved for backward compatibility but the behavior
+    has been changed to NOT add extensions.
+    """
+    # Previously this function added .js extensions, but this caused module resolution
+    # errors in TypeScript projects. We now preserve paths without adding extensions.
     return module_path
 
 

codeflash/verification/verifier.py

@@ -67,6 +67,7 @@ def generate_tests(
             from codeflash.languages.javascript.instrument import (
                 TestingMode,
                 instrument_generated_js_test,
+                strip_js_extensions,
                 validate_and_fix_import_style,
             )
             from codeflash.languages.javascript.module_system import ensure_module_system_compatibility
@@ -75,7 +76,10 @@ def generate_tests(
             func_name = function_to_optimize.function_name
             qualified_name = function_to_optimize.qualified_name
 
-            # First validate and fix import styles
+            # Strip incorrect file extensions from import paths (LLMs sometimes add .js to .ts imports)
+            generated_test_source = strip_js_extensions(generated_test_source)
+
+            # Validate and fix import styles (default vs named exports)
             generated_test_source = validate_and_fix_import_style(generated_test_source, source_file, func_name)
 
             # Convert module system if needed (e.g., CommonJS -> ESM for ESM projects)

tests/test_languages/test_javascript_instrumentation.py

@@ -652,3 +652,111 @@ def test_this_method_call_exact_output(self):
         expected = "        return codeflash.capture('Class.fibonacci', '1', this.fibonacci.bind(this), n - 1);"
         assert transformed == expected, f"Expected:\n{expected}\nGot:\n{transformed}"
         assert counter == 1
+
+
+class TestStripJsExtensions:
+    """Tests for stripping file extensions from import paths."""
+
+    def test_strip_js_extension_from_esm_import(self):
+        """Test stripping .js from ES module imports."""
+        from codeflash.languages.javascript.instrument import strip_js_extensions
+
+        code = "import { getDifferences } from '../src/utils/DynamicBindingUtils.js';"
+        expected = "import { getDifferences } from '../src/utils/DynamicBindingUtils';"
+
+        result = strip_js_extensions(code)
+        assert result == expected
+
+    def test_strip_ts_extension_from_esm_import(self):
+        """Test stripping .ts from ES module imports."""
+        from codeflash.languages.javascript.instrument import strip_js_extensions
+
+        code = "import { func } from './module.ts';"
+        expected = "import { func } from './module';"
+
+        result = strip_js_extensions(code)
+        assert result == expected
+
+    def test_strip_extension_from_require(self):
+        """Test stripping extensions from require() calls."""
+        from codeflash.languages.javascript.instrument import strip_js_extensions
+
+        code = "const { func } = require('../utils/helper.js');"
+        expected = "const { func } = require('../utils/helper');"
+
+        result = strip_js_extensions(code)
+        assert result == expected
+
+    def test_strip_extension_from_jest_mock(self):
+        """Test stripping extensions from jest.mock() calls."""
+        from codeflash.languages.javascript.instrument import strip_js_extensions
+
+        code = "jest.mock('../src/utils/DynamicBindingUtils.js');"
+        expected = "jest.mock('../src/utils/DynamicBindingUtils');"
+
+        result = strip_js_extensions(code)
+        assert result == expected
+
+    def test_strip_extension_from_jest_doMock(self):
+        """Test stripping extensions from jest.doMock() calls."""
+        from codeflash.languages.javascript.instrument import strip_js_extensions
+
+        code = "jest.doMock('./helper.ts');"
+        expected = "jest.doMock('./helper');"
+
+        result = strip_js_extensions(code)
+        assert result == expected
+
+    def test_preserve_external_package_imports(self):
+        """Test that external package imports are not modified."""
+        from codeflash.languages.javascript.instrument import strip_js_extensions
+
+        code = "import lodash from 'lodash';"
+
+        result = strip_js_extensions(code)
+        assert result == code  # Should be unchanged
+
+    def test_preserve_absolute_imports(self):
+        """Test that absolute/alias imports are not modified."""
+        from codeflash.languages.javascript.instrument import strip_js_extensions
+
+        code = "import { Component } from '@/components/Button';"
+
+        result = strip_js_extensions(code)
+        assert result == code  # Should be unchanged
+
+    def test_strip_multiple_extensions_in_file(self):
+        """Test stripping multiple extensions in a single file."""
+        from codeflash.languages.javascript.instrument import strip_js_extensions
+
+        code = """
+import { func1 } from '../utils/helper.js';
+import { func2 } from './local.ts';
+const { func3 } = require('../lib/util.tsx');
+jest.mock('../mocks/mock.jsx');
+"""
+        expected = """
+import { func1 } from '../utils/helper';
+import { func2 } from './local';
+const { func3 } = require('../lib/util');
+jest.mock('../mocks/mock');
+"""
+
+        result = strip_js_extensions(code)
+        assert result == expected
+
+    def test_strip_mjs_mts_extensions(self):
+        """Test stripping .mjs and .mts extensions."""
+        from codeflash.languages.javascript.instrument import strip_js_extensions
+
+        code = """
+import { a } from './module.mjs';
+import { b } from '../util.mts';
+"""
+        expected = """
+import { a } from './module';
+import { b } from '../util';
+"""
+
+        result = strip_js_extensions(code)
+        assert result == expected