diff --git a/agent/memory_manager.py b/agent/memory_manager.py index 7eda64fba4..c3ea0a2612 100644 --- a/agent/memory_manager.py +++ b/agent/memory_manager.py @@ -91,10 +91,12 @@ class StreamingContextScrubber: def __init__(self) -> None: self._in_span: bool = False self._buf: str = "" + self._at_block_boundary: bool = True def reset(self) -> None: self._in_span = False self._buf = "" + self._at_block_boundary = True def feed(self, text: str) -> str: """Return the visible portion of ``text`` after scrubbing. @@ -121,19 +123,19 @@ class StreamingContextScrubber: buf = buf[idx + len(self._CLOSE_TAG):] self._in_span = False else: - idx = buf.lower().find(self._OPEN_TAG) + 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) if held: - out.append(buf[:-held]) + self._append_visible(out, buf[:-held]) self._buf = buf[-held:] else: - out.append(buf) + self._append_visible(out, buf) return "".join(out) # Emit text before the tag, enter span if idx > 0: - out.append(buf[:idx]) + self._append_visible(out, buf[:idx]) buf = buf[idx + len(self._OPEN_TAG):] self._in_span = True @@ -169,6 +171,40 @@ class StreamingContextScrubber: return i return 0 + def _find_boundary_open_tag(self, buf: str) -> int: + """Find an opening fence only when it starts a block-like span.""" + buf_lower = buf.lower() + search_start = 0 + while True: + idx = buf_lower.find(self._OPEN_TAG, search_start) + if idx == -1: + return -1 + if self._is_block_boundary(buf, idx): + return idx + search_start = idx + 1 + + def _is_block_boundary(self, buf: str, idx: int) -> bool: + if idx == 0: + return self._at_block_boundary + preceding = buf[:idx] + last_newline = preceding.rfind("\n") + if last_newline == -1: + return self._at_block_boundary and preceding.strip() == "" + return preceding[last_newline + 1:].strip() == "" + + def _append_visible(self, out: list[str], text: str) -> None: + if not text: + return + out.append(text) + self._update_block_boundary(text) + + def _update_block_boundary(self, text: str) -> None: + last_newline = text.rfind("\n") + if last_newline != -1: + self._at_block_boundary = text[last_newline + 1:].strip() == "" + else: + self._at_block_boundary = self._at_block_boundary and text.strip() == "" + def build_memory_context_block(raw_context: str) -> str: """Wrap prefetched memory in a fenced block with system note.""" diff --git a/tests/agent/test_streaming_context_scrubber.py b/tests/agent/test_streaming_context_scrubber.py index 99f33e7ce9..94ca221dba 100644 --- a/tests/agent/test_streaming_context_scrubber.py +++ b/tests/agent/test_streaming_context_scrubber.py @@ -37,13 +37,13 @@ class TestStreamingContextScrubberBasics: """The real streaming case: tag pair split across deltas.""" s = StreamingContextScrubber() deltas = [ - "Hello ", + "Hello\n", "\npayload ", "more payload\n", " world", ] out = "".join(s.feed(d) for d in deltas) + s.flush() - assert out == "Hello world" + assert out == "Hello\n world" assert "payload" not in out def test_realistic_fragmented_chunks_strip_memory_payload(self): @@ -72,22 +72,22 @@ class TestStreamingContextScrubberBasics: """The open tag itself arriving in two fragments.""" s = StreamingContextScrubber() out = ( - s.feed("pre leak post") + s.flush() ) - assert out == "pre post" + assert out == "pre \n post" assert "leak" not in out def test_close_tag_split_across_two_deltas(self): """The close tag arriving in two fragments.""" s = StreamingContextScrubber() out = ( - s.feed("pre leakleak post") + s.flush() ) - assert out == "pre post" + assert out == "pre \n post" assert "leak" not in out @@ -105,13 +105,30 @@ class TestStreamingContextScrubberPartialTagFalsePositives: out = s.feed("price < ") + s.feed("10 dollars") + s.flush() assert out == "price < 10 dollars" + def test_inline_memory_context_tag_mention_is_not_scrubbed(self): + """A prose mention of the fence tag must not swallow the answer.""" + s = StreamingContextScrubber() + out = ( + s.feed("In that previous `` block, ") + + s.feed("there was no matching fact.") + + s.flush() + ) + assert out == "In that previous `` block, there was no matching fact." + + def test_mid_sentence_memory_context_pair_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." + 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 secret never closed") + s.flush() - assert out == "pre " + out = s.feed("pre \nsecret never closed") + s.flush() + assert out == "pre \n" assert "secret" not in out def test_reset_clears_hung_span(self): @@ -171,7 +188,7 @@ class TestStreamingContextScrubberCrossTurn: def test_reset_clears_in_span_state(self): s = StreamingContextScrubber() - s.feed("textsecret-tail") + s.feed("text\nsecret-tail") # Mid-span state held — without reset, subsequent text would be # discarded until we see . s.reset()