Converted all external tool calling to async

This commit is contained in:
Timmy
2026-06-29 16:59:54 +01:00
parent 75f41088b1
commit feb35ddd10
49 changed files with 482 additions and 496 deletions
+2 -2
View File
@@ -159,7 +159,7 @@ pub(super) async fn detect_conflicts(
our_claims.remove(&story_id); our_claims.remove(&story_id);
// Stop any local agent for this story by looking up its name. // Stop any local agent for this story by looking up its name.
if let Ok(agent_list) = agents.list_agents() { if let Ok(agent_list) = agents.list_agents().await {
for info in agent_list { for info in agent_list {
if info.story_id == story_id { if info.story_id == story_id {
let _ = agents let _ = agents
@@ -219,7 +219,7 @@ pub(super) fn reclaim_timed_out_work(_project_root: &Path) {
/// Check for completed agents, push their feature branches to the remote, /// Check for completed agents, push their feature branches to the remote,
/// and report completion via CRDT. /// and report completion via CRDT.
pub(super) async fn check_completions_and_push(agents: &AgentPool, _project_root: &Path) { pub(super) async fn check_completions_and_push(agents: &AgentPool, _project_root: &Path) {
let Ok(agent_list) = agents.list_agents() else { let Ok(agent_list) = agents.list_agents().await else {
return; return;
}; };
@@ -73,7 +73,7 @@ mod tests {
// task eventually fails. // task eventually fails.
pool.auto_assign_available_work(tmp.path()).await; pool.auto_assign_available_work(tmp.path()).await;
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
let has_pending = agents.values().any(|a| { let has_pending = agents.values().any(|a| {
a.agent_name == "coder-1" a.agent_name == "coder-1"
&& matches!(a.status, AgentStatus::Pending | AgentStatus::Running) && matches!(a.status, AgentStatus::Pending | AgentStatus::Running)
@@ -115,7 +115,7 @@ mod tests {
pool.auto_assign_available_work(root).await; pool.auto_assign_available_work(root).await;
// No agent should have been started for the spike. // No agent should have been started for the spike.
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
assert!( assert!(
agents.is_empty(), agents.is_empty(),
"No agents should be assigned to a spike with review_hold" "No agents should be assigned to a spike with review_hold"
@@ -155,7 +155,7 @@ mod tests {
pool.auto_assign_available_work(tmp.path()).await; pool.auto_assign_available_work(tmp.path()).await;
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
// coder-1 must NOT have been assigned to the QA story (wrong stage). // coder-1 must NOT have been assigned to the QA story (wrong stage).
let coder_assigned_to_qa = agents.iter().any(|(key, a)| { let coder_assigned_to_qa = agents.iter().any(|(key, a)| {
key.contains("9930_story_qa1") key.contains("9930_story_qa1")
@@ -209,7 +209,7 @@ mod tests {
pool.auto_assign_available_work(tmp.path()).await; pool.auto_assign_available_work(tmp.path()).await;
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
// coder-1 should have been picked (it matches the stage and is preferred). // coder-1 should have been picked (it matches the stage and is preferred).
let coder1_assigned = agents.values().any(|a| { let coder1_assigned = agents.values().any(|a| {
a.agent_name == "coder-1" a.agent_name == "coder-1"
@@ -262,7 +262,7 @@ mod tests {
// Must not panic. // Must not panic.
pool.auto_assign_available_work(tmp.path()).await; pool.auto_assign_available_work(tmp.path()).await;
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
// No agent should be assigned to the specific QA story (coder-1 may // No agent should be assigned to the specific QA story (coder-1 may
// be assigned to leaked 2_current items from the global CRDT store). // be assigned to leaked 2_current items from the global CRDT store).
let assigned_to_qa_story = agents.iter().any(|(key, a)| { let assigned_to_qa_story = agents.iter().any(|(key, a)| {
@@ -301,7 +301,7 @@ mod tests {
let pool = AgentPool::new_test(3001); let pool = AgentPool::new_test(3001);
pool.auto_assign_available_work(root).await; pool.auto_assign_available_work(root).await;
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
// Filter to only agents assigned to our specific story to avoid // Filter to only agents assigned to our specific story to avoid
// interference from other tests sharing the global CRDT store. // interference from other tests sharing the global CRDT store.
let assigned_to_our_story = agents.iter().any(|(key, a)| { let assigned_to_our_story = agents.iter().any(|(key, a)| {
@@ -347,7 +347,7 @@ mod tests {
let pool = AgentPool::new_test(3001); let pool = AgentPool::new_test(3001);
pool.auto_assign_available_work(root).await; pool.auto_assign_available_work(root).await;
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
let has_pending = agents.values().any(|a| { let has_pending = agents.values().any(|a| {
matches!( matches!(
a.status, a.status,
@@ -553,7 +553,7 @@ mod tests {
let _ = tokio::join!(t1, t2); let _ = tokio::join!(t1, t2);
// At most one Pending/Running entry should exist for coder-1. // At most one Pending/Running entry should exist for coder-1.
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
let active_coder_count = agents let active_coder_count = agents
.values() .values()
.filter(|a| { .filter(|a| {
@@ -602,7 +602,7 @@ mod tests {
pool.auto_assign_available_work(tmp.path()).await; pool.auto_assign_available_work(tmp.path()).await;
let count_after_first = { let count_after_first = {
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
agents agents
.iter() .iter()
.filter(|(key, a)| { .filter(|(key, a)| {
@@ -616,7 +616,7 @@ mod tests {
pool.auto_assign_available_work(tmp.path()).await; pool.auto_assign_available_work(tmp.path()).await;
let count_after_second = { let count_after_second = {
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
agents agents
.iter() .iter()
.filter(|(key, a)| { .filter(|(key, a)| {
+1 -7
View File
@@ -99,13 +99,7 @@ impl AgentPool {
// Skip if an explicit mergemaster LLM agent is already running // Skip if an explicit mergemaster LLM agent is already running
// (operator-driven failure recovery path). // (operator-driven failure recovery path).
let has_mergemaster = { let has_mergemaster = {
let agents = match self.agents.lock() { let agents = self.agents.lock().await;
Ok(a) => a,
Err(e) => {
slog_error!("[auto-assign] Failed to lock agents: {e}");
break;
}
};
is_story_assigned_for_stage(config, &agents, story_id, &PipelineStage::Mergemaster) is_story_assigned_for_stage(config, &agents, story_id, &PipelineStage::Mergemaster)
}; };
if has_mergemaster { if has_mergemaster {
@@ -54,7 +54,7 @@ pub(crate) fn spawn_merge_failure_block_subscriber(pool: Arc<AgentPool>, project
match rx.recv().await { match rx.recv().await {
Ok(fired) => { Ok(fired) => {
let recovery_running = let recovery_running =
is_mergemaster_running(&pool, &project_root, &fired.story_id.0); is_mergemaster_running(&pool, &project_root, &fired.story_id.0).await;
on_transition(&project_root, &fired, &mut counters, recovery_running); on_transition(&project_root, &fired, &mut counters, recovery_running);
} }
Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => { Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => {
@@ -72,15 +72,12 @@ pub(crate) fn spawn_merge_failure_block_subscriber(pool: Arc<AgentPool>, project
/// Return true if a mergemaster agent is currently in the pool for `story_id`. /// Return true if a mergemaster agent is currently in the pool for `story_id`.
/// Used to suppress counter increments while recovery is actively iterating /// Used to suppress counter increments while recovery is actively iterating
/// (bug 1025). /// (bug 1025).
fn is_mergemaster_running(pool: &AgentPool, project_root: &Path, story_id: &str) -> bool { async fn is_mergemaster_running(pool: &AgentPool, project_root: &Path, story_id: &str) -> bool {
let config = match crate::config::ProjectConfig::load(project_root) { let config = match crate::config::ProjectConfig::load(project_root) {
Ok(c) => c, Ok(c) => c,
Err(_) => return false, Err(_) => return false,
}; };
let agents = match pool.agents.lock() { let agents = pool.agents.lock().await;
Ok(a) => a,
Err(_) => return false,
};
is_story_assigned_for_stage(&config, &agents, story_id, &PipelineStage::Mergemaster) is_story_assigned_for_stage(&config, &agents, story_id, &PipelineStage::Mergemaster)
} }
@@ -100,15 +100,7 @@ async fn on_merge_failure_transition(
}; };
let agent_name = { let agent_name = {
let agents = match pool.agents.lock() { let agents = pool.agents.lock().await;
Ok(a) => a,
Err(e) => {
slog_warn!(
"[merge-failure-sub] Failed to lock agent pool for '{story_id}': {e}"
);
return;
}
};
if is_story_assigned_for_stage( if is_story_assigned_for_stage(
&config, &config,
&agents, &agents,
@@ -228,7 +220,7 @@ mod tests {
); );
on_merge_failure_transition(&pool, tmp.path(), &fired).await; on_merge_failure_transition(&pool, tmp.path(), &fired).await;
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.lock().await;
assert!( assert!(
agents.iter().any(|(key, a)| { agents.iter().any(|(key, a)| {
key.contains(story_id) key.contains(story_id)
@@ -259,7 +251,7 @@ mod tests {
// Give the subscriber time to run (it should do nothing). // Give the subscriber time to run (it should do nothing).
tokio::time::sleep(std::time::Duration::from_millis(100)).await; tokio::time::sleep(std::time::Duration::from_millis(100)).await;
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.lock().await;
let spawned = agents.iter().any(|(key, a)| { let spawned = agents.iter().any(|(key, a)| {
key.contains(story_id) key.contains(story_id)
&& a.agent_name == "mergemaster" && a.agent_name == "mergemaster"
@@ -287,7 +279,7 @@ mod tests {
tokio::time::sleep(std::time::Duration::from_millis(100)).await; tokio::time::sleep(std::time::Duration::from_millis(100)).await;
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.lock().await;
let spawned = agents.iter().any(|(key, a)| { let spawned = agents.iter().any(|(key, a)| {
key.contains(story_id) key.contains(story_id)
&& a.agent_name == "mergemaster" && a.agent_name == "mergemaster"
@@ -315,7 +307,7 @@ mod tests {
tokio::time::sleep(std::time::Duration::from_millis(100)).await; tokio::time::sleep(std::time::Duration::from_millis(100)).await;
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.lock().await;
let spawned = agents.iter().any(|(key, a)| { let spawned = agents.iter().any(|(key, a)| {
key.contains(story_id) key.contains(story_id)
&& a.agent_name == "mergemaster" && a.agent_name == "mergemaster"
@@ -343,7 +335,7 @@ mod tests {
tokio::time::sleep(std::time::Duration::from_millis(100)).await; tokio::time::sleep(std::time::Duration::from_millis(100)).await;
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.lock().await;
let spawned = agents.iter().any(|(key, a)| { let spawned = agents.iter().any(|(key, a)| {
key.contains(story_id) key.contains(story_id)
&& a.agent_name == "mergemaster" && a.agent_name == "mergemaster"
@@ -374,7 +366,7 @@ mod tests {
// First call — spawns mergemaster (agent enters Pending). // First call — spawns mergemaster (agent enters Pending).
on_merge_failure_transition(&pool, tmp.path(), &fired).await; on_merge_failure_transition(&pool, tmp.path(), &fired).await;
{ {
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.lock().await;
assert!( assert!(
agents.iter().any(|(key, a)| { agents.iter().any(|(key, a)| {
key.contains(story_id) key.contains(story_id)
@@ -388,7 +380,7 @@ mod tests {
// Second call (self-loop) — agent is still Pending; guard must prevent double-spawn. // Second call (self-loop) — agent is still Pending; guard must prevent double-spawn.
on_merge_failure_transition(&pool, tmp.path(), &fired).await; on_merge_failure_transition(&pool, tmp.path(), &fired).await;
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.lock().await;
let active_count = agents let active_count = agents
.iter() .iter()
.filter(|(key, a)| { .filter(|(key, a)| {
+2 -15
View File
@@ -5,7 +5,6 @@ use std::path::Path;
use crate::config::ProjectConfig; use crate::config::ProjectConfig;
use crate::pipeline_state::Stage; use crate::pipeline_state::Stage;
use crate::slog; use crate::slog;
use crate::slog_error;
use super::super::super::PipelineStage; use super::super::super::PipelineStage;
use super::super::AgentPool; use super::super::AgentPool;
@@ -80,13 +79,7 @@ impl AgentPool {
if *stage == PipelineStage::Coder if *stage == PipelineStage::Coder
&& let Some(max) = config.max_coders && let Some(max) = config.max_coders
{ {
let agents_lock = match self.agents.lock() { let agents_lock = self.agents.lock().await;
Ok(a) => a,
Err(e) => {
slog_error!("[auto-assign] Failed to lock agents: {e}");
break;
}
};
let active = count_active_agents_for_stage(config, &agents_lock, stage); let active = count_active_agents_for_stage(config, &agents_lock, stage);
if active >= max { if active >= max {
slog!( slog!(
@@ -102,13 +95,7 @@ impl AgentPool {
// stage_mismatch=true means the preferred agent's stage doesn't match the // stage_mismatch=true means the preferred agent's stage doesn't match the
// pipeline stage, so we fell back to a generic stage agent. // pipeline stage, so we fell back to a generic stage agent.
let (already_assigned, free_agent, preferred_busy, stage_mismatch) = { let (already_assigned, free_agent, preferred_busy, stage_mismatch) = {
let agents = match self.agents.lock() { let agents = self.agents.lock().await;
Ok(a) => a,
Err(e) => {
slog_error!("[auto-assign] Failed to lock agents: {e}");
break;
}
};
let assigned = is_story_assigned_for_stage(config, &agents, story_id, stage); let assigned = is_story_assigned_for_stage(config, &agents, story_id, stage);
if assigned { if assigned {
(true, None, false, false) (true, None, false, false)
+6 -6
View File
@@ -256,7 +256,7 @@ mod tests {
let pool = AgentPool::new_test(3001); let pool = AgentPool::new_test(3001);
pool.inject_test_agent("42_story_foo", "coder-1", AgentStatus::Running); pool.inject_test_agent("42_story_foo", "coder-1", AgentStatus::Running);
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
assert!(is_story_assigned_for_stage( assert!(is_story_assigned_for_stage(
&config, &config,
&agents, &agents,
@@ -285,7 +285,7 @@ mod tests {
let pool = AgentPool::new_test(3001); let pool = AgentPool::new_test(3001);
pool.inject_test_agent("42_story_foo", "coder-1", AgentStatus::Completed); pool.inject_test_agent("42_story_foo", "coder-1", AgentStatus::Completed);
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
// Completed agents don't count as assigned // Completed agents don't count as assigned
assert!(!is_story_assigned_for_stage( assert!(!is_story_assigned_for_stage(
&config, &config,
@@ -309,7 +309,7 @@ stage = "qa"
let pool = AgentPool::new_test(3001); let pool = AgentPool::new_test(3001);
pool.inject_test_agent("42_story_foo", "qa-2", AgentStatus::Running); pool.inject_test_agent("42_story_foo", "qa-2", AgentStatus::Running);
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
// qa-2 with stage=qa should be recognised as a QA agent // qa-2 with stage=qa should be recognised as a QA agent
assert!( assert!(
is_story_assigned_for_stage(&config, &agents, "42_story_foo", &PipelineStage::Qa), is_story_assigned_for_stage(&config, &agents, "42_story_foo", &PipelineStage::Qa),
@@ -338,7 +338,7 @@ name = "coder-2"
pool.inject_test_agent("s1", "coder-1", AgentStatus::Running); pool.inject_test_agent("s1", "coder-1", AgentStatus::Running);
pool.inject_test_agent("s2", "coder-2", AgentStatus::Running); pool.inject_test_agent("s2", "coder-2", AgentStatus::Running);
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
let free = find_free_agent_for_stage(&config, &agents, &PipelineStage::Coder); let free = find_free_agent_for_stage(&config, &agents, &PipelineStage::Coder);
assert!(free.is_none(), "no free coders should be available"); assert!(free.is_none(), "no free coders should be available");
} }
@@ -361,7 +361,7 @@ name = "coder-3"
// coder-1 is busy, coder-2 is free // coder-1 is busy, coder-2 is free
pool.inject_test_agent("s1", "coder-1", AgentStatus::Running); pool.inject_test_agent("s1", "coder-1", AgentStatus::Running);
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
let free = find_free_agent_for_stage(&config, &agents, &PipelineStage::Coder); let free = find_free_agent_for_stage(&config, &agents, &PipelineStage::Coder);
assert_eq!( assert_eq!(
free, free,
@@ -384,7 +384,7 @@ name = "coder-1"
// coder-1 completed its previous story — it's free for a new one // coder-1 completed its previous story — it's free for a new one
pool.inject_test_agent("s1", "coder-1", AgentStatus::Completed); pool.inject_test_agent("s1", "coder-1", AgentStatus::Completed);
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
let free = find_free_agent_for_stage(&config, &agents, &PipelineStage::Coder); let free = find_free_agent_for_stage(&config, &agents, &PipelineStage::Coder);
assert_eq!(free, Some("coder-1"), "completed coder-1 should be free"); assert_eq!(free, Some("coder-1"), "completed coder-1 should be free");
} }
@@ -2,7 +2,7 @@
use std::collections::HashMap; use std::collections::HashMap;
use std::path::Path; use std::path::Path;
use std::sync::Mutex; use tokio::sync::Mutex;
use tokio::sync::broadcast; use tokio::sync::broadcast;
use crate::agents::pool::StoryAgent; use crate::agents::pool::StoryAgent;
@@ -83,7 +83,7 @@ pub(crate) fn count_turns_in_log(path: &Path) -> u64 {
/// Turns and budget are counted from the **current session's** log file /// Turns and budget are counted from the **current session's** log file
/// only — prior sessions are excluded so that restart counts from earlier /// only — prior sessions are excluded so that restart counts from earlier
/// runs do not accumulate against the limits. /// runs do not accumulate against the limits.
pub(super) fn check_agent_limits( pub(super) async fn check_agent_limits(
agents: &Mutex<HashMap<String, StoryAgent>>, agents: &Mutex<HashMap<String, StoryAgent>>,
project_root: &Path, project_root: &Path,
) -> Vec<(String, TerminationReason)> { ) -> Vec<(String, TerminationReason)> {
@@ -94,10 +94,7 @@ pub(super) fn check_agent_limits(
// Snapshot running agents: (key, story_id, agent_name, tx, log_session_id). // Snapshot running agents: (key, story_id, agent_name, tx, log_session_id).
let running: Vec<RunningAgentSnapshot> = { let running: Vec<RunningAgentSnapshot> = {
let lock = match agents.lock() { let lock = agents.lock().await;
Ok(l) => l,
Err(_) => return Vec::new(),
};
lock.iter() lock.iter()
.filter(|(_, agent)| agent.status == AgentStatus::Running) .filter(|(_, agent)| agent.status == AgentStatus::Running)
.map(|(key, agent)| { .map(|(key, agent)| {
@@ -25,8 +25,8 @@ pub(crate) use limits::{count_turns_in_log, resolve_session_log};
impl AgentPool { impl AgentPool {
/// Run a single watchdog pass synchronously (test helper). /// Run a single watchdog pass synchronously (test helper).
#[cfg(test)] #[cfg(test)]
pub fn run_watchdog_once(&self) { pub async fn run_watchdog_once(&self) {
check_orphaned_agents(&self.agents); check_orphaned_agents(&self.agents).await;
} }
/// Run one watchdog pass: detect orphans, enforce limits, kill offenders. /// Run one watchdog pass: detect orphans, enforce limits, kill offenders.
@@ -39,21 +39,22 @@ impl AgentPool {
/// `retry_count` is incremented and the story stays in `2_current/` for /// `retry_count` is incremented and the story stays in `2_current/` for
/// re-attempt. This prevents the original kill-respawn loop (bug 646) /// re-attempt. This prevents the original kill-respawn loop (bug 646)
/// while restoring the `max_retries` semantic for turn/budget overruns. /// while restoring the `max_retries` semantic for turn/budget overruns.
pub fn run_watchdog_pass(&self, project_root: Option<&Path>) -> usize { pub async fn run_watchdog_pass(&self, project_root: Option<&Path>) -> usize {
let orphaned = check_orphaned_agents(&self.agents); let orphaned = check_orphaned_agents(&self.agents).await;
if let Some(root) = project_root { if let Some(root) = project_root {
let terminated = check_agent_limits(&self.agents, root); let terminated = check_agent_limits(&self.agents, root).await;
let config = ProjectConfig::load(root).unwrap_or_default(); let config = ProjectConfig::load(root).unwrap_or_default();
for (key, reason) in &terminated { for (key, reason) in &terminated {
// Step 1: snapshot the agent's worktree path so we can find every // Step 1: snapshot the agent's worktree path so we can find every
// process running in it (claude + any subprocesses). This must // process running in it (claude + any subprocesses). This must
// happen BEFORE we mutate the agent record so we can read the // happen BEFORE we mutate the agent record so we can read the
// worktree info safely. // worktree info safely.
let worktree_path = self.agents.lock().ok().and_then(|lock| { let worktree_path = {
let lock = self.agents.lock().await;
lock.get(key) lock.get(key)
.and_then(|a| a.worktree_info.as_ref().map(|wt| wt.path.clone())) .and_then(|a| a.worktree_info.as_ref().map(|wt| wt.path.clone()))
}); };
// Step 2: SIGKILL every process running in the worktree and // Step 2: SIGKILL every process running in the worktree and
// BLOCK until verified gone. The previous mechanism — portable_pty's // BLOCK until verified gone. The previous mechanism — portable_pty's
@@ -85,16 +86,16 @@ impl AgentPool {
"[watchdog] No worktree path recorded for '{key}'; cannot tree-kill, \ "[watchdog] No worktree path recorded for '{key}'; cannot tree-kill, \
falling back to portable_pty SIGHUP (likely no-op for claude-code)." falling back to portable_pty SIGHUP (likely no-op for claude-code)."
); );
self.kill_child_for_key(key); self.kill_child_for_key(key).await;
} }
// Step 3: NOW update the agent record. The process is verified // Step 3: NOW update the agent record. The process is verified
// gone (or we logged that SIGKILL didn't take effect, which is // gone (or we logged that SIGKILL didn't take effect, which is
// exceptional), so flipping status away from Running can no // exceptional), so flipping status away from Running can no
// longer open a window for a concurrent spawn. // longer open a window for a concurrent spawn.
if let Ok(mut lock) = self.agents.lock()
&& let Some(agent) = lock.get_mut(key)
{ {
let mut lock = self.agents.lock().await;
if let Some(agent) = lock.get_mut(key) {
agent.status = AgentStatus::Failed; agent.status = AgentStatus::Failed;
agent.termination_reason = Some(reason.clone()); agent.termination_reason = Some(reason.clone());
if let Some(handle) = agent.task_handle.take() { if let Some(handle) = agent.task_handle.take() {
@@ -104,6 +105,7 @@ impl AgentPool {
handle.abort(); handle.abort();
} }
} }
}
// Use the retry mechanism: increment retry_count and only block // Use the retry mechanism: increment retry_count and only block
// when the limit is exceeded, matching the pipeline's behaviour. // when the limit is exceeded, matching the pipeline's behaviour.
@@ -1,7 +1,7 @@
//! Orphan detection: marks running agents whose backing task has exited. //! Orphan detection: marks running agents whose backing task has exited.
use std::collections::HashMap; use std::collections::HashMap;
use std::sync::Mutex; use tokio::sync::Mutex;
use tokio::sync::broadcast; use tokio::sync::broadcast;
use crate::agents::pool::StoryAgent; use crate::agents::pool::StoryAgent;
@@ -15,11 +15,8 @@ use crate::slog;
/// without updating the agent status — for example when the process is killed /// without updating the agent status — for example when the process is killed
/// externally and the PTY master fd returns EOF before our inactivity timeout /// externally and the PTY master fd returns EOF before our inactivity timeout
/// fires, but some other edge case prevents the normal cleanup path from running. /// fires, but some other edge case prevents the normal cleanup path from running.
pub(super) fn check_orphaned_agents(agents: &Mutex<HashMap<String, StoryAgent>>) -> usize { pub(super) async fn check_orphaned_agents(agents: &Mutex<HashMap<String, StoryAgent>>) -> usize {
let mut lock = match agents.lock() { let mut lock = agents.lock().await;
Ok(l) => l,
Err(_) => return 0,
};
// Collect orphaned entries: Running or Pending agents whose task handle is finished. // Collect orphaned entries: Running or Pending agents whose task handle is finished.
// Pending agents can be orphaned if worktree creation panics before setting status. // Pending agents can be orphaned if worktree creation panics before setting status.
@@ -10,8 +10,8 @@ use crate::agents::{AgentEvent, AgentStatus, TerminationReason};
// ── Limit enforcement integration tests (bug 624) ──────────────────────── // ── Limit enforcement integration tests (bug 624) ────────────────────────
#[test] #[tokio::test]
fn watchdog_terminates_agent_exceeding_turn_limit() { async fn watchdog_terminates_agent_exceeding_turn_limit() {
let tmp = tempfile::tempdir().unwrap(); let tmp = tempfile::tempdir().unwrap();
let root = tmp.path(); let root = tmp.path();
@@ -37,12 +37,12 @@ max_turns = 10
); );
let mut rx = tx.subscribe(); let mut rx = tx.subscribe();
let found = pool.run_watchdog_pass(Some(root)); let found = pool.run_watchdog_pass(Some(root)).await;
assert!(found >= 1, "watchdog should detect the over-limit agent"); assert!(found >= 1, "watchdog should detect the over-limit agent");
// Agent should now be Failed with TurnLimit reason. // Agent should now be Failed with TurnLimit reason.
{ {
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
let key = composite_key("story_a", "coder-1"); let key = composite_key("story_a", "coder-1");
let agent = agents.get(&key).unwrap(); let agent = agents.get(&key).unwrap();
assert_eq!(agent.status, AgentStatus::Failed); assert_eq!(agent.status, AgentStatus::Failed);
@@ -60,8 +60,8 @@ max_turns = 10
); );
} }
#[test] #[tokio::test]
fn watchdog_terminates_agent_exceeding_budget_limit() { async fn watchdog_terminates_agent_exceeding_budget_limit() {
let tmp = tempfile::tempdir().unwrap(); let tmp = tempfile::tempdir().unwrap();
let root = tmp.path(); let root = tmp.path();
@@ -87,11 +87,11 @@ max_budget_usd = 5.00
); );
let mut rx = tx.subscribe(); let mut rx = tx.subscribe();
let found = pool.run_watchdog_pass(Some(root)); let found = pool.run_watchdog_pass(Some(root)).await;
assert!(found >= 1, "watchdog should detect the over-budget agent"); assert!(found >= 1, "watchdog should detect the over-budget agent");
{ {
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
let key = composite_key("story_b", "coder-1"); let key = composite_key("story_b", "coder-1");
let agent = agents.get(&key).unwrap(); let agent = agents.get(&key).unwrap();
assert_eq!(agent.status, AgentStatus::Failed); assert_eq!(agent.status, AgentStatus::Failed);
@@ -106,8 +106,8 @@ max_budget_usd = 5.00
assert!(matches!(event, AgentEvent::Error { .. })); assert!(matches!(event, AgentEvent::Error { .. }));
} }
#[test] #[tokio::test]
fn watchdog_does_not_terminate_agent_under_limits() { async fn watchdog_does_not_terminate_agent_under_limits() {
let tmp = tempfile::tempdir().unwrap(); let tmp = tempfile::tempdir().unwrap();
let root = tmp.path(); let root = tmp.path();
@@ -133,11 +133,11 @@ max_budget_usd = 10.00
// has 25 turns < 50 so no violation). // has 25 turns < 50 so no violation).
pool.inject_test_agent_with_session("story_c", "coder-1", AgentStatus::Running, "sess-ok"); pool.inject_test_agent_with_session("story_c", "coder-1", AgentStatus::Running, "sess-ok");
let found = pool.run_watchdog_pass(Some(root)); let found = pool.run_watchdog_pass(Some(root)).await;
assert_eq!(found, 0, "agent under limits should not be terminated"); assert_eq!(found, 0, "agent under limits should not be terminated");
{ {
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
let key = composite_key("story_c", "coder-1"); let key = composite_key("story_c", "coder-1");
let agent = agents.get(&key).unwrap(); let agent = agents.get(&key).unwrap();
assert_eq!( assert_eq!(
@@ -153,8 +153,8 @@ max_budget_usd = 10.00
/// coder-1 with max_turns=50, max_budget_usd=5.00 ran 5.6× over the turn /// coder-1 with max_turns=50, max_budget_usd=5.00 ran 5.6× over the turn
/// limit (280 turns). The watchdog must terminate at the turn limit (turns /// limit (280 turns). The watchdog must terminate at the turn limit (turns
/// hit first in the observed trace), with reason TurnLimit. /// hit first in the observed trace), with reason TurnLimit.
#[test] #[tokio::test]
fn regression_bug624_coder1_story623_trajectory() { async fn regression_bug624_coder1_story623_trajectory() {
let tmp = tempfile::tempdir().unwrap(); let tmp = tempfile::tempdir().unwrap();
let root = tmp.path(); let root = tmp.path();
@@ -183,11 +183,11 @@ max_budget_usd = 5.00
); );
let mut rx = tx.subscribe(); let mut rx = tx.subscribe();
let found = pool.run_watchdog_pass(Some(root)); let found = pool.run_watchdog_pass(Some(root)).await;
assert!(found >= 1, "watchdog must catch the turn-limit violation"); assert!(found >= 1, "watchdog must catch the turn-limit violation");
{ {
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
let key = composite_key("story_623", "coder-1"); let key = composite_key("story_623", "coder-1");
let agent = agents.get(&key).unwrap(); let agent = agents.get(&key).unwrap();
assert_eq!(agent.status, AgentStatus::Failed); assert_eq!(agent.status, AgentStatus::Failed);
@@ -218,8 +218,8 @@ max_budget_usd = 5.00
/// ///
/// This test seeds a single session that legitimately exceeds the limit /// This test seeds a single session that legitimately exceeds the limit
/// and uses `max_retries = 1` so that the first violation blocks. /// and uses `max_retries = 1` so that the first violation blocks.
#[test] #[tokio::test]
fn watchdog_marks_story_blocked_after_limit_termination() { async fn watchdog_marks_story_blocked_after_limit_termination() {
crate::db::ensure_content_store(); crate::db::ensure_content_store();
crate::crdt_state::init_for_test(); crate::crdt_state::init_for_test();
@@ -263,7 +263,7 @@ max_turns = 10
"sess-runaway", "sess-runaway",
); );
let found = pool.run_watchdog_pass(Some(root)); let found = pool.run_watchdog_pass(Some(root)).await;
assert!(found >= 1, "watchdog should detect the over-limit agent"); assert!(found >= 1, "watchdog should detect the over-limit agent");
// With max_retries=1, the first violation blocks immediately via the state machine. // With max_retries=1, the first violation blocks immediately via the state machine.
@@ -278,7 +278,7 @@ max_turns = 10
// Sanity: the agent itself is also Failed with the right reason. // Sanity: the agent itself is also Failed with the right reason.
{ {
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
let key = composite_key(story_id, "coder-1"); let key = composite_key(story_id, "coder-1");
let agent = agents.get(&key).unwrap(); let agent = agents.get(&key).unwrap();
assert_eq!(agent.status, AgentStatus::Failed); assert_eq!(agent.status, AgentStatus::Failed);
@@ -297,8 +297,8 @@ max_turns = 10
/// fresh session_id whose log has fewer events than `max_turns`. /// fresh session_id whose log has fewer events than `max_turns`.
/// Assert the agent is NOT terminated (per-session count is under the /// Assert the agent is NOT terminated (per-session count is under the
/// limit) AND the story is NOT marked blocked. /// limit) AND the story is NOT marked blocked.
#[test] #[tokio::test]
fn per_session_counting_does_not_terminate_under_limit() { async fn per_session_counting_does_not_terminate_under_limit() {
let tmp = tempfile::tempdir().unwrap(); let tmp = tempfile::tempdir().unwrap();
let root = tmp.path(); let root = tmp.path();
@@ -323,14 +323,14 @@ max_turns = 10
let pool = AgentPool::new_test(3001); let pool = AgentPool::new_test(3001);
pool.inject_test_agent_with_session("story_d", "coder-1", AgentStatus::Running, "new-sess"); pool.inject_test_agent_with_session("story_d", "coder-1", AgentStatus::Running, "new-sess");
let found = pool.run_watchdog_pass(Some(root)); let found = pool.run_watchdog_pass(Some(root)).await;
assert_eq!( assert_eq!(
found, 0, found, 0,
"agent under per-session limit should NOT be terminated" "agent under per-session limit should NOT be terminated"
); );
{ {
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
let key = composite_key("story_d", "coder-1"); let key = composite_key("story_d", "coder-1");
let agent = agents.get(&key).unwrap(); let agent = agents.get(&key).unwrap();
assert_eq!( assert_eq!(
@@ -345,8 +345,8 @@ max_turns = 10
/// Same setup as per_session_counting_does_not_terminate_under_limit, but /// Same setup as per_session_counting_does_not_terminate_under_limit, but
/// the new agent's own session log exceeds `max_turns`. Assert the agent /// the new agent's own session log exceeds `max_turns`. Assert the agent
/// IS terminated AND (with max_retries=1) the story IS marked blocked. /// IS terminated AND (with max_retries=1) the story IS marked blocked.
#[test] #[tokio::test]
fn per_session_counting_terminates_over_limit() { async fn per_session_counting_terminates_over_limit() {
crate::db::ensure_content_store(); crate::db::ensure_content_store();
crate::crdt_state::init_for_test(); crate::crdt_state::init_for_test();
@@ -390,14 +390,14 @@ max_turns = 10
pool.inject_test_agent_with_session(story_id, "coder-1", AgentStatus::Running, "new-sess"); pool.inject_test_agent_with_session(story_id, "coder-1", AgentStatus::Running, "new-sess");
let mut rx = tx.subscribe(); let mut rx = tx.subscribe();
let found = pool.run_watchdog_pass(Some(root)); let found = pool.run_watchdog_pass(Some(root)).await;
assert!( assert!(
found >= 1, found >= 1,
"agent over per-session limit must be terminated" "agent over per-session limit must be terminated"
); );
{ {
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
let key = composite_key(story_id, "coder-1"); let key = composite_key(story_id, "coder-1");
let agent = agents.get(&key).unwrap(); let agent = agents.get(&key).unwrap();
assert_eq!(agent.status, AgentStatus::Failed); assert_eq!(agent.status, AgentStatus::Failed);
@@ -423,8 +423,8 @@ max_turns = 10
/// `max_turns`. After session 1: retry_count=1, NOT blocked. After /// `max_turns`. After session 1: retry_count=1, NOT blocked. After
/// session 2: retry_count=2, NOT blocked. After session 3: /// session 2: retry_count=2, NOT blocked. After session 3:
/// retry_count=3 >= max_retries, story IS blocked. /// retry_count=3 >= max_retries, story IS blocked.
#[test] #[tokio::test]
fn watchdog_retry_semantic_blocks_after_max_retries() { async fn watchdog_retry_semantic_blocks_after_max_retries() {
crate::db::ensure_content_store(); crate::db::ensure_content_store();
let tmp = tempfile::tempdir().unwrap(); let tmp = tempfile::tempdir().unwrap();
@@ -453,7 +453,7 @@ max_turns = 10
write_fake_session_log(root, story_id, "coder-1", "session-1", 12); write_fake_session_log(root, story_id, "coder-1", "session-1", 12);
let pool = AgentPool::new_test(3001); let pool = AgentPool::new_test(3001);
pool.inject_test_agent_with_session(story_id, "coder-1", AgentStatus::Running, "session-1"); pool.inject_test_agent_with_session(story_id, "coder-1", AgentStatus::Running, "session-1");
pool.run_watchdog_pass(Some(root)); pool.run_watchdog_pass(Some(root)).await;
let item = crate::crdt_state::read_item(story_id).expect("story must be in CRDT"); let item = crate::crdt_state::read_item(story_id).expect("story must be in CRDT");
assert_eq!( assert_eq!(
@@ -473,7 +473,7 @@ max_turns = 10
write_fake_session_log(root, story_id, "coder-1", "session-2", 12); write_fake_session_log(root, story_id, "coder-1", "session-2", 12);
let pool = AgentPool::new_test(3001); let pool = AgentPool::new_test(3001);
pool.inject_test_agent_with_session(story_id, "coder-1", AgentStatus::Running, "session-2"); pool.inject_test_agent_with_session(story_id, "coder-1", AgentStatus::Running, "session-2");
pool.run_watchdog_pass(Some(root)); pool.run_watchdog_pass(Some(root)).await;
let item = crate::crdt_state::read_item(story_id).expect("story must be in CRDT"); let item = crate::crdt_state::read_item(story_id).expect("story must be in CRDT");
assert_eq!( assert_eq!(
@@ -493,7 +493,7 @@ max_turns = 10
write_fake_session_log(root, story_id, "coder-1", "session-3", 12); write_fake_session_log(root, story_id, "coder-1", "session-3", 12);
let pool = AgentPool::new_test(3001); let pool = AgentPool::new_test(3001);
pool.inject_test_agent_with_session(story_id, "coder-1", AgentStatus::Running, "session-3"); pool.inject_test_agent_with_session(story_id, "coder-1", AgentStatus::Running, "session-3");
pool.run_watchdog_pass(Some(root)); pool.run_watchdog_pass(Some(root)).await;
let item = crate::crdt_state::read_item(story_id).expect("story must be in CRDT"); let item = crate::crdt_state::read_item(story_id).expect("story must be in CRDT");
assert_eq!( assert_eq!(
@@ -518,8 +518,8 @@ max_turns = 10
/// must not count against the watchdog's turn budget. A session log with /// must not count against the watchdog's turn budget. A session log with
/// 5 tool turns and 30 narration turns reports turns_used == 5, so an /// 5 tool turns and 30 narration turns reports turns_used == 5, so an
/// agent with max_tool_turns = 10 stays Running. /// agent with max_tool_turns = 10 stays Running.
#[test] #[tokio::test]
fn watchdog_does_not_count_narration_only_turns() { async fn watchdog_does_not_count_narration_only_turns() {
let tmp = tempfile::tempdir().unwrap(); let tmp = tempfile::tempdir().unwrap();
let root = tmp.path(); let root = tmp.path();
@@ -542,13 +542,13 @@ max_turns = 200
let pool = AgentPool::new_test(3001); let pool = AgentPool::new_test(3001);
pool.inject_test_agent_with_session("story_923", "coder-1", AgentStatus::Running, "sess-narr"); pool.inject_test_agent_with_session("story_923", "coder-1", AgentStatus::Running, "sess-narr");
let found = pool.run_watchdog_pass(Some(root)); let found = pool.run_watchdog_pass(Some(root)).await;
assert_eq!( assert_eq!(
found, 0, found, 0,
"agent must not be terminated: only 5 tool turns of a 10-turn budget" "agent must not be terminated: only 5 tool turns of a 10-turn budget"
); );
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
let key = composite_key("story_923", "coder-1"); let key = composite_key("story_923", "coder-1");
let agent = agents.get(&key).unwrap(); let agent = agents.get(&key).unwrap();
assert_eq!(agent.status, AgentStatus::Running); assert_eq!(agent.status, AgentStatus::Running);
@@ -558,8 +558,8 @@ max_turns = 200
/// Story 923: max_tool_turns takes precedence over max_turns when both are /// Story 923: max_tool_turns takes precedence over max_turns when both are
/// set. With max_tool_turns = 3 and max_turns = 200, an agent that has 4 /// set. With max_tool_turns = 3 and max_turns = 200, an agent that has 4
/// tool turns is killed even though total turns (4) is far below max_turns. /// tool turns is killed even though total turns (4) is far below max_turns.
#[test] #[tokio::test]
fn watchdog_max_tool_turns_overrides_max_turns() { async fn watchdog_max_tool_turns_overrides_max_turns() {
let tmp = tempfile::tempdir().unwrap(); let tmp = tempfile::tempdir().unwrap();
let root = tmp.path(); let root = tmp.path();
@@ -585,13 +585,13 @@ max_turns = 200
); );
let mut rx = tx.subscribe(); let mut rx = tx.subscribe();
let found = pool.run_watchdog_pass(Some(root)); let found = pool.run_watchdog_pass(Some(root)).await;
assert!( assert!(
found >= 1, found >= 1,
"watchdog must terminate when tool turns exceed max_tool_turns" "watchdog must terminate when tool turns exceed max_tool_turns"
); );
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
let key = composite_key("story_923b", "coder-1"); let key = composite_key("story_923b", "coder-1");
let agent = agents.get(&key).unwrap(); let agent = agents.get(&key).unwrap();
assert_eq!(agent.status, AgentStatus::Failed); assert_eq!(agent.status, AgentStatus::Failed);
@@ -20,18 +20,18 @@ async fn check_orphaned_agents_returns_count_of_orphaned_agents() {
pool.inject_test_agent_with_handle("story_a", "coder", AgentStatus::Running, h1); pool.inject_test_agent_with_handle("story_a", "coder", AgentStatus::Running, h1);
pool.inject_test_agent_with_handle("story_b", "coder", AgentStatus::Running, h2); pool.inject_test_agent_with_handle("story_b", "coder", AgentStatus::Running, h2);
let found = check_orphaned_agents(&pool.agents); let found = check_orphaned_agents(&pool.agents).await;
assert_eq!(found, 2, "should detect both orphaned agents"); assert_eq!(found, 2, "should detect both orphaned agents");
} }
#[test] #[tokio::test]
fn check_orphaned_agents_returns_zero_when_no_orphans() { async fn check_orphaned_agents_returns_zero_when_no_orphans() {
let pool = AgentPool::new_test(3001); let pool = AgentPool::new_test(3001);
// Inject agents in terminal states — not orphaned. // Inject agents in terminal states — not orphaned.
pool.inject_test_agent("story_a", "coder", AgentStatus::Completed); pool.inject_test_agent("story_a", "coder", AgentStatus::Completed);
pool.inject_test_agent("story_b", "qa", AgentStatus::Failed); pool.inject_test_agent("story_b", "qa", AgentStatus::Failed);
let found = check_orphaned_agents(&pool.agents); let found = check_orphaned_agents(&pool.agents).await;
assert_eq!( assert_eq!(
found, 0, found, 0,
"no orphans should be detected for terminal agents" "no orphans should be detected for terminal agents"
@@ -53,10 +53,10 @@ async fn watchdog_detects_orphaned_running_agent() {
pool.inject_test_agent_with_handle("orphan_story", "coder", AgentStatus::Running, handle); pool.inject_test_agent_with_handle("orphan_story", "coder", AgentStatus::Running, handle);
let mut rx = tx.subscribe(); let mut rx = tx.subscribe();
pool.run_watchdog_once(); pool.run_watchdog_once().await;
{ {
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
let key = composite_key("orphan_story", "coder"); let key = composite_key("orphan_story", "coder");
let agent = agents.get(&key).unwrap(); let agent = agents.get(&key).unwrap();
assert_eq!( assert_eq!(
@@ -87,13 +87,13 @@ async fn watchdog_orphan_detection_returns_nonzero_enabling_auto_assign() {
// Before watchdog: agent is Running. // Before watchdog: agent is Running.
{ {
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
let key = composite_key("orphan_story", "coder"); let key = composite_key("orphan_story", "coder");
assert_eq!(agents.get(&key).unwrap().status, AgentStatus::Running); assert_eq!(agents.get(&key).unwrap().status, AgentStatus::Running);
} }
// Run watchdog pass — should return 1 (orphan found). // Run watchdog pass — should return 1 (orphan found).
let found = check_orphaned_agents(&pool.agents); let found = check_orphaned_agents(&pool.agents).await;
assert_eq!( assert_eq!(
found, 1, found, 1,
"watchdog must return 1 for a single orphaned agent" "watchdog must return 1 for a single orphaned agent"
@@ -101,7 +101,7 @@ async fn watchdog_orphan_detection_returns_nonzero_enabling_auto_assign() {
// After watchdog: agent is Failed. // After watchdog: agent is Failed.
{ {
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
let key = composite_key("orphan_story", "coder"); let key = composite_key("orphan_story", "coder");
assert_eq!( assert_eq!(
agents.get(&key).unwrap().status, agents.get(&key).unwrap().status,
+7 -5
View File
@@ -19,8 +19,8 @@ mod test_helpers;
use crate::io::watcher::WatcherEvent; use crate::io::watcher::WatcherEvent;
use crate::service::status::StatusBroadcaster; use crate::service::status::StatusBroadcaster;
use std::collections::HashMap; use std::collections::HashMap;
use std::sync::{Arc, Mutex}; use std::sync::Arc;
use tokio::sync::broadcast; use tokio::sync::{Mutex, broadcast};
// Bring pool-internal types into pool's namespace so that sub-modules // Bring pool-internal types into pool's namespace so that sub-modules
// (auto_assign, pipeline, etc.) can access them via `use super::...`. // (auto_assign, pipeline, etc.) can access them via `use super::...`.
@@ -87,10 +87,12 @@ impl AgentPool {
_ => continue, _ => continue,
}; };
let key = composite_key(&story_id, &agent_name); let key = composite_key(&story_id, &agent_name);
if let Ok(mut agents) = agents_clone.lock()
&& let Some(agent) = agents.get_mut(&key)
{ {
agent.throttled = Some(crate::agents::AgentExecution::Throttled { until }); let mut agents = agents_clone.lock().await;
if let Some(agent) = agents.get_mut(&key) {
agent.throttled =
Some(crate::agents::AgentExecution::Throttled { until });
}
} }
let _ = watcher_tx_clone.send(WatcherEvent::AgentStateChanged); let _ = watcher_tx_clone.send(WatcherEvent::AgentStateChanged);
} }
@@ -3,7 +3,7 @@
use std::collections::HashMap; use std::collections::HashMap;
use std::path::PathBuf; use std::path::PathBuf;
use std::sync::{Arc, Mutex}; use std::sync::Arc;
use tokio::sync::broadcast; use tokio::sync::broadcast;
@@ -16,7 +16,7 @@ use std::path::Path;
/// type cycle between `start_agent` and `run_server_owned_completion`. /// type cycle between `start_agent` and `run_server_owned_completion`.
#[allow(clippy::too_many_arguments)] #[allow(clippy::too_many_arguments)]
pub(crate) fn spawn_pipeline_advance( pub(crate) fn spawn_pipeline_advance(
agents: Arc<Mutex<HashMap<String, StoryAgent>>>, agents: Arc<tokio::sync::Mutex<HashMap<String, StoryAgent>>>,
port: u16, port: u16,
story_id: &str, story_id: &str,
agent_name: &str, agent_name: &str,
@@ -694,7 +694,7 @@ impl AgentPool {
if let Err(e) = crate::agents::lifecycle::move_story_to_done(story_id) { if let Err(e) = crate::agents::lifecycle::move_story_to_done(story_id) {
slog_error!("[pipeline] Failed to move '{story_id}' to done: {e}"); slog_error!("[pipeline] Failed to move '{story_id}' to done: {e}");
} }
self.remove_agents_for_story(story_id); self.remove_agents_for_story(story_id).await;
crate::crdt_state::delete_merge_job(story_id); crate::crdt_state::delete_merge_job(story_id);
// TODO: Re-enable worktree cleanup once we have persistent agent logs. // TODO: Re-enable worktree cleanup once we have persistent agent logs.
// Removing worktrees destroys evidence needed to debug empty-commit agents. // Removing worktrees destroys evidence needed to debug empty-commit agents.
@@ -104,7 +104,7 @@ async fn mergemaster_blocks_and_sends_story_blocked_when_no_commits_ahead() {
); );
// No mergemaster agent should have been started. // No mergemaster agent should have been started.
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
let mergemaster_started = agents let mergemaster_started = agents
.values() .values()
.any(|a| a.agent_name.contains("mergemaster")); .any(|a| a.agent_name.contains("mergemaster"));
@@ -162,7 +162,7 @@ stage = "qa"
// Verify that 293 cannot get a QA agent right now (QA is busy). // Verify that 293 cannot get a QA agent right now (QA is busy).
{ {
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
assert!( assert!(
!is_agent_free(&agents, "qa"), !is_agent_free(&agents, "qa"),
"qa should be busy on story 292" "qa should be busy on story 292"
@@ -172,7 +172,7 @@ stage = "qa"
// Simulate QA completing on story 292: remove the agent from the pool // Simulate QA completing on story 292: remove the agent from the pool
// (as run_server_owned_completion does) then run pipeline advance. // (as run_server_owned_completion does) then run pipeline advance.
{ {
let mut agents = pool.agents.lock().unwrap(); let mut agents = pool.agents.try_lock().unwrap();
agents.remove(&composite_key("292_story_first", "qa")); agents.remove(&composite_key("292_story_first", "qa"));
} }
@@ -193,7 +193,7 @@ stage = "qa"
.await; .await;
// After pipeline advance, auto_assign should have started QA on story 293. // After pipeline advance, auto_assign should have started QA on story 293.
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
let qa_on_293 = agents.values().any(|a| { let qa_on_293 = agents.values().any(|a| {
a.agent_name == "qa" && matches!(a.status, AgentStatus::Pending | AgentStatus::Running) a.agent_name == "qa" && matches!(a.status, AgentStatus::Pending | AgentStatus::Running)
}); });
@@ -278,7 +278,7 @@ async fn stale_mergemaster_advance_for_done_story_is_noop() {
.await; .await;
// No agents should have been started. // No agents should have been started.
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
assert!( assert!(
agents.is_empty(), agents.is_empty(),
"No agents should be started for a stale advance on a done story. \ "No agents should be started for a stale advance on a done story. \
@@ -871,7 +871,7 @@ stage = "coder"
.await; .await;
// The coder must be re-spawned — Pending or Running. // The coder must be re-spawned — Pending or Running.
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
let coder_restarted = agents.values().any(|a| { let coder_restarted = agents.values().any(|a| {
a.agent_name == "coder-1" && matches!(a.status, AgentStatus::Pending | AgentStatus::Running) a.agent_name == "coder-1" && matches!(a.status, AgentStatus::Pending | AgentStatus::Running)
}); });
@@ -957,7 +957,7 @@ stage = "coder"
.await; .await;
// The recovery respawn must have been issued — coder-1 should be Pending/Running. // The recovery respawn must have been issued — coder-1 should be Pending/Running.
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
let coder_restarted = agents.values().any(|a| { let coder_restarted = agents.values().any(|a| {
a.agent_name == "coder-1" && matches!(a.status, AgentStatus::Pending | AgentStatus::Running) a.agent_name == "coder-1" && matches!(a.status, AgentStatus::Pending | AgentStatus::Running)
}); });
@@ -1328,7 +1328,7 @@ async fn coder_completion_with_test_evidence_and_zero_commits_does_not_advance()
); );
// No QA or merge agent should have been started. // No QA or merge agent should have been started.
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
let qa_or_merge_started = agents let qa_or_merge_started = agents
.values() .values()
.any(|a| a.agent_name.contains("qa") || a.agent_name.contains("merge")); .any(|a| a.agent_name.contains("qa") || a.agent_name.contains("merge"));
@@ -28,7 +28,7 @@ impl AgentPool {
// Verify agent exists, is Running, and grab its worktree path. // Verify agent exists, is Running, and grab its worktree path.
let worktree_path = { let worktree_path = {
let agents = self.agents.lock().map_err(|e| e.to_string())?; let agents = self.agents.lock().await;
let agent = agents let agent = agents
.get(&key) .get(&key)
.ok_or_else(|| format!("No agent '{agent_name}' for story '{story_id}'"))?; .ok_or_else(|| format!("No agent '{agent_name}' for story '{story_id}'"))?;
@@ -82,7 +82,7 @@ impl AgentPool {
merge_failure_reported_for_advance, merge_failure_reported_for_advance,
session_id_for_advance, session_id_for_advance,
) = { ) = {
let mut agents = self.agents.lock().map_err(|e| e.to_string())?; let mut agents = self.agents.lock().await;
let agent = agents.get_mut(&key).ok_or_else(|| { let agent = agents.get_mut(&key).ok_or_else(|| {
format!("Agent '{agent_name}' for story '{story_id}' disappeared during gate check") format!("Agent '{agent_name}' for story '{story_id}' disappeared during gate check")
})?; })?;
@@ -2,7 +2,8 @@
use crate::io::watcher::WatcherEvent; use crate::io::watcher::WatcherEvent;
use crate::slog; use crate::slog;
use std::collections::HashMap; use std::collections::HashMap;
use std::sync::{Arc, Mutex}; use std::sync::Arc;
use tokio::sync::Mutex;
use tokio::sync::broadcast; use tokio::sync::broadcast;
use super::super::super::super::{AgentEvent, CompletionReport, PipelineStage, pipeline_stage}; use super::super::super::super::{AgentEvent, CompletionReport, PipelineStage, pipeline_stage};
@@ -45,10 +46,7 @@ pub(in crate::agents::pool) async fn run_server_owned_completion(
// Guard: skip if completion was already recorded (legacy path). // Guard: skip if completion was already recorded (legacy path).
{ {
let lock = match agents.lock() { let lock = agents.lock().await;
Ok(a) => a,
Err(_) => return,
};
match lock.get(&key) { match lock.get(&key) {
Some(agent) if agent.completion.is_some() => { Some(agent) if agent.completion.is_some() => {
slog!( slog!(
@@ -64,10 +62,7 @@ pub(in crate::agents::pool) async fn run_server_owned_completion(
// Get worktree path for running gates. // Get worktree path for running gates.
let worktree_path = { let worktree_path = {
let lock = match agents.lock() { let lock = agents.lock().await;
Ok(a) => a,
Err(_) => return,
};
lock.get(&key) lock.get(&key)
.and_then(|a| a.worktree_info.as_ref().map(|wt| wt.path.clone())) .and_then(|a| a.worktree_info.as_ref().map(|wt| wt.path.clone()))
}; };
@@ -192,10 +187,7 @@ pub(in crate::agents::pool) async fn run_server_owned_completion(
// Store completion report, extract data for pipeline advance, then // Store completion report, extract data for pipeline advance, then
// remove the entry so completed agents never appear in list_agents. // remove the entry so completed agents never appear in list_agents.
let (tx, project_root_for_advance, wt_path_for_advance, merge_failure_reported_for_advance) = { let (tx, project_root_for_advance, wt_path_for_advance, merge_failure_reported_for_advance) = {
let mut lock = match agents.lock() { let mut lock = agents.lock().await;
Ok(a) => a,
Err(_) => return,
};
let agent = match lock.get_mut(&key) { let agent = match lock.get_mut(&key) {
Some(a) => a, Some(a) => a,
None => return, None => return,
@@ -108,7 +108,7 @@ async fn server_owned_completion_skips_when_already_completed() {
); );
// Subscribe before calling so we can check if Done event was emitted. // Subscribe before calling so we can check if Done event was emitted.
let mut rx = pool.subscribe("s10", "coder-1").unwrap(); let mut rx = pool.subscribe("s10", "coder-1").await.unwrap();
run_server_owned_completion( run_server_owned_completion(
&pool.agents, &pool.agents,
@@ -121,7 +121,7 @@ async fn server_owned_completion_skips_when_already_completed() {
.await; .await;
// Status should remain Completed (unchanged) — no gate re-run. // Status should remain Completed (unchanged) — no gate re-run.
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
let key = super::super::super::composite_key("s10", "coder-1"); let key = super::super::super::composite_key("s10", "coder-1");
let agent = agents.get(&key).unwrap(); let agent = agents.get(&key).unwrap();
assert_eq!(agent.status, AgentStatus::Completed); assert_eq!(agent.status, AgentStatus::Completed);
@@ -147,7 +147,7 @@ async fn server_owned_completion_runs_gates_on_clean_worktree() {
let pool = AgentPool::new_test(3001); let pool = AgentPool::new_test(3001);
pool.inject_test_agent_with_path("s11", "coder-1", AgentStatus::Running, repo.to_path_buf()); pool.inject_test_agent_with_path("s11", "coder-1", AgentStatus::Running, repo.to_path_buf());
let mut rx = pool.subscribe("s11", "coder-1").unwrap(); let mut rx = pool.subscribe("s11", "coder-1").await.unwrap();
run_server_owned_completion( run_server_owned_completion(
&pool.agents, &pool.agents,
@@ -160,7 +160,7 @@ async fn server_owned_completion_runs_gates_on_clean_worktree() {
.await; .await;
// Agent entry should be removed from the map after completion. // Agent entry should be removed from the map after completion.
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
let key = super::super::super::composite_key("s11", "coder-1"); let key = super::super::super::composite_key("s11", "coder-1");
assert!( assert!(
agents.get(&key).is_none(), agents.get(&key).is_none(),
@@ -192,7 +192,7 @@ async fn server_owned_completion_fails_on_dirty_worktree() {
let pool = AgentPool::new_test(3001); let pool = AgentPool::new_test(3001);
pool.inject_test_agent_with_path("s12", "coder-1", AgentStatus::Running, repo.to_path_buf()); pool.inject_test_agent_with_path("s12", "coder-1", AgentStatus::Running, repo.to_path_buf());
let mut rx = pool.subscribe("s12", "coder-1").unwrap(); let mut rx = pool.subscribe("s12", "coder-1").await.unwrap();
run_server_owned_completion( run_server_owned_completion(
&pool.agents, &pool.agents,
@@ -205,7 +205,7 @@ async fn server_owned_completion_fails_on_dirty_worktree() {
.await; .await;
// Agent entry should be removed from the map after completion (even on failure). // Agent entry should be removed from the map after completion (even on failure).
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
let key = super::super::super::composite_key("s12", "coder-1"); let key = super::super::super::composite_key("s12", "coder-1");
assert!( assert!(
agents.get(&key).is_none(), agents.get(&key).is_none(),
@@ -307,7 +307,7 @@ async fn server_owned_completion_is_noop_for_mergemaster() {
// The agent entry should remain in the pool (lifecycle cleanup is the // The agent entry should remain in the pool (lifecycle cleanup is the
// caller's responsibility, not run_server_owned_completion's). // caller's responsibility, not run_server_owned_completion's).
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
let key = super::super::super::composite_key("99_story_merge445", "mergemaster"); let key = super::super::super::composite_key("99_story_merge445", "mergemaster");
assert!( assert!(
agents.get(&key).is_some(), agents.get(&key).is_some(),
@@ -361,7 +361,7 @@ async fn server_owned_completion_preserves_dirty_worktree_with_committed_work()
let pool = AgentPool::new_test(3001); let pool = AgentPool::new_test(3001);
pool.inject_test_agent_with_path("645_test", "coder-1", AgentStatus::Running, wt_path.clone()); pool.inject_test_agent_with_path("645_test", "coder-1", AgentStatus::Running, wt_path.clone());
let mut rx = pool.subscribe("645_test", "coder-1").unwrap(); let mut rx = pool.subscribe("645_test", "coder-1").await.unwrap();
run_server_owned_completion( run_server_owned_completion(
&pool.agents, &pool.agents,
@@ -34,9 +34,8 @@ impl AgentPool {
/// If the agent was already removed from the pool (race: `remove_agents_for_story` /// If the agent was already removed from the pool (race: `remove_agents_for_story`
/// ran first) this is a no-op; the `ContentKey::MergeSuccess` DB key written /// ran first) this is a no-op; the `ContentKey::MergeSuccess` DB key written
/// by the caller acts as the authoritative fallback in that case. /// by the caller acts as the authoritative fallback in that case.
pub fn set_merge_success_reported(&self, story_id: &str) { pub async fn set_merge_success_reported(&self, story_id: &str) {
match self.agents.lock() { let mut lock = self.agents.lock().await;
Ok(mut lock) => {
let found = lock.iter_mut().find(|(key, agent)| { let found = lock.iter_mut().find(|(key, agent)| {
let key_story_id = key let key_story_id = key
.rsplit_once(':') .rsplit_once(':')
@@ -61,11 +60,6 @@ impl AgentPool {
} }
} }
} }
Err(e) => {
slog_error!("[pipeline] set_merge_success_reported: could not lock agents: {e}");
}
}
}
/// Record that the mergemaster agent for `story_id` explicitly reported a /// Record that the mergemaster agent for `story_id` explicitly reported a
/// merge failure via the `report_merge_failure` MCP tool. /// merge failure via the `report_merge_failure` MCP tool.
@@ -74,9 +68,8 @@ impl AgentPool {
/// that `run_pipeline_advance` can block advancement to `5_done/` even when /// that `run_pipeline_advance` can block advancement to `5_done/` even when
/// the server-owned gate check returns `gates_passed=true` (those gates run /// the server-owned gate check returns `gates_passed=true` (those gates run
/// in the feature-branch worktree, not on master). /// in the feature-branch worktree, not on master).
pub fn set_merge_failure_reported(&self, story_id: &str) { pub async fn set_merge_failure_reported(&self, story_id: &str) {
match self.agents.lock() { let mut lock = self.agents.lock().await;
Ok(mut lock) => {
let found = lock.iter_mut().find(|(key, agent)| { let found = lock.iter_mut().find(|(key, agent)| {
let key_story_id = key let key_story_id = key
.rsplit_once(':') .rsplit_once(':')
@@ -101,9 +94,4 @@ impl AgentPool {
} }
} }
} }
Err(e) => {
slog_error!("[pipeline] set_merge_failure_reported: could not lock agents: {e}");
}
}
}
} }
@@ -283,7 +283,7 @@ impl AgentPool {
&& let Ok(ref r) = report && let Ok(ref r) = report
&& r.story_archived && r.story_archived
{ {
pool.set_merge_success_reported(&sid); pool.set_merge_success_reported(&sid).await;
crate::db::write_content(crate::db::ContentKey::MergeSuccess(&sid), "1"); crate::db::write_content(crate::db::ContentKey::MergeSuccess(&sid), "1");
} }
@@ -350,7 +350,7 @@ impl AgentPool {
let story_archived = crate::agents::lifecycle::move_story_to_done(story_id).is_ok(); let story_archived = crate::agents::lifecycle::move_story_to_done(story_id).is_ok();
if story_archived { if story_archived {
self.remove_agents_for_story(story_id); self.remove_agents_for_story(story_id).await;
} }
let worktree_cleaned_up = if wt_path.exists() { let worktree_cleaned_up = if wt_path.exists() {
@@ -149,7 +149,7 @@ async fn reap_stale_merge_jobs_removes_old_running_entry_without_merge() {
); );
// No agents must have been spawned (no merge was triggered). // No agents must have been spawned (no merge was triggered).
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
assert!( assert!(
agents.is_empty(), agents.is_empty(),
"reap must not spawn any agents; got {} agent(s)", "reap must not spawn any agents; got {} agent(s)",
@@ -811,7 +811,7 @@ async fn server_side_merge_happy_path_advances_to_done() {
} }
// Verify no LLM agent was spawned. // Verify no LLM agent was spawned.
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.try_lock().unwrap();
assert!( assert!(
agents.is_empty(), agents.is_empty(),
"no LLM agents should be spawned for deterministic merge; pool has {} agents", "no LLM agents should be spawned for deterministic merge; pool has {} agents",
+13 -17
View File
@@ -25,11 +25,9 @@ impl AgentPool {
/// continuing to run after the server exits. Collects each agent's worktree /// continuing to run after the server exits. Collects each agent's worktree
/// path, then SIGKILLs every process running inside that path and verifies /// path, then SIGKILLs every process running inside that path and verifies
/// termination before returning. /// termination before returning.
pub fn kill_all_children(&self) { pub async fn kill_all_children(&self) {
let worktree_paths: Vec<(String, std::path::PathBuf)> = { let worktree_paths: Vec<(String, std::path::PathBuf)> = {
let Ok(agents) = self.agents.lock() else { let agents = self.agents.lock().await;
return;
};
agents agents
.iter() .iter()
.filter_map(|(key, agent)| { .filter_map(|(key, agent)| {
@@ -69,11 +67,9 @@ impl AgentPool {
/// Fallback used by `stop_agent` when no worktree path is recorded for the /// Fallback used by `stop_agent` when no worktree path is recorded for the
/// agent. Also the primary kill path for any caller that has only a composite /// agent. Also the primary kill path for any caller that has only a composite
/// key and not a worktree path directly. /// key and not a worktree path directly.
pub(super) fn kill_child_for_key(&self, key: &str) { pub(super) async fn kill_child_for_key(&self, key: &str) {
let worktree_path = { let worktree_path = {
let Ok(agents) = self.agents.lock() else { let agents = self.agents.lock().await;
return;
};
agents agents
.get(key) .get(key)
.and_then(|a| a.worktree_info.as_ref().map(|wt| wt.path.clone())) .and_then(|a| a.worktree_info.as_ref().map(|wt| wt.path.clone()))
@@ -124,18 +120,18 @@ mod tests {
.unwrap_or(false) .unwrap_or(false)
} }
#[test] #[tokio::test]
fn kill_all_children_is_safe_on_empty_pool() { async fn kill_all_children_is_safe_on_empty_pool() {
let pool = AgentPool::new_test(3001); let pool = AgentPool::new_test(3001);
pool.kill_all_children(); // must not panic pool.kill_all_children().await; // must not panic
} }
/// AC 4 — `kill_child_for_key` SIGKILLs the single agent's process and /// AC 4 — `kill_child_for_key` SIGKILLs the single agent's process and
/// verifies it is gone within 2 s. The sleeper has the worktree path in /// verifies it is gone within 2 s. The sleeper has the worktree path in
/// its argv[0] so `pgrep -f` can locate it, mirroring how claude-code is /// its argv[0] so `pgrep -f` can locate it, mirroring how claude-code is
/// launched with `--directory <worktree>` in production. /// launched with `--directory <worktree>` in production.
#[test] #[tokio::test]
fn kill_child_for_key_kills_real_process() { async fn kill_child_for_key_kills_real_process() {
use std::os::unix::process::CommandExt; use std::os::unix::process::CommandExt;
let pool = AgentPool::new_test(3002); let pool = AgentPool::new_test(3002);
@@ -165,7 +161,7 @@ mod tests {
"sleeper pid {pid} should be running before kill_child_for_key" "sleeper pid {pid} should be running before kill_child_for_key"
); );
pool.kill_child_for_key("story-1090-kill:coder"); pool.kill_child_for_key("story-1090-kill:coder").await;
let _ = child.wait(); // reap zombie so ps -p returns false let _ = child.wait(); // reap zombie so ps -p returns false
assert!( assert!(
@@ -176,8 +172,8 @@ mod tests {
/// AC 5 — `kill_all_children` SIGKILLs all agents' processes. Two agents /// AC 5 — `kill_all_children` SIGKILLs all agents' processes. Two agents
/// with distinct worktree paths are injected; both must be gone after the call. /// with distinct worktree paths are injected; both must be gone after the call.
#[test] #[tokio::test]
fn kill_all_children_kills_multiple_real_processes() { async fn kill_all_children_kills_multiple_real_processes() {
use std::os::unix::process::CommandExt; use std::os::unix::process::CommandExt;
let pool = AgentPool::new_test(3003); let pool = AgentPool::new_test(3003);
@@ -213,7 +209,7 @@ mod tests {
); );
} }
pool.kill_all_children(); pool.kill_all_children().await;
for (pid, child, _tmp) in &mut sleepers { for (pid, child, _tmp) in &mut sleepers {
let _ = child.wait(); // reap zombie let _ = child.wait(); // reap zombie
+45 -16
View File
@@ -10,12 +10,12 @@ use super::types::{agent_info_from_entry, composite_key};
impl AgentPool { impl AgentPool {
/// Return the names of configured agents for `stage` that are not currently /// Return the names of configured agents for `stage` that are not currently
/// running or pending. /// running or pending.
pub fn available_agents_for_stage( pub async fn available_agents_for_stage(
&self, &self,
config: &ProjectConfig, config: &ProjectConfig,
stage: &PipelineStage, stage: &PipelineStage,
) -> Result<Vec<String>, String> { ) -> Result<Vec<String>, String> {
let agents = self.agents.lock().map_err(|e| e.to_string())?; let agents = self.agents.lock().await;
Ok(config Ok(config
.agent .agent
.iter() .iter()
@@ -44,8 +44,8 @@ impl AgentPool {
} }
/// List all agents with their status. /// List all agents with their status.
pub fn list_agents(&self) -> Result<Vec<AgentInfo>, String> { pub async fn list_agents(&self) -> Result<Vec<AgentInfo>, String> {
let agents = self.agents.lock().map_err(|e| e.to_string())?; let agents = self.agents.lock().await;
Ok(agents Ok(agents
.iter() .iter()
.map(|(key, agent)| { .map(|(key, agent)| {
@@ -59,14 +59,35 @@ impl AgentPool {
.collect()) .collect())
} }
/// Best-effort agent list for sync callers (chat commands, rendering).
///
/// Uses `try_lock()` so it never blocks the thread. Returns an empty list
/// if the lock is currently held — callers that refresh periodically (htop,
/// pipeline board) tolerate this gracefully.
pub fn list_agents_nonblocking(&self) -> Vec<AgentInfo> {
let Ok(agents) = self.agents.try_lock() else {
return Vec::new();
};
agents
.iter()
.map(|(key, agent)| {
let story_id = key
.rsplit_once(':')
.map(|(sid, _)| sid.to_string())
.unwrap_or_else(|| key.clone());
agent_info_from_entry(&story_id, agent)
})
.collect()
}
/// Subscribe to events for a story agent. /// Subscribe to events for a story agent.
pub fn subscribe( pub async fn subscribe(
&self, &self,
story_id: &str, story_id: &str,
agent_name: &str, agent_name: &str,
) -> Result<broadcast::Receiver<AgentEvent>, String> { ) -> Result<broadcast::Receiver<AgentEvent>, String> {
let key = composite_key(story_id, agent_name); let key = composite_key(story_id, agent_name);
let agents = self.agents.lock().map_err(|e| e.to_string())?; let agents = self.agents.lock().await;
let agent = agents let agent = agents
.get(&key) .get(&key)
.ok_or_else(|| format!("No agent '{agent_name}' for story '{story_id}'"))?; .ok_or_else(|| format!("No agent '{agent_name}' for story '{story_id}'"))?;
@@ -74,13 +95,13 @@ impl AgentPool {
} }
/// Drain accumulated events for polling. Returns all events since the last drain. /// Drain accumulated events for polling. Returns all events since the last drain.
pub fn drain_events( pub async fn drain_events(
&self, &self,
story_id: &str, story_id: &str,
agent_name: &str, agent_name: &str,
) -> Result<Vec<AgentEvent>, String> { ) -> Result<Vec<AgentEvent>, String> {
let key = composite_key(story_id, agent_name); let key = composite_key(story_id, agent_name);
let agents = self.agents.lock().map_err(|e| e.to_string())?; let agents = self.agents.lock().await;
let agent = agents let agent = agents
.get(&key) .get(&key)
.ok_or_else(|| format!("No agent '{agent_name}' for story '{story_id}'"))?; .ok_or_else(|| format!("No agent '{agent_name}' for story '{story_id}'"))?;
@@ -91,9 +112,13 @@ impl AgentPool {
/// Get the log session ID and project root for an agent, if available. /// Get the log session ID and project root for an agent, if available.
/// ///
/// Used by MCP tools to find the persistent log file for a completed agent. /// Used by MCP tools to find the persistent log file for a completed agent.
pub fn get_log_info(&self, story_id: &str, agent_name: &str) -> Option<(String, PathBuf)> { pub async fn get_log_info(
&self,
story_id: &str,
agent_name: &str,
) -> Option<(String, PathBuf)> {
let key = composite_key(story_id, agent_name); let key = composite_key(story_id, agent_name);
let agents = self.agents.lock().ok()?; let agents = self.agents.lock().await;
let agent = agents.get(&key)?; let agent = agents.get(&key)?;
let session_id = agent.log_session_id.clone()?; let session_id = agent.log_session_id.clone()?;
let project_root = agent.project_root.clone()?; let project_root = agent.project_root.clone()?;
@@ -111,8 +136,8 @@ mod tests {
ProjectConfig::parse(toml_str).unwrap() ProjectConfig::parse(toml_str).unwrap()
} }
#[test] #[tokio::test]
fn available_agents_for_stage_returns_idle_agents() { async fn available_agents_for_stage_returns_idle_agents() {
let config = make_config( let config = make_config(
r#" r#"
[[agent]] [[agent]]
@@ -133,17 +158,19 @@ stage = "qa"
let available = pool let available = pool
.available_agents_for_stage(&config, &PipelineStage::Coder) .available_agents_for_stage(&config, &PipelineStage::Coder)
.await
.unwrap(); .unwrap();
assert_eq!(available, vec!["coder-2"]); assert_eq!(available, vec!["coder-2"]);
let available_qa = pool let available_qa = pool
.available_agents_for_stage(&config, &PipelineStage::Qa) .available_agents_for_stage(&config, &PipelineStage::Qa)
.await
.unwrap(); .unwrap();
assert_eq!(available_qa, vec!["qa"]); assert_eq!(available_qa, vec!["qa"]);
} }
#[test] #[tokio::test]
fn available_agents_for_stage_returns_empty_when_all_busy() { async fn available_agents_for_stage_returns_empty_when_all_busy() {
let config = make_config( let config = make_config(
r#" r#"
[[agent]] [[agent]]
@@ -156,12 +183,13 @@ stage = "coder"
let available = pool let available = pool
.available_agents_for_stage(&config, &PipelineStage::Coder) .available_agents_for_stage(&config, &PipelineStage::Coder)
.await
.unwrap(); .unwrap();
assert!(available.is_empty()); assert!(available.is_empty());
} }
#[test] #[tokio::test]
fn available_agents_for_stage_ignores_completed_agents() { async fn available_agents_for_stage_ignores_completed_agents() {
let config = make_config( let config = make_config(
r#" r#"
[[agent]] [[agent]]
@@ -174,6 +202,7 @@ stage = "coder"
let available = pool let available = pool
.available_agents_for_stage(&config, &PipelineStage::Coder) .available_agents_for_stage(&config, &PipelineStage::Coder)
.await
.unwrap(); .unwrap();
assert_eq!(available, vec!["coder-1"]); assert_eq!(available, vec!["coder-1"]);
} }
+27 -24
View File
@@ -3,8 +3,10 @@
use crate::agent_log::AgentLogWriter; use crate::agent_log::AgentLogWriter;
use crate::config::ProjectConfig; use crate::config::ProjectConfig;
use crate::slog_error; use crate::slog_error;
use std::future::Future;
use std::path::Path; use std::path::Path;
use std::sync::{Arc, Mutex}; use std::pin::Pin;
use std::sync::Arc;
use tokio::sync::broadcast; use tokio::sync::broadcast;
use super::super::runtime::{ use super::super::runtime::{
@@ -38,48 +40,48 @@ impl AgentPool {
/// `resume_context` (if any) is sent as the new message. This lets /// `resume_context` (if any) is sent as the new message. This lets
/// the agent re-enter the previous conversation without re-reading /// the agent re-enter the previous conversation without re-reading
/// CLAUDE.md and README, satisfying story 543. /// CLAUDE.md and README, satisfying story 543.
pub async fn start_agent( pub fn start_agent<'a>(
&self, &'a self,
project_root: &Path, project_root: &'a Path,
story_id: &str, story_id: &'a str,
agent_name: Option<&str>, agent_name: Option<&'a str>,
resume_context: Option<&str>, resume_context: Option<&'a str>,
session_id_to_resume: Option<String>, session_id_to_resume: Option<String>,
) -> Result<AgentInfo, String> { ) -> Pin<Box<dyn Future<Output = Result<AgentInfo, String>> + Send + 'a>> {
self.start_agent_inner( Box::pin(self.start_agent_inner(
project_root, project_root,
story_id, story_id,
agent_name, agent_name,
resume_context, resume_context,
session_id_to_resume, session_id_to_resume,
None, None,
) ))
} }
/// Start an agent with an `AppContext` for direct MCP tool dispatch. /// Start an agent with an `AppContext` for direct MCP tool dispatch.
/// ///
/// API-based runtimes (Gemini, OpenAI) need the `AppContext` to invoke MCP /// API-based runtimes (Gemini, OpenAI) need the `AppContext` to invoke MCP
/// tools without an HTTP round-trip. CLI-based runtimes (Claude Code) do not. /// tools without an HTTP round-trip. CLI-based runtimes (Claude Code) do not.
pub fn start_agent_with_ctx( pub fn start_agent_with_ctx<'a>(
&self, &'a self,
project_root: &Path, project_root: &'a Path,
story_id: &str, story_id: &'a str,
agent_name: Option<&str>, agent_name: Option<&'a str>,
resume_context: Option<&str>, resume_context: Option<&'a str>,
session_id_to_resume: Option<String>, session_id_to_resume: Option<String>,
app_ctx: Arc<crate::http::context::AppContext>, app_ctx: Arc<crate::http::context::AppContext>,
) -> Result<AgentInfo, String> { ) -> Pin<Box<dyn Future<Output = Result<AgentInfo, String>> + Send + 'a>> {
self.start_agent_inner( Box::pin(self.start_agent_inner(
project_root, project_root,
story_id, story_id,
agent_name, agent_name,
resume_context, resume_context,
session_id_to_resume, session_id_to_resume,
Some(app_ctx), Some(app_ctx),
) ))
} }
fn start_agent_inner( async fn start_agent_inner(
&self, &self,
project_root: &Path, project_root: &Path,
story_id: &str, story_id: &str,
@@ -100,7 +102,8 @@ impl AgentPool {
// Create name-independent shared resources before the lock so they are // Create name-independent shared resources before the lock so they are
// ready for the atomic check-and-insert (story 132). // ready for the atomic check-and-insert (story 132).
let (tx, _) = broadcast::channel::<AgentEvent>(1024); let (tx, _) = broadcast::channel::<AgentEvent>(1024);
let event_log: Arc<Mutex<Vec<AgentEvent>>> = Arc::new(Mutex::new(Vec::new())); let event_log: Arc<std::sync::Mutex<Vec<AgentEvent>>> =
Arc::new(std::sync::Mutex::new(Vec::new()));
let log_session_id = uuid::Uuid::new_v4().to_string(); let log_session_id = uuid::Uuid::new_v4().to_string();
// Create the per-session status buffer subscribed to this project's // Create the per-session status buffer subscribed to this project's
@@ -149,7 +152,7 @@ impl AgentPool {
// agent turn (story 736). // agent turn (story 736).
let prior_events: Option<String>; let prior_events: Option<String>;
{ {
let mut agents = self.agents.lock().map_err(|e| e.to_string())?; let mut agents = self.agents.lock().await;
resolved_name = match agent_name { resolved_name = match agent_name {
Some(name) => name.to_string(), Some(name) => name.to_string(),
@@ -371,7 +374,7 @@ impl AgentPool {
// the atomic resolution above). // the atomic resolution above).
let log_writer = let log_writer =
match AgentLogWriter::new(project_root, story_id, &resolved_name, &log_session_id) { match AgentLogWriter::new(project_root, story_id, &resolved_name, &log_session_id) {
Ok(w) => Some(Arc::new(Mutex::new(w))), Ok(w) => Some(Arc::new(std::sync::Mutex::new(w))),
Err(e) => { Err(e) => {
eprintln!( eprintln!(
"[agents] Failed to create log writer for {story_id}:{resolved_name}: {e}" "[agents] Failed to create log writer for {story_id}:{resolved_name}: {e}"
@@ -436,7 +439,7 @@ impl AgentPool {
// Store the task handle while the agent is still Pending. // Store the task handle while the agent is still Pending.
{ {
let mut agents = self.agents.lock().map_err(|e| e.to_string())?; let mut agents = self.agents.lock().await;
if let Some(agent) = agents.get_mut(&key) { if let Some(agent) = agents.get_mut(&key) {
agent.task_handle = Some(handle); agent.task_handle = Some(handle);
} }
+50 -49
View File
@@ -146,7 +146,7 @@ pub(super) async fn run_agent_spawn(
story_id: String, story_id: String,
agent_name: String, agent_name: String,
tx: broadcast::Sender<AgentEvent>, tx: broadcast::Sender<AgentEvent>,
agents: Arc<Mutex<HashMap<String, StoryAgent>>>, agents: Arc<tokio::sync::Mutex<HashMap<String, StoryAgent>>>,
key: String, key: String,
event_log: Arc<Mutex<Vec<AgentEvent>>>, event_log: Arc<Mutex<Vec<AgentEvent>>>,
port: u16, port: u16,
@@ -218,11 +218,12 @@ pub(super) async fn run_agent_spawn(
log.push(event.clone()); log.push(event.clone());
} }
let _ = tx_clone.send(event); let _ = tx_clone.send(event);
if let Ok(mut agents) = agents_ref.lock()
&& let Some(agent) = agents.get_mut(&key_clone)
{ {
let mut agents = agents_ref.lock().await;
if let Some(agent) = agents.get_mut(&key_clone) {
agent.status = AgentStatus::Failed; agent.status = AgentStatus::Failed;
} }
}
AgentPool::notify_agent_state_changed(&watcher_tx_clone); AgentPool::notify_agent_state_changed(&watcher_tx_clone);
return; return;
} }
@@ -233,16 +234,27 @@ pub(super) async fn run_agent_spawn(
// Step 1.1: Install the pre-commit quality-gate hook in the worktree. // Step 1.1: Install the pre-commit quality-gate hook in the worktree.
// Non-fatal — if installation fails the agent can still run; the failure // Non-fatal — if installation fails the agent can still run; the failure
// is logged so the operator can investigate. // is logged so the operator can investigate.
if let Err(e) = crate::worktree::install_pre_commit_hook(&wt_info.path) { // Runs in spawn_blocking because install_pre_commit_hook executes
// synchronous git-config subprocesses that would otherwise pin a
// tokio worker thread and contribute to runtime starvation under
// concurrent agent spawns.
{
let hook_path = wt_info.path.clone();
let hook_result = tokio::task::spawn_blocking(move || {
crate::worktree::install_pre_commit_hook(&hook_path)
})
.await
.unwrap_or_else(|e| Err(format!("spawn_blocking panicked: {e}")));
if let Err(e) = hook_result {
slog_error!("[agents] pre-commit hook install failed for {sid}: {e}"); slog_error!("[agents] pre-commit hook install failed for {sid}: {e}");
} }
}
// Step 2: store worktree info and render agent command/args/prompt. // Step 2: store worktree info and render agent command/args/prompt.
let wt_path_str = wt_info.path.to_string_lossy().to_string(); let wt_path_str = wt_info.path.to_string_lossy().to_string();
{ {
if let Ok(mut agents) = agents_ref.lock() let mut agents = agents_ref.lock().await;
&& let Some(agent) = agents.get_mut(&key_clone) if let Some(agent) = agents.get_mut(&key_clone) {
{
agent.worktree_info = Some(wt_info.clone()); agent.worktree_info = Some(wt_info.clone());
} }
} }
@@ -266,11 +278,12 @@ pub(super) async fn run_agent_spawn(
log.push(event.clone()); log.push(event.clone());
} }
let _ = tx_clone.send(event); let _ = tx_clone.send(event);
if let Ok(mut agents) = agents_ref.lock()
&& let Some(agent) = agents.get_mut(&key_clone)
{ {
let mut agents = agents_ref.lock().await;
if let Some(agent) = agents.get_mut(&key_clone) {
agent.status = AgentStatus::Failed; agent.status = AgentStatus::Failed;
} }
}
AgentPool::notify_agent_state_changed(&watcher_tx_clone); AgentPool::notify_agent_state_changed(&watcher_tx_clone);
return; return;
} }
@@ -358,9 +371,8 @@ pub(super) async fn run_agent_spawn(
// Step 3: transition to Running now that the worktree is ready. // Step 3: transition to Running now that the worktree is ready.
{ {
if let Ok(mut agents) = agents_ref.lock() let mut agents = agents_ref.lock().await;
&& let Some(agent) = agents.get_mut(&key_clone) if let Some(agent) = agents.get_mut(&key_clone) {
{
agent.status = AgentStatus::Running; agent.status = AgentStatus::Running;
} }
} }
@@ -457,9 +469,9 @@ pub(super) async fn run_agent_spawn(
match run_result { match run_result {
Ok(result) => { Ok(result) => {
// Persist token usage if the agent reported it. // Persist token usage if the agent reported it.
if let Some(ref usage) = result.token_usage if let Some(ref usage) = result.token_usage {
&& let Ok(agents) = agents_ref.lock() let agents = agents_ref.lock().await;
&& let Some(agent) = agents.get(&key_clone) if let Some(agent) = agents.get(&key_clone)
&& let Some(ref pr) = agent.project_root && let Some(ref pr) = agent.project_root
{ {
let model_for_record = config_clone let model_for_record = config_clone
@@ -478,6 +490,7 @@ pub(super) async fn run_agent_spawn(
); );
} }
} }
}
// Persist session_id so respawns can resume prior reasoning. // Persist session_id so respawns can resume prior reasoning.
if let Some(ref sess_id) = result.session_id { if let Some(ref sess_id) = result.session_id {
@@ -526,10 +539,7 @@ pub(super) async fn run_agent_spawn(
// Remove the agent entry from the pool and emit Done so that // Remove the agent entry from the pool and emit Done so that
// any caller blocked on wait_for_agent is unblocked. // any caller blocked on wait_for_agent is unblocked.
let tx_done = { let tx_done = {
let mut lock = match agents_ref.lock() { let mut lock = agents_ref.lock().await;
Ok(a) => a,
Err(_) => return,
};
if let Some(agent) = lock.remove(&key_clone) { if let Some(agent) = lock.remove(&key_clone) {
agent.tx agent.tx
} else { } else {
@@ -608,10 +618,7 @@ pub(super) async fn run_agent_spawn(
if stage == PipelineStage::Mergemaster { if stage == PipelineStage::Mergemaster {
let (tx_done, done_session_id, merge_failure_reported, merge_success_reported) = { let (tx_done, done_session_id, merge_failure_reported, merge_success_reported) = {
let mut lock = match agents_ref.lock() { let mut lock = agents_ref.lock().await;
Ok(a) => a,
Err(_) => return,
};
if let Some(agent) = lock.remove(&key_clone) { if let Some(agent) = lock.remove(&key_clone) {
( (
agent.tx, agent.tx,
@@ -648,15 +655,14 @@ pub(super) async fn run_agent_spawn(
// Do NOT send WorkItem/reassign — story is already Done. // Do NOT send WorkItem/reassign — story is already Done.
// Drain one queued ConflictDetected story now that this // Drain one queued ConflictDetected story now that this
// mergemaster slot is free (story 1044). // mergemaster slot is free (story 1044).
if let Some((candidate_id, candidate_agent)) = let candidate =
crate::config::ProjectConfig::load(&project_root_clone) if let Ok(cfg) = crate::config::ProjectConfig::load(&project_root_clone) {
.ok() let agts = agents_ref.lock().await;
.and_then(|cfg| { pick_queued_conflict_detected(&cfg, &agts, &sid)
agents_ref.lock().ok().as_ref().and_then(|agts| { } else {
pick_queued_conflict_detected(&cfg, agts, &sid) None
}) };
}) if let Some((candidate_id, candidate_agent)) = candidate {
{
slog!( slog!(
"[agents] Mergemaster exit for '{sid}' (success): \ "[agents] Mergemaster exit for '{sid}' (success): \
queued ConflictDetected story '{candidate_id}' found; \ queued ConflictDetected story '{candidate_id}' found; \
@@ -766,17 +772,14 @@ pub(super) async fn run_agent_spawn(
}); });
// Drain one queued ConflictDetected story now that this // Drain one queued ConflictDetected story now that this
// mergemaster slot is free (story 1044). // mergemaster slot is free (story 1044).
if let Some((candidate_id, candidate_agent)) = let candidate =
crate::config::ProjectConfig::load(&project_root_clone) if let Ok(cfg) = crate::config::ProjectConfig::load(&project_root_clone) {
.ok() let agts = agents_ref.lock().await;
.and_then(|cfg| { pick_queued_conflict_detected(&cfg, &agts, &sid)
agents_ref } else {
.lock() None
.ok() };
.as_ref() if let Some((candidate_id, candidate_agent)) = candidate {
.and_then(|agts| pick_queued_conflict_detected(&cfg, agts, &sid))
})
{
slog!( slog!(
"[agents] Mergemaster exit for '{sid}': queued ConflictDetected \ "[agents] Mergemaster exit for '{sid}': queued ConflictDetected \
story '{candidate_id}' found; spawning '{candidate_agent}'." story '{candidate_id}' found; spawning '{candidate_agent}'."
@@ -833,10 +836,7 @@ pub(super) async fn run_agent_spawn(
// Remove agent from the pool and unblock any wait_for_agent callers. // Remove agent from the pool and unblock any wait_for_agent callers.
let tx_done = { let tx_done = {
let mut lock = match agents_ref.lock() { let mut lock = agents_ref.lock().await;
Ok(a) => a,
Err(_) => return,
};
if let Some(agent) = lock.remove(&key_clone) { if let Some(agent) = lock.remove(&key_clone) {
agent.tx agent.tx
} else { } else {
@@ -931,11 +931,12 @@ pub(super) async fn run_agent_spawn(
log.push(event.clone()); log.push(event.clone());
} }
let _ = tx_clone.send(event); let _ = tx_clone.send(event);
if let Ok(mut agents) = agents_ref.lock()
&& let Some(agent) = agents.get_mut(&key_clone)
{ {
let mut agents = agents_ref.lock().await;
if let Some(agent) = agents.get_mut(&key_clone) {
agent.status = AgentStatus::Failed; agent.status = AgentStatus::Failed;
} }
}
AgentPool::notify_agent_state_changed(&watcher_tx_clone); AgentPool::notify_agent_state_changed(&watcher_tx_clone);
} }
} }
@@ -109,7 +109,7 @@ async fn start_agent_cleans_up_pending_entry_on_failure() {
"agent must transition to Failed after worktree creation error" "agent must transition to Failed after worktree creation error"
); );
let agents = pool.agents.lock().unwrap(); let agents = pool.agents.lock().await;
let failed_entry = agents let failed_entry = agents
.values() .values()
.find(|a| a.agent_name == "coder-1" && a.status == AgentStatus::Failed); .find(|a| a.agent_name == "coder-1" && a.status == AgentStatus::Failed);
@@ -121,6 +121,7 @@ async fn start_agent_cleans_up_pending_entry_on_failure() {
let events = pool let events = pool
.drain_events("50_story_test", "coder-1") .drain_events("50_story_test", "coder-1")
.await
.expect("drain_events should succeed"); .expect("drain_events should succeed");
let has_error_event = events.iter().any(|e| matches!(e, AgentEvent::Error { .. })); let has_error_event = events.iter().any(|e| matches!(e, AgentEvent::Error { .. }));
assert!( assert!(
@@ -736,7 +737,7 @@ async fn reconcile_canonical_agents_stops_stale_coder_in_qa_stage() {
let pool = AgentPool::new_test(3099); let pool = AgentPool::new_test(3099);
pool.inject_test_agent("777_story_reconcile", "coder-1", AgentStatus::Running); pool.inject_test_agent("777_story_reconcile", "coder-1", AgentStatus::Running);
let before = pool.list_agents().unwrap(); let before = pool.list_agents().await.unwrap();
assert!( assert!(
before.iter().any(|a| a.agent_name == "coder-1" before.iter().any(|a| a.agent_name == "coder-1"
&& matches!(a.status, AgentStatus::Running | AgentStatus::Pending)), && matches!(a.status, AgentStatus::Running | AgentStatus::Pending)),
@@ -745,7 +746,7 @@ async fn reconcile_canonical_agents_stops_stale_coder_in_qa_stage() {
pool.reconcile_canonical_agents(root).await; pool.reconcile_canonical_agents(root).await;
let after = pool.list_agents().unwrap(); let after = pool.list_agents().await.unwrap();
let still_active = after.iter().any(|a| { let still_active = after.iter().any(|a| {
a.story_id == "777_story_reconcile" a.story_id == "777_story_reconcile"
&& a.agent_name == "coder-1" && a.agent_name == "coder-1"
@@ -786,7 +787,7 @@ async fn reconcile_canonical_agents_leaves_correct_stage_agent_alone() {
pool.reconcile_canonical_agents(root).await; pool.reconcile_canonical_agents(root).await;
let after = pool.list_agents().unwrap(); let after = pool.list_agents().await.unwrap();
let still_active = after.iter().any(|a| { let still_active = after.iter().any(|a| {
a.story_id == "555_story_correct" a.story_id == "555_story_correct"
&& a.agent_name == "coder-1" && a.agent_name == "coder-1"
@@ -851,7 +852,7 @@ async fn regression_1100_stale_coder_blocks_mergemaster_then_reconciler_clears()
pool.reconcile_canonical_agents(root).await; pool.reconcile_canonical_agents(root).await;
// coder-1 must be gone from the active pool. // coder-1 must be gone from the active pool.
let remaining = pool.list_agents().unwrap(); let remaining = pool.list_agents().await.unwrap();
assert!( assert!(
!remaining.iter().any(|a| { !remaining.iter().any(|a| {
a.story_id == "1100_reg" a.story_id == "1100_reg"
+18 -27
View File
@@ -1,7 +1,6 @@
//! Agent stop — terminates a running agent while preserving its worktree. //! Agent stop — terminates a running agent while preserving its worktree.
use crate::process_kill::{pids_matching, sigkill_pids_and_verify}; use crate::process_kill::{pids_matching, sigkill_pids_and_verify};
use crate::slog; use crate::slog;
use crate::slog_error;
use crate::slog_warn; use crate::slog_warn;
use std::path::Path; use std::path::Path;
@@ -40,7 +39,7 @@ impl AgentPool {
// Step 1: snapshot the worktree path (no status mutation yet). // Step 1: snapshot the worktree path (no status mutation yet).
let worktree_info = { let worktree_info = {
let agents = self.agents.lock().map_err(|e| e.to_string())?; let agents = self.agents.lock().await;
let agent = agents let agent = agents
.get(&key) .get(&key)
.ok_or_else(|| format!("No agent '{agent_name}' for story '{story_id}'"))?; .ok_or_else(|| format!("No agent '{agent_name}' for story '{story_id}'"))?;
@@ -71,12 +70,12 @@ impl AgentPool {
"[stop_agent] No worktree path recorded for '{key}'; cannot tree-kill, \ "[stop_agent] No worktree path recorded for '{key}'; cannot tree-kill, \
falling back to portable_pty SIGHUP (likely no-op for claude-code)." falling back to portable_pty SIGHUP (likely no-op for claude-code)."
); );
self.kill_child_for_key(&key); self.kill_child_for_key(&key).await;
} }
// Step 3: now safe to mutate. Status flip and handle abort. // Step 3: now safe to mutate. Status flip and handle abort.
let (task_handle, tx) = { let (task_handle, tx) = {
let mut agents = self.agents.lock().map_err(|e| e.to_string())?; let mut agents = self.agents.lock().await;
let agent = agents let agent = agents
.get_mut(&key) .get_mut(&key)
.ok_or_else(|| format!("No agent '{agent_name}' for story '{story_id}'"))?; .ok_or_else(|| format!("No agent '{agent_name}' for story '{story_id}'"))?;
@@ -107,7 +106,7 @@ impl AgentPool {
// Remove from map. // Remove from map.
{ {
let mut agents = self.agents.lock().map_err(|e| e.to_string())?; let mut agents = self.agents.lock().await;
agents.remove(&key); agents.remove(&key);
} }
@@ -138,9 +137,7 @@ impl AgentPool {
// Snapshot active LLM agents without holding the lock during async stops. // Snapshot active LLM agents without holding the lock during async stops.
let snapshot: Vec<(String, String, PipelineStage)> = { let snapshot: Vec<(String, String, PipelineStage)> = {
let Ok(agents) = self.agents.lock() else { let agents = self.agents.lock().await;
return;
};
agents agents
.iter() .iter()
.filter_map(|(key, a)| { .filter_map(|(key, a)| {
@@ -197,14 +194,8 @@ impl AgentPool {
/// ///
/// Called when a story is archived so that stale entries don't accumulate. /// Called when a story is archived so that stale entries don't accumulate.
/// Returns the number of entries removed. /// Returns the number of entries removed.
pub fn remove_agents_for_story(&self, story_id: &str) -> usize { pub async fn remove_agents_for_story(&self, story_id: &str) -> usize {
let mut agents = match self.agents.lock() { let mut agents = self.agents.lock().await;
Ok(a) => a,
Err(e) => {
slog_error!("[agents] Failed to lock pool for cleanup of '{story_id}': {e}");
return 0;
}
};
let prefix = format!("{story_id}:"); let prefix = format!("{story_id}:");
let keys_to_remove: Vec<String> = agents let keys_to_remove: Vec<String> = agents
.keys() .keys()
@@ -229,30 +220,30 @@ mod tests {
// ── remove_agents_for_story tests ──────────────────────────────────────── // ── remove_agents_for_story tests ────────────────────────────────────────
#[test] #[tokio::test]
fn remove_agents_for_story_removes_all_entries() { async fn remove_agents_for_story_removes_all_entries() {
let pool = AgentPool::new_test(3001); let pool = AgentPool::new_test(3001);
pool.inject_test_agent("story_a", "coder-1", AgentStatus::Completed); pool.inject_test_agent("story_a", "coder-1", AgentStatus::Completed);
pool.inject_test_agent("story_a", "qa", AgentStatus::Failed); pool.inject_test_agent("story_a", "qa", AgentStatus::Failed);
pool.inject_test_agent("story_b", "coder-1", AgentStatus::Running); pool.inject_test_agent("story_b", "coder-1", AgentStatus::Running);
let removed = pool.remove_agents_for_story("story_a"); let removed = pool.remove_agents_for_story("story_a").await;
assert_eq!(removed, 2, "should remove both agents for story_a"); assert_eq!(removed, 2, "should remove both agents for story_a");
let agents = pool.list_agents().unwrap(); let agents = pool.list_agents().await.unwrap();
assert_eq!(agents.len(), 1, "only story_b agent should remain"); assert_eq!(agents.len(), 1, "only story_b agent should remain");
assert_eq!(agents[0].story_id, "story_b"); assert_eq!(agents[0].story_id, "story_b");
} }
#[test] #[tokio::test]
fn remove_agents_for_story_returns_zero_when_no_match() { async fn remove_agents_for_story_returns_zero_when_no_match() {
let pool = AgentPool::new_test(3001); let pool = AgentPool::new_test(3001);
pool.inject_test_agent("story_a", "coder-1", AgentStatus::Running); pool.inject_test_agent("story_a", "coder-1", AgentStatus::Running);
let removed = pool.remove_agents_for_story("nonexistent"); let removed = pool.remove_agents_for_story("nonexistent").await;
assert_eq!(removed, 0); assert_eq!(removed, 0);
let agents = pool.list_agents().unwrap(); let agents = pool.list_agents().await.unwrap();
assert_eq!(agents.len(), 1, "existing agents should not be affected"); assert_eq!(agents.len(), 1, "existing agents should not be affected");
} }
@@ -283,12 +274,12 @@ mod tests {
pool.inject_test_agent("60_story_cleanup", "qa", AgentStatus::Completed); pool.inject_test_agent("60_story_cleanup", "qa", AgentStatus::Completed);
pool.inject_test_agent("61_story_other", "coder-1", AgentStatus::Running); pool.inject_test_agent("61_story_other", "coder-1", AgentStatus::Running);
assert_eq!(pool.list_agents().unwrap().len(), 3); assert_eq!(pool.list_agents().await.unwrap().len(), 3);
move_story_to_done("60_story_cleanup").unwrap(); move_story_to_done("60_story_cleanup").unwrap();
pool.remove_agents_for_story("60_story_cleanup"); pool.remove_agents_for_story("60_story_cleanup").await;
let remaining = pool.list_agents().unwrap(); let remaining = pool.list_agents().await.unwrap();
assert_eq!( assert_eq!(
remaining.len(), remaining.len(),
1, 1,
+8 -8
View File
@@ -20,7 +20,7 @@ impl AgentPool {
) -> broadcast::Sender<AgentEvent> { ) -> broadcast::Sender<AgentEvent> {
let (tx, _) = broadcast::channel::<AgentEvent>(64); let (tx, _) = broadcast::channel::<AgentEvent>(64);
let key = composite_key(story_id, agent_name); let key = composite_key(story_id, agent_name);
let mut agents = self.agents.lock().unwrap(); let mut agents = self.agents.try_lock().unwrap();
agents.insert( agents.insert(
key, key,
StoryAgent { StoryAgent {
@@ -55,7 +55,7 @@ impl AgentPool {
) -> broadcast::Sender<AgentEvent> { ) -> broadcast::Sender<AgentEvent> {
let (tx, _) = broadcast::channel::<AgentEvent>(64); let (tx, _) = broadcast::channel::<AgentEvent>(64);
let key = composite_key(story_id, agent_name); let key = composite_key(story_id, agent_name);
let mut agents = self.agents.lock().unwrap(); let mut agents = self.agents.try_lock().unwrap();
agents.insert( agents.insert(
key, key,
StoryAgent { StoryAgent {
@@ -95,7 +95,7 @@ impl AgentPool {
) -> broadcast::Sender<AgentEvent> { ) -> broadcast::Sender<AgentEvent> {
let (tx, _) = broadcast::channel::<AgentEvent>(64); let (tx, _) = broadcast::channel::<AgentEvent>(64);
let key = composite_key(story_id, agent_name); let key = composite_key(story_id, agent_name);
let mut agents = self.agents.lock().unwrap(); let mut agents = self.agents.try_lock().unwrap();
agents.insert( agents.insert(
key, key,
StoryAgent { StoryAgent {
@@ -130,7 +130,7 @@ impl AgentPool {
) -> broadcast::Sender<AgentEvent> { ) -> broadcast::Sender<AgentEvent> {
let (tx, _) = broadcast::channel::<AgentEvent>(64); let (tx, _) = broadcast::channel::<AgentEvent>(64);
let key = composite_key(story_id, agent_name); let key = composite_key(story_id, agent_name);
let mut agents = self.agents.lock().unwrap(); let mut agents = self.agents.try_lock().unwrap();
agents.insert( agents.insert(
key, key,
StoryAgent { StoryAgent {
@@ -165,7 +165,7 @@ impl AgentPool {
) -> broadcast::Sender<AgentEvent> { ) -> broadcast::Sender<AgentEvent> {
let (tx, _) = broadcast::channel::<AgentEvent>(64); let (tx, _) = broadcast::channel::<AgentEvent>(64);
let key = composite_key(story_id, agent_name); let key = composite_key(story_id, agent_name);
let mut agents = self.agents.lock().unwrap(); let mut agents = self.agents.try_lock().unwrap();
agents.insert( agents.insert(
key, key,
StoryAgent { StoryAgent {
@@ -197,7 +197,7 @@ impl AgentPool {
story_id: &str, story_id: &str,
agent_name: &str, agent_name: &str,
) -> Option<Vec<BufferedItem>> { ) -> Option<Vec<BufferedItem>> {
let agents = self.agents.lock().unwrap(); let agents = self.agents.try_lock().unwrap();
let key = composite_key(story_id, agent_name); let key = composite_key(story_id, agent_name);
agents agents
.get(&key) .get(&key)
@@ -219,7 +219,7 @@ impl AgentPool {
) -> broadcast::Sender<AgentEvent> { ) -> broadcast::Sender<AgentEvent> {
let (tx, _) = broadcast::channel::<AgentEvent>(64); let (tx, _) = broadcast::channel::<AgentEvent>(64);
let key = composite_key(story_id, agent_name); let key = composite_key(story_id, agent_name);
let mut agents = self.agents.lock().unwrap(); let mut agents = self.agents.try_lock().unwrap();
agents.insert( agents.insert(
key, key,
StoryAgent { StoryAgent {
@@ -258,7 +258,7 @@ impl AgentPool {
) -> broadcast::Sender<AgentEvent> { ) -> broadcast::Sender<AgentEvent> {
let (tx, _) = broadcast::channel::<AgentEvent>(64); let (tx, _) = broadcast::channel::<AgentEvent>(64);
let key = composite_key(story_id, agent_name); let key = composite_key(story_id, agent_name);
let mut agents = self.agents.lock().unwrap(); let mut agents = self.agents.try_lock().unwrap();
agents.insert( agents.insert(
key, key,
StoryAgent { StoryAgent {
+7 -5
View File
@@ -3,8 +3,8 @@ use crate::slog;
use crate::worktree::WorktreeInfo; use crate::worktree::WorktreeInfo;
use std::collections::HashMap; use std::collections::HashMap;
use std::path::PathBuf; use std::path::PathBuf;
use std::sync::{Arc, Mutex}; use std::sync::Arc;
use tokio::sync::broadcast; use tokio::sync::{Mutex, broadcast};
use super::super::{AgentEvent, AgentInfo, AgentStatus, CompletionReport}; use super::super::{AgentEvent, AgentInfo, AgentStatus, CompletionReport};
@@ -45,8 +45,10 @@ impl PendingGuard {
impl Drop for PendingGuard { impl Drop for PendingGuard {
fn drop(&mut self) { fn drop(&mut self) {
if self.armed if !self.armed {
&& let Ok(mut agents) = self.agents.lock() return;
}
if let Ok(mut agents) = self.agents.try_lock()
&& agents && agents
.get(&self.key) .get(&self.key)
.is_some_and(|a| a.status == AgentStatus::Pending) .is_some_and(|a| a.status == AgentStatus::Pending)
@@ -68,7 +70,7 @@ pub(super) struct StoryAgent {
pub(super) tx: broadcast::Sender<AgentEvent>, pub(super) tx: broadcast::Sender<AgentEvent>,
pub(super) task_handle: Option<tokio::task::JoinHandle<()>>, pub(super) task_handle: Option<tokio::task::JoinHandle<()>>,
/// Accumulated events for polling via get_agent_output. /// Accumulated events for polling via get_agent_output.
pub(super) event_log: Arc<Mutex<Vec<AgentEvent>>>, pub(super) event_log: Arc<std::sync::Mutex<Vec<AgentEvent>>>,
/// Set when the agent calls report_completion. /// Set when the agent calls report_completion.
pub(super) completion: Option<CompletionReport>, pub(super) completion: Option<CompletionReport>,
/// Project root, stored for pipeline advancement after completion. /// Project root, stored for pipeline advancement after completion.
+5 -5
View File
@@ -17,11 +17,11 @@ impl AgentPool {
) -> Result<AgentInfo, String> { ) -> Result<AgentInfo, String> {
// Subscribe before checking status so we don't miss the terminal event // Subscribe before checking status so we don't miss the terminal event
// if the agent completes in the window between the two operations. // if the agent completes in the window between the two operations.
let mut rx = self.subscribe(story_id, agent_name)?; let mut rx = self.subscribe(story_id, agent_name).await?;
// Return immediately if already in a terminal state. // Return immediately if already in a terminal state.
{ {
let agents = self.agents.lock().map_err(|e| e.to_string())?; let agents = self.agents.lock().await;
let key = composite_key(story_id, agent_name); let key = composite_key(story_id, agent_name);
if let Some(agent) = agents.get(&key) if let Some(agent) = agents.get(&key)
&& matches!(agent.status, AgentStatus::Completed | AgentStatus::Failed) && matches!(agent.status, AgentStatus::Completed | AgentStatus::Failed)
@@ -48,7 +48,7 @@ impl AgentPool {
_ => false, _ => false,
}; };
if is_terminal { if is_terminal {
let agents = self.agents.lock().map_err(|e| e.to_string())?; let agents = self.agents.lock().await;
let key = composite_key(story_id, agent_name); let key = composite_key(story_id, agent_name);
return Ok(if let Some(agent) = agents.get(&key) { return Ok(if let Some(agent) = agents.get(&key) {
agent_info_from_entry(story_id, agent) agent_info_from_entry(story_id, agent)
@@ -78,7 +78,7 @@ impl AgentPool {
} }
Ok(Err(broadcast::error::RecvError::Lagged(_))) => { Ok(Err(broadcast::error::RecvError::Lagged(_))) => {
// Missed some buffered events — check current status before resuming. // Missed some buffered events — check current status before resuming.
let agents = self.agents.lock().map_err(|e| e.to_string())?; let agents = self.agents.lock().await;
let key = composite_key(story_id, agent_name); let key = composite_key(story_id, agent_name);
if let Some(agent) = agents.get(&key) if let Some(agent) = agents.get(&key)
&& matches!(agent.status, AgentStatus::Completed | AgentStatus::Failed) && matches!(agent.status, AgentStatus::Completed | AgentStatus::Failed)
@@ -89,7 +89,7 @@ impl AgentPool {
} }
Ok(Err(broadcast::error::RecvError::Closed)) => { Ok(Err(broadcast::error::RecvError::Closed)) => {
// Channel closed: no more events will arrive. Return current state. // Channel closed: no more events will arrive. Return current state.
let agents = self.agents.lock().map_err(|e| e.to_string())?; let agents = self.agents.lock().await;
let key = composite_key(story_id, agent_name); let key = composite_key(story_id, agent_name);
if let Some(agent) = agents.get(&key) { if let Some(agent) = agents.get(&key) {
return Ok(agent_info_from_entry(story_id, agent)); return Ok(agent_info_from_entry(story_id, agent));
+1 -1
View File
@@ -90,7 +90,7 @@ pub(crate) fn build_status_from_items(
items: &[PipelineItem], items: &[PipelineItem],
) -> String { ) -> String {
// Build a map from story_id → active AgentInfo for quick lookup. // Build a map from story_id → active AgentInfo for quick lookup.
let active_agents = agents.list_agents().unwrap_or_default(); let active_agents = agents.list_agents_nonblocking();
let active_map: HashMap<String, &crate::agents::AgentInfo> = active_agents let active_map: HashMap<String, &crate::agents::AgentInfo> = active_agents
.iter() .iter()
.filter(|a| matches!(a.status, AgentStatus::Running | AgentStatus::Pending)) .filter(|a| matches!(a.status, AgentStatus::Running | AgentStatus::Pending))
@@ -111,6 +111,7 @@ pub async fn handle_assign(
// Check whether a coder is already running on this story. // Check whether a coder is already running on this story.
let running_coders: Vec<_> = agents let running_coders: Vec<_> = agents
.list_agents() .list_agents()
.await
.unwrap_or_default() .unwrap_or_default()
.into_iter() .into_iter()
.filter(|a| { .filter(|a| {
+1 -1
View File
@@ -191,7 +191,7 @@ pub fn build_htop_message(agents: &AgentPool, tick: u32, total_duration_secs: u6
String::new(), String::new(),
]; ];
let all_agents = agents.list_agents().unwrap_or_default(); let all_agents = agents.list_agents_nonblocking();
let active: Vec<_> = all_agents let active: Vec<_> = all_agents
.iter() .iter()
.filter(|a| matches!(a.status, AgentStatus::Running | AgentStatus::Pending)) .filter(|a| matches!(a.status, AgentStatus::Running | AgentStatus::Pending))
@@ -85,6 +85,7 @@ pub async fn handle_rmtree(
// Stop any running or pending agents for this story. // Stop any running or pending agents for this story.
let running_agents: Vec<(String, String)> = agents let running_agents: Vec<(String, String)> = agents
.list_agents() .list_agents()
.await
.unwrap_or_default() .unwrap_or_default()
.into_iter() .into_iter()
.filter(|a| { .filter(|a| {
+1 -1
View File
@@ -18,7 +18,7 @@ pub async fn agent_stream(
Path((story_id, agent_name)): Path<(String, String)>, Path((story_id, agent_name)): Path<(String, String)>,
ctx: Data<&Arc<AppContext>>, ctx: Data<&Arc<AppContext>>,
) -> impl IntoResponse { ) -> impl IntoResponse {
let mut rx = match ctx.services.agents.subscribe(&story_id, &agent_name) { let mut rx = match ctx.services.agents.subscribe(&story_id, &agent_name).await {
Ok(rx) => rx, Ok(rx) => rx,
Err(e) => { Err(e) => {
return Response::builder() return Response::builder()
+26 -21
View File
@@ -54,7 +54,7 @@ pub(crate) async fn tool_get_agent_output(
// writer failed and nothing was persisted to disk. // writer failed and nothing was persisted to disk.
if log_files.is_empty() if log_files.is_empty()
&& let Some(agent_name) = agent_name_filter && let Some(agent_name) = agent_name_filter
&& let Ok(live_events) = ctx.services.agents.drain_events(story_id, agent_name) && let Ok(live_events) = ctx.services.agents.drain_events(story_id, agent_name).await
&& !live_events.is_empty() && !live_events.is_empty()
{ {
all_lines.push(format!("=== {agent_name} (live) ===")); all_lines.push(format!("=== {agent_name} (live) ==="));
@@ -99,7 +99,7 @@ pub(crate) async fn tool_get_agent_output(
Ok(output) Ok(output)
} }
pub(crate) fn tool_get_agent_config(ctx: &AppContext) -> Result<String, String> { pub(crate) async fn tool_get_agent_config(ctx: &AppContext) -> Result<String, String> {
let project_root = ctx.services.agents.get_project_root(&ctx.state)?; let project_root = ctx.services.agents.get_project_root(&ctx.state)?;
let config = ProjectConfig::load(&project_root)?; let config = ProjectConfig::load(&project_root)?;
@@ -116,6 +116,7 @@ pub(crate) fn tool_get_agent_config(ctx: &AppContext) -> Result<String, String>
.services .services
.agents .agents
.available_agents_for_stage(&config, stage) .available_agents_for_stage(&config, stage)
.await
{ {
available_names.extend(names); available_names.extend(names);
} }
@@ -144,7 +145,7 @@ pub(crate) fn tool_get_agent_config(ctx: &AppContext) -> Result<String, String>
/// Returns turns used, max turns, remaining turns, budget used, max budget, /// Returns turns used, max turns, remaining turns, budget used, max budget,
/// and remaining budget for the named agent. Fails if the agent is not /// and remaining budget for the named agent. Fails if the agent is not
/// currently running or pending. /// currently running or pending.
pub(crate) fn tool_get_agent_remaining_turns_and_budget( pub(crate) async fn tool_get_agent_remaining_turns_and_budget(
args: &Value, args: &Value,
ctx: &AppContext, ctx: &AppContext,
) -> Result<String, String> { ) -> Result<String, String> {
@@ -158,7 +159,7 @@ pub(crate) fn tool_get_agent_remaining_turns_and_budget(
.ok_or("Missing required argument: agent_name")?; .ok_or("Missing required argument: agent_name")?;
// Verify the agent exists and is running/pending. // Verify the agent exists and is running/pending.
let agents = ctx.services.agents.list_agents()?; let agents = ctx.services.agents.list_agents().await?;
let agent_info = agents let agent_info = agents
.iter() .iter()
.find(|a| a.story_id == story_id && a.agent_name == agent_name) .find(|a| a.story_id == story_id && a.agent_name == agent_name)
@@ -270,12 +271,12 @@ mod tests {
use crate::http::test_helpers::test_ctx; use crate::http::test_helpers::test_ctx;
use serde_json::json; use serde_json::json;
#[test] #[tokio::test]
fn tool_get_agent_config_no_project_toml_returns_default_agent() { async fn tool_get_agent_config_no_project_toml_returns_default_agent() {
let tmp = tempfile::tempdir().unwrap(); let tmp = tempfile::tempdir().unwrap();
let ctx = test_ctx(tmp.path()); let ctx = test_ctx(tmp.path());
// No project.toml → default config with one fallback agent // No project.toml → default config with one fallback agent
let result = tool_get_agent_config(&ctx).unwrap(); let result = tool_get_agent_config(&ctx).await.unwrap();
let parsed: Vec<Value> = serde_json::from_str(&result).unwrap(); let parsed: Vec<Value> = serde_json::from_str(&result).unwrap();
// Default config contains one agent entry with default values // Default config contains one agent entry with default values
assert_eq!( assert_eq!(
@@ -457,34 +458,36 @@ mod tests {
// ── get_agent_remaining_turns_and_budget tests ────────────────────────── // ── get_agent_remaining_turns_and_budget tests ──────────────────────────
#[test] #[tokio::test]
fn tool_get_agent_remaining_turns_and_budget_missing_story_id() { async fn tool_get_agent_remaining_turns_and_budget_missing_story_id() {
let tmp = tempfile::tempdir().unwrap(); let tmp = tempfile::tempdir().unwrap();
let ctx = test_ctx(tmp.path()); let ctx = test_ctx(tmp.path());
let result = let result =
tool_get_agent_remaining_turns_and_budget(&json!({"agent_name": "coder-1"}), &ctx); tool_get_agent_remaining_turns_and_budget(&json!({"agent_name": "coder-1"}), &ctx)
.await;
assert!(result.is_err()); assert!(result.is_err());
assert!(result.unwrap_err().contains("story_id")); assert!(result.unwrap_err().contains("story_id"));
} }
#[test] #[tokio::test]
fn tool_get_agent_remaining_turns_and_budget_missing_agent_name() { async fn tool_get_agent_remaining_turns_and_budget_missing_agent_name() {
let tmp = tempfile::tempdir().unwrap(); let tmp = tempfile::tempdir().unwrap();
let ctx = test_ctx(tmp.path()); let ctx = test_ctx(tmp.path());
let result = let result =
tool_get_agent_remaining_turns_and_budget(&json!({"story_id": "1_test"}), &ctx); tool_get_agent_remaining_turns_and_budget(&json!({"story_id": "1_test"}), &ctx).await;
assert!(result.is_err()); assert!(result.is_err());
assert!(result.unwrap_err().contains("agent_name")); assert!(result.unwrap_err().contains("agent_name"));
} }
#[test] #[tokio::test]
fn tool_get_agent_remaining_turns_and_budget_no_agent_returns_error() { async fn tool_get_agent_remaining_turns_and_budget_no_agent_returns_error() {
let tmp = tempfile::tempdir().unwrap(); let tmp = tempfile::tempdir().unwrap();
let ctx = test_ctx(tmp.path()); let ctx = test_ctx(tmp.path());
let result = tool_get_agent_remaining_turns_and_budget( let result = tool_get_agent_remaining_turns_and_budget(
&json!({"story_id": "99_nope", "agent_name": "coder-1"}), &json!({"story_id": "99_nope", "agent_name": "coder-1"}),
&ctx, &ctx,
); )
.await;
assert!(result.is_err()); assert!(result.is_err());
let err = result.unwrap_err(); let err = result.unwrap_err();
assert!( assert!(
@@ -493,8 +496,8 @@ mod tests {
); );
} }
#[test] #[tokio::test]
fn tool_get_agent_remaining_turns_and_budget_completed_agent_returns_error() { async fn tool_get_agent_remaining_turns_and_budget_completed_agent_returns_error() {
use crate::agents::AgentStatus; use crate::agents::AgentStatus;
let tmp = tempfile::tempdir().unwrap(); let tmp = tempfile::tempdir().unwrap();
let ctx = test_ctx(tmp.path()); let ctx = test_ctx(tmp.path());
@@ -505,7 +508,8 @@ mod tests {
let result = tool_get_agent_remaining_turns_and_budget( let result = tool_get_agent_remaining_turns_and_budget(
&json!({"story_id": "42_story", "agent_name": "coder-1"}), &json!({"story_id": "42_story", "agent_name": "coder-1"}),
&ctx, &ctx,
); )
.await;
assert!(result.is_err()); assert!(result.is_err());
let err = result.unwrap_err(); let err = result.unwrap_err();
assert!( assert!(
@@ -514,8 +518,8 @@ mod tests {
); );
} }
#[test] #[tokio::test]
fn tool_get_agent_remaining_turns_and_budget_running_agent_returns_data() { async fn tool_get_agent_remaining_turns_and_budget_running_agent_returns_data() {
use crate::agents::AgentStatus; use crate::agents::AgentStatus;
use crate::store::StoreOps; use crate::store::StoreOps;
@@ -531,6 +535,7 @@ mod tests {
&json!({"story_id": "42_story", "agent_name": "coder-1"}), &json!({"story_id": "42_story", "agent_name": "coder-1"}),
&ctx, &ctx,
) )
.await
.unwrap(); .unwrap();
let parsed: Value = serde_json::from_str(&result).unwrap(); let parsed: Value = serde_json::from_str(&result).unwrap();
+5 -5
View File
@@ -67,9 +67,9 @@ pub(crate) async fn tool_stop_agent(args: &Value, ctx: &AppContext) -> Result<St
)) ))
} }
pub(crate) fn tool_list_agents(ctx: &AppContext) -> Result<String, String> { pub(crate) async fn tool_list_agents(ctx: &AppContext) -> Result<String, String> {
let project_root = ctx.services.agents.get_project_root(&ctx.state).ok(); let project_root = ctx.services.agents.get_project_root(&ctx.state).ok();
let agents = ctx.services.agents.list_agents()?; let agents = ctx.services.agents.list_agents().await?;
let mut entries: Vec<serde_json::Value> = agents let mut entries: Vec<serde_json::Value> = agents
.iter() .iter()
.filter(|a| { .filter(|a| {
@@ -156,11 +156,11 @@ mod tests {
use crate::http::test_helpers::test_ctx; use crate::http::test_helpers::test_ctx;
use serde_json::json; use serde_json::json;
#[test] #[tokio::test]
fn tool_list_agents_empty() { async fn tool_list_agents_empty() {
let tmp = tempfile::tempdir().unwrap(); let tmp = tempfile::tempdir().unwrap();
let ctx = test_ctx(tmp.path()); let ctx = test_ctx(tmp.path());
let result = tool_list_agents(&ctx).unwrap(); let result = tool_list_agents(&ctx).await.unwrap();
let parsed: Vec<serde_json::Value> = serde_json::from_str(&result).unwrap(); let parsed: Vec<serde_json::Value> = serde_json::from_str(&result).unwrap();
assert!(parsed.is_empty()); assert!(parsed.is_empty());
} }
@@ -362,8 +362,8 @@ mod tests {
// then exec() will be called — which would replace our test process. // then exec() will be called — which would replace our test process.
// So we only test that the function *runs* without panicking up to // So we only test that the function *runs* without panicking up to
// the agent-kill step. We do this by checking the pool is empty. // the agent-kill step. We do this by checking the pool is empty.
assert_eq!(ctx.services.agents.list_agents().unwrap().len(), 0); assert_eq!(ctx.services.agents.list_agents().await.unwrap().len(), 0);
ctx.services.agents.kill_all_children(); // should not panic on empty pool ctx.services.agents.kill_all_children().await; // should not panic on empty pool
} }
#[test] #[test]
+7 -7
View File
@@ -30,13 +30,13 @@ pub async fn dispatch_tool_call(
// Agent tools (async) // Agent tools (async)
"start_agent" => agent_tools::tool_start_agent(&args, ctx).await, "start_agent" => agent_tools::tool_start_agent(&args, ctx).await,
"stop_agent" => agent_tools::tool_stop_agent(&args, ctx).await, "stop_agent" => agent_tools::tool_stop_agent(&args, ctx).await,
"list_agents" => agent_tools::tool_list_agents(ctx), "list_agents" => agent_tools::tool_list_agents(ctx).await,
"get_agent_config" => agent_tools::tool_get_agent_config(ctx), "get_agent_config" => agent_tools::tool_get_agent_config(ctx).await,
"reload_agent_config" => agent_tools::tool_get_agent_config(ctx), "reload_agent_config" => agent_tools::tool_get_agent_config(ctx).await,
"get_agent_output" => agent_tools::tool_get_agent_output(&args, ctx).await, "get_agent_output" => agent_tools::tool_get_agent_output(&args, ctx).await,
"wait_for_agent" => agent_tools::tool_wait_for_agent(&args, ctx).await, "wait_for_agent" => agent_tools::tool_wait_for_agent(&args, ctx).await,
"get_agent_remaining_turns_and_budget" => { "get_agent_remaining_turns_and_budget" => {
agent_tools::tool_get_agent_remaining_turns_and_budget(&args, ctx) agent_tools::tool_get_agent_remaining_turns_and_budget(&args, ctx).await
} }
// Worktree tools // Worktree tools
"create_worktree" => agent_tools::tool_create_worktree(&args, ctx).await, "create_worktree" => agent_tools::tool_create_worktree(&args, ctx).await,
@@ -46,7 +46,7 @@ pub async fn dispatch_tool_call(
// Editor tools // Editor tools
"get_editor_command" => agent_tools::tool_get_editor_command(&args, ctx), "get_editor_command" => agent_tools::tool_get_editor_command(&args, ctx),
// Lifecycle tools // Lifecycle tools
"accept_story" => story_tools::tool_accept_story(&args, ctx), "accept_story" => story_tools::tool_accept_story(&args, ctx).await,
// Story mutation tools (auto-commit to master) // Story mutation tools (auto-commit to master)
"check_criterion" => story_tools::tool_check_criterion(&args, ctx), "check_criterion" => story_tools::tool_check_criterion(&args, ctx),
"edit_criterion" => story_tools::tool_edit_criterion(&args, ctx), "edit_criterion" => story_tools::tool_edit_criterion(&args, ctx),
@@ -58,7 +58,7 @@ pub async fn dispatch_tool_call(
// Bug lifecycle tools // Bug lifecycle tools
"create_bug" => story_tools::tool_create_bug(&args, ctx), "create_bug" => story_tools::tool_create_bug(&args, ctx),
"list_bugs" => story_tools::tool_list_bugs(ctx), "list_bugs" => story_tools::tool_list_bugs(ctx),
"close_bug" => story_tools::tool_close_bug(&args, ctx), "close_bug" => story_tools::tool_close_bug(&args, ctx).await,
// Refactor lifecycle tools // Refactor lifecycle tools
"create_refactor" => story_tools::tool_create_refactor(&args, ctx), "create_refactor" => story_tools::tool_create_refactor(&args, ctx),
"list_refactors" => story_tools::tool_list_refactors(ctx), "list_refactors" => story_tools::tool_list_refactors(ctx),
@@ -70,7 +70,7 @@ pub async fn dispatch_tool_call(
"merge_agent_work" => merge_tools::tool_merge_agent_work(&args, ctx).await, "merge_agent_work" => merge_tools::tool_merge_agent_work(&args, ctx).await,
"get_merge_status" => merge_tools::tool_get_merge_status(&args, ctx), "get_merge_status" => merge_tools::tool_get_merge_status(&args, ctx),
"move_story_to_merge" => merge_tools::tool_move_story_to_merge(&args, ctx).await, "move_story_to_merge" => merge_tools::tool_move_story_to_merge(&args, ctx).await,
"report_merge_failure" => merge_tools::tool_report_merge_failure(&args, ctx), "report_merge_failure" => merge_tools::tool_report_merge_failure(&args, ctx).await,
// QA tools // QA tools
"request_qa" => qa_tools::tool_request_qa(&args, ctx).await, "request_qa" => qa_tools::tool_request_qa(&args, ctx).await,
"approve_qa" => qa_tools::tool_approve_qa(&args, ctx).await, "approve_qa" => qa_tools::tool_approve_qa(&args, ctx).await,
+18 -11
View File
@@ -171,7 +171,10 @@ pub(super) async fn tool_move_story_to_merge(
.map_err(|e| format!("Serialization error: {e}")) .map_err(|e| format!("Serialization error: {e}"))
} }
pub(super) fn tool_report_merge_failure(args: &Value, ctx: &AppContext) -> Result<String, String> { pub(super) async fn tool_report_merge_failure(
args: &Value,
ctx: &AppContext,
) -> Result<String, String> {
let story_id = args let story_id = args
.get("story_id") .get("story_id")
.and_then(|v| v.as_str()) .and_then(|v| v.as_str())
@@ -182,7 +185,10 @@ pub(super) fn tool_report_merge_failure(args: &Value, ctx: &AppContext) -> Resul
.ok_or("Missing required argument: reason")?; .ok_or("Missing required argument: reason")?;
slog!("[mergemaster] Merge failure reported for '{story_id}': {reason}"); slog!("[mergemaster] Merge failure reported for '{story_id}': {reason}");
ctx.services.agents.set_merge_failure_reported(story_id); ctx.services
.agents
.set_merge_failure_reported(story_id)
.await;
// The mergemaster provides a freeform reason string; use Other so the // The mergemaster provides a freeform reason string; use Other so the
// auto-assigner does not re-spawn another mergemaster after this one fails. // auto-assigner does not re-spawn another mergemaster after this one fails.
@@ -412,26 +418,26 @@ mod tests {
assert!(req_names.contains(&"reason")); assert!(req_names.contains(&"reason"));
} }
#[test] #[tokio::test]
fn tool_report_merge_failure_missing_story_id() { async fn tool_report_merge_failure_missing_story_id() {
let tmp = tempfile::tempdir().unwrap(); let tmp = tempfile::tempdir().unwrap();
let ctx = test_ctx(tmp.path()); let ctx = test_ctx(tmp.path());
let result = tool_report_merge_failure(&json!({"reason": "conflicts"}), &ctx); let result = tool_report_merge_failure(&json!({"reason": "conflicts"}), &ctx).await;
assert!(result.is_err()); assert!(result.is_err());
assert!(result.unwrap_err().contains("story_id")); assert!(result.unwrap_err().contains("story_id"));
} }
#[test] #[tokio::test]
fn tool_report_merge_failure_missing_reason() { async fn tool_report_merge_failure_missing_reason() {
let tmp = tempfile::tempdir().unwrap(); let tmp = tempfile::tempdir().unwrap();
let ctx = test_ctx(tmp.path()); let ctx = test_ctx(tmp.path());
let result = tool_report_merge_failure(&json!({"story_id": "42_story_foo"}), &ctx); let result = tool_report_merge_failure(&json!({"story_id": "42_story_foo"}), &ctx).await;
assert!(result.is_err()); assert!(result.is_err());
assert!(result.unwrap_err().contains("reason")); assert!(result.unwrap_err().contains("reason"));
} }
#[test] #[tokio::test]
fn tool_report_merge_failure_returns_confirmation() { async fn tool_report_merge_failure_returns_confirmation() {
let tmp = tempfile::tempdir().unwrap(); let tmp = tempfile::tempdir().unwrap();
let ctx = test_ctx(tmp.path()); let ctx = test_ctx(tmp.path());
let result = tool_report_merge_failure( let result = tool_report_merge_failure(
@@ -440,7 +446,8 @@ mod tests {
"reason": "Unresolvable merge conflicts in src/main.rs" "reason": "Unresolvable merge conflicts in src/main.rs"
}), }),
&ctx, &ctx,
); )
.await;
assert!(result.is_ok()); assert!(result.is_ok());
let msg = result.unwrap(); let msg = result.unwrap();
assert!(msg.contains("42_story_foo")); assert!(msg.contains("42_story_foo"));
+1 -1
View File
@@ -81,7 +81,7 @@ pub(super) async fn tool_approve_qa(args: &Value, ctx: &AppContext) -> Result<St
move_story_to_done(story_id)?; move_story_to_done(story_id)?;
let pool = std::sync::Arc::clone(&ctx.services.agents); let pool = std::sync::Arc::clone(&ctx.services.agents);
pool.remove_agents_for_story(story_id); pool.remove_agents_for_story(story_id).await;
let wt_path = crate::worktree::worktree_path(&project_root, story_id); let wt_path = crate::worktree::worktree_path(&project_root, story_id);
if wt_path.exists() { if wt_path.exists() {
+10 -8
View File
@@ -66,14 +66,14 @@ pub(crate) fn tool_list_bugs(ctx: &AppContext) -> Result<String, String> {
.map_err(|e| format!("Serialization error: {e}")) .map_err(|e| format!("Serialization error: {e}"))
} }
pub(crate) fn tool_close_bug(args: &Value, ctx: &AppContext) -> Result<String, String> { pub(crate) async fn tool_close_bug(args: &Value, ctx: &AppContext) -> Result<String, String> {
let bug_id = args let bug_id = args
.get("bug_id") .get("bug_id")
.and_then(|v| v.as_str()) .and_then(|v| v.as_str())
.ok_or("Missing required argument: bug_id")?; .ok_or("Missing required argument: bug_id")?;
close_bug_to_archive(bug_id)?; close_bug_to_archive(bug_id)?;
ctx.services.agents.remove_agents_for_story(bug_id); ctx.services.agents.remove_agents_for_story(bug_id).await;
Ok(format!( Ok(format!(
"Bug '{bug_id}' closed, moved to bugs/archive/, and committed to master." "Bug '{bug_id}' closed, moved to bugs/archive/, and committed to master."
@@ -422,17 +422,17 @@ mod tests {
assert!(result.unwrap().contains("Created bug:")); assert!(result.unwrap().contains("Created bug:"));
} }
#[test] #[tokio::test]
fn tool_close_bug_missing_bug_id() { async fn tool_close_bug_missing_bug_id() {
let tmp = tempfile::tempdir().unwrap(); let tmp = tempfile::tempdir().unwrap();
let ctx = test_ctx(tmp.path()); let ctx = test_ctx(tmp.path());
let result = tool_close_bug(&json!({}), &ctx); let result = tool_close_bug(&json!({}), &ctx).await;
assert!(result.is_err()); assert!(result.is_err());
assert!(result.unwrap_err().contains("bug_id")); assert!(result.unwrap_err().contains("bug_id"));
} }
#[test] #[tokio::test]
fn tool_close_bug_moves_to_archive() { async fn tool_close_bug_moves_to_archive() {
let tmp = tempfile::tempdir().unwrap(); let tmp = tempfile::tempdir().unwrap();
setup_git_repo_in(tmp.path()); setup_git_repo_in(tmp.path());
let backlog_dir = tmp.path().join(".huskies/work/1_backlog"); let backlog_dir = tmp.path().join(".huskies/work/1_backlog");
@@ -460,7 +460,9 @@ mod tests {
.unwrap(); .unwrap();
let ctx = test_ctx(tmp.path()); let ctx = test_ctx(tmp.path());
let result = tool_close_bug(&json!({"bug_id": "9901_bug_crash"}), &ctx).unwrap(); let result = tool_close_bug(&json!({"bug_id": "9901_bug_crash"}), &ctx)
.await
.unwrap();
assert!(result.contains("9901_bug_crash")); assert!(result.contains("9901_bug_crash"));
assert!( assert!(
crate::db::read_content(crate::db::ContentKey::Story("9901_bug_crash")).is_some(), crate::db::read_content(crate::db::ContentKey::Story("9901_bug_crash")).is_some(),
+14 -14
View File
@@ -5,7 +5,7 @@ use crate::http::context::AppContext;
use crate::pipeline_state::{Stage, read_typed}; use crate::pipeline_state::{Stage, read_typed};
use serde_json::Value; use serde_json::Value;
pub(crate) fn tool_accept_story(args: &Value, ctx: &AppContext) -> Result<String, String> { pub(crate) async fn tool_accept_story(args: &Value, ctx: &AppContext) -> Result<String, String> {
let story_id = args let story_id = args
.get("story_id") .get("story_id")
.and_then(|v| v.as_str()) .and_then(|v| v.as_str())
@@ -33,7 +33,7 @@ pub(crate) fn tool_accept_story(args: &Value, ctx: &AppContext) -> Result<String
} }
move_story_to_done(story_id)?; move_story_to_done(story_id)?;
ctx.services.agents.remove_agents_for_story(story_id); ctx.services.agents.remove_agents_for_story(story_id).await;
Ok(format!( Ok(format!(
"Story '{story_id}' accepted, moved to done/, and committed to master." "Story '{story_id}' accepted, moved to done/, and committed to master."
@@ -146,27 +146,27 @@ mod tests {
assert!(!story_file.exists(), "story file should be deleted"); assert!(!story_file.exists(), "story file should be deleted");
} }
#[test] #[tokio::test]
fn tool_accept_story_missing_story_id() { async fn tool_accept_story_missing_story_id() {
let tmp = tempfile::tempdir().unwrap(); let tmp = tempfile::tempdir().unwrap();
let ctx = test_ctx(tmp.path()); let ctx = test_ctx(tmp.path());
let result = tool_accept_story(&json!({}), &ctx); let result = tool_accept_story(&json!({}), &ctx).await;
assert!(result.is_err()); assert!(result.is_err());
assert!(result.unwrap_err().contains("story_id")); assert!(result.unwrap_err().contains("story_id"));
} }
#[test] #[tokio::test]
fn tool_accept_story_nonexistent_story_returns_error() { async fn tool_accept_story_nonexistent_story_returns_error() {
let tmp = tempfile::tempdir().unwrap(); let tmp = tempfile::tempdir().unwrap();
setup_git_repo_in(tmp.path()); setup_git_repo_in(tmp.path());
let ctx = test_ctx(tmp.path()); let ctx = test_ctx(tmp.path());
// No story file in current/ — should fail // No story file in current/ — should fail
let result = tool_accept_story(&json!({"story_id": "99_nonexistent"}), &ctx); let result = tool_accept_story(&json!({"story_id": "99_nonexistent"}), &ctx).await;
assert!(result.is_err()); assert!(result.is_err());
} }
#[test] #[tokio::test]
fn tool_accept_story_refuses_when_feature_branch_has_unmerged_code() { async fn tool_accept_story_refuses_when_feature_branch_has_unmerged_code() {
let tmp = tempfile::tempdir().unwrap(); let tmp = tempfile::tempdir().unwrap();
setup_git_repo_in(tmp.path()); setup_git_repo_in(tmp.path());
@@ -203,7 +203,7 @@ mod tests {
.unwrap(); .unwrap();
let ctx = test_ctx(tmp.path()); let ctx = test_ctx(tmp.path());
let result = tool_accept_story(&json!({"story_id": "50_story_test"}), &ctx); let result = tool_accept_story(&json!({"story_id": "50_story_test"}), &ctx).await;
assert!( assert!(
result.is_err(), result.is_err(),
"should refuse when feature branch has unmerged code" "should refuse when feature branch has unmerged code"
@@ -215,8 +215,8 @@ mod tests {
); );
} }
#[test] #[tokio::test]
fn tool_accept_story_succeeds_when_no_feature_branch() { async fn tool_accept_story_succeeds_when_no_feature_branch() {
let tmp = tempfile::tempdir().unwrap(); let tmp = tempfile::tempdir().unwrap();
setup_git_repo_in(tmp.path()); setup_git_repo_in(tmp.path());
@@ -234,7 +234,7 @@ mod tests {
); );
let ctx = test_ctx(tmp.path()); let ctx = test_ctx(tmp.path());
let result = tool_accept_story(&json!({"story_id": "51_story_no_branch"}), &ctx); let result = tool_accept_story(&json!({"story_id": "51_story_no_branch"}), &ctx).await;
assert!( assert!(
result.is_ok(), result.is_ok(),
"should succeed when no feature branch: {result:?}" "should succeed when no feature branch: {result:?}"
+1 -4
View File
@@ -230,10 +230,7 @@ pub fn load_pipeline_state(ctx: &AppContext) -> Result<PipelineState, String> {
/// Build a map from story_id → AgentAssignment for all pending/running agents. /// Build a map from story_id → AgentAssignment for all pending/running agents.
fn build_active_agent_map(ctx: &AppContext) -> HashMap<String, AgentAssignment> { fn build_active_agent_map(ctx: &AppContext) -> HashMap<String, AgentAssignment> {
let agents = match ctx.services.agents.list_agents() { let agents = ctx.services.agents.list_agents_nonblocking();
Ok(a) => a,
Err(_) => return HashMap::new(),
};
let config_opt = ctx let config_opt = ctx
.state .state
+1 -1
View File
@@ -492,7 +492,7 @@ async fn main() -> Result<(), std::io::Error> {
tokio::time::sleep(std::time::Duration::from_millis(1500)).await; tokio::time::sleep(std::time::Duration::from_millis(1500)).await;
// Kill all active PTY child processes before exiting. // Kill all active PTY child processes before exiting.
agents_for_shutdown.kill_all_children(); agents_for_shutdown.kill_all_children().await;
if let Some(ref path) = port_file { if let Some(ref path) = port_file {
remove_port_file(path); remove_port_file(path);
+2 -1
View File
@@ -95,6 +95,7 @@ pub async fn rebuild_and_restart(
// 1. Gracefully stop all running agents. // 1. Gracefully stop all running agents.
let running_count = agents let running_count = agents
.list_agents() .list_agents()
.await
.unwrap_or_default() .unwrap_or_default()
.iter() .iter()
.filter(|a| a.status == crate::agents::AgentStatus::Running) .filter(|a| a.status == crate::agents::AgentStatus::Running)
@@ -102,7 +103,7 @@ pub async fn rebuild_and_restart(
if running_count > 0 { if running_count > 0 {
slog!("[rebuild] Stopping {running_count} running agent(s) before rebuild"); slog!("[rebuild] Stopping {running_count} running agent(s) before rebuild");
} }
agents.kill_all_children(); agents.kill_all_children().await;
// 2. Find the workspace root (parent of the server binary's source). // 2. Find the workspace root (parent of the server binary's source).
// CARGO_MANIFEST_DIR at compile time points to the `server/` crate; // CARGO_MANIFEST_DIR at compile time points to the `server/` crate;
+3 -2
View File
@@ -95,7 +95,7 @@ pub async fn delete_work_item(
// 1. Stop any running/pending agents (best-effort). // 1. Stop any running/pending agents (best-effort).
let mut agents_stopped: Vec<String> = Vec::new(); let mut agents_stopped: Vec<String> = Vec::new();
if let Ok(agent_list) = agents.list_agents() { if let Ok(agent_list) = agents.list_agents().await {
for agent in agent_list.iter().filter(|a| a.story_id == story_id) { for agent in agent_list.iter().filter(|a| a.story_id == story_id) {
match agents match agents
.stop_agent(project_root, story_id, &agent.agent_name) .stop_agent(project_root, story_id, &agent.agent_name)
@@ -120,7 +120,7 @@ pub async fn delete_work_item(
} }
// 2. Remove agent pool entries. // 2. Remove agent pool entries.
let removed_count = agents.remove_agents_for_story(story_id); let removed_count = agents.remove_agents_for_story(story_id).await;
slog_warn!("[delete_work_item] Removed {removed_count} agent pool entries for '{story_id}'"); slog_warn!("[delete_work_item] Removed {removed_count} agent pool entries for '{story_id}'");
// 3. Remove worktree (best-effort). // 3. Remove worktree (best-effort).
@@ -240,6 +240,7 @@ mod tests {
// Agent pool must have no entries for this story. // Agent pool must have no entries for this story.
let pool_entries = agents let pool_entries = agents
.list_agents() .list_agents()
.await
.unwrap_or_default() .unwrap_or_default()
.into_iter() .into_iter()
.filter(|a| a.story_id == story_id) .filter(|a| a.story_id == story_id)
+1 -1
View File
@@ -198,7 +198,7 @@ pub(crate) fn spawn_tick_loop(
// timeout window). A `TransitionFired` subscriber cannot observe the // timeout window). A `TransitionFired` subscriber cannot observe the
// absence of events, so this must remain on a periodic tick. // absence of events, so this must remain on a periodic tick.
if tick_count.is_multiple_of(30) { if tick_count.is_multiple_of(30) {
let found = agents.run_watchdog_pass(root.as_deref()); let found = agents.run_watchdog_pass(root.as_deref()).await;
if found > 0 { if found > 0 {
crate::slog!( crate::slog!(
"[tick] {found} orphaned agent(s) detected; triggering auto-assign." "[tick] {found} orphaned agent(s) detected; triggering auto-assign."