From 002012dfcd886692a521105d6441221ec10c9797 Mon Sep 17 00:00:00 2001 From: Ali Khokhar Date: Sun, 28 Jun 2026 01:32:34 -0700 Subject: [PATCH] Refactor messaging conversation state (#931) --- ARCHITECTURE.md | 19 +- api/runtime.py | 21 +- messaging/commands.py | 7 +- messaging/node_event_pipeline.py | 4 +- messaging/node_runner.py | 2 +- messaging/session.py | 334 ------------ messaging/session/__init__.py | 5 + messaging/session/message_log.py | 116 +++++ messaging/session/persistence.py | 131 +++++ messaging/session/store.py | 148 ++++++ messaging/trees/__init__.py | 6 +- messaging/trees/data.py | 482 ------------------ messaging/trees/graph.py | 156 ++++++ messaging/trees/manager.py | 18 +- messaging/trees/node.py | 57 +++ messaging/trees/processor.py | 3 +- messaging/trees/queue.py | 46 ++ messaging/trees/repository.py | 25 +- messaging/trees/runtime.py | 181 +++++++ messaging/trees/snapshot.py | 150 ++++++ messaging/turn_intake.py | 6 +- messaging/workflow.py | 4 +- pyproject.toml | 2 +- smoke/lib/e2e.py | 9 +- smoke/product/test_messaging_product_live.py | 11 +- tests/api/test_app_lifespan_and_errors.py | 35 +- tests/contracts/test_import_boundaries.py | 44 ++ tests/messaging/test_handler.py | 23 +- .../test_handler_context_isolation.py | 2 +- tests/messaging/test_handler_integration.py | 2 +- .../test_handler_markdown_and_status_edges.py | 22 +- tests/messaging/test_messaging.py | 48 +- tests/messaging/test_restart_reply_restore.py | 28 +- .../test_session_store_edge_cases.py | 187 ++++--- tests/messaging/test_tree_concurrency.py | 16 +- tests/messaging/test_tree_processor.py | 16 +- tests/messaging/test_tree_queue.py | 79 ++- tests/messaging/test_tree_repository.py | 29 +- uv.lock | 2 +- 39 files changed, 1389 insertions(+), 1087 deletions(-) delete mode 100644 messaging/session.py create mode 100644 messaging/session/__init__.py create mode 100644 messaging/session/message_log.py create mode 100644 messaging/session/persistence.py create mode 100644 messaging/session/store.py delete mode 100644 messaging/trees/data.py create mode 100644 messaging/trees/graph.py create mode 100644 messaging/trees/node.py create mode 100644 messaging/trees/queue.py create mode 100644 messaging/trees/runtime.py create mode 100644 messaging/trees/snapshot.py diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index f8215f5b..746b3ca8 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -582,12 +582,21 @@ each tree processes one node at a time while separate trees can progress independently. [messaging/trees/repository.py](messaging/trees/repository.py) owns the in-memory tree/node index, and [messaging/trees/processor.py](messaging/trees/processor.py) owns async queue -processing. [messaging/trees/data.py](messaging/trees/data.py) owns the persisted -`MessageNode` and `MessageTree` JSON shape. +processing. [messaging/trees/node.py](messaging/trees/node.py) owns +`MessageNode` and `MessageState`, +[messaging/trees/graph.py](messaging/trees/graph.py) owns parent/child and +status-message lookup state, [messaging/trees/runtime.py](messaging/trees/runtime.py) +owns locks/current-task/processing state, and +[messaging/trees/snapshot.py](messaging/trees/snapshot.py) owns typed persisted +conversation snapshots. -[messaging/session.py](messaging/session.py) persists trees, node-to-tree -mappings, and message IDs to a JSON file under the managed agent workspace. It -uses debounced atomic writes and flushes pending saves on shutdown. +[messaging/session/](messaging/session/) persists typed conversation snapshots +and message IDs to a JSON file under the managed agent workspace. +`SessionStore` reads existing `sessions.json` files but exposes typed snapshot +APIs to runtime code. Debounced atomic writes live in +[messaging/session/persistence.py](messaging/session/persistence.py), and +per-chat message ID tracking for `/clear` lives in +[messaging/session/message_log.py](messaging/session/message_log.py). ```mermaid sequenceDiagram diff --git a/api/runtime.py b/api/runtime.py index c4d8dc73..33be21e1 100644 --- a/api/runtime.py +++ b/api/runtime.py @@ -281,29 +281,28 @@ class AppRuntime: logger.info("{} platform started with messaging workflow", components.name) def _restore_tree_state(self, session_store: SessionStore) -> None: - saved_trees = session_store.get_all_trees() - if not saved_trees: + conversation_snapshot = session_store.load_conversation_snapshot() + if conversation_snapshot.is_empty: return if self.messaging_workflow is None: return - logger.info(f"Restoring {len(saved_trees)} conversation trees...") + logger.info( + "Restoring {} conversation trees...", + len(conversation_snapshot.trees), + ) from messaging.trees import TreeQueueManager self.messaging_workflow.replace_tree_queue( - TreeQueueManager.from_dict( - { - "trees": saved_trees, - "node_to_tree": session_store.get_node_mapping(), - }, + TreeQueueManager.from_snapshot( + conversation_snapshot, queue_update_callback=self.messaging_workflow.update_queue_positions, node_started_callback=self.messaging_workflow.mark_node_processing, ) ) if self.messaging_workflow.tree_queue.cleanup_stale_nodes() > 0: - tree_data = self.messaging_workflow.tree_queue.to_dict() - session_store.sync_from_tree_data( - tree_data["trees"], tree_data["node_to_tree"] + session_store.save_conversation_snapshot( + self.messaging_workflow.tree_queue.snapshot() ) def _publish_state(self) -> None: diff --git a/messaging/commands.py b/messaging/commands.py index 2808ebb5..26f9c9ec 100644 --- a/messaging/commands.py +++ b/messaging/commands.py @@ -160,19 +160,18 @@ async def _handle_clear_branch( await _delete_message_ids(handler, incoming.chat_id, msg_ids) # 4) Remove branch from tree - removed, root_id, removed_entire_tree = await handler.tree_queue.remove_branch( + _removed, root_id, removed_entire_tree = await handler.tree_queue.remove_branch( branch_root_id ) # 5) Update session store try: - handler.session_store.remove_node_mappings([n.node_id for n in removed]) if removed_entire_tree: - handler.session_store.remove_tree(root_id) + handler.session_store.remove_tree_snapshot(root_id) else: updated_tree = handler.tree_queue.get_tree(root_id) if updated_tree: - handler.session_store.save_tree(root_id, updated_tree.to_dict()) + handler.session_store.save_tree_snapshot(updated_tree.snapshot()) except Exception as e: logger.warning(f"Failed to update session store after branch clear: {e}") diff --git a/messaging/node_event_pipeline.py b/messaging/node_event_pipeline.py index 5f71851f..6d0dbdf1 100644 --- a/messaging/node_event_pipeline.py +++ b/messaging/node_event_pipeline.py @@ -51,7 +51,7 @@ async def handle_session_info_event( MessageState.IN_PROGRESS, session_id=real_session_id, ) - session_store.save_tree(tree.root_id, tree.to_dict()) + session_store.save_tree_snapshot(tree.snapshot()) return real_session_id, None @@ -101,7 +101,7 @@ async def process_parsed_cli_event( MessageState.COMPLETED, session_id=captured_session_id, ) - session_store.save_tree(tree.root_id, tree.to_dict()) + session_store.save_tree_snapshot(tree.snapshot()) elif ptype == "error": error_msg = parsed.get("message", "Unknown error") em = error_msg if isinstance(error_msg, str) else str(error_msg) diff --git a/messaging/node_runner.py b/messaging/node_runner.py index 350a4a8f..6cf75877 100644 --- a/messaging/node_runner.py +++ b/messaging/node_runner.py @@ -68,7 +68,7 @@ class MessagingNodeRunner: def _save_tree(self, tree: MessageTree | None) -> None: """Persist tree state after runner-owned mutations.""" if tree: - self.session_store.save_tree(tree.root_id, tree.to_dict()) + self.session_store.save_tree_snapshot(tree.snapshot()) async def process_node( self, diff --git a/messaging/session.py b/messaging/session.py deleted file mode 100644 index 89297726..00000000 --- a/messaging/session.py +++ /dev/null @@ -1,334 +0,0 @@ -""" -Session Store for Messaging Platforms - -Provides persistent storage for mapping platform messages to Claude CLI session IDs -and message trees for conversation continuation. -""" - -import contextlib -import json -import os -import tempfile -import threading -from datetime import UTC, datetime -from typing import Any - -from loguru import logger - - -class SessionStore: - """ - Persistent storage for message ↔ Claude session mappings and message trees. - - Uses a JSON file for storage with thread-safe operations. - Platform-agnostic: works with any messaging platform. - """ - - def __init__( - self, - storage_path: str = "sessions.json", - *, - message_log_cap: int | None = None, - ): - self.storage_path = storage_path - self._lock = threading.Lock() - self._trees: dict[str, dict] = {} # root_id -> tree data - self._node_to_tree: dict[str, str] = {} # node_id -> root_id - # Per-chat message ID log used to support best-effort UI clearing (/clear). - # Key: "{platform}:{chat_id}" -> list of records - self._message_log: dict[str, list[dict[str, Any]]] = {} - self._message_log_ids: dict[str, set[str]] = {} - self._dirty = False - self._save_timer: threading.Timer | None = None - self._save_debounce_secs = 0.5 - self._message_log_cap: int | None = message_log_cap - self._load() - - def _make_chat_key(self, platform: str, chat_id: str) -> str: - return f"{platform}:{chat_id}" - - def _load(self) -> None: - """Load sessions and trees from disk.""" - if not os.path.exists(self.storage_path): - return - - try: - with open(self.storage_path, encoding="utf-8") as f: - data = json.load(f) - - # Load trees - self._trees = data.get("trees", {}) - self._node_to_tree = data.get("node_to_tree", {}) - - # Load message log (optional/backward compatible) - raw_log = data.get("message_log", {}) or {} - if isinstance(raw_log, dict): - self._message_log = {} - self._message_log_ids = {} - for chat_key, items in raw_log.items(): - if not isinstance(chat_key, str) or not isinstance(items, list): - continue - cleaned: list[dict[str, Any]] = [] - seen: set[str] = set() - for it in items: - if not isinstance(it, dict): - continue - mid = it.get("message_id") - if mid is None: - continue - mid_s = str(mid) - if mid_s in seen: - continue - seen.add(mid_s) - cleaned.append( - { - "message_id": mid_s, - "ts": str(it.get("ts") or ""), - "direction": str(it.get("direction") or ""), - "kind": str(it.get("kind") or ""), - } - ) - self._message_log[chat_key] = cleaned - self._message_log_ids[chat_key] = seen - - logger.info( - f"Loaded {len(self._trees)} trees and " - f"{sum(len(v) for v in self._message_log.values())} msg_ids from {self.storage_path}" - ) - except Exception as e: - logger.error(f"Failed to load sessions: {e}") - - def _snapshot(self) -> dict: - """Snapshot current state for serialization. Caller must hold self._lock.""" - return { - "trees": dict(self._trees), - "node_to_tree": dict(self._node_to_tree), - "message_log": {k: list(v) for k, v in self._message_log.items()}, - } - - def _write_data(self, data: dict) -> None: - """Atomically write data dict to disk. Must be called WITHOUT holding self._lock.""" - abs_target = os.path.abspath(self.storage_path) - dir_name = os.path.dirname(abs_target) or "." - fd, tmp_path = tempfile.mkstemp( - dir=dir_name, prefix=".sessions.", suffix=".tmp.json" - ) - try: - with os.fdopen(fd, "w", encoding="utf-8") as f: - json.dump(data, f, indent=2) - f.flush() - os.fsync(f.fileno()) - os.replace(tmp_path, abs_target) - except BaseException: - with contextlib.suppress(OSError): - os.unlink(tmp_path) - raise - - def _schedule_save(self) -> None: - """Schedule a debounced save. Caller must hold self._lock.""" - self._dirty = True - if self._save_timer is not None: - self._save_timer.cancel() - self._save_timer = None - self._save_timer = threading.Timer( - self._save_debounce_secs, self._save_from_timer - ) - self._save_timer.daemon = True - self._save_timer.start() - - def _save_from_timer(self) -> None: - """Timer callback: save if dirty. Runs in timer thread.""" - with self._lock: - if not self._dirty: - self._save_timer = None - return - snapshot = self._snapshot() - self._dirty = False - self._save_timer = None - try: - self._write_data(snapshot) - except Exception as e: - logger.error(f"Failed to save sessions: {e}") - with self._lock: - self._dirty = True - - def _flush_save(self) -> dict: - """Cancel pending timer and snapshot current state. Caller must hold self._lock. - Returns snapshot dict; caller must call _write_data(snapshot) outside the lock.""" - if self._save_timer is not None: - self._save_timer.cancel() - self._save_timer = None - self._dirty = False - return self._snapshot() - - def flush_pending_save(self) -> None: - """Flush any pending debounced save. Call on shutdown to avoid losing data.""" - with self._lock: - snapshot = self._flush_save() - try: - self._write_data(snapshot) - except Exception as e: - logger.error(f"Failed to save sessions: {e}") - with self._lock: - self._dirty = True - - def record_message_id( - self, - platform: str, - chat_id: str, - message_id: str, - direction: str, - kind: str, - ) -> None: - """Record a message_id for later best-effort deletion (/clear).""" - if message_id is None: - return - - chat_key = self._make_chat_key(str(platform), str(chat_id)) - mid = str(message_id) - - with self._lock: - seen = self._message_log_ids.setdefault(chat_key, set()) - if mid in seen: - return - - rec = { - "message_id": mid, - "ts": datetime.now(UTC).isoformat(), - "direction": str(direction), - "kind": str(kind), - } - self._message_log.setdefault(chat_key, []).append(rec) - seen.add(mid) - - # Optional cap to prevent unbounded growth if configured. - if self._message_log_cap is not None and self._message_log_cap > 0: - items = self._message_log.get(chat_key, []) - if len(items) > self._message_log_cap: - self._message_log[chat_key] = items[-self._message_log_cap :] - self._message_log_ids[chat_key] = { - str(x.get("message_id")) for x in self._message_log[chat_key] - } - - self._schedule_save() - - def get_message_ids_for_chat(self, platform: str, chat_id: str) -> list[str]: - """Get all recorded message IDs for a chat (in insertion order).""" - chat_key = self._make_chat_key(str(platform), str(chat_id)) - with self._lock: - items = self._message_log.get(chat_key, []) - return [ - str(x.get("message_id")) - for x in items - if x.get("message_id") is not None - ] - - def clear_all(self) -> None: - """Clear all stored sessions/trees/mappings and persist an empty store.""" - with self._lock: - self._trees.clear() - self._node_to_tree.clear() - self._message_log.clear() - self._message_log_ids.clear() - snapshot = self._flush_save() - try: - self._write_data(snapshot) - except Exception as e: - logger.error(f"Failed to save sessions: {e}") - with self._lock: - self._dirty = True - - # ==================== Tree Methods ==================== - - @staticmethod - def _tree_lookup_ids(tree_data: dict) -> set[str]: - """Return lookup IDs represented by serialized tree nodes.""" - lookup_ids: set[str] = set() - nodes = tree_data.get("nodes", {}) - if not isinstance(nodes, dict): - return lookup_ids - - for node_key, node_data in nodes.items(): - lookup_ids.add(str(node_key)) - if not isinstance(node_data, dict): - continue - node_id = node_data.get("node_id") - if node_id is not None: - lookup_ids.add(str(node_id)) - status_message_id = node_data.get("status_message_id") - if status_message_id is not None: - lookup_ids.add(str(status_message_id)) - return lookup_ids - - def _remove_tree_lookup_ids_unlocked(self, root_id: str) -> None: - """Remove all lookup IDs currently pointing at a root. Caller holds lock.""" - stale_lookup_ids = [ - lookup_id - for lookup_id, mapped_root_id in self._node_to_tree.items() - if mapped_root_id == root_id - ] - for lookup_id in stale_lookup_ids: - self._node_to_tree.pop(lookup_id, None) - - def save_tree(self, root_id: str, tree_data: dict) -> None: - """ - Save a message tree. - - Args: - root_id: Root node ID of the tree - tree_data: Serialized tree data from tree.to_dict() - """ - with self._lock: - self._trees[root_id] = tree_data - - self._remove_tree_lookup_ids_unlocked(root_id) - for lookup_id in self._tree_lookup_ids(tree_data): - self._node_to_tree[lookup_id] = root_id - - self._schedule_save() - logger.debug(f"Saved tree {root_id}") - - def get_tree(self, root_id: str) -> dict | None: - """Get a tree by its root ID.""" - with self._lock: - return self._trees.get(root_id) - - def register_node(self, node_id: str, root_id: str) -> None: - """Register a node ID to a tree root.""" - with self._lock: - self._node_to_tree[node_id] = root_id - self._schedule_save() - - def remove_node_mappings(self, node_ids: list[str]) -> None: - """Remove node IDs from the node-to-tree mapping.""" - with self._lock: - for nid in node_ids: - self._node_to_tree.pop(nid, None) - self._schedule_save() - - def remove_tree(self, root_id: str) -> None: - """Remove a tree and all its node mappings from the store.""" - with self._lock: - tree_data = self._trees.pop(root_id, None) - if tree_data: - self._remove_tree_lookup_ids_unlocked(root_id) - self._schedule_save() - - def get_all_trees(self) -> dict[str, dict]: - """Get all stored trees (public accessor).""" - with self._lock: - return dict(self._trees) - - def get_node_mapping(self) -> dict[str, str]: - """Get the node-to-tree mapping (public accessor).""" - with self._lock: - return dict(self._node_to_tree) - - def sync_from_tree_data( - self, trees: dict[str, dict], node_to_tree: dict[str, str] - ) -> None: - """Sync internal tree state from external data and persist.""" - with self._lock: - self._trees = trees - self._node_to_tree = node_to_tree - self._schedule_save() diff --git a/messaging/session/__init__.py b/messaging/session/__init__.py new file mode 100644 index 00000000..cdf1f367 --- /dev/null +++ b/messaging/session/__init__.py @@ -0,0 +1,5 @@ +"""Public messaging session persistence API.""" + +from .store import SessionStore + +__all__ = ["SessionStore"] diff --git a/messaging/session/message_log.py b/messaging/session/message_log.py new file mode 100644 index 00000000..98d93e1b --- /dev/null +++ b/messaging/session/message_log.py @@ -0,0 +1,116 @@ +"""Per-chat message ID log used by messaging clear commands.""" + +from __future__ import annotations + +from datetime import UTC, datetime +from typing import Any + + +class MessageLog: + """Track inbound/outbound platform message IDs in insertion order.""" + + def __init__(self, *, cap: int | None = None) -> None: + self._items: dict[str, list[dict[str, Any]]] = {} + self._ids: dict[str, set[str]] = {} + self._cap = cap + + @property + def cap(self) -> int | None: + return self._cap + + @classmethod + def from_json(cls, raw_log: Any, *, cap: int | None = None) -> MessageLog: + log = cls(cap=cap) + if not isinstance(raw_log, dict): + return log + for chat_key, items in raw_log.items(): + if not isinstance(chat_key, str) or not isinstance(items, list): + continue + for item in items: + if not isinstance(item, dict): + continue + message_id = item.get("message_id") + if message_id is None: + continue + log._append( + chat_key, + str(message_id), + ts=str(item.get("ts") or ""), + direction=str(item.get("direction") or ""), + kind=str(item.get("kind") or ""), + ) + return log + + def to_json(self) -> dict[str, list[dict[str, Any]]]: + return {chat_key: list(items) for chat_key, items in self._items.items()} + + def record( + self, + *, + platform: str, + chat_id: str, + message_id: str, + direction: str, + kind: str, + ) -> bool: + chat_key = make_chat_key(platform, chat_id) + return self._append( + chat_key, + str(message_id), + ts=datetime.now(UTC).isoformat(), + direction=str(direction), + kind=str(kind), + ) + + def get_message_ids_for_chat(self, platform: str, chat_id: str) -> list[str]: + chat_key = make_chat_key(platform, chat_id) + return [ + str(item.get("message_id")) + for item in self._items.get(chat_key, []) + if item.get("message_id") is not None + ] + + def clear(self) -> None: + self._items.clear() + self._ids.clear() + + def _append( + self, + chat_key: str, + message_id: str, + *, + ts: str, + direction: str, + kind: str, + ) -> bool: + seen = self._ids.setdefault(chat_key, set()) + if message_id in seen: + return False + self._items.setdefault(chat_key, []).append( + { + "message_id": message_id, + "ts": ts, + "direction": direction, + "kind": kind, + } + ) + seen.add(message_id) + self._trim(chat_key) + return True + + def _trim(self, chat_key: str) -> None: + if self._cap is None or self._cap <= 0: + return + items = self._items.get(chat_key, []) + if len(items) <= self._cap: + return + self._items[chat_key] = items[-self._cap :] + self._ids[chat_key] = { + str(item.get("message_id")) + for item in self._items[chat_key] + if item.get("message_id") is not None + } + + +def make_chat_key(platform: str, chat_id: str) -> str: + return f"{platform}:{chat_id}" diff --git a/messaging/session/persistence.py b/messaging/session/persistence.py new file mode 100644 index 00000000..142c7db2 --- /dev/null +++ b/messaging/session/persistence.py @@ -0,0 +1,131 @@ +"""Atomic JSON persistence for messaging session state.""" + +from __future__ import annotations + +import contextlib +import json +import os +import tempfile +import threading +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +from loguru import logger + + +@dataclass(frozen=True) +class _PendingWrite: + generation: int + snapshot: dict[str, Any] + + +class DebouncedJsonPersistence: + """Thread-safe debounced JSON writer with atomic replace semantics.""" + + def __init__( + self, + storage_path: str, + *, + snapshot: Callable[[], dict[str, Any]], + on_dirty: Callable[[bool], None], + debounce_secs: float = 0.5, + ) -> None: + self.storage_path = storage_path + self._snapshot = snapshot + self._on_dirty = on_dirty + self._debounce_secs = debounce_secs + self._save_timer: threading.Timer | None = None + self._timer_lock = threading.Lock() + self._save_generation = 0 + + def load_json(self) -> dict[str, Any]: + if not os.path.exists(self.storage_path): + return {} + with open(self.storage_path, encoding="utf-8") as file: + data = json.load(file) + return data if isinstance(data, dict) else {} + + def schedule_save(self) -> None: + self._on_dirty(True) + with self._timer_lock: + if self._save_timer is not None: + self._save_timer.cancel() + self._save_generation += 1 + generation = self._save_generation + timer = threading.Timer( + self._debounce_secs, + self._save_from_timer, + args=(generation,), + ) + timer.daemon = True + self._save_timer = timer + timer.start() + + def flush(self) -> None: + pending = self._snapshot_for_write() + if pending is None: + return + self._write_pending(pending) + + def _save_from_timer(self, generation: int) -> None: + pending = self._snapshot_for_write(expected_generation=generation) + if pending is None: + return + self._write_pending(pending) + + def _write_pending(self, pending: _PendingWrite) -> None: + try: + self.write_data(pending.snapshot) + except Exception as e: + logger.error("Failed to save sessions: {}", e) + self._on_dirty(True) + return + self._mark_clean_if_current(pending.generation) + + def _snapshot_for_write( + self, *, expected_generation: int | None = None + ) -> _PendingWrite | None: + generation = self._claim_timer(expected_generation) + if generation is None: + return None + snapshot = self._snapshot() + return _PendingWrite(generation=generation, snapshot=snapshot) + + def _claim_timer(self, expected_generation: int | None) -> int | None: + with self._timer_lock: + if expected_generation is not None and ( + expected_generation != self._save_generation or self._save_timer is None + ): + return None + if self._save_timer is not None: + self._save_timer.cancel() + self._save_timer = None + return self._save_generation + + def _mark_clean_if_current(self, generation: int) -> None: + with self._timer_lock: + is_current = ( + self._save_timer is None and generation == self._save_generation + ) + if is_current: + self._on_dirty(False) + + def write_data(self, data: dict[str, Any]) -> None: + abs_target = os.path.abspath(self.storage_path) + dir_name = os.path.dirname(abs_target) or "." + fd, tmp_path = tempfile.mkstemp( + dir=dir_name, + prefix=".sessions.", + suffix=".tmp.json", + ) + try: + with os.fdopen(fd, "w", encoding="utf-8") as file: + json.dump(data, file, indent=2) + file.flush() + os.fsync(file.fileno()) + os.replace(tmp_path, abs_target) + except BaseException: + with contextlib.suppress(OSError): + os.unlink(tmp_path) + raise diff --git a/messaging/session/store.py b/messaging/session/store.py new file mode 100644 index 00000000..55ec158b --- /dev/null +++ b/messaging/session/store.py @@ -0,0 +1,148 @@ +"""Persistent messaging conversation state store.""" + +from __future__ import annotations + +import threading + +from loguru import logger + +from messaging.trees import ConversationSnapshot, TreeSnapshot + +from .message_log import MessageLog +from .persistence import DebouncedJsonPersistence + + +class SessionStore: + """ + Persistent storage for conversation snapshots and message IDs. + + The store reads both the old raw ``trees``/``node_to_tree`` shape and the + current typed ``conversation`` snapshot shape. Runtime callers deal in typed + snapshots only. + """ + + def __init__( + self, + storage_path: str = "sessions.json", + *, + message_log_cap: int | None = None, + ) -> None: + self.storage_path = storage_path + self._lock = threading.RLock() + self._conversation = ConversationSnapshot() + self._message_log = MessageLog(cap=message_log_cap) + self._dirty = False + self._persistence = DebouncedJsonPersistence( + storage_path, + snapshot=self._snapshot_for_persistence, + on_dirty=self._set_dirty, + ) + self._load() + + @property + def dirty(self) -> bool: + return self._dirty + + def _set_dirty(self, dirty: bool) -> None: + with self._lock: + self._dirty = dirty + + def _load(self) -> None: + try: + data = self._persistence.load_json() + except Exception as e: + logger.error("Failed to load sessions: {}", e) + return + + conversation_data = data.get("conversation") if isinstance(data, dict) else None + if not isinstance(conversation_data, dict): + conversation_data = data + + with self._lock: + self._conversation = ConversationSnapshot.from_json(conversation_data) + self._message_log = MessageLog.from_json( + data.get("message_log", {}) if isinstance(data, dict) else {}, + cap=self._message_log.cap, + ) + message_count = sum( + len(items) for items in self._message_log.to_json().values() + ) + logger.info( + "Loaded {} trees and {} msg_ids from {}", + len(self._conversation.trees), + message_count, + self.storage_path, + ) + + def _snapshot_for_persistence(self) -> dict: + with self._lock: + return { + "conversation": self._conversation.to_json(), + "message_log": self._message_log.to_json(), + } + + def load_conversation_snapshot(self) -> ConversationSnapshot: + with self._lock: + return ConversationSnapshot(trees=dict(self._conversation.trees)) + + def save_conversation_snapshot(self, snapshot: ConversationSnapshot) -> None: + with self._lock: + self._conversation = snapshot + self._persistence.schedule_save() + + def save_tree_snapshot(self, snapshot: TreeSnapshot) -> None: + with self._lock: + self._conversation = self._conversation.with_tree(snapshot) + self._persistence.schedule_save() + logger.debug("Saved tree {}", snapshot.root_id) + + def get_tree_snapshot(self, root_id: str) -> TreeSnapshot | None: + with self._lock: + return self._conversation.trees.get(root_id) + + def remove_tree_snapshot(self, root_id: str) -> None: + with self._lock: + self._conversation = self._conversation.without_tree(root_id) + self._persistence.schedule_save() + + def flush_pending_save(self) -> None: + self._persistence.flush() + + def record_message_id( + self, + platform: str, + chat_id: str, + message_id: str, + direction: str, + kind: str, + ) -> None: + if message_id is None: + return + with self._lock: + recorded = self._message_log.record( + platform=str(platform), + chat_id=str(chat_id), + message_id=str(message_id), + direction=str(direction), + kind=str(kind), + ) + if recorded: + self._persistence.schedule_save() + + def get_message_ids_for_chat(self, platform: str, chat_id: str) -> list[str]: + with self._lock: + return self._message_log.get_message_ids_for_chat( + str(platform), str(chat_id) + ) + + def clear_all(self) -> None: + with self._lock: + self._conversation = ConversationSnapshot() + self._message_log.clear() + snapshot = self._snapshot_for_persistence() + self._set_dirty(False) + try: + self._persistence.write_data(snapshot) + except Exception as e: + logger.error("Failed to save sessions: {}", e) + self._set_dirty(True) diff --git a/messaging/trees/__init__.py b/messaging/trees/__init__.py index cc96cb2e..f444b3cf 100644 --- a/messaging/trees/__init__.py +++ b/messaging/trees/__init__.py @@ -1,15 +1,19 @@ """Message tree data structures and queue management.""" -from .data import MessageNode, MessageState, MessageTree from .manager import TreeQueueManager +from .node import MessageNode, MessageState from .processor import TreeQueueProcessor from .repository import TreeRepository +from .runtime import MessageTree +from .snapshot import ConversationSnapshot, TreeSnapshot __all__ = [ + "ConversationSnapshot", "MessageNode", "MessageState", "MessageTree", "TreeQueueManager", "TreeQueueProcessor", "TreeRepository", + "TreeSnapshot", ] diff --git a/messaging/trees/data.py b/messaging/trees/data.py deleted file mode 100644 index d5db01db..00000000 --- a/messaging/trees/data.py +++ /dev/null @@ -1,482 +0,0 @@ -"""Tree data structures for message queue. - -Contains MessageState, MessageNode, and MessageTree classes. -""" - -import asyncio -from collections import deque -from contextlib import asynccontextmanager -from dataclasses import dataclass, field -from datetime import UTC, datetime -from enum import Enum -from typing import Any - -from loguru import logger - -from ..models import IncomingMessage - - -class _SnapshotQueue: - """Queue with snapshot/remove helpers, backed by a deque and a set index.""" - - def __init__(self) -> None: - self._deque: deque[str] = deque() - self._set: set[str] = set() - - async def put(self, item: str) -> None: - self._deque.append(item) - self._set.add(item) - - def put_nowait(self, item: str) -> None: - self._deque.append(item) - self._set.add(item) - - def get_nowait(self) -> str: - if not self._deque: - raise asyncio.QueueEmpty() - item = self._deque.popleft() - self._set.discard(item) - return item - - def qsize(self) -> int: - return len(self._deque) - - def get_snapshot(self) -> list[str]: - """Return current queue contents in FIFO order (read-only copy).""" - return list(self._deque) - - def remove_if_present(self, item: str) -> bool: - """Remove item from queue if present (O(1) membership check). Returns True if removed.""" - if item not in self._set: - return False - self._set.discard(item) - self._deque = deque(x for x in self._deque if x != item) - return True - - -class MessageState(Enum): - """State of a message node in the tree.""" - - PENDING = "pending" # Queued, waiting to be processed - IN_PROGRESS = "in_progress" # Currently being processed by Claude - COMPLETED = "completed" # Processing finished successfully - ERROR = "error" # Processing failed - - -@dataclass -class MessageNode: - """ - A node in the message tree. - - Each node represents a single message and tracks: - - Its relationship to parent/children - - Its processing state - - Claude session information - """ - - node_id: str # Unique ID (typically message_id) - incoming: IncomingMessage # The original message - status_message_id: str # Bot's status message ID - state: MessageState = MessageState.PENDING - parent_id: str | None = None # Parent node ID (None for root) - session_id: str | None = None # Claude session ID (forked from parent) - children_ids: list[str] = field(default_factory=list) - created_at: datetime = field(default_factory=lambda: datetime.now(UTC)) - completed_at: datetime | None = None - error_message: str | None = None - context: Any = None # Additional context if needed - - def set_context(self, context: Any) -> None: - self.context = context - - def to_dict(self) -> dict: - """Convert to dictionary for JSON serialization.""" - return { - "node_id": self.node_id, - "incoming": { - "text": self.incoming.text, - "chat_id": self.incoming.chat_id, - "user_id": self.incoming.user_id, - "message_id": self.incoming.message_id, - "platform": self.incoming.platform, - "reply_to_message_id": self.incoming.reply_to_message_id, - "message_thread_id": self.incoming.message_thread_id, - "username": self.incoming.username, - }, - "status_message_id": self.status_message_id, - "state": self.state.value, - "parent_id": self.parent_id, - "session_id": self.session_id, - "children_ids": self.children_ids, - "created_at": self.created_at.isoformat(), - "completed_at": self.completed_at.isoformat() - if self.completed_at - else None, - "error_message": self.error_message, - } - - @classmethod - def from_dict(cls, data: dict) -> MessageNode: - """Create from dictionary (JSON deserialization).""" - incoming_data = data["incoming"] - incoming = IncomingMessage( - text=incoming_data["text"], - chat_id=incoming_data["chat_id"], - user_id=incoming_data["user_id"], - message_id=incoming_data["message_id"], - platform=incoming_data["platform"], - reply_to_message_id=incoming_data.get("reply_to_message_id"), - message_thread_id=incoming_data.get("message_thread_id"), - username=incoming_data.get("username"), - ) - return cls( - node_id=data["node_id"], - incoming=incoming, - status_message_id=data["status_message_id"], - state=MessageState(data["state"]), - parent_id=data.get("parent_id"), - session_id=data.get("session_id"), - children_ids=data.get("children_ids", []), - created_at=datetime.fromisoformat(data["created_at"]), - completed_at=datetime.fromisoformat(data["completed_at"]) - if data.get("completed_at") - else None, - error_message=data.get("error_message"), - ) - - -class MessageTree: - """ - A tree of message nodes with queue functionality. - - Provides: - - O(1) node lookup via hashmap - - Per-tree message queue - - Thread-safe operations via asyncio.Lock - """ - - def __init__(self, root_node: MessageNode): - """ - Initialize tree with a root node. - - Args: - root_node: The root message node - """ - self.root_id = root_node.node_id - self._nodes: dict[str, MessageNode] = {root_node.node_id: root_node} - self._status_to_node: dict[str, str] = { - root_node.status_message_id: root_node.node_id - } - self._queue: _SnapshotQueue = _SnapshotQueue() - self._lock = asyncio.Lock() - self._is_processing = False - self._current_node_id: str | None = None - self._current_task: asyncio.Task | None = None - - logger.debug(f"Created MessageTree with root {self.root_id}") - - def set_current_task(self, task: asyncio.Task | None) -> None: - """Set the current processing task. Caller must hold lock.""" - self._current_task = task - - @property - def is_processing(self) -> bool: - """Check if tree is currently processing a message.""" - return self._is_processing - - async def add_node( - self, - node_id: str, - incoming: IncomingMessage, - status_message_id: str, - parent_id: str, - ) -> MessageNode: - """ - Add a child node to the tree. - - Args: - node_id: Unique ID for the new node - incoming: The incoming message - status_message_id: Bot's status message ID - parent_id: Parent node ID - - Returns: - The created MessageNode - """ - async with self._lock: - if parent_id not in self._nodes: - raise ValueError(f"Parent node {parent_id} not found in tree") - - node = MessageNode( - node_id=node_id, - incoming=incoming, - status_message_id=status_message_id, - parent_id=parent_id, - state=MessageState.PENDING, - ) - - self._nodes[node_id] = node - self._status_to_node[status_message_id] = node_id - self._nodes[parent_id].children_ids.append(node_id) - - logger.debug(f"Added node {node_id} as child of {parent_id}") - return node - - def get_node(self, node_id: str) -> MessageNode | None: - """Get a node by ID (O(1) lookup).""" - return self._nodes.get(node_id) - - def get_root(self) -> MessageNode: - """Get the root node.""" - return self._nodes[self.root_id] - - def get_children(self, node_id: str) -> list[MessageNode]: - """Get all child nodes of a given node.""" - node = self._nodes.get(node_id) - if not node: - return [] - return [self._nodes[cid] for cid in node.children_ids if cid in self._nodes] - - def get_parent(self, node_id: str) -> MessageNode | None: - """Get the parent node.""" - node = self._nodes.get(node_id) - if not node or not node.parent_id: - return None - return self._nodes.get(node.parent_id) - - def get_parent_session_id(self, node_id: str) -> str | None: - """ - Get the parent's session ID for forking. - - Returns None for root nodes. - """ - parent = self.get_parent(node_id) - return parent.session_id if parent else None - - async def update_state( - self, - node_id: str, - state: MessageState, - session_id: str | None = None, - error_message: str | None = None, - ) -> None: - """Update a node's state.""" - async with self._lock: - node = self._nodes.get(node_id) - if not node: - logger.warning(f"Node {node_id} not found for state update") - return - - node.state = state - if session_id: - node.session_id = session_id - if error_message: - node.error_message = error_message - if state in (MessageState.COMPLETED, MessageState.ERROR): - node.completed_at = datetime.now(UTC) - - logger.debug(f"Node {node_id} state -> {state.value}") - - async def enqueue(self, node_id: str) -> int: - """ - Add a node to the processing queue. - - Returns: - Queue position (1-indexed) - """ - async with self._lock: - await self._queue.put(node_id) - position = self._queue.qsize() - logger.debug(f"Enqueued node {node_id}, position {position}") - return position - - async def dequeue(self) -> str | None: - """ - Get the next node ID from the queue. - - Returns None if queue is empty. - """ - try: - return self._queue.get_nowait() - except asyncio.QueueEmpty: - return None - - async def get_queue_snapshot(self) -> list[str]: - """ - Get a snapshot of the current queue order. - - Returns: - List of node IDs in FIFO order. - """ - async with self._lock: - return self._queue.get_snapshot() - - def get_queue_size(self) -> int: - """Get number of messages waiting in queue.""" - return self._queue.qsize() - - def remove_from_queue(self, node_id: str) -> bool: - """ - Remove node_id from the internal queue if present. - - Caller must hold the tree lock (e.g. via with_lock). - Returns True if node was removed, False if not in queue. - """ - return self._queue.remove_if_present(node_id) - - @asynccontextmanager - async def with_lock(self): - """Async context manager for tree lock. Use when multiple operations need atomicity.""" - async with self._lock: - yield - - def set_processing_state(self, node_id: str | None, is_processing: bool) -> None: - """Set processing state. Caller must hold lock for consistency with queue operations.""" - self._is_processing = is_processing - self._current_node_id = node_id if is_processing else None - - def clear_current_node(self) -> None: - """Clear the currently processing node ID. Caller must hold lock.""" - self._current_node_id = None - - def is_current_node(self, node_id: str) -> bool: - """Check if node_id is the currently processing node.""" - return self._current_node_id == node_id - - def put_queue_unlocked(self, node_id: str) -> None: - """Add node to queue. Caller must hold lock (e.g. via with_lock).""" - self._queue.put_nowait(node_id) - - def cancel_current_task(self) -> bool: - """Cancel the currently running task. Returns True if a task was cancelled.""" - if self._current_task and not self._current_task.done(): - self._current_task.cancel() - return True - return False - - def set_node_error_sync(self, node: MessageNode, error_message: str) -> None: - """Synchronously mark a node as ERROR. Caller must ensure no concurrent access.""" - node.state = MessageState.ERROR - node.error_message = error_message - node.completed_at = datetime.now(UTC) - - def drain_queue_and_mark_cancelled( - self, error_message: str = "Cancelled by user" - ) -> list[MessageNode]: - """ - Drain the queue, mark each node as ERROR, and return affected nodes. - Does not acquire lock; caller must ensure no concurrent queue access. - """ - nodes: list[MessageNode] = [] - while True: - try: - node_id = self._queue.get_nowait() - except asyncio.QueueEmpty: - break - node = self._nodes.get(node_id) - if node: - self.set_node_error_sync(node, error_message) - nodes.append(node) - return nodes - - def reset_processing_state(self) -> None: - """Reset processing flags after cancel/cleanup.""" - self._is_processing = False - self._current_node_id = None - - @property - def current_node_id(self) -> str | None: - """Get the ID of the node currently being processed.""" - return self._current_node_id - - def to_dict(self) -> dict: - """Serialize tree to dictionary.""" - return { - "root_id": self.root_id, - "nodes": {nid: node.to_dict() for nid, node in self._nodes.items()}, - } - - def _add_node_from_dict(self, node: MessageNode) -> None: - """Register a deserialized node into the tree's internal indices.""" - self._nodes[node.node_id] = node - self._status_to_node[node.status_message_id] = node.node_id - - @classmethod - def from_dict(cls, data: dict) -> MessageTree: - """Deserialize tree from dictionary.""" - root_id = data["root_id"] - nodes_data = data["nodes"] - - # Create root node first - root_node = MessageNode.from_dict(nodes_data[root_id]) - tree = cls(root_node) - - # Add remaining nodes and build status->node index - for node_id, node_data in nodes_data.items(): - if node_id != root_id: - node = MessageNode.from_dict(node_data) - tree._add_node_from_dict(node) - - return tree - - def all_nodes(self) -> list[MessageNode]: - """Get all nodes in the tree.""" - return list(self._nodes.values()) - - def has_node(self, node_id: str) -> bool: - """Check if a node exists in this tree.""" - return node_id in self._nodes - - def find_node_by_status_message(self, status_msg_id: str) -> MessageNode | None: - """Find the node that has this status message ID (O(1) lookup).""" - node_id = self._status_to_node.get(status_msg_id) - return self._nodes.get(node_id) if node_id else None - - def get_descendants(self, node_id: str) -> list[str]: - """ - Get node_id and all descendant IDs (subtree). - - Returns: - List of node IDs including the given node. - """ - if node_id not in self._nodes: - return [] - result: list[str] = [] - stack = [node_id] - while stack: - nid = stack.pop() - result.append(nid) - node = self._nodes.get(nid) - if node: - stack.extend(node.children_ids) - return result - - def remove_branch(self, branch_root_id: str) -> list[MessageNode]: - """ - Remove a subtree (branch_root and all descendants) from the tree. - - Updates parent's children_ids. Caller must hold lock for consistency. - Does not acquire lock internally. - - Returns: - List of removed nodes. - """ - if branch_root_id not in self._nodes: - return [] - - parent = self.get_parent(branch_root_id) - removed = [] - for nid in self.get_descendants(branch_root_id): - node = self._nodes.get(nid) - if node: - removed.append(node) - del self._nodes[nid] - del self._status_to_node[node.status_message_id] - - if parent and branch_root_id in parent.children_ids: - parent.children_ids = [ - c for c in parent.children_ids if c != branch_root_id - ] - - logger.debug(f"Removed branch {branch_root_id} ({len(removed)} nodes)") - return removed diff --git a/messaging/trees/graph.py b/messaging/trees/graph.py new file mode 100644 index 00000000..85323c7b --- /dev/null +++ b/messaging/trees/graph.py @@ -0,0 +1,156 @@ +"""In-memory graph for one messaging conversation tree.""" + +from __future__ import annotations + +from loguru import logger + +from ..models import IncomingMessage +from .node import MessageNode, MessageState +from .snapshot import TreeSnapshot, node_from_snapshot, node_to_snapshot + + +class MessageTreeGraph: + """Own parent/child links, node lookup, and status-message lookup.""" + + def __init__(self, root_node: MessageNode) -> None: + self.root_id = root_node.node_id + self._nodes: dict[str, MessageNode] = {root_node.node_id: root_node} + self._status_to_node: dict[str, str] = { + root_node.status_message_id: root_node.node_id + } + + def add_node( + self, + *, + node_id: str, + incoming: IncomingMessage, + status_message_id: str, + parent_id: str, + ) -> MessageNode: + if parent_id not in self._nodes: + raise ValueError(f"Parent node {parent_id} not found in tree") + + node = MessageNode( + node_id=node_id, + incoming=incoming, + status_message_id=status_message_id, + parent_id=parent_id, + state=MessageState.PENDING, + ) + self._nodes[node_id] = node + self._status_to_node[status_message_id] = node_id + self._nodes[parent_id].children_ids.append(node_id) + logger.debug("Added node {} as child of {}", node_id, parent_id) + return node + + def get_node(self, node_id: str) -> MessageNode | None: + return self._nodes.get(node_id) + + def get_root(self) -> MessageNode: + return self._nodes[self.root_id] + + def get_children(self, node_id: str) -> list[MessageNode]: + node = self._nodes.get(node_id) + if not node: + return [] + return [self._nodes[cid] for cid in node.children_ids if cid in self._nodes] + + def get_parent(self, node_id: str) -> MessageNode | None: + node = self._nodes.get(node_id) + if not node or not node.parent_id: + return None + return self._nodes.get(node.parent_id) + + def get_parent_session_id(self, node_id: str) -> str | None: + parent = self.get_parent(node_id) + return parent.session_id if parent else None + + def update_node_state( + self, + node_id: str, + state: MessageState, + *, + session_id: str | None = None, + error_message: str | None = None, + ) -> bool: + node = self._nodes.get(node_id) + if not node: + logger.warning("Node {} not found for state update", node_id) + return False + node.update_state( + state, + session_id=session_id, + error_message=error_message, + ) + logger.debug("Node {} state -> {}", node_id, state.value) + return True + + def has_node(self, node_id: str) -> bool: + return node_id in self._nodes + + def find_node_by_status_message(self, status_msg_id: str) -> MessageNode | None: + node_id = self._status_to_node.get(status_msg_id) + return self._nodes.get(node_id) if node_id else None + + def all_nodes(self) -> list[MessageNode]: + return list(self._nodes.values()) + + def get_descendants(self, node_id: str) -> list[str]: + if node_id not in self._nodes: + return [] + result: list[str] = [] + stack = [node_id] + while stack: + current_id = stack.pop() + result.append(current_id) + node = self._nodes.get(current_id) + if node: + stack.extend(node.children_ids) + return result + + def remove_branch(self, branch_root_id: str) -> list[MessageNode]: + if branch_root_id not in self._nodes: + return [] + + parent = self.get_parent(branch_root_id) + removed: list[MessageNode] = [] + for node_id in self.get_descendants(branch_root_id): + node = self._nodes.get(node_id) + if not node: + continue + removed.append(node) + del self._nodes[node_id] + self._status_to_node.pop(node.status_message_id, None) + + if parent and branch_root_id in parent.children_ids: + parent.children_ids = [ + child_id + for child_id in parent.children_ids + if child_id != branch_root_id + ] + + logger.debug("Removed branch {} ({} nodes)", branch_root_id, len(removed)) + return removed + + def snapshot(self) -> TreeSnapshot: + return TreeSnapshot( + root_id=self.root_id, + nodes={ + node_id: node_to_snapshot(node) for node_id, node in self._nodes.items() + }, + ) + + @classmethod + def from_snapshot(cls, snapshot: TreeSnapshot) -> MessageTreeGraph: + root_data = snapshot.nodes[snapshot.root_id] + root_node = node_from_snapshot(root_data) + graph = cls(root_node) + for node_id, node_data in snapshot.nodes.items(): + if node_id == snapshot.root_id: + continue + if not isinstance(node_data, dict): + continue + node = node_from_snapshot(node_data) + graph._nodes[node.node_id] = node + graph._status_to_node[node.status_message_id] = node.node_id + return graph diff --git a/messaging/trees/manager.py b/messaging/trees/manager.py index dda4dcec..05722f64 100644 --- a/messaging/trees/manager.py +++ b/messaging/trees/manager.py @@ -6,9 +6,11 @@ from collections.abc import Awaitable, Callable from loguru import logger from ..models import IncomingMessage -from .data import MessageNode, MessageState, MessageTree +from .node import MessageNode, MessageState from .processor import TreeQueueProcessor from .repository import TreeRepository +from .runtime import MessageTree +from .snapshot import ConversationSnapshot class TreeQueueManager: @@ -408,23 +410,23 @@ class TreeQueueManager: """Get all message IDs for a given platform/chat.""" return self._repository.get_message_ids_for_chat(platform, chat_id) - def to_dict(self) -> dict: - """Serialize all trees.""" - return self._repository.to_dict() + def snapshot(self) -> ConversationSnapshot: + """Serialize all trees into a typed conversation snapshot.""" + return self._repository.snapshot() @classmethod - def from_dict( + def from_snapshot( cls, - data: dict, + snapshot: ConversationSnapshot, queue_update_callback: Callable[[MessageTree], Awaitable[None]] | None = None, node_started_callback: Callable[[MessageTree, str], Awaitable[None]] | None = None, ) -> TreeQueueManager: - """Deserialize from dictionary.""" + """Restore a manager from a typed conversation snapshot.""" return cls( queue_update_callback=queue_update_callback, node_started_callback=node_started_callback, - _repository=TreeRepository.from_dict(data), + _repository=TreeRepository.from_snapshot(snapshot), ) diff --git a/messaging/trees/node.py b/messaging/trees/node.py new file mode 100644 index 00000000..7c9e7993 --- /dev/null +++ b/messaging/trees/node.py @@ -0,0 +1,57 @@ +"""Message tree node model.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from datetime import UTC, datetime +from enum import Enum +from typing import Any + +from ..models import IncomingMessage + + +class MessageState(Enum): + """State of a message node in the tree.""" + + PENDING = "pending" + IN_PROGRESS = "in_progress" + COMPLETED = "completed" + ERROR = "error" + + +@dataclass +class MessageNode: + """A single user prompt/status node in a messaging conversation tree.""" + + node_id: str + incoming: IncomingMessage + status_message_id: str + state: MessageState = MessageState.PENDING + parent_id: str | None = None + session_id: str | None = None + children_ids: list[str] = field(default_factory=list) + created_at: datetime = field(default_factory=lambda: datetime.now(UTC)) + completed_at: datetime | None = None + error_message: str | None = None + context: Any = None + + def set_context(self, context: Any) -> None: + self.context = context + + def update_state( + self, + state: MessageState, + *, + session_id: str | None = None, + error_message: str | None = None, + ) -> None: + self.state = state + if session_id: + self.session_id = session_id + if error_message: + self.error_message = error_message + if state in (MessageState.COMPLETED, MessageState.ERROR): + self.completed_at = datetime.now(UTC) + + def mark_error(self, error_message: str) -> None: + self.update_state(MessageState.ERROR, error_message=error_message) diff --git a/messaging/trees/processor.py b/messaging/trees/processor.py index b6315957..10c28714 100644 --- a/messaging/trees/processor.py +++ b/messaging/trees/processor.py @@ -9,7 +9,8 @@ from config.settings import get_settings from core.anthropic import get_user_facing_error_message from ..safe_diagnostics import format_exception_for_log -from .data import MessageNode, MessageState, MessageTree +from .node import MessageNode, MessageState +from .runtime import MessageTree class TreeQueueProcessor: diff --git a/messaging/trees/queue.py b/messaging/trees/queue.py new file mode 100644 index 00000000..8ef3a07e --- /dev/null +++ b/messaging/trees/queue.py @@ -0,0 +1,46 @@ +"""FIFO queue state for one messaging conversation tree.""" + +from __future__ import annotations + +import asyncio +from collections import deque + + +class MessageNodeQueue: + """Queue with snapshot/remove helpers, backed by a deque and a set index.""" + + def __init__(self, items: list[str] | None = None) -> None: + self._deque: deque[str] = deque() + self._set: set[str] = set() + for item in items or []: + self.put_nowait(item) + + def put_nowait(self, item: str) -> None: + self._deque.append(item) + self._set.add(item) + + def get_nowait(self) -> str: + if not self._deque: + raise asyncio.QueueEmpty() + item = self._deque.popleft() + self._set.discard(item) + return item + + def qsize(self) -> int: + return len(self._deque) + + def snapshot(self) -> list[str]: + return list(self._deque) + + def remove_if_present(self, item: str) -> bool: + if item not in self._set: + return False + self._set.discard(item) + self._deque = deque(x for x in self._deque if x != item) + return True + + def drain(self) -> list[str]: + items = list(self._deque) + self._deque.clear() + self._set.clear() + return items diff --git a/messaging/trees/repository.py b/messaging/trees/repository.py index b4267c78..5c9dfa7f 100644 --- a/messaging/trees/repository.py +++ b/messaging/trees/repository.py @@ -2,7 +2,9 @@ from loguru import logger -from .data import MessageNode, MessageState, MessageTree +from .node import MessageNode, MessageState +from .runtime import MessageTree +from .snapshot import ConversationSnapshot class TreeRepository: @@ -157,20 +159,19 @@ class TreeRepository: msg_ids.add(str(node.status_message_id)) return msg_ids - def to_dict(self) -> dict: - """Serialize all trees.""" - return { - "trees": {rid: tree.to_dict() for rid, tree in self._trees.items()}, - "node_to_tree": self._node_to_tree.copy(), - } + def snapshot(self) -> ConversationSnapshot: + """Serialize all trees into a typed conversation snapshot.""" + return ConversationSnapshot( + trees={root_id: tree.snapshot() for root_id, tree in self._trees.items()} + ) @classmethod - def from_dict(cls, data: dict) -> TreeRepository: - """Deserialize from dictionary.""" + def from_snapshot(cls, snapshot: ConversationSnapshot) -> TreeRepository: + """Restore repository state from a typed conversation snapshot.""" repo = cls() - for root_id, tree_data in data.get("trees", {}).items(): - repo._trees[root_id] = MessageTree.from_dict(tree_data) - repo._node_to_tree = data.get("node_to_tree", {}) + for root_id, tree_snapshot in snapshot.trees.items(): + repo._trees[root_id] = MessageTree.from_snapshot(tree_snapshot) + repo._node_to_tree = snapshot.derive_node_to_tree() return repo diff --git a/messaging/trees/runtime.py b/messaging/trees/runtime.py new file mode 100644 index 00000000..48b83bfb --- /dev/null +++ b/messaging/trees/runtime.py @@ -0,0 +1,181 @@ +"""Runtime state for one messaging conversation tree.""" + +from __future__ import annotations + +import asyncio +from contextlib import asynccontextmanager + +from loguru import logger + +from ..models import IncomingMessage +from .graph import MessageTreeGraph +from .node import MessageNode, MessageState +from .queue import MessageNodeQueue +from .snapshot import TreeSnapshot + + +class MessageTree: + """Runtime aggregate for one ordered messaging conversation tree.""" + + def __init__( + self, + root_node: MessageNode, + *, + queue: MessageNodeQueue | None = None, + graph: MessageTreeGraph | None = None, + ) -> None: + self._graph = graph or MessageTreeGraph(root_node) + self._queue = queue or MessageNodeQueue() + self._lock = asyncio.Lock() + self._is_processing = False + self._current_node_id: str | None = None + self._current_task: asyncio.Task | None = None + logger.debug("Created MessageTree with root {}", self.root_id) + + @property + def root_id(self) -> str: + return self._graph.root_id + + @property + def is_processing(self) -> bool: + return self._is_processing + + @property + def current_node_id(self) -> str | None: + return self._current_node_id + + async def add_node( + self, + node_id: str, + incoming: IncomingMessage, + status_message_id: str, + parent_id: str, + ) -> MessageNode: + async with self._lock: + return self._graph.add_node( + node_id=node_id, + incoming=incoming, + status_message_id=status_message_id, + parent_id=parent_id, + ) + + def get_node(self, node_id: str) -> MessageNode | None: + return self._graph.get_node(node_id) + + def get_root(self) -> MessageNode: + return self._graph.get_root() + + def get_children(self, node_id: str) -> list[MessageNode]: + return self._graph.get_children(node_id) + + def get_parent(self, node_id: str) -> MessageNode | None: + return self._graph.get_parent(node_id) + + def get_parent_session_id(self, node_id: str) -> str | None: + return self._graph.get_parent_session_id(node_id) + + async def update_state( + self, + node_id: str, + state: MessageState, + session_id: str | None = None, + error_message: str | None = None, + ) -> None: + async with self._lock: + self._graph.update_node_state( + node_id, + state, + session_id=session_id, + error_message=error_message, + ) + + async def enqueue(self, node_id: str) -> int: + async with self._lock: + self._queue.put_nowait(node_id) + position = self._queue.qsize() + logger.debug("Enqueued node {}, position {}", node_id, position) + return position + + async def dequeue(self) -> str | None: + try: + return self._queue.get_nowait() + except asyncio.QueueEmpty: + return None + + async def get_queue_snapshot(self) -> list[str]: + async with self._lock: + return self._queue.snapshot() + + def get_queue_size(self) -> int: + return self._queue.qsize() + + def remove_from_queue(self, node_id: str) -> bool: + return self._queue.remove_if_present(node_id) + + @asynccontextmanager + async def with_lock(self): + async with self._lock: + yield + + def set_processing_state(self, node_id: str | None, is_processing: bool) -> None: + self._is_processing = is_processing + self._current_node_id = node_id if is_processing else None + + def clear_current_node(self) -> None: + self._current_node_id = None + + def is_current_node(self, node_id: str) -> bool: + return self._current_node_id == node_id + + def put_queue_unlocked(self, node_id: str) -> None: + self._queue.put_nowait(node_id) + + def set_current_task(self, task: asyncio.Task | None) -> None: + self._current_task = task + + def cancel_current_task(self) -> bool: + if self._current_task and not self._current_task.done(): + self._current_task.cancel() + return True + return False + + def set_node_error_sync(self, node: MessageNode, error_message: str) -> None: + node.mark_error(error_message) + + def drain_queue_and_mark_cancelled( + self, error_message: str = "Cancelled by user" + ) -> list[MessageNode]: + nodes: list[MessageNode] = [] + for node_id in self._queue.drain(): + node = self._graph.get_node(node_id) + if node: + self.set_node_error_sync(node, error_message) + nodes.append(node) + return nodes + + def reset_processing_state(self) -> None: + self._is_processing = False + self._current_node_id = None + + def all_nodes(self) -> list[MessageNode]: + return self._graph.all_nodes() + + def has_node(self, node_id: str) -> bool: + return self._graph.has_node(node_id) + + def find_node_by_status_message(self, status_msg_id: str) -> MessageNode | None: + return self._graph.find_node_by_status_message(status_msg_id) + + def get_descendants(self, node_id: str) -> list[str]: + return self._graph.get_descendants(node_id) + + def remove_branch(self, branch_root_id: str) -> list[MessageNode]: + return self._graph.remove_branch(branch_root_id) + + def snapshot(self) -> TreeSnapshot: + return self._graph.snapshot() + + @classmethod + def from_snapshot(cls, snapshot: TreeSnapshot) -> MessageTree: + graph = MessageTreeGraph.from_snapshot(snapshot) + return cls(graph.get_root(), graph=graph) diff --git a/messaging/trees/snapshot.py b/messaging/trees/snapshot.py new file mode 100644 index 00000000..41a02f91 --- /dev/null +++ b/messaging/trees/snapshot.py @@ -0,0 +1,150 @@ +"""Serializable messaging conversation snapshots.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from datetime import datetime +from typing import Any + +from ..models import IncomingMessage +from .node import MessageNode, MessageState + + +@dataclass(frozen=True) +class TreeSnapshot: + """Persisted representation of one conversation tree.""" + + root_id: str + nodes: dict[str, dict[str, Any]] + + def to_json(self) -> dict[str, Any]: + return {"root_id": self.root_id, "nodes": dict(self.nodes)} + + @classmethod + def from_json(cls, data: Any) -> TreeSnapshot | None: + if not isinstance(data, dict): + return None + root_id = data.get("root_id") + nodes = data.get("nodes") + if root_id is None or not isinstance(nodes, dict): + return None + return cls(root_id=str(root_id), nodes=dict(nodes)) + + def lookup_ids(self) -> set[str]: + lookup_ids: set[str] = set() + for node_key, node_data in self.nodes.items(): + lookup_ids.add(str(node_key)) + if not isinstance(node_data, dict): + continue + node_id = node_data.get("node_id") + if node_id is not None: + lookup_ids.add(str(node_id)) + status_message_id = node_data.get("status_message_id") + if status_message_id is not None: + lookup_ids.add(str(status_message_id)) + return lookup_ids + + +@dataclass(frozen=True) +class ConversationSnapshot: + """Persisted conversation trees plus derived lookup helpers.""" + + trees: dict[str, TreeSnapshot] = field(default_factory=dict) + + @property + def is_empty(self) -> bool: + return not self.trees + + def to_json(self) -> dict[str, Any]: + return { + "trees": { + root_id: tree_snapshot.to_json() + for root_id, tree_snapshot in self.trees.items() + } + } + + @classmethod + def from_json(cls, data: Any) -> ConversationSnapshot: + if not isinstance(data, dict): + return cls() + raw_trees = data.get("trees", {}) + if not isinstance(raw_trees, dict): + return cls() + + trees: dict[str, TreeSnapshot] = {} + for raw_root_id, raw_tree in raw_trees.items(): + snapshot = TreeSnapshot.from_json(raw_tree) + if snapshot is None: + continue + root_id = str(raw_root_id) if raw_root_id is not None else snapshot.root_id + trees[root_id] = snapshot + return cls(trees=trees) + + def derive_node_to_tree(self) -> dict[str, str]: + mapping: dict[str, str] = {} + for root_id, tree_snapshot in self.trees.items(): + for lookup_id in tree_snapshot.lookup_ids(): + mapping[lookup_id] = root_id + return mapping + + def with_tree(self, tree_snapshot: TreeSnapshot) -> ConversationSnapshot: + trees = dict(self.trees) + trees[tree_snapshot.root_id] = tree_snapshot + return ConversationSnapshot(trees=trees) + + def without_tree(self, root_id: str) -> ConversationSnapshot: + trees = dict(self.trees) + trees.pop(root_id, None) + return ConversationSnapshot(trees=trees) + + +def node_to_snapshot(node: MessageNode) -> dict[str, Any]: + return { + "node_id": node.node_id, + "incoming": { + "text": node.incoming.text, + "chat_id": node.incoming.chat_id, + "user_id": node.incoming.user_id, + "message_id": node.incoming.message_id, + "platform": node.incoming.platform, + "reply_to_message_id": node.incoming.reply_to_message_id, + "message_thread_id": node.incoming.message_thread_id, + "username": node.incoming.username, + }, + "status_message_id": node.status_message_id, + "state": node.state.value, + "parent_id": node.parent_id, + "session_id": node.session_id, + "children_ids": list(node.children_ids), + "created_at": node.created_at.isoformat(), + "completed_at": node.completed_at.isoformat() if node.completed_at else None, + "error_message": node.error_message, + } + + +def node_from_snapshot(data: dict[str, Any]) -> MessageNode: + incoming_data = data["incoming"] + incoming = IncomingMessage( + text=incoming_data["text"], + chat_id=incoming_data["chat_id"], + user_id=incoming_data["user_id"], + message_id=incoming_data["message_id"], + platform=incoming_data["platform"], + reply_to_message_id=incoming_data.get("reply_to_message_id"), + message_thread_id=incoming_data.get("message_thread_id"), + username=incoming_data.get("username"), + ) + return MessageNode( + node_id=data["node_id"], + incoming=incoming, + status_message_id=data["status_message_id"], + state=MessageState(data["state"]), + parent_id=data.get("parent_id"), + session_id=data.get("session_id"), + children_ids=list(data.get("children_ids", [])), + created_at=datetime.fromisoformat(data["created_at"]), + completed_at=datetime.fromisoformat(data["completed_at"]) + if data.get("completed_at") + else None, + error_message=data.get("error_message"), + ) diff --git a/messaging/turn_intake.py b/messaging/turn_intake.py index cc8a7c52..01802dfd 100644 --- a/messaging/turn_intake.py +++ b/messaging/turn_intake.py @@ -129,8 +129,6 @@ class MessagingTurnIntake: status_message_id=status_msg_id, ) tree_queue.register_node(status_msg_id, tree.root_id) - self.session_store.register_node(status_msg_id, tree.root_id) - self.session_store.register_node(node_id, tree.root_id) elif status_msg_id: tree = await tree_queue.create_tree( node_id=node_id, @@ -138,11 +136,9 @@ class MessagingTurnIntake: status_message_id=status_msg_id, ) tree_queue.register_node(status_msg_id, tree.root_id) - self.session_store.register_node(node_id, tree.root_id) - self.session_store.register_node(status_msg_id, tree.root_id) if tree: - self.session_store.save_tree(tree.root_id, tree.to_dict()) + self.session_store.save_tree_snapshot(tree.snapshot()) was_queued = await tree_queue.enqueue( node_id=node_id, diff --git a/messaging/workflow.py b/messaging/workflow.py index ae56bef5..f94a4c59 100644 --- a/messaging/workflow.py +++ b/messaging/workflow.py @@ -199,8 +199,8 @@ class MessagingWorkflow: tree = self.tree_queue.get_tree_for_node(node.node_id) if tree: trees_to_save[tree.root_id] = tree - for root_id, tree in trees_to_save.items(): - self.session_store.save_tree(root_id, tree.to_dict()) + for tree in trees_to_save.values(): + self.session_store.save_tree_snapshot(tree.snapshot()) __all__ = ["MessagingWorkflow"] diff --git a/pyproject.toml b/pyproject.toml index 95a0ba11..a033e445 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "free-claude-code" -version = "2.3.21" +version = "2.3.22" description = "Middleware between Claude Code CLI (Anthropic API) and NVIDIA NIM" readme = "README.md" requires-python = ">=3.14.0" diff --git a/smoke/lib/e2e.py b/smoke/lib/e2e.py index 0a9538b0..d12bde4a 100644 --- a/smoke/lib/e2e.py +++ b/smoke/lib/e2e.py @@ -552,12 +552,9 @@ class FakePlatformDriver: raise AssertionError("fake platform did not become idle") def _all_tree_nodes_terminal(self) -> bool: - data = self.workflow.tree_queue.to_dict() - for tree in data.get("trees", {}).values(): - nodes = tree.get("nodes", {}) if isinstance(tree, dict) else {} - for node in nodes.values(): - if not isinstance(node, dict): - continue + snapshot = self.workflow.tree_queue.snapshot() + for tree in snapshot.trees.values(): + for node in tree.nodes.values(): if node.get("state") in {"pending", "in_progress"}: return False return True diff --git a/smoke/product/test_messaging_product_live.py b/smoke/product/test_messaging_product_live.py index b14cc6bf..79711a65 100644 --- a/smoke/product/test_messaging_product_live.py +++ b/smoke/product/test_messaging_product_live.py @@ -77,7 +77,7 @@ async def test_messaging_commands_stop_clear_stats_e2e( assert "Stats" in sent_text assert "Stopped" in sent_text assert driver.platform.deletes - assert driver.session_store.get_all_trees() == {} + assert driver.session_store.load_conversation_snapshot().trees == {} @pytest.mark.asyncio @@ -108,14 +108,13 @@ async def test_restart_restore_and_session_persistence_e2e(tmp_path) -> None: session_file = tmp_path / "telegram-sessions.json" payload = json.loads(session_file.read_text(encoding="utf-8")) - assert payload["trees"] - assert payload["node_to_tree"] + assert payload["conversation"]["trees"] assert payload["message_log"] restored = FakePlatformDriver("telegram", tmp_path) - saved = restored.session_store.get_all_trees() - assert saved - assert root.message_id in restored.session_store.get_node_mapping() + saved = restored.session_store.load_conversation_snapshot() + assert saved.trees + assert root.message_id in saved.derive_node_to_tree() @pytest.mark.asyncio diff --git a/tests/api/test_app_lifespan_and_errors.py b/tests/api/test_app_lifespan_and_errors.py index 83fa8a00..bf239c84 100644 --- a/tests/api/test_app_lifespan_and_errors.py +++ b/tests/api/test_app_lifespan_and_errors.py @@ -9,6 +9,7 @@ from fastapi import FastAPI from fastapi.testclient import TestClient from config.settings import Settings +from messaging.trees import ConversationSnapshot, TreeSnapshot from providers.exceptions import ServiceUnavailableError from providers.runtime import ProviderRuntime @@ -318,17 +319,18 @@ def test_app_lifespan_sets_state_and_cleans_up(tmp_path, messaging_enabled): fake_platform.stop = AsyncMock() fake_components = _fake_messaging_components(fake_platform) + snapshot = ( + ConversationSnapshot(trees={"t": TreeSnapshot(root_id="t", nodes={})}) + if messaging_enabled + else ConversationSnapshot() + ) session_store = MagicMock() - session_store.get_all_trees.return_value = [{"t": 1}] if messaging_enabled else [] - session_store.get_node_mapping.return_value = {"n": "t"} - session_store.sync_from_tree_data = MagicMock() + session_store.load_conversation_snapshot.return_value = snapshot + session_store.save_conversation_snapshot = MagicMock() fake_queue = MagicMock() fake_queue.cleanup_stale_nodes.return_value = 1 - fake_queue.to_dict.return_value = { - "trees": [{"t": 1}], - "node_to_tree": {"n": "t"}, - } + fake_queue.snapshot.return_value = snapshot cli_manager = MagicMock() cli_manager.stop_all = AsyncMock() @@ -346,7 +348,7 @@ def test_app_lifespan_sets_state_and_cleans_up(tmp_path, messaging_enabled): patch("messaging.session.SessionStore", return_value=session_store), patch("cli.managed.ManagedClaudeSessionManager", return_value=cli_manager), patch( - "messaging.trees.TreeQueueManager.from_dict", + "messaging.trees.TreeQueueManager.from_snapshot", return_value=fake_queue, ), TestClient(app), @@ -360,10 +362,7 @@ def test_app_lifespan_sets_state_and_cleans_up(tmp_path, messaging_enabled): fake_platform.stop.assert_awaited_once() cli_manager.stop_all.assert_awaited_once() assert getattr(app.state, "messaging_workflow", None) is not None - session_store.sync_from_tree_data.assert_called_once_with( - [{"t": 1}], - {"n": "t"}, - ) + session_store.save_conversation_snapshot.assert_called_once_with(snapshot) else: fake_platform.start.assert_not_awaited() fake_platform.stop.assert_not_awaited() @@ -398,9 +397,7 @@ def test_app_lifespan_cleanup_continues_if_platform_stop_raises(tmp_path): fake_components = _fake_messaging_components(fake_platform) session_store = MagicMock() - session_store.get_all_trees.return_value = [] - session_store.get_node_mapping.return_value = {} - session_store.sync_from_tree_data = MagicMock() + session_store.load_conversation_snapshot.return_value = ConversationSnapshot() cli_manager = MagicMock() cli_manager.stop_all = AsyncMock() @@ -576,9 +573,7 @@ def test_app_lifespan_platform_start_exception_cleanup_still_runs(tmp_path): fake_components = _fake_messaging_components(fake_platform) session_store = MagicMock() - session_store.get_all_trees.return_value = [] - session_store.get_node_mapping.return_value = {} - session_store.sync_from_tree_data = MagicMock() + session_store.load_conversation_snapshot.return_value = ConversationSnapshot() cli_manager = MagicMock() cli_manager.stop_all = AsyncMock() @@ -627,9 +622,7 @@ def test_app_lifespan_flush_pending_save_exception_warning_only(tmp_path): fake_components = _fake_messaging_components(fake_platform) session_store = MagicMock() - session_store.get_all_trees.return_value = [] - session_store.get_node_mapping.return_value = {} - session_store.sync_from_tree_data = MagicMock() + session_store.load_conversation_snapshot.return_value = ConversationSnapshot() session_store.flush_pending_save = MagicMock(side_effect=OSError("disk full")) cli_manager = MagicMock() diff --git a/tests/contracts/test_import_boundaries.py b/tests/contracts/test_import_boundaries.py index 4014fdc9..21a248db 100644 --- a/tests/contracts/test_import_boundaries.py +++ b/tests/contracts/test_import_boundaries.py @@ -429,6 +429,50 @@ def test_messaging_transcript_uses_package_owners() -> None: assert "RenderCtx" in init_text +def test_messaging_conversation_state_uses_package_owners() -> None: + repo_root = Path(__file__).resolve().parents[2] + messaging_root = repo_root / "messaging" + trees_root = messaging_root / "trees" + session_root = messaging_root / "session" + + assert not (messaging_root / "session.py").exists() + assert not (trees_root / "data.py").exists() + for filename in { + "__init__.py", + "graph.py", + "manager.py", + "node.py", + "processor.py", + "queue.py", + "repository.py", + "runtime.py", + "snapshot.py", + }: + assert (trees_root / filename).exists() + for filename in { + "__init__.py", + "message_log.py", + "persistence.py", + "store.py", + }: + assert (session_root / filename).exists() + + offenders = _imports_matching( + [messaging_root, repo_root / "api", repo_root / "tests"], + forbidden_prefixes=("messaging.trees.data",), + ) + assert offenders == [] + + runtime_text = (repo_root / "api" / "runtime.py").read_text(encoding="utf-8") + for removed_api in { + "get_all_trees", + "get_node_mapping", + "sync_from_tree_data", + "TreeQueueManager.from_dict", + }: + assert removed_api not in runtime_text + + def test_messaging_workflow_uses_split_runtime_owners() -> None: repo_root = Path(__file__).resolve().parents[2] messaging_root = repo_root / "messaging" diff --git a/tests/messaging/test_handler.py b/tests/messaging/test_handler.py index fd0593ac..774fe28e 100644 --- a/tests/messaging/test_handler.py +++ b/tests/messaging/test_handler.py @@ -3,8 +3,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from messaging.models import IncomingMessage -from messaging.trees import MessageState -from messaging.trees.data import MessageNode, MessageTree +from messaging.trees import MessageNode, MessageState, MessageTree from messaging.workflow import MessagingWorkflow @@ -212,14 +211,14 @@ async def test_handle_message_new_conversation( ): mock_tree = MagicMock() mock_tree.root_id = "root_1" - mock_tree.to_dict.return_value = {"data": "tree"} + mock_tree.snapshot.return_value = {"data": "tree"} mock_create.return_value = mock_tree await handler.handle_message(incoming) mock_create.assert_called_once() mock_enqueue.assert_called_once() - mock_session_store.save_tree.assert_called_once_with("root_1", {"data": "tree"}) + mock_session_store.save_tree_snapshot.assert_called_once_with({"data": "tree"}) @pytest.mark.asyncio @@ -234,7 +233,7 @@ async def test_handle_message_queued(handler, mock_platform, incoming_message_fa ): mock_tree = MagicMock() mock_tree.root_id = "root_1" - mock_tree.to_dict.return_value = {} + mock_tree.snapshot.return_value = {} mock_create.return_value = mock_tree await handler.handle_message(incoming) @@ -407,7 +406,7 @@ async def test_node_runner_process_node_success_flow( mock_tree = MagicMock() mock_tree.update_state = AsyncMock() mock_tree.root_id = "root_1" - mock_tree.to_dict.return_value = {} + mock_tree.snapshot.return_value = {} with patch.object( handler.tree_queue, "get_tree_for_node", MagicMock(return_value=mock_tree) @@ -458,7 +457,7 @@ async def test_node_runner_process_node_reply_uses_parent_session_for_manager_an mock_tree = MagicMock() mock_tree.update_state = AsyncMock() mock_tree.root_id = "root_msg" - mock_tree.to_dict.return_value = {} + mock_tree.snapshot.return_value = {} mock_tree.get_parent_session_id = MagicMock(return_value=parent_claude_session) with patch.object( @@ -497,7 +496,7 @@ async def test_node_runner_process_node_error_flow( mock_tree = MagicMock() mock_tree.root_id = "root_1" - mock_tree.to_dict.return_value = {"data": "tree"} + mock_tree.snapshot.return_value = {"data": "tree"} mock_tree.update_state = AsyncMock() with ( @@ -513,8 +512,8 @@ async def test_node_runner_process_node_error_flow( handler.tree_queue.mark_node_error.assert_called_once_with( node_id, "CLI crashed", propagate_to_children=True ) - handler.session_store.save_tree.assert_called_once_with( - "root_1", {"data": "tree"} + handler.session_store.save_tree_snapshot.assert_called_once_with( + {"data": "tree"} ) last_call = mock_platform.queue_edit_message.call_args_list[-1] @@ -659,7 +658,7 @@ async def test_handle_message_clear_command_reply_clears_branch( assert set(deleted_ids) == {"102", "103", "150"} assert "100" not in deleted_ids assert "101" not in deleted_ids - mock_session_store.remove_node_mappings.assert_called() + mock_session_store.save_tree_snapshot.assert_called() assert handler.tree_queue.get_tree_for_node("102") is None assert handler.tree_queue.get_tree_for_node("100") is not None @@ -711,7 +710,7 @@ async def test_handle_message_clear_command_reply_to_root_clears_tree( await handler.handle_message(incoming) assert set(deleted_ids) == {"100", "101", "150"} - mock_session_store.remove_tree.assert_called_once_with("100") + mock_session_store.remove_tree_snapshot.assert_called_once_with("100") assert handler.tree_queue.get_tree_count() == 0 diff --git a/tests/messaging/test_handler_context_isolation.py b/tests/messaging/test_handler_context_isolation.py index 4bf31f5f..4933bf8e 100644 --- a/tests/messaging/test_handler_context_isolation.py +++ b/tests/messaging/test_handler_context_isolation.py @@ -2,7 +2,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest -from messaging.trees.data import MessageState +from messaging.trees import MessageState from messaging.workflow import MessagingWorkflow diff --git a/tests/messaging/test_handler_integration.py b/tests/messaging/test_handler_integration.py index 1727b462..daa87a1c 100644 --- a/tests/messaging/test_handler_integration.py +++ b/tests/messaging/test_handler_integration.py @@ -3,7 +3,7 @@ from unittest.mock import MagicMock import pytest -from messaging.trees.data import MessageState +from messaging.trees import MessageState from messaging.workflow import MessagingWorkflow diff --git a/tests/messaging/test_handler_markdown_and_status_edges.py b/tests/messaging/test_handler_markdown_and_status_edges.py index 6594ef11..b59515c9 100644 --- a/tests/messaging/test_handler_markdown_and_status_edges.py +++ b/tests/messaging/test_handler_markdown_and_status_edges.py @@ -6,7 +6,7 @@ import pytest from messaging.models import IncomingMessage from messaging.node_event_pipeline import process_parsed_cli_event from messaging.rendering.telegram_markdown import render_markdown_to_mdv2 -from messaging.trees.data import MessageNode, MessageState +from messaging.trees import MessageNode, MessageState from messaging.workflow import MessagingWorkflow @@ -153,7 +153,7 @@ async def test_node_runner_process_node_session_limit_marks_error_and_updates_ui fake_tree = MagicMock() fake_tree.root_id = "root" - fake_tree.to_dict = MagicMock(return_value={"root": "error"}) + fake_tree.snapshot = MagicMock(return_value={"root": "error"}) fake_tree.update_state = AsyncMock() with patch.object( handler.tree_queue, "get_tree_for_node", MagicMock(return_value=fake_tree) @@ -170,7 +170,7 @@ async def test_node_runner_process_node_session_limit_marks_error_and_updates_ui await handler.node_runner.process_node("n1", node) assert platform.queue_edit_message.await_count >= 1 fake_tree.update_state.assert_awaited() - session_store.save_tree.assert_called_once_with("root", {"root": "error"}) + session_store.save_tree_snapshot.assert_called_once_with({"root": "error"}) @pytest.mark.asyncio @@ -199,7 +199,7 @@ async def test_node_runner_cancellation_marks_error_and_saves_tree(): fake_tree = MagicMock() fake_tree.root_id = "root" - fake_tree.to_dict = MagicMock(return_value={"root": "cancelled"}) + fake_tree.snapshot = MagicMock(return_value={"root": "cancelled"}) fake_tree.update_state = AsyncMock() with patch.object( handler.tree_queue, "get_tree_for_node", MagicMock(return_value=fake_tree) @@ -218,7 +218,7 @@ async def test_node_runner_cancellation_marks_error_and_saves_tree(): fake_tree.update_state.assert_any_await( "n1", MessageState.ERROR, error_message="Cancelled by user" ) - session_store.save_tree.assert_called_once_with("root", {"root": "cancelled"}) + session_store.save_tree_snapshot.assert_called_once_with({"root": "cancelled"}) @pytest.mark.asyncio @@ -247,7 +247,7 @@ async def test_stop_all_tasks_saves_tree_for_cancelled_nodes(): tree = MagicMock() tree.root_id = "root" - tree.to_dict = MagicMock(return_value={"root": "ok"}) + tree.snapshot = MagicMock(return_value={"root": "ok"}) with ( patch.object(handler.tree_queue, "cancel_all", AsyncMock(return_value=[node])), patch.object( @@ -257,7 +257,7 @@ async def test_stop_all_tasks_saves_tree_for_cancelled_nodes(): count = await handler.stop_all_tasks() assert count == 1 cli_manager.stop_all.assert_awaited_once() - session_store.save_tree.assert_called_once_with("root", {"root": "ok"}) + session_store.save_tree_snapshot.assert_called_once_with({"root": "ok"}) @pytest.mark.asyncio @@ -277,7 +277,9 @@ async def test_handle_message_reply_with_tree_but_no_parent_treated_as_new(): mock_queue.get_tree_for_node.return_value = object() mock_queue.resolve_parent_node_id.return_value = None mock_queue.create_tree = AsyncMock( - return_value=MagicMock(root_id="root", to_dict=MagicMock(return_value={"t": 1})) + return_value=MagicMock( + root_id="root", snapshot=MagicMock(return_value={"t": 1}) + ) ) mock_queue.register_node = MagicMock() mock_queue.enqueue = AsyncMock(return_value=False) @@ -367,7 +369,9 @@ async def test_handle_message_incoming_text_none_safe(): mock_queue.get_tree_for_node.return_value = None mock_queue.resolve_parent_node_id.return_value = None mock_queue.create_tree = AsyncMock( - return_value=MagicMock(root_id="root", to_dict=MagicMock(return_value={"t": 1})) + return_value=MagicMock( + root_id="root", snapshot=MagicMock(return_value={"t": 1}) + ) ) mock_queue.register_node = MagicMock() mock_queue.enqueue = AsyncMock(return_value=True) diff --git a/tests/messaging/test_messaging.py b/tests/messaging/test_messaging.py index c9d49440..03deab44 100644 --- a/tests/messaging/test_messaging.py +++ b/tests/messaging/test_messaging.py @@ -79,38 +79,35 @@ class TestSessionStore: from messaging.session import SessionStore store = SessionStore(storage_path=str(tmp_path / "sessions.json")) - assert store._trees == {} + assert store.load_conversation_snapshot().is_empty # --- Tree Tests --- def test_save_and_get_tree(self, tmp_path): """Test saving and retrieving trees.""" from messaging.session import SessionStore + from messaging.trees import TreeSnapshot store = SessionStore(storage_path=str(tmp_path / "sessions.json")) tree_data = { - "root": "r1", - "nodes": {"r1": {"content": "root"}, "n1": {"content": "child"}}, + "root_id": "r1", + "nodes": { + "r1": {"node_id": "r1", "status_message_id": "s1"}, + "n1": {"node_id": "n1", "status_message_id": "s2"}, + }, } - store.save_tree("r1", tree_data) + snapshot = TreeSnapshot.from_json(tree_data) + assert snapshot is not None + store.save_tree_snapshot(snapshot) - loaded = store.get_tree("r1") - assert loaded == tree_data + loaded = store.get_tree_snapshot("r1") + assert loaded == snapshot - # Verify node mapping - node_map = store.get_node_mapping() + node_map = store.load_conversation_snapshot().derive_node_to_tree() assert node_map["r1"] == "r1" assert node_map["n1"] == "r1" - def test_register_node(self, tmp_path): - """Test manual node registration.""" - from messaging.session import SessionStore - - store = SessionStore(storage_path=str(tmp_path / "sessions.json")) - store.register_node("n_manual", "r_manual") - assert store.get_node_mapping()["n_manual"] == "r_manual" - # --- Persistence & Edge Cases --- def test_load_existing_file_with_trees(self, tmp_path): @@ -129,7 +126,7 @@ class TestSessionStore: json.dump(data, f) store = SessionStore(storage_path=str(p)) - assert store.get_tree("r1") is not None + assert store.get_tree_snapshot("r1") is not None def test_load_corrupt_file(self, tmp_path): """Test loading corrupt/invalid json file.""" @@ -141,21 +138,24 @@ class TestSessionStore: # Should log error and start empty, avoiding crash store = SessionStore(storage_path=str(p)) - assert store._trees == {} + assert store.load_conversation_snapshot().is_empty def test_save_error_handling(self, tmp_path): """Test error during save.""" from messaging.session import SessionStore + from messaging.trees import TreeSnapshot store = SessionStore(storage_path=str(tmp_path / "sessions.json")) - store.save_tree("r1", {"root_id": "r1", "nodes": {"r1": {}}}) + snapshot = TreeSnapshot(root_id="r1", nodes={"r1": {}}) + store.save_tree_snapshot(snapshot) - # Mock open to raise exception - with patch("builtins.open", side_effect=OSError("Disk full")): - store.save_tree("r2", {"root_id": "r2", "nodes": {"r2": {}}}) + with patch( + "messaging.session.persistence.os.replace", side_effect=OSError("Disk full") + ): + store.flush_pending_save() - # Should log error but not crash. Tree should be in memory. - assert "r2" in store._trees + assert store.dirty is True + assert store.get_tree_snapshot("r1") is not None class TestTreeQueueManager: diff --git a/tests/messaging/test_restart_reply_restore.py b/tests/messaging/test_restart_reply_restore.py index 82342783..c761e033 100644 --- a/tests/messaging/test_restart_reply_restore.py +++ b/tests/messaging/test_restart_reply_restore.py @@ -28,19 +28,15 @@ async def test_reply_to_old_status_message_after_restore_routes_to_parent( "A", a_incoming, status_message_id="status_A" ) handler1.tree_queue.register_node("status_A", tree.root_id) - store.register_node("status_A", tree.root_id) - store.save_tree(tree.root_id, tree.to_dict()) + store.save_tree_snapshot(tree.snapshot()) store.flush_pending_save() # "Restart": new store instance loads from disk, and we restore TreeQueueManager. store2 = SessionStore(storage_path=str(store_path)) handler2 = MessagingWorkflow(mock_platform, mock_cli_manager, store2) handler2.replace_tree_queue( - TreeQueueManager.from_dict( - { - "trees": store2.get_all_trees(), - "node_to_tree": store2.get_node_mapping(), - }, + TreeQueueManager.from_snapshot( + store2.load_conversation_snapshot(), queue_update_callback=handler2.update_queue_positions, node_started_callback=handler2.mark_node_processing, ) @@ -86,17 +82,14 @@ async def test_save_tree_persists_status_message_mapping_without_manual_register tree = await handler1.tree_queue.create_tree( "A", a_incoming, status_message_id="status_A" ) - store.save_tree(tree.root_id, tree.to_dict()) + store.save_tree_snapshot(tree.snapshot()) store.flush_pending_save() store2 = SessionStore(storage_path=str(store_path)) handler2 = MessagingWorkflow(mock_platform, mock_cli_manager, store2) handler2.replace_tree_queue( - TreeQueueManager.from_dict( - { - "trees": store2.get_all_trees(), - "node_to_tree": store2.get_node_mapping(), - }, + TreeQueueManager.from_snapshot( + store2.load_conversation_snapshot(), queue_update_callback=handler2.update_queue_positions, node_started_callback=handler2.mark_node_processing, ) @@ -154,7 +147,7 @@ async def test_reply_clear_purges_removed_status_mapping_from_persisted_store( "root", "child", child_incoming, status_message_id="child_status" ) handler.tree_queue.register_node("child_status", tree.root_id) - store.save_tree(tree.root_id, tree.to_dict()) + store.save_tree_snapshot(tree.snapshot()) clear_reply = IncomingMessage( text="/clear", @@ -169,11 +162,8 @@ async def test_reply_clear_purges_removed_status_mapping_from_persisted_store( store.flush_pending_save() restored_store = SessionStore(storage_path=str(store_path)) - restored_tree_queue = TreeQueueManager.from_dict( - { - "trees": restored_store.get_all_trees(), - "node_to_tree": restored_store.get_node_mapping(), - } + restored_tree_queue = TreeQueueManager.from_snapshot( + restored_store.load_conversation_snapshot() ) assert restored_tree_queue.get_tree_for_node("root") is not None diff --git a/tests/messaging/test_session_store_edge_cases.py b/tests/messaging/test_session_store_edge_cases.py index ca3d8b83..ddb962ca 100644 --- a/tests/messaging/test_session_store_edge_cases.py +++ b/tests/messaging/test_session_store_edge_cases.py @@ -1,11 +1,16 @@ -"""Edge case tests for messaging/session.py SessionStore.""" +"""Edge case tests for the messaging session store.""" import json +from collections.abc import Callable +from typing import Any, ClassVar from unittest.mock import patch import pytest +import messaging.session.persistence as persistence_module from messaging.session import SessionStore +from messaging.session.persistence import DebouncedJsonPersistence +from messaging.trees import TreeSnapshot @pytest.fixture @@ -22,6 +27,52 @@ def _tree_node(node_id: str, status_message_id: str) -> dict: } +class FakeTimer: + instances: ClassVar[list[FakeTimer]] = [] + + def __init__( + self, + interval: float, + function: Callable[..., None], + args: tuple[Any, ...] | None = None, + kwargs: dict[str, Any] | None = None, + ) -> None: + self.interval = interval + self.function = function + self.args = args or () + self.kwargs = kwargs or {} + self.daemon = False + self.canceled = False + self.started = False + self.instances.append(self) + + def cancel(self) -> None: + self.canceled = True + + def start(self) -> None: + self.started = True + + def fire(self, *, force: bool = False) -> None: + if self.canceled and not force: + return + self.function(*self.args, **self.kwargs) + + +class RecordingPersistence(DebouncedJsonPersistence): + def __init__( + self, + storage_path: str, + *, + snapshot: Callable[[], dict[str, Any]], + on_dirty: Callable[[bool], None], + ) -> None: + self.writes: list[dict[str, Any]] = [] + super().__init__(storage_path, snapshot=snapshot, on_dirty=on_dirty) + + def write_data(self, data: dict[str, Any]) -> None: + self.writes.append(data) + + class TestSessionStoreLoadEdgeCases: """Tests for loading corrupted/malformed data.""" @@ -32,7 +83,7 @@ class TestSessionStoreLoadEdgeCases: f.write("{invalid json") store = SessionStore(storage_path=path) - assert len(store._trees) == 0 + assert store.load_conversation_snapshot().is_empty def test_load_truncated_json(self, tmp_path): """Truncated JSON file is handled gracefully.""" @@ -41,7 +92,7 @@ class TestSessionStoreLoadEdgeCases: f.write('{"sessions": {"s1": {"session_id": "s1"') store = SessionStore(storage_path=path) - assert len(store._trees) == 0 + assert store.load_conversation_snapshot().is_empty def test_load_empty_file(self, tmp_path): """Empty file is handled gracefully.""" @@ -50,13 +101,13 @@ class TestSessionStoreLoadEdgeCases: f.write("") store = SessionStore(storage_path=path) - assert len(store._trees) == 0 + assert store.load_conversation_snapshot().is_empty def test_load_nonexistent_file(self, tmp_path): """Non-existent file starts with empty state.""" path = str(tmp_path / "nonexistent.json") store = SessionStore(storage_path=path) - assert len(store._trees) == 0 + assert store.load_conversation_snapshot().is_empty def test_load_legacy_sessions_ignored(self, tmp_path): """Legacy sessions in file are ignored; trees and message_log load.""" @@ -81,87 +132,112 @@ class TestSessionStoreLoadEdgeCases: json.dump(data, f) store = SessionStore(storage_path=path) - assert store.get_tree("r1") is not None + assert store.get_tree_snapshot("r1") is not None class TestSessionStoreSaveEdgeCases: """Tests for save failure handling.""" def test_save_io_error_handled(self, tmp_store): - """Write failure during atomic replace is surfaced to callers.""" - tmp_store.save_tree("r1", {"root_id": "r1", "nodes": {"r1": {}}}) - with ( - patch("messaging.session.os.replace", side_effect=OSError("disk full")), - pytest.raises(OSError), + """Write failure marks pending state dirty without crashing callers.""" + tmp_store.save_tree_snapshot(TreeSnapshot(root_id="r1", nodes={"r1": {}})) + with patch( + "messaging.session.persistence.os.replace", + side_effect=OSError("disk full"), ): - tmp_store._write_data(tmp_store._snapshot()) + tmp_store.flush_pending_save() + assert tmp_store.dirty is True + + def test_stale_timer_callback_cannot_clear_newer_timer(self, tmp_path, monkeypatch): + """An already-running old timer cannot consume the newest save.""" + FakeTimer.instances = [] + monkeypatch.setattr(persistence_module.threading, "Timer", FakeTimer) + + dirty_states: list[bool] = [] + snapshot_count = 0 + + def snapshot() -> dict[str, Any]: + nonlocal snapshot_count + snapshot_count += 1 + return {"snapshot": snapshot_count} + + persistence = RecordingPersistence( + str(tmp_path / "sessions.json"), + snapshot=snapshot, + on_dirty=dirty_states.append, + ) + + persistence.schedule_save() + first_timer = FakeTimer.instances[0] + persistence.schedule_save() + second_timer = FakeTimer.instances[1] + + first_timer.fire(force=True) + assert persistence.writes == [] + assert dirty_states[-1] is True + assert second_timer.canceled is False + + second_timer.fire() + assert persistence.writes == [{"snapshot": 1}] + assert dirty_states[-1] is False class TestSessionStoreTreeMappings: def test_save_tree_rebuilds_lookup_ids_for_that_root(self, tmp_path): path = str(tmp_path / "sessions.json") store = SessionStore(storage_path=path) - store.register_node("unrelated_status", "other_root") - store.save_tree( - "root", - { - "root_id": "root", - "nodes": { + store.save_tree_snapshot( + TreeSnapshot( + root_id="root", + nodes={ "root": _tree_node("root", "root_status"), "child": _tree_node("child", "child_status"), }, - }, + ) ) - mapping = store.get_node_mapping() + mapping = store.load_conversation_snapshot().derive_node_to_tree() assert mapping["root"] == "root" assert mapping["root_status"] == "root" assert mapping["child"] == "root" assert mapping["child_status"] == "root" - store.save_tree( - "root", - { - "root_id": "root", - "nodes": { + store.save_tree_snapshot( + TreeSnapshot( + root_id="root", + nodes={ "root": _tree_node("root", "root_status"), }, - }, + ) ) - mapping = store.get_node_mapping() + mapping = store.load_conversation_snapshot().derive_node_to_tree() assert mapping["root"] == "root" assert mapping["root_status"] == "root" assert "child" not in mapping assert "child_status" not in mapping - assert mapping["unrelated_status"] == "other_root" def test_remove_tree_removes_all_lookup_ids_for_that_root(self, tmp_path): path = str(tmp_path / "sessions.json") store = SessionStore(storage_path=path) - store.register_node("old_status", "root") - store.register_node("unrelated_status", "other_root") - store.save_tree( - "root", - { - "root_id": "root", - "nodes": { + store.save_tree_snapshot( + TreeSnapshot( + root_id="root", + nodes={ "root": _tree_node("root", "root_status"), "child": _tree_node("child", "child_status"), }, - }, + ) ) - store.remove_tree("root") + store.remove_tree_snapshot("root") - mapping = store.get_node_mapping() + mapping = store.load_conversation_snapshot().derive_node_to_tree() assert "root" not in mapping assert "root_status" not in mapping assert "child" not in mapping assert "child_status" not in mapping - assert "old_status" not in mapping - assert mapping["unrelated_status"] == "other_root" class TestSessionStoreAtomicWrites: @@ -170,23 +246,24 @@ class TestSessionStoreAtomicWrites: def test_failed_replace_keeps_prior_bytes_and_marks_dirty(self, tmp_path): path = str(tmp_path / "sessions.json") store = SessionStore(storage_path=path) - store.save_tree("r1", {"root_id": "r1", "nodes": {"r1": {}}}) + store.save_tree_snapshot(TreeSnapshot(root_id="r1", nodes={"r1": {}})) store.flush_pending_save() with open(path, encoding="utf-8") as f: disk_after_first = f.read() - store.save_tree("r2", {"root_id": "r2", "nodes": {"r2": {}}}) + store.save_tree_snapshot(TreeSnapshot(root_id="r2", nodes={"r2": {}})) with patch( - "messaging.session.os.replace", side_effect=OSError("replace failed") + "messaging.session.persistence.os.replace", + side_effect=OSError("replace failed"), ): store.flush_pending_save() with open(path, encoding="utf-8") as f: disk_after_failed = f.read() assert disk_after_failed == disk_after_first - assert store._dirty is True - assert store.get_tree("r2") is not None + assert store.dirty is True + assert store.get_tree_snapshot("r2") is not None class TestSessionStoreClearAll: @@ -194,11 +271,10 @@ class TestSessionStoreClearAll: path = str(tmp_path / "sessions.json") store = SessionStore(storage_path=path) - store.save_tree( - "root1", - { - "root_id": "root1", - "nodes": { + store.save_tree_snapshot( + TreeSnapshot( + root_id="root1", + nodes={ "root1": { "node_id": "root1", "incoming": { @@ -208,7 +284,6 @@ class TestSessionStoreClearAll: "message_id": "m1", "platform": "telegram", "reply_to_message_id": None, - "username": None, }, "status_message_id": "status1", "state": "pending", @@ -220,22 +295,20 @@ class TestSessionStoreClearAll: "error_message": None, } }, - }, + ) ) store.clear_all() - assert store.get_all_trees() == {} - assert store.get_node_mapping() == {} + assert store.load_conversation_snapshot().is_empty with open(path, encoding="utf-8") as f: data = json.load(f) - assert data["trees"] == {} - assert data["node_to_tree"] == {} + assert data["conversation"]["trees"] == {} assert data["message_log"] == {} store2 = SessionStore(storage_path=path) - assert len(store2._trees) == 0 + assert store2.load_conversation_snapshot().is_empty def test_message_log_persists_and_dedups(self, tmp_path): path = str(tmp_path / "sessions.json") diff --git a/tests/messaging/test_tree_concurrency.py b/tests/messaging/test_tree_concurrency.py index 281843cc..209e7431 100644 --- a/tests/messaging/test_tree_concurrency.py +++ b/tests/messaging/test_tree_concurrency.py @@ -5,8 +5,8 @@ import asyncio import pytest from messaging.models import IncomingMessage -from messaging.trees import TreeQueueManager -from messaging.trees.data import MessageNode, MessageState, MessageTree +from messaging.trees import MessageNode, MessageState, MessageTree, TreeQueueManager +from messaging.trees.snapshot import node_from_snapshot, node_to_snapshot def _make_incoming(text: str = "hello", msg_id: str = "m1") -> IncomingMessage: @@ -254,8 +254,8 @@ class TestMessageTreeSerialization: await tree.add_node("c2", _make_incoming(msg_id="c2"), "s2", "root") await tree.update_state("root", MessageState.COMPLETED, session_id="sess1") - data = tree.to_dict() - restored = MessageTree.from_dict(data) + snapshot = tree.snapshot() + restored = MessageTree.from_snapshot(snapshot) assert restored.root_id == "root" assert len(restored.all_nodes()) == 3 @@ -276,8 +276,8 @@ class TestMessageTreeSerialization: session_id="sess_test", error_message="test error", ) - data = node.to_dict() - restored = MessageNode.from_dict(data) + data = node_to_snapshot(node) + restored = node_from_snapshot(data) assert restored.node_id == "n1" assert restored.state == MessageState.COMPLETED @@ -507,8 +507,8 @@ class TestTreeQueueManagerConcurrency: await mgr.create_tree("root", _make_incoming(msg_id="root"), "s_root") _, _ = await mgr.add_to_tree("root", "c1", _make_incoming(msg_id="c1"), "s1") - data = mgr.to_dict() - restored = TreeQueueManager.from_dict(data) + snapshot = mgr.snapshot() + restored = TreeQueueManager.from_snapshot(snapshot) assert restored.get_tree_count() == 1 assert restored.get_node("c1") is not None diff --git a/tests/messaging/test_tree_processor.py b/tests/messaging/test_tree_processor.py index b3a00128..0f597579 100644 --- a/tests/messaging/test_tree_processor.py +++ b/tests/messaging/test_tree_processor.py @@ -5,7 +5,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest from messaging.models import IncomingMessage -from messaging.trees.data import MessageNode, MessageState, MessageTree +from messaging.trees import MessageNode, MessageState, MessageTree from messaging.trees.processor import TreeQueueProcessor @@ -142,7 +142,7 @@ async def test_process_next_queue_empty(tree_processor, sample_tree): @pytest.mark.asyncio async def test_process_next_with_item(tree_processor, sample_tree): processor = AsyncMock() - await sample_tree._queue.put("next_node") + sample_tree.put_queue_unlocked("next_node") node = MagicMock(spec=MessageNode) sample_tree.get_node = MagicMock(return_value=node) @@ -182,8 +182,8 @@ async def test_process_next_skips_stale_id_and_runs_next_valid_node(sample_tree) "valid_status", sample_tree.root_id, ) - await sample_tree._queue.put("missing_node") - await sample_tree._queue.put("valid_node") + sample_tree.put_queue_unlocked("missing_node") + sample_tree.put_queue_unlocked("valid_node") async def node_processor(node_id, node): await release.wait() @@ -209,8 +209,8 @@ async def test_process_next_drains_all_stale_ids_without_wedging(sample_tree): node_started_callback=node_started, ) sample_tree._is_processing = True - await sample_tree._queue.put("missing_one") - await sample_tree._queue.put("missing_two") + sample_tree.put_queue_unlocked("missing_one") + sample_tree.put_queue_unlocked("missing_two") await processor._process_next(sample_tree, AsyncMock()) @@ -226,7 +226,7 @@ async def test_process_next_triggers_queue_update(sample_tree): callback = AsyncMock() processor = TreeQueueProcessor(queue_update_callback=callback) - await sample_tree._queue.put("next_node") + sample_tree.put_queue_unlocked("next_node") sample_tree.get_node = MagicMock(return_value=None) await processor._process_next(sample_tree, AsyncMock()) @@ -250,7 +250,7 @@ async def test_process_next_triggers_node_started(sample_tree): await sample_tree.add_node( "next_node", incoming, "next_status", sample_tree.root_id ) - await sample_tree._queue.put("next_node") + sample_tree.put_queue_unlocked("next_node") await processor._process_next(sample_tree, AsyncMock()) diff --git a/tests/messaging/test_tree_queue.py b/tests/messaging/test_tree_queue.py index cbece5d6..5f852d82 100644 --- a/tests/messaging/test_tree_queue.py +++ b/tests/messaging/test_tree_queue.py @@ -11,7 +11,10 @@ from messaging.trees import ( MessageState, MessageTree, TreeQueueManager, + TreeSnapshot, ) +from messaging.trees.graph import MessageTreeGraph +from messaging.trees.snapshot import node_from_snapshot, node_to_snapshot class TestMessageState: @@ -49,7 +52,7 @@ class TestMessageNode: assert node.children_ids == [] assert node.session_id is None - def test_node_to_dict(self): + def test_node_to_snapshot(self): """Test serializing a node.""" incoming = IncomingMessage( text="Test", @@ -66,12 +69,12 @@ class TestMessageNode: session_id="sess_123", ) - data = node.to_dict() + data = node_to_snapshot(node) assert data["node_id"] == "3" assert data["state"] == "completed" assert data["session_id"] == "sess_123" - def test_node_from_dict(self): + def test_node_from_snapshot(self): """Test deserializing a node.""" data = { "node_id": "n1", @@ -90,7 +93,7 @@ class TestMessageNode: "created_at": "2025-01-01T00:00:00", } - node = MessageNode.from_dict(data) + node = node_from_snapshot(data) assert node.node_id == "n1" assert node.state == MessageState.IN_PROGRESS assert node.parent_id == "parent_1" @@ -258,8 +261,8 @@ class TestMessageTree: snapshot = await tree.get_queue_snapshot() assert snapshot == ["child_1", "child_2"] - def test_tree_serialization(self): - """Test tree to_dict and from_dict.""" + def test_tree_snapshot_round_trip(self): + """Test tree snapshot round-trip.""" incoming = IncomingMessage( text="Test", chat_id="1", @@ -276,14 +279,44 @@ class TestMessageTree: ) tree = MessageTree(root) - data = tree.to_dict() - restored = MessageTree.from_dict(data) + snapshot = tree.snapshot() + restored = MessageTree.from_snapshot(snapshot) assert restored.root_id == "m1" node = restored.get_node("m1") assert node is not None assert node.session_id == "sess_1" + def test_tree_from_snapshot_uses_one_graph_construction(self, monkeypatch): + """Restore should not build a temporary graph and replace it.""" + incoming = IncomingMessage( + text="Test", + chat_id="1", + user_id="1", + message_id="m1", + platform="test", + ) + root = MessageNode( + node_id="m1", + incoming=incoming, + status_message_id="s1", + ) + snapshot = MessageTree(root).snapshot() + original_init = MessageTreeGraph.__init__ + init_calls = 0 + + def counting_init(self: MessageTreeGraph, root_node: MessageNode) -> None: + nonlocal init_calls + init_calls += 1 + original_init(self, root_node) + + monkeypatch.setattr(MessageTreeGraph, "__init__", counting_init) + + restored = MessageTree.from_snapshot(snapshot) + + assert restored.root_id == "m1" + assert init_calls == 1 + @pytest.mark.asyncio async def test_get_descendants(self): """Test get_descendants returns node and all descendants.""" @@ -777,11 +810,13 @@ class TestSessionStoreTrees: }, } - store.save_tree("root_1", tree_data) + snapshot = TreeSnapshot.from_json(tree_data) + assert snapshot is not None + store.save_tree_snapshot(snapshot) - retrieved = store.get_tree("root_1") + retrieved = store.get_tree_snapshot("root_1") assert retrieved is not None - assert retrieved["root_id"] == "root_1" + assert retrieved.root_id == "root_1" def test_get_tree_by_root_id(self, tmp_path): """Test getting tree by root ID and node mapping.""" @@ -797,19 +832,13 @@ class TestSessionStoreTrees: }, } - store.save_tree("root", tree_data) + snapshot = TreeSnapshot.from_json(tree_data) + assert snapshot is not None + store.save_tree_snapshot(snapshot) - retrieved = store.get_tree("root") + retrieved = store.get_tree_snapshot("root") assert retrieved is not None - assert retrieved["root_id"] == "root" - assert store.get_node_mapping()["child"] == "root" - - def test_register_node(self, tmp_path): - """Test registering a node to a tree.""" - from messaging.session import SessionStore - - store = SessionStore(storage_path=str(tmp_path / "sessions.json")) - - store.register_node("new_node", "root_tree") - - assert store.get_node_mapping()["new_node"] == "root_tree" + assert retrieved.root_id == "root" + assert ( + store.load_conversation_snapshot().derive_node_to_tree()["child"] == "root" + ) diff --git a/tests/messaging/test_tree_repository.py b/tests/messaging/test_tree_repository.py index 2860ae6e..4ed66392 100644 --- a/tests/messaging/test_tree_repository.py +++ b/tests/messaging/test_tree_repository.py @@ -3,7 +3,7 @@ from unittest.mock import MagicMock import pytest from messaging.models import IncomingMessage -from messaging.trees.data import MessageNode, MessageState, MessageTree +from messaging.trees import MessageNode, MessageTree from messaging.trees.repository import TreeRepository @@ -93,10 +93,10 @@ def test_resolve_parent_node_id(repository, sample_tree): assert repository.resolve_parent_node_id("unknown") is None -def test_get_pending_children(repository, sample_tree): +@pytest.mark.asyncio +async def test_get_pending_children(repository, sample_tree): repository.add_tree("root_id", sample_tree) - # Create a child node child_incoming = IncomingMessage( text="child", chat_id="c1", @@ -104,16 +104,7 @@ def test_get_pending_children(repository, sample_tree): message_id="child_id", platform="telegram", ) - child_node = MessageNode( - node_id="child_id", - incoming=child_incoming, - status_message_id="s2", - parent_id="root_id", - state=MessageState.PENDING, - ) - - sample_tree._nodes["child_id"] = child_node - sample_tree.get_node("root_id").children_ids.append("child_id") + await sample_tree.add_node("child_id", child_incoming, "s2", "root_id") repository.register_node("child_id", "root_id") pending = repository.get_pending_children("root_id") @@ -121,16 +112,14 @@ def test_get_pending_children(repository, sample_tree): assert pending[0].node_id == "child_id" -def test_to_from_dict(repository, sample_tree): +def test_snapshot_round_trip(repository, sample_tree): repository.add_tree("root_id", sample_tree) - data = repository.to_dict() + snapshot = repository.snapshot() - assert "trees" in data - assert "root_id" in data["trees"] - assert "node_to_tree" in data - assert data["node_to_tree"]["root_id"] == "root_id" + assert "root_id" in snapshot.trees + assert snapshot.derive_node_to_tree()["root_id"] == "root_id" - new_repo = TreeRepository.from_dict(data) + new_repo = TreeRepository.from_snapshot(snapshot) tree = new_repo.get_tree("root_id") assert tree is not None assert tree.root_id == "root_id" diff --git a/uv.lock b/uv.lock index 82d62331..03bf8455 100644 --- a/uv.lock +++ b/uv.lock @@ -561,7 +561,7 @@ wheels = [ [[package]] name = "free-claude-code" -version = "2.3.21" +version = "2.3.22" source = { editable = "." } dependencies = [ { name = "aiohttp" },