refactor(agent): require runtime for consolidation

This commit is contained in:
chengyongru
2026-07-10 17:54:34 +08:00
committed by Xubin Ren
parent 5bd3d1e0af
commit c9d3e74342
18 changed files with 486 additions and 214 deletions
+29 -10
View File
@@ -76,7 +76,10 @@ async def test_prompt_above_threshold_archives_until_next_user_boundary(tmp_path
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(session)
await loop.consolidator.maybe_consolidate_by_tokens(
session,
runtime=loop.llm_runtime(),
)
archived_chunk = loop.consolidator.archive.await_args.args[0]
assert [message["content"] for message in archived_chunk] == ["u1", "a1", "u2", "a2"]
@@ -102,7 +105,7 @@ async def test_consolidation_loops_until_target_met(tmp_path, monkeypatch) -> No
loop.sessions.save(session)
call_count = [0]
def mock_estimate(_session, *, session_summary=None):
def mock_estimate(_session, *, runtime):
call_count[0] += 1
if call_count[0] == 1:
return (500, "test")
@@ -113,7 +116,10 @@ async def test_consolidation_loops_until_target_met(tmp_path, monkeypatch) -> No
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)
await loop.consolidator.maybe_consolidate_by_tokens(
session,
runtime=loop.llm_runtime(),
)
assert loop.consolidator.archive.await_count == 2
assert session.last_consolidated == 6
@@ -139,7 +145,7 @@ async def test_consolidation_continues_below_trigger_until_half_target(tmp_path,
call_count = [0]
def mock_estimate(_session, *, session_summary=None):
def mock_estimate(_session, *, runtime):
call_count[0] += 1
if call_count[0] == 1:
return (500, "test")
@@ -150,7 +156,10 @@ async def test_consolidation_continues_below_trigger_until_half_target(tmp_path,
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)
await loop.consolidator.maybe_consolidate_by_tokens(
session,
runtime=loop.llm_runtime(),
)
assert loop.consolidator.archive.await_count == 2
assert session.last_consolidated == 6
@@ -171,7 +180,7 @@ async def test_consolidation_persists_summary_for_next_prepare_session(tmp_path,
call_count = [0]
def mock_estimate(_session, *, session_summary=None):
def mock_estimate(_session, *, runtime):
call_count[0] += 1
if call_count[0] == 1:
return (500, "test")
@@ -180,7 +189,10 @@ async def test_consolidation_persists_summary_for_next_prepare_session(tmp_path,
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(session)
await loop.consolidator.maybe_consolidate_by_tokens(
session,
runtime=loop.llm_runtime(),
)
reloaded = loop.sessions.get_or_create("cli:test")
meta = reloaded.metadata.get("_last_summary")
@@ -204,12 +216,19 @@ async def test_preflight_consolidation_receives_pending_summary(tmp_path) -> Non
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=None) # type: ignore[method-assign]
loop._schedule_background = lambda coro: coro.close() # type: ignore[method-assign]
await loop.process_direct("hello", session_key="cli:test")
runtime = loop.llm_runtime()
await loop.process_direct("hello", session_key="cli:test", runtime=runtime)
loop.consolidator.maybe_consolidate_by_tokens.assert_any_await(
session,
runtime=runtime,
replay_max_messages=loop._max_messages,
)
assert len(loop.consolidator.maybe_consolidate_by_tokens.call_args_list) == 2
assert all(
call.kwargs["runtime"] is runtime
for call in loop.consolidator.maybe_consolidate_by_tokens.call_args_list
)
@pytest.mark.asyncio
@@ -221,7 +240,7 @@ async def test_preflight_consolidation_before_llm_call(tmp_path, monkeypatch) ->
archived_session_keys: list[str | None] = []
async def track_consolidate(messages, *, session_key=None):
async def track_consolidate(messages, *, runtime, session_key=None):
order.append("consolidate")
archived_session_keys.append(session_key)
return True
@@ -244,7 +263,7 @@ async def test_preflight_consolidation_before_llm_call(tmp_path, monkeypatch) ->
monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda _m: 500)
call_count = [0]
def mock_estimate(_session, *, session_summary=None):
def mock_estimate(_session, *, runtime):
call_count[0] += 1
return (1000 if call_count[0] <= 1 else 80, "test")
loop.consolidator.estimate_session_prompt_tokens = mock_estimate # type: ignore[method-assign]