Merge branch 'main' of https://github.com/HKUDS/nanobot into codex/webui-performance

This commit is contained in:
Xubin Ren
2026-05-18 01:18:28 +08:00
8 changed files with 635 additions and 363 deletions
+140 -163
View File
@@ -45,6 +45,73 @@ def _add_turns(session, turns: int, *, prefix: str = "msg") -> None:
session.add_message("assistant", f"{prefix} assistant {i}")
def _make_fake_compact(
loop: AgentLoop,
*,
summary: str = "Summary.",
on_archive=None,
track_archived: list | None = None,
track_count: bool = False,
):
"""Return a fake compact_idle_session that mirrors the real method's session mutation."""
from nanobot.session.manager import Session as _Session
state = {"count": 0}
async def _fake_compact(key: str, max_suffix: int = 8) -> str:
state["count"] += 1
session = loop.sessions.get_or_create(key)
tail = list(session.messages[session.last_consolidated:])
if not tail:
session.updated_at = datetime.now()
loop.sessions.save(session)
return ""
probe = _Session(
key=session.key,
messages=tail.copy(),
created_at=session.created_at,
updated_at=session.updated_at,
metadata={},
last_consolidated=0,
)
probe.retain_recent_legal_suffix(max_suffix)
kept = probe.messages
cut = len(tail) - len(kept)
archive_msgs = tail[:cut]
if not archive_msgs and not kept:
session.updated_at = datetime.now()
loop.sessions.save(session)
return ""
last_active = session.updated_at
s = summary
if archive_msgs:
if on_archive:
result = on_archive(archive_msgs)
s = result if isinstance(result, str) else summary
if track_archived is not None:
track_archived.extend(archive_msgs)
if s and s != "(nothing)":
session.metadata["_last_summary"] = {
"text": s,
"last_active": last_active.isoformat(),
}
session.messages = kept
session.last_consolidated = 0
session.updated_at = datetime.now()
loop.sessions.save(session)
return s
# Attach state for count access
_fake_compact.state = state # type: ignore[attr-defined]
return _fake_compact
class TestSessionTTLConfig:
"""Test session TTL configuration."""
@@ -201,10 +268,7 @@ class TestAutoCompact:
s2.add_message("user", "recent")
loop.sessions.save(s2)
async def _fake_archive(messages):
return "Summary."
loop.consolidator.archive = _fake_archive
loop.consolidator.compact_idle_session = _make_fake_compact(loop)
loop.auto_compact.check_expired(loop._schedule_background)
await asyncio.sleep(0.1)
@@ -222,12 +286,9 @@ class TestAutoCompact:
loop.sessions.save(session)
archived_messages = []
async def _fake_archive(messages):
archived_messages.extend(messages)
return "Summary."
loop.consolidator.archive = _fake_archive
loop.consolidator.compact_idle_session = _make_fake_compact(
loop, track_archived=archived_messages,
)
await loop.auto_compact._archive("cli:test")
@@ -246,10 +307,9 @@ class TestAutoCompact:
_add_turns(session, 6, prefix="hello")
loop.sessions.save(session)
async def _fake_archive(messages):
return "User said hello."
loop.consolidator.archive = _fake_archive
loop.consolidator.compact_idle_session = _make_fake_compact(
loop, summary="User said hello.",
)
await loop.auto_compact._archive("cli:test")
@@ -262,23 +322,16 @@ class TestAutoCompact:
@pytest.mark.asyncio
async def test_auto_compact_empty_session(self, tmp_path):
"""_archive on empty session should not archive."""
"""_archive on empty session should not store a summary."""
loop = _make_loop(tmp_path, session_ttl_minutes=15)
archive_called = False
async def _fake_archive(messages):
nonlocal archive_called
archive_called = True
return "Summary."
loop.consolidator.archive = _fake_archive
loop.consolidator.compact_idle_session = _make_fake_compact(loop)
await loop.auto_compact._archive("cli:test")
assert not archive_called
session_after = loop.sessions.get_or_create("cli:test")
assert len(session_after.messages) == 0
assert "cli:test" not in loop.auto_compact._summaries
await loop.close_mcp()
@pytest.mark.asyncio
@@ -290,18 +343,14 @@ class TestAutoCompact:
session.last_consolidated = 18
loop.sessions.save(session)
archived_count = 0
async def _fake_archive(messages):
nonlocal archived_count
archived_count = len(messages)
return "Summary."
loop.consolidator.archive = _fake_archive
archived_messages = []
loop.consolidator.compact_idle_session = _make_fake_compact(
loop, track_archived=archived_messages,
)
await loop.auto_compact._archive("cli:test")
assert archived_count == 2
assert len(archived_messages) == 2
await loop.close_mcp()
@@ -334,12 +383,9 @@ class TestAutoCompactIdleDetection:
loop.sessions.save(session)
archived_messages = []
async def _fake_archive(messages):
archived_messages.extend(messages)
return "Summary."
loop.consolidator.archive = _fake_archive
loop.consolidator.compact_idle_session = _make_fake_compact(
loop, track_archived=archived_messages,
)
# Simulate proactive archive completing before message arrives
await loop.auto_compact._archive("cli:test")
@@ -402,10 +448,7 @@ class TestAutoCompactIdleDetection:
session.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(session)
async def _fake_archive(messages):
return "Summary."
loop.consolidator.archive = _fake_archive
loop.consolidator.compact_idle_session = _make_fake_compact(loop)
msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="/new")
response = await loop._process_message(msg)
@@ -466,10 +509,7 @@ class TestAutoCompactSystemMessages:
session.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(session)
async def _fake_archive(messages):
return "Summary."
loop.consolidator.archive = _fake_archive
loop.consolidator.compact_idle_session = _make_fake_compact(loop)
# Simulate proactive archive completing before system message arrives
await loop.auto_compact._archive("cli:test")
@@ -547,12 +587,9 @@ class TestAutoCompactEdgeCases:
loop.sessions.save(session)
archived_messages = []
async def _fake_archive(messages):
archived_messages.extend(messages)
return "Summary."
loop.consolidator.archive = _fake_archive
loop.consolidator.compact_idle_session = _make_fake_compact(
loop, track_archived=archived_messages,
)
# Simulate proactive archive completing before message arrives
await loop.auto_compact._archive("cli:test")
@@ -644,10 +681,7 @@ class TestAutoCompactIntegration:
session.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(session)
async def _fake_archive(messages):
return "Summary."
loop.consolidator.archive = _fake_archive
loop.consolidator.compact_idle_session = _make_fake_compact(loop)
# Simulate proactive archive completing before message arrives
await loop.auto_compact._archive("cli:test")
@@ -704,12 +738,9 @@ class TestProactiveAutoCompact:
loop.sessions.save(session)
archived_messages = []
async def _fake_archive(messages):
archived_messages.extend(messages)
return "User chatted about old things."
loop.consolidator.archive = _fake_archive
loop.consolidator.compact_idle_session = _make_fake_compact(
loop, summary="User chatted about old things.", track_archived=archived_messages,
)
await self._run_check_expired(loop)
@@ -748,14 +779,14 @@ class TestProactiveAutoCompact:
started = asyncio.Event()
block_forever = asyncio.Event()
async def _slow_archive(messages):
async def _slow_compact(key, max_suffix=8):
nonlocal archive_count
archive_count += 1
started.set()
await block_forever.wait()
return "Summary."
loop.consolidator.archive = _slow_archive
loop.consolidator.compact_idle_session = _slow_compact
# First call starts archiving via callback
loop.auto_compact.check_expired(loop._schedule_background)
@@ -781,10 +812,10 @@ class TestProactiveAutoCompact:
session.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(session)
async def _failing_archive(messages):
async def _failing_compact(key, max_suffix=8):
raise RuntimeError("LLM down")
loop.consolidator.archive = _failing_archive
loop.consolidator.compact_idle_session = _failing_compact
# Should not raise
await self._run_check_expired(loop)
@@ -795,24 +826,18 @@ class TestProactiveAutoCompact:
@pytest.mark.asyncio
async def test_proactive_archive_skips_empty_sessions(self, tmp_path):
"""Proactive archive should not call LLM for sessions with no un-consolidated messages."""
"""Proactive archive should not produce a summary for sessions with no messages."""
loop = _make_loop(tmp_path, session_ttl_minutes=15)
session = loop.sessions.get_or_create("cli:test")
session.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(session)
archive_called = False
async def _fake_archive(messages):
nonlocal archive_called
archive_called = True
return "Summary."
loop.consolidator.archive = _fake_archive
loop.consolidator.compact_idle_session = _make_fake_compact(loop)
await self._run_check_expired(loop)
assert not archive_called
# Empty session should not produce a summary
assert "cli:test" not in loop.auto_compact._summaries
await loop.close_mcp()
@pytest.mark.asyncio
@@ -824,18 +849,12 @@ class TestProactiveAutoCompact:
session.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(session)
archive_count = 0
async def _fake_archive(messages):
nonlocal archive_count
archive_count += 1
return "Summary."
loop.consolidator.archive = _fake_archive
_fake_compact = _make_fake_compact(loop)
loop.consolidator.compact_idle_session = _fake_compact
# Simulate an active agent task for this session
await self._run_check_expired(loop, active_session_keys={"cli:test"})
assert archive_count == 0
assert _fake_compact.state["count"] == 0
session_after = loop.sessions.get_or_create("cli:test")
assert len(session_after.messages) == 12 # All messages preserved
@@ -851,22 +870,16 @@ class TestProactiveAutoCompact:
session.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(session)
archive_count = 0
async def _fake_archive(messages):
nonlocal archive_count
archive_count += 1
return "Summary."
loop.consolidator.archive = _fake_archive
_fake_compact = _make_fake_compact(loop)
loop.consolidator.compact_idle_session = _fake_compact
# First tick: active task, skip
await self._run_check_expired(loop, active_session_keys={"cli:test"})
assert archive_count == 0
assert _fake_compact.state["count"] == 0
# Second tick: task completed, should archive
await self._run_check_expired(loop)
assert archive_count == 1
assert _fake_compact.state["count"] == 1
await loop.close_mcp()
@pytest.mark.asyncio
@@ -888,18 +901,12 @@ class TestProactiveAutoCompact:
s3.add_message("user", "recent")
loop.sessions.save(s3)
archive_count = 0
async def _fake_archive(messages):
nonlocal archive_count
archive_count += 1
return "Summary."
loop.consolidator.archive = _fake_archive
_fake_compact = _make_fake_compact(loop)
loop.consolidator.compact_idle_session = _fake_compact
await self._run_check_expired(loop, active_session_keys={"cli:expired_active"})
assert archive_count == 1
assert _fake_compact.state["count"] == 1
s1_after = loop.sessions.get_or_create("cli:expired_idle")
assert len(s1_after.messages) == loop.auto_compact._RECENT_SUFFIX_MESSAGES
s2_after = loop.sessions.get_or_create("cli:expired_active")
@@ -917,22 +924,16 @@ class TestProactiveAutoCompact:
session.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(session)
archive_count = 0
async def _fake_archive(messages):
nonlocal archive_count
archive_count += 1
return "Summary."
loop.consolidator.archive = _fake_archive
_fake_compact = _make_fake_compact(loop)
loop.consolidator.compact_idle_session = _fake_compact
# First tick: archives the session
await self._run_check_expired(loop)
assert archive_count == 1
assert _fake_compact.state["count"] == 1
# Second tick: should NOT re-schedule (updated_at is fresh after clear)
await self._run_check_expired(loop)
assert archive_count == 1 # Still 1, not re-scheduled
assert _fake_compact.state["count"] == 1 # Still 1, not re-scheduled
await loop.close_mcp()
@pytest.mark.asyncio
@@ -943,22 +944,15 @@ class TestProactiveAutoCompact:
session.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(session)
archive_count = 0
async def _fake_archive(messages):
nonlocal archive_count
archive_count += 1
return "Summary."
loop.consolidator.archive = _fake_archive
loop.consolidator.compact_idle_session = _make_fake_compact(loop)
# First tick: skips (no messages), refreshes updated_at
await self._run_check_expired(loop)
assert archive_count == 0
assert "cli:test" not in loop.auto_compact._summaries
# Second tick: should NOT re-schedule because updated_at is fresh
await self._run_check_expired(loop)
assert archive_count == 0
assert "cli:test" not in loop.auto_compact._summaries
await loop.close_mcp()
@pytest.mark.asyncio
@@ -970,18 +964,12 @@ class TestProactiveAutoCompact:
session.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(session)
archive_count = 0
async def _fake_archive(messages):
nonlocal archive_count
archive_count += 1
return "Summary."
loop.consolidator.archive = _fake_archive
_fake_compact = _make_fake_compact(loop)
loop.consolidator.compact_idle_session = _fake_compact
# First compact cycle
await loop.auto_compact._archive("cli:test")
assert archive_count == 1
assert _fake_compact.state["count"] == 1
# User returns, sends new messages
msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="second topic")
@@ -995,7 +983,7 @@ class TestProactiveAutoCompact:
# Second compact cycle should succeed
await loop.auto_compact._archive("cli:test")
assert archive_count == 2
assert _fake_compact.state["count"] == 2
await loop.close_mcp()
@@ -1011,10 +999,9 @@ class TestSummaryPersistence:
session.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(session)
async def _fake_archive(messages):
return "User said hello."
loop.consolidator.archive = _fake_archive
loop.consolidator.compact_idle_session = _make_fake_compact(
loop, summary="User said hello.",
)
await loop.auto_compact._archive("cli:test")
@@ -1036,10 +1023,9 @@ class TestSummaryPersistence:
session.updated_at = last_active
loop.sessions.save(session)
async def _fake_archive(messages):
return "User said hello."
loop.consolidator.archive = _fake_archive
loop.consolidator.compact_idle_session = _make_fake_compact(
loop, summary="User said hello.",
)
# Archive
await loop.auto_compact._archive("cli:test")
@@ -1069,10 +1055,7 @@ class TestSummaryPersistence:
session.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(session)
async def _fake_archive(messages):
return "Summary."
loop.consolidator.archive = _fake_archive
loop.consolidator.compact_idle_session = _make_fake_compact(loop)
await loop.auto_compact._archive("cli:test")
@@ -1100,10 +1083,7 @@ class TestSummaryPersistence:
session.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(session)
async def _fake_archive(messages):
return "Summary."
loop.consolidator.archive = _fake_archive
loop.consolidator.compact_idle_session = _make_fake_compact(loop)
await loop.auto_compact._archive("cli:test")
@@ -1129,10 +1109,9 @@ class TestSummaryPersistence:
session.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(session)
async def _fake_archive(messages):
return "First summary."
loop.consolidator.archive = _fake_archive
loop.consolidator.compact_idle_session = _make_fake_compact(
loop, summary="First summary.",
)
await loop.auto_compact._archive("cli:test")
# Consume the first summary via hot path
@@ -1148,10 +1127,9 @@ class TestSummaryPersistence:
session.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(session)
async def _fake_archive2(messages):
return "Second summary."
loop.consolidator.archive = _fake_archive2
loop.consolidator.compact_idle_session = _make_fake_compact(
loop, summary="Second summary.",
)
await loop.auto_compact._archive("cli:test")
# The second archive writes a new summary
@@ -1173,10 +1151,9 @@ class TestSummaryPersistence:
session.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(session)
async def _fake_archive(messages):
return "Old summary."
loop.consolidator.archive = _fake_archive
loop.consolidator.compact_idle_session = _make_fake_compact(
loop, summary="Old summary.",
)
await loop.auto_compact._archive("cli:test")
# Verify summary exists before /new
+32 -143
View File
@@ -38,7 +38,7 @@ def _make_autocompact(
sessions = MagicMock(spec=SessionManager)
if consolidator is None:
consolidator = MagicMock()
consolidator.archive = AsyncMock(return_value="Summary.")
consolidator.compact_idle_session = AsyncMock(return_value="Summary.")
return AutoCompact(
sessions=sessions,
consolidator=consolidator,
@@ -178,62 +178,6 @@ class TestFormatSummary:
assert result.startswith("Previous conversation summary (last active ")
# ---------------------------------------------------------------------------
# _split_unconsolidated
# ---------------------------------------------------------------------------
class TestSplitUnconsolidated:
"""Test AutoCompact._split_unconsolidated splitting logic."""
def test_empty_session_returns_both_empty(self):
"""Empty session should return ([], [])."""
ac = _make_autocompact()
session = _make_session(messages=[])
archive, kept = ac._split_unconsolidated(session)
assert archive == []
assert kept == []
def test_all_messages_archivable_when_more_than_suffix(self):
"""Session with many messages should archive a prefix and keep suffix."""
ac = _make_autocompact()
msgs = [{"role": "user", "content": f"u{i}"} for i in range(20)]
session = _make_session(messages=msgs)
archive, kept = ac._split_unconsolidated(session)
assert len(archive) > 0
assert len(kept) <= AutoCompact._RECENT_SUFFIX_MESSAGES
def test_fewer_messages_than_suffix_returns_empty_archive(self):
"""Session with fewer messages than suffix should have empty archive."""
ac = _make_autocompact()
msgs = [{"role": "user", "content": f"u{i}"} for i in range(3)]
session = _make_session(messages=msgs)
archive, kept = ac._split_unconsolidated(session)
assert archive == []
assert len(kept) == len(msgs)
def test_respects_last_consolidated_offset(self):
"""Only messages after last_consolidated should be considered."""
ac = _make_autocompact()
msgs = [{"role": "user", "content": f"u{i}"} for i in range(20)]
# First 10 are already consolidated
session = _make_session(messages=msgs, last_consolidated=10)
archive, kept = ac._split_unconsolidated(session)
# Only the tail of 10 messages is considered for splitting
assert all(m["content"] in [f"u{i}" for i in range(10, 20)] for m in kept)
assert all(m["content"] in [f"u{i}" for i in range(10, 20)] for m in archive)
def test_retain_recent_legal_suffix_keeps_last_n(self):
"""The kept suffix should be at most _RECENT_SUFFIX_MESSAGES long."""
ac = _make_autocompact()
# 20 user messages = 20 messages total, all after last_consolidated=0
msgs = [{"role": "user", "content": f"u{i}"} for i in range(20)]
session = _make_session(messages=msgs)
archive, kept = ac._split_unconsolidated(session)
assert len(kept) <= AutoCompact._RECENT_SUFFIX_MESSAGES
assert len(archive) == len(msgs) - len(kept)
# ---------------------------------------------------------------------------
# check_expired
# ---------------------------------------------------------------------------
@@ -313,126 +257,71 @@ class TestCheckExpired:
# ---------------------------------------------------------------------------
class TestArchive:
"""Test AutoCompact._archive async method."""
class TestArchiveDelegates:
"""_archive should delegate all session mutation to Consolidator."""
@pytest.mark.asyncio
async def test_empty_session_updates_timestamp_no_archive_call(self):
"""Empty session should refresh updated_at and not call consolidator.archive."""
async def test_calls_compact_idle_session(self):
ac = _make_autocompact()
mock_sm = MagicMock(spec=SessionManager)
empty_session = _make_session(messages=[])
mock_sm.get_or_create.return_value = empty_session
ac.sessions = mock_sm
ac.consolidator.archive = AsyncMock(return_value="Summary.")
ac.consolidator.compact_idle_session = AsyncMock(return_value="Summary.")
await ac._archive("cli:test")
ac.consolidator.archive.assert_not_called()
mock_sm.save.assert_called_once_with(empty_session)
# updated_at was refreshed
assert empty_session.updated_at > datetime.now() - timedelta(seconds=5)
ac.consolidator.compact_idle_session.assert_awaited_once_with(
"cli:test", ac._RECENT_SUFFIX_MESSAGES,
)
@pytest.mark.asyncio
async def test_archive_returns_empty_string_no_summary_stored(self):
"""If archive returns empty string, no summary should be stored."""
async def test_populates_summaries_from_metadata(self):
ac = _make_autocompact()
mock_sm = MagicMock(spec=SessionManager)
msgs = [{"role": "user", "content": f"u{i}"} for i in range(20)]
session = _make_session(messages=msgs)
session = _make_session(
metadata={"_last_summary": {"text": "Hello.", "last_active": "2026-05-13T10:00:00"}}
)
mock_sm.get_or_create.return_value = session
ac.sessions = mock_sm
ac.consolidator.archive = AsyncMock(return_value="")
ac.consolidator.compact_idle_session = AsyncMock(return_value="Hello.")
await ac._archive("cli:test")
assert "cli:test" not in ac._summaries
@pytest.mark.asyncio
async def test_archive_returns_nothing_no_summary_stored(self):
"""If archive returns '(nothing)', no summary should be stored."""
ac = _make_autocompact()
mock_sm = MagicMock(spec=SessionManager)
msgs = [{"role": "user", "content": f"u{i}"} for i in range(20)]
session = _make_session(messages=msgs)
mock_sm.get_or_create.return_value = session
ac.sessions = mock_sm
ac.consolidator.archive = AsyncMock(return_value="(nothing)")
await ac._archive("cli:test")
assert "cli:test" not in ac._summaries
@pytest.mark.asyncio
async def test_archive_exception_caught_key_removed_from_archiving(self):
"""If archive raises, exception is caught and key removed from _archiving."""
ac = _make_autocompact()
mock_sm = MagicMock(spec=SessionManager)
msgs = [{"role": "user", "content": f"u{i}"} for i in range(20)]
session = _make_session(messages=msgs)
mock_sm.get_or_create.return_value = session
ac.sessions = mock_sm
ac.consolidator.archive = AsyncMock(side_effect=RuntimeError("LLM down"))
# Should not raise
await ac._archive("cli:test")
assert "cli:test" not in ac._archiving
@pytest.mark.asyncio
async def test_successful_archive_stores_summary_in_summaries_and_metadata(self):
"""Successful archive should store summary in _summaries dict and metadata."""
ac = _make_autocompact()
mock_sm = MagicMock(spec=SessionManager)
msgs = [{"role": "user", "content": f"u{i}"} for i in range(20)]
last_active = datetime(2026, 5, 13, 10, 0, 0)
session = _make_session(messages=msgs, updated_at=last_active)
mock_sm.get_or_create.return_value = session
ac.sessions = mock_sm
ac.consolidator.archive = AsyncMock(return_value="User discussed AI.")
await ac._archive("cli:test")
# _summaries
entry = ac._summaries.get("cli:test")
assert entry is not None
assert entry[0] == "User discussed AI."
assert entry[1] == last_active
# metadata
meta = session.metadata.get("_last_summary")
assert meta is not None
assert meta["text"] == "User discussed AI."
assert "last_active" in meta
assert entry[0] == "Hello."
@pytest.mark.asyncio
async def test_finally_block_always_removes_from_archiving(self):
"""Finally block should always remove key from _archiving, even on error."""
async def test_no_summary_when_compact_returns_empty(self):
ac = _make_autocompact()
mock_sm = MagicMock(spec=SessionManager)
msgs = [{"role": "user", "content": f"u{i}"} for i in range(20)]
session = _make_session(messages=msgs)
mock_sm.get_or_create.return_value = session
ac.sessions = mock_sm
ac.consolidator.archive = AsyncMock(side_effect=RuntimeError("fail"))
ac.consolidator.compact_idle_session = AsyncMock(return_value="")
# Pre-add key to archiving to verify it gets removed
ac._archiving.add("cli:test")
await ac._archive("cli:test")
assert "cli:test" not in ac._archiving
assert "cli:test" not in ac._summaries
@pytest.mark.asyncio
async def test_finally_removes_from_archiving_on_success(self):
"""Finally block should remove key from _archiving on success too."""
async def test_no_summary_when_compact_returns_nothing(self):
ac = _make_autocompact()
mock_sm = MagicMock(spec=SessionManager)
msgs = [{"role": "user", "content": f"u{i}"} for i in range(20)]
session = _make_session(messages=msgs)
mock_sm.get_or_create.return_value = session
ac.sessions = mock_sm
ac.consolidator.archive = AsyncMock(return_value="Summary.")
ac.consolidator.compact_idle_session = AsyncMock(return_value="(nothing)")
await ac._archive("cli:test")
assert "cli:test" not in ac._summaries
@pytest.mark.asyncio
async def test_exception_still_removes_from_archiving(self):
ac = _make_autocompact()
mock_sm = MagicMock(spec=SessionManager)
ac.sessions = mock_sm
ac.consolidator.compact_idle_session = AsyncMock(side_effect=RuntimeError("fail"))
ac._archiving.add("cli:test")
await ac._archive("cli:test")
assert "cli:test" not in ac._archiving
+267
View File
@@ -28,6 +28,12 @@ def mock_provider():
def consolidator(store, mock_provider):
sessions = MagicMock()
sessions.save = MagicMock()
# When maybe_consolidate_by_tokens refreshes the session reference via
# get_or_create(session.key), it should get back the same object the test
# passed in. Store sessions by key so the lookup is transparent.
_session_cache: dict[str, MagicMock] = {}
sessions.get_or_create = MagicMock(side_effect=lambda key: _session_cache.get(key, MagicMock()))
sessions._session_cache = _session_cache
return Consolidator(
store=store,
provider=mock_provider,
@@ -117,6 +123,7 @@ class TestConsolidatorTokenBudget:
session.last_consolidated = 0
session.messages = [{"role": "user", "content": "hi"}]
session.key = "test:key"
consolidator.sessions._session_cache[session.key] = session
consolidator.estimate_session_prompt_tokens = MagicMock(return_value=(100, "tiktoken"))
consolidator.archive = AsyncMock(return_value=True)
await consolidator.maybe_consolidate_by_tokens(session)
@@ -152,6 +159,7 @@ class TestConsolidatorTokenBudget:
session.add_message("user", f"u{i}")
session.add_message("assistant", f"a{i}")
consolidator.sessions._session_cache[session.key] = session
consolidator.estimate_session_prompt_tokens = MagicMock(return_value=(100, "tiktoken"))
consolidator.archive = AsyncMock(return_value="old conversation summary")
@@ -184,6 +192,7 @@ class TestConsolidatorTokenBudget:
session.add_message("tool", "tool result", tool_call_id="call-1", name="x")
session.add_message("assistant", "final answer")
consolidator.sessions._session_cache[session.key] = session
consolidator.estimate_session_prompt_tokens = MagicMock(return_value=(100, "tiktoken"))
consolidator.archive = AsyncMock(return_value="tool turn summary")
@@ -210,6 +219,7 @@ class TestConsolidatorTokenBudget:
}
for i in range(70)
]
consolidator.sessions._session_cache[session.key] = session
consolidator.estimate_session_prompt_tokens = MagicMock(
side_effect=[(1200, "tiktoken"), (400, "tiktoken")]
)
@@ -238,6 +248,7 @@ class TestConsolidatorTokenBudget:
for i in range(70)
]
session.metadata = {}
consolidator.sessions._session_cache[session.key] = session
consolidator.estimate_session_prompt_tokens = MagicMock(
side_effect=[(1200, "tiktoken"), (400, "tiktoken")]
)
@@ -263,6 +274,7 @@ class TestConsolidatorTokenBudget:
for i in range(70)
]
session.metadata = {}
consolidator.sessions._session_cache[session.key] = session
# Keep estimates high so the loop would otherwise run multiple rounds.
consolidator.estimate_session_prompt_tokens = MagicMock(
return_value=(1200, "tiktoken")
@@ -287,6 +299,7 @@ class TestConsolidatorTokenBudget:
}
for i in range(70)
]
consolidator.sessions._session_cache[session.key] = session
consolidator.estimate_session_prompt_tokens = MagicMock(
side_effect=[(1200, "tiktoken"), (400, "tiktoken")]
)
@@ -299,6 +312,260 @@ class TestConsolidatorTokenBudget:
assert session.last_consolidated == 61
class TestCompactIdleSession:
"""Tests for Consolidator.compact_idle_session — lock-protected idle truncation."""
@pytest.fixture
def real_consolidator(self, store, mock_provider):
"""Create a Consolidator with a real SessionManager (not a mock)."""
from nanobot.session.manager import SessionManager
sessions = SessionManager(store.workspace)
return Consolidator(
store=store,
provider=mock_provider,
model="test-model",
sessions=sessions,
context_window_tokens=1000,
build_messages=MagicMock(return_value=[]),
get_tool_definitions=MagicMock(return_value=[]),
max_completion_tokens=100,
)
@pytest.mark.asyncio
async def test_archives_prefix_keeps_suffix(self, real_consolidator, mock_provider):
"""20 user/assistant turns → compact with max_suffix=8 → messages ≤ 8,
last_consolidated=0, _last_summary stored."""
mock_provider.chat_with_retry.return_value = MagicMock(
content="Summary of old conversation.", finish_reason="stop"
)
sessions = real_consolidator.sessions
session = sessions.get_or_create("cli:test")
for i in range(20):
session.add_message("user", f"user msg {i}")
session.add_message("assistant", f"assistant msg {i}")
sessions.save(session)
result = await real_consolidator.compact_idle_session("cli:test", max_suffix=8)
assert result == "Summary of old conversation."
reloaded = sessions.get_or_create("cli:test")
assert len(reloaded.messages) <= 8
assert reloaded.last_consolidated == 0
meta = reloaded.metadata.get("_last_summary")
assert meta is not None
assert meta["text"] == "Summary of old conversation."
assert "last_active" in meta
@pytest.mark.asyncio
async def test_empty_session_refreshes_timestamp(self, real_consolidator):
"""Empty session with old updated_at → refreshed after call, returns ''."""
from datetime import datetime, timedelta
sessions = real_consolidator.sessions
session = sessions.get_or_create("cli:empty")
old_ts = datetime.now() - timedelta(hours=2)
session.updated_at = old_ts
sessions.save(session)
result = await real_consolidator.compact_idle_session("cli:empty")
assert result == ""
reloaded = sessions.get_or_create("cli:empty")
assert reloaded.updated_at > old_ts
@pytest.mark.asyncio
async def test_nothing_summary_not_stored(self, real_consolidator, mock_provider):
"""LLM returns '(nothing)' → _last_summary NOT in metadata."""
mock_provider.chat_with_retry.return_value = MagicMock(
content="(nothing)", finish_reason="stop"
)
sessions = real_consolidator.sessions
session = sessions.get_or_create("cli:nothing")
for i in range(10):
session.add_message("user", f"u{i}")
session.add_message("assistant", f"a{i}")
sessions.save(session)
result = await real_consolidator.compact_idle_session("cli:nothing", max_suffix=4)
assert result == "(nothing)"
reloaded = sessions.get_or_create("cli:nothing")
assert "_last_summary" not in reloaded.metadata
@pytest.mark.asyncio
async def test_llm_failure_still_truncates(self, real_consolidator, mock_provider, store):
"""LLM raises RuntimeError → raw_archive fires, session still truncated, returns None."""
mock_provider.chat_with_retry.side_effect = RuntimeError("LLM unavailable")
sessions = real_consolidator.sessions
session = sessions.get_or_create("cli:fail")
for i in range(10):
session.add_message("user", f"u{i}")
session.add_message("assistant", f"a{i}")
sessions.save(session)
result = await real_consolidator.compact_idle_session("cli:fail", max_suffix=4)
assert result is None
# raw_archive should have been called (history.jsonl gets an entry)
entries = store.read_unprocessed_history(since_cursor=0)
assert any("[RAW]" in e["content"] for e in entries)
# Session should still be truncated
reloaded = sessions.get_or_create("cli:fail")
assert len(reloaded.messages) <= 4
@pytest.mark.asyncio
async def test_respects_last_consolidated(self, real_consolidator, mock_provider):
"""30 turns with last_consolidated=50 → only unconsolidated tail considered."""
mock_provider.chat_with_retry.return_value = MagicMock(
content="Tail summary.", finish_reason="stop"
)
sessions = real_consolidator.sessions
session = sessions.get_or_create("cli:offset")
for i in range(30):
session.add_message("user", f"u{i}")
session.add_message("assistant", f"a{i}")
session.last_consolidated = 50 # Only 10 messages unconsolidated
sessions.save(session)
result = await real_consolidator.compact_idle_session("cli:offset", max_suffix=4)
assert result == "Tail summary."
# Verify only the unconsolidated tail was processed:
# 10 unconsolidated messages (50-59), keep suffix of 4 → archive 6
archived_call = mock_provider.chat_with_retry.call_args
user_content = archived_call.kwargs["messages"][1]["content"]
# Should contain only tail messages, not early ones
assert "u0" not in user_content
assert "u25" in user_content or "a25" in user_content
@pytest.mark.asyncio
async def test_acquires_consolidation_lock(self, real_consolidator, mock_provider):
"""Verify lock is held during execution."""
import asyncio
# Use a slow LLM response to ensure the lock is held while we check
started = asyncio.Event()
async def slow_chat(**kwargs):
started.set()
await asyncio.sleep(0.1)
return MagicMock(content="Summary.", finish_reason="stop")
mock_provider.chat_with_retry = slow_chat
sessions = real_consolidator.sessions
session = sessions.get_or_create("cli:lock")
for i in range(10):
session.add_message("user", f"u{i}")
session.add_message("assistant", f"a{i}")
sessions.save(session)
lock = real_consolidator.get_lock("cli:lock")
assert not lock.locked()
task = asyncio.ensure_future(
real_consolidator.compact_idle_session("cli:lock", max_suffix=4)
)
await started.wait()
assert lock.locked()
await task
assert not lock.locked()
class TestConsolidatorSessionRefresh:
"""Background consolidation must detect stale session references."""
@pytest.mark.asyncio
async def test_reloads_before_empty_session_guard(self, tmp_path):
"""A stale empty reference must not skip a non-empty cached session."""
from nanobot.agent.memory import Consolidator, MemoryStore
from nanobot.session.manager import Session, SessionManager
store = MemoryStore(tmp_path)
provider = MagicMock()
provider.chat_with_retry = AsyncMock(
return_value=MagicMock(content="summary", finish_reason="stop")
)
provider.generation.max_tokens = 4096
provider.estimate_prompt_tokens = MagicMock(return_value=(10, "test"))
sessions = SessionManager(tmp_path)
consolidator = Consolidator(
store=store,
provider=provider,
model="test-model",
sessions=sessions,
context_window_tokens=128_000,
build_messages=MagicMock(return_value=[]),
get_tool_definitions=MagicMock(return_value=[]),
)
fresh = sessions.get_or_create("cli:test")
fresh.add_message("user", "fresh message")
sessions.save(fresh)
stale_empty = Session(key="cli:test")
seen: dict[str, Session] = {}
def estimate(session: Session):
seen["session"] = session
return 10, "test"
consolidator.estimate_session_prompt_tokens = MagicMock(side_effect=estimate)
await consolidator.maybe_consolidate_by_tokens(stale_empty)
assert seen["session"] is fresh
@pytest.mark.asyncio
async def test_reloads_stale_session_after_compact(self, tmp_path):
"""After compact_idle_session replaces the session, a concurrent
maybe_consolidate_by_tokens with the old reference should use the
fresh session from cache instead of overwriting."""
from nanobot.agent.memory import Consolidator, MemoryStore
from nanobot.session.manager import SessionManager
store = MemoryStore(tmp_path)
provider = MagicMock()
provider.chat_with_retry = AsyncMock(
return_value=MagicMock(content="summary", finish_reason="stop")
)
provider.generation.max_tokens = 4096
provider.estimate_prompt_tokens = MagicMock(return_value=(10, "test"))
sessions = SessionManager(tmp_path)
consolidator = Consolidator(
store=store,
provider=provider,
model="test-model",
sessions=sessions,
context_window_tokens=128_000,
build_messages=MagicMock(return_value=[]),
get_tool_definitions=MagicMock(return_value=[]),
)
# Populate session with many messages
session = sessions.get_or_create("cli:test")
for i in range(20):
session.add_message("user", f"u{i}")
session.add_message("assistant", f"a{i}")
sessions.save(session)
# Simulate: background consolidation captures old reference
old_ref = session
# AutoCompact runs first and truncates to 8
await consolidator.compact_idle_session("cli:test", max_suffix=8)
# Background consolidation runs with stale reference —
# should detect the session was replaced and not undo the compact.
await consolidator.maybe_consolidate_by_tokens(old_ref)
session_after = sessions.get_or_create("cli:test")
# Messages should still be truncated (not restored to 40)
assert len(session_after.messages) <= 8
class TestRawArchiveTruncation:
"""raw_archive() must cap entry size to avoid bloating history.jsonl."""
+1
View File
@@ -387,6 +387,7 @@ class TestConsolidationUnaffectedByUnifiedSession:
session = Session(key="unified:default")
session.messages = [{"role": "user", "content": "msg"}]
sessions.get_or_create.return_value = session
# Simulate over-budget: estimated > budget
consolidator.estimate_session_prompt_tokens = MagicMock(return_value=(950, "tiktoken"))