mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(content-router): protect custom-tag blocks before mixed-content section split
Protect custom-tag blocks during mixed-content routing.
This commit is contained in:
parent
e4904e23a6
commit
d7bc1e275f
2 changed files with 112 additions and 2 deletions
|
|
@ -2406,7 +2406,24 @@ class ContentRouter(Transform):
|
|||
Returns:
|
||||
RouterCompressionResult with reassembled content.
|
||||
"""
|
||||
sections = split_into_sections(content)
|
||||
from .tag_protector import protect_tags, restore_tags
|
||||
|
||||
# Protect custom-tag blocks BEFORE splitting into sections. Section
|
||||
# boundaries (code fences, blank lines) split a
|
||||
# ``<system-reminder>...</system-reminder>`` pair across sections, so
|
||||
# the per-section tag protection inside ``_try_ml_compressor`` never
|
||||
# sees a matched pair (an unmatched tag protects nothing) and
|
||||
# instruction blocks — Claude Code ships CLAUDE.md inside
|
||||
# <system-reminder> — leak into lossy ML compression and arrive
|
||||
# word-dropped. Protecting here keeps the whole block as one
|
||||
# placeholder that spans sections intact.
|
||||
cleaned, protected = protect_tags(
|
||||
content,
|
||||
compress_tagged_content=self.config.compress_tagged_content,
|
||||
)
|
||||
sections_source = cleaned if protected else content
|
||||
|
||||
sections = split_into_sections(sections_source)
|
||||
if logger.isEnabledFor(logging.DEBUG):
|
||||
_log_router_debug(
|
||||
"content_router_mixed_sections",
|
||||
|
|
@ -2422,10 +2439,32 @@ class ContentRouter(Transform):
|
|||
strategy_used=CompressionStrategy.PASSTHROUGH,
|
||||
)
|
||||
|
||||
# Placeholders must survive byte-exact: ``restore_tags`` DISCARDS a
|
||||
# protected block whose placeholder was stripped or rewritten
|
||||
# (Hotfix-A9), so a compressor eating a placeholder would silently
|
||||
# drop the whole tag block — worse than the mangling this fixes.
|
||||
# Any section carrying a placeholder is passed through verbatim
|
||||
# instead of ever entering a compressor.
|
||||
placeholders = [placeholder for placeholder, _ in protected]
|
||||
|
||||
compressed_sections: list[str] = []
|
||||
routing_log: list[RoutingDecision] = []
|
||||
|
||||
for i, section in enumerate(sections):
|
||||
if placeholders and any(ph in section.content for ph in placeholders):
|
||||
section_tokens = _estimate_tokens(section.content)
|
||||
compressed_sections.append(section.content)
|
||||
routing_log.append(
|
||||
RoutingDecision(
|
||||
content_type=section.content_type,
|
||||
strategy=CompressionStrategy.PASSTHROUGH,
|
||||
original_tokens=section_tokens,
|
||||
compressed_tokens=section_tokens,
|
||||
section_index=i,
|
||||
)
|
||||
)
|
||||
continue
|
||||
|
||||
# Get strategy for this section
|
||||
strategy = self._strategy_from_detection_type(section.content_type)
|
||||
|
||||
|
|
@ -2455,8 +2494,12 @@ class ContentRouter(Transform):
|
|||
)
|
||||
)
|
||||
|
||||
compressed = "\n\n".join(compressed_sections)
|
||||
if protected:
|
||||
compressed = restore_tags(compressed, protected)
|
||||
|
||||
return RouterCompressionResult(
|
||||
compressed="\n\n".join(compressed_sections),
|
||||
compressed=compressed,
|
||||
original=content,
|
||||
strategy_used=CompressionStrategy.MIXED,
|
||||
routing_log=routing_log,
|
||||
|
|
|
|||
|
|
@ -1699,3 +1699,70 @@ class TestCompressBlockContent:
|
|||
assert any("router:tool_result" in t for t in transforms_applied), (
|
||||
f"Expected router:tool_result:* in transforms, got: {transforms_applied}"
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Mixed content: custom-tag protection (system-reminder mangling regression)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestMixedContentTagProtection:
|
||||
"""_compress_mixed must protect custom-tag blocks BEFORE section split.
|
||||
|
||||
Splitting first lands the open/close tags of a
|
||||
``<system-reminder>...</system-reminder>`` pair in different sections;
|
||||
per-section protection then sees only unmatched tags (which protect
|
||||
nothing) and the block's content — Claude Code ships CLAUDE.md this way —
|
||||
is lossy-compressed and arrives word-dropped.
|
||||
"""
|
||||
|
||||
REMINDER = (
|
||||
"<system-reminder>\n"
|
||||
"Instruction prose that must survive byte-exact.\n\n"
|
||||
"```bash\nrtk gain\n```\n\n"
|
||||
"More instructions after the fence, also byte-exact.\n"
|
||||
"</system-reminder>"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _mangling_router() -> ContentRouter:
|
||||
"""Router whose per-section compressor visibly mangles everything."""
|
||||
router = ContentRouter(ContentRouterConfig(min_section_tokens=1))
|
||||
|
||||
def mangle(content, strategy, context, language=None, question=None, bias=1.0):
|
||||
return "MANGLED", 1, None
|
||||
|
||||
router._apply_strategy_to_content = mangle # type: ignore[method-assign]
|
||||
return router
|
||||
|
||||
def test_reminder_block_survives_mixed_compression_verbatim(self):
|
||||
router = self._mangling_router()
|
||||
content = (
|
||||
"Prose before the reminder that may compress.\n\n"
|
||||
+ self.REMINDER
|
||||
+ "\n\nProse after the reminder that may compress."
|
||||
)
|
||||
|
||||
result = router._compress_mixed(content, context="")
|
||||
|
||||
# The tag block (fence and all) is byte-exact in the output...
|
||||
assert self.REMINDER in result.compressed
|
||||
# ...while content outside it still went through the compressor.
|
||||
assert "MANGLED" in result.compressed
|
||||
|
||||
def test_reminder_only_content_passes_through(self):
|
||||
router = self._mangling_router()
|
||||
|
||||
result = router._compress_mixed(self.REMINDER, context="")
|
||||
|
||||
assert self.REMINDER in result.compressed
|
||||
assert "MANGLED" not in result.compressed
|
||||
|
||||
def test_untagged_mixed_content_still_compresses(self):
|
||||
router = self._mangling_router()
|
||||
content = "Plain prose section.\n\n```python\nprint('hi')\n```\n\nMore prose."
|
||||
|
||||
result = router._compress_mixed(content, context="")
|
||||
|
||||
assert "MANGLED" in result.compressed
|
||||
assert result.strategy_used == CompressionStrategy.MIXED
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue