fix(agent): bound remaining memory/history pollution paths from #3412
#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:
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user