refactor how async decorators are applied at source site

[KRRT7]

User description

this is a superficial refactor, all the logic of the instrumentation for the deco @ call site is the same, my original implementation was awkward, this is in preparation for the synchronous portion;

following up with runtime timings fixes for async

---

PR Type

Enhancement, Tests

---

Description

• Refactor async decorator application to in-place

• Prioritize async throughput in reporting

• Include async throughput in PR comments

• Adjust thresholds and skip async opt review

---

Diagram Walkthrough

flowchart LR A["instrument_existing_tests.add_async_decorator_to_function (in-place)"] -- "used by" --> B["function_optimizer establish/run candidate"] C["PrComment supports async throughput"] -- "adds fields" --> D["PR JSON output"] E["Explanation"] -- "prefers throughput metric" --> F["perf_improvement_line/speedup"] G["config_consts"] -- "lower threshold" --> H["MIN_THROUGHPUT_IMPROVEMENT_THRESHOLD 0.05"] I["Tests"] -- "updated for in-place API" --> A 

Loading

File Walkthrough

	Relevant files
Enhancement	7 files config_consts.py Lower async throughput improvement threshold to 5%              +1/-1      instrument_existing_tests.py Rewrite async decorator API to write file in-place              +19/-29  PrComment.py Add async throughput fields to PR JSON report                        +9/-1      function_optimizer.py Prefer throughput for async metrics and PR flow tweaks      +51/-59  create_pr.py Pass async throughput metrics into PR creation                      +4/-0      explanation.py Centralize speedup logic; throughput-aware improvement text +10/-25  concolic_testing.py Restrict concolic generation to sync functions                      +1/-1
Tests	2 files test_async_run_and_parse_tests.py Update async tests for in-place instrumentation API            +23/-26  test_instrument_async_tests.py Adapt unit tests to file-writing decorator and temp files +71/-58

---

[github-actions]

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

⏱️ Estimated effort to review: 3 🔵🔵🔵⚪⚪

🧪 PR contains tests

🔒 No security concerns identified

