Spaces:
Sleeping
Sleeping
| """Atomic runtime aggregate for one messaging conversation tree.""" | |
| from __future__ import annotations | |
| import asyncio | |
| from dataclasses import dataclass | |
| from uuid import uuid4 | |
| from loguru import logger | |
| from ..models import MessageScope | |
| from .graph import MessageTreeGraph | |
| from .identity import TreeIdentity | |
| from .node import MessageNode, MessageReferenceKind, MessageState | |
| from .queue import MessageNodeQueue | |
| from .snapshot import TreeSnapshot | |
| from .transitions import ( | |
| AdmissionRejection, | |
| CompletionResult, | |
| FailureResult, | |
| MessageSubtreeRemoval, | |
| NodeClaim, | |
| NodeUiTarget, | |
| NodeView, | |
| QueueDecision, | |
| QueueEntry, | |
| ReplyTarget, | |
| TreeCancellation, | |
| ) | |
| class _ActiveClaim: | |
| """Runtime execution identity kept separate from the node's UI state.""" | |
| claim: NodeClaim | |
| cancellation_requested: bool = False | |
| class MessageTree: | |
| """Own graph, queue, claim identity, and every concurrency invariant.""" | |
| def __init__( | |
| self, | |
| root_node: MessageNode, | |
| *, | |
| graph: MessageTreeGraph | None = None, | |
| ) -> None: | |
| self._graph = graph or MessageTreeGraph(root_node) | |
| self._queue = MessageNodeQueue() | |
| self._lock = asyncio.Lock() | |
| self._active: _ActiveClaim | None = None | |
| self._restored_snapshot: TreeSnapshot | None = None | |
| self._restored_stale_targets: tuple[NodeUiTarget, ...] = () | |
| logger.debug("Created MessageTree with root {}", self.root_id) | |
| def root_id(self) -> str: | |
| return self._graph.root_id | |
| def identity(self) -> TreeIdentity: | |
| return self._graph.identity | |
| def restored_snapshot(self) -> TreeSnapshot | None: | |
| """Normalized startup snapshot captured before the tree is published.""" | |
| return self._restored_snapshot | |
| def restored_stale_targets(self) -> tuple[NodeUiTarget, ...]: | |
| """UI targets normalized from runnable to interrupted on restore.""" | |
| return self._restored_stale_targets | |
| def _ui_target(self, node: MessageNode) -> NodeUiTarget: | |
| if node.status_message_id is None: | |
| raise ValueError("Runnable node has no status message") | |
| return NodeUiTarget( | |
| scope=node.scope, | |
| node_id=node.node_id, | |
| status_message_id=node.status_message_id, | |
| ) | |
| def _queue_entries(self) -> tuple[QueueEntry, ...]: | |
| entries: list[QueueEntry] = [] | |
| for node_id in self._queue.items(): | |
| node = self._graph.get_node(node_id) | |
| if node is None or node.state is not MessageState.PENDING: | |
| continue | |
| entries.append( | |
| QueueEntry(node=self._ui_target(node), position=len(entries) + 1) | |
| ) | |
| return tuple(entries) | |
| def _claim(self, node: MessageNode) -> NodeClaim: | |
| node.update_state(MessageState.IN_PROGRESS) | |
| claim = NodeClaim( | |
| identity=self.identity, | |
| claim_id=uuid4().hex, | |
| node=self._ui_target(node), | |
| prompt=node.prompt, | |
| parent_session_id=self._graph.get_parent_session_id(node.node_id), | |
| ) | |
| self._active = _ActiveClaim(claim=claim) | |
| return claim | |
| def _enqueue_or_claim(self, node_id: str) -> QueueDecision: | |
| node = self._graph.get_node(node_id) | |
| if node is None or node.state is not MessageState.PENDING: | |
| return QueueDecision( | |
| claim=None, | |
| position=None, | |
| snapshot=None, | |
| rejection=AdmissionRejection.DUPLICATE, | |
| ) | |
| if self._active is None: | |
| claim = self._claim(node) | |
| return QueueDecision( | |
| claim=claim, | |
| position=None, | |
| snapshot=self._graph.snapshot(), | |
| ) | |
| if not self._queue.put(node_id): | |
| return QueueDecision( | |
| claim=None, | |
| position=None, | |
| snapshot=None, | |
| rejection=AdmissionRejection.DUPLICATE, | |
| ) | |
| position = self._queue.qsize() | |
| logger.info("Queued node {}, position {}", node_id, position) | |
| return QueueDecision( | |
| claim=None, | |
| position=position, | |
| snapshot=self._graph.snapshot(), | |
| ) | |
| async def enqueue_or_claim(self, node_id: str) -> QueueDecision: | |
| """Atomically reject, queue, or exclusively claim an existing node.""" | |
| async with self._lock: | |
| return self._enqueue_or_claim(node_id) | |
| async def add_and_enqueue( | |
| self, | |
| node_id: str, | |
| scope: MessageScope, | |
| prompt: str, | |
| status_message_id: str, | |
| parent_id: str, | |
| parent_reference_id: str, | |
| ) -> QueueDecision: | |
| """Atomically add a reply and admit it to this tree.""" | |
| async with self._lock: | |
| self._graph.add_node( | |
| node_id=node_id, | |
| scope=scope, | |
| prompt=prompt, | |
| status_message_id=status_message_id, | |
| parent_id=parent_id, | |
| parent_reference_id=parent_reference_id, | |
| ) | |
| return self._enqueue_or_claim(node_id) | |
| async def finish_and_claim_next(self, claim_id: str) -> CompletionResult: | |
| """Release only the matching claim and atomically select its successor.""" | |
| async with self._lock: | |
| if self._active is None or self._active.claim.claim_id != claim_id: | |
| return CompletionResult( | |
| next_claim=None, | |
| queue=self._queue_entries(), | |
| ) | |
| self._active = None | |
| next_claim: NodeClaim | None = None | |
| while node_id := self._queue.pop(): | |
| node = self._graph.get_node(node_id) | |
| if node is not None and node.state is MessageState.PENDING: | |
| next_claim = self._claim(node) | |
| break | |
| return CompletionResult( | |
| next_claim=next_claim, | |
| queue=self._queue_entries(), | |
| ) | |
| async def cancel_node( | |
| self, | |
| node_id: str, | |
| ) -> TreeCancellation: | |
| """Atomically cancel one active, queued, or stale runnable node.""" | |
| async with self._lock: | |
| node = self._graph.get_node(node_id) | |
| active_claim = ( | |
| self._active.claim | |
| if self._active is not None | |
| and self._active.claim.node.node_id == node_id | |
| else None | |
| ) | |
| if active_claim is not None: | |
| active = self._active | |
| if active is not None: | |
| active.cancellation_requested = True | |
| if node is None: | |
| return TreeCancellation( | |
| nodes=(), | |
| active_claim=active_claim, | |
| queue_update=None, | |
| ) | |
| queue_changed = self._queue.remove(node_id) | |
| cancelled_nodes: tuple[NodeUiTarget, ...] = () | |
| if node.state in (MessageState.PENDING, MessageState.IN_PROGRESS): | |
| node.mark_error() | |
| cancelled_nodes = (self._ui_target(node),) | |
| elif node.state is MessageState.ERROR and active_claim is not None: | |
| cancelled_nodes = (self._ui_target(node),) | |
| return TreeCancellation( | |
| nodes=cancelled_nodes, | |
| active_claim=active_claim, | |
| queue_update=self._queue_entries() if queue_changed else None, | |
| ) | |
| async def cancel_all( | |
| self, | |
| ) -> TreeCancellation: | |
| """Atomically cancel every runnable node present at the transition.""" | |
| async with self._lock: | |
| cancelled_nodes: list[NodeUiTarget] = [] | |
| seen: set[str] = set() | |
| active_claim: NodeClaim | None = None | |
| if self._active is not None: | |
| active_claim = self._active.claim | |
| self._active.cancellation_requested = True | |
| active_node = self._graph.get_node(active_claim.node.node_id) | |
| if active_node is not None and active_node.state in ( | |
| MessageState.PENDING, | |
| MessageState.IN_PROGRESS, | |
| ): | |
| active_node.mark_error() | |
| seen.add(active_node.node_id) | |
| cancelled_nodes.append(self._ui_target(active_node)) | |
| elif active_node is not None: | |
| seen.add(active_node.node_id) | |
| if active_node.state is MessageState.ERROR: | |
| cancelled_nodes.append(self._ui_target(active_node)) | |
| queued_ids = self._queue.drain() | |
| for node_id in queued_ids: | |
| node = self._graph.get_node(node_id) | |
| if node is None or node.state not in ( | |
| MessageState.PENDING, | |
| MessageState.IN_PROGRESS, | |
| ): | |
| continue | |
| node.mark_error() | |
| seen.add(node_id) | |
| cancelled_nodes.append(self._ui_target(node)) | |
| for node in self._graph.all_nodes(): | |
| if node.node_id in seen or node.state not in ( | |
| MessageState.PENDING, | |
| MessageState.IN_PROGRESS, | |
| ): | |
| continue | |
| node.mark_error() | |
| cancelled_nodes.append(self._ui_target(node)) | |
| return TreeCancellation( | |
| nodes=tuple(cancelled_nodes), | |
| active_claim=active_claim, | |
| queue_update=() if queued_ids else None, | |
| ) | |
| async def remove_message_subtree( | |
| self, | |
| reference_id: str, | |
| ) -> MessageSubtreeRemoval: | |
| """Atomically cancel and detach one literal platform reply subtree.""" | |
| async with self._lock: | |
| resolved = self._graph.resolve_reference(reference_id) | |
| reference_ids = tuple(self._graph.get_reference_descendants(reference_id)) | |
| if resolved is None or not reference_ids: | |
| empty = TreeCancellation( | |
| nodes=(), | |
| active_claim=None, | |
| queue_update=None, | |
| ) | |
| return MessageSubtreeRemoval( | |
| cancellation=empty, | |
| removed_message_ids=frozenset(), | |
| removed_entire_tree=False, | |
| ) | |
| owner, reference_kind = resolved | |
| removed_node_ids = { | |
| candidate | |
| for candidate in reference_ids | |
| if self._graph.get_node(candidate) is not None | |
| } | |
| affected_node_ids = set(removed_node_ids) | |
| if reference_kind is MessageReferenceKind.STATUS: | |
| affected_node_ids.add(owner.node_id) | |
| active_claim = ( | |
| self._active.claim | |
| if self._active is not None | |
| and self._active.claim.node.node_id in affected_node_ids | |
| else None | |
| ) | |
| if active_claim is not None: | |
| active = self._active | |
| if active is not None: | |
| active.cancellation_requested = True | |
| cancelled_nodes: list[NodeUiTarget] = [] | |
| queue_changed = False | |
| for node_id in affected_node_ids: | |
| node = self._graph.get_node(node_id) | |
| if node is None: | |
| continue | |
| queue_changed = self._queue.remove(node_id) or queue_changed | |
| if node.state in (MessageState.PENDING, MessageState.IN_PROGRESS): | |
| target = self._ui_target(node) | |
| node.mark_error() | |
| cancelled_nodes.append(target) | |
| elif ( | |
| node.state is MessageState.ERROR | |
| and active_claim is not None | |
| and active_claim.node.node_id == node_id | |
| ): | |
| cancelled_nodes.append(self._ui_target(node)) | |
| if reference_kind is MessageReferenceKind.STATUS: | |
| self._graph.clear_status(owner.node_id) | |
| removed_entire_tree = self.root_id in removed_node_ids | |
| self._graph.remove_nodes(removed_node_ids) | |
| cancellation = TreeCancellation( | |
| nodes=tuple(cancelled_nodes), | |
| active_claim=active_claim, | |
| queue_update=self._queue_entries() if queue_changed else None, | |
| ) | |
| return MessageSubtreeRemoval( | |
| cancellation=cancellation, | |
| removed_message_ids=frozenset(reference_ids), | |
| removed_entire_tree=removed_entire_tree, | |
| ) | |
| async def record_session( | |
| self, claim_id: str, session_id: str | |
| ) -> TreeSnapshot | None: | |
| """Record a real CLI session only for the currently active claim.""" | |
| async with self._lock: | |
| if ( | |
| self._active is None | |
| or self._active.claim.claim_id != claim_id | |
| or self._active.cancellation_requested | |
| ): | |
| return None | |
| node = self._graph.get_node(self._active.claim.node.node_id) | |
| if node is None or node.state is not MessageState.IN_PROGRESS: | |
| return None | |
| node.update_state(MessageState.IN_PROGRESS, session_id=session_id) | |
| return self._graph.snapshot() | |
| async def complete_claim( | |
| self, claim_id: str, session_id: str | None | |
| ) -> TreeSnapshot | None: | |
| """Mark the currently active claim complete.""" | |
| async with self._lock: | |
| if ( | |
| self._active is None | |
| or self._active.claim.claim_id != claim_id | |
| or self._active.cancellation_requested | |
| ): | |
| return None | |
| node = self._graph.get_node(self._active.claim.node.node_id) | |
| if node is None or node.state not in ( | |
| MessageState.IN_PROGRESS, | |
| MessageState.ERROR, | |
| ): | |
| return None | |
| node.update_state(MessageState.COMPLETED, session_id=session_id) | |
| return self._graph.snapshot() | |
| async def fail_claim( | |
| self, | |
| claim_id: str, | |
| *, | |
| propagate: bool, | |
| ) -> FailureResult: | |
| """Atomically fail the active claim and its pending descendants.""" | |
| async with self._lock: | |
| if ( | |
| self._active is None | |
| or self._active.claim.claim_id != claim_id | |
| or self._active.cancellation_requested | |
| ): | |
| return FailureResult(affected=(), queue_update=None, snapshot=None) | |
| node = self._graph.get_node(self._active.claim.node.node_id) | |
| if node is None: | |
| return FailureResult(affected=(), queue_update=None, snapshot=None) | |
| affected: list[NodeUiTarget] = [] | |
| queue_changed = False | |
| if node.state is not MessageState.COMPLETED: | |
| if node.state is not MessageState.ERROR: | |
| node.mark_error() | |
| affected.append(self._ui_target(node)) | |
| if propagate: | |
| for descendant_id in self._graph.get_descendants(node.node_id)[1:]: | |
| child = self._graph.get_node(descendant_id) | |
| if child is None or child.state is not MessageState.PENDING: | |
| continue | |
| child.mark_error() | |
| queue_changed = ( | |
| self._queue.remove(child.node_id) or queue_changed | |
| ) | |
| affected.append(self._ui_target(child)) | |
| return FailureResult( | |
| affected=tuple(affected), | |
| queue_update=self._queue_entries() if queue_changed else None, | |
| snapshot=self._graph.snapshot(), | |
| ) | |
| async def resolve_reply(self, reference_id: str) -> ReplyTarget | None: | |
| """Resolve a node/status reference without exposing the mutable graph.""" | |
| async with self._lock: | |
| resolved = self._graph.resolve_reference(reference_id) | |
| if resolved is None: | |
| return None | |
| node, reference_kind = resolved | |
| return ReplyTarget( | |
| node_id=node.node_id, | |
| reference_id=reference_id, | |
| reference_kind=reference_kind, | |
| queue_position=(self._queue.qsize() + 1) | |
| if self._active is not None | |
| else None, | |
| ) | |
| async def node_view(self, node_id: str) -> NodeView | None: | |
| """Return a copied node read model.""" | |
| async with self._lock: | |
| node = self._graph.get_node(node_id) | |
| if node is None: | |
| return None | |
| return NodeView( | |
| identity=self.identity, | |
| node_id=node.node_id, | |
| state=node.state, | |
| parent_id=node.parent_id, | |
| session_id=node.session_id, | |
| ) | |
| async def snapshot(self) -> TreeSnapshot: | |
| """Capture a detached persistence snapshot under the aggregate lock.""" | |
| async with self._lock: | |
| return self._graph.snapshot() | |
| async def message_ids_for_chat(self, platform: str, chat_id: str) -> set[str]: | |
| """Copy every prompt and FCC status belonging to one platform chat.""" | |
| async with self._lock: | |
| if self.identity.scope.platform != str(platform) or ( | |
| self.identity.scope.chat_id != str(chat_id) | |
| ): | |
| return set() | |
| return self._graph.all_reference_ids() | |
| def from_snapshot(cls, snapshot: TreeSnapshot) -> "MessageTree": | |
| """Restore and reconcile interrupted nodes before publishing the tree.""" | |
| graph = MessageTreeGraph.from_snapshot(snapshot) | |
| tree = cls(graph.get_root(), graph=graph) | |
| stale_targets: list[NodeUiTarget] = [] | |
| for node in graph.all_nodes(): | |
| if node.state in (MessageState.PENDING, MessageState.IN_PROGRESS): | |
| stale_targets.append(tree._ui_target(node)) | |
| node.mark_error() | |
| tree._restored_stale_targets = tuple(stale_targets) | |
| tree._restored_snapshot = graph.snapshot() | |
| return tree | |
| __all__ = ["MessageTree"] | |