diff --git a/src/airflow/model/common/gantt.rs b/src/airflow/model/common/gantt.rs index 13c7641e..c992a0a3 100644 --- a/src/airflow/model/common/gantt.rs +++ b/src/airflow/model/common/gantt.rs @@ -80,7 +80,7 @@ impl GanttData { } /// Recalculate `window_start` and `window_end` from all tries. - fn recompute_window(&mut self) { + pub fn recompute_window(&mut self) { let mut min_start: Option = None; let mut max_end: Option = None; let mut any_running = false; diff --git a/src/app/model/taskinstances.rs b/src/app/model/taskinstances.rs index ca62a96f..fcc41227 100644 --- a/src/app/model/taskinstances.rs +++ b/src/app/model/taskinstances.rs @@ -11,7 +11,8 @@ use ratatui::widgets::{Block, BorderType, Borders, Row, StatefulWidget, Table, W use crate::airflow::graph::{sort_task_instances, TaskGraph}; use crate::airflow::model::common::{ - calculate_duration, format_duration, GanttData, TaskId, TaskInstance, TaskInstanceState, + calculate_duration, format_duration, DagId, DagRunId, GanttData, TaskId, TaskInstance, + TaskInstanceState, }; use crate::app::events::custom::FlowrsEvent; use crate::ui::common::{create_headers, state_to_colored_square}; @@ -32,6 +33,9 @@ pub struct TaskInstanceModel { pub popup: Popup, /// Gantt chart data computed from task instances and their tries pub gantt_data: GanttData, + /// Tracks which DAG + run the cached `gantt_data` belongs to, so we can + /// invalidate it when the user navigates to a different DAG or run. + current_gantt_key: Option<(DagId, DagRunId)>, ticks: u32, event_buffer: Vec, pub task_graph: Option, @@ -43,6 +47,7 @@ impl Default for TaskInstanceModel { table: FilterableTable::new(), popup: Popup::None, gantt_data: GanttData::default(), + current_gantt_key: None, ticks: 0, event_buffer: Vec::new(), task_graph: None, @@ -55,6 +60,17 @@ impl TaskInstanceModel { Self::default() } + /// Notify the model which DAG + run is now active. + /// Resets `gantt_data` when either the DAG or the run changes so cached + /// retry history cannot leak into a different context. + pub fn set_gantt_context(&mut self, dag_id: &DagId, dag_run_id: &DagRunId) { + let key = (dag_id.clone(), dag_run_id.clone()); + if self.current_gantt_key.as_ref() != Some(&key) { + self.gantt_data = GanttData::default(); + self.current_gantt_key = Some(key); + } + } + /// Sort task instances by topological order (or timestamp fallback) pub fn sort_task_instances(&mut self) { if let Some(graph) = &self.task_graph { @@ -64,15 +80,38 @@ impl TaskInstanceModel { /// Rebuild Gantt data from the current task instance list. /// Returns task IDs that have retries (`try_number` > 1) for fetching detailed tries. + /// + /// For tasks with retries, previously cached try history is preserved so the + /// Gantt chart does not flicker while fresh retry data is being fetched. pub fn rebuild_gantt(&mut self) -> Vec { - self.gantt_data = GanttData::from_task_instances(&self.table.all); + let mut new_gantt = GanttData::from_task_instances(&self.table.all); + + // Collect retried task IDs and carry over cached retry data let mut seen = HashSet::new(); - self.table + let retried: Vec = self + .table .all .iter() .filter(|ti| ti.try_number > 1 && seen.insert(ti.task_id.clone())) .map(|ti| ti.task_id.clone()) - .collect() + .collect(); + + for task_id in &retried { + if let Some(cached_tries) = self.gantt_data.task_tries.get(task_id) { + // Only carry over if the cache has more tries than the fresh data + // (i.e. it already includes the detailed retry history) + let new_tries = new_gantt.task_tries.get(task_id); + if cached_tries.len() > new_tries.map_or(0, Vec::len) { + new_gantt + .task_tries + .insert(task_id.clone(), cached_tries.clone()); + } + } + } + + new_gantt.recompute_window(); + self.gantt_data = new_gantt; + retried } /// Mark a task instance with a new status (optimistic update) diff --git a/src/app/state.rs b/src/app/state.rs index 6b74d091..69caf37e 100644 --- a/src/app/state.rs +++ b/src/app/state.rs @@ -370,6 +370,7 @@ impl App { if let (Some(dag_id), Some(dag_run_id)) = (self.nav_context.dag_id(), self.nav_context.dag_run_id()) { + self.task_instances.set_gantt_context(dag_id, dag_run_id); self.task_instances.table.all = self .environment_state .get_active_task_instances(dag_id, dag_run_id);