huskies/server/src/chat/transport/matrix/assign.rs

//! Assign command: pre-assign or re-assign a coder model to a story.
//!
//! `{bot_name} assign {number} {model}` finds the story by number, updates the
//! `agent` field in its front matter, and — when a coder is already running on
//! the story — stops the current coder and starts the newly-assigned one.
//!
//! When no coder is running (the story has not been started yet), the command
//! behaves as before: it simply persists the assignment in the front matter so
//! that the next `start` invocation picks it up automatically.

use crate::agents::{AgentPool, AgentStatus};
use crate::chat::util::strip_bot_mention;
use crate::io::story_metadata::{parse_front_matter, set_front_matter_field};
use std::path::Path;

/// All pipeline stage directories to search when finding a work item by number.
const STAGES: &[&str] = &[
    "1_backlog",
    "2_current",
    "3_qa",
    "4_merge",
    "5_done",
    "6_archived",
];

/// A parsed assign command from a Matrix message body.
#[derive(Debug, PartialEq)]
pub enum AssignCommand {
    /// Assign the story with this number to the given model.
    Assign {
        story_number: String,
        model: String,
    },
    /// The user typed `assign` but without valid arguments.
    BadArgs,
}

/// Parse an assign command from a raw Matrix message body.
///
/// Strips the bot mention prefix and checks whether the first word is `assign`.
/// Returns `None` when the message is not an assign command at all.
pub fn extract_assign_command(
    message: &str,
    bot_name: &str,
    bot_user_id: &str,
) -> Option<AssignCommand> {
    let stripped = strip_bot_mention(message, bot_name, bot_user_id);
    let trimmed = stripped
        .trim()
        .trim_start_matches(|c: char| !c.is_alphanumeric());

    let (cmd, args) = match trimmed.split_once(char::is_whitespace) {
        Some((c, a)) => (c, a.trim()),
        None => (trimmed, ""),
    };

    if !cmd.eq_ignore_ascii_case("assign") {
        return None;
    }

    // Split args into story number and model.
    let (number_str, model_str) = match args.split_once(char::is_whitespace) {
        Some((n, m)) => (n.trim(), m.trim()),
        None => (args, ""),
    };

    if number_str.is_empty()
        || !number_str.chars().all(|c| c.is_ascii_digit())
        || model_str.is_empty()
    {
        return Some(AssignCommand::BadArgs);
    }

    Some(AssignCommand::Assign {
        story_number: number_str.to_string(),
        model: model_str.to_string(),
    })
}

/// Resolve a model name hint (e.g. `"opus"`) to a full agent name
/// (e.g. `"coder-opus"`). If the hint already starts with `"coder-"`,
/// it is returned unchanged to prevent double-prefixing.
pub fn resolve_agent_name(model: &str) -> String {
    if model.starts_with("coder-") {
        model.to_string()
    } else {
        format!("coder-{model}")
    }
}

