refactor(agent): require runtime for consolidation
This commit is contained in:
@@ -1,5 +1,6 @@
|
||||
"""Tests for the lightweight Consolidator — append-only to HISTORY.md."""
|
||||
|
||||
from dataclasses import replace
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
@@ -9,8 +10,9 @@ from nanobot.agent.memory import (
|
||||
Consolidator,
|
||||
MemoryStore,
|
||||
)
|
||||
from nanobot.providers.base import LLMResponse
|
||||
from nanobot.providers.base import GenerationSettings, LLMResponse
|
||||
from nanobot.session.manager import Session
|
||||
from nanobot.utils.llm_runtime import LLMRuntime
|
||||
from nanobot.utils.prompt_templates import render_template
|
||||
|
||||
|
||||
@@ -23,11 +25,21 @@ def store(tmp_path):
|
||||
def mock_provider():
|
||||
p = MagicMock()
|
||||
p.chat_with_retry = AsyncMock()
|
||||
p.generation = GenerationSettings(max_tokens=100)
|
||||
return p
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def consolidator(store, mock_provider):
|
||||
def runtime(mock_provider):
|
||||
return LLMRuntime.capture(
|
||||
mock_provider,
|
||||
"test-model",
|
||||
context_window_tokens=1000,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def consolidator(store):
|
||||
sessions = MagicMock()
|
||||
sessions.save = MagicMock()
|
||||
# When maybe_consolidate_by_tokens refreshes the session reference via
|
||||
@@ -38,13 +50,9 @@ def consolidator(store, mock_provider):
|
||||
sessions._session_cache = _session_cache
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
@@ -62,7 +70,41 @@ def _tool_round(call_id: str) -> list[dict]:
|
||||
|
||||
|
||||
class TestConsolidatorSummarize:
|
||||
async def test_summarize_appends_to_history(self, consolidator, mock_provider, store):
|
||||
async def test_archive_uses_captured_generation(
|
||||
self, consolidator, mock_provider, runtime
|
||||
):
|
||||
admitted = replace(
|
||||
runtime,
|
||||
generation=GenerationSettings(
|
||||
temperature=0.25,
|
||||
max_tokens=321,
|
||||
reasoning_effort="medium",
|
||||
),
|
||||
)
|
||||
mock_provider.generation = GenerationSettings(
|
||||
temperature=0.9,
|
||||
max_tokens=999,
|
||||
reasoning_effort="high",
|
||||
)
|
||||
mock_provider.chat_with_retry.return_value = MagicMock(
|
||||
content="Summary.",
|
||||
finish_reason="stop",
|
||||
)
|
||||
|
||||
await consolidator.archive(
|
||||
[{"role": "user", "content": "hello"}],
|
||||
runtime=admitted,
|
||||
)
|
||||
|
||||
call = mock_provider.chat_with_retry.call_args.kwargs
|
||||
assert call["model"] == admitted.model
|
||||
assert call["temperature"] == 0.25
|
||||
assert call["max_tokens"] == 321
|
||||
assert call["reasoning_effort"] == "medium"
|
||||
|
||||
async def test_summarize_appends_to_history(
|
||||
self, consolidator, mock_provider, store, runtime
|
||||
):
|
||||
"""Consolidator should call LLM to summarize, then append to HISTORY.md."""
|
||||
mock_provider.chat_with_retry.return_value = MagicMock(
|
||||
content="User fixed a bug in the auth module."
|
||||
@@ -71,7 +113,7 @@ class TestConsolidatorSummarize:
|
||||
{"role": "user", "content": "fix the auth bug"},
|
||||
{"role": "assistant", "content": "Done, fixed the race condition."},
|
||||
]
|
||||
result = await consolidator.archive(messages)
|
||||
result = await consolidator.archive(messages, runtime=runtime)
|
||||
assert result == "User fixed a bug in the auth module."
|
||||
entries = store.read_unprocessed_history(since_cursor=0)
|
||||
assert len(entries) == 1
|
||||
@@ -81,6 +123,7 @@ class TestConsolidatorSummarize:
|
||||
consolidator,
|
||||
mock_provider,
|
||||
store,
|
||||
runtime,
|
||||
):
|
||||
mock_provider.chat_with_retry.return_value = MagicMock(
|
||||
content="User fixed a bug in the auth module.",
|
||||
@@ -88,16 +131,22 @@ class TestConsolidatorSummarize:
|
||||
)
|
||||
messages = [{"role": "user", "content": "fix the auth bug"}]
|
||||
|
||||
await consolidator.archive(messages, session_key="telegram:chat-1")
|
||||
await consolidator.archive(
|
||||
messages,
|
||||
runtime=runtime,
|
||||
session_key="telegram:chat-1",
|
||||
)
|
||||
|
||||
entries = store.read_unprocessed_history(since_cursor=0)
|
||||
assert entries[0]["session_key"] == "telegram:chat-1"
|
||||
|
||||
async def test_summarize_raw_dumps_on_llm_failure(self, consolidator, mock_provider, store):
|
||||
async def test_summarize_raw_dumps_on_llm_failure(
|
||||
self, consolidator, mock_provider, store, runtime
|
||||
):
|
||||
"""On LLM failure, raw-dump messages to HISTORY.md."""
|
||||
mock_provider.chat_with_retry.side_effect = Exception("API error")
|
||||
messages = [{"role": "user", "content": "hello"}]
|
||||
result = await consolidator.archive(messages)
|
||||
result = await consolidator.archive(messages, runtime=runtime)
|
||||
assert result is None # no summary on raw dump fallback
|
||||
entries = store.read_unprocessed_history(since_cursor=0)
|
||||
assert len(entries) == 1
|
||||
@@ -108,17 +157,22 @@ class TestConsolidatorSummarize:
|
||||
consolidator,
|
||||
mock_provider,
|
||||
store,
|
||||
runtime,
|
||||
):
|
||||
mock_provider.chat_with_retry.side_effect = Exception("API error")
|
||||
messages = [{"role": "user", "content": "hello"}]
|
||||
|
||||
await consolidator.archive(messages, session_key="slack:chat-2")
|
||||
await consolidator.archive(
|
||||
messages,
|
||||
runtime=runtime,
|
||||
session_key="slack:chat-2",
|
||||
)
|
||||
|
||||
entries = store.read_unprocessed_history(since_cursor=0)
|
||||
assert entries[0]["session_key"] == "slack:chat-2"
|
||||
|
||||
async def test_summarize_skips_empty_messages(self, consolidator):
|
||||
result = await consolidator.archive([])
|
||||
async def test_summarize_skips_empty_messages(self, consolidator, runtime):
|
||||
result = await consolidator.archive([], runtime=runtime)
|
||||
assert result is None
|
||||
|
||||
|
||||
@@ -139,7 +193,9 @@ class TestConsolidatorArchiveErrorHandling:
|
||||
See https://github.com/HKUDS/nanobot/issues/3244
|
||||
"""
|
||||
|
||||
async def test_archive_falls_back_on_error_finish_reason(self, consolidator, mock_provider, store):
|
||||
async def test_archive_falls_back_on_error_finish_reason(
|
||||
self, consolidator, mock_provider, store, runtime
|
||||
):
|
||||
"""LLM returning finish_reason='error' should trigger raw_archive, not write error text."""
|
||||
mock_provider.chat_with_retry.return_value = MagicMock(
|
||||
content="Error: {'type': 'error', 'error': {'type': 'overloaded_error', 'message': 'overloaded_error (529)'}}",
|
||||
@@ -149,14 +205,16 @@ class TestConsolidatorArchiveErrorHandling:
|
||||
{"role": "user", "content": "fix the auth bug"},
|
||||
{"role": "assistant", "content": "Done, fixed the race condition."},
|
||||
]
|
||||
result = await consolidator.archive(messages)
|
||||
result = await consolidator.archive(messages, runtime=runtime)
|
||||
assert result is None
|
||||
entries = store.read_unprocessed_history(since_cursor=0)
|
||||
assert len(entries) == 1
|
||||
assert "[RAW]" in entries[0]["content"]
|
||||
assert "Error:" not in entries[0]["content"]
|
||||
|
||||
async def test_archive_preserves_summary_on_success(self, consolidator, mock_provider, store):
|
||||
async def test_archive_preserves_summary_on_success(
|
||||
self, consolidator, mock_provider, store, runtime
|
||||
):
|
||||
"""Normal LLM response should still produce a proper summary entry."""
|
||||
mock_provider.chat_with_retry.return_value = MagicMock(
|
||||
content="User fixed a bug in the auth module.",
|
||||
@@ -166,7 +224,7 @@ class TestConsolidatorArchiveErrorHandling:
|
||||
{"role": "user", "content": "fix the auth bug"},
|
||||
{"role": "assistant", "content": "Done."},
|
||||
]
|
||||
result = await consolidator.archive(messages)
|
||||
result = await consolidator.archive(messages, runtime=runtime)
|
||||
assert result == "User fixed a bug in the auth module."
|
||||
entries = store.read_unprocessed_history(since_cursor=0)
|
||||
assert len(entries) == 1
|
||||
@@ -174,7 +232,9 @@ class TestConsolidatorArchiveErrorHandling:
|
||||
|
||||
|
||||
class TestConsolidatorTokenBudget:
|
||||
async def test_prompt_below_threshold_does_not_consolidate(self, consolidator):
|
||||
async def test_prompt_below_threshold_does_not_consolidate(
|
||||
self, consolidator, runtime
|
||||
):
|
||||
"""No consolidation when tokens are within budget."""
|
||||
session = MagicMock()
|
||||
session.last_consolidated = 0
|
||||
@@ -183,10 +243,10 @@ class TestConsolidatorTokenBudget:
|
||||
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)
|
||||
await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime)
|
||||
consolidator.archive.assert_not_called()
|
||||
|
||||
async def test_estimate_uses_full_unconsolidated_tail(self, consolidator):
|
||||
async def test_estimate_uses_full_unconsolidated_tail(self, consolidator, runtime):
|
||||
"""Consolidation pressure must see messages hidden by the replay window."""
|
||||
session = Session(key="test:full-tail")
|
||||
for i in range(160):
|
||||
@@ -200,7 +260,7 @@ class TestConsolidatorTokenBudget:
|
||||
|
||||
consolidator._build_messages = build_messages
|
||||
|
||||
consolidator.estimate_session_prompt_tokens(session)
|
||||
consolidator.estimate_session_prompt_tokens(session, runtime=runtime)
|
||||
|
||||
assert len(captured["history"]) == 160
|
||||
assert captured["history"][0]["content"].endswith("msg-0")
|
||||
@@ -208,6 +268,7 @@ class TestConsolidatorTokenBudget:
|
||||
async def test_replay_window_overflow_is_archived_even_under_token_budget(
|
||||
self,
|
||||
consolidator,
|
||||
runtime,
|
||||
):
|
||||
"""Old messages that cannot be replayed should be materialized first."""
|
||||
consolidator._SAFETY_BUFFER = 0
|
||||
@@ -222,6 +283,7 @@ class TestConsolidatorTokenBudget:
|
||||
|
||||
await consolidator.maybe_consolidate_by_tokens(
|
||||
session,
|
||||
runtime=runtime,
|
||||
replay_max_messages=6,
|
||||
)
|
||||
|
||||
@@ -235,6 +297,7 @@ class TestConsolidatorTokenBudget:
|
||||
async def test_replay_window_overflow_extends_to_long_recent_user_turn(
|
||||
self,
|
||||
consolidator,
|
||||
runtime,
|
||||
):
|
||||
"""Replay-window consolidation must not cut into the latest user turn."""
|
||||
session = Session(key="test:replay-tool-boundary")
|
||||
@@ -251,6 +314,7 @@ class TestConsolidatorTokenBudget:
|
||||
|
||||
await consolidator.maybe_consolidate_by_tokens(
|
||||
session,
|
||||
runtime=runtime,
|
||||
replay_max_messages=4,
|
||||
)
|
||||
|
||||
@@ -266,6 +330,7 @@ class TestConsolidatorTokenBudget:
|
||||
async def test_replay_window_overflow_uses_newer_user_inside_window(
|
||||
self,
|
||||
consolidator,
|
||||
runtime,
|
||||
):
|
||||
"""Do not extend to an older long turn when the hard window has a newer user."""
|
||||
session = Session(key="test:replay-newer-user")
|
||||
@@ -284,6 +349,7 @@ class TestConsolidatorTokenBudget:
|
||||
|
||||
await consolidator.maybe_consolidate_by_tokens(
|
||||
session,
|
||||
runtime=runtime,
|
||||
replay_max_messages=6,
|
||||
)
|
||||
|
||||
@@ -295,7 +361,7 @@ class TestConsolidatorTokenBudget:
|
||||
history = session.get_history(max_messages=6, extend_to_user=True)
|
||||
assert [m["content"] for m in history] == ["new question", "new answer"]
|
||||
|
||||
async def test_large_chunk_archived_without_cap(self, consolidator):
|
||||
async def test_large_chunk_archived_without_cap(self, consolidator, runtime):
|
||||
"""Without chunk cap, the full range from pick_consolidation_boundary is archived."""
|
||||
consolidator._SAFETY_BUFFER = 0
|
||||
session = MagicMock()
|
||||
@@ -316,14 +382,16 @@ class TestConsolidatorTokenBudget:
|
||||
# (user message at 50, token budget met)
|
||||
consolidator.archive = AsyncMock(return_value=True)
|
||||
|
||||
await consolidator.maybe_consolidate_by_tokens(session)
|
||||
await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime)
|
||||
|
||||
archived_chunk = consolidator.archive.await_args.args[0]
|
||||
# pick_consolidation_boundary returns (50, tokens) — user turn at idx 50
|
||||
assert archived_chunk[0]["content"] == "m0"
|
||||
assert session.last_consolidated > 0
|
||||
|
||||
async def test_raw_archive_fallback_advances_last_consolidated(self, consolidator):
|
||||
async def test_raw_archive_fallback_advances_last_consolidated(
|
||||
self, consolidator, runtime
|
||||
):
|
||||
"""When archive() falls back to raw-archive (LLM failed), the cursor
|
||||
must still advance. Otherwise the same chunk gets raw-archived again
|
||||
on every subsequent maybe_consolidate_by_tokens() call, spamming
|
||||
@@ -344,14 +412,16 @@ class TestConsolidatorTokenBudget:
|
||||
# LLM consolidation fails — archive() returns None (raw_archive fired).
|
||||
consolidator.archive = AsyncMock(return_value=None)
|
||||
|
||||
await consolidator.maybe_consolidate_by_tokens(session)
|
||||
await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime)
|
||||
|
||||
consolidator.archive.assert_awaited_once()
|
||||
# The chunk is considered "materialized" (as a raw-archive breadcrumb),
|
||||
# so last_consolidated must have moved past it.
|
||||
assert session.last_consolidated == 50
|
||||
|
||||
async def test_raw_archive_fallback_breaks_round_loop(self, consolidator):
|
||||
async def test_raw_archive_fallback_breaks_round_loop(
|
||||
self, consolidator, runtime
|
||||
):
|
||||
"""A degraded LLM should not trigger more archive() calls within the
|
||||
same maybe_consolidate_by_tokens invocation — bail after one fallback."""
|
||||
consolidator._SAFETY_BUFFER = 0
|
||||
@@ -370,12 +440,14 @@ class TestConsolidatorTokenBudget:
|
||||
)
|
||||
consolidator.archive = AsyncMock(return_value=None)
|
||||
|
||||
await consolidator.maybe_consolidate_by_tokens(session)
|
||||
await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime)
|
||||
|
||||
# Exactly one fallback per call — not _MAX_CONSOLIDATION_ROUNDS.
|
||||
assert consolidator.archive.await_count == 1
|
||||
|
||||
async def test_boundary_respected_when_no_intermediate_user_turn(self, consolidator):
|
||||
async def test_boundary_respected_when_no_intermediate_user_turn(
|
||||
self, consolidator, runtime
|
||||
):
|
||||
"""When boundary points past a long tool chain, the full chunk is archived."""
|
||||
consolidator._SAFETY_BUFFER = 0
|
||||
session = MagicMock()
|
||||
@@ -394,7 +466,7 @@ class TestConsolidatorTokenBudget:
|
||||
)
|
||||
consolidator.archive = AsyncMock(return_value=True)
|
||||
|
||||
await consolidator.maybe_consolidate_by_tokens(session)
|
||||
await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime)
|
||||
|
||||
consolidator.archive.assert_awaited_once()
|
||||
# pick_consolidation_boundary finds the only boundary at idx=61
|
||||
@@ -412,17 +484,15 @@ class TestCompactIdleSession:
|
||||
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):
|
||||
async def test_archives_prefix_keeps_suffix(
|
||||
self, real_consolidator, mock_provider, runtime
|
||||
):
|
||||
"""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(
|
||||
@@ -437,7 +507,9 @@ class TestCompactIdleSession:
|
||||
session.updated_at = old_ts
|
||||
sessions.save(session)
|
||||
|
||||
result = await real_consolidator.compact_idle_session("cli:test", max_suffix=8)
|
||||
result = await real_consolidator.compact_idle_session(
|
||||
"cli:test", runtime=runtime, max_suffix=8
|
||||
)
|
||||
assert result == "Summary of old conversation."
|
||||
|
||||
reloaded = sessions.get_or_create("cli:test")
|
||||
@@ -451,7 +523,7 @@ class TestCompactIdleSession:
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_summarizes_retained_suffix_not_just_dropped_prefix(
|
||||
self, real_consolidator, mock_provider
|
||||
self, real_consolidator, mock_provider, runtime
|
||||
):
|
||||
"""idleCompact must summarize over the full unconsolidated tail, including
|
||||
the recent suffix it retains. Otherwise a late user correction / final
|
||||
@@ -470,14 +542,16 @@ class TestCompactIdleSession:
|
||||
session.add_message("assistant", "CORRECTED_FINAL_RESULT_alpha")
|
||||
sessions.save(session)
|
||||
|
||||
await real_consolidator.compact_idle_session("cli:correction", max_suffix=8)
|
||||
await real_consolidator.compact_idle_session(
|
||||
"cli:correction", runtime=runtime, max_suffix=8
|
||||
)
|
||||
|
||||
summarized = mock_provider.chat_with_retry.call_args.kwargs["messages"][1]["content"]
|
||||
assert "CORRECTED_FINAL_RESULT_alpha" in summarized
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_raw_dumps_only_dropped_messages_on_llm_failure(
|
||||
self, real_consolidator, mock_provider, store
|
||||
self, real_consolidator, mock_provider, store, runtime
|
||||
):
|
||||
"""Summarizing over the full tail must not widen what gets raw-dumped on
|
||||
LLM failure: the breadcrumb should contain only the removed prefix, not
|
||||
@@ -492,7 +566,9 @@ class TestCompactIdleSession:
|
||||
session.add_message("assistant", "RETAINED_SUFFIX_marker")
|
||||
sessions.save(session)
|
||||
|
||||
await real_consolidator.compact_idle_session("cli:rawdrop", max_suffix=8)
|
||||
await real_consolidator.compact_idle_session(
|
||||
"cli:rawdrop", runtime=runtime, max_suffix=8
|
||||
)
|
||||
|
||||
raw = "\n".join(e["content"] for e in store.read_unprocessed_history(since_cursor=0))
|
||||
assert "[RAW]" in raw
|
||||
@@ -505,6 +581,7 @@ class TestCompactIdleSession:
|
||||
real_consolidator,
|
||||
mock_provider,
|
||||
store,
|
||||
runtime,
|
||||
):
|
||||
mock_provider.chat_with_retry.return_value = MagicMock(
|
||||
content="Summary of old conversation.", finish_reason="stop"
|
||||
@@ -515,14 +592,16 @@ class TestCompactIdleSession:
|
||||
session.add_message("assistant", f"assistant msg {i}")
|
||||
real_consolidator.sessions.save(session)
|
||||
|
||||
await real_consolidator.compact_idle_session("cli:test", max_suffix=4)
|
||||
await real_consolidator.compact_idle_session(
|
||||
"cli:test", runtime=runtime, max_suffix=4
|
||||
)
|
||||
|
||||
entries = store.read_unprocessed_history(since_cursor=0)
|
||||
assert entries[0]["session_key"] == "cli:test"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_session_does_not_refresh_timestamp(
|
||||
self, real_consolidator
|
||||
self, real_consolidator, runtime
|
||||
):
|
||||
"""Empty session with old updated_at does not look active after compaction."""
|
||||
from datetime import datetime, timedelta
|
||||
@@ -533,7 +612,9 @@ class TestCompactIdleSession:
|
||||
session.updated_at = old_ts
|
||||
sessions.save(session)
|
||||
|
||||
result = await real_consolidator.compact_idle_session("cli:empty")
|
||||
result = await real_consolidator.compact_idle_session(
|
||||
"cli:empty", runtime=runtime
|
||||
)
|
||||
assert result == ""
|
||||
|
||||
reloaded = sessions.get_or_create("cli:empty")
|
||||
@@ -541,7 +622,9 @@ class TestCompactIdleSession:
|
||||
assert reloaded.metadata == {}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_nothing_summary_not_stored(self, real_consolidator, mock_provider):
|
||||
async def test_nothing_summary_not_stored(
|
||||
self, real_consolidator, mock_provider, runtime
|
||||
):
|
||||
"""LLM returns '(nothing)' → _last_summary NOT in metadata."""
|
||||
mock_provider.chat_with_retry.return_value = MagicMock(
|
||||
content="(nothing)", finish_reason="stop"
|
||||
@@ -553,14 +636,18 @@ class TestCompactIdleSession:
|
||||
session.add_message("assistant", f"a{i}")
|
||||
sessions.save(session)
|
||||
|
||||
result = await real_consolidator.compact_idle_session("cli:nothing", max_suffix=4)
|
||||
result = await real_consolidator.compact_idle_session(
|
||||
"cli:nothing", runtime=runtime, 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):
|
||||
async def test_llm_failure_still_truncates(
|
||||
self, real_consolidator, mock_provider, store, runtime
|
||||
):
|
||||
"""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
|
||||
@@ -570,7 +657,9 @@ class TestCompactIdleSession:
|
||||
session.add_message("assistant", f"a{i}")
|
||||
sessions.save(session)
|
||||
|
||||
result = await real_consolidator.compact_idle_session("cli:fail", max_suffix=4)
|
||||
result = await real_consolidator.compact_idle_session(
|
||||
"cli:fail", runtime=runtime, max_suffix=4
|
||||
)
|
||||
assert result is None
|
||||
|
||||
# raw_archive should have been called (history.jsonl gets an entry)
|
||||
@@ -582,7 +671,9 @@ class TestCompactIdleSession:
|
||||
assert len(reloaded.messages) <= 4
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_respects_last_consolidated(self, real_consolidator, mock_provider):
|
||||
async def test_respects_last_consolidated(
|
||||
self, real_consolidator, mock_provider, runtime
|
||||
):
|
||||
"""30 turns with last_consolidated=50 → only unconsolidated tail considered."""
|
||||
mock_provider.chat_with_retry.return_value = MagicMock(
|
||||
content="Tail summary.", finish_reason="stop"
|
||||
@@ -595,7 +686,9 @@ class TestCompactIdleSession:
|
||||
session.last_consolidated = 50 # Only 10 messages unconsolidated
|
||||
sessions.save(session)
|
||||
|
||||
result = await real_consolidator.compact_idle_session("cli:offset", max_suffix=4)
|
||||
result = await real_consolidator.compact_idle_session(
|
||||
"cli:offset", runtime=runtime, max_suffix=4
|
||||
)
|
||||
assert result == "Tail summary."
|
||||
|
||||
# Verify only the unconsolidated tail was processed:
|
||||
@@ -611,6 +704,7 @@ class TestCompactIdleSession:
|
||||
self,
|
||||
real_consolidator,
|
||||
mock_provider,
|
||||
runtime,
|
||||
):
|
||||
"""Assistant-only tails extend back to the latest user turn, so archive
|
||||
the actual dropped messages rather than a computed prefix."""
|
||||
@@ -625,7 +719,9 @@ class TestCompactIdleSession:
|
||||
session.add_message("assistant", f"assistant-{i:02d}")
|
||||
sessions.save(session)
|
||||
|
||||
result = await real_consolidator.compact_idle_session("cli:noncontiguous", max_suffix=6)
|
||||
result = await real_consolidator.compact_idle_session(
|
||||
"cli:noncontiguous", runtime=runtime, max_suffix=6
|
||||
)
|
||||
assert result == "Tail summary."
|
||||
|
||||
reloaded = sessions.get_or_create("cli:noncontiguous")
|
||||
@@ -653,7 +749,9 @@ class TestCompactIdleSession:
|
||||
assert "user-14" in user_content
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acquires_consolidation_lock(self, real_consolidator, mock_provider):
|
||||
async def test_acquires_consolidation_lock(
|
||||
self, real_consolidator, mock_provider, runtime
|
||||
):
|
||||
"""Verify lock is held during execution."""
|
||||
import asyncio
|
||||
|
||||
@@ -679,7 +777,9 @@ class TestCompactIdleSession:
|
||||
assert not lock.locked()
|
||||
|
||||
task = asyncio.ensure_future(
|
||||
real_consolidator.compact_idle_session("cli:lock", max_suffix=4)
|
||||
real_consolidator.compact_idle_session(
|
||||
"cli:lock", runtime=runtime, max_suffix=4
|
||||
)
|
||||
)
|
||||
await started.wait()
|
||||
assert lock.locked()
|
||||
@@ -702,15 +802,17 @@ class TestConsolidatorSessionRefresh:
|
||||
provider.chat_with_retry = AsyncMock(
|
||||
return_value=MagicMock(content="summary", finish_reason="stop")
|
||||
)
|
||||
provider.generation.max_tokens = 4096
|
||||
provider.generation = GenerationSettings(max_tokens=4096)
|
||||
provider.estimate_prompt_tokens = MagicMock(return_value=(10, "test"))
|
||||
runtime = LLMRuntime.capture(
|
||||
provider,
|
||||
"test-model",
|
||||
context_window_tokens=128_000,
|
||||
)
|
||||
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=[]),
|
||||
)
|
||||
@@ -722,13 +824,16 @@ class TestConsolidatorSessionRefresh:
|
||||
|
||||
seen: dict[str, Session] = {}
|
||||
|
||||
def estimate(session: Session):
|
||||
def estimate(session: Session, *, runtime):
|
||||
seen["session"] = session
|
||||
return 10, "test"
|
||||
|
||||
consolidator.estimate_session_prompt_tokens = MagicMock(side_effect=estimate)
|
||||
|
||||
await consolidator.maybe_consolidate_by_tokens(stale_empty)
|
||||
await consolidator.maybe_consolidate_by_tokens(
|
||||
stale_empty,
|
||||
runtime=runtime,
|
||||
)
|
||||
|
||||
assert seen["session"] is fresh
|
||||
|
||||
@@ -745,15 +850,17 @@ class TestConsolidatorSessionRefresh:
|
||||
provider.chat_with_retry = AsyncMock(
|
||||
return_value=MagicMock(content="summary", finish_reason="stop")
|
||||
)
|
||||
provider.generation.max_tokens = 4096
|
||||
provider.generation = GenerationSettings(max_tokens=4096)
|
||||
provider.estimate_prompt_tokens = MagicMock(return_value=(10, "test"))
|
||||
runtime = LLMRuntime.capture(
|
||||
provider,
|
||||
"test-model",
|
||||
context_window_tokens=128_000,
|
||||
)
|
||||
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=[]),
|
||||
)
|
||||
@@ -769,11 +876,18 @@ class TestConsolidatorSessionRefresh:
|
||||
old_ref = session
|
||||
|
||||
# AutoCompact runs first and truncates to 8
|
||||
await consolidator.compact_idle_session("cli:test", max_suffix=8)
|
||||
await consolidator.compact_idle_session(
|
||||
"cli:test",
|
||||
runtime=runtime,
|
||||
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)
|
||||
await consolidator.maybe_consolidate_by_tokens(
|
||||
old_ref,
|
||||
runtime=runtime,
|
||||
)
|
||||
|
||||
session_after = sessions.get_or_create("cli:test")
|
||||
# Messages should still be truncated (not restored to 40)
|
||||
@@ -818,7 +932,9 @@ class TestRawArchiveTruncation:
|
||||
class TestArchiveTruncation:
|
||||
"""archive() must truncate formatted text before sending to consolidation LLM."""
|
||||
|
||||
async def test_archive_truncates_large_formatted_text(self, consolidator, mock_provider, store):
|
||||
async def test_archive_truncates_large_formatted_text(
|
||||
self, consolidator, mock_provider, store, runtime
|
||||
):
|
||||
"""Large formatted text should be truncated to token budget before LLM call."""
|
||||
# context_window_tokens=1000, max_completion_tokens=100, _SAFETY_BUFFER=1024
|
||||
# budget = 1000 - 100 - 1024 = -124 → fallback via truncate_text(budget*4)
|
||||
@@ -826,21 +942,23 @@ class TestArchiveTruncation:
|
||||
mock_provider.chat_with_retry.return_value = MagicMock(
|
||||
content="Summary of large input.", finish_reason="stop"
|
||||
)
|
||||
await consolidator.archive(big_messages)
|
||||
await consolidator.archive(big_messages, runtime=runtime)
|
||||
|
||||
call_args = mock_provider.chat_with_retry.call_args
|
||||
user_content = call_args.kwargs["messages"][1]["content"]
|
||||
# Should be significantly shorter than 100K
|
||||
assert len(user_content) < 50_000
|
||||
|
||||
async def test_archive_truncates_with_small_token_budget(self, consolidator, mock_provider, store):
|
||||
async def test_archive_truncates_with_small_token_budget(
|
||||
self, consolidator, mock_provider, store, runtime
|
||||
):
|
||||
"""Small context window: truncation uses actual tokenizer count."""
|
||||
consolidator.context_window_tokens = 500
|
||||
runtime = replace(runtime, context_window_tokens=500)
|
||||
big_messages = [{"role": "user", "content": "word " * 50_000}]
|
||||
mock_provider.chat_with_retry.return_value = MagicMock(
|
||||
content="Summary.", finish_reason="stop"
|
||||
)
|
||||
await consolidator.archive(big_messages)
|
||||
await consolidator.archive(big_messages, runtime=runtime)
|
||||
|
||||
sent_messages = mock_provider.chat_with_retry.call_args.kwargs["messages"]
|
||||
user_content = sent_messages[1]["content"]
|
||||
@@ -848,7 +966,9 @@ class TestArchiveTruncation:
|
||||
# Should be truncated
|
||||
assert len(user_content) < 250_000
|
||||
|
||||
async def test_oversized_summary_is_capped_before_append(self, consolidator, mock_provider, store):
|
||||
async def test_oversized_summary_is_capped_before_append(
|
||||
self, consolidator, mock_provider, store, runtime
|
||||
):
|
||||
"""A pathologically large LLM summary must not land full-length in
|
||||
history.jsonl — that would re-open the #3412 bloat vector from the
|
||||
*success* path instead of the fallback path."""
|
||||
@@ -856,21 +976,26 @@ class TestArchiveTruncation:
|
||||
content="S" * (_ARCHIVE_SUMMARY_MAX_CHARS * 10),
|
||||
finish_reason="stop",
|
||||
)
|
||||
await consolidator.archive([{"role": "user", "content": "hi"}])
|
||||
await consolidator.archive(
|
||||
[{"role": "user", "content": "hi"}],
|
||||
runtime=runtime,
|
||||
)
|
||||
|
||||
entry = store.read_unprocessed_history(since_cursor=0)[0]
|
||||
assert len(entry["content"]) <= _ARCHIVE_SUMMARY_MAX_CHARS + 50
|
||||
|
||||
async def test_archive_truncates_via_tiktoken_with_positive_budget(self, consolidator, mock_provider, store):
|
||||
async def test_archive_truncates_via_tiktoken_with_positive_budget(
|
||||
self, consolidator, mock_provider, store, runtime
|
||||
):
|
||||
"""Positive token budget should use tiktoken for precise truncation."""
|
||||
consolidator.context_window_tokens = 10_000
|
||||
runtime = replace(runtime, context_window_tokens=10_000)
|
||||
consolidator._SAFETY_BUFFER = 0
|
||||
# budget = 10000 - 100 - 0 = 9900 tokens
|
||||
big_messages = [{"role": "user", "content": "word " * 50_000}]
|
||||
mock_provider.chat_with_retry.return_value = MagicMock(
|
||||
content="Summary.", finish_reason="stop"
|
||||
)
|
||||
await consolidator.archive(big_messages)
|
||||
await consolidator.archive(big_messages, runtime=runtime)
|
||||
|
||||
import tiktoken
|
||||
enc = tiktoken.get_encoding("cl100k_base")
|
||||
|
||||
Reference in New Issue
Block a user