From 6b50670f2ec07ee248eb17b945da8e4a54983199 Mon Sep 17 00:00:00 2001
From: XianBW <36835909+XianBW@users.noreply.github.com>
Date: Wed, 6 Aug 2025 15:43:28 +0800
Subject: [PATCH] fix loop step time log problem (#1158)

---
 rdagent/utils/workflow/loop.py | 20 ++++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/rdagent/utils/workflow/loop.py b/rdagent/utils/workflow/loop.py
index 5f43a02d..577ad5d4 100644
--- a/rdagent/utils/workflow/loop.py
+++ b/rdagent/utils/workflow/loop.py
@@ -235,16 +235,6 @@ class LoopBase:
                     # Store result in the nested dictionary
                     self.loop_prev_out[li][name] = result
 
-                    # Record the trace
-                    end = datetime.now(timezone.utc)
-                    self.loop_trace[li].append(LoopTrace(start, end, step_idx=si))
-                    logger.log_object(
-                        {
-                            "start_time": start,
-                            "end_time": end,
-                        },
-                        tag="time_info",
-                    )
                     # Save snapshot after completing the step
                     self.dump(self.session_folder / f"{li}" / f"{si}_{name}")
                 except Exception as e:
@@ -264,6 +254,16 @@ class LoopBase:
                     else:
                         raise  # re-raise unhandled exceptions
                 finally:
+                    # Record the trace
+                    end = datetime.now(timezone.utc)
+                    self.loop_trace[li].append(LoopTrace(start, end, step_idx=si))
+                    logger.log_object(
+                        {
+                            "start_time": start,
+                            "end_time": end,
+                        },
+                        tag="time_info",
+                    )
                     if step_forward:
                         # Increment step index
                         self.step_idx[li] = next_step_idx