zeph-skills 0.12.2

// SPDX-FileCopyrightText: 2026 Andrei G <bug-ops>
// SPDX-License-Identifier: MIT OR Apache-2.0

use std::collections::HashMap;
use std::fmt::Write;

use crate::loader::Skill;
use crate::resource::discover_resources;
use crate::trust::TrustLevel;

// XML tag patterns (lowercase) that could break prompt structure if injected verbatim.
// Matching is case-insensitive; the replacement is always the canonical escaped form.
const SANITIZE_PATTERNS: &[(&str, &str)] = &[
    ("</skill>", "&lt;/skill&gt;"),
    ("<skill", "&lt;skill"),
    ("</instructions>", "&lt;/instructions&gt;"),
    ("<instructions", "&lt;instructions"),
    ("</available_skills>", "&lt;/available_skills&gt;"),
    ("<available_skills", "&lt;available_skills"),
];

/// Case-insensitive replacement of `pattern` (given in lowercase) with `replacement` in `src`.
fn replace_case_insensitive(src: &str, pattern: &str, replacement: &str) -> String {
    let lower = src.to_ascii_lowercase();
    let mut out = String::with_capacity(src.len());
    let mut pos = 0;
    while pos < src.len() {
        if lower[pos..].starts_with(pattern) {
            out.push_str(replacement);
            pos += pattern.len();
        } else {
            // Safety: pos is always at a char boundary because ascii_lowercase preserves boundaries
            let ch = src[pos..].chars().next().unwrap();
            out.push(ch);
            pos += ch.len_utf8();
        }
    }
    out
}

/// Escape XML tags that could break prompt structure when emitted verbatim.
///
/// Matching is case-insensitive so mixed-case variants like `</Skill>` are also escaped.
/// Applied only to untrusted (non-`Trusted`) skill bodies before prompt injection.
#[must_use]
pub fn sanitize_skill_body(body: &str) -> String {
    let mut out = body.to_string();
    for (pattern, replacement) in SANITIZE_PATTERNS {
        out = replace_case_insensitive(&out, pattern, replacement);
    }
    out
}

#[must_use]
pub fn format_skills_prompt<S: std::hash::BuildHasher>(
    skills: &[Skill],
    trust_levels: &HashMap<String, TrustLevel, S>,
) -> String {
    if skills.is_empty() {
        return String::new();
    }

    let mut out = String::from("<available_skills>\n");

    for skill in skills {
        let trust = trust_levels
            .get(skill.name())
            .copied()
            .unwrap_or(TrustLevel::Trusted);
        let raw_body = if trust == TrustLevel::Trusted {
            skill.body.clone()
        } else {
            sanitize_skill_body(&skill.body)
        };
        let body = if trust == TrustLevel::Quarantined {
            wrap_quarantined(skill.name(), &raw_body)
        } else {
            raw_body
        };
        let _ = write!(
            out,
            "  <skill name=\"{}\">\n    <description>{}</description>\n    <instructions>\n{}",
            skill.name(),
            skill.description(),
            body,
        );

        let resources = discover_resources(&skill.meta.skill_dir);

        let ref_names: Vec<&str> = resources
            .references
            .iter()
            .filter_map(|p| p.file_name()?.to_str())
            .collect();
        if !ref_names.is_empty() {
            let _ = write!(out, "\nAvailable references: {}", ref_names.join(", "));
        }

        let script_names: Vec<&str> = resources
            .scripts
            .iter()
            .filter_map(|p| p.file_name()?.to_str())
            .collect();
        if !script_names.is_empty() {
            let _ = write!(out, "\nAvailable scripts: {}", script_names.join(", "));
        }

        let asset_names: Vec<&str> = resources
            .assets
            .iter()
            .filter_map(|p| p.file_name()?.to_str())
            .collect();
        if !asset_names.is_empty() {
            let _ = write!(out, "\nAvailable assets: {}", asset_names.join(", "));
        }

        out.push_str("\n    </instructions>\n  </skill>\n");
    }

    out.push_str("</available_skills>");
    out
}

