refactor(memory): remove consolidation ratio (#5575)

* refactor(memory): remove consolidation ratio

* docs(memory): document fixed consolidation policy

* docs(memory): simplify consolidation overview

* docs(memory): rely on soft wrapping
This commit is contained in:
chengyongru
2026-08-28 13:20:09 +08:00
committed by GitHub
parent 29025f5a8b
commit cace42af14
8 changed files with 71 additions and 315 deletions
+1 -3
View File
@@ -29,9 +29,7 @@ Memory moves through nanobot in two stages.
### Stage 1: Consolidator ### Stage 1: Consolidator
When a conversation grows large enough to pressure the context window, nanobot does not try to carry every old message forever. When a conversation grows large, the `Consolidator` summarizes older turns and appends the result to `memory/history.jsonl`, while keeping recent conversation available. Each summary preserves useful long-term facts and a short handoff for active work.
Instead, the `Consolidator` summarizes the oldest safe slice of the conversation and appends that summary to `memory/history.jsonl`.
This file is: This file is:
-3
View File
@@ -273,7 +273,6 @@ class AgentLoop:
channels_config: ChannelsConfig | None = None, channels_config: ChannelsConfig | None = None,
timezone: str | None = None, timezone: str | None = None,
session_ttl_minutes: int = 0, session_ttl_minutes: int = 0,
consolidation_ratio: float = 0.5,
hooks: list[AgentHook] | None = None, hooks: list[AgentHook] | None = None,
hook_factories: list[AgentTurnHookFactory] | None = None, hook_factories: list[AgentTurnHookFactory] | None = None,
unified_session: bool = False, unified_session: bool = False,
@@ -444,7 +443,6 @@ class AgentLoop:
workspace_scopes=self.workspace_scopes, workspace_scopes=self.workspace_scopes,
unified_session=unified_session, unified_session=unified_session,
), ),
consolidation_ratio=consolidation_ratio,
unified_session=unified_session, unified_session=unified_session,
) )
self.auto_compact = AutoCompact( self.auto_compact = AutoCompact(
@@ -517,7 +515,6 @@ class AgentLoop:
disabled_skills=defaults.disabled_skills, disabled_skills=defaults.disabled_skills,
session_ttl_minutes=defaults.session_ttl_minutes, session_ttl_minutes=defaults.session_ttl_minutes,
idle_compact_check_interval_seconds=defaults.idle_compact_check_interval_seconds, idle_compact_check_interval_seconds=defaults.idle_compact_check_interval_seconds,
consolidation_ratio=defaults.consolidation_ratio,
tools_config=config.tools, tools_config=config.tools,
model_presets=preset_helpers.configured_model_presets(config), model_presets=preset_helpers.configured_model_presets(config),
model_preset=defaults.model_preset, model_preset=defaults.model_preset,
+41 -73
View File
@@ -32,7 +32,6 @@ from nanobot.utils.gitstore import GitStore
from nanobot.utils.helpers import ( from nanobot.utils.helpers import (
content_with_media_breadcrumbs, content_with_media_breadcrumbs,
ensure_dir, ensure_dir,
estimate_message_tokens,
estimate_prompt_tokens_chain, estimate_prompt_tokens_chain,
strip_think, strip_think,
truncate_text, truncate_text,
@@ -956,8 +955,6 @@ class MemoryArchiver:
class Consolidator: class Consolidator:
"""Legacy context-pressure coordinator backed by a MemoryArchiver.""" """Legacy context-pressure coordinator backed by a MemoryArchiver."""
_MAX_CONSOLIDATION_ROUNDS = 5
_SAFETY_BUFFER = 1024 # extra headroom for tokenizer estimation drift _SAFETY_BUFFER = 1024 # extra headroom for tokenizer estimation drift
def __init__( def __init__(
@@ -967,12 +964,10 @@ class Consolidator:
build_messages: Callable[..., list[dict[str, Any]]], build_messages: Callable[..., list[dict[str, Any]]],
get_tool_definitions: Callable[[], list[dict[str, Any]]], get_tool_definitions: Callable[[], list[dict[str, Any]]],
resolve_prompt_context: Callable[[Session], tuple[str | None, Path | None]] | None = None, resolve_prompt_context: Callable[[Session], tuple[str | None, Path | None]] | None = None,
consolidation_ratio: float = 0.5,
unified_session: bool = False, unified_session: bool = False,
): ):
self.store = store self.store = store
self.sessions = sessions self.sessions = sessions
self.consolidation_ratio = consolidation_ratio
self.unified_session = unified_session self.unified_session = unified_session
self._build_messages = build_messages self._build_messages = build_messages
self._get_tool_definitions = get_tool_definitions self._get_tool_definitions = get_tool_definitions
@@ -995,24 +990,19 @@ class Consolidator:
def pick_consolidation_boundary( def pick_consolidation_boundary(
self, self,
session: Session, session: Session,
tokens_to_remove: int, ) -> int | None:
) -> tuple[int, int] | None: """Return the fixed user-led boundary before the recent replay tail."""
"""Pick a user-turn boundary that removes enough old prompt tokens.""" if not session.messages:
start = session.last_archived
if start >= len(session.messages) or tokens_to_remove <= 0:
return None return None
boundary = max(0, len(session.messages) - MIN_COMPACTED_REPLAY_MESSAGES)
removed_tokens = 0 while boundary > 0 and session.messages[boundary].get("role") != "user":
last_boundary: tuple[int, int] | None = None boundary -= 1
for idx in range(start, len(session.messages)): if (
message = session.messages[idx] boundary <= session.last_archived
if idx > start and message.get("role") == "user": or session.messages[boundary].get("role") != "user"
last_boundary = (idx, removed_tokens) ):
if removed_tokens >= tokens_to_remove: return None
return last_boundary return boundary
removed_tokens += estimate_message_tokens(message)
return last_boundary
@staticmethod @staticmethod
def _full_replay_history( def _full_replay_history(
@@ -1106,7 +1096,7 @@ class Consolidator:
*, *,
runtime: LLMRuntime, runtime: LLMRuntime,
) -> None: ) -> None:
"""Loop: archive old messages until prompt fits within safe budget. """Archive one fixed old prefix when the prompt exceeds the safe budget.
The budget reserves space for completion tokens and a safety buffer The budget reserves space for completion tokens and a safety buffer
so the LLM request never exceeds the context window. so the LLM request never exceeds the context window.
@@ -1124,7 +1114,6 @@ class Consolidator:
return return
budget = self._input_token_budget(runtime) budget = self._input_token_budget(runtime)
target = int(budget * self.consolidation_ratio)
last_summary: str | None = None last_summary: str | None = None
estimated, source = self.estimate_session_prompt_tokens( estimated, source = self.estimate_session_prompt_tokens(
session, session,
@@ -1146,58 +1135,37 @@ class Consolidator:
self._persist_last_summary(session, last_summary) self._persist_last_summary(session, last_summary)
return return
for round_num in range(self._MAX_CONSOLIDATION_ROUNDS): end_idx = self.pick_consolidation_boundary(session)
if estimated <= target: if end_idx is None:
break logger.debug(
"Token consolidation: no safe fixed boundary for {}",
boundary = self.pick_consolidation_boundary(session, max(1, estimated - target))
if boundary is None:
logger.debug(
"Token consolidation: no safe boundary for {} (round {})",
session.key,
round_num,
)
break
end_idx = boundary[0]
chunk = session.messages[session.last_archived:end_idx]
if not chunk:
break
logger.info(
"Token consolidation round {} for {}: {}/{} via {}, chunk={} msgs",
round_num,
session.key, session.key,
estimated,
runtime.context_window_tokens,
source,
len(chunk),
) )
summary = await self.archive_session( return
session,
archive_end=end_idx,
runtime=runtime,
)
# Advance the cursor either way: on success the chunk was
# summarized; on failure archive_session() 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
session.last_archived = 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
estimated, source = self.estimate_session_prompt_tokens( chunk = session.messages[session.last_archived:end_idx]
session, if not chunk:
runtime=runtime, return
)
if estimated <= 0: logger.info(
break "Token consolidation for {}: {}/{} via {}, chunk={} msgs",
session.key,
estimated,
runtime.context_window_tokens,
source,
len(chunk),
)
summary = await self.archive_session(
session,
archive_end=end_idx,
runtime=runtime,
)
# Advance either way: archive_session raw-archives on degradation,
# and replaying the same chunk would duplicate Memory material.
if summary:
last_summary = summary
session.last_archived = end_idx
self.sessions.save(session)
# Persist the last summary to session metadata so it can be injected # Persist the last summary to session metadata so it can be injected
# into the runtime context on the next prepare_session() call, aligning # into the runtime context on the next prepare_session() call, aligning
-7
View File
@@ -155,13 +155,6 @@ class AgentDefaults(Base):
default=60, default=60,
ge=0, ge=0,
) # Minimum interval in seconds between scans for idle sessions ) # Minimum interval in seconds between scans for idle sessions
consolidation_ratio: float = Field(
default=0.5,
ge=0.1,
le=0.95,
validation_alias=AliasChoices("consolidationRatio"),
serialization_alias="consolidationRatio",
) # Consolidation target ratio (0.5 = 50% of budget retained after compression)
dream: DreamConfig = Field(default_factory=DreamConfig) dream: DreamConfig = Field(default_factory=DreamConfig)
@model_validator(mode="before") @model_validator(mode="before")
@@ -6,6 +6,8 @@ Use [skip] unless a fact meets all SNIP criteria:
- Important: prevents rework or captures preferences / rules - Important: prevents rework or captures preferences / rules
- Persistent: still relevant after 2 weeks - Persistent: still relevant after 2 weeks
Also preserve a compact working-state handoff even when it is not Persistent: the active objective, current status, completed steps, unresolved blockers, next action, and exact identifiers needed to continue without rework. Mark these facts [ephemeral].
Format each fact as: Format each fact as:
- [mark] fact content - [mark] fact content
-112
View File
@@ -1,112 +0,0 @@
"""Tests for configurable consolidation_ratio."""
from unittest.mock import AsyncMock, MagicMock
import pytest
from pydantic import ValidationError
import nanobot.agent.memory as memory_module
from nanobot.agent.loop import AgentLoop
from nanobot.bus.queue import MessageBus
from nanobot.config.schema import AgentDefaults
from nanobot.providers.base import GenerationSettings, LLMResponse
def _make_loop(
tmp_path,
*,
estimated_tokens: int = 0,
context_window_tokens: int = 200,
consolidation_ratio: float = 0.5,
) -> AgentLoop:
provider = MagicMock()
provider.get_default_model.return_value = "test-model"
provider.generation = GenerationSettings(max_tokens=0)
provider.estimate_prompt_tokens.return_value = (estimated_tokens, "test-counter")
_response = LLMResponse(content="ok", tool_calls=[])
provider.chat_with_retry = AsyncMock(return_value=_response)
provider.chat_stream_with_retry = AsyncMock(return_value=_response)
loop = AgentLoop(
bus=MessageBus(),
provider=provider,
workspace=tmp_path,
model="test-model",
context_window_tokens=context_window_tokens,
consolidation_ratio=consolidation_ratio,
)
loop.tools.get_definitions = MagicMock(return_value=[])
loop.consolidator._SAFETY_BUFFER = 0
return loop
def _session_with_turns(loop: AgentLoop, *, turns: int):
session = loop.sessions.get_or_create("cli:test")
session.messages = []
for i in range(turns):
session.messages.append({"role": "user", "content": f"u{i}", "timestamp": f"2026-01-01T00:00:{i:02d}"})
session.messages.append({"role": "assistant", "content": f"a{i}", "timestamp": f"2026-01-01T00:01:{i:02d}"})
loop.sessions.save(session)
return session
@pytest.mark.asyncio
@pytest.mark.parametrize(
("ratio", "context_window_tokens", "estimates", "expected_archives"),
[
(0.5, 200, [250, 90], 1),
(0.1, 1000, [1200, 800, 400, 50], 2),
(0.9, 200, [300, 175], 1),
],
)
async def test_consolidation_ratio_controls_target(
tmp_path,
monkeypatch,
ratio: float,
context_window_tokens: int,
estimates: list[int],
expected_archives: int,
) -> None:
loop = _make_loop(
tmp_path,
context_window_tokens=context_window_tokens,
consolidation_ratio=ratio,
)
loop.consolidator.archive_session = AsyncMock(return_value=True) # type: ignore[method-assign]
session = _session_with_turns(loop, turns=10)
remaining_estimates = list(estimates)
runtime = loop.llm_runtime()
def mock_estimate(_session, *, runtime):
return (remaining_estimates.pop(0), "test")
loop.consolidator.estimate_session_prompt_tokens = mock_estimate # type: ignore[method-assign]
monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda _m: 100)
await loop.consolidator.maybe_consolidate_by_tokens(
session,
runtime=runtime,
)
assert loop.consolidator.archive_session.await_count == expected_archives
def test_ratio_propagated_from_config_schema() -> None:
defaults = AgentDefaults()
assert defaults.consolidation_ratio == 0.5
defaults = AgentDefaults.model_validate({"consolidationRatio": 0.3})
assert defaults.consolidation_ratio == 0.3
dumped = defaults.model_dump(by_alias=True)
assert dumped["consolidationRatio"] == 0.3
def test_ratio_validation_rejects_out_of_range() -> None:
with pytest.raises(ValidationError):
AgentDefaults(consolidation_ratio=0.05)
with pytest.raises(ValidationError):
AgentDefaults(consolidation_ratio=1.0)
+7 -5
View File
@@ -232,17 +232,19 @@ class TestConsolidatorSummarize:
class TestConsolidatorPromptContract: class TestConsolidatorPromptContract:
def test_archive_prompt_outputs_attribute_tags_without_missing_context_claims(self): def test_archive_prompt_preserves_working_state_with_memory_facts(self):
prompt = render_template("agent/consolidator_archive.md", strip=True, archive_count=4) prompt = render_template("agent/consolidator_archive.md", strip=True, archive_count=4)
assert "SNIP" in prompt assert "SNIP" in prompt
assert "final 4 conversation messages" in prompt assert "final 4 conversation messages" in prompt
for mark in ("[permanent]", "[durable]", "[ephemeral]", "[correction]", "[skip]"): for mark in ("[permanent]", "[durable]", "[ephemeral]", "[correction]", "[skip]"):
assert mark in prompt assert mark in prompt
assert "check context below" not in prompt.lower() assert "working-state handoff" in prompt
assert "exact identifiers needed to continue without rework" in prompt
assert "Do not output facts already present in the system prompt's Recent History" in prompt assert "Do not output facts already present in the system prompt's Recent History" in prompt
assert "Do not mark something [skip] merely because it might already exist" in prompt assert "Do not mark something [skip] merely because it might already exist" in prompt
class TestConsolidatorArchiveErrorHandling: class TestConsolidatorArchiveErrorHandling:
"""archive() must fall back when the LLM does not complete its overview. """archive() must fall back when the LLM does not complete its overview.
@@ -420,7 +422,7 @@ class TestConsolidatorTokenBudget:
consolidator.estimate_session_prompt_tokens = MagicMock( consolidator.estimate_session_prompt_tokens = MagicMock(
side_effect=[(1200, "tiktoken"), (400, "tiktoken")] side_effect=[(1200, "tiktoken"), (400, "tiktoken")]
) )
consolidator.pick_consolidation_boundary = MagicMock(return_value=(50, 800)) consolidator.pick_consolidation_boundary = MagicMock(return_value=50)
consolidator.archiver._build_messages = MagicMock(side_effect=_build_test_messages) consolidator.archiver._build_messages = MagicMock(side_effect=_build_test_messages)
mock_provider.estimate_prompt_tokens.return_value = (100, "test-counter") mock_provider.estimate_prompt_tokens.return_value = (100, "test-counter")
mock_provider.chat_with_retry.return_value = LLMResponse( mock_provider.chat_with_retry.return_value = LLMResponse(
@@ -493,7 +495,7 @@ class TestConsolidatorTokenBudget:
await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime) await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime)
# Exactly one fallback per call — not _MAX_CONSOLIDATION_ROUNDS. # The fixed policy archives at most one prefix per call.
assert consolidator.archive_session.await_count == 1 assert consolidator.archive_session.await_count == 1
async def test_boundary_respected_when_no_intermediate_user_turn( async def test_boundary_respected_when_no_intermediate_user_turn(
@@ -520,7 +522,7 @@ class TestConsolidatorTokenBudget:
await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime) await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime)
consolidator.archive_session.assert_awaited_once() consolidator.archive_session.assert_awaited_once()
# pick_consolidation_boundary finds the only boundary at idx=61 # The fixed recent tail expands backward to the user at idx=61.
assert session.last_archived == 61 assert session.last_archived == 61
+20 -112
View File
@@ -2,7 +2,6 @@ from unittest.mock import AsyncMock, MagicMock
import pytest import pytest
import nanobot.agent.memory as memory_module
from nanobot.agent.loop import AgentLoop from nanobot.agent.loop import AgentLoop
from nanobot.bus.queue import MessageBus from nanobot.bus.queue import MessageBus
from nanobot.providers.base import LLMResponse from nanobot.providers.base import LLMResponse
@@ -41,17 +40,16 @@ async def test_prompt_below_threshold_does_not_consolidate(tmp_path) -> None:
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_prompt_above_threshold_triggers_consolidation(tmp_path, monkeypatch) -> None: async def test_prompt_above_threshold_triggers_consolidation(tmp_path) -> None:
loop = _make_loop(tmp_path, estimated_tokens=1000, context_window_tokens=200) loop = _make_loop(tmp_path, estimated_tokens=1000, context_window_tokens=200)
loop.consolidator.archive_session = AsyncMock(return_value=True) # type: ignore[method-assign] loop.consolidator.archive_session = AsyncMock(return_value=True) # type: ignore[method-assign]
session = loop.sessions.get_or_create("cli:test") session = loop.sessions.get_or_create("cli:test")
session.messages = [ session.messages = [
{"role": "user", "content": "u1", "timestamp": "2026-01-01T00:00:00"}, {"role": role, "content": f"{role[0]}{turn}"}
{"role": "assistant", "content": "a1", "timestamp": "2026-01-01T00:00:01"}, for turn in range(10)
{"role": "user", "content": "u2", "timestamp": "2026-01-01T00:00:02"}, for role in ("user", "assistant")
] ]
loop.sessions.save(session) loop.sessions.save(session)
monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda _message: 500)
await loop.process_direct("hello", session_key="cli:test") await loop.process_direct("hello", session_key="cli:test")
@@ -59,23 +57,18 @@ async def test_prompt_above_threshold_triggers_consolidation(tmp_path, monkeypat
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_prompt_above_threshold_archives_until_next_user_boundary(tmp_path, monkeypatch) -> None: async def test_prompt_above_threshold_uses_fixed_recent_tail(tmp_path) -> None:
loop = _make_loop(tmp_path, estimated_tokens=1000, context_window_tokens=200) loop = _make_loop(tmp_path, estimated_tokens=1000, context_window_tokens=200)
loop.consolidator.archive_session = AsyncMock(return_value=True) # type: ignore[method-assign] loop.consolidator.archive_session = AsyncMock(return_value=True) # type: ignore[method-assign]
session = loop.sessions.get_or_create("cli:test") session = loop.sessions.get_or_create("cli:test")
session.messages = [ session.messages = [
{"role": "user", "content": "u1", "timestamp": "2026-01-01T00:00:00"}, {"role": role, "content": f"{role[0]}{turn}"}
{"role": "assistant", "content": "a1", "timestamp": "2026-01-01T00:00:01"}, for turn in range(10)
{"role": "user", "content": "u2", "timestamp": "2026-01-01T00:00:02"}, for role in ("user", "assistant")
{"role": "assistant", "content": "a2", "timestamp": "2026-01-01T00:00:03"},
{"role": "user", "content": "u3", "timestamp": "2026-01-01T00:00:04"},
] ]
loop.sessions.save(session) loop.sessions.save(session)
token_map = {"u1": 120, "a1": 120, "u2": 120, "a2": 120, "u3": 120}
monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda message: token_map[message["content"]])
await loop.consolidator.maybe_consolidate_by_tokens( await loop.consolidator.maybe_consolidate_by_tokens(
session, session,
runtime=loop.llm_runtime(), runtime=loop.llm_runtime(),
@@ -83,112 +76,29 @@ async def test_prompt_above_threshold_archives_until_next_user_boundary(tmp_path
archive_end = loop.consolidator.archive_session.await_args.kwargs["archive_end"] archive_end = loop.consolidator.archive_session.await_args.kwargs["archive_end"]
archived_chunk = session.messages[:archive_end] archived_chunk = session.messages[:archive_end]
assert [message["content"] for message in archived_chunk] == ["u1", "a1", "u2", "a2"] assert [message["content"] for message in archived_chunk] == [
assert session.last_archived == 4 "u0", "a0", "u1", "a1", "u2", "a2", "u3", "a3", "u4", "a4", "u5", "a5",
@pytest.mark.asyncio
async def test_consolidation_loops_until_target_met(tmp_path, monkeypatch) -> None:
"""Verify maybe_consolidate_by_tokens keeps looping until under threshold."""
loop = _make_loop(tmp_path, estimated_tokens=0, context_window_tokens=200)
loop.consolidator.archive_session = AsyncMock(return_value=True) # type: ignore[method-assign]
session = loop.sessions.get_or_create("cli:test")
session.messages = [
{"role": "user", "content": "u1", "timestamp": "2026-01-01T00:00:00"},
{"role": "assistant", "content": "a1", "timestamp": "2026-01-01T00:00:01"},
{"role": "user", "content": "u2", "timestamp": "2026-01-01T00:00:02"},
{"role": "assistant", "content": "a2", "timestamp": "2026-01-01T00:00:03"},
{"role": "user", "content": "u3", "timestamp": "2026-01-01T00:00:04"},
{"role": "assistant", "content": "a3", "timestamp": "2026-01-01T00:00:05"},
{"role": "user", "content": "u4", "timestamp": "2026-01-01T00:00:06"},
] ]
loop.sessions.save(session) assert session.last_archived == 12
call_count = [0]
def mock_estimate(_session, *, runtime):
call_count[0] += 1
if call_count[0] == 1:
return (500, "test")
if call_count[0] == 2:
return (300, "test")
return (80, "test")
loop.consolidator.estimate_session_prompt_tokens = mock_estimate # type: ignore[method-assign]
monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda _m: 100)
await loop.consolidator.maybe_consolidate_by_tokens(
session,
runtime=loop.llm_runtime(),
)
assert loop.consolidator.archive_session.await_count == 2
assert session.last_archived == 6
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_consolidation_continues_below_trigger_until_half_target(tmp_path, monkeypatch) -> None: async def test_consolidation_persists_summary_for_next_prepare_session(tmp_path) -> None:
"""Once triggered, consolidation should continue until it drops below half threshold."""
loop = _make_loop(tmp_path, estimated_tokens=0, context_window_tokens=200)
loop.consolidator.archive_session = AsyncMock(return_value=True) # type: ignore[method-assign]
session = loop.sessions.get_or_create("cli:test")
session.messages = [
{"role": "user", "content": "u1", "timestamp": "2026-01-01T00:00:00"},
{"role": "assistant", "content": "a1", "timestamp": "2026-01-01T00:00:01"},
{"role": "user", "content": "u2", "timestamp": "2026-01-01T00:00:02"},
{"role": "assistant", "content": "a2", "timestamp": "2026-01-01T00:00:03"},
{"role": "user", "content": "u3", "timestamp": "2026-01-01T00:00:04"},
{"role": "assistant", "content": "a3", "timestamp": "2026-01-01T00:00:05"},
{"role": "user", "content": "u4", "timestamp": "2026-01-01T00:00:06"},
]
loop.sessions.save(session)
call_count = [0]
def mock_estimate(_session, *, runtime):
call_count[0] += 1
if call_count[0] == 1:
return (500, "test")
if call_count[0] == 2:
return (150, "test")
return (80, "test")
loop.consolidator.estimate_session_prompt_tokens = mock_estimate # type: ignore[method-assign]
monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda _m: 100)
await loop.consolidator.maybe_consolidate_by_tokens(
session,
runtime=loop.llm_runtime(),
)
assert loop.consolidator.archive_session.await_count == 2
assert session.last_archived == 6
@pytest.mark.asyncio
async def test_consolidation_persists_summary_for_next_prepare_session(tmp_path, monkeypatch) -> None:
loop = _make_loop(tmp_path, estimated_tokens=0, context_window_tokens=200) loop = _make_loop(tmp_path, estimated_tokens=0, context_window_tokens=200)
loop.consolidator.archive_session = AsyncMock(return_value="User discussed project status.") # type: ignore[method-assign] loop.consolidator.archive_session = AsyncMock(return_value="User discussed project status.") # type: ignore[method-assign]
session = loop.sessions.get_or_create("cli:test") session = loop.sessions.get_or_create("cli:test")
session.messages = [ session.messages = [
{"role": "user", "content": "u1", "timestamp": "2026-01-01T00:00:00"}, {"role": role, "content": f"{role[0]}{turn}"}
{"role": "assistant", "content": "a1", "timestamp": "2026-01-01T00:00:01"}, for turn in range(5)
{"role": "user", "content": "u2", "timestamp": "2026-01-01T00:00:02"}, for role in ("user", "assistant")
] ]
loop.sessions.save(session) loop.sessions.save(session)
call_count = [0]
def mock_estimate(_session, *, runtime): def mock_estimate(_session, *, runtime):
call_count[0] += 1 return (500, "test")
if call_count[0] == 1:
return (500, "test")
return (80, "test")
loop.consolidator.estimate_session_prompt_tokens = mock_estimate # type: ignore[method-assign] loop.consolidator.estimate_session_prompt_tokens = mock_estimate # type: ignore[method-assign]
monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda _m: 150)
await loop.consolidator.maybe_consolidate_by_tokens( await loop.consolidator.maybe_consolidate_by_tokens(
session, session,
@@ -235,7 +145,7 @@ async def test_preflight_consolidation_receives_pending_summary(tmp_path) -> Non
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_preflight_consolidation_before_llm_call(tmp_path, monkeypatch) -> None: async def test_preflight_consolidation_before_llm_call(tmp_path) -> None:
"""Verify preflight consolidation runs before the LLM call in process_direct.""" """Verify preflight consolidation runs before the LLM call in process_direct."""
order: list[str] = [] order: list[str] = []
@@ -258,13 +168,11 @@ async def test_preflight_consolidation_before_llm_call(tmp_path, monkeypatch) ->
session = loop.sessions.get_or_create("cli:test") session = loop.sessions.get_or_create("cli:test")
session.messages = [ session.messages = [
{"role": "user", "content": "u1", "timestamp": "2026-01-01T00:00:00"}, {"role": role, "content": f"{role[0]}{turn}"}
{"role": "assistant", "content": "a1", "timestamp": "2026-01-01T00:00:01"}, for turn in range(10)
{"role": "user", "content": "u2", "timestamp": "2026-01-01T00:00:02"}, for role in ("user", "assistant")
] ]
loop.sessions.save(session) loop.sessions.save(session)
monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda _m: 500)
call_count = [0] call_count = [0]
def mock_estimate(_session, *, runtime): def mock_estimate(_session, *, runtime):
call_count[0] += 1 call_count[0] += 1