fix: clean up async helper file and combine all decorators into single file

Write all three async decorator implementations into one helper file to
avoid overwrite issues when switching modes. Clean up the helper file in
revert_code_and_helpers and early-exit paths so it doesn't persist in the
user's project root after optimization.

Author: KRRT7

codeflash/code_utils/instrument_existing_tests.py

@@ -1497,11 +1497,11 @@ def _is_target_decorator(self, decorator_node: cst.Name | cst.Attribute | cst.Ca
         return False
 
 
-def get_behavior_async_inline_code() -> str:
-    return """import asyncio
+ASYNC_HELPER_INLINE_CODE = """import asyncio
 import gc
 import os
 import sqlite3
+import time
 from functools import wraps
 from pathlib import Path
 from tempfile import TemporaryDirectory
@@ -1590,25 +1590,6 @@ async def async_wrapper(*args, **kwargs):
             raise exception
         return return_value
     return async_wrapper
-"""
-
-
-def get_performance_async_inline_code() -> str:
-    return """import asyncio
-import gc
-import os
-from functools import wraps
-
-
-def extract_test_context_from_env():
-    test_module = os.environ["CODEFLASH_TEST_MODULE"]
-    test_class = os.environ.get("CODEFLASH_TEST_CLASS", None)
-    test_function = os.environ["CODEFLASH_TEST_FUNCTION"]
-    if test_module and test_function:
-        return (test_module, test_class if test_class else None, test_function)
-    raise RuntimeError(
-        "Test context environment variables not set - ensure tests are run through codeflash test runner"
-    )
 
 
 def codeflash_performance_async(func):
@@ -1649,15 +1630,6 @@ async def async_wrapper(*args, **kwargs):
             raise exception
         return return_value
     return async_wrapper
-"""
-
-
-def get_concurrency_async_inline_code() -> str:
-    return """import asyncio
-import gc
-import os
-import time
-from functools import wraps
 
 
 def codeflash_concurrency_async(func):
@@ -1691,15 +1663,6 @@ async def async_wrapper(*args, **kwargs):
     return async_wrapper
 """
 
-
-def get_async_inline_code(mode: TestingMode) -> str:
-    if mode == TestingMode.BEHAVIOR:
-        return get_behavior_async_inline_code()
-    if mode == TestingMode.CONCURRENCY:
-        return get_concurrency_async_inline_code()
-    return get_performance_async_inline_code()
-
-
 ASYNC_HELPER_FILENAME = "codeflash_async_wrapper.py"
 
 
@@ -1711,14 +1674,11 @@ def get_decorator_name_for_mode(mode: TestingMode) -> str:
     return "codeflash_performance_async"
 
 
-def write_async_helper_file(target_dir: Path, mode: TestingMode) -> Path:
+def write_async_helper_file(target_dir: Path) -> 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")
+    if not helper_path.exists():
+        helper_path.write_text(ASYNC_HELPER_INLINE_CODE, "utf-8")
     return helper_path
 
 
@@ -1750,7 +1710,7 @@ def add_async_decorator_to_function(
         if decorator_transformer.added_decorator:
             # 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)
+            write_async_helper_file(helper_dir)
             # 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}")

codeflash/optimization/function_optimizer.py

@@ -1897,6 +1897,7 @@ def setup_and_establish_baseline(
             if self.args.override_fixtures:
                 restore_conftest(original_conftest_content)
             cleanup_paths(paths_to_cleanup)
+            self.cleanup_async_helper_file()
             return Failure(baseline_result.failure())
 
         original_code_baseline, test_functions_to_remove = baseline_result.unwrap()
@@ -1908,6 +1909,7 @@ def setup_and_establish_baseline(
             if self.args.override_fixtures:
                 restore_conftest(original_conftest_content)
             cleanup_paths(paths_to_cleanup)
+            self.cleanup_async_helper_file()
             return Failure("The threshold for test confidence was not met.")
 
         return Success(
@@ -2279,6 +2281,13 @@ def revert_code_and_helpers(self, original_helper_code: dict[Path, str]) -> None
         self.write_code_and_helpers(
             self.function_to_optimize_source_code, original_helper_code, self.function_to_optimize.file_path
         )
+        self.cleanup_async_helper_file()
+
+    def cleanup_async_helper_file(self) -> None:
+        from codeflash.code_utils.instrument_existing_tests import ASYNC_HELPER_FILENAME
+
+        helper_path = self.project_root / ASYNC_HELPER_FILENAME
+        helper_path.unlink(missing_ok=True)
 
     def establish_original_code_baseline(
         self,