/// Wrap a quarantined skill's prompt with warning markers.
#[must_use]
pub fn wrap_quarantined(skill_name: &str, body: &str) -> String {
    format!(
        "[QUARANTINED SKILL: {skill_name}] The following skill is quarantined. \
         It has restricted tool access (no bash, file_write, web_scrape).\n\n{body}"
    )
}

/// Format skills as a compact single-line XML list (name + description + path only).
///
/// Used when the model context window is small (< 8192 tokens) to save space.
#[must_use]
pub fn format_skills_prompt_compact(skills: &[Skill]) -> String {
    if skills.is_empty() {
        return String::new();
    }

    let mut out = String::from("<available_skills mode=\"compact\">\n");
    for skill in skills {
        let _ = writeln!(
            out,
            "  <skill name=\"{}\" description=\"{}\" />",
            skill.name(),
            skill.description(),
        );
    }
    out.push_str("</available_skills>");
    out
}

#[must_use]
pub fn format_skills_catalog(skills: &[Skill]) -> String {
    if skills.is_empty() {
        return String::new();
    }

    let mut out = String::from("<other_skills>\n");
    for skill in skills {
        let _ = writeln!(
            out,
            "  <skill name=\"{}\" description=\"{}\" />",
            skill.name(),
            skill.description(),
        );
    }
    out.push_str("</other_skills>");
    out
}

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

    use crate::loader::SkillMeta;

    fn make_skill(name: &str, description: &str, body: &str) -> Skill {
        Skill {
            meta: SkillMeta {
                name: name.into(),
                description: description.into(),
                compatibility: None,
                license: None,
                metadata: Vec::new(),
                allowed_tools: Vec::new(),
                requires_secrets: Vec::new(),
                skill_dir: PathBuf::new(),
            },
            body: body.into(),
        }
    }

    fn make_skill_with_dir(name: &str, description: &str, body: &str, dir: PathBuf) -> Skill {
        Skill {
            meta: SkillMeta {
                name: name.into(),
                description: description.into(),
                compatibility: None,
                license: None,
                metadata: Vec::new(),
                allowed_tools: Vec::new(),
                requires_secrets: Vec::new(),
                skill_dir: dir,
            },
            body: body.into(),
        }
    }

    #[test]
    fn empty_skills_returns_empty_string() {
        let empty: &[Skill] = &[];
        assert_eq!(format_skills_prompt(empty, &HashMap::new()), "");
    }

    #[test]
    fn single_skill_format() {
        let skills = vec![make_skill("test", "A test.", "# Hello\nworld")];

        let output = format_skills_prompt(&skills, &HashMap::new());
        assert!(output.starts_with("<available_skills>"));
        assert!(output.ends_with("</available_skills>"));
        assert!(output.contains("<skill name=\"test\">"));
        assert!(output.contains("<description>A test.</description>"));
        assert!(output.contains("# Hello\nworld"));
    }

    #[test]
    fn multiple_skills() {
        let skills = vec![
            make_skill("a", "desc a", "body a"),
            make_skill("b", "desc b", "body b"),
        ];

        let output = format_skills_prompt(&skills, &HashMap::new());
        assert!(output.contains("<skill name=\"a\">"));
        assert!(output.contains("<skill name=\"b\">"));
    }

    #[test]
    fn references_listed_not_injected() {
        let dir = tempfile::tempdir().unwrap();
        let refs = dir.path().join("references");
        std::fs::create_dir(&refs).unwrap();
        std::fs::write(refs.join("api-guide.md"), "# API Guide content").unwrap();
        std::fs::write(refs.join("common.md"), "# Common docs content").unwrap();

        let skills = vec![make_skill_with_dir(
            "test",
            "desc",
            "body",
            dir.path().to_path_buf(),
        )];

        let output = format_skills_prompt(&skills, &HashMap::new());
        // filenames listed
        assert!(output.contains("Available references:"));
        assert!(output.contains("api-guide.md"));
        assert!(output.contains("common.md"));
        // content NOT injected
        assert!(!output.contains("# API Guide content"));
        assert!(!output.contains("# Common docs content"));
        assert!(!output.contains("<reference"));
    }

    #[test]
    fn scripts_listed_not_injected() {
        let dir = tempfile::tempdir().unwrap();
        let scripts = dir.path().join("scripts");
        std::fs::create_dir(&scripts).unwrap();
        std::fs::write(scripts.join("extract.py"), "print('hi')").unwrap();

        let skills = vec![make_skill_with_dir(
            "test",
            "desc",
            "body",
            dir.path().to_path_buf(),
        )];

        let output = format_skills_prompt(&skills, &HashMap::new());
        assert!(output.contains("Available scripts: extract.py"));
        assert!(!output.contains("print('hi')"));
    }

    #[test]
    fn assets_listed_not_injected() {
        let dir = tempfile::tempdir().unwrap();
        let assets = dir.path().join("assets");
        std::fs::create_dir(&assets).unwrap();
        std::fs::write(assets.join("logo.png"), &[0u8; 4]).unwrap();

        let skills = vec![make_skill_with_dir(
            "test",
            "desc",
            "body",
            dir.path().to_path_buf(),
        )];

        let output = format_skills_prompt(&skills, &HashMap::new());
        assert!(output.contains("Available assets: logo.png"));
    }

    #[test]
    fn no_resources_dir_produces_body_only() {
        let dir = tempfile::tempdir().unwrap();
        let skills = vec![make_skill_with_dir(
            "test",
            "desc",
            "skill body",
            dir.path().to_path_buf(),
        )];

        let output = format_skills_prompt(&skills, &HashMap::new());
        assert!(output.contains("skill body"));
        assert!(!output.contains("Available references"));
        assert!(!output.contains("Available scripts"));
        assert!(!output.contains("Available assets"));
    }

    #[test]
    fn quarantined_skill_gets_wrapped() {
        let skills = vec![make_skill("untrusted", "desc", "do stuff")];
        let mut trust = HashMap::new();
        trust.insert("untrusted".into(), TrustLevel::Quarantined);
        let output = format_skills_prompt(&skills, &trust);
        assert!(output.contains("[QUARANTINED SKILL: untrusted]"));
        assert!(output.contains("restricted tool access"));
    }

    #[test]
    fn trusted_skill_not_wrapped() {
        let skills = vec![make_skill("safe", "desc", "do stuff")];
        let mut trust = HashMap::new();
        trust.insert("safe".into(), TrustLevel::Trusted);
        let output = format_skills_prompt(&skills, &trust);
        assert!(!output.contains("QUARANTINED"));
        assert!(output.contains("do stuff"));
    }

    #[test]
    fn sanitize_case_insensitive() {
        let body = "Close </Skill> and </INSTRUCTIONS> and </Available_Skills>.";
        let sanitized = sanitize_skill_body(body);
        assert!(!sanitized.contains("</Skill>"));
        assert!(!sanitized.contains("</INSTRUCTIONS>"));
        assert!(!sanitized.contains("</Available_Skills>"));
        assert!(sanitized.contains("&lt;/skill&gt;"));
        assert!(sanitized.contains("&lt;/instructions&gt;"));
        assert!(sanitized.contains("&lt;/available_skills&gt;"));
    }

    #[test]
    fn sanitize_escapes_xml_tags() {
        let body = "Do not close </skill> or </instructions> tags.";
        let sanitized = sanitize_skill_body(body);
        assert!(!sanitized.contains("</skill>"));
        assert!(!sanitized.contains("</instructions>"));
        assert!(sanitized.contains("&lt;/skill&gt;"));
        assert!(sanitized.contains("&lt;/instructions&gt;"));
    }

    #[test]
    fn sanitize_escapes_opening_xml_tags() {
        let body = "Inject <skill name=\"evil\"> and <instructions> here.";
        let sanitized = sanitize_skill_body(body);
        assert!(!sanitized.contains("<skill"));
        assert!(!sanitized.contains("<instructions"));
        assert!(sanitized.contains("&lt;skill"));
        assert!(sanitized.contains("&lt;instructions"));
    }

    #[test]
    fn trusted_skill_not_sanitized() {
        let body = "Some </skill> content.";
        let skills = vec![make_skill("safe", "desc", body)];
        let mut trust = HashMap::new();
        trust.insert("safe".into(), TrustLevel::Trusted);
        let output = format_skills_prompt(&skills, &trust);
        assert!(output.contains("Some </skill> content."));
    }

    #[test]
    fn verified_skill_is_sanitized() {
        let body = "Inject </skill> here.";
        let skills = vec![make_skill("ver", "desc", body)];
        let mut trust = HashMap::new();
        trust.insert("ver".into(), TrustLevel::Verified);
        let output = format_skills_prompt(&skills, &trust);
        assert!(output.contains("&lt;/skill&gt;"));
        assert!(!output.contains("Inject </skill> here."));
    }

    #[test]
    fn quarantined_skill_is_sanitized_and_wrapped() {
        let body = "Inject </instructions> and </skill>.";
        let skills = vec![make_skill("evil", "desc", body)];
        let mut trust = HashMap::new();
        trust.insert("evil".into(), TrustLevel::Quarantined);
        let output = format_skills_prompt(&skills, &trust);
        assert!(output.contains("[QUARANTINED SKILL: evil]"));
        assert!(output.contains("&lt;/instructions&gt;"));
        assert!(output.contains("&lt;/skill&gt;"));
        assert!(!output.contains("Inject </instructions>"));
    }

    #[test]
    fn compact_empty_returns_empty_string() {
        let empty: &[Skill] = &[];
        assert_eq!(format_skills_prompt_compact(empty), "");
    }

    #[test]
    fn compact_single_skill_no_path() {
        let skills = vec![make_skill("my-skill", "Does things.", "body")];
        let output = format_skills_prompt_compact(&skills);
        assert!(output.starts_with("<available_skills mode=\"compact\">"));
        assert!(output.ends_with("</available_skills>"));
        assert!(output.contains("name=\"my-skill\""));
        assert!(output.contains("description=\"Does things.\""));
        assert!(!output.contains("path="), "path must not be present");
    }

    #[test]
    fn compact_multiple_skills() {
        let skills = vec![
            make_skill("a", "desc a", "body a"),
            make_skill("b", "desc b", "body b"),
        ];
        let output = format_skills_prompt_compact(&skills);
        assert!(output.contains("name=\"a\""));
        assert!(output.contains("name=\"b\""));
        assert!(!output.contains("path="));
    }

    #[test]
    fn compact_mode_attribute_present() {
        let skills = vec![make_skill("x", "y", "z")];
        let output = format_skills_prompt_compact(&skills);
        assert!(output.contains("mode=\"compact\""));
    }

    #[test]
    fn format_skills_catalog_empty() {
        let empty: &[Skill] = &[];
        assert_eq!(format_skills_catalog(empty), "");
    }

    #[test]
    fn format_skills_catalog_produces_other_skills_tag() {
        let skills = vec![make_skill("test", "A test skill.", "body")];
        let output = format_skills_catalog(&skills);
        assert!(output.starts_with("<other_skills>"));
        assert!(output.ends_with("</other_skills>"));
        assert!(output.contains("name=\"test\""));
        assert!(output.contains("description=\"A test skill.\""));
        assert!(!output.contains("body"));
    }
}