@hej / sigit / commits / f543385

Show model download and loading progress in Zed agent

Add download and loading progress notifications when switching models in the Zed agent. Show a progress bar during downloads and a spinner during model loading. Progress pollers are stopped once the model is loaded or on error.

paydii committed Apr 26, 2026 at 08:29 UTC f543385723ccfc8b04a2d4f7f349aea5590f81ca

2 files changed +254 -26

src/chat.rs

+35 -25

@@ -162,6 +162,9 @@ struct App {

162 blink_counter: u8,

163 /// True while a model switch is in progress.

164 switching_model: bool,

165 + /// Tool-calling flag for the model currently being loaded in the background.

166 + /// Applied to `app.tool_calling` when `ModelLoadUpdate::Loaded` arrives.

167 + pending_tool_calling: Option<bool>,

168

169 // ── Loading-phase state ───────────────────────────────────────────────────

170 /// True while the model is still loading; switches to false on completion.

@@ -226,6 +229,7 @@ impl App {

229 blink_on: true,

230 blink_counter: 0,

231 switching_model: false,

232 + pending_tool_calling: None,

233 is_loading: true,

234 load_tick: 0,

235 load_error: None,

@@ -1107,7 +1111,7 @@ fn handle_key(app: &mut App, key: KeyEvent) -> Option<String> {

1111 async fn exec_slash<B: ratatui::backend::Backend>(

1112 app: &mut App,

1113 cmd: SlashCommand,

1110 - engine: &ChatEngine,

1114 + engine: Arc<ChatEngine>,

1115 terminal: &mut ratatui::Terminal<B>,

1116 ) {

1117 match cmd {

@@ -1130,7 +1134,7 @@ async fn exec_slash<B: ratatui::backend::Backend>(

1134 )));

1135 }

1136 SlashCommand::Status => {

1133 - let info = engine.info().await;

1137 + let info = engine.as_ref().info().await;

1138 let model = info.model_name.as_deref().unwrap_or("(none)");

1139 let mem = info.approx_memory.as_deref().unwrap_or("unknown");

1140 app.messages.push(ChatMessage::system(format!(

@@ -1140,7 +1144,7 @@ async fn exec_slash<B: ratatui::backend::Backend>(

1144 }

1145 SlashCommand::Models(selection) => match selection {

1146 None => {

1143 - app.open_model_picker(engine);

1147 + app.open_model_picker(&engine);

1148 }

1149 Some(n) => {

1150 let idx = n.saturating_sub(1);

@@ -1182,28 +1186,30 @@ async fn exec_slash<B: ratatui::backend::Backend>(

1186 ..SamplingConfig::default()

1187 };

1188

1185 - // load_gguf_model unloads any existing model internally before

1186 - // loading the new one. Calling unload_model() explicitly first

1187 - // would create a window where no model is loaded — if a message

1188 - // arrived in that gap it would fail with NoModelLoaded.

1189 + // Spawn onto a background task so the event loop keeps

1190 + // running (and the spinner keeps animating) during the

1191 + // download + load — which can take several minutes for

1192 + // a large model fetched from HuggingFace for the first time.

1193 let system_prompt = crate::system_prompt_for_model(model.tool_calling);

1190 - let update = match engine

1191 - .load_gguf_model(

1192 - model.config.clone(),

1193 - Some(system_prompt.to_string()),

1194 - Some(sampling),

1195 - )

1196 - .await

1197 - {

1198 - Ok(_) => {

1199 - engine.clear_history().await;

1200 - app.tool_calling = model.tool_calling;

1201 - ModelLoadUpdate::Loaded(model.display_name.clone())

1202 - }

1203 - Err(err) => ModelLoadUpdate::Error(err.to_string()),

1204 - };

1205 -

1206 - let _ = tx.send(update).await;

1194 + let engine_handle = Arc::clone(&engine);

1195 + let tool_calling = model.tool_calling;

1196 + tokio::spawn(async move {

1197 + let update = match engine_handle

1198 + .load_gguf_model(

1199 + model.config.clone(),

1200 + Some(system_prompt.to_string()),

1201 + Some(sampling),

1202 + )

1203 + .await

1204 + {

1205 + Ok(_) => ModelLoadUpdate::Loaded(model.display_name.clone()),

1206 + Err(err) => ModelLoadUpdate::Error(err.to_string()),

1207 + };

1208 + let _ = tx.send(update).await;

1209 + });

1210 + // tool_calling is applied when ModelLoadUpdate::Loaded

1211 + // arrives in the event loop (see model_load_rx handler).

1212 + app.pending_tool_calling = Some(tool_calling);

1213 }

1214 }

1215 }

@@ -1377,6 +1383,10 @@ async fn event_loop<B: ratatui::backend::Backend>(

1383 if let Some(rx) = app.model_load_rx.as_mut() {

1384 match rx.try_recv() {

1385 Ok(ModelLoadUpdate::Loaded(model_name)) => {

1386 + engine.clear_history().await;

1387 + if let Some(tc) = app.pending_tool_calling.take() {

1388 + app.tool_calling = tc;

1389 + }

1390 app.switching_model = false;

1391 app.model_load_rx = None;

1392 app.current_model_name = model_name.clone();

@@ -1548,7 +1558,7 @@ async fn event_loop<B: ratatui::backend::Backend>(

1558

1559 if let Some(text) = handle_key(&mut app, key) {

1560 if let Some(cmd) = parse_slash(&text) {

1551 - exec_slash(&mut app, cmd, &engine, terminal).await;

1561 + exec_slash(&mut app, cmd, Arc::clone(&engine), terminal).await;

1562 continue;

1563 }

1564

src/main.rs

+219 -1

@@ -66,6 +66,7 @@ use agent_client_protocol::{

66 use futures::future::LocalBoxFuture;

67 use onde::inference::{ChatEngine, GgufModelConfig, ToolDefinition, ToolResult};

68 use std::path::PathBuf;

69 +use std::sync::atomic::{AtomicBool, Ordering};

70 use tokio::sync::mpsc;

71 use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt};

72 use tracing_subscriber::{EnvFilter, fmt as tracing_fmt};

@@ -1163,7 +1164,179 @@ impl Agent for SiGitAgent {

1164 }

1165

1166 let model_id = args.value.0.as_ref();

1166 - let _new_config = self.switch_model_by_id(model_id).await?;

1167 +

1168 + // Check if this model needs to be downloaded first so we can show

1169 + // a progress indicator in Zed while the download + load is happening.

1170 + let needs_download = models::build_model_picker_items()

1171 + .into_iter()

1172 + .find(|item| item.config.model_id == model_id)

1173 + .map(|item| item.cache_health == setup::ModelCacheHealth::NotDownloaded)

1174 + .unwrap_or(false);

1175 +

1176 + // Spawn a progress-poller task that sends periodic download status

1177 + // messages to Zed via the notification channel. A shared flag lets

1178 + // us stop the poller once the load finishes.

1179 + let stop_flag = Arc::new(AtomicBool::new(false));

1180 +

1181 + if needs_download {

1182 + // Send the initial "downloading" banner immediately.

1183 + let model_id_owned = model_id.to_string();

1184 + let expected_bytes = onde::inference::models::SUPPORTED_MODEL_INFO

1185 + .iter()

1186 + .find(|m| m.id == model_id_owned)

1187 + .map(|m| m.expected_size_bytes)

1188 + .unwrap_or(0);

1189 +

1190 + let display_name = models::build_model_picker_items()

1191 + .into_iter()

1192 + .find(|item| item.config.model_id == model_id_owned)

1193 + .map(|item| item.display_name.clone())

1194 + .unwrap_or_else(|| model_id_owned.clone());

1195 +

1196 + let size_hint = if expected_bytes > 0 {

1197 + format!(" (~{})", format_size_human(expected_bytes))

1198 + } else {

1199 + String::new()

1200 + };

1201 +

1202 + self.send_assistant_message(

1203 + args.session_id.clone(),

1204 + format!("⏬ Downloading {display_name}{size_hint}… this may take a few minutes."),

1205 + )

1206 + .await;

1207 +

1208 + // Poller: every 4 seconds report bytes-on-disk / expected.

1209 + let poller_tx = self.notification_tx.clone();

1210 + let poller_session = args.session_id.clone();

1211 + let poller_model_id = model_id_owned.clone();

1212 + let poller_stop = Arc::clone(&stop_flag);

1213 +

1214 + tokio::spawn(async move {

1215 + let cache_path = onde::hf_cache::model_cache_path(&poller_model_id);

1216 + let mut interval = tokio::time::interval(std::time::Duration::from_secs(4));

1217 + interval.tick().await; // consume the immediate first tick

1218 +

1219 + while !poller_stop.load(Ordering::Relaxed) {

1220 + interval.tick().await;

1221 +

1222 + if poller_stop.load(Ordering::Relaxed) {

1223 + break;

1224 + }

1225 +

1226 + let downloaded = cache_path

1227 + .as_ref()

1228 + .filter(|p| p.exists())

1229 + .map(|p| dir_size_recursive(p))

1230 + .unwrap_or(0);

1231 +

1232 + let msg = if expected_bytes > 0 {

1233 + let pct =

1234 + ((downloaded as f64 / expected_bytes as f64) * 100.0).min(99.0) as u8;

1235 + let bar = progress_bar(pct, 20);

1236 + format!(

1237 + "⏬ {display_name} — {bar} {pct}% ({} / {})",

1238 + format_size_human(downloaded),

1239 + format_size_human(expected_bytes),

1240 + )

1241 + } else {

1242 + format!(

1243 + "⏬ {display_name} — {} downloaded…",

1244 + format_size_human(downloaded)

1245 + )

1246 + };

1247 +

1248 + let notification = SessionNotification::new(

1249 + poller_session.clone(),

1250 + SessionUpdate::AgentMessageChunk(ContentChunk::new(ContentBlock::from(

1251 + msg,

1252 + ))),

1253 + );

1254 + if poller_tx.send(notification).await.is_err() {

1255 + break;

1256 + }

1257 + }

1258 + });

1259 + }

1260 +

1261 + // For already-cached models, send a "loading" message and a spinner

1262 + // so the user sees activity while mistralrs loads the weights (~10-30 s).

1263 + if !needs_download {

1264 + let cached_display_name = models::build_model_picker_items()

1265 + .into_iter()

1266 + .find(|item| item.config.model_id == model_id)

1267 + .map(|item| item.display_name.clone())

1268 + .unwrap_or_else(|| model_id.to_string());

1269 +

1270 + self.send_assistant_message(

1271 + args.session_id.clone(),

1272 + format!("⏳ Loading {cached_display_name}…"),

1273 + )

1274 + .await;

1275 +

1276 + // Spinner poller: send an elapsed-time update every 5 seconds so

1277 + // the user can tell siGit is still working.

1278 + let spinner_tx = self.notification_tx.clone();

1279 + let spinner_session = args.session_id.clone();

1280 + let spinner_name = cached_display_name.clone();

1281 + let spinner_stop = Arc::clone(&stop_flag);

1282 + let load_start = std::time::Instant::now();

1283 +

1284 + tokio::spawn(async move {

1285 + const SPINNER: &[char] = &['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '⠇', '⠏'];

1286 + let mut tick: usize = 0;

1287 + let mut interval = tokio::time::interval(std::time::Duration::from_secs(5));

1288 + interval.tick().await; // consume the immediate first tick

1289 +

1290 + while !spinner_stop.load(Ordering::Relaxed) {

1291 + interval.tick().await;

1292 +

1293 + if spinner_stop.load(Ordering::Relaxed) {

1294 + break;

1295 + }

1296 +

1297 + let elapsed = load_start.elapsed();

1298 + let elapsed_str = if elapsed.as_secs() >= 60 {

1299 + format!("{}m {:02}s", elapsed.as_secs() / 60, elapsed.as_secs() % 60)

1300 + } else {

1301 + format!("{}s", elapsed.as_secs())

1302 + };

1303 + let frame = SPINNER[tick % SPINNER.len()];

1304 + tick += 1;

1305 +

1306 + let msg = format!("{frame} Loading {spinner_name}… ({elapsed_str})");

1307 + let notification = SessionNotification::new(

1308 + spinner_session.clone(),

1309 + SessionUpdate::AgentMessageChunk(ContentChunk::new(ContentBlock::from(

1310 + msg,

1311 + ))),

1312 + );

1313 + if spinner_tx.send(notification).await.is_err() {

1314 + break;

1315 + }

1316 + }

1317 + });

1318 + }

1319 +

1320 + let switch_result = self.switch_model_by_id(model_id).await;

1321 +

1322 + // Stop the progress / spinner poller regardless of success/failure.

1323 + stop_flag.store(true, Ordering::Relaxed);

1324 +

1325 + let new_config = switch_result?;

1326 +

1327 + if needs_download {

1328 + self.send_assistant_message(

1329 + args.session_id.clone(),

1330 + format!("✓ {} downloaded and loaded.", new_config.display_name),

1331 + )

1332 + .await;

1333 + } else {

1334 + self.send_assistant_message(

1335 + args.session_id.clone(),

1336 + format!("✓ Switched to {}.", new_config.display_name),

1337 + )

1338 + .await;

1339 + }

1340

1341 let config_options = {

1342 let guard = self.current_model.lock().unwrap();

@@ -1175,6 +1348,51 @@ impl Agent for SiGitAgent {

1348 }

1349 }

1350

1351 +// ── Download progress helpers ─────────────────────────────────────────────────

1352 +

1353 +/// Recursively sum the sizes of all files under `path`, following symlinks.

1354 +/// Used by the ACP download-progress poller to report bytes-on-disk before

1355 +/// hf-hub renames the staging files to their final blob names.

1356 +fn dir_size_recursive(path: &std::path::Path) -> u64 {

1357 + let mut total: u64 = 0;

1358 + let Ok(entries) = std::fs::read_dir(path) else {

1359 + return 0;

1360 + };

1361 + for entry in entries.flatten() {

1362 + let entry_path = entry.path();

1363 + if entry_path.is_dir() {

1364 + total += dir_size_recursive(&entry_path);

1365 + } else if let Ok(meta) = entry_path.metadata() {

1366 + total += meta.len();

1367 + }

1368 + }

1369 + total

1370 +}

1371 +

1372 +/// Format a byte count as a human-readable string (B / KB / MB / GB).

1373 +fn format_size_human(bytes: u64) -> String {

1374 + const GB: u64 = 1_073_741_824;

1375 + const MB: u64 = 1_048_576;

1376 + const KB: u64 = 1_024;

1377 + if bytes >= GB {

1378 + format!("{:.2} GB", bytes as f64 / GB as f64)

1379 + } else if bytes >= MB {

1380 + format!("{:.1} MB", bytes as f64 / MB as f64)

1381 + } else if bytes >= KB {

1382 + format!("{:.0} KB", bytes as f64 / KB as f64)

1383 + } else {

1384 + format!("{bytes} B")

1385 + }

1386 +}

1387 +

1388 +/// Build a simple ASCII progress bar string of the given width.

1389 +/// e.g. `[████████░░░░░░░░░░░░]` at 40 %

1390 +fn progress_bar(pct: u8, width: usize) -> String {

1391 + let filled = ((pct as usize) * width) / 100;

1392 + let empty = width.saturating_sub(filled);

1393 + format!("[{}{}]", "█".repeat(filled), "░".repeat(empty))

1394 +}

1395 +

1396 // ── Output capture ────────────────────────────────────────────────────────────

1397

1398 /// Redirect **both** stdout and stderr to `$TMPDIR/sigit.log` at the