Merge branch 'main' of github.com:codeflash-ai/codeflash into fix/js-jest30-loop-runner

Author: mohammedahmed18

.github/workflows/claude.yml

@@ -48,7 +48,7 @@ jobs:
         with:
           use_foundry: "true"
           use_sticky_comment: true
-          allowed_bots: "claude[bot]"
+          allowed_bots: "claude[bot],codeflash-ai[bot]"
           prompt: |
             REPO: ${{ github.repository }}
             PR NUMBER: ${{ github.event.pull_request.number }}

CLAUDE.md

@@ -27,12 +27,29 @@ uv run ruff format codeflash/    # Format
 # Linting (run before committing)
 uv run prek run --from-ref origin/main
 
+# Mypy type checking (run on changed files before committing)
+uv run mypy --non-interactive --config-file pyproject.toml <changed_files>
+
 # Running the CLI
 uv run codeflash --help
 uv run codeflash init            # Initialize in a project
 uv run codeflash --all           # Optimize entire codebase
 ```
 
+## Mypy Type Checking
+
+When modifying code, fix any mypy type errors in the files you changed. Run mypy on changed files:
+
+```bash
+uv run mypy --non-interactive --config-file pyproject.toml <changed_files>
+```
+
+Rules:
+- Fix type annotation issues: missing return types, incorrect types, Optional/None unions, import errors for type hints
+- Do NOT add `# type: ignore` comments — always fix the root cause
+- Do NOT fix type errors that require logic changes, complex generic type rework, or anything that could change runtime behavior
+- Files in `mypy_allowlist.txt` are checked in CI — ensure they remain error-free
+
 <!-- Section below is auto-generated by `tessl install` - do not edit manually -->
 
 # Agent Rules <!-- tessl-managed -->

MULTI_LANGUAGE_ARCHITECTURE.md

@@ -386,7 +386,7 @@ class JavaScriptTransformer:
 
 from pathlib import Path
 from codeflash.languages.base import LanguageSupport, FunctionInfo, CodeContext
-from codeflash.languages.treesitter_utils import TreeSitterAnalyzer
+from codeflash.languages.javascript.treesitter import TreeSitterAnalyzer
 from codeflash.languages.javascript.transformer import JavaScriptTransformer
 
 class JavaScriptSupport(LanguageSupport):
@@ -523,7 +523,7 @@ class JavaScriptSupport(LanguageSupport):
 # codeflash/languages/javascript/test_discovery.py
 
 from pathlib import Path
-from codeflash.languages.treesitter_utils import TreeSitterAnalyzer
+from codeflash.languages.javascript.treesitter import TreeSitterAnalyzer
 
 class JestTestDiscovery:
     """Static analysis-based test discovery for Jest."""

codeflash/code_utils/code_extractor.py

