1use std::sync::{
7 atomic::{AtomicBool, Ordering},
8 Arc,
9};
10
11use chrono::{DateTime, Utc};
12use eframe::egui;
13
14use crate::{
15 auto_register::RegistrationState,
16 config::Config,
17 runtime::{HeartbeatOutcome, HeartbeatStatus},
18};
19
20#[derive(Debug, Clone, PartialEq)]
25pub enum StatusView {
26 Initialising { api_base_url: String },
30 Pending {
34 api_base_url: String,
35 request_id: String,
36 since: DateTime<Utc>,
37 },
38 Rejected {
42 api_base_url: String,
43 reason: String,
44 },
45 Registered {
46 worker_id: String,
47 api_base_url: String,
48 vram_total_gb: f32,
49 vram_threshold_gb: f32,
50 paused: bool,
51 busy: bool,
52 last_heartbeat: Option<HeartbeatSummary>,
53 },
54}
55
56#[derive(Debug, Clone, PartialEq)]
57pub struct HeartbeatSummary {
58 pub when: DateTime<Utc>,
59 pub ok: bool,
60 pub reason: Option<String>,
61}
62
63impl HeartbeatSummary {
64 pub fn from(status: &HeartbeatStatus) -> Self {
65 match &status.outcome {
66 HeartbeatOutcome::Ok => Self {
67 when: status.last_attempt_at,
68 ok: true,
69 reason: None,
70 },
71 HeartbeatOutcome::Err { reason } => Self {
72 when: status.last_attempt_at,
73 ok: false,
74 reason: Some(reason.clone()),
75 },
76 }
77 }
78}
79
80impl StatusView {
81 pub fn build(
82 cfg: &Config,
83 registration: &RegistrationState,
84 busy: bool,
85 paused: bool,
86 last_heartbeat: Option<&HeartbeatStatus>,
87 vram_total_gb: f32,
88 ) -> Self {
89 let registered = cfg.worker_id.is_some() && cfg.auth_token.is_some();
90 if registered {
91 return Self::Registered {
92 worker_id: cfg.worker_id.clone().unwrap_or_default(),
93 api_base_url: cfg.api_base_url.clone(),
94 vram_total_gb,
95 vram_threshold_gb: cfg.vram_threshold_gb,
96 paused,
97 busy,
98 last_heartbeat: last_heartbeat.map(HeartbeatSummary::from),
99 };
100 }
101 match registration {
102 RegistrationState::Pending { request_id, since } => Self::Pending {
103 api_base_url: cfg.api_base_url.clone(),
104 request_id: request_id.clone(),
105 since: *since,
106 },
107 RegistrationState::Rejected { reason } => Self::Rejected {
108 api_base_url: cfg.api_base_url.clone(),
109 reason: reason.clone(),
110 },
111 RegistrationState::Pristine | RegistrationState::Approved => Self::Initialising {
114 api_base_url: cfg.api_base_url.clone(),
115 },
116 }
117 }
118}
119
120pub fn format_age(now: DateTime<Utc>, when: DateTime<Utc>) -> String {
122 let delta = now.signed_duration_since(when);
123 let secs = delta.num_seconds();
124 if secs < 0 {
125 return "just now".into();
126 }
127 if secs < 60 {
128 return format!("{secs}s ago");
129 }
130 let mins = secs / 60;
131 if mins < 60 {
132 let rem = secs % 60;
133 return format!("{mins}m {rem:02}s ago");
134 }
135 let hours = mins / 60;
136 let rem_min = mins % 60;
137 format!("{hours}h {rem_min:02}m ago")
138}
139
140pub fn render(ui: &mut egui::Ui, view: &StatusView, paused_flag: &Arc<AtomicBool>) {
145 match view {
146 StatusView::Initialising { api_base_url } => render_initialising(ui, api_base_url),
147 StatusView::Pending {
148 api_base_url,
149 request_id,
150 since,
151 } => render_pending(ui, api_base_url, request_id, *since),
152 StatusView::Rejected {
153 api_base_url,
154 reason,
155 } => render_rejected(ui, api_base_url, reason),
156 StatusView::Registered { .. } => render_registered(ui, view, paused_flag),
157 }
158}
159
160fn render_initialising(ui: &mut egui::Ui, api_base_url: &str) {
161 ui.heading("Initialising");
162 ui.add_space(4.0);
163 ui.horizontal(|ui| {
164 ui.spinner();
165 ui.label(format!(
166 "Asking {api_base_url} for a registration slot\u{2026}"
167 ));
168 });
169 ui.add_space(8.0);
170 ui.label(
171 egui::RichText::new(
172 "No action needed. The worker will keep retrying until it gets through.",
173 )
174 .italics()
175 .color(egui::Color32::from_gray(160)),
176 );
177}
178
179fn render_pending(ui: &mut egui::Ui, api_base_url: &str, request_id: &str, since: DateTime<Utc>) {
180 ui.heading("Waiting for approval");
181 ui.add_space(4.0);
182 ui.label(format!(
183 "This worker has registered with {api_base_url} and is waiting for the \
184 studio operator to approve it. You can keep this window open or close \
185 it \u{2014} the worker keeps polling in the background."
186 ));
187 ui.add_space(12.0);
188 egui::Grid::new("pending_grid")
189 .num_columns(2)
190 .spacing([12.0, 6.0])
191 .show(ui, |ui| {
192 ui.label("Request ID");
193 ui.horizontal(|ui| {
194 ui.monospace(request_id);
195 if ui.button("Copy").clicked() {
196 ui.ctx().copy_text(request_id.to_string());
197 }
198 });
199 ui.end_row();
200
201 ui.label("Waiting");
202 ui.label(format_age(Utc::now(), since));
203 ui.end_row();
204 });
205 ui.add_space(8.0);
206 ui.label(
207 egui::RichText::new(
208 "Share the Request ID with the studio operator if you want them to \
209 find your pending row quickly.",
210 )
211 .italics()
212 .color(egui::Color32::from_gray(160)),
213 );
214}
215
216fn render_rejected(ui: &mut egui::Ui, api_base_url: &str, reason: &str) {
217 ui.heading("Registration rejected");
218 ui.add_space(4.0);
219 ui.colored_label(
220 egui::Color32::LIGHT_RED,
221 if reason.is_empty() {
222 "The studio operator rejected this worker's registration.".to_string()
223 } else {
224 format!("The studio operator rejected this worker's registration: {reason}")
225 },
226 );
227 ui.add_space(12.0);
228 ui.label(format!(
229 "To try again, contact the operator of {api_base_url} to understand why, then run:"
230 ));
231 ui.add_space(4.0);
232 ui.monospace("studio-worker register --reset");
233 ui.add_space(4.0);
234 ui.label(
235 egui::RichText::new(
236 "This clears the local request state and submits a fresh request on \
237 the next launch.",
238 )
239 .italics()
240 .color(egui::Color32::from_gray(160)),
241 );
242}
243
244fn render_registered(ui: &mut egui::Ui, view: &StatusView, paused_flag: &Arc<AtomicBool>) {
245 let StatusView::Registered {
246 worker_id,
247 api_base_url,
248 vram_total_gb,
249 vram_threshold_gb,
250 paused,
251 busy,
252 last_heartbeat,
253 } = view
254 else {
255 unreachable!();
256 };
257
258 ui.heading("Worker status");
259 ui.add_space(4.0);
260
261 let badge = if *busy {
262 ("BUSY", egui::Color32::from_rgb(232, 168, 56))
263 } else if *paused {
264 ("PAUSED", egui::Color32::LIGHT_GRAY)
265 } else {
266 ("IDLE", egui::Color32::LIGHT_GREEN)
267 };
268 ui.horizontal(|ui| {
269 ui.label(egui::RichText::new(badge.0).color(badge.1).strong());
270 ui.label("\u{2014}");
271 ui.label(if *busy {
272 "running a job"
273 } else if *paused {
274 "claiming paused by operator"
275 } else {
276 "waiting for work"
277 });
278 });
279 ui.add_space(8.0);
280
281 ui.horizontal(|ui| {
282 let (label, hint) = if *paused {
283 ("Resume", "start accepting new job offers again")
284 } else {
285 (
286 "Pause",
287 "stop accepting new job offers (in-flight job, if any, will finish)",
288 )
289 };
290 if ui.button(label).on_hover_text(hint).clicked() {
291 toggle_pause(paused_flag);
292 }
293 });
294 ui.add_space(8.0);
295
296 egui::Grid::new("status_grid")
297 .num_columns(2)
298 .spacing([12.0, 6.0])
299 .show(ui, |ui| {
300 ui.label("Worker ID");
301 ui.monospace(worker_id);
302 ui.end_row();
303
304 ui.label("API base URL");
305 ui.monospace(api_base_url);
306 ui.end_row();
307
308 ui.label("VRAM total");
309 ui.label(format!("{vram_total_gb:.1} GB"));
310 ui.end_row();
311
312 ui.label("VRAM threshold");
313 ui.label(format!("{vram_threshold_gb:.1} GB per claim"));
314 ui.end_row();
315
316 ui.label("Last heartbeat");
317 match last_heartbeat {
318 None => ui.label("never"),
319 Some(h) => {
320 let when = format_age(Utc::now(), h.when);
321 if h.ok {
322 ui.colored_label(egui::Color32::LIGHT_GREEN, format!("ok \u{00b7} {when}"))
323 } else {
324 let reason = h.reason.as_deref().unwrap_or("unknown");
325 ui.colored_label(
326 egui::Color32::LIGHT_RED,
327 format!("error \u{00b7} {when} \u{00b7} {reason}"),
328 )
329 }
330 }
331 };
332 ui.end_row();
333 });
334}
335
336fn toggle_pause(paused_flag: &Arc<AtomicBool>) -> bool {
345 let now_paused = !paused_flag.fetch_xor(true, Ordering::SeqCst);
346 tracing::info!(
347 target: "studio_worker::ui::status",
348 paused = now_paused,
349 "pause toggled from status tab"
350 );
351 now_paused
352}
353
354#[cfg(test)]
355mod tests {
356 use super::*;
357 use crate::config::Config;
358 use crate::runtime::HeartbeatStatus;
359 use chrono::TimeZone;
360
361 fn registered_cfg() -> Config {
362 Config {
363 worker_id: Some("w-abc".into()),
364 auth_token: Some("tok-xyz".into()),
365 api_base_url: "https://studio.example".into(),
366 vram_threshold_gb: 12.0,
367 ..Config::default()
368 }
369 }
370
371 #[test]
372 fn build_initialising_when_pristine_and_unregistered() {
373 let cfg = Config::default();
374 let view = StatusView::build(&cfg, &RegistrationState::Pristine, false, false, None, 0.0);
375 match view {
376 StatusView::Initialising { api_base_url } => {
377 assert_eq!(api_base_url, cfg.api_base_url);
378 }
379 other => panic!("expected Initialising, got {other:?}"),
380 }
381 }
382
383 #[test]
384 fn build_pending_when_state_pending() {
385 let cfg = Config::default();
386 let since = Utc::now();
387 let view = StatusView::build(
388 &cfg,
389 &RegistrationState::Pending {
390 request_id: "rr-42".into(),
391 since,
392 },
393 false,
394 false,
395 None,
396 0.0,
397 );
398 match view {
399 StatusView::Pending {
400 request_id,
401 since: s,
402 ..
403 } => {
404 assert_eq!(request_id, "rr-42");
405 assert_eq!(s, since);
406 }
407 other => panic!("expected Pending, got {other:?}"),
408 }
409 }
410
411 #[test]
412 fn build_rejected_when_state_rejected() {
413 let cfg = Config::default();
414 let view = StatusView::build(
415 &cfg,
416 &RegistrationState::Rejected {
417 reason: "unknown contributor".into(),
418 },
419 false,
420 false,
421 None,
422 0.0,
423 );
424 match view {
425 StatusView::Rejected { reason, .. } => assert_eq!(reason, "unknown contributor"),
426 other => panic!("expected Rejected, got {other:?}"),
427 }
428 }
429
430 #[test]
431 fn build_registered_takes_precedence_over_registration_state() {
432 let cfg = registered_cfg();
435 let view = StatusView::build(
436 &cfg,
437 &RegistrationState::Pending {
438 request_id: "rr-stale".into(),
439 since: Utc::now(),
440 },
441 false,
442 false,
443 None,
444 24.0,
445 );
446 assert!(matches!(view, StatusView::Registered { .. }));
447 }
448
449 #[test]
450 fn build_registered_when_worker_id_and_token_present() {
451 let cfg = registered_cfg();
452 let view = StatusView::build(&cfg, &RegistrationState::Approved, false, false, None, 24.0);
453 match view {
454 StatusView::Registered {
455 worker_id,
456 api_base_url,
457 vram_total_gb,
458 vram_threshold_gb,
459 paused,
460 busy,
461 last_heartbeat,
462 } => {
463 assert_eq!(worker_id, "w-abc");
464 assert_eq!(api_base_url, "https://studio.example");
465 assert!((vram_total_gb - 24.0).abs() < f32::EPSILON);
466 assert!((vram_threshold_gb - 12.0).abs() < f32::EPSILON);
467 assert!(!paused);
468 assert!(!busy);
469 assert!(last_heartbeat.is_none());
470 }
471 _ => panic!("expected Registered"),
472 }
473 }
474
475 #[test]
476 fn build_registered_propagates_paused() {
477 let cfg = registered_cfg();
478 let view = StatusView::build(&cfg, &RegistrationState::Approved, false, true, None, 24.0);
479 match view {
480 StatusView::Registered { paused, .. } => assert!(paused),
481 _ => panic!("expected Registered"),
482 }
483 }
484
485 #[test]
486 fn build_propagates_heartbeat_ok() {
487 let cfg = registered_cfg();
488 let hb = HeartbeatStatus {
489 last_attempt_at: Utc::now(),
490 outcome: HeartbeatOutcome::Ok,
491 };
492 let view = StatusView::build(
493 &cfg,
494 &RegistrationState::Approved,
495 false,
496 false,
497 Some(&hb),
498 24.0,
499 );
500 match view {
501 StatusView::Registered {
502 last_heartbeat: Some(s),
503 ..
504 } => {
505 assert!(s.ok);
506 assert!(s.reason.is_none());
507 }
508 _ => panic!("expected Registered with heartbeat"),
509 }
510 }
511
512 #[test]
513 fn build_propagates_heartbeat_err() {
514 let cfg = registered_cfg();
515 let hb = HeartbeatStatus {
516 last_attempt_at: Utc::now(),
517 outcome: HeartbeatOutcome::Err {
518 reason: "5xx".into(),
519 },
520 };
521 let view = StatusView::build(
522 &cfg,
523 &RegistrationState::Approved,
524 true,
525 false,
526 Some(&hb),
527 24.0,
528 );
529 match view {
530 StatusView::Registered {
531 busy,
532 last_heartbeat: Some(s),
533 ..
534 } => {
535 assert!(busy);
536 assert!(!s.ok);
537 assert_eq!(s.reason.as_deref(), Some("5xx"));
538 }
539 _ => panic!("expected Registered with err heartbeat"),
540 }
541 }
542
543 #[test]
544 fn format_age_sub_minute() {
545 let now = Utc.with_ymd_and_hms(2026, 5, 25, 12, 0, 30).unwrap();
546 let then = Utc.with_ymd_and_hms(2026, 5, 25, 12, 0, 18).unwrap();
547 assert_eq!(format_age(now, then), "12s ago");
548 }
549
550 #[test]
551 fn format_age_sub_hour() {
552 let now = Utc.with_ymd_and_hms(2026, 5, 25, 12, 5, 30).unwrap();
553 let then = Utc.with_ymd_and_hms(2026, 5, 25, 12, 0, 18).unwrap();
554 assert_eq!(format_age(now, then), "5m 12s ago");
555 }
556
557 #[test]
558 fn format_age_multi_hour() {
559 let now = Utc.with_ymd_and_hms(2026, 5, 25, 14, 5, 0).unwrap();
560 let then = Utc.with_ymd_and_hms(2026, 5, 25, 12, 0, 0).unwrap();
561 assert_eq!(format_age(now, then), "2h 05m ago");
562 }
563
564 #[test]
565 fn format_age_future_clamps_to_just_now() {
566 let now = Utc.with_ymd_and_hms(2026, 5, 25, 12, 0, 0).unwrap();
567 let then = Utc.with_ymd_and_hms(2026, 5, 25, 12, 0, 5).unwrap();
568 assert_eq!(format_age(now, then), "just now");
569 }
570
571 #[test]
572 fn toggle_pause_flips_flag_and_logs_both_directions() {
573 let flag = Arc::new(AtomicBool::new(false));
578
579 let out = crate::test_support::capture({
580 let flag = flag.clone();
581 move || assert!(toggle_pause(&flag), "first toggle must pause")
582 });
583 assert!(
584 flag.load(Ordering::SeqCst),
585 "flag is paused after first toggle"
586 );
587 assert!(out.contains("INFO"), "expected INFO level, got: {out}");
588 assert!(
589 out.contains("studio_worker::ui::status"),
590 "expected the status target, got: {out}"
591 );
592 assert!(
593 out.contains("pause toggled from status tab"),
594 "expected the toggle message, got: {out}"
595 );
596 assert!(
597 out.contains("paused=true"),
598 "expected paused=true, got: {out}"
599 );
600
601 let out = crate::test_support::capture({
602 let flag = flag.clone();
603 move || assert!(!toggle_pause(&flag), "second toggle must resume")
604 });
605 assert!(
606 !flag.load(Ordering::SeqCst),
607 "flag is resumed after second toggle"
608 );
609 assert!(
610 out.contains("paused=false"),
611 "expected paused=false, got: {out}"
612 );
613 }
614}