fix code revert after optimization

Author: mohammedahmed18

codeflash/optimization/function_optimizer.py

@@ -1189,7 +1189,9 @@ def process_review(
             "coverage_message": coverage_message,
         }
 
-        if not self.args.no_pr and not self.args.staging_review:
+        raise_pr = not self.args.no_pr
+
+        if raise_pr and not self.args.staging_review:
             data["git_remote"] = self.args.git_remote
             check_create_pr(**data)
         elif self.args.staging_review:
@@ -1200,12 +1202,24 @@ def process_review(
                 trace_id=self.function_trace_id, is_optimization_found=best_optimization is not None
             )
 
-        if ((not self.args.no_pr) or not self.args.staging_review) and (
-            self.args.all or env_utils.get_pr_number() or (self.args.file and not self.args.function)
+        if raise_pr and (
+            self.args.all
+            or env_utils.get_pr_number()
+            or self.args.replay_test
+            or (self.args.file and not self.args.function)
         ):
-            self.write_code_and_helpers(
-                self.function_to_optimize_source_code, original_helper_code, self.function_to_optimize.file_path
-            )
+            self.revert_code_and_helpers(original_helper_code)
+            return
+
+        if self.args.staging_review:
+            # always revert code and helpers when staging review
+            self.revert_code_and_helpers(original_helper_code)
+            return
+
+    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
+        )
 
     def establish_original_code_baseline(
         self,