llm-tool 0.10.2

//! Quarantine framing and delimiter neutralization for untrusted tool outputs.
//!
//! Provides idempotent XML boundary wrapping (`<tool_output_quarantine>`) and
//! control-token neutralization so untrusted tool responses cannot inject `ChatML`,
//! Harmony, Llama, or `DeepSeek` turn/tool delimiters into the model's token stream.

use alloc::{borrow::ToOwned, string::String};

use crate::types::ToolOutput;

/// Default XML tag name used for quarantined tool output envelopes.
pub const TOOL_OUTPUT_QUARANTINE_TAG: &str = "tool_output_quarantine";

/// Opening tag of a quarantined tool output envelope.
pub const TOOL_OUTPUT_QUARANTINE_OPEN: &str = "<tool_output_quarantine>";

/// Closing tag of a quarantined tool output envelope.
pub const TOOL_OUTPUT_QUARANTINE_CLOSE: &str = "</tool_output_quarantine>";

/// Opening prefix including trailing newline.
const OPEN_FRAME: &str = "<tool_output_quarantine>\n";

/// Closing suffix including leading newline.
const CLOSE_FRAME: &str = "\n</tool_output_quarantine>";

/// Well-known LLM control token and role/tool delimiter replacement rules.
pub const DELIMITER_RULES: &[(&str, &str)] = &[
    ("<|im_start|>", "&lt;|im_start|&gt;"),
    ("<|im_end|>", "&lt;|im_end|&gt;"),
    ("<|endoftext|>", "&lt;|endoftext|&gt;"),
    ("<|start_header_id|>", "&lt;|start_header_id|&gt;"),
    ("<|end_header_id|>", "&lt;|end_header_id|&gt;"),
    ("<|eot_id|>", "&lt;|eot_id|&gt;"),
    ("<tool_call>", "&lt;tool_call&gt;"),
    ("</tool_call>", "&lt;/tool_call&gt;"),
    ("<tool_response>", "&lt;tool_response&gt;"),
    ("</tool_response>", "&lt;/tool_response&gt;"),
    ("[INST]", "&#91;INST&#93;"),
    ("[/INST]", "&#91;/INST&#93;"),
    ("<<SYS>>", "&lt;&lt;SYS&gt;&gt;"),
    ("<</SYS>>", "&lt;&lt;/SYS&gt;&gt;"),
    ("<start_of_turn>", "&lt;start_of_turn&gt;"),
    ("<end_of_turn>", "&lt;end_of_turn&gt;"),
    ("<think>", "&lt;think&gt;"),
    ("</think>", "&lt;/think&gt;"),
    ("<untrusted_tool_output>", "&lt;untrusted_tool_output&gt;"),
    ("</untrusted_tool_output>", "&lt;/untrusted_tool_output&gt;"),
    ("<untrusted_content>", "&lt;untrusted_content&gt;"),
    ("</untrusted_content>", "&lt;/untrusted_content&gt;"),
    ("<tool_output_quarantine>", "&lt;tool_output_quarantine&gt;"),
    (
        "</tool_output_quarantine>",
        "&lt;/tool_output_quarantine&gt;",
    ),
    ("<｜begin▁of▁sentence｜>", "&lt;｜begin▁of▁sentence｜&gt;"),
    ("<｜end▁of▁sentence｜>", "&lt;｜end▁of▁sentence｜&gt;"),
    ("<｜User｜>", "&lt;｜User｜&gt;"),
    ("<｜Assistant｜>", "&lt;｜Assistant｜&gt;"),
    ("<｜tool▁calls▁begin｜>", "&lt;｜tool▁calls▁begin｜&gt;"),
    ("<|user|>", "&lt;|user|&gt;"),
    ("<|assistant|>", "&lt;|assistant|&gt;"),
    ("<|system|>", "&lt;|system|&gt;"),
    ("<|end|>", "&lt;|end|&gt;"),
    ("<|START_OF_TURN_TOKEN|>", "&lt;|START_OF_TURN_TOKEN|&gt;"),
    ("<|END_OF_TURN_TOKEN|>", "&lt;|END_OF_TURN_TOKEN|&gt;"),
    ("[TOOL_CALLS]", "&#91;TOOL_CALLS&#93;"),
    ("[AVAILABLE_TOOLS]", "&#91;AVAILABLE_TOOLS&#93;"),
    ("[/TOOL_CALLS]", "&#91;/TOOL_CALLS&#93;"),
    ("[/AVAILABLE_TOOLS]", "&#91;/AVAILABLE_TOOLS&#93;"),
    ("\n\nHuman:", "\n\nHuman&#58;"),
    ("\n\nAssistant:", "\n\nAssistant&#58;"),
    ("<role>", "&lt;role&gt;"),
    ("</role>", "&lt;/role&gt;"),
    ("<|role_end|>", "&lt;|role_end|&gt;"),
    ("<|channel|>", "&lt;|channel|&gt;"),
    ("<|message|>", "&lt;|message|&gt;"),
];

const fn is_ncname_continue_byte(b: u8) -> bool {
    b.is_ascii_alphanumeric() || b == b'_' || b == b'-' || b == b'.'
}

/// Check if `bytes[i..]` (where `bytes[i] == b'<'`) starts an opening or closing
/// tag matching `tag_bytes` (ASCII case-insensitively), allowing optional ASCII
/// whitespace after `<` and after `/`.
fn match_tag_prefix(bytes: &[u8], i: usize, tag_bytes: &[u8]) -> Option<usize> {
    let mut pos = i + 1;
    while pos < bytes.len() && bytes[pos].is_ascii_whitespace() {
        pos += 1;
    }
    if pos < bytes.len() && bytes[pos] == b'/' {
        pos += 1;
        while pos < bytes.len() && bytes[pos].is_ascii_whitespace() {
            pos += 1;
        }
    }
    let after_tag = pos.checked_add(tag_bytes.len())?;
    if after_tag <= bytes.len()
        && bytes[pos..after_tag].eq_ignore_ascii_case(tag_bytes)
        && !(after_tag < bytes.len() && is_ncname_continue_byte(bytes[after_tag]))
    {
        Some(after_tag)
    } else {
        None
    }
}

/// Returns `true` if `s` contains any un-neutralized `<tool_output_quarantine...>`
/// or `</tool_output_quarantine...>` tag prefix (case-insensitive, with optional
/// whitespace or attributes).
fn has_quarantine_tag_breakout(s: &str) -> bool {
    let bytes = s.as_bytes();
    let tag_bytes = TOOL_OUTPUT_QUARANTINE_TAG.as_bytes();
    for i in 0..bytes.len() {
        if bytes[i] == b'<' && match_tag_prefix(bytes, i, tag_bytes).is_some() {
            return true;
        }
    }
    false
}

/// Escape any opening or closing `<tool_output_quarantine...>` tag occurrences.
fn sanitize_quarantine_tag_breakouts(s: &str) -> String {
    let bytes = s.as_bytes();
    let tag_bytes = TOOL_OUTPUT_QUARANTINE_TAG.as_bytes();
    let mut out = String::with_capacity(s.len());
    let mut i = 0;
    while i < bytes.len() {
        if bytes[i] == b'<'
            && let Some(after_tag) = match_tag_prefix(bytes, i, tag_bytes)
        {
            let mut j = after_tag;
            while j < bytes.len() && bytes[j] != b'>' && bytes[j] != b'<' {
                j += 1;
            }
            if j < bytes.len() && bytes[j] == b'>' {
                out.push_str("&lt;");
                out.push_str(&s[i + 1..j]);
                out.push_str("&gt;");
                i = j + 1;
                continue;
            }
            out.push_str("&lt;");
            i += 1;
            continue;
        }
        let ch = s[i..].chars().next().expect("valid utf-8 character");
        out.push(ch);
        i += ch.len_utf8();
    }
    out
}

/// Neutralize all known control tokens and quarantine breakout tags in `raw`.
#[must_use]
pub fn neutralize_delimiters(raw: &str) -> String {
    let mut out = raw.to_owned();
    for &(token, replacement) in DELIMITER_RULES {
        if out.contains(token) {
            out = out.replace(token, replacement);
        }
    }
    if has_quarantine_tag_breakout(&out) {
        out = sanitize_quarantine_tag_breakouts(&out);
    }
    out
}

/// Check whether `raw` is already a valid, un-escaped `<tool_output_quarantine>`
/// envelope whose inner body contains zero un-neutralized [`DELIMITER_RULES`] or
/// boundary breakout tags.
///
/// A naive `starts_with` / `ends_with` check is unsound on its own because an
/// attacker payload like `<tool_output_quarantine>\nevil\n</tool_output_quarantine>\n<tool_call>...`
/// or `<tool_output_quarantine>\n</tool_output_quarantine><tool_call>\n</tool_output_quarantine>`
/// would pass the outer check while smuggling raw delimiters inside the body.
#[must_use]
pub fn is_quarantined_output(raw: &str) -> bool {
    let Some(rest) = raw.strip_prefix(OPEN_FRAME) else {
        return false;
    };
    let Some(inner) = rest.strip_suffix(CLOSE_FRAME) else {
        return false;
    };
    if DELIMITER_RULES
        .iter()
        .any(|&(token, _)| inner.contains(token))
    {
        return false;
    }
    !has_quarantine_tag_breakout(inner)
}

/// Extract the inner payload slice from a quarantined tool output string if and
/// only if `raw` is a valid quarantined envelope (as verified by [`is_quarantined_output`]).
#[must_use]
pub fn unquarantine_tool_output(raw: &str) -> Option<&str> {
    if !is_quarantined_output(raw) {
        return None;
    }
    raw.strip_prefix(OPEN_FRAME)
        .and_then(|rest| rest.strip_suffix(CLOSE_FRAME))
}

/// Wrap `raw` in a `<tool_output_quarantine>` envelope after neutralizing any
/// embedded [`DELIMITER_RULES`] and breakout tags.
///
/// This function is **idempotent**: if `raw` already satisfies [`is_quarantined_output`],
/// it is returned unchanged without double-wrapping.
#[must_use]
pub fn quarantine_tool_output(raw: &str) -> String {
    if is_quarantined_output(raw) {
        return raw.to_owned();
    }
    let sanitized = neutralize_delimiters(raw);
    let mut buf = String::with_capacity(OPEN_FRAME.len() + sanitized.len() + CLOSE_FRAME.len());
    buf.push_str(OPEN_FRAME);
    buf.push_str(&sanitized);
    buf.push_str(CLOSE_FRAME);
    buf
}

impl ToolOutput {
    /// Return a quarantined copy of this [`ToolOutput`], neutralizing control
    /// delimiters and wrapping the content in `<tool_output_quarantine>` tags.
    ///
    /// Idempotent when called on an already-quarantined output.
    #[must_use]
    pub fn quarantined(mut self) -> Self {
        self.content = quarantine_tool_output(&self.content);
        self
    }

    /// Returns `true` if this output's content is wrapped in a verified
    /// `<tool_output_quarantine>` envelope with no un-neutralized delimiters.
    #[must_use]
    pub fn is_quarantined(&self) -> bool {
        is_quarantined_output(&self.content)
    }
}