diff --git a/skyrl-train/examples/terminal_bench/terminal_bench_generator.py b/skyrl-train/examples/terminal_bench/terminal_bench_generator.py
index e0fa6648b4..3ce044bd62 100644
--- a/skyrl-train/examples/terminal_bench/terminal_bench_generator.py
+++ b/skyrl-train/examples/terminal_bench/terminal_bench_generator.py
@@ -159,6 +159,7 @@ async def terminal_bench_agent_loop(
                         "key": "fake_key",
                         "max_episodes": self.max_episodes,
                         "session_id": session_id,
+                        "store_all_messages": True,
                         "enable_summarize": self.enable_summarize,
                     },
                 ),
@@ -176,7 +177,7 @@ async def terminal_bench_agent_loop(
         else:
             raise ValueError(f"Invalid agent name: {self.agent_name}")
 
-        trial = Trial(trial_config)
+        
 
         # Run the trial to get `rewards`, `chat_history`, and `summarization_count`
         successful = False
@@ -184,6 +185,7 @@ async def terminal_bench_agent_loop(
         chat_history = None
         summarization_count = None
         for i in range(MAX_NUM_RETRIES_PER_TRIAL):
+            trial = Trial(trial_config)
             prefix = f"Trajectory {trajectory_id} attempt {i+1}/{MAX_NUM_RETRIES_PER_TRIAL}"
             results = None
             try: