From 3ce0cd972eee57318bb7eb472494c00367cb22a1 Mon Sep 17 00:00:00 2001 From: chengyongru Date: Mon, 15 Jun 2026 18:03:14 +0800 Subject: [PATCH] fix(session): keep auto compact suffix on user turn --- nanobot/agent/memory.py | 2 +- nanobot/session/manager.py | 32 +++++++---- tests/agent/test_auto_compact.py | 61 ++++++++++++++++++++- tests/agent/test_consolidator.py | 14 +++-- tests/agent/test_session_manager_history.py | 18 ++++++ 5 files changed, 110 insertions(+), 17 deletions(-) diff --git a/nanobot/agent/memory.py b/nanobot/agent/memory.py index f8d93f7e..fb08e371 100644 --- a/nanobot/agent/memory.py +++ b/nanobot/agent/memory.py @@ -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:] diff --git a/nanobot/session/manager.py b/nanobot/session/manager.py index 5f0af684..622e5e55 100644 --- a/nanobot/session/manager.py +++ b/nanobot/session/manager.py @@ -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: diff --git a/tests/agent/test_auto_compact.py b/tests/agent/test_auto_compact.py index ceada74c..e6293ebe 100644 --- a/tests/agent/test_auto_compact.py +++ b/tests/agent/test_auto_compact.py @@ -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.""" diff --git a/tests/agent/test_consolidator.py b/tests/agent/test_consolidator.py index a87445a6..33754eb7 100644 --- a/tests/agent/test_consolidator.py +++ b/tests/agent/test_consolidator.py @@ -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 diff --git a/tests/agent/test_session_manager_history.py b/tests/agent/test_session_manager_history.py index 3441c483..58c87e6a 100644 --- a/tests/agent/test_session_manager_history.py +++ b/tests/agent/test_session_manager_history.py @@ -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) ---