From 774548c04cf458ac62a14b930ca304f5869c2579 Mon Sep 17 00:00:00 2001 From: Dave Date: Tue, 24 Feb 2026 15:50:34 +0000 Subject: [PATCH] story-kit: merge 153_bug_auto_assign_broken_after_stage_field_was_added_to_agent_config Co-Authored-By: Claude Opus 4.6 --- server/src/agents.rs | 118 +++++++++++++++++++++++++++++++++++++++++-- server/src/config.rs | 6 +++ 2 files changed, 119 insertions(+), 5 deletions(-) diff --git a/server/src/agents.rs b/server/src/agents.rs index b64cbe3..847d75c 100644 --- a/server/src/agents.rs +++ b/server/src/agents.rs @@ -2,7 +2,7 @@ use crate::agent_log::AgentLogWriter; use crate::slog; use crate::slog_error; use crate::slog_warn; -use crate::config::ProjectConfig; +use crate::config::{AgentConfig, ProjectConfig}; use crate::worktree::{self, WorktreeInfo}; use portable_pty::{CommandBuilder, PtySize, native_pty_system}; use serde::Serialize; @@ -160,6 +160,21 @@ pub fn pipeline_stage(agent_name: &str) -> PipelineStage { } } +/// Determine the pipeline stage for a configured agent. +/// +/// Prefers the explicit `stage` config field (added in Bug 150) over the +/// legacy name-based heuristic so that agents with non-standard names +/// (e.g. `qa-2`, `coder-opus`) are assigned to the correct stage. +fn agent_config_stage(cfg: &AgentConfig) -> PipelineStage { + match cfg.stage.as_deref() { + Some("coder") => PipelineStage::Coder, + Some("qa") => PipelineStage::Qa, + Some("mergemaster") => PipelineStage::Mergemaster, + Some(_) => PipelineStage::Other, + None => pipeline_stage(&cfg.name), + } +} + /// Completion report produced when acceptance gates are run. /// /// Created automatically by the server when an agent process exits normally, @@ -721,7 +736,11 @@ impl AgentPool { } }; - let stage = pipeline_stage(agent_name); + let config = ProjectConfig::load(&project_root).unwrap_or_default(); + let stage = config + .find_agent(agent_name) + .map(agent_config_stage) + .unwrap_or_else(|| pipeline_stage(agent_name)); match stage { PipelineStage::Other => { @@ -1222,7 +1241,7 @@ impl AgentPool { break; } }; - let assigned = is_story_assigned_for_stage(&agents, story_id, stage); + let assigned = is_story_assigned_for_stage(&config, &agents, story_id, stage); let free = if assigned { None } else { @@ -1689,7 +1708,11 @@ fn scan_stage_items(project_root: &Path, stage_dir: &str) -> Vec { } /// Return `true` if `story_id` has any active (pending/running) agent matching `stage`. +/// +/// Uses the explicit `stage` config field when the agent is found in `config`; +/// falls back to the legacy name-based heuristic for unlisted agents. fn is_story_assigned_for_stage( + config: &ProjectConfig, agents: &HashMap, story_id: &str, stage: &PipelineStage, @@ -1697,21 +1720,26 @@ fn is_story_assigned_for_stage( agents.iter().any(|(key, agent)| { // Composite key format: "{story_id}:{agent_name}" let key_story_id = key.rsplit_once(':').map(|(sid, _)| sid).unwrap_or(key); + let agent_stage = config + .find_agent(&agent.agent_name) + .map(agent_config_stage) + .unwrap_or_else(|| pipeline_stage(&agent.agent_name)); key_story_id == story_id - && pipeline_stage(&agent.agent_name) == *stage + && agent_stage == *stage && matches!(agent.status, AgentStatus::Running | AgentStatus::Pending) }) } /// Find the first configured agent for `stage` that has no active (pending/running) assignment. /// Returns `None` if all agents for that stage are busy or none are configured. +/// Uses the agent's explicit `stage` config field (preferred) or falls back to name-based detection. fn find_free_agent_for_stage<'a>( config: &'a ProjectConfig, agents: &HashMap, stage: &PipelineStage, ) -> Option<&'a str> { for agent_config in &config.agent { - if pipeline_stage(&agent_config.name) != *stage { + if agent_config_stage(agent_config) != *stage { continue; } let is_busy = agents.values().any(|a| { @@ -4117,23 +4145,28 @@ mod tests { #[test] fn is_story_assigned_returns_true_for_running_coder() { + use crate::config::ProjectConfig; + let config = ProjectConfig::default(); let pool = AgentPool::new(3001); pool.inject_test_agent("42_story_foo", "coder-1", AgentStatus::Running); let agents = pool.agents.lock().unwrap(); assert!(is_story_assigned_for_stage( + &config, &agents, "42_story_foo", &PipelineStage::Coder )); // Same story but wrong stage — should be false assert!(!is_story_assigned_for_stage( + &config, &agents, "42_story_foo", &PipelineStage::Qa )); // Different story — should be false assert!(!is_story_assigned_for_stage( + &config, &agents, "99_story_other", &PipelineStage::Coder @@ -4142,12 +4175,15 @@ mod tests { #[test] fn is_story_assigned_returns_false_for_completed_agent() { + use crate::config::ProjectConfig; + let config = ProjectConfig::default(); let pool = AgentPool::new(3001); pool.inject_test_agent("42_story_foo", "coder-1", AgentStatus::Completed); let agents = pool.agents.lock().unwrap(); // Completed agents don't count as assigned assert!(!is_story_assigned_for_stage( + &config, &agents, "42_story_foo", &PipelineStage::Coder @@ -4240,6 +4276,78 @@ name = "qa" assert_eq!(free_qa, Some("qa")); } + // ── agent_config_stage / stage field tests ──────────────────────────────── + + #[test] + fn find_free_agent_uses_config_stage_field_not_name() { + // Agents named "qa-2" and "coder-opus" don't match the legacy name heuristic + // but should be picked up via their explicit stage field. + use crate::config::ProjectConfig; + let config = ProjectConfig::parse( + r#" +[[agent]] +name = "qa-2" +stage = "qa" + +[[agent]] +name = "coder-opus" +stage = "coder" +"#, + ) + .unwrap(); + + let agents: HashMap = HashMap::new(); + + // qa-2 should be found for PipelineStage::Qa via config stage field + let free_qa = find_free_agent_for_stage(&config, &agents, &PipelineStage::Qa); + assert_eq!(free_qa, Some("qa-2"), "qa-2 with stage=qa should be found"); + + // coder-opus should be found for PipelineStage::Coder via config stage field + let free_coder = find_free_agent_for_stage(&config, &agents, &PipelineStage::Coder); + assert_eq!( + free_coder, + Some("coder-opus"), + "coder-opus with stage=coder should be found" + ); + + // Neither should match the other stage + let free_merge = find_free_agent_for_stage(&config, &agents, &PipelineStage::Mergemaster); + assert!(free_merge.is_none()); + } + + #[test] + fn is_story_assigned_uses_config_stage_field_for_nonstandard_names() { + use crate::config::ProjectConfig; + let config = ProjectConfig::parse( + r#" +[[agent]] +name = "qa-2" +stage = "qa" +"#, + ) + .unwrap(); + + let pool = AgentPool::new(3001); + pool.inject_test_agent("42_story_foo", "qa-2", AgentStatus::Running); + + let agents = pool.agents.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), + "qa-2 should be detected as assigned to QA stage" + ); + // Should NOT appear as a coder + assert!( + !is_story_assigned_for_stage( + &config, + &agents, + "42_story_foo", + &PipelineStage::Coder + ), + "qa-2 should not be detected as a coder" + ); + } + // ── find_active_story_stage tests ───────────────────────────────────────── #[test] diff --git a/server/src/config.rs b/server/src/config.rs index a9667e1..6e634e1 100644 --- a/server/src/config.rs +++ b/server/src/config.rs @@ -45,6 +45,11 @@ pub struct AgentConfig { pub max_budget_usd: Option, #[serde(default)] pub system_prompt: Option, + /// Pipeline stage this agent belongs to. Supported values: "coder", "qa", + /// "mergemaster", "other". When set, overrides the legacy name-based + /// detection used by `pipeline_stage()`. + #[serde(default)] + pub stage: Option, /// Inactivity timeout in seconds for the PTY read loop. /// If no output is received within this duration, the agent process is killed /// and marked as Failed. Default: 300 (5 minutes). Set to 0 to disable. @@ -99,6 +104,7 @@ impl Default for ProjectConfig { max_turns: None, max_budget_usd: None, system_prompt: None, + stage: None, inactivity_timeout_secs: default_inactivity_timeout_secs(), }], }