@@ -1772,7 +1772,7 @@ def _extract_calling_function_js(source_code: str, function_name: str, ref_line:
 
     """
     try:
-        from codeflash.languages.treesitter_utils import TreeSitterAnalyzer, TreeSitterLanguage
+        from codeflash.languages.javascript.treesitter import TreeSitterAnalyzer, TreeSitterLanguage
 
         # Try TypeScript first, fall back to JavaScript
         for lang in [TreeSitterLanguage.TYPESCRIPT, TreeSitterLanguage.TSX, TreeSitterLanguage.JAVASCRIPT]:

codeflash/code_utils/code_replacer.py

@@ -26,7 +26,7 @@
 
     from codeflash.discovery.functions_to_optimize import FunctionToOptimize
     from codeflash.languages.base import Language, LanguageSupport
-    from codeflash.languages.treesitter_utils import TreeSitterAnalyzer
+    from codeflash.languages.javascript.treesitter import TreeSitterAnalyzer
     from codeflash.models.models import CodeOptimizationContext, CodeStringsMarkdown, OptimizedCandidate, ValidCode
 
 ASTNodeT = TypeVar("ASTNodeT", bound=ast.AST)
@@ -640,7 +640,7 @@ def _add_global_declarations_for_language(
         return original_source
 
     try:
-        from codeflash.languages.treesitter_utils import get_analyzer_for_file
+        from codeflash.languages.javascript.treesitter import get_analyzer_for_file
 
         analyzer = get_analyzer_for_file(module_abspath)
 

codeflash/code_utils/concolic_utils.py

@@ -105,12 +105,12 @@ def clean_concolic_tests(test_suite_code: str) -> str:
         can_parse = False
         tree = None
 
-    if not can_parse:
+    if not can_parse or tree is None:
         return AssertCleanup().transform_asserts(test_suite_code)
 
     for node in ast.walk(tree):
         if isinstance(node, ast.FunctionDef) and node.name.startswith("test_"):
-            new_body = []
+            new_body: list[ast.stmt] = []
             for stmt in node.body:
                 if isinstance(stmt, ast.Assert):
                     if isinstance(stmt.test, ast.Compare) and isinstance(stmt.test.left, ast.Call):

codeflash/code_utils/coverage_utils.py

@@ -13,14 +13,28 @@
 def extract_dependent_function(main_function: str, code_context: CodeOptimizationContext) -> str | Literal[False]:
     """Extract the single dependent function from the code context excluding the main function."""
     dependent_functions = set()
+
+    # Compare using bare name since AST extracts bare function names
+    bare_main = main_function.rsplit(".", 1)[-1] if "." in main_function else main_function
+
     for code_string in code_context.testgen_context.code_strings:
-        ast_tree = ast.parse(code_string.code)
-        dependent_functions.update(
-            {node.name for node in ast_tree.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))}
-        )
+        # Quick heuristic: skip parsing entirely if there is no 'def' token,
+        # since no function definitions can be present without it.
+        if "def" not in code_string.code:
+            continue
 
-    if main_function in dependent_functions:
-        dependent_functions.discard(main_function)
+        ast_tree = ast.parse(code_string.code)
+        # Add function names directly, skipping the bare main name.
+        for node in ast_tree.body:
+            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+                name = node.name
+                if name == bare_main:
+                    continue
+                dependent_functions.add(name)
+                # If more than one dependent function (other than the main) is found,
+                # we can return False early since the final result cannot be a single name.
+                if len(dependent_functions) > 1:
+                    return False
 
     if not dependent_functions:
         return False
@@ -32,6 +46,9 @@ def extract_dependent_function(main_function: str, code_context: CodeOptimizatio
 
 
 def build_fully_qualified_name(function_name: str, code_context: CodeOptimizationContext) -> str:
+    # If the name is already qualified (contains a dot), return as-is
+    if "." in function_name:
+        return function_name
     full_name = function_name
     for obj_name, parents in code_context.preexisting_objects:
         if obj_name == function_name:

codeflash/code_utils/normalizers/javascript.py

@@ -233,7 +233,7 @@ def normalize(self, code: str) -> str:
 
         """
         try:
-            from codeflash.languages.treesitter_utils import TreeSitterAnalyzer, TreeSitterLanguage
+            from codeflash.languages.javascript.treesitter import TreeSitterAnalyzer, TreeSitterLanguage
 
             lang_map = {"javascript": TreeSitterLanguage.JAVASCRIPT, "typescript": TreeSitterLanguage.TYPESCRIPT}
             lang = lang_map.get(self._get_tree_sitter_language(), TreeSitterLanguage.JAVASCRIPT)

codeflash/discovery/functions_to_optimize.py

@@ -201,7 +201,7 @@ def _is_js_ts_function_exported(file_path: Path, function_name: str) -> tuple[bo
         Tuple of (is_exported, export_name). export_name may be 'default' for default exports.
 
     """
-    from codeflash.languages.treesitter_utils import get_analyzer_for_file
+    from codeflash.languages.javascript.treesitter import get_analyzer_for_file
 
     try:
         source = file_path.read_text(encoding="utf-8")

codeflash/github/PrComment.py

@@ -26,10 +26,10 @@ class PrComment:
 
     def to_json(self) -> dict[str, Union[str, int, dict[str, dict[str, int]], list[BenchmarkDetail], None]]:
         report_table: dict[str, dict[str, int]] = {}
-        for test_type, test_result in self.winning_behavior_test_results.get_test_pass_fail_report_by_type().items():
+        for test_type, counts in self.winning_behavior_test_results.get_test_pass_fail_report_by_type().items():
             name = test_type.to_name()
             if name:
-                report_table[name] = test_result
+                report_table[name] = counts
 
         result: dict[str, Union[str, int, dict[str, dict[str, int]], list[BenchmarkDetail], None]] = {
             "optimization_explanation": self.optimization_explanation,