Skip to main content

agentic_tools_utils/
prompt.rs

1//! Prompt hardening utilities.
2
3/// Truncate a string for safe embedding into LLM prompts.
4///
5/// Returns `(truncated_string, was_truncated)`. Truncation is UTF-8 safe (char-boundary slicing)
6/// and appends `...[truncated]` when the input exceeds `max_chars`.
7pub fn truncate_for_prompt(s: &str, max_chars: usize) -> (String, bool) {
8    const SUFFIX: &str = "...[truncated]";
9
10    match s.char_indices().nth(max_chars) {
11        None => (s.to_string(), false),
12        Some((byte_idx, _)) => {
13            let mut out = String::with_capacity(byte_idx + SUFFIX.len());
14            out.push_str(&s[..byte_idx]);
15            out.push_str(SUFFIX);
16            (out, true)
17        }
18    }
19}
20
21/// Wrap untrusted text in explicit XML-like tags to make the security boundary obvious to the model.
22pub fn wrap_untrusted(tag: &str, body: &str) -> String {
23    format!("<{tag}>\n{body}\n</{tag}>")
24}
25
26#[cfg(test)]
27mod tests {
28    use super::*;
29
30    #[test]
31    fn truncate_does_not_modify_short_string() {
32        let s = "hello";
33        let (out, truncated) = truncate_for_prompt(s, 10);
34        assert_eq!(out, "hello");
35        assert!(!truncated);
36    }
37
38    #[test]
39    fn truncate_adds_suffix_to_long_string() {
40        let s = "hello world this is a long string";
41        let (out, truncated) = truncate_for_prompt(s, 10);
42        assert!(truncated);
43        assert!(out.ends_with("...[truncated]"));
44        // First 10 chars + suffix
45        assert_eq!(&out[..10], "hello worl");
46    }
47
48    #[test]
49    fn truncate_handles_exact_length() {
50        let s = "12345";
51        let (out, truncated) = truncate_for_prompt(s, 5);
52        assert_eq!(out, "12345");
53        assert!(!truncated);
54    }
55
56    #[test]
57    fn truncate_handles_unicode() {
58        // 3 chars: emoji, emoji, "a"
59        let s = "🎉🎊a";
60        let (out, truncated) = truncate_for_prompt(s, 2);
61        assert!(truncated);
62        // Should keep first 2 chars (🎉🎊)
63        assert!(out.starts_with("🎉🎊"));
64        assert!(out.ends_with("...[truncated]"));
65    }
66
67    #[test]
68    fn wrap_untrusted_produces_xml_tags() {
69        let result = wrap_untrusted("untrusted_input", "some data");
70        assert_eq!(result, "<untrusted_input>\nsome data\n</untrusted_input>");
71    }
72
73    #[test]
74    fn wrap_untrusted_preserves_body() {
75        let body = "line1\nline2\nline3";
76        let result = wrap_untrusted("data", body);
77        assert!(result.contains("line1\nline2\nline3"));
78    }
79}