fix(session): keep auto compact suffix on user turn

This commit is contained in:
chengyongru
2026-06-15 19:03:37 +08:00
committed by Xubin Ren
parent 9814a3b9fe
commit 3ce0cd972e
5 changed files with 110 additions and 17 deletions
+1 -1
View File
@@ -1006,7 +1006,7 @@ class Consolidator:
metadata={},
last_consolidated=0,
)
dropped, already_consolidated = probe.retain_recent_legal_suffix(max_suffix)
dropped, already_consolidated = probe.retain_recent_legal_suffix(max_suffix, extend_to_user=True)
messages_to_keep = probe.messages
messages_to_remove = dropped[already_consolidated:]
+22 -10
View File
@@ -287,8 +287,13 @@ class Session:
self.updated_at = datetime.now()
self.metadata.pop("_last_summary", None)
def retain_recent_legal_suffix(self, max_messages: int) -> tuple[list[dict], int]:
"""Keep a legal recent suffix constrained by a hard message cap.
def retain_recent_legal_suffix(
self,
max_messages: int,
*,
extend_to_user: bool = False,
) -> tuple[list[dict], int]:
"""Keep a legal recent suffix, optionally extending it back to a user turn.
Returns ``(dropped, already_consolidated_count)`` where *dropped* is
the list of removed messages (in original order) and
@@ -307,30 +312,37 @@ class Session:
original = list(self.messages)
before_lc = self.last_consolidated
retained = list(self.messages[-max_messages:])
start_idx = max(0, len(self.messages) - max_messages)
if extend_to_user:
start_idx = next(
(i for i in range(start_idx, -1, -1) if self.messages[i].get("role") == "user"),
start_idx,
)
# Prefer starting at a user turn when one exists within the tail.
retained = self.messages[start_idx:]
# Prefer starting at a user turn when one exists within the retained window.
first_user = next((i for i, m in enumerate(retained) if m.get("role") == "user"), None)
if first_user is not None:
retained = retained[first_user:]
else:
# If the tail is assistant/tool-only, anchor to the latest user in
# the full session and take a capped forward window from there.
elif not extend_to_user:
# If the hard-capped tail is assistant/tool-only, anchor to the
# latest user in the full session and take a capped forward window.
latest_user = next(
(i for i in range(len(self.messages) - 1, -1, -1)
if self.messages[i].get("role") == "user"),
None,
)
if latest_user is not None:
retained = list(self.messages[latest_user: latest_user + max_messages])
retained = self.messages[latest_user: latest_user + max_messages]
# Mirror get_history(): avoid persisting orphan tool results at the front.
start = find_legal_message_start(retained)
if start:
retained = retained[start:]
# Hard-cap guarantee: never keep more than max_messages.
if len(retained) > max_messages:
# Hard-cap guarantee unless the caller requested user-turn extension.
if not extend_to_user and len(retained) > max_messages:
retained = retained[-max_messages:]
start = find_legal_message_start(retained)
if start:
+60 -1
View File
@@ -45,6 +45,33 @@ def _add_turns(session, turns: int, *, prefix: str = "msg") -> None:
session.add_message("assistant", f"{prefix} assistant {i}")
def _add_tool_turn(session, prefix: str, idx: int) -> None:
call_id = f"{prefix}_{idx}"
session.messages.append(
{
"role": "assistant",
"content": "",
"tool_calls": [
{
"id": call_id,
"type": "function",
"function": {"name": "exec", "arguments": "{}"},
}
],
"timestamp": datetime.now().isoformat(),
}
)
session.messages.append(
{
"role": "tool",
"tool_call_id": call_id,
"name": "exec",
"content": "ok",
"timestamp": datetime.now().isoformat(),
}
)
def _make_fake_compact(
loop: AgentLoop,
*,
@@ -76,7 +103,10 @@ def _make_fake_compact(
metadata={},
last_consolidated=0,
)
dropped, already_consolidated = probe.retain_recent_legal_suffix(max_suffix)
dropped, already_consolidated = probe.retain_recent_legal_suffix(
max_suffix,
extend_to_user=True,
)
kept = probe.messages
archive_msgs = dropped[already_consolidated:]
@@ -305,6 +335,35 @@ class TestAutoCompact:
assert session_after.messages[-1]["content"] == "msg assistant 5"
await loop.close_mcp()
@pytest.mark.asyncio
async def test_auto_compact_extends_recent_suffix_to_user_turn(self, tmp_path):
loop = _make_loop(tmp_path, session_ttl_minutes=15)
session = loop.sessions.get_or_create("cli:test")
_add_turns(session, 2, prefix="old")
session.add_message("user", "record this")
for i in range(8):
_add_tool_turn(session, "recent", i)
session.add_message("assistant", "done")
loop.sessions.save(session)
await loop.auto_compact._archive("cli:test")
session_after = loop.sessions.get_or_create("cli:test")
assert len(session_after.messages) > loop.auto_compact._RECENT_SUFFIX_MESSAGES
assert session_after.messages[0]["content"] == "record this"
assert session_after.messages[-1]["content"] == "done"
tool_results = {
m.get("tool_call_id")
for m in session_after.messages
if m.get("role") == "tool"
}
assert all(
tc["id"] in tool_results
for m in session_after.messages
for tc in (m.get("tool_calls") or [])
)
await loop.close_mcp()
@pytest.mark.asyncio
async def test_auto_compact_stores_summary(self, tmp_path):
"""_archive should store the summary in _summaries."""
+9 -5
View File
@@ -561,8 +561,8 @@ class TestCompactIdleSession:
real_consolidator,
mock_provider,
):
"""Assistant-only tails retain a non-contiguous slice, so archive the
actual dropped messages rather than a computed prefix."""
"""Assistant-only tails extend back to the latest user turn, so archive
the actual dropped messages rather than a computed prefix."""
mock_provider.chat_with_retry.return_value = MagicMock(
content="Tail summary.", finish_reason="stop"
)
@@ -585,12 +585,16 @@ class TestCompactIdleSession:
"assistant-02",
"assistant-03",
"assistant-04",
"assistant-05",
"assistant-06",
"assistant-07",
"assistant-08",
"assistant-09",
]
# #4264: idle compaction now summarizes the full unconsolidated tail, so
# the dropped head (user-00), the non-contiguous dropped tail
# (assistant-09), and the retained suffix (user-14) are all summarized.
# Retention above still proves the non-contiguous suffix is handled.
# the dropped head (user-00) and retained suffix (user-14 through
# assistant-09) are all summarized.
archived_call = mock_provider.chat_with_retry.call_args
user_content = archived_call.kwargs["messages"][1]["content"]
assert "user-00" in user_content
@@ -623,6 +623,24 @@ def test_retain_recent_legal_suffix_hard_cap_with_long_non_user_chain():
assert len(session.messages) <= 6
def test_retain_recent_legal_suffix_can_extend_to_user_for_long_recent_turn():
session = Session(key="test:extend-to-user")
session.messages.append({"role": "user", "content": "old"})
session.messages.append({"role": "assistant", "content": "old answer"})
session.messages.append({"role": "user", "content": "record this"})
for i in range(4):
session.messages.extend(_tool_turn("recent", i))
session.messages.append({"role": "assistant", "content": "done"})
session.retain_recent_legal_suffix(8, extend_to_user=True)
assert len(session.messages) > 8
assert session.messages[0]["content"] == "record this"
assert session.messages[-1]["content"] == "done"
history = session.get_history(max_messages=500)
_assert_no_orphans(history)
# --- enforce_file_cap archive correctness (issue #4128) ---