studio-worker 0.4.7

Pull-based image-generation worker for the minis.gg studio.
Documentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
//! Long-running auto-update task + one-shot CLI helpers.
//!
//! After the WS migration the runtime owns just two background
//! tasks: the WebSocket session (`ws::session::spawn_ws_session`,
//! which subsumes heartbeats, claim/accept/complete, fail, and log
//! shipping) and the auto-updater (`spawn_auto_updater`).  Per-tick
//! helpers from the old polling loops are gone.
use crate::{
    config::{self, Config, SharedConfig},
    engine::Engine,
    sys,
    types::*,
    update, AGENT_VERSION,
};
use anyhow::{anyhow, Result};
use chrono::{DateTime, SecondsFormat, Utc};
use parking_lot::Mutex;
use std::{
    collections::VecDeque,
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc,
    },
    time::Duration,
};
use tracing::{info, warn};

/// Tracing target for runtime-level events (startup, state mutations).
/// Stable so operators can filter with `RUST_LOG=studio_worker::runtime=debug`.
const TRACE_TARGET: &str = "studio_worker::runtime";

/// Maximum number of finished jobs kept in `WorkerObservers::recent_jobs`.
/// Older entries fall off the back of the ring.
pub const RECENT_JOBS_CAP: usize = 50;

/// Maximum number of log entries kept in `WorkerObservers::recent_logs`
/// for the UI's Logs tab.  The shipping queue (`logs: Arc<Mutex<Vec<…>>>`)
/// is drained on every WS tick — the display ring is what the UI reads.
pub const RECENT_LOGS_CAP: usize = 1000;

/// Prompt previews stored in `CurrentJob` / `RecentJob` are clipped to
/// this many chars so the in-memory state stays bounded even when LLM
/// prompts are huge.
pub const PROMPT_PREVIEW_CHARS: usize = 200;

/// Maximum number of entries the WS ship queue (`logs:
/// Arc<Mutex<Vec<LogEntry>>>`) may hold.  The shipper pump only drains
/// while a session is connected, so a long approval wait or reconnect
/// backoff would otherwise grow the queue without bound.  On overflow
/// the oldest entries are dropped and a warn-level marker records the
/// loss.
pub const LOG_SHIP_QUEUE_CAP: usize = 5_000;

/// Job in flight right now.  Populated by the WS session before
/// dispatch, cleared once the job finishes (success or failure).
#[derive(Debug, Clone)]
pub struct CurrentJob {
    pub job_id: String,
    pub kind: TaskKind,
    pub model: String,
    pub prompt: String,
    pub started_at: DateTime<Utc>,
}

/// Outcome a finished job ended with.  Failures carry the human
/// reason (already surfaced to logs + Sentry).
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum JobOutcome {
    Completed,
    Failed { reason: String },
}

/// One finished job, retained in the recent-jobs ring for the UI.
#[derive(Debug, Clone)]
pub struct RecentJob {
    pub job_id: String,
    pub kind: TaskKind,
    pub model: String,
    pub prompt: String,
    pub outcome: JobOutcome,
    pub started_at: DateTime<Utc>,
    pub finished_at: DateTime<Utc>,
}

/// Result of the most recent heartbeat the WS session sent.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum HeartbeatOutcome {
    Ok,
    Err { reason: String },
}

#[derive(Debug, Clone)]
pub struct HeartbeatStatus {
    pub last_attempt_at: DateTime<Utc>,
    pub outcome: HeartbeatOutcome,
}

/// Bundle of in-process observation slots the WS session writes to and
/// the optional native UI reads from.  `Default` gives empty slots so
/// existing (headless) call sites stay one-liners.  Cheap to clone —
/// every field is an `Arc`.
#[derive(Clone, Default)]
pub struct WorkerObservers {
    pub current_job: Arc<Mutex<Option<CurrentJob>>>,
    pub recent_jobs: Arc<Mutex<VecDeque<RecentJob>>>,
    pub last_heartbeat: Arc<Mutex<Option<HeartbeatStatus>>>,
    /// Bounded ring of every log entry the worker has emitted, kept
    /// for the UI's Logs tab.  Separate from the WS ship queue
    /// (which is drained every second) so the display doesn't blank
    /// out between ticks.
    pub recent_logs: Arc<Mutex<VecDeque<LogEntry>>>,
}

pub fn truncate_prompt(s: &str) -> String {
    if s.chars().count() <= PROMPT_PREVIEW_CHARS {
        return s.to_string();
    }
    let mut out: String = s.chars().take(PROMPT_PREVIEW_CHARS).collect();
    out.push('');
    out
}

pub fn record_recent_job(observers: &WorkerObservers, entry: RecentJob) {
    let mut ring = observers.recent_jobs.lock();
    ring.push_front(entry);
    while ring.len() > RECENT_JOBS_CAP {
        ring.pop_back();
    }
}

/// Test-only helper to populate the recent-jobs ring without driving a
/// full claim cycle.  Lives in the library surface so integration
/// tests can pin the ring-capacity contract cheaply.
#[doc(hidden)]
pub fn push_recent_job_for_tests(observers: &WorkerObservers, job_id: &str) {
    let now = Utc::now();
    record_recent_job(
        observers,
        RecentJob {
            job_id: job_id.to_string(),
            kind: TaskKind::Image,
            model: "synthetic".into(),
            prompt: String::new(),
            outcome: JobOutcome::Completed,
            started_at: now,
            finished_at: now,
        },
    );
}

pub const AUTO_UPDATE_TICK: Duration = Duration::from_secs(60);
/// Cadence at which the auto-updater's idle wait re-checks the `stop`
/// flag.  Mirrors the WS session's shutdown tick so a SIGTERM / SIGINT
/// landing during the (up to `AUTO_UPDATE_TICK`-long) idle window wakes
/// the loop within ~250 ms instead of leaving `run_loops`' join blocked
/// for a whole tick.
pub const AUTO_UPDATE_SHUTDOWN_TICK: Duration = Duration::from_millis(250);
/// Default WS heartbeat interval, re-exported here so the native UI
/// (and any other downstream readers) get a stable constant without
/// reaching into `ws::session`.
pub const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(5);

/// Schedule for the long-running loops.
#[derive(Debug, Clone, Copy)]
pub struct LoopSchedule {
    pub ws_session: crate::ws::session::SessionSchedule,
    pub auto_update_tick: Duration,
    /// How often the idle wait between update checks re-polls the
    /// `stop` flag, so a shutdown request isn't deferred for a whole
    /// `auto_update_tick`.
    pub shutdown_tick: Duration,
}

impl Default for LoopSchedule {
    fn default() -> Self {
        Self {
            ws_session: crate::ws::session::SessionSchedule::default(),
            auto_update_tick: AUTO_UPDATE_TICK,
            shutdown_tick: AUTO_UPDATE_SHUTDOWN_TICK,
        }
    }
}

impl LoopSchedule {
    /// Schedule with 1 ms intervals — used by tests to exercise the
    /// loop wrappers without blocking.
    pub fn fast_for_tests() -> Self {
        Self {
            ws_session: crate::ws::session::SessionSchedule::fast_for_tests(),
            auto_update_tick: Duration::from_millis(1),
            shutdown_tick: Duration::from_millis(1),
        }
    }
}

