storkit: merge 418_refactor_split_pool_auto_assign_rs_into_submodules

2026-03-28 08:44:27 +00:00
parent 4ea4be1462
commit d375c4b1d3
7 changed files with 1901 additions and 1813 deletions
@@ -0,0 +1,552 @@
+//! Scanning pipeline stages for work items and querying agent pool state.
+
+use crate::config::ProjectConfig;
+use std::collections::HashMap;
+use std::path::Path;
+
+use super::super::super::{AgentStatus, PipelineStage, agent_config_stage, pipeline_stage};
+use super::super::StoryAgent;
+
+/// Return `true` if `agent_name` has no active (pending/running) entry in the pool.
+pub(in crate::agents::pool) fn is_agent_free(
+    agents: &HashMap<String, StoryAgent>,
+    agent_name: &str,
+) -> bool {
+    !agents.values().any(|a| {
+        a.agent_name == agent_name
+            && matches!(a.status, AgentStatus::Running | AgentStatus::Pending)
+    })
+}
+
+pub(super) fn scan_stage_items(project_root: &Path, stage_dir: &str) -> Vec<String> {
+    let dir = project_root.join(".storkit").join("work").join(stage_dir);
+    if !dir.is_dir() {
+        return Vec::new();
+    }
+    let mut items = Vec::new();
+    if let Ok(entries) = std::fs::read_dir(&dir) {
+        for entry in entries.flatten() {
+            let path = entry.path();
+            if path.extension().and_then(|e| e.to_str()) == Some("md")
+                && let Some(stem) = path.file_stem().and_then(|s| s.to_str())
+            {
+                items.push(stem.to_string());
+            }
+        }
+    }
+    items.sort();
+    items
+}
+
+/// 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.
+pub(super) fn is_story_assigned_for_stage(
+    config: &ProjectConfig,
+    agents: &HashMap<String, StoryAgent>,
+    story_id: &str,
+    stage: &PipelineStage,
+) -> bool {
+    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
+            && agent_stage == *stage
+            && matches!(agent.status, AgentStatus::Running | AgentStatus::Pending)
+    })
+}
+
+/// Count active (pending/running) agents for a given pipeline stage.
+pub(super) fn count_active_agents_for_stage(
+    config: &ProjectConfig,
+    agents: &HashMap<String, StoryAgent>,
+    stage: &PipelineStage,
+) -> usize {
+    agents
+        .values()
+        .filter(|a| {
+            matches!(a.status, AgentStatus::Running | AgentStatus::Pending)
+                && config
+                    .find_agent(&a.agent_name)
+                    .map(|cfg| agent_config_stage(cfg) == *stage)
+                    .unwrap_or_else(|| pipeline_stage(&a.agent_name) == *stage)
+        })
+        .count()
+}
+
+/// Find the first configured agent for `stage` that has no active (pending/running) assignment.
+/// Returns `None` if all agents for that stage are busy, none are configured,
+/// or the `max_coders` limit has been reached (for the Coder stage).
+///
+/// For the Coder stage, when `default_coder_model` is set, only considers agents whose
+/// model matches the default. This ensures opus-class agents are reserved for explicit
+/// front-matter requests.
+pub(in crate::agents::pool) fn find_free_agent_for_stage<'a>(
+    config: &'a ProjectConfig,
+    agents: &HashMap<String, StoryAgent>,
+    stage: &PipelineStage,
+) -> Option<&'a str> {
+    // Enforce max_coders limit for the Coder stage.
+    if *stage == PipelineStage::Coder
+        && let Some(max) = config.max_coders
+    {
+        let active = count_active_agents_for_stage(config, agents, stage);
+        if active >= max {
+            return None;
+        }
+    }
+
+    for agent_config in &config.agent {
+        if agent_config_stage(agent_config) != *stage {
+            continue;
+        }
+        // When default_coder_model is set, only auto-assign coder agents whose
+        // model matches. This keeps opus agents reserved for explicit requests.
+        if *stage == PipelineStage::Coder
+            && let Some(ref default_model) = config.default_coder_model
+        {
+            let agent_model = agent_config.model.as_deref().unwrap_or("");
+            if agent_model != default_model {
+                continue;
+            }
+        }
+        let is_busy = agents.values().any(|a| {
+            a.agent_name == agent_config.name
+                && matches!(a.status, AgentStatus::Running | AgentStatus::Pending)
+        });
+        if !is_busy {
+            return Some(&agent_config.name);
+        }
+    }
+    None
+}
+
+// ── Tests ──────────────────────────────────────────────────────────────────
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::config::ProjectConfig;
+    use std::sync::{Arc, Mutex};
+    use tokio::sync::broadcast;
+
+    use super::super::super::AgentPool;
+
+    fn make_config(toml_str: &str) -> ProjectConfig {
+        ProjectConfig::parse(toml_str).unwrap()
+    }
+
+    fn make_test_story_agent(agent_name: &str, status: AgentStatus) -> StoryAgent {
+        StoryAgent {
+            agent_name: agent_name.to_string(),
+            status,
+            worktree_info: None,
+            session_id: None,
+            tx: broadcast::channel(1).0,
+            task_handle: None,
+            event_log: Arc::new(Mutex::new(Vec::new())),
+            completion: None,
+            project_root: None,
+            log_session_id: None,
+            merge_failure_reported: false,
+        }
+    }
+
+    #[test]
+    fn scan_stage_items_returns_empty_for_missing_dir() {
+        let tmp = tempfile::tempdir().unwrap();
+        let items = scan_stage_items(tmp.path(), "2_current");
+        assert!(items.is_empty());
+    }
+
+    #[test]
+    fn scan_stage_items_returns_sorted_story_ids() {
+        use std::fs;
+        let tmp = tempfile::tempdir().unwrap();
+        let stage_dir = tmp.path().join(".storkit").join("work").join("2_current");
+        fs::create_dir_all(&stage_dir).unwrap();
+        fs::write(stage_dir.join("42_story_foo.md"), "---\nname: foo\n---").unwrap();
+        fs::write(stage_dir.join("10_story_bar.md"), "---\nname: bar\n---").unwrap();
+        fs::write(stage_dir.join("5_story_baz.md"), "---\nname: baz\n---").unwrap();
+        // non-md file should be ignored
+        fs::write(stage_dir.join("README.txt"), "ignore me").unwrap();
+
+        let items = scan_stage_items(tmp.path(), "2_current");
+        assert_eq!(items, vec!["10_story_bar", "42_story_foo", "5_story_baz"]);
+    }
+
+    #[test]
+    fn is_story_assigned_returns_true_for_running_coder() {
+        let config = ProjectConfig::default();
+        let pool = AgentPool::new_test(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
+        ));
+    }
+
+    #[test]
+    fn is_story_assigned_returns_false_for_completed_agent() {
+        let config = ProjectConfig::default();
+        let pool = AgentPool::new_test(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
+        ));
+    }
+
+    #[test]
+    fn is_story_assigned_uses_config_stage_field_for_nonstandard_names() {
+        let config = ProjectConfig::parse(
+            r#"
+[[agent]]
+name = "qa-2"
+stage = "qa"
+"#,
+        )
+        .unwrap();
+
+        let pool = AgentPool::new_test(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"
+        );
+    }
+
+    #[test]
+    fn find_free_agent_returns_none_when_all_busy() {
+        let config = ProjectConfig::parse(
+            r#"
+[[agent]]
+name = "coder-1"
+[[agent]]
+name = "coder-2"
+"#,
+        )
+        .unwrap();
+
+        let pool = AgentPool::new_test(3001);
+        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 free = find_free_agent_for_stage(&config, &agents, &PipelineStage::Coder);
+        assert!(free.is_none(), "no free coders should be available");
+    }
+
+    #[test]
+    fn find_free_agent_returns_first_free_coder() {
+        let config = ProjectConfig::parse(
+            r#"
+[[agent]]
+name = "coder-1"
+[[agent]]
+name = "coder-2"
+[[agent]]
+name = "coder-3"
+"#,
+        )
+        .unwrap();
+
+        let pool = AgentPool::new_test(3001);
+        // coder-1 is busy, coder-2 is free
+        pool.inject_test_agent("s1", "coder-1", AgentStatus::Running);
+
+        let agents = pool.agents.lock().unwrap();
+        let free = find_free_agent_for_stage(&config, &agents, &PipelineStage::Coder);
+        assert_eq!(
+            free,
+            Some("coder-2"),
+            "coder-2 should be the first free coder"
+        );
+    }
+
+    #[test]
+    fn find_free_agent_ignores_completed_agents() {
+        let config = ProjectConfig::parse(
+            r#"
+[[agent]]
+name = "coder-1"
+"#,
+        )
+        .unwrap();
+
+        let pool = AgentPool::new_test(3001);
+        // 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 free = find_free_agent_for_stage(&config, &agents, &PipelineStage::Coder);
+        assert_eq!(free, Some("coder-1"), "completed coder-1 should be free");
+    }
+
+    #[test]
+    fn find_free_agent_returns_none_for_wrong_stage() {
+        let config = ProjectConfig::parse(
+            r#"
+[[agent]]
+name = "qa"
+"#,
+        )
+        .unwrap();
+
+        let agents: HashMap<String, StoryAgent> = HashMap::new();
+        // Looking for a Coder but only QA is configured
+        let free = find_free_agent_for_stage(&config, &agents, &PipelineStage::Coder);
+        assert!(free.is_none());
+        // Looking for QA should find it
+        let free_qa = find_free_agent_for_stage(&config, &agents, &PipelineStage::Qa);
+        assert_eq!(free_qa, Some("qa"));
+    }
+
+    #[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.
+        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());
+    }
+
+    // ── find_free_agent_for_stage: default_coder_model filtering ─────────
+
+    #[test]
+    fn find_free_agent_skips_opus_when_default_coder_model_set() {
+        let config = make_config(
+            r#"
+default_coder_model = "sonnet"
+
+[[agent]]
+name = "coder-1"
+stage = "coder"
+model = "sonnet"
+
+[[agent]]
+name = "coder-opus"
+stage = "coder"
+model = "opus"
+"#,
+        );
+
+        let agents = HashMap::new();
+        let free = find_free_agent_for_stage(&config, &agents, &PipelineStage::Coder);
+        assert_eq!(free, Some("coder-1"));
+    }
+
+    #[test]
+    fn find_free_agent_returns_opus_when_no_default_coder_model() {
+        let config = make_config(
+            r#"
+[[agent]]
+name = "coder-opus"
+stage = "coder"
+model = "opus"
+"#,
+        );
+
+        let agents = HashMap::new();
+        let free = find_free_agent_for_stage(&config, &agents, &PipelineStage::Coder);
+        assert_eq!(free, Some("coder-opus"));
+    }
+
+    #[test]
+    fn find_free_agent_returns_none_when_all_sonnet_coders_busy() {
+        let config = make_config(
+            r#"
+default_coder_model = "sonnet"
+
+[[agent]]
+name = "coder-1"
+stage = "coder"
+model = "sonnet"
+
+[[agent]]
+name = "coder-opus"
+stage = "coder"
+model = "opus"
+"#,
+        );
+
+        let mut agents = HashMap::new();
+        agents.insert(
+            "story1:coder-1".to_string(),
+            make_test_story_agent("coder-1", AgentStatus::Running),
+        );
+
+        let free = find_free_agent_for_stage(&config, &agents, &PipelineStage::Coder);
+        assert_eq!(free, None, "opus agent should not be auto-assigned");
+    }
+
+    // ── find_free_agent_for_stage: max_coders limit ─────────────────────
+
+    #[test]
+    fn find_free_agent_respects_max_coders() {
+        let config = make_config(
+            r#"
+max_coders = 1
+
+[[agent]]
+name = "coder-1"
+stage = "coder"
+model = "sonnet"
+
+[[agent]]
+name = "coder-2"
+stage = "coder"
+model = "sonnet"
+"#,
+        );
+
+        let mut agents = HashMap::new();
+        agents.insert(
+            "story1:coder-1".to_string(),
+            make_test_story_agent("coder-1", AgentStatus::Running),
+        );
+
+        let free = find_free_agent_for_stage(&config, &agents, &PipelineStage::Coder);
+        assert_eq!(free, None, "max_coders=1 should block second coder");
+    }
+
+    #[test]
+    fn find_free_agent_allows_within_max_coders() {
+        let config = make_config(
+            r#"
+max_coders = 2
+
+[[agent]]
+name = "coder-1"
+stage = "coder"
+model = "sonnet"
+
+[[agent]]
+name = "coder-2"
+stage = "coder"
+model = "sonnet"
+"#,
+        );
+
+        let mut agents = HashMap::new();
+        agents.insert(
+            "story1:coder-1".to_string(),
+            make_test_story_agent("coder-1", AgentStatus::Running),
+        );
+
+        let free = find_free_agent_for_stage(&config, &agents, &PipelineStage::Coder);
+        assert_eq!(free, Some("coder-2"));
+    }
+
+    #[test]
+    fn max_coders_does_not_affect_qa_stage() {
+        let config = make_config(
+            r#"
+max_coders = 1
+
+[[agent]]
+name = "qa"
+stage = "qa"
+model = "sonnet"
+"#,
+        );
+
+        let agents = HashMap::new();
+        let free = find_free_agent_for_stage(&config, &agents, &PipelineStage::Qa);
+        assert_eq!(free, Some("qa"));
+    }
+
+    // ── count_active_agents_for_stage ────────────────────────────────────
+
+    #[test]
+    fn count_active_agents_counts_running_and_pending() {
+        let config = make_config(
+            r#"
+[[agent]]
+name = "coder-1"
+stage = "coder"
+
+[[agent]]
+name = "coder-2"
+stage = "coder"
+"#,
+        );
+
+        let mut agents = HashMap::new();
+        agents.insert(
+            "s1:coder-1".to_string(),
+            make_test_story_agent("coder-1", AgentStatus::Running),
+        );
+        agents.insert(
+            "s2:coder-2".to_string(),
+            make_test_story_agent("coder-2", AgentStatus::Completed),
+        );
+
+        let count = count_active_agents_for_stage(&config, &agents, &PipelineStage::Coder);
+        assert_eq!(count, 1, "Only Running coder should be counted, not Completed");
+    }
+
+}