studio_worker/ui/tabs/

status.rs

//! Status tab — surfaces who the worker is, who it's talking to, and
//! how recently it last successfully heartbeat.  When the worker
//! hasn't registered yet, this tab shows the in-window Register form
//! (fork #2 of plans/native-ui.md, default A).

use chrono::{DateTime, Utc};
use eframe::egui;

use crate::{
    auto_register::RegistrationState,
    config::Config,
    runtime::{GpuRuntimeStatus, HeartbeatOutcome, HeartbeatStatus, SessionState},
};

/// Pure-data view of the Status tab.  Constructed each frame from
/// the live shared state; no egui types in scope so it's
/// unit-testable.  Maps onto `auto_register::RegistrationState` for
/// the pre-registration phases.
#[derive(Debug, Clone, PartialEq)]
pub enum StatusView {
    /// Worker is auto-registering but hasn't received a request_id
    /// from the studio yet.  Transient — normally only seen for a
    /// frame or two on first launch.
    Initialising { api_base_url: String },
    /// Studio has a Pending Workers row for this install; UI shows
    /// the request id + a copy button so the operator can match it
    /// in the dashboard.
    Pending {
        api_base_url: String,
        request_id: String,
        since: DateTime<Utc>,
    },
    /// Operator rejected the request.  Worker stops trying; the
    /// hint instructs the user to run `studio-worker register --reset`
    /// to clear state and try again.
    Rejected {
        api_base_url: String,
        reason: String,
    },
    Registered {
        worker_id: String,
        api_base_url: String,
        vram_total_gb: f32,
        vram_threshold_gb: f32,
        paused: bool,
        busy: bool,
        last_heartbeat: Option<HeartbeatSummary>,
        /// One-line WS lifecycle summary (connected / reconnecting /
        /// auth-failed + recovery action).
        session: String,
        /// GPU-runtime readiness: `None` = not probed, else (ok, detail).
        gpu: Option<(bool, String)>,
    },
}

#[derive(Debug, Clone, PartialEq)]
pub struct HeartbeatSummary {
    pub when: DateTime<Utc>,
    pub ok: bool,
    pub reason: Option<String>,
}

impl HeartbeatSummary {
    pub fn from(status: &HeartbeatStatus) -> Self {
        match &status.outcome {
            HeartbeatOutcome::Ok => Self {
                when: status.last_attempt_at,
                ok: true,
                reason: None,
            },
            HeartbeatOutcome::Err { reason } => Self {
                when: status.last_attempt_at,
                ok: false,
                reason: Some(reason.clone()),
            },
        }
    }
}

impl StatusView {
    // A pure-data view assembled from the several independent runtime
    // signals the Status tab shows; grouping them would only rename the
    // list, not shorten it.
    #[allow(clippy::too_many_arguments)]
    pub fn build(
        cfg: &Config,
        registered: bool,
        registration: &RegistrationState,
        busy: bool,
        paused: bool,
        last_heartbeat: Option<&HeartbeatStatus>,
        vram_total_gb: f32,
        session_state: &SessionState,
        gpu_runtime: Option<&GpuRuntimeStatus>,
    ) -> Self {
        if registered {
            return Self::Registered {
                worker_id: cfg.worker_id.clone().unwrap_or_default(),
                api_base_url: cfg.api_base_url.clone(),
                vram_total_gb,
                vram_threshold_gb: cfg.vram_threshold_gb,
                paused,
                busy,
                last_heartbeat: last_heartbeat.map(HeartbeatSummary::from),
                session: session_state.summary(),
                gpu: gpu_runtime.map(|g| (g.ok, g.detail.clone())),
            };
        }
        match registration {
            RegistrationState::Pending { request_id, since } => Self::Pending {
                api_base_url: cfg.api_base_url.clone(),
                request_id: request_id.clone(),
                since: *since,
            },
            RegistrationState::Rejected { reason } => Self::Rejected {
                api_base_url: cfg.api_base_url.clone(),
                reason: reason.clone(),
            },
            // Pristine — cold start, between requests, or operator
            // bootstrap path that hasn't completed yet.
            RegistrationState::Pristine | RegistrationState::Approved => Self::Initialising {
                api_base_url: cfg.api_base_url.clone(),
            },
        }
    }
}

/// Human-friendly "5s ago" formatting for a heartbeat timestamp.
pub fn format_age(now: DateTime<Utc>, when: DateTime<Utc>) -> String {
    let delta = now.signed_duration_since(when);
    let secs = delta.num_seconds();
    if secs < 0 {
        return "just now".into();
    }
    if secs < 60 {
        return format!("{secs}s ago");
    }
    let mins = secs / 60;
    if mins < 60 {
        let rem = secs % 60;
        return format!("{mins}m {rem:02}s ago");
    }
    let hours = mins / 60;
    let rem_min = mins % 60;
    format!("{hours}h {rem_min:02}m ago")
}

