From 4531167c1209e2be2e52913dd294e2285fad88f9 Mon Sep 17 00:00:00 2001 From: Xubin Ren Date: Thu, 23 Apr 2026 20:11:44 +0000 Subject: [PATCH] fix(agent): bound remaining memory/history pollution paths from #3412 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit #3412 stopped the headline raw_archive bloat but left four adjacent leaks on the same pollution chain: - archive() success path appended uncapped LLM summaries to history.jsonl, so a misbehaving LLM could re-open the #3412 bug from the happy path. - maybe_consolidate_by_tokens did not advance last_consolidated when archive() fell back to raw_archive, causing duplicate [RAW] dumps of the same chunk on every subsequent call. - Dream's Phase 1/2 prompt injected MEMORY.md / SOUL.md / USER.md and each history entry without caps, so any legacy oversized record (or an unbounded user edit) would blow past the context window every dream. - append_history itself had no default cap, leaving future new callers one forgotten-cap-away from the same vector. Changes: - Cap LLM-produced summaries at 8K chars (_ARCHIVE_SUMMARY_MAX_CHARS) before writing to history.jsonl. - Advance session.last_consolidated after archive() regardless of whether it summarized or raw-archived — both outcomes materialize the chunk; still break the round loop on fallback so a degraded LLM isn't hammered. - Truncate MEMORY.md / SOUL.md / USER.md and each history entry in Dream's Phase 1 prompt preview (Phase 2 still reaches full files via read_file). - Add _HISTORY_ENTRY_HARD_CAP (64K) as belt-and-suspenders default in append_history with a once-per-store warning, so any new caller that forgets its own tighter cap gets caught and observable. Layer the caps by scope: raw_archive=16K, archive summary=8K, append_history default=64K. Tight per-caller values cover expected payloads; the wide default only catches regressions. Tests: +9 regression tests covering each fix. Full suite: 2372 passed. Made-with: Cursor --- nanobot/agent/memory.py | 69 ++++++++++++++++++++++++++----- tests/agent/test_consolidator.py | 70 +++++++++++++++++++++++++++++++- tests/agent/test_dream.py | 51 +++++++++++++++++++++++ tests/agent/test_memory_store.py | 45 +++++++++++++++++++- 4 files changed, 222 insertions(+), 13 deletions(-) diff --git a/nanobot/agent/memory.py b/nanobot/agent/memory.py index 6a23227f..16c01d31 100644 --- a/nanobot/agent/memory.py +++ b/nanobot/agent/memory.py @@ -51,6 +51,7 @@ class MemoryStore: self._cursor_file = self.memory_dir / ".cursor" self._dream_cursor_file = self.memory_dir / ".dream_cursor" self._corruption_logged = False # rate-limit non-int cursor warning + self._oversize_logged = False # rate-limit oversized-entry warning self._git = GitStore(workspace, tracked_files=[ "SOUL.md", "USER.md", "memory/MEMORY.md", ]) @@ -222,7 +223,7 @@ class MemoryStore: # -- history.jsonl — append-only, JSONL format --------------------------- - def append_history(self, entry: str) -> int: + def append_history(self, entry: str, *, max_chars: int | None = None) -> int: """Append *entry* to history.jsonl and return its auto-incrementing cursor. Entries are passed through `strip_think` to drop template-level leaks @@ -231,10 +232,26 @@ class MemoryStore: the record is persisted with an empty string rather than falling back to the raw leak — otherwise `strip_think`'s guarantees would be undone by history replay / consolidation downstream. + + A defensive cap (*max_chars*, default ``_HISTORY_ENTRY_HARD_CAP``) is + applied as a final safety net: individual callers should cap their own + content more tightly; this default only exists to catch unintentional + large writes (e.g. an LLM echoing its input back as a "summary"). """ + limit = max_chars if max_chars is not None else _HISTORY_ENTRY_HARD_CAP cursor = self._next_cursor() ts = datetime.now().strftime("%Y-%m-%d %H:%M") raw = entry.rstrip() + if len(raw) > limit: + if not self._oversize_logged: + self._oversize_logged = True + logger.warning( + "history entry exceeds {} chars ({}); truncating. " + "Usually means a caller forgot its own cap; " + "further occurrences suppressed.", + limit, len(raw), + ) + raw = truncate_text(raw, limit) content = strip_think(raw) if raw and not content: logger.debug( @@ -393,7 +410,12 @@ class MemoryStore: # --------------------------------------------------------------------------- -_RAW_ARCHIVE_MAX_CHARS = 16_000 # cap raw_archive entries to avoid bloating history.jsonl +# Individual history.jsonl writers cap their own payloads tightly; the +# _HISTORY_ENTRY_HARD_CAP at append_history() is a belt-and-suspenders default +# that catches any new caller that forgot to set its own cap. +_RAW_ARCHIVE_MAX_CHARS = 16_000 # fallback dump (LLM failed) +_ARCHIVE_SUMMARY_MAX_CHARS = 8_000 # LLM-produced consolidation summary +_HISTORY_ENTRY_HARD_CAP = 64_000 # emergency cap in append_history class Consolidator: @@ -522,7 +544,7 @@ class Consolidator: if response.finish_reason == "error": raise RuntimeError(f"LLM returned error: {response.content}") summary = response.content or "[no summary]" - self.store.append_history(summary) + self.store.append_history(summary, max_chars=_ARCHIVE_SUMMARY_MAX_CHARS) return summary except Exception: logger.warning("Consolidation LLM call failed, raw-dumping to history") @@ -599,12 +621,18 @@ class Consolidator: len(chunk), ) summary = await self.archive(chunk) + # Advance the cursor either way: on success the chunk was + # summarized; on failure archive() already raw-archived it as + # a breadcrumb. Re-archiving the same chunk on the next call + # would just emit duplicate [RAW] entries. if summary: last_summary = summary - else: - break session.last_consolidated = end_idx self.sessions.save(session) + if not summary: + # LLM is degraded — stop hammering it this call; + # the next invocation can retry a fresh chunk. + break try: estimated, source = self.estimate_session_prompt_tokens( @@ -648,6 +676,15 @@ class Dream: LLM can make targeted, incremental edits instead of replacing entire files. """ + # Caps on prompt-bound inputs so Dream's LLM calls never exceed the model's + # context window just because a file (or a legacy large history entry) grew + # unexpectedly. Each file still appears in full via read_file when the agent + # needs it in Phase 2 — these caps only bound the Phase 1/2 prompt preview. + _MEMORY_FILE_MAX_CHARS = 32_000 + _SOUL_FILE_MAX_CHARS = 16_000 + _USER_FILE_MAX_CHARS = 16_000 + _HISTORY_ENTRY_PREVIEW_MAX_CHARS = 4_000 + def __init__( self, store: MemoryStore, @@ -786,21 +823,31 @@ class Dream: len(entries), last_cursor, batch[-1]["cursor"], len(batch), ) - # Build history text for LLM + # Build history text for LLM — cap each entry so a legacy oversized + # record (e.g. pre-#3412 raw_archive dump) can't blow up the prompt. history_text = "\n".join( - f"[{e['timestamp']}] {e['content']}" for e in batch + f"[{e['timestamp']}] " + f"{truncate_text(e['content'], self._HISTORY_ENTRY_PREVIEW_MAX_CHARS)}" + for e in batch ) - # Current file contents + per-line age annotations (MEMORY.md only) + # Current file contents + per-line age annotations (MEMORY.md only). + # Each file is capped in the *prompt preview* only; Phase 2 still sees + # the full file via the read_file tool. current_date = datetime.now().strftime("%Y-%m-%d") raw_memory = self.store.read_memory() or "(empty)" - current_memory = ( + annotated_memory = ( self._annotate_with_ages(raw_memory) if self.annotate_line_ages else raw_memory ) - current_soul = self.store.read_soul() or "(empty)" - current_user = self.store.read_user() or "(empty)" + current_memory = truncate_text(annotated_memory, self._MEMORY_FILE_MAX_CHARS) + current_soul = truncate_text( + self.store.read_soul() or "(empty)", self._SOUL_FILE_MAX_CHARS, + ) + current_user = truncate_text( + self.store.read_user() or "(empty)", self._USER_FILE_MAX_CHARS, + ) file_context = ( f"## Current Date\n{current_date}\n\n" diff --git a/tests/agent/test_consolidator.py b/tests/agent/test_consolidator.py index 58dce3c6..75a6d1f7 100644 --- a/tests/agent/test_consolidator.py +++ b/tests/agent/test_consolidator.py @@ -4,7 +4,12 @@ import pytest import asyncio from unittest.mock import AsyncMock, MagicMock, patch -from nanobot.agent.memory import Consolidator, MemoryStore, _RAW_ARCHIVE_MAX_CHARS +from nanobot.agent.memory import ( + Consolidator, + MemoryStore, + _ARCHIVE_SUMMARY_MAX_CHARS, + _RAW_ARCHIVE_MAX_CHARS, +) @pytest.fixture @@ -144,6 +149,56 @@ class TestConsolidatorTokenBudget: assert archived_chunk[0]["content"] == "m0" assert session.last_consolidated > 0 + async def test_raw_archive_fallback_advances_last_consolidated(self, consolidator): + """When archive() falls back to raw-archive (LLM failed), the cursor + must still advance. Otherwise the same chunk gets raw-archived again + on every subsequent maybe_consolidate_by_tokens() call, spamming + duplicate [RAW] entries into history.jsonl.""" + consolidator._SAFETY_BUFFER = 0 + session = MagicMock() + session.last_consolidated = 0 + session.key = "test:key" + session.messages = [ + {"role": "user" if i in {0, 50} else "assistant", "content": f"m{i}"} + for i in range(70) + ] + session.metadata = {} + consolidator.estimate_session_prompt_tokens = MagicMock( + side_effect=[(1200, "tiktoken"), (400, "tiktoken")] + ) + # LLM consolidation fails — archive() returns None (raw_archive fired). + consolidator.archive = AsyncMock(return_value=None) + + await consolidator.maybe_consolidate_by_tokens(session) + + consolidator.archive.assert_awaited_once() + # The chunk is considered "materialized" (as a raw-archive breadcrumb), + # so last_consolidated must have moved past it. + assert session.last_consolidated == 50 + + async def test_raw_archive_fallback_breaks_round_loop(self, consolidator): + """A degraded LLM should not trigger more archive() calls within the + same maybe_consolidate_by_tokens invocation — bail after one fallback.""" + consolidator._SAFETY_BUFFER = 0 + session = MagicMock() + session.last_consolidated = 0 + session.key = "test:key" + session.messages = [ + {"role": "user" if i in {0, 20, 40, 60} else "assistant", "content": f"m{i}"} + for i in range(70) + ] + session.metadata = {} + # Keep estimates high so the loop would otherwise run multiple rounds. + consolidator.estimate_session_prompt_tokens = MagicMock( + return_value=(1200, "tiktoken") + ) + consolidator.archive = AsyncMock(return_value=None) + + await consolidator.maybe_consolidate_by_tokens(session) + + # Exactly one fallback per call — not _MAX_CONSOLIDATION_ROUNDS. + assert consolidator.archive.await_count == 1 + async def test_boundary_respected_when_no_intermediate_user_turn(self, consolidator): """When boundary points past a long tool chain, the full chunk is archived.""" consolidator._SAFETY_BUFFER = 0 @@ -231,6 +286,19 @@ class TestArchiveTruncation: # Should be truncated assert len(user_content) < 250_000 + async def test_oversized_summary_is_capped_before_append(self, consolidator, mock_provider, store): + """A pathologically large LLM summary must not land full-length in + history.jsonl — that would re-open the #3412 bloat vector from the + *success* path instead of the fallback path.""" + mock_provider.chat_with_retry.return_value = MagicMock( + content="S" * (_ARCHIVE_SUMMARY_MAX_CHARS * 10), + finish_reason="stop", + ) + await consolidator.archive([{"role": "user", "content": "hi"}]) + + entry = store.read_unprocessed_history(since_cursor=0)[0] + assert len(entry["content"]) <= _ARCHIVE_SUMMARY_MAX_CHARS + 50 + async def test_archive_truncates_via_tiktoken_with_positive_budget(self, consolidator, mock_provider, store): """Positive token budget should use tiktoken for precise truncation.""" consolidator.context_window_tokens = 10_000 diff --git a/tests/agent/test_dream.py b/tests/agent/test_dream.py index cb6c8de7..27e49fda 100644 --- a/tests/agent/test_dream.py +++ b/tests/agent/test_dream.py @@ -1,5 +1,7 @@ """Tests for the Dream class — two-phase memory consolidation via AgentRunner.""" +import json + import pytest from unittest.mock import AsyncMock, MagicMock, patch @@ -256,3 +258,52 @@ class TestDreamRun: # The template renders with stale_threshold_days=14 → LLM must see "N>14" assert "N>14" in system_msg + +class TestDreamPromptCaps: + """Dream's Phase 1/2 prompt must not be poisoned by a legacy oversized + history entry or a runaway MEMORY.md. Without caps, a single pre-#3412 + raw_archive dump in history.jsonl would make every subsequent Dream run + exceed the context window and silently advance the cursor past real work. + """ + + async def test_phase1_caps_huge_memory_file( + self, dream, mock_provider, mock_runner, store, + ): + """A MEMORY.md much larger than _MEMORY_FILE_MAX_CHARS must be truncated + in the prompt preview (full content is still reachable via read_file).""" + store.write_memory("M" * (dream._MEMORY_FILE_MAX_CHARS * 5)) + store.append_history("some event") + mock_provider.chat_with_retry.return_value = MagicMock(content="[SKIP]") + mock_runner.run = AsyncMock(return_value=_make_run_result()) + + await dream.run() + + user_msg = mock_provider.chat_with_retry.call_args.kwargs["messages"][1]["content"] + memory_section = user_msg.split("## Current MEMORY.md")[1].split("## Current SOUL.md")[0] + assert len(memory_section) < dream._MEMORY_FILE_MAX_CHARS + 500 + + async def test_phase1_caps_huge_history_entry( + self, dream, mock_provider, mock_runner, store, + ): + """A legacy oversized history entry (e.g. pre-#3412 raw_archive dump) + must not explode the Phase 1 prompt — each entry is capped in the + preview, even though the JSONL record itself stays full-size.""" + # Bypass the append_history cap by writing directly, simulating a + # record that was written by an older nanobot build before any caps. + store.history_file.write_text( + json.dumps({ + "cursor": 1, + "timestamp": "2026-04-01 10:00", + "content": "H" * (dream._HISTORY_ENTRY_PREVIEW_MAX_CHARS * 8), + }) + "\n", + encoding="utf-8", + ) + mock_provider.chat_with_retry.return_value = MagicMock(content="[SKIP]") + mock_runner.run = AsyncMock(return_value=_make_run_result()) + + await dream.run() + + user_msg = mock_provider.chat_with_retry.call_args.kwargs["messages"][1]["content"] + history_section = user_msg.split("## Conversation History\n")[1].split("\n\n## Current Date")[0] + assert len(history_section) < dream._HISTORY_ENTRY_PREVIEW_MAX_CHARS + 500 + diff --git a/tests/agent/test_memory_store.py b/tests/agent/test_memory_store.py index 94adbf37..8f322045 100644 --- a/tests/agent/test_memory_store.py +++ b/tests/agent/test_memory_store.py @@ -5,7 +5,7 @@ from datetime import datetime import pytest -from nanobot.agent.memory import MemoryStore +from nanobot.agent.memory import MemoryStore, _HISTORY_ENTRY_HARD_CAP @pytest.fixture @@ -142,6 +142,49 @@ class TestHistoryWithCursor: assert entries[0]["cursor"] in {4, 5} +class TestAppendHistoryHardCap: + """append_history has a defensive cap that catches new callers who forgot + to set their own tighter cap. The default is intentionally larger than + any current caller's per-call cap, so normal operation never trips it.""" + + def test_oversized_entry_is_truncated(self, store): + """An entry above _HISTORY_ENTRY_HARD_CAP is truncated before being persisted.""" + huge = "x" * (_HISTORY_ENTRY_HARD_CAP + 10_000) + store.append_history(huge) + entry = store.read_unprocessed_history(since_cursor=0)[0] + assert len(entry["content"]) <= _HISTORY_ENTRY_HARD_CAP + 50 + + def test_oversize_warning_is_emitted_once(self, store, caplog): + """Repeated oversized writes should warn only on the first occurrence.""" + from loguru import logger as loguru_logger + + records: list[str] = [] + handler_id = loguru_logger.add(lambda m: records.append(m), level="WARNING") + try: + huge = "x" * (_HISTORY_ENTRY_HARD_CAP + 1) + store.append_history(huge) + store.append_history(huge) + store.append_history(huge) + finally: + loguru_logger.remove(handler_id) + + oversize_warnings = [r for r in records if "exceeds" in r and "chars" in r] + assert len(oversize_warnings) == 1 + + def test_custom_max_chars_overrides_default(self, store): + """Callers that pass max_chars should get their tighter cap applied.""" + store.append_history("a" * 500, max_chars=100) + entry = store.read_unprocessed_history(since_cursor=0)[0] + assert len(entry["content"]) <= 150 # 100 + "\n... (truncated)" + + def test_normal_sized_entries_unaffected(self, store): + """The hard cap must not alter entries that fit within it.""" + msg = "normal short entry" + store.append_history(msg) + entry = store.read_unprocessed_history(since_cursor=0)[0] + assert entry["content"] == msg + + class TestDreamCursor: def test_initial_cursor_is_zero(self, store): assert store.get_last_dream_cursor() == 0