mirror/nanobot
1
0
Fork 0
mirror of https://github.com/HKUDS/nanobot.git synced 2026-08-19 02:26:12 +03:00
Code Issues Packages Projects Releases Wiki Activity

refactor: move MCP lifecycle out of AgentLoop (#5343)

Browse Source
This commit is contained in:
chengyongru
2026-08-12 17:51:04 +08:00
committed by GitHub
parent 686dd0603e
commit 19997d20bb
39 changed files with 1192 additions and 846 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/agent/conftest.py
-3
View File
@@ -46,7 +46,6 @@ def make_loop(
context_window_tokens: int = 128_000,
session_ttl_minutes: int = 0,
unified_session: bool = False,
mcp_servers: dict | None = None,
tools_config=None,
model_presets: dict | None = None,
hooks: list | None = None,
@@ -72,8 +71,6 @@ def make_loop(
session_ttl_minutes=session_ttl_minutes,
unified_session=unified_session,
)
if mcp_servers is not None:
kwargs["mcp_servers"] = mcp_servers
if tools_config is not None:
kwargs["tools_config"] = tools_config
if model_presets is not None:
tests/agent/test_auto_compact.py
+46 -41
View File
@@ -8,6 +8,7 @@ from unittest.mock import AsyncMock, MagicMock
import pytest
from nanobot.agent.loop import AgentLoop
from nanobot.agent.tools.registry import ToolRegistry
from nanobot.bus.events import InboundMessage
from nanobot.bus.queue import MessageBus
from nanobot.command import CommandContext
@@ -193,7 +194,11 @@ class TestIdleScanThrottling:
})
provider = MagicMock()
provider.get_default_model.return_value = "test-model"
loop = AgentLoop.from_config(config, provider=provider)
loop = AgentLoop.from_config(
config,
tool_registry=ToolRegistry(),
provider=provider,
)
loop.auto_compact.check_expired = MagicMock()
loop._check_expired_sessions_if_due()
@@ -310,7 +315,7 @@ class TestAutoCompact:
assert loop.auto_compact._is_expired(ts) is True
ts2 = datetime.now() - timedelta(minutes=14, seconds=59)
assert loop.auto_compact._is_expired(ts2) is False
await loop.close_mcp()
await loop.aclose()
@pytest.mark.asyncio
async def test_is_expired_string_timestamp(self, tmp_path):
@@ -320,7 +325,7 @@ class TestAutoCompact:
assert loop.auto_compact._is_expired(ts) is True
assert loop.auto_compact._is_expired(None) is False
assert loop.auto_compact._is_expired("") is False
await loop.close_mcp()
await loop.aclose()
@pytest.mark.asyncio
async def test_check_expired_only_archives_expired_sessions(self, tmp_path):
@@ -343,7 +348,7 @@ class TestAutoCompact:
active_after = loop.sessions.get_or_create("cli:active")
assert len(active_after.messages) == 1
assert active_after.messages[0]["content"] == "recent"
await loop.close_mcp()
await loop.aclose()
@pytest.mark.asyncio
async def test_auto_compact_archives_full_tail_without_deleting_history(self, tmp_path):
@@ -367,7 +372,7 @@ class TestAutoCompact:
assert len(visible) == loop.auto_compact._RECENT_SUFFIX_MESSAGES
assert visible[0]["content"] == "msg user 2"
assert visible[-1]["content"] == "msg assistant 5"
await loop.close_mcp()
await loop.aclose()
@pytest.mark.asyncio
async def test_auto_compact_extends_recent_suffix_to_user_turn(self, tmp_path):
@@ -398,7 +403,7 @@ class TestAutoCompact:
for m in visible
for tc in (m.get("tool_calls") or [])
)
await loop.close_mcp()
await loop.aclose()
@pytest.mark.asyncio
async def test_auto_compact_stores_summary(self, tmp_path):
@@ -422,7 +427,7 @@ class TestAutoCompact:
assert len(session_after.get_history(max_messages=12)) == (
loop.auto_compact._RECENT_SUFFIX_MESSAGES
)
await loop.close_mcp()
await loop.aclose()
@pytest.mark.asyncio
async def test_auto_compact_empty_session(self, tmp_path):
@@ -436,7 +441,7 @@ class TestAutoCompact:
session_after = loop.sessions.get_or_create("cli:test")
assert len(session_after.messages) == 0
assert "cli:test" not in loop.auto_compact._summaries
await loop.close_mcp()
await loop.aclose()
@pytest.mark.asyncio
async def test_auto_compact_respects_last_consolidated(self, tmp_path):
@@ -455,7 +460,7 @@ class TestAutoCompact:
await loop.auto_compact._archive("cli:test", runtime=loop.llm_runtime())
assert len(archived_messages) == 10
await loop.close_mcp()
await loop.aclose()
class TestAutoCompactIdleDetection:
@@ -474,7 +479,7 @@ class TestAutoCompactIdleDetection:
session_after = loop.sessions.get_or_create("cli:test")
assert any(m["content"] == "old message" for m in session_after.messages)
await loop.close_mcp()
await loop.aclose()
@pytest.mark.asyncio
async def test_auto_compact_triggers_on_idle(self, tmp_path):
@@ -503,7 +508,7 @@ class TestAutoCompactIdleDetection:
for m in session_after.get_history(max_messages=len(session_after.messages))
)
assert any(m["content"] == "new msg" for m in session_after.messages)
await loop.close_mcp()
await loop.aclose()
@pytest.mark.asyncio
async def test_no_auto_compact_when_active(self, tmp_path):
@@ -517,7 +522,7 @@ class TestAutoCompactIdleDetection:
session_after = loop.sessions.get_or_create("cli:test")
assert any(m["content"] == "recent message" for m in session_after.messages)
await loop.close_mcp()
await loop.aclose()
@pytest.mark.asyncio
async def test_auto_compact_does_not_affect_priority_commands(self, tmp_path):
@@ -540,7 +545,7 @@ class TestAutoCompactIdleDetection:
# Session should be untouched since priority commands skip _process_message
session_after = loop.sessions.get_or_create("cli:test")
assert any(m["content"] == "old message" for m in session_after.messages)
await loop.close_mcp()
await loop.aclose()
@pytest.mark.asyncio
async def test_auto_compact_with_slash_new(self, tmp_path):
@@ -562,7 +567,7 @@ class TestAutoCompactIdleDetection:
session_after = loop.sessions.get_or_create("cli:test")
assert len(session_after.messages) == 0
await loop.close_mcp()
await loop.aclose()
@pytest.mark.asyncio
async def test_shortcut_command_persisted_with_command_flag(self, tmp_path):
@@ -581,7 +586,7 @@ class TestAutoCompactIdleDetection:
assert session_after.messages[1]["role"] == "assistant"
assert session_after.messages[1].get("_command") is True
assert AgentLoop._PENDING_USER_TURN_KEY not in session_after.metadata
await loop.close_mcp()
await loop.aclose()
@pytest.mark.asyncio
async def test_shortcut_command_excluded_from_get_history(self, tmp_path):
@@ -597,7 +602,7 @@ class TestAutoCompactIdleDetection:
assert len(history) == 2
assert all(m["content"] != "/help" for m in history)
assert all(m["content"] != "help text" for m in history)
await loop.close_mcp()
await loop.aclose()
class TestAutoCompactSystemMessages:
@@ -628,7 +633,7 @@ class TestAutoCompactSystemMessages:
m["content"] == "old user 0"
for m in session_after.get_history(max_messages=len(session_after.messages))
)
await loop.close_mcp()
await loop.aclose()
class TestAutoCompactEdgeCases:
@@ -656,7 +661,7 @@ class TestAutoCompactEdgeCases:
# "(nothing)" summary should not be stored
assert "cli:test" not in loop.auto_compact._summaries
await loop.close_mcp()
await loop.aclose()
@pytest.mark.asyncio
async def test_auto_compact_archive_failure_preserves_raw_history(self, tmp_path):
@@ -677,7 +682,7 @@ class TestAutoCompactEdgeCases:
loop.auto_compact._RECENT_SUFFIX_MESSAGES
)
await loop.close_mcp()
await loop.aclose()
@pytest.mark.asyncio
async def test_auto_compact_preserves_runtime_checkpoint_before_check(self, tmp_path):
@@ -709,7 +714,7 @@ class TestAutoCompactEdgeCases:
assert any(m["content"] == "previous message" for m in session_after.messages)
assert any(m["content"] == "interrupted response" for m in session_after.messages)
await loop.close_mcp()
await loop.aclose()
class TestAutoCompactIntegration:
@@ -779,7 +784,7 @@ class TestAutoCompactIntegration:
# The new message should be processed (response exists)
assert response is not None
await loop.close_mcp()
await loop.aclose()
@pytest.mark.asyncio
async def test_runtime_context_markers_not_persisted_for_multi_paragraph_turn(self, tmp_path):
@@ -807,7 +812,7 @@ class TestAutoCompactIntegration:
content = str(persisted.get("content", ""))
assert "[Runtime Context" not in content
assert "[/Runtime Context]" not in content
await loop.close_mcp()
await loop.aclose()
class TestProactiveAutoCompact:
@@ -870,7 +875,7 @@ class TestProactiveAutoCompact:
session_after = loop.sessions.get_or_create("cli:test")
assert len(session_after.messages) == 1
await loop.close_mcp()
await loop.aclose()
@pytest.mark.asyncio
async def test_proactive_archive_on_idle_tick(self, tmp_path):
@@ -897,7 +902,7 @@ class TestProactiveAutoCompact:
entry = loop.auto_compact._summaries.get("cli:test")
assert entry is not None
assert entry[0] == "User chatted about old things."
await loop.close_mcp()
await loop.aclose()
@pytest.mark.asyncio
async def test_proactive_archive_skips_dream_sessions(self, tmp_path):
@@ -918,7 +923,7 @@ class TestProactiveAutoCompact:
assert _fake_compact.state["count"] == 0
assert "dream:20260602-155256" not in loop.auto_compact._archiving
assert "dream:20260602-155256" not in loop.auto_compact._summaries
await loop.close_mcp()
await loop.aclose()
@pytest.mark.asyncio
async def test_no_proactive_archive_when_active(self, tmp_path):
@@ -932,7 +937,7 @@ class TestProactiveAutoCompact:
session_after = loop.sessions.get_or_create("cli:test")
assert len(session_after.messages) == 1
await loop.close_mcp()
await loop.aclose()
@pytest.mark.asyncio
async def test_no_duplicate_archive(self, tmp_path):
@@ -968,7 +973,7 @@ class TestProactiveAutoCompact:
# Clean up
block_forever.set()
await _drain_background_tasks(loop)
await loop.close_mcp()
await loop.aclose()
@pytest.mark.asyncio
async def test_proactive_archive_error_does_not_block(self, tmp_path):
@@ -989,7 +994,7 @@ class TestProactiveAutoCompact:
# Key should be removed from _archiving (finally block)
assert "cli:test" not in loop.auto_compact._archiving
await loop.close_mcp()
await loop.aclose()
@pytest.mark.asyncio
async def test_proactive_archive_skips_empty_sessions(self, tmp_path):
@@ -1005,7 +1010,7 @@ class TestProactiveAutoCompact:
# Empty session should not produce a summary
assert "cli:test" not in loop.auto_compact._summaries
await loop.close_mcp()
await loop.aclose()
@pytest.mark.asyncio
async def test_skip_expired_session_with_active_agent_task(self, tmp_path):
@@ -1026,7 +1031,7 @@ class TestProactiveAutoCompact:
session_after = loop.sessions.get_or_create("cli:test")
assert len(session_after.messages) == 12 # All messages preserved
await loop.close_mcp()
await loop.aclose()
@pytest.mark.asyncio
async def test_archive_after_active_task_completes(self, tmp_path):
@@ -1047,7 +1052,7 @@ class TestProactiveAutoCompact:
# Second tick: task completed, should archive
await self._run_check_expired(loop)
assert _fake_compact.state["count"] == 1
await loop.close_mcp()
await loop.aclose()
@pytest.mark.asyncio
async def test_partial_active_set_only_archives_inactive_expired(self, tmp_path):
@@ -1083,7 +1088,7 @@ class TestProactiveAutoCompact:
assert len(s2_after.messages) == 12 # Preserved
s3_after = loop.sessions.get_or_create("cli:recent")
assert len(s3_after.messages) == 1 # Preserved
await loop.close_mcp()
await loop.aclose()
@pytest.mark.asyncio
async def test_no_reschedule_after_successful_archive(self, tmp_path):
@@ -1104,7 +1109,7 @@ class TestProactiveAutoCompact:
# Second tick: should NOT re-schedule because the session has no removable tail.
await self._run_check_expired(loop)
assert _fake_compact.state["count"] == 1 # Still 1, not re-scheduled
await loop.close_mcp()
await loop.aclose()
@pytest.mark.asyncio
async def test_empty_session_does_not_schedule_idle_compact(self, tmp_path):
@@ -1124,7 +1129,7 @@ class TestProactiveAutoCompact:
await self._run_check_expired(loop)
assert _fake_compact.state["count"] == 0
assert "cli:test" not in loop.auto_compact._summaries
await loop.close_mcp()
await loop.aclose()
@pytest.mark.asyncio
async def test_session_can_be_compacted_again_after_new_messages(self, tmp_path):
@@ -1155,7 +1160,7 @@ class TestProactiveAutoCompact:
# Second compact cycle should succeed
await loop.auto_compact._archive("cli:test", runtime=loop.llm_runtime())
assert _fake_compact.state["count"] == 2
await loop.close_mcp()
await loop.aclose()
class TestSummaryPersistence:
@@ -1182,7 +1187,7 @@ class TestSummaryPersistence:
assert meta is not None
assert meta["text"] == "User said hello."
assert "last_active" in meta
await loop.close_mcp()
await loop.aclose()
@pytest.mark.asyncio
async def test_summary_recovered_after_restart(self, tmp_path):
@@ -1218,7 +1223,7 @@ class TestSummaryPersistence:
assert "Previous conversation summary" in summary
# _last_summary persists in metadata for restart survival.
assert "_last_summary" in reloaded.metadata
await loop.close_mcp()
await loop.aclose()
@pytest.mark.asyncio
async def test_metadata_persists_for_restart(self, tmp_path):
@@ -1246,7 +1251,7 @@ class TestSummaryPersistence:
assert "Summary." in summary2
# _last_summary persists in metadata for restart survival.
assert "_last_summary" in reloaded.metadata
await loop.close_mcp()
await loop.aclose()
@pytest.mark.asyncio
async def test_metadata_cleanup_on_inmemory_path(self, tmp_path):
@@ -1272,7 +1277,7 @@ class TestSummaryPersistence:
assert summary is not None
# _last_summary persists in metadata for restart survival.
assert "_last_summary" in reloaded.metadata
await loop.close_mcp()
await loop.aclose()
@pytest.mark.asyncio
async def test_new_summary_overrides_old(self, tmp_path):
@@ -1314,7 +1319,7 @@ class TestSummaryPersistence:
_, summary2 = loop.auto_compact.prepare_session(reloaded, "cli:test")
assert summary2 is not None
assert "Second summary." in summary2
await loop.close_mcp()
await loop.aclose()
@pytest.mark.asyncio
async def test_new_command_clears_last_summary(self, tmp_path):
@@ -1342,4 +1347,4 @@ class TestSummaryPersistence:
# After /new, metadata should no longer contain _last_summary
fresh = loop.sessions.get_or_create("cli:test")
assert "_last_summary" not in fresh.metadata
await loop.close_mcp()
await loop.aclose()
tests/agent/test_consolidate_offset.py
+5 -5
View File
@@ -538,7 +538,7 @@ class TestNewCommandArchival:
session_after = loop.sessions.get_or_create("cli:test")
assert len(session_after.messages) == 0
await loop.close_mcp()
await loop.aclose()
assert call_count == 1
@pytest.mark.asyncio
@@ -572,7 +572,7 @@ class TestNewCommandArchival:
assert response is not None
assert "new session started" in response.content.lower()
await loop.close_mcp()
await loop.aclose()
assert archived_count == 3
assert archived_session_key == "cli:test"
@@ -603,8 +603,8 @@ class TestNewCommandArchival:
assert loop.sessions.get_or_create("cli:test").messages == []
@pytest.mark.asyncio
async def test_close_mcp_drains_background_tasks(self, tmp_path: Path) -> None:
"""close_mcp waits for background tasks to complete."""
async def test_aclose_drains_background_tasks(self, tmp_path: Path) -> None:
"""aclose waits for background tasks to complete."""
from nanobot.bus.events import InboundMessage
loop = self._make_loop(tmp_path)
@@ -632,5 +632,5 @@ class TestNewCommandArchival:
assert not archived.is_set()
release_archive.set()
await loop.close_mcp()
await loop.aclose()
assert archived.is_set()
tests/agent/test_loop_direct_websocket_status.py
-1
View File
@@ -93,7 +93,6 @@ async def test_process_direct_websocket_clears_run_status(tmp_path) -> None:
@pytest.mark.asyncio
async def test_process_direct_reuses_existing_session_lock(tmp_path) -> None:
loop = _make_loop(tmp_path)
loop._connect_mcp = AsyncMock()
session_key = "api:fixed"
lock = loop._session_locks.setdefault(session_key, asyncio.Lock())
await lock.acquire()
tests/agent/test_loop_save_turn.py
-3
View File
@@ -1519,13 +1519,11 @@ async def test_run_agent_loop_goal_continue_message_reads_latest_metadata(
@pytest.mark.asyncio
async def test_process_direct_rejects_reserved_system_channel(tmp_path: Path) -> None:
loop = _make_full_loop(tmp_path)
loop._connect_mcp = AsyncMock() # type: ignore[method-assign]
loop._process_message = AsyncMock(return_value=None) # type: ignore[method-assign]
with pytest.raises(ValueError, match="reserved for internal messages"):
await loop.process_direct("external input", channel="system")
loop._connect_mcp.assert_not_awaited()
loop._process_message.assert_not_awaited()
@@ -1534,7 +1532,6 @@ async def test_process_direct_skip_user_persist_does_not_save_retry_user(
tmp_path: Path,
) -> None:
loop = _make_full_loop(tmp_path)
loop._connect_mcp = AsyncMock()
session = loop.sessions.get_or_create("api:default")
session.add_message("user", "hello")
session.add_message("assistant", "previous empty-response attempt")
tests/agent/test_loop_session_policy.py
+1 -2
View File
@@ -123,8 +123,7 @@ async def test_session_discard_control_cancels_active_turn(tmp_path, monkeypatch
await asyncio.sleep(0)
loop.provider.chat_with_retry = AsyncMock(side_effect=block_provider)
monkeypatch.setattr(loop, "_connect_mcp", AsyncMock())
monkeypatch.setattr(loop, "close_mcp", AsyncMock())
monkeypatch.setattr(loop, "aclose", AsyncMock())
terminate_exec_sessions = AsyncMock(return_value=1)
monkeypatch.setattr(
loop._exec_session_manager,
tests/agent/test_loop_tool_context.py
+48
View File
@@ -1,4 +1,5 @@
import asyncio
import inspect
from pathlib import Path
from unittest.mock import AsyncMock, MagicMock
@@ -11,8 +12,10 @@ from nanobot.agent.tools.context import (
current_request_context,
reset_request_context,
)
from nanobot.agent.tools.registry import ToolRegistry
from nanobot.bus.events import InboundMessage
from nanobot.bus.queue import MessageBus
from nanobot.config.schema import Config
from nanobot.providers.base import LLMResponse, ToolCallRequest
from nanobot.session.turn_continuation import INTERNAL_CONTINUATION_META
@@ -56,6 +59,51 @@ class _Tools:
return (self.tool, arguments, None) if name == "cron" else (None, arguments, None)
def test_loop_registers_default_tools_in_injected_registry(tmp_path: Path) -> None:
provider = MagicMock()
provider.get_default_model.return_value = "test-model"
registry = ToolRegistry()
loop = AgentLoop(
bus=MessageBus(),
provider=provider,
workspace=tmp_path,
tool_registry=registry,
)
assert loop.tools is registry
assert registry.has("read_file")
def _config_for_loop(tmp_path: Path) -> Config:
return Config.model_validate({"agents": {"defaults": {"workspace": str(tmp_path)}}})
def _provider_for_loop() -> MagicMock:
provider = MagicMock()
provider.get_default_model.return_value = "test-model"
return provider
def test_loop_from_config_requires_caller_owned_registry(tmp_path: Path) -> None:
signature = inspect.signature(AgentLoop.from_config)
with pytest.raises(TypeError, match="tool_registry"):
signature.bind(_config_for_loop(tmp_path))
def test_loop_from_config_uses_caller_owned_registry(tmp_path: Path) -> None:
registry = ToolRegistry()
loop = AgentLoop.from_config(
_config_for_loop(tmp_path),
tool_registry=registry,
provider=_provider_for_loop(),
)
assert loop.tools is registry
assert loop.tools.has("read_file")
@pytest.mark.asyncio
async def test_loop_binds_request_context_for_tool_execution(tmp_path: Path) -> None:
provider = MagicMock()
tests/agent/test_mcp_connection.py
+165 -97
View File
@@ -1,4 +1,4 @@
"""Tests for MCP connection lifecycle in AgentLoop."""
"""Tests for the application-owned MCP provider lifecycle."""
from __future__ import annotations
@@ -6,7 +6,7 @@ import asyncio
from contextlib import AsyncExitStack
from types import SimpleNamespace
from typing import Any
from unittest.mock import AsyncMock, MagicMock
from unittest.mock import AsyncMock
import anyio
import pytest
@@ -15,11 +15,10 @@ from mcp.shared.exceptions import McpError
from mcp.shared.message import SessionMessage
from mcp.types import ErrorData
from nanobot.agent.loop import AgentLoop
from nanobot.agent.tools import mcp as mcp_runtime
from nanobot.agent.tools.base import Tool
from nanobot.agent.tools.mcp import MCPResourceWrapper, MCPToolWrapper
from nanobot.bus.queue import MessageBus
from nanobot.agent.tools.mcp import MCPProvider, MCPResourceWrapper, MCPToolWrapper
from nanobot.agent.tools.registry import ToolRegistry
from nanobot.config.loader import load_config, save_config
from nanobot.config.schema import MCPServerConfig
@@ -74,18 +73,20 @@ class _FakeMcpTool(Tool):
return "ok"
def _make_loop(tmp_path, *, mcp_servers: dict | None = None) -> AgentLoop:
bus = MessageBus()
provider = MagicMock()
provider.get_default_model.return_value = "test-model"
provider.generation.max_tokens = 4096
return AgentLoop(
bus=bus,
provider=provider,
workspace=tmp_path,
model="test-model",
mcp_servers=mcp_servers or {"test": object()},
def _stdio_server(command: str = "test-mcp") -> MCPServerConfig:
return MCPServerConfig(type="stdio", command=command)
def _make_provider(
*,
mcp_servers: dict[str, MCPServerConfig] | None = None,
) -> tuple[MCPProvider, ToolRegistry]:
registry = ToolRegistry()
provider = MCPProvider(
mcp_servers if mcp_servers is not None else {"test": _stdio_server()},
registry,
)
return provider, registry
@pytest.mark.asyncio
@@ -140,7 +141,7 @@ async def test_owned_mcp_connection_closes_from_its_owner_task():
@pytest.mark.asyncio
async def test_connect_mcp_retries_when_no_servers_connect(tmp_path, monkeypatch: pytest.MonkeyPatch):
loop = _make_loop(tmp_path)
provider, _registry = _make_provider()
attempts = 0
async def _fake_connect(_servers, _registry):
@@ -150,12 +151,12 @@ async def test_connect_mcp_retries_when_no_servers_connect(tmp_path, monkeypatch
monkeypatch.setattr("nanobot.agent.tools.mcp.connect_mcp_servers", _fake_connect)
await loop._connect_mcp()
await loop._connect_mcp()
await provider.connect()
await provider.connect()
assert attempts == 2
assert loop._mcp_stacks == {}
assert loop.mcp_runtime_status() == {"test": "failed"}
assert provider.connected_server_names == set()
assert provider.runtime_status() == {"test": "failed"}
@pytest.mark.asyncio
@@ -168,7 +169,7 @@ async def test_connect_mcp_does_not_report_failure_before_oauth_authorization(
auth="oauth",
url="https://mcp.example.com/mcp",
)
loop = _make_loop(tmp_path, mcp_servers={"oauth-app": cfg})
provider, _registry = _make_provider(mcp_servers={"oauth-app": cfg})
connect = AsyncMock()
monkeypatch.setattr("nanobot.agent.tools.mcp.connect_mcp_servers", connect)
monkeypatch.setattr(
@@ -176,19 +177,20 @@ async def test_connect_mcp_does_not_report_failure_before_oauth_authorization(
lambda _name, _url: False,
)
await loop._connect_mcp()
await provider.connect()
connect.assert_not_awaited()
assert loop.mcp_runtime_status() == {}
assert provider.runtime_status() == {}
@pytest.mark.asyncio
async def test_agent_loop_run_closes_mcp_from_connection_owner_task(
async def test_mcp_provider_closes_connections_independently_from_agent_loop(
tmp_path,
monkeypatch: pytest.MonkeyPatch,
):
loop = _make_loop(tmp_path, mcp_servers={"playwright": object()})
connected = asyncio.Event()
provider, registry = _make_provider(
mcp_servers={"playwright": _stdio_server("playwright")}
)
owner_tasks: list[asyncio.Task | None] = []
closed_tasks: list[asyncio.Task | None] = []
@@ -203,40 +205,38 @@ async def test_agent_loop_run_closes_mcp_from_connection_owner_task(
async def _fake_connect(servers, _registry):
stacks = {name: _OwnerCheckedStack() for name in servers}
connected.set()
return stacks
monkeypatch.setattr("nanobot.agent.tools.mcp.connect_mcp_servers", _fake_connect)
task = asyncio.create_task(loop.run())
await asyncio.wait_for(connected.wait(), timeout=1)
loop.stop()
task.cancel()
await asyncio.gather(task, return_exceptions=True)
await provider.connect()
registry.register(_FakeMcpTool("mcp_playwright_search"))
await provider.aclose()
assert owner_tasks
assert closed_tasks == owner_tasks
assert loop._mcp_stacks == {}
assert provider.connected_server_names == set()
assert registry.get("mcp_playwright_search") is None
@pytest.mark.asyncio
async def test_close_server_ignores_server_cancelled_error(tmp_path):
loop = _make_loop(tmp_path)
provider, _registry = _make_provider()
class _ServerCancelledStack:
async def aclose(self) -> None:
raise asyncio.CancelledError()
loop._mcp_stacks = {"test": _ServerCancelledStack()}
provider._connections = {"test": _ServerCancelledStack()}
await mcp_runtime._close_server(loop, "test")
await provider._close_server("test")
assert loop._mcp_stacks == {}
assert provider.connected_server_names == set()
@pytest.mark.asyncio
async def test_close_mcp_servers_continues_after_server_cancelled_error(tmp_path):
loop = _make_loop(tmp_path)
async def test_provider_close_continues_after_server_cancelled_error(tmp_path):
provider, _registry = _make_provider()
closed: list[str] = []
class _ServerCancelledStack:
@@ -247,21 +247,53 @@ async def test_close_mcp_servers_continues_after_server_cancelled_error(tmp_path
async def aclose(self) -> None:
closed.append("second")
loop._mcp_stacks = {
provider._connections = {
"first": _ServerCancelledStack(),
"second": _TrackedStack(),
}
await mcp_runtime.close_mcp_servers(loop)
await provider.aclose()
assert closed == ["second"]
assert loop._mcp_stacks == {}
assert provider.connected_server_names == set()
@pytest.mark.asyncio
async def test_provider_close_finishes_other_connections_before_propagating_cancellation(
tmp_path,
):
provider, _registry = _make_provider()
started = asyncio.Event()
closed: list[str] = []
class _BlockingStack:
async def aclose(self) -> None:
started.set()
await asyncio.Event().wait()
class _TrackedStack:
async def aclose(self) -> None:
closed.append("second")
provider._connections = {
"first": _BlockingStack(),
"second": _TrackedStack(),
}
task = asyncio.create_task(provider.aclose())
await asyncio.wait_for(started.wait(), timeout=1)
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
assert closed == ["second"]
assert provider.connected_server_names == set()
@pytest.mark.asyncio
@pytest.mark.parametrize("close_all", [False, True], ids=["single", "all"])
async def test_mcp_cleanup_re_raises_external_cancellation(tmp_path, close_all: bool):
loop = _make_loop(tmp_path)
provider, _registry = _make_provider()
started = asyncio.Event()
class _BlockingStack:
@@ -269,12 +301,12 @@ async def test_mcp_cleanup_re_raises_external_cancellation(tmp_path, close_all:
started.set()
await asyncio.Event().wait()
loop._mcp_stacks = {"test": _BlockingStack()}
provider._connections = {"test": _BlockingStack()}
if close_all:
task = asyncio.create_task(mcp_runtime.close_mcp_servers(loop))
task = asyncio.create_task(provider.aclose())
else:
task = asyncio.create_task(mcp_runtime._close_server(loop, "test"))
task = asyncio.create_task(provider._close_server("test"))
await asyncio.wait_for(started.wait(), timeout=1)
task.cancel()
@@ -312,41 +344,38 @@ async def test_reload_mcp_servers_adds_and_removes_tools_without_restart(
return stacks
monkeypatch.setattr("nanobot.agent.tools.mcp.connect_mcp_servers", _fake_connect)
loop = _make_loop(tmp_path, mcp_servers={})
provider, registry = _make_provider(mcp_servers={})
added = await mcp_runtime.reload_servers(loop, loop.tools)
added = await provider.reload()
assert added["ok"] is True
assert added["added"] == ["browserbase"]
assert loop.tools.has("mcp_browserbase_navigate")
assert "browserbase" in loop._mcp_stacks
assert registry.has("mcp_browserbase_navigate")
assert provider.connected_server_names == {"browserbase"}
config = load_config()
del config.tools.mcp_servers["browserbase"]
save_config(config)
removed = await mcp_runtime.reload_servers(loop, loop.tools)
removed = await provider.reload()
assert removed["ok"] is True
assert removed["removed"] == ["browserbase"]
assert not loop.tools.has("mcp_browserbase_navigate")
assert "browserbase" not in loop._mcp_stacks
assert not registry.has("mcp_browserbase_navigate")
assert provider.connected_server_names == set()
assert closed == ["browserbase"]
@pytest.mark.asyncio
async def test_request_mcp_reload_reaches_runtime_control_without_restart(
async def test_reload_is_a_direct_provider_operation_without_an_agent_loop(
tmp_path,
monkeypatch: pytest.MonkeyPatch,
):
config_path = tmp_path / "config.json"
monkeypatch.setattr("nanobot.config.loader._current_config_path", config_path)
config = load_config()
config.tools.mcp_servers["browserbase"] = MCPServerConfig(
browserbase = MCPServerConfig(
type="stdio",
command="browserbase-mcp",
)
save_config(config)
configured: dict[str, MCPServerConfig] = {"browserbase": browserbase}
closed: list[str] = []
@@ -364,37 +393,68 @@ async def test_request_mcp_reload_reaches_runtime_control_without_restart(
return stacks
monkeypatch.setattr("nanobot.agent.tools.mcp.connect_mcp_servers", _fake_connect)
loop = _make_loop(tmp_path, mcp_servers={})
registry = ToolRegistry()
provider = MCPProvider({}, registry, server_loader=lambda: configured)
async def _handle_one_runtime_control() -> None:
msg = await loop.bus.consume_inbound()
handled = await mcp_runtime.handle_runtime_control(loop, msg, loop.tools)
assert handled is True
consumer = asyncio.create_task(_handle_one_runtime_control())
result = await mcp_runtime.request_mcp_reload(loop.bus, timeout=2.0)
await consumer
result = await provider.reload()
assert result["ok"] is True
assert result["added"] == ["browserbase"]
assert result["requires_restart"] is False
assert loop.tools.has("mcp_browserbase_navigate")
assert registry.has("mcp_browserbase_navigate")
config = load_config()
del config.tools.mcp_servers["browserbase"]
save_config(config)
configured = {}
consumer = asyncio.create_task(_handle_one_runtime_control())
result = await mcp_runtime.request_mcp_reload(loop.bus, timeout=2.0)
await consumer
result = await provider.reload()
assert result["ok"] is True
assert result["removed"] == ["browserbase"]
assert result["requires_restart"] is False
assert not loop.tools.has("mcp_browserbase_navigate")
assert not registry.has("mcp_browserbase_navigate")
assert closed == ["browserbase"]
@pytest.mark.asyncio
async def test_reload_timeout_marks_attempted_server_failed_and_allows_retry(
monkeypatch: pytest.MonkeyPatch,
):
server = _stdio_server("slow-mcp")
started = asyncio.Event()
attempts = 0
async def _fake_connect(servers, _registry):
nonlocal attempts
attempts += 1
if attempts == 1:
started.set()
await asyncio.Event().wait()
stack = AsyncExitStack()
await stack.__aenter__()
return {name: stack for name in servers}
monkeypatch.setattr("nanobot.agent.tools.mcp.connect_mcp_servers", _fake_connect)
provider = MCPProvider(
{"test": server},
ToolRegistry(),
server_loader=lambda: {"test": server},
)
reload_task = asyncio.create_task(provider.reload())
await asyncio.wait_for(started.wait(), timeout=1.0)
with pytest.raises(asyncio.TimeoutError):
await asyncio.wait_for(reload_task, timeout=0.01)
assert provider.connected_server_names == set()
assert provider.runtime_status() == {"test": "failed"}
result = await provider.reload()
assert result["ok"] is True
assert provider.connected_server_names == {"test"}
assert provider.runtime_status() == {"test": "connected"}
await provider.aclose()
@pytest.mark.asyncio
async def test_reload_mcp_servers_retries_configured_server_without_live_stack(
tmp_path,
@@ -419,16 +479,18 @@ async def test_reload_mcp_servers_retries_configured_server_without_live_stack(
return stacks
monkeypatch.setattr("nanobot.agent.tools.mcp.connect_mcp_servers", _fake_connect)
loop = _make_loop(tmp_path, mcp_servers={"browserbase": config.tools.mcp_servers["browserbase"]})
provider, registry = _make_provider(
mcp_servers={"browserbase": config.tools.mcp_servers["browserbase"]}
)
result = await mcp_runtime.reload_servers(loop, loop.tools)
result = await provider.reload()
assert result["ok"] is True
assert result["added"] == []
assert result["changed"] == []
assert result["retried"] == ["browserbase"]
assert loop.tools.has("mcp_browserbase_navigate")
await loop.close_mcp()
assert registry.has("mcp_browserbase_navigate")
await provider.aclose()
@pytest.mark.asyncio
@@ -465,16 +527,16 @@ async def test_reload_mcp_servers_skips_oauth_server_waiting_for_authorization(
"nanobot.agent.tools.mcp_oauth.mcp_oauth_has_credentials",
lambda name, _url: name == "linear",
)
loop = _make_loop(tmp_path, mcp_servers={"notion": notion})
provider, _registry = _make_provider(mcp_servers={"notion": notion})
result = await mcp_runtime.reload_servers(loop, loop.tools)
result = await provider.reload()
assert attempted == ["linear"]
assert result["ok"] is True
assert result["failed"] == []
assert result["retried"] == []
assert result["connected"] == ["linear"]
await loop.close_mcp()
await provider.aclose()
@pytest.mark.asyncio
@@ -482,7 +544,9 @@ async def test_mcp_tool_reconnects_after_session_terminated(
tmp_path,
monkeypatch: pytest.MonkeyPatch,
):
loop = _make_loop(tmp_path, mcp_servers={"remote": object()})
provider, registry = _make_provider(
mcp_servers={"remote": _stdio_server("remote")}
)
closed: list[str] = []
sessions: list[Any] = []
connect_count = 0
@@ -525,8 +589,8 @@ async def test_mcp_tool_reconnects_after_session_terminated(
monkeypatch.setattr("nanobot.agent.tools.mcp.connect_mcp_servers", _fake_connect)
await loop._connect_mcp()
old_tool = loop.tools.get("mcp_remote_quote")
await provider.connect()
old_tool = registry.get("mcp_remote_quote")
assert isinstance(old_tool, MCPToolWrapper)
output = await old_tool.execute(symbol="AAPL")
@@ -536,8 +600,8 @@ async def test_mcp_tool_reconnects_after_session_terminated(
assert closed == ["remote"]
assert sessions[0].call_count == 1
assert sessions[1].call_count == 1
assert "remote" in loop._mcp_stacks
assert loop.tools.get("mcp_remote_quote") is not old_tool
assert provider.connected_server_names == {"remote"}
assert registry.get("mcp_remote_quote") is not old_tool
@pytest.mark.asyncio
@@ -545,7 +609,9 @@ async def test_mcp_reconnect_handler_uses_sanitized_server_prefix(
tmp_path,
monkeypatch: pytest.MonkeyPatch,
):
loop = _make_loop(tmp_path, mcp_servers={"remote_": object()})
provider, registry = _make_provider(
mcp_servers={"remote_": _stdio_server("remote")}
)
connect_count = 0
class _FakeSession:
@@ -578,15 +644,15 @@ async def test_mcp_reconnect_handler_uses_sanitized_server_prefix(
monkeypatch.setattr("nanobot.agent.tools.mcp.connect_mcp_servers", _fake_connect)
await loop._connect_mcp()
old_tool = loop.tools.get("mcp_remote_quote")
await provider.connect()
old_tool = registry.get("mcp_remote_quote")
assert isinstance(old_tool, MCPToolWrapper)
output = await old_tool.execute()
assert output == "recovered"
assert connect_count == 2
assert loop.tools.get("mcp_remote_quote") is not old_tool
assert registry.get("mcp_remote_quote") is not old_tool
@pytest.mark.asyncio
@@ -594,7 +660,9 @@ async def test_concurrent_mcp_reconnect_reuses_fresh_session(
tmp_path,
monkeypatch: pytest.MonkeyPatch,
):
loop = _make_loop(tmp_path, mcp_servers={"remote": object()})
provider, registry = _make_provider(
mcp_servers={"remote": _stdio_server("remote")}
)
closed: list[str] = []
connect_count = 0
@@ -638,9 +706,9 @@ async def test_concurrent_mcp_reconnect_reuses_fresh_session(
monkeypatch.setattr("nanobot.agent.tools.mcp.connect_mcp_servers", _fake_connect)
await loop._connect_mcp()
old_alpha = loop.tools.get("mcp_remote_resource_alpha")
old_beta = loop.tools.get("mcp_remote_resource_beta")
await provider.connect()
old_alpha = registry.get("mcp_remote_resource_alpha")
old_beta = registry.get("mcp_remote_resource_beta")
assert isinstance(old_alpha, MCPResourceWrapper)
assert isinstance(old_beta, MCPResourceWrapper)
tests/agent/test_mcp_reconnect_crash.py
+15 -26
View File
@@ -15,15 +15,13 @@ import asyncio
import multiprocessing
import socket
import time
from unittest.mock import MagicMock
import httpx
import pytest
from nanobot.agent.loop import AgentLoop
from nanobot.agent.tools import mcp as mcp_module
from nanobot.agent.tools.mcp import MCPToolWrapper
from nanobot.bus.queue import MessageBus
from nanobot.agent.tools.mcp import MCPProvider, MCPToolWrapper
from nanobot.agent.tools.registry import ToolRegistry
from nanobot.config.schema import MCPServerConfig
from nanobot.security import network as security_network
@@ -113,18 +111,9 @@ def mcp_server_url():
process.join(timeout=2.0)
def _make_loop(tmp_path, *, mcp_servers: dict) -> AgentLoop:
bus = MessageBus()
provider = MagicMock()
provider.get_default_model.return_value = "test-model"
provider.generation.max_tokens = 4096
return AgentLoop(
bus=bus,
provider=provider,
workspace=tmp_path,
model="test-model",
mcp_servers=mcp_servers,
)
def _make_provider(*, mcp_servers: dict) -> tuple[MCPProvider, ToolRegistry]:
registry = ToolRegistry()
return MCPProvider(mcp_servers, registry), registry
@pytest.fixture(autouse=True)
@@ -170,12 +159,12 @@ async def test_mcp_reconnect_after_session_timeout(tmp_path, mcp_server_url):
tool_timeout=_TOOL_TIMEOUT_SECONDS,
enabled_tools=["*"],
)
loop = _make_loop(tmp_path, mcp_servers={"repro": cfg})
provider, registry = _make_provider(mcp_servers={"repro": cfg})
await asyncio.create_task(loop._connect_mcp())
assert "repro" in loop._mcp_stacks
await asyncio.create_task(provider.connect())
assert provider.connected_server_names == {"repro"}
tool = loop.tools.get("mcp_repro_greet")
tool = registry.get("mcp_repro_greet")
assert isinstance(tool, MCPToolWrapper)
output = await asyncio.create_task(tool.execute(name="first"))
@@ -187,7 +176,7 @@ async def test_mcp_reconnect_after_session_timeout(tmp_path, mcp_server_url):
output = await asyncio.create_task(tool.execute(name="second"))
assert "Hello, second" in output
await asyncio.create_task(loop.close_mcp())
await asyncio.create_task(provider.aclose())
@pytest.mark.asyncio
@@ -203,10 +192,10 @@ async def test_mcp_reconnect_during_shutdown_does_not_crash(
tool_timeout=_TOOL_TIMEOUT_SECONDS,
enabled_tools=["*"],
)
loop = _make_loop(tmp_path, mcp_servers={"repro": cfg})
provider, registry = _make_provider(mcp_servers={"repro": cfg})
await asyncio.create_task(loop._connect_mcp())
tool = loop.tools.get("mcp_repro_greet")
await asyncio.create_task(provider.connect())
tool = registry.get("mcp_repro_greet")
assert isinstance(tool, MCPToolWrapper)
await asyncio.create_task(tool.execute(name="first"))
@@ -224,7 +213,7 @@ async def test_mcp_reconnect_during_shutdown_does_not_crash(
monkeypatch.setattr(mcp_module, "connect_mcp_servers", gated_connect)
call_task = asyncio.create_task(tool.execute(name="second"))
await asyncio.wait_for(reconnect_started.wait(), timeout=5)
close_task = asyncio.create_task(loop.close_mcp())
close_task = asyncio.create_task(provider.aclose())
await asyncio.sleep(0)
finish_reconnect.set()
@@ -245,4 +234,4 @@ async def test_mcp_reconnect_during_shutdown_does_not_crash(
unhandled.append(exc)
assert not unhandled, f"Unhandled exception leaked during reconnect/shutdown: {unhandled[0]}"
assert loop._mcp_stacks == {}
assert provider.connected_server_names == set()
tests/agent/test_runtime_refresh.py
+6 -1
View File
@@ -7,6 +7,7 @@ from unittest.mock import MagicMock
import pytest
from nanobot.agent.loop import AgentLoop
from nanobot.agent.tools.registry import ToolRegistry
from nanobot.bus.queue import MessageBus
from nanobot.bus.runtime_events import RuntimeModelChanged
from nanobot.config.errors import ConfigLoadError
@@ -312,7 +313,11 @@ def test_settings_context_window_refreshes_runtime_state(
def loader(*, preset_name: str | None = None) -> ProviderSnapshot:
return load_provider_snapshot(config_path, preset_name=preset_name)
loop = AgentLoop.from_config(config, provider_snapshot_loader=loader)
loop = AgentLoop.from_config(
config,
tool_registry=ToolRegistry(),
provider_snapshot_loader=loader,
)
payload = update_agent_settings({"context_window_tokens": ["262144"]})
loop.runtime_resolver.invalidate()
tests/agent/test_self_model_preset.py
+3 -2
View File
@@ -5,6 +5,7 @@ import pytest
from nanobot.agent.loop import AgentLoop
from nanobot.agent.tools.context import RequestContext, request_context
from nanobot.agent.tools.registry import ToolRegistry
from nanobot.agent.tools.runtime_control import AgentRuntimeControl
from nanobot.agent.tools.self import MyTool
from nanobot.bus.queue import MessageBus
@@ -390,7 +391,7 @@ def test_from_config_injects_default_preset(tmp_path) -> None:
})
fake_provider = _provider("openai/gpt-4.1")
with patch("nanobot.providers.factory.make_provider", return_value=fake_provider):
loop = AgentLoop.from_config(config)
loop = AgentLoop.from_config(config, tool_registry=ToolRegistry())
assert loop.model == "openai/gpt-4.1"
assert loop.model_preset is None
assert "default" in loop.model_presets
@@ -407,7 +408,7 @@ def test_from_config_static_preset_loader_does_not_enable_hot_reload(tmp_path) -
})
fake_provider = _provider("openai/gpt-4.1")
with patch("nanobot.providers.factory.make_provider", return_value=fake_provider):
loop = AgentLoop.from_config(config)
loop = AgentLoop.from_config(config, tool_registry=ToolRegistry())
default_runtime = loop.runtime_resolver.runtime
resolved = loop.runtime_resolver.resolve_preset("fast")
assert resolved.model == "openai/gpt-4.1-mini"
tests/agent/test_task_cancel.py
+11 -14
View File
@@ -56,7 +56,7 @@ class TestHandleStop:
assert "No active task" in out.content
@pytest.mark.asyncio
async def test_close_mcp_cancels_active_turn_before_resources(self):
async def test_aclose_cancels_active_turn_before_resources(self):
loop, _bus = _make_loop()
events: list[str] = []
@@ -76,14 +76,13 @@ class TestHandleStop:
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()
await loop.aclose()
assert events == ["turn_cancelled", "resources_closed"]
assert task.cancelled()
@pytest.mark.asyncio
async def test_close_mcp_serializes_duplicate_cleanup(self):
async def test_aclose_serializes_duplicate_cleanup(self):
loop, _bus = _make_loop()
entered = asyncio.Event()
release = asyncio.Event()
@@ -100,14 +99,13 @@ class TestHandleStop:
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)
first = asyncio.create_task(loop.aclose())
await entered.wait()
second = asyncio.create_task(loop.aclose())
await asyncio.sleep(0)
assert not second.done()
release.set()
await asyncio.gather(first, second)
assert max_concurrent == 1
@@ -172,8 +170,7 @@ class TestDispatch:
@pytest.mark.asyncio
async def test_run_logs_and_continues_after_leaked_cancelled_error(self, monkeypatch):
loop, bus = _make_loop()
loop._connect_mcp = AsyncMock()
loop.close_mcp = AsyncMock()
loop.aclose = AsyncMock()
loop.auto_compact.check_expired = MagicMock()
warnings: list[str] = []
calls = 0
tests/agent/tools/test_self_tool.py
-21
View File
@@ -493,20 +493,6 @@ class TestModifyOpen:
assert "Set workspace" in result
assert tool._runtime_control.snapshot().workspace == "/new/path"
@pytest.mark.asyncio
async def test_modify_mcp_servers_blocked(self):
"""_mcp_servers contains API credentials — must be blocked."""
tool = _make_tool()
result = await tool.execute(action="set", key="_mcp_servers", value={"evil": "leaked"})
assert "protected" in result
@pytest.mark.asyncio
async def test_modify_mcp_stacks_blocked(self):
"""_mcp_stacks holds connection handles — must be blocked."""
tool = _make_tool()
result = await tool.execute(action="set", key="_mcp_stacks", value={})
assert "protected" in result
@pytest.mark.asyncio
async def test_modify_pending_queues_blocked(self):
"""_pending_queues controls message routing — must be blocked."""
@@ -535,13 +521,6 @@ class TestModifyOpen:
result = await tool.execute(action="set", key="_background_tasks", value=[])
assert "protected" in result
@pytest.mark.asyncio
async def test_inspect_mcp_servers_blocked(self):
"""_mcp_servers contains credentials — check must be blocked too."""
tool = _make_tool()
result = await tool.execute(action="check", key="_mcp_servers")
assert "not accessible" in result
@pytest.mark.asyncio
async def test_modify_wrapped_denied(self):
"""__wrapped__ allows decorator bypass — must be denied."""
tests/cli/test_agent_interactive.py
+4 -4
View File
@@ -64,7 +64,7 @@ def test_interactive_agent_routes_a_complete_user_turn(
def __init__(self, bus) -> None:
self.bus = bus
self.stopped = asyncio.Event()
self.close_mcp_calls = 0
self.aclose_calls = 0
async def run(self) -> None:
message = await self.bus.consume_inbound()
@@ -97,8 +97,8 @@ def test_interactive_agent_routes_a_complete_user_turn(
def stop(self) -> None:
self.stopped.set()
async def close_mcp(self) -> None:
self.close_mcp_calls += 1
async def aclose(self) -> None:
self.aclose_calls += 1
read_input = AsyncMock(side_effect=["hello nanobot", "exit"])
print_response = MagicMock()
@@ -136,7 +136,7 @@ def test_interactive_agent_routes_a_complete_user_turn(
assert inbound.metadata == {"_wants_stream": True}
loop = seen["loop"]
assert isinstance(loop, _AgentLoop)
assert loop.close_mcp_calls == 1
assert loop.aclose_calls == 1
assert len(renderers) == 1
renderer = renderers[0]
assert isinstance(renderer, _Renderer)
tests/cli/test_commands.py
+50 -19
View File
@@ -1543,7 +1543,7 @@ def mock_agent_runtime(tmp_path):
agent_loop.process_direct = AsyncMock(
return_value=OutboundMessage(channel="cli", chat_id="direct", content="mock-response"),
)
agent_loop.close_mcp = AsyncMock(return_value=None)
agent_loop.aclose = AsyncMock(return_value=None)
mock_from_config.return_value = agent_loop
yield {
@@ -1621,7 +1621,7 @@ def test_agent_config_sets_active_path(monkeypatch, tmp_path: Path) -> None:
async def process_direct(self, *_args, **_kwargs):
return OutboundMessage(channel="cli", chat_id="direct", content="ok")
async def close_mcp(self) -> None:
async def aclose(self) -> None:
return None
monkeypatch.setattr("nanobot.cli.agent.AgentLoop", _FakeAgentLoop)
@@ -1662,7 +1662,7 @@ def test_agent_uses_workspace_directory_for_cron_store(monkeypatch, tmp_path: Pa
async def process_direct(self, *_args, **_kwargs):
return OutboundMessage(channel="cli", chat_id="direct", content="ok")
async def close_mcp(self) -> None:
async def aclose(self) -> None:
return None
monkeypatch.setattr("nanobot.cron.service.CronService", _FakeCron)
@@ -1712,7 +1712,7 @@ def test_agent_workspace_override_does_not_migrate_legacy_cron(
async def process_direct(self, *_args, **_kwargs):
return OutboundMessage(channel="cli", chat_id="direct", content="ok")
async def close_mcp(self) -> None:
async def aclose(self) -> None:
return None
monkeypatch.setattr("nanobot.cron.service.CronService", _FakeCron)
@@ -1768,7 +1768,7 @@ def test_agent_custom_config_workspace_does_not_migrate_legacy_cron(
async def process_direct(self, *_args, **_kwargs):
return OutboundMessage(channel="cli", chat_id="direct", content="ok")
async def close_mcp(self) -> None:
async def aclose(self) -> None:
return None
monkeypatch.setattr("nanobot.cron.service.CronService", _FakeCron)
@@ -2062,7 +2062,7 @@ def test_heartbeat_empty_response_still_retains_recent_messages(
async def process_direct(self, *_args, **_kwargs):
return SimpleNamespace(content="")
async def close_mcp(self) -> None:
async def aclose(self) -> None:
return None
async def run(self) -> None:
@@ -2738,10 +2738,7 @@ def _patch_serve_runtime(monkeypatch, config: Config, seen: dict[str, object]) -
def __init__(self, **kwargs) -> None:
seen["workspace"] = kwargs["workspace"]
async def _connect_mcp(self) -> None:
return None
async def close_mcp(self) -> None:
async def aclose(self) -> None:
return None
def _fake_create_app(
@@ -2749,11 +2746,13 @@ def _patch_serve_runtime(monkeypatch, config: Config, seen: dict[str, object]) -
model_name: str,
request_timeout: float,
api_key: str = "",
prepare_agent=None,
):
seen["agent_loop"] = agent_loop
seen["model_name"] = model_name
seen["request_timeout"] = request_timeout
seen["api_key"] = api_key
seen["prepare_agent"] = prepare_agent
return _FakeApiApp()
def _fake_run_app(api_app, host: str, port: int, print):
@@ -2914,7 +2913,7 @@ def test_gateway_unbound_agent_cron_is_skipped(
async def submit_cron_turn(self, _msg: InboundMessage):
raise AssertionError("unbound cron job must not run as a bound cron turn")
async def close_mcp(self) -> None:
async def aclose(self) -> None:
return None
async def run(self) -> None:
@@ -3033,7 +3032,7 @@ def test_gateway_bound_cron_runs_as_session_turn(
content="Checked the repo.",
)
async def close_mcp(self) -> None:
async def aclose(self) -> None:
return None
async def run(self) -> None:
@@ -3253,7 +3252,7 @@ def test_gateway_local_trigger_queue_submits_agent_turns(
self.runtime_resolver.invalidate.assert_called_once_with()
await asyncio.Event().wait()
async def close_mcp(self) -> None:
async def aclose(self) -> None:
return None
def stop(self) -> None:
@@ -3499,7 +3498,7 @@ def test_gateway_health_endpoint_binds_and_serves_expected_responses(
async def run(self) -> None:
await asyncio.Event().wait()
async def close_mcp(self) -> None:
async def aclose(self) -> None:
return None
def stop(self) -> None:
@@ -3668,7 +3667,7 @@ def test_gateway_health_endpoint_binds_and_serves_expected_responses(
assert timed_out_writer.output == b""
def test_gateway_shutdown_lets_agent_task_own_mcp_cleanup(
def test_gateway_agent_task_owns_initial_mcp_provider_close(
monkeypatch,
tmp_path: Path,
) -> None:
@@ -3696,17 +3695,41 @@ def test_gateway_shutdown_lets_agent_task_own_mcp_cleanup(
return None
async def run(self) -> None:
seen["agent_task"] = asyncio.current_task()
try:
await asyncio.Event().wait()
finally:
seen["agent_task_cleaned_up"] = True
async def close_mcp(self) -> None:
raise AssertionError("gateway must not close MCP from the outer task")
async def aclose(self) -> None:
seen["agent_closed"] = True
def stop(self) -> None:
seen["agent_stopped"] = True
class _FakeMCPProvider:
def __init__(self) -> None:
self.connect_task: asyncio.Task | None = None
self.close_tasks: list[asyncio.Task | None] = []
@classmethod
def from_config(cls, _config, _registry):
provider = cls()
seen["mcp_provider"] = provider
return provider
async def connect(self) -> None:
self.connect_task = asyncio.current_task()
async def aclose(self) -> None:
self.close_tasks.append(asyncio.current_task())
def runtime_status(self) -> dict[str, str]:
return {}
async def reload(self) -> dict[str, object]:
return {"ok": True}
class _FakeChannelManager:
def __init__(self, _config, _bus, **_kwargs) -> None:
self.enabled_channels = ["telegram"]
@@ -3753,6 +3776,7 @@ def test_gateway_shutdown_lets_agent_task_own_mcp_cleanup(
session_manager=lambda _workspace: object(),
)
monkeypatch.setattr("nanobot.cli.gateway_runtime.AgentLoop", _FakeAgentLoop)
monkeypatch.setattr("nanobot.cli.gateway_runtime.MCPProvider", _FakeMCPProvider)
monkeypatch.setattr("nanobot.channels.manager.ChannelManager", _FakeChannelManager)
monkeypatch.setattr("nanobot.cron.service.CronService", _FakeCronService)
monkeypatch.setattr("asyncio.start_server", _fake_start_server)
@@ -3761,9 +3785,15 @@ def test_gateway_shutdown_lets_agent_task_own_mcp_cleanup(
assert result.exit_code == 0
assert seen["agent_stopped"] is True
assert seen["agent_closed"] is True
assert seen["agent_task_cleaned_up"] is True
assert seen["channels_stopped"] is True
assert seen["cron_stopped"] is True
mcp_provider = seen["mcp_provider"]
assert isinstance(mcp_provider, _FakeMCPProvider)
assert mcp_provider.connect_task is seen["agent_task"]
assert mcp_provider.close_tasks[0] is mcp_provider.connect_task
assert len(mcp_provider.close_tasks) == 2
def test_gateway_shutdown_event_exits_forever_runtime_tasks(
@@ -3800,8 +3830,8 @@ def test_gateway_shutdown_event_exits_forever_runtime_tasks(
finally:
seen["agent_task_cleaned_up"] = True
async def close_mcp(self) -> None:
raise AssertionError("gateway must not close MCP from the outer task")
async def aclose(self) -> None:
seen["agent_closed"] = True
def stop(self) -> None:
seen["agent_stopped"] = True
@@ -3881,6 +3911,7 @@ def test_gateway_shutdown_event_exits_forever_runtime_tasks(
assert result.exit_code == 0
assert seen["agent_stopped"] is True
assert seen["agent_closed"] is True
assert seen["agent_task_cleaned_up"] is True
assert seen["channel_task_cleaned_up"] is True
assert seen["channels_stopped"] is True
tests/cli/test_gateway_runtime.py
+43 -11
View File
@@ -22,7 +22,7 @@ class _FakeAgent:
self.raise_on_close = False
self.background: asyncio.Task[None] | None = None
async def close_mcp(self) -> None:
async def aclose(self) -> None:
self.close_calls += 1
if self.hang_on_close:
await asyncio.sleep(3600)
@@ -30,7 +30,7 @@ class _FakeAgent:
raise RuntimeError("cleanup exploded")
if self.background is not None:
await self.background
self.events.append("close_mcp")
self.events.append("aclose")
class _FakeChannels:
@@ -43,6 +43,16 @@ class _FakeChannels:
self.events.append("channels_stopped")
class _FakeMCPProvider:
def __init__(self, events: list[str] | None = None) -> None:
self.close_calls = 0
self.events = events if events is not None else []
async def aclose(self) -> None:
self.close_calls += 1
self.events.append("mcp_closed")
async def _cancellable_task(events: list[str]) -> None:
try:
await asyncio.sleep(3600)
@@ -64,13 +74,14 @@ async def _stubborn_task(events: list[str]) -> None:
async def test_runtime_tasks_cancelled_before_resources_closed() -> None:
events: list[str] = []
agent = _FakeAgent(events)
provider = _FakeMCPProvider(events)
channels = _FakeChannels()
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)
await _close_gateway_runtime(agent, provider, channels, [task], None)
assert events == ["cancelled", "close_mcp"] # cancel happens before close
assert events == ["cancelled", "aclose", "mcp_closed"]
assert channels.stopped == 1
assert agent.close_calls == 1
assert task.cancelled()
@@ -78,6 +89,7 @@ async def test_runtime_tasks_cancelled_before_resources_closed() -> None:
async def test_pending_background_work_is_drained_before_close_returns() -> None:
agent = _FakeAgent()
provider = _FakeMCPProvider()
channels = _FakeChannels()
done: dict[str, bool] = {"done": False}
@@ -87,7 +99,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)
await _close_gateway_runtime(agent, provider, channels, [], None)
assert done["done"] is True
assert agent.close_calls == 1
@@ -95,6 +107,7 @@ async def test_pending_background_work_is_drained_before_close_returns() -> None
async def test_stubborn_task_does_not_block_past_wait_timeout() -> None:
agent = _FakeAgent()
provider = _FakeMCPProvider()
channels = _FakeChannels()
events: list[str] = []
task = asyncio.create_task(_stubborn_task(events))
@@ -104,6 +117,7 @@ async def test_stubborn_task_does_not_block_past_wait_timeout() -> None:
start = time.monotonic()
await _close_gateway_runtime(
agent,
provider,
channels,
[task],
runtime_tasks,
@@ -117,70 +131,88 @@ async def test_stubborn_task_does_not_block_past_wait_timeout() -> None:
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 provider.close_calls == 1
assert elapsed < 1.0 # bounded, not held open by the stubborn task
async def test_hanging_close_is_bounded_and_does_not_raise() -> None:
agent = _FakeAgent()
provider = _FakeMCPProvider()
agent.hang_on_close = True
channels = _FakeChannels()
start = time.monotonic()
await _close_gateway_runtime(agent, channels, [], None, close_timeout=0.05)
await _close_gateway_runtime(
agent,
provider,
channels,
[],
None,
close_timeout=0.05,
)
elapsed = time.monotonic() - start
assert agent.close_calls == 1
assert provider.close_calls == 1
assert channels.stopped == 1
assert elapsed < 1.0
async def test_failing_close_is_logged_but_shutdown_proceeds() -> None:
agent = _FakeAgent()
provider = _FakeMCPProvider()
agent.raise_on_close = True
channels = _FakeChannels()
await _close_gateway_runtime(agent, channels, [], None)
await _close_gateway_runtime(agent, provider, channels, [], None)
assert agent.close_calls == 1
assert provider.close_calls == 1
assert channels.stopped == 1 # teardown continued past the failure
async def test_duplicate_cleanup_is_idempotent() -> None:
agent = _FakeAgent()
provider = _FakeMCPProvider()
channels = _FakeChannels()
task = asyncio.create_task(_cancellable_task([]))
await _close_gateway_runtime(agent, channels, [task], None)
await _close_gateway_runtime(agent, channels, [task], None)
await _close_gateway_runtime(agent, provider, channels, [task], None)
await _close_gateway_runtime(agent, provider, channels, [task], None)
assert agent.close_calls == 2 # second pass is a clean no-op
assert provider.close_calls == 2
assert channels.stopped == 2
assert task.cancelled()
async def test_finished_runtime_tasks_gather_is_retrieved() -> None:
agent = _FakeAgent()
provider = _FakeMCPProvider()
channels = _FakeChannels()
finished = asyncio.get_running_loop().create_future()
finished.set_result(None)
runtime_tasks = asyncio.gather(finished)
await asyncio.sleep(0) # let the gather observe the finished child
await _close_gateway_runtime(agent, channels, [], runtime_tasks)
await _close_gateway_runtime(agent, provider, channels, [], runtime_tasks)
assert runtime_tasks.done()
assert agent.close_calls == 1
assert provider.close_calls == 1
async def test_cancelled_runtime_tasks_gather_does_not_raise() -> None:
agent = _FakeAgent()
provider = _FakeMCPProvider()
channels = _FakeChannels()
runtime_tasks = asyncio.gather(asyncio.sleep(3600))
runtime_tasks.cancel()
await _close_gateway_runtime(agent, channels, [], runtime_tasks)
await _close_gateway_runtime(agent, provider, channels, [], runtime_tasks)
with suppress(asyncio.CancelledError):
await runtime_tasks # settle the cancelled gather without raising
assert runtime_tasks.done() # the cancelled gather was awaited without raising
assert agent.close_calls == 1
assert provider.close_calls == 1
tests/test_api_attachment.py
+1 -2
View File
@@ -32,8 +32,7 @@ AUTH_HEADERS = {"Authorization": f"Bearer {API_KEY}"}
def _make_mock_agent(response_text: str = "mock response") -> MagicMock:
agent = MagicMock()
agent.process_direct = AsyncMock(return_value=response_text)
agent._connect_mcp = AsyncMock()
agent.close_mcp = AsyncMock()
agent.aclose = AsyncMock()
agent._last_usage = {}
return agent
tests/test_api_stream.py
+8 -16
View File
@@ -64,8 +64,7 @@ def test_sse_done_format() -> None:
def _make_streaming_agent(tokens: list[str]) -> MagicMock:
"""Create a mock agent that streams tokens via on_stream callback."""
agent = MagicMock()
agent._connect_mcp = AsyncMock()
agent.close_mcp = AsyncMock()
agent.aclose = AsyncMock()
async def fake_process_direct(*, content="", media=None, session_key="",
channel="", chat_id="", on_stream=None,
@@ -136,8 +135,7 @@ async def test_stream_false_returns_json(aiohttp_client) -> None:
"""stream=false should still return regular JSON response."""
agent = MagicMock()
agent.process_direct = AsyncMock(return_value="normal reply")
agent._connect_mcp = AsyncMock()
agent.close_mcp = AsyncMock()
agent.aclose = AsyncMock()
agent._last_usage = {}
app = create_app(agent, model_name="m", api_key=API_KEY)
@@ -160,8 +158,7 @@ async def test_stream_default_is_false(aiohttp_client) -> None:
"""Omitting stream should behave like stream=false."""
agent = MagicMock()
agent.process_direct = AsyncMock(return_value="default reply")
agent._connect_mcp = AsyncMock()
agent.close_mcp = AsyncMock()
agent.aclose = AsyncMock()
agent._last_usage = {}
app = create_app(agent, model_name="m", api_key=API_KEY)
@@ -217,8 +214,7 @@ async def test_stream_passes_on_stream_callbacks(aiohttp_client) -> None:
agent = MagicMock()
agent.process_direct = fake_process_direct
agent._connect_mcp = AsyncMock()
agent.close_mcp = AsyncMock()
agent.aclose = AsyncMock()
agent._last_usage = {}
app = create_app(agent, model_name="m", api_key=API_KEY)
@@ -251,8 +247,7 @@ async def test_stream_segment_end_does_not_close_sse(aiohttp_client) -> None:
return "planning final"
agent.process_direct = fake_process_direct
agent._connect_mcp = AsyncMock()
agent.close_mcp = AsyncMock()
agent.aclose = AsyncMock()
agent._last_usage = {}
app = create_app(agent, model_name="m", api_key=API_KEY)
@@ -291,8 +286,7 @@ async def test_stream_uses_final_response_when_no_deltas(aiohttp_client) -> None
return "plain final"
agent.process_direct = fake_process_direct
agent._connect_mcp = AsyncMock()
agent.close_mcp = AsyncMock()
agent.aclose = AsyncMock()
agent._last_usage = {}
app = create_app(agent, model_name="m", api_key=API_KEY)
@@ -334,8 +328,7 @@ async def test_stream_with_session_id(aiohttp_client) -> None:
agent = MagicMock()
agent.process_direct = fake_process_direct
agent._connect_mcp = AsyncMock()
agent.close_mcp = AsyncMock()
agent.aclose = AsyncMock()
agent._last_usage = {}
app = create_app(agent, model_name="m", api_key=API_KEY)
@@ -364,8 +357,7 @@ async def test_streaming_backend_failure_does_not_emit_success_terminator(aiohtt
raise RuntimeError("backend blew up")
agent.process_direct = boom
agent._connect_mcp = AsyncMock()
agent.close_mcp = AsyncMock()
agent.aclose = AsyncMock()
agent._last_usage = {}
app = create_app(agent, model_name="m", api_key=API_KEY)
tests/test_nanobot_facade.py
+23 -7
View File
@@ -101,6 +101,19 @@ def test_from_config_creates_instance(tmp_path):
assert bot._loop.workspace == tmp_path
def test_from_config_composes_configured_mcp_outside_agent_loop(tmp_path):
config_path = _write_config(
tmp_path,
{"tools": {"mcpServers": {"demo": {"command": "fake-mcp"}}}},
)
bot = Nanobot.from_config(config_path, workspace=tmp_path)
assert bot._mcp_provider is not None
assert bot._mcp_provider.configured_server_names == {"demo"}
assert bot._mcp_provider._registry is bot._loop.tools
def test_from_config_accepts_default_model_override(tmp_path):
config_path = _write_config(tmp_path)
@@ -1637,37 +1650,40 @@ async def test_runtime_helpers_expose_model_workspace_and_compact(tmp_path):
@pytest.mark.asyncio
async def test_aclose_delegates_to_loop_close_mcp(tmp_path):
async def test_aclose_releases_loop_and_mcp_provider(tmp_path):
config_path = _write_config(tmp_path)
bot = Nanobot.from_config(config_path, workspace=tmp_path)
bot._loop.close_mcp = AsyncMock()
bot._loop.aclose = AsyncMock()
assert bot._mcp_provider is not None
bot._mcp_provider.aclose = AsyncMock()
await bot.aclose()
bot._loop.close_mcp.assert_awaited_once()
bot._loop.aclose.assert_awaited_once()
bot._mcp_provider.aclose.assert_awaited_once()
@pytest.mark.asyncio
async def test_context_manager_calls_aclose_on_exit(tmp_path):
config_path = _write_config(tmp_path)
bot = Nanobot.from_config(config_path, workspace=tmp_path)
bot._loop.close_mcp = AsyncMock()
bot._loop.aclose = AsyncMock()
async with bot as b:
assert b is bot
bot._loop.close_mcp.assert_awaited_once()
bot._loop.aclose.assert_awaited_once()
@pytest.mark.asyncio
async def test_context_manager_does_not_swallow_exceptions(tmp_path):
config_path = _write_config(tmp_path)
bot = Nanobot.from_config(config_path, workspace=tmp_path)
bot._loop.close_mcp = AsyncMock()
bot._loop.aclose = AsyncMock()
with pytest.raises(ValueError):
async with bot as b:
assert b is bot
raise ValueError("boom")
bot._loop.close_mcp.assert_awaited_once()
bot._loop.aclose.assert_awaited_once()
tests/test_openai_api.py
+57 -9
View File
@@ -34,8 +34,7 @@ AUTH_HEADERS = {"Authorization": f"Bearer {API_KEY}"}
def _make_mock_agent(response_text: str = "mock response") -> MagicMock:
agent = MagicMock()
agent.process_direct = AsyncMock(return_value=response_text)
agent._connect_mcp = AsyncMock()
agent.close_mcp = AsyncMock()
agent.aclose = AsyncMock()
agent._last_usage = {"prompt_tokens": 100, "completion_tokens": 50}
return agent
@@ -149,6 +148,59 @@ async def test_api_routes_allow_requests_without_configured_api_key(aiohttp_clie
mock_agent.process_direct.assert_called_once()
@pytest.mark.skipif(not HAS_AIOHTTP, reason="aiohttp not installed")
@pytest.mark.asyncio
async def test_api_prepares_application_resources_before_each_turn(aiohttp_client) -> None:
events: list[str] = []
agent = _make_mock_agent()
async def prepare_agent() -> None:
events.append("prepare")
async def process_direct(**_kwargs):
events.append("process")
return "ready"
agent.process_direct = process_direct
app = create_app(agent, prepare_agent=prepare_agent)
client = await aiohttp_client(app)
response = await client.post(
"/v1/chat/completions",
json={"messages": [{"role": "user", "content": "hello"}]},
)
assert response.status == 200
assert events == ["prepare", "process"]
@pytest.mark.skipif(not HAS_AIOHTTP, reason="aiohttp not installed")
@pytest.mark.asyncio
async def test_api_preparation_is_bounded_by_request_timeout(aiohttp_client) -> None:
agent = _make_mock_agent()
started = asyncio.Event()
async def prepare_agent() -> None:
started.set()
await asyncio.Event().wait()
app = create_app(
agent,
request_timeout=0.01,
prepare_agent=prepare_agent,
)
client = await aiohttp_client(app)
response = await client.post(
"/v1/chat/completions",
json={"messages": [{"role": "user", "content": "hello"}]},
)
assert started.is_set()
assert response.status == 504
agent.process_direct.assert_not_awaited()
@pytest.mark.skipif(not HAS_AIOHTTP, reason="aiohttp not installed")
@pytest.mark.asyncio
async def test_no_user_message_returns_400(aiohttp_client, app) -> None:
@@ -275,8 +327,7 @@ async def test_followup_requests_share_same_session_key(aiohttp_client) -> None:
agent = MagicMock()
agent.process_direct = fake_process
agent._connect_mcp = AsyncMock()
agent.close_mcp = AsyncMock()
agent.aclose = AsyncMock()
agent._last_usage = {}
app = create_app(agent, model_name="m", api_key=API_KEY)
@@ -315,8 +366,7 @@ async def test_fixed_session_requests_are_serialized(aiohttp_client) -> None:
agent = MagicMock()
agent.process_direct = slow_process
agent._connect_mcp = AsyncMock()
agent.close_mcp = AsyncMock()
agent.aclose = AsyncMock()
agent._last_usage = {}
app = create_app(agent, model_name="m", api_key=API_KEY)
@@ -433,8 +483,7 @@ async def test_empty_response_falls_back_without_retry(aiohttp_client) -> None:
agent = MagicMock()
agent.process_direct = always_empty
agent._connect_mcp = AsyncMock()
agent.close_mcp = AsyncMock()
agent.aclose = AsyncMock()
agent._last_usage = {}
app = create_app(agent, model_name="m", api_key=API_KEY)
@@ -457,7 +506,6 @@ async def test_process_direct_accepts_media() -> None:
from nanobot.bus.runtime_events import RuntimeEventPublisher
loop = AgentLoop.__new__(AgentLoop)
loop._connect_mcp = AsyncMock()
loop._session_locks = {}
loop.runtime_event_publisher = RuntimeEventPublisher()
tests/tools/test_exec_session_tools.py
+7 -17
View File
@@ -12,7 +12,6 @@ from unittest.mock import AsyncMock
import pytest
from nanobot.agent import context as agent_context
from nanobot.agent.loop import AgentLoop
from nanobot.agent.tools.context import RequestContext, bind_request_context, reset_request_context
from nanobot.agent.tools.exec_session import (
@@ -712,7 +711,7 @@ def test_exec_session_manager_preserves_single_cleanup_error():
asyncio.run(run())
def test_agent_loop_shutdown_closes_exec_sessions(tmp_path, monkeypatch):
def test_agent_loop_shutdown_closes_exec_sessions(tmp_path):
async def run() -> None:
manager = ExecSessionManager()
tool = ExecTool(working_dir=str(tmp_path), timeout=30, session_manager=manager)
@@ -723,14 +722,13 @@ def test_agent_loop_shutdown_closes_exec_sessions(tmp_path, monkeypatch):
sid = _session_id(initial)
process = manager._sessions[sid].process
monkeypatch.setattr(agent_context, "close_mcp", lambda _state: asyncio.sleep(0))
loop = object.__new__(AgentLoop)
loop._background_tasks = set()
loop._exec_session_manager = manager
loop.subagents = SimpleNamespace(close=AsyncMock())
await loop.close_mcp()
await loop.close_mcp()
await loop.aclose()
await loop.aclose()
assert process.returncode is not None
assert manager._sessions == {}
@@ -739,7 +737,7 @@ def test_agent_loop_shutdown_closes_exec_sessions(tmp_path, monkeypatch):
asyncio.run(run())
def test_agent_loop_shutdown_attempts_all_cleanup_after_errors(monkeypatch):
def test_agent_loop_shutdown_attempts_all_cleanup_after_errors():
async def run() -> None:
loop = object.__new__(AgentLoop)
loop._background_tasks = set()
@@ -749,16 +747,12 @@ def test_agent_loop_shutdown_attempts_all_cleanup_after_errors(monkeypatch):
loop._exec_session_manager = SimpleNamespace(
close_all=AsyncMock(side_effect=OSError("exec cleanup failed")),
)
close_mcp = AsyncMock()
monkeypatch.setattr(agent_context, "close_mcp", close_mcp)
with pytest.raises(BaseExceptionGroup) as exc_info:
await loop.close_mcp()
await loop.aclose()
assert len(exc_info.value.exceptions) == 2
loop.subagents.close.assert_awaited_once()
loop._exec_session_manager.close_all.assert_awaited_once()
close_mcp.assert_awaited_once_with(loop)
asyncio.run(run())
@@ -892,7 +886,7 @@ def test_terminate_by_owner_skips_sessions_without_owner_key(tmp_path):
asyncio.run(run())
def test_agent_loop_shutdown_preserves_single_cleanup_error(monkeypatch):
def test_agent_loop_shutdown_preserves_single_cleanup_error():
async def run() -> None:
loop = object.__new__(AgentLoop)
loop._background_tasks = set()
@@ -900,13 +894,9 @@ def test_agent_loop_shutdown_preserves_single_cleanup_error(monkeypatch):
close=AsyncMock(side_effect=RuntimeError("subagent cleanup failed")),
)
loop._exec_session_manager = SimpleNamespace(close_all=AsyncMock())
close_mcp = AsyncMock()
monkeypatch.setattr(agent_context, "close_mcp", close_mcp)
with pytest.raises(RuntimeError, match="subagent cleanup failed"):
await loop.close_mcp()
await loop.aclose()
loop._exec_session_manager.close_all.assert_awaited_once()
close_mcp.assert_awaited_once_with(loop)
asyncio.run(run())
tests/tools/test_mcp_tool.py
+76 -26
View File
@@ -13,6 +13,7 @@ import pytest
import nanobot.agent.tools.mcp as mcp_mod
from nanobot.agent.tools.mcp import (
MCPPromptWrapper,
MCPProvider,
MCPResourceWrapper,
MCPToolWrapper,
_normalize_windows_stdio_command,
@@ -153,7 +154,7 @@ def _make_wrapper(session: object, *, timeout: float = 0.1) -> MCPToolWrapper:
@pytest.mark.asyncio
async def test_connect_missing_servers_propagates_external_cancellation(monkeypatch) -> None:
async def test_mcp_provider_connect_propagates_external_cancellation(monkeypatch) -> None:
started = asyncio.Event()
async def connect_mcp_servers(_servers: dict, _registry: ToolRegistry) -> dict:
@@ -161,24 +162,21 @@ async def test_connect_missing_servers_propagates_external_cancellation(monkeypa
await asyncio.sleep(60)
return {}
class State:
pass
state = State()
state._mcp_closing = False
state._mcp_servers = {"test": MCPServerConfig(command="fake")}
state._mcp_stacks = {}
state._mcp_connecting = False
provider = MCPProvider(
{"test": MCPServerConfig(command="fake")},
ToolRegistry(),
)
monkeypatch.setattr(mcp_mod, "connect_mcp_servers", connect_mcp_servers)
task = asyncio.create_task(mcp_mod.connect_missing_servers(state, ToolRegistry()))
task = asyncio.create_task(provider.connect())
await asyncio.wait_for(started.wait(), timeout=1.0)
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
assert state._mcp_connecting is False
assert provider.connected_server_names == set()
assert provider.runtime_status() == {"test": "failed"}
@pytest.mark.asyncio
@@ -223,25 +221,20 @@ async def test_saved_oauth_http_403_projects_failed_runtime_without_details(
rejected_streamable_http,
)
class State:
pass
state = State()
state._mcp_closing = False
state._mcp_servers = {
"xmind": MCPServerConfig(
provider = MCPProvider(
{
"xmind": MCPServerConfig(
type="streamableHttp",
auth="oauth",
url="https://app.xmind.com/api/mcp",
)
}
state._mcp_stacks = {}
state._mcp_runtime_statuses = {}
state._mcp_connecting = False
)
},
ToolRegistry(),
)
await mcp_mod.connect_missing_servers(state, ToolRegistry())
await provider.connect()
snapshot = mcp_mod.runtime_status(state)
snapshot = provider.runtime_status()
assert snapshot == {"xmind": "failed"}
assert "saved-oauth-secret" not in str(snapshot)
assert "app.xmind.com" not in str(snapshot)
@@ -1263,6 +1256,59 @@ async def test_connect_mcp_servers_propagates_external_cancellation(
await asyncio.wait_for(closed.wait(), timeout=1.0)
@pytest.mark.asyncio
async def test_connect_mcp_servers_rolls_back_completed_batch_on_cancellation(
monkeypatch: pytest.MonkeyPatch,
) -> None:
slow_started = asyncio.Event()
closed: list[str] = []
sessions = {"fast": _make_fake_session(["demo"])}
class _SelectiveClientSession:
def __init__(self, read: object, _write: object) -> None:
self._session = sessions[str(read)]
async def __aenter__(self) -> object:
return self._session
async def __aexit__(self, exc_type, exc, tb) -> bool:
return False
@asynccontextmanager
async def _selective_stdio_client(params: object):
command = str(params.command)
try:
if command == "slow":
slow_started.set()
await asyncio.Event().wait()
yield command, object()
finally:
closed.append(command)
monkeypatch.setattr(sys.modules["mcp"], "ClientSession", _SelectiveClientSession)
monkeypatch.setattr(sys.modules["mcp.client.stdio"], "stdio_client", _selective_stdio_client)
registry = ToolRegistry()
task = asyncio.create_task(
connect_mcp_servers(
{
"fast": MCPServerConfig(command="fast"),
"slow": MCPServerConfig(command="slow"),
},
registry,
)
)
await asyncio.wait_for(slow_started.wait(), timeout=1.0)
assert registry.tool_names == ["mcp_fast_demo"]
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
assert registry.tool_names == []
assert sorted(closed) == ["fast", "slow"]
@pytest.mark.asyncio
async def test_connect_mcp_servers_streamable_http_uses_finite_timeout(
fake_mcp_runtime: dict[str, object | None],
@@ -1900,7 +1946,11 @@ def test_long_server_name_tools_are_matched_by_server_name() -> None:
assert len(wrapper.name) == 64
assert not wrapper.name.startswith(mcp_mod._tool_prefix(server_name))
mcp_mod._attach_reconnect_handlers(SimpleNamespace(), registry, {server_name})
provider = MCPProvider(
{server_name: MCPServerConfig(command="fake")},
registry,
)
provider._attach_reconnect_handlers({server_name})
assert wrapper._reconnect is not None
assert other_wrapper._reconnect is None
tests/webui/test_settings_routes.py
+31 -1
View File
@@ -1,7 +1,8 @@
from __future__ import annotations
import asyncio
import json
from collections.abc import Callable, Mapping
from collections.abc import Awaitable, Callable, Mapping
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import ANY, AsyncMock, MagicMock
@@ -22,6 +23,7 @@ def _router(
authorized: bool = True,
config_path: Path | None = None,
mcp_runtime_status: Callable[[], Mapping[str, str]] | None = None,
mcp_reload: Callable[[], Awaitable[dict[str, object]]] | None = None,
) -> WebUISettingsRouter:
return WebUISettingsRouter(
settings=WebUISettingsServices.create(config_path or get_config_path()),
@@ -37,6 +39,7 @@ def _router(
runtime_surface="browser",
runtime_capabilities={},
mcp_runtime_status=mcp_runtime_status,
mcp_reload=mcp_reload,
mcp_oauth_redirect_uri=lambda _request: "https://gateway.example/auth/mcp/callback",
)
@@ -89,6 +92,33 @@ async def test_mcp_list_serializes_local_runtime_failure_snapshot(tmp_path) -> N
assert snapshot_calls == 1
@pytest.mark.asyncio
async def test_mcp_reload_callback_is_bounded(
monkeypatch: pytest.MonkeyPatch,
) -> None:
started = asyncio.Event()
async def reload_mcp() -> dict[str, object]:
started.set()
await asyncio.Event().wait()
return {"ok": True}
monkeypatch.setattr(
"nanobot.webui.settings_routes._MCP_RELOAD_TIMEOUT_SECONDS",
0.01,
)
router = _router(mcp_reload=reload_mcp)
result = await router._reload_mcp_runtime()
assert started.is_set()
assert result == {
"ok": False,
"message": "MCP hot reload timed out. Restart nanobot to pick up changes.",
"requires_restart": True,
}
@pytest.mark.asyncio
async def test_mcp_oauth_start_uses_gateway_callback_and_requires_api_auth(monkeypatch) -> None:
config = SimpleNamespace(