1use std::sync::{Arc, RwLock};
14
15use super::{SubagentRegistry, SubagentState, SubagentStatus};
16
17#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
25pub struct CarriedWorker {
26 pub id: String,
27 #[serde(default)]
28 pub subagent_id: u64,
29 #[serde(default)]
30 pub agent: String,
31 pub status: String,
33 #[serde(default)]
35 pub archive_path: Option<std::path::PathBuf>,
36 #[serde(default)]
40 pub stopped_by_reload: bool,
41 #[serde(default)]
45 pub wake_lost: bool,
46 #[serde(default)]
51 pub wake: bool,
52 #[serde(default)]
53 pub elapsed_secs: u64,
54 #[serde(default)]
56 pub preview: String,
57}
58
59impl CarriedWorker {
60 pub fn needs_wake(&self) -> bool {
62 self.wake
63 }
64}
65
66pub const MAX_CARRIED_ARCHIVES: usize = 1024;
68
69pub fn handle_number(id: &str) -> Option<u64> {
71 id.strip_prefix("sa_")?.parse().ok()
72}
73
74impl SubagentRegistry {
75 pub fn reload_carry(
81 &self,
82 undelivered_wakes: &std::collections::HashSet<String>,
83 ) -> Vec<CarriedWorker> {
84 let mut carried: Vec<CarriedWorker> = self
85 .handles
86 .values()
87 .map(|h| {
88 let s = h.state.read().unwrap_or_else(|p| p.into_inner());
89 let terminal = s.finalization_ready || !matches!(s.status, SubagentStatus::Running);
90 let stopped_by_reload =
94 s.stopped_by_reload || (!terminal && !s.cancel_requested);
95 let wake_lost = !stopped_by_reload
100 && !h.collected
101 && (undelivered_wakes.contains(&h.id)
102 || (!s.completion_settled && !s.cancel_requested));
103 let latest = s.last_response_text().trim();
104 let source = if latest.is_empty() { s.partial_text.trim() } else { latest };
105 CarriedWorker {
106 id: h.id.clone(),
107 subagent_id: h.subagent_id,
108 agent: h.agent_name.clone(),
109 status: if terminal {
110 s.status.as_str().to_string()
111 } else {
112 "cancelled".to_string()
115 },
116 archive_path: s.archive_path.clone(),
117 stopped_by_reload,
118 wake_lost,
119 wake: stopped_by_reload || wake_lost,
120 elapsed_secs: h.elapsed_secs() as u64,
121 preview: source.chars().take(300).collect(),
122 }
123 })
124 .collect();
125 let live: std::collections::HashSet<&String> = self.handles.keys().collect();
126 let mut reaped: Vec<(&String, &std::path::PathBuf)> = self
127 .archived
128 .iter()
129 .filter(|(id, _)| !live.contains(id))
130 .collect();
131 reaped.sort_by_key(|(id, _)| std::cmp::Reverse(handle_number(id)));
133 carried.extend(reaped.into_iter().take(MAX_CARRIED_ARCHIVES).map(|(id, path)| {
134 CarriedWorker {
135 id: id.clone(),
136 subagent_id: handle_number(id).unwrap_or(0),
137 agent: String::new(),
138 status: "archived".into(),
139 archive_path: Some(path.clone()),
140 stopped_by_reload: self.stopped_by_reload.contains(id),
141 wake_lost: false,
142 wake: false,
143 elapsed_secs: 0,
144 preview: String::new(),
145 }
146 }));
147 for (id, posted) in &self.reload_woken {
151 let pending = undelivered_wakes.contains(id);
152 match carried.iter_mut().find(|w| &w.id == id) {
153 Some(w) if pending && !self.handles.contains_key(id) => {
154 *w = CarriedWorker { wake: true, ..posted.clone() };
155 }
156 Some(w) => w.wake = pending,
157 None if pending => carried.push(CarriedWorker { wake: true, ..posted.clone() }),
158 None => {}
159 }
160 }
161 carried.sort_by_key(|w| (handle_number(&w.id), w.id.clone()));
162 carried
163 }
164
165 pub fn adopt_carried(&mut self, worker: &CarriedWorker) {
169 if self.handles.contains_key(&worker.id) {
170 return;
171 }
172 if let Some(path) = &worker.archive_path {
173 self.archived.insert(worker.id.clone(), path.clone());
174 }
175 if worker.stopped_by_reload {
176 self.stopped_by_reload.insert(worker.id.clone());
177 }
178 }
179
180 pub fn note_reload_wake(&mut self, worker: &CarriedWorker) {
183 self.reload_woken.insert(worker.id.clone(), worker.clone());
184 }
185
186 pub fn rewake_after_failed_exec(&self, carried: Vec<CarriedWorker>) -> Vec<CarriedWorker> {
194 carried
195 .into_iter()
196 .map(|w| {
197 let wake = w.wake && w.stopped_by_reload && !self.reload_woken.contains_key(&w.id);
198 CarriedWorker { wake, ..w }
199 })
200 .collect()
201 }
202
203 pub fn was_stopped_by_reload(&self, id: &str) -> bool {
205 self.stopped_by_reload.contains(id)
206 || self.handles.get(id).is_some_and(|h| {
207 h.state.read().unwrap_or_else(|p| p.into_inner()).stopped_by_reload
208 })
209 }
210
211 pub fn lift_reload_fence(&mut self) {
215 if let Some(previous) = self.reload_fence.take() {
216 self.set_spawn_cancellation(previous);
217 }
218 }
219
220 pub fn stop_for_reload(&mut self) -> Vec<(String, Arc<RwLock<SubagentState>>)> {
231 if self.reload_fence.is_none() {
232 let fence = crate::CancellationToken::new();
233 fence.cancel();
234 let previous = self.spawn_cancellation.take();
235 self.set_spawn_cancellation(Some(fence));
236 self.reload_fence = Some(previous);
237 }
238 let mut awaited = Vec::new();
239 for h in self.handles.values_mut() {
240 let stop = {
241 let mut s = h.state.write().unwrap_or_else(|p| p.into_inner());
246 if s.completion_settled {
247 continue;
248 }
249 let running = matches!(s.status, SubagentStatus::Running) && !s.finalization_started;
250 let stop = running && !s.cancel_requested;
251 if stop {
252 s.stopped_by_reload = true;
253 s.cancel_requested = true;
254 }
255 stop
256 };
257 if stop {
258 h.cancel();
259 }
260 awaited.push((h.id.clone(), Arc::clone(&h.state)));
261 }
262 awaited.sort_by(|a, b| a.0.cmp(&b.0));
263 awaited
264 }
265}
266
267#[cfg(test)]
268mod tests {
269 use super::super::SubagentHandle;
270 use super::*;
271
272 fn make_handle(id: &str) -> SubagentHandle {
275 let (steer_tx, _) = tokio::sync::mpsc::unbounded_channel();
276 let (shutdown_tx, _) = tokio::sync::oneshot::channel();
277 let (_, result_rx) = tokio::sync::oneshot::channel();
278 SubagentHandle::new(
279 id.to_string(),
280 handle_number(id).unwrap_or(0),
281 "test-agent".to_string(),
282 "test task".to_string(),
283 "claude-sonnet-4-6".to_string(),
284 "You are a test agent.".to_string(),
285 300,
286 Arc::new(RwLock::new(SubagentState::new())),
287 Some(steer_tx),
288 Some(shutdown_tx),
289 Some(result_rx),
290 )
291 }
292
293 fn make_finished_handle(id: &str) -> SubagentHandle {
294 let h = make_handle(id);
295 {
296 let mut s = h.state.write().unwrap();
297 s.status = SubagentStatus::Completed;
298 s.finished_at = Some(std::time::Instant::now());
299 }
300 h
301 }
302
303 fn finalize_as(h: &SubagentHandle, status: SubagentStatus, archive: Option<&str>) {
304 let mut s = h.state.write().unwrap();
305 if s.cancel_requested {
306 s.status = SubagentStatus::Cancelled;
307 } else {
308 s.status = status;
309 }
310 s.finished_at = Some(std::time::Instant::now());
311 s.finalization_started = true;
312 s.finalization_ready = true;
313 s.completion_settled = true;
314 s.archive_path = archive.map(std::path::PathBuf::from);
315 s.partial_text = "latest words".into();
316 }
317
318 #[test]
324 fn stop_for_reload_cancels_running_workers_flagged_and_skips_finished() {
325 let mut registry = SubagentRegistry::new();
326 let running = make_handle("sa_4");
327 let user_cancelled = make_handle("sa_5");
328 let finished = make_finished_handle("sa_6");
329 finalize_as(&finished, SubagentStatus::Completed, Some("/a/sa_6.json"));
330 let finishing = make_finished_handle("sa_7");
331 let (running_state, user_state, finished_state, finishing_state) = (
332 Arc::clone(&running.state),
333 Arc::clone(&user_cancelled.state),
334 Arc::clone(&finished.state),
335 Arc::clone(&finishing.state),
336 );
337 registry.register(running);
338 registry.register(user_cancelled);
339 registry.register(finished);
340 registry.register(finishing);
341 registry.get_mut("sa_5").unwrap().cancel();
342
343 let stopping: Vec<String> = registry.stop_for_reload().into_iter().map(|(id, _)| id).collect();
344 assert_eq!(stopping, ["sa_4", "sa_5", "sa_7"]);
345 let r = running_state.read().unwrap();
346 assert!(r.cancel_requested && r.stopped_by_reload);
347 let u = user_state.read().unwrap();
348 assert!(u.cancel_requested && !u.stopped_by_reload, "a user cancel stays a user cancel");
349 let f = finished_state.read().unwrap();
350 assert!(!f.cancel_requested && !f.stopped_by_reload);
351 let g = finishing_state.read().unwrap();
352 assert!(!g.cancel_requested && !g.stopped_by_reload, "its completion stands");
353 assert_eq!(g.status, SubagentStatus::Completed);
354 drop((r, u, f, g));
355 assert_eq!(registry.stop_for_reload().len(), 3);
357 }
358
359 #[test]
364 fn a_worker_finishing_during_the_checkpoint_keeps_its_wake() {
365 let mut registry = SubagentRegistry::new();
366 registry.register(make_finished_handle("sa_20"));
367 let cancelled = make_finished_handle("sa_21");
368 cancelled.state.write().unwrap().cancel_requested = true;
369 registry.register(cancelled);
370 let none = std::collections::HashSet::new();
371 let carried = registry.reload_carry(&none);
372 let by = |id: &str| carried.iter().find(|w| w.id == id).unwrap().clone();
373 let w = by("sa_20");
374 assert!(w.wake_lost && w.wake && !w.stopped_by_reload, "{w:?}");
375 assert_eq!(w.status, "completed");
376 assert!(!by("sa_21").wake, "a cancel owes no wake");
377
378 finalize_as(registry.get("sa_20").unwrap(), SubagentStatus::Completed, Some("/a/sa_20.json"));
380 assert!(!registry.reload_carry(&none).iter().any(|w| w.wake));
381 let queued: std::collections::HashSet<String> = ["sa_20".to_string()].into();
383 let w = registry.reload_carry(&queued).into_iter().find(|w| w.id == "sa_20").unwrap();
384 assert!(w.wake_lost && w.wake);
385 }
386
387 #[test]
390 fn reload_carry_then_adopt_resolves_old_ids_to_their_archives() {
391 let mut old = SubagentRegistry::new();
392 let stopped = make_handle("sa_10");
394 let stuck = make_handle("sa_11");
396 let owed = make_finished_handle("sa_12");
398 let delivered = make_finished_handle("sa_13");
400 let mut collected = make_finished_handle("sa_14");
402 collected.mark_collected();
403 for (h, path) in [
404 (&owed, "/a/sa_12.json"),
405 (&delivered, "/a/sa_13.json"),
406 (&collected, "/a/sa_14.json"),
407 ] {
408 finalize_as(h, SubagentStatus::Completed, Some(path));
409 }
410 let stopped_state = Arc::clone(&stopped.state);
411 for h in [stopped, stuck, owed, delivered, collected] {
412 old.register(h);
413 }
414 old.archived.insert("sa_3".into(), "/a/sa_3.json".into());
415 assert_eq!(old.stop_for_reload().len(), 2);
416 {
417 let h = old.get("sa_10").unwrap();
418 finalize_as(h, SubagentStatus::Completed, Some("/a/sa_10.json"));
419 }
420 assert_eq!(stopped_state.read().unwrap().status, SubagentStatus::Cancelled);
421
422 let undelivered: std::collections::HashSet<String> = ["sa_12".to_string()].into();
423 let carried = old.reload_carry(&undelivered);
424 let ids: Vec<&str> = carried.iter().map(|w| w.id.as_str()).collect();
425 assert_eq!(ids, ["sa_3", "sa_10", "sa_11", "sa_12", "sa_13", "sa_14"]);
426 let by = |id: &str| carried.iter().find(|w| w.id == id).unwrap().clone();
427
428 let w = by("sa_10");
429 assert!(w.stopped_by_reload && !w.wake_lost && w.needs_wake());
430 assert_eq!(w.status, "cancelled");
431 assert_eq!(w.archive_path.as_deref(), Some(std::path::Path::new("/a/sa_10.json")));
432 assert_eq!(w.preview, "latest words");
433 assert_eq!(w.subagent_id, 10);
434
435 let w = by("sa_11");
436 assert!(w.stopped_by_reload && w.needs_wake(), "the exec destroys it: still announced");
437 assert_eq!(w.archive_path, None);
438
439 let w = by("sa_12");
440 assert!(!w.stopped_by_reload && w.wake_lost && w.needs_wake());
441 assert_eq!(w.status, "completed");
442 assert!(!by("sa_13").needs_wake(), "its wake was delivered: no duplicate");
443 assert!(!by("sa_14").needs_wake());
444 let w = by("sa_3");
445 assert!(!w.needs_wake());
446 assert_eq!(w.status, "archived");
447
448 let mut new = SubagentRegistry::new();
450 for w in &carried {
451 new.adopt_carried(w);
452 if w.needs_wake() {
453 new.note_reload_wake(w);
454 }
455 }
456 assert_eq!(new.archive_path("sa_10"), Some("/a/sa_10.json".into()));
457 assert_eq!(new.archive_path("sa_3"), Some("/a/sa_3.json".into()));
458 assert_eq!(new.archive_path("sa_11"), None, "no archive: nothing to resolve");
459 assert!(new.was_stopped_by_reload("sa_10"));
460 assert!(new.was_stopped_by_reload("sa_11"));
461 assert!(!new.was_stopped_by_reload("sa_12"));
462 assert!(new.list_active().is_empty(), "carried workers are not live");
463 assert!(new.display_rows().is_empty(), "nor shown as running");
464
465 let mut live = SubagentRegistry::new();
467 let h = make_handle("sa_10");
468 finalize_as(&h, SubagentStatus::Completed, Some("/live/sa_10.json"));
469 live.register(h);
470 live.adopt_carried(&by("sa_10"));
471 assert_eq!(live.archive_path("sa_10"), Some("/live/sa_10.json".into()));
472
473 let again = new.reload_carry(&Default::default());
477 assert!(again.iter().all(|w| !w.needs_wake()), "{again:?}");
478 let labelled: Vec<&str> =
479 again.iter().filter(|w| w.stopped_by_reload).map(|w| w.id.as_str()).collect();
480 assert_eq!(labelled, ["sa_10"]);
482
483 let queued: std::collections::HashSet<String> =
486 ["sa_11".to_string(), "sa_12".to_string()].into();
487 let again = new.reload_carry(&queued);
488 let woken: Vec<&CarriedWorker> = again.iter().filter(|w| w.needs_wake()).collect();
489 assert_eq!(woken.iter().map(|w| w.id.as_str()).collect::<Vec<_>>(), ["sa_11", "sa_12"]);
490 assert_eq!(*woken[0], by("sa_11"), "the stuck worker's stop, as posted");
491 assert_eq!(*woken[1], by("sa_12"), "the lost completion, as posted");
492 assert!(!again.iter().any(|w| w.id == "sa_10" && w.needs_wake()));
493 }
494
495 #[test]
501 fn a_worker_registered_after_the_reload_stop_is_stopped_by_the_reload() {
502 let mut registry = SubagentRegistry::new();
503 registry.register(make_handle("sa_1"));
504 assert_eq!(registry.stop_for_reload().len(), 1);
505
506 let late = make_handle("sa_2");
507 let late_state = Arc::clone(&late.state);
508 registry.register_with_cancellation(late, None);
509 {
510 let s = late_state.read().unwrap();
511 assert!(s.cancel_requested && s.stopped_by_reload, "stopped at birth, by the reload");
512 }
513 let turn = crate::CancellationToken::new();
516 turn.cancel();
517 let racing = make_handle("sa_3");
518 let racing_state = Arc::clone(&racing.state);
519 registry.register_with_cancellation(racing, Some(&turn));
520 assert!(racing_state.read().unwrap().stopped_by_reload);
521 assert!(registry.cancel_running(super::super::WorkerCancelScope::All).is_empty());
523 let awaited: Vec<String> = registry.stop_for_reload().into_iter().map(|(id, _)| id).collect();
525 assert_eq!(awaited, ["sa_1", "sa_2", "sa_3"]);
526 let carried = registry.reload_carry(&Default::default());
527 for id in ["sa_2", "sa_3"] {
528 let w = carried.iter().find(|w| w.id == id).unwrap();
529 assert!(w.stopped_by_reload && w.needs_wake(), "{w:?}");
530 }
531
532 registry.lift_reload_fence();
534 let after = make_handle("sa_4");
535 let after_state = Arc::clone(&after.state);
536 registry.register(after);
537 assert!(!after_state.read().unwrap().cancel_requested);
538
539 let mut fenced = SubagentRegistry::new();
542 let driver = crate::CancellationToken::new();
543 driver.cancel();
544 fenced.set_spawn_cancellation(Some(driver));
545 fenced.stop_for_reload();
546 fenced.lift_reload_fence();
547 let late = make_handle("sa_9");
548 let late_state = Arc::clone(&late.state);
549 fenced.register(late);
550 let s = late_state.read().unwrap();
551 assert!(s.cancel_requested && !s.stopped_by_reload, "the driver's fence, not the reload's");
552 }
553
554 #[test]
558 fn a_failed_exec_rewakes_only_the_workers_the_reload_stopped() {
559 let mut registry = SubagentRegistry::new();
560 registry.register(make_handle("sa_1"));
561 registry.register(make_finished_handle("sa_2"));
562 let stopped_earlier = CarriedWorker {
563 id: "sa_3".into(),
564 subagent_id: 3,
565 agent: String::new(),
566 status: "cancelled".into(),
567 archive_path: Some("/a/sa_3.json".into()),
568 stopped_by_reload: true,
569 wake_lost: false,
570 wake: true,
571 elapsed_secs: 0,
572 preview: String::new(),
573 };
574 registry.adopt_carried(&stopped_earlier);
575 registry.note_reload_wake(&stopped_earlier);
576 registry.stop_for_reload();
577 let queued: std::collections::HashSet<String> = ["sa_3".to_string()].into();
578 let carried = registry.reload_carry(&queued);
579 assert_eq!(carried.iter().filter(|w| w.needs_wake()).count(), 3, "{carried:?}");
580
581 let rewake = registry.rewake_after_failed_exec(carried.clone());
582 let woken: Vec<&str> =
583 rewake.iter().filter(|w| w.needs_wake()).map(|w| w.id.as_str()).collect();
584 assert_eq!(woken, ["sa_1"], "only this reload's own stop: {rewake:?}");
585 assert_eq!(rewake.len(), carried.len(), "every record is kept");
586 }
587
588 #[test]
589 fn reload_carry_caps_reaped_archives_keeping_the_newest() {
590 let mut registry = SubagentRegistry::new();
591 for n in 0..(MAX_CARRIED_ARCHIVES as u64 + 10) {
592 registry
593 .archived
594 .insert(format!("sa_{n}"), format!("/a/sa_{n}.json").into());
595 }
596 let carried = registry.reload_carry(&Default::default());
597 assert_eq!(carried.len(), MAX_CARRIED_ARCHIVES);
598 assert_eq!(carried.first().unwrap().id, "sa_10");
599 assert_eq!(carried.last().unwrap().id, format!("sa_{}", MAX_CARRIED_ARCHIVES as u64 + 9));
600 }
601
602 #[test]
603 fn carried_worker_reads_records_without_the_optional_fields() {
604 let w: CarriedWorker =
605 serde_json::from_value(serde_json::json!({"id": "sa_7", "status": "cancelled"})).unwrap();
606 assert_eq!(w.id, "sa_7");
607 assert!(!w.needs_wake());
608 assert_eq!(handle_number("sa_7"), Some(7));
609 assert_eq!(handle_number("x_7"), None);
610 }
611}