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

Make messaging cancellation terminal (#995)

Browse source 
## Problem

Messaging `/clear` and `/stop` could return before cancelled node tasks
finished cleanup. Late cleanup could save stale conversation state after
`/clear` reset FCC state.

## Changes

| Before | After |
| --- | --- |
| Tree cancellation called `task.cancel()` and returned immediately. |
Tree cancellation awaits cancelled task cleanup outside tree locks with
a bounded timeout. |
| Node runners saved snapshots even after their tree was removed or
replaced. | Node runners save only when their node still belongs to the
active tree queue. |
| `/clear` deletion stopped at a batch failure and left tracking vague.
| `/clear` attempts each tracked delete independently and clears
FCC-owned tracking state. |
| Architecture docs did not state terminal cancellation ownership. |
Architecture docs assign terminal cancellation to `messaging/trees` and
guarded cleanup persistence to node runners. |
| Package metadata stayed at `3.4.0`. | Package metadata and lockfile
move to `3.4.1`. |

<!-- greptile_comment -->

<details open><summary><h3>Greptile Summary</h3></summary>

This PR makes messaging cancellation wait for node cleanup before
command cleanup continues. The main changes are:

- Drains cancelled tree tasks with a bounded timeout outside tree locks.
- Guards node-runner snapshot saves so removed or replaced trees are not
restored by late cleanup.
- Changes `/clear` deletion to try each tracked platform message
independently.
- Removes cleared branch message IDs from FCC-owned tracking state.
- Adds cancellation, stale-save, and clear-delete regression tests.
- Bumps package metadata and lockfile version to `3.4.1`.
</details>

<h3>Confidence Score: 5/5</h3>

Safe to merge with minimal risk.

The cancellation paths now drain outside tree locks with a bounded wait,
stale snapshot persistence is guarded, and `/clear` continues through
individual delete failures. Tests cover the key cancellation cleanup,
timeout, stale-save, and clear-delete cases. No blocking correctness or
security issues were found in the changed files.

No files require special attention.

<details><summary><h3><a href="https://www.greptile.com/trex"><img
alt="T-Rex"
src="https://greptile-static-assets.s3.amazonaws.com/trex/trex_green.svg"
height="20" align="absmiddle"></a> T-Rex Logs</h3></summary>

**What T-Rex did**
- Ran the messaging cancellation test suite with verbose output to
capture the command, current working directory, timestamps, and verbose
test names.
- Verified the targeted cancellation tests passed, including
test\_cancel\_tree\_waits\_for\_current\_task\_cleanup,
test\_cancel\_node\_waits\_for\_current\_task\_cleanup,
test\_cancel\_branch\_waits\_for\_current\_task\_cleanup,
test\_cancel\_all\_waits\_for\_current\_task\_cleanup\_across\_trees,
and test\_cancel\_task\_drain\_timeout\_is\_bounded.
- Verified that the coverage tests for /clear handling and stale
persistence guard also passed, including
test\_handle\_message\_clear\_command\_stops\_deletes\_and\_wipes\_state
and
test\_cancelled\_node\_runner\_does\_not\_save\_after\_clear\_replaces\_queue.

<a
href="https://app.greptile.com/trex/runs/13359514/artifacts"><picture><source
media="(prefers-color-scheme: dark)"
srcset="https://greptile-static-assets.s3.amazonaws.com/badges/ViewAllArtifactsDark.svg?v=4"><source
media="(prefers-color-scheme: light)"
srcset="https://greptile-static-assets.s3.amazonaws.com/badges/ViewAllArtifacts.svg?v=4"><img
alt="View all artifacts"
src="https://greptile-static-assets.s3.amazonaws.com/badges/ViewAllArtifacts.svg?v=4"></picture></a>

<sub><a href="https://www.greptile.com/trex"><img alt="T-Rex"
src="https://greptile-static-assets.s3.amazonaws.com/trex/trex_green.svg"
height="14" align="absmiddle"></a> Ran code and verified through
T-Rex</sub>
</details>

<details open><summary><h3>Important Files Changed</h3></summary>

| Filename | Overview |
|----------|----------|
| messaging/trees/manager.py | Makes tree, node, branch, and all-tree
cancellation await cancelled task cleanup outside tree locks with a
bounded drain helper. |
| messaging/node_runner.py | Guards runner-owned snapshot saves by
confirming the node still belongs to the active tree queue. |
| messaging/commands.py | Updates `/clear` deletion to attempt message
deletes individually and forget branch-owned message IDs after branch
clears. |
| messaging/session/message_log.py | Adds targeted removal of tracked
message IDs while keeping the per-chat ID cache synchronized. |
| tests/messaging/test_tree_queue.py | Adds cancellation-drain tests for
tree, node, branch, all-tree, and timeout-bounded cleanup paths. |
| tests/messaging/test_handler.py | Adds coverage for resilient clear
deletion, branch message-log cleanup, and stale cancellation persistence
prevention. |

</details>

<details open><summary><h3>Sequence Diagram</h3></summary>

<a href="#gh-light-mode-only">

```mermaid
%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
participant User
participant Commands as messaging/commands.py
participant Manager as TreeQueueManager
participant Tree as MessageTree
participant Runner as MessagingNodeRunner
participant Store as SessionStore
participant Outbound as OutboundMessenger

User->>Commands: /clear or /stop
Commands->>Manager: cancel_tree/cancel_branch/cancel_all
Manager->>Tree: cancel_current_task()
Tree-->>Manager: cancelled asyncio.Task
Manager->>Tree: mark queued/current nodes ERROR
Manager->>Runner: task cancellation propagates
Manager->>Manager: await _drain_cancelled_tasks(timeout)
Runner->>Runner: cancellation cleanup/update UI
Runner->>Manager: check active tree for node
alt node still belongs to active queue
    Runner->>Store: save_tree_snapshot(snapshot)
else tree removed or queue replaced
    Runner-->>Store: skip stale save
end
Commands->>Outbound: delete tracked messages individually
Commands->>Store: clear_all or forget_message_ids
```

</a>
<a href="#gh-dark-mode-only">

```mermaid
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
participant User
participant Commands as messaging/commands.py
participant Manager as TreeQueueManager
participant Tree as MessageTree
participant Runner as MessagingNodeRunner
participant Store as SessionStore
participant Outbound as OutboundMessenger

User->>Commands: /clear or /stop
Commands->>Manager: cancel_tree/cancel_branch/cancel_all
Manager->>Tree: cancel_current_task()
Tree-->>Manager: cancelled asyncio.Task
Manager->>Tree: mark queued/current nodes ERROR
Manager->>Runner: task cancellation propagates
Manager->>Manager: await _drain_cancelled_tasks(timeout)
Runner->>Runner: cancellation cleanup/update UI
Runner->>Manager: check active tree for node
alt node still belongs to active queue
    Runner->>Store: save_tree_snapshot(snapshot)
else tree removed or queue replaced
    Runner-->>Store: skip stale save
end
Commands->>Outbound: delete tracked messages individually
Commands->>Store: clear_all or forget_message_ids
```

</a>
</details>

<sub>Reviews (2): Last reviewed commit: ["Bound messaging cancellation
drain"](d63b5b43f5)
| [Re-trigger
Greptile](https://app.greptile.com/api/retrigger?id=41990665)</sub>

<!-- /greptile_comment -->
This commit is contained in:
Ali Khokhar 2026-07-05 21:24:11 -07:00 • committed by GitHub
parent 05dae97248
commit 28e8996121
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
15 changed files with 446 additions and 65 deletions
Download patch file Download diff file Expand all files Collapse all files

14
ARCHITECTURE.md
View file

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

45
messaging/commands.py
View file

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

12
messaging/node_event_pipeline.py
View file

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

26
messaging/node_runner.py
View file

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

24
messaging/session/message_log.py
View file

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

12
messaging/session/store.py
View file

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

56
messaging/trees/manager.py
View file

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

2
messaging/trees/processor.py
View file

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

10
messaging/trees/runtime.py
View file

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

2
pyproject.toml
View file

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

76
tests/messaging/test_handler.py
View file

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

2
tests/messaging/test_tree_concurrency.py
View file

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

58
tests/messaging/test_tree_processor.py
View file

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

170
tests/messaging/test_tree_queue.py
View file

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

2
uv.lock generated
View file

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