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 tokio::process::Command;
use uuid::Uuid;

use crate::config::TaskMode;
use crate::just::model::{
    JustDump, JustRecipe, Recipe, RecipeSource, 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 is based entirely on the JSON output of `just --dump`:
/// 1. `[group('allow-agent')]` attribute (preferred, deterministic), or
/// 2. `[allow-agent]` marker inside the recipe's doc comment (legacy fallback).
pub async fn list_recipes(
    justfile_path: &Path,
    mode: &TaskMode,
    workdir: Option<&Path>,
) -> Result<Vec<Recipe>, JustError> {
    list_recipes_with_source(justfile_path, mode, workdir, RecipeSource::Project).await
}

/// Discover and merge recipes from a project justfile and an optional global justfile.
///
/// Merge semantics:
/// - Both sides are filtered by `mode` independently.
/// - Project recipes have `source = Project`; global recipes have `source = Global`.
/// - On name collision, the project recipe wins (global is hidden).
/// - Result order: project recipes (alphabetical) followed by global-only recipes (alphabetical).
///
/// When `global_path` is `None`, this is equivalent to `list_recipes`.
pub async fn list_recipes_merged(
    project_path: &Path,
    global_path: Option<&Path>,
    mode: &TaskMode,
    project_workdir: Option<&Path>,
) -> Result<Vec<Recipe>, JustError> {
    // Collect project recipes, tagging them as Project source.
    // If the project justfile does not exist (e.g. new project with global-only setup),
    // treat it as an empty list rather than propagating a "file not found" error.
    let project_recipes = if tokio::fs::metadata(project_path).await.is_ok() {
        list_recipes_with_source(project_path, mode, project_workdir, RecipeSource::Project).await?
    } else {
        Vec::new()
    };

    let global_path = match global_path {
        Some(p) => p,
        None => return Ok(project_recipes),
    };

    // Collect global recipes, tagging them as Global source.
    // Use project_workdir as workdir so global recipes run in project context.
    let global_recipes =
        list_recipes_with_source(global_path, mode, project_workdir, RecipeSource::Global).await?;

    // Build name set from project recipes for override detection.
    let project_names: std::collections::HashSet<&str> =
        project_recipes.iter().map(|r| r.name.as_str()).collect();

    // Global-only recipes (not overridden by project).
    let global_only: Vec<Recipe> = global_recipes
        .into_iter()
        .filter(|r| !project_names.contains(r.name.as_str()))
        .collect();

    // Result: project first, then global-only.
    let mut merged = project_recipes;
    merged.extend(global_only);
    Ok(merged)
}

/// Internal helper: like `list_recipes` but tags each recipe with `source`.
async fn list_recipes_with_source(
    justfile_path: &Path,
    mode: &TaskMode,
    workdir: Option<&Path>,
    source: RecipeSource,
) -> Result<Vec<Recipe>, JustError> {
    let dump = dump_json(justfile_path, workdir).await?;

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

    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('allow-agent')]` attribute.
fn has_allow_agent_group_attribute(recipe: &JustRecipe) -> bool {
    recipe
        .attributes
        .iter()
        .any(|a| a.group() == Some("allow-agent"))
}

/// Pattern B: legacy `# [allow-agent]` marker embedded in the recipe doc comment.
///
/// `just --dump` preserves the comment line immediately above a recipe as its
/// `doc` field. When users tag recipes with a `# [allow-agent]` line, that text
/// shows up here, so we can detect it without re-reading the source file.
///
/// Note: `just` only keeps the closest comment line as the doc, so combining
/// `# [allow-agent]` with a descriptive doc comment causes one of the two to be
/// dropped. Prefer the `[group('allow-agent')]` attribute for new recipes.
fn has_allow_agent_doc(recipe: &JustRecipe) -> bool {
    recipe
        .doc
        .as_deref()
        .is_some_and(|d| d.split_whitespace().any(|t| t == "[allow-agent]"))
}

/// Determine if a recipe is agent-safe via group attribute or legacy doc marker.
fn is_allow_agent(recipe: &JustRecipe) -> bool {
    has_allow_agent_group_attribute(recipe) || has_allow_agent_doc(recipe)
}

/// 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 control characters.
///
/// `tokio::process::Command` bypasses the shell, and arguments are passed to
/// `just` as OS-level argv entries without shell interpretation.  Shell
/// metacharacter validation is therefore unnecessary and would block legitimate
/// values (URLs with `&`, jq filters with `|`, template values with `${}`, etc.).
///
/// Control characters (`\n`, `\r`) are still rejected as they are invalid in
/// single-value arguments and can cause log injection or unexpected behavior
/// in recipe body text substitution.
pub fn validate_arg_value(value: &str) -> Result<(), TaskError> {
    const REJECTED_CONTROL_CHARS: &[&str] = &["\n", "\r"];
    for pattern in REJECTED_CONTROL_CHARS {
        if value.contains(pattern) {
            return Err(TaskError::DangerousArgument(value.to_string()));
        }
    }
    Ok(())
}

/// Environment variable prefix for content arguments.
const CONTENT_ENV_PREFIX: &str = "TASK_MCP_CONTENT_";

/// Validate that a content key is a valid environment variable suffix.
///
/// Must match `[A-Za-z][A-Za-z0-9_]*`: start with an ASCII letter, followed by
/// zero or more ASCII letters, digits, or underscores.  Keys are uppercased and
/// prepended with `TASK_MCP_CONTENT_` before being set as environment variables.
pub fn validate_content_key(key: &str) -> Result<(), TaskError> {
    let mut chars = key.chars();
    match chars.next() {
        None => return Err(TaskError::InvalidContentKey(key.to_string())),
        Some(c) if !c.is_ascii_alphabetic() => {
            return Err(TaskError::InvalidContentKey(key.to_string()));
        }
        _ => {}
    }
    for c in chars {
        if !c.is_ascii_alphanumeric() && c != '_' {
            return Err(TaskError::InvalidContentKey(key.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>,
    content: &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)?;
    }

    // 3. Content key validation (early, before I/O)
    for key in content.keys() {
        validate_content_key(key)?;
    }

    execute_with_justfile(recipe, args, content, justfile_path, workdir, timeout).await
}

/// Execute a recipe resolved from a merged recipe list (project + optional global).
///
/// Lookup order: project first, then global. When the recipe is found in the global
/// justfile, `global_justfile_path` is used as `--justfile` but the cwd remains
/// `project_workdir` so that recipes that write to `./` target the project directory.
#[allow(clippy::too_many_arguments)]
pub async fn execute_recipe_merged(
    recipe_name: &str,
    args: &HashMap<String, String>,
    content: &HashMap<String, String>,
    project_justfile_path: &Path,
    global_justfile_path: Option<&Path>,
    timeout: Duration,
    mode: &TaskMode,
    project_workdir: Option<&Path>,
) -> Result<TaskExecution, TaskError> {
    // Build merged recipe list to find the target recipe and its source.
    let recipes = list_recipes_merged(
        project_justfile_path,
        global_justfile_path,
        mode,
        project_workdir,
    )
    .await?;

    let recipe = recipes
        .iter()
        .find(|r| r.name == recipe_name)
        .ok_or_else(|| TaskError::RecipeNotFound(recipe_name.to_string()))?;

    // Validate arguments.
    for value in args.values() {
        validate_arg_value(value)?;
    }

    // Content key validation (early, before I/O)
    for key in content.keys() {
        validate_content_key(key)?;
    }

    // Determine which justfile to invoke based on recipe source.
    let effective_justfile = match recipe.source {
        RecipeSource::Global => global_justfile_path
            .ok_or_else(|| TaskError::RecipeNotFound(recipe_name.to_string()))?,
        RecipeSource::Project => project_justfile_path,
    };

    execute_with_justfile(
        recipe,
        args,
        content,
        effective_justfile,
        project_workdir,
        timeout,
    )
    .await
}

/// Internal helper that constructs and runs a `just` command for the given recipe.
///
/// Handles argument ordering, timeout, output capture/truncation, and
/// `TaskExecution` assembly.  Both `execute_recipe` and `execute_recipe_merged`
/// delegate here after performing their recipe-lookup and validation steps.
async fn execute_with_justfile(
    recipe: &Recipe,
    args: &HashMap<String, String>,
    content: &HashMap<String, String>,
    effective_justfile: &Path,
    project_workdir: Option<&Path>,
    timeout: Duration,
) -> Result<TaskExecution, TaskError> {
    // 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();

    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(effective_justfile);
    if let Some(dir) = project_workdir {
        cmd.arg("--working-directory").arg(dir);
        cmd.current_dir(dir);
    }
    cmd.arg(&recipe.name);
    for arg in &positional {
        cmd.arg(arg);
    }

    // Set content entries as environment variables.
    // Keys are uppercased and prefixed with TASK_MCP_CONTENT_.
    // validate_content_key was already called in execute_recipe* callers, but
    // we call it again here to be safe in case execute_with_justfile is called
    // directly in the future.
    for (key, value) in content {
        validate_content_key(key)?;
        let env_name = format!("{}{}", CONTENT_ENV_PREFIX, key.to_uppercase());
        cmd.env(&env_name, value);
    }

    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();
    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;

    Ok(TaskExecution {
        id: Uuid::new_v4().to_string(),
        task_name: recipe.name.clone(),
        args: args.clone(),
        content: content.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 {
        make_recipe_with_doc(name, attributes, None)
    }

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

    #[test]
    fn has_allow_agent_group_attribute_true() {
        let recipe = make_recipe(
            "build",
            vec![RecipeAttribute::Object(
                [("group".to_string(), Some("allow-agent".to_string()))]
                    .into_iter()
                    .collect(),
            )],
        );
        assert!(has_allow_agent_group_attribute(&recipe));
    }

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

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

    #[test]
    fn has_allow_agent_group_attribute_false_legacy_agent_literal() {
        // The bare 'agent' literal is no longer recognized; only 'allow-agent'
        // matches Pattern A. This guards against accidental regressions.
        let recipe = make_recipe(
            "build",
            vec![RecipeAttribute::Object(
                [("group".to_string(), Some("agent".to_string()))]
                    .into_iter()
                    .collect(),
            )],
        );
        assert!(!has_allow_agent_group_attribute(&recipe));
    }

    #[test]
    fn has_allow_agent_doc_true() {
        let recipe = make_recipe_with_doc("build", vec![], Some("[allow-agent]"));
        assert!(has_allow_agent_doc(&recipe));
    }

    #[test]
    fn has_allow_agent_doc_false_no_doc() {
        let recipe = make_recipe("build", vec![]);
        assert!(!has_allow_agent_doc(&recipe));
    }

    #[test]
    fn has_allow_agent_doc_false_other_doc() {
        let recipe = make_recipe_with_doc("build", vec![], Some("Build the project"));
        assert!(!has_allow_agent_doc(&recipe));
    }

    #[test]
    fn has_allow_agent_doc_false_substring_in_prose() {
        // Prose mentioning the marker as part of a sentence must not match.
        let recipe = make_recipe_with_doc(
            "build",
            vec![],
            Some("do not add-[allow-agent]-here casually"),
        );
        assert!(!has_allow_agent_doc(&recipe));
    }

    #[test]
    fn has_allow_agent_doc_true_with_surrounding_whitespace() {
        let recipe = make_recipe_with_doc("build", vec![], Some("  [allow-agent]  "));
        assert!(has_allow_agent_doc(&recipe));
    }

    #[test]
    fn is_allow_agent_pattern_a() {
        let recipe = make_recipe(
            "build",
            vec![RecipeAttribute::Object(
                [("group".to_string(), Some("allow-agent".to_string()))]
                    .into_iter()
                    .collect(),
            )],
        );
        assert!(is_allow_agent(&recipe));
    }

    #[test]
    fn is_allow_agent_pattern_b() {
        let recipe = make_recipe_with_doc("build", vec![], Some("[allow-agent]"));
        assert!(is_allow_agent(&recipe));
    }

    #[test]
    fn is_allow_agent_pattern_a_plus_other_groups() {
        // A recipe with multiple stacked group attributes including
        // `allow-agent` stays agent-safe. This is the just-native form
        // (one `[group(...)]` attribute per line).
        let recipe = make_recipe(
            "build",
            vec![
                RecipeAttribute::Object(
                    [("group".to_string(), Some("allow-agent".to_string()))]
                        .into_iter()
                        .collect(),
                ),
                RecipeAttribute::Object(
                    [("group".to_string(), Some("profile".to_string()))]
                        .into_iter()
                        .collect(),
                ),
            ],
        );
        assert!(is_allow_agent(&recipe));
    }

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

    #[test]
    fn is_allow_agent_non_agent_group_only() {
        // Previously this case would be mistakenly filtered out when the user had
        // both a `# [allow-agent]` comment and a `[group('profile')]` attribute,
        // because the source-text scanner tripped on the attribute line.
        // The doc-based fallback now handles it correctly.
        let recipe = make_recipe_with_doc(
            "foo",
            vec![RecipeAttribute::Object(
                [("group".to_string(), Some("profile".to_string()))]
                    .into_iter()
                    .collect(),
            )],
            Some("[allow-agent]"),
        );
        assert!(is_allow_agent(&recipe));
    }

    #[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_n(char_3bytes, 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_allowed() {
        assert!(validate_arg_value("foo; rm -rf /").is_ok());
    }

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

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

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

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

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

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

    #[test]
    fn validate_arg_value_shell_metacharacters_allowed() {
        // Shell metacharacters are no longer rejected because task-mcp uses
        // Command::new("just").arg() which bypasses shell interpretation.
        assert!(validate_arg_value("https://example.com?a=1&b=2").is_ok());
        assert!(validate_arg_value("value with ${VAR} reference").is_ok());
        assert!(validate_arg_value(".items[] | select(.name)").is_ok());
        assert!(validate_arg_value("echo hello; echo world").is_ok());
        assert!(validate_arg_value("foo || bar").is_ok());
        assert!(validate_arg_value("result=$(cmd)").is_ok());
        assert!(validate_arg_value("hello `world`").is_ok());
    }

    // -------------------------------------------------------------------------
    // validate_content_key tests
    // -------------------------------------------------------------------------

    #[test]
    fn validate_content_key_single_letter() {
        assert!(validate_content_key("a").is_ok());
        assert!(validate_content_key("Z").is_ok());
    }

    #[test]
    fn validate_content_key_alphanumeric_and_underscore() {
        assert!(validate_content_key("body").is_ok());
        assert!(validate_content_key("my_content").is_ok());
        assert!(validate_content_key("Content123").is_ok());
        assert!(validate_content_key("A_B_C").is_ok());
        assert!(validate_content_key("myKey_2").is_ok());
    }

    #[test]
    fn validate_content_key_empty_rejected() {
        assert!(validate_content_key("").is_err());
    }

    #[test]
    fn validate_content_key_digit_start_rejected() {
        assert!(validate_content_key("1key").is_err());
        assert!(validate_content_key("0abc").is_err());
    }

    #[test]
    fn validate_content_key_hyphen_rejected() {
        assert!(validate_content_key("my-key").is_err());
    }

    #[test]
    fn validate_content_key_space_rejected() {
        assert!(validate_content_key("my key").is_err());
        assert!(validate_content_key(" key").is_err());
    }

    #[test]
    fn validate_content_key_newline_rejected() {
        assert!(validate_content_key("key\nvalue").is_err());
        assert!(validate_content_key("\nkey").is_err());
    }

    #[test]
    fn validate_content_key_underscore_start_rejected() {
        // Underscore-only start is not a letter — env var convention
        assert!(validate_content_key("_key").is_err());
    }

    #[test]
    fn validate_content_key_dot_rejected() {
        assert!(validate_content_key("my.key").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(),
            content: 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);
    }
}