1use std::collections::BTreeSet;
2use std::sync::OnceLock;
3
4use crate::llm::tools::{
5 TEXT_TOOL_CALL_CLOSE, TEXT_TOOL_CALL_CLOSE_COMPACT, TEXT_TOOL_CALL_OPEN,
6 TEXT_TOOL_CALL_OPEN_COMPACT,
7};
8use regex::Regex;
9
10#[derive(Default, Clone, Debug, PartialEq, Eq)]
11pub struct VisibleTextState {
12 raw_text: String,
13 last_visible_text: String,
14}
15
16impl VisibleTextState {
17 pub fn push(&mut self, delta: &str, partial: bool) -> (String, String) {
18 self.raw_text.push_str(delta);
19 let visible_text = sanitize_visible_assistant_text(&self.raw_text, partial);
20 let visible_delta = visible_text
21 .strip_prefix(&self.last_visible_text)
22 .unwrap_or(visible_text.as_str())
23 .to_string();
24 self.last_visible_text = visible_text.clone();
25 (visible_text, visible_delta)
26 }
27
28 pub fn clear(&mut self) {
29 self.raw_text.clear();
30 self.last_visible_text.clear();
31 }
32}
33
34fn internal_block_patterns() -> &'static [Regex] {
35 static PATTERNS: OnceLock<Vec<Regex>> = OnceLock::new();
36 PATTERNS.get_or_init(|| {
37 [
38 r"(?s)<think>.*?</think>",
39 r"(?s)<think>.*$",
40 r"(?s)<\|tool_call\|>.*?</\|tool_call\|>",
41 r"(?s)<tool_?call>.*?</tool_?call>",
45 r"(?s)<done>.*?</done>",
46 r"(?s)<tool_result[^>]*>.*?</tool_result>",
47 r"(?s)\[result of [^\]]+\].*?\[end of [^\]]+\]",
48 r"(?m)^\s*(##DONE##|DONE|PLAN_READY)\s*$",
49 r"(?s)\s*(##DONE##|PLAN_READY)\s*$",
50 ]
51 .into_iter()
52 .map(|pattern| Regex::new(pattern).expect("valid assistant sanitization regex"))
53 .collect()
54 })
55}
56
57fn assistant_prose_regex() -> &'static Regex {
58 static RE: OnceLock<Regex> = OnceLock::new();
59 RE.get_or_init(|| {
60 Regex::new(r"(?ms)^[ \t]*<assistant_?prose>\s*(.*?)\s*</assistant_?prose>")
61 .expect("valid assistant_prose regex")
62 })
63}
64
65fn user_response_regex() -> &'static Regex {
66 static RE: OnceLock<Regex> = OnceLock::new();
67 RE.get_or_init(|| {
68 Regex::new(r"(?ms)^[ \t]*<user_?response>\s*(.*?)\s*</user_?response>")
69 .expect("valid user_response regex")
70 })
71}
72
73fn inside_markdown_fence(text: &str, idx: usize) -> bool {
74 let mut count = 0;
75 let mut cursor = 0;
76 while cursor < idx {
77 let Some(pos) = text[cursor..idx].find("```") else {
78 break;
79 };
80 count += 1;
81 cursor += pos + 3;
82 }
83 count % 2 == 1
84}
85
86fn is_top_level_tag_position(text: &str, idx: usize) -> bool {
87 let line_start = text[..idx].rfind('\n').map(|pos| pos + 1).unwrap_or(0);
88 text[line_start..idx]
89 .chars()
90 .all(|ch| matches!(ch, ' ' | '\t' | '\r'))
91}
92
93fn is_protocol_tag_position(text: &str, idx: usize) -> bool {
94 is_top_level_tag_position(text, idx) && !inside_markdown_fence(text, idx)
95}
96
97fn extract_user_response(text: &str) -> Option<String> {
98 let sections: Vec<String> = user_response_regex()
99 .captures_iter(text)
100 .filter(|caps| {
101 caps.get(0)
102 .is_some_and(|m| is_protocol_tag_position(text, m.start()))
103 })
104 .filter_map(|caps| caps.get(1).map(|m| m.as_str().trim().to_string()))
105 .filter(|section| !section.is_empty())
106 .collect();
107 if sections.is_empty() {
108 None
109 } else {
110 Some(sections.join("\n\n"))
111 }
112}
113
114fn unwrap_assistant_prose(text: &str) -> String {
115 let mut out = String::with_capacity(text.len());
116 let mut last = 0;
117 for caps in assistant_prose_regex().captures_iter(text) {
118 let Some(block) = caps.get(0) else {
119 continue;
120 };
121 if !is_protocol_tag_position(text, block.start()) {
122 continue;
123 }
124 out.push_str(&text[last..block.start()]);
125 if let Some(body) = caps.get(1) {
126 out.push_str(body.as_str().trim());
127 }
128 last = block.end();
129 }
130 out.push_str(&text[last..]);
131 out
132}
133
134fn extract_visible_prose(text: &str) -> String {
139 if let Some(user_response) = extract_user_response(text) {
140 return user_response;
141 }
142 unwrap_assistant_prose(text)
143}
144
145fn json_fence_regex() -> &'static Regex {
146 static JSON_FENCE: OnceLock<Regex> = OnceLock::new();
147 JSON_FENCE
148 .get_or_init(|| Regex::new(r"(?s)```json[^\n]*\n(.*?)```").expect("valid json fence regex"))
149}
150
151fn inline_planner_json_regex() -> &'static Regex {
152 static INLINE_PLANNER_JSON: OnceLock<Regex> = OnceLock::new();
153 INLINE_PLANNER_JSON.get_or_init(|| {
154 Regex::new(r#"(?s)\{\s*"mode"\s*:\s*"(?:fast_execute|plan_then_execute|ask_user)".*?\}"#)
155 .expect("valid inline planner json regex")
156 })
157}
158
159fn partial_inline_planner_json_regex() -> &'static Regex {
160 static PARTIAL_INLINE_PLANNER_JSON: OnceLock<Regex> = OnceLock::new();
161 PARTIAL_INLINE_PLANNER_JSON.get_or_init(|| {
162 Regex::new(r#"(?s)\{\s*"mode"\s*:\s*"(?:fast_execute|plan_then_execute|ask_user)".*$"#)
163 .expect("valid partial inline planner json regex")
164 })
165}
166
167fn looks_like_internal_planning_json(source: &str) -> bool {
168 let trimmed = source.trim();
169 if !(trimmed.starts_with('{') || trimmed.starts_with('[')) {
170 return false;
171 }
172
173 fn collect_keys(value: &serde_json::Value, keys: &mut BTreeSet<String>) {
174 match value {
175 serde_json::Value::Object(map) => {
176 for (key, child) in map {
177 keys.insert(key.clone());
178 collect_keys(child, keys);
179 }
180 }
181 serde_json::Value::Array(items) => {
182 for item in items {
183 collect_keys(item, keys);
184 }
185 }
186 _ => {}
187 }
188 }
189
190 if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(trimmed) {
191 let mut keys = BTreeSet::new();
192 collect_keys(&parsed, &mut keys);
193 let has_planner_mode = match &parsed {
194 serde_json::Value::Object(map) => map
195 .get("mode")
196 .and_then(|value| value.as_str())
197 .is_some_and(|mode| {
198 matches!(mode, "fast_execute" | "plan_then_execute" | "ask_user")
199 }),
200 _ => false,
201 };
202 let has_internal_keys = [
203 "plan",
204 "steps",
205 "tool_calls",
206 "tool_name",
207 "verification",
208 "execution_mode",
209 "required_outputs",
210 "files_to_edit",
211 "next_action",
212 "reasoning",
213 "direction",
214 "targets",
215 "tasks",
216 "unknowns",
217 ]
218 .into_iter()
219 .any(|key| keys.contains(key));
220 return has_planner_mode || has_internal_keys;
221 }
222
223 false
224}
225
226fn strip_internal_json_fences(text: &str) -> String {
227 json_fence_regex()
228 .replace_all(text, |caps: ®ex::Captures| {
229 let body = caps
230 .get(1)
231 .map(|match_| match_.as_str())
232 .unwrap_or_default();
233 if looks_like_internal_planning_json(body) {
234 String::new()
235 } else {
236 caps.get(0)
237 .map(|match_| match_.as_str().to_string())
238 .unwrap_or_default()
239 }
240 })
241 .to_string()
242}
243
244fn strip_unclosed_internal_blocks(text: &str) -> String {
245 if let Some(open_idx) = text.rfind("<|tool_call|>") {
246 let close_idx = text.rfind("</|tool_call|>");
247 if close_idx.is_none_or(|idx| idx < open_idx) {
248 return text[..open_idx].to_string();
249 }
250 }
251
252 if let Some(open_idx) = text.rfind(TEXT_TOOL_CALL_OPEN) {
253 let close_idx = text.rfind(TEXT_TOOL_CALL_CLOSE);
254 if is_protocol_tag_position(text, open_idx) && close_idx.is_none_or(|idx| idx < open_idx) {
255 return text[..open_idx].to_string();
256 }
257 }
258
259 if let Some(open_idx) = text.rfind(TEXT_TOOL_CALL_OPEN_COMPACT) {
260 let close_idx = text.rfind(TEXT_TOOL_CALL_CLOSE_COMPACT);
261 if is_protocol_tag_position(text, open_idx) && close_idx.is_none_or(|idx| idx < open_idx) {
262 return text[..open_idx].to_string();
263 }
264 }
265
266 if let Some(open_idx) = text.rfind("<done>") {
267 let close_idx = text.rfind("</done>");
268 if is_protocol_tag_position(text, open_idx) && close_idx.is_none_or(|idx| idx < open_idx) {
269 return text[..open_idx].to_string();
270 }
271 }
272
273 if let Some(open_idx) = text.rfind("<user_response>") {
274 let close_idx = text.rfind("</user_response>");
275 if is_protocol_tag_position(text, open_idx) && close_idx.is_none_or(|idx| idx < open_idx) {
276 return text[..open_idx].to_string();
277 }
278 }
279
280 if let Some(open_idx) = text.rfind("<userresponse>") {
281 let close_idx = text.rfind("</userresponse>");
282 if is_protocol_tag_position(text, open_idx) && close_idx.is_none_or(|idx| idx < open_idx) {
283 return text[..open_idx].to_string();
284 }
285 }
286
287 if let Some(open_idx) = text.rfind("[result of ") {
288 let close_idx = text.rfind("[end of ");
289 if close_idx.is_none_or(|idx| idx < open_idx) {
290 return text[..open_idx].to_string();
291 }
292 }
293
294 if let Some(open_idx) = text.rfind("<tool_result") {
295 let close_idx = text.rfind("</tool_result>");
296 if is_protocol_tag_position(text, open_idx) && close_idx.is_none_or(|idx| idx < open_idx) {
297 return text[..open_idx].to_string();
298 }
299 }
300
301 text.to_string()
302}
303
304fn strip_inline_internal_planning_json(text: &str, partial: bool) -> String {
305 let mut stripped = inline_planner_json_regex()
306 .replace_all(text, "")
307 .to_string();
308 if partial {
309 stripped = partial_inline_planner_json_regex()
310 .replace_all(&stripped, "")
311 .to_string();
312 }
313 stripped
314}
315
316fn protocol_residue_regex() -> &'static Regex {
317 static RE: OnceLock<Regex> = OnceLock::new();
337 RE.get_or_init(|| {
338 Regex::new(r"<?/?\|?(?:t?o?o?l?)_call\|?>|<?/?\|?[a-z]*_prose>")
339 .expect("valid protocol residue regex")
340 })
341}
342
343fn strip_protocol_residue(text: &str) -> String {
344 protocol_residue_regex()
348 .replace_all(text, |caps: ®ex::Captures| {
349 let matched = caps.get(0).expect("capture group 0 always present");
350 if inside_markdown_fence(text, matched.start()) {
351 matched.as_str().to_string()
352 } else {
353 String::new()
354 }
355 })
356 .to_string()
357}
358
359fn looks_like_bare_internal_verdict_json(source: &str) -> bool {
360 let trimmed = source.trim();
361 if !trimmed.starts_with('{') {
362 return false;
363 }
364
365 let Ok(serde_json::Value::Object(map)) = serde_json::from_str::<serde_json::Value>(trimmed)
366 else {
367 return false;
368 };
369
370 let Some(verdict) = map
371 .get("verdict")
372 .and_then(|value| value.as_str())
373 .map(str::trim)
374 .filter(|value| !value.is_empty())
375 else {
376 return false;
377 };
378
379 let verdict = verdict.to_ascii_lowercase();
380 let has_completion_explanation = map.contains_key("reasoning")
381 || map.contains_key("reason")
382 || map.contains_key("next_step")
383 || map.contains_key("nextStep");
384 let has_judge_metadata = map.contains_key("critique")
385 || map.contains_key("confidence")
386 || map.contains_key("category")
387 || map.contains_key("error");
388
389 let known_internal_verdict = matches!(verdict.as_str(), "done" | "continue")
390 && has_completion_explanation
391 || matches!(verdict.as_str(), "revise" | "pass" | "fail" | "unclear") && has_judge_metadata
392 || matches!(verdict.as_str(), "allow" | "warn" | "block") && has_judge_metadata;
393 if !known_internal_verdict {
394 return false;
395 }
396
397 map.keys().all(|key| {
398 matches!(
399 key.as_str(),
400 "verdict"
401 | "reasoning"
402 | "reason"
403 | "next_step"
404 | "nextStep"
405 | "critique"
406 | "confidence"
407 | "category"
408 | "error"
409 )
410 })
411}
412
413fn strip_bare_internal_json(text: &str) -> String {
414 if looks_like_bare_internal_verdict_json(text) {
423 return String::new();
424 }
425 text.to_string()
426}
427
428fn strip_partial_marker_suffix(text: &str) -> String {
429 const MARKERS: [&str; 13] = [
430 "<|tool_call|>",
431 TEXT_TOOL_CALL_OPEN,
432 TEXT_TOOL_CALL_OPEN_COMPACT,
433 "<assistant_prose>",
434 "<assistantprose>",
435 "<user_response>",
436 "<userresponse>",
437 "<done>",
438 "<tool_result",
439 "[result of ",
440 "##DONE##",
441 "DONE",
442 "PLAN_READY",
443 ];
444 for marker in MARKERS {
445 for len in (1..marker.len()).rev() {
446 let prefix = &marker[..len];
447 if let Some(stripped) = text.strip_suffix(prefix) {
448 if is_protocol_tag_position(text, stripped.len()) {
449 return stripped.to_string();
450 }
451 }
452 }
453 }
454 text.to_string()
455}
456
457fn normalize_visible_whitespace(text: &str) -> String {
458 text.replace("\r\n", "\n")
459 .replace("\n\n\n", "\n\n")
460 .trim()
461 .to_string()
462}
463
464pub fn sanitize_visible_assistant_text(text: &str, partial: bool) -> String {
465 let mut sanitized = text.to_string();
466 for pattern in internal_block_patterns() {
467 sanitized = pattern.replace_all(&sanitized, "").to_string();
468 }
469 sanitized = extract_visible_prose(&sanitized);
473 sanitized = strip_internal_json_fences(&sanitized);
474 sanitized = strip_inline_internal_planning_json(&sanitized, partial);
475 sanitized = strip_protocol_residue(&sanitized);
480 sanitized = strip_bare_internal_json(sanitized.trim());
481 if partial {
482 sanitized = strip_unclosed_internal_blocks(&sanitized);
483 sanitized = strip_partial_marker_suffix(&sanitized);
484 }
485 normalize_visible_whitespace(&sanitized)
486}
487
488#[cfg(test)]
489mod tests {
490 use super::{sanitize_visible_assistant_text, VisibleTextState};
491
492 #[test]
493 fn push_emits_incremental_visible_delta_for_plain_chunks() {
494 let mut state = VisibleTextState::default();
495 let (visible, delta) = state.push("Hello", true);
496 assert_eq!(visible, "Hello");
497 assert_eq!(delta, "Hello");
498
499 let (visible, delta) = state.push(" world", true);
500 assert_eq!(visible, "Hello world");
501 assert_eq!(delta, " world");
502 }
503
504 #[test]
505 fn push_hides_open_think_block_until_closed() {
506 let mut state = VisibleTextState::default();
507 let (visible, delta) = state.push("Hi <think>secret", true);
508 assert_eq!(visible, "Hi");
509 assert_eq!(delta, "Hi");
510
511 let (visible, delta) = state.push(" plan</think> bye", true);
512 assert_eq!(visible, "Hi bye");
513 assert_eq!(delta, " bye");
514 }
515
516 #[test]
517 fn push_emits_full_visible_text_when_sanitization_shrinks_output() {
518 let mut state = VisibleTextState::default();
519 let (visible, _) = state.push("ok", true);
520 assert_eq!(visible, "ok");
521
522 let (visible, delta) = state.push(" <think>", true);
523 assert_eq!(visible, "ok");
524 assert_eq!(delta, "");
526 }
527
528 #[test]
529 fn push_partial_marker_suffix_is_held_back_until_resolved() {
530 let mut state = VisibleTextState::default();
531 let (visible, delta) = state.push("Hello\n##DON", true);
532 assert_eq!(visible, "Hello");
533 assert_eq!(delta, "Hello");
534
535 let (visible, delta) = state.push("E##\nmore", true);
536 assert_eq!(visible, "Hello\n\nmore");
537 assert_eq!(delta, "\n\nmore");
538 }
539
540 #[test]
541 fn clear_resets_streaming_state() {
542 let mut state = VisibleTextState::default();
543 let _ = state.push("Hello world", true);
544 state.clear();
545 let (visible, delta) = state.push("fresh", true);
546 assert_eq!(visible, "fresh");
547 assert_eq!(delta, "fresh");
548 }
549
550 #[test]
551 fn sanitize_drops_inline_planner_json_only_with_planner_mode() {
552 let raw = r#"{"mode":"plan_then_execute","plan":[]}"#;
553 assert_eq!(sanitize_visible_assistant_text(raw, false), "");
554 let raw = r#"{"status":"ok","message":"hello"}"#;
555 assert_eq!(sanitize_visible_assistant_text(raw, false), raw);
556 }
557
558 #[test]
559 fn sanitize_strips_orphan_tool_call_residue_and_truncations() {
560 assert_eq!(sanitize_visible_assistant_text("_call>", false), "");
563 assert_eq!(sanitize_visible_assistant_text("l_call>l_call>", false), "");
564 assert_eq!(
565 sanitize_visible_assistant_text("Done.\n})\n</tool_call>_call>", false),
566 "Done.\n})"
567 );
568 assert_eq!(
569 sanitize_visible_assistant_text("Implemented.</assistant_prose>", false),
570 "Implemented."
571 );
572 assert_eq!(
574 sanitize_visible_assistant_text("Implemented.\nnt_prose>", false),
575 "Implemented."
576 );
577 let fenced = "```\n</tool_call>\n```\nDone.";
579 assert_eq!(sanitize_visible_assistant_text(fenced, false), fenced);
580 }
581
582 #[test]
583 fn sanitize_does_not_touch_ordinary_prose_or_inequalities() {
584 let raw = "Use a_call> only as— wait, compare x > y and y > z here.";
586 let out = sanitize_visible_assistant_text(raw, false);
588 assert!(out.contains("compare x > y and y > z here."), "got: {out}");
589 assert_eq!(
590 sanitize_visible_assistant_text("The phrase tool call is normal prose.", false),
591 "The phrase tool call is normal prose."
592 );
593 }
594
595 #[test]
596 fn sanitize_drops_bare_completion_judge_verdict_json() {
597 let raw = r#"{"verdict":"done","reasoning":"All tests pass.","next_step":""}"#;
598 assert_eq!(sanitize_visible_assistant_text(raw, false), "");
599 let padded = "\n {\"verdict\":\"continue\",\"reasoning\":\"does not compile\"} \n";
601 assert_eq!(sanitize_visible_assistant_text(padded, false), "");
602 let keep = r#"{"status":"ok","message":"hello"}"#;
604 assert_eq!(sanitize_visible_assistant_text(keep, false), keep);
605 let visible_answer =
609 r#"{"tasks":["ship"],"steps":["test"],"reasoning":"user-visible rationale"}"#;
610 assert_eq!(
611 sanitize_visible_assistant_text(visible_answer, false),
612 visible_answer
613 );
614 let visible_verdict = r#"{"verdict":"pass","summary":"public result"}"#;
615 assert_eq!(
616 sanitize_visible_assistant_text(visible_verdict, false),
617 visible_verdict
618 );
619 let visible_verdict_rationale = r#"{"verdict":"pass","reasoning":"public rationale"}"#;
620 assert_eq!(
621 sanitize_visible_assistant_text(visible_verdict_rationale, false),
622 visible_verdict_rationale
623 );
624 }
625
626 #[test]
627 fn sanitize_prefers_user_response_blocks_over_other_prose() {
628 let raw = "Working...\n<assistant_prose>internal narration</assistant_prose>\n<user_response>Visible answer.</user_response>\n##DONE##";
629 assert_eq!(
630 sanitize_visible_assistant_text(raw, false),
631 "Visible answer."
632 );
633 }
634
635 #[test]
636 fn sanitize_strips_trailing_runtime_sentinel_after_answer_text() {
637 assert_eq!(
638 sanitize_visible_assistant_text("HARN_LOCAL_TOOL_OK##DONE##", false),
639 "HARN_LOCAL_TOOL_OK"
640 );
641 assert_eq!(
642 sanitize_visible_assistant_text("Done.\nPLAN_READY", false),
643 "Done."
644 );
645 }
646
647 #[test]
648 fn sanitize_accepts_compact_protocol_tag_aliases_without_hiding_plain_words() {
649 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>";
650 assert_eq!(
651 sanitize_visible_assistant_text(raw, false),
652 "Visible answer."
653 );
654
655 assert_eq!(
656 sanitize_visible_assistant_text("A tool call summary is fine.", false),
657 "A tool call summary is fine."
658 );
659 }
660
661 #[test]
662 fn sanitize_ignores_inline_user_response_placeholder() {
663 let raw = "Wrap final answers in `<user_response>...</user_response>`.\nAudit: real answer";
664 assert_eq!(sanitize_visible_assistant_text(raw, false), raw);
665 }
666
667 #[test]
668 fn sanitize_prefers_top_level_user_response_over_inline_placeholder() {
669 let raw =
670 "Remember `<user_response>...</user_response>` is the wrapper.\n<user_response>Visible answer.</user_response>";
671 assert_eq!(
672 sanitize_visible_assistant_text(raw, false),
673 "Visible answer."
674 );
675 }
676
677 #[test]
678 fn sanitize_ignores_user_response_inside_markdown_fence() {
679 let raw = "```xml\n<user_response>example only</user_response>\n```\nFinal plain answer.";
680 assert_eq!(sanitize_visible_assistant_text(raw, false), raw);
681 }
682
683 #[test]
684 fn sanitize_partial_keeps_inline_protocol_prefixes() {
685 let raw = "Mention `<user_resp";
686 assert_eq!(sanitize_visible_assistant_text(raw, true), raw);
687 }
688
689 #[test]
690 fn sanitize_partial_hides_top_level_protocol_prefixes() {
691 assert_eq!(sanitize_visible_assistant_text("<user_resp", true), "");
692 }
693}