diff --git a/code_puppy/plugins/session_tree/__init__.py b/code_puppy/plugins/session_tree/__init__.py new file mode 100644 index 000000000..52230a858 --- /dev/null +++ b/code_puppy/plugins/session_tree/__init__.py @@ -0,0 +1 @@ +"""Session tree navigation plugin.""" diff --git a/code_puppy/plugins/session_tree/register_callbacks.py b/code_puppy/plugins/session_tree/register_callbacks.py new file mode 100644 index 000000000..cbb382100 --- /dev/null +++ b/code_puppy/plugins/session_tree/register_callbacks.py @@ -0,0 +1,254 @@ +"""Session navigation commands for Code Puppy.""" + +from __future__ import annotations + +import asyncio +import concurrent.futures +from datetime import datetime +from pathlib import Path +from typing import Any + +from code_puppy.callbacks import register_callback + +from .tree_menu import TreeSelectionMenu +from .tree_model import ( + FilterMode, + HistoryNode, + message_text, + resolve_node, + selectable_nodes, + user_nodes, +) +from .tree_store import SessionTree, TreeStore + + +def _emit(kind: str, message: str) -> None: + from code_puppy import messaging + + getattr(messaging, f"emit_{kind}")(message) + + +def _current_history() -> tuple[Any, list[Any]]: + from code_puppy.agents.agent_manager import get_current_agent + + agent = get_current_agent() + return agent, list(agent.get_message_history()) + + +def _tree_store() -> TreeStore: + from code_puppy.config import AUTOSAVE_DIR, get_current_autosave_id + + path = Path(AUTOSAVE_DIR) / "session_trees" / f"{get_current_autosave_id()}.pkl" + return TreeStore(path) + + +def _load_tree(history: list[Any]) -> tuple[TreeStore, SessionTree]: + store = _tree_store() + tree = store.load() + tree.sync_history(history) + store.save(tree) + return store, tree + + +def _help_entries() -> list[tuple[str, str]]: + return [ + ("tree", "Navigate the persistent conversation tree and switch branches"), + ("fork", "Re-edit a previous user message to create a branch"), + ("clone", "Duplicate the active branch into a new autosave session"), + ] + + +def _tree_args(command: str) -> tuple[bool, str]: + parts = command.split(maxsplit=2) + if len(parts) >= 2 and parts[1].lower() in {"summary", "summarize"}: + return True, parts[2].strip() if len(parts) == 3 else "" + return False, parts[1].strip() if len(parts) >= 2 else "" + + +def _markdown_result(prompt: str) -> Any: + from code_puppy.plugins.customizable_commands.register_callbacks import ( + MarkdownCommandResult, + ) + + return MarkdownCommandResult(prompt) + + +def _run_tree_menu(tree: SessionTree, *, user_only: bool = False) -> Any: + mode = FilterMode.USER_ONLY if user_only else FilterMode.DEFAULT + with concurrent.futures.ThreadPoolExecutor() as executor: + future = executor.submit( + lambda: asyncio.run(TreeSelectionMenu(tree=tree, mode=mode).run_async()) + ) + return future.result(timeout=300) + + +def _resolve_or_menu( + command: str, tree: SessionTree, nodes: list[HistoryNode] +) -> tuple[HistoryNode | None, bool]: + summarize, selection = _tree_args(command) + if selection: + node = resolve_node(selection, nodes) + if node is None: + _emit("warning", f"/tree: no unique history point matches '{selection}'") + return node, summarize + result = _run_tree_menu(tree) + if result.cancelled: + return None, False + return result.node, result.summarize + + +def _summary_messages(history: list[Any], abandoned: list[Any]) -> list[Any]: + """Summarize only the abandoned tail and return generated context messages.""" + from code_puppy.agents._compaction import _run_summarization_core + + if not abandoned: + return [] + system = history[:1] + compacted, _ = _run_summarization_core( + [*system, *abandoned], + protected_tokens=0, + with_protection=False, + model_name=None, + ) + return list(compacted[len(system) :]) + + +def _navigate( + agent: Any, + store: TreeStore, + tree: SessionTree, + node: HistoryNode, + *, + summarize: bool, +) -> bool | Any: + old_leaf_id = tree.active_leaf_id + target_id = node.parent_id if node.is_userish else node.node_id + prompt = ( + message_text(tree.nodes[node.node_id].message).strip() + if node.is_userish + else "" + ) + target_history = tree.history_through(target_id) + + if summarize: + try: + target_history.extend( + _summary_messages( + tree.history_through(old_leaf_id), + tree.abandoned_history(node.node_id), + ) + ) + except Exception as exc: # noqa: BLE001 - command boundary must fail soft + _emit("error", f"/tree: failed to summarize abandoned branch - {exc}") + return True + + try: + agent.set_message_history(target_history) + tree.sync_history(target_history) + store.save(tree) + except Exception as exc: # noqa: BLE001 + _emit("error", f"/tree: failed to navigate - {exc}") + return True + + if node.is_userish: + _emit( + "success", "Moved before selected user message; edit it to create a branch." + ) + return _markdown_result(prompt) if prompt else True + _emit("success", f"Navigated to {node.role} message [{node.node_id}]") + return True + + +def _handle_tree(command: str) -> bool | Any: + try: + agent, history = _current_history() + store, tree = _load_tree(history) + except Exception as exc: # noqa: BLE001 + _emit("error", f"/tree: could not read session tree - {exc}") + return True + + nodes = selectable_nodes(tree) + if not nodes: + _emit("warning", "/tree: conversation history is empty") + return True + try: + node, summarize = _resolve_or_menu(command, tree, nodes) + # Label changes made in the selector survive selection and cancellation. + store.save(tree) + except Exception as exc: # noqa: BLE001 + _emit("error", f"/tree: selector failed - {exc}") + return True + return ( + True + if node is None + else _navigate(agent, store, tree, node, summarize=summarize) + ) + + +def _handle_fork(command: str) -> bool | Any: + try: + agent, history = _current_history() + store, tree = _load_tree(history) + except Exception as exc: # noqa: BLE001 + _emit("error", f"/fork: could not read session tree - {exc}") + return True + nodes = user_nodes(tree) + if not nodes: + _emit("warning", "/fork: no previous user messages are available") + return True + selection = command.split(maxsplit=1)[1].strip() if " " in command else "" + if selection: + node = resolve_node(selection, nodes) + else: + try: + result = _run_tree_menu(tree, user_only=True) + node = result.node if not result.cancelled else None + except Exception as exc: # noqa: BLE001 + _emit("error", f"/fork: selector failed - {exc}") + return True + if node is None or not node.is_userish: + if selection: + _emit("warning", f"/fork: no unique user message matches '{selection}'") + return True + return _navigate(agent, store, tree, node, summarize=False) + + +def _handle_clone() -> bool: + try: + agent, history = _current_history() + if not history: + _emit("warning", "/clone: no active branch to clone") + return True + from code_puppy.config import AUTOSAVE_DIR + from code_puppy.session_storage import save_session + + # A clone is a new saved artifact, not a mutation of the current session identity. + clone_id = datetime.now().strftime("%Y%m%d_%H%M%S_%f") + metadata = save_session( + history=history, + session_name=f"auto_session_{clone_id}", + base_dir=Path(AUTOSAVE_DIR), + timestamp=datetime.now().isoformat(), + token_estimator=agent.estimate_tokens_for_message, + auto_saved=True, + ) + _emit("success", f"Cloned active branch to {metadata.session_name}.") + except Exception as exc: # noqa: BLE001 + _emit("error", f"/clone: failed to clone active branch - {exc}") + return True + + +def _handle_custom_command(command: str, name: str) -> bool | str | None: + if name == "tree": + return _handle_tree(command) + if name == "fork": + return _handle_fork(command) + if name == "clone": + return _handle_clone() + return None + + +register_callback("custom_command_help", _help_entries) +register_callback("custom_command", _handle_custom_command) + +__all__ = ["_handle_custom_command", "_handle_fork", "_handle_tree", "_help_entries"] diff --git a/code_puppy/plugins/session_tree/tree_menu.py b/code_puppy/plugins/session_tree/tree_menu.py new file mode 100644 index 000000000..0855cfd68 --- /dev/null +++ b/code_puppy/plugins/session_tree/tree_menu.py @@ -0,0 +1,295 @@ +"""Prompt-toolkit TUI for session tree selection.""" + +from __future__ import annotations + +import sys +from dataclasses import dataclass, field + +from prompt_toolkit import Application +from prompt_toolkit.key_binding import KeyBindings +from prompt_toolkit.layout import Layout, Window +from prompt_toolkit.layout.controls import FormattedTextControl + +from code_puppy.command_line.pagination import ( + ensure_visible_page, + get_page_bounds, + get_page_for_index, + get_total_pages, +) + +from .tree_model import ( + FILTER_MODES, + FilterMode, + HistoryNode, + build_nodes, + visible_nodes, +) +from .tree_store import SessionTree + +TREE_PAGE_SIZE = 14 + + +@dataclass(slots=True) +class TreeMenuResult: + node: HistoryNode | None = None + cancelled: bool = False + summarize: bool = False + custom_summary_focus: str = "" + + +@dataclass(slots=True) +class TreeSelectionMenu: + tree: SessionTree + mode: FilterMode = FilterMode.DEFAULT + selected_index: int = 0 + page: int = 0 + page_size: int = TREE_PAGE_SIZE + result: TreeMenuResult = field(default_factory=TreeMenuResult) + show_label_timestamps: bool = False + search_query: str = "" + + def __post_init__(self) -> None: + self._clamp_selection() + + @property + def nodes(self) -> list[HistoryNode]: + return build_nodes(self.tree) + + @property + def visible(self) -> list[HistoryNode]: + return visible_nodes(self.nodes, self.mode, self.search_query) + + @property + def total_pages(self) -> int: + return get_total_pages(len(self.visible), self.page_size) + + @property + def page_start(self) -> int: + start, _ = get_page_bounds(self.page, len(self.visible), self.page_size) + return start + + @property + def page_end(self) -> int: + _, end = get_page_bounds(self.page, len(self.visible), self.page_size) + return end + + @property + def page_nodes(self) -> list[HistoryNode]: + return self.visible[self.page_start : self.page_end] + + def selected_node(self) -> HistoryNode | None: + if 0 <= self.selected_index < len(self.visible): + return self.visible[self.selected_index] + return None + + def _clamp_selection(self) -> None: + if not self.visible: + self.selected_index = 0 + self.page = 0 + return + self.selected_index = max(0, min(self.selected_index, len(self.visible) - 1)) + self.page = ensure_visible_page( + self.selected_index, self.page, len(self.visible), self.page_size + ) + + def _move(self, delta: int) -> None: + self.selected_index += delta + self._clamp_selection() + + def _page(self, delta: int) -> None: + if not self.visible: + return + self.page = max(0, min(self.page + delta, self.total_pages - 1)) + self.selected_index = self.page_start + + def _cycle_filter(self) -> None: + modes = list(FILTER_MODES) + next_index = (modes.index(self.mode) + 1) % len(modes) + selected = self.selected_node() + self.mode = modes[next_index] + if selected and selected in self.visible: + self.selected_index = self.visible.index(selected) + else: + self.selected_index = 0 + self._clamp_selection() + + def _toggle_label_timestamp(self) -> None: + self.show_label_timestamps = not self.show_label_timestamps + + def _toggle_label(self) -> None: + selected = self.selected_node() + if selected is None: + return + self.tree.set_label(selected.node_id, None if selected.label else "label") + + def _branch_prefix(self, node: HistoryNode) -> str: + ancestors = "".join( + "│ " if has_next else " " for has_next in node.ancestor_has_next + ) + if node.depth == 0: + return "" + return f"{ancestors}{'└─ ' if node.is_last else '├─ '}" + + def _render(self): + lines = [("bold cyan", " Session Tree")] + lines.append(("fg:ansibrightblack", f"\n Filter: {self.mode.value}")) + lines.append( + ( + "fg:ansibrightblack", + f" Search: {self.search_query or '(type to search)'}", + ) + ) + if self.total_pages > 1: + lines.append( + ("fg:ansibrightblack", f" Page {self.page + 1}/{self.total_pages}") + ) + lines.append(("", "\n")) + + if not self.visible: + lines.append(("fg:ansiyellow", "\n No entries match this filter.\n\n")) + for offset, node in enumerate(self.page_nodes): + absolute_index = self.page_start + offset + is_selected = absolute_index == self.selected_index + is_active = node.is_active + cursor = " › " if is_selected else " " + style = "fg:ansiwhite bold" if is_selected else "fg:ansibrightblack" + timestamp = "" + if self.show_label_timestamps and node.label_timestamp: + timestamp = f" @{node.label_timestamp[11:16]}" + label = f" [{node.label}{timestamp}]" if node.label else "" + active = " ← active" if is_active else "" + path_marker = "• " if node.is_on_active_path else " " + lines.append( + ( + style, + f'{cursor}{self._branch_prefix(node)}{path_marker}{node.role}: "{node.preview}"', + ) + ) + lines.append(("fg:ansiblue", f" #{node.node_id}{label}{active}\n")) + + lines.append(("", "\n")) + lines.append(("fg:ansibrightblack", " ↑/↓ ")) + lines.append(("", "Navigate\n")) + lines.append(("fg:ansibrightblack", " ←/→ PgUp/PgDn ")) + lines.append(("", "Page\n")) + lines.append(("fg:ansibrightblack", " Ctrl+O ")) + lines.append(("", "Cycle filter\n")) + lines.append(("fg:ansibrightblack", " L ")) + lines.append(("", "Set/clear label\n")) + lines.append(("fg:ansibrightblack", " T ")) + lines.append(("", "Toggle label timestamps\n")) + lines.append(("fg:ansigreen", " Enter ")) + lines.append(("", "Select/restore\n")) + lines.append(("fg:ansigreen", " S ")) + lines.append(("", "Summarize through selection as new conversation\n")) + lines.append(("fg:ansiyellow", " Esc/Ctrl+C ")) + lines.append(("", "Cancel\n")) + return lines + + async def run_async(self) -> TreeMenuResult: + control = FormattedTextControl(lambda: self._render()) + kb = KeyBindings() + + def refresh(event) -> None: + control.text = self._render() + event.app.invalidate() + + @kb.add("up") + @kb.add("c-p") + def _(event): + self._move(-1) + refresh(event) + + @kb.add("down") + @kb.add("c-n") + def _(event): + self._move(1) + refresh(event) + + @kb.add("pageup") + @kb.add("left") + def _(event): + self._page(-1) + refresh(event) + + @kb.add("pagedown") + @kb.add("right") + def _(event): + self._page(1) + refresh(event) + + @kb.add("c-o") + def _(event): + self._cycle_filter() + refresh(event) + + @kb.add("L") + @kb.add("l") + def _(event): + self._toggle_label() + refresh(event) + + @kb.add("T") + @kb.add("t") + def _(event): + self._toggle_label_timestamp() + refresh(event) + + @kb.add("enter") + def _(event): + self.result.node = self.selected_node() + event.app.exit() + + @kb.add("S") + @kb.add("s") + def _(event): + self.result.node = self.selected_node() + self.result.summarize = True + event.app.exit() + + @kb.add("backspace") + def _(event): + if self.search_query: + self.search_query = self.search_query[:-1] + self.selected_index = 0 + self._clamp_selection() + refresh(event) + + @kb.add("") + def _(event): + text = event.data + if text and text.isprintable(): + self.search_query += text + self.selected_index = 0 + self._clamp_selection() + refresh(event) + + @kb.add("escape") + @kb.add("c-c") + def _(event): + if self.search_query: + self.search_query = "" + self.selected_index = 0 + self._clamp_selection() + refresh(event) + return + self.result.cancelled = True + event.app.exit() + + sys.stdout.write("\033[?1049h\033[2J\033[H") + sys.stdout.flush() + try: + app = Application( + layout=Layout(Window(content=control, wrap_lines=True)), + key_bindings=kb, + full_screen=False, + ) + await app.run_async() + finally: + sys.stdout.write("\033[?1049l") + sys.stdout.flush() + return self.result + + +def selected_index_to_page(index: int, page_size: int = TREE_PAGE_SIZE) -> int: + return get_page_for_index(index, page_size) diff --git a/code_puppy/plugins/session_tree/tree_model.py b/code_puppy/plugins/session_tree/tree_model.py new file mode 100644 index 000000000..6bd86090f --- /dev/null +++ b/code_puppy/plugins/session_tree/tree_model.py @@ -0,0 +1,178 @@ +"""View-model helpers for a persistent conversation tree.""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import StrEnum +from typing import Any, Iterable, Sequence + +from .tree_store import SessionTree, StoredNode + + +class FilterMode(StrEnum): + DEFAULT = "default" + NO_TOOLS = "no-tools" + USER_ONLY = "user-only" + LABELED_ONLY = "labeled-only" + ALL = "all" + + +FILTER_MODES = tuple(FilterMode) +TOOL_ROLES = {"tool", "toolresult", "tool_result", "bashexecution"} +USERISH_ROLES = {"user", "custom"} + + +@dataclass(frozen=True, slots=True) +class HistoryNode: + node_id: str + parent_id: str | None + role: str + preview: str + depth: int + is_last: bool + ancestor_has_next: tuple[bool, ...] + is_system: bool = False + label: str = "" + label_timestamp: str | None = None + is_active: bool = False + is_on_active_path: bool = False + + @property + def is_userish(self) -> bool: + return self.role.lower() in USERISH_ROLES + + @property + def is_toolish(self) -> bool: + return self.role.lower() in TOOL_ROLES + + +@dataclass(frozen=True, slots=True) +class _FlatNode: + stored: StoredNode + depth: int + is_last: bool + ancestor_has_next: tuple[bool, ...] + + +def message_role(message: Any) -> str: + role = getattr(message, "role", None) + if isinstance(role, str): + return role + name = type(message).__name__.lower() + if "response" in name: + return "assistant" + if "request" in name: + return "user" + return name or "message" + + +def message_text(message: Any) -> str: + content = getattr(message, "content", None) + if isinstance(content, str): + return content + parts = getattr(message, "parts", None) + if parts is None: + parts = content if isinstance(content, list) else [] + chunks: list[str] = [] + for part in parts or []: + text = getattr(part, "content", None) + if text is None: + text = getattr(part, "text", None) + if isinstance(part, dict): + text = part.get("content") or part.get("text") + if isinstance(text, str): + chunks.append(text) + return " ".join(chunks) if chunks else str(message) + + +def _preview(text: str, max_len: int = 96) -> str: + clean = " ".join(text.split()) or "(no text)" + return clean if len(clean) <= max_len else f"{clean[: max_len - 1]}…" + + +def _flatten(tree: SessionTree) -> list[_FlatNode]: + result: list[_FlatNode] = [] + roots = tree.children(None) + stack: list[tuple[StoredNode, int, bool, tuple[bool, ...]]] = [] + for index in range(len(roots) - 1, -1, -1): + stack.append((roots[index], 0, index == len(roots) - 1, ())) + while stack: + node, depth, is_last, ancestors = stack.pop() + result.append(_FlatNode(node, depth, is_last, ancestors)) + children = tree.children(node.node_id) + for index in range(len(children) - 1, -1, -1): + child_is_last = index == len(children) - 1 + stack.append( + (children[index], depth + 1, child_is_last, (*ancestors, not is_last)) + ) + return result + + +def build_nodes(tree: SessionTree) -> list[HistoryNode]: + active_path = set(tree.path_ids(tree.active_leaf_id)) + nodes: list[HistoryNode] = [] + for item in _flatten(tree): + role = message_role(item.stored.message) + nodes.append( + HistoryNode( + node_id=item.stored.node_id, + parent_id=item.stored.parent_id, + role=role, + preview=_preview(message_text(item.stored.message)), + depth=item.depth, + is_last=item.is_last, + ancestor_has_next=item.ancestor_has_next, + is_system=item.stored.parent_id is None and role == "system", + label=item.stored.label, + label_timestamp=item.stored.label_timestamp, + is_active=item.stored.node_id == tree.active_leaf_id, + is_on_active_path=item.stored.node_id in active_path, + ) + ) + return nodes + + +def selectable_nodes(tree: SessionTree) -> list[HistoryNode]: + return [node for node in build_nodes(tree) if not node.is_system] + + +def user_nodes(tree: SessionTree) -> list[HistoryNode]: + return [node for node in selectable_nodes(tree) if node.is_userish] + + +def visible_nodes( + nodes: Sequence[HistoryNode], mode: FilterMode, query: str = "" +) -> list[HistoryNode]: + if mode == FilterMode.ALL: + visible = list(nodes) + elif mode == FilterMode.NO_TOOLS: + visible = [node for node in nodes if not node.is_toolish] + elif mode == FilterMode.USER_ONLY: + visible = [node for node in nodes if node.is_userish] + elif mode == FilterMode.LABELED_ONLY: + visible = [node for node in nodes if node.label] + else: + visible = [node for node in nodes if not node.is_system and not node.is_toolish] + tokens = query.lower().split() + if not tokens: + return visible + return [ + node + for node in visible + if all( + token in f"{node.role} {node.preview} {node.label}".lower() + for token in tokens + ) + ] + + +def resolve_node(selection: str, nodes: Iterable[HistoryNode]) -> HistoryNode | None: + value = selection.strip() + if not value: + return None + node_list = list(nodes) + if value.isdigit() and 1 <= int(value) <= len(node_list): + return node_list[int(value) - 1] + lowered = value.lower() + matches = [node for node in node_list if node.node_id.startswith(lowered)] + return matches[0] if len(matches) == 1 else None diff --git a/code_puppy/plugins/session_tree/tree_store.py b/code_puppy/plugins/session_tree/tree_store.py new file mode 100644 index 000000000..c9b9044b5 --- /dev/null +++ b/code_puppy/plugins/session_tree/tree_store.py @@ -0,0 +1,135 @@ +"""Persistent append-only conversation graph for the session tree plugin.""" + +from __future__ import annotations + +import hashlib +import pickle +import uuid +from dataclasses import dataclass, field +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Sequence + + +def _message_fingerprint(message: Any) -> str: + """Return a stable-enough identity for matching a message on one path.""" + try: + payload = pickle.dumps(message) + except Exception: # pragma: no cover - exotic third-party messages + payload = repr(message).encode("utf-8", errors="replace") + return hashlib.sha256(payload).hexdigest() + + +@dataclass(slots=True) +class StoredNode: + node_id: str + parent_id: str | None + message: Any + fingerprint: str + sequence: int + label: str = "" + label_timestamp: str | None = None + + +@dataclass(slots=True) +class SessionTree: + """All known messages and branches for one Code Puppy autosave session.""" + + nodes: dict[str, StoredNode] = field(default_factory=dict) + active_leaf_id: str | None = None + next_sequence: int = 0 + + def sync_history(self, history: Sequence[Any]) -> str | None: + """Merge a live linear history into the graph and mark its final node active.""" + parent_id: str | None = None + for message in history: + fingerprint = _message_fingerprint(message) + matching = next( + ( + node + for node in self.nodes.values() + if node.parent_id == parent_id and node.fingerprint == fingerprint + ), + None, + ) + if matching is None: + matching = StoredNode( + node_id=uuid.uuid4().hex[:12], + parent_id=parent_id, + message=message, + fingerprint=fingerprint, + sequence=self.next_sequence, + ) + self.nodes[matching.node_id] = matching + self.next_sequence += 1 + else: + # Keep the newest object shape after dependency upgrades/deserialization. + matching.message = message + parent_id = matching.node_id + self.active_leaf_id = parent_id + return parent_id + + def children(self, parent_id: str | None) -> list[StoredNode]: + return sorted( + (node for node in self.nodes.values() if node.parent_id == parent_id), + key=lambda node: node.sequence, + ) + + def path_ids(self, node_id: str | None) -> list[str]: + path: list[str] = [] + seen: set[str] = set() + while node_id is not None and node_id not in seen: + seen.add(node_id) + node = self.nodes.get(node_id) + if node is None: + break + path.append(node_id) + node_id = node.parent_id + path.reverse() + return path + + def history_through(self, node_id: str | None) -> list[Any]: + return [self.nodes[item].message for item in self.path_ids(node_id)] + + def common_ancestor(self, left_id: str | None, right_id: str | None) -> str | None: + common: str | None = None + for left, right in zip(self.path_ids(left_id), self.path_ids(right_id)): + if left != right: + break + common = left + return common + + def abandoned_history(self, target_id: str) -> list[Any]: + ancestor = self.common_ancestor(self.active_leaf_id, target_id) + active_path = self.path_ids(self.active_leaf_id) + start = active_path.index(ancestor) + 1 if ancestor in active_path else 0 + return [self.nodes[item].message for item in active_path[start:]] + + def set_label(self, node_id: str, label: str | None) -> None: + node = self.nodes[node_id] + node.label = (label or "").strip() + node.label_timestamp = ( + datetime.now(timezone.utc).isoformat() if node.label else None + ) + + +class TreeStore: + """Atomic pickle persistence for a session graph.""" + + def __init__(self, path: Path) -> None: + self.path = path + + def load(self) -> SessionTree: + if not self.path.exists(): + return SessionTree() + try: + value = pickle.loads(self.path.read_bytes()) # noqa: S301 + except Exception: + return SessionTree() + return value if isinstance(value, SessionTree) else SessionTree() + + def save(self, tree: SessionTree) -> None: + self.path.parent.mkdir(parents=True, exist_ok=True) + temporary = self.path.with_suffix(".tmp") + temporary.write_bytes(pickle.dumps(tree)) + temporary.replace(self.path) diff --git a/tests/plugins/test_session_tree.py b/tests/plugins/test_session_tree.py new file mode 100644 index 000000000..d3a9fec87 --- /dev/null +++ b/tests/plugins/test_session_tree.py @@ -0,0 +1,257 @@ +"""Tests for persistent /tree, /fork, and /clone navigation.""" + +from __future__ import annotations + +import importlib +import sys +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +from code_puppy.plugins.session_tree.tree_model import ( + FilterMode, + build_nodes, + resolve_node, + selectable_nodes, + visible_nodes, +) +from code_puppy.plugins.session_tree.tree_store import SessionTree, TreeStore + + +class Message(SimpleNamespace): + pass + + +def _plugin_module(): + sys.modules.setdefault("dbos", MagicMock()) + return importlib.import_module("code_puppy.plugins.session_tree.register_callbacks") + + +def _history() -> list[Message]: + return [ + Message(role="system", content="system prompt"), + Message(role="user", content="first task"), + Message(role="assistant", content="first answer"), + Message(role="toolResult", content="tool output"), + Message(role="user", content="second task"), + Message(role="assistant", content="second answer"), + ] + + +def _tree(history=None) -> SessionTree: + tree = SessionTree() + tree.sync_history(history or _history()) + return tree + + +def _command_context(tmp_path: Path, history: list[Message]): + agent = MagicMock() + agent.get_message_history.return_value = history + manager = SimpleNamespace(get_current_agent=lambda: agent) + config = SimpleNamespace( + AUTOSAVE_DIR=str(tmp_path), get_current_autosave_id=lambda: "current" + ) + return agent, patch.dict( + sys.modules, + { + "code_puppy.agents.agent_manager": manager, + "code_puppy.config": config, + }, + ) + + +def test_sync_history_is_stable_and_builds_real_branches(): + history = _history() + tree = _tree(history) + original_ids = tree.path_ids(tree.active_leaf_id) + + tree.sync_history(history) + assert tree.path_ids(tree.active_leaf_id) == original_ids + + tree.sync_history([*history[:4], Message(role="user", content="alternate")]) + branch_parent = original_ids[3] + assert len(tree.children(branch_parent)) == 2 + assert len(tree.nodes) == len(history) + 1 + + +def test_duplicate_messages_on_one_path_remain_distinct(): + repeated = Message(role="user", content="same") + tree = _tree([repeated, repeated]) + path = tree.path_ids(tree.active_leaf_id) + + assert len(path) == 2 + assert path[0] != path[1] + + +def test_store_persists_branches_labels_and_active_leaf(tmp_path): + store = TreeStore(tmp_path / "tree.pkl") + tree = _tree() + selected = selectable_nodes(tree)[0] + tree.set_label(selected.node_id, "keeper") + store.save(tree) + + restored = store.load() + assert restored.nodes[selected.node_id].label == "keeper" + assert restored.active_leaf_id == tree.active_leaf_id + + +def test_model_marks_active_path_and_renders_branch_depth(): + history = _history() + tree = _tree(history) + first_path = tree.path_ids(tree.active_leaf_id) + tree.sync_history([*history[:3], Message(role="user", content="alternate")]) + nodes = build_nodes(tree) + + assert sum(node.is_active for node in nodes) == 1 + assert all( + node.is_on_active_path + for node in nodes + if node.node_id in tree.path_ids(tree.active_leaf_id) + ) + assert any(node.depth > 0 and not node.is_on_active_path for node in nodes) + assert first_path[-1] in tree.nodes + + +def test_filters_search_and_unique_resolution(): + tree = _tree() + nodes = build_nodes(tree) + assistant = next(node for node in nodes if node.role == "assistant") + tree.set_label(assistant.node_id, "keeper") + nodes = build_nodes(tree) + + assert {node.role for node in visible_nodes(nodes, FilterMode.USER_ONLY)} == { + "user" + } + assert all( + not node.is_toolish for node in visible_nodes(nodes, FilterMode.NO_TOOLS) + ) + assert [node.node_id for node in visible_nodes(nodes, FilterMode.LABELED_ONLY)] == [ + assistant.node_id + ] + assert ( + visible_nodes(nodes, FilterMode.ALL, "first answer")[0].node_id + == assistant.node_id + ) + selectable = selectable_nodes(tree) + assert resolve_node("1", selectable) == selectable[0] + assert resolve_node(assistant.node_id[:6], selectable).node_id == assistant.node_id + + +def test_abandoned_history_starts_after_common_ancestor(): + history = _history() + tree = _tree(history) + old_leaf = tree.active_leaf_id + target = tree.path_ids(old_leaf)[2] + + abandoned = tree.abandoned_history(target) + assert abandoned == history[3:] + + +def test_tree_user_selection_restores_parent_and_preserves_old_branch(tmp_path): + plugin = _plugin_module() + history = _history() + agent, context = _command_context(tmp_path, history) + with context, patch.object(plugin, "_emit"): + result = plugin._handle_custom_command("/tree 4", "tree") + + assert str(result) == "second task" + agent.set_message_history.assert_called_once_with(history[:4]) + stored = TreeStore(tmp_path / "session_trees/current.pkl").load() + assert len(stored.nodes) == len(history) + assert stored.history_through(stored.active_leaf_id) == history[:4] + + +def test_new_turn_after_navigation_creates_sibling_branch(tmp_path): + plugin = _plugin_module() + history = _history() + agent, context = _command_context(tmp_path, history) + with context, patch.object(plugin, "_emit"): + plugin._handle_custom_command("/tree 4", "tree") + + alternate = [*history[:4], Message(role="user", content="replacement")] + agent, context = _command_context(tmp_path, alternate) + with context, patch.object(plugin, "_emit"): + plugin._handle_custom_command("/tree nope", "tree") + + stored = TreeStore(tmp_path / "session_trees/current.pkl").load() + parent = stored.path_ids(stored.active_leaf_id)[3] + assert len(stored.children(parent)) == 2 + + +def test_tree_non_user_selection_restores_through_selected_node(tmp_path): + plugin = _plugin_module() + history = _history() + agent, context = _command_context(tmp_path, history) + with context, patch.object(plugin, "_emit"): + result = plugin._handle_custom_command("/tree 2", "tree") + + assert result is True + agent.set_message_history.assert_called_once_with(history[:3]) + + +def test_summary_receives_only_abandoned_tail_and_attaches_at_target(tmp_path): + plugin = _plugin_module() + history = _history() + summary = Message(role="user", content="branch summary") + agent, context = _command_context(tmp_path, history) + with ( + context, + patch.object(plugin, "_emit"), + patch.object(plugin, "_summary_messages", return_value=[summary]) as summarize, + ): + plugin._handle_custom_command("/tree summary 2", "tree") + + assert summarize.call_args.args[1] == history[3:] + agent.set_message_history.assert_called_once_with([*history[:3], summary]) + + +def test_summary_failure_does_not_mutate_history(tmp_path): + plugin = _plugin_module() + history = _history() + agent, context = _command_context(tmp_path, history) + with ( + context, + patch.object(plugin, "_emit"), + patch.object( + plugin, "_summary_messages", side_effect=RuntimeError("no biscuits") + ), + ): + result = plugin._handle_custom_command("/tree summary 2", "tree") + + assert result is True + agent.set_message_history.assert_not_called() + + +def test_fork_without_argument_uses_tree_selector(tmp_path): + plugin = _plugin_module() + history = _history() + agent, context = _command_context(tmp_path, history) + + def select_last_user(tree, **_kwargs): + selected = [node for node in selectable_nodes(tree) if node.is_userish][-1] + return SimpleNamespace(node=selected, cancelled=False) + + with ( + context, + patch.object(plugin, "_emit"), + patch.object(plugin, "_run_tree_menu", side_effect=select_last_user), + ): + result = plugin._handle_custom_command("/fork", "fork") + + assert str(result) == "second task" + agent.set_message_history.assert_called_once_with(history[:4]) + + +def test_clone_does_not_rotate_current_session_id(tmp_path): + plugin = _plugin_module() + history = _history() + agent, context = _command_context(tmp_path, history) + agent.estimate_tokens_for_message.return_value = 1 + with context, patch.object(plugin, "_emit"): + assert plugin._handle_custom_command("/clone", "clone") is True + + assert len(list(tmp_path.glob("auto_session_*.pkl"))) == 1 + + +def test_unknown_command_is_ignored(): + assert _plugin_module()._handle_custom_command("/nope", "nope") is None