From feb35ddd10965248affdb3bcb0a3dcac18904f52 Mon Sep 17 00:00:00 2001 From: Timmy Date: Mon, 29 Jun 2026 16:59:54 +0100 Subject: [PATCH] Converted all external tool calling to async --- server/src/agent_mode/loop_ops.rs | 4 +- .../agents/pool/auto_assign/auto_assign.rs | 20 +-- server/src/agents/pool/auto_assign/merge.rs | 8 +- .../merge_failure_block_subscriber.rs | 9 +- .../auto_assign/merge_failure_subscriber.rs | 24 +-- .../src/agents/pool/auto_assign/pipeline.rs | 17 +-- server/src/agents/pool/auto_assign/scan.rs | 12 +- .../pool/auto_assign/watchdog/limits.rs | 9 +- .../agents/pool/auto_assign/watchdog/mod.rs | 36 ++--- .../pool/auto_assign/watchdog/orphan.rs | 9 +- .../watchdog/tests/limits_tests.rs | 82 +++++------ .../watchdog/tests/orphan_tests.rs | 18 +-- server/src/agents/pool/mod.rs | 12 +- .../agents/pool/pipeline/advance/helpers.rs | 4 +- .../src/agents/pool/pipeline/advance/mod.rs | 2 +- .../pool/pipeline/advance/tests_regression.rs | 16 +- .../agents/pool/pipeline/completion/legacy.rs | 4 +- .../agents/pool/pipeline/completion/server.rs | 18 +-- .../agents/pool/pipeline/completion/tests.rs | 16 +- .../src/agents/pool/pipeline/merge/control.rs | 100 ++++++------- .../src/agents/pool/pipeline/merge/runner.rs | 4 +- .../src/agents/pool/pipeline/merge/tests.rs | 4 +- server/src/agents/pool/process.rs | 30 ++-- server/src/agents/pool/query.rs | 61 ++++++-- server/src/agents/pool/start/mod.rs | 51 ++++--- server/src/agents/pool/start/spawn.rs | 137 +++++++++--------- .../agents/pool/start/tests_concurrency.rs | 11 +- server/src/agents/pool/stop.rs | 45 +++--- server/src/agents/pool/test_helpers.rs | 16 +- server/src/agents/pool/types.rs | 12 +- server/src/agents/pool/wait.rs | 10 +- server/src/chat/commands/status/render.rs | 2 +- server/src/chat/transport/matrix/assign.rs | 1 + server/src/chat/transport/matrix/htop.rs | 2 +- server/src/chat/transport/matrix/rmtree.rs | 1 + server/src/http/agents_sse.rs | 2 +- server/src/http/mcp/agent_tools/inspection.rs | 47 +++--- server/src/http/mcp/agent_tools/lifecycle.rs | 10 +- server/src/http/mcp/diagnostics/permission.rs | 4 +- server/src/http/mcp/dispatch.rs | 14 +- server/src/http/mcp/merge_tools.rs | 29 ++-- server/src/http/mcp/qa_tools.rs | 2 +- server/src/http/mcp/story_tools/bug.rs | 18 ++- .../src/http/mcp/story_tools/story/delete.rs | 28 ++-- server/src/http/workflow/pipeline.rs | 5 +- server/src/main.rs | 2 +- server/src/rebuild.rs | 3 +- server/src/service/work_item/delete.rs | 5 +- server/src/startup/tick_loop.rs | 2 +- 49 files changed, 482 insertions(+), 496 deletions(-) diff --git a/server/src/agent_mode/loop_ops.rs b/server/src/agent_mode/loop_ops.rs index bee86110..4a6f138e 100644 --- a/server/src/agent_mode/loop_ops.rs +++ b/server/src/agent_mode/loop_ops.rs @@ -159,7 +159,7 @@ pub(super) async fn detect_conflicts( our_claims.remove(&story_id); // 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 { if info.story_id == story_id { 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, /// and report completion via CRDT. 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; }; diff --git a/server/src/agents/pool/auto_assign/auto_assign.rs b/server/src/agents/pool/auto_assign/auto_assign.rs index ba2fa7b2..b239c2bd 100644 --- a/server/src/agents/pool/auto_assign/auto_assign.rs +++ b/server/src/agents/pool/auto_assign/auto_assign.rs @@ -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)| { diff --git a/server/src/agents/pool/auto_assign/merge.rs b/server/src/agents/pool/auto_assign/merge.rs index 47113028..46e02493 100644 --- a/server/src/agents/pool/auto_assign/merge.rs +++ b/server/src/agents/pool/auto_assign/merge.rs @@ -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 { diff --git a/server/src/agents/pool/auto_assign/merge_failure_block_subscriber.rs b/server/src/agents/pool/auto_assign/merge_failure_block_subscriber.rs index d67758cf..47534810 100644 --- a/server/src/agents/pool/auto_assign/merge_failure_block_subscriber.rs +++ b/server/src/agents/pool/auto_assign/merge_failure_block_subscriber.rs @@ -54,7 +54,7 @@ pub(crate) fn spawn_merge_failure_block_subscriber(pool: Arc, 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, 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) } diff --git a/server/src/agents/pool/auto_assign/merge_failure_subscriber.rs b/server/src/agents/pool/auto_assign/merge_failure_subscriber.rs index 81604cec..b5c50647 100644 --- a/server/src/agents/pool/auto_assign/merge_failure_subscriber.rs +++ b/server/src/agents/pool/auto_assign/merge_failure_subscriber.rs @@ -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)| { diff --git a/server/src/agents/pool/auto_assign/pipeline.rs b/server/src/agents/pool/auto_assign/pipeline.rs index f46d79e9..8276b006 100644 --- a/server/src/agents/pool/auto_assign/pipeline.rs +++ b/server/src/agents/pool/auto_assign/pipeline.rs @@ -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) diff --git a/server/src/agents/pool/auto_assign/scan.rs b/server/src/agents/pool/auto_assign/scan.rs index 29212932..ab405764 100644 --- a/server/src/agents/pool/auto_assign/scan.rs +++ b/server/src/agents/pool/auto_assign/scan.rs @@ -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"); } diff --git a/server/src/agents/pool/auto_assign/watchdog/limits.rs b/server/src/agents/pool/auto_assign/watchdog/limits.rs index b741d1b0..92422251 100644 --- a/server/src/agents/pool/auto_assign/watchdog/limits.rs +++ b/server/src/agents/pool/auto_assign/watchdog/limits.rs @@ -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>, 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 = { - 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)| { diff --git a/server/src/agents/pool/auto_assign/watchdog/mod.rs b/server/src/agents/pool/auto_assign/watchdog/mod.rs index e7d4709d..7376fbda 100644 --- a/server/src/agents/pool/auto_assign/watchdog/mod.rs +++ b/server/src/agents/pool/auto_assign/watchdog/mod.rs @@ -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(); + } } } diff --git a/server/src/agents/pool/auto_assign/watchdog/orphan.rs b/server/src/agents/pool/auto_assign/watchdog/orphan.rs index 5f7b3358..e05e29be 100644 --- a/server/src/agents/pool/auto_assign/watchdog/orphan.rs +++ b/server/src/agents/pool/auto_assign/watchdog/orphan.rs @@ -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>) -> usize { - let mut lock = match agents.lock() { - Ok(l) => l, - Err(_) => return 0, - }; +pub(super) async fn check_orphaned_agents(agents: &Mutex>) -> 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. diff --git a/server/src/agents/pool/auto_assign/watchdog/tests/limits_tests.rs b/server/src/agents/pool/auto_assign/watchdog/tests/limits_tests.rs index 501c0071..822896d0 100644 --- a/server/src/agents/pool/auto_assign/watchdog/tests/limits_tests.rs +++ b/server/src/agents/pool/auto_assign/watchdog/tests/limits_tests.rs @@ -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); diff --git a/server/src/agents/pool/auto_assign/watchdog/tests/orphan_tests.rs b/server/src/agents/pool/auto_assign/watchdog/tests/orphan_tests.rs index 1a1d9f7a..d97975e0 100644 --- a/server/src/agents/pool/auto_assign/watchdog/tests/orphan_tests.rs +++ b/server/src/agents/pool/auto_assign/watchdog/tests/orphan_tests.rs @@ -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, diff --git a/server/src/agents/pool/mod.rs b/server/src/agents/pool/mod.rs index 8c0cd43d..8c2bcedc 100644 --- a/server/src/agents/pool/mod.rs +++ b/server/src/agents/pool/mod.rs @@ -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); } diff --git a/server/src/agents/pool/pipeline/advance/helpers.rs b/server/src/agents/pool/pipeline/advance/helpers.rs index 2ac440a2..c9ca5d25 100644 --- a/server/src/agents/pool/pipeline/advance/helpers.rs +++ b/server/src/agents/pool/pipeline/advance/helpers.rs @@ -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>>, + agents: Arc>>, port: u16, story_id: &str, agent_name: &str, diff --git a/server/src/agents/pool/pipeline/advance/mod.rs b/server/src/agents/pool/pipeline/advance/mod.rs index c71332ab..bb235e2b 100644 --- a/server/src/agents/pool/pipeline/advance/mod.rs +++ b/server/src/agents/pool/pipeline/advance/mod.rs @@ -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. diff --git a/server/src/agents/pool/pipeline/advance/tests_regression.rs b/server/src/agents/pool/pipeline/advance/tests_regression.rs index 2defcd6d..dcf14766 100644 --- a/server/src/agents/pool/pipeline/advance/tests_regression.rs +++ b/server/src/agents/pool/pipeline/advance/tests_regression.rs @@ -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")); diff --git a/server/src/agents/pool/pipeline/completion/legacy.rs b/server/src/agents/pool/pipeline/completion/legacy.rs index 6f3351fb..25d9db5f 100644 --- a/server/src/agents/pool/pipeline/completion/legacy.rs +++ b/server/src/agents/pool/pipeline/completion/legacy.rs @@ -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") })?; diff --git a/server/src/agents/pool/pipeline/completion/server.rs b/server/src/agents/pool/pipeline/completion/server.rs index aaf5a371..56223fc0 100644 --- a/server/src/agents/pool/pipeline/completion/server.rs +++ b/server/src/agents/pool/pipeline/completion/server.rs @@ -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, diff --git a/server/src/agents/pool/pipeline/completion/tests.rs b/server/src/agents/pool/pipeline/completion/tests.rs index 39f3dbda..8164bc4b 100644 --- a/server/src/agents/pool/pipeline/completion/tests.rs +++ b/server/src/agents/pool/pipeline/completion/tests.rs @@ -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, diff --git a/server/src/agents/pool/pipeline/merge/control.rs b/server/src/agents/pool/pipeline/merge/control.rs index f2155a63..97304412 100644 --- a/server/src/agents/pool/pipeline/merge/control.rs +++ b/server/src/agents/pool/pipeline/merge/control.rs @@ -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" + ); } } } diff --git a/server/src/agents/pool/pipeline/merge/runner.rs b/server/src/agents/pool/pipeline/merge/runner.rs index f04026a2..3ef15129 100644 --- a/server/src/agents/pool/pipeline/merge/runner.rs +++ b/server/src/agents/pool/pipeline/merge/runner.rs @@ -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() { diff --git a/server/src/agents/pool/pipeline/merge/tests.rs b/server/src/agents/pool/pipeline/merge/tests.rs index 51652ef3..cd7193c4 100644 --- a/server/src/agents/pool/pipeline/merge/tests.rs +++ b/server/src/agents/pool/pipeline/merge/tests.rs @@ -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", diff --git a/server/src/agents/pool/process.rs b/server/src/agents/pool/process.rs index fc767925..fbc14dc4 100644 --- a/server/src/agents/pool/process.rs +++ b/server/src/agents/pool/process.rs @@ -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 ` 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 diff --git a/server/src/agents/pool/query.rs b/server/src/agents/pool/query.rs index 6e76edc0..7d564618 100644 --- a/server/src/agents/pool/query.rs +++ b/server/src/agents/pool/query.rs @@ -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, 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, String> { - let agents = self.agents.lock().map_err(|e| e.to_string())?; + pub async fn list_agents(&self) -> Result, 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 { + 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, 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, 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"]); } diff --git a/server/src/agents/pool/start/mod.rs b/server/src/agents/pool/start/mod.rs index 820a9ab8..4a45bc1a 100644 --- a/server/src/agents/pool/start/mod.rs +++ b/server/src/agents/pool/start/mod.rs @@ -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, - ) -> Result { - self.start_agent_inner( + ) -> Pin> + 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, app_ctx: Arc, - ) -> Result { - self.start_agent_inner( + ) -> Pin> + 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::(1024); - let event_log: Arc>> = Arc::new(Mutex::new(Vec::new())); + let event_log: Arc>> = + 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; { - 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); } diff --git a/server/src/agents/pool/start/spawn.rs b/server/src/agents/pool/start/spawn.rs index 9e2cad3e..3acd008a 100644 --- a/server/src/agents/pool/start/spawn.rs +++ b/server/src/agents/pool/start/spawn.rs @@ -146,7 +146,7 @@ pub(super) async fn run_agent_spawn( story_id: String, agent_name: String, tx: broadcast::Sender, - agents: Arc>>, + agents: Arc>>, key: String, event_log: Arc>>, 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); } diff --git a/server/src/agents/pool/start/tests_concurrency.rs b/server/src/agents/pool/start/tests_concurrency.rs index d34e49c1..486c98ce 100644 --- a/server/src/agents/pool/start/tests_concurrency.rs +++ b/server/src/agents/pool/start/tests_concurrency.rs @@ -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" diff --git a/server/src/agents/pool/stop.rs b/server/src/agents/pool/stop.rs index 64446f32..a7321c0d 100644 --- a/server/src/agents/pool/stop.rs +++ b/server/src/agents/pool/stop.rs @@ -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 = 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, diff --git a/server/src/agents/pool/test_helpers.rs b/server/src/agents/pool/test_helpers.rs index 42383cbf..4e2f6f38 100644 --- a/server/src/agents/pool/test_helpers.rs +++ b/server/src/agents/pool/test_helpers.rs @@ -20,7 +20,7 @@ impl AgentPool { ) -> broadcast::Sender { let (tx, _) = broadcast::channel::(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 { let (tx, _) = broadcast::channel::(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 { let (tx, _) = broadcast::channel::(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 { let (tx, _) = broadcast::channel::(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 { let (tx, _) = broadcast::channel::(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> { - 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 { let (tx, _) = broadcast::channel::(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 { let (tx, _) = broadcast::channel::(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 { diff --git a/server/src/agents/pool/types.rs b/server/src/agents/pool/types.rs index 2a708fb4..9a6ee054 100644 --- a/server/src/agents/pool/types.rs +++ b/server/src/agents/pool/types.rs @@ -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, pub(super) task_handle: Option>, /// Accumulated events for polling via get_agent_output. - pub(super) event_log: Arc>>, + pub(super) event_log: Arc>>, /// Set when the agent calls report_completion. pub(super) completion: Option, /// Project root, stored for pipeline advancement after completion. diff --git a/server/src/agents/pool/wait.rs b/server/src/agents/pool/wait.rs index f347a63c..0d26f0b2 100644 --- a/server/src/agents/pool/wait.rs +++ b/server/src/agents/pool/wait.rs @@ -17,11 +17,11 @@ impl AgentPool { ) -> Result { // 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)); diff --git a/server/src/chat/commands/status/render.rs b/server/src/chat/commands/status/render.rs index a0183534..3978213b 100644 --- a/server/src/chat/commands/status/render.rs +++ b/server/src/chat/commands/status/render.rs @@ -90,7 +90,7 @@ pub(crate) fn build_status_from_items( items: &[PipelineItem], ) -> String { // 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 = active_agents .iter() .filter(|a| matches!(a.status, AgentStatus::Running | AgentStatus::Pending)) diff --git a/server/src/chat/transport/matrix/assign.rs b/server/src/chat/transport/matrix/assign.rs index 704ae676..417cfa55 100644 --- a/server/src/chat/transport/matrix/assign.rs +++ b/server/src/chat/transport/matrix/assign.rs @@ -111,6 +111,7 @@ pub async fn handle_assign( // Check whether a coder is already running on this story. let running_coders: Vec<_> = agents .list_agents() + .await .unwrap_or_default() .into_iter() .filter(|a| { diff --git a/server/src/chat/transport/matrix/htop.rs b/server/src/chat/transport/matrix/htop.rs index ef7ef89d..d459c0b6 100644 --- a/server/src/chat/transport/matrix/htop.rs +++ b/server/src/chat/transport/matrix/htop.rs @@ -191,7 +191,7 @@ pub fn build_htop_message(agents: &AgentPool, tick: u32, total_duration_secs: u6 String::new(), ]; - let all_agents = agents.list_agents().unwrap_or_default(); + let all_agents = agents.list_agents_nonblocking(); let active: Vec<_> = all_agents .iter() .filter(|a| matches!(a.status, AgentStatus::Running | AgentStatus::Pending)) diff --git a/server/src/chat/transport/matrix/rmtree.rs b/server/src/chat/transport/matrix/rmtree.rs index 5f190d3d..226e9972 100644 --- a/server/src/chat/transport/matrix/rmtree.rs +++ b/server/src/chat/transport/matrix/rmtree.rs @@ -85,6 +85,7 @@ pub async fn handle_rmtree( // Stop any running or pending agents for this story. let running_agents: Vec<(String, String)> = agents .list_agents() + .await .unwrap_or_default() .into_iter() .filter(|a| { diff --git a/server/src/http/agents_sse.rs b/server/src/http/agents_sse.rs index 6bdbc7ed..579b3835 100644 --- a/server/src/http/agents_sse.rs +++ b/server/src/http/agents_sse.rs @@ -18,7 +18,7 @@ pub async fn agent_stream( Path((story_id, agent_name)): Path<(String, String)>, ctx: Data<&Arc>, ) -> 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, Err(e) => { return Response::builder() diff --git a/server/src/http/mcp/agent_tools/inspection.rs b/server/src/http/mcp/agent_tools/inspection.rs index bcb39d08..2089e4e0 100644 --- a/server/src/http/mcp/agent_tools/inspection.rs +++ b/server/src/http/mcp/agent_tools/inspection.rs @@ -54,7 +54,7 @@ pub(crate) async fn tool_get_agent_output( // writer failed and nothing was persisted to disk. if log_files.is_empty() && 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() { all_lines.push(format!("=== {agent_name} (live) ===")); @@ -99,7 +99,7 @@ pub(crate) async fn tool_get_agent_output( Ok(output) } -pub(crate) fn tool_get_agent_config(ctx: &AppContext) -> Result { +pub(crate) async fn tool_get_agent_config(ctx: &AppContext) -> Result { let project_root = ctx.services.agents.get_project_root(&ctx.state)?; let config = ProjectConfig::load(&project_root)?; @@ -116,6 +116,7 @@ pub(crate) fn tool_get_agent_config(ctx: &AppContext) -> Result .services .agents .available_agents_for_stage(&config, stage) + .await { available_names.extend(names); } @@ -144,7 +145,7 @@ pub(crate) fn tool_get_agent_config(ctx: &AppContext) -> Result /// Returns turns used, max turns, remaining turns, budget used, max budget, /// and remaining budget for the named agent. Fails if the agent is not /// 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, ctx: &AppContext, ) -> Result { @@ -158,7 +159,7 @@ pub(crate) fn tool_get_agent_remaining_turns_and_budget( .ok_or("Missing required argument: agent_name")?; // 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 .iter() .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 serde_json::json; - #[test] - fn tool_get_agent_config_no_project_toml_returns_default_agent() { + #[tokio::test] + async fn tool_get_agent_config_no_project_toml_returns_default_agent() { let tmp = tempfile::tempdir().unwrap(); let ctx = test_ctx(tmp.path()); // 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 = serde_json::from_str(&result).unwrap(); // Default config contains one agent entry with default values assert_eq!( @@ -457,34 +458,36 @@ mod tests { // ── get_agent_remaining_turns_and_budget tests ────────────────────────── - #[test] - fn tool_get_agent_remaining_turns_and_budget_missing_story_id() { + #[tokio::test] + async fn tool_get_agent_remaining_turns_and_budget_missing_story_id() { let tmp = tempfile::tempdir().unwrap(); let ctx = test_ctx(tmp.path()); 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.unwrap_err().contains("story_id")); } - #[test] - fn tool_get_agent_remaining_turns_and_budget_missing_agent_name() { + #[tokio::test] + async fn tool_get_agent_remaining_turns_and_budget_missing_agent_name() { let tmp = tempfile::tempdir().unwrap(); let ctx = test_ctx(tmp.path()); 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.unwrap_err().contains("agent_name")); } - #[test] - fn tool_get_agent_remaining_turns_and_budget_no_agent_returns_error() { + #[tokio::test] + async fn tool_get_agent_remaining_turns_and_budget_no_agent_returns_error() { let tmp = tempfile::tempdir().unwrap(); let ctx = test_ctx(tmp.path()); let result = tool_get_agent_remaining_turns_and_budget( &json!({"story_id": "99_nope", "agent_name": "coder-1"}), &ctx, - ); + ) + .await; assert!(result.is_err()); let err = result.unwrap_err(); assert!( @@ -493,8 +496,8 @@ mod tests { ); } - #[test] - fn tool_get_agent_remaining_turns_and_budget_completed_agent_returns_error() { + #[tokio::test] + async fn tool_get_agent_remaining_turns_and_budget_completed_agent_returns_error() { use crate::agents::AgentStatus; let tmp = tempfile::tempdir().unwrap(); let ctx = test_ctx(tmp.path()); @@ -505,7 +508,8 @@ mod tests { let result = tool_get_agent_remaining_turns_and_budget( &json!({"story_id": "42_story", "agent_name": "coder-1"}), &ctx, - ); + ) + .await; assert!(result.is_err()); let err = result.unwrap_err(); assert!( @@ -514,8 +518,8 @@ mod tests { ); } - #[test] - fn tool_get_agent_remaining_turns_and_budget_running_agent_returns_data() { + #[tokio::test] + async fn tool_get_agent_remaining_turns_and_budget_running_agent_returns_data() { use crate::agents::AgentStatus; use crate::store::StoreOps; @@ -531,6 +535,7 @@ mod tests { &json!({"story_id": "42_story", "agent_name": "coder-1"}), &ctx, ) + .await .unwrap(); let parsed: Value = serde_json::from_str(&result).unwrap(); diff --git a/server/src/http/mcp/agent_tools/lifecycle.rs b/server/src/http/mcp/agent_tools/lifecycle.rs index 5b1e627e..adae5a03 100644 --- a/server/src/http/mcp/agent_tools/lifecycle.rs +++ b/server/src/http/mcp/agent_tools/lifecycle.rs @@ -67,9 +67,9 @@ pub(crate) async fn tool_stop_agent(args: &Value, ctx: &AppContext) -> Result Result { +pub(crate) async fn tool_list_agents(ctx: &AppContext) -> Result { 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 = agents .iter() .filter(|a| { @@ -156,11 +156,11 @@ mod tests { use crate::http::test_helpers::test_ctx; use serde_json::json; - #[test] - fn tool_list_agents_empty() { + #[tokio::test] + async fn tool_list_agents_empty() { let tmp = tempfile::tempdir().unwrap(); 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::from_str(&result).unwrap(); assert!(parsed.is_empty()); } diff --git a/server/src/http/mcp/diagnostics/permission.rs b/server/src/http/mcp/diagnostics/permission.rs index 8c103b12..2dc4440d 100644 --- a/server/src/http/mcp/diagnostics/permission.rs +++ b/server/src/http/mcp/diagnostics/permission.rs @@ -362,8 +362,8 @@ mod tests { // then exec() will be called — which would replace our test process. // 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. - assert_eq!(ctx.services.agents.list_agents().unwrap().len(), 0); - ctx.services.agents.kill_all_children(); // should not panic on empty pool + assert_eq!(ctx.services.agents.list_agents().await.unwrap().len(), 0); + ctx.services.agents.kill_all_children().await; // should not panic on empty pool } #[test] diff --git a/server/src/http/mcp/dispatch.rs b/server/src/http/mcp/dispatch.rs index 88dc0732..3be98581 100644 --- a/server/src/http/mcp/dispatch.rs +++ b/server/src/http/mcp/dispatch.rs @@ -30,13 +30,13 @@ pub async fn dispatch_tool_call( // Agent tools (async) "start_agent" => agent_tools::tool_start_agent(&args, ctx).await, "stop_agent" => agent_tools::tool_stop_agent(&args, ctx).await, - "list_agents" => agent_tools::tool_list_agents(ctx), - "get_agent_config" => agent_tools::tool_get_agent_config(ctx), - "reload_agent_config" => agent_tools::tool_get_agent_config(ctx), + "list_agents" => agent_tools::tool_list_agents(ctx).await, + "get_agent_config" => agent_tools::tool_get_agent_config(ctx).await, + "reload_agent_config" => agent_tools::tool_get_agent_config(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, "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 "create_worktree" => agent_tools::tool_create_worktree(&args, ctx).await, @@ -46,7 +46,7 @@ pub async fn dispatch_tool_call( // Editor tools "get_editor_command" => agent_tools::tool_get_editor_command(&args, ctx), // 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) "check_criterion" => story_tools::tool_check_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 "create_bug" => story_tools::tool_create_bug(&args, 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 "create_refactor" => story_tools::tool_create_refactor(&args, 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, "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, - "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 "request_qa" => qa_tools::tool_request_qa(&args, ctx).await, "approve_qa" => qa_tools::tool_approve_qa(&args, ctx).await, diff --git a/server/src/http/mcp/merge_tools.rs b/server/src/http/mcp/merge_tools.rs index be4165f1..9850f02d 100644 --- a/server/src/http/mcp/merge_tools.rs +++ b/server/src/http/mcp/merge_tools.rs @@ -171,7 +171,10 @@ pub(super) async fn tool_move_story_to_merge( .map_err(|e| format!("Serialization error: {e}")) } -pub(super) fn tool_report_merge_failure(args: &Value, ctx: &AppContext) -> Result { +pub(super) async fn tool_report_merge_failure( + args: &Value, + ctx: &AppContext, +) -> Result { let story_id = args .get("story_id") .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")?; 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 // auto-assigner does not re-spawn another mergemaster after this one fails. @@ -412,26 +418,26 @@ mod tests { assert!(req_names.contains(&"reason")); } - #[test] - fn tool_report_merge_failure_missing_story_id() { + #[tokio::test] + async fn tool_report_merge_failure_missing_story_id() { let tmp = tempfile::tempdir().unwrap(); 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.unwrap_err().contains("story_id")); } - #[test] - fn tool_report_merge_failure_missing_reason() { + #[tokio::test] + async fn tool_report_merge_failure_missing_reason() { let tmp = tempfile::tempdir().unwrap(); 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.unwrap_err().contains("reason")); } - #[test] - fn tool_report_merge_failure_returns_confirmation() { + #[tokio::test] + async fn tool_report_merge_failure_returns_confirmation() { let tmp = tempfile::tempdir().unwrap(); let ctx = test_ctx(tmp.path()); let result = tool_report_merge_failure( @@ -440,7 +446,8 @@ mod tests { "reason": "Unresolvable merge conflicts in src/main.rs" }), &ctx, - ); + ) + .await; assert!(result.is_ok()); let msg = result.unwrap(); assert!(msg.contains("42_story_foo")); diff --git a/server/src/http/mcp/qa_tools.rs b/server/src/http/mcp/qa_tools.rs index 920266f2..d2f57b18 100644 --- a/server/src/http/mcp/qa_tools.rs +++ b/server/src/http/mcp/qa_tools.rs @@ -81,7 +81,7 @@ pub(super) async fn tool_approve_qa(args: &Value, ctx: &AppContext) -> Result Result { .map_err(|e| format!("Serialization error: {e}")) } -pub(crate) fn tool_close_bug(args: &Value, ctx: &AppContext) -> Result { +pub(crate) async fn tool_close_bug(args: &Value, ctx: &AppContext) -> Result { let bug_id = args .get("bug_id") .and_then(|v| v.as_str()) .ok_or("Missing required argument: 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!( "Bug '{bug_id}' closed, moved to bugs/archive/, and committed to master." @@ -422,17 +422,17 @@ mod tests { assert!(result.unwrap().contains("Created bug:")); } - #[test] - fn tool_close_bug_missing_bug_id() { + #[tokio::test] + async fn tool_close_bug_missing_bug_id() { let tmp = tempfile::tempdir().unwrap(); 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.unwrap_err().contains("bug_id")); } - #[test] - fn tool_close_bug_moves_to_archive() { + #[tokio::test] + async fn tool_close_bug_moves_to_archive() { let tmp = tempfile::tempdir().unwrap(); setup_git_repo_in(tmp.path()); let backlog_dir = tmp.path().join(".huskies/work/1_backlog"); @@ -460,7 +460,9 @@ mod tests { .unwrap(); 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!( crate::db::read_content(crate::db::ContentKey::Story("9901_bug_crash")).is_some(), diff --git a/server/src/http/mcp/story_tools/story/delete.rs b/server/src/http/mcp/story_tools/story/delete.rs index 7ccc9ae8..3541b364 100644 --- a/server/src/http/mcp/story_tools/story/delete.rs +++ b/server/src/http/mcp/story_tools/story/delete.rs @@ -5,7 +5,7 @@ use crate::http::context::AppContext; use crate::pipeline_state::{Stage, read_typed}; use serde_json::Value; -pub(crate) fn tool_accept_story(args: &Value, ctx: &AppContext) -> Result { +pub(crate) async fn tool_accept_story(args: &Value, ctx: &AppContext) -> Result { let story_id = args .get("story_id") .and_then(|v| v.as_str()) @@ -33,7 +33,7 @@ pub(crate) fn tool_accept_story(args: &Value, ctx: &AppContext) -> Result Result { /// Build a map from story_id → AgentAssignment for all pending/running agents. fn build_active_agent_map(ctx: &AppContext) -> HashMap { - let agents = match ctx.services.agents.list_agents() { - Ok(a) => a, - Err(_) => return HashMap::new(), - }; + let agents = ctx.services.agents.list_agents_nonblocking(); let config_opt = ctx .state diff --git a/server/src/main.rs b/server/src/main.rs index 61318e08..abce1605 100644 --- a/server/src/main.rs +++ b/server/src/main.rs @@ -492,7 +492,7 @@ async fn main() -> Result<(), std::io::Error> { tokio::time::sleep(std::time::Duration::from_millis(1500)).await; // 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 { remove_port_file(path); diff --git a/server/src/rebuild.rs b/server/src/rebuild.rs index 761e9405..e8d96828 100644 --- a/server/src/rebuild.rs +++ b/server/src/rebuild.rs @@ -95,6 +95,7 @@ pub async fn rebuild_and_restart( // 1. Gracefully stop all running agents. let running_count = agents .list_agents() + .await .unwrap_or_default() .iter() .filter(|a| a.status == crate::agents::AgentStatus::Running) @@ -102,7 +103,7 @@ pub async fn rebuild_and_restart( if running_count > 0 { 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). // CARGO_MANIFEST_DIR at compile time points to the `server/` crate; diff --git a/server/src/service/work_item/delete.rs b/server/src/service/work_item/delete.rs index 5a0ac9d9..ee28c4bd 100644 --- a/server/src/service/work_item/delete.rs +++ b/server/src/service/work_item/delete.rs @@ -95,7 +95,7 @@ pub async fn delete_work_item( // 1. Stop any running/pending agents (best-effort). let mut agents_stopped: Vec = 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) { match agents .stop_agent(project_root, story_id, &agent.agent_name) @@ -120,7 +120,7 @@ pub async fn delete_work_item( } // 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}'"); // 3. Remove worktree (best-effort). @@ -240,6 +240,7 @@ mod tests { // Agent pool must have no entries for this story. let pool_entries = agents .list_agents() + .await .unwrap_or_default() .into_iter() .filter(|a| a.story_id == story_id) diff --git a/server/src/startup/tick_loop.rs b/server/src/startup/tick_loop.rs index 67c97ebc..614c176e 100644 --- a/server/src/startup/tick_loop.rs +++ b/server/src/startup/tick_loop.rs @@ -198,7 +198,7 @@ pub(crate) fn spawn_tick_loop( // timeout window). A `TransitionFired` subscriber cannot observe the // absence of events, so this must remain on a periodic tick. 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 { crate::slog!( "[tick] {found} orphaned agent(s) detected; triggering auto-assign."