// ---------------------------------------------------------------------------
// One-shot helpers used by the CLI subcommands
// ---------------------------------------------------------------------------

/// Bundle of flags from `studio-worker register`.
#[derive(Debug, Clone, Default)]
pub struct RegisterArgs {
    pub api_base_url: Option<String>,
    pub reset: bool,
}

/// Persist registration metadata for the next launch.  No HTTP — the
/// auto-register orchestration inside `run` / `ui` is the only thing
/// that talks to the studio.
pub async fn register(config_path: Option<&str>, args: RegisterArgs) -> Result<()> {
    let (mut cfg, path) = config::load(config_path)?;

    if args.reset {
        cfg.worker_id = None;
        cfg.auth_token = None;
        cfg.registration_request_id = None;
        cfg.registration_secret = None;
        cfg.install_id = None;
    }
    if let Some(url) = args.api_base_url {
        cfg.api_base_url = url;
    }

    config::save(&cfg, &path)?;
    if args.reset {
        info!(
            config_path = %path.display(),
            "local registration state cleared; next launch will auto-register"
        );
        println!(
            "local registration state cleared; run `studio-worker run` or \
             `studio-worker ui` to auto-register"
        );
    } else {
        info!(
            config_path = %path.display(),
            "register flags persisted; next launch will auto-register"
        );
        println!(
            "saved; run `studio-worker run` or `studio-worker ui` to auto-register against {}",
            cfg.api_base_url
        );
    }
    Ok(())
}

pub async fn status(config_path: Option<&str>) -> Result<()> {
    let (cfg, path) = config::load(config_path)?;
    println!("{}", format_status(&cfg, &path));
    Ok(())
}

pub fn format_status(cfg: &Config, path: &std::path::Path) -> String {
    let mut out = String::new();
    use std::fmt::Write as _;
    let _ = writeln!(out, "config path:        {}", path.display());
    let _ = writeln!(out, "api_base_url:       {}", cfg.api_base_url);
    let registration_line = if cfg.worker_id.is_some() && cfg.auth_token.is_some() {
        format!("approved as {}", cfg.worker_id.as_deref().unwrap_or(""))
    } else if let Some(rid) = cfg.registration_request_id.as_deref() {
        format!("pending operator approval (request {rid})")
    } else {
        "not registered (will auto-register on next launch)".into()
    };
    let _ = writeln!(out, "registration:       {registration_line}");
    let _ = writeln!(out, "vram_threshold_gb:  {}", cfg.vram_threshold_gb);
    let _ = writeln!(out, "auto_start:         {}", cfg.auto_start);
    let _ = writeln!(out, "models_root:        {}", cfg.models_root.display());
    let _ = writeln!(out, "auto_update:        {}", cfg.auto_update_enabled);
    let _ = writeln!(
        out,
        "update_interval:    {}s",
        cfg.auto_update_interval_secs
    );
    out
}

pub fn set_threshold(config_path: Option<&str>, gb: f32) -> Result<()> {
    if gb < 0.0 {
        return Err(anyhow!("threshold must be >= 0"));
    }
    let (mut cfg, path) = config::load(config_path)?;
    cfg.vram_threshold_gb = gb;
    config::save(&cfg, &path)?;
    info!(
        target: TRACE_TARGET,
        op = "set_threshold",
        vram_threshold_gb = gb,
        config_path = path.display().to_string(),
        "VRAM threshold persisted"
    );
    println!("vram_threshold_gb = {gb}");
    Ok(())
}

/// Emit a one-shot startup banner so operators can confirm which
/// config the worker actually loaded.  Without this the only thing in
/// `journalctl -u studio-worker` on a healthy boot is whatever the
/// loops happen to log on their first tick.
pub fn log_startup_banner(cfg: &Config, path: &std::path::Path) {
    info!(
        target: TRACE_TARGET,
        op = "startup",
        version = AGENT_VERSION,
        config_path = path.display().to_string(),
        api_base_url = cfg.api_base_url.as_str(),
        vram_threshold_gb = cfg.vram_threshold_gb,
        auto_start = cfg.auto_start,
        auto_update_enabled = cfg.auto_update_enabled,
        auto_update_interval_secs = cfg.auto_update_interval_secs,
        models_root = cfg.models_root.display().to_string(),
        worker_id = cfg.worker_id.as_deref().unwrap_or("(unregistered)"),
        "studio-worker booting"
    );
}

pub fn show_config(config_path: Option<&str>) -> Result<()> {
    let (cfg, path) = config::load(config_path)?;
    println!("# {}", path.display());
    print!("{}", toml::to_string_pretty(&cfg)?);
    Ok(())
}

pub async fn check_update(config_path: Option<&str>) -> Result<()> {
    let (cfg, _) = config::load(config_path)?;
    let current = semver::Version::parse(AGENT_VERSION)
        .map_err(|e| anyhow!("invalid current version {AGENT_VERSION}: {e}"))?;
    let outcome = tokio::task::spawn_blocking(move || {
        update::check(&cfg.auto_update_feed, &current, cfg.auto_update_prerelease)
    })
    .await??;
    println!("{}", format_check_outcome(&outcome));
    Ok(())
}

pub fn format_check_outcome(outcome: &update::CheckOutcome) -> String {
    match outcome {
        update::CheckOutcome::UpToDate { current } => format!("up to date: {current}"),
        update::CheckOutcome::NewerAvailable { current, latest } => {
            format!("update available: {current} -> {latest}")
        }
    }
}

// ---------------------------------------------------------------------------
// Long-running run loop
// ---------------------------------------------------------------------------

pub async fn run(config_path: Option<&str>) -> Result<()> {
    let (cfg, path) = config::load(config_path)?;
    log_startup_banner(&cfg, &path);

    let cfg = config::shared(cfg);
    let stop = Arc::new(AtomicBool::new(false));
    let busy = Arc::new(AtomicBool::new(false));
    // Operator pause toggle.  Runtime-only — never persisted, so the
    // worker comes up unpaused after every restart.
    let paused = Arc::new(AtomicBool::new(false));
    let logs: Arc<Mutex<Vec<LogEntry>>> = Arc::new(Mutex::new(Vec::new()));
    let observers = WorkerObservers::default();
    let registration = crate::auto_register::shared_initial();

    let stop_clone = stop.clone();
    tokio::spawn(async move {
        let signal = wait_for_shutdown_signal().await;
        request_shutdown(&stop_clone, signal);
    });

    // Block on auto-register until the operator approves (or rejects).
    // Polls every 30s; aborts on Ctrl-C.  A stop signal that arrives
    // before approval is a clean shutdown (the pre-approval wait is the
    // normal state of a fresh worker), so exit Ok rather than letting
    // `run_cli` log it at error and exit non-zero.
    if ensure_registered(&cfg, &path, &registration, &stop).await? == RegistrationGate::Stopped {
        info!(
            target: TRACE_TARGET,
            op = "shutdown",
            "stopped before registration completed; exiting cleanly"
        );
        return Ok(());
    }

    run_loops(
        cfg,
        stop,
        logs,
        busy,
        paused,
        observers,
        LoopSchedule::default(),
    )
    .await
}

