diff --git a/backend/packages/harness/deerflow/subagents/capacity.py b/backend/packages/harness/deerflow/subagents/capacity.py index 4ed80541b..2f35242aa 100644 --- a/backend/packages/harness/deerflow/subagents/capacity.py +++ b/backend/packages/harness/deerflow/subagents/capacity.py @@ -105,13 +105,32 @@ class SubagentExecutionCapacity: async with self._lock: self._release_locked() + async def _release_cancellation_safe(self) -> None: + """Release an acquired slot before propagating repeated cancellation.""" + release_task = asyncio.create_task(self._release()) + cancellation: asyncio.CancelledError | None = None + while not release_task.done(): + try: + await asyncio.shield(release_task) + except asyncio.CancelledError as exc: + if cancellation is None: + cancellation = exc + + if cancellation is not None: + try: + release_task.result() + except Exception as exc: + raise cancellation from exc + raise cancellation + release_task.result() + @asynccontextmanager async def slot(self) -> AsyncIterator[None]: await self._acquire() try: yield finally: - await self._release() + await self._release_cancellation_safe() _config = SubagentRuntimeConfig() diff --git a/backend/tests/test_subagent_capacity.py b/backend/tests/test_subagent_capacity.py index 294d8806a..4e0253581 100644 --- a/backend/tests/test_subagent_capacity.py +++ b/backend/tests/test_subagent_capacity.py @@ -57,3 +57,43 @@ async def test_snapshot_reports_running_and_queued_waiters(): queued.cancel() holder.cancel() await asyncio.gather(queued, holder, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_repeated_cancellation_cannot_leak_running_slot(monkeypatch): + """A second cancellation during cleanup must not strand process capacity.""" + capacity = SubagentExecutionCapacity(SubagentRuntimeConfig(max_running=1, max_queued=1, queue_timeout_seconds=1)) + entered = asyncio.Event() + release_started = asyncio.Event() + never = asyncio.Event() + original_release = capacity._release + + async def observed_release(): + release_started.set() + await original_release() + + monkeypatch.setattr(capacity, "_release", observed_release) + + async def hold_slot(): + async with capacity.slot(): + entered.set() + await never.wait() + + holder = asyncio.create_task(hold_slot()) + await asyncio.wait_for(entered.wait(), timeout=1) + await capacity._lock.acquire() + try: + holder.cancel() + await asyncio.wait_for(release_started.wait(), timeout=1) + holder.cancel() + await asyncio.sleep(0) + finally: + capacity._lock.release() + + with pytest.raises(asyncio.CancelledError): + await holder + + assert capacity.snapshot().running == 0 + async with asyncio.timeout(0.5): + async with capacity.slot(): + pass