1828: feat(TUI): enable various plan rendering formats for stages

.cursor/rules.md

@@ -0,0 +1,5 @@
+The AI agents should never link to any issue or a pull request
+in any GitHub repository in the code reviews!
+
+The AI agents should not review AI agents' config files like CLAUDE.md or AGENTS.md!
+

AGENTS.md

@@ -0,0 +1,5 @@
+The AI agents should never link to any issue or a pull request
+in any GitHub repository in the code reviews!
+
+The AI agents should not review AI agents' config files like CLAUDE.md or AGENTS.md!
+

CLAUDE.md

@@ -0,0 +1,5 @@
+The AI agents should never link to any issue or a pull request
+in any GitHub repository in the code reviews!
+
+The AI agents should not review AI agents' config files like CLAUDE.md or AGENTS.md!
+

ballista-cli/src/tui/app.rs

@@ -52,7 +52,7 @@ use crate::tui::http_client::HttpClient;
 #[cfg(not(feature = "web"))]
 use crate::tui::ui::{
     load_executor_details_popup, load_executors_data, load_job_details, load_job_dot,
-    load_job_stages_popup, load_jobs_data, load_metrics_data,
+    load_job_stages_popup, load_jobs_data, load_metrics_data, load_stage_plan,
 };
 
 const INVALID_DATE: &str = "Invalid date";
