diff --git a/ui/easydiffusion/task_manager.py b/ui/easydiffusion/task_manager.py
index 31fdaa6f..91adc04b 100644
--- a/ui/easydiffusion/task_manager.py
+++ b/ui/easydiffusion/task_manager.py
@@ -317,6 +317,9 @@ def thread_render(device):
             def step_callback():
                 global current_state_error
 
+                task_cache.keep(id(task), TASK_TTL)
+                session_cache.keep(task.task_data.session_id, TASK_TTL)
+
                 if (
                     isinstance(current_state_error, SystemExit)
                     or isinstance(current_state_error, StopAsyncIteration)