mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-12 15:19:16 +03:00
feat(dream): single-phase consolidation (cherry-pick from feat/single-phase-dream)
This commit is contained in:
chengyongru
+397
-216
@@ -1,19 +1,32 @@
|
||||
"""Tests for the Dream class — two-phase memory consolidation via AgentRunner."""
|
||||
"""Tests for Dream driven through AgentLoop._process_system_message."""
|
||||
|
||||
import json
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from nanobot.agent.memory import Dream, MemoryStore
|
||||
from nanobot.agent.loop import AgentLoop
|
||||
from nanobot.agent.runner import AgentRunResult
|
||||
from nanobot.agent.skills import BUILTIN_SKILLS_DIR
|
||||
from nanobot.bus.events import InboundMessage
|
||||
from nanobot.bus.queue import MessageBus
|
||||
from nanobot.utils.gitstore import LineAge
|
||||
|
||||
|
||||
def _provider(default_model: str, max_tokens: int = 123) -> MagicMock:
|
||||
provider = MagicMock()
|
||||
provider.get_default_model.return_value = default_model
|
||||
provider.generation = SimpleNamespace(
|
||||
max_tokens=max_tokens, temperature=0.1, reasoning_effort=None
|
||||
)
|
||||
return provider
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def store(tmp_path):
|
||||
from nanobot.agent.memory import MemoryStore
|
||||
|
||||
s = MemoryStore(tmp_path)
|
||||
s.write_soul("# Soul\n- Helpful")
|
||||
s.write_user("# User\n- Developer")
|
||||
@@ -23,9 +36,7 @@ def store(tmp_path):
|
||||
|
||||
@pytest.fixture
|
||||
def mock_provider():
|
||||
p = MagicMock()
|
||||
p.chat_with_retry = AsyncMock()
|
||||
return p
|
||||
return _provider("test-model")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
@@ -34,10 +45,16 @@ def mock_runner():
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def dream(store, mock_provider, mock_runner):
|
||||
d = Dream(store=store, provider=mock_provider, model="test-model", max_batch_size=5)
|
||||
d._runner = mock_runner
|
||||
return d
|
||||
def loop(tmp_path, mock_provider, mock_runner):
|
||||
loop = AgentLoop(
|
||||
bus=MessageBus(),
|
||||
provider=mock_provider,
|
||||
workspace=tmp_path,
|
||||
model="test-model",
|
||||
context_window_tokens=1000,
|
||||
)
|
||||
loop.dream._runner = mock_runner
|
||||
return loop
|
||||
|
||||
|
||||
def _make_run_result(
|
||||
@@ -56,254 +73,418 @@ def _make_run_result(
|
||||
)
|
||||
|
||||
|
||||
class TestDreamRun:
|
||||
async def test_noop_when_no_unprocessed_history(self, dream, mock_provider, mock_runner, store):
|
||||
"""Dream should not call LLM when there's nothing to process."""
|
||||
result = await dream.run()
|
||||
assert result is False
|
||||
mock_provider.chat_with_retry.assert_not_called()
|
||||
class TestDreamAgentLoopIntegration:
|
||||
async def test_completes_goal_state_after_full_backlog(self, loop, mock_runner, store):
|
||||
"""Goal should be completed after processing all backlog in internal loop."""
|
||||
for i in range(6):
|
||||
store.append_history(f"event {i}")
|
||||
mock_runner.run = AsyncMock(return_value=_make_run_result())
|
||||
msg = InboundMessage(
|
||||
channel="system", sender_id="dream", chat_id="dream", content=""
|
||||
)
|
||||
await loop._process_system_message(msg)
|
||||
session = loop.sessions.get_or_create("system:dream")
|
||||
goal = session.metadata.get("goal_state")
|
||||
assert isinstance(goal, dict)
|
||||
assert goal["status"] == "completed"
|
||||
assert store.get_last_dream_cursor() == 6
|
||||
|
||||
async def test_completes_goal_state_on_finish(self, loop, mock_runner, store):
|
||||
"""Goal should be marked completed when backlog is fully processed."""
|
||||
store.append_history("event 1")
|
||||
mock_runner.run = AsyncMock(return_value=_make_run_result())
|
||||
msg = InboundMessage(
|
||||
channel="system", sender_id="dream", chat_id="dream", content=""
|
||||
)
|
||||
await loop._process_system_message(msg)
|
||||
session = loop.sessions.get_or_create("system:dream")
|
||||
goal = session.metadata.get("goal_state")
|
||||
assert goal["status"] == "completed"
|
||||
assert "completed_at" in goal
|
||||
assert "recap" in goal
|
||||
|
||||
async def test_noop_when_no_unprocessed_history(self, loop, mock_runner):
|
||||
"""Dream should not call runner when there's nothing to process."""
|
||||
msg = InboundMessage(
|
||||
channel="system", sender_id="dream", chat_id="dream", content=""
|
||||
)
|
||||
result = await loop._process_system_message(msg)
|
||||
assert result is None
|
||||
mock_runner.run.assert_not_called()
|
||||
|
||||
async def test_calls_runner_for_unprocessed_entries(self, dream, mock_provider, mock_runner, store):
|
||||
async def test_calls_runner_for_unprocessed_entries(self, loop, mock_runner, store):
|
||||
"""Dream should call AgentRunner when there are unprocessed history entries."""
|
||||
store.append_history("User prefers dark mode")
|
||||
mock_provider.chat_with_retry.return_value = MagicMock(content="New fact")
|
||||
mock_runner.run = AsyncMock(return_value=_make_run_result(
|
||||
tool_events=[{"name": "edit_file", "status": "ok", "detail": "memory/MEMORY.md"}],
|
||||
))
|
||||
result = await dream.run()
|
||||
assert result is True
|
||||
mock_runner.run = AsyncMock(
|
||||
return_value=_make_run_result(
|
||||
tool_events=[
|
||||
{"name": "edit_file", "status": "ok", "detail": "memory/MEMORY.md"}
|
||||
],
|
||||
)
|
||||
)
|
||||
msg = InboundMessage(
|
||||
channel="system", sender_id="dream", chat_id="dream", content=""
|
||||
)
|
||||
await loop._process_system_message(msg)
|
||||
mock_runner.run.assert_called_once()
|
||||
spec = mock_runner.run.call_args[0][0]
|
||||
assert spec.max_iterations == 10
|
||||
assert spec.fail_on_tool_error is False
|
||||
|
||||
async def test_advances_dream_cursor(self, dream, mock_provider, mock_runner, store):
|
||||
async def test_advances_dream_cursor(self, loop, mock_runner, store):
|
||||
"""Dream should advance the cursor after processing."""
|
||||
store.append_history("event 1")
|
||||
store.append_history("event 2")
|
||||
mock_provider.chat_with_retry.return_value = MagicMock(content="Nothing new")
|
||||
mock_runner.run = AsyncMock(return_value=_make_run_result())
|
||||
await dream.run()
|
||||
msg = InboundMessage(
|
||||
channel="system", sender_id="dream", chat_id="dream", content=""
|
||||
)
|
||||
await loop._process_system_message(msg)
|
||||
assert store.get_last_dream_cursor() == 2
|
||||
|
||||
async def test_compacts_processed_history(self, dream, mock_provider, mock_runner, store):
|
||||
async def test_compacts_processed_history(self, loop, mock_runner, store):
|
||||
"""Dream should compact history after processing."""
|
||||
store.append_history("event 1")
|
||||
store.append_history("event 2")
|
||||
store.append_history("event 3")
|
||||
mock_provider.chat_with_retry.return_value = MagicMock(content="Nothing new")
|
||||
mock_runner.run = AsyncMock(return_value=_make_run_result())
|
||||
await dream.run()
|
||||
# After Dream, cursor is advanced and 3, compact keeps last max_history_entries
|
||||
msg = InboundMessage(
|
||||
channel="system", sender_id="dream", chat_id="dream", content=""
|
||||
)
|
||||
await loop._process_system_message(msg)
|
||||
entries = store.read_unprocessed_history(since_cursor=0)
|
||||
assert all(e["cursor"] > 0 for e in entries)
|
||||
|
||||
async def test_skill_phase_uses_builtin_skill_creator_path(self, dream, mock_provider, mock_runner, store):
|
||||
"""Dream should point skill creation guidance at the builtin skill-creator template."""
|
||||
async def test_processes_full_backlog_in_one_call(self, loop, mock_runner, store):
|
||||
"""Backlog larger than max_batch_size should be fully processed in one call."""
|
||||
for i in range(12):
|
||||
store.append_history(f"event {i}")
|
||||
mock_runner.run = AsyncMock(return_value=_make_run_result())
|
||||
msg = InboundMessage(
|
||||
channel="system", sender_id="dream", chat_id="dream", content=""
|
||||
)
|
||||
await loop._process_system_message(msg)
|
||||
assert store.get_last_dream_cursor() == 12
|
||||
assert mock_runner.run.call_count == 3 # 5 + 5 + 2
|
||||
|
||||
async def test_single_git_commit_for_multi_batch(self, loop, mock_runner, store):
|
||||
"""Multi-batch run should collapse into exactly one git commit."""
|
||||
store.git.init()
|
||||
store.git.auto_commit("initial")
|
||||
for i in range(12):
|
||||
store.append_history(f"event {i}")
|
||||
mock_runner.run = AsyncMock(
|
||||
return_value=_make_run_result(
|
||||
tool_events=[
|
||||
{"name": "edit_file", "status": "ok", "detail": "memory/MEMORY.md"}
|
||||
],
|
||||
)
|
||||
)
|
||||
msg = InboundMessage(
|
||||
channel="system", sender_id="dream", chat_id="dream", content=""
|
||||
)
|
||||
await loop._process_system_message(msg)
|
||||
commits = store.git.log()
|
||||
dream_commits = [c for c in commits if c.message.startswith("dream:")]
|
||||
assert len(dream_commits) == 1
|
||||
|
||||
async def test_system_prompt_cached(self, loop, mock_runner, store):
|
||||
"""Batches within one run should reuse cached system prompt when template mtime unchanged."""
|
||||
for i in range(6):
|
||||
store.append_history(f"event {i}")
|
||||
mock_runner.run = AsyncMock(return_value=_make_run_result())
|
||||
msg = InboundMessage(
|
||||
channel="system", sender_id="dream", chat_id="dream", content=""
|
||||
)
|
||||
await loop._process_system_message(msg)
|
||||
# Two batches (5 + 1), both should use the same cached prompt
|
||||
assert mock_runner.run.call_count == 2
|
||||
first_prompt = mock_runner.run.call_args_list[0][0][0].initial_messages[0]["content"]
|
||||
second_prompt = mock_runner.run.call_args_list[1][0][0].initial_messages[0]["content"]
|
||||
assert second_prompt is first_prompt
|
||||
|
||||
async def test_noop_when_empty_backlog(self, loop, mock_runner, store):
|
||||
"""Empty backlog should not advance cursor or create a commit."""
|
||||
store.git.init()
|
||||
msg = InboundMessage(
|
||||
channel="system", sender_id="dream", chat_id="dream", content=""
|
||||
)
|
||||
await loop._process_system_message(msg)
|
||||
assert store.get_last_dream_cursor() == 0
|
||||
commits = store.git.log()
|
||||
assert len([c for c in commits if c.message.startswith("dream:")]) == 0
|
||||
|
||||
|
||||
class TestDreamPrompt:
|
||||
async def test_prompt_contains_mece_rules(self, loop, mock_runner, store):
|
||||
store.append_history("some event")
|
||||
mock_runner.run = AsyncMock(return_value=_make_run_result())
|
||||
msg = InboundMessage(
|
||||
channel="system", sender_id="dream", chat_id="dream", content=""
|
||||
)
|
||||
await loop._process_system_message(msg)
|
||||
spec = mock_runner.run.call_args[0][0]
|
||||
system_prompt = spec.initial_messages[0]["content"]
|
||||
assert "Do NOT guess paths" in system_prompt
|
||||
assert "SOUL.md" in system_prompt
|
||||
assert "USER.md" in system_prompt
|
||||
assert "MEMORY.md" in system_prompt
|
||||
|
||||
async def test_skill_phase_uses_builtin_skill_creator_path(self, loop, mock_runner, store):
|
||||
store.append_history("Repeated workflow one")
|
||||
store.append_history("Repeated workflow two")
|
||||
mock_provider.chat_with_retry.return_value = MagicMock(content="[SKILL] test-skill: test description")
|
||||
mock_runner.run = AsyncMock(return_value=_make_run_result())
|
||||
|
||||
await dream.run()
|
||||
|
||||
msg = InboundMessage(
|
||||
channel="system", sender_id="dream", chat_id="dream", content=""
|
||||
)
|
||||
await loop._process_system_message(msg)
|
||||
spec = mock_runner.run.call_args[0][0]
|
||||
system_prompt = spec.initial_messages[0]["content"]
|
||||
expected = str(BUILTIN_SKILLS_DIR / "skill-creator" / "SKILL.md")
|
||||
assert expected in system_prompt
|
||||
|
||||
async def test_skill_write_tool_accepts_workspace_relative_skill_path(self, dream, store):
|
||||
"""Dream skill creation should allow skills/<name>/SKILL.md relative to workspace root."""
|
||||
write_tool = dream._tools.get("write_file")
|
||||
assert write_tool is not None
|
||||
|
||||
result = await write_tool.execute(
|
||||
path="skills/test-skill/SKILL.md",
|
||||
content="---\nname: test-skill\ndescription: Test\n---\n",
|
||||
async def test_system_prompt_uses_threshold_from_template_var(self, loop, mock_runner, store):
|
||||
store.append_history("some event")
|
||||
mock_runner.run = AsyncMock(return_value=_make_run_result())
|
||||
msg = InboundMessage(
|
||||
channel="system", sender_id="dream", chat_id="dream", content=""
|
||||
)
|
||||
|
||||
assert "Successfully wrote" in result
|
||||
assert (store.workspace / "skills" / "test-skill" / "SKILL.md").exists()
|
||||
|
||||
async def test_phase1_prompt_includes_line_age_annotations(self, dream, mock_provider, mock_runner, store):
|
||||
"""Phase 1 prompt should have per-line age suffixes in MEMORY.md when git is available."""
|
||||
store.append_history("some event")
|
||||
mock_provider.chat_with_retry.return_value = MagicMock(content="[SKIP]")
|
||||
mock_runner.run = AsyncMock(return_value=_make_run_result())
|
||||
|
||||
# Init git so line_ages works
|
||||
store.git.init()
|
||||
store.git.auto_commit("initial memory state")
|
||||
|
||||
await dream.run()
|
||||
|
||||
# The MEMORY.md section should not crash and should contain the memory content
|
||||
call_args = mock_provider.chat_with_retry.call_args
|
||||
user_msg = call_args.kwargs.get("messages", call_args[1].get("messages"))[1]["content"]
|
||||
assert "## Current MEMORY.md" in user_msg
|
||||
|
||||
async def test_phase1_annotates_only_memory_not_soul_or_user(self, dream, mock_provider, mock_runner, store):
|
||||
"""SOUL.md and USER.md should never have age annotations — they are permanent."""
|
||||
store.append_history("some event")
|
||||
mock_provider.chat_with_retry.return_value = MagicMock(content="[SKIP]")
|
||||
mock_runner.run = AsyncMock(return_value=_make_run_result())
|
||||
|
||||
store.git.init()
|
||||
store.git.auto_commit("initial state")
|
||||
|
||||
await dream.run()
|
||||
|
||||
call_args = mock_provider.chat_with_retry.call_args
|
||||
user_msg = call_args.kwargs.get("messages", call_args[1].get("messages"))[1]["content"]
|
||||
# The ← suffix should only appear in MEMORY.md section
|
||||
memory_section = user_msg.split("## Current MEMORY.md")[1].split("## Current SOUL.md")[0]
|
||||
soul_section = user_msg.split("## Current SOUL.md")[1].split("## Current USER.md")[0]
|
||||
user_section = user_msg.split("## Current USER.md")[1]
|
||||
# SOUL and USER should not contain age arrows
|
||||
assert "\u2190" not in soul_section
|
||||
assert "\u2190" not in user_section
|
||||
|
||||
async def test_phase1_prompt_works_without_git(self, dream, mock_provider, mock_runner, store):
|
||||
"""Phase 1 should work fine even if git is not initialized (no age annotations)."""
|
||||
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()
|
||||
|
||||
# Should still succeed — just without age annotations
|
||||
mock_provider.chat_with_retry.assert_called_once()
|
||||
call_args = mock_provider.chat_with_retry.call_args
|
||||
user_msg = call_args.kwargs.get("messages", call_args[1].get("messages"))[1]["content"]
|
||||
assert "## Current MEMORY.md" in user_msg
|
||||
|
||||
async def test_phase1_prompt_carries_age_suffix_for_stale_lines(
|
||||
self, dream, mock_provider, mock_runner, store,
|
||||
):
|
||||
"""End-to-end: ages >14d must appear verbatim in the LLM prompt, ages ≤14d must not."""
|
||||
# MEMORY.md fixture has 2 non-blank lines ("# Memory" and "- Project X active").
|
||||
# Inject four ages to cover threshold boundaries: >14 suffix, ==14 no suffix, <14 no suffix.
|
||||
store.write_memory("# Memory\n- Project X active\n- fresh item\n- edge case line")
|
||||
store.append_history("some event")
|
||||
mock_provider.chat_with_retry.return_value = MagicMock(content="[SKIP]")
|
||||
mock_runner.run = AsyncMock(return_value=_make_run_result())
|
||||
|
||||
fake_ages = [
|
||||
LineAge(age_days=30), # "# Memory" → should get ← 30d
|
||||
LineAge(age_days=20), # "- Project X..." → should get ← 20d
|
||||
LineAge(age_days=14), # "- fresh item" → ==14, threshold is strictly >14, no suffix
|
||||
LineAge(age_days=5), # "- edge case..." → no suffix
|
||||
]
|
||||
with patch.object(store.git, "line_ages", return_value=fake_ages):
|
||||
await dream.run()
|
||||
|
||||
call_args = mock_provider.chat_with_retry.call_args
|
||||
user_msg = call_args.kwargs.get("messages", call_args[1].get("messages"))[1]["content"]
|
||||
memory_section = user_msg.split("## Current MEMORY.md")[1].split("## Current SOUL.md")[0]
|
||||
assert "\u2190 30d" in memory_section
|
||||
assert "\u2190 20d" in memory_section
|
||||
assert "\u2190 14d" not in memory_section
|
||||
assert "\u2190 5d" not in memory_section
|
||||
|
||||
async def test_phase1_skips_annotation_when_disabled(
|
||||
self, dream, mock_provider, mock_runner, store,
|
||||
):
|
||||
"""`annotate_line_ages=False` must bypass the git lookup entirely and keep MEMORY.md raw."""
|
||||
store.append_history("some event")
|
||||
mock_provider.chat_with_retry.return_value = MagicMock(content="[SKIP]")
|
||||
mock_runner.run = AsyncMock(return_value=_make_run_result())
|
||||
|
||||
dream.annotate_line_ages = False
|
||||
# line_ages must be bypassed entirely — verify with a spy rather than a
|
||||
# raising side_effect, because _annotate_with_ages catches Exception
|
||||
# (which swallows AssertionError) and would hide an accidental call.
|
||||
with patch.object(store.git, "line_ages") as mock_line_ages:
|
||||
await dream.run()
|
||||
mock_line_ages.assert_not_called()
|
||||
|
||||
call_args = mock_provider.chat_with_retry.call_args
|
||||
user_msg = call_args.kwargs.get("messages", call_args[1].get("messages"))[1]["content"]
|
||||
assert "\u2190" not in user_msg
|
||||
|
||||
async def test_phase1_skips_annotation_on_line_ages_length_mismatch(
|
||||
self, dream, mock_provider, mock_runner, store,
|
||||
):
|
||||
"""If ages length != lines length (dirty working tree), skip annotation instead of mis-tagging."""
|
||||
# MEMORY.md has 2 non-blank lines but we hand back only 1 age → mismatch.
|
||||
store.append_history("some event")
|
||||
mock_provider.chat_with_retry.return_value = MagicMock(content="[SKIP]")
|
||||
mock_runner.run = AsyncMock(return_value=_make_run_result())
|
||||
|
||||
with patch.object(store.git, "line_ages", return_value=[LineAge(age_days=999)]):
|
||||
await dream.run()
|
||||
|
||||
call_args = mock_provider.chat_with_retry.call_args
|
||||
user_msg = call_args.kwargs.get("messages", call_args[1].get("messages"))[1]["content"]
|
||||
memory_section = user_msg.split("## Current MEMORY.md")[1].split("## Current SOUL.md")[0]
|
||||
# No age arrow at all — we refused to annotate rather than tag the wrong line.
|
||||
assert "\u2190" not in memory_section
|
||||
|
||||
async def test_phase1_prompt_uses_threshold_from_template_var(
|
||||
self, dream, mock_provider, mock_runner, store,
|
||||
):
|
||||
"""System prompt should reference the stale-threshold constant, not a hardcoded 14."""
|
||||
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()
|
||||
|
||||
system_msg = mock_provider.chat_with_retry.call_args.kwargs["messages"][0]["content"]
|
||||
# The template renders with stale_threshold_days=14 → LLM must see "N>14"
|
||||
await loop._process_system_message(msg)
|
||||
spec = mock_runner.run.call_args[0][0]
|
||||
system_msg = spec.initial_messages[0]["content"]
|
||||
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))
|
||||
async def test_caps_huge_memory_file(self, loop, mock_runner, store):
|
||||
store.write_memory("M" * (loop.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())
|
||||
msg = InboundMessage(
|
||||
channel="system", sender_id="dream", chat_id="dream", content=""
|
||||
)
|
||||
await loop._process_system_message(msg)
|
||||
spec = mock_runner.run.call_args[0][0]
|
||||
user_msg = spec.initial_messages[1]["content"]
|
||||
memory_section = user_msg.split("## Current MEMORY.md")[1].split(
|
||||
"## Current SOUL.md"
|
||||
)[0]
|
||||
assert len(memory_section) < loop.dream._MEMORY_FILE_MAX_CHARS + 500
|
||||
|
||||
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.
|
||||
async def test_caps_huge_history_entry(self, loop, mock_runner, store):
|
||||
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",
|
||||
json.dumps(
|
||||
{
|
||||
"cursor": 1,
|
||||
"timestamp": "2026-04-01 10:00",
|
||||
"content": "H" * (loop.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())
|
||||
msg = InboundMessage(
|
||||
channel="system", sender_id="dream", chat_id="dream", content=""
|
||||
)
|
||||
await loop._process_system_message(msg)
|
||||
spec = mock_runner.run.call_args[0][0]
|
||||
user_msg = spec.initial_messages[1]["content"]
|
||||
history_section = user_msg.split("## Conversation History\n")[1].split(
|
||||
"\n\n## Current Date"
|
||||
)[0]
|
||||
assert len(history_section) < loop.dream._HISTORY_ENTRY_PREVIEW_MAX_CHARS + 500
|
||||
|
||||
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
|
||||
class TestDreamTools:
|
||||
def test_apply_patch_tool_registered(self, loop):
|
||||
tool = loop.dream._tools.get("apply_patch")
|
||||
assert tool is not None
|
||||
|
||||
|
||||
class TestDreamCaps:
|
||||
def test_batch_size_default_is_5(self):
|
||||
from nanobot.config.schema import DreamConfig
|
||||
|
||||
assert DreamConfig().max_batch_size == 5
|
||||
|
||||
def test_memory_cap_is_16k(self, loop):
|
||||
assert loop.dream._MEMORY_FILE_MAX_CHARS == 16_000
|
||||
|
||||
|
||||
class TestDreamSkipFiltering:
|
||||
async def test_skip_entries_removed_from_prompt(self, loop, mock_runner, store):
|
||||
store.append_history("- [skip] greeting\n- [permanent] User prefers dark mode")
|
||||
mock_runner.run = AsyncMock(return_value=_make_run_result())
|
||||
msg = InboundMessage(
|
||||
channel="system", sender_id="dream", chat_id="dream", content=""
|
||||
)
|
||||
await loop._process_system_message(msg)
|
||||
spec = mock_runner.run.call_args[0][0]
|
||||
user_msg = spec.initial_messages[1]["content"]
|
||||
assert "User prefers dark mode" in user_msg
|
||||
assert "[skip]" not in user_msg
|
||||
assert "greeting" not in user_msg
|
||||
|
||||
|
||||
class TestDreamAgeAnnotations:
|
||||
async def test_prompt_includes_line_age_annotations(self, loop, mock_runner, store):
|
||||
store.append_history("some event")
|
||||
mock_runner.run = AsyncMock(return_value=_make_run_result())
|
||||
store.git.init()
|
||||
store.git.auto_commit("initial memory state")
|
||||
msg = InboundMessage(
|
||||
channel="system", sender_id="dream", chat_id="dream", content=""
|
||||
)
|
||||
await loop._process_system_message(msg)
|
||||
spec = mock_runner.run.call_args[0][0]
|
||||
user_msg = spec.initial_messages[1]["content"]
|
||||
assert "## Current MEMORY.md" in user_msg
|
||||
|
||||
async def test_annotates_only_memory_not_soul_or_user(self, loop, mock_runner, store):
|
||||
store.append_history("some event")
|
||||
mock_runner.run = AsyncMock(return_value=_make_run_result())
|
||||
store.git.init()
|
||||
store.git.auto_commit("initial state")
|
||||
msg = InboundMessage(
|
||||
channel="system", sender_id="dream", chat_id="dream", content=""
|
||||
)
|
||||
await loop._process_system_message(msg)
|
||||
spec = mock_runner.run.call_args[0][0]
|
||||
user_msg = spec.initial_messages[1]["content"]
|
||||
soul_section = user_msg.split("## Current SOUL.md")[1].split(
|
||||
"## Current USER.md"
|
||||
)[0]
|
||||
user_section = user_msg.split("## Current USER.md")[1]
|
||||
assert "←" not in soul_section
|
||||
assert "←" not in user_section
|
||||
|
||||
async def test_prompt_works_without_git(self, loop, mock_runner, store):
|
||||
store.append_history("some event")
|
||||
mock_runner.run = AsyncMock(return_value=_make_run_result())
|
||||
msg = InboundMessage(
|
||||
channel="system", sender_id="dream", chat_id="dream", content=""
|
||||
)
|
||||
await loop._process_system_message(msg)
|
||||
mock_runner.run.assert_called_once()
|
||||
spec = mock_runner.run.call_args[0][0]
|
||||
user_msg = spec.initial_messages[1]["content"]
|
||||
assert "## Current MEMORY.md" in user_msg
|
||||
|
||||
async def test_prompt_carries_age_suffix_for_stale_lines(self, loop, mock_runner, store):
|
||||
store.write_memory(
|
||||
"# Memory\n- Project X active\n- fresh item\n- edge case line"
|
||||
)
|
||||
store.append_history("some event")
|
||||
mock_runner.run = AsyncMock(return_value=_make_run_result())
|
||||
fake_ages = [
|
||||
LineAge(age_days=30),
|
||||
LineAge(age_days=20),
|
||||
LineAge(age_days=14),
|
||||
LineAge(age_days=5),
|
||||
]
|
||||
with patch.object(loop.dream.store.git, "line_ages", return_value=fake_ages):
|
||||
msg = InboundMessage(
|
||||
channel="system", sender_id="dream", chat_id="dream", content=""
|
||||
)
|
||||
await loop._process_system_message(msg)
|
||||
spec = mock_runner.run.call_args[0][0]
|
||||
user_msg = spec.initial_messages[1]["content"]
|
||||
memory_section = user_msg.split("## Current MEMORY.md")[1].split(
|
||||
"## Current SOUL.md"
|
||||
)[0]
|
||||
assert "← 30d" in memory_section
|
||||
assert "← 20d" in memory_section
|
||||
assert "← 14d" not in memory_section
|
||||
assert "← 5d" not in memory_section
|
||||
|
||||
async def test_skips_annotation_when_disabled(self, loop, mock_runner, store):
|
||||
store.append_history("some event")
|
||||
mock_runner.run = AsyncMock(return_value=_make_run_result())
|
||||
loop.dream.annotate_line_ages = False
|
||||
with patch.object(loop.dream.store.git, "line_ages") as mock_line_ages:
|
||||
msg = InboundMessage(
|
||||
channel="system", sender_id="dream", chat_id="dream", content=""
|
||||
)
|
||||
await loop._process_system_message(msg)
|
||||
mock_line_ages.assert_not_called()
|
||||
spec = mock_runner.run.call_args[0][0]
|
||||
user_msg = spec.initial_messages[1]["content"]
|
||||
assert "←" not in user_msg
|
||||
|
||||
async def test_skips_annotation_on_line_ages_length_mismatch(self, loop, mock_runner, store):
|
||||
store.append_history("some event")
|
||||
mock_runner.run = AsyncMock(return_value=_make_run_result())
|
||||
with patch.object(
|
||||
loop.dream.store.git, "line_ages", return_value=[LineAge(age_days=999)]
|
||||
):
|
||||
msg = InboundMessage(
|
||||
channel="system", sender_id="dream", chat_id="dream", content=""
|
||||
)
|
||||
await loop._process_system_message(msg)
|
||||
spec = mock_runner.run.call_args[0][0]
|
||||
user_msg = spec.initial_messages[1]["content"]
|
||||
memory_section = user_msg.split("## Current MEMORY.md")[1].split(
|
||||
"## Current SOUL.md"
|
||||
)[0]
|
||||
assert "←" not in memory_section
|
||||
|
||||
|
||||
class TestDreamSessionPersistence:
|
||||
async def test_writes_session_on_success(self, loop, mock_runner, store):
|
||||
store.append_history("event one")
|
||||
store.append_history("event two")
|
||||
mock_runner.run = AsyncMock(
|
||||
return_value=_make_run_result(
|
||||
tool_events=[
|
||||
{"name": "edit_file", "status": "ok", "detail": "memory/MEMORY.md"}
|
||||
],
|
||||
)
|
||||
)
|
||||
msg = InboundMessage(
|
||||
channel="system", sender_id="dream", chat_id="dream", content=""
|
||||
)
|
||||
await loop._process_system_message(msg)
|
||||
session_path = store.memory_dir / ".dream_session.json"
|
||||
assert session_path.exists()
|
||||
data = json.loads(session_path.read_text(encoding="utf-8"))
|
||||
assert data["batch"]["from_cursor"] == 0
|
||||
assert data["batch"]["to_cursor"] == 2
|
||||
assert data["batch"]["count"] == 2
|
||||
assert data["stop_reason"] == "completed"
|
||||
assert data["changelog"] == ["edit_file: memory/MEMORY.md"]
|
||||
assert "timestamp" in data
|
||||
assert "elapsed_seconds" in data
|
||||
assert "messages" in data
|
||||
|
||||
async def test_no_session_record_on_failure(self, loop, mock_runner, store):
|
||||
"""Failed batch should not write a session record (cursor stays put for retry)."""
|
||||
store.append_history("event one")
|
||||
mock_runner.run = AsyncMock(side_effect=RuntimeError("LLM error"))
|
||||
msg = InboundMessage(
|
||||
channel="system", sender_id="dream", chat_id="dream", content=""
|
||||
)
|
||||
await loop._process_system_message(msg)
|
||||
session_path = store.memory_dir / ".dream_session.json"
|
||||
assert not session_path.exists()
|
||||
assert store.get_last_dream_cursor() == 0
|
||||
|
||||
async def test_session_contains_full_messages(self, loop, mock_runner, store):
|
||||
store.append_history("event one")
|
||||
messages = [
|
||||
{"role": "system", "content": "you are a memory bot"},
|
||||
{"role": "user", "content": "history here"},
|
||||
{"role": "assistant", "content": "I will edit MEMORY.md"},
|
||||
]
|
||||
result = _make_run_result()
|
||||
result.messages = messages
|
||||
mock_runner.run = AsyncMock(return_value=result)
|
||||
msg = InboundMessage(
|
||||
channel="system", sender_id="dream", chat_id="dream", content=""
|
||||
)
|
||||
await loop._process_system_message(msg)
|
||||
session_path = store.memory_dir / ".dream_session.json"
|
||||
data = json.loads(session_path.read_text(encoding="utf-8"))
|
||||
assert data["messages"] == messages
|
||||
assert data["prompt_chars"] > 0
|
||||
assert data["commit_sha"] is None
|
||||
|
||||
@@ -0,0 +1,125 @@
|
||||
"""Tests for memory system: Consolidator, token estimation, truncation."""
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from nanobot.agent.memory import _TIKTOKEN_ENC, Consolidator, MemoryStore, _estimate_tokens
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def store(tmp_path):
|
||||
s = MemoryStore(tmp_path)
|
||||
s.write_soul("# Soul\n- Helpful")
|
||||
s.write_user("# User\n- Developer")
|
||||
s.write_memory("# Memory\n- Project X active")
|
||||
return s
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_provider():
|
||||
p = MagicMock()
|
||||
p.chat_with_retry = AsyncMock()
|
||||
p.generation.max_tokens = 4096
|
||||
return p
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_sessions():
|
||||
return MagicMock()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_build_messages():
|
||||
return MagicMock(return_value=[])
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_get_tool_definitions():
|
||||
return MagicMock(return_value=[])
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def consolidator(store, mock_provider, mock_sessions, mock_build_messages, mock_get_tool_definitions):
|
||||
return Consolidator(
|
||||
store=store,
|
||||
provider=mock_provider,
|
||||
model="test-model",
|
||||
sessions=mock_sessions,
|
||||
context_window_tokens=128_000,
|
||||
build_messages=mock_build_messages,
|
||||
get_tool_definitions=mock_get_tool_definitions,
|
||||
)
|
||||
|
||||
|
||||
class TestEstimateTokens:
|
||||
def test_estimate_tokens_returns_positive(self):
|
||||
assert _estimate_tokens("hello world") > 0
|
||||
|
||||
def test_estimate_tokens_english_approximate(self):
|
||||
# English is roughly 1 token per 4 chars as fallback
|
||||
text = "a " * 100
|
||||
if _TIKTOKEN_ENC is not None:
|
||||
expected = len(_TIKTOKEN_ENC.encode(text))
|
||||
else:
|
||||
expected = len(text) // 4
|
||||
assert _estimate_tokens(text) == expected
|
||||
|
||||
|
||||
class TestTruncateToTokenBudget:
|
||||
def test_reserve_tokens_reduces_budget(self, consolidator):
|
||||
long_text = "word " * 200_000
|
||||
# Without reserve, more text survives
|
||||
no_reserve = consolidator._truncate_to_token_budget(long_text, reserve_tokens=0)
|
||||
with_reserve = consolidator._truncate_to_token_budget(long_text, reserve_tokens=500)
|
||||
assert len(with_reserve) < len(no_reserve)
|
||||
|
||||
def test_reserve_tokens_zero_default(self, consolidator):
|
||||
text = "hello world"
|
||||
result = consolidator._truncate_to_token_budget(text)
|
||||
assert result == text
|
||||
|
||||
|
||||
class TestConsolidatorPrompt:
|
||||
def test_prompt_contains_snip(self):
|
||||
from nanobot.utils.prompt_templates import render_template
|
||||
text = render_template("agent/consolidator_archive.md", strip=True)
|
||||
assert "SNIP" in text
|
||||
assert "[permanent]" in text
|
||||
assert "[skip]" in text
|
||||
|
||||
|
||||
class TestConsolidatorArchive:
|
||||
async def test_archive_injects_dedup_context(self, consolidator, mock_provider, store):
|
||||
store.write_memory("- User prefers dark mode")
|
||||
store.write_user("- Developer")
|
||||
messages = [{"role": "user", "content": "hello", "timestamp": "2026-01-01 10:00"}]
|
||||
|
||||
mock_provider.chat_with_retry.return_value = MagicMock(
|
||||
content="(nothing)", finish_reason="stop"
|
||||
)
|
||||
await consolidator.archive(messages)
|
||||
|
||||
call_args = mock_provider.chat_with_retry.call_args
|
||||
user_msg = call_args.kwargs["messages"][1]["content"]
|
||||
assert "## Current MEMORY.md (for dedup)" in user_msg
|
||||
assert "User prefers dark mode" in user_msg
|
||||
assert "## Current USER.md (for dedup)" in user_msg
|
||||
assert "Developer" in user_msg
|
||||
|
||||
async def test_archive_skips_dedup_when_budget_exhausted(self, consolidator, mock_provider, store):
|
||||
# Shrink token budget so dedup context (always capped at ~6000 chars)
|
||||
# exceeds the available room.
|
||||
consolidator.context_window_tokens = 6_000
|
||||
store.write_memory("word " * 10_000)
|
||||
messages = [{"role": "user", "content": "hello", "timestamp": "2026-01-01 10:00"}]
|
||||
|
||||
mock_provider.chat_with_retry.return_value = MagicMock(
|
||||
content="(nothing)", finish_reason="stop"
|
||||
)
|
||||
await consolidator.archive(messages)
|
||||
|
||||
call_args = mock_provider.chat_with_retry.call_args
|
||||
user_msg = call_args.kwargs["messages"][1]["content"]
|
||||
# Should not contain dedup context when budget is exhausted
|
||||
assert "## Current MEMORY.md (for dedup)" not in user_msg
|
||||
@@ -292,3 +292,95 @@ def test_from_config_static_preset_loader_does_not_enable_hot_reload(tmp_path) -
|
||||
loop = AgentLoop.from_config(config)
|
||||
assert loop._provider_snapshot_loader is None
|
||||
assert loop._preset_snapshot_loader is not None
|
||||
|
||||
|
||||
class TestDreamModelOverride:
|
||||
def test_dream_follows_main_when_no_override(self, tmp_path) -> None:
|
||||
provider = _provider("base-model")
|
||||
loop = AgentLoop(
|
||||
bus=MessageBus(),
|
||||
provider=provider,
|
||||
workspace=tmp_path,
|
||||
model="base-model",
|
||||
context_window_tokens=1000,
|
||||
)
|
||||
assert loop.dream.model == "base-model"
|
||||
assert loop.dream.provider is provider
|
||||
|
||||
def test_dream_raw_model_override(self, tmp_path) -> None:
|
||||
provider = _provider("base-model")
|
||||
loop = AgentLoop(
|
||||
bus=MessageBus(),
|
||||
provider=provider,
|
||||
workspace=tmp_path,
|
||||
model="base-model",
|
||||
context_window_tokens=1000,
|
||||
dream_model_override="custom-model-v2",
|
||||
)
|
||||
assert loop.dream.model == "custom-model-v2"
|
||||
assert loop.dream.provider is provider
|
||||
|
||||
def test_dream_preset_override(self, tmp_path) -> None:
|
||||
cheap_provider = _provider("openai/gpt-4.1-mini", max_tokens=2048)
|
||||
preset = ModelPresetConfig(
|
||||
model="openai/gpt-4.1-mini",
|
||||
provider="openai",
|
||||
max_tokens=2048,
|
||||
context_window_tokens=128_000,
|
||||
)
|
||||
loop = AgentLoop(
|
||||
bus=MessageBus(),
|
||||
provider=_provider("base-model"),
|
||||
workspace=tmp_path,
|
||||
model="base-model",
|
||||
context_window_tokens=1000,
|
||||
model_presets={"cheap": preset},
|
||||
dream_model_override="cheap",
|
||||
preset_snapshot_loader=lambda _name: ProviderSnapshot(
|
||||
provider=cheap_provider,
|
||||
model=preset.model,
|
||||
context_window_tokens=preset.context_window_tokens,
|
||||
signature=("cheap", preset.model),
|
||||
),
|
||||
)
|
||||
assert loop.dream.model == "openai/gpt-4.1-mini"
|
||||
assert loop.dream.provider is cheap_provider
|
||||
assert loop.dream._runner.provider is cheap_provider
|
||||
|
||||
def test_dream_override_survives_main_preset_switch(self, tmp_path) -> None:
|
||||
base_provider = _provider("base-model")
|
||||
fast_provider = _provider("openai/gpt-4.1", max_tokens=4096)
|
||||
cheap_provider = _provider("openai/gpt-4.1-mini", max_tokens=2048)
|
||||
loop = AgentLoop(
|
||||
bus=MessageBus(),
|
||||
provider=base_provider,
|
||||
workspace=tmp_path,
|
||||
model="base-model",
|
||||
context_window_tokens=1000,
|
||||
model_presets={
|
||||
"fast": ModelPresetConfig(model="openai/gpt-4.1"),
|
||||
"cheap": ModelPresetConfig(model="openai/gpt-4.1-mini"),
|
||||
},
|
||||
dream_model_override="cheap",
|
||||
preset_snapshot_loader=lambda name: ProviderSnapshot(
|
||||
provider=fast_provider if name == "fast" else cheap_provider,
|
||||
model="openai/gpt-4.1" if name == "fast" else "openai/gpt-4.1-mini",
|
||||
context_window_tokens=32_768 if name == "fast" else 128_000,
|
||||
signature=(name, "model"),
|
||||
),
|
||||
)
|
||||
# Initially dream is on cheap
|
||||
assert loop.dream.model == "openai/gpt-4.1-mini"
|
||||
assert loop.dream.provider is cheap_provider
|
||||
|
||||
# Switch main preset to fast
|
||||
loop.set_model_preset("fast")
|
||||
|
||||
# Main agent should be on fast
|
||||
assert loop.model == "openai/gpt-4.1"
|
||||
assert loop.provider is fast_provider
|
||||
|
||||
# Dream should still be on cheap override
|
||||
assert loop.dream.model == "openai/gpt-4.1-mini"
|
||||
assert loop.dream.provider is cheap_provider
|
||||
assert loop.dream._runner.provider is cheap_provider
|
||||
|
||||
Reference in New Issue
Block a user