huskies/server/src/agents/pool/auto_assign/scan.rs

//! 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,
            throttled: 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");
    }

}