/// Flip the `stop` flag and emit a shutdown breadcrumb so an operator
/// tailing the journal sees a clean stop, mirroring
/// [`log_startup_banner`].  Pulled out of the signal task so the
/// shutdown decision is unit-testable without delivering a real OS
/// signal.  `signal` names whatever woke us (e.g. `"SIGTERM"`).
pub fn request_shutdown(stop: &AtomicBool, signal: &str) {
    let already_stopping = stop.swap(true, Ordering::SeqCst);
    info!(
        target: TRACE_TARGET,
        op = "shutdown",
        signal,
        already_stopping,
        "shutdown signal received; stopping worker gracefully"
    );
}

/// Block until the OS asks the worker to stop, returning the name of
/// the signal that fired.
///
/// On Unix we wait on **both** SIGINT (interactive Ctrl-C) and SIGTERM.
/// SIGTERM is the signal `systemctl stop` / `launchctl unload` / host
/// shutdown deliver by default, and the worker ships as a `Type=simple`
/// systemd unit (see `service::render_service`).  Listening for Ctrl-C
/// alone meant the service manager's stop never reached the graceful
/// path: the WS session was killed mid-`close`, the studio saw an
/// abrupt disconnect, and the final log batch never flushed.  If the
/// SIGTERM handler can't be installed we degrade to Ctrl-C only rather
/// than abort the shutdown task.
///
/// On non-Unix we wait on Ctrl-C, which tokio maps to the console
/// Ctrl-C / close events.
#[cfg_attr(coverage_nightly, coverage(off))]
async fn wait_for_shutdown_signal() -> &'static str {
    #[cfg(unix)]
    {
        use tokio::signal::unix::{signal, SignalKind};
        let mut sigterm = match signal(SignalKind::terminate()) {
            Ok(s) => s,
            Err(e) => {
                warn!(
                    target: TRACE_TARGET,
                    op = "shutdown",
                    error = %e,
                    "could not install SIGTERM handler; falling back to Ctrl-C only"
                );
                let _ = tokio::signal::ctrl_c().await;
                return "SIGINT";
            }
        };
        tokio::select! {
            _ = tokio::signal::ctrl_c() => "SIGINT",
            _ = sigterm.recv() => "SIGTERM",
        }
    }
    #[cfg(not(unix))]
    {
        let _ = tokio::signal::ctrl_c().await;
        "ctrl-c"
    }
}

/// Outcome of the startup registration gate ([`ensure_registered`]).
///
/// A clean stop signal (Ctrl-C / SIGTERM) that arrives **before** the
/// studio approves the worker is a routine shutdown, not a failure:
/// the pre-approval wait is the normal state of a freshly-installed
/// worker sitting in the studio's approval queue.  Surfacing it as a
/// distinct [`Stopped`](RegistrationGate::Stopped) outcome lets `run`
/// exit 0 — so `systemctl stop` doesn't mark the unit failed — and
/// skip the top-level `tracing::error!` that would otherwise ship a
/// spurious Sentry event on every clean stop of an unapproved worker.
/// An operator *rejection*, by contrast, stays a hard `Err`: it's a
/// terminal state the operator must act on (`register --reset`).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RegistrationGate {
    /// The worker is registered (already, or just approved); proceed
    /// to open the WS session.
    Ready,
    /// A stop signal arrived before approval; shut down cleanly.
    Stopped,
}

/// Loop auto_register::tick on a 30s cadence until `worker_id` +
/// `auth_token` are populated (Approved → [`RegistrationGate::Ready`]),
/// a stop signal arrives (→ [`RegistrationGate::Stopped`]), or the
/// operator rejects the worker (→ `Err` with recovery guidance).
pub async fn ensure_registered(
    cfg: &SharedConfig,
    path: &std::path::Path,
    registration: &crate::auto_register::SharedRegistration,
    stop: &Arc<AtomicBool>,
) -> Result<RegistrationGate> {
    use std::time::Duration;
    loop {
        if stop.load(Ordering::SeqCst) {
            return Ok(RegistrationGate::Stopped);
        }
        {
            let snap = cfg.lock();
            if snap.worker_id.is_some() && snap.auth_token.is_some() {
                return Ok(RegistrationGate::Ready);
            }
        }
        let state = crate::auto_register::tick(cfg, path, registration).await;
        match state {
            crate::auto_register::RegistrationState::Approved => {
                return Ok(RegistrationGate::Ready)
            }
            crate::auto_register::RegistrationState::Rejected { reason } => {
                return Err(anyhow!(
                    "registration rejected by the studio operator: {reason}.  \
                     Run `studio-worker register --reset` to clear local state \
                     and submit a fresh request."
                ));
            }
            _ => {}
        }
        // Sleep with a fast-cancel on stop.
        for _ in 0..30 {
            if stop.load(Ordering::SeqCst) {
                return Ok(RegistrationGate::Stopped);
            }
            tokio::time::sleep(Duration::from_secs(1)).await;
        }
    }
}

/// Spawn the WS session + auto-updater, wait for them.  Pulled out of
/// `run` so tests can drive with a different schedule.
///
/// `paused` is the runtime-only Pause / Resume toggle the UI flips.
/// When set, the WS session advertises `auto_enabled = false` in
/// heartbeats and refuses new job offers without restarting the
/// session.
pub async fn run_loops(
    cfg: SharedConfig,
    stop: Arc<AtomicBool>,
    logs: Arc<Mutex<Vec<LogEntry>>>,
    busy: Arc<AtomicBool>,
    paused: Arc<AtomicBool>,
    observers: WorkerObservers,
    schedule: LoopSchedule,
) -> Result<()> {
    let session = crate::ws::session::spawn_ws_session(
        cfg.clone(),
        stop.clone(),
        logs.clone(),
        busy.clone(),
        paused.clone(),
        observers.clone(),
        schedule.ws_session,
    );
    let auto_updater = spawn_auto_updater(
        cfg.clone(),
        stop.clone(),
        logs.clone(),
        busy.clone(),
        schedule,
    );
    let (session_result, _) = tokio::join!(session, auto_updater);
    session_result
}

// ---------------------------------------------------------------------------
// Per-tick helpers — pure async fns, easy to drive from unit tests.
// ---------------------------------------------------------------------------

// (The old per-tick HTTP helpers — heartbeat_tick, claim_tick, log_shipper_tick,
//  run_job, ClaimOutcome — lived here.  They are gone with the WS migration.
//  See `ws::session::spawn_ws_session` for the replacement that runs the
//  whole session in one connected loop.)

/// What the auto-updater decided this tick.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AutoUpdateDecision {
    /// Auto-update is turned off — do nothing.
    Disabled,
    /// Worker is currently running a job — skip.
    SkippedBusy,
    /// Local version is already the latest.
    UpToDate,
    /// Check failed (network etc.) — leave a log entry, try again later.
    CheckError(String),
    /// A newer version was applied successfully.  Caller should restart.
    Updated,
    /// A newer version was found but the install failed.
    UpdateError(String),
}

