harn_vm/
visible_text.rs

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            // Tagged response protocol: hide tool-call bodies (executed as
42            // structured data, never surfaced as narration) and done
43            // blocks (runtime signal, not user-facing).
44            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
134/// Strip the wrapper tags around `<assistant_prose>` blocks so the
135/// surfaced visible text reads as plain narration. When a
136/// `<user_response>` block is present, it becomes the authoritative
137/// host-facing surface and supersedes generic assistant prose.
138fn 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: &regex::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 strip_partial_marker_suffix(text: &str) -> String {
317    const MARKERS: [&str; 13] = [
318        "<|tool_call|>",
319        TEXT_TOOL_CALL_OPEN,
320        TEXT_TOOL_CALL_OPEN_COMPACT,
321        "<assistant_prose>",
322        "<assistantprose>",
323        "<user_response>",
324        "<userresponse>",
325        "<done>",
326        "<tool_result",
327        "[result of ",
328        "##DONE##",
329        "DONE",
330        "PLAN_READY",
331    ];
332    for marker in MARKERS {
333        for len in (1..marker.len()).rev() {
334            let prefix = &marker[..len];
335            if let Some(stripped) = text.strip_suffix(prefix) {
336                if is_protocol_tag_position(text, stripped.len()) {
337                    return stripped.to_string();
338                }
339            }
340        }
341    }
342    text.to_string()
343}
344
345fn normalize_visible_whitespace(text: &str) -> String {
346    text.replace("\r\n", "\n")
347        .replace("\n\n\n", "\n\n")
348        .trim()
349        .to_string()
350}
351
352pub fn sanitize_visible_assistant_text(text: &str, partial: bool) -> String {
353    let mut sanitized = text.to_string();
354    for pattern in internal_block_patterns() {
355        sanitized = pattern.replace_all(&sanitized, "").to_string();
356    }
357    // After runtime tags are stripped, surface only the explicit
358    // user-facing response when one exists; otherwise unwrap
359    // <assistant_prose> into plain narration.
360    sanitized = extract_visible_prose(&sanitized);
361    sanitized = strip_internal_json_fences(&sanitized);
362    sanitized = strip_inline_internal_planning_json(&sanitized, partial);
363    if partial {
364        sanitized = strip_unclosed_internal_blocks(&sanitized);
365        sanitized = strip_partial_marker_suffix(&sanitized);
366    }
367    normalize_visible_whitespace(&sanitized)
368}
369
370#[cfg(test)]
371mod tests {
372    use super::{sanitize_visible_assistant_text, VisibleTextState};
373
374    #[test]
375    fn push_emits_incremental_visible_delta_for_plain_chunks() {
376        let mut state = VisibleTextState::default();
377        let (visible, delta) = state.push("Hello", true);
378        assert_eq!(visible, "Hello");
379        assert_eq!(delta, "Hello");
380
381        let (visible, delta) = state.push(" world", true);
382        assert_eq!(visible, "Hello world");
383        assert_eq!(delta, " world");
384    }
385
386    #[test]
387    fn push_hides_open_think_block_until_closed() {
388        let mut state = VisibleTextState::default();
389        let (visible, delta) = state.push("Hi <think>secret", true);
390        assert_eq!(visible, "Hi");
391        assert_eq!(delta, "Hi");
392
393        let (visible, delta) = state.push(" plan</think> bye", true);
394        assert_eq!(visible, "Hi  bye");
395        assert_eq!(delta, "  bye");
396    }
397
398    #[test]
399    fn push_emits_full_visible_text_when_sanitization_shrinks_output() {
400        let mut state = VisibleTextState::default();
401        let (visible, _) = state.push("ok", true);
402        assert_eq!(visible, "ok");
403
404        let (visible, delta) = state.push(" <think>", true);
405        assert_eq!(visible, "ok");
406        // No prefix change so delta is empty.
407        assert_eq!(delta, "");
408    }
409
410    #[test]
411    fn push_partial_marker_suffix_is_held_back_until_resolved() {
412        let mut state = VisibleTextState::default();
413        let (visible, delta) = state.push("Hello\n##DON", true);
414        assert_eq!(visible, "Hello");
415        assert_eq!(delta, "Hello");
416
417        let (visible, delta) = state.push("E##\nmore", true);
418        assert_eq!(visible, "Hello\n\nmore");
419        assert_eq!(delta, "\n\nmore");
420    }
421
422    #[test]
423    fn clear_resets_streaming_state() {
424        let mut state = VisibleTextState::default();
425        let _ = state.push("Hello world", true);
426        state.clear();
427        let (visible, delta) = state.push("fresh", true);
428        assert_eq!(visible, "fresh");
429        assert_eq!(delta, "fresh");
430    }
431
432    #[test]
433    fn sanitize_drops_inline_planner_json_only_with_planner_mode() {
434        let raw = r#"{"mode":"plan_then_execute","plan":[]}"#;
435        assert_eq!(sanitize_visible_assistant_text(raw, false), "");
436        let raw = r#"{"status":"ok","message":"hello"}"#;
437        assert_eq!(sanitize_visible_assistant_text(raw, false), raw);
438    }
439
440    #[test]
441    fn sanitize_prefers_user_response_blocks_over_other_prose() {
442        let raw = "Working...\n<assistant_prose>internal narration</assistant_prose>\n<user_response>Visible answer.</user_response>\n##DONE##";
443        assert_eq!(
444            sanitize_visible_assistant_text(raw, false),
445            "Visible answer."
446        );
447    }
448
449    #[test]
450    fn sanitize_strips_trailing_runtime_sentinel_after_answer_text() {
451        assert_eq!(
452            sanitize_visible_assistant_text("HARN_LOCAL_TOOL_OK##DONE##", false),
453            "HARN_LOCAL_TOOL_OK"
454        );
455        assert_eq!(
456            sanitize_visible_assistant_text("Done.\nPLAN_READY", false),
457            "Done."
458        );
459    }
460
461    #[test]
462    fn sanitize_accepts_compact_protocol_tag_aliases_without_hiding_plain_words() {
463        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>";
464        assert_eq!(
465            sanitize_visible_assistant_text(raw, false),
466            "Visible answer."
467        );
468
469        assert_eq!(
470            sanitize_visible_assistant_text("A tool call summary is fine.", false),
471            "A tool call summary is fine."
472        );
473    }
474
475    #[test]
476    fn sanitize_ignores_inline_user_response_placeholder() {
477        let raw = "Wrap final answers in `<user_response>...</user_response>`.\nAudit: real answer";
478        assert_eq!(sanitize_visible_assistant_text(raw, false), raw);
479    }
480
481    #[test]
482    fn sanitize_prefers_top_level_user_response_over_inline_placeholder() {
483        let raw =
484            "Remember `<user_response>...</user_response>` is the wrapper.\n<user_response>Visible answer.</user_response>";
485        assert_eq!(
486            sanitize_visible_assistant_text(raw, false),
487            "Visible answer."
488        );
489    }
490
491    #[test]
492    fn sanitize_ignores_user_response_inside_markdown_fence() {
493        let raw = "```xml\n<user_response>example only</user_response>\n```\nFinal plain answer.";
494        assert_eq!(sanitize_visible_assistant_text(raw, false), raw);
495    }
496
497    #[test]
498    fn sanitize_partial_keeps_inline_protocol_prefixes() {
499        let raw = "Mention `<user_resp";
500        assert_eq!(sanitize_visible_assistant_text(raw, true), raw);
501    }
502
503    #[test]
504    fn sanitize_partial_hides_top_level_protocol_prefixes() {
505        assert_eq!(sanitize_visible_assistant_text("<user_resp", true), "");
506    }
507}
harn_vm/visible_text.rs

harn_vm/
visible_text.rs
