fix(runtime): complete reconnect recovery paths

This commit is contained in:
Junyan Qin
2026-07-23 17:28:31 +08:00
parent 9cbbaf617b
commit 06e03af994
14 changed files with 223 additions and 34 deletions
+1 -1
View File
@@ -70,7 +70,7 @@ dependencies = [
"chromadb>=1.0.0,<2.0.0", "chromadb>=1.0.0,<2.0.0",
"qdrant-client (>=1.15.1,<2.0.0)", "qdrant-client (>=1.15.1,<2.0.0)",
"pyseekdb==1.1.0.post3", "pyseekdb==1.1.0.post3",
"langbot-plugin==0.4.16", "langbot-plugin==0.4.17",
"asyncpg>=0.30.0", "asyncpg>=0.30.0",
"line-bot-sdk>=3.19.0", "line-bot-sdk>=3.19.0",
"matrix-nio>=0.25.2", "matrix-nio>=0.25.2",
+3
View File
@@ -189,6 +189,9 @@ class BoxRuntimeConnector(ManagedRuntimeConnector):
await self._close_managed_subprocess() await self._close_managed_subprocess()
raise raise
if self._heartbeat_task is None or self._heartbeat_task.done():
self._heartbeat_task = asyncio.create_task(self._heartbeat_loop())
# -- heartbeat ----------------------------------------------------------- # -- heartbeat -----------------------------------------------------------
async def _heartbeat_loop(self) -> None: async def _heartbeat_loop(self) -> None:
+6
View File
@@ -143,8 +143,14 @@ class BoxService:
await asyncio.sleep(delay) await asyncio.sleep(delay)
try: try:
await connector.reconnect() await connector.reconnect()
self._ensure_default_workspace()
await self._purge_attachment_dirs()
self._available = True self._available = True
self._connector_error = '' 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.') self.ap.logger.info('Box runtime reconnected, sandbox features restored.')
return return
except Exception as exc: except Exception as exc:
+17 -5
View File
@@ -307,10 +307,10 @@ class RuntimeMCPSession:
await self._box_stdio_runtime.initialize() await self._box_stdio_runtime.initialize()
return return
# Box is configured (ap.box_service exists) but currently unavailable # Box is configured but explicitly disabled. Refuse stdio MCP rather
# (disabled by config or connection failed). Refuse stdio MCP rather # than silently falling through to host-stdio — the operator asked for
# than silently falling through to host-stdio — the operator asked # the sandbox and the failure mode should be visible. An enabled Box
# for the sandbox and the failure mode should be visible. # that is reconnecting is handled above and waits for availability.
# #
# Set ``error_phase = BOX_UNAVAILABLE`` BEFORE raising so the retry # Set ``error_phase = BOX_UNAVAILABLE`` BEFORE raising so the retry
# wrapper can short-circuit (retrying is pointless when Box is # wrapper can short-circuit (retrying is pointless when Box is
@@ -1823,11 +1823,23 @@ class MCPLoader(loader.ToolLoader):
async def shutdown(self): async def shutdown(self):
"""关闭所有工具""" """关闭所有工具"""
self.ap.logger.info('Shutting down all MCP sessions...') 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: try:
await session.shutdown() await session.shutdown()
self.ap.logger.debug(f'Shutdown MCP session: {server_name}') self.ap.logger.debug(f'Shutdown MCP session: {server_name}')
except Exception as e: except Exception as e:
self.ap.logger.error(f'Error shutting down MCP session {server_name}: {e}\n{traceback.format_exc()}') 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.sessions.clear()
self.ap.logger.info('All MCP sessions shutdown complete') self.ap.logger.info('All MCP sessions shutdown complete')
@@ -185,11 +185,10 @@ class BoxStdioSessionRuntime:
box_service = getattr(self.ap, 'box_service', None) box_service = getattr(self.ap, 'box_service', None)
if box_service is None: if box_service is None:
return False return False
# When Box is configured but currently unavailable (disabled or # An enabled Box service remains the required transport while it is
# connection failed), do NOT silently fall through to host-stdio — # reconnecting. initialize() waits for availability instead of
# that would bypass the sandbox the operator asked for. The caller # permanently failing the MCP server or falling through to host stdio.
# is expected to refuse the stdio MCP server with a clear error. return bool(getattr(box_service, 'enabled', True))
return bool(getattr(box_service, 'available', False))
async def initialize(self) -> None: async def initialize(self) -> None:
await self._wait_for_box_runtime() await self._wait_for_box_runtime()
@@ -488,7 +487,7 @@ class BoxStdioSessionRuntime:
) )
warned = True warned = True
if asyncio.get_running_loop().time() >= deadline: 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') raise Exception(f'Box runtime is not available after {int(timeout_sec)} seconds')
await asyncio.sleep(1) await asyncio.sleep(1)
@@ -57,7 +57,7 @@ class NativeToolLoader(loader.ToolLoader):
return await is_box_backend_available(self.ap) return await is_box_backend_available(self.ap)
async def get_tools(self, bound_plugins: list[str] | None = None) -> list[resource_tool.LLMTool]: 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 [] return []
if self._tools is None: if self._tools is None:
self._tools = [ self._tools = [
@@ -71,7 +71,7 @@ class NativeToolLoader(loader.ToolLoader):
return list(self._tools) return list(self._tools)
async def has_tool(self, name: str) -> bool: 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): async def invoke_tool(self, name: str, parameters: dict, query: pipeline_query.Query):
if name == EXEC_TOOL_NAME: if name == EXEC_TOOL_NAME:
@@ -652,12 +652,9 @@ else:
if callable(refresh_skill): if callable(refresh_skill):
refresh_skill(selected_skill.get('name', '')) refresh_skill(selected_skill.get('name', ''))
def _is_sandbox_available(self) -> bool: async def _is_sandbox_available(self) -> bool:
"""Check if sandbox backend is available. """Refresh backend availability so Box reconnects restore tool exposure."""
self._backend_available = await self._check_backend_available()
This checks the cached backend availability from initialization,
not just whether the box_service process is running.
"""
return bool(self._backend_available) return bool(self._backend_available)
def _build_exec_tool(self) -> resource_tool.LLMTool: def _build_exec_tool(self) -> resource_tool.LLMTool:
@@ -49,19 +49,27 @@ class SkillToolLoader(loader.ToolLoader):
return await is_box_backend_available(self.ap) return await is_box_backend_available(self.ap)
async def get_tools(self, bound_plugins: list[str] | None = None) -> list[resource_tool.LLMTool]: 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 [] return []
if not self._tools:
self._tools = [
self._build_activate_skill_tool(),
self._build_register_skill_tool(),
]
return list(self._tools) return list(self._tools)
async def has_tool(self, name: str) -> bool: 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. """Check if skill tools should be available.
Skill tools require both a skill manager and a sandbox backend. 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: async def invoke_tool(self, name: str, parameters: dict, query) -> typing.Any:
if name == ACTIVATE_SKILL_TOOL_NAME: if name == ACTIVATE_SKILL_TOOL_NAME:
@@ -126,6 +126,29 @@ async def test_box_runtime_connector_cleans_partial_transport_on_connect_failure
connector._close_managed_subprocess.assert_awaited_once() 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 @pytest.mark.asyncio
async def test_box_stdio_connection_does_not_capture_unconsumed_stderr( async def test_box_stdio_connection_does_not_capture_unconsumed_stderr(
monkeypatch: pytest.MonkeyPatch, monkeypatch: pytest.MonkeyPatch,
+22
View File
@@ -347,6 +347,28 @@ async def test_box_service_shutdown_reaps_connector_when_runtime_rpc_is_offline(
connector.aclose.assert_awaited_once() 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 @pytest.mark.asyncio
async def test_box_runtime_reuses_request_session(): async def test_box_runtime_reuses_request_session():
logger = Mock() logger = Mock()
@@ -725,14 +725,10 @@ class TestGetRuntimeInfoDict:
) )
assert writable._build_box_session_id() != default._build_box_session_id() assert writable._build_box_session_id() != default._build_box_session_id()
def test_stdio_session_refuses_when_box_unavailable(self, mcp_module): def test_stdio_session_waits_for_enabled_box_when_temporarily_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."""
ap = _make_ap() ap = _make_ap()
ap.box_service.available = False ap.box_service.available = False
ap.box_service.enabled = True
s = _make_session( s = _make_session(
mcp_module, mcp_module,
{ {
@@ -745,9 +741,58 @@ class TestGetRuntimeInfoDict:
ap=ap, ap=ap,
) )
info = s.get_runtime_info_dict() 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_session_id' not in info
assert 'box_enabled' 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): def test_stdio_session_without_box_service_uses_local_stdio(self, mcp_module):
ap = _make_ap() ap = _make_ap()
del ap.box_service del ap.box_service
@@ -1,5 +1,6 @@
from __future__ import annotations from __future__ import annotations
import asyncio
import base64 import base64
from types import SimpleNamespace from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock 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 assert query.variables[MCP_RESOURCE_CONTEXT_QUERY_KEY]['resource_count'] == 1
docs.session.read_resource.assert_awaited_once() docs.session.read_resource.assert_awaited_once()
other.session.read_resource.assert_not_called() 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 == {}
@@ -458,6 +458,25 @@ class TestSkillToolLoader:
assert await loader.has_tool('activate') is True assert await loader.has_tool('activate') is True
assert await loader.has_tool('register_skill') 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: class TestNativeToolLoaderSkillPaths:
@pytest.mark.asyncio @pytest.mark.asyncio
@@ -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 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 ───────────────────────────── # ── read/write/edit file tool tests ─────────────────────────────
Generated
+4 -4
View File
@@ -2124,7 +2124,7 @@ requires-dist = [
{ name = "ebooklib", specifier = ">=0.18" }, { name = "ebooklib", specifier = ">=0.18" },
{ name = "gewechat-client", specifier = ">=0.1.5" }, { name = "gewechat-client", specifier = ">=0.1.5" },
{ name = "html2text", specifier = ">=2024.2.26" }, { 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", specifier = ">=1.3.9" },
{ name = "langchain-core", specifier = ">=1.3.3" }, { name = "langchain-core", specifier = ">=1.3.3" },
{ name = "langchain-text-splitters", specifier = ">=1.1.2" }, { name = "langchain-text-splitters", specifier = ">=1.1.2" },
@@ -2189,7 +2189,7 @@ dev = [
[[package]] [[package]]
name = "langbot-plugin" name = "langbot-plugin"
version = "0.4.16" version = "0.4.17"
source = { registry = "https://pypi.org/simple" } source = { registry = "https://pypi.org/simple" }
dependencies = [ dependencies = [
{ name = "aiofiles" }, { name = "aiofiles" },
@@ -2210,9 +2210,9 @@ dependencies = [
{ name = "watchdog" }, { name = "watchdog" },
{ name = "websockets" }, { 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 = [ 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]] [[package]]