@sigitcodereview-old / sigit / commits / f543385

Fork 0

A modular local inference LLM coding agent. Core component of siGit Factory.

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