From 92fba55f700632e09d71db0d0b9d65c21b3b481e Mon Sep 17 00:00:00 2001 From: Joey Yakimowich-Payne Date: Fri, 13 Mar 2026 21:07:52 -0600 Subject: [PATCH] fix: accurate context reduction stats and SSE cleanup tag filter Measure incoming_bytes before _preprocess() so bytes_saved reflects true reduction. Add SSECleanupFilter that intercepts memory_cleanup/yuyay-response tags in streaming responses, strips them from output, and executes ops (drops, collapses, releases) in real-time. Handles partial tags split across SSE chunks with a safety valve to flush stale buffers for prose. --- src/mnemosyne/benchmark.py | 5 +- src/mnemosyne/gateway.py | 282 ++++++++++++++++++++++++++++++++++++- 2 files changed, 282 insertions(+), 5 deletions(-) diff --git a/src/mnemosyne/benchmark.py b/src/mnemosyne/benchmark.py index ce50978..d7d5301 100644 --- a/src/mnemosyne/benchmark.py +++ b/src/mnemosyne/benchmark.py @@ -477,7 +477,10 @@ class BenchmarkCollector: "max_ms": round(max_time, 2), } - context_reduction = total_outgoing / total_incoming if total_incoming > 0 else 1.0 + total_bytes_saved = sum(s.tokens.total_bytes_saved for s in sessions) + context_reduction = ( + min(total_bytes_saved / total_incoming, 1.0) if total_incoming > 0 else 0.0 + ) return { "sessions": n, diff --git a/src/mnemosyne/gateway.py b/src/mnemosyne/gateway.py index af7fed0..89fa156 100644 --- a/src/mnemosyne/gateway.py +++ b/src/mnemosyne/gateway.py @@ -57,8 +57,274 @@ from mnemosyne.object_store import ObjectStoreBackend from mnemosyne.pager import PageStore, compact_messages from mnemosyne.message_store import MessageStore from mnemosyne.providers import adapters +from mnemosyne.tags import ( + parse_cleanup_tags, + parse_yuyay_response, + strip_cleanup_tags, + strip_yuyay_tags, +) from mnemosyne.telemetry import Telemetry +import re as _re + +# --------------------------------------------------------------------------- +# SSE Cleanup Filter — strips / tags from +# streaming responses and executes the contained ops in real-time. +# --------------------------------------------------------------------------- + +_TAG_OPEN_RE = _re.compile(r"<(memory_cleanup|yuyay-response)") +_TAG_CLOSE_RE = _re.compile(r"") + +# Matches a trailing '<' optionally followed by a prefix of a known tag name +# or ' bool: + """Return True if *buf* ends with a partial tag opener we care about. + + Only checks the last 25 characters — a complete tag name is at most + ```` (19 chars). This prevents false positives when + the model writes prose containing '<' earlier in the buffer. + """ + # Only look at the tail of the buffer for partial tags + window = buf[-25:] if len(buf) > 25 else buf + idx = window.rfind("<") + if idx == -1: + return False + tail = window[idx + 1 :] # everything after the last '<' + if not tail: + return True # bare '<' at the very end — could be start of any tag + # If the tail contains '>' then the tag is already closed — not partial + if ">" in tail: + return False + return any(p.startswith(tail) for p in _KNOWN_TAG_PREFIXES) + + +class SSECleanupFilter: + """Intercept SSE text deltas and strip/execute cleanup tags. + + Anthropic SSE streams emit ``content_block_delta`` events with + ``{"delta": {"type": "text_delta", "text": "..."}}``. This filter + accumulates those text fragments, detects complete + ```` / ```` blocks, executes the + operations against the session's BlockStore / PageStore, and rewrites + the SSE ``data:`` lines with the tags removed. + + Non-text events (ping, message_start, content_block_start, etc.) + pass through unchanged. + """ + + # Max text deltas to buffer while waiting for a tag to complete. + # Real tags complete within 2-3 deltas. If we exceed this, it's prose. + _MAX_BUFFERED_DELTAS = 6 + + def __init__(self, block_store: "BlockStore | None", page_store: "PageStore | None"): + self._bs = block_store + self._ps = page_store + self._buf = "" # accumulated text not yet flushed + self._inside_tag = False # currently buffering a tag body + self._buffered_deltas = 0 # how many deltas we've buffered without flushing + self._stats: list[str] = [] # executed ops log + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def filter_chunk(self, raw_chunk: bytes) -> bytes: + """Process a raw SSE chunk (may contain multiple lines). + + Returns the (possibly rewritten) chunk to forward downstream. + """ + lines = raw_chunk.split(b"\n") + out_lines: list[bytes] = [] + + for line in lines: + if not line.startswith(b"data: "): + out_lines.append(line) + continue + + json_bytes = line[6:] # strip "data: " prefix + if json_bytes.strip() in (b"", b"[DONE]"): + out_lines.append(line) + continue + + try: + event = json.loads(json_bytes) + except (json.JSONDecodeError, UnicodeDecodeError): + out_lines.append(line) + continue + + if event.get("type") != "content_block_delta": + out_lines.append(line) + continue + + delta = event.get("delta", {}) + if delta.get("type") != "text_delta": + out_lines.append(line) + continue + + text = delta.get("text", "") + if not text: + out_lines.append(line) + continue + + # Feed text into buffer and get cleaned output + cleaned = self._feed(text) + + if cleaned is None: + # Entire chunk is being buffered (inside a tag) — suppress + continue + elif cleaned == text: + # No change — pass through original bytes + out_lines.append(line) + elif not cleaned: + # Text was entirely a cleanup tag — suppress this event + continue + else: + # Rewrite the delta with cleaned text + delta["text"] = cleaned + out_lines.append(b"data: " + json.dumps(event, ensure_ascii=False).encode("utf-8")) + + return b"\n".join(out_lines) + + @property + def stats(self) -> str: + return "; ".join(self._stats) if self._stats else "" + + def flush(self) -> str: + """Flush any remaining buffered text (call at stream end).""" + if self._buf: + result = self._buf + self._buf = "" + self._inside_tag = False + return result + return "" + + # ------------------------------------------------------------------ + # Internal + # ------------------------------------------------------------------ + + def _feed(self, text: str) -> str | None: + """Feed a text delta fragment. Returns cleaned text or None to suppress. + + Returns: + str — cleaned text to emit (may be empty string to skip event) + None — suppress entirely (buffering inside a tag) + """ + self._buf += text + self._buffered_deltas += 1 + + has_open = bool(_TAG_OPEN_RE.search(self._buf)) + has_close = bool(_TAG_CLOSE_RE.search(self._buf)) + partial = _has_partial_tag(self._buf) + + # Safety valve: if we've buffered too many deltas without resolving + # a tag, this is just prose containing '<' — flush everything. + if self._buffered_deltas > self._MAX_BUFFERED_DELTAS and not has_close: + result = self._buf + self._buf = "" + self._inside_tag = False + self._buffered_deltas = 0 + return result + + # Fast path: no tag markers and no partial tag at the end + if not self._inside_tag and not has_open and not partial: + result = self._buf + self._buf = "" + self._buffered_deltas = 0 + return result + + # Partial tag opener at the end (e.g. " 0: + emit = self._buf[:idx] + self._buf = self._buf[idx:] + return emit + return None + + # We see a complete tag opening + if not self._inside_tag and has_open: + self._inside_tag = True + m = _TAG_OPEN_RE.search(self._buf) + if m and m.start() > 0: + emit = self._buf[: m.start()] + self._buf = self._buf[m.start() :] + return emit + + # Check if we have a complete tag + if self._inside_tag and has_close: + self._execute_ops(self._buf) + cleaned = strip_yuyay_tags(strip_cleanup_tags(self._buf)) + self._buf = "" + self._inside_tag = False + self._buffered_deltas = 0 + + if _TAG_OPEN_RE.search(cleaned): + self._buf = cleaned + self._inside_tag = True + return None + if _has_partial_tag(cleaned): + idx = cleaned.rfind("<") + if idx > 0: + emit = cleaned[:idx] + self._buf = cleaned[idx:] + return emit + self._buf = cleaned + return None + + return cleaned + + # Still inside an incomplete tag — keep buffering + if self._inside_tag: + return None + + # Shouldn't reach here, but safety + result = self._buf + self._buf = "" + self._buffered_deltas = 0 + return result + + def _execute_ops(self, text: str) -> None: + """Parse and execute cleanup/yuyay ops from buffered text.""" + ops_list = [] + + cleanup_ops = parse_cleanup_tags(text) + if not cleanup_ops.empty: + ops_list.append(cleanup_ops) + + yuyay_ops = parse_yuyay_response(text) + if not yuyay_ops.empty: + ops_list.append(yuyay_ops) + + for ops in ops_list: + if self._bs is not None: + for block_id in ops.drops: + if self._bs.drop(block_id): + self._stats.append(f"dropped {block_id}") + for block_id, summary in ops.summaries: + if self._bs.summarize(block_id, summary): + self._stats.append(f"summarized {block_id}") + for block_id in ops.anchors: + if self._bs.anchor(block_id): + self._stats.append(f"anchored {block_id}") + for collapse in ops.collapses: + collapsed = self._bs.collapse_range( + collapse.start_turn, collapse.end_turn, collapse.summary + ) + if collapsed: + self._stats.append( + f"collapsed turns {collapse.start_turn}-{collapse.end_turn} " + f"({len(collapsed)} blocks)" + ) + if self._ps is not None and ops.releases: + for path in ops.releases: + self._ps.mark_released(path) + self._stats.append(f"released {len(ops.releases)} path(s)") + # ANSI for stderr status lines _DIM = "\033[2m" @@ -2011,6 +2277,9 @@ def create_app( # Cleanup now runs inside _preprocess (before manifest injection) + # Measure raw incoming size BEFORE any preprocessing + incoming_bytes = len(json.dumps(payload, default=str).encode("utf-8")) + # Pre-process: system status, block labeling preprocess_bytes_saved = 0 if endpoint == "messages": @@ -2027,7 +2296,6 @@ def create_app( request_id = str(uuid.uuid4()) started = time.perf_counter() - incoming_bytes = len(json.dumps(payload, default=str).encode("utf-8")) session_id = session.id req = adapter.normalize_request(payload) @@ -2087,6 +2355,10 @@ def create_app( bytes_out = len(body) yield body else: + cleanup_filter = SSECleanupFilter( + block_store=session.block_store, + page_store=session.page_store, + ) for chunk in resp.iter_bytes(): bytes_out += len(chunk) chunk_count += 1 @@ -2099,7 +2371,9 @@ def create_app( provider=provider, usage_accumulator=usage, ) - yield chunk + filtered = cleanup_filter.filter_chunk(chunk) + if filtered: + yield filtered except Exception as e: stream_error = True emit_event( @@ -2121,7 +2395,7 @@ def create_app( session, incoming_bytes=incoming_bytes, outgoing_bytes=outgoing_bytes, - bytes_saved=preprocess_bytes_saved if endpoint == "messages" else 0, + bytes_saved=max(0, incoming_bytes - outgoing_bytes), ) _update_fidelity_pressure(usage, session) emit_event( @@ -2183,7 +2457,7 @@ def create_app( session, incoming_bytes=incoming_bytes, outgoing_bytes=outgoing_bytes, - bytes_saved=preprocess_bytes_saved, + bytes_saved=max(0, incoming_bytes - outgoing_bytes), ) _update_fidelity_pressure(usage, session)