[FIX] Ensure cleanup after each candidate testing (CF-724)

[mohammedahmed18]

If a continue statement is hit inside the try block, the files are not reset, causing the next candidate’s code to be compared against the previous candidate instead of the original code.

---

PR Type

Bug fix

---

Description

• Always reset code after each candidate

• Move cleanup into finally block

• Improve interrupt logging and re-raise

---

Diagram Walkthrough

flowchart LR
  TryBlock["Candidate evaluation try"] -- may raise --> Except["KeyboardInterrupt handler"]
  TryBlock -- always --> Finally["Finally: reset code/helpers"]
  Except -- log & re-raise --> Propagate["Propagation"]
  Finally -- prepare --> NextCandidate["Next candidate run"]

Loading

File Walkthrough

	Relevant files
Bug fix	function_optimizer.py Always reset state via finally; improve interrupt handling codeflash/optimization/function_optimizer.py Move code/helper reset into a finally block. Add exception logging for KeyboardInterrupt, then re-raise. Remove duplicated reset from normal flow/except path. +4/-6

---

[github-actions]

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

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

🧪 No relevant tests

🔒 No security concerns identified

⚡ Recommended focus areas for review Logging Level Using logger.exception in a KeyboardInterrupt path will include a full traceback; consider whether logger.info/warning is preferable to avoid noisy logs for intentional interrupts. logger.exception(f"Optimization interrupted: {e}") raise Cleanup Timing Ensure write_code_and_helpers in finally does not overwrite persisted best-candidate changes if later code expects to keep an accepted optimization applied beyond the candidate loop. finally: # reset for the next candidate self.write_code_and_helpers( self.function_to_optimize_source_code, original_helper_code, self.function_to_optimize.file_path )

[github-actions]

PR Code Suggestions ✨

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Make cleanup failure-safe Guard the cleanup to ensure it never raises and masks the original error or interrupt. Wrap the write in its own try/except and log failures, so KeyboardInterrupt or other exceptions aren't hidden by cleanup errors. codeflash/optimization/function_optimizer.py [655-659] finally: # reset for the next candidate - self.write_code_and_helpers( - self.function_to_optimize_source_code, original_helper_code, self.function_to_optimize.file_path - ) + try: + self.write_code_and_helpers( + self.function_to_optimize_source_code, original_helper_code, self.function_to_optimize.file_path + ) + except Exception as ce: + logger.error(f"Cleanup failed while restoring original code: {ce}") Suggestion importance[1-10]: 7 __ Why: Wrapping cleanup in a try/except prevents masking the original KeyboardInterrupt or other exceptions, improving robustness. It's accurate to the new hunk and beneficial, though not a critical bug fix.	Medium
General	Avoid misleading exception logs Logging inside except and then re-raising will also trigger the finally, which still performs cleanup—good—but the current message uses logger.exception for an expected interrupt. Use logger.info (or warning) in the except to avoid noisy stack traces, and keep finally for cleanup. This preserves cleanup semantics without misleading error logs. codeflash/optimization/function_optimizer.py [652-659] except KeyboardInterrupt as e: - logger.exception(f"Optimization interrupted: {e}") + logger.info(f"Optimization interrupted by user: {e}") raise finally: # reset for the next candidate self.write_code_and_helpers( self.function_to_optimize_source_code, original_helper_code, self.function_to_optimize.file_path ) Suggestion importance[1-10]: 6 __ Why: The existing_code exactly matches the new hunk and the change from logger.exception to logger.info reduces noisy stack traces for an expected KeyboardInterrupt without altering control flow. It's a reasonable logging-level improvement but not critical.	Low