fix(tui): use the shared API wire types instead of local copies

ballista-cli/Cargo.toml

@@ -34,6 +34,7 @@ crate-type = ["cdylib", "rlib"]
 [dependencies]
 # CLI-only deps (DataFusion/Ballista) — gated by `cli` feature
 ballista = { path = "../ballista/client", version = "54.0.0", features = ["standalone"], optional = true }
+ballista-api-types = { path = "../ballista/api-types", version = "54.0.0", optional = true }
 datafusion = { workspace = true, optional = true }
 datafusion-cli = { workspace = true, optional = true }
 mimalloc = { workspace = true, optional = true }
@@ -89,6 +90,7 @@ cli = [
 ]
 
 tui = [
+    "dep:ballista-api-types",
     "dep:chrono",
     "dep:config",
     "dep:crossterm",
@@ -111,6 +113,7 @@ tui = [
 ]
 
 web = [
+    "dep:ballista-api-types",
     "dep:chrono",
     "dep:config",
     "dep:critical-section",

ballista-cli/src/tui/app.rs

@@ -18,6 +18,7 @@
 #[cfg(not(feature = "web"))]
 use crate::tui::TuiError;
 use crate::tui::TuiResult;
+use crate::tui::domain::jobs::JobStatusExt;
 use crate::tui::event::Event;
 #[cfg(feature = "web")]
 use crate::tui::event::web::Sender;
@@ -1335,6 +1336,9 @@ mod tests {
             num_stages: 1,
             completed_stages: 0,
             percent_complete: 0,
+            logical_plan: None,
+            physical_plan: None,
+            stage_plan: None,
         }
     }
 

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

@@ -21,33 +21,38 @@ use ratatui::widgets::{ScrollbarState, TableState};
 use serde::Deserialize;
 use std::collections::BTreeMap;
 
-#[derive(Deserialize, Clone, Debug)]
-pub struct Job {
-    pub job_id: String,
-    pub job_name: String,
-    pub status: String,     // Running, Completed, Failed, Canceled
-    pub job_status: String, // human-readable status/failure detail, e.g. "Failed: <reason>"
-    pub start_time: i64,
-    pub end_time: i64,
-    pub num_stages: usize,
-    pub completed_stages: usize,
-    pub percent_complete: u8,
+// Shared with the scheduler that serves it, so a field change there is a
+// compile error here rather than a runtime deserialization failure (#2257).
+pub use ballista_api_types::dto::JobResponse as Job;
+
+/// Status predicates over the shared [`Job`] type. An extension trait because
+/// `Job` is defined in `ballista-api-types`, which is deliberately free of any
+/// TUI concerns.
+pub trait JobStatusExt {
+    /// Job is accepted but not yet scheduled.
+    fn is_queued(&self) -> bool;
+    /// Job is currently executing.
+    fn is_running(&self) -> bool;
+    /// Job finished successfully.
+    fn is_completed(&self) -> bool;
+    /// Job terminated with an error.
+    fn is_failed(&self) -> bool;
 }
 
-impl Job {
-    pub fn is_queued(&self) -> bool {
+impl JobStatusExt for Job {
+    fn is_queued(&self) -> bool {
         self.status == "Queued"
     }
 
-    pub fn is_running(&self) -> bool {
+    fn is_running(&self) -> bool {
         self.status == "Running"
     }
 
-    pub fn is_completed(&self) -> bool {
+    fn is_completed(&self) -> bool {
         self.status == "Completed"
     }
 
-    pub fn is_failed(&self) -> bool {
+    fn is_failed(&self) -> bool {
         self.status == "Failed"
     }
 }
@@ -454,8 +459,8 @@ mod tests {
         id: &str,
         name: &str,
         status: &str,
-        start_time: i64,
-        end_time: i64,
+        start_time: u64,
+        end_time: u64,
         num_stages: usize,
         completed_stages: usize,
         percent_complete: u8,
@@ -470,6 +475,9 @@ mod tests {
             num_stages,
             completed_stages,
             percent_complete,
+            logical_plan: None,
+            physical_plan: None,
+            stage_plan: None,
         }
     }
 
@@ -748,6 +756,9 @@ mod tests {
             num_stages: 1,
             completed_stages: 0,
             percent_complete: 0,
+            logical_plan: None,
+            physical_plan: None,
+            stage_plan: None,
         };
         assert_eq!(job.status, "Failed");
         assert_eq!(job.job_status, "Failed: division by zero");

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

@@ -16,65 +16,16 @@
 // under the License.
 
 use ratatui::widgets::{ScrollbarState, TableState};
-use serde::Deserialize;
 use std::fmt::Display;
 
-#[derive(Deserialize, Clone, Debug)]
-pub struct JobStagesResponse {
-    pub stages: Vec<JobStageResponse>,
-}
-
-#[derive(Deserialize, Clone, Debug)]
-pub struct JobStageResponse {
-    #[serde(rename = "stage_id")]
-    pub id: String,
-    #[serde(rename = "stage_status")]
-    pub status: String,
-    #[serde(rename = "stage_plan", default)]
-    pub plan: String,
-    pub input_rows: usize,
-    pub output_rows: usize,
-    pub elapsed_compute: Option<String>,
-    #[serde(default)]
-    pub task_duration_percentiles: Option<TaskPercentiles>,
-    #[serde(default)]
-    pub task_input_percentiles: Option<TaskPercentiles>,
-    #[serde(default)]
-    pub tasks: Vec<Option<StageTaskResponse>>,
-}
-
-// TaskStatus
-#[derive(Deserialize, Clone, Debug)]
-pub enum StageTaskStatus {
-    Running,
-    Successful,
-    Failed { reason: String },
-}
-
-// TaskSummary
-#[derive(Deserialize, Clone, Debug)]
-pub struct StageTaskResponse {
-    pub id: usize,
-    pub status: StageTaskStatus,
-    pub partition_id: u32,
-    pub input_rows: usize,
-    pub output_rows: usize,
-    pub scheduled_time: u64,
-    pub launch_time: u64,
-    pub start_exec_time: u64,
-    pub end_exec_time: u64,
-    pub finish_time: u64,
-}
-
-// Percentiles
-#[derive(Deserialize, Clone, Debug)]
-pub struct TaskPercentiles {
-    pub min: u64,
-    pub max: u64,
-    pub median: u64,
-    pub p25: u64,
-    pub p75: u64,
-}
+// The `/api/*` response shapes are defined once in `ballista-api-types` and
+// shared with the scheduler that serves them. Re-exported here under the names
+// the TUI already used, so a scheduler-side field change is a compile error
+// rather than a deserialization failure at runtime (#2257).
+pub use ballista_api_types::dto::{
+    QueryStageSummary as JobStageResponse, QueryStagesResponse as JobStagesResponse,
+    TaskStatus as StageTaskStatus, TaskSummary as StageTaskResponse,
+};
 
 #[derive(Debug, Default)]
 pub struct PlanCache {
@@ -372,11 +323,12 @@ impl StagesGraph {
 mod tests {
     use super::{
         JobStageResponse, JobStagesPopup, JobStagesResponse, StageTaskResponse,
-        StageTaskStatus, StagesGraph, TaskPercentiles,
+        StageTaskStatus, StagesGraph,
     };
+    use ballista_api_types::dto::Percentiles;
 
-    fn make_percentiles() -> TaskPercentiles {
-        TaskPercentiles {
+    fn make_percentiles() -> Percentiles {
+        Percentiles {
             min: 0,
             max: 100,
             median: 50,
@@ -387,9 +339,9 @@ mod tests {
 
     fn make_stage(id: &str) -> JobStageResponse {
         JobStageResponse {
-            id: id.to_string(),
-            status: "Completed".to_string(),
-            plan: String::new(),
+            stage_id: id.to_string(),
+            stage_status: "Completed".to_string(),
+            stage_plan: None,
             input_rows: 0,
             output_rows: 0,
             elapsed_compute: Some("1ns".to_string()),
@@ -568,7 +520,7 @@ mod tests {
         let mut popup = make_popup(3);
         popup.table_state.select(Some(1));
         let stage = popup.selected_stage().unwrap();
-        assert_eq!(stage.id, "1");
+        assert_eq!(stage.stage_id, "1");
     }
 
     #[test]
@@ -611,19 +563,69 @@ mod tests {
         assert_eq!(graph.scroll_position, 0);
     }
 
+    /// Regression test for #2257. The scheduler serves `partition_id` as an
+    /// array of the global partitions a task owns, and the TUI used to declare
+    /// it as a scalar `u32`, so the whole stages response failed to parse. This
+    /// pins the shape against a payload shaped like the real one.
+    #[test]
+    fn deserializes_multi_partition_task_payload() {
+        let payload = r#"{
+            "stages": [{
+                "stage_id": "1",
+                "stage_status": "Successful",
+                "input_rows": 10,
+                "output_rows": 10,
+                "elapsed_compute": "1ms",
+                "tasks": [{
+                    "id": 0,
+                    "status": "Successful",
+                    "partition_id": [0, 1, 2],
+                    "scheduled_time": 0,
+                    "launch_time": 0,
+                    "start_exec_time": 0,
+                    "end_exec_time": 5,
+                    "exec_duration": 5,
+                    "finish_time": 5,
+                    "input_rows": 10,
+                    "output_rows": 10
+                }]
+            }]
+        }"#;
+
+        let parsed: JobStagesResponse = serde_json::from_str(payload).unwrap();
+        let task = parsed.stages[0].tasks[0].as_ref().unwrap();
+        assert_eq!(task.partition_id, vec![0, 1, 2]);
+    }
+
+    /// The scheduler's failed-task status carries an `error` alongside
+    /// `reason`; the TUI's old copy only had `reason`.
+    #[test]
+    fn deserializes_failed_task_status_with_error() {
+        let payload = r#"{"Failed": {"reason": "ExecutionError", "error": "boom"}}"#;
+        let status: StageTaskStatus = serde_json::from_str(payload).unwrap();
+        match status {
+            StageTaskStatus::Failed { reason, error } => {
+                assert_eq!(reason, "ExecutionError");
+                assert_eq!(error, "boom");
+            }
+            other => panic!("expected Failed, got {other:?}"),
+        }
+    }
+
     // --- Helpers for task-bearing stages ---
 
     fn make_task(id: usize) -> StageTaskResponse {
         StageTaskResponse {
             id,
             status: StageTaskStatus::Successful,
-            partition_id: id as u32,
+            partition_id: vec![id as u32],
             input_rows: 0,
             output_rows: 0,
             scheduled_time: 0,
             launch_time: 0,
             start_exec_time: 0,
             end_exec_time: 0,
+            exec_duration: 0,
             finish_time: 0,
         }
     }

ballista-cli/src/tui/mod.rs

@@ -297,9 +297,9 @@ pub(crate) mod web {
                     .await
                 {
                     Ok(mut stages) => {
-                        stages
-                            .stages
-                            .sort_by_key(|s| s.id.parse::<u64>().unwrap_or(u64::MAX));
+                        stages.stages.sort_by_key(|s| {
+                            s.stage_id.parse::<u64>().unwrap_or(u64::MAX)
+                        });
                         send_data(UiData::JobStagesData(id, stages), tx).await;
                     }
                     Err(e) => {
@@ -383,9 +383,9 @@ pub(crate) mod web {
             WebKeyAsyncAction::LoadStagePlan(job_id, tab) => {
                 match http_client.get_job_stages(&job_id, &tab).await {
                     Ok(mut stages) => {
-                        stages
-                            .stages
-                            .sort_by_key(|s| s.id.parse::<u64>().unwrap_or(u64::MAX));
+                        stages.stages.sort_by_key(|s| {
+                            s.stage_id.parse::<u64>().unwrap_or(u64::MAX)
+                        });
                         send_data(UiData::JobStagesPlanData(tab, stages), tx).await;
                     }
                     Err(e) => {

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

@@ -88,7 +88,7 @@ fn build_stage_row(i: usize, stage: &JobStageResponse, app: &App) -> Row<'static
         app.theme.row_odd
     };
 
-    let status_style = match stage.status.as_str() {
+    let status_style = match stage.stage_status.as_str() {
         "Running" => app.theme.status_running,
         "Queued" => app.theme.status_queued,
         "Successful" | "Completed" => app.theme.status_completed,
@@ -124,9 +124,9 @@ fn build_stage_row(i: usize, stage: &JobStageResponse, app: &App) -> Row<'static
     );
 
     Row::new(vec![
-        Cell::from(Text::from(stage.id.clone()).right_aligned()),
+        Cell::from(Text::from(stage.stage_id.clone()).right_aligned()),
         Cell::from(
-            Text::from(stage.status.clone())
+            Text::from(stage.stage_status.clone())
                 .style(status_style)
                 .centered(),
         ),