// ---------------------------------------------------------------------------
// Rendering
// ---------------------------------------------------------------------------

/// What the operator asked for on the Status tab.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StatusAction {
    /// Pause (`true`) or resume claiming studio jobs.
    SetPaused(bool),
    /// Clear a rejected registration and ask the studio again.
    ResetRegistration,
}

pub fn render(ui: &mut egui::Ui, view: &StatusView) -> Option<StatusAction> {
    match view {
        StatusView::Initialising { api_base_url } => render_initialising(ui, api_base_url),
        StatusView::Pending {
            api_base_url,
            request_id,
            since,
        } => render_pending(ui, api_base_url, request_id, *since),
        StatusView::Rejected {
            api_base_url,
            reason,
        } => return render_rejected(ui, api_base_url, reason),
        StatusView::Registered { .. } => return render_registered(ui, view),
    }
    None
}

fn render_initialising(ui: &mut egui::Ui, api_base_url: &str) {
    ui.heading("Initialising");
    ui.add_space(4.0);
    ui.horizontal(|ui| {
        ui.spinner();
        ui.label(format!(
            "Asking {api_base_url} for a registration slot\u{2026}"
        ));
    });
    ui.add_space(8.0);
    ui.label(
        egui::RichText::new(
            "No action needed.  The worker will keep retrying until it gets through.",
        )
        .italics()
        .color(egui::Color32::from_gray(160)),
    );
}

fn render_pending(ui: &mut egui::Ui, api_base_url: &str, request_id: &str, since: DateTime<Utc>) {
    ui.heading("Waiting for approval");
    ui.add_space(4.0);
    ui.label(format!(
        "This worker has registered with {api_base_url} and is waiting for the \
         studio operator to approve it.  You can keep this window open or close \
         it \u{2014} the worker keeps polling in the background."
    ));
    ui.add_space(12.0);
    egui::Grid::new("pending_grid")
        .num_columns(2)
        .spacing([12.0, 6.0])
        .show(ui, |ui| {
            ui.label("Request ID");
            ui.horizontal(|ui| {
                ui.monospace(request_id);
                if ui.button("Copy").clicked() {
                    ui.ctx().copy_text(request_id.to_string());
                }
            });
            ui.end_row();

            ui.label("Waiting");
            ui.label(format_age(Utc::now(), since));
            ui.end_row();
        });
    ui.add_space(8.0);
    ui.label(
        egui::RichText::new(
            "Share the Request ID with the studio operator if you want them to \
             find your pending row quickly.",
        )
        .italics()
        .color(egui::Color32::from_gray(160)),
    );
}

fn render_rejected(ui: &mut egui::Ui, api_base_url: &str, reason: &str) -> Option<StatusAction> {
    ui.heading("Registration rejected");
    ui.add_space(4.0);
    ui.colored_label(
        egui::Color32::LIGHT_RED,
        if reason.is_empty() {
            "The studio operator rejected this worker's registration.".to_string()
        } else {
            format!("The studio operator rejected this worker's registration: {reason}")
        },
    );
    ui.add_space(12.0);
    ui.label(format!(
        "Local models and the local API keep working.  To ask {api_base_url} again, \
         check with its operator why, then reset the registration:"
    ));
    ui.add_space(4.0);
    let clicked = ui
        .button("Reset registration")
        .on_hover_text("clear the local registration state and submit a fresh request")
        .clicked();
    ui.add_space(4.0);
    ui.label(
        egui::RichText::new("Same as `studio-worker register --reset`, without a restart.")
            .italics()
            .color(egui::Color32::from_gray(160)),
    );
    clicked.then_some(StatusAction::ResetRegistration)
}

