vrr/free-claude-code
2
0
Fork 0
mirror of https://github.com/Alishahryar1/free-claude-code.git synced 2026-07-09 16:00:45 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

Refactor messaging conversation state (#931)

Browse source
This commit is contained in:
Ali Khokhar 2026-06-28 01:32:34 -07:00 • committed by GitHub
parent 3afdd98bd1
commit 002012dfcd
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
39 changed files with 1389 additions and 1087 deletions
Download patch file Download diff file Expand all files Collapse all files

19
ARCHITECTURE.md
View file

@ -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

21
api/runtime.py
View file

@ -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:

7
messaging/commands.py
View file

@ -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}")

4
messaging/node_event_pipeline.py
View file

@ -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)

2
messaging/node_runner.py
View file

@ -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,

334
messaging/session.py
View file

@ -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()

5
messaging/session/__init__.py Normal file
View file

@ -0,0 +1,5 @@
"""Public messaging session persistence API."""
from .store import SessionStore
__all__ = ["SessionStore"]

116
messaging/session/message_log.py Normal file
View file

@ -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}"

131
messaging/session/persistence.py Normal file
View file

@ -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

148
messaging/session/store.py Normal file
View file

@ -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)

6
messaging/trees/__init__.py
View file

@ -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",
]

482
messaging/trees/data.py
View file

@ -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

156
messaging/trees/graph.py Normal file
View file

@ -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

18
messaging/trees/manager.py
View file

@ -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),
)

57
messaging/trees/node.py Normal file
View file

@ -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)

3
messaging/trees/processor.py
View file

@ -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:

46
messaging/trees/queue.py Normal file
View file

@ -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

25
messaging/trees/repository.py
View file

@ -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

181
messaging/trees/runtime.py Normal file
View file

@ -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)

150
messaging/trees/snapshot.py Normal file
View file

@ -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"),
)

6
messaging/turn_intake.py
View file

@ -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,

4
messaging/workflow.py
View file

@ -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"]

2
pyproject.toml
View file

@ -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"

9
smoke/lib/e2e.py
View file

@ -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

11
smoke/product/test_messaging_product_live.py
View file

@ -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

35
tests/api/test_app_lifespan_and_errors.py
View file

@ -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()

44
tests/contracts/test_import_boundaries.py
View file

@ -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"

23
tests/messaging/test_handler.py
View file

@ -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

2
tests/messaging/test_handler_context_isolation.py
View file

@ -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

2
tests/messaging/test_handler_integration.py
View file

@ -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

22
tests/messaging/test_handler_markdown_and_status_edges.py
View file

@ -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)

48
tests/messaging/test_messaging.py
View file

@ -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:

28
tests/messaging/test_restart_reply_restore.py
View file

@ -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

187
tests/messaging/test_session_store_edge_cases.py
View file

@ -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")

16
tests/messaging/test_tree_concurrency.py
View file

@ -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

16
tests/messaging/test_tree_processor.py
View file

@ -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())

79
tests/messaging/test_tree_queue.py
View file

@ -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"
)

29
tests/messaging/test_tree_repository.py
View file

@ -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"

2
uv.lock generated
View file

@ -561,7 +561,7 @@ wheels = [
[[package]]
name = "free-claude-code"
version = "2.3.21"
version = "2.3.22"
source = { editable = "." }
dependencies = [
{ name = "aiohttp" },