refix

Author: KRRT7

codeflash/optimization/optimizer.py

@@ -409,6 +409,12 @@ def run(self) -> None:
         if self.args.worktree:
             self.worktree_mode()
 
+        if not self.args.replay_test and self.test_cfg.tests_root.exists():
+            leftover_trace_files = list(self.test_cfg.tests_root.glob("*.trace"))
+            if leftover_trace_files:
+                logger.debug(f"Cleaning up {len(leftover_trace_files)} leftover trace file(s) from previous runs")
+                cleanup_paths(leftover_trace_files)
+
         cleanup_paths(Optimizer.find_leftover_instrumented_test_files(self.test_cfg.tests_root))
 
         function_optimizer = None
@@ -519,6 +525,8 @@ def run(self) -> None:
                 )
             if self.functions_checkpoint:
                 self.functions_checkpoint.cleanup()
+            if hasattr(self.args, "command") and self.args.command == "optimize":
+                self.cleanup_replay_tests()
             if optimizations_found == 0:
                 logger.info("❌ No optimizations found.")
             elif self.args.all:
@@ -554,6 +562,17 @@ def find_leftover_instrumented_test_files(test_root: Path) -> list[Path]:
             file_path for file_path in test_root.rglob("*") if file_path.is_file() and pattern.match(file_path.name)
         ]
 
+    def cleanup_replay_tests(self) -> None:
+        paths_to_cleanup = []
+        if self.replay_tests_dir and self.replay_tests_dir.exists():
+            logger.debug(f"Cleaning up replay tests directory: {self.replay_tests_dir}")
+            paths_to_cleanup.append(self.replay_tests_dir)
+        if self.trace_file and self.trace_file.exists():
+            logger.debug(f"Cleaning up trace file: {self.trace_file}")
+            paths_to_cleanup.append(self.trace_file)
+        if paths_to_cleanup:
+            cleanup_paths(paths_to_cleanup)
+
     def cleanup_temporary_paths(self) -> None:
         if hasattr(get_run_tmp_file, "tmpdir"):
             get_run_tmp_file.tmpdir.cleanup()
@@ -568,6 +587,14 @@ def cleanup_temporary_paths(self) -> None:
 
         if self.current_function_optimizer:
             self.current_function_optimizer.cleanup_generated_files()
+        paths_to_cleanup = [self.replay_tests_dir]
+        if self.trace_file:
+            paths_to_cleanup.append(self.trace_file)
+        if self.test_cfg.tests_root.exists():
+            for trace_file in self.test_cfg.tests_root.glob("*.trace"):
+                if trace_file not in paths_to_cleanup:
+                    paths_to_cleanup.append(trace_file)
+        cleanup_paths(paths_to_cleanup)
 
     def worktree_mode(self) -> None:
         if self.current_worktree:

codeflash/tracer.py

@@ -86,6 +86,7 @@ def main(args: Namespace | None = None) -> ArgumentParser:
     # to the output file at startup.
     if parsed_args.outfile is not None:
         parsed_args.outfile = Path(parsed_args.outfile).resolve()
+    outfile = parsed_args.outfile
     config, found_config_path = parse_config_file(parsed_args.codeflash_config)
     project_root = project_root_from_module_root(Path(config["module_root"]), found_config_path)
     if len(unknown_args) > 0:
@@ -215,6 +216,12 @@ def main(args: Namespace | None = None) -> ArgumentParser:
 
                 optimizer.run_with_args(args)
 
+                # Delete the trace file and the replay test file if they exist
+                if outfile:
+                    outfile.unlink(missing_ok=True)
+                for replay_test_path in replay_test_paths:
+                    Path(replay_test_path).unlink(missing_ok=True)
+
         except BrokenPipeError as exc:
             # Prevent "Exception ignored" during interpreter shutdown.
             sys.stdout = None