mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-07 21:08:34 +03:00
fix(gateway): make resource teardown cancellation-safe
This commit is contained in:
@@ -55,6 +55,62 @@ class TestHandleStop:
|
||||
out = await cmd_stop(ctx)
|
||||
assert "No active task" in out.content
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_close_mcp_cancels_active_turn_before_resources(self):
|
||||
loop, _bus = _make_loop()
|
||||
events: list[str] = []
|
||||
|
||||
async def active_turn():
|
||||
try:
|
||||
await asyncio.sleep(60)
|
||||
except asyncio.CancelledError:
|
||||
events.append("turn_cancelled")
|
||||
raise
|
||||
|
||||
task = asyncio.create_task(active_turn())
|
||||
await asyncio.sleep(0)
|
||||
loop._active_tasks["test:c1"] = {task}
|
||||
|
||||
async def close_subagents():
|
||||
events.append("resources_closed")
|
||||
|
||||
loop.subagents.close = close_subagents
|
||||
loop._exec_session_manager.close_all = AsyncMock()
|
||||
with patch("nanobot.agent.loop.agent_context.close_mcp", AsyncMock()):
|
||||
await loop.close_mcp()
|
||||
|
||||
assert events == ["turn_cancelled", "resources_closed"]
|
||||
assert task.cancelled()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_close_mcp_serializes_duplicate_cleanup(self):
|
||||
loop, _bus = _make_loop()
|
||||
entered = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
concurrent = 0
|
||||
max_concurrent = 0
|
||||
|
||||
async def close_subagents():
|
||||
nonlocal concurrent, max_concurrent
|
||||
concurrent += 1
|
||||
max_concurrent = max(max_concurrent, concurrent)
|
||||
entered.set()
|
||||
await release.wait()
|
||||
concurrent -= 1
|
||||
|
||||
loop.subagents.close = close_subagents
|
||||
loop._exec_session_manager.close_all = AsyncMock()
|
||||
with patch("nanobot.agent.loop.agent_context.close_mcp", AsyncMock()):
|
||||
first = asyncio.create_task(loop.close_mcp())
|
||||
await entered.wait()
|
||||
second = asyncio.create_task(loop.close_mcp())
|
||||
await asyncio.sleep(0)
|
||||
assert not second.done()
|
||||
release.set()
|
||||
await asyncio.gather(first, second)
|
||||
|
||||
assert max_concurrent == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stop_cancels_active_task(self):
|
||||
from nanobot.bus.events import InboundMessage
|
||||
|
||||
@@ -67,7 +67,7 @@ async def test_runtime_tasks_cancelled_before_resources_closed() -> None:
|
||||
task = asyncio.create_task(_cancellable_task(events))
|
||||
await asyncio.sleep(0) # let the task start (cancellation pre-start skips its body)
|
||||
|
||||
await _close_gateway_runtime(agent, channels, [task], None, False)
|
||||
await _close_gateway_runtime(agent, channels, [task], None)
|
||||
|
||||
assert events == ["cancelled", "close_mcp"] # cancel happens before close
|
||||
assert channels.stopped == 1
|
||||
@@ -86,7 +86,7 @@ async def test_pending_background_work_is_drained_before_close_returns() -> None
|
||||
|
||||
agent.background = asyncio.create_task(background_work())
|
||||
|
||||
await _close_gateway_runtime(agent, channels, [], None, False)
|
||||
await _close_gateway_runtime(agent, channels, [], None)
|
||||
|
||||
assert done["done"] is True
|
||||
assert agent.close_calls == 1
|
||||
@@ -98,14 +98,14 @@ async def test_stubborn_task_does_not_block_past_wait_timeout() -> None:
|
||||
events: list[str] = []
|
||||
task = asyncio.create_task(_stubborn_task(events))
|
||||
await asyncio.sleep(0) # let the task start (cancellation pre-start skips its body)
|
||||
runtime_tasks = asyncio.gather(task)
|
||||
|
||||
start = time.monotonic()
|
||||
await _close_gateway_runtime(
|
||||
agent,
|
||||
channels,
|
||||
[task],
|
||||
None,
|
||||
False,
|
||||
runtime_tasks,
|
||||
task_wait_timeout=0.05,
|
||||
)
|
||||
elapsed = time.monotonic() - start
|
||||
@@ -113,7 +113,8 @@ async def test_stubborn_task_does_not_block_past_wait_timeout() -> None:
|
||||
await asyncio.sleep(0) # let the swallowed cancellation handler run
|
||||
|
||||
assert "swallowed" in events # task was cancelled, then refused to die
|
||||
assert not task.done() # still running despite cancellation
|
||||
assert task.done() # the timed-out task received a second cancellation
|
||||
assert runtime_tasks.done()
|
||||
assert agent.close_calls == 1 # resources still closed underneath it
|
||||
assert elapsed < 1.0 # bounded, not held open by the stubborn task
|
||||
|
||||
@@ -124,7 +125,7 @@ async def test_hanging_close_is_bounded_and_does_not_raise() -> None:
|
||||
channels = _FakeChannels()
|
||||
|
||||
start = time.monotonic()
|
||||
await _close_gateway_runtime(agent, channels, [], None, False, close_timeout=0.05)
|
||||
await _close_gateway_runtime(agent, channels, [], None, close_timeout=0.05)
|
||||
elapsed = time.monotonic() - start
|
||||
|
||||
assert agent.close_calls == 1
|
||||
@@ -137,7 +138,7 @@ async def test_failing_close_is_logged_but_shutdown_proceeds() -> None:
|
||||
agent.raise_on_close = True
|
||||
channels = _FakeChannels()
|
||||
|
||||
await _close_gateway_runtime(agent, channels, [], None, False)
|
||||
await _close_gateway_runtime(agent, channels, [], None)
|
||||
|
||||
assert agent.close_calls == 1
|
||||
assert channels.stopped == 1 # teardown continued past the failure
|
||||
@@ -148,20 +149,20 @@ async def test_duplicate_cleanup_is_idempotent() -> None:
|
||||
channels = _FakeChannels()
|
||||
task = asyncio.create_task(_cancellable_task([]))
|
||||
|
||||
await _close_gateway_runtime(agent, channels, [task], None, False)
|
||||
await _close_gateway_runtime(agent, channels, [task], None, False)
|
||||
await _close_gateway_runtime(agent, channels, [task], None)
|
||||
await _close_gateway_runtime(agent, channels, [task], None)
|
||||
|
||||
assert agent.close_calls == 2 # second pass is a clean no-op
|
||||
assert channels.stopped == 2
|
||||
assert task.cancelled()
|
||||
|
||||
|
||||
async def test_runtime_tasks_gather_is_awaited_when_not_drained() -> None:
|
||||
async def test_finished_runtime_tasks_gather_is_retrieved() -> None:
|
||||
agent = _FakeAgent()
|
||||
channels = _FakeChannels()
|
||||
runtime_tasks = asyncio.gather(asyncio.sleep(0))
|
||||
|
||||
await _close_gateway_runtime(agent, channels, [], runtime_tasks, False)
|
||||
await _close_gateway_runtime(agent, channels, [], runtime_tasks)
|
||||
|
||||
assert runtime_tasks.done()
|
||||
assert agent.close_calls == 1
|
||||
@@ -173,7 +174,7 @@ async def test_cancelled_runtime_tasks_gather_does_not_raise() -> None:
|
||||
runtime_tasks = asyncio.gather(asyncio.sleep(3600))
|
||||
runtime_tasks.cancel()
|
||||
|
||||
await _close_gateway_runtime(agent, channels, [], runtime_tasks, False)
|
||||
await _close_gateway_runtime(agent, channels, [], runtime_tasks)
|
||||
|
||||
assert runtime_tasks.done() # the cancelled gather was awaited without raising
|
||||
assert agent.close_calls == 1
|
||||
|
||||
Reference in New Issue
Block a user