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
//! First-run **provisioning plumbing** — the event channel and status lines
//! for the on-host summarizer download on first launch.
//!
//! The binary (newt-cli) provisions the model on a background thread and hands
//! `run_code` a [`SetupHandle`]. The ONE splash screen renders that work as an
//! extra spinner line (`crate::splash::SetupWait`) — there is deliberately no
//! second, separate setup screen (field note: two consecutive splashes read as
//! a bug). [`SetupEvent`]s stream progress from the download thread; the
//! no-splash fallback is [`run_setup_inline`].
//!
//! This is the UI of the setup step. The *logic* lives elsewhere: [`crate::setup`]
//! is the human-driven config wizard (`newt setup`), and [`crate::wizard`] is the
//! silent auto-prober behind `newt init`. Extracted from `lib.rs` in the #1096
//! functional-cohesion pass; the module name says what it is — the UI, not the
//! logic.
use std::io::{self, Write as _};
use crossterm::event::{Event, KeyCode, KeyEvent, KeyModifiers};
/// Progress events from the background provisioning thread to the setup screen.
pub enum SetupEvent {
/// A named step began (e.g. `weights`, `tokenizer`).
Step(String),
/// Byte progress on the current step (`total` unknown until headers arrive).
Progress { done: u64, total: Option<u64> },
/// Every step finished — the model is provisioned.
Done,
/// Setup failed (offline / firewalled / disk); the session continues degraded.
Failed(String),
}
/// Handed to [`run_code`] by the binary when a first-run provision is needed.
/// `rx` streams [`SetupEvent`]s from the download thread; setting `cancel` asks
/// that thread to stop (triple-Esc / Ctrl+C).
pub struct SetupHandle {
/// Human label for what's being set up (e.g. `on-host summarizer (qwen2.5-0.5b)`).
pub what: String,
/// Progress stream from the background thread.
pub rx: std::sync::mpsc::Receiver<SetupEvent>,
/// Set to request the background thread abort.
pub cancel: std::sync::Arc<std::sync::atomic::AtomicBool>,
}
/// Is this key an abort press (Esc or Ctrl+C)? Three in a row abort setup.
pub(crate) fn is_abort_key(ev: &Event) -> bool {
match ev {
Event::Key(KeyEvent {
code: KeyCode::Esc, ..
}) => true,
Event::Key(KeyEvent {
code: KeyCode::Char('c'),
modifiers,
..
}) => modifiers.contains(KeyModifiers::CONTROL),
_ => false,
}
}
/// The status line from the latest provisioning state.
pub(crate) fn setup_status_line(what: &str, step: &str, done: u64, total: Option<u64>) -> String {
let mb = |b: u64| b / 1_048_576;
match total.filter(|&t| t > 0) {
Some(t) => format!(
"setting up {what} — {step} {}/{} MB {}%",
mb(done),
mb(t),
done.saturating_mul(100) / t
),
None => format!("setting up {what} — {step} {} MB", mb(done)),
}
}
/// Poll `rx` non-blocking, folding events into the running state. Returns
/// `Some(Ok/Err)` once setup is finished (or the sender died), else `None`.
pub(crate) fn drain_setup(
rx: &std::sync::mpsc::Receiver<SetupEvent>,
step: &mut String,
done: &mut u64,
total: &mut Option<u64>,
) -> Option<Result<(), String>> {
use std::sync::mpsc::TryRecvError;
loop {
match rx.try_recv() {
Ok(SetupEvent::Step(s)) => {
*step = s;
*done = 0;
*total = None;
}
Ok(SetupEvent::Progress { done: d, total: t }) => {
*done = d;
*total = t;
}
Ok(SetupEvent::Done) => return Some(Ok(())),
Ok(SetupEvent::Failed(e)) => return Some(Err(e)),
Err(TryRecvError::Empty) => return None,
// Sender dropped without a terminal event → the thread died.
Err(TryRecvError::Disconnected) => return Some(Err("interrupted".into())),
}
}
}
/// The no-splash analog: a single carriage-return-updating spinner line (the
/// header is already printed by the caller). No input capture — Ctrl+C is the
/// terminal's SIGINT.
pub(crate) fn run_setup_inline(setup: &SetupHandle) {
let mut frame = 0usize;
let mut step = "starting".to_string();
let (mut done, mut total) = (0u64, None);
loop {
match drain_setup(&setup.rx, &mut step, &mut done, &mut total) {
None => {
eprint!(
"\r {} {} ",
newt_core::tty::SPINNER_FRAMES[frame % newt_core::tty::SPINNER_FRAMES.len()],
setup_status_line(&setup.what, &step, done, total)
);
let _ = io::stderr().flush();
}
Some(r) => {
let msg = match r {
Ok(()) => format!("✓ {} ready", setup.what),
Err(e) => format!("⚠ setup skipped ({e}) — using the session model"),
};
eprintln!("\r {msg} ");
return;
}
}
std::thread::sleep(std::time::Duration::from_millis(100));
frame += 1;
}
}
#[cfg(test)]
mod setup_screen_tests {
use super::{drain_setup, is_abort_key, setup_status_line, SetupEvent};
use crossterm::event::{Event, KeyCode, KeyEvent, KeyModifiers};
const MB: u64 = 1_048_576;
#[test]
fn status_line_shows_pct_with_total_and_bare_mb_without() {
let known = setup_status_line("summarizer", "weights", 50 * MB, Some(100 * MB));
assert!(known.contains("50/100 MB"), "{known}");
assert!(known.contains("50%"), "{known}");
let unknown = setup_status_line("summarizer", "weights", 7 * MB, None);
assert!(
unknown.contains("7 MB") && !unknown.contains('%'),
"{unknown}"
);
}
#[test]
fn abort_key_is_esc_or_ctrl_c_only() {
assert!(is_abort_key(&Event::Key(KeyEvent::new(
KeyCode::Esc,
KeyModifiers::NONE
))));
assert!(is_abort_key(&Event::Key(KeyEvent::new(
KeyCode::Char('c'),
KeyModifiers::CONTROL
))));
// A bare 'c' or Enter must NOT abort — the splash blocks ordinary input.
assert!(!is_abort_key(&Event::Key(KeyEvent::new(
KeyCode::Char('c'),
KeyModifiers::NONE
))));
assert!(!is_abort_key(&Event::Key(KeyEvent::new(
KeyCode::Enter,
KeyModifiers::NONE
))));
}
#[test]
fn drain_folds_progress_and_reports_done() {
let (tx, rx) = std::sync::mpsc::channel();
let (mut step, mut done, mut total) = (String::new(), 0u64, None);
assert!(drain_setup(&rx, &mut step, &mut done, &mut total).is_none());
tx.send(SetupEvent::Step("weights".into())).unwrap();
tx.send(SetupEvent::Progress {
done: 42,
total: Some(100),
})
.unwrap();
assert!(drain_setup(&rx, &mut step, &mut done, &mut total).is_none());
assert_eq!((step.as_str(), done, total), ("weights", 42, Some(100)));
tx.send(SetupEvent::Done).unwrap();
assert!(matches!(
drain_setup(&rx, &mut step, &mut done, &mut total),
Some(Ok(()))
));
}
#[test]
fn drain_reports_sender_death_as_error() {
// If the download thread dies without a terminal event, the screen must
// stop waiting (else it spins forever) — Disconnected → Err.
let (tx, rx) = std::sync::mpsc::channel::<SetupEvent>();
drop(tx);
let (mut step, mut done, mut total) = (String::new(), 0u64, None);
assert!(matches!(
drain_setup(&rx, &mut step, &mut done, &mut total),
Some(Err(_))
));
}
}