🐛 fix(memory): keep inline memory-context mentions visible

This commit is contained in:
墨綠BG 2026-05-01 20:29:22 +08:00 committed by kshitij
parent 956dd44625
commit 341c8d3030
2 changed files with 66 additions and 13 deletions

View file

@ -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."""