wehub-resource-sync
195 lines
6.6 KiB
Python
195 lines
6.6 KiB
Python
"""Persistence and restored-manager contracts for messaging trees."""
|
|
|
|
import asyncio
|
|
|
|
import pytest
|
|
|
|
from free_claude_code.messaging.models import IncomingMessage, MessageScope
|
|
from free_claude_code.messaging.trees import (
|
|
ConversationSnapshot,
|
|
MessageState,
|
|
NodeClaim,
|
|
TreeIdentity,
|
|
TreeQueueManager,
|
|
TreeSnapshot,
|
|
)
|
|
from free_claude_code.messaging.trees.node import MessageNode
|
|
from free_claude_code.messaging.trees.snapshot import node_to_snapshot
|
|
|
|
TELEGRAM_CHAT = MessageScope(platform="telegram", chat_id="chat")
|
|
ROOT_IDENTITY = TreeIdentity(scope=TELEGRAM_CHAT, root_id="root")
|
|
|
|
|
|
def _incoming(node_id: str, *, reply_to: str | None = None) -> IncomingMessage:
|
|
return IncomingMessage(
|
|
text=f"prompt {node_id}",
|
|
chat_id="chat",
|
|
user_id="user",
|
|
message_id=node_id,
|
|
platform="telegram",
|
|
reply_to_message_id=reply_to,
|
|
)
|
|
|
|
|
|
async def _wait_for_no_tasks(manager: TreeQueueManager) -> None:
|
|
loop = asyncio.get_running_loop()
|
|
for _ in range(20):
|
|
if manager.task_count() == 0:
|
|
return
|
|
checkpoint = asyncio.Event()
|
|
loop.call_soon(checkpoint.set)
|
|
await checkpoint.wait()
|
|
assert manager.task_count() == 0
|
|
|
|
|
|
def _interrupted_conversation() -> ConversationSnapshot:
|
|
root = MessageNode(
|
|
node_id="root",
|
|
scope=TELEGRAM_CHAT,
|
|
prompt="prompt root",
|
|
status_message_id="status-root",
|
|
state=MessageState.COMPLETED,
|
|
session_id="session-root",
|
|
children_ids=["pending", "failed"],
|
|
)
|
|
pending = MessageNode(
|
|
node_id="pending",
|
|
scope=TELEGRAM_CHAT,
|
|
prompt="prompt pending",
|
|
status_message_id="status-pending",
|
|
parent_id="root",
|
|
state=MessageState.PENDING,
|
|
children_ids=["running"],
|
|
)
|
|
running = MessageNode(
|
|
node_id="running",
|
|
scope=TELEGRAM_CHAT,
|
|
prompt="prompt running",
|
|
status_message_id="status-running",
|
|
parent_id="pending",
|
|
state=MessageState.IN_PROGRESS,
|
|
)
|
|
failed = MessageNode(
|
|
node_id="failed",
|
|
scope=TELEGRAM_CHAT,
|
|
prompt="prompt failed",
|
|
status_message_id="status-failed",
|
|
parent_id="root",
|
|
state=MessageState.ERROR,
|
|
)
|
|
snapshot = TreeSnapshot(
|
|
scope=TELEGRAM_CHAT,
|
|
root_id="root",
|
|
nodes={
|
|
node.node_id: node_to_snapshot(node)
|
|
for node in (root, pending, running, failed)
|
|
},
|
|
)
|
|
return ConversationSnapshot(trees={snapshot.identity: snapshot})
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_restore_reconciles_interrupted_nodes_before_manager_exposure() -> None:
|
|
processed: list[str] = []
|
|
|
|
async def process(claim: NodeClaim) -> None:
|
|
processed.append(claim.node.node_id)
|
|
|
|
manager = TreeQueueManager.from_snapshot(_interrupted_conversation(), process)
|
|
|
|
assert len(manager.restored_stale_targets) == 2
|
|
assert manager.restored_snapshot is not None
|
|
assert manager.restored_snapshot == await manager.snapshot()
|
|
assert manager.task_count() == 0
|
|
assert processed == []
|
|
|
|
root = await manager.get_node(TELEGRAM_CHAT, "root")
|
|
pending = await manager.get_node(TELEGRAM_CHAT, "pending")
|
|
running = await manager.get_node(TELEGRAM_CHAT, "status-running")
|
|
failed = await manager.get_node(TELEGRAM_CHAT, "failed")
|
|
assert root is not None and root.state is MessageState.COMPLETED
|
|
assert root.session_id == "session-root"
|
|
assert pending is not None and pending.state is MessageState.ERROR
|
|
assert running is not None and running.state is MessageState.ERROR
|
|
assert failed is not None and failed.state is MessageState.ERROR
|
|
assert await manager.resolve_node_id(TELEGRAM_CHAT, "status-pending") == "pending"
|
|
assert {
|
|
(target.scope, target.node_id) for target in manager.restored_stale_targets
|
|
} == {
|
|
(TELEGRAM_CHAT, "pending"),
|
|
(TELEGRAM_CHAT, "running"),
|
|
}
|
|
|
|
normalized = await manager.snapshot()
|
|
normalized_tree = normalized.get_tree(ROOT_IDENTITY)
|
|
assert normalized_tree is not None
|
|
assert normalized_tree.nodes["pending"]["state"] == "error"
|
|
assert normalized_tree.nodes["running"]["state"] == "error"
|
|
assert normalized_tree.nodes["failed"]["state"] == "error"
|
|
assert all("error_message" not in node for node in normalized_tree.nodes.values())
|
|
|
|
|
|
@pytest.mark.parametrize("corruption", ["cross_scope", "duplicate_reference"])
|
|
def test_restore_rejects_tree_that_violates_scoped_reference_invariants(
|
|
corruption: str,
|
|
) -> None:
|
|
snapshot = _interrupted_conversation()
|
|
tree = snapshot.get_tree(ROOT_IDENTITY)
|
|
assert tree is not None
|
|
if corruption == "cross_scope":
|
|
tree.nodes["pending"]["incoming"] = {
|
|
"platform": "telegram",
|
|
"chat_id": "other-chat",
|
|
}
|
|
else:
|
|
tree.nodes["failed"]["status_message_id"] = "status-pending"
|
|
|
|
manager = TreeQueueManager.from_snapshot(snapshot, lambda _claim: asyncio.sleep(0))
|
|
|
|
assert manager.get_tree_count() == 0
|
|
assert manager.restored_snapshot == ConversationSnapshot()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_claim_state_and_session_updates_produce_detached_snapshots() -> None:
|
|
release = asyncio.Event()
|
|
started = asyncio.Event()
|
|
|
|
async def process(_claim: NodeClaim) -> None:
|
|
started.set()
|
|
await release.wait()
|
|
|
|
manager = TreeQueueManager(process)
|
|
decision = await manager.admit(_incoming("root"), "status-root")
|
|
assert decision.claim is not None
|
|
await started.wait()
|
|
|
|
session_snapshot = await manager.record_session(decision.claim, "session-root")
|
|
completed_snapshot = await manager.complete_claim(
|
|
decision.claim,
|
|
"session-root",
|
|
)
|
|
|
|
assert session_snapshot is not None
|
|
assert session_snapshot.nodes["root"]["state"] == "in_progress"
|
|
assert session_snapshot.nodes["root"]["session_id"] == "session-root"
|
|
assert completed_snapshot is not None
|
|
assert completed_snapshot.nodes["root"]["state"] == "completed"
|
|
assert completed_snapshot.nodes["root"]["session_id"] == "session-root"
|
|
|
|
assert decision.snapshot is not None
|
|
decision.snapshot.nodes["root"]["state"] = "mutated-copy"
|
|
view = await manager.get_node(TELEGRAM_CHAT, "status-root")
|
|
assert view is not None and view.state is MessageState.COMPLETED
|
|
assert view.session_id == "session-root"
|
|
assert await manager.get_message_ids_for_chat("telegram", "chat") == {
|
|
"root",
|
|
"status-root",
|
|
}
|
|
|
|
release.set()
|
|
await _wait_for_no_tasks(manager)
|
|
assert await manager.snapshot() == ConversationSnapshot(
|
|
trees={ROOT_IDENTITY: completed_snapshot}
|
|
)
|