diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index ce391e77..32af75d9 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -589,7 +589,9 @@ creation/extension, initial status messages, persistence, and enqueueing. [messaging/node_runner.py](messaging/node_runner.py) owns managed CLI session lifecycle for queued nodes: parent-session fork/resume, session registration, CLI event parsing, transcript/status updates, cancellation, error propagation, -and session cleanup. +and session cleanup. Runner persistence must be guarded by active tree +membership so late cancellation cleanup cannot restore state after `/clear` +removed a tree or reset the queue. [messaging/event_parser.py](messaging/event_parser.py) normalizes managed Claude JSON events into low-level transcript events. @@ -605,7 +607,10 @@ depend on the concrete workflow object or on platform SDK runtimes. [messaging/trees/manager.py](messaging/trees/manager.py) preserves per-conversation ordering with tree-aware queues. Replies become child nodes, and each tree processes one node at a time while separate trees can progress -independently. [messaging/trees/repository.py](messaging/trees/repository.py) +independently. Tree cancellation is terminal: cancelling `/stop`, reply +`/stop`, reply `/clear`, or global `/clear` awaits active task cleanup outside +tree locks before command state cleanup continues. +[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/node.py](messaging/trees/node.py) owns @@ -622,7 +627,10 @@ and message IDs to a JSON file under the managed agent workspace. 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). +[messaging/session/message_log.py](messaging/session/message_log.py). `/clear` +guarantees FCC state cleanup and tries each tracked platform delete, but +Discord/Telegram can still reject individual message deletions for platform +reasons such as permissions, age, or missing messages. ```mermaid sequenceDiagram diff --git a/messaging/commands.py b/messaging/commands.py index 3991f7dd..5965296e 100644 --- a/messaging/commands.py +++ b/messaging/commands.py @@ -3,12 +3,21 @@ Commands depend on MessagingCommandContext instead of the concrete workflow. """ +from dataclasses import dataclass + from loguru import logger from .command_context import MessagingCommandContext from .models import IncomingMessage +@dataclass(frozen=True, slots=True) +class MessageDeleteReport: + attempted: int + deleted: int + failed: int + + async def handle_stop_command( handler: MessagingCommandContext, incoming: IncomingMessage ) -> None: @@ -86,10 +95,10 @@ async def handle_stats_command( async def _delete_message_ids( handler: MessagingCommandContext, chat_id: str, msg_ids: set[str] -) -> None: +) -> MessageDeleteReport: """Best-effort delete messages by ID. Sorts numeric IDs descending.""" if not msg_ids: - return + return MessageDeleteReport(attempted=0, deleted=0, failed=0) def _as_int(s: str) -> int | None: try: @@ -108,15 +117,28 @@ async def _delete_message_ids( numeric.sort(reverse=True) ordered = [mid for _, mid in numeric] + non_numeric - try: - CHUNK = 100 - for i in range(0, len(ordered), CHUNK): - chunk = ordered[i : i + CHUNK] - await handler.outbound.queue_delete_messages( - chat_id, chunk, fire_and_forget=False + deleted = 0 + failed = 0 + for mid in ordered: + try: + await handler.outbound.queue_delete_message( + chat_id, mid, fire_and_forget=False ) - except Exception as e: - logger.debug(f"Batch delete failed: {type(e).__name__}: {e}") + deleted += 1 + except Exception as e: + failed += 1 + logger.debug( + "Message delete failed for chat {}: {}", chat_id, type(e).__name__ + ) + + if ordered: + logger.info( + "Clear delete attempted={} deleted={} failed={}", + len(ordered), + deleted, + failed, + ) + return MessageDeleteReport(attempted=len(ordered), deleted=deleted, failed=failed) async def _handle_clear_branch( @@ -166,6 +188,9 @@ async def _handle_clear_branch( updated_tree = handler.tree_queue.get_tree(root_id) if updated_tree: handler.session_store.save_tree_snapshot(updated_tree.snapshot()) + handler.session_store.forget_message_ids( + incoming.platform, incoming.chat_id, msg_ids + ) except Exception as e: logger.warning(f"Failed to update session store after branch clear: {e}") diff --git a/messaging/node_event_pipeline.py b/messaging/node_event_pipeline.py index c654d80b..f68c35a4 100644 --- a/messaging/node_event_pipeline.py +++ b/messaging/node_event_pipeline.py @@ -24,6 +24,7 @@ async def handle_session_info_event( *, cli_manager: ManagedClaudeSessionManagerProtocol, session_store: SessionStore, + save_tree_snapshot: Callable[[MessageTree], None] | None = None, ) -> tuple[str | None, str | None]: """Handle session_info event; return updated (captured_session_id, temp_session_id).""" if event_data.get("type") != "session_info": @@ -49,7 +50,10 @@ async def handle_session_info_event( MessageState.IN_PROGRESS, session_id=real_session_id, ) - session_store.save_tree_snapshot(tree.snapshot()) + if save_tree_snapshot is None: + session_store.save_tree_snapshot(tree.snapshot()) + else: + save_tree_snapshot(tree) return real_session_id, None @@ -65,6 +69,7 @@ async def process_parsed_cli_event( captured_session_id: str | None, *, session_store: SessionStore, + save_tree_snapshot: Callable[[MessageTree], None] | None = None, format_status: Callable[..., str], propagate_error_to_children: Callable[[str, str, str], Awaitable[None]], log_messaging_error_details: bool = False, @@ -99,7 +104,10 @@ async def process_parsed_cli_event( MessageState.COMPLETED, session_id=captured_session_id, ) - session_store.save_tree_snapshot(tree.snapshot()) + if save_tree_snapshot is None: + session_store.save_tree_snapshot(tree.snapshot()) + else: + save_tree_snapshot(tree) elif ptype == "error": error_msg = parsed.get("message", "Unknown error") em = error_msg if isinstance(error_msg, str) else str(error_msg) diff --git a/messaging/node_runner.py b/messaging/node_runner.py index 0d7de158..bb9397f8 100644 --- a/messaging/node_runner.py +++ b/messaging/node_runner.py @@ -63,10 +63,18 @@ class MessagingNodeRunner: ) return transcript, self._get_render_ctx() - def _save_tree(self, tree: MessageTree | None) -> None: + def _save_tree(self, tree: MessageTree | None, node_id: str) -> None: """Persist tree state after runner-owned mutations.""" - if tree: - self.session_store.save_tree_snapshot(tree.snapshot()) + if not tree: + return + active_tree = self._get_tree_queue().get_tree_for_node(node_id) + if active_tree is not tree: + logger.debug( + "Skipping stale tree save for node {} after cancellation/clear", + node_id, + ) + return + self.session_store.save_tree_snapshot(tree.snapshot()) async def process_node( self, @@ -186,7 +194,7 @@ class MessagingNodeRunner: MessageState.ERROR, error_message=error_message, ) - self._save_tree(tree) + self._save_tree(tree, node_id) trace_event( stage="claude_cli", event="claude_cli.session.limit_reached", @@ -218,6 +226,9 @@ class MessagingNodeRunner: temp_session_id, cli_manager=self.cli_manager, session_store=self.session_store, + save_tree_snapshot=lambda updated_tree: self._save_tree( + updated_tree, node_id + ), ) if event_data.get("type") == "session_info": continue @@ -240,6 +251,9 @@ class MessagingNodeRunner: node_id, captured_session_id, session_store=self.session_store, + save_tree_snapshot=lambda updated_tree: self._save_tree( + updated_tree, node_id + ), format_status=self._format_status, propagate_error_to_children=self.propagate_error_to_children, log_messaging_error_details=self._log_messaging_error_details, @@ -270,7 +284,7 @@ class MessagingNodeRunner: await tree.update_state( node_id, MessageState.ERROR, error_message="Cancelled by user" ) - self._save_tree(tree) + self._save_tree(tree, node_id) except Exception as e: trace_event( stage="claude_cli", @@ -328,7 +342,7 @@ class MessagingNodeRunner: node_id, error_msg, propagate_to_children=True ) if affected: - self._save_tree(tree_queue.get_tree_for_node(node_id)) + self._save_tree(tree_queue.get_tree_for_node(node_id), node_id) for child in affected[1:]: self.outbound.fire_and_forget( self.outbound.queue_edit_message( diff --git a/messaging/session/message_log.py b/messaging/session/message_log.py index eece6851..60d18a9b 100644 --- a/messaging/session/message_log.py +++ b/messaging/session/message_log.py @@ -68,6 +68,30 @@ class MessageLog: if item.get("message_id") is not None ] + def remove_message_ids( + self, platform: str, chat_id: str, message_ids: set[str] + ) -> bool: + chat_key = make_chat_key(platform, chat_id) + if not message_ids or chat_key not in self._items: + return False + + before_count = len(self._items[chat_key]) + self._items[chat_key] = [ + item + for item in self._items[chat_key] + if str(item.get("message_id")) not in message_ids + ] + if not self._items[chat_key]: + self._items.pop(chat_key, None) + self._ids.pop(chat_key, None) + else: + self._ids[chat_key] = { + str(item.get("message_id")) + for item in self._items[chat_key] + if item.get("message_id") is not None + } + return len(self._items.get(chat_key, [])) != before_count + def clear(self) -> None: self._items.clear() self._ids.clear() diff --git a/messaging/session/store.py b/messaging/session/store.py index b8f0b0af..4bb17fb7 100644 --- a/messaging/session/store.py +++ b/messaging/session/store.py @@ -133,6 +133,18 @@ class SessionStore: str(platform), str(chat_id) ) + def forget_message_ids( + self, platform: str, chat_id: str, message_ids: set[str] + ) -> None: + with self._lock: + removed = self._message_log.remove_message_ids( + str(platform), + str(chat_id), + {str(message_id) for message_id in message_ids}, + ) + if removed: + self._persistence.schedule_save() + def clear_all(self) -> None: with self._lock: self._conversation = ConversationSnapshot() diff --git a/messaging/trees/manager.py b/messaging/trees/manager.py index 05722f64..e025ff18 100644 --- a/messaging/trees/manager.py +++ b/messaging/trees/manager.py @@ -12,6 +12,35 @@ from .repository import TreeRepository from .runtime import MessageTree from .snapshot import ConversationSnapshot +CANCEL_TASK_DRAIN_TIMEOUT_S = 5.0 + + +async def _drain_cancelled_tasks(tasks: list[asyncio.Task]) -> None: + """Wait briefly for cancelled node tasks to finish their cleanup.""" + if not tasks: + return + + done, pending = await asyncio.wait( + set(tasks), + timeout=CANCEL_TASK_DRAIN_TIMEOUT_S, + ) + if pending: + logger.warning( + "Timed out waiting for {} cancelled messaging task(s) to finish cleanup", + len(pending), + ) + + for task in done: + if task.cancelled(): + continue + try: + task.result() + except Exception as exc: + logger.debug( + "Cancelled messaging task finished with {}", + type(exc).__name__, + ) + class TreeQueueManager: """ @@ -218,10 +247,13 @@ class TreeQueueManager: return [] cancelled_nodes = [] + cancelled_tasks: list[asyncio.Task] = [] cleanup_count = 0 async with tree.with_lock(): - if tree.cancel_current_task(): + cancelled_task = tree.cancel_current_task() + if cancelled_task: + cancelled_tasks.append(cancelled_task) current_id = tree.current_node_id if current_id: node = tree.get_node(current_id) @@ -246,6 +278,8 @@ class TreeQueueManager: tree.reset_processing_state() + await _drain_cancelled_tasks(cancelled_tasks) + if cancelled_nodes: logger.info( f"Cancelled {len(cancelled_nodes)} active nodes in tree {root_id}" @@ -274,8 +308,11 @@ class TreeQueueManager: if node.state in (MessageState.COMPLETED, MessageState.ERROR): return [] + cancelled_tasks: list[asyncio.Task] = [] if tree.is_current_node(node_id): - self._processor.cancel_current(tree) + cancelled_task = self._processor.cancel_current(tree) + if cancelled_task: + cancelled_tasks.append(cancelled_task) removed_from_queue = False try: @@ -289,6 +326,7 @@ class TreeQueueManager: if removed_from_queue: await self._processor.notify_queue_updated(tree) + await _drain_cancelled_tasks(cancelled_tasks) return [node] @@ -296,10 +334,10 @@ class TreeQueueManager: """Cancel all messages in all trees.""" async with self._lock: root_ids = list(self._repository.tree_ids()) - all_cancelled: list[MessageNode] = [] - for root_id in root_ids: - all_cancelled.extend(await self.cancel_tree(root_id)) - return all_cancelled + all_cancelled: list[MessageNode] = [] + for root_id in root_ids: + all_cancelled.extend(await self.cancel_tree(root_id)) + return all_cancelled def cleanup_stale_nodes(self) -> int: """ @@ -348,6 +386,7 @@ class TreeQueueManager: branch_ids = set(tree.get_descendants(branch_root_id)) cancelled: list[MessageNode] = [] + cancelled_tasks: list[asyncio.Task] = [] removed_from_queue = False async with tree.with_lock(): @@ -360,7 +399,9 @@ class TreeQueueManager: continue if tree.is_current_node(nid): - self._processor.cancel_current(tree) + cancelled_task = self._processor.cancel_current(tree) + if cancelled_task: + cancelled_tasks.append(cancelled_task) tree.set_node_error_sync(node, "Cancelled by user") cancelled.append(node) else: @@ -374,6 +415,7 @@ class TreeQueueManager: logger.info(f"Cancelled {len(cancelled)} nodes in branch {branch_root_id}") if removed_from_queue: await self._processor.notify_queue_updated(tree) + await _drain_cancelled_tasks(cancelled_tasks) return cancelled async def remove_branch( diff --git a/messaging/trees/processor.py b/messaging/trees/processor.py index 10c28714..d07c5f5e 100644 --- a/messaging/trees/processor.py +++ b/messaging/trees/processor.py @@ -172,7 +172,7 @@ class TreeQueueProcessor: ) return False - def cancel_current(self, tree: MessageTree) -> bool: + def cancel_current(self, tree: MessageTree) -> asyncio.Task | None: """Cancel the currently running task in a tree.""" return tree.cancel_current_task() diff --git a/messaging/trees/runtime.py b/messaging/trees/runtime.py index 65adac77..33095f4c 100644 --- a/messaging/trees/runtime.py +++ b/messaging/trees/runtime.py @@ -131,11 +131,15 @@ class MessageTree: def set_current_task(self, task: asyncio.Task | None) -> None: self._current_task = task - def cancel_current_task(self) -> bool: + @property + def current_task(self) -> asyncio.Task | None: + return self._current_task + + def cancel_current_task(self) -> asyncio.Task | None: if self._current_task and not self._current_task.done(): self._current_task.cancel() - return True - return False + return self._current_task + return None def set_node_error_sync(self, node: MessageNode, error_message: str) -> None: node.mark_error(error_message) diff --git a/pyproject.toml b/pyproject.toml index 1bdc760d..1f75a1e6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "free-claude-code" -version = "3.4.0" +version = "3.4.1" description = "Middleware between Claude Code CLI (Anthropic API) and NVIDIA NIM" readme = "README.md" requires-python = ">=3.14.0" diff --git a/tests/messaging/test_handler.py b/tests/messaging/test_handler.py index 774fe28e..708b8b92 100644 --- a/tests/messaging/test_handler.py +++ b/tests/messaging/test_handler.py @@ -1,9 +1,10 @@ +import asyncio from unittest.mock import AsyncMock, MagicMock, patch import pytest from messaging.models import IncomingMessage -from messaging.trees import MessageNode, MessageState, MessageTree +from messaging.trees import MessageNode, MessageState, MessageTree, TreeQueueManager from messaging.workflow import MessagingWorkflow @@ -614,6 +615,29 @@ async def test_handle_message_clear_command_deletes_message_log_ids( assert deleted == {"42", "43", "150"} +@pytest.mark.asyncio +async def test_handle_message_clear_command_continues_after_delete_failure( + handler, mock_platform, mock_session_store, incoming_message_factory +): + handler.stop_all_tasks = AsyncMock(return_value=0) + mock_session_store.get_message_ids_for_chat.return_value = ["41", "42", "43"] + + async def delete_message(chat_id, message_id, fire_and_forget=True): + if message_id == "42": + raise RuntimeError("platform rejected delete") + + mock_platform.queue_delete_message = AsyncMock(side_effect=delete_message) + + incoming = incoming_message_factory( + text="/clear", chat_id="chat_1", message_id="150" + ) + await handler.handle_message(incoming) + + deleted = [c.args[1] for c in mock_platform.queue_delete_message.call_args_list] + assert deleted == ["150", "43", "42", "41"] + mock_session_store.clear_all.assert_called_once() + + @pytest.mark.asyncio async def test_handle_message_clear_command_reply_clears_branch( handler, mock_platform, mock_session_store, incoming_message_factory @@ -659,6 +683,9 @@ async def test_handle_message_clear_command_reply_clears_branch( assert "100" not in deleted_ids assert "101" not in deleted_ids mock_session_store.save_tree_snapshot.assert_called() + mock_session_store.forget_message_ids.assert_called_once_with( + "telegram", "chat_1", {"102", "103", "150"} + ) assert handler.tree_queue.get_tree_for_node("102") is None assert handler.tree_queue.get_tree_for_node("100") is not None @@ -711,9 +738,56 @@ async def test_handle_message_clear_command_reply_to_root_clears_tree( assert set(deleted_ids) == {"100", "101", "150"} mock_session_store.remove_tree_snapshot.assert_called_once_with("100") + mock_session_store.forget_message_ids.assert_called_once_with( + "telegram", "chat_1", {"100", "101", "150"} + ) assert handler.tree_queue.get_tree_count() == 0 +@pytest.mark.asyncio +async def test_cancelled_node_runner_does_not_save_after_clear_replaces_queue( + handler, mock_cli_manager, mock_session_store, incoming_message_factory +): + """Late cancellation cleanup must not restore a tree after /clear reset.""" + started = asyncio.Event() + + async def start_task(*args, **kwargs): + started.set() + try: + await asyncio.sleep(60) + except asyncio.CancelledError: + raise + if False: + yield {} + + cli_session = MagicMock() + cli_session.start_task = start_task + mock_cli_manager.get_or_create_session.return_value = ( + cli_session, + "pending_1", + True, + ) + + incoming = incoming_message_factory(text="work", chat_id="chat_1", message_id="100") + tree = await handler.tree_queue.create_tree("100", incoming, "101") + node = tree.get_node("100") + assert node is not None + + task = asyncio.create_task(handler.node_runner.process_node("100", node)) + await started.wait() + handler.replace_tree_queue( + TreeQueueManager( + queue_update_callback=handler.update_queue_positions, + node_started_callback=handler.mark_node_processing, + ) + ) + + task.cancel() + await task + + mock_session_store.save_tree_snapshot.assert_not_called() + + @pytest.mark.asyncio async def test_handle_message_clear_command_reply_pending_voice_cancels( handler, mock_platform, mock_session_store, incoming_message_factory diff --git a/tests/messaging/test_tree_concurrency.py b/tests/messaging/test_tree_concurrency.py index 209e7431..aa584aa3 100644 --- a/tests/messaging/test_tree_concurrency.py +++ b/tests/messaging/test_tree_concurrency.py @@ -412,7 +412,7 @@ class TestTreeQueueManagerConcurrency: assert len(cancelled) >= 1 # At least the current + queued # Tree should no longer be processing - assert tree._is_processing is False + assert tree.is_processing is False @pytest.mark.asyncio async def test_cancel_nonexistent_tree(self): diff --git a/tests/messaging/test_tree_processor.py b/tests/messaging/test_tree_processor.py index 0f597579..3a2ee394 100644 --- a/tests/messaging/test_tree_processor.py +++ b/tests/messaging/test_tree_processor.py @@ -44,7 +44,7 @@ async def test_process_node_success(tree_processor, sample_tree, sample_node): await tree_processor.process_node(sample_tree, sample_node, processor) processor.assert_called_once_with(sample_node.node_id, sample_node) - assert sample_tree._current_node_id is None + assert sample_tree.current_node_id is None @pytest.mark.asyncio @@ -54,7 +54,7 @@ async def test_process_node_cancelled(tree_processor, sample_tree, sample_node): with pytest.raises(asyncio.CancelledError): await tree_processor.process_node(sample_tree, sample_node, processor) - assert sample_tree._current_node_id is None + assert sample_tree.current_node_id is None @pytest.mark.asyncio @@ -69,7 +69,7 @@ async def test_process_node_exception(tree_processor, sample_tree, sample_node): sample_tree.update_state.assert_called_once_with( sample_node.node_id, MessageState.ERROR, error_message="Test error" ) - assert sample_tree._current_node_id is None + assert sample_tree.current_node_id is None @pytest.mark.asyncio @@ -84,20 +84,20 @@ async def test_enqueue_and_start_when_free(tree_processor, sample_tree): was_queued = await tree_processor.enqueue_and_start(sample_tree, node_id, processor) assert was_queued is False - assert sample_tree._is_processing is True - assert sample_tree._current_node_id == node_id - assert sample_tree._current_task is not None + assert sample_tree.is_processing is True + assert sample_tree.current_node_id == node_id + assert sample_tree.current_task is not None # Clean up task - sample_tree._current_task.cancel() + sample_tree.current_task.cancel() with contextlib.suppress(asyncio.CancelledError): - await sample_tree._current_task + await sample_tree.current_task @pytest.mark.asyncio async def test_enqueue_and_start_when_busy(tree_processor, sample_tree): processor = AsyncMock() - sample_tree._is_processing = True + sample_tree.set_processing_state("busy", True) node_id = "node1" was_queued = await tree_processor.enqueue_and_start(sample_tree, node_id, processor) @@ -110,33 +110,33 @@ async def test_enqueue_and_start_when_busy(tree_processor, sample_tree): def test_cancel_current_task(tree_processor, sample_tree): mock_task = MagicMock(spec=asyncio.Task) mock_task.done.return_value = False - sample_tree._current_task = mock_task + sample_tree.set_current_task(mock_task) cancelled = tree_processor.cancel_current(sample_tree) - assert cancelled is True + assert cancelled is mock_task mock_task.cancel.assert_called_once() def test_cancel_current_task_already_done(tree_processor, sample_tree): mock_task = MagicMock(spec=asyncio.Task) mock_task.done.return_value = True - sample_tree._current_task = mock_task + sample_tree.set_current_task(mock_task) cancelled = tree_processor.cancel_current(sample_tree) - assert cancelled is False + assert cancelled is None mock_task.cancel.assert_not_called() @pytest.mark.asyncio async def test_process_next_queue_empty(tree_processor, sample_tree): processor = AsyncMock() - sample_tree._is_processing = True + sample_tree.set_processing_state("busy", True) await tree_processor._process_next(sample_tree, processor) - assert sample_tree._is_processing is False + assert sample_tree.is_processing is False @pytest.mark.asyncio @@ -149,13 +149,13 @@ async def test_process_next_with_item(tree_processor, sample_tree): await tree_processor._process_next(sample_tree, processor) - assert sample_tree._current_node_id == "next_node" - assert sample_tree._current_task is not None + assert sample_tree.current_node_id == "next_node" + assert sample_tree.current_task is not None # Clean up - sample_tree._current_task.cancel() + sample_tree.current_task.cancel() with contextlib.suppress(asyncio.CancelledError): - await sample_tree._current_task + await sample_tree.current_task @pytest.mark.asyncio @@ -190,14 +190,14 @@ async def test_process_next_skips_stale_id_and_runs_next_valid_node(sample_tree) await processor._process_next(sample_tree, node_processor) - assert sample_tree._is_processing is True - assert sample_tree._current_node_id == "valid_node" + assert sample_tree.is_processing is True + assert sample_tree.current_node_id == "valid_node" node_started.assert_awaited_once_with(sample_tree, "valid_node") queue_updated.assert_awaited_once_with(sample_tree) release.set() - if sample_tree._current_task: - await sample_tree._current_task + if sample_tree.current_task: + await sample_tree.current_task @pytest.mark.asyncio @@ -208,14 +208,14 @@ async def test_process_next_drains_all_stale_ids_without_wedging(sample_tree): queue_update_callback=queue_updated, node_started_callback=node_started, ) - sample_tree._is_processing = True + sample_tree.set_processing_state("busy", True) sample_tree.put_queue_unlocked("missing_one") sample_tree.put_queue_unlocked("missing_two") await processor._process_next(sample_tree, AsyncMock()) - assert sample_tree._is_processing is False - assert sample_tree._current_node_id is None + assert sample_tree.is_processing is False + assert sample_tree.current_node_id is None assert sample_tree._queue.qsize() == 0 node_started.assert_not_awaited() queue_updated.assert_awaited_once_with(sample_tree) @@ -256,7 +256,7 @@ async def test_process_next_triggers_node_started(sample_tree): node_started.assert_awaited_once_with(sample_tree, "next_node") - if sample_tree._current_task: - sample_tree._current_task.cancel() + if sample_tree.current_task: + sample_tree.current_task.cancel() with contextlib.suppress(asyncio.CancelledError): - await sample_tree._current_task + await sample_tree.current_task diff --git a/tests/messaging/test_tree_queue.py b/tests/messaging/test_tree_queue.py index 5f852d82..2c521989 100644 --- a/tests/messaging/test_tree_queue.py +++ b/tests/messaging/test_tree_queue.py @@ -1,6 +1,7 @@ """Tests for tree-based message queue system.""" import asyncio +import contextlib from unittest.mock import AsyncMock import pytest @@ -13,6 +14,7 @@ from messaging.trees import ( TreeQueueManager, TreeSnapshot, ) +from messaging.trees import manager as tree_manager_module from messaging.trees.graph import MessageTreeGraph from messaging.trees.snapshot import node_from_snapshot, node_to_snapshot @@ -553,6 +555,68 @@ class TestTreeQueueManager: processing_complete.set() + @pytest.mark.asyncio + async def test_cancel_tree_waits_for_current_task_cleanup(self): + """cancel_tree returns only after the cancelled task cleanup runs.""" + manager = TreeQueueManager() + started = asyncio.Event() + cleanup_done = asyncio.Event() + + async def slow_processor(node_id, node): + started.set() + try: + await asyncio.sleep(60) + except asyncio.CancelledError: + cleanup_done.set() + raise + + incoming = IncomingMessage( + text="Test", + chat_id="1", + user_id="1", + message_id="m1", + platform="test", + ) + await manager.create_tree("m1", incoming, "s1") + await manager.enqueue("m1", slow_processor) + await started.wait() + + cancelled = await manager.cancel_tree("m1") + + assert [node.node_id for node in cancelled] == ["m1"] + assert cleanup_done.is_set() + + @pytest.mark.asyncio + async def test_cancel_node_waits_for_current_task_cleanup(self): + """cancel_node waits when the target node is actively running.""" + manager = TreeQueueManager() + started = asyncio.Event() + cleanup_done = asyncio.Event() + + async def slow_processor(node_id, node): + started.set() + try: + await asyncio.sleep(60) + except asyncio.CancelledError: + cleanup_done.set() + raise + + incoming = IncomingMessage( + text="Test", + chat_id="1", + user_id="1", + message_id="m1", + platform="test", + ) + await manager.create_tree("m1", incoming, "s1") + await manager.enqueue("m1", slow_processor) + await started.wait() + + cancelled = await manager.cancel_node("m1") + + assert [node.node_id for node in cancelled] == ["m1"] + assert cleanup_done.is_set() + @pytest.mark.asyncio async def test_cancel_branch(self): """Test cancel_branch cancels only nodes in subtree.""" @@ -595,6 +659,112 @@ class TestTreeQueueManager: assert sibling_node is not None assert sibling_node.state == MessageState.PENDING + @pytest.mark.asyncio + async def test_cancel_branch_waits_for_current_task_cleanup(self): + """cancel_branch waits when the active node is inside the branch.""" + manager = TreeQueueManager() + started = asyncio.Event() + cleanup_done = asyncio.Event() + + async def slow_processor(node_id, node): + started.set() + try: + await asyncio.sleep(60) + except asyncio.CancelledError: + cleanup_done.set() + raise + + root_incoming = IncomingMessage( + text="Root", chat_id="1", user_id="1", message_id="root", platform="test" + ) + await manager.create_tree("root", root_incoming, "s1") + child_incoming = IncomingMessage( + text="Child", + chat_id="1", + user_id="1", + message_id="child", + platform="test", + reply_to_message_id="root", + ) + await manager.add_to_tree("root", "child", child_incoming, "s2") + await manager.enqueue("child", slow_processor) + await started.wait() + + cancelled = await manager.cancel_branch("child") + + assert [node.node_id for node in cancelled] == ["child"] + assert cleanup_done.is_set() + + @pytest.mark.asyncio + async def test_cancel_all_waits_for_current_task_cleanup_across_trees(self): + """cancel_all waits for active task cleanup in every tree.""" + manager = TreeQueueManager() + started: set[str] = set() + cleanup_done: set[str] = set() + all_started = asyncio.Event() + + async def slow_processor(node_id, node): + started.add(node_id) + if started == {"a", "b"}: + all_started.set() + try: + await asyncio.sleep(60) + except asyncio.CancelledError: + cleanup_done.add(node_id) + raise + + for node_id in ("a", "b"): + incoming = IncomingMessage( + text="Test", + chat_id="1", + user_id="1", + message_id=node_id, + platform="test", + ) + await manager.create_tree(node_id, incoming, f"s_{node_id}") + await manager.enqueue(node_id, slow_processor) + await all_started.wait() + + cancelled = await manager.cancel_all() + + assert {node.node_id for node in cancelled} == {"a", "b"} + assert cleanup_done == {"a", "b"} + + @pytest.mark.asyncio + async def test_cancel_task_drain_timeout_is_bounded(self, monkeypatch): + """Cancellation drain returns when a task refuses to finish cleanup.""" + monkeypatch.setattr(tree_manager_module, "CANCEL_TASK_DRAIN_TIMEOUT_S", 0.01) + started = asyncio.Event() + cleanup_started = asyncio.Event() + cleanup_release = asyncio.Event() + + async def stubborn_task(): + started.set() + try: + await asyncio.sleep(60) + except asyncio.CancelledError: + cleanup_started.set() + await cleanup_release.wait() + raise + + task = asyncio.create_task(stubborn_task()) + await started.wait() + task.cancel() + await cleanup_started.wait() + + start = asyncio.get_running_loop().time() + try: + await tree_manager_module._drain_cancelled_tasks([task]) + elapsed = asyncio.get_running_loop().time() - start + assert not task.done() + finally: + cleanup_release.set() + if not task.done(): + with contextlib.suppress(asyncio.CancelledError): + await task + + assert elapsed < 0.5 + @pytest.mark.asyncio async def test_cancel_node_refreshes_queue_positions_for_remaining_nodes(self): """Reply-scoped stop refreshes queued siblings after removing one queued node.""" diff --git a/uv.lock b/uv.lock index 0e68181c..9d676953 100644 --- a/uv.lock +++ b/uv.lock @@ -561,7 +561,7 @@ wheels = [ [[package]] name = "free-claude-code" -version = "3.4.0" +version = "3.4.1" source = { editable = "." } dependencies = [ { name = "aiohttp" },