⚡ Recommended focus areas for review API Change Risk The signature and behavior of add_async_decorator_to_function changed from returning modified code to performing in-place writes and returning a boolean. Verify all call sites outside the shown diff are updated to the new API to avoid breakages. source_path: Path, function: FunctionToOptimize, mode: TestingMode = TestingMode.BEHAVIOR ) -> 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. """ if not function.is_async: return False try: # Read source code with source_path.open(encoding="utf8") as f: source_code = f.read() module = cst.parse_module(source_code) # Add the decorator to the 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 = AsyncDecoratorImportAdder(mode) module = module.visit(import_transformer) modified_code = sort_imports(code=module.code, float_to_top=True) except Exception as e: logger.exception(f"Error adding async decorator to function {function.qualified_name}: {e}") return False else: if decorator_transformer.added_decorator: with source_path.open("w", encoding="utf8") as f: f.write(modified_code) logger.debug(f"Applied async {mode.value} instrumentation to {source_path}") return True return False Semantics Change perf_improvement_line now delegates to speedup_pct/speedup_x, while speedup gained logic to prefer throughput. Confirm that all consumers expect speedup to sometimes reflect throughput rather than runtime and that formatting remains consistent across UI/PR comments. # speedup property already handles choosing between runtime and throughput return f"{self.speedup_pct} improvement ({self.speedup_x} faster)." @property def speedup(self) -> float: runtime_improvement = (self.original_runtime_ns / self.best_runtime_ns) - 1 # Use throughput improvement if we have async metrics and throughput is better if ( self.original_async_throughput is not None and self.best_async_throughput is not None and self.original_async_throughput > 0 ): throughput_improvement = throughput_gain( original_throughput=self.original_async_throughput, optimized_throughput=self.best_async_throughput ) # Use throughput metrics if throughput improvement is better or runtime got worse if throughput_improvement > runtime_improvement or runtime_improvement <= 0: return throughput_improvement return runtime_improvement @property def speedup_x(self) -> str: return f"{self.speedup:,.2f}x" Feature Regression Concolic test generation condition now excludes ast.AsyncFunctionDef. Ensure async targets intentionally skip concolic generation, or restore support if this was not intended. if ( test_cfg.concolic_test_root_dir and isinstance(function_to_optimize_ast, ast.FunctionDef) and has_typed_parameters(function_to_optimize_ast, function_to_optimize.parents) ):

[github-actions]

PR Code Suggestions ✨

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Avoid unintended async coverage loss This change excludes async functions (ast.AsyncFunctionDef) from concolic test generation, which likely reduces coverage for async targets. If intentional, guard with the function’s is_async to keep previous behavior for sync functions while explicitly skipping only when async. codeflash/verification/concolic_testing.py [38-44] if ( test_cfg.concolic_test_root_dir - and isinstance(function_to_optimize_ast, ast.FunctionDef) + and isinstance(function_to_optimize_ast, (ast.FunctionDef, ast.AsyncFunctionDef)) + and (not function_to_optimize.is_async) and has_typed_parameters(function_to_optimize_ast, function_to_optimize.parents) ): logger.info("Generating concolic opcode coverage tests for the original code…") console.rule() Suggestion importance[1-10]: 7 __ Why: The PR narrowed concolic test generation to only FunctionDef. Restoring AsyncFunctionDef while explicitly skipping when is_async preserves sync behavior and makes the intent clear; it’s a reasonable correction to prevent accidental coverage regression.	Medium
	Ensure PRs for async proceed When skipping optimization review for async functions, opt_review_response remains empty and bypasses PR creation due to the != "low" check. Explicitly set a non-blocking default (e.g., "ok") for async paths so valid PRs are not accidentally suppressed. codeflash/optimization/function_optimizer.py [1504-1529] # Skip optimization review for async functions for now if (raise_pr or staging_review) and not self.function_to_optimize.is_async: data["root_dir"] = git_root_dir() try: opt_review_response = self.aiservice_client.get_optimization_review( **data, calling_fn_details=function_references ) except Exception as e: logger.debug(f"optimization review response failed, investigate {e}") -# Always set optimization_review in data (empty string for async functions) +else: + # Provide a neutral default so PR creation isn't suppressed for async functions + opt_review_response = "ok" +# Always set optimization_review in data data["optimization_review"] = opt_review_response if raise_pr and not staging_review and opt_review_response != "low": - # Ensure root_dir is set for PR creation (needed for async functions that skip opt_review) if "root_dir" not in data: data["root_dir"] = git_root_dir() data["git_remote"] = self.args.git_remote check_create_pr(**data) Suggestion importance[1-10]: 5 __ Why: Setting a neutral default like "ok" ensures async paths don't block PR creation due to an empty opt_review_response. This is a small logic tweak with potential workflow impact; however, changing review semantics may be policy-dependent, so impact is moderate.	Low
General	Prevent unnecessary file writes Avoid writing back to disk when no changes are needed to prevent unnecessary file churn and timestamp updates. Also, short-circuit early if the computed modified_code is identical to the original to skip writes even when added_decorator is true due to formatting-only changes. codeflash/code_utils/instrument_existing_tests.py [1271-1315] def add_async_decorator_to_function( source_path: Path, function: FunctionToOptimize, mode: TestingMode = TestingMode.BEHAVIOR ) -> bool: - """Add async decorator to an async function definition and write back to file. -... + """Add async decorator to an async function definition and write back to file.""" if not function.is_async: return False try: # Read source code with source_path.open(encoding="utf8") as f: - source_code = f.read() -... + original_source = f.read() + + module = cst.parse_module(original_source) + + # Add the decorator to the 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 = AsyncDecoratorImportAdder(mode) + module = module.visit(import_transformer) + modified_code = sort_imports(code=module.code, float_to_top=True) except Exception as e: logger.exception(f"Error adding async decorator to function {function.qualified_name}: {e}") return False else: - if decorator_transformer.added_decorator: - with source_path.open("w", encoding="utf8") as f: - f.write(modified_code) - logger.debug(f"Applied async {mode.value} instrumentation to {source_path}") - return True - return False + if not decorator_transformer.added_decorator: + return False + if modified_code == original_source: + return False + with source_path.open("w", encoding="utf8") as f: + f.write(modified_code) + logger.debug(f"Applied async {mode.value} instrumentation to {source_path}") + return True Suggestion importance[1-10]: 6 __ Why: The suggestion correctly targets the new add_async_decorator_to_function implementation to avoid no-op writes by comparing modified_code with the original; this can reduce file churn. It’s a maintainability optimization, not critical, but accurate and aligned with the PR’s changes.	Low