mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-10 14:28:38 +03:00
fix(agent): stop session-owned exec processes
This commit is contained in:
@@ -788,9 +788,9 @@ class AgentLoop:
|
|||||||
logger.warning("Command '{}' matched but dispatch returned None", raw)
|
logger.warning("Command '{}' matched but dispatch returned None", raw)
|
||||||
|
|
||||||
async def _cancel_active_tasks(self, key: str) -> int:
|
async def _cancel_active_tasks(self, key: str) -> int:
|
||||||
"""Cancel and await all active tasks and subagents for *key*.
|
"""Cancel and await all active work for *key*.
|
||||||
|
|
||||||
Returns the total number of cancelled tasks + subagents.
|
Returns the total number of cancelled tasks, subagents, and exec sessions.
|
||||||
"""
|
"""
|
||||||
tasks = tuple(self._active_tasks.pop(key, set()))
|
tasks = tuple(self._active_tasks.pop(key, set()))
|
||||||
cancelled = sum(1 for t in tasks if not t.done() and t.cancel())
|
cancelled = sum(1 for t in tasks if not t.done() and t.cancel())
|
||||||
@@ -798,7 +798,8 @@ class AgentLoop:
|
|||||||
with suppress(asyncio.CancelledError, Exception):
|
with suppress(asyncio.CancelledError, Exception):
|
||||||
await t
|
await t
|
||||||
sub_cancelled = await self.subagents.cancel_by_session(key)
|
sub_cancelled = await self.subagents.cancel_by_session(key)
|
||||||
return cancelled + sub_cancelled
|
exec_cancelled = await self._exec_session_manager.terminate_by_owner(key)
|
||||||
|
return cancelled + sub_cancelled + exec_cancelled
|
||||||
|
|
||||||
async def discard_session(self, key: str) -> None:
|
async def discard_session(self, key: str) -> None:
|
||||||
"""Stop active work for *key* and forget its cached session."""
|
"""Stop active work for *key* and forget its cached session."""
|
||||||
|
|||||||
@@ -119,12 +119,18 @@ async def test_session_discard_control_cancels_active_turn(tmp_path, monkeypatch
|
|||||||
loop = _loop(tmp_path, [])
|
loop = _loop(tmp_path, [])
|
||||||
|
|
||||||
async def wait_for_discard(key: str) -> None:
|
async def wait_for_discard(key: str) -> None:
|
||||||
while loop.sessions.get_cached(key) is not None:
|
while loop.sessions.get_cached(key) is not None or key in loop._discarding_sessions:
|
||||||
await asyncio.sleep(0)
|
await asyncio.sleep(0)
|
||||||
|
|
||||||
loop.provider.chat_with_retry = AsyncMock(side_effect=block_provider)
|
loop.provider.chat_with_retry = AsyncMock(side_effect=block_provider)
|
||||||
monkeypatch.setattr(loop, "_connect_mcp", AsyncMock())
|
monkeypatch.setattr(loop, "_connect_mcp", AsyncMock())
|
||||||
monkeypatch.setattr(loop, "close_mcp", AsyncMock())
|
monkeypatch.setattr(loop, "close_mcp", AsyncMock())
|
||||||
|
terminate_exec_sessions = AsyncMock(return_value=1)
|
||||||
|
monkeypatch.setattr(
|
||||||
|
loop._exec_session_manager,
|
||||||
|
"terminate_by_owner",
|
||||||
|
terminate_exec_sessions,
|
||||||
|
)
|
||||||
key = "websocket:transient-cancelled"
|
key = "websocket:transient-cancelled"
|
||||||
loop.sessions.get_or_create_transient(
|
loop.sessions.get_or_create_transient(
|
||||||
key,
|
key,
|
||||||
@@ -152,6 +158,7 @@ async def test_session_discard_control_cancels_active_turn(tmp_path, monkeypatch
|
|||||||
await asyncio.wait_for(active_task, timeout=2)
|
await asyncio.wait_for(active_task, timeout=2)
|
||||||
await asyncio.wait_for(wait_for_discard(key), timeout=2)
|
await asyncio.wait_for(wait_for_discard(key), timeout=2)
|
||||||
assert loop.sessions.get_cached(key) is None
|
assert loop.sessions.get_cached(key) is None
|
||||||
|
terminate_exec_sessions.assert_awaited_once_with(key)
|
||||||
|
|
||||||
loop.stop()
|
loop.stop()
|
||||||
await loop.bus.publish_inbound(_message(key, "wake"))
|
await loop.bus.publish_inbound(_message(key, "wake"))
|
||||||
|
|||||||
Reference in New Issue
Block a user