refactor: use helper file for async decorator instrumentation

Replace inline code injection with a helper file approach that writes
decorator implementations to a separate codeflash_async_wrapper.py file.
This removes the codeflash package import dependency from instrumented
source files while keeping line numbers stable (only 1 import + 1
decorator line added, same as before).

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

Author: KRRT7

code_to_optimize/code_directories/async_e2e/main.py

@@ -1,16 +1,18 @@
-import time
 import asyncio
 
 
 async def retry_with_backoff(func, max_retries=3):
     if max_retries < 1:
         raise ValueError("max_retries must be at least 1")
     last_exception = None
+    _sleep = asyncio.sleep
     for attempt in range(max_retries):
         try:
             return await func()
         except Exception as e:
             last_exception = e
             if attempt < max_retries - 1:
-                time.sleep(0.0001 * attempt)
+                delay = 0.0001 * attempt
+                if delay:
+                    await _sleep(delay)
     raise last_exception

codeflash/code_utils/instrument_existing_tests.py

@@ -1700,69 +1700,44 @@ def get_async_inline_code(mode: TestingMode) -> str:
     return get_performance_async_inline_code()
 
 
-class AsyncInlineCodeInjector(cst.CSTTransformer):
-    """Injects async decorator function definitions inline instead of importing from codeflash."""
+ASYNC_HELPER_FILENAME = "codeflash_async_wrapper.py"
 
-    def __init__(self, mode: TestingMode = TestingMode.BEHAVIOR) -> None:
-        self.mode = mode
-        self.has_inline_definition = False
-        self.has_old_import = False
-
-    def _get_decorator_name(self) -> str:
-        if self.mode == TestingMode.BEHAVIOR:
-            return "codeflash_behavior_async"
-        if self.mode == TestingMode.CONCURRENCY:
-            return "codeflash_concurrency_async"
-        return "codeflash_performance_async"
-
-    def visit_ImportFrom(self, node: cst.ImportFrom) -> None:
-        if (
-            isinstance(node.module, cst.Attribute)
-            and isinstance(node.module.value, cst.Attribute)
-            and isinstance(node.module.value.value, cst.Name)
-            and node.module.value.value.value == "codeflash"
-            and node.module.value.attr.value == "code_utils"
-            and node.module.attr.value == "codeflash_wrap_decorator"
-            and not isinstance(node.names, cst.ImportStar)
-        ):
-            decorator_name = self._get_decorator_name()
-            for import_alias in node.names:
-                if import_alias.name.value == decorator_name:
-                    self.has_old_import = True
 
-    def visit_FunctionDef(self, node: cst.FunctionDef) -> None:
-        if node.name.value == self._get_decorator_name():
-            self.has_inline_definition = True
+def get_decorator_name_for_mode(mode: TestingMode) -> str:
+    if mode == TestingMode.BEHAVIOR:
+        return "codeflash_behavior_async"
+    if mode == TestingMode.CONCURRENCY:
+        return "codeflash_concurrency_async"
+    return "codeflash_performance_async"
+
 
-    def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> cst.Module:
-        if self.has_inline_definition or self.has_old_import:
-            return updated_node
-        inline_code = get_async_inline_code(self.mode)
-        inline_stmts = cst.parse_module(inline_code).body
-        return updated_node.with_changes(body=[*inline_stmts, *list(updated_node.body)])
+def write_async_helper_file(target_dir: Path, mode: TestingMode) -> Path:
+    """Write the async decorator helper file to the target directory."""
+    helper_path = target_dir / ASYNC_HELPER_FILENAME
+    if helper_path.exists():
+        decorator_name = get_decorator_name_for_mode(mode)
+        if f"def {decorator_name}" in helper_path.read_text("utf-8"):
+            return helper_path
+    helper_path.write_text(get_async_inline_code(mode), "utf-8")
+    return helper_path
 
 
 def add_async_decorator_to_function(
-    source_path: Path, function: FunctionToOptimize, mode: TestingMode = TestingMode.BEHAVIOR
+    source_path: Path,
+    function: FunctionToOptimize,
+    mode: TestingMode = TestingMode.BEHAVIOR,
+    project_root: Path | None = None,
 ) -> bool:
     """Add async decorator to an async function definition and write back to file.
 
-    Args:
-    ----
-        source_path: Path to the source file to modify in-place.
-        function: The FunctionToOptimize object representing the target async function.
-        mode: The testing mode to determine which decorator to apply.
-
-    Returns:
-    -------
-        Boolean indicating whether the decorator was successfully added.
+    Writes a helper file containing the decorator implementation to project_root (or source directory
+    as fallback) and adds a standard import + decorator to the source file.
 
     """
     if not function.is_async:
         return False
 
     try:
