refactor(agent): require runtime for consolidation
This commit is contained in:
@@ -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]
|
||||
|
||||
Reference in New Issue
Block a user