fix(session): keep auto compact suffix on user turn
This commit is contained in:
@@ -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
@@ -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:
|
||||
|
||||
@@ -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."""
|
||||
|
||||
@@ -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) ---
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user