use std::collections::BTreeSet;
use std::sync::OnceLock;
use crate::llm::tools::{
TEXT_TOOL_CALL_CLOSE, TEXT_TOOL_CALL_CLOSE_COMPACT, TEXT_TOOL_CALL_OPEN,
TEXT_TOOL_CALL_OPEN_COMPACT,
};
use regex::Regex;
#[derive(Default, Clone, Debug, PartialEq, Eq)]
pub struct VisibleTextState {
raw_text: String,
last_visible_text: String,
}
impl VisibleTextState {
pub fn push(&mut self, delta: &str, partial: bool) -> (String, String) {
self.raw_text.push_str(delta);
let visible_text = sanitize_visible_assistant_text(&self.raw_text, partial);
let visible_delta = visible_text
.strip_prefix(&self.last_visible_text)
.unwrap_or(visible_text.as_str())
.to_string();
self.last_visible_text = visible_text.clone();
(visible_text, visible_delta)
}
pub fn clear(&mut self) {
self.raw_text.clear();
self.last_visible_text.clear();
}
}
fn internal_block_patterns() -> &'static [Regex] {
static PATTERNS: OnceLock<Vec<Regex>> = OnceLock::new();
PATTERNS.get_or_init(|| {
[
r"(?s)<think>.*?</think>",
r"(?s)<think>.*$",
r"(?s)<\|tool_call\|>.*?</\|tool_call\|>",
r"(?s)<tool_?call>.*?</tool_?call>",
r"(?s)<done>.*?</done>",
r"(?s)<tool_result[^>]*>.*?</tool_result>",
r"(?s)\[result of [^\]]+\].*?\[end of [^\]]+\]",
r"(?m)^\s*(##DONE##|DONE|PLAN_READY)\s*$",
r"(?s)\s*(##DONE##|PLAN_READY)\s*$",
]
.into_iter()
.map(|pattern| Regex::new(pattern).expect("valid assistant sanitization regex"))
.collect()
})
}
fn assistant_prose_regex() -> &'static Regex {
static RE: OnceLock<Regex> = OnceLock::new();
RE.get_or_init(|| {
Regex::new(r"(?ms)^[ \t]*<assistant_?prose>\s*(.*?)\s*</assistant_?prose>")
.expect("valid assistant_prose regex")
})
}
fn user_response_regex() -> &'static Regex {
static RE: OnceLock<Regex> = OnceLock::new();
RE.get_or_init(|| {
Regex::new(r"(?ms)^[ \t]*<user_?response>\s*(.*?)\s*</user_?response>")
.expect("valid user_response regex")
})
}
fn inside_markdown_fence(text: &str, idx: usize) -> bool {
let mut count = 0;
let mut cursor = 0;
while cursor < idx {
let Some(pos) = text[cursor..idx].find("```") else {
break;
};
count += 1;
cursor += pos + 3;
}
count % 2 == 1
}
fn is_top_level_tag_position(text: &str, idx: usize) -> bool {
let line_start = text[..idx].rfind('\n').map(|pos| pos + 1).unwrap_or(0);
text[line_start..idx]
.chars()
.all(|ch| matches!(ch, ' ' | '\t' | '\r'))
}
fn is_protocol_tag_position(text: &str, idx: usize) -> bool {
is_top_level_tag_position(text, idx) && !inside_markdown_fence(text, idx)
}
fn extract_user_response(text: &str) -> Option<String> {
let sections: Vec<String> = user_response_regex()
.captures_iter(text)
.filter(|caps| {
caps.get(0)
.is_some_and(|m| is_protocol_tag_position(text, m.start()))
})
.filter_map(|caps| caps.get(1).map(|m| m.as_str().trim().to_string()))
.filter(|section| !section.is_empty())
.collect();
if sections.is_empty() {
None
} else {
Some(sections.join("\n\n"))
}
}
fn unwrap_assistant_prose(text: &str) -> String {
let mut out = String::with_capacity(text.len());
let mut last = 0;
for caps in assistant_prose_regex().captures_iter(text) {
let Some(block) = caps.get(0) else {
continue;
};
if !is_protocol_tag_position(text, block.start()) {
continue;
}
out.push_str(&text[last..block.start()]);
if let Some(body) = caps.get(1) {
out.push_str(body.as_str().trim());
}
last = block.end();
}
out.push_str(&text[last..]);
out
}
fn extract_visible_prose(text: &str) -> String {
if let Some(user_response) = extract_user_response(text) {
return user_response;
}
unwrap_assistant_prose(text)
}
fn json_fence_regex() -> &'static Regex {
static JSON_FENCE: OnceLock<Regex> = OnceLock::new();
JSON_FENCE
.get_or_init(|| Regex::new(r"(?s)```json[^\n]*\n(.*?)```").expect("valid json fence regex"))
}
fn inline_planner_json_regex() -> &'static Regex {
static INLINE_PLANNER_JSON: OnceLock<Regex> = OnceLock::new();
INLINE_PLANNER_JSON.get_or_init(|| {
Regex::new(r#"(?s)\{\s*"mode"\s*:\s*"(?:fast_execute|plan_then_execute|ask_user)".*?\}"#)
.expect("valid inline planner json regex")
})
}
fn partial_inline_planner_json_regex() -> &'static Regex {
static PARTIAL_INLINE_PLANNER_JSON: OnceLock<Regex> = OnceLock::new();
PARTIAL_INLINE_PLANNER_JSON.get_or_init(|| {
Regex::new(r#"(?s)\{\s*"mode"\s*:\s*"(?:fast_execute|plan_then_execute|ask_user)".*$"#)
.expect("valid partial inline planner json regex")
})
}
fn looks_like_internal_planning_json(source: &str) -> bool {
let trimmed = source.trim();
if !(trimmed.starts_with('{') || trimmed.starts_with('[')) {
return false;
}
fn collect_keys(value: &serde_json::Value, keys: &mut BTreeSet<String>) {
match value {
serde_json::Value::Object(map) => {
for (key, child) in map {
keys.insert(key.clone());
collect_keys(child, keys);
}
}
serde_json::Value::Array(items) => {
for item in items {
collect_keys(item, keys);
}
}
_ => {}
}
}
if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(trimmed) {
let mut keys = BTreeSet::new();
collect_keys(&parsed, &mut keys);
let has_planner_mode = match &parsed {
serde_json::Value::Object(map) => map
.get("mode")
.and_then(|value| value.as_str())
.is_some_and(|mode| {
matches!(mode, "fast_execute" | "plan_then_execute" | "ask_user")
}),
_ => false,
};
let has_internal_keys = [
"plan",
"steps",
"tool_calls",
"tool_name",
"verification",
"execution_mode",
"required_outputs",
"files_to_edit",
"next_action",
"reasoning",
"direction",
"targets",
"tasks",
"unknowns",
]
.into_iter()
.any(|key| keys.contains(key));
return has_planner_mode || has_internal_keys;
}
false
}
fn strip_internal_json_fences(text: &str) -> String {
json_fence_regex()
.replace_all(text, |caps: ®ex::Captures| {
let body = caps
.get(1)
.map(|match_| match_.as_str())
.unwrap_or_default();
if looks_like_internal_planning_json(body) {
String::new()
} else {
caps.get(0)
.map(|match_| match_.as_str().to_string())
.unwrap_or_default()
}
})
.to_string()
}
fn strip_unclosed_internal_blocks(text: &str) -> String {
if let Some(open_idx) = text.rfind("<|tool_call|>") {
let close_idx = text.rfind("</|tool_call|>");
if close_idx.is_none_or(|idx| idx < open_idx) {
return text[..open_idx].to_string();
}
}
if let Some(open_idx) = text.rfind(TEXT_TOOL_CALL_OPEN) {
let close_idx = text.rfind(TEXT_TOOL_CALL_CLOSE);
if is_protocol_tag_position(text, open_idx) && close_idx.is_none_or(|idx| idx < open_idx) {
return text[..open_idx].to_string();
}
}
if let Some(open_idx) = text.rfind(TEXT_TOOL_CALL_OPEN_COMPACT) {
let close_idx = text.rfind(TEXT_TOOL_CALL_CLOSE_COMPACT);
if is_protocol_tag_position(text, open_idx) && close_idx.is_none_or(|idx| idx < open_idx) {
return text[..open_idx].to_string();
}
}
if let Some(open_idx) = text.rfind("<done>") {
let close_idx = text.rfind("</done>");
if is_protocol_tag_position(text, open_idx) && close_idx.is_none_or(|idx| idx < open_idx) {
return text[..open_idx].to_string();
}
}
if let Some(open_idx) = text.rfind("<user_response>") {
let close_idx = text.rfind("</user_response>");
if is_protocol_tag_position(text, open_idx) && close_idx.is_none_or(|idx| idx < open_idx) {
return text[..open_idx].to_string();
}
}
if let Some(open_idx) = text.rfind("<userresponse>") {
let close_idx = text.rfind("</userresponse>");
if is_protocol_tag_position(text, open_idx) && close_idx.is_none_or(|idx| idx < open_idx) {
return text[..open_idx].to_string();
}
}
if let Some(open_idx) = text.rfind("[result of ") {
let close_idx = text.rfind("[end of ");
if close_idx.is_none_or(|idx| idx < open_idx) {
return text[..open_idx].to_string();
}
}
if let Some(open_idx) = text.rfind("<tool_result") {
let close_idx = text.rfind("</tool_result>");
if is_protocol_tag_position(text, open_idx) && close_idx.is_none_or(|idx| idx < open_idx) {
return text[..open_idx].to_string();
}
}
text.to_string()
}
fn strip_inline_internal_planning_json(text: &str, partial: bool) -> String {
let mut stripped = inline_planner_json_regex()
.replace_all(text, "")
.to_string();
if partial {
stripped = partial_inline_planner_json_regex()
.replace_all(&stripped, "")
.to_string();
}
stripped
}
fn strip_partial_marker_suffix(text: &str) -> String {
const MARKERS: [&str; 13] = [
"<|tool_call|>",
TEXT_TOOL_CALL_OPEN,
TEXT_TOOL_CALL_OPEN_COMPACT,
"<assistant_prose>",
"<assistantprose>",
"<user_response>",
"<userresponse>",
"<done>",
"<tool_result",
"[result of ",
"##DONE##",
"DONE",
"PLAN_READY",
];
for marker in MARKERS {
for len in (1..marker.len()).rev() {
let prefix = &marker[..len];
if let Some(stripped) = text.strip_suffix(prefix) {
if is_protocol_tag_position(text, stripped.len()) {
return stripped.to_string();
}
}
}
}
text.to_string()
}
fn normalize_visible_whitespace(text: &str) -> String {
text.replace("\r\n", "\n")
.replace("\n\n\n", "\n\n")
.trim()
.to_string()
}
pub fn sanitize_visible_assistant_text(text: &str, partial: bool) -> String {
let mut sanitized = text.to_string();
for pattern in internal_block_patterns() {
sanitized = pattern.replace_all(&sanitized, "").to_string();
}
sanitized = extract_visible_prose(&sanitized);
sanitized = strip_internal_json_fences(&sanitized);
sanitized = strip_inline_internal_planning_json(&sanitized, partial);
if partial {
sanitized = strip_unclosed_internal_blocks(&sanitized);
sanitized = strip_partial_marker_suffix(&sanitized);
}
normalize_visible_whitespace(&sanitized)
}
#[cfg(test)]
mod tests {
use super::{sanitize_visible_assistant_text, VisibleTextState};
#[test]
fn push_emits_incremental_visible_delta_for_plain_chunks() {
let mut state = VisibleTextState::default();
let (visible, delta) = state.push("Hello", true);
assert_eq!(visible, "Hello");
assert_eq!(delta, "Hello");
let (visible, delta) = state.push(" world", true);
assert_eq!(visible, "Hello world");
assert_eq!(delta, " world");
}
#[test]
fn push_hides_open_think_block_until_closed() {
let mut state = VisibleTextState::default();
let (visible, delta) = state.push("Hi <think>secret", true);
assert_eq!(visible, "Hi");
assert_eq!(delta, "Hi");
let (visible, delta) = state.push(" plan</think> bye", true);
assert_eq!(visible, "Hi bye");
assert_eq!(delta, " bye");
}
#[test]
fn push_emits_full_visible_text_when_sanitization_shrinks_output() {
let mut state = VisibleTextState::default();
let (visible, _) = state.push("ok", true);
assert_eq!(visible, "ok");
let (visible, delta) = state.push(" <think>", true);
assert_eq!(visible, "ok");
assert_eq!(delta, "");
}
#[test]
fn push_partial_marker_suffix_is_held_back_until_resolved() {
let mut state = VisibleTextState::default();
let (visible, delta) = state.push("Hello\n##DON", true);
assert_eq!(visible, "Hello");
assert_eq!(delta, "Hello");
let (visible, delta) = state.push("E##\nmore", true);
assert_eq!(visible, "Hello\n\nmore");
assert_eq!(delta, "\n\nmore");
}
#[test]
fn clear_resets_streaming_state() {
let mut state = VisibleTextState::default();
let _ = state.push("Hello world", true);
state.clear();
let (visible, delta) = state.push("fresh", true);
assert_eq!(visible, "fresh");
assert_eq!(delta, "fresh");
}
#[test]
fn sanitize_drops_inline_planner_json_only_with_planner_mode() {
let raw = r#"{"mode":"plan_then_execute","plan":[]}"#;
assert_eq!(sanitize_visible_assistant_text(raw, false), "");
let raw = r#"{"status":"ok","message":"hello"}"#;
assert_eq!(sanitize_visible_assistant_text(raw, false), raw);
}
#[test]
fn sanitize_prefers_user_response_blocks_over_other_prose() {
let raw = "Working...\n<assistant_prose>internal narration</assistant_prose>\n<user_response>Visible answer.</user_response>\n##DONE##";
assert_eq!(
sanitize_visible_assistant_text(raw, false),
"Visible answer."
);
}
#[test]
fn sanitize_strips_trailing_runtime_sentinel_after_answer_text() {
assert_eq!(
sanitize_visible_assistant_text("HARN_LOCAL_TOOL_OK##DONE##", false),
"HARN_LOCAL_TOOL_OK"
);
assert_eq!(
sanitize_visible_assistant_text("Done.\nPLAN_READY", false),
"Done."
);
}
#[test]
fn sanitize_accepts_compact_protocol_tag_aliases_without_hiding_plain_words() {
let raw = "The phrase tool call is normal prose.\n<assistantprose>hidden</assistantprose>\n<toolcall>\nrun({ command: \"git status\" })\n</toolcall>\n<userresponse>Visible answer.</userresponse>\n<done>##DONE##</done>";
assert_eq!(
sanitize_visible_assistant_text(raw, false),
"Visible answer."
);
assert_eq!(
sanitize_visible_assistant_text("A tool call summary is fine.", false),
"A tool call summary is fine."
);
}
#[test]
fn sanitize_ignores_inline_user_response_placeholder() {
let raw = "Wrap final answers in `<user_response>...</user_response>`.\nAudit: real answer";
assert_eq!(sanitize_visible_assistant_text(raw, false), raw);
}
#[test]
fn sanitize_prefers_top_level_user_response_over_inline_placeholder() {
let raw =
"Remember `<user_response>...</user_response>` is the wrapper.\n<user_response>Visible answer.</user_response>";
assert_eq!(
sanitize_visible_assistant_text(raw, false),
"Visible answer."
);
}
#[test]
fn sanitize_ignores_user_response_inside_markdown_fence() {
let raw = "```xml\n<user_response>example only</user_response>\n```\nFinal plain answer.";
assert_eq!(sanitize_visible_assistant_text(raw, false), raw);
}
#[test]
fn sanitize_partial_keeps_inline_protocol_prefixes() {
let raw = "Mention `<user_resp";
assert_eq!(sanitize_visible_assistant_text(raw, true), raw);
}
#[test]
fn sanitize_partial_hides_top_level_protocol_prefixes() {
assert_eq!(sanitize_visible_assistant_text("<user_resp", true), "");
}
}