pub async fn auto_update_tick(
    cfg: &Config,
    busy: bool,
    logs: &Arc<Mutex<Vec<LogEntry>>>,
) -> AutoUpdateDecision {
    if !cfg.auto_update_enabled {
        return AutoUpdateDecision::Disabled;
    }
    if busy {
        push_log(
            logs,
            "info",
            "auto-update",
            "skipping check: worker is busy on a job",
            None,
        );
        return AutoUpdateDecision::SkippedBusy;
    }
    let feed = cfg.auto_update_feed.clone();
    let prerelease = cfg.auto_update_prerelease;
    let logs_for_task = logs.clone();
    let outcome = tokio::task::spawn_blocking(move || -> Result<AutoUpdateDecision> {
        let current = semver::Version::parse(AGENT_VERSION)
            .map_err(|e| anyhow!("invalid AGENT_VERSION {AGENT_VERSION}: {e}"))?;
        match update::check(&feed, &current, prerelease) {
            Ok(update::CheckOutcome::UpToDate { current }) => {
                push_log(
                    &logs_for_task,
                    "info",
                    "auto-update",
                    &format!("up to date at {current}"),
                    None,
                );
                Ok(AutoUpdateDecision::UpToDate)
            }
            Ok(update::CheckOutcome::NewerAvailable { current, latest }) => {
                push_log(
                    &logs_for_task,
                    "info",
                    "auto-update",
                    &format!("update available {current} -> {latest}; applying"),
                    None,
                );
                match update::apply(&feed, &latest) {
                    Ok(()) => {
                        push_log(
                            &logs_for_task,
                            "info",
                            "auto-update",
                            "binary replaced; restart pending",
                            None,
                        );
                        Ok(AutoUpdateDecision::Updated)
                    }
                    Err(e) => {
                        push_log(
                            &logs_for_task,
                            "error",
                            "auto-update",
                            &format!("update failed: {e}"),
                            None,
                        );
                        Ok(AutoUpdateDecision::UpdateError(e.to_string()))
                    }
                }
            }
            Err(e) => {
                push_log(
                    &logs_for_task,
                    "warn",
                    "auto-update",
                    &format!("check failed: {e}"),
                    None,
                );
                Ok(AutoUpdateDecision::CheckError(e.to_string()))
            }
        }
    })
    .await;
    match outcome {
        Ok(Ok(decision)) => decision,
        Ok(Err(e)) => AutoUpdateDecision::CheckError(e.to_string()),
        Err(e) => AutoUpdateDecision::CheckError(e.to_string()),
    }
}

// ---------------------------------------------------------------------------
// Long-running task wrappers — they exist solely to call the ticks in a
// loop on a schedule.  All real logic lives in the ticks.
// ---------------------------------------------------------------------------

// (`spawn_heartbeat`, `spawn_claim_loop`, `spawn_log_shipper`, and
//  `next_delay_for` lived here.  Their behaviour is now carried by the
//  WS-driven tasks in `ws::session`.)

/// Sleep up to `total`, re-checking `stop` every `tick` and returning
/// the instant a shutdown is requested.  Keeps long idle waits (the
/// auto-update tick here, reconnect backoff in the WS session)
/// responsive to SIGTERM / SIGINT without busy-looping.  Shared by the
/// runtime auto-updater and `ws::session`.
pub(crate) async fn wait_with_stop(total: Duration, stop: &Arc<AtomicBool>, tick: Duration) {
    let mut elapsed = Duration::ZERO;
    while elapsed < total {
        if stop.load(Ordering::SeqCst) {
            return;
        }
        let next = tick.min(total - elapsed);
        tokio::time::sleep(next).await;
        elapsed += next;
    }
}

pub fn spawn_auto_updater(
    cfg: SharedConfig,
    stop: Arc<AtomicBool>,
    logs: Arc<Mutex<Vec<LogEntry>>>,
    busy: Arc<AtomicBool>,
    schedule: LoopSchedule,
) -> tokio::task::JoinHandle<()> {
    tokio::spawn(async move {
        let mut elapsed = Duration::from_secs(0);
        while !stop.load(Ordering::SeqCst) {
            // Stop-aware idle wait: a shutdown signal during this window
            // wakes the loop within `schedule.shutdown_tick` instead of
            // leaving `run_loops`' join() blocked for a full
            // `auto_update_tick`.
            wait_with_stop(schedule.auto_update_tick, &stop, schedule.shutdown_tick).await;
            if stop.load(Ordering::SeqCst) {
                break;
            }
            elapsed += schedule.auto_update_tick;
            let snapshot = cfg.lock().clone();
            if elapsed < Duration::from_secs(snapshot.auto_update_interval_secs) {
                continue;
            }
            elapsed = Duration::from_secs(0);
            let busy_now = busy.load(Ordering::SeqCst);
            let decision = auto_update_tick(&snapshot, busy_now, &logs).await;
            if matches!(decision, AutoUpdateDecision::Updated) {
                stop.store(true, Ordering::SeqCst);
                update::restart_self();
            }
        }
    })
}

// (`run_job` lived here.  See `ws::session::run_offered_job` for the
//  WS-driven replacement.)

pub fn prompt_for(task: &Task) -> String {
    match task {
        Task::Image(p) => p.prompt.clone(),
        Task::Llm(p) => p
            .messages
            .last()
            .map(|m| m.content.clone())
            .unwrap_or_default(),
        Task::AudioStt(p) => p.input_url.clone(),
        Task::AudioTts(p) => p.text.clone(),
        Task::Video(p) => p.prompt.clone(),
    }
}