@@ -221,13 +221,47 @@ impl App {
                     _ => {}
                 }
             } else if popup.is_plan_view() {
-                match key.code {
-                    KeyCode::Esc => popup.set_no_details_view(),
-                    KeyCode::Up => popup.scroll_up(),
-                    KeyCode::Down => popup.scroll_down(),
-                    KeyCode::Left => popup.scroll_left(),
-                    KeyCode::Right => popup.scroll_right(),
-                    _ => {}
+                use crate::tui::domain::jobs::stages::StagePlanTab;
+
+                // Collect job_id + tab to fetch while popup is still borrowed,
+                // then drop the borrow before the async call.
+                let fetch = match key.code {
+                    KeyCode::Char('d') => popup
+                        .set_tab(StagePlanTab::Default)
+                        .map(|tab| (popup.job_id.clone(), tab)),
+                    KeyCode::Char('t') => popup
+                        .set_tab(StagePlanTab::Tree)
+                        .map(|tab| (popup.job_id.clone(), tab)),
+                    KeyCode::Char('m') => popup
+                        .set_tab(StagePlanTab::Metrics)
+                        .map(|tab| (popup.job_id.clone(), tab)),
+                    KeyCode::Esc => {
+                        popup.set_no_details_view();
+                        None
+                    }
+                    KeyCode::Up => {
+                        popup.scroll_up();
+                        None
+                    }
+                    KeyCode::Down => {
+                        popup.scroll_down();
+                        None
+                    }
+                    KeyCode::Left => {
+                        popup.scroll_left();
+                        None
+                    }
+                    KeyCode::Right => {
+                        popup.scroll_right();
+                        None
+                    }
+                    _ => None,
+                };
+
+                if let Some((job_id, tab)) = fetch
+                    && let Err(e) = load_stage_plan(self, &job_id, tab).await
+                {
+                    tracing::error!("Failed to load stage plan: {e:?}");
                 }
             } else if popup.is_no_details_view() {
                 match key.code {
@@ -797,6 +831,11 @@ impl App {
             UiData::JobStagesData(job_id, stages) => {
                 self.job_stages_popup = Some(JobStagesPopup::new(job_id, stages));
             }
+            UiData::JobStagesPlanData(_job_id, tab, stages) => {
+                if let Some(popup) = &mut self.job_stages_popup {
+                    popup.cache_plan_response(tab, stages);
+                }
+            }
             UiData::ExecutorDetails(executor) => {
                 self.executor_details_popup = Some(ExecutorDetailsPopup::new(executor));
             }

ballista-cli/src/tui/domain/jobs/stages.rs

@@ -14,7 +14,7 @@
 // KIND, either express or implied.  See the License for the
 // specific language governing permissions and limitations
 // under the License.
-
+#![allow(unfulfilled_lint_expectations)]
 use ratatui::widgets::{ScrollbarState, TableState};
 use serde::Deserialize;
 
@@ -67,17 +67,33 @@ pub struct TaskPercentiles {
     pub p75: u64,
 }
 
+#[derive(Debug, Default)]
+pub struct PlanCache {
+    pub default: Option<JobStagesResponse>,
+    pub tree: Option<JobStagesResponse>,
+    pub metrics: Option<JobStagesResponse>,
+}
+
 #[derive(Debug, PartialEq)]
 pub enum StageDetailsView {
     None,
     Tasks,
-    Plan,
+    Plan(StagePlanTab),
+}
+
+#[allow(dead_code)]
+#[derive(Debug, Clone, PartialEq)]
+pub enum StagePlanTab {
+    Default,
+    Tree,
+    Metrics,
 }
 
 #[derive(Debug)]
 pub struct JobStagesPopup {
     pub job_id: String,
     pub stages: JobStagesResponse,
+    pub plan_cache: PlanCache,
     pub table_state: TableState,
     pub scrollbar_state: ScrollbarState,
     pub tasks_table_state: TableState,
@@ -92,6 +108,10 @@ impl JobStagesPopup {
         Self {
             job_id,
             scrollbar_state: ScrollbarState::new(stages.stages.len()),
+            plan_cache: PlanCache {
+                default: Some(stages.clone()),
+                ..Default::default()
+            },
             stages,
             table_state: TableState::default(),
             tasks_table_state: TableState::default(),
@@ -102,6 +122,36 @@ impl JobStagesPopup {
         }
     }
 
+    pub fn cache_plan_response(&mut self, fmt: StagePlanTab, resp: JobStagesResponse) {
+        match fmt {
+            StagePlanTab::Default => self.plan_cache.default = Some(resp.clone()),
+            StagePlanTab::Tree => self.plan_cache.tree = Some(resp.clone()),
+            StagePlanTab::Metrics => self.plan_cache.metrics = Some(resp.clone()),
+        }
+        // If we're currently on that tab, update the live stages too so the
+        // selection / scroll state is preserved.
+        let active_fmt = self.active_plan_format();
+        if active_fmt == Some(fmt) {
+            self.stages = resp;
+        }
+    }
+
+    #[allow(dead_code)]
+    pub fn cached_response(&self, tab: &StagePlanTab) -> Option<JobStagesResponse> {
+        match tab {
+            StagePlanTab::Default => self.plan_cache.default.clone(),
+            StagePlanTab::Tree => self.plan_cache.tree.clone(),
+            StagePlanTab::Metrics => self.plan_cache.metrics.clone(),
+        }
+    }
+
+    pub fn active_plan_format(&self) -> Option<StagePlanTab> {
+        match &self.details_view {
+            StageDetailsView::Plan(tab) => Some(tab.clone()),
+            _ => None,
+        }
+    }
+
     pub fn plan_vertical_scroll_position(&self) -> u16 {
         self.plan_vertical_scroll_position
     }
@@ -117,7 +167,7 @@ impl JobStagesPopup {
     }
 
     pub fn set_plan_view(&mut self) {
-        self.details_view = StageDetailsView::Plan;
+        self.details_view = StageDetailsView::Plan(StagePlanTab::Default);
         self.plan_vertical_scroll_position = 0;
         self.plan_horizontal_scroll_position = 0;
     }
@@ -135,7 +185,21 @@ impl JobStagesPopup {
     }
 
     pub fn is_plan_view(&self) -> bool {
-        self.details_view == StageDetailsView::Plan
+        matches!(self.details_view, StageDetailsView::Plan(_))
+    }
+
+    #[allow(dead_code)]
+    pub fn set_tab(&mut self, tab: StagePlanTab) -> Option<StagePlanTab> {
+        self.details_view = StageDetailsView::Plan(tab.clone());
+        self.plan_vertical_scroll_position = 0;
+        self.plan_horizontal_scroll_position = 0;
+
+        if let Some(cached) = self.cached_response(&tab) {
+            self.stages = cached;
+            None
+        } else {
+            Some(tab)
+        }
     }
 
     pub fn scroll_down(&mut self) {

ballista-cli/src/tui/event.rs

@@ -20,7 +20,7 @@ use crate::tui::domain::{
     executors::Executor,
     jobs::{
         CancelJobResult, Job, JobDetails,
-        stages::{JobStagesResponse, StagesGraph},
+        stages::{JobStagesResponse, StagePlanTab, StagesGraph},
     },
     metrics::Metric,
 };
@@ -38,6 +38,8 @@ pub enum UiData {
     JobDetails(JobDetails),
     JobStagesGraph(StagesGraph),
     JobStagesData(String, JobStagesResponse),
+    #[allow(dead_code)]
+    JobStagesPlanData(String, StagePlanTab, JobStagesResponse),
     ExecutorDetails(Executor),
     CancelJobResult(CancelJobResult),
 }

ballista-cli/src/tui/http_client.rs

@@ -126,14 +126,26 @@ impl HttpClient {
         self.text(&url).await
     }
 
-    pub async fn get_job_stages(&self, job_id: &str) -> TuiResult<JobStagesResponse> {
+    pub async fn get_job_stages(
+        &self,
+        job_id: &str,
+        plan_format: Option<&str>,
+    ) -> TuiResult<JobStagesResponse> {
+        let fmt = plan_format.unwrap_or({
+            if self.config.job.stage.plan.tree {
+                "tree"
+            } else {
+                ""
+            }
+        });
+
         let url = self.url(&format!(
             "job/{}/stages{}",
             self.url_encode(job_id),
-            if self.config.job.stage.plan.tree {
-                "?plan_format=tree"
+            if fmt.is_empty() {
+                String::new()
             } else {
-                ""
+                format!("?plan_format={fmt}")
             }
         ));
         self.json::<JobStagesResponse>(&url).await

ballista-cli/src/tui/mod.rs

@@ -291,7 +291,7 @@ pub(crate) mod web {
 
         match action {
             WebKeyAsyncAction::LoadJobStages(id) => {
-                match http_client.get_job_stages(&id).await {
+                match http_client.get_job_stages(&id, None).await {
                     Ok(mut stages) => {
                         stages
                             .stages

ballista-cli/src/tui/ui/footer.rs

@@ -80,6 +80,12 @@ pub(super) fn render_footer(f: &mut Frame, area: Rect, app: &App) {
                         } else if app.is_job_stage_plan_popup_open() {
                             current_view_key_bindings
                                 .push(Span::from("[↑↓] Scroll up/down, "));
+                            current_view_key_bindings
+                                .push(Span::from("[d] Default format, "));
+                            current_view_key_bindings
+                                .push(Span::from("[t] Tree format, "));
+                            current_view_key_bindings
+                                .push(Span::from("[m] Show metrics, "));
                             current_view_key_bindings
                                 .push(Span::from("[Esc] Close popup, "));
                         } else if app.is_job_stage_tasks_popup_open() {

ballista-cli/src/tui/ui/main/jobs/mod.rs

@@ -81,11 +81,12 @@ pub async fn load_job_dot(app: &App, job_id: &str) -> TuiResult<()> {
     }
 }
 
+/// Loading whole job's stages to render the popup window
 #[cfg(not(feature = "web"))]
 pub async fn load_job_stages_popup(app: &App, job_id: &str) -> TuiResult<()> {
     let mut stages = app
         .http_client
-        .get_job_stages(job_id)
+        .get_job_stages(job_id, None)
         .await
         .inspect(|stages| tracing::trace!("Loaded stages for job '{job_id}': {stages:?}"))
         .inspect_err(|e| {
@@ -102,6 +103,40 @@ pub async fn load_job_stages_popup(app: &App, job_id: &str) -> TuiResult<()> {
     .await
 }
 
+/// Loading stage's plan to render the popup window
+#[cfg(not(feature = "web"))]
+pub async fn load_stage_plan(
+    app: &App,
+    job_id: &str,
+    tab: crate::tui::domain::jobs::stages::StagePlanTab,
+) -> TuiResult<()> {
+    use crate::tui::domain::jobs::stages::StagePlanTab;
+
+    let fmt = match tab {
+        StagePlanTab::Default => None,
+        StagePlanTab::Tree => Some("tree"),
+        StagePlanTab::Metrics => Some("metrics"),
+    };
+
+    let mut stages = app
+        .http_client
+        .get_job_stages(job_id, fmt)
+        .await
+        .inspect(|s| tracing::trace!("Loaded {fmt:?} plan for job '{job_id}': {s:?}"))
+        .inspect_err(|e| {
+            tracing::error!("Failed to load {fmt:?} plan for job '{job_id}': {e:?}")
+        })?;
+
+    stages
+        .stages
+        .sort_by_key(|s| s.id.parse::<u64>().unwrap_or(u64::MAX));
+
+    app.send_event(Event::DataLoaded {
+        data: UiData::JobStagesPlanData(job_id.to_owned(), tab, stages),
+    })
+    .await
+}
+
 #[cfg(not(feature = "web"))]
 pub async fn load_job_details(app: &App, job_id: &str) -> TuiResult<()> {
     let details = match app.http_client.get_job_details(job_id).await {

ballista-cli/src/tui/ui/main/mod.rs

@@ -23,7 +23,10 @@ pub use executors::{executor_details_popup, render_executors};
 #[cfg(not(feature = "web"))]
 pub use executors::{load_executor_details_popup, load_executors_data};
 #[cfg(not(feature = "web"))]
-pub use jobs::{load_job_details, load_job_dot, load_job_stages_popup, load_jobs_data};
+pub use jobs::{
+    load_job_details, load_job_dot, load_job_stages_popup, load_jobs_data,
+    load_stage_plan,
+};
 
 #[cfg(feature = "web")]
 pub(crate) use jobs::dot_parser;

ballista-cli/src/tui/ui/mod.rs

@@ -36,7 +36,7 @@ pub use main::{
 #[cfg(not(feature = "web"))]
 pub use main::{
     load_executor_details_popup, load_executors_data, load_job_details, load_job_dot,
-    load_job_stages_popup, load_jobs_data, load_metrics_data,
+    load_job_stages_popup, load_jobs_data, load_metrics_data, load_stage_plan,
 };
 
 use ratatui::{