diff --git a/pyproject.toml b/pyproject.toml index 031157a53..53425aa01 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -70,7 +70,7 @@ dependencies = [ "chromadb>=1.0.0,<2.0.0", "qdrant-client (>=1.15.1,<2.0.0)", "pyseekdb==1.1.0.post3", - "langbot-plugin==0.4.16", + "langbot-plugin==0.4.17", "asyncpg>=0.30.0", "line-bot-sdk>=3.19.0", "matrix-nio>=0.25.2", diff --git a/src/langbot/pkg/box/connector.py b/src/langbot/pkg/box/connector.py index d56270105..f24287477 100644 --- a/src/langbot/pkg/box/connector.py +++ b/src/langbot/pkg/box/connector.py @@ -189,6 +189,9 @@ class BoxRuntimeConnector(ManagedRuntimeConnector): await self._close_managed_subprocess() raise + if self._heartbeat_task is None or self._heartbeat_task.done(): + self._heartbeat_task = asyncio.create_task(self._heartbeat_loop()) + # -- heartbeat ----------------------------------------------------------- async def _heartbeat_loop(self) -> None: diff --git a/src/langbot/pkg/box/service.py b/src/langbot/pkg/box/service.py index 2b900df89..809a94130 100644 --- a/src/langbot/pkg/box/service.py +++ b/src/langbot/pkg/box/service.py @@ -143,8 +143,14 @@ class BoxService: await asyncio.sleep(delay) try: await connector.reconnect() + self._ensure_default_workspace() + await self._purge_attachment_dirs() self._available = True self._connector_error = '' + skill_mgr = getattr(self.ap, 'skill_mgr', None) + reload_skills = getattr(skill_mgr, 'reload_skills', None) + if callable(reload_skills): + await reload_skills() self.ap.logger.info('Box runtime reconnected, sandbox features restored.') return except Exception as exc: diff --git a/src/langbot/pkg/provider/tools/loaders/mcp.py b/src/langbot/pkg/provider/tools/loaders/mcp.py index edb25b03b..07f4349a9 100644 --- a/src/langbot/pkg/provider/tools/loaders/mcp.py +++ b/src/langbot/pkg/provider/tools/loaders/mcp.py @@ -307,10 +307,10 @@ class RuntimeMCPSession: await self._box_stdio_runtime.initialize() return - # Box is configured (ap.box_service exists) but currently unavailable - # (disabled by config or connection failed). Refuse stdio MCP rather - # than silently falling through to host-stdio — the operator asked - # for the sandbox and the failure mode should be visible. + # Box is configured but explicitly disabled. Refuse stdio MCP rather + # than silently falling through to host-stdio — the operator asked for + # the sandbox and the failure mode should be visible. An enabled Box + # that is reconnecting is handled above and waits for availability. # # Set ``error_phase = BOX_UNAVAILABLE`` BEFORE raising so the retry # wrapper can short-circuit (retrying is pointless when Box is @@ -1823,11 +1823,23 @@ class MCPLoader(loader.ToolLoader): async def shutdown(self): """关闭所有工具""" self.ap.logger.info('Shutting down all MCP sessions...') - for server_name, session in list(self.sessions.items()): + + hosted_tasks = [task for task in self._hosted_mcp_tasks if not task.done()] + for task in hosted_tasks: + task.cancel() + if hosted_tasks: + await asyncio.gather(*hosted_tasks, return_exceptions=True) + self._hosted_mcp_tasks.clear() + + async def shutdown_session(server_name: str, session: RuntimeMCPSession) -> None: try: await session.shutdown() self.ap.logger.debug(f'Shutdown MCP session: {server_name}') except Exception as e: self.ap.logger.error(f'Error shutting down MCP session {server_name}: {e}\n{traceback.format_exc()}') + + await asyncio.gather( + *(shutdown_session(server_name, session) for server_name, session in list(self.sessions.items())) + ) self.sessions.clear() self.ap.logger.info('All MCP sessions shutdown complete') diff --git a/src/langbot/pkg/provider/tools/loaders/mcp_stdio.py b/src/langbot/pkg/provider/tools/loaders/mcp_stdio.py index 134c60a88..8c956fda3 100644 --- a/src/langbot/pkg/provider/tools/loaders/mcp_stdio.py +++ b/src/langbot/pkg/provider/tools/loaders/mcp_stdio.py @@ -185,11 +185,10 @@ class BoxStdioSessionRuntime: box_service = getattr(self.ap, 'box_service', None) if box_service is None: return False - # When Box is configured but currently unavailable (disabled or - # connection failed), do NOT silently fall through to host-stdio — - # that would bypass the sandbox the operator asked for. The caller - # is expected to refuse the stdio MCP server with a clear error. - return bool(getattr(box_service, 'available', False)) + # An enabled Box service remains the required transport while it is + # reconnecting. initialize() waits for availability instead of + # permanently failing the MCP server or falling through to host stdio. + return bool(getattr(box_service, 'enabled', True)) async def initialize(self) -> None: await self._wait_for_box_runtime() @@ -488,7 +487,7 @@ class BoxStdioSessionRuntime: ) warned = True if asyncio.get_running_loop().time() >= deadline: - self.owner.error_phase = MCPSessionErrorPhase.SESSION_CREATE + self.owner.error_phase = MCPSessionErrorPhase.BOX_UNAVAILABLE raise Exception(f'Box runtime is not available after {int(timeout_sec)} seconds') await asyncio.sleep(1) diff --git a/src/langbot/pkg/provider/tools/loaders/native.py b/src/langbot/pkg/provider/tools/loaders/native.py index 7f5ee4226..a1a8c027b 100644 --- a/src/langbot/pkg/provider/tools/loaders/native.py +++ b/src/langbot/pkg/provider/tools/loaders/native.py @@ -57,7 +57,7 @@ class NativeToolLoader(loader.ToolLoader): return await is_box_backend_available(self.ap) async def get_tools(self, bound_plugins: list[str] | None = None) -> list[resource_tool.LLMTool]: - if not self._is_sandbox_available(): + if not await self._is_sandbox_available(): return [] if self._tools is None: self._tools = [ @@ -71,7 +71,7 @@ class NativeToolLoader(loader.ToolLoader): return list(self._tools) async def has_tool(self, name: str) -> bool: - return name in _ALL_TOOL_NAMES and self._is_sandbox_available() + return name in _ALL_TOOL_NAMES and await self._is_sandbox_available() async def invoke_tool(self, name: str, parameters: dict, query: pipeline_query.Query): if name == EXEC_TOOL_NAME: @@ -652,12 +652,9 @@ else: if callable(refresh_skill): refresh_skill(selected_skill.get('name', '')) - def _is_sandbox_available(self) -> bool: - """Check if sandbox backend is available. - - This checks the cached backend availability from initialization, - not just whether the box_service process is running. - """ + async def _is_sandbox_available(self) -> bool: + """Refresh backend availability so Box reconnects restore tool exposure.""" + self._backend_available = await self._check_backend_available() return bool(self._backend_available) def _build_exec_tool(self) -> resource_tool.LLMTool: diff --git a/src/langbot/pkg/provider/tools/loaders/skill_authoring.py b/src/langbot/pkg/provider/tools/loaders/skill_authoring.py index d53721785..fe0f65620 100644 --- a/src/langbot/pkg/provider/tools/loaders/skill_authoring.py +++ b/src/langbot/pkg/provider/tools/loaders/skill_authoring.py @@ -49,19 +49,27 @@ class SkillToolLoader(loader.ToolLoader): return await is_box_backend_available(self.ap) async def get_tools(self, bound_plugins: list[str] | None = None) -> list[resource_tool.LLMTool]: - if not self._is_available(): + if not await self._is_available(): return [] + if not self._tools: + self._tools = [ + self._build_activate_skill_tool(), + self._build_register_skill_tool(), + ] return list(self._tools) async def has_tool(self, name: str) -> bool: - return self._is_available() and name in SKILL_TOOL_NAMES + return await self._is_available() and name in SKILL_TOOL_NAMES - def _is_available(self) -> bool: + async def _is_available(self) -> bool: """Check if skill tools should be available. Skill tools require both a skill manager and a sandbox backend. """ - return self._has_skill_manager() and self._sandbox_available + if not self._has_skill_manager(): + return False + self._sandbox_available = await self._check_sandbox_available() + return self._sandbox_available async def invoke_tool(self, name: str, parameters: dict, query) -> typing.Any: if name == ACTIVATE_SKILL_TOOL_NAME: diff --git a/tests/unit_tests/box/test_box_connector.py b/tests/unit_tests/box/test_box_connector.py index 9cb7d6249..91016291d 100644 --- a/tests/unit_tests/box/test_box_connector.py +++ b/tests/unit_tests/box/test_box_connector.py @@ -126,6 +126,29 @@ async def test_box_runtime_connector_cleans_partial_transport_on_connect_failure connector._close_managed_subprocess.assert_awaited_once() +@pytest.mark.asyncio +async def test_box_runtime_connector_starts_heartbeat_after_reconnect( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr('langbot.pkg.utils.platform.get_platform', lambda: 'linux') + monkeypatch.setattr('langbot.pkg.utils.platform.standalone_box', False) + connector = BoxRuntimeConnector(make_app(Mock())) + connector._start_local_stdio = AsyncMock(side_effect=[RuntimeError('bind failed'), None]) + connector._stop_transport = AsyncMock() + connector._close_managed_subprocess = AsyncMock() + + with pytest.raises(RuntimeError, match='bind failed'): + await connector.initialize() + + assert connector._heartbeat_task is None + + await connector.reconnect() + + assert connector._heartbeat_task is not None + assert not connector._heartbeat_task.done() + await connector.aclose() + + @pytest.mark.asyncio async def test_box_stdio_connection_does_not_capture_unconsumed_stderr( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unit_tests/box/test_box_service.py b/tests/unit_tests/box/test_box_service.py index 5c8c1da85..4737dab15 100644 --- a/tests/unit_tests/box/test_box_service.py +++ b/tests/unit_tests/box/test_box_service.py @@ -347,6 +347,28 @@ async def test_box_service_shutdown_reaps_connector_when_runtime_rpc_is_offline( connector.aclose.assert_awaited_once() +@pytest.mark.asyncio +async def test_box_service_reconnect_restores_workspace_and_runs_cleanup( + monkeypatch: pytest.MonkeyPatch, +): + app = make_app(Mock()) + app.skill_mgr = SimpleNamespace(reload_skills=AsyncMock()) + service = BoxService(app, client=Mock(spec=BoxRuntimeClient)) + connector = Mock() + connector.reconnect = AsyncMock() + service._ensure_default_workspace = Mock() + service._purge_attachment_dirs = AsyncMock() + monkeypatch.setattr('langbot.pkg.box.service.asyncio.sleep', AsyncMock()) + + await service._reconnect_loop(connector) + + connector.reconnect.assert_awaited_once() + service._ensure_default_workspace.assert_called_once() + service._purge_attachment_dirs.assert_awaited_once() + app.skill_mgr.reload_skills.assert_awaited_once() + assert service.available is True + + @pytest.mark.asyncio async def test_box_runtime_reuses_request_session(): logger = Mock() diff --git a/tests/unit_tests/provider/test_mcp_box_integration.py b/tests/unit_tests/provider/test_mcp_box_integration.py index 4a790908b..b234218bf 100644 --- a/tests/unit_tests/provider/test_mcp_box_integration.py +++ b/tests/unit_tests/provider/test_mcp_box_integration.py @@ -725,14 +725,10 @@ class TestGetRuntimeInfoDict: ) assert writable._build_box_session_id() != default._build_box_session_id() - def test_stdio_session_refuses_when_box_unavailable(self, mcp_module): - """Policy: when Box is configured but unavailable (disabled in config - OR connection failed), stdio MCP servers are NOT treated as box-stdio. - ``_init_stdio_python_server`` will raise a clear refusal at start - time; until then, the runtime info simply omits box_session_id so the - UI can render the disabled state cleanly.""" + def test_stdio_session_waits_for_enabled_box_when_temporarily_unavailable(self, mcp_module): ap = _make_ap() ap.box_service.available = False + ap.box_service.enabled = True s = _make_session( mcp_module, { @@ -745,9 +741,58 @@ class TestGetRuntimeInfoDict: ap=ap, ) info = s.get_runtime_info_dict() + assert info['box_session_id'] == 'mcp-shared' + assert info['box_enabled'] is True + + def test_stdio_session_refuses_when_box_is_disabled(self, mcp_module): + ap = _make_ap() + ap.box_service.available = False + ap.box_service.enabled = False + s = _make_session( + mcp_module, + { + 'name': 'test', + 'uuid': 'test-uuid', + 'mode': 'stdio', + 'command': 'python', + 'args': [], + }, + ap=ap, + ) + + info = s.get_runtime_info_dict() + assert 'box_session_id' not in info assert 'box_enabled' not in info + @pytest.mark.asyncio + async def test_stdio_session_waits_until_box_reconnects(self, mcp_module, monkeypatch): + mcp_stdio_module = sys.modules['langbot.pkg.provider.tools.loaders.mcp_stdio'] + ap = _make_ap() + ap.box_service.available = False + ap.box_service.enabled = True + session = _make_session( + mcp_module, + { + 'name': 'test', + 'uuid': 'test-uuid', + 'mode': 'stdio', + 'command': 'python', + 'args': [], + 'box': {'startup_timeout_sec': 1}, + }, + ap=ap, + ) + + async def reconnect_box(_delay): + ap.box_service.available = True + + monkeypatch.setattr(mcp_stdio_module.asyncio, 'sleep', reconnect_box) + + await session._box_stdio_runtime._wait_for_box_runtime() + + assert ap.box_service.available is True + def test_stdio_session_without_box_service_uses_local_stdio(self, mcp_module): ap = _make_ap() del ap.box_service diff --git a/tests/unit_tests/provider/test_mcp_resources.py b/tests/unit_tests/provider/test_mcp_resources.py index 773efe71c..05f74b0f0 100644 --- a/tests/unit_tests/provider/test_mcp_resources.py +++ b/tests/unit_tests/provider/test_mcp_resources.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import base64 from types import SimpleNamespace from unittest.mock import AsyncMock, Mock @@ -321,3 +322,42 @@ async def test_build_resource_context_for_query_uses_only_bound_attached_text_re assert query.variables[MCP_RESOURCE_CONTEXT_QUERY_KEY]['resource_count'] == 1 docs.session.read_resource.assert_awaited_once() other.session.read_resource.assert_not_called() + + +@pytest.mark.asyncio +async def test_mcp_loader_shutdown_cancels_startup_tasks_and_closes_sessions_concurrently(): + loader = MCPLoader(_app()) + hosted_cancelled = asyncio.Event() + + async def pending_host(): + try: + await asyncio.Event().wait() + finally: + hosted_cancelled.set() + + hosted_task = asyncio.create_task(pending_host()) + await asyncio.sleep(0) + loader._hosted_mcp_tasks = [hosted_task] + + started: set[str] = set() + all_started = asyncio.Event() + + class Session: + def __init__(self, name: str): + self.name = name + + async def shutdown(self): + started.add(self.name) + if len(started) == 2: + all_started.set() + await all_started.wait() + + loader.sessions = {'one': Session('one'), 'two': Session('two')} + + await asyncio.wait_for(loader.shutdown(), timeout=1) + + assert hosted_cancelled.is_set() + assert hosted_task.cancelled() + assert started == {'one', 'two'} + assert loader._hosted_mcp_tasks == [] + assert loader.sessions == {} diff --git a/tests/unit_tests/provider/test_skill_tools.py b/tests/unit_tests/provider/test_skill_tools.py index 9db7b945e..405f11d17 100644 --- a/tests/unit_tests/provider/test_skill_tools.py +++ b/tests/unit_tests/provider/test_skill_tools.py @@ -458,6 +458,25 @@ class TestSkillToolLoader: assert await loader.has_tool('activate') is True assert await loader.has_tool('register_skill') is True + @pytest.mark.asyncio + async def test_tools_reappear_after_box_backend_recovers(self): + from langbot.pkg.provider.tools.loaders.skill_authoring import SkillToolLoader + + ap = _make_ap() + ap.skill_mgr = SimpleNamespace(skills={'demo': _make_skill_data(name='demo')}) + ap.box_service = SimpleNamespace( + available=False, + get_status=AsyncMock(return_value={'backend': {'available': True}}), + ) + + loader = SkillToolLoader(ap) + await loader.initialize() + assert await loader.get_tools() == [] + + ap.box_service.available = True + + assert sorted(tool.name for tool in await loader.get_tools()) == ['activate', 'register_skill'] + class TestNativeToolLoaderSkillPaths: @pytest.mark.asyncio diff --git a/tests/unit_tests/provider/test_tool_manager_native.py b/tests/unit_tests/provider/test_tool_manager_native.py index 2085a9e8c..d73c66e4f 100644 --- a/tests/unit_tests/provider/test_tool_manager_native.py +++ b/tests/unit_tests/provider/test_tool_manager_native.py @@ -151,6 +151,21 @@ async def test_native_tool_loader_exposes_all_tools_when_box_available(): assert await loader.has_tool(tool_name) is True +@pytest.mark.asyncio +async def test_native_tool_loader_refreshes_after_box_recovers(): + box_service = SimpleNamespace( + available=False, + get_status=AsyncMock(return_value={'backend': {'available': True}}), + ) + loader = NativeToolLoader(SimpleNamespace(box_service=box_service, logger=Mock())) + await loader.initialize() + assert await loader.get_tools() == [] + + box_service.available = True + + assert [tool.name for tool in await loader.get_tools()] == ['exec', 'read', 'write', 'edit', 'glob', 'grep'] + + # ── read/write/edit file tool tests ───────────────────────────── diff --git a/uv.lock b/uv.lock index 078aa24db..cc71d0888 100644 --- a/uv.lock +++ b/uv.lock @@ -2124,7 +2124,7 @@ requires-dist = [ { name = "ebooklib", specifier = ">=0.18" }, { name = "gewechat-client", specifier = ">=0.1.5" }, { name = "html2text", specifier = ">=2024.2.26" }, - { name = "langbot-plugin", specifier = "==0.4.16" }, + { name = "langbot-plugin", specifier = "==0.4.17" }, { name = "langchain", specifier = ">=1.3.9" }, { name = "langchain-core", specifier = ">=1.3.3" }, { name = "langchain-text-splitters", specifier = ">=1.1.2" }, @@ -2189,7 +2189,7 @@ dev = [ [[package]] name = "langbot-plugin" -version = "0.4.16" +version = "0.4.17" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "aiofiles" }, @@ -2210,9 +2210,9 @@ dependencies = [ { name = "watchdog" }, { name = "websockets" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/a6/7e/b01ab601b25dcff1e6a1a713ccfa212c13fdadcf2149092ded66c1d4daf9/langbot_plugin-0.4.16.tar.gz", hash = "sha256:9c587ff30bca4b3e718d1b2b1692a54f9a12f1be57a2f160844b308006a95a4b", size = 346875, upload-time = "2026-07-23T08:45:50.950Z" } +sdist = { url = "https://files.pythonhosted.org/packages/3a/96/336d7ac97ff5a9c413e7aaf0c329a7047ec91dec5fa4f9e0b1d0abb57d0e/langbot_plugin-0.4.17.tar.gz", hash = "sha256:b1539444f16568c0b3244f68c498ac9498c56e2c20f12e220e6ee735b8d6b5c2", size = 347080, upload-time = "2026-07-23T09:23:55.941Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/70/9f/67e80a60d33cc38173edbaa2c31e984b3e9e838b7b14a76abeb0f0e7a281/langbot_plugin-0.4.16-py3-none-any.whl", hash = "sha256:330d4d930b601a6fb0f49b6360963a17134bd6fe96f23f16855112faaf1b5a0d", size = 229855, upload-time = "2026-07-23T08:45:49.479Z" }, + { url = "https://files.pythonhosted.org/packages/e1/98/ac372a238e59894ac54189cb220a3fb75da48f319881eeb49d583712549a/langbot_plugin-0.4.17-py3-none-any.whl", hash = "sha256:6fd6c9d7e0583f04659f023a233c4201cb4c450d1ef696e79fa1e338933ee156", size = 229881, upload-time = "2026-07-23T09:23:54.582Z" }, ] [[package]]