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

Converted all external tool calling to async

Browse Source
This commit is contained in:
Timmy
2026-06-29 16:59:54 +01:00
parent 75f41088b1
commit feb35ddd10
49 changed files with 482 additions and 496 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server/src/agents/pool/start/mod.rs
+27 -24
View File
@@ -3,8 +3,10 @@
use crate::agent_log::AgentLogWriter;
use crate::config::ProjectConfig;
use crate::slog_error;
use std::future::Future;
use std::path::Path;
use std::sync::{Arc, Mutex};
use std::pin::Pin;
use std::sync::Arc;
use tokio::sync::broadcast;
use super::super::runtime::{
@@ -38,48 +40,48 @@ impl AgentPool {
/// `resume_context` (if any) is sent as the new message. This lets
/// the agent re-enter the previous conversation without re-reading
/// CLAUDE.md and README, satisfying story 543.
pub async fn start_agent(
&self,
project_root: &Path,
story_id: &str,
agent_name: Option<&str>,
resume_context: Option<&str>,
pub fn start_agent<'a>(
&'a self,
project_root: &'a Path,
story_id: &'a str,
agent_name: Option<&'a str>,
resume_context: Option<&'a str>,
session_id_to_resume: Option<String>,
) -> Result<AgentInfo, String> {
self.start_agent_inner(
) -> Pin<Box<dyn Future<Output = Result<AgentInfo, String>> + Send + 'a>> {
Box::pin(self.start_agent_inner(
project_root,
story_id,
agent_name,
resume_context,
session_id_to_resume,
None,
)
))
}
/// Start an agent with an `AppContext` for direct MCP tool dispatch.
///
/// API-based runtimes (Gemini, OpenAI) need the `AppContext` to invoke MCP
/// tools without an HTTP round-trip. CLI-based runtimes (Claude Code) do not.
pub fn start_agent_with_ctx(
&self,
project_root: &Path,
story_id: &str,
agent_name: Option<&str>,
resume_context: Option<&str>,
pub fn start_agent_with_ctx<'a>(
&'a self,
project_root: &'a Path,
story_id: &'a str,
agent_name: Option<&'a str>,
resume_context: Option<&'a str>,
session_id_to_resume: Option<String>,
app_ctx: Arc<crate::http::context::AppContext>,
) -> Result<AgentInfo, String> {
self.start_agent_inner(
) -> Pin<Box<dyn Future<Output = Result<AgentInfo, String>> + Send + 'a>> {
Box::pin(self.start_agent_inner(
project_root,
story_id,
agent_name,
resume_context,
session_id_to_resume,
Some(app_ctx),
)
))
}
fn start_agent_inner(
async fn start_agent_inner(
&self,
project_root: &Path,
story_id: &str,
@@ -100,7 +102,8 @@ impl AgentPool {
// Create name-independent shared resources before the lock so they are
// ready for the atomic check-and-insert (story 132).
let (tx, _) = broadcast::channel::<AgentEvent>(1024);
let event_log: Arc<Mutex<Vec<AgentEvent>>> = Arc::new(Mutex::new(Vec::new()));
let event_log: Arc<std::sync::Mutex<Vec<AgentEvent>>> =
Arc::new(std::sync::Mutex::new(Vec::new()));
let log_session_id = uuid::Uuid::new_v4().to_string();
// Create the per-session status buffer subscribed to this project's
@@ -149,7 +152,7 @@ impl AgentPool {
// agent turn (story 736).
let prior_events: Option<String>;
{
let mut agents = self.agents.lock().map_err(|e| e.to_string())?;
let mut agents = self.agents.lock().await;
resolved_name = match agent_name {
Some(name) => name.to_string(),
@@ -371,7 +374,7 @@ impl AgentPool {
// the atomic resolution above).
let log_writer =
match AgentLogWriter::new(project_root, story_id, &resolved_name, &log_session_id) {
Ok(w) => Some(Arc::new(Mutex::new(w))),
Ok(w) => Some(Arc::new(std::sync::Mutex::new(w))),
Err(e) => {
eprintln!(
"[agents] Failed to create log writer for {story_id}:{resolved_name}: {e}"
@@ -436,7 +439,7 @@ impl AgentPool {
// Store the task handle while the agent is still Pending.
{
let mut agents = self.agents.lock().map_err(|e| e.to_string())?;
let mut agents = self.agents.lock().await;
if let Some(agent) = agents.get_mut(&key) {
agent.task_handle = Some(handle);
}
server/src/agents/pool/start/spawn.rs
+69 -68
View File
@@ -146,7 +146,7 @@ pub(super) async fn run_agent_spawn(
story_id: String,
agent_name: String,
tx: broadcast::Sender<AgentEvent>,
agents: Arc<Mutex<HashMap<String, StoryAgent>>>,
agents: Arc<tokio::sync::Mutex<HashMap<String, StoryAgent>>>,
key: String,
event_log: Arc<Mutex<Vec<AgentEvent>>>,
port: u16,
@@ -218,10 +218,11 @@ pub(super) async fn run_agent_spawn(
log.push(event.clone());
}
let _ = tx_clone.send(event);
if let Ok(mut agents) = agents_ref.lock()
&& let Some(agent) = agents.get_mut(&key_clone)
{
agent.status = AgentStatus::Failed;
let mut agents = agents_ref.lock().await;
if let Some(agent) = agents.get_mut(&key_clone) {
agent.status = AgentStatus::Failed;
}
}
AgentPool::notify_agent_state_changed(&watcher_tx_clone);
return;
@@ -233,16 +234,27 @@ pub(super) async fn run_agent_spawn(
// Step 1.1: Install the pre-commit quality-gate hook in the worktree.
// Non-fatal — if installation fails the agent can still run; the failure
// is logged so the operator can investigate.
if let Err(e) = crate::worktree::install_pre_commit_hook(&wt_info.path) {
slog_error!("[agents] pre-commit hook install failed for {sid}: {e}");
// Runs in spawn_blocking because install_pre_commit_hook executes
// synchronous git-config subprocesses that would otherwise pin a
// tokio worker thread and contribute to runtime starvation under
// concurrent agent spawns.
{
let hook_path = wt_info.path.clone();
let hook_result = tokio::task::spawn_blocking(move || {
crate::worktree::install_pre_commit_hook(&hook_path)
})
.await
.unwrap_or_else(|e| Err(format!("spawn_blocking panicked: {e}")));
if let Err(e) = hook_result {
slog_error!("[agents] pre-commit hook install failed for {sid}: {e}");
}
}
// Step 2: store worktree info and render agent command/args/prompt.
let wt_path_str = wt_info.path.to_string_lossy().to_string();
{
if let Ok(mut agents) = agents_ref.lock()
&& let Some(agent) = agents.get_mut(&key_clone)
{
let mut agents = agents_ref.lock().await;
if let Some(agent) = agents.get_mut(&key_clone) {
agent.worktree_info = Some(wt_info.clone());
}
}
@@ -266,10 +278,11 @@ pub(super) async fn run_agent_spawn(
log.push(event.clone());
}
let _ = tx_clone.send(event);
if let Ok(mut agents) = agents_ref.lock()
&& let Some(agent) = agents.get_mut(&key_clone)
{
agent.status = AgentStatus::Failed;
let mut agents = agents_ref.lock().await;
if let Some(agent) = agents.get_mut(&key_clone) {
agent.status = AgentStatus::Failed;
}
}
AgentPool::notify_agent_state_changed(&watcher_tx_clone);
return;
@@ -358,9 +371,8 @@ pub(super) async fn run_agent_spawn(
// Step 3: transition to Running now that the worktree is ready.
{
if let Ok(mut agents) = agents_ref.lock()
&& let Some(agent) = agents.get_mut(&key_clone)
{
let mut agents = agents_ref.lock().await;
if let Some(agent) = agents.get_mut(&key_clone) {
agent.status = AgentStatus::Running;
}
}
@@ -457,25 +469,26 @@ pub(super) async fn run_agent_spawn(
match run_result {
Ok(result) => {
// Persist token usage if the agent reported it.
if let Some(ref usage) = result.token_usage
&& let Ok(agents) = agents_ref.lock()
&& let Some(agent) = agents.get(&key_clone)
&& let Some(ref pr) = agent.project_root
{
let model_for_record = config_clone
.find_agent(&aname)
.and_then(|a| a.model.clone());
let record = crate::agents::token_usage::build_record(
&sid,
&aname,
model_for_record,
usage.clone(),
);
if let Err(e) = crate::agents::token_usage::append_record(pr, &record) {
slog_error!(
"[agents] Failed to persist token usage for \
{sid}:{aname}: {e}"
if let Some(ref usage) = result.token_usage {
let agents = agents_ref.lock().await;
if let Some(agent) = agents.get(&key_clone)
&& let Some(ref pr) = agent.project_root
{
let model_for_record = config_clone
.find_agent(&aname)
.and_then(|a| a.model.clone());
let record = crate::agents::token_usage::build_record(
&sid,
&aname,
model_for_record,
usage.clone(),
);
if let Err(e) = crate::agents::token_usage::append_record(pr, &record) {
slog_error!(
"[agents] Failed to persist token usage for \
{sid}:{aname}: {e}"
);
}
}
}
@@ -526,10 +539,7 @@ pub(super) async fn run_agent_spawn(
// Remove the agent entry from the pool and emit Done so that
// any caller blocked on wait_for_agent is unblocked.
let tx_done = {
let mut lock = match agents_ref.lock() {
Ok(a) => a,
Err(_) => return,
};
let mut lock = agents_ref.lock().await;
if let Some(agent) = lock.remove(&key_clone) {
agent.tx
} else {
@@ -608,10 +618,7 @@ pub(super) async fn run_agent_spawn(
if stage == PipelineStage::Mergemaster {
let (tx_done, done_session_id, merge_failure_reported, merge_success_reported) = {
let mut lock = match agents_ref.lock() {
Ok(a) => a,
Err(_) => return,
};
let mut lock = agents_ref.lock().await;
if let Some(agent) = lock.remove(&key_clone) {
(
agent.tx,
@@ -648,15 +655,14 @@ pub(super) async fn run_agent_spawn(
// Do NOT send WorkItem/reassign — story is already Done.
// Drain one queued ConflictDetected story now that this
// mergemaster slot is free (story 1044).
if let Some((candidate_id, candidate_agent)) =
crate::config::ProjectConfig::load(&project_root_clone)
.ok()
.and_then(|cfg| {
agents_ref.lock().ok().as_ref().and_then(|agts| {
pick_queued_conflict_detected(&cfg, agts, &sid)
})
})
{
let candidate =
if let Ok(cfg) = crate::config::ProjectConfig::load(&project_root_clone) {
let agts = agents_ref.lock().await;
pick_queued_conflict_detected(&cfg, &agts, &sid)
} else {
None
};
if let Some((candidate_id, candidate_agent)) = candidate {
slog!(
"[agents] Mergemaster exit for '{sid}' (success): \
queued ConflictDetected story '{candidate_id}' found; \
@@ -766,17 +772,14 @@ pub(super) async fn run_agent_spawn(
});
// Drain one queued ConflictDetected story now that this
// mergemaster slot is free (story 1044).
if let Some((candidate_id, candidate_agent)) =
crate::config::ProjectConfig::load(&project_root_clone)
.ok()
.and_then(|cfg| {
agents_ref
.lock()
.ok()
.as_ref()
.and_then(|agts| pick_queued_conflict_detected(&cfg, agts, &sid))
})
{
let candidate =
if let Ok(cfg) = crate::config::ProjectConfig::load(&project_root_clone) {
let agts = agents_ref.lock().await;
pick_queued_conflict_detected(&cfg, &agts, &sid)
} else {
None
};
if let Some((candidate_id, candidate_agent)) = candidate {
slog!(
"[agents] Mergemaster exit for '{sid}': queued ConflictDetected \
story '{candidate_id}' found; spawning '{candidate_agent}'."
@@ -833,10 +836,7 @@ pub(super) async fn run_agent_spawn(
// Remove agent from the pool and unblock any wait_for_agent callers.
let tx_done = {
let mut lock = match agents_ref.lock() {
Ok(a) => a,
Err(_) => return,
};
let mut lock = agents_ref.lock().await;
if let Some(agent) = lock.remove(&key_clone) {
agent.tx
} else {
@@ -931,10 +931,11 @@ pub(super) async fn run_agent_spawn(
log.push(event.clone());
}
let _ = tx_clone.send(event);
if let Ok(mut agents) = agents_ref.lock()
&& let Some(agent) = agents.get_mut(&key_clone)
{
agent.status = AgentStatus::Failed;
let mut agents = agents_ref.lock().await;
if let Some(agent) = agents.get_mut(&key_clone) {
agent.status = AgentStatus::Failed;
}
}
AgentPool::notify_agent_state_changed(&watcher_tx_clone);
}
server/src/agents/pool/start/tests_concurrency.rs
+6 -5
View File
@@ -109,7 +109,7 @@ async fn start_agent_cleans_up_pending_entry_on_failure() {
"agent must transition to Failed after worktree creation error"
);
let agents = pool.agents.lock().unwrap();
let agents = pool.agents.lock().await;
let failed_entry = agents
.values()
.find(|a| a.agent_name == "coder-1" && a.status == AgentStatus::Failed);
@@ -121,6 +121,7 @@ async fn start_agent_cleans_up_pending_entry_on_failure() {
let events = pool
.drain_events("50_story_test", "coder-1")
.await
.expect("drain_events should succeed");
let has_error_event = events.iter().any(|e| matches!(e, AgentEvent::Error { .. }));
assert!(
@@ -736,7 +737,7 @@ async fn reconcile_canonical_agents_stops_stale_coder_in_qa_stage() {
let pool = AgentPool::new_test(3099);
pool.inject_test_agent("777_story_reconcile", "coder-1", AgentStatus::Running);
let before = pool.list_agents().unwrap();
let before = pool.list_agents().await.unwrap();
assert!(
before.iter().any(|a| a.agent_name == "coder-1"
&& matches!(a.status, AgentStatus::Running | AgentStatus::Pending)),
@@ -745,7 +746,7 @@ async fn reconcile_canonical_agents_stops_stale_coder_in_qa_stage() {
pool.reconcile_canonical_agents(root).await;
let after = pool.list_agents().unwrap();
let after = pool.list_agents().await.unwrap();
let still_active = after.iter().any(|a| {
a.story_id == "777_story_reconcile"
&& a.agent_name == "coder-1"
@@ -786,7 +787,7 @@ async fn reconcile_canonical_agents_leaves_correct_stage_agent_alone() {
pool.reconcile_canonical_agents(root).await;
let after = pool.list_agents().unwrap();
let after = pool.list_agents().await.unwrap();
let still_active = after.iter().any(|a| {
a.story_id == "555_story_correct"
&& a.agent_name == "coder-1"
@@ -851,7 +852,7 @@ async fn regression_1100_stale_coder_blocks_mergemaster_then_reconciler_clears()
pool.reconcile_canonical_agents(root).await;
// coder-1 must be gone from the active pool.
let remaining = pool.list_agents().unwrap();
let remaining = pool.list_agents().await.unwrap();
assert!(
!remaining.iter().any(|a| {
a.story_id == "1100_reg"