diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index 050f0cf7..7b281df1 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -1172,6 +1172,7 @@ class AgentLoop: "max_messages": self._max_messages, "max_tokens": self._replay_token_budget(), "include_timestamps": True, + "extend_to_user": True, } history = session.get_history(**_hist_kwargs) current_role = "assistant" if is_subagent else "user" @@ -1447,6 +1448,7 @@ class AgentLoop: "max_messages": self._max_messages, "max_tokens": self._replay_token_budget(), "include_timestamps": True, + "extend_to_user": True, } ctx.history = ctx.session.get_history(**_hist_kwargs) self._runtime_events().record_turn_runtime( diff --git a/nanobot/agent/memory.py b/nanobot/agent/memory.py index 56588bba..02453a04 100644 --- a/nanobot/agent/memory.py +++ b/nanobot/agent/memory.py @@ -22,6 +22,7 @@ from nanobot.utils.helpers import ( estimate_message_tokens, estimate_prompt_tokens_chain, find_legal_message_start, + recent_message_start_index, strip_think, truncate_text, truncate_text_to_tokens, @@ -717,7 +718,13 @@ class Consolidator: if len(tail) <= replay_max_messages: return None - sliced = tail[-replay_max_messages:] + tail_messages = [message for _idx, message in tail] + start_idx = recent_message_start_index( + tail_messages, + replay_max_messages, + extend_to_user=True, + ) + sliced = tail[start_idx:] for i, (_idx, message) in enumerate(sliced): if message.get("role") == "user": start = i diff --git a/nanobot/session/manager.py b/nanobot/session/manager.py index 622e5e55..dbf79757 100644 --- a/nanobot/session/manager.py +++ b/nanobot/session/manager.py @@ -19,6 +19,7 @@ from nanobot.utils.helpers import ( estimate_message_tokens, find_legal_message_start, image_placeholder_text, + recent_message_start_index, safe_filename, strip_think, ) @@ -153,6 +154,7 @@ class Session: *, max_tokens: int = 0, include_timestamps: bool = False, + extend_to_user: bool = False, ) -> list[dict[str, Any]]: """Return unconsolidated messages for LLM input. @@ -161,7 +163,12 @@ class Session: """ unconsolidated = self.messages[self.last_consolidated:] max_messages = max_messages if max_messages > 0 else 120 - sliced = unconsolidated[-max_messages:] + start_idx = recent_message_start_index( + unconsolidated, + max_messages, + extend_to_user=extend_to_user, + ) + sliced = unconsolidated[start_idx:] # Avoid starting mid-turn when possible, except for proactive # assistant deliveries that the user may be replying to. diff --git a/nanobot/utils/helpers.py b/nanobot/utils/helpers.py index c7a91821..34177bef 100644 --- a/nanobot/utils/helpers.py +++ b/nanobot/utils/helpers.py @@ -270,6 +270,30 @@ def truncate_text_to_tokens(text: str, max_tokens: int) -> str: return truncate_text(text, max_chars - suffix_chars) +def recent_message_start_index( + messages: list[dict[str, Any]], + max_messages: int, + *, + extend_to_user: bool = False, +) -> int: + """Return the start index for a recent replay window.""" + if max_messages <= 0: + return len(messages) + start_idx = max(0, len(messages) - max_messages) + if not extend_to_user or len(messages) <= max_messages: + return start_idx + + recovered_user = next( + (i for i in range(start_idx, -1, -1) if messages[i].get("role") == "user"), + None, + ) + if recovered_user is None: + return start_idx + if recovered_user > 0 and messages[recovered_user - 1].get("_channel_delivery"): + return recovered_user - 1 + return recovered_user + + def find_legal_message_start(messages: list[dict[str, Any]]) -> int: """Find the first index whose tool results have matching assistant calls.""" declared: set[str] = set() diff --git a/tests/agent/test_consolidator.py b/tests/agent/test_consolidator.py index dafb1d40..da79f600 100644 --- a/tests/agent/test_consolidator.py +++ b/tests/agent/test_consolidator.py @@ -48,6 +48,19 @@ def consolidator(store, mock_provider): ) +def _tool_round(call_id: str) -> list[dict]: + return [ + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": call_id, "type": "function", "function": {"name": "x", "arguments": "{}"}} + ], + }, + {"role": "tool", "tool_call_id": call_id, "name": "x", "content": "ok"}, + ] + + class TestConsolidatorSummarize: async def test_summarize_appends_to_history(self, consolidator, mock_provider, store): """Consolidator should call LLM to summarize, then append to HISTORY.md.""" @@ -219,21 +232,17 @@ class TestConsolidatorTokenBudget: assert session.metadata["_last_summary"]["text"] == "old conversation summary" consolidator.sessions.save.assert_called() - async def test_replay_window_overflow_matches_history_tool_boundary( + async def test_replay_window_overflow_extends_to_long_recent_user_turn( self, consolidator, ): - """Archive the exact prefix hidden by get_history's legal-start trimming.""" + """Replay-window consolidation must not cut into the latest user turn.""" session = Session(key="test:replay-tool-boundary") - session.add_message("user", "run the tool") - session.add_message( - "assistant", - "", - tool_calls=[ - {"id": "call-1", "type": "function", "function": {"name": "x", "arguments": "{}"}} - ], - ) - session.add_message("tool", "tool result", tool_call_id="call-1", name="x") + session.add_message("user", "old") + session.add_message("assistant", "old answer") + session.add_message("user", "record this") + for i in range(4): + session.messages.extend(_tool_round(f"call-{i}")) session.add_message("assistant", "final answer") consolidator.sessions._session_cache[session.key] = session @@ -242,13 +251,17 @@ class TestConsolidatorTokenBudget: await consolidator.maybe_consolidate_by_tokens( session, - replay_max_messages=2, + replay_max_messages=4, ) archived_chunk = consolidator.archive.await_args.args[0] - assert [m["role"] for m in archived_chunk] == ["user", "assistant", "tool"] - assert session.last_consolidated == 3 - assert session.get_history(max_messages=2) == [{"role": "assistant", "content": "final answer"}] + assert [m["content"] for m in archived_chunk] == ["old", "old answer"] + assert session.last_consolidated == 2 + + history = session.get_history(max_messages=4, extend_to_user=True) + assert len(history) > 4 + assert history[0]["content"] == "record this" + assert history[-1]["content"] == "final answer" async def test_large_chunk_archived_without_cap(self, consolidator): """Without chunk cap, the full range from pick_consolidation_boundary is archived.""" diff --git a/tests/agent/test_max_messages_config.py b/tests/agent/test_max_messages_config.py index 491ff16e..2a598ce4 100644 --- a/tests/agent/test_max_messages_config.py +++ b/tests/agent/test_max_messages_config.py @@ -111,6 +111,7 @@ class TestMaxMessagesIntegration: assert result is not None assert mock_hist.call_count == 1 assert mock_hist.call_args.kwargs["max_messages"] == 25 + assert mock_hist.call_args.kwargs["extend_to_user"] is True @pytest.mark.asyncio async def test_zero_config_passes_builtin_limit_to_history_call(self, tmp_path: Path) -> None: @@ -129,6 +130,7 @@ class TestMaxMessagesIntegration: assert result is not None assert mock_hist.call_args.kwargs["max_messages"] == DEFAULT_MAX_MESSAGES + assert mock_hist.call_args.kwargs["extend_to_user"] is True class TestSchemaConfig: diff --git a/tests/agent/test_session_manager_history.py b/tests/agent/test_session_manager_history.py index 58c87e6a..4b8fbbdf 100644 --- a/tests/agent/test_session_manager_history.py +++ b/tests/agent/test_session_manager_history.py @@ -641,6 +641,25 @@ def test_retain_recent_legal_suffix_can_extend_to_user_for_long_recent_turn(): _assert_no_orphans(history) +def test_get_history_can_extend_to_user_for_long_recent_turn(): + session = Session(key="test:history-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"}) + + hard_capped = session.get_history(max_messages=8) + extended = session.get_history(max_messages=8, extend_to_user=True) + + assert len(hard_capped) <= 8 + assert len(extended) > 8 + assert extended[0]["content"] == "record this" + assert extended[-1]["content"] == "done" + _assert_no_orphans(extended) + + # --- enforce_file_cap archive correctness (issue #4128) ---