mirror/nanobot
1
0
Fork 0
mirror of https://github.com/HKUDS/nanobot.git synced 2026-08-15 00:29:52 +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."""