diff --git a/agent/memory_manager.py b/agent/memory_manager.py
index c3ea0a26128..79547139086 100644
--- a/agent/memory_manager.py
+++ b/agent/memory_manager.py
@@ -126,7 +126,10 @@ class StreamingContextScrubber:
idx = self._find_boundary_open_tag(buf)
if idx == -1:
# No open tag — hold back a potential partial open tag
- held = self._max_partial_suffix(buf, self._OPEN_TAG)
+ held = (
+ self._max_pending_open_suffix(buf)
+ or self._max_partial_suffix(buf, self._OPEN_TAG)
+ )
if held:
self._append_visible(out, buf[:-held])
self._buf = buf[-held:]
@@ -179,10 +182,25 @@ class StreamingContextScrubber:
idx = buf_lower.find(self._OPEN_TAG, search_start)
if idx == -1:
return -1
- if self._is_block_boundary(buf, idx):
+ if self._is_block_boundary(buf, idx) and self._has_block_opener_suffix(buf, idx):
return idx
search_start = idx + 1
+ def _max_pending_open_suffix(self, buf: str) -> int:
+ """Hold a complete boundary tag until the following char confirms it."""
+ if not buf.lower().endswith(self._OPEN_TAG):
+ return 0
+ idx = len(buf) - len(self._OPEN_TAG)
+ if not self._is_block_boundary(buf, idx):
+ return 0
+ return len(self._OPEN_TAG)
+
+ def _has_block_opener_suffix(self, buf: str, idx: int) -> bool:
+ after_idx = idx + len(self._OPEN_TAG)
+ if after_idx >= len(buf):
+ return False
+ return buf[after_idx] in "\r\n"
+
def _is_block_boundary(self, buf: str, idx: int) -> bool:
if idx == 0:
return self._at_block_boundary
diff --git a/tests/agent/test_streaming_context_scrubber.py b/tests/agent/test_streaming_context_scrubber.py
index 94ca221dba5..ed633b6b19f 100644
--- a/tests/agent/test_streaming_context_scrubber.py
+++ b/tests/agent/test_streaming_context_scrubber.py
@@ -73,7 +73,18 @@ class TestStreamingContextScrubberBasics:
s = StreamingContextScrubber()
out = (
s.feed("pre \nleak post")
+ + s.feed("-context>\nleak post")
+ + s.flush()
+ )
+ assert out == "pre \n post"
+ assert "leak" not in out
+
+ def test_open_tag_waits_for_newline_confirmation_across_deltas(self):
+ """A boundary tag is only a leaked block when the next char is a newline."""
+ s = StreamingContextScrubber()
+ out = (
+ s.feed("pre \n")
+ + s.feed("\nleak post")
+ s.flush()
)
assert out == "pre \n post"
@@ -83,7 +94,7 @@ class TestStreamingContextScrubberBasics:
"""The close tag arriving in two fragments."""
s = StreamingContextScrubber()
out = (
- s.feed("pre \nleak\nleak post")
+ s.flush()
)
@@ -116,18 +127,28 @@ class TestStreamingContextScrubberPartialTagFalsePositives:
)
assert out == "In that previous `` block, there was no matching fact."
- def test_mid_sentence_memory_context_pair_is_not_scrubbed(self):
+ def test_mid_sentence_memory_context_mention_is_not_scrubbed(self):
"""Only block-like memory-context spans are treated as leaked context."""
s = StreamingContextScrubber()
out = s.feed("The tag name is documented here.") + s.flush()
assert out == "The tag name is documented here."
+ def test_line_start_memory_context_mention_without_close_is_not_scrubbed(self):
+ """A plain-text line that starts with the tag name must be preserved."""
+ s = StreamingContextScrubber()
+ out = (
+ s.feed("Visible intro\n")
+ + s.feed(" is the literal tag name mentioned here.")
+ + s.flush()
+ )
+ assert out == "Visible intro\n is the literal tag name mentioned here."
+
class TestStreamingContextScrubberUnterminatedSpan:
def test_unterminated_span_drops_payload(self):
"""Provider drops close tag — better to lose output than to leak."""
s = StreamingContextScrubber()
- out = s.feed("pre \nsecret never closed") + s.flush()
+ out = s.feed("pre \n\nsecret never closed") + s.flush()
assert out == "pre \n"
assert "secret" not in out
@@ -144,7 +165,7 @@ class TestStreamingContextScrubberCaseInsensitivity:
def test_uppercase_tags_still_scrubbed(self):
s = StreamingContextScrubber()
out = (
- s.feed("secret")
+ s.feed("\nsecret")
+ s.feed("visible")
+ s.flush()
)