From 543fb2c1efd7a73f136df407f192e944d42a53be Mon Sep 17 00:00:00 2001
From: Tim <illking@foxmail.com>
Date: Wed, 4 Jun 2025 18:30:29 +0800
Subject: [PATCH] fix: ds trace (#929)

* fix: remove_ansi_codes
* chore: rename function
---
 rdagent/log/ui/ds_trace.py | 4 ++--
 rdagent/utils/workflow.py  | 8 ++++----
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/rdagent/log/ui/ds_trace.py b/rdagent/log/ui/ds_trace.py
index 22cf4d8e..dd3aeff5 100644
--- a/rdagent/log/ui/ds_trace.py
+++ b/rdagent/log/ui/ds_trace.py
@@ -14,12 +14,12 @@ from rdagent.app.data_science.loop import DataScienceRDLoop
 from rdagent.log.storage import FileStorage
 from rdagent.log.ui.utils import load_times
 from rdagent.log.utils import (
+    LogColors,
     extract_evoid,
     extract_json,
     extract_loopid_func_name,
     is_valid_session,
 )
-from rdagent.utils import remove_ansi_codes
 from rdagent.utils.repo.diff import generate_diff_from_dict
 
 if "show_stdout" not in state:
@@ -598,7 +598,7 @@ def stdout_win(loop_id: int):
         return
     start_index = stdout.find(f"Start Loop {loop_id}")
     end_index = stdout.find(f"Start Loop {loop_id + 1}")
-    loop_stdout = remove_ansi_codes(stdout[start_index:end_index])
+    loop_stdout = LogColors.remove_ansi_codes(stdout[start_index:end_index])
     with st.container(border=True):
         st.subheader(f"Loop {loop_id} stdout")
         pattern = f"Start Loop {loop_id}, " + r"Step \d+: \w+"
diff --git a/rdagent/utils/workflow.py b/rdagent/utils/workflow.py
index 1081d9ab..5f9fa215 100644
--- a/rdagent/utils/workflow.py
+++ b/rdagent/utils/workflow.py
@@ -189,7 +189,7 @@ class LoopBase:
                         elif isinstance(e, self.withdraw_loop_error):
                             logger.warning(f"Withdraw loop {li} due to {e}")
                             # Back to previous loop
-                            self.step_backward(li - 1)
+                            self.withdraw_loop(li)
                             continue
                         else:
                             raise
@@ -215,8 +215,8 @@ class LoopBase:
 
                 self.dump(self.session_folder / f"{li}" / f"{si}_{name}")  # save a snapshot after the session
 
-    def step_backward(self, li: int) -> None:
-        prev_session_dir = self.session_folder / str(li)
+    def withdraw_loop(self, loop_idx: int) -> None:
+        prev_session_dir = self.session_folder / str(loop_idx - 1)
         prev_path = min(
             (p for p in prev_session_dir.glob("*_*") if p.is_file()),
             key=lambda item: int(item.name.split("_", 1)[0]),
@@ -233,7 +233,7 @@ class LoopBase:
             # Overwrite current instance state
             self.__dict__ = loaded.__dict__
         else:
-            logger.error(f"No previous dump found at {prev_session_dir}, cannot withdraw loop {li}")
+            logger.error(f"No previous dump found at {prev_session_dir}, cannot withdraw loop {loop_idx}")
             raise
 
     def dump(self, path: str | Path) -> None: