llm_tool/types/
quarantine.rs

1//! Quarantine framing and delimiter neutralization for untrusted tool outputs.
2//!
3//! Provides idempotent XML boundary wrapping (`<tool_output_quarantine>`) and
4//! control-token neutralization so untrusted tool responses cannot inject `ChatML`,
5//! Harmony, Llama, or `DeepSeek` turn/tool delimiters into the model's token stream.
6
7use alloc::{borrow::ToOwned, string::String};
8
9use crate::types::ToolOutput;
10
11/// Default XML tag name used for quarantined tool output envelopes.
12pub const TOOL_OUTPUT_QUARANTINE_TAG: &str = "tool_output_quarantine";
13
14/// Opening tag of a quarantined tool output envelope.
15pub const TOOL_OUTPUT_QUARANTINE_OPEN: &str = "<tool_output_quarantine>";
16
17/// Closing tag of a quarantined tool output envelope.
18pub const TOOL_OUTPUT_QUARANTINE_CLOSE: &str = "</tool_output_quarantine>";
19
20/// Opening prefix including trailing newline.
21const OPEN_FRAME: &str = "<tool_output_quarantine>\n";
22
23/// Closing suffix including leading newline.
24const CLOSE_FRAME: &str = "\n</tool_output_quarantine>";
25
26/// Well-known LLM control token and role/tool delimiter replacement rules.
27pub const DELIMITER_RULES: &[(&str, &str)] = &[
28    ("<|im_start|>", "&lt;|im_start|&gt;"),
29    ("<|im_end|>", "&lt;|im_end|&gt;"),
30    ("<|endoftext|>", "&lt;|endoftext|&gt;"),
31    ("<|start_header_id|>", "&lt;|start_header_id|&gt;"),
32    ("<|end_header_id|>", "&lt;|end_header_id|&gt;"),
33    ("<|eot_id|>", "&lt;|eot_id|&gt;"),
34    ("<tool_call>", "&lt;tool_call&gt;"),
35    ("</tool_call>", "&lt;/tool_call&gt;"),
36    ("<tool_response>", "&lt;tool_response&gt;"),
37    ("</tool_response>", "&lt;/tool_response&gt;"),
38    ("[INST]", "&#91;INST&#93;"),
39    ("[/INST]", "&#91;/INST&#93;"),
40    ("<<SYS>>", "&lt;&lt;SYS&gt;&gt;"),
41    ("<</SYS>>", "&lt;&lt;/SYS&gt;&gt;"),
42    ("<start_of_turn>", "&lt;start_of_turn&gt;"),
43    ("<end_of_turn>", "&lt;end_of_turn&gt;"),
44    ("<think>", "&lt;think&gt;"),
45    ("</think>", "&lt;/think&gt;"),
46    ("<untrusted_tool_output>", "&lt;untrusted_tool_output&gt;"),
47    ("</untrusted_tool_output>", "&lt;/untrusted_tool_output&gt;"),
48    ("<untrusted_content>", "&lt;untrusted_content&gt;"),
49    ("</untrusted_content>", "&lt;/untrusted_content&gt;"),
50    ("<tool_output_quarantine>", "&lt;tool_output_quarantine&gt;"),
51    (
52        "</tool_output_quarantine>",
53        "&lt;/tool_output_quarantine&gt;",
54    ),
55    ("<｜begin▁of▁sentence｜>", "&lt;｜begin▁of▁sentence｜&gt;"),
56    ("<｜end▁of▁sentence｜>", "&lt;｜end▁of▁sentence｜&gt;"),
57    ("<｜User｜>", "&lt;｜User｜&gt;"),
58    ("<｜Assistant｜>", "&lt;｜Assistant｜&gt;"),
59    ("<｜tool▁calls▁begin｜>", "&lt;｜tool▁calls▁begin｜&gt;"),
60    ("<|user|>", "&lt;|user|&gt;"),
61    ("<|assistant|>", "&lt;|assistant|&gt;"),
62    ("<|system|>", "&lt;|system|&gt;"),
63    ("<|end|>", "&lt;|end|&gt;"),
64    ("<|START_OF_TURN_TOKEN|>", "&lt;|START_OF_TURN_TOKEN|&gt;"),
65    ("<|END_OF_TURN_TOKEN|>", "&lt;|END_OF_TURN_TOKEN|&gt;"),
66    ("[TOOL_CALLS]", "&#91;TOOL_CALLS&#93;"),
67    ("[AVAILABLE_TOOLS]", "&#91;AVAILABLE_TOOLS&#93;"),
68    ("[/TOOL_CALLS]", "&#91;/TOOL_CALLS&#93;"),
69    ("[/AVAILABLE_TOOLS]", "&#91;/AVAILABLE_TOOLS&#93;"),
70    ("\n\nHuman:", "\n\nHuman&#58;"),
71    ("\n\nAssistant:", "\n\nAssistant&#58;"),
72    ("<role>", "&lt;role&gt;"),
73    ("</role>", "&lt;/role&gt;"),
74    ("<|role_end|>", "&lt;|role_end|&gt;"),
75    ("<|channel|>", "&lt;|channel|&gt;"),
76    ("<|message|>", "&lt;|message|&gt;"),
77];
78
79const fn is_ncname_continue_byte(b: u8) -> bool {
80    b.is_ascii_alphanumeric() || b == b'_' || b == b'-' || b == b'.'
81}
82
83/// Check if `bytes[i..]` (where `bytes[i] == b'<'`) starts an opening or closing
84/// tag matching `tag_bytes` (ASCII case-insensitively), allowing optional ASCII
85/// whitespace after `<` and after `/`.
86fn match_tag_prefix(bytes: &[u8], i: usize, tag_bytes: &[u8]) -> Option<usize> {
87    let mut pos = i + 1;
88    while pos < bytes.len() && bytes[pos].is_ascii_whitespace() {
89        pos += 1;
90    }
91    if pos < bytes.len() && bytes[pos] == b'/' {
92        pos += 1;
93        while pos < bytes.len() && bytes[pos].is_ascii_whitespace() {
94            pos += 1;
95        }
96    }
97    let after_tag = pos.checked_add(tag_bytes.len())?;
98    if after_tag <= bytes.len()
99        && bytes[pos..after_tag].eq_ignore_ascii_case(tag_bytes)
100        && !(after_tag < bytes.len() && is_ncname_continue_byte(bytes[after_tag]))
101    {
102        Some(after_tag)
103    } else {
104        None
105    }
106}
107
108/// Returns `true` if `s` contains any un-neutralized `<tool_output_quarantine...>`
109/// or `</tool_output_quarantine...>` tag prefix (case-insensitive, with optional
110/// whitespace or attributes).
111fn has_quarantine_tag_breakout(s: &str) -> bool {
112    let bytes = s.as_bytes();
113    let tag_bytes = TOOL_OUTPUT_QUARANTINE_TAG.as_bytes();
114    for i in 0..bytes.len() {
115        if bytes[i] == b'<' && match_tag_prefix(bytes, i, tag_bytes).is_some() {
116            return true;
117        }
118    }
119    false
120}
121
122/// Escape any opening or closing `<tool_output_quarantine...>` tag occurrences.
123fn sanitize_quarantine_tag_breakouts(s: &str) -> String {
124    let bytes = s.as_bytes();
125    let tag_bytes = TOOL_OUTPUT_QUARANTINE_TAG.as_bytes();
126    let mut out = String::with_capacity(s.len());
127    let mut i = 0;
128    while i < bytes.len() {
129        if bytes[i] == b'<'
130            && let Some(after_tag) = match_tag_prefix(bytes, i, tag_bytes)
131        {
132            let mut j = after_tag;
133            while j < bytes.len() && bytes[j] != b'>' && bytes[j] != b'<' {
134                j += 1;
135            }
136            if j < bytes.len() && bytes[j] == b'>' {
137                out.push_str("&lt;");
138                out.push_str(&s[i + 1..j]);
139                out.push_str("&gt;");
140                i = j + 1;
141                continue;
142            }
143            out.push_str("&lt;");
144            i += 1;
145            continue;
146        }
147        let ch = s[i..].chars().next().expect("valid utf-8 character");
148        out.push(ch);
149        i += ch.len_utf8();
150    }
151    out
152}
153
154/// Neutralize all known control tokens and quarantine breakout tags in `raw`.
155#[must_use]
156pub fn neutralize_delimiters(raw: &str) -> String {
157    let mut out = raw.to_owned();
158    for &(token, replacement) in DELIMITER_RULES {
159        if out.contains(token) {
160            out = out.replace(token, replacement);
161        }
162    }
163    if has_quarantine_tag_breakout(&out) {
164        out = sanitize_quarantine_tag_breakouts(&out);
165    }
166    out
167}
168
169/// Check whether `raw` is already a valid, un-escaped `<tool_output_quarantine>`
170/// envelope whose inner body contains zero un-neutralized [`DELIMITER_RULES`] or
171/// boundary breakout tags.
172///
173/// A naive `starts_with` / `ends_with` check is unsound on its own because an
174/// attacker payload like `<tool_output_quarantine>\nevil\n</tool_output_quarantine>\n<tool_call>...`
175/// or `<tool_output_quarantine>\n</tool_output_quarantine><tool_call>\n</tool_output_quarantine>`
176/// would pass the outer check while smuggling raw delimiters inside the body.
177#[must_use]
178pub fn is_quarantined_output(raw: &str) -> bool {
179    let Some(rest) = raw.strip_prefix(OPEN_FRAME) else {
180        return false;
181    };
182    let Some(inner) = rest.strip_suffix(CLOSE_FRAME) else {
183        return false;
184    };
185    if DELIMITER_RULES
186        .iter()
187        .any(|&(token, _)| inner.contains(token))
188    {
189        return false;
190    }
191    !has_quarantine_tag_breakout(inner)
192}
193
194/// Extract the inner payload slice from a quarantined tool output string if and
195/// only if `raw` is a valid quarantined envelope (as verified by [`is_quarantined_output`]).
196#[must_use]
197pub fn unquarantine_tool_output(raw: &str) -> Option<&str> {
198    if !is_quarantined_output(raw) {
199        return None;
200    }
201    raw.strip_prefix(OPEN_FRAME)
202        .and_then(|rest| rest.strip_suffix(CLOSE_FRAME))
203}
204
205/// Wrap `raw` in a `<tool_output_quarantine>` envelope after neutralizing any
206/// embedded [`DELIMITER_RULES`] and breakout tags.
207///
208/// This function is **idempotent**: if `raw` already satisfies [`is_quarantined_output`],
209/// it is returned unchanged without double-wrapping.
210#[must_use]
211pub fn quarantine_tool_output(raw: &str) -> String {
212    if is_quarantined_output(raw) {
213        return raw.to_owned();
214    }
215    let sanitized = neutralize_delimiters(raw);
216    let mut buf = String::with_capacity(OPEN_FRAME.len() + sanitized.len() + CLOSE_FRAME.len());
217    buf.push_str(OPEN_FRAME);
218    buf.push_str(&sanitized);
219    buf.push_str(CLOSE_FRAME);
220    buf
221}
222
223impl ToolOutput {
224    /// Return a quarantined copy of this [`ToolOutput`], neutralizing control
225    /// delimiters and wrapping the content in `<tool_output_quarantine>` tags.
226    ///
227    /// Idempotent when called on an already-quarantined output.
228    #[must_use]
229    pub fn quarantined(mut self) -> Self {
230        self.content = quarantine_tool_output(&self.content);
231        self
232    }
233
234    /// Returns `true` if this output's content is wrapped in a verified
235    /// `<tool_output_quarantine>` envelope with no un-neutralized delimiters.
236    #[must_use]
237    pub fn is_quarantined(&self) -> bool {
238        is_quarantined_output(&self.content)
239    }
240}
llm_tool/types/quarantine.rs

llm_tool/types/
quarantine.rs