/// Handle an assign command asynchronously.
///
/// Finds the work item by `story_number` across all pipeline stages, updates
/// the `agent` field in its front matter, and — if a coder is currently
/// running on the story — stops it and starts the newly-assigned agent.
/// Returns a markdown-formatted response string.
pub async fn handle_assign(
    bot_name: &str,
    story_number: &str,
    model_str: &str,
    project_root: &Path,
    agents: &AgentPool,
) -> String {
    // Find the story file across all pipeline stages.
    // Try the content store / CRDT state first, then fall back to filesystem.
    let mut found: Option<(std::path::PathBuf, String)> = None;

    // --- DB-first lookup ---
    for id in crate::db::all_content_ids() {
        let file_num = id.split('_').next().unwrap_or("");
        if file_num == story_number && let Ok(Some(item)) = crate::pipeline_state::read_typed(&id) {
            let path = project_root
                .join(".huskies")
                .join("work")
                .join(item.stage.dir_name())
                .join(format!("{id}.md"));
            found = Some((path, id));
            break;
        }
    }

    // --- Filesystem fallback ---
    if found.is_none() {
        'outer: for stage in STAGES {
            let dir = project_root.join(".huskies").join("work").join(stage);
            if !dir.exists() {
                continue;
            }
            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") {
                        continue;
                    }
                    if let Some(stem) = path
                        .file_stem()
                        .and_then(|s| s.to_str())
                        .map(|s| s.to_string())
                    {
                        let file_num = stem
                            .split('_')
                            .next()
                            .filter(|s| !s.is_empty() && s.chars().all(|c| c.is_ascii_digit()))
                            .unwrap_or("")
                            .to_string();
                        if file_num == story_number {
                            found = Some((path, stem));
                            break 'outer;
                        }
                    }
                }
            }
        }
    }

    let (path, story_id) = match found {
        Some(f) => f,
        None => {
            return format!(
                "No story, bug, or spike with number **{story_number}** found."
            );
        }
    };

    // Read the human-readable name from front matter for the response.
    // Try the content store first, then fall back to reading from disk.
    let story_name = crate::db::read_content(&story_id)
        .or_else(|| std::fs::read_to_string(&path).ok())
        .and_then(|contents| {
            parse_front_matter(&contents)
                .ok()
                .and_then(|m| m.name)
        })
        .unwrap_or_else(|| story_id.clone());

    let agent_name = resolve_agent_name(model_str);

    // Write `agent: <agent_name>` into the story's front matter.
    let write_result = std::fs::read_to_string(&path)
        .map_err(|e| format!("Failed to read story file: {e}"))
        .and_then(|contents| {
            let updated = set_front_matter_field(&contents, "agent", &agent_name);
            std::fs::write(&path, &updated)
                .map_err(|e| format!("Failed to write story file: {e}"))
        });

    if let Err(e) = write_result {
        return format!("Failed to assign model to **{story_name}**: {e}");
    }

    // Check whether a coder is already running on this story.
    let running_coders: Vec<_> = agents
        .list_agents()
        .unwrap_or_default()
        .into_iter()
        .filter(|a| {
            a.story_id == story_id
                && a.agent_name.starts_with("coder")
                && matches!(a.status, AgentStatus::Running | AgentStatus::Pending)
        })
        .collect();

    if running_coders.is_empty() {
        // No coder running — just persist the assignment.
        return format!(
            "Assigned **{agent_name}** to **{story_name}** (story {story_number}). \
             The model will be used when the story starts."
        );
    }

    // Stop each running coder, then start the newly assigned one.
    let stopped: Vec<String> = running_coders
        .iter()
        .map(|a| a.agent_name.clone())
        .collect();

    for coder in &running_coders {
        if let Err(e) = agents
            .stop_agent(project_root, &story_id, &coder.agent_name)
            .await
        {
            crate::slog!(
                "[matrix-bot] assign: failed to stop agent {} for {}: {e}",
                coder.agent_name,
                story_id
            );
        }
    }

    crate::slog!(
        "[matrix-bot] assign (bot={bot_name}): stopped {:?} for {}; starting {agent_name}",
        stopped,
        story_id
    );

    match agents
        .start_agent(project_root, &story_id, Some(&agent_name), None)
        .await
    {
        Ok(info) => {
            format!(
                "Reassigned **{story_name}** (story {story_number}): \
                 stopped **{}** and started **{}**.",
                stopped.join(", "),
                info.agent_name
            )
        }
        Err(e) => {
            format!(
                "Assigned **{agent_name}** to **{story_name}** (story {story_number}): \
                 stopped **{}** but failed to start the new agent: {e}",
                stopped.join(", ")
            )
        }
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;

    // -- extract_assign_command -----------------------------------------------

    #[test]
    fn extract_with_full_user_id() {
        let cmd = extract_assign_command(
            "@timmy:home.local assign 42 opus",
            "Timmy",
            "@timmy:home.local",
        );
        assert_eq!(
            cmd,
            Some(AssignCommand::Assign {
                story_number: "42".to_string(),
                model: "opus".to_string()
            })
        );
    }

    #[test]
    fn extract_with_display_name() {
        let cmd = extract_assign_command("Timmy assign 42 sonnet", "Timmy", "@timmy:home.local");
        assert_eq!(
            cmd,
            Some(AssignCommand::Assign {
                story_number: "42".to_string(),
                model: "sonnet".to_string()
            })
        );
    }

    #[test]
    fn extract_with_localpart() {
        let cmd = extract_assign_command("@timmy assign 7 opus", "Timmy", "@timmy:home.local");
        assert_eq!(
            cmd,
            Some(AssignCommand::Assign {
                story_number: "7".to_string(),
                model: "opus".to_string()
            })
        );
    }

    #[test]
    fn extract_case_insensitive_command() {
        let cmd = extract_assign_command("Timmy ASSIGN 99 opus", "Timmy", "@timmy:home.local");
        assert_eq!(
            cmd,
            Some(AssignCommand::Assign {
                story_number: "99".to_string(),
                model: "opus".to_string()
            })
        );
    }

    #[test]
    fn extract_no_args_is_bad_args() {
        let cmd = extract_assign_command("Timmy assign", "Timmy", "@timmy:home.local");
        assert_eq!(cmd, Some(AssignCommand::BadArgs));
    }

    #[test]
    fn extract_missing_model_is_bad_args() {
        let cmd = extract_assign_command("Timmy assign 42", "Timmy", "@timmy:home.local");
        assert_eq!(cmd, Some(AssignCommand::BadArgs));
    }

    #[test]
    fn extract_non_numeric_number_is_bad_args() {
        let cmd = extract_assign_command("Timmy assign abc opus", "Timmy", "@timmy:home.local");
        assert_eq!(cmd, Some(AssignCommand::BadArgs));
    }

    #[test]
    fn extract_non_assign_command_returns_none() {
        let cmd = extract_assign_command("Timmy help", "Timmy", "@timmy:home.local");
        assert_eq!(cmd, None);
    }

    #[test]
    fn extract_assign_command_multibyte_prefix_no_panic() {
        // "xxxx⏺ assign 42 opus" — ⏺ (U+23FA) is 3 bytes, starting at byte 4.
        // "@timmy" has len 6 so text[..6] lands inside ⏺ — panics without the fix.
        let cmd = extract_assign_command(
            "xxxx\u{23FA} assign 42 opus",
            "Timmy",
            "@timmy:home.local",
        );
        assert_eq!(cmd, None);
    }

    // -- resolve_agent_name --------------------------------------------------

    #[test]
    fn resolve_agent_name_prefixes_bare_model() {
        assert_eq!(resolve_agent_name("opus"), "coder-opus");
        assert_eq!(resolve_agent_name("sonnet"), "coder-sonnet");
        assert_eq!(resolve_agent_name("haiku"), "coder-haiku");
    }

    #[test]
    fn resolve_agent_name_does_not_double_prefix() {
        assert_eq!(resolve_agent_name("coder-opus"), "coder-opus");
        assert_eq!(resolve_agent_name("coder-sonnet"), "coder-sonnet");
    }

    // -- handle_assign (no running coder) ------------------------------------

    use crate::chat::test_helpers::write_story_file;

    #[tokio::test]
    async fn handle_assign_returns_not_found_for_unknown_number() {
        let tmp = tempfile::tempdir().unwrap();
        for stage in STAGES {
            std::fs::create_dir_all(tmp.path().join(".huskies/work").join(stage)).unwrap();
        }
        let agents = std::sync::Arc::new(AgentPool::new_test(3000));
        let response = handle_assign("Timmy", "999", "opus", tmp.path(), &agents).await;
        assert!(
            response.contains("No story") && response.contains("999"),
            "unexpected response: {response}"
        );
    }

    #[tokio::test]
    async fn handle_assign_writes_front_matter_when_no_coder_running() {
        let tmp = tempfile::tempdir().unwrap();
        write_story_file(
            tmp.path(),
            "1_backlog",
            "42_story_test.md",
            "---\nname: Test Feature\n---\n\n# Story 42\n",
        );

        let agents = std::sync::Arc::new(AgentPool::new_test(3000));
        let response = handle_assign("Timmy", "42", "opus", tmp.path(), &agents).await;

        assert!(
            response.contains("coder-opus"),
            "response should mention agent: {response}"
        );
        assert!(
            response.contains("Test Feature"),
            "response should mention story name: {response}"
        );
        // Should say "will be used when the story starts" (no restart)
        assert!(
            response.contains("start"),
            "response should indicate assignment for future start: {response}"
        );

        let contents = std::fs::read_to_string(
            tmp.path().join(".huskies/work/1_backlog/42_story_test.md"),
        )
        .unwrap();
        assert!(
            contents.contains("agent: coder-opus"),
            "front matter should contain agent field: {contents}"
        );
    }

    #[tokio::test]
    async fn handle_assign_with_already_prefixed_name_does_not_double_prefix() {
        let tmp = tempfile::tempdir().unwrap();
        write_story_file(
            tmp.path(),
            "1_backlog",
            "7_story_small.md",
            "---\nname: Small Story\n---\n",
        );

        let agents = std::sync::Arc::new(AgentPool::new_test(3000));
        let response = handle_assign("Timmy", "7", "coder-opus", tmp.path(), &agents).await;

        assert!(
            response.contains("coder-opus"),
            "should not double-prefix: {response}"
        );
        assert!(
            !response.contains("coder-coder-opus"),
            "must not double-prefix: {response}"
        );

        let contents = std::fs::read_to_string(
            tmp.path().join(".huskies/work/1_backlog/7_story_small.md"),
        )
        .unwrap();
        assert!(
            contents.contains("agent: coder-opus"),
            "must write coder-opus, not coder-coder-opus: {contents}"
        );
    }

    #[tokio::test]
    async fn handle_assign_overwrites_existing_agent_field() {
        let tmp = tempfile::tempdir().unwrap();
        write_story_file(
            tmp.path(),
            "1_backlog",
            "5_story_existing.md",
            "---\nname: Existing\nagent: coder-sonnet\n---\n",
        );

        let agents = std::sync::Arc::new(AgentPool::new_test(3000));
        handle_assign("Timmy", "5", "opus", tmp.path(), &agents).await;

        let contents = std::fs::read_to_string(
            tmp.path().join(".huskies/work/1_backlog/5_story_existing.md"),
        )
        .unwrap();
        assert!(
            contents.contains("agent: coder-opus"),
            "should overwrite old agent: {contents}"
        );
        assert!(
            !contents.contains("coder-sonnet"),
            "old agent should no longer appear: {contents}"
        );
    }

    #[tokio::test]
    async fn handle_assign_finds_story_in_any_stage() {
        let tmp = tempfile::tempdir().unwrap();
        write_story_file(
            tmp.path(),
            "3_qa",
            "99_story_in_qa.md",
            "---\nname: In QA\n---\n",
        );

        let agents = std::sync::Arc::new(AgentPool::new_test(3000));
        let response = handle_assign("Timmy", "99", "opus", tmp.path(), &agents).await;
        assert!(
            response.contains("coder-opus"),
            "should find story in qa stage: {response}"
        );
    }

    // -- handle_assign (with running coder) ----------------------------------

    #[tokio::test]
    async fn handle_assign_stops_running_coder_and_reports_reassignment() {
        let tmp = tempfile::tempdir().unwrap();
        write_story_file(
            tmp.path(),
            "2_current",
            "10_story_current.md",
            "---\nname: Current Story\nagent: coder-sonnet\n---\n",
        );

        let agents = std::sync::Arc::new(AgentPool::new_test(3000));
        // Inject a running coder for this story.
        agents.inject_test_agent("10_story_current", "coder-sonnet", AgentStatus::Running);

        let response = handle_assign("Timmy", "10", "opus", tmp.path(), &agents).await;

        // The response should mention both stopped and started agents.
        assert!(
            response.contains("coder-sonnet"),
            "response should mention the stopped agent: {response}"
        );
        // Should indicate a restart occurred (not just "will be used when starts")
        assert!(
            response.to_lowercase().contains("stop") || response.to_lowercase().contains("reassign"),
            "response should indicate stop/reassign: {response}"
        );
    }
}