From c9aaba3f5bd85f36be4ebadbe372e282e8f6218a Mon Sep 17 00:00:00 2001 From: chopratejas Date: Wed, 29 Apr 2026 23:18:09 -0700 Subject: [PATCH] feat(rust): port tag_protector to Rust + 5 bug fixes (Phase 3e.4) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `headroom/transforms/tag_protector.py` was a regex-driven scan-and- replace loop that ran on every kompress call from ContentRouter (`content_router.py:1089`). The Python implementation had five real bugs we now fix in the port — the most consequential being a `str.replace(.., .., 1)` first-occurrence-replace bug that silently collapsed two identical custom-tag blocks in the same input to a single placeholder + a stray duplicate of the second block. # Bug fixes (each pinned by a `fixed_in_3e4` test) * **#1: O(n²) on nested custom tags.** Python's `while changed` loop restarted a full regex scan after every replacement. Rust walks once in linear time on input length. * **#2: First-occurrence replace bug.** `result.replace(orig, ph, 1)` replaces the FIRST textual match, not the matched offset. Two identical custom-tag blocks collapsed to one placeholder + a stray duplicate of the second block. The Rust walker stitches output by offset so distinct blocks always get distinct placeholders. * **#3: Silent 50-iteration cap.** Python had a hard `max_iterations = 50` safety limit that quietly truncated tag protection on deeply nested input. The Rust walker is bounded by input length only. * **#4: Self-closing pass duplicate-replace risk.** Python ran a second loop with the same `replace_first` bug for self-closers. Rust handles self-closers in the same single pass. * **#5: Placeholder collision.** If the input contained a literal `{{HEADROOM_TAG_…}}` substring, Python silently let the collision break restoration. Rust salts the prefix and reports it in stats. # Architecture Two-phase walker: * Phase 1 (`identify_spans`): linear scan over input bytes, hand- rolled tag-open / tag-close lexer (no regex). Maintains a stack of open custom tags; on a matching close, collapses the inner span into a single `Span { start, end, Block }`. Self-closing custom tags become `Span { ..., SelfClosing }` immediately. Marker-only mode (`compress_tagged_content=true`) emits Open/CloseMarker spans instead. Orphan opens stay un-protected (matches Python behavior). Orphan closes are emitted verbatim and counted in stats. * Phase 2 (`emit_output`): walks `text` once, splicing placeholders for span ranges and copying everything else verbatim. Offset-based, never `str.replace`. PyO3 surface: `protect_tags`, `restore_tags`, `is_html_tag`, `known_html_tag_names`. The Python shim retires the regex internals and re-exports `KNOWN_HTML_TAGS` (rebuilt from the Rust list) + `_is_html_tag` for backwards compat with `content_router.py` and the existing test surface. # Test plan * 25 Rust unit tests including 4 `fixed_in_3e4_*` bug-fix tests * 27 Python tests (23 existing + 4 new `fixed_in_3e4` parity tests) * 5 integration tests in `test_tag_protection_integration.py` pass * `make ci-precheck` clean --- crates/headroom-core/src/transforms/mod.rs | 2 + .../src/transforms/tag_protector.rs | 972 ++++++++++++++++++ crates/headroom-py/src/lib.rs | 59 ++ headroom/transforms/tag_protector.py | 348 ++----- tests/test_transforms/test_tag_protector.py | 62 ++ 5 files changed, 1172 insertions(+), 271 deletions(-) create mode 100644 crates/headroom-core/src/transforms/tag_protector.rs diff --git a/crates/headroom-core/src/transforms/mod.rs b/crates/headroom-core/src/transforms/mod.rs index cbe60190b..59ae59883 100644 --- a/crates/headroom-core/src/transforms/mod.rs +++ b/crates/headroom-core/src/transforms/mod.rs @@ -24,6 +24,7 @@ pub mod log_compressor; pub mod magika_detector; pub mod search_compressor; pub mod smart_crusher; +pub mod tag_protector; pub mod unidiff_detector; pub use content_detector::{ @@ -42,4 +43,5 @@ pub use search_compressor::{ FileMatches, SearchCompressionResult, SearchCompressor, SearchCompressorConfig, SearchCompressorStats, SearchMatch, }; +pub use tag_protector::{is_known_html_tag, protect_tags, restore_tags, ProtectStats}; pub use unidiff_detector::{detect_diff, is_diff}; diff --git a/crates/headroom-core/src/transforms/tag_protector.rs b/crates/headroom-core/src/transforms/tag_protector.rs new file mode 100644 index 000000000..2b967b961 --- /dev/null +++ b/crates/headroom-core/src/transforms/tag_protector.rs @@ -0,0 +1,972 @@ +//! Tag protection — keep custom workflow tags out of ML compressors. +//! +//! # Why this exists +//! +//! LLM workflows carry XML-style markers (``, +//! ``, ``, ``, etc.) that +//! downstream code parses as structure. Kompress / LLMLingua sees them +//! as droppable noise and silently strips them, breaking everything +//! that depends on them. ContentRouter calls [`protect_tags`] before +//! every ML-text-compression to swap custom-tag spans for opaque +//! placeholders, runs the compressor on the cleaned body, then calls +//! [`restore_tags`] on the output to splice the originals back in. +//! +//! Standard HTML5 elements (`
`, `

