fix: parallel subagent cancellation + register task before lock

- cancel_by_session: use asyncio.gather for parallel cancellation
  instead of sequential await per task
- _dispatch: register in _active_tasks before acquiring lock so /stop
  can find queued tasks (synced from #1179)
This commit is contained in:
coldxiangyu
2026-02-25 18:21:46 +08:00
parent 2466b8b843
commit 4768b9a09d
2 changed files with 32 additions and 29 deletions

View File

@@ -304,9 +304,14 @@ class AgentLoop:
))
async def _dispatch(self, msg: InboundMessage) -> None:
"""Dispatch a message for processing under the global lock."""
async with self._processing_lock:
"""Dispatch a message for processing under the global lock.
The task is registered in _active_tasks *before* acquiring the lock
so that /stop can find (and cancel) tasks that are still queued.
"""
self._active_tasks[msg.session_key] = asyncio.current_task() # type: ignore[arg-type]
try:
async with self._processing_lock:
try:
response = await self._process_message(msg)
if response is not None:

View File

@@ -268,17 +268,15 @@ When you have completed the task, provide a clear summary of your findings or ac
async def cancel_by_session(self, session_key: str) -> int:
"""Cancel all subagents spawned under the given session. Returns count cancelled."""
task_ids = list(self._session_tasks.get(session_key, []))
cancelled = 0
to_cancel: list[asyncio.Task] = []
for tid in task_ids:
t = self._running_tasks.get(tid)
if t and not t.done():
t.cancel()
try:
await t
except (asyncio.CancelledError, Exception):
pass
cancelled += 1
return cancelled
to_cancel.append(t)
if to_cancel:
await asyncio.gather(*to_cancel, return_exceptions=True)
return len(to_cancel)
def get_running_count(self) -> int:
"""Return the number of currently running subagents."""