story-kit: merge 153_bug_auto_assign_broken_after_stage_field_was_added_to_agent_config

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

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<String> {
 }
 
 /// 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<String, StoryAgent>,
     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<String, StoryAgent>,
     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<String, StoryAgent> = 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]

server/src/config.rs

@@ -45,6 +45,11 @@ pub struct AgentConfig {
     pub max_budget_usd: Option<f64>,
     #[serde(default)]
     pub system_prompt: Option<String>,
+    /// 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<String>,
     /// 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(),
             }],
         }