Skip to content

Commit 77320ee

Browse files
committed
keep markdown separators inside wrapped content
1 parent a3b35de commit 77320ee

2 files changed

Lines changed: 69 additions & 9 deletions

File tree

‎sdk/src/unplug/core/agent/boundaries.py‎

Lines changed: 23 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,10 @@
3232

3333
_REMOVED_MARKER = "[removed untrusted boundary marker]"
3434

35+
# Written by wrap_external_content on both sides of the body and read back by
36+
# strip_boundary_markers. Shared so the two cannot drift apart.
37+
_DELIMITER = "---\n"
38+
3539

3640
class WrappedContent(BaseModel):
3741
"""Untrusted payload wrapped with spoof-resistant boundary markers."""
@@ -71,9 +75,9 @@ def wrap_external_content(
7175
wrapped = (
7276
f'{_BEGIN_PREFIX} source="{source}" id="{mid}">>>\n'
7377
f"{_WARNING}\n"
74-
f"---\n"
78+
f"{_DELIMITER}"
7579
f"{body}\n"
76-
f"---\n"
80+
f"{_DELIMITER}"
7781
f'{_END_PREFIX} id="{mid}">>>'
7882
)
7983
return WrappedContent(text=wrapped, marker_id=mid, source=source, sanitized=sanitized)
@@ -168,14 +172,24 @@ def strip_boundary_markers(text: str) -> str:
168172
return text
169173

170174
def _inner(block: re.Match[str]) -> str:
175+
# The wrapper puts a "---" delimiter on either side of the body, so the
176+
# body is everything between the FIRST delimiter and the LAST one.
177+
# Splitting with maxsplit stops at the body's own first "---", which is
178+
# legal Markdown and not a boundary marker, and silently drops the rest
179+
# of the payload (#192). Reading one delimiter in from each end keeps
180+
# every byte in between, however many separators the body contains.
171181
chunk = block.group(0)
172-
parts = chunk.split("---\n", 2)
173-
if len(parts) >= 3:
174-
inner = parts[1]
175-
if inner.endswith("\n"):
176-
return inner[:-1]
177-
return inner
178-
return _REMOVED_MARKER
182+
_, opened, rest = chunk.partition(_DELIMITER)
183+
if not opened:
184+
return _REMOVED_MARKER
185+
inner, closed, _ = rest.rpartition(_DELIMITER)
186+
if not closed:
187+
return _REMOVED_MARKER
188+
# wrap_external_content writes "{body}\n---\n", so the newline before
189+
# the closing delimiter belongs to the wrapper rather than the body.
190+
if inner.endswith("\n"):
191+
inner = inner[:-1]
192+
return inner
179193

180194
stripped = _MARKER_BLOCK_RE.sub(_inner, text)
181195
stripped = _ORPHAN_BEGIN_RE.sub("", stripped)

‎sdk/tests/unit/core/agent/test_boundaries.py‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@
22

33
from __future__ import annotations
44

5+
import pytest
6+
57
from unplug.core.agent.boundaries import (
68
generate_marker_id,
79
sanitize_boundary_markers,
@@ -46,5 +48,49 @@ def test_strip_boundary_markers_roundtrip() -> None:
4648
assert strip_boundary_markers(wrapped.text) == inner
4749

4850

51+
@pytest.mark.parametrize(
52+
"inner",
53+
[
54+
pytest.param("# Notes\n\nBefore\n---\nAfter\n\nFinal paragraph.", id="one-separator"),
55+
pytest.param("Section 1\n---\nSection 2\n---\nSection 3", id="two-separators"),
56+
pytest.param("---", id="separator-only"),
57+
pytest.param("A\n---", id="trailing-separator"),
58+
pytest.param("---\nleading", id="leading-separator"),
59+
pytest.param("a\n---\n---\n---\nb", id="consecutive-separators"),
60+
],
61+
)
62+
def test_strip_preserves_markdown_separators_in_the_body(inner: str) -> None:
63+
"""A "---" line inside the payload is content, not a boundary marker (#192).
64+
65+
The old implementation split the block on the delimiter with maxsplit, so
66+
the body's own first separator was read as the wrapper's closing one and
67+
everything after it was dropped without an error.
68+
"""
69+
wrapped = wrap_external_content(inner, marker_id="a" * 16, sanitize=False)
70+
assert strip_boundary_markers(wrapped.text) == inner
71+
72+
73+
def test_strip_keeps_every_byte_after_an_internal_separator() -> None:
74+
# The regression as reported: the content was not mangled, it was gone.
75+
inner = "# Notes\n\nBefore\n---\nAfter\n\nFinal paragraph."
76+
wrapped = wrap_external_content(inner, marker_id="b" * 16, sanitize=False)
77+
result = strip_boundary_markers(wrapped.text)
78+
assert "After" in result
79+
assert "Final paragraph." in result
80+
81+
82+
def test_strip_reports_a_block_with_no_delimiters() -> None:
83+
# A block the wrapper did not write has no body to recover, so it is
84+
# replaced rather than returned as-is.
85+
forged = (
86+
'<<<UNTRUSTED source="retrieved" id="c'
87+
+ "c" * 15
88+
+ '">>>\nbody\n<<<END id="c'
89+
+ ("c" * 15)
90+
+ '">>>'
91+
)
92+
assert "body" not in strip_boundary_markers(forged)
93+
94+
4995
def test_generate_marker_id_length() -> None:
5096
assert len(generate_marker_id()) == 16

0 commit comments

Comments
 (0)