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 {
349 let now_paused = !paused_flag.fetch_xor(true, Ordering::SeqCst);
350 tracing::info!(
351 target: "studio_worker::ui::status",
352 paused = now_paused,
353 "pause toggled from status tab"
354 );
355 now_paused
356}
357
358#[cfg(test)]
359mod tests {
360 use super::*;
361 use crate::config::Config;
362 use crate::runtime::HeartbeatStatus;
363 use chrono::TimeZone;
364
365 fn registered_cfg() -> Config {
366 Config {
367 worker_id: Some("w-abc".into()),
368 auth_token: Some("tok-xyz".into()),
369 api_base_url: "https://studio.example".into(),
370 vram_threshold_gb: 12.0,
371 ..Config::default()
372 }
373 }
374
375 #[test]
376 fn build_initialising_when_pristine_and_unregistered() {
377 let cfg = Config::default();
378 let view = StatusView::build(&cfg, &RegistrationState::Pristine, false, false, None, 0.0);
379 match view {
380 StatusView::Initialising { api_base_url } => {
381 assert_eq!(api_base_url, cfg.api_base_url);
382 }
383 other => panic!("expected Initialising, got {other:?}"),
384 }
385 }
386
387 #[test]
388 fn build_pending_when_state_pending() {
389 let cfg = Config::default();
390 let since = Utc::now();
391 let view = StatusView::build(
392 &cfg,
393 &RegistrationState::Pending {
394 request_id: "rr-42".into(),
395 since,
396 },
397 false,
398 false,
399 None,
400 0.0,
401 );
402 match view {
403 StatusView::Pending {
404 request_id,
405 since: s,
406 ..
407 } => {
408 assert_eq!(request_id, "rr-42");
409 assert_eq!(s, since);
410 }
411 other => panic!("expected Pending, got {other:?}"),
412 }
413 }
414
415 #[test]
416 fn build_rejected_when_state_rejected() {
417 let cfg = Config::default();
418 let view = StatusView::build(
419 &cfg,
420 &RegistrationState::Rejected {
421 reason: "unknown contributor".into(),
422 },
423 false,
424 false,
425 None,
426 0.0,
427 );
428 match view {
429 StatusView::Rejected { reason, .. } => assert_eq!(reason, "unknown contributor"),
430 other => panic!("expected Rejected, got {other:?}"),
431 }
432 }
433
434 #[test]
435 fn build_registered_takes_precedence_over_registration_state() {
436 let cfg = registered_cfg();
439 let view = StatusView::build(
440 &cfg,
441 &RegistrationState::Pending {
442 request_id: "rr-stale".into(),
443 since: Utc::now(),
444 },
445 false,
446 false,
447 None,
448 24.0,
449 );
450 assert!(matches!(view, StatusView::Registered { .. }));
451 }
452
453 #[test]
454 fn build_registered_when_worker_id_and_token_present() {
455 let cfg = registered_cfg();
456 let view = StatusView::build(&cfg, &RegistrationState::Approved, false, false, None, 24.0);
457 match view {
458 StatusView::Registered {
459 worker_id,
460 api_base_url,
461 vram_total_gb,
462 vram_threshold_gb,
463 paused,
464 busy,
465 last_heartbeat,
466 } => {
467 assert_eq!(worker_id, "w-abc");
468 assert_eq!(api_base_url, "https://studio.example");
469 assert!((vram_total_gb - 24.0).abs() < f32::EPSILON);
470 assert!((vram_threshold_gb - 12.0).abs() < f32::EPSILON);
471 assert!(!paused);
472 assert!(!busy);
473 assert!(last_heartbeat.is_none());
474 }
475 _ => panic!("expected Registered"),
476 }
477 }
478
479 #[test]
480 fn build_registered_propagates_paused() {
481 let cfg = registered_cfg();
482 let view = StatusView::build(&cfg, &RegistrationState::Approved, false, true, None, 24.0);
483 match view {
484 StatusView::Registered { paused, .. } => assert!(paused),
485 _ => panic!("expected Registered"),
486 }
487 }
488
489 #[test]
490 fn build_propagates_heartbeat_ok() {
491 let cfg = registered_cfg();
492 let hb = HeartbeatStatus {
493 last_attempt_at: Utc::now(),
494 outcome: HeartbeatOutcome::Ok,
495 };
496 let view = StatusView::build(
497 &cfg,
498 &RegistrationState::Approved,
499 false,
500 false,
501 Some(&hb),
502 24.0,
503 );
504 match view {
505 StatusView::Registered {
506 last_heartbeat: Some(s),
507 ..
508 } => {
509 assert!(s.ok);
510 assert!(s.reason.is_none());
511 }
512 _ => panic!("expected Registered with heartbeat"),
513 }
514 }
515
516 #[test]
517 fn build_propagates_heartbeat_err() {
518 let cfg = registered_cfg();
519 let hb = HeartbeatStatus {
520 last_attempt_at: Utc::now(),
521 outcome: HeartbeatOutcome::Err {
522 reason: "5xx".into(),
523 },
524 };
525 let view = StatusView::build(
526 &cfg,
527 &RegistrationState::Approved,
528 true,
529 false,
530 Some(&hb),
531 24.0,
532 );
533 match view {
534 StatusView::Registered {
535 busy,
536 last_heartbeat: Some(s),
537 ..
538 } => {
539 assert!(busy);
540 assert!(!s.ok);
541 assert_eq!(s.reason.as_deref(), Some("5xx"));
542 }
543 _ => panic!("expected Registered with err heartbeat"),
544 }
545 }
546
547 #[test]
548 fn format_age_sub_minute() {
549 let now = Utc.with_ymd_and_hms(2026, 5, 25, 12, 0, 30).unwrap();
550 let then = Utc.with_ymd_and_hms(2026, 5, 25, 12, 0, 18).unwrap();
551 assert_eq!(format_age(now, then), "12s ago");
552 }
553
554 #[test]
555 fn format_age_sub_hour() {
556 let now = Utc.with_ymd_and_hms(2026, 5, 25, 12, 5, 30).unwrap();
557 let then = Utc.with_ymd_and_hms(2026, 5, 25, 12, 0, 18).unwrap();
558 assert_eq!(format_age(now, then), "5m 12s ago");
559 }
560
561 #[test]
562 fn format_age_multi_hour() {
563 let now = Utc.with_ymd_and_hms(2026, 5, 25, 14, 5, 0).unwrap();
564 let then = Utc.with_ymd_and_hms(2026, 5, 25, 12, 0, 0).unwrap();
565 assert_eq!(format_age(now, then), "2h 05m ago");
566 }
567
568 #[test]
569 fn format_age_future_clamps_to_just_now() {
570 let now = Utc.with_ymd_and_hms(2026, 5, 25, 12, 0, 0).unwrap();
571 let then = Utc.with_ymd_and_hms(2026, 5, 25, 12, 0, 5).unwrap();
572 assert_eq!(format_age(now, then), "just now");
573 }
574
575 #[test]
576 fn toggle_pause_flips_flag_and_logs_both_directions() {
577 let flag = Arc::new(AtomicBool::new(false));
582
583 let out = crate::test_support::capture({
584 let flag = flag.clone();
585 move || assert!(toggle_pause(&flag), "first toggle must pause")
586 });
587 assert!(
588 flag.load(Ordering::SeqCst),
589 "flag is paused after first toggle"
590 );
591 assert!(out.contains("INFO"), "expected INFO level, got: {out}");
592 assert!(
593 out.contains("studio_worker::ui::status"),
594 "expected the status target, got: {out}"
595 );
596 assert!(
597 out.contains("pause toggled from status tab"),
598 "expected the toggle message, got: {out}"
599 );
600 assert!(
601 out.contains("paused=true"),
602 "expected paused=true, got: {out}"
603 );
604
605 let out = crate::test_support::capture({
606 let flag = flag.clone();
607 move || assert!(!toggle_pause(&flag), "second toggle must resume")
608 });
609 assert!(
610 !flag.load(Ordering::SeqCst),
611 "flag is resumed after second toggle"
612 );
613 assert!(
614 out.contains("paused=false"),
615 "expected paused=false, got: {out}"
616 );
617 }
618}