fix: stop adding .js extensions when converting CommonJS to ESM

The `add_js_extension()` function in `module_system.py` was adding .js
extensions when converting CommonJS requires to ESM imports. This caused
"Cannot find module" errors because TypeScript/Jest module resolution
doesn't require explicit extensions.

This change modifies `add_js_extension()` to NOT add extensions, as
TypeScript and modern bundlers handle extension resolution automatically.

Note: The aiservice now handles stripping any extensions that LLMs might
add (see codeflash-internal PR #2341), so CLI-side stripping is removed.

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

Author: misrasaurabh1

codeflash/discovery/functions_to_optimize.py

@@ -216,6 +216,32 @@ def get_files_for_language(
     return files
 
 
+def _is_js_ts_function_exported(file_path: Path, function_name: str) -> tuple[bool, str | None]:
+    """Check if a JavaScript/TypeScript function is exported from its module.
+
+    For JS/TS, functions that are not exported cannot be imported by tests,
+    making them impossible to optimize.
+
+    Args:
+        file_path: Path to the source file.
+        function_name: Name of the function to check.
+
+    Returns:
+        Tuple of (is_exported, export_name). export_name may be 'default' for default exports.
+
+    """
+    from codeflash.languages.treesitter_utils import get_analyzer_for_file
+
+    try:
+        source = file_path.read_text(encoding="utf-8")
+        analyzer = get_analyzer_for_file(file_path)
+        return analyzer.is_function_exported(source, function_name)
+    except Exception as e:
+        logger.debug(f"Failed to check export status for {function_name}: {e}")
+        # Return True to avoid blocking in case of errors
+        return True, None
+
+
 def _find_all_functions_in_python_file(file_path: Path) -> dict[Path, list[FunctionToOptimize]]:
     """Find all optimizable functions in a Python file using AST parsing.
 
@@ -338,6 +364,36 @@ def get_functions_to_optimize(
                     exit_with_message(
                         f"Function {only_get_this_function} not found in file {file}\nor the function does not have a 'return' statement or is a property"
                     )
+
+                # For JavaScript/TypeScript, verify that the function (or its parent class) is exported
+                # Non-exported functions cannot be imported by tests
+                if found_function.language in ("javascript", "typescript"):
+                    # For class methods, check if the parent class is exported
+                    # For standalone functions, check if the function itself is exported
+                    if found_function.parents:
+                        # It's a class method - check if the class is exported
+                        name_to_check = found_function.top_level_parent_name
+                    else:
+                        # It's a standalone function - check if the function is exported
+                        name_to_check = found_function.function_name
+
+                    is_exported, export_name = _is_js_ts_function_exported(file, name_to_check)
+                    if not is_exported:
+                        if found_function.parents:
+                            logger.debug(
+                                f"Class '{name_to_check}' containing method '{found_function.function_name}' "
+                                f"is not exported from {file}. "
+                                f"In JavaScript/TypeScript, only exported classes/functions can be optimized "
+                                f"because tests need to import them."
+                            )
+                        else:
+                            logger.debug(
+                                f"Function '{found_function.function_name}' is not exported from {file}. "
+                                f"In JavaScript/TypeScript, only exported functions can be optimized because "
+                                f"tests need to import them."
+                            )
+                        return {}, 0, None
+
                 functions[file] = [found_function]
         else:
             logger.info("Finding all functions modified in the current git diff ...")

codeflash/languages/javascript/instrument.py

@@ -18,43 +18,6 @@
     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/verification/verifier.py

@@ -67,7 +67,6 @@ 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
@@ -76,9 +75,6 @@ def generate_tests(
             func_name = function_to_optimize.function_name
             qualified_name = function_to_optimize.qualified_name
 
-            # 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)
 

tests/test_javascript_function_discovery.py

@@ -346,11 +346,11 @@ def test_get_specific_function(self, tmp_path):
         """Test getting a specific function by name."""
         js_file = tmp_path / "math_utils.js"
         js_file.write_text("""
-function add(a, b) {
+export function add(a, b) {
     return a + b;
 }
 
-function subtract(a, b) {
+export function subtract(a, b) {
     return a - b;
 }
 """)
@@ -378,7 +378,7 @@ def test_get_class_method(self, tmp_path):
         """Test getting a specific class method."""
         js_file = tmp_path / "calculator.js"
         js_file.write_text("""
-class Calculator {
+export class Calculator {
     add(a, b) {
         return a + b;
     }
@@ -388,7 +388,7 @@ class Calculator {
     }
 }
 
-function standaloneFunc() {
+export function standaloneFunc() {
     return 42;
 }
 """)

tests/test_languages/test_javascript_instrumentation.py

@@ -651,112 +651,4 @@ 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
+        assert counter == 1