feedback

Author: KRRT7

codeflash.code-workspace

@@ -70,7 +70,7 @@
                 "request": "launch",
                 "program": "${workspaceFolder:codeflash}/codeflash/main.py",
                 "args": [
-                    "--all",
+                    "--file", "src/async_examples/shocker.py", "--verbose"
                 ],
                 "cwd": "${input:chooseCwd}",
                 "console": "integratedTerminal",

codeflash/code_utils/codeflash_wrap_decorator.py

@@ -108,12 +108,12 @@ async def async_wrapper(*args: Any, **kwargs: Any) -> Any:  # noqa: ANN401
 
 
 def codeflash_performance_async(func: F) -> F:
-    function_name = func.__name__
-    line_id = f"{func.__name__}_{func.__code__.co_firstlineno}"
-    loop_index = int(os.environ.get("CODEFLASH_LOOP_INDEX", "1"))
-
     @wraps(func)
     async def async_wrapper(*args: Any, **kwargs: Any) -> Any:  # noqa: ANN401
+        function_name = func.__name__
+        line_id = f"{func.__name__}_{func.__code__.co_firstlineno}"
+        loop_index = os.environ["CODEFLASH_LOOP_INDEX"]
+
         test_module_name, test_class_name, test_name = extract_test_context_from_frame()
 
         test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{loop_index}"

codeflash/optimization/function_optimizer.py

@@ -1541,11 +1541,25 @@ def run_optimized_candidate(
             candidate_helper_code = {}
             for module_abspath in original_helper_code:
                 candidate_helper_code[module_abspath] = Path(module_abspath).read_text("utf-8")
+            if self.function_to_optimize.is_async:
+                from codeflash.code_utils.instrument_existing_tests import (
+                    instrument_source_module_with_async_decorators,
+                )
+
+                success, instrumented_source = instrument_source_module_with_async_decorators(
+                    self.function_to_optimize.file_path, self.function_to_optimize, TestingMode.BEHAVIOR
+                )
+                if success and instrumented_source:
+                    with self.function_to_optimize.file_path.open("w", encoding="utf8") as f:
+                        f.write(instrumented_source)
+                    logger.debug(
+                        f"Applied async behavioral instrumentation to {self.function_to_optimize.file_path} for candidate {optimization_candidate_index}"
+                    )
+
             try:
                 instrument_codeflash_capture(
                     self.function_to_optimize, file_path_to_helper_classes, self.test_cfg.tests_root
                 )
-
                 candidate_behavior_results, _ = self.run_and_parse_tests(
                     testing_type=TestingMode.BEHAVIOR,
                     test_env=test_env,

tests/test_instrument_async_tests.py

@@ -248,13 +248,11 @@ async def test_async_function():
         func,
         temp_dir,
         "pytest",
-        mode=TestingMode.BEHAVIOR,
-        source_module_path=source_file
+        mode=TestingMode.BEHAVIOR
     )
 
-    assert success
-    assert instrumented_test_code is not None
-    assert 'codeflash_wrap(' in instrumented_test_code
+    assert not success
+    assert instrumented_test_code is None
 
 
 def test_inject_profiling_async_function_performance_mode(temp_dir):
@@ -316,14 +314,11 @@ async def test_async_function():
         func,
         temp_dir,
         "pytest",
-        mode=TestingMode.PERFORMANCE,
-        source_module_path=source_file
+        mode=TestingMode.PERFORMANCE
     )
 
-    assert success
-    assert instrumented_test_code is not None
-    # Test should use wrapper calls but source module should have async decorators
-    assert 'codeflash_wrap(' in instrumented_test_code
+    assert not success
+    assert instrumented_test_code is None
 
 
 def test_mixed_sync_async_instrumentation(temp_dir):
@@ -390,14 +385,12 @@ async def test_mixed_functions():
         async_func,
         temp_dir,
         "pytest",
-        mode=TestingMode.BEHAVIOR,
-        source_module_path=source_file
+        mode=TestingMode.BEHAVIOR
     )
 
-    assert success
-    assert instrumented_test_code is not None
-    # Both sync and async function calls should use wrappers in test file
-    assert 'codeflash_wrap(' in instrumented_test_code
+    # Async functions should not be instrumented at the test level
+    assert not success
+    assert instrumented_test_code is None