Multiple fixes around vitest

codeflash/code_utils/edit_generated_tests.py

@@ -361,21 +361,27 @@ def normalize_codeflash_imports(source: str) -> str:
     return _CODEFLASH_IMPORT_PATTERN.sub(r"import \1 from 'codeflash'", source)
 
 
-def inject_test_globals(generated_tests: GeneratedTestsList) -> GeneratedTestsList:
+def inject_test_globals(generated_tests: GeneratedTestsList, test_framework: str = "jest") -> GeneratedTestsList:
     # TODO: inside the prompt tell the llm if it should import jest functions or it's already injected in the global window
     """Inject test globals into all generated tests.
 
     Args:
         generated_tests: List of generated tests.
+        test_framework: The test framework being used ("jest", "vitest", or "mocha").
 
     Returns:
         Generated tests with test globals injected.
 
     """
     # we only inject test globals for esm modules
-    global_import = (
-        "import { jest, describe, it, expect, beforeEach, afterEach, beforeAll, test } from '@jest/globals'\n"
-    )
+    # Use vitest imports for vitest projects, jest imports for jest projects
+    if test_framework == "vitest":
+        global_import = "import { vi, describe, it, expect, beforeEach, afterEach, beforeAll, test } from 'vitest'\n"
+    else:
+        # Default to jest imports for jest and other frameworks
+        global_import = (
+            "import { jest, describe, it, expect, beforeEach, afterEach, beforeAll, test } from '@jest/globals'\n"
+        )
 
     for test in generated_tests.generated_tests:
         test.generated_original_test_source = global_import + test.generated_original_test_source

codeflash/languages/javascript/import_resolver.py