pub fn is_unsupported_kind(e: &anyhow::Error) -> bool {
    // Typed check first — survives context wrapping and rewording.
    // The string check remains as a fallback for error paths that
    // haven't migrated to `engine::UnsupportedTask` yet.
    e.chain().any(|cause| {
        cause
            .downcast_ref::<crate::engine::UnsupportedTask>()
            .is_some()
    }) || e.to_string().contains("cannot serve")
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

pub fn build_capabilities(cfg: &Config, engine: &dyn Engine) -> WorkerCapabilities {
    build_capabilities_with(cfg, engine, true)
}

/// Same as [`build_capabilities`] but lets the caller drive
/// `auto_enabled` from a runtime pause flag (the UI's Pause/Resume
/// button).  The persisted [`Config`] no longer carries that bit —
/// it's an in-process toggle.
pub fn build_capabilities_with(
    cfg: &Config,
    engine: &dyn Engine,
    auto_enabled: bool,
) -> WorkerCapabilities {
    let vram = sys::detect_vram_gb().unwrap_or(0.0);
    let caps = engine.capabilities();
    let supported_models_per_kind = caps.supported_models_per_kind.clone();
    let task_kinds = caps.kinds();
    // Legacy `supported_models` is a flat list across all kinds so the
    // studio API's claim filter (which only knows about this field) can
    // match jobs of any modality this worker can serve.
    let supported_models = {
        let mut all = caps.flat_models();
        all.sort();
        all.dedup();
        all
    };

    WorkerCapabilities {
        machine_name: sys::machine_name(),
        username: sys::username(),
        agent_version: AGENT_VERSION.to_string(),
        engine: engine.name().to_string(),
        vram_total_gb: vram,
        vram_threshold_gb: cfg.vram_threshold_gb,
        auto_enabled,
        auto_start: cfg.auto_start,
        supported_models,
        task_kinds,
        supported_models_per_kind,
    }
}

/// One-line, operator-facing summary of what this worker advertises to
/// the studio on the WS handshake.  Logged once per session attempt so
/// the worker's own logs (and the studio's shipped-log view) record
/// exactly which task kinds, models, and VRAM budget were offered — the
/// missing complement to [`log_startup_banner`], which only covers the
/// loaded config.  Without it, an operator chasing "why won't my worker
/// claim image jobs" has no record of what the worker told the studio
/// it could do.  Pure so the formatting is unit-tested without a live
/// session.
pub fn summarize_capabilities(caps: &WorkerCapabilities) -> String {
    let kinds = caps
        .task_kinds
        .iter()
        .map(|k| k.as_str())
        .collect::<Vec<_>>()
        .join(", ");
    format!(
        "advertising engine={}, vram={:.1}/{:.1}GB threshold, auto_enabled={}, \
         kinds=[{}], {} model(s)=[{}]",
        caps.engine,
        caps.vram_total_gb,
        caps.vram_threshold_gb,
        caps.auto_enabled,
        kinds,
        caps.supported_models.len(),
        caps.supported_models.join(", "),
    )
}

/// Operator-facing warning when the configured VRAM threshold exceeds
/// the GPU VRAM the worker actually detected.
///
/// The studio matches jobs to a worker purely by its advertised
/// `vram_threshold_gb`, so a threshold set above the card's real
/// capacity — e.g. the default 12 GB on an 8 GB consumer GPU — makes the
/// worker accept jobs its GPU can't fit: they load, exhaust VRAM, and
/// fail with an OOM the operator then has to trace back to a config
/// value.  Surfacing it on the handshake (one line next to the
/// capability summary) turns a silent OOM-on-claim into an actionable
/// "lower your threshold" breadcrumb.
///
/// Only fires when the VRAM probe returned a real positive total: a
/// detected 0 GB means the probe failed (no `nvidia-smi` / sysfs tree,
/// or a non-NVIDIA GPU we can't size), where the threshold is the only
/// capacity signal we have and second-guessing it would be wrong.  The
/// boundary is strict (`threshold > total`), so a threshold that exactly
/// matches the card stays silent.  Pure so the wording + boundary are
/// unit-tested without a live GPU.
pub fn vram_threshold_warning(caps: &WorkerCapabilities) -> Option<String> {
    if caps.vram_total_gb > 0.0 && caps.vram_threshold_gb > caps.vram_total_gb {
        Some(format!(
            "configured VRAM threshold {:.1}GB exceeds detected GPU VRAM {:.1}GB; \
             the studio may offer jobs larger than this card can fit and they will \
             OOM on load — lower vram_threshold_gb to at or below {:.1}GB",
            caps.vram_threshold_gb, caps.vram_total_gb, caps.vram_total_gb
        ))
    } else {
        None
    }
}

pub fn push_log(
    logs: &Arc<Mutex<Vec<LogEntry>>>,
    level: &str,
    category: &str,
    message: &str,
    job_id: Option<String>,
) {
    push_log_with_observers(logs, None, level, category, message, job_id);
}

/// Same as [`push_log`] but also appends to
/// [`WorkerObservers::recent_logs`] so the UI's Logs tab keeps a
/// rolling display window.  The WS session uses this variant so
/// operators don't see the Logs tab blank out every second when the
/// shipping queue gets drained.
pub fn push_log_with_observers(
    logs: &Arc<Mutex<Vec<LogEntry>>>,
    observers: Option<&WorkerObservers>,
    level: &str,
    category: &str,
    message: &str,
    job_id: Option<String>,
) {
    let entry = LogEntry {
        ts: Utc::now().to_rfc3339_opts(SecondsFormat::Millis, true),
        level: level.to_string(),
        category: category.to_string(),
        message: message.to_string(),
        job_id,
    };
    // Carry the job id as a structured field so operators can pivot
    // shipped studio logs / Sentry breadcrumbs on it. `Option<&str>`
    // only records the field when `Some`, so jobless breadcrumbs stay
    // free of a noisy empty `job_id`.
    let job_id = entry.job_id.as_deref();
    if level == "error" {
        tracing::error!(target: "studio_worker", job_id, "[{category}] {message}");
    } else if level == "warn" {
        tracing::warn!(target: "studio_worker", job_id, "[{category}] {message}");
    } else {
        info!(target: "studio_worker", job_id, "[{category}] {message}");
    }
    {
        let mut queue = logs.lock();
        if queue.len() >= LOG_SHIP_QUEUE_CAP {
            // +1 for the entry below, +1 for the drop marker.
            let overflow = queue.len() + 2 - LOG_SHIP_QUEUE_CAP;
            queue.drain(0..overflow);
            queue.push(LogEntry {
                ts: Utc::now().to_rfc3339_opts(SecondsFormat::Millis, true),
                level: "warn".to_string(),
                category: "logs".to_string(),
                message: format!(
                    "ship queue full ({LOG_SHIP_QUEUE_CAP} entries); dropped {overflow} oldest"
                ),
                job_id: None,
            });
        }
        queue.push(entry.clone());
    }
    if let Some(o) = observers {
        let mut ring = o.recent_logs.lock();
        ring.push_back(entry);
        while ring.len() > RECENT_LOGS_CAP {
            ring.pop_front();
        }
    }
}

/// Put a drained-but-unsent batch back at the front of the ship queue
/// so it survives for the next session attempt.  Entries that arrived
/// while the batch was in flight stay behind it (newest last).  The
/// combined queue is clipped to [`LOG_SHIP_QUEUE_CAP`], dropping the
/// oldest entries first.
pub fn restore_unshipped(logs: &Arc<Mutex<Vec<LogEntry>>>, mut batch: Vec<LogEntry>) {
    let mut queue = logs.lock();
    batch.append(&mut queue);
    *queue = batch;
    if queue.len() > LOG_SHIP_QUEUE_CAP {
        let overflow = queue.len() - LOG_SHIP_QUEUE_CAP;
        queue.drain(0..overflow);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::Config;
    use crate::engine::SyntheticEngine;

    #[test]
    fn is_unsupported_kind_detects_typed_unsupported_task() {
        let err: anyhow::Error =
            crate::engine::UnsupportedTask::new("synthetic", TaskKind::Llm).into();
        assert!(is_unsupported_kind(&err));
        // The message keeps the legacy operator-facing shape.
        assert!(err.to_string().contains("cannot serve llm"));
    }

    #[test]
    fn is_unsupported_kind_survives_context_wrapping() {
        // String sniffing broke as soon as a caller added context (the
        // outer message no longer contains "cannot serve"); the typed
        // downcast searches the whole chain.
        let err = anyhow::Error::from(crate::engine::UnsupportedTask::new(
            "sdcpp",
            TaskKind::AudioTts,
        ))
        .context("dispatching job j-1");
        assert!(is_unsupported_kind(&err));
    }

    fn entry(message: &str) -> LogEntry {
        LogEntry {
            ts: Utc::now().to_rfc3339_opts(SecondsFormat::Millis, true),
            level: "info".into(),
            category: "test".into(),
            message: message.into(),
            job_id: None,
        }
    }

    #[test]
    fn restore_unshipped_requeues_batch_ahead_of_newer_entries() {
        // A batch the shipper drained but failed to send must survive
        // for the next session, ordered before entries that arrived
        // while it was in flight.
        let logs: Arc<Mutex<Vec<LogEntry>>> = Arc::new(Mutex::new(vec![entry("newer")]));
        restore_unshipped(&logs, vec![entry("batch-1"), entry("batch-2")]);
        let queue = logs.lock();
        let order: Vec<&str> = queue.iter().map(|e| e.message.as_str()).collect();
        assert_eq!(order, vec!["batch-1", "batch-2", "newer"]);
    }

    #[test]
    fn restore_unshipped_respects_the_queue_cap() {
        // Requeueing must never grow the queue past the ship cap; the
        // oldest (front) entries give way so the newest survive.
        let logs: Arc<Mutex<Vec<LogEntry>>> =
            Arc::new(Mutex::new(vec![entry("newest"); LOG_SHIP_QUEUE_CAP]));
        restore_unshipped(&logs, vec![entry("old-batch"); 100]);
        let queue = logs.lock();
        assert_eq!(queue.len(), LOG_SHIP_QUEUE_CAP);
        assert_eq!(
            queue.last().map(|e| e.message.as_str()),
            Some("newest"),
            "newest entries must survive the cap"
        );
    }

    #[test]
    fn ship_queue_is_bounded_and_records_dropped_entries() {
        // The WS shipper only drains while a session is connected; a
        // long approval wait / reconnect backoff must not grow the
        // queue without bound.
        let logs: Arc<Mutex<Vec<LogEntry>>> = Arc::new(Mutex::new(Vec::new()));
        for i in 0..(LOG_SHIP_QUEUE_CAP + 100) {
            push_log_with_observers(&logs, None, "info", "test", &format!("entry {i}"), None);
        }
        let queue = logs.lock();
        assert!(
            queue.len() <= LOG_SHIP_QUEUE_CAP,
            "ship queue exceeded its cap: {}",
            queue.len()
        );
        // The newest entry always survives.
        assert_eq!(
            queue.last().map(|e| e.message.as_str()),
            Some(format!("entry {}", LOG_SHIP_QUEUE_CAP + 99).as_str())
        );
        // Loss is visible: a marker entry names how many were dropped.
        assert!(
            queue
                .iter()
                .any(|e| e.level == "warn" && e.message.contains("dropped")),
            "overflow must leave a visible drop marker"
        );
    }

    #[test]
    fn recent_logs_ring_is_bounded_at_recent_logs_cap() {
        // The observer ring backing the UI Logs tab is never drained
        // (unlike the ship queue, which the WS shipper empties every
        // second), so this cap is its only bound.  A regression that
        // dropped the eviction loop would leak memory for the lifetime
        // of a long-running worker; one that flipped `pop_front` for
        // `pop_back` would silently retain the *oldest* entries and show
        // a stale Logs tab.  Mirrors `recent_jobs_ring_caps_at_*`.
        let logs: Arc<Mutex<Vec<LogEntry>>> = Arc::new(Mutex::new(Vec::new()));
        let observers = WorkerObservers::default();
        let overflow = 25;
        for i in 0..(RECENT_LOGS_CAP + overflow) {
            push_log_with_observers(
                &logs,
                Some(&observers),
                "info",
                "test",
                &format!("entry {i}"),
                None,
            );
        }
        let ring = observers.recent_logs.lock();
        assert_eq!(
            ring.len(),
            RECENT_LOGS_CAP,
            "the recent-logs ring must cap at RECENT_LOGS_CAP"
        );
        // Newest entries go to the back; the oldest `overflow` entries
        // must have been evicted from the front.
        assert_eq!(
            ring.back().map(|e| e.message.as_str()),
            Some(format!("entry {}", RECENT_LOGS_CAP + overflow - 1).as_str()),
            "the newest entry must survive at the back of the ring"
        );
        assert_eq!(
            ring.front().map(|e| e.message.as_str()),
            Some(format!("entry {overflow}").as_str()),
            "the oldest surviving entry must be entry #overflow (older evicted)"
        );
    }

    #[test]
    fn capabilities_advertises_all_synthetic_kinds() {
        let cfg = Config::default();
        let engine = SyntheticEngine::new();
        let cap = build_capabilities(&cfg, &engine);
        assert_eq!(cap.engine, "synthetic");
        assert_eq!(cap.task_kinds.len(), TaskKind::ALL.len());
        assert!(cap.auto_enabled, "default capability snapshot is unpaused");
        for kind in TaskKind::ALL {
            assert!(cap.supported_models_per_kind.contains_key(&kind));
        }
    }

    #[test]
    fn capabilities_with_paused_flag_drives_auto_enabled() {
        let cfg = Config::default();
        let engine = SyntheticEngine::new();
        let paused_caps = build_capabilities_with(&cfg, &engine, false);
        assert!(!paused_caps.auto_enabled);
    }

    #[test]
    fn summarize_capabilities_lists_engine_kinds_models_vram_and_pause_state() {
        let cfg = Config {
            vram_threshold_gb: 6.0,
            ..Config::default()
        };
        let engine = SyntheticEngine::new();
        let caps = build_capabilities_with(&cfg, &engine, true);
        let summary = summarize_capabilities(&caps);
        // Engine name + every advertised kind is present.
        assert!(summary.contains("engine=synthetic"), "got: {summary}");
        for kind in &caps.task_kinds {
            assert!(
                summary.contains(kind.as_str()),
                "missing kind {} in: {summary}",
                kind.as_str()
            );
        }
        // Model count + an actual advertised model id are present.
        assert!(
            summary.contains(&format!("{} model(s)", caps.supported_models.len())),
            "missing model count in: {summary}"
        );
        assert!(
            summary.contains("synthetic"),
            "missing model id in: {summary}"
        );
        // VRAM budget (total/threshold) + unpaused state are visible.
        assert!(
            summary.contains("6.0"),
            "missing vram threshold in: {summary}"
        );
        assert!(summary.contains("auto_enabled=true"), "got: {summary}");
    }

    #[test]
    fn summarize_capabilities_reflects_paused_state() {
        let cfg = Config::default();
        let engine = SyntheticEngine::new();
        let caps = build_capabilities_with(&cfg, &engine, false);
        assert!(
            summarize_capabilities(&caps).contains("auto_enabled=false"),
            "paused worker must advertise auto_enabled=false"
        );
    }

    /// Build a capability snapshot, then override the two VRAM fields so
    /// the threshold/total relationship is deterministic regardless of
    /// the host's real GPU (the probe is `0.0` on CI).
    fn caps_with_vram(total_gb: f32, threshold_gb: f32) -> WorkerCapabilities {
        let mut caps = build_capabilities_with(&Config::default(), &SyntheticEngine::new(), true);
        caps.vram_total_gb = total_gb;
        caps.vram_threshold_gb = threshold_gb;
        caps
    }

    #[test]
    fn vram_threshold_warning_flags_threshold_above_detected_vram() {
        // The default 12 GB threshold on an 8 GB card: the studio will
        // offer up-to-12 GB jobs this GPU can't fit, and they OOM on
        // load.  The breadcrumb must name both numbers and the config
        // key the operator has to lower.
        let warning = vram_threshold_warning(&caps_with_vram(8.0, 12.0))
            .expect("threshold above detected VRAM must warn");
        assert!(warning.contains("12.0"), "missing threshold in: {warning}");
        assert!(
            warning.contains("8.0"),
            "missing detected VRAM in: {warning}"
        );
        assert!(
            warning.contains("vram_threshold_gb"),
            "must name the config key to change: {warning}"
        );
    }

    #[test]
    fn vram_threshold_warning_silent_when_threshold_within_detected_vram() {
        // A 24 GB card with a 12 GB threshold is correctly conservative.
        assert!(vram_threshold_warning(&caps_with_vram(24.0, 12.0)).is_none());
    }

    #[test]
    fn vram_threshold_warning_silent_when_threshold_equals_detected() {
        // The boundary is strict: a threshold that exactly matches the
        // card fits, so it stays silent.
        assert!(vram_threshold_warning(&caps_with_vram(12.0, 12.0)).is_none());
    }

    #[test]
    fn vram_threshold_warning_silent_when_vram_undetected() {
        // A detected 0 GB means the probe failed (no nvidia-smi / sysfs)
        // or it's a non-NVIDIA GPU we can't size; the threshold is then
        // the only capacity signal we have, so second-guessing it with a
        // spurious OOM warning would be wrong.
        assert!(vram_threshold_warning(&caps_with_vram(0.0, 12.0)).is_none());
    }

    #[test]
    fn prompt_for_extracts_per_kind() {
        let image = Task::Image(ImageParams {
            prompt: "a stone golem".into(),
            ..Default::default()
        });
        assert_eq!(prompt_for(&image), "a stone golem");

        let llm = Task::Llm(LlmParams {
            messages: vec![
                ChatMessage {
                    role: "system".into(),
                    content: "be helpful".into(),
                },
                ChatMessage {
                    role: "user".into(),
                    content: "hi".into(),
                },
            ],
            max_tokens: 32,
            temperature: 0.5,
            ..Default::default()
        });
        assert_eq!(prompt_for(&llm), "hi");

        let llm_empty = Task::Llm(LlmParams {
            messages: vec![],
            ..Default::default()
        });
        assert_eq!(prompt_for(&llm_empty), "");

        let stt = Task::AudioStt(AudioSttParams {
            input_url: "https://example.com/clip.wav".into(),
            ..Default::default()
        });
        assert_eq!(prompt_for(&stt), "https://example.com/clip.wav");

        let tts = Task::AudioTts(AudioTtsParams {
            text: "hi there".into(),
            voice: "v".into(),
            ext: "wav".into(),
            ..Default::default()
        });
        assert_eq!(prompt_for(&tts), "hi there");

        let video = Task::Video(VideoParams {
            prompt: "a tiny dragon".into(),
            seconds: 1.0,
            width: 256,
            height: 256,
            ext: "mp4".into(),
            ..Default::default()
        });
        assert_eq!(prompt_for(&video), "a tiny dragon");
    }

    #[test]
    fn truncate_prompt_passes_short_through_and_clips_long_prompts() {
        // Under the cap → returned verbatim, no ellipsis.
        let short = "a stone golem";
        assert_eq!(truncate_prompt(short), short);

        // Exactly at the cap is the boundary: still untouched.
        let exactly = "x".repeat(PROMPT_PREVIEW_CHARS);
        assert_eq!(
            truncate_prompt(&exactly),
            exactly,
            "a prompt exactly at the cap must not be clipped"
        );

        // One past the cap → clipped to PROMPT_PREVIEW_CHARS chars plus
        // the single ellipsis terminator.
        let over = "y".repeat(PROMPT_PREVIEW_CHARS + 1);
        let clipped = truncate_prompt(&over);
        assert_eq!(
            clipped.chars().count(),
            PROMPT_PREVIEW_CHARS + 1,
            "clipped preview is the cap plus one ellipsis char"
        );
        assert!(
            clipped.ends_with('\u{2026}'),
            "a clipped preview ends with an ellipsis"
        );
        assert_eq!(
            clipped
                .chars()
                .take(PROMPT_PREVIEW_CHARS)
                .collect::<String>(),
            "y".repeat(PROMPT_PREVIEW_CHARS),
            "the kept prefix is the first PROMPT_PREVIEW_CHARS chars"
        );
    }

    #[test]
    fn truncate_prompt_clips_on_char_boundaries_for_multibyte_text() {
        // Each char here is 3 bytes, so the cap-th *byte* lands
        // mid-codepoint: a naive `&s[..PROMPT_PREVIEW_CHARS]` byte slice
        // would panic.  `truncate_prompt` counts chars, so a one-over
        // multibyte prompt clips cleanly to the cap plus the ellipsis.
        let multibyte = "\u{3042}".repeat(PROMPT_PREVIEW_CHARS + 1);
        let clipped = truncate_prompt(&multibyte);
        assert_eq!(clipped.chars().count(), PROMPT_PREVIEW_CHARS + 1);
        assert!(clipped.ends_with('\u{2026}'));
        assert_eq!(
            clipped.chars().filter(|c| *c == '\u{3042}').count(),
            PROMPT_PREVIEW_CHARS,
            "exactly PROMPT_PREVIEW_CHARS multibyte chars survive the clip"
        );
    }

    #[test]
    fn is_unsupported_kind_matches_engine_message() {
        let err = anyhow!("multi engine cannot serve llm tasks");
        assert!(is_unsupported_kind(&err));
        let other = anyhow!("network timeout");
        assert!(!is_unsupported_kind(&other));
    }

    #[test]
    fn format_status_includes_every_field() {
        let cfg = Config::default();
        let out = format_status(&cfg, std::path::Path::new("/tmp/x.toml"));
        assert!(out.contains("config path:"));
        assert!(out.contains("api_base_url:"));
        assert!(out.contains("registration:"));
        assert!(out.contains("not registered"));
        assert!(out.contains("models_root:"));
        assert!(out.contains("auto_update:"));
        assert!(out.contains("update_interval:"));
    }

    #[test]
    fn format_status_shows_worker_id_when_registered() {
        let cfg = Config {
            worker_id: Some("w-abc".into()),
            auth_token: Some("tok".into()),
            ..Config::default()
        };
        let out = format_status(&cfg, std::path::Path::new("/tmp/x.toml"));
        assert!(out.contains("w-abc"));
        assert!(out.contains("approved"));
    }

    #[test]
    fn format_status_shows_pending_request_id() {
        let cfg = Config {
            registration_request_id: Some("rr-7".into()),
            ..Config::default()
        };
        let out = format_status(&cfg, std::path::Path::new("/tmp/x.toml"));
        assert!(out.contains("pending operator approval"));
        assert!(out.contains("rr-7"));
    }

    #[test]
    fn format_check_outcome_handles_both_branches() {
        let up = update::CheckOutcome::UpToDate {
            current: semver::Version::new(1, 2, 3),
        };
        assert!(format_check_outcome(&up).contains("up to date"));
        let newer = update::CheckOutcome::NewerAvailable {
            current: semver::Version::new(1, 2, 3),
            latest: semver::Version::new(1, 3, 0),
        };
        let s = format_check_outcome(&newer);
        assert!(s.contains("1.2.3 -> 1.3.0"));
    }

    #[test]
    fn push_log_appends_an_entry() {
        let logs: Arc<Mutex<Vec<LogEntry>>> = Arc::new(Mutex::new(Vec::new()));
        push_log(&logs, "info", "test", "hi", None);
        push_log(&logs, "warn", "test", "wat", Some("j-1".into()));
        push_log(&logs, "error", "test", "boom", None);
        let v = logs.lock();
        assert_eq!(v.len(), 3);
        assert_eq!(v[0].level, "info");
        assert_eq!(v[1].level, "warn");
        assert_eq!(v[1].job_id.as_deref(), Some("j-1"));
        assert_eq!(v[2].level, "error");
    }

    #[test]
    fn push_log_emits_job_id_as_a_structured_tracing_field() {
        // Operators correlating shipped studio logs / Sentry
        // breadcrumbs by job need the job id as a *field*, not just
        // buried in the message text, so `RUST_LOG` filters and Sentry
        // tag search can pivot on it.
        use crate::test_support::capture;
        let logs = capture(|| {
            let logs: Arc<Mutex<Vec<LogEntry>>> = Arc::new(Mutex::new(Vec::new()));
            push_log(
                &logs,
                "info",
                "ws",
                "binary upload ok",
                Some("job-42".into()),
            );
        });
        assert!(
            logs.contains("job_id=\"job-42\""),
            "expected structured job_id field, got: {logs}"
        );
        assert!(
            logs.contains("[ws] binary upload ok"),
            "expected the human-readable message to survive, got: {logs}"
        );
    }

    #[test]
    fn push_log_omits_job_id_field_when_absent() {
        // Jobless breadcrumbs (startup banners, heartbeats, auto-update
        // ticks) must not gain a noisy empty `job_id` field.
        use crate::test_support::capture;
        let logs = capture(|| {
            let logs: Arc<Mutex<Vec<LogEntry>>> = Arc::new(Mutex::new(Vec::new()));
            push_log(&logs, "info", "auto-update", "up to date", None);
        });
        assert!(
            !logs.contains("job_id"),
            "expected no job_id field for a jobless log, got: {logs}"
        );
    }

    // --- async tick tests ---

    #[test]
    fn request_shutdown_sets_the_stop_flag() {
        let stop = AtomicBool::new(false);
        request_shutdown(&stop, "SIGTERM");
        assert!(stop.load(Ordering::SeqCst));
    }

    #[test]
    fn request_shutdown_reconfirms_when_already_stopping() {
        // A second signal (or a race with another shutdown path) must
        // not panic or clear the flag — it just re-confirms the stop.
        let stop = AtomicBool::new(true);
        request_shutdown(&stop, "SIGINT");
        assert!(stop.load(Ordering::SeqCst));
    }

    #[test]
    fn request_shutdown_emits_a_named_shutdown_breadcrumb() {
        use crate::test_support::capture;
        let logs = capture(|| {
            let stop = AtomicBool::new(false);
            request_shutdown(&stop, "SIGTERM");
        });
        assert!(logs.contains("INFO"), "expected INFO event, got: {logs}");
        assert!(
            logs.contains("studio_worker::runtime"),
            "expected runtime target, got: {logs}"
        );
        assert!(
            logs.contains("op=\"shutdown\""),
            "expected op field, got: {logs}"
        );
        assert!(
            logs.contains("signal=\"SIGTERM\""),
            "expected signal field, got: {logs}"
        );
    }

    #[tokio::test]
    async fn auto_update_tick_disabled_when_flag_off() {
        let cfg = Config {
            auto_update_enabled: false,
            ..Config::default()
        };
        let logs = Arc::new(Mutex::new(Vec::new()));
        let decision = auto_update_tick(&cfg, false, &logs).await;
        assert_eq!(decision, AutoUpdateDecision::Disabled);
    }

    #[tokio::test]
    async fn auto_update_tick_skipped_when_busy() {
        let cfg = Config {
            auto_update_enabled: true,
            ..Config::default()
        };
        let logs = Arc::new(Mutex::new(Vec::new()));
        let decision = auto_update_tick(&cfg, true, &logs).await;
        assert_eq!(decision, AutoUpdateDecision::SkippedBusy);
        let entries = logs.lock();
        assert!(entries.iter().any(|e| e.message.contains("busy on a job")));
    }

    #[tokio::test]
    async fn wait_with_stop_short_circuits_when_already_stopped() {
        let stop = Arc::new(AtomicBool::new(true));
        let start = std::time::Instant::now();
        wait_with_stop(Duration::from_secs(60), &stop, Duration::from_millis(10)).await;
        assert!(
            start.elapsed() < Duration::from_millis(100),
            "an already-set stop must return without sleeping the full duration"
        );
    }

    #[tokio::test]
    async fn auto_updater_stops_promptly_during_idle_wait() {
        // A huge auto_update_tick means a non-cancellable idle sleep
        // would pin the JoinHandle — and thus `run_loops`' join() — for
        // the whole tick after stop is set, defeating graceful
        // shutdown.  The stop-aware wait must let the task finish well
        // inside the tick.
        let cfg = crate::config::shared(Config {
            auto_update_enabled: false,
            ..Config::default()
        });
        let stop = Arc::new(AtomicBool::new(false));
        let logs: Arc<Mutex<Vec<LogEntry>>> = Arc::new(Mutex::new(Vec::new()));
        let busy = Arc::new(AtomicBool::new(false));
        let schedule = LoopSchedule {
            ws_session: crate::ws::session::SessionSchedule::fast_for_tests(),
            auto_update_tick: Duration::from_secs(3600),
            shutdown_tick: Duration::from_millis(1),
        };
        let handle = spawn_auto_updater(cfg, stop.clone(), logs, busy, schedule);
        // Let the loop reach its idle wait, then request shutdown.
        tokio::time::sleep(Duration::from_millis(10)).await;
        stop.store(true, Ordering::SeqCst);
        tokio::time::timeout(Duration::from_millis(250), handle)
            .await
            .expect("auto-updater did not observe stop promptly")
            .expect("auto-updater task panicked");
    }
}