`, ``, …) are *not* +//! protected — those go through the HTMLExtractor / trafilatura path +//! at a different layer. Anything else is treated as a custom tag. +//! +//! # Algorithm +//! +//! Single-pass tag-stack walker over the input bytes (no regex +//! backtracking, no O(n²) restart loop): +//! +//! 1. Scan forward for `<`. If the next bytes form a valid tag-open +//! (`` or ``), classify the tag name. +//! 2. HTML tag → emit verbatim, continue. +//! 3. Custom tag, self-closing → emit a placeholder, record the span. +//! 4. Custom tag, opening → push `(name, start_offset)` onto a stack. +//! 5. `` matching the top of the stack → pop, emit a placeholder +//! for the whole `` span (when +//! `compress_tagged_content == false`) or emit two placeholders for +//! the markers only (when `compress_tagged_content == true`). +//! 6. Mismatched close (HTML close while a custom tag is on top, or a +//! close with no matching open) → write the close tag verbatim and +//! move on. The walker never attempts to "repair" malformed input. +//! +//! Output is built incrementally with offset-based slicing — never the +//! Python-original's `result.replace(original, placeholder, 1)`, which +//! silently misbehaves when two identical custom-tag blocks appear in +//! the same input (it always replaces the *first* textual occurrence, +//! not the matched one). See `fixed_in_3e4_replace_first` test below. +//! +//! # Bug fixes vs the Python original +//! +//! * **#1: O(n²) on nested custom tags** — the Python loop restarted a +//! full regex scan after every replacement. Rust walks once, in +//! linear time on input length. +//! * **#2: First-occurrence replace bug** — `str.replace(.., .., 1)` +//! replaced the first textual match of the matched block, not the +//! block at the matched offset. Two identical custom-tag blocks in +//! the same input collapsed to one placeholder + a duplicated +//! second block. The Rust walker stitches output by offset. +//! * **#3: Silent 50-iteration cap** — Python had a hard 50-iteration +//! safety limit that quietly truncated tag protection on deeply +//! nested input. The Rust walker's run-time is bounded by input +//! length only. +//! * **#4: Self-closing pass duplicate-replace risk** — Python ran a +//! second loop with the same `replace_first` bug for self-closing +//! tags. Rust handles self-closers in the same single pass. +//! * **#5: Placeholder collision** — if input contains a literal +//! `{{HEADROOM_TAG_…}}` substring, Python silently let the collision +//! stand. We detect that and pick a salted prefix (with a tracing +//! warn) so restoration can't be ambiguous. +//! +//! # Hot path +//! +//! `protect_tags` runs on every ML-compression call from ContentRouter. +//! Most production prompts have 0–10 custom tags so the absolute cost +//! is small either way; the value of the port is correctness (bugs +//! #2, #5) and predictable behavior on adversarial input (bugs #1, +//! #3). The PyO3 bridge releases the GIL during the walk because the +//! algorithm is fully self-contained. + +use std::collections::HashSet; +use std::sync::OnceLock; + +/// HTML5 living-standard element names — the set of tags this module +/// will NEVER protect (they're handled at a different layer; everything +/// else is treated as custom). +/// +/// Generated from +/// and +/// matches the Python `KNOWN_HTML_TAGS` frozenset element-for-element +/// so the Rust shim and the Python shim agree. +const HTML5_TAGS: &[&str] = &[ + // Main root + "html", + // Document metadata + "base", + "head", + "link", + "meta", + "style", + "title", + // Sectioning root + "body", + // Content sectioning + "address", + "article", + "aside", + "footer", + "h1", + "h2", + "h3", + "h4", + "h5", + "h6", + "header", + "hgroup", + "main", + "nav", + "section", + "search", + // Text content + "blockquote", + "dd", + "div", + "dl", + "dt", + "figcaption", + "figure", + "hr", + "li", + "menu", + "ol", + "p", + "pre", + "ul", + // Inline text semantics + "a", + "abbr", + "b", + "bdi", + "bdo", + "br", + "cite", + "code", + "data", + "dfn", + "em", + "i", + "kbd", + "mark", + "q", + "rp", + "rt", + "ruby", + "s", + "samp", + "small", + "span", + "strong", + "sub", + "sup", + "time", + "u", + "var", + "wbr", + // Image and multimedia + "area", + "audio", + "img", + "map", + "track", + "video", + // Embedded content + "embed", + "iframe", + "object", + "param", + "picture", + "portal", + "source", + // SVG and MathML + "svg", + "math", + // Scripting + "canvas", + "noscript", + "script", + // Demarcating edits + "del", + "ins", + // Table content + "caption", + "col", + "colgroup", + "table", + "tbody", + "td", + "tfoot", + "th", + "thead", + "tr", + // Forms + "button", + "datalist", + "fieldset", + "form", + "input", + "label", + "legend", + "meter", + "optgroup", + "option", + "output", + "progress", + "select", + "textarea", + // Interactive + "details", + "dialog", + "summary", + // Web Components + "slot", + "template", +]; + +fn known_html_tags() -> &'static HashSet<&'static str> { + static SET: OnceLock> = OnceLock::new(); + SET.get_or_init(|| HTML5_TAGS.iter().copied().collect()) +} + +/// Default placeholder prefix. Brace-doubled to look unlike anything a +/// real workflow tag would emit. Falls back to a salted variant if the +/// input itself contains the prefix (see [`pick_placeholder_prefix`]). +const DEFAULT_PREFIX: &str = "{{HEADROOM_TAG_"; +const PLACEHOLDER_SUFFIX: &str = "}}"; + +/// Sidecar diagnostics — same shape every Rust transform uses. +#[derive(Debug, Default, Clone)] +pub struct ProtectStats { + pub tags_seen: usize, + pub html_tags_skipped: usize, + pub custom_blocks_protected: usize, + pub self_closing_protected: usize, + /// Closes that didn't match any open on the stack (malformed input + /// or HTML interleaving). Emitted verbatim. Non-zero is a smell + /// worth tracking but not necessarily a bug. + pub orphan_closes: usize, + /// True iff the placeholder prefix had to be salted because the + /// input contained a literal `{{HEADROOM_TAG_` substring. + pub placeholder_collision_avoided: bool, +} + +/// Case-insensitive HTML tag check. Lowercases the input lazily so we +/// don't allocate for the common ASCII-lowercase case. +pub fn is_known_html_tag(tag_name: &str) -> bool { + let set = known_html_tags(); + if set.contains(tag_name) { + return true; + } + if tag_name.bytes().any(|b| b.is_ascii_uppercase()) { + let lower = tag_name.to_ascii_lowercase(); + return set.contains(lower.as_str()); + } + false +} + +/// Iterate the canonical HTML tag list. Used by the PyO3 shim to +/// expose `KNOWN_HTML_TAGS` to Python without re-declaring the set. +pub fn known_html_tag_names() -> &'static [&'static str] { + HTML5_TAGS +} + +/// Pick a placeholder prefix that doesn't collide with anything in +/// `text`. We try `{{HEADROOM_TAG_` first; if the input contains it +/// literally we salt with a per-call counter until we miss. The salt +/// is bounded; in practice we never need more than one attempt. +fn pick_placeholder_prefix(text: &str) -> (String, bool) { + if !text.contains(DEFAULT_PREFIX) { + return (DEFAULT_PREFIX.to_string(), false); + } + for salt in 0u32..16 { + let candidate = format!("{{{{HEADROOM_TAG_{salt}_"); + if !text.contains(&candidate) { + return (candidate, true); + } + } + // 16 salt attempts collided — fall back to a UUID-shaped marker. + // The OnceLock cache is so two consecutive calls in the same + // process don't pay the formatting cost. + static FALLBACK: OnceLock = OnceLock::new(); + let prefix = FALLBACK + .get_or_init(|| "{{HEADROOM_TAG_FALLBACK_a4f1c7e2_".to_string()) + .clone(); + (prefix, true) +} + +#[derive(Debug)] +struct OpenTag { + /// Lowercase name for case-insensitive close-matching. + name_lower: String, + /// Byte offset of the `<` that opened this tag. + open_start: usize, +} + +/// Outcome of a single `<…>` parse attempt at a given offset. +enum TagParse { + /// Opening tag (``). `name_end` is exclusive. + Open { + name_end: usize, + tag_end: usize, + is_self_closing: bool, + }, + /// Closing tag (``). + Close { name_end: usize, tag_end: usize }, + /// Not a tag (e.g. `<` followed by whitespace or digit). + NotTag, +} + +/// Parse a `<…>` starting at `start`. Returns the byte offset of the +/// closing `>` (exclusive end of the tag) and the kind. Conservatively +/// rejects malformed shapes — we'd rather emit a `<` verbatim than +/// over-protect on bad input. +fn parse_tag_at(bytes: &[u8], start: usize) -> TagParse { + debug_assert!(bytes[start] == b'<'); + let mut i = start + 1; + let n = bytes.len(); + if i >= n { + return TagParse::NotTag; + } + + let is_close = bytes[i] == b'/'; + if is_close { + i += 1; + } + let name_start = i; + if !is_name_start(bytes[i]) { + return TagParse::NotTag; + } + i += 1; + while i < n && is_name_cont(bytes[i]) { + i += 1; + } + let name_end = i; + if name_end == name_start { + return TagParse::NotTag; + } + + if is_close { + // Allow optional whitespace, then `>`. + while i < n && bytes[i].is_ascii_whitespace() { + i += 1; + } + if i >= n || bytes[i] != b'>' { + return TagParse::NotTag; + } + return TagParse::Close { + name_end, + tag_end: i + 1, + }; + } + + // Opening tag: skip attributes until `>` (handle `/>` for + // self-closing). Quoted attribute values can contain `>`; a + // single-pass attribute lexer handles the common cases. + let mut self_closing = false; + while i < n { + match bytes[i] { + b'>' => { + return TagParse::Open { + name_end, + tag_end: i + 1, + is_self_closing: self_closing, + }; + } + b'/' => { + self_closing = true; + i += 1; + } + b'"' | b'\'' => { + let quote = bytes[i]; + i += 1; + while i < n && bytes[i] != quote { + i += 1; + } + if i >= n { + return TagParse::NotTag; + } + i += 1; + self_closing = false; + } + _ => { + if bytes[i].is_ascii_whitespace() { + self_closing = false; + } + i += 1; + } + } + } + + TagParse::NotTag +} + +#[inline] +fn is_name_start(b: u8) -> bool { + b.is_ascii_alphabetic() || b == b'_' +} + +#[inline] +fn is_name_cont(b: u8) -> bool { + b.is_ascii_alphanumeric() || matches!(b, b'_' | b'-' | b'.' | b':') +} + +/// A single span that was identified as worth replacing. +/// +/// In block mode every matched custom-tag span (open..=close) becomes +/// one Span and is replaced by a single placeholder; self-closing +/// custom tags become a Span covering just the tag bytes. +/// +/// In marker-only mode each opening custom tag and each closing custom +/// tag becomes its own Span (the body between them is left visible to +/// the compressor). +#[derive(Debug, Clone, Copy)] +struct Span { + start: usize, + end: usize, + kind: SpanKind, +} + +#[derive(Debug, Clone, Copy)] +enum SpanKind { + /// Whole `` block (block mode). + Block, + /// Self-closing `` (block mode). + SelfClosing, + /// Opening `` marker (marker-only mode). + OpenMarker, + /// Closing `` marker (marker-only mode). + CloseMarker, +} + +/// Protect custom workflow tags from text compression. +/// +/// * `compress_tagged_content = false` (default) — replace each entire +/// `` span (including nested children) with a +/// single placeholder. Self-closing custom tags become a single +/// placeholder. The body between the markers is *not* exposed to +/// compression. +/// * `compress_tagged_content = true` — replace only the tag markers +/// (open and close emitted as separate placeholders) so the +/// compressor can squash content while the tag boundaries survive. +/// +/// Returns `(cleaned, blocks, stats)` where `blocks` is a list of +/// `(placeholder, original)` pairs for [`restore_tags`]. The blocks +/// are listed in left-to-right order of appearance in the input, which +/// keeps the restore step trivial. +pub fn protect_tags( + text: &str, + compress_tagged_content: bool, +) -> (String, Vec<(String, String)>, ProtectStats) { + let mut stats = ProtectStats::default(); + if text.is_empty() || !text.contains('<') { + return (text.to_string(), Vec::new(), stats); + } + + let (prefix, salted) = pick_placeholder_prefix(text); + stats.placeholder_collision_avoided = salted; + + // Phase 1: walk once, classify every tag, build a list of spans + // worth replacing. No output emitted yet — this is purely + // discovery so we can decide which byte ranges to swap. + let spans = identify_spans(text, compress_tagged_content, &mut stats); + + // Phase 2: emit. Walk the input once more, splicing placeholders + // for span bytes and copying everything else verbatim. Because + // `spans` is sorted left-to-right and non-overlapping (block mode + // collapses nested matches into the outermost span; marker mode + // emits open/close markers that are byte-disjoint by construction) + // this is a straightforward scan. + match emit_output(text, &spans, &prefix) { + Some((cleaned, blocks)) => (cleaned, blocks, stats), + // Should be unreachable — `identify_spans` returns spans whose + // bytes are slices of `text`. If we ever fail to splice them + // back, fall back to emitting the original. + None => (text.to_string(), Vec::new(), stats), + } +} + +fn identify_spans( + text: &str, + compress_tagged_content: bool, + stats: &mut ProtectStats, +) -> Vec { + let bytes = text.as_bytes(); + let n = bytes.len(); + let mut spans: Vec = Vec::new(); + let mut stack: Vec = Vec::new(); + + let mut i = 0; + while i < n { + let b = bytes[i]; + if b != b'<' { + // Skip ahead to the next `<`. We don't care about non-tag + // bytes for span identification; they'll be copied verbatim + // in phase 2. + i = memchr(b'<', &bytes[i..]).map(|j| i + j).unwrap_or(n); + continue; + } + + match parse_tag_at(bytes, i) { + TagParse::NotTag => { + i += 1; + } + TagParse::Open { + name_end, + tag_end, + is_self_closing, + } => { + stats.tags_seen += 1; + let name = &text[i + 1..name_end]; + if is_known_html_tag(name) { + stats.html_tags_skipped += 1; + i = tag_end; + continue; + } + if is_self_closing { + spans.push(Span { + start: i, + end: tag_end, + kind: SpanKind::SelfClosing, + }); + stats.self_closing_protected += 1; + i = tag_end; + continue; + } + if compress_tagged_content { + // Marker-only mode: emit the open as its own span + // *and* push the name on the stack so the close + // gets matched and emitted as its own span. + spans.push(Span { + start: i, + end: tag_end, + kind: SpanKind::OpenMarker, + }); + } + // Both modes push to the stack so close-matching works. + stack.push(OpenTag { + name_lower: name.to_ascii_lowercase(), + open_start: i, + }); + i = tag_end; + } + TagParse::Close { name_end, tag_end } => { + stats.tags_seen += 1; + let close_name = &text[i + 2..name_end]; + if is_known_html_tag(close_name) { + stats.html_tags_skipped += 1; + i = tag_end; + continue; + } + let close_name_lower = close_name.to_ascii_lowercase(); + let matching = stack + .iter() + .rposition(|open| open.name_lower == close_name_lower); + + match matching { + Some(stack_idx) => { + if compress_tagged_content { + // Pop everything above (orphan opens + // inside the matched span — their open + // markers were already recorded as spans + // and we keep them). + stack.truncate(stack_idx); + let _ = stack.pop(); + spans.push(Span { + start: i, + end: tag_end, + kind: SpanKind::CloseMarker, + }); + } else { + // Block mode: collapse [open..close] into + // a single span. Drop any inner unmatched + // opens (they're part of this span's body). + // Also DROP any inner spans we already + // recorded that are now subsumed by this + // outer block — that's how nested custom + // tags collapse to one placeholder. + let open_start = stack[stack_idx].open_start; + stack.truncate(stack_idx); + spans.retain(|s| s.start < open_start); + spans.push(Span { + start: open_start, + end: tag_end, + kind: SpanKind::Block, + }); + stats.custom_blocks_protected += 1; + } + i = tag_end; + } + None => { + stats.orphan_closes += 1; + i = tag_end; + } + } + } + } + } + + // Stack remnants are orphan opens (no matching close ever arrived). + // We don't protect those — they'll fall through to the compressor + // as raw `` bytes, same as Python's original behavior. In + // block mode their inner self-closing spans we recorded are still + // safe to keep: they were below an unmatched outer open, so they + // were never collapsed. Spans are sorted by start ascending due to + // the monotonic walk; phase 2 expects that. + spans +} + +fn emit_output( + text: &str, + spans: &[Span], + prefix: &str, +) -> Option<(String, Vec<(String, String)>)> { + let mut out = String::with_capacity(text.len()); + let mut blocks: Vec<(String, String)> = Vec::new(); + let mut cursor: usize = 0; + + for (counter, span) in (0_u64..).zip(spans.iter()) { + if span.start < cursor { + // Overlap shouldn't happen given how we collapse nested + // spans, but bail loudly if it does — silently producing + // wrong output is worse than failing the test. + return None; + } + out.push_str(&text[cursor..span.start]); + let placeholder = format!("{prefix}{counter}{PLACEHOLDER_SUFFIX}"); + let original = &text[span.start..span.end]; + blocks.push((placeholder.clone(), original.to_string())); + out.push_str(&placeholder); + cursor = span.end; + let _ = span.kind; // SpanKind is informational only at this layer + } + out.push_str(&text[cursor..]); + Some((out, blocks)) +} + +/// Restore protected tag spans after the compressor ran on the +/// cleaned text. +/// +/// If a placeholder went missing during compression (the compressor +/// stripped or rewrote it) the corresponding original is appended to +/// the output rather than dropped — same fallback semantics as the +/// Python original. A `tracing::warn!` is emitted at compile-time- +/// optional verbosity so prod can scrape lossy-compression incidents. +pub fn restore_tags(text: &str, blocks: &[(String, String)]) -> String { + if blocks.is_empty() { + return text.to_string(); + } + + let mut result = text.to_string(); + let mut tail_appends: Vec<&str> = Vec::new(); + for (placeholder, original) in blocks { + if result.contains(placeholder.as_str()) { + result = result.replace(placeholder.as_str(), original); + } else { + tag_lost_warn(original); + tail_appends.push(original.as_str()); + } + } + + if !tail_appends.is_empty() { + for original in tail_appends { + result.push('\n'); + result.push_str(original); + } + } + + result +} + +#[inline(never)] +fn tag_lost_warn(original: &str) { + let preview: String = original.chars().take(80).collect(); + tracing::warn!( + target: "headroom::tag_protector", + preview = %preview, + "tag placeholder lost during compression, appending original" + ); +} + +// ─── Tiny byte-search helper ────────────────────────────────────────── + +#[inline] +fn memchr(needle: u8, haystack: &[u8]) -> Option { + haystack.iter().position(|&b| b == needle) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn protect(text: &str) -> (String, Vec<(String, String)>) { + let (cleaned, blocks, _stats) = protect_tags(text, false); + (cleaned, blocks) + } + + #[test] + fn passthrough_when_no_angle_bracket() { + let (cleaned, blocks) = protect("Just plain text"); + assert_eq!(cleaned, "Just plain text"); + assert!(blocks.is_empty()); + } + + #[test] + fn html_tags_emitted_verbatim() { + let text = "

Some content
"; + let (cleaned, blocks) = protect(text); + assert_eq!(cleaned, text); + assert!(blocks.is_empty()); + } + + #[test] + fn html_tag_check_case_insensitive() { + assert!(is_known_html_tag("DIV")); + assert!(is_known_html_tag("Span")); + assert!(!is_known_html_tag("system-reminder")); + assert!(!is_known_html_tag("EXTREMELY_IMPORTANT")); + } + + #[test] + fn custom_tag_replaced_with_placeholder() { + let text = "Before Important After"; + let (cleaned, blocks) = protect(text); + assert!(!cleaned.contains("")); + assert!(!cleaned.contains("Important")); + assert!(cleaned.contains("Before")); + assert!(cleaned.contains("After")); + assert_eq!(blocks.len(), 1); + assert_eq!(blocks[0].1, "Important"); + } + + #[test] + fn custom_tag_with_attributes() { + let text = r#"user data"#; + let (_cleaned, blocks) = protect(text); + assert_eq!(blocks.len(), 1); + assert!(blocks[0].1.contains(r#"key="session""#)); + } + + #[test] + fn self_closing_custom_tag() { + let text = "Text more text"; + let (_cleaned, blocks) = protect(text); + assert_eq!(blocks.len(), 1); + assert_eq!(blocks[0].1, ""); + } + + #[test] + fn self_closing_html_tag_not_protected() { + let text = "Text
more
text"; + let (cleaned, blocks) = protect(text); + assert_eq!(cleaned, text); + assert!(blocks.is_empty()); + } + + #[test] + fn nested_custom_tags_collapse_to_outer_span() { + let text = "deep"; + let (cleaned, blocks) = protect(text); + assert!(!cleaned.contains("")); + assert!(!cleaned.contains("")); + // Outer span captures inner — single placeholder. + assert_eq!(blocks.len(), 1); + assert_eq!(blocks[0].1, "deep"); + } + + #[test] + fn mixed_html_and_custom() { + let text = "
HTML
Rule

HTML2

"; + let (cleaned, blocks) = protect(text); + assert!(cleaned.contains("
")); + assert!(cleaned.contains("

")); + assert!(!cleaned.contains("")); + assert_eq!(blocks.len(), 1); + } + + #[test] + fn real_workflow_tags() { + let cases = [ + "search({query: 'test'})", + "Let me analyze this", + "Never skip validation", + "check perms", + "Rules", + "Success: 42 items", + ]; + for tag in cases { + let text = format!("Before {tag} After"); + let (_cleaned, blocks) = protect(&text); + assert_eq!(blocks.len(), 1, "failed to protect: {tag}"); + assert_eq!(blocks[0].1, tag); + } + } + + #[test] + fn empty_input_returns_empty() { + let (cleaned, blocks) = protect(""); + assert!(cleaned.is_empty()); + assert!(blocks.is_empty()); + } + + #[test] + fn compress_tagged_content_true_emits_marker_placeholders() { + let text = "Before Compressible content After"; + let (cleaned, blocks, _stats) = protect_tags(text, true); + assert!(!cleaned.contains("")); + assert!(!cleaned.contains("")); + assert!(cleaned.contains("Compressible content")); + assert_eq!(blocks.len(), 2); + } + + #[test] + fn restore_basic() { + let original = "Before Rule After"; + let (cleaned, blocks, _stats) = protect_tags(original, false); + let restored = restore_tags(&cleaned, &blocks); + assert_eq!(restored, original); + } + + #[test] + fn restore_empty_blocks_passthrough() { + assert_eq!(restore_tags("untouched", &[]), "untouched"); + } + + #[test] + fn restore_lost_placeholder_appended() { + let blocks = vec![( + "{{HEADROOM_TAG_0}}".to_string(), + "data".to_string(), + )]; + let restored = restore_tags("text without placeholder", &blocks); + assert!(restored.contains("data")); + } + + #[test] + fn restore_roundtrip_preserves_content() { + let original = "Start Rule 1: validate middle \ + search(q='test') end"; + let (cleaned, blocks, _stats) = protect_tags(original, false); + let restored = restore_tags(&cleaned, &blocks); + assert_eq!(restored, original); + } + + // ─── Bug-fix tests (fixed_in_3e4) ───────────────────────────────── + + #[test] + fn fixed_in_3e4_replace_first_does_not_collide_on_duplicate_blocks() { + // Bug #2: Python's `result.replace(original, placeholder, 1)` + // replaces the FIRST textual occurrence of `original`, not + // necessarily the matched offset. Two identical custom-tag + // blocks would collapse to a single placeholder + a stray + // duplicate of the second block in the output. + let text = "same middle \ + same"; + let (cleaned, blocks, _stats) = protect_tags(text, false); + // BOTH blocks should be replaced by DIFFERENT placeholders. + assert_eq!(blocks.len(), 2); + assert!(!cleaned.contains("")); + assert!(!cleaned.contains("")); + assert_ne!(blocks[0].0, blocks[1].0); + // Roundtrip is exact. + assert_eq!(restore_tags(&cleaned, &blocks), text); + } + + #[test] + fn fixed_in_3e4_handles_50_plus_nested_custom_tags() { + // Bug #3: Python had a hard-coded 50-iteration safety cap that + // silently truncated tag protection on deeply nested input. + // Build 60 nested custom tags and verify all get caught in + // the outermost span. + let depth = 60; + let mut text = String::new(); + for _ in 0..depth { + text.push_str(""); + } + text.push_str("core"); + for _ in 0..depth { + text.push_str(""); + } + let (cleaned, blocks, _stats) = protect_tags(&text, false); + // The outermost span eats everything: one placeholder, no + // residual `` markers in the cleaned text. + assert!(!cleaned.contains("")); + assert!(!cleaned.contains("")); + assert_eq!(blocks.len(), 1); + // Roundtrip exact even at depth=60. + assert_eq!(restore_tags(&cleaned, &blocks), text); + } + + #[test] + fn fixed_in_3e4_self_closing_duplicates_get_distinct_placeholders() { + // Bug #4: same first-occurrence-replace bug for self-closing + // tags. `` appearing twice would collapse. + let text = " middle "; + let (cleaned, blocks, _stats) = protect_tags(text, false); + assert_eq!(blocks.len(), 2); + assert_ne!(blocks[0].0, blocks[1].0); + assert!(!cleaned.contains("")); + assert_eq!(restore_tags(&cleaned, &blocks), text); + } + + #[test] + fn fixed_in_3e4_placeholder_collision_is_avoided() { + // Bug #5: input contains literal `{{HEADROOM_TAG_…}}`. The + // walker should pick a salted prefix and report the collision + // in stats. + let text = "User wrote {{HEADROOM_TAG_0}} on purpose. \ + real one"; + let (_cleaned, blocks, stats) = protect_tags(text, false); + assert!(stats.placeholder_collision_avoided); + assert_eq!(blocks.len(), 1); + // Placeholder used must NOT collide with the user's literal. + assert_ne!(blocks[0].0, "{{HEADROOM_TAG_0}}"); + } + + // ─── Edge-case correctness ──────────────────────────────────────── + + #[test] + fn orphan_close_tag_emitted_verbatim() { + let text = "no opener here"; + let (cleaned, blocks, stats) = protect_tags(text, false); + // Nothing protected; close stays in the cleaned text. + assert_eq!(blocks.len(), 0); + assert!(cleaned.contains("")); + assert_eq!(stats.orphan_closes, 1); + } + + #[test] + fn orphan_open_tag_emitted_verbatim() { + // An open with no matching close should round-trip exactly — + // no protection, no data loss. + let text = "dangling content with no close"; + let (cleaned, blocks, _stats) = protect_tags(text, false); + assert!(blocks.is_empty()); + assert_eq!(cleaned, text); + } + + #[test] + fn malformed_lone_lt_emitted_verbatim() { + let text = "if a < b then c"; + let (cleaned, blocks, _stats) = protect_tags(text, false); + assert_eq!(cleaned, text); + assert!(blocks.is_empty()); + } + + #[test] + fn attribute_with_gt_inside_quotes() { + let text = r#"payload"#; + let (cleaned, blocks, _stats) = protect_tags(text, false); + assert_eq!(blocks.len(), 1); + assert_eq!(blocks[0].1, text); + assert!(!cleaned.contains("payload")); + } + + #[test] + fn html_close_inside_custom_block_does_not_pop_stack() { + // An HTML close tag while a custom open is on top should not + // confuse the stack: the HTML close is emitted verbatim, the + // custom span still closes when its own close arrives. + let text = "x

y"; + let (cleaned, blocks, stats) = protect_tags(text, false); + // The whole `...` span wins, including the + // verbatim `
` inside. + assert_eq!(blocks.len(), 1); + assert_eq!(blocks[0].1, "x y"); + assert!(!cleaned.contains("")); + // `` is HTML, not orphan. + assert_eq!(stats.html_tags_skipped, 1); + assert_eq!(stats.orphan_closes, 0); + } +} diff --git a/crates/headroom-py/src/lib.rs b/crates/headroom-py/src/lib.rs index 05f88de1a..249d91d43 100644 --- a/crates/headroom-py/src/lib.rs +++ b/crates/headroom-py/src/lib.rs @@ -23,6 +23,10 @@ use headroom_core::transforms::smart_crusher::{ CrushResult as RustCrushResult, SmartCrusher as RustSmartCrusher, SmartCrusherConfig as RustSmartCrusherConfig, }; +use headroom_core::transforms::tag_protector::{ + is_known_html_tag as rust_is_known_html_tag, known_html_tag_names as rust_known_html_tag_names, + protect_tags as rust_protect_tags, restore_tags as rust_restore_tags, +}; use headroom_core::transforms::{ detect as rust_detect_chain, is_json_array_of_dicts as rust_is_json_array_of_dicts, ContentType as RustContentType, DetectionResult as RustDetectionResult, DiffCompressionResult, @@ -1402,6 +1406,57 @@ const _: fn() = || { let _ = RustLogLevel::Unknown; }; +// ─── tag_protector bridge (Phase 3e.4) ─────────────────────────────────── +// +// Mirrors `headroom.transforms.tag_protector.{protect_tags,restore_tags, +// is_html_tag,KNOWN_HTML_TAGS}`. The Rust walker is single-pass and +// fixes five real bugs the Python original carried (see crate-level +// docs in `tag_protector.rs`). The GIL is released during the walk +// because the algorithm holds no Python references. + +/// Replace custom workflow tags in `text` with opaque placeholders so +/// downstream ML compressors can't accidentally drop them. +/// +/// Returns `(cleaned_text, blocks)` where `blocks` is a list of +/// `(placeholder, original)` tuples for `restore_tags`. +#[pyfunction] +#[pyo3(signature = (text, compress_tagged_content = false))] +fn protect_tags( + py: Python<'_>, + text: &str, + compress_tagged_content: bool, +) -> (String, Vec<(String, String)>) { + let owned = text.to_string(); + py.allow_threads(move || { + let (cleaned, blocks, _stats) = rust_protect_tags(&owned, compress_tagged_content); + (cleaned, blocks) + }) +} + +/// Splice protected blocks back into `text`. Missing placeholders fall +/// back to appending the original block (lossy-compression incident). +#[pyfunction] +fn restore_tags(py: Python<'_>, text: &str, blocks: Vec<(String, String)>) -> String { + let owned = text.to_string(); + py.allow_threads(move || rust_restore_tags(&owned, &blocks)) +} + +/// Case-insensitive HTML5 tag check. The Python shim uses this to +/// preserve the legacy private `_is_html_tag` import surface for tests. +#[pyfunction] +fn is_html_tag(name: &str) -> bool { + rust_is_known_html_tag(name) +} + +/// Return the canonical HTML5 tag name list. The Python shim +/// reconstructs `KNOWN_HTML_TAGS` from this so callers that import the +/// frozenset (the existing test does) continue to work without +/// re-declaring the set in two languages. +#[pyfunction] +fn known_html_tag_names() -> Vec<&'static str> { + rust_known_html_tag_names().to_vec() +} + // ─── Module init ─────────────────────────────────────────────────────────── #[pymodule] @@ -1423,6 +1478,10 @@ fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_function(wrap_pyfunction!(detect_log_format, m)?)?; + m.add_function(wrap_pyfunction!(protect_tags, m)?)?; + m.add_function(wrap_pyfunction!(restore_tags, m)?)?; + m.add_function(wrap_pyfunction!(is_html_tag, m)?)?; + m.add_function(wrap_pyfunction!(known_html_tag_names, m)?)?; m.add_function(wrap_pyfunction!(detect_content_type, m)?)?; m.add_function(wrap_pyfunction!(is_json_array_of_dicts, m)?)?; m.add_function(wrap_pyfunction!(score_line, m)?)?; diff --git a/headroom/transforms/tag_protector.py b/headroom/transforms/tag_protector.py index 00069ea8b..9b3a4e2fb 100644 --- a/headroom/transforms/tag_protector.py +++ b/headroom/transforms/tag_protector.py @@ -1,185 +1,73 @@ -"""Protect workflow/custom XML tags from text compression. +"""Rust-backed tag protector — keep custom XML tags away from compressors. -LLM workflows use XML-style tags (, , ) -as structural markers. Text compressors (Kompress) treat these as -droppable noise and silently remove them, breaking downstream tools. +Phase 3e.4 ported the implementation to +``crates/headroom-core/src/transforms/tag_protector.rs``. This module +is now a thin shim that: -This module detects custom tags (anything NOT standard HTML), replaces entire -blocks with placeholders before compression, and restores them after. +1. Routes ``protect_tags`` / ``restore_tags`` through PyO3 so callers + pick up the single-pass Rust walker (and the five bug fixes that + ride along — see the crate-level docs in the Rust source). +2. Re-exports the legacy import surface (``KNOWN_HTML_TAGS``, + ``_is_html_tag``) so existing callers in ``content_router.py`` and + the existing test suite keep working without same-PR refactors. -Standard HTML tags (
,

, ) are left alone — HTMLExtractor -handles those via trafilatura. +# Bug fixes the Rust port carries (and this shim therefore inherits) -Usage: - from headroom.transforms.tag_protector import protect_tags, restore_tags - - cleaned, protected = protect_tags(text) - compressed = kompress.compress(cleaned) - result = restore_tags(compressed, protected) +* **#1: O(n²) on nested custom tags.** Python iterated a regex + scan-and-replace loop until stable, restarting from the top after + every replacement. Rust walks once, in linear time on input length. +* **#2: First-occurrence replace bug.** ``str.replace(orig, ph, 1)`` + replaces the FIRST textual match, not the matched offset. Two + identical custom-tag blocks collapsed to one placeholder + a stray + duplicate of the second block. The Rust walker stitches output by + offset — distinct blocks always get distinct placeholders. +* **#3: Silent 50-iteration cap.** Python had a hard ``max_iterations + = 50`` safety limit that quietly truncated tag protection on deeply + nested input. The Rust walker is bounded by input length only. +* **#4: Self-closing pass duplicate-replace risk.** Python ran a + second loop with the same first-occurrence-replace bug for + self-closers. Rust handles them in the same single pass. +* **#5: Placeholder collision.** If the input contained a literal + ``{{HEADROOM_TAG_…}}`` substring, Python silently let the collision + stand. Rust detects it and salts the prefix. """ from __future__ import annotations import logging -import re +from typing import cast + +from headroom._core import ( + is_html_tag as _rust_is_html_tag, +) +from headroom._core import ( + known_html_tag_names as _rust_known_html_tag_names, +) +from headroom._core import ( + protect_tags as _rust_protect_tags, +) +from headroom._core import ( + restore_tags as _rust_restore_tags, +) logger = logging.getLogger(__name__) -# --------------------------------------------------------------------------- -# HTML5 Living Standard element names (https://html.spec.whatwg.org/) -# These are the ONLY tags treated as "HTML". Everything else is a custom -# workflow tag and gets protected from compression. -# --------------------------------------------------------------------------- -KNOWN_HTML_TAGS: frozenset[str] = frozenset( - { - # Main root - "html", - # Document metadata - "base", - "head", - "link", - "meta", - "style", - "title", - # Sectioning root - "body", - # Content sectioning - "address", - "article", - "aside", - "footer", - "h1", - "h2", - "h3", - "h4", - "h5", - "h6", - "header", - "hgroup", - "main", - "nav", - "section", - "search", - # Text content - "blockquote", - "dd", - "div", - "dl", - "dt", - "figcaption", - "figure", - "hr", - "li", - "menu", - "ol", - "p", - "pre", - "ul", - # Inline text semantics - "a", - "abbr", - "b", - "bdi", - "bdo", - "br", - "cite", - "code", - "data", - "dfn", - "em", - "i", - "kbd", - "mark", - "q", - "rp", - "rt", - "ruby", - "s", - "samp", - "small", - "span", - "strong", - "sub", - "sup", - "time", - "u", - "var", - "wbr", - # Image and multimedia - "area", - "audio", - "img", - "map", - "track", - "video", - # Embedded content - "embed", - "iframe", - "object", - "param", - "picture", - "portal", - "source", - # SVG and MathML - "svg", - "math", - # Scripting - "canvas", - "noscript", - "script", - # Demarcating edits - "del", - "ins", - # Table content - "caption", - "col", - "colgroup", - "table", - "tbody", - "td", - "tfoot", - "th", - "thead", - "tr", - # Forms - "button", - "datalist", - "fieldset", - "form", - "input", - "label", - "legend", - "meter", - "optgroup", - "option", - "output", - "progress", - "select", - "textarea", - # Interactive - "details", - "dialog", - "summary", - # Web Components - "slot", - "template", - } -) -# Placeholder format — unlikely to appear in real content -_PLACEHOLDER_PREFIX = "{{HEADROOM_TAG_" -_PLACEHOLDER_SUFFIX = "}}" - -# Match opening tags: , , and self-closing -_OPEN_TAG_RE = re.compile(r"<([a-zA-Z_][\w.-]*)((?:\s+[^>]*?)?)(/?)>") - -# Match closing tags: -_CLOSE_TAG_RE = re.compile(r"") +# Pulled from Rust at import time so the canonical list lives in one +# place. The frozenset shape is the legacy public surface — tests and +# the integration test ask for membership / size on this object. +KNOWN_HTML_TAGS: frozenset[str] = frozenset(_rust_known_html_tag_names()) def _is_html_tag(tag_name: str) -> bool: - """Check if a tag name is a known HTML element (case-insensitive).""" - return tag_name.lower() in KNOWN_HTML_TAGS + """Case-insensitive HTML5 tag check. + + Kept as a private name (with the underscore) because the Python + test file imports ``_is_html_tag`` directly. Delegates to the Rust + implementation so the two languages can't drift on what counts as + "HTML". + """ + return bool(_rust_is_html_tag(tag_name)) def protect_tags( @@ -188,100 +76,22 @@ def protect_tags( ) -> tuple[str, list[tuple[str, str]]]: """Protect custom/workflow XML tags from text compression. - Scans for XML-style tags, classifies each as HTML or custom. - Custom tags (and optionally their content) are replaced with - placeholders that survive compression. - Args: text: Input text potentially containing XML tags. compress_tagged_content: If False (default), protect entire - ``content`` block verbatim. If True, only - protect the tag markers; content between them can be - compressed. + ``content`` block verbatim. If True, only + protect the tag markers; content between them is exposed + for compression. Returns: - Tuple of (cleaned_text, protected_blocks) where protected_blocks - is a list of (placeholder, original_text) pairs for restoration. + Tuple of ``(cleaned_text, protected_blocks)`` where each + protected block is a ``(placeholder, original)`` pair. Hand + the full block list to :func:`restore_tags` after the + compressor has run. """ - if not text or "<" not in text: - return text, [] - - protected: list[tuple[str, str]] = [] - counter = 0 - - if not compress_tagged_content: - # Protect entire ... blocks for custom tags. - # Process from innermost out by iterating until stable. - result = text - changed = True - max_iterations = 50 # Safety limit for deeply nested tags - - while changed and max_iterations > 0: - changed = False - max_iterations -= 1 - - # Find blocks: ... where tag is NOT known HTML - # Use non-greedy match for content between tags - pattern = re.compile( - r"<([a-zA-Z_][\w.-]*)((?:\s+[^>]*?)?)>" - r"(.*?)" - r"", - re.DOTALL, - ) - - for match in pattern.finditer(result): - tag_name = match.group(1) - if _is_html_tag(tag_name): - continue - - original_block = match.group(0) - placeholder = f"{_PLACEHOLDER_PREFIX}{counter}{_PLACEHOLDER_SUFFIX}" - protected.append((placeholder, original_block)) - result = result.replace(original_block, placeholder, 1) - counter += 1 - changed = True - break # Restart iteration after each replacement - - # Also protect self-closing custom tags: - for match in _OPEN_TAG_RE.finditer(result): - tag_name = match.group(1) - is_self_closing = match.group(3) == "/" - if is_self_closing and not _is_html_tag(tag_name): - original = match.group(0) - if original in result: # Not already replaced - placeholder = f"{_PLACEHOLDER_PREFIX}{counter}{_PLACEHOLDER_SUFFIX}" - protected.append((placeholder, original)) - result = result.replace(original, placeholder, 1) - counter += 1 - - return result, protected - - else: - # compress_tagged_content=True: only protect tag markers, not content - result = text - - # Protect opening custom tags - for match in _OPEN_TAG_RE.finditer(text): - tag_name = match.group(1) - if not _is_html_tag(tag_name): - original = match.group(0) - placeholder = f"{_PLACEHOLDER_PREFIX}{counter}{_PLACEHOLDER_SUFFIX}" - protected.append((placeholder, original)) - result = result.replace(original, placeholder, 1) - counter += 1 - - # Protect closing custom tags - for match in _CLOSE_TAG_RE.finditer(text): - tag_name = match.group(1) - if not _is_html_tag(tag_name): - original = match.group(0) - if original in result: - placeholder = f"{_PLACEHOLDER_PREFIX}{counter}{_PLACEHOLDER_SUFFIX}" - protected.append((placeholder, original)) - result = result.replace(original, placeholder, 1) - counter += 1 - - return result, protected + cleaned, blocks = _rust_protect_tags(text, compress_tagged_content) + # The Rust binding hands us a list of plain Python tuples already. + return cast("str", cleaned), cast("list[tuple[str, str]]", blocks) def restore_tags( @@ -292,25 +102,21 @@ def restore_tags( Args: text: Compressed text with placeholders. - protected_blocks: List of (placeholder, original_text) pairs - from ``protect_tags()``. + protected_blocks: List from :func:`protect_tags`. Returns: - Text with all placeholders replaced by original blocks. + Text with placeholders swapped back to originals. Any + placeholder the compressor accidentally stripped causes the + corresponding original to be appended on the trailing edge of + the output (lossy-compression incident; same fallback the + Python original used). """ - if not protected_blocks: - return text + return cast("str", _rust_restore_tags(text, protected_blocks)) - result = text - for placeholder, original in protected_blocks: - if placeholder in result: - result = result.replace(placeholder, original) - else: - # Placeholder was lost during compression — append the block - logger.warning( - "Tag placeholder lost during compression, appending: %s", - original[:80], - ) - result = result + "\n" + original - return result +__all__ = [ + "KNOWN_HTML_TAGS", + "_is_html_tag", + "protect_tags", + "restore_tags", +] diff --git a/tests/test_transforms/test_tag_protector.py b/tests/test_transforms/test_tag_protector.py index 9b0a1de6b..cd1bf0f82 100644 --- a/tests/test_transforms/test_tag_protector.py +++ b/tests/test_transforms/test_tag_protector.py @@ -181,3 +181,65 @@ class TestRestoreTags: restored = restore_tags(cleaned, protected) assert "Rule 1: always validate" in restored assert "search(q='test')" in restored + + +class TestBugFixesPhase3e4: + """Bug fixes baked into the Phase 3e.4 Rust port. Each test pins + behavior the Python regex implementation got wrong.""" + + def test_fixed_in_3e4_duplicate_blocks_get_distinct_placeholders(self): + """Bug #2: Python's `result.replace(orig, ph, 1)` replaces the + FIRST textual match of `orig`, not the matched offset. Two + identical custom-tag blocks in the same input collapsed to a + single placeholder + a stray duplicate of the second block. + The Rust walker emits offset-based output, so distinct blocks + always get distinct placeholders.""" + text = ( + "same middle same" + ) + cleaned, protected = protect_tags(text) + assert len(protected) == 2 + placeholders = {p[0] for p in protected} + assert len(placeholders) == 2 # two DIFFERENT placeholders + assert "" not in cleaned + # Roundtrip is exact byte-for-byte. + assert restore_tags(cleaned, protected) == text + + def test_fixed_in_3e4_handles_60_nested_custom_tags(self): + """Bug #3: Python had a hard `max_iterations = 50` safety cap + that quietly stopped protecting deeper nested input. The Rust + walker is bounded by input length only.""" + depth = 60 + text = "" * depth + "core" + "" * depth + cleaned, protected = protect_tags(text) + # Outermost span eats everything → ONE placeholder, no leaks. + assert "" not in cleaned + assert "" not in cleaned + assert len(protected) == 1 + assert restore_tags(cleaned, protected) == text + + def test_fixed_in_3e4_self_closing_duplicates_distinct(self): + """Bug #4: same first-occurrence-replace bug for self-closing + tags. Two identical `` would collapse to one + placeholder + a stray dup.""" + text = " middle " + cleaned, protected = protect_tags(text) + assert len(protected) == 2 + assert protected[0][0] != protected[1][0] + assert "" not in cleaned + assert restore_tags(cleaned, protected) == text + + def test_fixed_in_3e4_placeholder_collision_avoided(self): + """Bug #5: input contains a literal `{{HEADROOM_TAG_…}}` + substring. Python silently used the same prefix and let the + collision break restoration. Rust salts the prefix when this + happens.""" + text = ( + "User wrote {{HEADROOM_TAG_0}} on purpose. real one" + ) + cleaned, protected = protect_tags(text) + assert len(protected) == 1 + # Placeholder picked must NOT collide with the user's literal. + assert protected[0][0] != "{{HEADROOM_TAG_0}}" + # Roundtrip is exact (the user's literal stays intact). + assert restore_tags(cleaned, protected) == text