@@ -558,6 +558,7 @@ def _find_helpers_recursive(
 
         """
         from codeflash.discovery.functions_to_optimize import FunctionToOptimize
+        from codeflash.languages.registry import get_language_support
         from codeflash.languages.treesitter_utils import get_analyzer_for_file
 
         if context.current_depth >= context.max_depth:
@@ -578,12 +579,15 @@ def _find_helpers_recursive(
         imports = analyzer.find_imports(source)
 
         # Create FunctionToOptimize for the helper
+        # Get language from the language support registry
+        lang_support = get_language_support(file_path)
         func_info = FunctionToOptimize(
             function_name=helper.name,
             file_path=file_path,
             parents=[],
             starting_line=helper.start_line,
             ending_line=helper.end_line,
+            language=str(lang_support.language),
         )
 
         # Recursively find helpers

codeflash/languages/javascript/module_system.py

@@ -416,8 +416,10 @@ def ensure_module_system_compatibility(code: str, target_module_system: str, pro
     is_esm = has_import or has_export
 
     # Convert if needed
-    if target_module_system == ModuleSystem.ES_MODULE and is_commonjs and not is_esm:
-        logger.debug("Converting CommonJS to ES Module syntax")
+    # For ESM target: convert any require statements, even if there are also import statements
+    # This handles generated tests that have ESM imports for test globals but CommonJS for the function
+    if target_module_system == ModuleSystem.ES_MODULE and has_require:
+        logger.debug("Converting CommonJS require statements to ES Module syntax")
         return convert_commonjs_to_esm(code)
 
     if target_module_system == ModuleSystem.COMMONJS and is_esm and not is_commonjs:

codeflash/languages/javascript/parse.py

@@ -16,12 +16,7 @@
 from junitparser.xunit2 import JUnitXml
 
 from codeflash.cli_cmds.console import logger
-from codeflash.models.models import (
-    FunctionTestInvocation,
-    InvocationId,
-    TestResults,
-    TestType,
-)
+from codeflash.models.models import FunctionTestInvocation, InvocationId, TestResults, TestType
 
 if TYPE_CHECKING:
     import subprocess
@@ -127,6 +122,7 @@ def parse_jest_test_xml(
     # This handles cases where instrumented files are in temp directories
     instrumented_path_lookup: dict[str, tuple[Path, TestType]] = {}
     for test_file in test_files.test_files:
+        # Add behavior instrumented file paths
         if test_file.instrumented_behavior_file_path:
             # Store both the absolute path and resolved path as keys
             abs_path = str(test_file.instrumented_behavior_file_path.resolve())
@@ -137,18 +133,35 @@ def parse_jest_test_xml(
                 test_file.test_type,
             )
             logger.debug(f"Jest XML lookup: registered {abs_path}")
+        # Also add benchmarking file paths (perf-only instrumented tests)
+        if test_file.benchmarking_file_path:
+            bench_abs_path = str(test_file.benchmarking_file_path.resolve())
+            instrumented_path_lookup[bench_abs_path] = (test_file.benchmarking_file_path, test_file.test_type)
+            instrumented_path_lookup[str(test_file.benchmarking_file_path)] = (
+                test_file.benchmarking_file_path,
+                test_file.test_type,
+            )
+            logger.debug(f"Jest XML lookup: registered benchmark {bench_abs_path}")
 
     # Also build a filename-only lookup for fallback matching
     # This handles cases where JUnit XML has relative paths that don't match absolute paths
     # e.g., JUnit has "test/utils__perfinstrumented.test.ts" but lookup has absolute paths
     filename_lookup: dict[str, tuple[Path, TestType]] = {}
     for test_file in test_files.test_files:
+        # Add instrumented_behavior_file_path (behavior tests)
         if test_file.instrumented_behavior_file_path:
             filename = test_file.instrumented_behavior_file_path.name
             # Only add if not already present (avoid overwrites in case of duplicate filenames)
             if filename not in filename_lookup:
                 filename_lookup[filename] = (test_file.instrumented_behavior_file_path, test_file.test_type)
                 logger.debug(f"Jest XML filename lookup: registered {filename}")
+        # Also add benchmarking_file_path (perf-only tests) - these have different filenames
+        # e.g., utils__perfonlyinstrumented.test.ts vs utils__perfinstrumented.test.ts
+        if test_file.benchmarking_file_path:
+            bench_filename = test_file.benchmarking_file_path.name
+            if bench_filename not in filename_lookup:
+                filename_lookup[bench_filename] = (test_file.benchmarking_file_path, test_file.test_type)
+                logger.debug(f"Jest XML filename lookup: registered benchmark file {bench_filename}")
 
     # Fallback: if JUnit XML doesn't have system-out, use subprocess stdout directly
     global_stdout = ""
@@ -184,6 +197,21 @@ def parse_jest_test_xml(
             key = match.groups()[:5]
             end_matches_dict[key] = match
 
+        # Also collect timing markers from testcase-level system-out (Vitest puts output at testcase level)
+        for tc in suite:
+            tc_system_out = tc._elem.find("system-out")  # noqa: SLF001
+            if tc_system_out is not None and tc_system_out.text:
+                tc_stdout = tc_system_out.text.strip()
+                logger.debug(f"Vitest testcase system-out found: {len(tc_stdout)} chars, first 200: {tc_stdout[:200]}")
+                end_marker_count = 0
+                for match in jest_end_pattern.finditer(tc_stdout):
+                    key = match.groups()[:5]
+                    end_matches_dict[key] = match
+                    end_marker_count += 1
+                if end_marker_count > 0:
+                    logger.debug(f"Found {end_marker_count} END timing markers in testcase system-out")
+                start_matches.extend(jest_start_pattern.finditer(tc_stdout))
+
         for testcase in suite:
             testcase_count += 1
             test_class_path = testcase.classname  # For Jest, this is the file path
@@ -311,7 +339,18 @@ def parse_jest_test_xml(
                         matching_ends_direct.append(end_match)
 
             if not matching_starts and not matching_ends_direct:
-                # No timing markers found - add basic result
+                # No timing markers found - use JUnit XML time attribute as fallback
+                # The time attribute is in seconds (e.g., "0.00077875"), convert to nanoseconds
+                runtime = None
+                try:
+                    time_attr = testcase._elem.attrib.get("time")  # noqa: SLF001
+                    if time_attr:
+                        time_seconds = float(time_attr)
+                        runtime = int(time_seconds * 1_000_000_000)  # Convert seconds to nanoseconds
+                        logger.debug(f"Jest XML: using time attribute for {test_name}: {time_seconds}s = {runtime}ns")
+                except (ValueError, TypeError) as e:
+                    logger.debug(f"Jest XML: could not parse time attribute: {e}")
+
                 test_results.add(
                     FunctionTestInvocation(
                         loop_index=1,
@@ -323,7 +362,7 @@ def parse_jest_test_xml(
                             iteration_id="",
                         ),
                         file_name=test_file_path,
-                        runtime=None,
+                        runtime=runtime,
                         test_framework=test_config.test_framework,
                         did_pass=result,
                         test_type=test_type,

codeflash/languages/javascript/support.py

@@ -2098,6 +2098,10 @@ def run_behavioral_tests(
             candidate_index=candidate_index,
         )
 
+    # JavaScript/TypeScript benchmarking uses high max_loops like Python (100,000)
+    # The actual loop count is limited by target_duration_seconds, not max_loops
+    JS_BENCHMARKING_MAX_LOOPS = 100_000
+
     def run_benchmarking_tests(
         self,
         test_paths: Any,
@@ -2131,6 +2135,9 @@ def run_benchmarking_tests(
 
         framework = test_framework or get_js_test_framework_or_default()
 
+        # Use JS-specific high max_loops - actual loop count is limited by target_duration
+        effective_max_loops = self.JS_BENCHMARKING_MAX_LOOPS
+
         if framework == "vitest":
             from codeflash.languages.javascript.vitest_runner import run_vitest_benchmarking_tests
 
@@ -2141,7 +2148,7 @@ def run_benchmarking_tests(
                 timeout=timeout,
                 project_root=project_root,
                 min_loops=min_loops,
-                max_loops=max_loops,
+                max_loops=effective_max_loops,
                 target_duration_ms=int(target_duration_seconds * 1000),
             )
 
@@ -2154,7 +2161,7 @@ def run_benchmarking_tests(
             timeout=timeout,
             project_root=project_root,
             min_loops=min_loops,
-            max_loops=max_loops,
+            max_loops=effective_max_loops,
             target_duration_ms=int(target_duration_seconds * 1000),
         )