1use std::collections::{HashMap, HashSet};
9use std::path::PathBuf;
10use std::sync::{Arc, Mutex};
11use std::time::{Duration, Instant};
12
13use futures::stream::BoxStream;
14use serde::{Deserialize, Serialize};
15use tokio::io::AsyncBufReadExt;
16use tokio::sync::{Mutex as AsyncMutex, RwLock, watch};
17
18use crate::bus::{EventBus, InProcessBus};
19use crate::caliband::client::CalibandClient;
20use crate::caliband::stream::{NormalizeOptions, Normalized, normalize_frame};
21use crate::caliband::wire::{AgentRecord, AttachInbound, Endpoint, SpawnSpec};
22use crate::config_store::{ConfigStore, SqliteConfigStore};
23use crate::discovery::{DiscoveryEnv, EnsureConfig, ensure_caliband};
24use crate::error::{CoreError, Result};
25use crate::event::{EventKind, FleetEvent};
26use crate::metrics::{Metrics, MetricsSnapshot};
27use crate::model::{
28 Agent, AgentId, AgentStatus, FleetChange, FleetSnapshot, Workspace, WorkspaceHealth,
29};
30use crate::ownership::{Ownership, SelfOwnsAll};
31use crate::registry::Registry;
32use crate::store::Store;
33
34#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
38pub struct SpawnRequest {
39 pub prompt: String,
41 pub label: Option<String>,
43 pub model: Option<String>,
45 pub isolation_worktree: bool,
48 pub tool_allowlist: Option<Vec<String>>,
50 pub interactive: bool,
53 pub frontmatter_path: Option<PathBuf>,
56 pub provider_ref: Option<String>,
61}
62
63impl SpawnRequest {
64 pub fn new(prompt: impl Into<String>) -> Self {
66 Self {
67 prompt: prompt.into(),
68 label: None,
69 model: None,
70 isolation_worktree: true,
71 tool_allowlist: None,
72 interactive: false,
73 frontmatter_path: None,
74 provider_ref: None,
75 }
76 }
77
78 fn into_spec(self) -> SpawnSpec {
79 SpawnSpec {
80 label: self.label,
81 frontmatter_path: self.frontmatter_path,
82 initial_prompt: self.prompt,
83 model: self.model,
84 provider: None,
87 tool_allowlist: self.tool_allowlist,
88 isolation_worktree: self.isolation_worktree,
89 inherit_hooks: true,
90 interactive: self.interactive,
91 }
92 }
93}
94
95#[derive(Debug, Clone)]
97pub struct FleetConfig {
98 pub host: String,
100 pub data_dir: PathBuf,
102 pub poll_interval: Duration,
104 pub discovery_env: DiscoveryEnv,
106 pub ensure: EnsureConfig,
108 pub normalize: NormalizeOptions,
110 pub event_buffer: usize,
112 pub default_env: std::collections::BTreeMap<String, String>,
114 pub attach_backoff: AttachBackoff,
116 pub caliband_network: Option<CalibandNetworkConfig>,
122}
123
124#[derive(Debug, Clone)]
129pub struct CalibandNetworkConfig {
130 pub addr: String,
132 pub ca_pem: Vec<u8>,
134 pub server_name: String,
136 pub token: Option<String>,
138}
139
140#[derive(Debug, Clone, Copy)]
149pub struct AttachBackoff {
150 pub base: Duration,
152 pub max: Duration,
154 pub max_retries: u32,
156}
157
158impl Default for AttachBackoff {
159 fn default() -> Self {
160 Self {
161 base: Duration::from_millis(200),
162 max: Duration::from_secs(10),
163 max_retries: 8,
164 }
165 }
166}
167
168impl AttachBackoff {
169 fn delay_for(&self, agent_id: &str, attempt: u32) -> Duration {
173 use std::hash::{Hash, Hasher};
174 let exp = self
175 .base
176 .saturating_mul(2u32.saturating_pow(attempt.min(31)));
177 let capped = exp.min(self.max);
178 let mut h = std::collections::hash_map::DefaultHasher::new();
179 agent_id.hash(&mut h);
180 attempt.hash(&mut h);
181 let frac = 0.5 + 0.5 * ((h.finish() % 1000) as f64 / 1000.0);
182 Duration::from_secs_f64(capped.as_secs_f64() * frac)
183 }
184}
185
186impl FleetConfig {
187 pub fn new(host: impl Into<String>, data_dir: impl Into<PathBuf>) -> Self {
189 Self {
190 host: host.into(),
191 data_dir: data_dir.into(),
192 poll_interval: Duration::from_secs(2),
193 discovery_env: DiscoveryEnv::from_process(),
194 ensure: EnsureConfig::default(),
195 normalize: NormalizeOptions::default(),
196 event_buffer: 1024,
197 default_env: std::collections::BTreeMap::new(),
198 attach_backoff: AttachBackoff::default(),
199 caliband_network: None,
200 }
201 }
202
203 fn network_client(&self) -> Result<Option<CalibandClient>> {
208 match &self.caliband_network {
209 None => Ok(None),
210 Some(net) => {
211 let tls =
212 crate::caliband::transport::tls_client_from_pem(&net.ca_pem, &net.server_name)?;
213 Ok(Some(CalibandClient::connect_tcp(
214 net.addr.clone(),
215 Some(tls),
216 net.token.clone(),
217 )))
218 }
219 }
220 }
221}
222
223#[derive(Clone)]
231pub(crate) struct Emitter {
232 store: Arc<dyn Store>,
233 bus: Arc<dyn EventBus>,
234 seqs: Arc<AsyncMutex<HashMap<String, u64>>>,
236 metrics: Arc<Metrics>,
237}
238
239impl Emitter {
240 #[cfg(feature = "k8s")]
247 pub(crate) fn new(bus: Arc<dyn EventBus>, store: Arc<dyn Store>) -> Self {
248 Self {
249 store,
250 bus,
251 seqs: Arc::new(AsyncMutex::new(HashMap::new())),
252 metrics: Arc::new(Metrics::default()),
253 }
254 }
255
256 #[cfg(feature = "k8s")]
259 pub(crate) fn store(&self) -> Arc<dyn Store> {
260 self.store.clone()
261 }
262
263 #[cfg(feature = "k8s")]
266 pub(crate) fn metrics_snapshot(&self, active: u64) -> MetricsSnapshot {
267 self.metrics.snapshot(active)
268 }
269
270 async fn next_event(&self, repo: &str, agent_id: &str, kind: EventKind) -> FleetEvent {
271 let stream_key = crate::event::stream_key_for(repo, agent_id);
272 let seq = self.next_seq(&stream_key).await;
273 FleetEvent {
274 seq,
275 ts: chrono::Utc::now().to_rfc3339(),
276 repo: repo.to_string(),
277 agent_id: agent_id.to_string(),
278 kind,
279 }
280 }
281
282 async fn next_seq(&self, stream_key: &str) -> u64 {
286 {
288 let mut seqs = self.seqs.lock().await;
289 if let Some(n) = seqs.get(stream_key) {
290 let next = n + 1;
291 seqs.insert(stream_key.to_string(), next);
292 return next;
293 }
294 }
295 let seeded = self.store.high_water(stream_key).await.unwrap_or_else(|e| {
297 tracing::warn!(
298 target: "prospero_fleet", stream = %stream_key, error = %e,
299 "high_water read failed seeding per-stream seq; starting at 0"
300 );
301 0
302 });
303 let mut seqs = self.seqs.lock().await;
308 let next = match seqs.get(stream_key) {
310 Some(n) => n + 1,
311 None => seeded + 1,
312 };
313 seqs.insert(stream_key.to_string(), next);
314 next
315 }
316
317 pub(crate) async fn emit(&self, repo: &str, agent_id: &str, kind: EventKind) {
323 let mut event = self.next_event(repo, agent_id, kind).await;
324 let append_err = self.append_with_seq_retry(&mut event).await;
325 let lost_seq = event.seq;
326 match &append_err {
327 None => self.metrics.record_append_ok(),
328 Some(_) => self.metrics.record_append_failure(),
329 }
330 self.bus.publish(event);
334 if let Some(e) = append_err {
335 tracing::warn!(target: "prospero_fleet", error = %e, "failed to persist event");
336 self.emit_persist_gap(repo, agent_id, lost_seq, e).await;
337 }
338 }
339
340 async fn append_with_seq_retry(&self, event: &mut FleetEvent) -> Option<CoreError> {
345 const SEQ_CONFLICT_RETRIES: u32 = 8;
346 let stream_key = event.stream_key();
347 let mut attempts = 0;
348 loop {
349 match self.store.append(event).await {
350 Ok(()) => return None,
351 Err(CoreError::SeqConflict) if attempts < SEQ_CONFLICT_RETRIES => {
352 attempts += 1;
353 self.reseed_seq(&stream_key).await;
356 event.seq = self.next_seq(&stream_key).await;
357 }
358 Err(e) => return Some(e),
359 }
360 }
361 }
362
363 async fn reseed_seq(&self, stream_key: &str) {
366 self.seqs.lock().await.remove(stream_key);
367 }
368
369 async fn emit_persist_gap(&self, repo: &str, agent_id: &str, lost_seq: u64, err: CoreError) {
375 let mut marker = self
376 .next_event(
377 repo,
378 agent_id,
379 EventKind::StorePersistFailed {
380 lost_seq,
381 detail: err.to_string(),
382 },
383 )
384 .await;
385 match self.append_with_seq_retry(&mut marker).await {
386 None => self.metrics.record_append_ok(),
387 Some(e) => {
388 self.metrics.record_append_failure();
391 tracing::warn!(target: "prospero_fleet", error = %e, "failed to persist store-gap marker");
392 }
393 }
394 self.bus.publish(marker);
395 }
396}
397
398struct Inner {
399 config: FleetConfig,
400 snapshot: RwLock<FleetSnapshot>,
401 registry: RwLock<Registry>,
402 config_store: Arc<dyn ConfigStore>,
403 clients: Mutex<HashMap<String, CalibandClient>>,
405 attached: Mutex<HashSet<String>>,
407 recent_spawns: Mutex<HashMap<String, String>>,
414 emitter: Emitter,
415 ownership: Arc<dyn Ownership>,
416 shutdown: watch::Sender<bool>,
419}
420
421#[derive(Clone)]
423pub struct FleetManager {
424 inner: Arc<Inner>,
425}
426
427impl FleetManager {
428 pub async fn new(config: FleetConfig, store: Arc<dyn Store>) -> Result<Self> {
433 let config_store = Arc::new(SqliteConfigStore::open(&config.data_dir).await?);
434 Self::with_config_store(config, store, config_store).await
435 }
436
437 pub async fn with_config_store(
440 config: FleetConfig,
441 store: Arc<dyn Store>,
442 config_store: Arc<dyn ConfigStore>,
443 ) -> Result<Self> {
444 let bus: Arc<dyn EventBus> = Arc::new(InProcessBus::new(config.event_buffer));
445 let ownership: Arc<dyn Ownership> = Arc::new(SelfOwnsAll);
446 Self::with_seams(config, store, config_store, bus, ownership).await
447 }
448
449 pub async fn with_seams(
453 config: FleetConfig,
454 store: Arc<dyn Store>,
455 config_store: Arc<dyn ConfigStore>,
456 bus: Arc<dyn EventBus>,
457 ownership: Arc<dyn Ownership>,
458 ) -> Result<Self> {
459 let registry = Registry {
460 workspaces: config_store.list_repos().await?,
461 };
462 let emitter = Emitter {
463 store,
464 bus,
465 seqs: Arc::new(AsyncMutex::new(HashMap::new())),
466 metrics: Arc::new(Metrics::default()),
467 };
468 let snapshot = FleetSnapshot {
469 host: config.host.clone(),
470 workspaces: registry
471 .workspaces
472 .iter()
473 .map(|r| Workspace {
474 name: r.name.clone(),
475 root: r.root.clone(),
476 sources: crate::caliband::sources::discover_sources(&r.root),
477 health: WorkspaceHealth::Healthy,
478 config: r.config.clone(),
479 agents: Vec::new(),
480 })
481 .collect(),
482 };
483 Ok(Self {
484 inner: Arc::new(Inner {
485 config,
486 snapshot: RwLock::new(snapshot),
487 registry: RwLock::new(registry),
488 config_store,
489 clients: Mutex::new(HashMap::new()),
490 attached: Mutex::new(HashSet::new()),
491 recent_spawns: Mutex::new(HashMap::new()),
492 emitter,
493 ownership,
494 shutdown: watch::channel(false).0,
495 }),
496 })
497 }
498
499 pub fn begin_shutdown(&self) {
505 self.inner.shutdown.send_replace(true);
506 }
507
508 pub fn subscribe(&self, stream_key: &str) -> crate::bus::BusSubscription {
512 self.inner.emitter.bus.subscribe(stream_key)
513 }
514
515 #[must_use]
519 pub fn store(&self) -> Arc<dyn Store> {
520 self.inner.emitter.store.clone()
521 }
522
523 #[must_use]
525 pub fn bus(&self) -> Arc<dyn EventBus> {
526 self.inner.emitter.bus.clone()
527 }
528
529 fn subscribe_all(&self) -> crate::bus::BusSubscription {
535 self.inner.emitter.bus.subscribe_all()
536 }
537
538 pub async fn snapshot(&self) -> FleetSnapshot {
541 let mut snap = self.inner.snapshot.read().await.clone();
542 let reg = self.inner.registry.read().await;
543 for repo in &mut snap.workspaces {
544 if let Some(r) = reg.get(&repo.name) {
545 repo.config = r.config.clone();
546 }
547 }
548 snap
549 }
550
551 pub fn metrics(&self) -> MetricsSnapshot {
554 let active = self.inner.attached.lock().unwrap().len() as u64;
555 self.inner.emitter.metrics.snapshot(active)
556 }
557
558 pub async fn readiness(&self) -> crate::model::Readiness {
562 let store_writable = self.inner.emitter.store.writable().await;
563 let snap = self.inner.snapshot.read().await;
564 let workspaces_total = snap.workspaces.len();
565 let workspaces_healthy = snap
566 .workspaces
567 .iter()
568 .filter(|r| matches!(r.health, WorkspaceHealth::Healthy))
569 .count();
570 crate::model::Readiness {
571 ready: store_writable,
572 store_writable,
573 workspaces_total,
574 workspaces_healthy,
575 workspaces_unreachable: workspaces_total - workspaces_healthy,
576 }
577 }
578
579 pub async fn history(&self, stream_key: &str, from_seq: u64) -> Result<Vec<FleetEvent>> {
584 self.inner.emitter.store.replay(stream_key, from_seq).await
585 }
586
587 pub async fn prune_older_than(&self, max_age: std::time::Duration) -> Result<u64> {
590 crate::store::prune_store_older_than(self.inner.emitter.store.as_ref(), max_age).await
591 }
592
593 pub async fn add_workspace(
595 &self,
596 name: impl Into<String>,
597 root: impl Into<PathBuf>,
598 ) -> Result<()> {
599 self.add_workspace_with_config(name, root, Default::default())
600 .await
601 }
602
603 pub async fn add_repo(&self, name: impl Into<String>, root: impl Into<PathBuf>) -> Result<()> {
605 self.add_workspace(name, root).await
606 }
607
608 pub async fn add_repo_with_config(
610 &self,
611 name: impl Into<String>,
612 root: impl Into<PathBuf>,
613 config: crate::registry::RepoProviderConfig,
614 ) -> Result<()> {
615 self.add_workspace_with_config(name, root, config).await
616 }
617
618 pub async fn add_workspace_with_config(
620 &self,
621 name: impl Into<String>,
622 root: impl Into<PathBuf>,
623 config: crate::registry::RepoProviderConfig,
624 ) -> Result<()> {
625 let name = name.into();
626 let root = root.into();
627 crate::provider_env::validate_provider_config(&config)
631 .map_err(CoreError::ProviderMisconfigured)?;
632 let root = std::fs::canonicalize(&root).unwrap_or(root);
638 let repo = {
639 let mut reg = self.inner.registry.write().await;
640 reg.add(name.clone(), root.clone())?;
641 reg.set_config(&name, config);
642 reg.get(&name)
643 .cloned()
644 .expect("repo just inserted must exist")
645 };
646 self.inner.config_store.upsert_repo(&repo).await?;
647 {
648 let mut snap = self.inner.snapshot.write().await;
649 if !snap.workspaces.iter().any(|r| r.name == name) {
650 snap.workspaces.push(Workspace {
651 name: name.clone(),
652 root: root.clone(),
653 sources: crate::caliband::sources::discover_sources(&root),
654 health: WorkspaceHealth::Healthy,
655 config: repo.config.clone(),
656 agents: Vec::new(),
657 });
658 }
659 }
660 self.poll_repo_once(&name).await;
661 Ok(())
662 }
663
664 pub async fn repo_config(&self, repo: &str) -> Option<crate::registry::RepoProviderConfig> {
666 self.inner
667 .registry
668 .read()
669 .await
670 .get(repo)
671 .map(|r| r.config.clone())
672 }
673
674 pub async fn remove_repo(&self, name: &str) -> Result<bool> {
676 let removed = {
677 let mut reg = self.inner.registry.write().await;
678 reg.remove(name)
679 };
680 if removed {
681 self.inner.config_store.delete_repo(name).await?;
683 self.inner
684 .snapshot
685 .write()
686 .await
687 .workspaces
688 .retain(|r| r.name != name);
689 self.inner.clients.lock().unwrap().remove(name);
690 }
691 Ok(removed)
692 }
693
694 pub async fn ensure_config_for(&self, repo: &str) -> Result<EnsureConfig> {
697 let cfg = {
698 let reg = self.inner.registry.read().await;
699 reg.get(repo)
700 .map(|r| r.config.clone())
701 .ok_or_else(|| CoreError::WorkspaceNotFound(repo.to_string()))?
702 };
703 let env = crate::provider_env::resolve_env(&self.inner.config.default_env, &cfg, &|k| {
704 std::env::var(k).ok()
705 });
706 let mut ensure = self.inner.config.ensure.clone();
707 ensure.env = env;
708 Ok(ensure)
709 }
710
711 pub async fn set_repo_config_registry_only(
713 &self,
714 repo: &str,
715 config: crate::registry::RepoProviderConfig,
716 ) -> Result<()> {
717 crate::provider_env::validate_provider_config(&config)
722 .map_err(CoreError::ProviderMisconfigured)?;
723 let mut reg = self.inner.registry.write().await;
729 if !reg.set_config(repo, config) {
730 return Err(CoreError::WorkspaceNotFound(repo.to_string()));
731 }
732 let record = reg
733 .get(repo)
734 .cloned()
735 .expect("repo exists after successful set_config");
736 self.inner.config_store.upsert_repo(&record).await?;
737 Ok(())
738 }
739
740 async fn client_for(&self, repo: &str) -> Result<CalibandClient> {
742 if let Some(c) = self.inner.clients.lock().unwrap().get(repo).cloned() {
743 return Ok(c);
744 }
745 if let Some(client) = self.inner.config.network_client()? {
749 self.inner
750 .clients
751 .lock()
752 .unwrap()
753 .insert(repo.to_string(), client.clone());
754 return Ok(client);
755 }
756 let root = {
757 let reg = self.inner.registry.read().await;
758 reg.get(repo)
759 .map(|r| r.root.clone())
760 .ok_or_else(|| CoreError::WorkspaceNotFound(repo.to_string()))?
761 };
762 let ensure = self.ensure_config_for(repo).await?;
763 let client = ensure_caliband(&root, &self.inner.config.discovery_env, &ensure).await?;
764 self.inner
765 .clients
766 .lock()
767 .unwrap()
768 .insert(repo.to_string(), client.clone());
769 Ok(client)
770 }
771
772 async fn validate_provider_env(&self, repo: &str) -> Result<()> {
777 let cfg = {
778 let reg = self.inner.registry.read().await;
779 reg.get(repo)
780 .map(|r| r.config.clone())
781 .ok_or_else(|| CoreError::WorkspaceNotFound(repo.to_string()))?
782 };
783 let env = crate::provider_env::resolve_env(&self.inner.config.default_env, &cfg, &|k| {
784 std::env::var(k).ok()
785 });
786 crate::provider_env::validate_provider_env(&cfg, &env)
787 .map_err(CoreError::ProviderMisconfigured)
788 }
789
790 pub async fn spawn_agent(&self, repo: &str, req: SpawnRequest) -> Result<String> {
792 Ok(self.spawn_agent_with_socket(repo, req).await?.0)
793 }
794
795 pub async fn spawn_agent_with_socket(
800 &self,
801 repo: &str,
802 req: SpawnRequest,
803 ) -> Result<(String, Endpoint)> {
804 self.validate_provider_env(repo).await?;
805 let client = self.client_for(repo).await?;
806 let mut spec = req.into_spec();
807 spec.provider = self.repo_config(repo).await.and_then(|c| c.provider);
812 let (id, endpoint) = client.spawn(spec).await?;
813 self.inner
817 .recent_spawns
818 .lock()
819 .unwrap()
820 .insert(id.clone(), repo.to_string());
821 self.inner
822 .emitter
823 .emit(repo, &id, EventKind::AgentSpawned)
824 .await;
825 self.start_attach(repo, &id, client).await;
826 Ok((id, endpoint))
827 }
828
829 pub async fn kill_agent(&self, agent_id: &str) -> Result<()> {
831 let repo = self.repo_of(agent_id).await?;
832 self.client_for(&repo).await?.kill(agent_id).await
833 }
834
835 pub async fn send_agent_input(&self, agent_id: &str, input: AttachInbound) -> Result<()> {
843 let (repo, interactive, terminal) = {
844 let snap = self.inner.snapshot.read().await;
845 let (repo, agent) = snap
846 .find_agent(agent_id)
847 .ok_or_else(|| CoreError::AgentNotFound(agent_id.to_string()))?;
848 (
849 repo.to_string(),
850 agent.interactive,
851 agent.status.is_terminal(),
852 )
853 };
854 if terminal {
855 return Err(CoreError::InvalidState {
856 op: "send_input".into(),
857 id: agent_id.to_string(),
858 status: "terminal".into(),
859 });
860 }
861 if !interactive {
862 return Err(CoreError::InvalidState {
863 op: "send_input".into(),
864 id: agent_id.to_string(),
865 status: "not interactive".into(),
866 });
867 }
868 let client = self.client_for(&repo).await?;
869 let endpoint = client.attach(agent_id).await?;
870 client.send_inbound(&endpoint, &input).await
871 }
872
873 pub async fn respawn_agent(&self, agent_id: &str) -> Result<String> {
875 let repo = self.repo_of(agent_id).await?;
876 self.client_for(&repo).await?.respawn(agent_id).await
877 }
878
879 pub async fn drain_agent(&self, agent_id: &str, timeout: Duration) -> Result<()> {
887 let _ = self
888 .send_agent_input(agent_id, AttachInbound::EndInput)
889 .await;
890
891 let poll_interval = Duration::from_millis(20).min(timeout);
892 let deadline = tokio::time::Instant::now() + timeout;
893 while tokio::time::Instant::now() < deadline {
894 let snap = self.snapshot().await;
895 match snap.find_agent(agent_id) {
896 Some((_, agent)) if agent.status.is_terminal() => break,
897 None => break,
898 _ => tokio::time::sleep(poll_interval).await,
899 }
900 }
901
902 self.kill_agent(agent_id).await
903 }
904
905 pub fn watch_changes(&self) -> BoxStream<'static, FleetChange> {
917 let live = self.subscribe_all();
918
919 let seed_mgr = self.clone();
920 let seed = futures::stream::once(async move {
921 let snap = seed_mgr.snapshot().await;
922 let mut items = Vec::new();
923 for repo in snap.workspaces {
924 items.push(FleetChange::WorkspaceHealth {
925 workspace: repo.name.clone(),
926 health: repo.health.clone(),
927 });
928 for agent in repo.agents {
929 items.push(FleetChange::Discovered {
930 id: AgentId::from(agent.id.clone()),
931 workspace: repo.name.clone(),
932 agent,
933 });
934 }
935 }
936 futures::stream::iter(items)
937 });
938 let seed = futures::stream::StreamExt::flatten(seed);
939
940 let live_mgr = self.clone();
941 let live_changes = futures::stream::StreamExt::filter_map(live, move |be| {
942 let mgr = live_mgr.clone();
943 async move {
944 match be {
945 crate::bus::BusEvent::Event(ev) => event_to_change(&mgr, ev).await,
946 crate::bus::BusEvent::Lagged(n) => {
951 tracing::warn!(
952 target: "prospero_fleet", skipped = n,
953 "watch_changes subscriber lagged; some live FleetChanges were dropped"
954 );
955 None
956 }
957 }
958 }
959 });
960
961 Box::pin(futures::stream::StreamExt::chain(seed, live_changes))
962 }
963
964 pub async fn rm_agent(&self, agent_id: &str, force: bool) -> Result<()> {
971 let repo = self.repo_of(agent_id).await?;
972 self.client_for(&repo).await?.rm(agent_id, force).await?;
973
974 self.inner.recent_spawns.lock().unwrap().remove(agent_id);
977
978 let removed = {
980 let mut snap = self.inner.snapshot.write().await;
981 match snap.workspaces.iter_mut().find(|r| r.name == repo) {
982 Some(r) => {
983 let before = r.agents.len();
984 r.agents.retain(|a| a.id != agent_id);
985 r.agents.len() != before
986 }
987 None => false,
988 }
989 };
990
991 if removed
999 && self
1000 .inner
1001 .ownership
1002 .try_acquire(&crate::event::stream_key_for(&repo, ""))
1003 .await
1004 .is_some()
1005 {
1006 self.inner
1007 .emitter
1008 .emit(&repo, agent_id, EventKind::AgentGone)
1009 .await;
1010 }
1011 Ok(())
1012 }
1013
1014 async fn repo_of(&self, agent_id: &str) -> Result<String> {
1015 if let Some(repo) = self
1016 .inner
1017 .snapshot
1018 .read()
1019 .await
1020 .find_agent(agent_id)
1021 .map(|(repo, _)| repo.to_string())
1022 {
1023 return Ok(repo);
1024 }
1025 self.inner
1029 .recent_spawns
1030 .lock()
1031 .unwrap()
1032 .get(agent_id)
1033 .cloned()
1034 .ok_or_else(|| CoreError::AgentNotFound(agent_id.to_string()))
1035 }
1036
1037 pub async fn poll_all_once(&self) {
1041 self.refresh_registry_from_store().await;
1042 let names: Vec<String> = {
1043 let reg = self.inner.registry.read().await;
1044 reg.workspaces.iter().map(|r| r.name.clone()).collect()
1045 };
1046 for name in names {
1047 self.poll_repo_once(&name).await;
1048 }
1049 }
1050
1051 async fn refresh_registry_from_store(&self) {
1058 let durable = {
1064 let mut reg = self.inner.registry.write().await;
1065 let durable = match self.inner.config_store.list_repos().await {
1066 Ok(repos) => repos,
1067 Err(e) => {
1068 tracing::warn!(
1069 target: "prospero_fleet", error = %e,
1070 "registry refresh from config store failed; keeping cached view"
1071 );
1072 return;
1073 }
1074 };
1075 reg.workspaces = durable.clone();
1076 durable
1077 };
1078 let mut snap = self.inner.snapshot.write().await;
1079 snap.workspaces
1080 .retain(|r| durable.iter().any(|d| d.name == r.name));
1081 for d in &durable {
1082 if !snap.workspaces.iter().any(|r| r.name == d.name) {
1083 snap.workspaces.push(Workspace {
1085 name: d.name.clone(),
1086 root: d.root.clone(),
1087 sources: crate::caliband::sources::discover_sources(&d.root),
1088 health: WorkspaceHealth::Healthy,
1089 config: d.config.clone(),
1090 agents: Vec::new(),
1091 });
1092 }
1093 }
1094 }
1095
1096 pub async fn poll_repo_once(&self, repo: &str) {
1100 self.inner.emitter.metrics.record_repo_poll();
1101 let own_lifecycle = self
1107 .inner
1108 .ownership
1109 .try_acquire(&crate::event::stream_key_for(repo, ""))
1110 .await
1111 .is_some();
1112 let client = match self.client_for(repo).await {
1113 Ok(c) => c,
1114 Err(e) => {
1115 self.mark_unreachable(repo, e.to_string(), own_lifecycle)
1116 .await;
1117 return;
1118 }
1119 };
1120 match client.list().await {
1121 Ok(records) => self.reconcile(repo, records, client, own_lifecycle).await,
1122 Err(e) => {
1123 self.inner.clients.lock().unwrap().remove(repo);
1126 self.mark_unreachable(repo, e.to_string(), own_lifecycle)
1127 .await;
1128 }
1129 }
1130 }
1131
1132 async fn mark_unreachable(&self, repo: &str, reason: String, own_lifecycle: bool) {
1133 let mut snap = self.inner.snapshot.write().await;
1134 if let Some(r) = snap.workspaces.iter_mut().find(|r| r.name == repo) {
1135 let new_health = WorkspaceHealth::Unreachable {
1136 reason: reason.clone(),
1137 };
1138 if r.health != new_health {
1139 r.health = new_health.clone();
1140 drop(snap);
1141 if own_lifecycle {
1144 self.inner
1145 .emitter
1146 .emit(repo, "", EventKind::RepoHealth { state: new_health })
1147 .await;
1148 }
1149 }
1150 }
1151 }
1152
1153 async fn reconcile(
1154 &self,
1155 repo: &str,
1156 records: Vec<AgentRecord>,
1157 client: CalibandClient,
1158 own_lifecycle: bool,
1159 ) {
1160 let prior: HashMap<String, AgentStatus> = {
1162 let snap = self.inner.snapshot.read().await;
1163 snap.workspaces
1164 .iter()
1165 .find(|r| r.name == repo)
1166 .map(|r| r.agents.iter().map(|a| (a.id.clone(), a.status)).collect())
1167 .unwrap_or_default()
1168 };
1169
1170 let mut new_agents = Vec::new();
1171 let mut to_attach: Vec<String> = Vec::new();
1172 let attached_now = self.inner.attached.lock().unwrap().clone();
1173
1174 for rec in &records {
1175 let agent = Agent {
1176 id: rec.id.clone(),
1177 name: rec.name.clone(),
1178 workspace: repo.to_string(),
1179 status: rec.status,
1180 started_at: rec.started_at.clone(),
1181 isolated: rec.spec.isolation_worktree,
1182 interactive: rec.spec.interactive,
1183 session_dir: rec.session_dir.clone(),
1184 };
1185 match prior.get(&rec.id) {
1186 None if own_lifecycle && !attached_now.contains(&rec.id) => {
1191 self.inner
1192 .emitter
1193 .emit(repo, &rec.id, EventKind::AgentDiscovered)
1194 .await;
1195 }
1196 None => {}
1197 Some(&old) if own_lifecycle && old != rec.status => {
1199 self.inner
1200 .emitter
1201 .emit(
1202 repo,
1203 &rec.id,
1204 EventKind::StatusChanged {
1205 from: old,
1206 to: rec.status,
1207 },
1208 )
1209 .await;
1210 }
1211 _ => {}
1212 }
1213 if !rec.status.is_terminal() && !attached_now.contains(&rec.id) {
1219 to_attach.push(rec.id.clone());
1220 }
1221 new_agents.push(agent);
1222 }
1223
1224 for old_id in prior.keys() {
1227 if own_lifecycle && !records.iter().any(|r| &r.id == old_id) {
1228 self.inner
1229 .emitter
1230 .emit(repo, old_id, EventKind::AgentGone)
1231 .await;
1232 }
1233 }
1234
1235 {
1236 let mut snap = self.inner.snapshot.write().await;
1237 if let Some(r) = snap.workspaces.iter_mut().find(|r| r.name == repo) {
1238 let was_unreachable = matches!(r.health, WorkspaceHealth::Unreachable { .. });
1239 r.health = WorkspaceHealth::Healthy;
1240 r.agents = new_agents;
1241 if was_unreachable {
1242 drop(snap);
1243 if own_lifecycle {
1246 self.inner
1247 .emitter
1248 .emit(
1249 repo,
1250 "",
1251 EventKind::RepoHealth {
1252 state: WorkspaceHealth::Healthy,
1253 },
1254 )
1255 .await;
1256 }
1257 }
1258 }
1259 }
1260
1261 if !records.is_empty() {
1265 let mut recent = self.inner.recent_spawns.lock().unwrap();
1266 recent.retain(|id, _| !records.iter().any(|r| &r.id == id));
1267 }
1268
1269 for id in to_attach {
1270 self.start_attach(repo, &id, client.clone()).await;
1271 }
1272 }
1273
1274 async fn start_attach(&self, repo: &str, agent_id: &str, client: CalibandClient) {
1278 if self.inner.attached.lock().unwrap().contains(agent_id) {
1281 return;
1282 }
1283 if self.inner.ownership.try_acquire(agent_id).await.is_none() {
1287 return;
1288 }
1289 {
1290 let mut attached = self.inner.attached.lock().unwrap();
1291 if !attached.insert(agent_id.to_string()) {
1292 return;
1296 }
1297 }
1298 let repo = repo.to_string();
1299 let agent_id = agent_id.to_string();
1300 let emitter = self.inner.emitter.clone();
1301 let normalize = self.inner.config.normalize;
1302 let backoff = self.inner.config.attach_backoff;
1303 let mut shutdown = self.inner.shutdown.subscribe();
1304 let attached = self.inner.clone();
1305
1306 tokio::spawn(async move {
1307 let result = attach_loop(
1308 &client,
1309 AttachTarget {
1312 repo: &repo,
1313 agent_id: &agent_id,
1314 attach_id: &agent_id,
1315 },
1316 &emitter,
1317 normalize,
1318 backoff,
1319 &mut shutdown,
1320 )
1321 .await;
1322 if let Err(e) = result {
1323 tracing::warn!(
1324 target: "prospero_fleet",
1325 %repo, %agent_id, error = %e,
1326 "attach task ended with error"
1327 );
1328 }
1329 attached.attached.lock().unwrap().remove(&agent_id);
1330 attached.ownership.release(&agent_id).await;
1332 });
1333 }
1334
1335 pub async fn cached_client_names(&self) -> Vec<String> {
1337 self.inner.clients.lock().unwrap().keys().cloned().collect()
1338 }
1339
1340 pub fn is_attached(&self, agent_id: &str) -> bool {
1342 self.inner.attached.lock().unwrap().contains(agent_id)
1343 }
1344
1345 pub async fn restart_caliband(&self, repo: &str) -> Result<()> {
1348 let client = self.inner.clients.lock().unwrap().get(repo).cloned();
1349 if let Some(client) = client {
1350 let res = client.shutdown().await;
1351 if let Err(e) = res {
1352 tracing::warn!(target: "prospero_fleet", repo, error = %e,
1353 "shutdown request to caliband failed (continuing)");
1354 }
1355 }
1356 self.inner.clients.lock().unwrap().remove(repo);
1357
1358 let root = {
1359 let reg = self.inner.registry.read().await;
1360 reg.get(repo).map(|r| r.root.clone())
1361 };
1362 if let Some(root) = root {
1363 let socket_res =
1364 crate::discovery::resolve_socket(&root, &self.inner.config.discovery_env);
1365 if let Ok(socket) = socket_res {
1366 let deadline =
1369 tokio::time::Instant::now() + self.inner.config.ensure.startup_timeout;
1370 while tokio::net::UnixStream::connect(&socket).await.is_ok() {
1371 if tokio::time::Instant::now() >= deadline {
1372 tracing::warn!(target: "prospero_fleet", repo,
1373 "old caliband socket still reachable after shutdown; proceeding");
1374 break;
1375 }
1376 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
1377 }
1378 }
1379 }
1380 self.poll_repo_once(repo).await;
1381 Ok(())
1382 }
1383
1384 pub async fn set_repo_config(
1386 &self,
1387 repo: &str,
1388 config: crate::registry::RepoProviderConfig,
1389 ) -> Result<()> {
1390 self.set_repo_config_registry_only(repo, config).await?;
1391 self.restart_caliband(repo).await
1392 }
1393
1394 pub async fn run(self) {
1401 let interval = self.inner.config.poll_interval;
1402 let mut shutdown = self.inner.shutdown.subscribe();
1403 if *shutdown.borrow_and_update() {
1404 return;
1405 }
1406 loop {
1407 self.poll_all_once().await;
1408 tokio::select! {
1409 _ = tokio::time::sleep(interval) => {}
1410 _ = shutdown.changed() => break,
1411 }
1412 }
1413 tracing::info!(target: "prospero_fleet", "poll loop drained on shutdown");
1414 }
1415}
1416
1417async fn event_to_change(mgr: &FleetManager, ev: FleetEvent) -> Option<FleetChange> {
1440 match ev.kind {
1441 EventKind::AgentDiscovered => {
1442 const RETRY_BUDGET: Duration = Duration::from_millis(250);
1443 const RETRY_INTERVAL: Duration = Duration::from_millis(5);
1444 let deadline = Instant::now() + RETRY_BUDGET;
1445 loop {
1446 {
1447 let snap = mgr.snapshot().await;
1448 if let Some((_, agent)) = snap.find_agent(&ev.agent_id) {
1449 return Some(FleetChange::Discovered {
1450 id: AgentId::from(ev.agent_id.clone()),
1451 workspace: ev.repo.clone(),
1452 agent: agent.clone(),
1453 });
1454 }
1455 }
1456 if Instant::now() >= deadline {
1457 break;
1458 }
1459 tokio::time::sleep(RETRY_INTERVAL).await;
1460 }
1461 tracing::warn!(
1462 target: "prospero_fleet",
1463 agent_id = %ev.agent_id, repo = %ev.repo,
1464 "watch_changes: AgentDiscovered fired but the agent never appeared in the \
1465 snapshot within the retry budget; dropping the FleetChange"
1466 );
1467 None
1468 }
1469 EventKind::StatusChanged { from, to } => Some(FleetChange::StatusChanged {
1470 id: AgentId::from(ev.agent_id.clone()),
1471 workspace: ev.repo.clone(),
1472 from,
1473 to,
1474 }),
1475 EventKind::AgentGone => Some(FleetChange::Gone {
1476 id: AgentId::from(ev.agent_id.clone()),
1477 workspace: ev.repo.clone(),
1478 }),
1479 EventKind::RepoHealth { state } => Some(FleetChange::WorkspaceHealth {
1480 workspace: ev.repo.clone(),
1481 health: state,
1482 }),
1483 EventKind::AgentSpawned
1486 | EventKind::AgentInit { .. }
1487 | EventKind::Output { .. }
1488 | EventKind::ToolStarted { .. }
1489 | EventKind::ToolFinished { .. }
1490 | EventKind::AgentFinished { .. }
1491 | EventKind::StorePersistFailed { .. } => None,
1492 }
1493}
1494
1495#[derive(Clone, Copy)]
1507pub(crate) struct AttachTarget<'a> {
1508 pub repo: &'a str,
1510 pub agent_id: &'a str,
1512 pub attach_id: &'a str,
1514}
1515
1516enum StreamOutcome {
1518 Finished,
1520 Disconnected,
1522}
1523
1524pub(crate) async fn attach_loop(
1547 client: &CalibandClient,
1548 target: AttachTarget<'_>,
1549 emitter: &Emitter,
1550 normalize: NormalizeOptions,
1551 backoff: AttachBackoff,
1552 shutdown: &mut watch::Receiver<bool>,
1553) -> Result<()> {
1554 let AttachTarget { repo, agent_id, .. } = target;
1555 let mut frames_seen: u64 = 0;
1556 let mut attempt: u32 = 0;
1557 loop {
1558 let before = frames_seen;
1559 let err = match attach_once(
1560 client,
1561 target,
1562 emitter,
1563 normalize,
1564 &mut frames_seen,
1565 shutdown,
1566 )
1567 .await
1568 {
1569 Ok(StreamOutcome::Finished) => return Ok(()),
1570 Ok(StreamOutcome::Disconnected) => None,
1571 Err(e) => Some(e),
1572 };
1573 if *shutdown.borrow() {
1575 return Ok(());
1576 }
1577 if frames_seen > before {
1579 attempt = 0;
1580 }
1581 if attempt >= backoff.max_retries {
1582 return match err {
1583 Some(e) => Err(e),
1584 None => {
1585 tracing::warn!(
1586 target: "prospero_fleet", %repo, %agent_id,
1587 "attach reconnection budget exhausted; poll loop will re-attach"
1588 );
1589 Ok(())
1590 }
1591 };
1592 }
1593 let delay = backoff.delay_for(agent_id, attempt);
1594 tracing::warn!(
1595 target: "prospero_fleet", %repo, %agent_id, attempt,
1596 delay_ms = delay.as_millis() as u64,
1597 reason = if err.is_some() { "error" } else { "premature-eof" },
1598 "attach stream dropped; reconnecting after backoff"
1599 );
1600 tokio::select! {
1601 _ = tokio::time::sleep(delay) => {}
1602 _ = shutdown.changed() => return Ok(()),
1603 }
1604 attempt += 1;
1605 }
1606}
1607
1608async fn attach_once(
1613 client: &CalibandClient,
1614 target: AttachTarget<'_>,
1615 emitter: &Emitter,
1616 normalize: NormalizeOptions,
1617 frames_seen: &mut u64,
1618 shutdown: &mut watch::Receiver<bool>,
1619) -> Result<StreamOutcome> {
1620 let AttachTarget {
1621 repo,
1622 agent_id,
1623 attach_id,
1624 } = target;
1625 let endpoint = client.attach(attach_id).await?;
1628 let mut reader = client.open_stream(&endpoint).await?;
1629 let mut line = String::new();
1630 let mut idx: u64 = 0;
1631 let mut saw_terminal = false;
1632 loop {
1633 line.clear();
1634 let n = tokio::select! {
1635 r = reader.read_line(&mut line) => r?,
1636 _ = shutdown.changed() => {
1637 return Ok(StreamOutcome::Finished);
1639 }
1640 };
1641 if n == 0 {
1642 return Ok(if saw_terminal {
1643 StreamOutcome::Finished
1644 } else {
1645 StreamOutcome::Disconnected
1646 });
1647 }
1648 let trimmed = line.trim_end();
1649 if trimmed.is_empty() {
1650 continue;
1651 }
1652 idx += 1;
1653 if idx <= *frames_seen {
1655 continue;
1656 }
1657 *frames_seen = idx;
1658 let frame: serde_json::Value = match serde_json::from_str(trimmed) {
1659 Ok(v) => v,
1660 Err(_) => {
1661 tracing::warn!(target: "prospero_fleet", %agent_id, "unparseable stream line");
1662 continue;
1663 }
1664 };
1665 match normalize_frame(&frame, normalize) {
1666 Normalized::Event(kind) => {
1667 if matches!(kind, EventKind::AgentFinished { .. }) {
1668 saw_terminal = true;
1669 }
1670 emitter.emit(repo, agent_id, kind).await;
1671 }
1672 Normalized::Dropped => {}
1673 Normalized::Unknown => {
1674 emitter.metrics.record_unknown_frame();
1675 tracing::debug!(target: "prospero_fleet", %agent_id, "unknown caliban frame type");
1676 }
1677 }
1678 }
1679}
1680
1681#[cfg(test)]
1682mod tests {
1683 use super::*;
1684
1685 struct SlowListConfigStore {
1690 repos: std::sync::Mutex<Vec<crate::registry::RegisteredWorkspace>>,
1691 read_delay: Duration,
1692 }
1693 impl SlowListConfigStore {
1694 fn new(read_delay: Duration) -> Self {
1695 Self {
1696 repos: std::sync::Mutex::new(Vec::new()),
1697 read_delay,
1698 }
1699 }
1700 }
1701 #[async_trait::async_trait]
1702 impl crate::config_store::ConfigStore for SlowListConfigStore {
1703 async fn list_repos(&self) -> Result<Vec<crate::registry::RegisteredWorkspace>> {
1704 let snapshot = self.repos.lock().unwrap().clone(); tokio::time::sleep(self.read_delay).await; Ok(snapshot)
1707 }
1708 async fn upsert_repo(&self, repo: &crate::registry::RegisteredWorkspace) -> Result<()> {
1709 let mut v = self.repos.lock().unwrap();
1710 if let Some(e) = v.iter_mut().find(|e| e.name == repo.name) {
1711 *e = repo.clone();
1712 } else {
1713 v.push(repo.clone());
1714 }
1715 Ok(())
1716 }
1717 async fn delete_repo(&self, name: &str) -> Result<bool> {
1718 let mut v = self.repos.lock().unwrap();
1719 let before = v.len();
1720 v.retain(|e| e.name != name);
1721 Ok(v.len() != before)
1722 }
1723 }
1724
1725 #[tokio::test]
1726 async fn concurrent_refresh_does_not_clobber_a_just_set_config() {
1727 let dir = tempfile::tempdir().unwrap();
1728 let mut config = FleetConfig::new("local", dir.path());
1729 config.ensure.autostart = false;
1730 let root = dir.path().join("r");
1731 std::fs::create_dir_all(&root).unwrap();
1732
1733 let store: Arc<dyn crate::store::Store> =
1734 Arc::new(crate::store::JsonlStore::open(dir.path()).unwrap());
1735 let cfg_store: Arc<dyn crate::config_store::ConfigStore> =
1736 Arc::new(SlowListConfigStore::new(Duration::from_millis(150)));
1737 let mgr = FleetManager::with_config_store(config, store, cfg_store)
1738 .await
1739 .unwrap();
1740 mgr.add_repo("r", &root).await.unwrap(); let m = mgr.clone();
1745 let refresh = tokio::spawn(async move { m.refresh_registry_from_store().await });
1746
1747 tokio::time::sleep(Duration::from_millis(30)).await;
1749 let cfg = crate::registry::RepoProviderConfig {
1750 provider: Some("ollama".to_string()),
1751 ..Default::default()
1752 };
1753 mgr.set_repo_config_registry_only("r", cfg).await.unwrap();
1754 refresh.await.unwrap();
1755
1756 let snap = mgr.snapshot().await;
1758 let repo = snap.workspaces.iter().find(|w| w.name == "r").unwrap();
1759 assert_eq!(
1760 repo.config.provider.as_deref(),
1761 Some("ollama"),
1762 "refresh clobbered a concurrent set_config back to durable; got {:?}",
1763 repo.config
1764 );
1765 }
1766
1767 #[test]
1768 fn fleet_config_network_yields_tcp_client() {
1769 let plain = FleetConfig::new("local", std::path::Path::new("/tmp/x"));
1772 assert!(plain.network_client().unwrap().is_none());
1773
1774 let cert = rcgen::generate_simple_self_signed(vec!["localhost".into()]).unwrap();
1775 let mut cfg = FleetConfig::new("local", std::path::Path::new("/tmp/x"));
1776 cfg.caliband_network = Some(CalibandNetworkConfig {
1777 addr: "h:9443".into(),
1778 ca_pem: cert.cert.pem().into_bytes(),
1779 server_name: "localhost".into(),
1780 token: Some("t".into()),
1781 });
1782 let client = cfg.network_client().unwrap().expect("network client");
1783 assert!(matches!(client.endpoint(), Endpoint::Tcp { .. }));
1784 }
1785
1786 #[test]
1787 fn spawn_request_forwards_frontmatter_to_spec() {
1788 let mut req = SpawnRequest::new("p");
1790 assert_eq!(req.clone().into_spec().frontmatter_path, None);
1791 req.frontmatter_path = Some(std::path::PathBuf::from("/tpl.md"));
1792 assert_eq!(
1793 req.into_spec().frontmatter_path,
1794 Some(std::path::PathBuf::from("/tpl.md"))
1795 );
1796 }
1797
1798 #[test]
1799 fn attach_backoff_is_exponential_capped_and_jittered() {
1800 let b = AttachBackoff {
1801 base: Duration::from_millis(100),
1802 max: Duration::from_millis(800),
1803 max_retries: 8,
1804 };
1805 for (attempt, exp_ms) in [(0u32, 100u64), (1, 200), (2, 400), (3, 800)] {
1807 let d = b.delay_for("agent-a", attempt).as_millis() as u64;
1808 assert!(
1809 d >= exp_ms / 2 && d <= exp_ms,
1810 "attempt {attempt}: {d}ms outside [{}, {exp_ms}]",
1811 exp_ms / 2
1812 );
1813 }
1814 let capped = b.delay_for("agent-a", 20).as_millis() as u64;
1816 assert!((400..=800).contains(&capped), "capped delay {capped}ms");
1817 assert_eq!(
1819 b.delay_for("agent-a", 2).as_millis(),
1820 b.delay_for("agent-a", 2).as_millis()
1821 );
1822 }
1823
1824 #[tokio::test]
1825 async fn restart_caliband_shuts_down_and_clears_client() {
1826 use crate::registry::RepoProviderConfig;
1827 use crate::testkit::FakeCaliband;
1828
1829 let dir = tempfile::tempdir().unwrap();
1830 let mut config = FleetConfig::new("local", dir.path());
1831 config.discovery_env.caliban_daemon_runtime_dir = Some(dir.path().to_path_buf());
1832 config.ensure.autostart = false; let root = dir.path().join("repo");
1834 std::fs::create_dir_all(&root).unwrap();
1835 let socket = crate::discovery::resolve_socket(&root, &config.discovery_env).unwrap();
1836
1837 let fake = FakeCaliband::start_at(&socket).await.unwrap();
1838 let store = std::sync::Arc::new(crate::store::JsonlStore::open(dir.path()).unwrap());
1839 let mgr = FleetManager::new(config, store).await.unwrap();
1840 mgr.add_repo("p", &root).await.unwrap();
1841
1842 mgr.poll_repo_once("p").await; mgr.set_repo_config("p", RepoProviderConfig::default())
1845 .await
1846 .unwrap();
1847
1848 assert_eq!(fake.shutdowns(), 1, "restart should send one Shutdown");
1849 assert!(
1850 mgr.cached_client_names().await.iter().all(|n| n != "p"),
1851 "cached client for the repo should be cleared after restart"
1852 );
1853 }
1854
1855 #[tokio::test]
1856 async fn send_agent_input_rejects_terminal_unknown_and_non_interactive() {
1857 use crate::caliband::wire::AttachInbound;
1858 use crate::model::AgentStatus;
1859 use crate::testkit::{FakeCaliband, test_record};
1860
1861 let dir = tempfile::tempdir().unwrap();
1862 let mut config = FleetConfig::new("local", dir.path());
1863 config.discovery_env.caliban_daemon_runtime_dir = Some(dir.path().to_path_buf());
1864 config.ensure.autostart = false;
1865 let root = dir.path().join("repo");
1866 std::fs::create_dir_all(&root).unwrap();
1867 let socket = crate::discovery::resolve_socket(&root, &config.discovery_env).unwrap();
1868
1869 let mut fake = FakeCaliband::start_at(&socket).await.unwrap();
1870 let mut done = test_record("ag-done", dir.path(), AgentStatus::Done, false);
1872 done.spec.interactive = true;
1873 fake.add_agent(done, vec![]).await;
1874 let idle = test_record("ag-idle", dir.path(), AgentStatus::Idle, false);
1876 fake.add_agent(idle, vec![]).await;
1877
1878 let store = std::sync::Arc::new(crate::store::JsonlStore::open(dir.path()).unwrap());
1879 let mgr = FleetManager::new(config, store).await.unwrap();
1880 mgr.add_repo("repo", &root).await.unwrap();
1881 mgr.poll_repo_once("repo").await;
1882
1883 let r1 = mgr
1884 .send_agent_input("ag-done", AttachInbound::EndInput)
1885 .await;
1886 assert!(
1887 matches!(r1, Err(CoreError::InvalidState { .. })),
1888 "terminal must reject"
1889 );
1890 let r2 = mgr
1891 .send_agent_input("ag-idle", AttachInbound::EndInput)
1892 .await;
1893 assert!(
1894 matches!(r2, Err(CoreError::InvalidState { .. })),
1895 "non-interactive must reject"
1896 );
1897 let r3 = mgr.send_agent_input("nope", AttachInbound::EndInput).await;
1898 assert!(
1899 matches!(r3, Err(CoreError::AgentNotFound(_))),
1900 "unknown id must 404"
1901 );
1902 }
1903
1904 #[tokio::test]
1905 async fn ownership_gates_the_attach_path() {
1906 use crate::testkit::FakeCaliband;
1909
1910 let dir = tempfile::tempdir().unwrap();
1911 let mut config = FleetConfig::new("local", dir.path());
1912 config.discovery_env.caliban_daemon_runtime_dir = Some(dir.path().to_path_buf());
1913 config.ensure.autostart = false;
1914 let root = dir.path().join("repo");
1915 std::fs::create_dir_all(&root).unwrap();
1916 let socket = crate::discovery::resolve_socket(&root, &config.discovery_env).unwrap();
1917 let _fake = FakeCaliband::start_at(&socket).await.unwrap();
1918
1919 let store = Arc::new(crate::store::JsonlStore::open(dir.path()).unwrap());
1920 let mgr = FleetManager::new(config, store).await.unwrap();
1921 mgr.add_repo("p", &root).await.unwrap();
1922
1923 let id = mgr.spawn_agent("p", SpawnRequest::new("hi")).await.unwrap();
1924 assert!(mgr.is_attached(&id), "owned agent must be attached");
1927 }
1928
1929 #[tokio::test]
1930 async fn refused_ownership_blocks_the_attach_path() {
1931 use crate::bus::InProcessBus;
1932 use crate::ownership::{Lease, Ownership};
1933 use crate::testkit::FakeCaliband;
1934 use async_trait::async_trait;
1935
1936 struct NeverOwns;
1938 #[async_trait]
1939 impl Ownership for NeverOwns {
1940 async fn try_acquire(&self, _: &str) -> Option<Lease> {
1941 None
1942 }
1943 async fn renew(&self, _: &Lease) -> crate::error::Result<()> {
1944 Ok(())
1945 }
1946 async fn release(&self, _: &str) {}
1947 fn owns(&self, _: &str) -> bool {
1948 false
1949 }
1950 }
1951
1952 let dir = tempfile::tempdir().unwrap();
1953 let mut config = FleetConfig::new("local", dir.path());
1954 config.discovery_env.caliban_daemon_runtime_dir = Some(dir.path().to_path_buf());
1955 config.ensure.autostart = false;
1956 let root = dir.path().join("repo");
1957 std::fs::create_dir_all(&root).unwrap();
1958 let socket = crate::discovery::resolve_socket(&root, &config.discovery_env).unwrap();
1959 let _fake = FakeCaliband::start_at(&socket).await.unwrap();
1960
1961 let store: Arc<dyn Store> = Arc::new(crate::store::JsonlStore::open(dir.path()).unwrap());
1962 let config_store: Arc<dyn ConfigStore> = Arc::new(
1963 crate::config_store::SqliteConfigStore::open(dir.path())
1964 .await
1965 .unwrap(),
1966 );
1967 let bus: Arc<dyn EventBus> = Arc::new(InProcessBus::new(config.event_buffer));
1968 let ownership: Arc<dyn Ownership> = Arc::new(NeverOwns);
1969 let mgr = FleetManager::with_seams(config, store, config_store, bus, ownership)
1970 .await
1971 .unwrap();
1972 mgr.add_repo("p", &root).await.unwrap();
1973
1974 let id = mgr.spawn_agent("p", SpawnRequest::new("hi")).await.unwrap();
1975 assert!(
1976 !mgr.is_attached(&id),
1977 "peer-owned agent must NOT be attached locally"
1978 );
1979 }
1980
1981 #[tokio::test]
1982 async fn spawn_passes_repo_provider_into_spawnspec() {
1983 use crate::registry::RepoProviderConfig;
1984 use crate::testkit::FakeCaliband;
1985
1986 let dir = tempfile::tempdir().unwrap();
1987 let mut config = FleetConfig::new("local", dir.path());
1988 config.discovery_env.caliban_daemon_runtime_dir = Some(dir.path().to_path_buf());
1989 config.ensure.autostart = false; let root = dir.path().join("repo");
1991 std::fs::create_dir_all(&root).unwrap();
1992 let socket = crate::discovery::resolve_socket(&root, &config.discovery_env).unwrap();
1993
1994 let fake = FakeCaliband::start_at(&socket).await.unwrap();
1995 let store = std::sync::Arc::new(crate::store::JsonlStore::open(dir.path()).unwrap());
1996 let mgr = FleetManager::new(config, store).await.unwrap();
1997 mgr.add_repo("p", &root).await.unwrap();
1998 mgr.set_repo_config_registry_only(
1999 "p",
2000 RepoProviderConfig {
2001 provider: Some("ollama".into()),
2002 ..Default::default()
2003 },
2004 )
2005 .await
2006 .unwrap();
2007
2008 mgr.spawn_agent("p", SpawnRequest::new("hi")).await.unwrap();
2009
2010 let specs = fake.received_specs();
2011 assert_eq!(specs.len(), 1, "exactly one spawn reached caliband");
2012 assert_eq!(
2013 specs[0].provider.as_deref(),
2014 Some("ollama"),
2015 "the repo's configured provider must be carried in SpawnSpec.provider (#93)"
2016 );
2017 }
2018
2019 #[tokio::test]
2020 async fn ensure_config_for_merges_default_and_repo_config() {
2021 use crate::registry::RepoProviderConfig;
2022 let dir = tempfile::tempdir().unwrap();
2023 let mut config = FleetConfig::new("local", dir.path());
2024 config.default_env.insert("KEEP".into(), "global".into());
2025 let store = std::sync::Arc::new(crate::store::JsonlStore::open(dir.path()).unwrap());
2026 let mgr = FleetManager::new(config, store).await.unwrap();
2027
2028 mgr.add_repo("p", "/tmp/p").await.ok(); let cfg = RepoProviderConfig {
2030 provider: Some("ollama".into()),
2031 base_url: Some("http://h:11434".into()),
2032 env: [("EXTRA".to_string(), "1".to_string())]
2033 .into_iter()
2034 .collect(),
2035 ..Default::default()
2036 };
2037 mgr.set_repo_config_registry_only("p", cfg).await.unwrap();
2038
2039 let ec = mgr.ensure_config_for("p").await.unwrap();
2040 assert_eq!(ec.env.get("KEEP").unwrap(), "global");
2041 assert_eq!(ec.env.get("CALIBAN_PROVIDER").unwrap(), "ollama");
2042 assert_eq!(ec.env.get("OLLAMA_BASE_URL").unwrap(), "http://h:11434");
2043 assert_eq!(ec.env.get("EXTRA").unwrap(), "1");
2044 }
2045
2046 struct FlakyStore {
2050 inner: crate::store::JsonlStore,
2051 fail_seqs: std::sync::Mutex<std::collections::HashSet<u64>>,
2052 }
2053
2054 impl FlakyStore {
2055 fn new(inner: crate::store::JsonlStore, fail: impl IntoIterator<Item = u64>) -> Self {
2056 Self {
2057 inner,
2058 fail_seqs: std::sync::Mutex::new(fail.into_iter().collect()),
2059 }
2060 }
2061 }
2062
2063 #[async_trait::async_trait]
2064 impl Store for FlakyStore {
2065 async fn append(&self, event: &FleetEvent) -> Result<()> {
2066 if self.fail_seqs.lock().unwrap().contains(&event.seq) {
2067 return Err(CoreError::Store("injected append failure".into()));
2068 }
2069 self.inner.append(event).await
2070 }
2071 async fn replay(&self, stream_key: &str, from_seq: u64) -> Result<Vec<FleetEvent>> {
2072 self.inner.replay(stream_key, from_seq).await
2073 }
2074 async fn high_water(&self, stream_key: &str) -> Result<u64> {
2075 self.inner.high_water(stream_key).await
2076 }
2077 async fn writable(&self) -> bool {
2078 self.inner.writable().await
2079 }
2080 async fn prune(&self, before_ts: &str) -> Result<u64> {
2081 self.inner.prune(before_ts).await
2082 }
2083 async fn usage(&self, since: &str, until: &str) -> Result<Vec<crate::store::UsageRow>> {
2084 self.inner.usage(since, until).await
2085 }
2086 }
2087
2088 fn emitter_with(store: Arc<dyn Store>) -> Emitter {
2089 Emitter {
2090 store,
2091 bus: Arc::new(InProcessBus::new(16)),
2092 seqs: Arc::new(AsyncMutex::new(HashMap::new())),
2093 metrics: Arc::new(Metrics::default()),
2094 }
2095 }
2096
2097 fn ev(seq: u64, agent: &str, chunk: &str) -> FleetEvent {
2098 use crate::event::OutputStream;
2099 FleetEvent {
2100 seq,
2101 ts: "t".into(),
2102 repo: "r".into(),
2103 agent_id: agent.into(),
2104 kind: EventKind::Output {
2105 stream: OutputStream::Stdout,
2106 chunk: chunk.into(),
2107 },
2108 }
2109 }
2110
2111 #[tokio::test]
2112 async fn seq_is_monotonic_per_stream_not_global() {
2113 let dir = tempfile::tempdir().unwrap();
2114 let store = Arc::new(crate::store::JsonlStore::open(dir.path()).unwrap());
2115 let emitter = emitter_with(store);
2116
2117 emitter.emit("r", "a1", EventKind::AgentSpawned).await;
2119 emitter.emit("r", "a2", EventKind::AgentSpawned).await;
2120 emitter.emit("r", "a1", EventKind::AgentGone).await;
2121
2122 let a1 = emitter
2123 .store
2124 .replay(&crate::event::stream_key_for("r", "a1"), 0)
2125 .await
2126 .unwrap();
2127 let a2 = emitter
2128 .store
2129 .replay(&crate::event::stream_key_for("r", "a2"), 0)
2130 .await
2131 .unwrap();
2132 assert_eq!(a1.iter().map(|e| e.seq).collect::<Vec<_>>(), vec![1, 2]);
2133 assert_eq!(a2.iter().map(|e| e.seq).collect::<Vec<_>>(), vec![1]);
2134 }
2135
2136 #[tokio::test]
2137 async fn seq_resumes_per_stream_from_high_water() {
2138 let dir = tempfile::tempdir().unwrap();
2139 {
2141 let store = crate::store::JsonlStore::open(dir.path()).unwrap();
2142 store.append(&ev(5, "a1", "old")).await.unwrap();
2143 }
2144 let store = Arc::new(crate::store::JsonlStore::open(dir.path()).unwrap());
2145 let emitter = emitter_with(store);
2146 emitter.emit("r", "a1", EventKind::AgentGone).await;
2147
2148 let a1 = emitter
2149 .store
2150 .replay(&crate::event::stream_key_for("r", "a1"), 0)
2151 .await
2152 .unwrap();
2153 assert_eq!(a1.last().unwrap().seq, 6);
2155 }
2156
2157 #[tokio::test]
2158 async fn emit_retries_on_seq_conflict_instead_of_dropping() {
2159 let dir = tempfile::tempdir().unwrap();
2164 let store = Arc::new(
2165 crate::sqlite_store::SqliteStore::open(dir.path())
2166 .await
2167 .unwrap(),
2168 );
2169 let emitter = emitter_with(store.clone());
2170
2171 emitter.emit("r", "a1", EventKind::AgentSpawned).await;
2173
2174 store.append(&ev(2, "a1", "from-peer")).await.unwrap();
2176
2177 emitter.emit("r", "a1", EventKind::AgentGone).await;
2180
2181 let hist = store
2182 .replay(&crate::event::stream_key_for("r", "a1"), 0)
2183 .await
2184 .unwrap();
2185 assert_eq!(
2186 hist.iter().map(|e| e.seq).collect::<Vec<_>>(),
2187 vec![1, 2, 3],
2188 "the racing event must be preserved at the next free seq, not dropped"
2189 );
2190 assert!(
2191 matches!(hist.last().unwrap().kind, EventKind::AgentGone),
2192 "the retried emitter event lands at seq 3: {:?}",
2193 hist.last().unwrap().kind
2194 );
2195 assert_eq!(
2196 emitter.metrics.snapshot(0).append_failures,
2197 0,
2198 "a resolved conflict must not count as a persist failure"
2199 );
2200 }
2201
2202 #[tokio::test]
2203 async fn append_failure_emits_persist_gap_marker_visible_to_history() {
2204 use crate::event::OutputStream;
2205
2206 let dir = tempfile::tempdir().unwrap();
2207 let inner = crate::store::JsonlStore::open(dir.path()).unwrap();
2208 let store = Arc::new(FlakyStore::new(inner, [1])); let emitter = emitter_with(store.clone());
2210 use tokio_stream::StreamExt;
2211 let mut sub = emitter
2212 .bus
2213 .subscribe(&crate::event::stream_key_for("repo", "a1"));
2214
2215 emitter
2216 .emit(
2217 "repo",
2218 "a1",
2219 EventKind::Output {
2220 stream: OutputStream::Stdout,
2221 chunk: "lost".into(),
2222 },
2223 )
2224 .await;
2225
2226 let ev = match sub.next().await {
2228 Some(crate::bus::BusEvent::Event(ev)) => ev,
2229 other => panic!("expected a live event, got {other:?}"),
2230 };
2231 assert_eq!(ev.seq, 1);
2232 assert!(matches!(ev.kind, EventKind::Output { .. }));
2233 let marker = match sub.next().await {
2235 Some(crate::bus::BusEvent::Event(ev)) => ev,
2236 other => panic!("expected a live event, got {other:?}"),
2237 };
2238 assert_eq!(marker.agent_id, "a1");
2239 assert!(matches!(
2240 marker.kind,
2241 EventKind::StorePersistFailed { lost_seq: 1, .. }
2242 ));
2243
2244 let history = store.replay("a1", 0).await.unwrap();
2247 assert!(
2248 history
2249 .iter()
2250 .any(|e| matches!(e.kind, EventKind::StorePersistFailed { lost_seq: 1, .. })),
2251 "history reader must see the persist-gap marker"
2252 );
2253 assert!(
2254 !history.iter().any(|e| e.seq == 1),
2255 "the un-persisted event must not appear in durable history"
2256 );
2257 }
2258
2259 #[tokio::test]
2260 async fn healthy_append_emits_no_gap_marker() {
2261 let dir = tempfile::tempdir().unwrap();
2262 let store = Arc::new(crate::store::JsonlStore::open(dir.path()).unwrap());
2263 let emitter = emitter_with(store);
2264 use tokio_stream::StreamExt;
2265 let mut sub = emitter
2266 .bus
2267 .subscribe(&crate::event::stream_key_for("repo", "a1"));
2268
2269 emitter.emit("repo", "a1", EventKind::AgentSpawned).await;
2270
2271 let ev = match sub.next().await {
2272 Some(crate::bus::BusEvent::Event(ev)) => ev,
2273 other => panic!("expected a live event, got {other:?}"),
2274 };
2275 assert!(matches!(ev.kind, EventKind::AgentSpawned));
2276 assert!(
2278 tokio::time::timeout(std::time::Duration::from_millis(100), sub.next())
2279 .await
2280 .is_err(),
2281 "a healthy append must not emit a gap marker"
2282 );
2283 }
2284
2285 #[tokio::test]
2286 async fn append_failure_and_success_advance_metrics() {
2287 let dir = tempfile::tempdir().unwrap();
2288 let inner = crate::store::JsonlStore::open(dir.path()).unwrap();
2289 let store = Arc::new(FlakyStore::new(inner, [1]));
2291 let emitter = emitter_with(store);
2292
2293 emitter.emit("repo", "a1", EventKind::AgentSpawned).await;
2294
2295 let m = emitter.metrics.snapshot(0);
2296 assert_eq!(m.append_failures, 1, "the failed append must be counted");
2297 assert_eq!(
2298 m.events_appended, 1,
2299 "the successful gap-marker append must be counted"
2300 );
2301 }
2302
2303 #[tokio::test]
2304 async fn run_drains_and_returns_on_shutdown() {
2305 let dir = tempfile::tempdir().unwrap();
2306 let mut config = FleetConfig::new("local", dir.path());
2307 config.poll_interval = Duration::from_millis(50);
2308 let store = Arc::new(crate::store::JsonlStore::open(dir.path()).unwrap());
2309 let mgr = FleetManager::new(config, store).await.unwrap();
2310
2311 let signaller = mgr.clone();
2312 let handle = tokio::spawn(mgr.run());
2313 signaller.begin_shutdown();
2314
2315 tokio::time::timeout(Duration::from_secs(2), handle)
2318 .await
2319 .expect("run() must return after begin_shutdown")
2320 .expect("run task panicked");
2321 }
2322
2323 #[tokio::test]
2324 async fn prune_older_than_removes_aged_events() {
2325 let dir = tempfile::tempdir().unwrap();
2326 let store = Arc::new(crate::store::JsonlStore::open(dir.path()).unwrap());
2327 store
2328 .append(&FleetEvent {
2329 seq: 1,
2330 ts: "2000-01-01T00:00:00+00:00".into(),
2331 repo: "r".into(),
2332 agent_id: "a".into(),
2333 kind: EventKind::AgentSpawned,
2334 })
2335 .await
2336 .unwrap();
2337 store
2338 .append(&FleetEvent {
2339 seq: 2,
2340 ts: chrono::Utc::now().to_rfc3339(),
2341 repo: "r".into(),
2342 agent_id: "a".into(),
2343 kind: EventKind::AgentGone,
2344 })
2345 .await
2346 .unwrap();
2347
2348 let config = FleetConfig::new("local", dir.path());
2349 let mgr = FleetManager::new(config, store).await.unwrap();
2350
2351 let removed = mgr
2352 .prune_older_than(std::time::Duration::from_secs(24 * 3600))
2353 .await
2354 .unwrap();
2355 assert_eq!(removed, 1);
2356 assert_eq!(mgr.history("a", 0).await.unwrap().len(), 1);
2357 }
2358
2359 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2360 async fn concurrent_seed_of_same_stream_assigns_unique_seqs() {
2361 let dir = tempfile::tempdir().unwrap();
2362 let store = Arc::new(crate::store::JsonlStore::open(dir.path()).unwrap());
2363 let emitter = emitter_with(store);
2364 let e1 = emitter.clone();
2366 let e2 = emitter.clone();
2367 let t1 = tokio::spawn(async move { e1.emit("r", "a1", EventKind::AgentSpawned).await });
2368 let t2 = tokio::spawn(async move { e2.emit("r", "a1", EventKind::AgentGone).await });
2369 t1.await.unwrap();
2370 t2.await.unwrap();
2371 let a1 = emitter
2372 .store
2373 .replay(&crate::event::stream_key_for("r", "a1"), 0)
2374 .await
2375 .unwrap();
2376 let seqs: Vec<u64> = a1.iter().map(|e| e.seq).collect();
2377 assert_eq!(seqs.len(), 2, "both events persisted");
2378 assert_ne!(seqs[0], seqs[1], "no duplicate seq within a stream");
2379 }
2380}