fn render_registered(ui: &mut egui::Ui, view: &StatusView) -> Option<StatusAction> {
    let StatusView::Registered {
        worker_id,
        api_base_url,
        vram_total_gb,
        vram_threshold_gb,
        paused,
        busy,
        last_heartbeat,
        session,
        gpu,
    } = view
    else {
        unreachable!();
    };

    ui.heading("Worker status");
    ui.add_space(4.0);

    let badge = if *busy {
        ("BUSY", egui::Color32::from_rgb(232, 168, 56))
    } else if *paused {
        ("PAUSED", egui::Color32::LIGHT_GRAY)
    } else {
        ("IDLE", egui::Color32::LIGHT_GREEN)
    };
    ui.horizontal(|ui| {
        ui.label(egui::RichText::new(badge.0).color(badge.1).strong());
        ui.label("\u{2014}");
        ui.label(if *busy {
            "running a job"
        } else if *paused {
            "claiming paused by operator"
        } else {
            "waiting for work"
        });
    });
    ui.add_space(8.0);

    let mut action = None;
    ui.horizontal(|ui| {
        let (label, hint) = if *paused {
            ("Resume", "start accepting new job offers again")
        } else {
            (
                "Pause",
                "stop accepting new job offers (in-flight job, if any, will finish)",
            )
        };
        if ui.button(label).on_hover_text(hint).clicked() {
            action = Some(StatusAction::SetPaused(!*paused));
        }
    });
    ui.add_space(8.0);

    egui::Grid::new("status_grid")
        .num_columns(2)
        .spacing([12.0, 6.0])
        .show(ui, |ui| {
            ui.label("Worker ID");
            ui.monospace(worker_id);
            ui.end_row();

            ui.label("API base URL");
            ui.monospace(api_base_url);
            ui.end_row();

            ui.label("VRAM total");
            ui.label(format!("{vram_total_gb:.1} GB"));
            ui.end_row();

            ui.label("VRAM threshold");
            ui.label(format!("{vram_threshold_gb:.1} GB per claim"));
            ui.end_row();

            ui.label("Connection");
            // Auth-failed / fatal states carry a recovery action — draw
            // them red so a stranded worker is impossible to miss.
            if session.contains("failed") || session.contains("ended") {
                ui.colored_label(egui::Color32::LIGHT_RED, session);
            } else {
                ui.label(session);
            }
            ui.end_row();

            // GPU runtime: red when missing, with the exact remedy, so a
            // box that can't run image jobs says why up front.
            if let Some((ok, detail)) = gpu {
                ui.label("GPU runtime");
                if *ok {
                    ui.colored_label(egui::Color32::LIGHT_GREEN, detail);
                } else {
                    ui.colored_label(egui::Color32::LIGHT_RED, detail);
                }
                ui.end_row();
            }

            ui.label("Last heartbeat");
            match last_heartbeat {
                None => ui.label("never"),
                Some(h) => {
                    let when = format_age(Utc::now(), h.when);
                    if h.ok {
                        ui.colored_label(egui::Color32::LIGHT_GREEN, format!("ok \u{00b7} {when}"))
                    } else {
                        let reason = h.reason.as_deref().unwrap_or("unknown");
                        ui.colored_label(
                            egui::Color32::LIGHT_RED,
                            format!("error \u{00b7} {when} \u{00b7} {reason}"),
                        )
                    }
                }
            };
            ui.end_row();
        });
    action
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::Config;
    use crate::runtime::HeartbeatStatus;
    use chrono::TimeZone;

    fn registered_of(cfg: &Config) -> bool {
        cfg.worker_id.is_some() && cfg.auth_token.is_some()
    }

    fn registered_cfg() -> Config {
        Config {
            worker_id: Some("w-abc".into()),
            auth_token: Some("tok-xyz".into()),
            api_base_url: "https://studio.example".into(),
            vram_threshold_gb: 12.0,
            ..Config::default()
        }
    }

    #[test]
    fn build_initialising_when_pristine_and_unregistered() {
        let cfg = Config::default();
        let view = StatusView::build(
            &cfg,
            registered_of(&cfg),
            &RegistrationState::Pristine,
            false,
            false,
            None,
            0.0,
            &SessionState::default(),
            None,
        );
        match view {
            StatusView::Initialising { api_base_url } => {
                assert_eq!(api_base_url, cfg.api_base_url);
            }
            other => panic!("expected Initialising, got {other:?}"),
        }
    }

    #[test]
    fn build_pending_when_state_pending() {
        let cfg = Config::default();
        let since = Utc::now();
        let view = StatusView::build(
            &cfg,
            registered_of(&cfg),
            &RegistrationState::Pending {
                request_id: "rr-42".into(),
                since,
            },
            false,
            false,
            None,
            0.0,
            &SessionState::default(),
            None,
        );
        match view {
            StatusView::Pending {
                request_id,
                since: s,
                ..
            } => {
                assert_eq!(request_id, "rr-42");
                assert_eq!(s, since);
            }
            other => panic!("expected Pending, got {other:?}"),
        }
    }

    #[test]
    fn build_rejected_when_state_rejected() {
        let cfg = Config::default();
        let view = StatusView::build(
            &cfg,
            registered_of(&cfg),
            &RegistrationState::Rejected {
                reason: "unknown contributor".into(),
            },
            false,
            false,
            None,
            0.0,
            &SessionState::default(),
            None,
        );
        match view {
            StatusView::Rejected { reason, .. } => assert_eq!(reason, "unknown contributor"),
            other => panic!("expected Rejected, got {other:?}"),
        }
    }

    #[test]
    fn build_registered_takes_precedence_over_registration_state() {
        // If worker_id + auth_token are set, the registration state
        // is irrelevant — we're operational.
        let cfg = registered_cfg();
        let view = StatusView::build(
            &cfg,
            registered_of(&cfg),
            &RegistrationState::Pending {
                request_id: "rr-stale".into(),
                since: Utc::now(),
            },
            false,
            false,
            None,
            24.0,
            &SessionState::default(),
            None,
        );
        assert!(matches!(view, StatusView::Registered { .. }));
    }

    #[test]
    fn build_registered_when_worker_id_and_token_present() {
        let cfg = registered_cfg();
        let view = StatusView::build(
            &cfg,
            registered_of(&cfg),
            &RegistrationState::Approved,
            false,
            false,
            None,
            24.0,
            &SessionState::default(),
            None,
        );
        match view {
            StatusView::Registered {
                worker_id,
                api_base_url,
                vram_total_gb,
                vram_threshold_gb,
                paused,
                busy,
                last_heartbeat,
                session: _,
                gpu: _,
            } => {
                assert_eq!(worker_id, "w-abc");
                assert_eq!(api_base_url, "https://studio.example");
                assert!((vram_total_gb - 24.0).abs() < f32::EPSILON);
                assert!((vram_threshold_gb - 12.0).abs() < f32::EPSILON);
                assert!(!paused);
                assert!(!busy);
                assert!(last_heartbeat.is_none());
            }
            _ => panic!("expected Registered"),
        }
    }

    #[test]
    fn build_registered_propagates_paused() {
        let cfg = registered_cfg();
        let view = StatusView::build(
            &cfg,
            registered_of(&cfg),
            &RegistrationState::Approved,
            false,
            true,
            None,
            24.0,
            &SessionState::default(),
            None,
        );
        match view {
            StatusView::Registered { paused, .. } => assert!(paused),
            _ => panic!("expected Registered"),
        }
    }

    #[test]
    fn build_propagates_heartbeat_ok() {
        let cfg = registered_cfg();
        let hb = HeartbeatStatus {
            last_attempt_at: Utc::now(),
            outcome: HeartbeatOutcome::Ok,
        };
        let view = StatusView::build(
            &cfg,
            registered_of(&cfg),
            &RegistrationState::Approved,
            false,
            false,
            Some(&hb),
            24.0,
            &SessionState::default(),
            None,
        );
        match view {
            StatusView::Registered {
                last_heartbeat: Some(s),
                ..
            } => {
                assert!(s.ok);
                assert!(s.reason.is_none());
            }
            _ => panic!("expected Registered with heartbeat"),
        }
    }

    #[test]
    fn build_propagates_heartbeat_err() {
        let cfg = registered_cfg();
        let hb = HeartbeatStatus {
            last_attempt_at: Utc::now(),
            outcome: HeartbeatOutcome::Err {
                reason: "5xx".into(),
            },
        };
        let view = StatusView::build(
            &cfg,
            registered_of(&cfg),
            &RegistrationState::Approved,
            true,
            false,
            Some(&hb),
            24.0,
            &SessionState::default(),
            None,
        );
        match view {
            StatusView::Registered {
                busy,
                last_heartbeat: Some(s),
                ..
            } => {
                assert!(busy);
                assert!(!s.ok);
                assert_eq!(s.reason.as_deref(), Some("5xx"));
            }
            _ => panic!("expected Registered with err heartbeat"),
        }
    }

    #[test]
    fn format_age_sub_minute() {
        let now = Utc.with_ymd_and_hms(2026, 5, 25, 12, 0, 30).unwrap();
        let then = Utc.with_ymd_and_hms(2026, 5, 25, 12, 0, 18).unwrap();
        assert_eq!(format_age(now, then), "12s ago");
    }

    #[test]
    fn format_age_sub_hour() {
        let now = Utc.with_ymd_and_hms(2026, 5, 25, 12, 5, 30).unwrap();
        let then = Utc.with_ymd_and_hms(2026, 5, 25, 12, 0, 18).unwrap();
        assert_eq!(format_age(now, then), "5m 12s ago");
    }

    #[test]
    fn format_age_multi_hour() {
        let now = Utc.with_ymd_and_hms(2026, 5, 25, 14, 5, 0).unwrap();
        let then = Utc.with_ymd_and_hms(2026, 5, 25, 12, 0, 0).unwrap();
        assert_eq!(format_age(now, then), "2h 05m ago");
    }

    #[test]
    fn format_age_future_clamps_to_just_now() {
        let now = Utc.with_ymd_and_hms(2026, 5, 25, 12, 0, 0).unwrap();
        let then = Utc.with_ymd_and_hms(2026, 5, 25, 12, 0, 5).unwrap();
        assert_eq!(format_age(now, then), "just now");
    }
}