diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py
index acae4b9fb..82cf4bc57 100644
--- a/codeflash/optimization/function_optimizer.py
+++ b/codeflash/optimization/function_optimizer.py
@@ -1457,11 +1457,9 @@ def submit_test_generation_tasks(
         ]
 
     def cleanup_generated_files(self) -> None:
-        paths_to_cleanup = [self.test_cfg.concolic_test_root_dir]
+        paths_to_cleanup = []
         for test_file in self.test_files:
             paths_to_cleanup.append(test_file.instrumented_behavior_file_path)
             paths_to_cleanup.append(test_file.benchmarking_file_path)
 
         cleanup_paths(paths_to_cleanup)
-        if hasattr(get_run_tmp_file, "tmpdir"):
-            get_run_tmp_file.tmpdir.cleanup()
diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py
index 70ba8cc12..819e61264 100644
--- a/codeflash/optimization/optimizer.py
+++ b/codeflash/optimization/optimizer.py
@@ -12,7 +12,7 @@
 from codeflash.api.aiservice import AiServiceClient, LocalAiServiceClient
 from codeflash.cli_cmds.console import console, logger, progress_bar
 from codeflash.code_utils import env_utils
-from codeflash.code_utils.code_utils import cleanup_paths
+from codeflash.code_utils.code_utils import cleanup_paths, get_run_tmp_file
 from codeflash.code_utils.env_utils import get_pr_number
 from codeflash.either import is_successful
 from codeflash.models.models import ValidCode
@@ -289,31 +289,36 @@ def run(self) -> None:
                         f"{function_to_optimize.qualified_name}"
                     )
                     console.rule()
+                    function_optimizer = None
+                    try:
+                        function_optimizer = self.create_function_optimizer(
+                            function_to_optimize,
+                            function_to_tests=function_to_tests,
+                            function_to_optimize_source_code=validated_original_code[original_module_path].source_code,
+                            function_benchmark_timings=function_benchmark_timings,
+                            total_benchmark_timings=total_benchmark_timings,
+                            original_module_ast=original_module_ast,
+                            original_module_path=original_module_path,
+                        )
 
-                    function_optimizer = self.create_function_optimizer(
-                        function_to_optimize,
-                        function_to_tests=function_to_tests,
-                        function_to_optimize_source_code=validated_original_code[original_module_path].source_code,
-                        function_benchmark_timings=function_benchmark_timings,
-                        total_benchmark_timings=total_benchmark_timings,
-                        original_module_ast=original_module_ast,
-                        original_module_path=original_module_path,
-                    )
-
-                    self.current_function_optimizer = (
-                        function_optimizer  # needed to clean up from the outside of this function
-                    )
-                    best_optimization = function_optimizer.optimize_function()
-                    if self.functions_checkpoint:
-                        self.functions_checkpoint.add_function_to_checkpoint(
-                            function_to_optimize.qualified_name_with_modules_from_root(self.args.project_root)
+                        self.current_function_optimizer = (
+                            function_optimizer  # needed to clean up from the outside of this function
                         )
-                    if is_successful(best_optimization):
-                        optimizations_found += 1
-                    else:
-                        logger.warning(best_optimization.failure())
-                        console.rule()
-                        continue
+                        best_optimization = function_optimizer.optimize_function()
+                        if self.functions_checkpoint:
+                            self.functions_checkpoint.add_function_to_checkpoint(
+                                function_to_optimize.qualified_name_with_modules_from_root(self.args.project_root)
+                            )
+                        if is_successful(best_optimization):
+                            optimizations_found += 1
+                        else:
+                            logger.warning(best_optimization.failure())
+                            console.rule()
+                            continue
+                    finally:
+                        if function_optimizer is not None:
+                            function_optimizer.cleanup_generated_files()
+
             ph("cli-optimize-run-finished", {"optimizations_found": optimizations_found})
             if self.functions_checkpoint:
                 self.functions_checkpoint.cleanup()
@@ -351,6 +356,9 @@ def cleanup_temporary_paths(self) -> None:
         if self.current_function_optimizer:
             self.current_function_optimizer.cleanup_generated_files()
 
+        if hasattr(get_run_tmp_file, "tmpdir"):
+            get_run_tmp_file.tmpdir.cleanup()
+            del get_run_tmp_file.tmpdir
         cleanup_paths([self.test_cfg.concolic_test_root_dir, self.replay_tests_dir])