-        # Read source code
         with source_path.open(encoding="utf8") as f:
             source_code = f.read()
 
@@ -1772,10 +1747,14 @@ def add_async_decorator_to_function(
         decorator_transformer = AsyncDecoratorAdder(function, mode)
         module = module.visit(decorator_transformer)
 
-        # Add the import if decorator was added
         if decorator_transformer.added_decorator:
-            import_transformer = AsyncInlineCodeInjector(mode)
-            module = module.visit(import_transformer)
+            # Write the helper file to project_root (on sys.path) or source dir as fallback
+            helper_dir = project_root if project_root is not None else source_path.parent
+            write_async_helper_file(helper_dir, mode)
+            # Add the import via CST so sort_imports can place it correctly
+            decorator_name = get_decorator_name_for_mode(mode)
+            import_node = cst.parse_statement(f"from codeflash_async_wrapper import {decorator_name}")
+            module = module.with_changes(body=[import_node, *list(module.body)])
 
         modified_code = sort_imports(code=module.code, float_to_top=True)
     except Exception as e:

codeflash/optimization/function_optimizer.py

@@ -2296,7 +2296,10 @@ def establish_original_code_baseline(
             from codeflash.code_utils.instrument_existing_tests import add_async_decorator_to_function
 
             success = add_async_decorator_to_function(
-                self.function_to_optimize.file_path, self.function_to_optimize, TestingMode.BEHAVIOR
+                self.function_to_optimize.file_path,
+                self.function_to_optimize,
+                TestingMode.BEHAVIOR,
+                project_root=self.project_root,
             )
 
         # Instrument codeflash capture
@@ -2361,7 +2364,10 @@ def establish_original_code_baseline(
                 from codeflash.code_utils.instrument_existing_tests import add_async_decorator_to_function
 
                 add_async_decorator_to_function(
-                    self.function_to_optimize.file_path, self.function_to_optimize, TestingMode.PERFORMANCE
+                    self.function_to_optimize.file_path,
+                    self.function_to_optimize,
+                    TestingMode.PERFORMANCE,
+                    project_root=self.project_root,
                 )
 
             try:
@@ -2535,7 +2541,10 @@ def run_optimized_candidate(
                 from codeflash.code_utils.instrument_existing_tests import add_async_decorator_to_function
 
                 add_async_decorator_to_function(
-                    self.function_to_optimize.file_path, self.function_to_optimize, TestingMode.BEHAVIOR
+                    self.function_to_optimize.file_path,
+                    self.function_to_optimize,
+                    TestingMode.BEHAVIOR,
+                    project_root=self.project_root,
                 )
 
             try:
@@ -2611,7 +2620,10 @@ def run_optimized_candidate(
                 from codeflash.code_utils.instrument_existing_tests import add_async_decorator_to_function
 
                 add_async_decorator_to_function(
-                    self.function_to_optimize.file_path, self.function_to_optimize, TestingMode.PERFORMANCE
+                    self.function_to_optimize.file_path,
+                    self.function_to_optimize,
+                    TestingMode.PERFORMANCE,
+                    project_root=self.project_root,
                 )
 
             try:
@@ -2974,7 +2986,10 @@ def run_concurrency_benchmark(
         try:
             # Add concurrency decorator to the source function
             add_async_decorator_to_function(
-                self.function_to_optimize.file_path, self.function_to_optimize, TestingMode.CONCURRENCY
+                self.function_to_optimize.file_path,
+                self.function_to_optimize,
+                TestingMode.CONCURRENCY,
+                project_root=self.project_root,
             )
 
             # Run the concurrency benchmark tests

tests/test_async_run_and_parse_tests.py

@@ -8,8 +8,9 @@
 import pytest
 
 from codeflash.code_utils.instrument_existing_tests import (
+    ASYNC_HELPER_FILENAME,
     add_async_decorator_to_function,
-    get_async_inline_code,
+    get_decorator_name_for_mode,
     inject_profiling_into_existing_test,
 )
 from codeflash.discovery.functions_to_optimize import FunctionToOptimize
@@ -56,19 +57,22 @@ async def test_async_sort():
         func = FunctionToOptimize(function_name="async_sorter", parents=[], file_path=Path(fto_path), is_async=True)
 
         # For async functions, instrument the source module directly with decorators
-        source_success = add_async_decorator_to_function(fto_path, func, TestingMode.BEHAVIOR)
+        source_success = add_async_decorator_to_function(
+            fto_path, func, TestingMode.BEHAVIOR, project_root=project_root_path
+        )
 
         assert source_success
 
         # Verify the file was modified with exact expected output
         instrumented_source = fto_path.read_text("utf-8")
         from codeflash.code_utils.formatter import sort_imports
 
-        inline_code = get_async_inline_code(TestingMode.BEHAVIOR)
+        decorator_name = get_decorator_name_for_mode(TestingMode.BEHAVIOR)
         decorated_original = original_code.replace(
-            "async def async_sorter", "@codeflash_behavior_async\nasync def async_sorter"
+            "async def async_sorter", f"@{decorator_name}\nasync def async_sorter"
         )
-        expected = sort_imports(code=inline_code + decorated_original, float_to_top=True)
+        code_with_import = f"from codeflash_async_wrapper import {decorator_name}\n{decorated_original}"
+        expected = sort_imports(code=code_with_import, float_to_top=True)
         assert instrumented_source.strip() == expected.strip()
 
         # Add codeflash capture
@@ -147,6 +151,9 @@ async def test_async_sort():
             test_path.unlink()
         if test_path_perf.exists():
             test_path_perf.unlink()
+        helper_path = project_root_path / ASYNC_HELPER_FILENAME
+        if helper_path.exists():
+            helper_path.unlink()
 
 
 @pytest.mark.skipif(sys.platform == "win32", reason="pending support for asyncio on windows")
@@ -187,7 +194,9 @@ async def test_async_class_sort():
             is_async=True,
         )
 
-        source_success = add_async_decorator_to_function(fto_path, func, TestingMode.BEHAVIOR)
+        source_success = add_async_decorator_to_function(
+            fto_path, func, TestingMode.BEHAVIOR, project_root=project_root_path
+        )
 
         assert source_success
 
@@ -269,6 +278,9 @@ async def test_async_class_sort():
             test_path.unlink()
         if test_path_perf.exists():
             test_path_perf.unlink()
+        helper_path = project_root_path / ASYNC_HELPER_FILENAME
+        if helper_path.exists():
+            helper_path.unlink()
 
 
 @pytest.mark.skipif(sys.platform == "win32", reason="pending support for asyncio on windows")
@@ -299,19 +311,22 @@ async def test_async_perf():
         func = FunctionToOptimize(function_name="async_sorter", parents=[], file_path=Path(fto_path), is_async=True)
 
         # Instrument the source module with async performance decorators
-        source_success = add_async_decorator_to_function(fto_path, func, TestingMode.PERFORMANCE)
+        source_success = add_async_decorator_to_function(
+            fto_path, func, TestingMode.PERFORMANCE, project_root=project_root_path
+        )
 
         assert source_success
 
         # Verify the file was modified
         instrumented_source = fto_path.read_text("utf-8")
         from codeflash.code_utils.formatter import sort_imports
 
-        inline_code = get_async_inline_code(TestingMode.PERFORMANCE)
+        decorator_name = get_decorator_name_for_mode(TestingMode.PERFORMANCE)
         decorated_original = original_code.replace(
-            "async def async_sorter", "@codeflash_performance_async\nasync def async_sorter"
+            "async def async_sorter", f"@{decorator_name}\nasync def async_sorter"
         )
-        expected = sort_imports(code=inline_code + decorated_original, float_to_top=True)
+        code_with_import = f"from codeflash_async_wrapper import {decorator_name}\n{decorated_original}"
+        expected = sort_imports(code=code_with_import, float_to_top=True)
         assert instrumented_source.strip() == expected.strip()
 
         instrument_codeflash_capture(func, {}, tests_root)
@@ -368,6 +383,9 @@ async def test_async_perf():
         # Clean up test files
         if test_path.exists():
             test_path.unlink()
+        helper_path = project_root_path / ASYNC_HELPER_FILENAME
+        if helper_path.exists():
+            helper_path.unlink()
 
 
 @pytest.mark.skipif(sys.platform == "win32", reason="pending support for asyncio on windows")
@@ -413,7 +431,9 @@ async def async_error_function(lst):
             function_name="async_error_function", parents=[], file_path=Path(fto_path), is_async=True
         )
 
-        source_success = add_async_decorator_to_function(fto_path, func, TestingMode.BEHAVIOR)
+        source_success = add_async_decorator_to_function(
+            fto_path, func, TestingMode.BEHAVIOR, project_root=project_root_path
+        )
 
         assert source_success
 
@@ -422,11 +442,12 @@ async def async_error_function(lst):
 
         from codeflash.code_utils.formatter import sort_imports
 
-        inline_code = get_async_inline_code(TestingMode.BEHAVIOR)
+        decorator_name = get_decorator_name_for_mode(TestingMode.BEHAVIOR)
         decorated_modified = modified_code.replace(
-            "async def async_error_function", "@codeflash_behavior_async\nasync def async_error_function"
+            "async def async_error_function", f"@{decorator_name}\nasync def async_error_function"
         )
-        expected = sort_imports(code=inline_code + decorated_modified, float_to_top=True)
+        code_with_import = f"from codeflash_async_wrapper import {decorator_name}\n{decorated_modified}"
+        expected = sort_imports(code=code_with_import, float_to_top=True)
         assert instrumented_source.strip() == expected.strip()
         instrument_codeflash_capture(func, {}, tests_root)
 
@@ -488,6 +509,9 @@ async def async_error_function(lst):
             test_path.unlink()
         if test_path_perf.exists():
             test_path_perf.unlink()
+        helper_path = project_root_path / ASYNC_HELPER_FILENAME
+        if helper_path.exists():
+            helper_path.unlink()
 
 
 @pytest.mark.skipif(sys.platform == "win32", reason="pending support for asyncio on windows")
@@ -525,7 +549,9 @@ async def test_async_multi():
 
         func = FunctionToOptimize(function_name="async_sorter", parents=[], file_path=Path(fto_path), is_async=True)
 
-        source_success = add_async_decorator_to_function(fto_path, func, TestingMode.BEHAVIOR)
+        source_success = add_async_decorator_to_function(
+            fto_path, func, TestingMode.BEHAVIOR, project_root=project_root_path
+        )
 
         assert source_success
         instrument_codeflash_capture(func, {}, tests_root)
@@ -598,6 +624,9 @@ async def test_async_multi():
             test_path.unlink()
         if test_path_perf.exists():
             test_path_perf.unlink()
+        helper_path = project_root_path / ASYNC_HELPER_FILENAME
+        if helper_path.exists():
+            helper_path.unlink()
 
 
 @pytest.mark.skipif(sys.platform == "win32", reason="pending support for asyncio on windows")
@@ -640,7 +669,9 @@ async def test_async_edge_cases():
 
         func = FunctionToOptimize(function_name="async_sorter", parents=[], file_path=Path(fto_path), is_async=True)
 
-        source_success = add_async_decorator_to_function(fto_path, func, TestingMode.BEHAVIOR)
+        source_success = add_async_decorator_to_function(
+            fto_path, func, TestingMode.BEHAVIOR, project_root=project_root_path
+        )
 
         assert source_success
         instrument_codeflash_capture(func, {}, tests_root)
@@ -715,6 +746,9 @@ async def test_async_edge_cases():
             test_path.unlink()
         if test_path_perf.exists():
             test_path_perf.unlink()
+        helper_path = project_root_path / ASYNC_HELPER_FILENAME
+        if helper_path.exists():
+            helper_path.unlink()
 
 
 @pytest.mark.skipif(sys.platform == "win32", reason="pending support for asyncio on windows")
@@ -949,7 +983,9 @@ async def test_mixed_sorting():
             function_name="async_merge_sort", parents=[], file_path=Path(mixed_fto_path), is_async=True
         )
 
-        source_success = add_async_decorator_to_function(mixed_fto_path, async_func, TestingMode.BEHAVIOR)
+        source_success = add_async_decorator_to_function(
+            mixed_fto_path, async_func, TestingMode.BEHAVIOR, project_root=project_root_path
+        )
 
         assert source_success
 
@@ -1022,3 +1058,6 @@ async def test_mixed_sorting():
             test_path.unlink()
         if test_path_perf.exists():
             test_path_perf.unlink()
+        helper_path = project_root_path / ASYNC_HELPER_FILENAME
+        if helper_path.exists():
+            helper_path.unlink()