use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use tokio::sync::Semaphore;
use super::harness::{AgentRunSummary, ForemanError, WorktreeAgent, WorktreeAgentRequest};
pub struct FleetWorker {
pub id: String,
pub agent: Arc<dyn WorktreeAgent>,
pub capacity: usize,
pub remote: bool,
}
impl FleetWorker {
pub fn local(id: impl Into<String>, agent: Arc<dyn WorktreeAgent>, capacity: usize) -> Self {
Self {
id: id.into(),
agent,
capacity: capacity.max(1),
remote: false,
}
}
pub fn remote(id: impl Into<String>, agent: Arc<dyn WorktreeAgent>, capacity: usize) -> Self {
Self {
id: id.into(),
agent,
capacity: capacity.max(1),
remote: true,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct FailedAttempt {
#[serde(rename = "worker")]
pub worker_id: String,
pub error: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Placement {
pub subtask_id: String,
#[serde(rename = "worker")]
pub worker_id: Option<String>,
pub remote: bool,
#[serde(rename = "failed_attempts", default)]
pub attempts: Vec<FailedAttempt>,
}
pub struct FleetPool {
slots: Vec<Slot>,
ledger: Mutex<Vec<Placement>>,
}
struct Slot {
worker: FleetWorker,
permits: Arc<Semaphore>,
quarantined: AtomicBool,
}
impl FleetPool {
pub fn new(workers: Vec<FleetWorker>) -> Self {
let slots = workers
.into_iter()
.map(|worker| {
let permits = Arc::new(Semaphore::new(worker.capacity.max(1)));
Slot {
worker,
permits,
quarantined: AtomicBool::new(false),
}
})
.collect();
Self {
slots,
ledger: Mutex::new(Vec::new()),
}
}
pub fn capacity(&self) -> usize {
self.slots.iter().map(|s| s.worker.capacity).sum()
}
pub fn worker_ids(&self) -> Vec<&str> {
self.slots.iter().map(|s| s.worker.id.as_str()).collect()
}
pub fn placements(&self) -> Vec<Placement> {
self.ledger
.lock()
.unwrap_or_else(|e| e.into_inner())
.clone()
}
fn record(&self, placement: Placement) {
self.ledger
.lock()
.unwrap_or_else(|e| e.into_inner())
.push(placement);
}
pub fn quarantined(&self) -> Vec<&str> {
self.slots
.iter()
.filter(|s| s.quarantined.load(Ordering::Relaxed))
.map(|s| s.worker.id.as_str())
.collect()
}
fn candidates(&self) -> Vec<usize> {
let mut order: Vec<usize> = (0..self.slots.len())
.filter(|&i| !self.slots[i].quarantined.load(Ordering::Relaxed))
.collect();
order.sort_by_key(|&i| {
(
std::cmp::Reverse(self.slots[i].permits.available_permits()),
i,
)
});
order
}
}
#[async_trait]
impl WorktreeAgent for FleetPool {
async fn run_in(
&self,
req: &WorktreeAgentRequest<'_>,
) -> Result<AgentRunSummary, ForemanError> {
if self.slots.is_empty() {
return Err(ForemanError::Agent(
"no workers in the fleet pool: nothing can run this subtask".into(),
));
}
let subtask_id = req.subtask.id.clone();
let mut attempts: Vec<FailedAttempt> = Vec::new();
let mut attempted = false;
for idx in self.candidates() {
let slot = &self.slots[idx];
if slot.quarantined.load(Ordering::Relaxed) {
continue;
}
if attempted {
if let Err(e) = reset_worktree(req.cwd) {
self.record(Placement {
subtask_id,
worker_id: None,
remote: false,
attempts,
});
return Err(ForemanError::Git(format!(
"cannot fail over to `{}`: {e}",
slot.worker.id
)));
}
}
let _permit = match slot.permits.clone().acquire_owned().await {
Ok(p) => p,
Err(_) => {
attempts.push(FailedAttempt {
worker_id: slot.worker.id.clone(),
error: "worker capacity closed".into(),
});
continue;
}
};
if slot.quarantined.load(Ordering::Relaxed) {
drop(_permit);
continue;
}
attempted = true;
match slot.worker.agent.run_in(req).await {
Ok(summary) => {
self.record(Placement {
subtask_id,
worker_id: Some(slot.worker.id.clone()),
remote: slot.worker.remote,
attempts,
});
return Ok(summary);
}
Err(e) => {
if slot.worker.remote
&& matches!(e, ForemanError::Worker(_))
&& !slot.quarantined.swap(true, Ordering::Relaxed)
{
tracing::warn!(
worker = %slot.worker.id,
error = %e,
"quarantining fleet worker for the rest of the run"
);
}
tracing::warn!(
subtask = %subtask_id,
worker = %slot.worker.id,
error = %e,
"fleet worker failed; trying the next"
);
attempts.push(FailedAttempt {
worker_id: slot.worker.id.clone(),
error: e.to_string(),
});
}
}
}
if !attempted {
return Err(ForemanError::Agent(format!(
"no worker could be offered this subtask; quarantined for this run: {}",
self.quarantined().join(", ")
)));
}
let detail = attempts
.iter()
.map(|a| format!("{}: {}", a.worker_id, a.error))
.collect::<Vec<_>>()
.join("; ");
self.record(Placement {
subtask_id,
worker_id: None,
remote: false,
attempts,
});
Err(ForemanError::Agent(format!(
"every fleet worker failed — {detail}"
)))
}
}
fn reset_worktree(cwd: &std::path::Path) -> Result<(), String> {
for args in [vec!["reset", "--hard", "--quiet"], vec!["clean", "-fdq"]] {
let out = std::process::Command::new("git")
.arg("-C")
.arg(cwd)
.args(&args)
.output()
.map_err(|e| format!("git {}: {e}", args.join(" ")))?;
if !out.status.success() {
return Err(format!(
"git {} failed: {}",
args.join(" "),
String::from_utf8_lossy(&out.stderr).trim()
));
}
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::patterns::foreman::harness::Subtask;
use std::path::{Path, PathBuf};
#[test]
fn placement_wire_shape_is_pinned() {
let json = serde_json::to_value(vec![Placement {
subtask_id: "s1".into(),
worker_id: Some("studio".into()),
remote: true,
attempts: vec![FailedAttempt {
worker_id: "laptop".into(),
error: "boom".into(),
}],
}])
.unwrap();
assert_eq!(
json,
serde_json::json!([{
"subtask_id": "s1",
"worker": "studio",
"remote": true,
"failed_attempts": [{ "worker": "laptop", "error": "boom" }],
}])
);
let round: Vec<Placement> = serde_json::from_value(json).unwrap();
assert_eq!(round[0].worker_id.as_deref(), Some("studio"));
assert_eq!(round[0].attempts[0].worker_id, "laptop");
}
#[test]
fn optional_keys_are_optional_and_required_ones_are_required() {
let p: Placement = serde_json::from_value(serde_json::json!({
"subtask_id": "s1",
"remote": false,
}))
.expect("worker and failed_attempts may both be absent");
assert_eq!(p.worker_id, None);
assert!(p.attempts.is_empty());
for missing in [
serde_json::json!({ "remote": false }),
serde_json::json!({ "subtask_id": "s1" }),
] {
assert!(
serde_json::from_value::<Placement>(missing.clone()).is_err(),
"{missing} must not deserialize"
);
}
}
use std::process::Command;
use std::sync::atomic::{AtomicUsize, Ordering};
fn git(cwd: &Path, args: &[&str]) {
let out = Command::new("git")
.args(args)
.current_dir(cwd)
.output()
.expect("git runs");
assert!(
out.status.success(),
"git {args:?}: {}",
String::from_utf8_lossy(&out.stderr)
);
}
fn worktree() -> tempfile::TempDir {
let dir = tempfile::tempdir().unwrap();
let root = dir.path();
git(root, &["init", "-q", "-b", "main"]);
git(root, &["config", "user.email", "t@t.t"]);
git(root, &["config", "user.name", "t"]);
git(root, &["config", "core.autocrlf", "false"]);
std::fs::write(root.join("seed.txt"), "seed\n").unwrap();
git(root, &["add", "-A"]);
git(root, &["commit", "-qm", "base"]);
dir
}
struct Counting {
id: &'static str,
inflight: Arc<AtomicUsize>,
peak: Arc<AtomicUsize>,
ran: Arc<Mutex<Vec<String>>>,
}
#[async_trait]
impl WorktreeAgent for Counting {
async fn run_in(
&self,
req: &WorktreeAgentRequest<'_>,
) -> Result<AgentRunSummary, ForemanError> {
let now = self.inflight.fetch_add(1, Ordering::SeqCst) + 1;
self.peak.fetch_max(now, Ordering::SeqCst);
tokio::time::sleep(std::time::Duration::from_millis(20)).await;
self.ran
.lock()
.unwrap()
.push(format!("{}:{}", self.id, req.subtask.id));
self.inflight.fetch_sub(1, Ordering::SeqCst);
Ok(AgentRunSummary {
answer: self.id.to_string(),
})
}
}
struct MessyFailure;
#[async_trait]
impl WorktreeAgent for MessyFailure {
async fn run_in(
&self,
req: &WorktreeAgentRequest<'_>,
) -> Result<AgentRunSummary, ForemanError> {
std::fs::write(req.cwd.join("half-done.txt"), "partial\n").unwrap();
std::fs::write(req.cwd.join("seed.txt"), "clobbered\n").unwrap();
Err(ForemanError::Agent("peer dropped mid-subtask".into()))
}
}
struct ExpectsCleanTree;
#[async_trait]
impl WorktreeAgent for ExpectsCleanTree {
async fn run_in(
&self,
req: &WorktreeAgentRequest<'_>,
) -> Result<AgentRunSummary, ForemanError> {
if req.cwd.join("half-done.txt").exists() {
return Err(ForemanError::Agent(
"inherited a failed worker's untracked file".into(),
));
}
let seed = std::fs::read_to_string(req.cwd.join("seed.txt")).unwrap();
if seed != "seed\n" {
return Err(ForemanError::Agent(
"inherited a failed worker's tracked edit".into(),
));
}
Ok(AgentRunSummary {
answer: "clean".into(),
})
}
}
#[tokio::test]
async fn a_failover_that_cannot_reset_the_tree_still_records_what_ran() {
let pool = FleetPool::new(vec![
FleetWorker::local("first", Arc::new(MessyFailure), 1),
FleetWorker::local("second", Arc::new(ExpectsCleanTree), 1),
]);
let dir = tempfile::tempdir().unwrap();
let cwd = dir.path().to_path_buf();
let a = Subtask::files_only("a", "a", vec![]);
let err = pool.run_in(&request(&a, &cwd)).await.unwrap_err();
assert!(
matches!(err, ForemanError::Git(_)),
"expected the reset failure to surface: {err:?}"
);
let placements = pool.placements();
assert_eq!(placements.len(), 1, "the subtask must appear in the ledger");
assert_eq!(placements[0].subtask_id, "a");
assert_eq!(placements[0].worker_id, None);
assert_eq!(placements[0].attempts.len(), 1);
assert_eq!(placements[0].attempts[0].worker_id, "first");
assert!(
!placements[0]
.attempts
.iter()
.any(|a| a.worker_id == "second"),
"a worker the failover never reached is not a failed attempt"
);
}
struct Failing {
calls: Arc<AtomicUsize>,
worker_level: bool,
}
#[async_trait]
impl WorktreeAgent for Failing {
async fn run_in(
&self,
_req: &WorktreeAgentRequest<'_>,
) -> Result<AgentRunSummary, ForemanError> {
self.calls.fetch_add(1, Ordering::SeqCst);
Err(if self.worker_level {
ForemanError::Worker("peer is not there".into())
} else {
ForemanError::Agent("the subtask failed".into())
})
}
}
fn always_ok(id: &'static str) -> Arc<Counting> {
Arc::new(Counting {
id,
inflight: Arc::new(AtomicUsize::new(0)),
peak: Arc::new(AtomicUsize::new(0)),
ran: Arc::new(Mutex::new(Vec::new())),
})
}
#[tokio::test]
async fn a_worker_level_failure_takes_the_peer_out_for_the_rest_of_the_run() {
let calls = Arc::new(AtomicUsize::new(0));
let pool = FleetPool::new(vec![
FleetWorker::remote(
"dead-peer",
Arc::new(Failing {
calls: Arc::clone(&calls),
worker_level: true,
}),
1,
),
FleetWorker::local("here", always_ok("here"), 1),
]);
let dir = worktree();
let cwd = dir.path().to_path_buf();
for id in ["a", "b", "c"] {
let sub = Subtask::files_only(id, id, vec![]);
pool.run_in(&request(&sub, &cwd))
.await
.expect("the local worker picks it up");
}
assert_eq!(
calls.load(Ordering::SeqCst),
1,
"the dead peer must be offered exactly one subtask, not one per subtask"
);
assert_eq!(pool.quarantined(), vec!["dead-peer"]);
let with_attempts = pool
.placements()
.iter()
.filter(|p| !p.attempts.is_empty())
.count();
assert_eq!(with_attempts, 1);
}
struct SlowFailing {
calls: Arc<AtomicUsize>,
}
#[async_trait]
impl WorktreeAgent for SlowFailing {
async fn run_in(
&self,
_req: &WorktreeAgentRequest<'_>,
) -> Result<AgentRunSummary, ForemanError> {
self.calls.fetch_add(1, Ordering::SeqCst);
tokio::time::sleep(std::time::Duration::from_millis(60)).await;
Err(ForemanError::Worker("peer is not there".into()))
}
}
#[tokio::test]
async fn a_whole_level_queued_behind_a_dying_peer_stops_dispatching_to_it() {
const CAPACITY: usize = 2;
const LEVEL: usize = 8;
let calls = Arc::new(AtomicUsize::new(0));
let pool = FleetPool::new(vec![
FleetWorker::remote(
"dying-peer",
Arc::new(SlowFailing {
calls: Arc::clone(&calls),
}),
CAPACITY,
),
FleetWorker::local("here", always_ok("here"), 1),
]);
let dirs: Vec<tempfile::TempDir> = (0..LEVEL).map(|_| worktree()).collect();
let cwds: Vec<PathBuf> = dirs.iter().map(|d| d.path().to_path_buf()).collect();
let subtasks: Vec<Subtask> = (0..LEVEL)
.map(|i| Subtask::files_only(format!("s{i}"), "go", vec![]))
.collect();
let reqs: Vec<WorktreeAgentRequest<'_>> = subtasks
.iter()
.zip(cwds.iter())
.map(|(sub, cwd)| request(sub, cwd))
.collect();
let results = futures::future::join_all(reqs.iter().map(|r| pool.run_in(r))).await;
for r in &results {
r.as_ref()
.expect("the local worker picks every one of them up");
}
assert_eq!(pool.quarantined(), vec!["dying-peer"]);
assert!(
calls.load(Ordering::SeqCst) <= CAPACITY,
"a level of {LEVEL} must not dispatch past the peer's capacity of \
{CAPACITY} once it is quarantined; got {}",
calls.load(Ordering::SeqCst)
);
}
#[tokio::test]
async fn a_subtask_that_failed_does_not_quarantine_the_worker_that_ran_it() {
let calls = Arc::new(AtomicUsize::new(0));
let pool = FleetPool::new(vec![
FleetWorker::remote(
"healthy-peer",
Arc::new(Failing {
calls: Arc::clone(&calls),
worker_level: false,
}),
1,
),
FleetWorker::local("here", always_ok("here"), 1),
]);
let dir = worktree();
let cwd = dir.path().to_path_buf();
for id in ["a", "b", "c"] {
let sub = Subtask::files_only(id, id, vec![]);
pool.run_in(&request(&sub, &cwd)).await.expect("local runs");
}
assert_eq!(
calls.load(Ordering::SeqCst),
3,
"an `Agent` failure must leave the worker in the pool"
);
assert!(pool.quarantined().is_empty());
}
#[tokio::test]
async fn the_local_worker_is_never_quarantined() {
let calls = Arc::new(AtomicUsize::new(0));
let pool = FleetPool::new(vec![FleetWorker::local(
"here",
Arc::new(Failing {
calls: Arc::clone(&calls),
worker_level: true,
}),
1,
)]);
let dir = worktree();
let cwd = dir.path().to_path_buf();
for id in ["a", "b"] {
let sub = Subtask::files_only(id, id, vec![]);
pool.run_in(&request(&sub, &cwd))
.await
.expect_err("it fails, but it stays in the pool");
}
assert_eq!(calls.load(Ordering::SeqCst), 2);
assert!(pool.quarantined().is_empty());
}
#[tokio::test]
async fn an_all_remote_pool_that_quarantines_everyone_says_so() {
let calls = Arc::new(AtomicUsize::new(0));
let pool = FleetPool::new(vec![FleetWorker::remote(
"only-peer",
Arc::new(Failing {
calls: Arc::clone(&calls),
worker_level: true,
}),
1,
)]);
let dir = worktree();
let cwd = dir.path().to_path_buf();
let a = Subtask::files_only("a", "a", vec![]);
pool.run_in(&request(&a, &cwd)).await.unwrap_err();
let b = Subtask::files_only("b", "b", vec![]);
let err = pool.run_in(&request(&b, &cwd)).await.unwrap_err();
assert_eq!(calls.load(Ordering::SeqCst), 1);
let text = err.to_string();
assert!(
text.contains("quarantined") && text.contains("only-peer"),
"the error must name the quarantine and the peer: {text}"
);
}
fn request<'a>(subtask: &'a Subtask, cwd: &'a PathBuf) -> WorktreeAgentRequest<'a> {
WorktreeAgentRequest {
subtask,
cwd,
allowed_tools: None,
mcp_endpoint: None,
}
}
fn counting(
id: &'static str,
ran: &Arc<Mutex<Vec<String>>>,
) -> (Arc<Counting>, Arc<AtomicUsize>) {
let peak = Arc::new(AtomicUsize::new(0));
(
Arc::new(Counting {
id,
inflight: Arc::new(AtomicUsize::new(0)),
peak: Arc::clone(&peak),
ran: Arc::clone(ran),
}),
peak,
)
}
#[tokio::test]
async fn work_spreads_across_workers_instead_of_queueing_on_one() {
let ran = Arc::new(Mutex::new(Vec::new()));
let (here, _) = counting("here", &ran);
let (studio, _) = counting("studio", &ran);
let pool = FleetPool::new(vec![
FleetWorker::local("here", here, 1),
FleetWorker::remote("studio", studio, 1),
]);
assert_eq!(pool.capacity(), 2);
let dir = worktree();
let cwd = dir.path().to_path_buf();
let a = Subtask::files_only("a", "a", vec![]);
let b = Subtask::files_only("b", "b", vec![]);
let (req_a, req_b) = (request(&a, &cwd), request(&b, &cwd));
let (ra, rb) = tokio::join!(pool.run_in(&req_a), pool.run_in(&req_b));
ra.unwrap();
rb.unwrap();
let placements = pool.placements();
assert_eq!(placements.len(), 2);
let used: std::collections::HashSet<String> = placements
.iter()
.filter_map(|p| p.worker_id.clone())
.collect();
assert_eq!(used.len(), 2, "both workers took a subtask: {placements:?}");
}
#[tokio::test]
async fn a_worker_never_exceeds_its_capacity() {
let ran = Arc::new(Mutex::new(Vec::new()));
let (only, peak) = counting("only", &ran);
let pool = FleetPool::new(vec![FleetWorker::local("only", only, 1)]);
let dir = worktree();
let cwd = dir.path().to_path_buf();
let a = Subtask::files_only("a", "a", vec![]);
let b = Subtask::files_only("b", "b", vec![]);
let c = Subtask::files_only("c", "c", vec![]);
let (req_a, req_b, req_c) = (request(&a, &cwd), request(&b, &cwd), request(&c, &cwd));
let _ = tokio::join!(
pool.run_in(&req_a),
pool.run_in(&req_b),
pool.run_in(&req_c)
);
assert_eq!(
peak.load(Ordering::SeqCst),
1,
"capacity 1 means one at a time, however many subtasks arrive"
);
assert_eq!(pool.placements().len(), 3);
}
#[tokio::test]
async fn failover_hands_the_next_worker_a_clean_tree() {
let pool = FleetPool::new(vec![
FleetWorker::remote("flaky", Arc::new(MessyFailure), 1),
FleetWorker::local("here", Arc::new(ExpectsCleanTree), 1),
]);
let dir = worktree();
let cwd = dir.path().to_path_buf();
let s = Subtask::files_only("s", "s", vec![]);
let summary = pool
.run_in(&request(&s, &cwd))
.await
.expect("second worker takes it");
assert_eq!(summary.answer, "clean");
let placements = pool.placements();
assert_eq!(placements[0].worker_id.as_deref(), Some("here"));
assert_eq!(placements[0].attempts.len(), 1);
assert_eq!(placements[0].attempts[0].worker_id, "flaky");
assert!(!cwd.join("half-done.txt").exists(), "reset cleaned up");
}
#[tokio::test]
async fn every_worker_failing_reports_all_of_them() {
let pool = FleetPool::new(vec![
FleetWorker::remote("a", Arc::new(MessyFailure), 1),
FleetWorker::remote("b", Arc::new(MessyFailure), 1),
]);
let dir = worktree();
let cwd = dir.path().to_path_buf();
let s = Subtask::files_only("s", "s", vec![]);
let err = pool.run_in(&request(&s, &cwd)).await.unwrap_err();
let msg = err.to_string();
assert!(msg.contains('a') && msg.contains('b'), "{msg}");
let placements = pool.placements();
assert!(placements[0].worker_id.is_none());
assert_eq!(placements[0].attempts.len(), 2);
}
#[tokio::test]
async fn an_empty_pool_says_so_instead_of_silently_doing_nothing() {
let pool = FleetPool::new(Vec::new());
let dir = worktree();
let cwd = dir.path().to_path_buf();
let s = Subtask::files_only("s", "s", vec![]);
let err = pool.run_in(&request(&s, &cwd)).await.unwrap_err();
assert!(err.to_string().contains("no workers"), "{err}");
}
struct PatchReturningWorker {
base: PathBuf,
}
#[async_trait]
impl WorktreeAgent for PatchReturningWorker {
async fn run_in(
&self,
req: &WorktreeAgentRequest<'_>,
) -> Result<AgentRunSummary, ForemanError> {
let remote = tempfile::tempdir().unwrap();
let clone = remote.path().join("checkout");
let out = Command::new("git")
.args(["clone", "-q"])
.arg(&self.base)
.arg(&clone)
.output()
.unwrap();
assert!(out.status.success(), "{out:?}");
std::fs::write(clone.join("new.rs"), "pub fn added() {}\n").unwrap();
std::fs::write(clone.join("seed.txt"), "edited\n").unwrap();
let patch = super::super::harness::capture_patch(&clone)?;
super::super::harness::git_apply(req.cwd, &patch)?;
Ok(AgentRunSummary {
answer: "remote".into(),
})
}
}
#[tokio::test]
async fn a_patch_made_on_another_checkout_lands_in_the_gated_worktree() {
let dir = worktree();
let cwd = dir.path().to_path_buf();
let pool = FleetPool::new(vec![FleetWorker::remote(
"studio",
Arc::new(PatchReturningWorker { base: cwd.clone() }),
1,
)]);
let s = Subtask::files_only("s", "s", vec![]);
pool.run_in(&request(&s, &cwd)).await.expect("applied");
assert_eq!(
std::fs::read_to_string(cwd.join("seed.txt")).unwrap(),
"edited\n",
"the peer's edit to a tracked file arrived"
);
assert!(cwd.join("new.rs").exists(), "and so did its new file");
let placement = &pool.placements()[0];
assert_eq!(placement.worker_id.as_deref(), Some("studio"));
assert!(placement.remote);
}
#[test]
fn a_zero_capacity_worker_is_clamped_rather_than_deadlocking() {
let pool = FleetPool::new(vec![FleetWorker::local(
"here",
Arc::new(ExpectsCleanTree),
0,
)]);
assert_eq!(pool.capacity(), 1);
assert_eq!(pool.worker_ids(), vec!["here"]);
}
}