From 2056061765895e8a3fddd9b98899eb6845307ba5 Mon Sep 17 00:00:00 2001 From: Xubin Ren Date: Mon, 23 Mar 2026 16:27:20 +0000 Subject: [PATCH] refine heartbeat session retention boundaries --- nanobot/cli/commands.py | 9 ++--- nanobot/config/schema.py | 1 + nanobot/session/manager.py | 26 ++++++++++++++ tests/test_commands.py | 6 ++++ tests/test_session_manager_history.py | 52 +++++++++++++++++++++++++++ 5 files changed, 90 insertions(+), 4 deletions(-) diff --git a/nanobot/cli/commands.py b/nanobot/cli/commands.py index 372056a..acea2db 100644 --- a/nanobot/cli/commands.py +++ b/nanobot/cli/commands.py @@ -619,12 +619,13 @@ def gateway( chat_id=chat_id, on_progress=_silent, ) - - # Clear the heartbeat session to prevent token overflow from accumulated tasks + + # Keep a small tail of heartbeat history so the loop stays bounded + # without losing all short-term context between runs. session = agent.sessions.get_or_create("heartbeat") - session.clear() + session.retain_recent_legal_suffix(hb_cfg.keep_recent_messages) agent.sessions.save(session) - + return resp.content if resp else "" async def on_heartbeat_notify(response: str) -> None: diff --git a/nanobot/config/schema.py b/nanobot/config/schema.py index 58ead15..7d8f5c8 100644 --- a/nanobot/config/schema.py +++ b/nanobot/config/schema.py @@ -90,6 +90,7 @@ class HeartbeatConfig(Base): enabled: bool = True interval_s: int = 30 * 60 # 30 minutes + keep_recent_messages: int = 8 class GatewayConfig(Base): diff --git a/nanobot/session/manager.py b/nanobot/session/manager.py index f8244e5..537ba42 100644 --- a/nanobot/session/manager.py +++ b/nanobot/session/manager.py @@ -98,6 +98,32 @@ class Session: self.last_consolidated = 0 self.updated_at = datetime.now() + def retain_recent_legal_suffix(self, max_messages: int) -> None: + """Keep a legal recent suffix, mirroring get_history boundary rules.""" + if max_messages <= 0: + self.clear() + return + if len(self.messages) <= max_messages: + return + + start_idx = max(0, len(self.messages) - max_messages) + + # If the cutoff lands mid-turn, extend backward to the nearest user turn. + while start_idx > 0 and self.messages[start_idx].get("role") != "user": + start_idx -= 1 + + retained = self.messages[start_idx:] + + # Mirror get_history(): avoid persisting orphan tool results at the front. + start = self._find_legal_start(retained) + if start: + retained = retained[start:] + + dropped = len(self.messages) - len(retained) + self.messages = retained + self.last_consolidated = max(0, self.last_consolidated - dropped) + self.updated_at = datetime.now() + class SessionManager: """ diff --git a/tests/test_commands.py b/tests/test_commands.py index 09b74f2..7d2c178 100644 --- a/tests/test_commands.py +++ b/tests/test_commands.py @@ -477,6 +477,12 @@ def test_agent_hints_about_deprecated_memory_window(mock_agent_runtime, tmp_path assert "no longer used" in result.stdout +def test_heartbeat_retains_recent_messages_by_default(): + config = Config() + + assert config.gateway.heartbeat.keep_recent_messages == 8 + + def test_gateway_uses_workspace_from_config_by_default(monkeypatch, tmp_path: Path) -> None: config_file = tmp_path / "instance" / "config.json" config_file.parent.mkdir(parents=True) diff --git a/tests/test_session_manager_history.py b/tests/test_session_manager_history.py index 4f56344..83036c8 100644 --- a/tests/test_session_manager_history.py +++ b/tests/test_session_manager_history.py @@ -64,6 +64,58 @@ def test_legitimate_tool_pairs_preserved_after_trim(): assert history[0]["role"] == "user" +def test_retain_recent_legal_suffix_keeps_recent_messages(): + session = Session(key="test:trim") + for i in range(10): + session.messages.append({"role": "user", "content": f"msg{i}"}) + + session.retain_recent_legal_suffix(4) + + assert len(session.messages) == 4 + assert session.messages[0]["content"] == "msg6" + assert session.messages[-1]["content"] == "msg9" + + +def test_retain_recent_legal_suffix_adjusts_last_consolidated(): + session = Session(key="test:trim-cons") + for i in range(10): + session.messages.append({"role": "user", "content": f"msg{i}"}) + session.last_consolidated = 7 + + session.retain_recent_legal_suffix(4) + + assert len(session.messages) == 4 + assert session.last_consolidated == 1 + + +def test_retain_recent_legal_suffix_zero_clears_session(): + session = Session(key="test:trim-zero") + for i in range(10): + session.messages.append({"role": "user", "content": f"msg{i}"}) + session.last_consolidated = 5 + + session.retain_recent_legal_suffix(0) + + assert session.messages == [] + assert session.last_consolidated == 0 + + +def test_retain_recent_legal_suffix_keeps_legal_tool_boundary(): + session = Session(key="test:trim-tools") + session.messages.append({"role": "user", "content": "old"}) + session.messages.extend(_tool_turn("old", 0)) + session.messages.append({"role": "user", "content": "keep"}) + session.messages.extend(_tool_turn("keep", 0)) + session.messages.append({"role": "assistant", "content": "done"}) + + session.retain_recent_legal_suffix(4) + + history = session.get_history(max_messages=500) + _assert_no_orphans(history) + assert history[0]["role"] == "user" + assert history[0]["content"] == "keep" + + # --- last_consolidated > 0 --- def test_orphan_trim_with_last_consolidated():