From 0ccb7234d4f01dbbefe1cfe8a21d8906e77ea420 Mon Sep 17 00:00:00 2001 From: chengyongru Date: Thu, 2 Jul 2026 16:21:06 +0800 Subject: [PATCH] refactor(session): extract turn history recovery --- nanobot/agent/loop.py | 234 +++------------------- nanobot/session/manager.py | 5 +- nanobot/session/turn_history.py | 266 ++++++++++++++++++++++++++ tests/test_truncate_text_shadowing.py | 13 +- 4 files changed, 302 insertions(+), 216 deletions(-) create mode 100644 nanobot/session/turn_history.py diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index b0e8114c2..fcf8c029f 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -55,7 +55,7 @@ from nanobot.security.workspace_access import ( bind_workspace_scope, reset_workspace_scope, ) -from nanobot.session import turn_continuation +from nanobot.session import turn_continuation, turn_history from nanobot.session.automation_turns import automation_history_overrides from nanobot.session.goal_state import ( goal_state_runtime_lines, @@ -71,8 +71,6 @@ from nanobot.session.manager import ( ) from nanobot.triggers.local_turns import LocalTriggerTurnCoordinator from nanobot.utils.document import extract_documents, reference_non_image_attachments -from nanobot.utils.helpers import image_placeholder_text -from nanobot.utils.helpers import truncate_text as truncate_text_fn from nanobot.utils.image_generation_intent import image_generation_prompt from nanobot.utils.llm_runtime import LLMRuntime from nanobot.utils.runtime import ( @@ -175,8 +173,8 @@ class AgentLoop: self._refresh_provider_snapshot() return LLMRuntime(self.provider, self.model) - _RUNTIME_CHECKPOINT_KEY = "runtime_checkpoint" - _PENDING_USER_TURN_KEY = "pending_user_turn" + _RUNTIME_CHECKPOINT_KEY = turn_history.RUNTIME_CHECKPOINT_KEY + _PENDING_USER_TURN_KEY = turn_history.PENDING_USER_TURN_KEY # Event-driven state transition table. # Handlers return an event string; the driver looks up the next state here. @@ -1654,38 +1652,13 @@ class AgentLoop: should_truncate_text: bool = False, drop_runtime: bool = False, ) -> list[dict[str, Any]]: - """Strip volatile multimodal payloads before writing session history.""" - filtered: list[dict[str, Any]] = [] - for block in content: - if not isinstance(block, dict): - filtered.append(block) - continue - - if ( - drop_runtime - and block.get("type") == "text" - and isinstance(block.get("text"), str) - and block["text"].startswith(ContextBuilder._RUNTIME_CONTEXT_TAG) - ): - continue - - if block.get("type") == "image_url" and block.get("image_url", {}).get( - "url", "" - ).startswith("data:image/"): - path = (block.get("_meta") or {}).get("path", "") - filtered.append({"type": "text", "text": image_placeholder_text(path)}) - continue - - if block.get("type") == "text" and isinstance(block.get("text"), str): - text = block["text"] - if should_truncate_text and len(text) > self.max_tool_result_chars: - text = truncate_text_fn(text, self.max_tool_result_chars) - filtered.append({**block, "text": text}) - continue - - filtered.append(block) - - return filtered + return turn_history.sanitize_persisted_blocks( + content, + max_tool_result_chars=self.max_tool_result_chars, + runtime_context_tag=ContextBuilder._RUNTIME_CONTEXT_TAG, + should_truncate_text=should_truncate_text, + drop_runtime=drop_runtime, + ) def _save_turn( self, @@ -1695,193 +1668,36 @@ class AgentLoop: *, turn_latency_ms: int | None = None, ) -> None: - """Save new-turn messages into session, truncating large tool results.""" - from datetime import datetime - - declared_tool_call_ids = { - str(tc["id"]) - for m in session.messages - if m.get("role") == "assistant" - for tc in m.get("tool_calls") or [] - if isinstance(tc, dict) and tc.get("id") - } - last_assistant_idx: int | None = None - for m in messages[skip:]: - entry = dict(m) - role, content = entry.get("role"), entry.get("content") - if role == "assistant" and not content and not entry.get("tool_calls"): - continue # skip empty assistant messages — they poison session context - if role == "tool": - tool_call_id = entry.get("tool_call_id") - if not tool_call_id or str(tool_call_id) not in declared_tool_call_ids: - # Undeclared tool results corrupt future provider requests. - logger.warning( - "Dropping orphaned tool result {} from session {} during persistence", - tool_call_id or "(missing id)", - session.key, - ) - continue - if isinstance(content, str) and len(content) > self.max_tool_result_chars: - entry["content"] = truncate_text_fn(content, self.max_tool_result_chars) - elif isinstance(content, list): - filtered = self._sanitize_persisted_blocks(content, should_truncate_text=True) - if not filtered: - # Preserve the tool_call/result pair after block filtering. - filtered = [ - {"type": "text", "text": "[tool result omitted during persistence]"} - ] - entry["content"] = filtered - elif role == "user": - if isinstance(content, str) and ContextBuilder._RUNTIME_CONTEXT_TAG in content: - # Strip the runtime-context block appended at the end. - tag_pos = content.find(ContextBuilder._RUNTIME_CONTEXT_TAG) - before = content[:tag_pos].rstrip("\n ") - if before: - entry["content"] = before - else: - continue - if isinstance(content, list): - filtered = self._sanitize_persisted_blocks(content, drop_runtime=True) - if not filtered: - continue - entry["content"] = filtered - entry.setdefault("timestamp", datetime.now().isoformat()) - session.messages.append(entry) - if role == "assistant": - last_assistant_idx = len(session.messages) - 1 - declared_tool_call_ids.update( - str(tc["id"]) - for tc in entry.get("tool_calls") or [] - if isinstance(tc, dict) and tc.get("id") - ) - if turn_latency_ms is not None and last_assistant_idx is not None: - session.messages[last_assistant_idx]["latency_ms"] = int(turn_latency_ms) - session.updated_at = datetime.now() + turn_history.save_turn( + session, + messages, + skip, + max_tool_result_chars=self.max_tool_result_chars, + runtime_context_tag=ContextBuilder._RUNTIME_CONTEXT_TAG, + turn_latency_ms=turn_latency_ms, + ) def _persist_subagent_followup(self, session: Session, msg: InboundMessage) -> bool: - """Persist subagent follow-ups before prompt assembly so history stays durable. - - Returns True if a new entry was appended; False if the follow-up was - deduped (same ``subagent_task_id`` already in session) or carries no - content worth persisting. - """ - if not msg.content: - return False - task_id = msg.metadata.get("subagent_task_id") if isinstance(msg.metadata, dict) else None - if task_id and any( - m.get("injected_event") == "subagent_result" and m.get("subagent_task_id") == task_id - for m in session.messages - ): - return False - session.add_message( - "assistant", - msg.content, - sender_id=msg.sender_id, - injected_event="subagent_result", - subagent_task_id=task_id, - ) - return True + return turn_history.persist_subagent_followup(session, msg) def _set_runtime_checkpoint(self, session: Session, payload: dict[str, Any]) -> None: - """Persist the latest in-flight turn state into session metadata.""" - session.metadata[self._RUNTIME_CHECKPOINT_KEY] = payload + turn_history.set_runtime_checkpoint(session, payload) self.sessions.save(session) def _mark_pending_user_turn(self, session: Session) -> None: - session.metadata[self._PENDING_USER_TURN_KEY] = True + turn_history.mark_pending_user_turn(session) def _clear_pending_user_turn(self, session: Session) -> None: - session.metadata.pop(self._PENDING_USER_TURN_KEY, None) + turn_history.clear_pending_user_turn(session) def _clear_runtime_checkpoint(self, session: Session) -> None: - if self._RUNTIME_CHECKPOINT_KEY in session.metadata: - session.metadata.pop(self._RUNTIME_CHECKPOINT_KEY, None) - - @staticmethod - def _checkpoint_message_key(message: dict[str, Any]) -> tuple[Any, ...]: - return ( - message.get("role"), - message.get("content"), - message.get("tool_call_id"), - message.get("name"), - message.get("tool_calls"), - message.get("reasoning_content"), - message.get("thinking_blocks"), - ) + turn_history.clear_runtime_checkpoint(session) def _restore_runtime_checkpoint(self, session: Session) -> bool: - """Materialize an unfinished turn into session history before a new request.""" - from datetime import datetime - - checkpoint = session.metadata.get(self._RUNTIME_CHECKPOINT_KEY) - if not isinstance(checkpoint, dict): - return False - - assistant_message = checkpoint.get("assistant_message") - completed_tool_results = checkpoint.get("completed_tool_results") or [] - pending_tool_calls = checkpoint.get("pending_tool_calls") or [] - - restored_messages: list[dict[str, Any]] = [] - if isinstance(assistant_message, dict): - restored = dict(assistant_message) - restored.setdefault("timestamp", datetime.now().isoformat()) - restored_messages.append(restored) - for message in completed_tool_results: - if isinstance(message, dict): - restored = dict(message) - restored.setdefault("timestamp", datetime.now().isoformat()) - restored_messages.append(restored) - for tool_call in pending_tool_calls: - if not isinstance(tool_call, dict): - continue - tool_id = tool_call.get("id") - name = ((tool_call.get("function") or {}).get("name")) or "tool" - restored_messages.append( - { - "role": "tool", - "tool_call_id": tool_id, - "name": name, - "content": "Error: Task interrupted before this tool finished.", - "timestamp": datetime.now().isoformat(), - } - ) - - overlap = 0 - max_overlap = min(len(session.messages), len(restored_messages)) - for size in range(max_overlap, 0, -1): - existing = session.messages[-size:] - restored = restored_messages[:size] - if all( - self._checkpoint_message_key(left) == self._checkpoint_message_key(right) - for left, right in zip(existing, restored) - ): - overlap = size - break - session.messages.extend(restored_messages[overlap:]) - - self._clear_pending_user_turn(session) - self._clear_runtime_checkpoint(session) - return True + return turn_history.restore_runtime_checkpoint(session) def _restore_pending_user_turn(self, session: Session) -> bool: - """Close a turn that only persisted the user message before crashing.""" - from datetime import datetime - - if not session.metadata.get(self._PENDING_USER_TURN_KEY): - return False - - if session.messages and session.messages[-1].get("role") == "user": - session.messages.append( - { - "role": "assistant", - "content": "Error: Task interrupted before a response was generated.", - "timestamp": datetime.now().isoformat(), - } - ) - session.updated_at = datetime.now() - - self._clear_pending_user_turn(session) - return True + return turn_history.restore_pending_user_turn(session) async def process_direct( self, diff --git a/nanobot/session/manager.py b/nanobot/session/manager.py index 533b5c9ce..e73136b98 100644 --- a/nanobot/session/manager.py +++ b/nanobot/session/manager.py @@ -15,6 +15,7 @@ from typing import Any from loguru import logger from nanobot.config.paths import get_legacy_sessions_dir +from nanobot.session.turn_history import PENDING_USER_TURN_KEY, RUNTIME_CHECKPOINT_KEY from nanobot.utils.helpers import ( ensure_dir, estimate_message_tokens, @@ -37,8 +38,8 @@ _SESSION_LIST_PREVIEW_MAX_RECORDS = 200 _SESSION_LIST_PREVIEW_MAX_CHARS = 1_000_000 _FORK_VOLATILE_METADATA_KEYS = { "goal_state", - "pending_user_turn", - "runtime_checkpoint", + PENDING_USER_TURN_KEY, + RUNTIME_CHECKPOINT_KEY, "thread_goal", "title", "title_user_edited", diff --git a/nanobot/session/turn_history.py b/nanobot/session/turn_history.py new file mode 100644 index 000000000..7c01174ce --- /dev/null +++ b/nanobot/session/turn_history.py @@ -0,0 +1,266 @@ +"""Turn history persistence and interrupted-turn recovery.""" + +from __future__ import annotations + +from datetime import datetime +from typing import TYPE_CHECKING, Any + +from loguru import logger + +from nanobot.utils.helpers import image_placeholder_text +from nanobot.utils.helpers import truncate_text as truncate_text_fn + +if TYPE_CHECKING: + from nanobot.session.manager import Session + +RUNTIME_CHECKPOINT_KEY = "runtime_checkpoint" +PENDING_USER_TURN_KEY = "pending_user_turn" + + +def sanitize_persisted_blocks( + content: list[dict[str, Any]], + *, + max_tool_result_chars: int, + runtime_context_tag: str, + should_truncate_text: bool = False, + drop_runtime: bool = False, +) -> list[dict[str, Any]]: + """Strip volatile multimodal payloads before writing session history.""" + filtered: list[dict[str, Any]] = [] + for block in content: + if not isinstance(block, dict): + filtered.append(block) + continue + + if ( + drop_runtime + and block.get("type") == "text" + and isinstance(block.get("text"), str) + and block["text"].startswith(runtime_context_tag) + ): + continue + + if block.get("type") == "image_url" and block.get("image_url", {}).get( + "url", "" + ).startswith("data:image/"): + path = (block.get("_meta") or {}).get("path", "") + filtered.append({"type": "text", "text": image_placeholder_text(path)}) + continue + + if block.get("type") == "text" and isinstance(block.get("text"), str): + text = block["text"] + if should_truncate_text and len(text) > max_tool_result_chars: + text = truncate_text_fn(text, max_tool_result_chars) + filtered.append({**block, "text": text}) + continue + + filtered.append(block) + + return filtered + + +def save_turn( + session: Session, + messages: list[dict], + skip: int, + *, + max_tool_result_chars: int, + runtime_context_tag: str, + turn_latency_ms: int | None = None, +) -> None: + """Save new-turn messages into session, truncating large tool results.""" + declared_tool_call_ids = { + str(tc["id"]) + for m in session.messages + if m.get("role") == "assistant" + for tc in m.get("tool_calls") or [] + if isinstance(tc, dict) and tc.get("id") + } + last_assistant_idx: int | None = None + for m in messages[skip:]: + entry = dict(m) + role, content = entry.get("role"), entry.get("content") + if role == "assistant" and not content and not entry.get("tool_calls"): + continue # skip empty assistant messages - they poison session context + if role == "tool": + tool_call_id = entry.get("tool_call_id") + if not tool_call_id or str(tool_call_id) not in declared_tool_call_ids: + # Undeclared tool results corrupt future provider requests. + logger.warning( + "Dropping orphaned tool result {} from session {} during persistence", + tool_call_id or "(missing id)", + session.key, + ) + continue + if isinstance(content, str) and len(content) > max_tool_result_chars: + entry["content"] = truncate_text_fn(content, max_tool_result_chars) + elif isinstance(content, list): + filtered = sanitize_persisted_blocks( + content, + max_tool_result_chars=max_tool_result_chars, + runtime_context_tag=runtime_context_tag, + should_truncate_text=True, + ) + if not filtered: + # Preserve the tool_call/result pair after block filtering. + filtered = [ + {"type": "text", "text": "[tool result omitted during persistence]"} + ] + entry["content"] = filtered + elif role == "user": + if isinstance(content, str) and runtime_context_tag in content: + # Strip the runtime-context block appended at the end. + tag_pos = content.find(runtime_context_tag) + before = content[:tag_pos].rstrip("\n ") + if before: + entry["content"] = before + else: + continue + if isinstance(content, list): + filtered = sanitize_persisted_blocks( + content, + max_tool_result_chars=max_tool_result_chars, + runtime_context_tag=runtime_context_tag, + drop_runtime=True, + ) + if not filtered: + continue + entry["content"] = filtered + entry.setdefault("timestamp", datetime.now().isoformat()) + session.messages.append(entry) + if role == "assistant": + last_assistant_idx = len(session.messages) - 1 + declared_tool_call_ids.update( + str(tc["id"]) + for tc in entry.get("tool_calls") or [] + if isinstance(tc, dict) and tc.get("id") + ) + if turn_latency_ms is not None and last_assistant_idx is not None: + session.messages[last_assistant_idx]["latency_ms"] = int(turn_latency_ms) + session.updated_at = datetime.now() + + +def persist_subagent_followup(session: Session, msg: Any) -> bool: + """Persist subagent follow-ups before prompt assembly so history stays durable. + + Returns True if a new entry was appended; False if the follow-up was + deduped (same ``subagent_task_id`` already in session) or carries no + content worth persisting. + """ + if not msg.content: + return False + task_id = msg.metadata.get("subagent_task_id") if isinstance(msg.metadata, dict) else None + if task_id and any( + m.get("injected_event") == "subagent_result" and m.get("subagent_task_id") == task_id + for m in session.messages + ): + return False + session.add_message( + "assistant", + msg.content, + sender_id=msg.sender_id, + injected_event="subagent_result", + subagent_task_id=task_id, + ) + return True + + +def set_runtime_checkpoint(session: Session, payload: dict[str, Any]) -> None: + """Persist the latest in-flight turn state into session metadata.""" + session.metadata[RUNTIME_CHECKPOINT_KEY] = payload + + +def mark_pending_user_turn(session: Session) -> None: + session.metadata[PENDING_USER_TURN_KEY] = True + + +def clear_pending_user_turn(session: Session) -> None: + session.metadata.pop(PENDING_USER_TURN_KEY, None) + + +def clear_runtime_checkpoint(session: Session) -> None: + session.metadata.pop(RUNTIME_CHECKPOINT_KEY, None) + + +def checkpoint_message_key(message: dict[str, Any]) -> tuple[Any, ...]: + return ( + message.get("role"), + message.get("content"), + message.get("tool_call_id"), + message.get("name"), + message.get("tool_calls"), + message.get("reasoning_content"), + message.get("thinking_blocks"), + ) + + +def restore_runtime_checkpoint(session: Session) -> bool: + """Materialize an unfinished turn into session history before a new request.""" + checkpoint = session.metadata.get(RUNTIME_CHECKPOINT_KEY) + if not isinstance(checkpoint, dict): + return False + + assistant_message = checkpoint.get("assistant_message") + completed_tool_results = checkpoint.get("completed_tool_results") or [] + pending_tool_calls = checkpoint.get("pending_tool_calls") or [] + + restored_messages: list[dict[str, Any]] = [] + if isinstance(assistant_message, dict): + restored = dict(assistant_message) + restored.setdefault("timestamp", datetime.now().isoformat()) + restored_messages.append(restored) + for message in completed_tool_results: + if isinstance(message, dict): + restored = dict(message) + restored.setdefault("timestamp", datetime.now().isoformat()) + restored_messages.append(restored) + for tool_call in pending_tool_calls: + if not isinstance(tool_call, dict): + continue + tool_id = tool_call.get("id") + name = ((tool_call.get("function") or {}).get("name")) or "tool" + restored_messages.append( + { + "role": "tool", + "tool_call_id": tool_id, + "name": name, + "content": "Error: Task interrupted before this tool finished.", + "timestamp": datetime.now().isoformat(), + } + ) + + overlap = 0 + max_overlap = min(len(session.messages), len(restored_messages)) + for size in range(max_overlap, 0, -1): + existing = session.messages[-size:] + restored = restored_messages[:size] + if all( + checkpoint_message_key(left) == checkpoint_message_key(right) + for left, right in zip(existing, restored) + ): + overlap = size + break + session.messages.extend(restored_messages[overlap:]) + + clear_pending_user_turn(session) + clear_runtime_checkpoint(session) + return True + + +def restore_pending_user_turn(session: Session) -> bool: + """Close a turn that only persisted the user message before crashing.""" + if not session.metadata.get(PENDING_USER_TURN_KEY): + return False + + if session.messages and session.messages[-1].get("role") == "user": + session.messages.append( + { + "role": "assistant", + "content": "Error: Task interrupted before a response was generated.", + "timestamp": datetime.now().isoformat(), + } + ) + session.updated_at = datetime.now() + + clear_pending_user_turn(session) + return True diff --git a/tests/test_truncate_text_shadowing.py b/tests/test_truncate_text_shadowing.py index 11132b511..635facac6 100644 --- a/tests/test_truncate_text_shadowing.py +++ b/tests/test_truncate_text_shadowing.py @@ -1,5 +1,4 @@ import inspect -from types import SimpleNamespace def test_sanitize_persisted_blocks_truncate_text_shadowing_regression() -> None: @@ -14,16 +13,20 @@ def test_sanitize_persisted_blocks_truncate_text_shadowing_regression() -> None: This test asserts the fixed API exists and truncation works without raising. """ - from nanobot.agent.loop import AgentLoop + from nanobot.session.turn_history import sanitize_persisted_blocks - sig = inspect.signature(AgentLoop._sanitize_persisted_blocks) + sig = inspect.signature(sanitize_persisted_blocks) assert "should_truncate_text" in sig.parameters assert "truncate_text" not in sig.parameters - dummy = SimpleNamespace(max_tool_result_chars=5) content = [{"type": "text", "text": "0123456789"}] - out = AgentLoop._sanitize_persisted_blocks(dummy, content, should_truncate_text=True) + out = sanitize_persisted_blocks( + content, + max_tool_result_chars=5, + runtime_context_tag="[runtime]", + should_truncate_text=True, + ) assert isinstance(out, list) assert out and out[0]["type"] == "text" assert isinstance(out[0]["text"], str)