task_mcp/just/

mod.rs

pub mod model;

use std::collections::{HashMap, VecDeque};
use std::path::{Path, PathBuf};
use std::sync::Mutex;
use std::time::{Duration, SystemTime, UNIX_EPOCH};

use regex::Regex;
use tokio::process::Command;
use uuid::Uuid;

use crate::config::TaskMode;
use crate::just::model::{
    JustDump, JustRecipe, Recipe, TaskError, TaskExecution, TaskExecutionSummary,
};

// =============================================================================
// Error
// =============================================================================

#[derive(Debug, thiserror::Error)]
pub enum JustError {
    #[error("just command not found: {0}")]
    NotFound(String),
    #[error("just command failed (exit {code}): {stderr}")]
    CommandFailed { code: i32, stderr: String },
    #[error("failed to parse just dump json: {0}")]
    ParseError(#[from] serde_json::Error),
    #[error("I/O error while reading justfile: {0}")]
    Io(#[from] std::io::Error),
}

// =============================================================================
// Public API
// =============================================================================

/// Discover recipes from the justfile at `justfile_path`.
///
/// Filtering behaviour depends on `mode`:
/// - `TaskMode::AgentOnly`: only recipes marked agent-safe are returned.
/// - `TaskMode::All`: all non-private recipes are returned.
///
/// Agent-safe detection: patten A (group attribute) first, pattern B (comment regex) as fallback.
pub async fn list_recipes(
    justfile_path: &Path,
    mode: &TaskMode,
    workdir: Option<&Path>,
) -> Result<Vec<Recipe>, JustError> {
    let dump = dump_json(justfile_path, workdir).await?;

    // Read justfile text for pattern B fallback
    let justfile_text = tokio::fs::read_to_string(justfile_path).await.ok();

    // Build comment-tag set from justfile text (pattern B)
    let comment_tagged = justfile_text
        .as_deref()
        .map(extract_comment_tagged_recipes)
        .unwrap_or_default();

    let mut recipes: Vec<Recipe> = dump
        .recipes
        .into_values()
        .filter(|r| !r.private)
        .map(|raw| {
            let allow_agent = is_allow_agent(&raw, &comment_tagged);
            Recipe::from_just_recipe(raw, allow_agent)
        })
        .collect();

    // Sort by name for deterministic output
    recipes.sort_by(|a, b| a.name.cmp(&b.name));

    match mode {
        TaskMode::AgentOnly => Ok(recipes.into_iter().filter(|r| r.allow_agent).collect()),
        TaskMode::All => Ok(recipes),
    }
}

// =============================================================================
// Internal helpers
// =============================================================================

/// Run `just --dump --dump-format json --unstable` and return parsed output.
async fn dump_json(justfile_path: &Path, workdir: Option<&Path>) -> Result<JustDump, JustError> {
    let mut cmd = Command::new("just");
    cmd.arg("--justfile")
        .arg(justfile_path)
        .arg("--dump")
        .arg("--dump-format")
        .arg("json")
        .arg("--unstable");
    if let Some(dir) = workdir {
        cmd.current_dir(dir);
    }
    let output = cmd
        .output()
        .await
        .map_err(|e| JustError::NotFound(e.to_string()))?;

    if !output.status.success() {
        let code = output.status.code().unwrap_or(-1);
        let stderr = String::from_utf8_lossy(&output.stderr).into_owned();
        return Err(JustError::CommandFailed { code, stderr });
    }

    let json_str = String::from_utf8_lossy(&output.stdout);
    let dump: JustDump = serde_json::from_str(&json_str)?;
    Ok(dump)
}

/// Pattern A: check if the recipe has a `[group: 'agent']` attribute.
fn has_group_agent_attribute(recipe: &JustRecipe) -> bool {
    recipe.attributes.iter().any(|a| {
        a.group
            .as_ref()
            .map(|g| g.as_str() == "agent")
            .unwrap_or(false)
    })
}

/// Pattern B: extract recipe names that are preceded by a `# [allow-agent]` comment.
///
/// Regex matches lines like (justfile syntax, not Rust):
/// ```text
/// # [allow-agent]
/// recipe_name:
/// ```
fn extract_comment_tagged_recipes(justfile_text: &str) -> std::collections::HashSet<String> {
    // Allow-agent comment followed (with optional blank/doc lines) by a recipe line
    // Simple approach: line-by-line scan with state
    let allow_agent_re = Regex::new(r"^\s*#\s*\[allow-agent\]").expect("valid regex");
    let recipe_name_re = Regex::new(r"^([a-zA-Z0-9_-]+)\s*(?:\S.*)?:").expect("valid regex");

    let mut tagged = std::collections::HashSet::new();
    let mut saw_allow_agent = false;

    for line in justfile_text.lines() {
        if allow_agent_re.is_match(line) {
            saw_allow_agent = true;
            continue;
        }

        if saw_allow_agent {
            // Skip doc comment lines and blank lines between tag and recipe
            let trimmed = line.trim();
            if trimmed.is_empty() || trimmed.starts_with('#') {
                continue;
            }

            // Try to extract recipe name
            if let Some(cap) = recipe_name_re.captures(line) {
                let name = cap[1].to_string();
                // Exclude attribute lines like [group: 'agent']
                if !name.starts_with('[') {
                    tagged.insert(name);
                }
            }
            // Once we see a non-comment, non-blank line, reset state
            saw_allow_agent = false;
        }
    }

    tagged
}

/// Determine if a recipe is agent-safe using pattern A first, then pattern B.
fn is_allow_agent(recipe: &JustRecipe, comment_tagged: &std::collections::HashSet<String>) -> bool {
    has_group_agent_attribute(recipe) || comment_tagged.contains(&recipe.name)
}

/// Resolve justfile path from an optional override, workdir, or the current directory.
pub fn resolve_justfile_path(override_path: Option<&str>, workdir: Option<&Path>) -> PathBuf {
    match override_path {
        Some(p) => PathBuf::from(p),
        None => match workdir {
            Some(dir) => dir.join("justfile"),
            None => PathBuf::from("justfile"),
        },
    }
}

// =============================================================================
// Output truncation
// =============================================================================

const MAX_OUTPUT_BYTES: usize = 100 * 1024; // 100 KB
const HEAD_BYTES: usize = 50 * 1024; // 50 KB
const TAIL_BYTES: usize = 50 * 1024; // 50 KB

/// Truncate output to at most `MAX_OUTPUT_BYTES`.
///
/// If truncation is necessary the result contains:
/// `{head}\n...[truncated {n} bytes]...\n{tail}`
///
/// UTF-8 multi-byte boundaries are respected — the slice points are adjusted
/// so that we never split a multi-byte character.
pub fn truncate_output(output: &str) -> (String, bool) {
    if output.len() <= MAX_OUTPUT_BYTES {
        return (output.to_string(), false);
    }

    // Find safe byte boundary for the head (≤ HEAD_BYTES)
    let head_end = safe_byte_boundary(output, HEAD_BYTES);
    // Find safe byte boundary for the tail (last TAIL_BYTES)
    let tail_start_raw = output.len().saturating_sub(TAIL_BYTES);
    let tail_start = safe_tail_start(output, tail_start_raw);

    let head = &output[..head_end];
    let tail = &output[tail_start..];
    let truncated_bytes = output.len() - head_end - (output.len() - tail_start);

    (
        format!("{head}\n...[truncated {truncated_bytes} bytes]...\n{tail}"),
        true,
    )
}

/// Find the largest byte index `<= limit` that lies on a UTF-8 character boundary.
fn safe_byte_boundary(s: &str, limit: usize) -> usize {
    if limit >= s.len() {
        return s.len();
    }
    // Walk backwards from `limit` until we hit a valid char boundary
    let mut idx = limit;
    while idx > 0 && !s.is_char_boundary(idx) {
        idx -= 1;
    }
    idx
}

/// Find the smallest byte index `>= hint` that lies on a UTF-8 character boundary.
fn safe_tail_start(s: &str, hint: usize) -> usize {
    if hint >= s.len() {
        return s.len();
    }
    let mut idx = hint;
    while idx < s.len() && !s.is_char_boundary(idx) {
        idx += 1;
    }
    idx
}

// =============================================================================
// Argument validation
// =============================================================================

/// Reject argument values that contain shell meta-characters.
///
/// `tokio::process::Command` bypasses the shell, but `just` itself invokes a
/// shell interpreter for recipe bodies.  Validating inputs here prevents
/// injection attacks in case a recipe passes an argument through to the shell.
pub fn validate_arg_value(value: &str) -> Result<(), TaskError> {
    const DANGEROUS: &[&str] = &[";", "|", "&&", "||", "`", "$(", "${", "\n", "\r"];
    for pattern in DANGEROUS {
        if value.contains(pattern) {
            return Err(TaskError::DangerousArgument(value.to_string()));
        }
    }
    Ok(())
}

// =============================================================================
// Recipe execution
// =============================================================================

/// Execute a recipe by name, passing `args` as positional parameters.
///
/// Steps:
/// 1. Confirm the recipe exists in `list_recipes(justfile_path, mode)`.
/// 2. Validate each argument value for dangerous characters.
/// 3. Run `just --justfile {path} {recipe_name} {arg_values...}` with a
///    timeout.
/// 4. Capture stdout/stderr and apply truncation.
/// 5. Return a `TaskExecution` record.
pub async fn execute_recipe(
    recipe_name: &str,
    args: &HashMap<String, String>,
    justfile_path: &Path,
    timeout: Duration,
    mode: &TaskMode,
    workdir: Option<&Path>,
) -> Result<TaskExecution, TaskError> {
    // 1. Whitelist check
    let recipes = list_recipes(justfile_path, mode, workdir).await?;
    let recipe = recipes
        .iter()
        .find(|r| r.name == recipe_name)
        .ok_or_else(|| TaskError::RecipeNotFound(recipe_name.to_string()))?;

    // 2. Argument validation
    for value in args.values() {
        validate_arg_value(value)?;
    }

    // Build positional argument list in parameter definition order
    let positional: Vec<&str> = recipe
        .parameters
        .iter()
        .filter_map(|p| args.get(&p.name).map(|v| v.as_str()))
        .collect();

    // 3. Construct and run the command
    let started_at = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs();
    let start_instant = std::time::Instant::now();

    let mut cmd = Command::new("just");
    cmd.arg("--justfile").arg(justfile_path).arg(recipe_name);
    for arg in &positional {
        cmd.arg(arg);
    }
    if let Some(dir) = workdir {
        cmd.current_dir(dir);
    }

    let run_result = tokio::time::timeout(timeout, cmd.output()).await;

    let duration_ms = start_instant.elapsed().as_millis() as u64;

    let output = match run_result {
        Err(_) => return Err(TaskError::Timeout),
        Ok(Err(io_err)) => return Err(TaskError::Io(io_err)),
        Ok(Ok(out)) => out,
    };

    let exit_code = output.status.code();

    // 4. Capture + truncate
    let raw_stdout = String::from_utf8_lossy(&output.stdout).into_owned();
    let raw_stderr = String::from_utf8_lossy(&output.stderr).into_owned();

    let (stdout, stdout_truncated) = truncate_output(&raw_stdout);
    let (stderr, stderr_truncated) = truncate_output(&raw_stderr);
    let truncated = stdout_truncated || stderr_truncated;

    // 5. Return execution record
    Ok(TaskExecution {
        id: Uuid::new_v4().to_string(),
        task_name: recipe_name.to_string(),
        args: args.clone(),
        exit_code,
        stdout,
        stderr,
        started_at,
        duration_ms,
        truncated,
    })
}

// =============================================================================
// Task log store
// =============================================================================

/// In-memory ring buffer of recent task executions.
///
/// `Arc<TaskLogStore>` is `Clone` because `Arc<T>` implements `Clone` for any
/// `T: ?Sized`.  `TaskLogStore` itself does not need to implement `Clone`.
pub struct TaskLogStore {
    logs: Mutex<VecDeque<TaskExecution>>,
    max_entries: usize,
}

impl TaskLogStore {
    pub fn new(max_entries: usize) -> Self {
        Self {
            logs: Mutex::new(VecDeque::new()),
            max_entries,
        }
    }

    /// Append an execution record, evicting the oldest entry when full.
    pub fn push(&self, execution: TaskExecution) {
        let mut guard = self.logs.lock().expect("log store lock poisoned");
        if guard.len() >= self.max_entries {
            guard.pop_front();
        }
        guard.push_back(execution);
    }

    /// Look up a specific execution by ID.  Returns a clone.
    pub fn get(&self, id: &str) -> Option<TaskExecution> {
        let guard = self.logs.lock().expect("log store lock poisoned");
        guard.iter().find(|e| e.id == id).cloned()
    }

    /// Return summaries of the most recent `n` executions (newest first).
    pub fn recent(&self, n: usize) -> Vec<TaskExecutionSummary> {
        let guard = self.logs.lock().expect("log store lock poisoned");
        guard
            .iter()
            .rev()
            .take(n)
            .map(TaskExecutionSummary::from_execution)
            .collect()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::just::model::RecipeAttribute;

    fn make_recipe(name: &str, attributes: Vec<RecipeAttribute>) -> JustRecipe {
        crate::just::model::JustRecipe {
            name: name.to_string(),
            namepath: name.to_string(),
            doc: None,
            attributes,
            parameters: vec![],
            private: false,
            quiet: false,
        }
    }

    #[test]
    fn has_group_agent_attribute_true() {
        let recipe = make_recipe(
            "build",
            vec![RecipeAttribute {
                group: Some("agent".to_string()),
            }],
        );
        assert!(has_group_agent_attribute(&recipe));
    }

    #[test]
    fn has_group_agent_attribute_false_no_attrs() {
        let recipe = make_recipe("deploy", vec![]);
        assert!(!has_group_agent_attribute(&recipe));
    }

    #[test]
    fn has_group_agent_attribute_false_other_group() {
        let recipe = make_recipe(
            "build",
            vec![RecipeAttribute {
                group: Some("ci".to_string()),
            }],
        );
        assert!(!has_group_agent_attribute(&recipe));
    }

    #[test]
    fn extract_comment_tagged_recipes_basic() {
        let text = "# [allow-agent]\nbuild:\n    cargo build\n\ndeploy:\n    ./deploy.sh\n";
        let tagged = extract_comment_tagged_recipes(text);
        assert!(tagged.contains("build"), "build should be tagged");
        assert!(!tagged.contains("deploy"), "deploy should not be tagged");
    }

    #[test]
    fn extract_comment_tagged_recipes_with_doc_comment() {
        let text = "# [allow-agent]\n# Run tests\ntest filter=\"\":\n    cargo test {{filter}}\n\ndeploy:\n    ./deploy.sh\n";
        let tagged = extract_comment_tagged_recipes(text);
        assert!(tagged.contains("test"), "test should be tagged");
        assert!(!tagged.contains("deploy"));
    }

    #[test]
    fn extract_comment_tagged_recipes_multiple() {
        let text = "# [allow-agent]\nbuild:\n    cargo build\n\n# [allow-agent]\ninfo:\n    echo info\n\ndeploy:\n    ./deploy.sh\n";
        let tagged = extract_comment_tagged_recipes(text);
        assert!(tagged.contains("build"));
        assert!(tagged.contains("info"));
        assert!(!tagged.contains("deploy"));
    }

    #[test]
    fn is_allow_agent_pattern_a() {
        let tagged = std::collections::HashSet::new();
        // not in comment set
        let recipe = make_recipe(
            "build",
            vec![RecipeAttribute {
                group: Some("agent".to_string()),
            }],
        );
        assert!(is_allow_agent(&recipe, &tagged));
    }

    #[test]
    fn is_allow_agent_pattern_b() {
        let mut tagged = std::collections::HashSet::new();
        tagged.insert("build".to_string());
        let recipe = make_recipe("build", vec![]);
        assert!(is_allow_agent(&recipe, &tagged));
    }

    #[test]
    fn is_allow_agent_neither() {
        let tagged = std::collections::HashSet::new();
        let recipe = make_recipe("deploy", vec![]);
        assert!(!is_allow_agent(&recipe, &tagged));
    }

    #[test]
    fn resolve_justfile_path_override() {
        let p = resolve_justfile_path(Some("/custom/justfile"), None);
        assert_eq!(p, PathBuf::from("/custom/justfile"));
    }

    #[test]
    fn resolve_justfile_path_default() {
        let p = resolve_justfile_path(None, None);
        assert_eq!(p, PathBuf::from("justfile"));
    }

    #[test]
    fn resolve_justfile_path_with_workdir() {
        let workdir = Path::new("/some/project");
        let p = resolve_justfile_path(None, Some(workdir));
        assert_eq!(p, PathBuf::from("/some/project/justfile"));
    }

    #[test]
    fn resolve_justfile_path_override_ignores_workdir() {
        // override_path takes precedence over workdir
        let workdir = Path::new("/some/project");
        let p = resolve_justfile_path(Some("/custom/justfile"), Some(workdir));
        assert_eq!(p, PathBuf::from("/custom/justfile"));
    }

    // -------------------------------------------------------------------------
    // truncate_output tests
    // -------------------------------------------------------------------------

    #[test]
    fn truncate_output_short_input_unchanged() {
        let input = "hello";
        let (result, truncated) = truncate_output(input);
        assert!(!truncated);
        assert_eq!(result, input);
    }

    #[test]
    fn truncate_output_long_input_truncated() {
        // Create a string longer than MAX_OUTPUT_BYTES (100 KB)
        let input = "x".repeat(200 * 1024);
        let (result, truncated) = truncate_output(&input);
        assert!(truncated);
        assert!(result.contains("...[truncated"));
        // Result should be smaller than the input
        assert!(result.len() < input.len());
    }

    #[test]
    fn truncate_output_utf8_boundary() {
        // Build a string that is just over HEAD_BYTES using multi-byte chars
        // Each '日' is 3 bytes; we need HEAD_BYTES+1 bytes to trigger truncation
        let char_3bytes = '日';
        // Fill slightly above MAX_OUTPUT_BYTES boundary
        let count = (MAX_OUTPUT_BYTES / 3) + 10;
        let input: String = std::iter::repeat(char_3bytes).take(count).collect();
        let (result, truncated) = truncate_output(&input);
        // Verify the result is valid UTF-8 (no panic = success)
        assert!(std::str::from_utf8(result.as_bytes()).is_ok());
        if truncated {
            assert!(result.contains("...[truncated"));
        }
    }

    // -------------------------------------------------------------------------
    // validate_arg_value tests
    // -------------------------------------------------------------------------

    #[test]
    fn validate_arg_value_safe_values() {
        assert!(validate_arg_value("hello world").is_ok());
        assert!(validate_arg_value("value_123-abc").is_ok());
        assert!(validate_arg_value("path/to/file.txt").is_ok());
    }

    #[test]
    fn validate_arg_value_semicolon_rejected() {
        assert!(validate_arg_value("foo; rm -rf /").is_err());
    }

    #[test]
    fn validate_arg_value_pipe_rejected() {
        assert!(validate_arg_value("foo | cat /etc/passwd").is_err());
    }

    #[test]
    fn validate_arg_value_and_and_rejected() {
        assert!(validate_arg_value("foo && evil").is_err());
    }

    #[test]
    fn validate_arg_value_backtick_rejected() {
        assert!(validate_arg_value("foo`id`").is_err());
    }

    #[test]
    fn validate_arg_value_dollar_paren_rejected() {
        assert!(validate_arg_value("$(id)").is_err());
    }

    #[test]
    fn validate_arg_value_newline_rejected() {
        assert!(validate_arg_value("foo\nbar").is_err());
    }

    // -------------------------------------------------------------------------
    // TaskLogStore tests
    // -------------------------------------------------------------------------

    fn make_execution(id: &str, task_name: &str) -> TaskExecution {
        TaskExecution {
            id: id.to_string(),
            task_name: task_name.to_string(),
            args: HashMap::new(),
            exit_code: Some(0),
            stdout: "".to_string(),
            stderr: "".to_string(),
            started_at: 0,
            duration_ms: 0,
            truncated: false,
        }
    }

    #[test]
    fn task_log_store_push_and_get() {
        let store = TaskLogStore::new(10);
        let exec = make_execution("id-1", "build");
        store.push(exec);
        let retrieved = store.get("id-1").expect("should find id-1");
        assert_eq!(retrieved.task_name, "build");
    }

    #[test]
    fn task_log_store_get_missing() {
        let store = TaskLogStore::new(10);
        assert!(store.get("nonexistent").is_none());
    }

    #[test]
    fn task_log_store_evicts_oldest_when_full() {
        let store = TaskLogStore::new(3);
        store.push(make_execution("id-1", "a"));
        store.push(make_execution("id-2", "b"));
        store.push(make_execution("id-3", "c"));
        store.push(make_execution("id-4", "d")); // evicts id-1
        assert!(store.get("id-1").is_none(), "id-1 should be evicted");
        assert!(store.get("id-4").is_some(), "id-4 should exist");
    }

    #[test]
    fn task_log_store_recent_newest_first() {
        let store = TaskLogStore::new(10);
        store.push(make_execution("id-1", "a"));
        store.push(make_execution("id-2", "b"));
        store.push(make_execution("id-3", "c"));
        let recent = store.recent(2);
        assert_eq!(recent.len(), 2);
        assert_eq!(recent[0].id, "id-3", "newest should be first");
        assert_eq!(recent[1].id, "id-2");
    }

    #[test]
    fn task_log_store_recent_n_larger_than_store() {
        let store = TaskLogStore::new(10);
        store.push(make_execution("id-1", "a"));
        let recent = store.recent(5);
        assert_eq!(recent.len(), 1);
    }
}