diff --git a/src/safety/mod.rs b/src/safety/mod.rs index 31fda95e..83d79627 100644 --- a/src/safety/mod.rs +++ b/src/safety/mod.rs @@ -247,26 +247,42 @@ fn escape_xml_attr(s: &str) -> String { } /// Neutralize closing ``. The leading `<` is replaced /// with `<\u{200B}` (zero-width space) so JSON and other content passes /// through unchanged. fn escape_tool_output_close(s: &str) -> String { - // Case-insensitive search for String { /// Reverse the escaping applied by [`escape_tool_output_close`] by removing /// the zero-width space inserted after `<` in ` String { - s.replace("<\u{200B}/", " suffix"; + let wrapped = safety.wrap_for_llm("t", malicious); + assert!(wrapped.contains("<\u{200B}/ Tool_Output>")); + let unwrapped = SafetyLayer::unwrap_tool_output(&wrapped).expect("should unwrap"); + assert_eq!(unwrapped, malicious); + } + #[test] fn test_escape_tool_output_close_only_targets_closing_tag() { // Regular content passes through unchanged @@ -434,6 +501,13 @@ mod tests { ); // Only ").contains("<\u{200B}/tool_output>")); + assert!(escape_tool_output_close("").contains("<\u{200B}/ Tool_Output>")); + } + + #[test] + fn test_unescape_tool_output_close_ignores_other_sequences() { + let untouched = "prefix <\u{200B}/not_tool_output> suffix"; + assert_eq!(unescape_tool_output_close(untouched), untouched); } #[test]