Grayscale/nanobot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(agent): bound remaining memory/history pollution paths from #3412

Browse Source 
#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
This commit is contained in:
Xubin Ren
2026-04-24 04:17:19 +08:00
committed by Xubin Ren
parent 81a5af2352
commit 4531167c12
4 changed files with 222 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nanobot/agent/memory.py
+58 -11
View File
@@ -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"
tests/agent/test_consolidator.py
+69 -1
View File
@@ -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
tests/agent/test_dream.py
+51
View File
@@ -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
tests/agent/test_memory_store.py
+44 -1
View File
@@ -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