crashlabs/huskies
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 37 Wiki Activity

Converted all external tool calling to async

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
server/src/agents/pool/auto_assign/auto_assign.rs
+10 -10
View File
@@ -73,7 +73,7 @@ mod tests {
// task eventually fails.
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| {
a.agent_name == "coder-1"
&& matches!(a.status, AgentStatus::Pending | AgentStatus::Running)
@@ -115,7 +115,7 @@ mod tests {
pool.auto_assign_available_work(root).await;
// No agent should have been started for the spike.
let agents = pool.agents.lock().unwrap();
let agents = pool.agents.try_lock().unwrap();
assert!(
agents.is_empty(),
"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;
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).
let coder_assigned_to_qa = agents.iter().any(|(key, a)| {
key.contains("9930_story_qa1")
@@ -209,7 +209,7 @@ mod tests {
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).
let coder1_assigned = agents.values().any(|a| {
a.agent_name == "coder-1"
@@ -262,7 +262,7 @@ mod tests {
// Must not panic.
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
// be assigned to leaked 2_current items from the global CRDT store).
let assigned_to_qa_story = agents.iter().any(|(key, a)| {
@@ -301,7 +301,7 @@ mod tests {
let pool = AgentPool::new_test(3001);
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
// interference from other tests sharing the global CRDT store.
let assigned_to_our_story = agents.iter().any(|(key, a)| {
@@ -347,7 +347,7 @@ mod tests {
let pool = AgentPool::new_test(3001);
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| {
matches!(
a.status,
@@ -553,7 +553,7 @@ mod tests {
let _ = tokio::join!(t1, t2);
// 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
.values()
.filter(|a| {
@@ -602,7 +602,7 @@ mod tests {
pool.auto_assign_available_work(tmp.path()).await;
let count_after_first = {
let agents = pool.agents.lock().unwrap();
let agents = pool.agents.try_lock().unwrap();
agents
.iter()
.filter(|(key, a)| {
@@ -616,7 +616,7 @@ mod tests {
pool.auto_assign_available_work(tmp.path()).await;
let count_after_second = {
let agents = pool.agents.lock().unwrap();
let agents = pool.agents.try_lock().unwrap();
agents
.iter()
.filter(|(key, a)| {
server/src/agents/pool/auto_assign/merge.rs
+1 -7
View File
@@ -99,13 +99,7 @@ impl AgentPool {
// Skip if an explicit mergemaster LLM agent is already running
// (operator-driven failure recovery path).
let has_mergemaster = {
let agents = match self.agents.lock() {
Ok(a) => a,
Err(e) => {
slog_error!("[auto-assign] Failed to lock agents: {e}");
break;
}
};
let agents = self.agents.lock().await;
is_story_assigned_for_stage(config, &agents, story_id, &PipelineStage::Mergemaster)
};
if has_mergemaster {
server/src/agents/pool/auto_assign/merge_failure_block_subscriber.rs
+3 -6
View File
@@ -54,7 +54,7 @@ pub(crate) fn spawn_merge_failure_block_subscriber(pool: Arc<AgentPool>, project
match rx.recv().await {
Ok(fired) => {
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);
}
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`.
/// Used to suppress counter increments while recovery is actively iterating
/// (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) {
Ok(c) => c,
Err(_) => return false,
};
let agents = match pool.agents.lock() {
Ok(a) => a,
Err(_) => return false,
};
let agents = pool.agents.lock().await;
is_story_assigned_for_stage(&config, &agents, story_id, &PipelineStage::Mergemaster)
}
server/src/agents/pool/auto_assign/merge_failure_subscriber.rs
+8 -16
View File
@@ -100,15 +100,7 @@ async fn on_merge_failure_transition(
};
let agent_name = {
let agents = match pool.agents.lock() {
Ok(a) => a,
Err(e) => {
slog_warn!(
"[merge-failure-sub] Failed to lock agent pool for '{story_id}': {e}"
);
return;
}
};
let agents = pool.agents.lock().await;
if is_story_assigned_for_stage(
&config,
&agents,
@@ -228,7 +220,7 @@ mod tests {
);
on_merge_failure_transition(&pool, tmp.path(), &fired).await;
let agents = pool.agents.lock().unwrap();
let agents = pool.agents.lock().await;
assert!(
agents.iter().any(|(key, a)| {
key.contains(story_id)
@@ -259,7 +251,7 @@ mod tests {
// Give the subscriber time to run (it should do nothing).
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)| {
key.contains(story_id)
&& a.agent_name == "mergemaster"
@@ -287,7 +279,7 @@ mod tests {
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)| {
key.contains(story_id)
&& a.agent_name == "mergemaster"
@@ -315,7 +307,7 @@ mod tests {
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)| {
key.contains(story_id)
&& a.agent_name == "mergemaster"
@@ -343,7 +335,7 @@ mod tests {
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)| {
key.contains(story_id)
&& a.agent_name == "mergemaster"
@@ -374,7 +366,7 @@ mod tests {
// First call — spawns mergemaster (agent enters Pending).
on_merge_failure_transition(&pool, tmp.path(), &fired).await;
{
let agents = pool.agents.lock().unwrap();
let agents = pool.agents.lock().await;
assert!(
agents.iter().any(|(key, a)| {
key.contains(story_id)
@@ -388,7 +380,7 @@ mod tests {
// Second call (self-loop) — agent is still Pending; guard must prevent double-spawn.
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
.iter()
.filter(|(key, a)| {
server/src/agents/pool/auto_assign/pipeline.rs
+2 -15
View File
@@ -5,7 +5,6 @@ use std::path::Path;
use crate::config::ProjectConfig;
use crate::pipeline_state::Stage;
use crate::slog;
use crate::slog_error;
use super::super::super::PipelineStage;
use super::super::AgentPool;
@@ -80,13 +79,7 @@ impl AgentPool {
if *stage == PipelineStage::Coder
&& let Some(max) = config.max_coders
{
let agents_lock = match self.agents.lock() {
Ok(a) => a,
Err(e) => {
slog_error!("[auto-assign] Failed to lock agents: {e}");
break;
}
};
let agents_lock = self.agents.lock().await;
let active = count_active_agents_for_stage(config, &agents_lock, stage);
if active >= max {
slog!(
@@ -102,13 +95,7 @@ impl AgentPool {
// stage_mismatch=true means the preferred agent's stage doesn't match the
// pipeline stage, so we fell back to a generic stage agent.
let (already_assigned, free_agent, preferred_busy, stage_mismatch) = {
let agents = match self.agents.lock() {
Ok(a) => a,
Err(e) => {
slog_error!("[auto-assign] Failed to lock agents: {e}");
break;
}
};
let agents = self.agents.lock().await;
let assigned = is_story_assigned_for_stage(config, &agents, story_id, stage);
if assigned {
(true, None, false, false)
server/src/agents/pool/auto_assign/scan.rs
+6 -6
View File
@@ -256,7 +256,7 @@ mod tests {
let pool = AgentPool::new_test(3001);
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(
&config,
&agents,
@@ -285,7 +285,7 @@ mod tests {
let pool = AgentPool::new_test(3001);
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
assert!(!is_story_assigned_for_stage(
&config,
@@ -309,7 +309,7 @@ stage = "qa"
let pool = AgentPool::new_test(3001);
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
assert!(
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("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);
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
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);
assert_eq!(
free,
@@ -384,7 +384,7 @@ name = "coder-1"
// coder-1 completed its previous story — it's free for a new one
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);
assert_eq!(free, Some("coder-1"), "completed coder-1 should be free");
}
server/src/agents/pool/auto_assign/watchdog/limits.rs
+3 -6
View File
@@ -2,7 +2,7 @@
use std::collections::HashMap;
use std::path::Path;
use std::sync::Mutex;
use tokio::sync::Mutex;
use tokio::sync::broadcast;
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
/// only — prior sessions are excluded so that restart counts from earlier
/// 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>>,
project_root: &Path,
) -> 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).
let running: Vec<RunningAgentSnapshot> = {
let lock = match agents.lock() {
Ok(l) => l,
Err(_) => return Vec::new(),
};
let lock = agents.lock().await;
lock.iter()
.filter(|(_, agent)| agent.status == AgentStatus::Running)
.map(|(key, agent)| {
server/src/agents/pool/auto_assign/watchdog/mod.rs
+19 -17
View File
@@ -25,8 +25,8 @@ pub(crate) use limits::{count_turns_in_log, resolve_session_log};
impl AgentPool {
/// Run a single watchdog pass synchronously (test helper).
#[cfg(test)]
pub fn run_watchdog_once(&self) {
check_orphaned_agents(&self.agents);
pub async fn run_watchdog_once(&self) {
check_orphaned_agents(&self.agents).await;
}
/// 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
/// re-attempt. This prevents the original kill-respawn loop (bug 646)
/// while restoring the `max_retries` semantic for turn/budget overruns.
pub fn run_watchdog_pass(&self, project_root: Option<&Path>) -> usize {
let orphaned = check_orphaned_agents(&self.agents);
pub async fn run_watchdog_pass(&self, project_root: Option<&Path>) -> usize {
let orphaned = check_orphaned_agents(&self.agents).await;
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();
for (key, reason) in &terminated {
// Step 1: snapshot the agent's worktree path so we can find every
// process running in it (claude + any subprocesses). This must
// happen BEFORE we mutate the agent record so we can read the
// 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)
.and_then(|a| a.worktree_info.as_ref().map(|wt| wt.path.clone()))
});
};
// Step 2: SIGKILL every process running in the worktree and
// BLOCK until verified gone. The previous mechanism — portable_pty's
@@ -85,23 +86,24 @@ impl AgentPool {
"[watchdog] No worktree path recorded for '{key}'; cannot tree-kill, \
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
// gone (or we logged that SIGKILL didn't take effect, which is
// exceptional), so flipping status away from Running can no
// longer open a window for a concurrent spawn.
if let Ok(mut lock) = self.agents.lock()
&& let Some(agent) = lock.get_mut(key)
{
agent.status = AgentStatus::Failed;
agent.termination_reason = Some(reason.clone());
if let Some(handle) = agent.task_handle.take() {
// Best-effort abort of the outer tokio task. The PTY
// blocking thread already returned (claude is dead),
// so this is bookkeeping rather than load-bearing.
handle.abort();
let mut lock = self.agents.lock().await;
if let Some(agent) = lock.get_mut(key) {
agent.status = AgentStatus::Failed;
agent.termination_reason = Some(reason.clone());
if let Some(handle) = agent.task_handle.take() {
// Best-effort abort of the outer tokio task. The PTY
// blocking thread already returned (claude is dead),
// so this is bookkeeping rather than load-bearing.
handle.abort();
}
}
}
server/src/agents/pool/auto_assign/watchdog/orphan.rs
+3 -6
View File
@@ -1,7 +1,7 @@
//! Orphan detection: marks running agents whose backing task has exited.
use std::collections::HashMap;
use std::sync::Mutex;
use tokio::sync::Mutex;
use tokio::sync::broadcast;
use crate::agents::pool::StoryAgent;
@@ -15,11 +15,8 @@ use crate::slog;
/// without updating the agent status — for example when the process is killed
/// 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.
pub(super) fn check_orphaned_agents(agents: &Mutex<HashMap<String, StoryAgent>>) -> usize {
let mut lock = match agents.lock() {
Ok(l) => l,
Err(_) => return 0,
};
pub(super) async fn check_orphaned_agents(agents: &Mutex<HashMap<String, StoryAgent>>) -> usize {
let mut lock = agents.lock().await;
// Collect orphaned entries: Running or Pending agents whose task handle is finished.
// Pending agents can be orphaned if worktree creation panics before setting status.
server/src/agents/pool/auto_assign/watchdog/tests/limits_tests.rs
+41 -41
View File
@@ -10,8 +10,8 @@ use crate::agents::{AgentEvent, AgentStatus, TerminationReason};
// ── Limit enforcement integration tests (bug 624) ────────────────────────
#[test]
fn watchdog_terminates_agent_exceeding_turn_limit() {
#[tokio::test]
async fn watchdog_terminates_agent_exceeding_turn_limit() {
let tmp = tempfile::tempdir().unwrap();
let root = tmp.path();
@@ -37,12 +37,12 @@ max_turns = 10
);
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");
// 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 agent = agents.get(&key).unwrap();
assert_eq!(agent.status, AgentStatus::Failed);
@@ -60,8 +60,8 @@ max_turns = 10
);
}
#[test]
fn watchdog_terminates_agent_exceeding_budget_limit() {
#[tokio::test]
async fn watchdog_terminates_agent_exceeding_budget_limit() {
let tmp = tempfile::tempdir().unwrap();
let root = tmp.path();
@@ -87,11 +87,11 @@ max_budget_usd = 5.00
);
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");
{
let agents = pool.agents.lock().unwrap();
let agents = pool.agents.try_lock().unwrap();
let key = composite_key("story_b", "coder-1");
let agent = agents.get(&key).unwrap();
assert_eq!(agent.status, AgentStatus::Failed);
@@ -106,8 +106,8 @@ max_budget_usd = 5.00
assert!(matches!(event, AgentEvent::Error { .. }));
}
#[test]
fn watchdog_does_not_terminate_agent_under_limits() {
#[tokio::test]
async fn watchdog_does_not_terminate_agent_under_limits() {
let tmp = tempfile::tempdir().unwrap();
let root = tmp.path();
@@ -133,11 +133,11 @@ max_budget_usd = 10.00
// has 25 turns < 50 so no violation).
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");
{
let agents = pool.agents.lock().unwrap();
let agents = pool.agents.try_lock().unwrap();
let key = composite_key("story_c", "coder-1");
let agent = agents.get(&key).unwrap();
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
/// limit (280 turns). The watchdog must terminate at the turn limit (turns
/// hit first in the observed trace), with reason TurnLimit.
#[test]
fn regression_bug624_coder1_story623_trajectory() {
#[tokio::test]
async fn regression_bug624_coder1_story623_trajectory() {
let tmp = tempfile::tempdir().unwrap();
let root = tmp.path();
@@ -183,11 +183,11 @@ max_budget_usd = 5.00
);
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");
{
let agents = pool.agents.lock().unwrap();
let agents = pool.agents.try_lock().unwrap();
let key = composite_key("story_623", "coder-1");
let agent = agents.get(&key).unwrap();
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
/// and uses `max_retries = 1` so that the first violation blocks.
#[test]
fn watchdog_marks_story_blocked_after_limit_termination() {
#[tokio::test]
async fn watchdog_marks_story_blocked_after_limit_termination() {
crate::db::ensure_content_store();
crate::crdt_state::init_for_test();
@@ -263,7 +263,7 @@ max_turns = 10
"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");
// 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.
{
let agents = pool.agents.lock().unwrap();
let agents = pool.agents.try_lock().unwrap();
let key = composite_key(story_id, "coder-1");
let agent = agents.get(&key).unwrap();
assert_eq!(agent.status, AgentStatus::Failed);
@@ -297,8 +297,8 @@ max_turns = 10
/// fresh session_id whose log has fewer events than `max_turns`.
/// Assert the agent is NOT terminated (per-session count is under the
/// limit) AND the story is NOT marked blocked.
#[test]
fn per_session_counting_does_not_terminate_under_limit() {
#[tokio::test]
async fn per_session_counting_does_not_terminate_under_limit() {
let tmp = tempfile::tempdir().unwrap();
let root = tmp.path();
@@ -323,14 +323,14 @@ max_turns = 10
let pool = AgentPool::new_test(3001);
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!(
found, 0,
"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 agent = agents.get(&key).unwrap();
assert_eq!(
@@ -345,8 +345,8 @@ max_turns = 10
/// 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
/// IS terminated AND (with max_retries=1) the story IS marked blocked.
#[test]
fn per_session_counting_terminates_over_limit() {
#[tokio::test]
async fn per_session_counting_terminates_over_limit() {
crate::db::ensure_content_store();
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");
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,
"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 agent = agents.get(&key).unwrap();
assert_eq!(agent.status, AgentStatus::Failed);
@@ -423,8 +423,8 @@ max_turns = 10
/// `max_turns`. After session 1: retry_count=1, NOT blocked. After
/// session 2: retry_count=2, NOT blocked. After session 3:
/// retry_count=3 >= max_retries, story IS blocked.
#[test]
fn watchdog_retry_semantic_blocks_after_max_retries() {
#[tokio::test]
async fn watchdog_retry_semantic_blocks_after_max_retries() {
crate::db::ensure_content_store();
let tmp = tempfile::tempdir().unwrap();
@@ -453,7 +453,7 @@ max_turns = 10
write_fake_session_log(root, story_id, "coder-1", "session-1", 12);
let pool = AgentPool::new_test(3001);
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");
assert_eq!(
@@ -473,7 +473,7 @@ max_turns = 10
write_fake_session_log(root, story_id, "coder-1", "session-2", 12);
let pool = AgentPool::new_test(3001);
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");
assert_eq!(
@@ -493,7 +493,7 @@ max_turns = 10
write_fake_session_log(root, story_id, "coder-1", "session-3", 12);
let pool = AgentPool::new_test(3001);
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");
assert_eq!(
@@ -518,8 +518,8 @@ max_turns = 10
/// 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
/// agent with max_tool_turns = 10 stays Running.
#[test]
fn watchdog_does_not_count_narration_only_turns() {
#[tokio::test]
async fn watchdog_does_not_count_narration_only_turns() {
let tmp = tempfile::tempdir().unwrap();
let root = tmp.path();
@@ -542,13 +542,13 @@ max_turns = 200
let pool = AgentPool::new_test(3001);
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!(
found, 0,
"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 agent = agents.get(&key).unwrap();
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
/// 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.
#[test]
fn watchdog_max_tool_turns_overrides_max_turns() {
#[tokio::test]
async fn watchdog_max_tool_turns_overrides_max_turns() {
let tmp = tempfile::tempdir().unwrap();
let root = tmp.path();
@@ -585,13 +585,13 @@ max_turns = 200
);
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 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 agent = agents.get(&key).unwrap();
assert_eq!(agent.status, AgentStatus::Failed);
server/src/agents/pool/auto_assign/watchdog/tests/orphan_tests.rs
+9 -9
View File
@@ -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_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");
}
#[test]
fn check_orphaned_agents_returns_zero_when_no_orphans() {
#[tokio::test]
async fn check_orphaned_agents_returns_zero_when_no_orphans() {
let pool = AgentPool::new_test(3001);
// Inject agents in terminal states — not orphaned.
pool.inject_test_agent("story_a", "coder", AgentStatus::Completed);
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!(
found, 0,
"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);
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 agent = agents.get(&key).unwrap();
assert_eq!(
@@ -87,13 +87,13 @@ async fn watchdog_orphan_detection_returns_nonzero_enabling_auto_assign() {
// 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");
assert_eq!(agents.get(&key).unwrap().status, AgentStatus::Running);
}
// 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!(
found, 1,
"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.
{
let agents = pool.agents.lock().unwrap();
let agents = pool.agents.try_lock().unwrap();
let key = composite_key("orphan_story", "coder");
assert_eq!(
agents.get(&key).unwrap().status,
server/src/agents/pool/mod.rs
+7 -5
View File
@@ -19,8 +19,8 @@ mod test_helpers;
use crate::io::watcher::WatcherEvent;
use crate::service::status::StatusBroadcaster;
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use tokio::sync::broadcast;
use std::sync::Arc;
use tokio::sync::{Mutex, broadcast};
// Bring pool-internal types into pool's namespace so that sub-modules
// (auto_assign, pipeline, etc.) can access them via `use super::...`.
@@ -87,10 +87,12 @@ impl AgentPool {
_ => continue,
};
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);
}
server/src/agents/pool/pipeline/advance/helpers.rs
+2 -2
View File
@@ -3,7 +3,7 @@
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::{Arc, Mutex};
use std::sync::Arc;
use tokio::sync::broadcast;
@@ -16,7 +16,7 @@ use std::path::Path;
/// type cycle between `start_agent` and `run_server_owned_completion`.
#[allow(clippy::too_many_arguments)]
pub(crate) fn spawn_pipeline_advance(
agents: Arc<Mutex<HashMap<String, StoryAgent>>>,
agents: Arc<tokio::sync::Mutex<HashMap<String, StoryAgent>>>,
port: u16,
story_id: &str,
agent_name: &str,
server/src/agents/pool/pipeline/advance/mod.rs
+1 -1
View File
@@ -694,7 +694,7 @@ impl AgentPool {
if let Err(e) = crate::agents::lifecycle::move_story_to_done(story_id) {
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);
// TODO: Re-enable worktree cleanup once we have persistent agent logs.
// Removing worktrees destroys evidence needed to debug empty-commit agents.
server/src/agents/pool/pipeline/advance/tests_regression.rs
+8 -8
View File
@@ -104,7 +104,7 @@ async fn mergemaster_blocks_and_sends_story_blocked_when_no_commits_ahead() {
);
// No mergemaster agent should have been started.
let agents = pool.agents.lock().unwrap();
let agents = pool.agents.try_lock().unwrap();
let mergemaster_started = agents
.values()
.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).
{
let agents = pool.agents.lock().unwrap();
let agents = pool.agents.try_lock().unwrap();
assert!(
!is_agent_free(&agents, "qa"),
"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
// (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"));
}
@@ -193,7 +193,7 @@ stage = "qa"
.await;
// 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| {
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;
// No agents should have been started.
let agents = pool.agents.lock().unwrap();
let agents = pool.agents.try_lock().unwrap();
assert!(
agents.is_empty(),
"No agents should be started for a stale advance on a done story. \
@@ -871,7 +871,7 @@ stage = "coder"
.await;
// 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| {
a.agent_name == "coder-1" && matches!(a.status, AgentStatus::Pending | AgentStatus::Running)
});
@@ -957,7 +957,7 @@ stage = "coder"
.await;
// 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| {
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.
let agents = pool.agents.lock().unwrap();
let agents = pool.agents.try_lock().unwrap();
let qa_or_merge_started = agents
.values()
.any(|a| a.agent_name.contains("qa") || a.agent_name.contains("merge"));
server/src/agents/pool/pipeline/completion/legacy.rs
+2 -2
View File
@@ -28,7 +28,7 @@ impl AgentPool {
// Verify agent exists, is Running, and grab its 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
.get(&key)
.ok_or_else(|| format!("No agent '{agent_name}' for story '{story_id}'"))?;
@@ -82,7 +82,7 @@ impl AgentPool {
merge_failure_reported_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(|| {
format!("Agent '{agent_name}' for story '{story_id}' disappeared during gate check")
})?;
server/src/agents/pool/pipeline/completion/server.rs
+5 -13
View File
@@ -2,7 +2,8 @@
use crate::io::watcher::WatcherEvent;
use crate::slog;
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::sync::Arc;
use tokio::sync::Mutex;
use tokio::sync::broadcast;
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).
{
let lock = match agents.lock() {
Ok(a) => a,
Err(_) => return,
};
let lock = agents.lock().await;
match lock.get(&key) {
Some(agent) if agent.completion.is_some() => {
slog!(
@@ -64,10 +62,7 @@ pub(in crate::agents::pool) async fn run_server_owned_completion(
// Get worktree path for running gates.
let worktree_path = {
let lock = match agents.lock() {
Ok(a) => a,
Err(_) => return,
};
let lock = agents.lock().await;
lock.get(&key)
.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
// 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 mut lock = match agents.lock() {
Ok(a) => a,
Err(_) => return,
};
let mut lock = agents.lock().await;
let agent = match lock.get_mut(&key) {
Some(a) => a,
None => return,
server/src/agents/pool/pipeline/completion/tests.rs
+8 -8
View File
@@ -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.
let mut rx = pool.subscribe("s10", "coder-1").unwrap();
let mut rx = pool.subscribe("s10", "coder-1").await.unwrap();
run_server_owned_completion(
&pool.agents,
@@ -121,7 +121,7 @@ async fn server_owned_completion_skips_when_already_completed() {
.await;
// 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 agent = agents.get(&key).unwrap();
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);
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(
&pool.agents,
@@ -160,7 +160,7 @@ async fn server_owned_completion_runs_gates_on_clean_worktree() {
.await;
// 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");
assert!(
agents.get(&key).is_none(),
@@ -192,7 +192,7 @@ async fn server_owned_completion_fails_on_dirty_worktree() {
let pool = AgentPool::new_test(3001);
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(
&pool.agents,
@@ -205,7 +205,7 @@ async fn server_owned_completion_fails_on_dirty_worktree() {
.await;
// 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");
assert!(
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
// 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");
assert!(
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);
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(
&pool.agents,
server/src/agents/pool/pipeline/merge/control.rs
+44 -56
View File
@@ -34,35 +34,29 @@ impl AgentPool {
/// 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
/// by the caller acts as the authoritative fallback in that case.
pub fn set_merge_success_reported(&self, story_id: &str) {
match self.agents.lock() {
Ok(mut lock) => {
let found = lock.iter_mut().find(|(key, agent)| {
let key_story_id = key
.rsplit_once(':')
.map(|(sid, _)| sid)
.unwrap_or(key.as_str());
key_story_id == story_id
&& pipeline_stage(&agent.agent_name) == PipelineStage::Mergemaster
});
match found {
Some((_, agent)) => {
agent.merge_success_reported = true;
slog!(
"[pipeline] Merge success flag set for '{story_id}:{}'",
agent.agent_name
);
}
None => {
slog!(
"[pipeline] set_merge_success_reported: no running mergemaster \
for '{story_id}' — DB key is the authoritative fallback"
);
}
}
pub async fn set_merge_success_reported(&self, story_id: &str) {
let mut lock = self.agents.lock().await;
let found = lock.iter_mut().find(|(key, agent)| {
let key_story_id = key
.rsplit_once(':')
.map(|(sid, _)| sid)
.unwrap_or(key.as_str());
key_story_id == story_id
&& pipeline_stage(&agent.agent_name) == PipelineStage::Mergemaster
});
match found {
Some((_, agent)) => {
agent.merge_success_reported = true;
slog!(
"[pipeline] Merge success flag set for '{story_id}:{}'",
agent.agent_name
);
}
Err(e) => {
slog_error!("[pipeline] set_merge_success_reported: could not lock agents: {e}");
None => {
slog!(
"[pipeline] set_merge_success_reported: no running mergemaster \
for '{story_id}' — DB key is the authoritative fallback"
);
}
}
}
@@ -74,35 +68,29 @@ impl AgentPool {
/// that `run_pipeline_advance` can block advancement to `5_done/` even when
/// the server-owned gate check returns `gates_passed=true` (those gates run
/// in the feature-branch worktree, not on master).
pub fn set_merge_failure_reported(&self, story_id: &str) {
match self.agents.lock() {
Ok(mut lock) => {
let found = lock.iter_mut().find(|(key, agent)| {
let key_story_id = key
.rsplit_once(':')
.map(|(sid, _)| sid)
.unwrap_or(key.as_str());
key_story_id == story_id
&& pipeline_stage(&agent.agent_name) == PipelineStage::Mergemaster
});
match found {
Some((_, agent)) => {
agent.merge_failure_reported = true;
slog!(
"[pipeline] Merge failure flag set for '{story_id}:{}'",
agent.agent_name
);
}
None => {
slog_warn!(
"[pipeline] set_merge_failure_reported: no running mergemaster found \
for story '{story_id}' — flag not set"
);
}
}
pub async fn set_merge_failure_reported(&self, story_id: &str) {
let mut lock = self.agents.lock().await;
let found = lock.iter_mut().find(|(key, agent)| {
let key_story_id = key
.rsplit_once(':')
.map(|(sid, _)| sid)
.unwrap_or(key.as_str());
key_story_id == story_id
&& pipeline_stage(&agent.agent_name) == PipelineStage::Mergemaster
});
match found {
Some((_, agent)) => {
agent.merge_failure_reported = true;
slog!(
"[pipeline] Merge failure flag set for '{story_id}:{}'",
agent.agent_name
);
}
Err(e) => {
slog_error!("[pipeline] set_merge_failure_reported: could not lock agents: {e}");
None => {
slog_warn!(
"[pipeline] set_merge_failure_reported: no running mergemaster found \
for story '{story_id}' — flag not set"
);
}
}
}
server/src/agents/pool/pipeline/merge/runner.rs
+2 -2
View File
@@ -283,7 +283,7 @@ impl AgentPool {
&& let Ok(ref r) = report
&& 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");
}
@@ -350,7 +350,7 @@ impl AgentPool {
let story_archived = crate::agents::lifecycle::move_story_to_done(story_id).is_ok();
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() {
server/src/agents/pool/pipeline/merge/tests.rs
+2 -2
View File
@@ -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).
let agents = pool.agents.lock().unwrap();
let agents = pool.agents.try_lock().unwrap();
assert!(
agents.is_empty(),
"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.
let agents = pool.agents.lock().unwrap();
let agents = pool.agents.try_lock().unwrap();
assert!(
agents.is_empty(),
"no LLM agents should be spawned for deterministic merge; pool has {} agents",
server/src/agents/pool/process.rs
+13 -17
View File
@@ -25,11 +25,9 @@ impl AgentPool {
/// continuing to run after the server exits. Collects each agent's worktree
/// path, then SIGKILLs every process running inside that path and verifies
/// 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 Ok(agents) = self.agents.lock() else {
return;
};
let agents = self.agents.lock().await;
agents
.iter()
.filter_map(|(key, agent)| {
@@ -69,11 +67,9 @@ impl AgentPool {
/// 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
/// 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 Ok(agents) = self.agents.lock() else {
return;
};
let agents = self.agents.lock().await;
agents
.get(key)
.and_then(|a| a.worktree_info.as_ref().map(|wt| wt.path.clone()))
@@ -124,18 +120,18 @@ mod tests {
.unwrap_or(false)
}
#[test]
fn kill_all_children_is_safe_on_empty_pool() {
#[tokio::test]
async fn kill_all_children_is_safe_on_empty_pool() {
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
/// 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
/// launched with `--directory <worktree>` in production.
#[test]
fn kill_child_for_key_kills_real_process() {
#[tokio::test]
async fn kill_child_for_key_kills_real_process() {
use std::os::unix::process::CommandExt;
let pool = AgentPool::new_test(3002);
@@ -165,7 +161,7 @@ mod tests {
"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
assert!(
@@ -176,8 +172,8 @@ mod tests {
/// AC 5 — `kill_all_children` SIGKILLs all agents' processes. Two agents
/// with distinct worktree paths are injected; both must be gone after the call.
#[test]
fn kill_all_children_kills_multiple_real_processes() {
#[tokio::test]
async fn kill_all_children_kills_multiple_real_processes() {
use std::os::unix::process::CommandExt;
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 {
let _ = child.wait(); // reap zombie
server/src/agents/pool/query.rs
+45 -16
View File
@@ -10,12 +10,12 @@ use super::types::{agent_info_from_entry, composite_key};
impl AgentPool {
/// Return the names of configured agents for `stage` that are not currently
/// running or pending.
pub fn available_agents_for_stage(
pub async fn available_agents_for_stage(
&self,
config: &ProjectConfig,
stage: &PipelineStage,
) -> Result<Vec<String>, String> {
let agents = self.agents.lock().map_err(|e| e.to_string())?;
let agents = self.agents.lock().await;
Ok(config
.agent
.iter()
@@ -44,8 +44,8 @@ impl AgentPool {
}
/// List all agents with their status.
pub fn list_agents(&self) -> Result<Vec<AgentInfo>, String> {
let agents = self.agents.lock().map_err(|e| e.to_string())?;
pub async fn list_agents(&self) -> Result<Vec<AgentInfo>, String> {
let agents = self.agents.lock().await;
Ok(agents
.iter()
.map(|(key, agent)| {
@@ -59,14 +59,35 @@ impl AgentPool {
.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.
pub fn subscribe(
pub async fn subscribe(
&self,
story_id: &str,
agent_name: &str,
) -> Result<broadcast::Receiver<AgentEvent>, String> {
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
.get(&key)
.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.
pub fn drain_events(
pub async fn drain_events(
&self,
story_id: &str,
agent_name: &str,
) -> Result<Vec<AgentEvent>, String> {
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
.get(&key)
.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.
///
/// 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 agents = self.agents.lock().ok()?;
let agents = self.agents.lock().await;
let agent = agents.get(&key)?;
let session_id = agent.log_session_id.clone()?;
let project_root = agent.project_root.clone()?;
@@ -111,8 +136,8 @@ mod tests {
ProjectConfig::parse(toml_str).unwrap()
}
#[test]
fn available_agents_for_stage_returns_idle_agents() {
#[tokio::test]
async fn available_agents_for_stage_returns_idle_agents() {
let config = make_config(
r#"
[[agent]]
@@ -133,17 +158,19 @@ stage = "qa"
let available = pool
.available_agents_for_stage(&config, &PipelineStage::Coder)
.await
.unwrap();
assert_eq!(available, vec!["coder-2"]);
let available_qa = pool
.available_agents_for_stage(&config, &PipelineStage::Qa)
.await
.unwrap();
assert_eq!(available_qa, vec!["qa"]);
}
#[test]
fn available_agents_for_stage_returns_empty_when_all_busy() {
#[tokio::test]
async fn available_agents_for_stage_returns_empty_when_all_busy() {
let config = make_config(
r#"
[[agent]]
@@ -156,12 +183,13 @@ stage = "coder"
let available = pool
.available_agents_for_stage(&config, &PipelineStage::Coder)
.await
.unwrap();
assert!(available.is_empty());
}
#[test]
fn available_agents_for_stage_ignores_completed_agents() {
#[tokio::test]
async fn available_agents_for_stage_ignores_completed_agents() {
let config = make_config(
r#"
[[agent]]
@@ -174,6 +202,7 @@ stage = "coder"
let available = pool
.available_agents_for_stage(&config, &PipelineStage::Coder)
.await
.unwrap();
assert_eq!(available, vec!["coder-1"]);
}
server/src/agents/pool/start/mod.rs
+27 -24
View File
@@ -3,8 +3,10 @@
use crate::agent_log::AgentLogWriter;
use crate::config::ProjectConfig;
use crate::slog_error;
use std::future::Future;
use std::path::Path;
use std::sync::{Arc, Mutex};
use std::pin::Pin;
use std::sync::Arc;
use tokio::sync::broadcast;
use super::super::runtime::{
@@ -38,48 +40,48 @@ impl AgentPool {
/// `resume_context` (if any) is sent as the new message. This lets
/// the agent re-enter the previous conversation without re-reading
/// CLAUDE.md and README, satisfying story 543.
pub async fn start_agent(
&self,
project_root: &Path,
story_id: &str,
agent_name: Option<&str>,
resume_context: Option<&str>,
pub fn start_agent<'a>(
&'a self,
project_root: &'a Path,
story_id: &'a str,
agent_name: Option<&'a str>,
resume_context: Option<&'a str>,
session_id_to_resume: Option<String>,
) -> Result<AgentInfo, String> {
self.start_agent_inner(
) -> Pin<Box<dyn Future<Output = Result<AgentInfo, String>> + Send + 'a>> {
Box::pin(self.start_agent_inner(
project_root,
story_id,
agent_name,
resume_context,
session_id_to_resume,
None,
)
))
}
/// Start an agent with an `AppContext` for direct MCP tool dispatch.
///
/// API-based runtimes (Gemini, OpenAI) need the `AppContext` to invoke MCP
/// tools without an HTTP round-trip. CLI-based runtimes (Claude Code) do not.
pub fn start_agent_with_ctx(
&self,
project_root: &Path,
story_id: &str,
agent_name: Option<&str>,
resume_context: Option<&str>,
pub fn start_agent_with_ctx<'a>(
&'a self,
project_root: &'a Path,
story_id: &'a str,
agent_name: Option<&'a str>,
resume_context: Option<&'a str>,
session_id_to_resume: Option<String>,
app_ctx: Arc<crate::http::context::AppContext>,
) -> Result<AgentInfo, String> {
self.start_agent_inner(
) -> Pin<Box<dyn Future<Output = Result<AgentInfo, String>> + Send + 'a>> {
Box::pin(self.start_agent_inner(
project_root,
story_id,
agent_name,
resume_context,
session_id_to_resume,
Some(app_ctx),
)
))
}
fn start_agent_inner(
async fn start_agent_inner(
&self,
project_root: &Path,
story_id: &str,
@@ -100,7 +102,8 @@ impl AgentPool {
// Create name-independent shared resources before the lock so they are
// ready for the atomic check-and-insert (story 132).
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();
// Create the per-session status buffer subscribed to this project's
@@ -149,7 +152,7 @@ impl AgentPool {
// agent turn (story 736).
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 {
Some(name) => name.to_string(),
@@ -371,7 +374,7 @@ impl AgentPool {
// the atomic resolution above).
let log_writer =
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) => {
eprintln!(
"[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.
{
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) {
agent.task_handle = Some(handle);
}
server/src/agents/pool/start/spawn.rs
+69 -68
View File
@@ -146,7 +146,7 @@ pub(super) async fn run_agent_spawn(
story_id: String,
agent_name: String,
tx: broadcast::Sender<AgentEvent>,
agents: Arc<Mutex<HashMap<String, StoryAgent>>>,
agents: Arc<tokio::sync::Mutex<HashMap<String, StoryAgent>>>,
key: String,
event_log: Arc<Mutex<Vec<AgentEvent>>>,
port: u16,
@@ -218,10 +218,11 @@ pub(super) async fn run_agent_spawn(
log.push(event.clone());
}
let _ = tx_clone.send(event);
if let Ok(mut agents) = agents_ref.lock()
&& let Some(agent) = agents.get_mut(&key_clone)
{
agent.status = AgentStatus::Failed;
let mut agents = agents_ref.lock().await;
if let Some(agent) = agents.get_mut(&key_clone) {
agent.status = AgentStatus::Failed;
}
}
AgentPool::notify_agent_state_changed(&watcher_tx_clone);
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.
// Non-fatal — if installation fails the agent can still run; the failure
// is logged so the operator can investigate.
if let Err(e) = crate::worktree::install_pre_commit_hook(&wt_info.path) {
slog_error!("[agents] pre-commit hook install failed for {sid}: {e}");
// 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}");
}
}
// Step 2: store worktree info and render agent command/args/prompt.
let wt_path_str = wt_info.path.to_string_lossy().to_string();
{
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.worktree_info = Some(wt_info.clone());
}
}
@@ -266,10 +278,11 @@ pub(super) async fn run_agent_spawn(
log.push(event.clone());
}
let _ = tx_clone.send(event);
if let Ok(mut agents) = agents_ref.lock()
&& let Some(agent) = agents.get_mut(&key_clone)
{
agent.status = AgentStatus::Failed;
let mut agents = agents_ref.lock().await;
if let Some(agent) = agents.get_mut(&key_clone) {
agent.status = AgentStatus::Failed;
}
}
AgentPool::notify_agent_state_changed(&watcher_tx_clone);
return;
@@ -358,9 +371,8 @@ pub(super) async fn run_agent_spawn(
// Step 3: transition to Running now that the worktree is ready.
{
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::Running;
}
}
@@ -457,25 +469,26 @@ pub(super) async fn run_agent_spawn(
match run_result {
Ok(result) => {
// Persist token usage if the agent reported it.
if let Some(ref usage) = result.token_usage
&& let Ok(agents) = agents_ref.lock()
&& let Some(agent) = agents.get(&key_clone)
&& let Some(ref pr) = agent.project_root
{
let model_for_record = config_clone
.find_agent(&aname)
.and_then(|a| a.model.clone());
let record = crate::agents::token_usage::build_record(
&sid,
&aname,
model_for_record,
usage.clone(),
);
if let Err(e) = crate::agents::token_usage::append_record(pr, &record) {
slog_error!(
"[agents] Failed to persist token usage for \
{sid}:{aname}: {e}"
if let Some(ref usage) = result.token_usage {
let agents = agents_ref.lock().await;
if let Some(agent) = agents.get(&key_clone)
&& let Some(ref pr) = agent.project_root
{
let model_for_record = config_clone
.find_agent(&aname)
.and_then(|a| a.model.clone());
let record = crate::agents::token_usage::build_record(
&sid,
&aname,
model_for_record,
usage.clone(),
);
if let Err(e) = crate::agents::token_usage::append_record(pr, &record) {
slog_error!(
"[agents] Failed to persist token usage for \
{sid}:{aname}: {e}"
);
}
}
}
@@ -526,10 +539,7 @@ pub(super) async fn run_agent_spawn(
// Remove the agent entry from the pool and emit Done so that
// any caller blocked on wait_for_agent is unblocked.
let tx_done = {
let mut lock = match agents_ref.lock() {
Ok(a) => a,
Err(_) => return,
};
let mut lock = agents_ref.lock().await;
if let Some(agent) = lock.remove(&key_clone) {
agent.tx
} else {
@@ -608,10 +618,7 @@ pub(super) async fn run_agent_spawn(
if stage == PipelineStage::Mergemaster {
let (tx_done, done_session_id, merge_failure_reported, merge_success_reported) = {
let mut lock = match agents_ref.lock() {
Ok(a) => a,
Err(_) => return,
};
let mut lock = agents_ref.lock().await;
if let Some(agent) = lock.remove(&key_clone) {
(
agent.tx,
@@ -648,15 +655,14 @@ pub(super) async fn run_agent_spawn(
// Do NOT send WorkItem/reassign — story is already Done.
// Drain one queued ConflictDetected story now that this
// mergemaster slot is free (story 1044).
if let Some((candidate_id, candidate_agent)) =
crate::config::ProjectConfig::load(&project_root_clone)
.ok()
.and_then(|cfg| {
agents_ref.lock().ok().as_ref().and_then(|agts| {
pick_queued_conflict_detected(&cfg, agts, &sid)
})
})
{
let candidate =
if let Ok(cfg) = crate::config::ProjectConfig::load(&project_root_clone) {
let agts = agents_ref.lock().await;
pick_queued_conflict_detected(&cfg, &agts, &sid)
} else {
None
};
if let Some((candidate_id, candidate_agent)) = candidate {
slog!(
"[agents] Mergemaster exit for '{sid}' (success): \
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
// mergemaster slot is free (story 1044).
if let Some((candidate_id, candidate_agent)) =
crate::config::ProjectConfig::load(&project_root_clone)
.ok()
.and_then(|cfg| {
agents_ref
.lock()
.ok()
.as_ref()
.and_then(|agts| pick_queued_conflict_detected(&cfg, agts, &sid))
})
{
let candidate =
if let Ok(cfg) = crate::config::ProjectConfig::load(&project_root_clone) {
let agts = agents_ref.lock().await;
pick_queued_conflict_detected(&cfg, &agts, &sid)
} else {
None
};
if let Some((candidate_id, candidate_agent)) = candidate {
slog!(
"[agents] Mergemaster exit for '{sid}': queued ConflictDetected \
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.
let tx_done = {
let mut lock = match agents_ref.lock() {
Ok(a) => a,
Err(_) => return,
};
let mut lock = agents_ref.lock().await;
if let Some(agent) = lock.remove(&key_clone) {
agent.tx
} else {
@@ -931,10 +931,11 @@ pub(super) async fn run_agent_spawn(
log.push(event.clone());
}
let _ = tx_clone.send(event);
if let Ok(mut agents) = agents_ref.lock()
&& let Some(agent) = agents.get_mut(&key_clone)
{
agent.status = AgentStatus::Failed;
let mut agents = agents_ref.lock().await;
if let Some(agent) = agents.get_mut(&key_clone) {
agent.status = AgentStatus::Failed;
}
}
AgentPool::notify_agent_state_changed(&watcher_tx_clone);
}
server/src/agents/pool/start/tests_concurrency.rs
+6 -5
View File
@@ -109,7 +109,7 @@ async fn start_agent_cleans_up_pending_entry_on_failure() {
"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
.values()
.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
.drain_events("50_story_test", "coder-1")
.await
.expect("drain_events should succeed");
let has_error_event = events.iter().any(|e| matches!(e, AgentEvent::Error { .. }));
assert!(
@@ -736,7 +737,7 @@ async fn reconcile_canonical_agents_stops_stale_coder_in_qa_stage() {
let pool = AgentPool::new_test(3099);
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!(
before.iter().any(|a| a.agent_name == "coder-1"
&& 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;
let after = pool.list_agents().unwrap();
let after = pool.list_agents().await.unwrap();
let still_active = after.iter().any(|a| {
a.story_id == "777_story_reconcile"
&& 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;
let after = pool.list_agents().unwrap();
let after = pool.list_agents().await.unwrap();
let still_active = after.iter().any(|a| {
a.story_id == "555_story_correct"
&& 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;
// coder-1 must be gone from the active pool.
let remaining = pool.list_agents().unwrap();
let remaining = pool.list_agents().await.unwrap();
assert!(
!remaining.iter().any(|a| {
a.story_id == "1100_reg"
server/src/agents/pool/stop.rs
+18 -27
View File
@@ -1,7 +1,6 @@
//! Agent stop — terminates a running agent while preserving its worktree.
use crate::process_kill::{pids_matching, sigkill_pids_and_verify};
use crate::slog;
use crate::slog_error;
use crate::slog_warn;
use std::path::Path;
@@ -40,7 +39,7 @@ impl AgentPool {
// Step 1: snapshot the worktree path (no status mutation yet).
let worktree_info = {
let agents = self.agents.lock().map_err(|e| e.to_string())?;
let agents = self.agents.lock().await;
let agent = agents
.get(&key)
.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, \
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.
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
.get_mut(&key)
.ok_or_else(|| format!("No agent '{agent_name}' for story '{story_id}'"))?;
@@ -107,7 +106,7 @@ impl AgentPool {
// 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);
}
@@ -138,9 +137,7 @@ impl AgentPool {
// Snapshot active LLM agents without holding the lock during async stops.
let snapshot: Vec<(String, String, PipelineStage)> = {
let Ok(agents) = self.agents.lock() else {
return;
};
let agents = self.agents.lock().await;
agents
.iter()
.filter_map(|(key, a)| {
@@ -197,14 +194,8 @@ impl AgentPool {
///
/// Called when a story is archived so that stale entries don't accumulate.
/// Returns the number of entries removed.
pub fn remove_agents_for_story(&self, story_id: &str) -> usize {
let mut agents = match self.agents.lock() {
Ok(a) => a,
Err(e) => {
slog_error!("[agents] Failed to lock pool for cleanup of '{story_id}': {e}");
return 0;
}
};
pub async fn remove_agents_for_story(&self, story_id: &str) -> usize {
let mut agents = self.agents.lock().await;
let prefix = format!("{story_id}:");
let keys_to_remove: Vec<String> = agents
.keys()
@@ -229,30 +220,30 @@ mod tests {
// ── remove_agents_for_story tests ────────────────────────────────────────
#[test]
fn remove_agents_for_story_removes_all_entries() {
#[tokio::test]
async fn remove_agents_for_story_removes_all_entries() {
let pool = AgentPool::new_test(3001);
pool.inject_test_agent("story_a", "coder-1", AgentStatus::Completed);
pool.inject_test_agent("story_a", "qa", AgentStatus::Failed);
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");
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[0].story_id, "story_b");
}
#[test]
fn remove_agents_for_story_returns_zero_when_no_match() {
#[tokio::test]
async fn remove_agents_for_story_returns_zero_when_no_match() {
let pool = AgentPool::new_test(3001);
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);
let agents = pool.list_agents().unwrap();
let agents = pool.list_agents().await.unwrap();
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("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();
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!(
remaining.len(),
1,
server/src/agents/pool/test_helpers.rs
+8 -8
View File
@@ -20,7 +20,7 @@ impl AgentPool {
) -> broadcast::Sender<AgentEvent> {
let (tx, _) = broadcast::channel::<AgentEvent>(64);
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(
key,
StoryAgent {
@@ -55,7 +55,7 @@ impl AgentPool {
) -> broadcast::Sender<AgentEvent> {
let (tx, _) = broadcast::channel::<AgentEvent>(64);
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(
key,
StoryAgent {
@@ -95,7 +95,7 @@ impl AgentPool {
) -> broadcast::Sender<AgentEvent> {
let (tx, _) = broadcast::channel::<AgentEvent>(64);
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(
key,
StoryAgent {
@@ -130,7 +130,7 @@ impl AgentPool {
) -> broadcast::Sender<AgentEvent> {
let (tx, _) = broadcast::channel::<AgentEvent>(64);
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(
key,
StoryAgent {
@@ -165,7 +165,7 @@ impl AgentPool {
) -> broadcast::Sender<AgentEvent> {
let (tx, _) = broadcast::channel::<AgentEvent>(64);
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(
key,
StoryAgent {
@@ -197,7 +197,7 @@ impl AgentPool {
story_id: &str,
agent_name: &str,
) -> Option<Vec<BufferedItem>> {
let agents = self.agents.lock().unwrap();
let agents = self.agents.try_lock().unwrap();
let key = composite_key(story_id, agent_name);
agents
.get(&key)
@@ -219,7 +219,7 @@ impl AgentPool {
) -> broadcast::Sender<AgentEvent> {
let (tx, _) = broadcast::channel::<AgentEvent>(64);
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(
key,
StoryAgent {
@@ -258,7 +258,7 @@ impl AgentPool {
) -> broadcast::Sender<AgentEvent> {
let (tx, _) = broadcast::channel::<AgentEvent>(64);
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(
key,
StoryAgent {
server/src/agents/pool/types.rs
+7 -5
View File
@@ -3,8 +3,8 @@ use crate::slog;
use crate::worktree::WorktreeInfo;
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::{Arc, Mutex};
use tokio::sync::broadcast;
use std::sync::Arc;
use tokio::sync::{Mutex, broadcast};
use super::super::{AgentEvent, AgentInfo, AgentStatus, CompletionReport};
@@ -45,8 +45,10 @@ impl PendingGuard {
impl Drop for PendingGuard {
fn drop(&mut self) {
if self.armed
&& let Ok(mut agents) = self.agents.lock()
if !self.armed {
return;
}
if let Ok(mut agents) = self.agents.try_lock()
&& agents
.get(&self.key)
.is_some_and(|a| a.status == AgentStatus::Pending)
@@ -68,7 +70,7 @@ pub(super) struct StoryAgent {
pub(super) tx: broadcast::Sender<AgentEvent>,
pub(super) task_handle: Option<tokio::task::JoinHandle<()>>,
/// 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.
pub(super) completion: Option<CompletionReport>,
/// Project root, stored for pipeline advancement after completion.
server/src/agents/pool/wait.rs
+5 -5
View File
@@ -17,11 +17,11 @@ impl AgentPool {
) -> Result<AgentInfo, String> {
// Subscribe before checking status so we don't miss the terminal event
// 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.
{
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);
if let Some(agent) = agents.get(&key)
&& matches!(agent.status, AgentStatus::Completed | AgentStatus::Failed)
@@ -48,7 +48,7 @@ impl AgentPool {
_ => false,
};
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);
return Ok(if let Some(agent) = agents.get(&key) {
agent_info_from_entry(story_id, agent)
@@ -78,7 +78,7 @@ impl AgentPool {
}
Ok(Err(broadcast::error::RecvError::Lagged(_))) => {
// 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);
if let Some(agent) = agents.get(&key)
&& matches!(agent.status, AgentStatus::Completed | AgentStatus::Failed)
@@ -89,7 +89,7 @@ impl AgentPool {
}
Ok(Err(broadcast::error::RecvError::Closed)) => {
// 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);
if let Some(agent) = agents.get(&key) {
return Ok(agent_info_from_entry(story_id, agent));