From 9a8cdde86cdec2a2adfc3e463a738b1981680e87 Mon Sep 17 00:00:00 2001 From: Hyu Date: Fri, 3 Jul 2026 13:14:14 +0800 Subject: [PATCH] refactor(mcp): make MCP test reuse the shared Box session (near-instant tests) (#2304) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * refactor(mcp): make MCP test reuse the shared Box session instead of a per-test session Testing an MCP server (config-page "test" button) previously spun up a fresh isolated mcp-test- Box session every time: cold-start the container, run the dependency bootstrap, probe, then tear the whole session down. That is slow (tens of seconds) and, on an already-hosted server, wasteful — the server is already running in the shared session. Change the test to reuse the shared session / live process: - _build_box_session_id: transient tests now use mcp-shared, the same Box session as live servers, so a test reuses the running container (and, for an existing server, its live managed process) instead of a cold per-test session. - cleanup_session: a transient test no longer deletes the whole session (which under the shared model would kill every other MCP server in the container). It stops only its own process_id, exactly like a live server. Isolation is now at the process level (distinct process_id per server/test), not the session level. - test_mcp_server (persisted server): reuse the live connection with a real list_tools refresh/probe; only fall back to a full start() when there is no live connection to probe or the refresh fails, instead of an ERROR->start() rebuild. Trade-off: a failing test now shares the container with live servers rather than a throwaway session. Accepted deliberately in favour of near-instant tests; process-level isolation keeps a test from stopping another server's process. * chore(deps): pin langbot-plugin 0.4.9 for the nsjail RLIMIT_AS node/npx MCP fix --------- Co-authored-by: dadachann <185672915+dadachann@users.noreply.github.com> --- pyproject.toml | 2 +- src/langbot/pkg/api/http/service/mcp.py | 16 ++++++++++++++-- src/langbot/pkg/provider/tools/loaders/mcp.py | 13 ++++++++----- .../pkg/provider/tools/loaders/mcp_stdio.py | 16 ++++++++++------ .../provider/test_mcp_box_integration.py | 19 ++++++++++++------- uv.lock | 8 ++++---- 6 files changed, 49 insertions(+), 25 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index c33ba8034..b76d21c83 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.8", + "langbot-plugin==0.4.9", "asyncpg>=0.30.0", "line-bot-sdk>=3.19.0", "matrix-nio>=0.25.2", diff --git a/src/langbot/pkg/api/http/service/mcp.py b/src/langbot/pkg/api/http/service/mcp.py index 157294d48..f04137535 100644 --- a/src/langbot/pkg/api/http/service/mcp.py +++ b/src/langbot/pkg/api/http/service/mcp.py @@ -188,10 +188,22 @@ class MCPService: persisted_session = runtime_mcp_session async def _refresh_and_report() -> None: - if persisted_session.status == MCPSessionStatus.ERROR: + # Testing a persisted server should REUSE its live shared-session + # process, not rebuild it. Try a lightweight refresh (a real + # list_tools probe over the existing connection) first; only fall + # back to a full start() when the session has no live connection + # to probe (never connected, or the process is actually gone). + needs_start = persisted_session.status == MCPSessionStatus.ERROR or persisted_session.session is None + if needs_start: await persisted_session.start() else: - await persisted_session.refresh() + try: + await persisted_session.refresh() + except Exception: + # The live connection was stale/dropped: reconnect once + # (reusing the live managed process where possible) and + # re-probe, instead of reporting a false failure. + await persisted_session.start() # Surface the discovered tools so the config page can render them # even for an already-hosted server. ctx.metadata['runtime_info'] = persisted_session.get_runtime_info_dict() diff --git a/src/langbot/pkg/provider/tools/loaders/mcp.py b/src/langbot/pkg/provider/tools/loaders/mcp.py index 9fb3907ab..539910ac7 100644 --- a/src/langbot/pkg/provider/tools/loaders/mcp.py +++ b/src/langbot/pkg/provider/tools/loaders/mcp.py @@ -986,11 +986,14 @@ class RuntimeMCPSession: return self._box_stdio_runtime.uses_box_stdio() def _build_box_session_id(self) -> str: - # Transient test sessions get their own isolated Box session so a - # failing/short-lived test can never disturb the shared session that - # hosts live, already-connected MCP servers. - if self.is_transient: - return f'mcp-test-{self.server_uuid}' + # Both live servers and transient config-page tests share ONE Box + # session ('mcp-shared'). A test therefore reuses the already-running + # container (and, for an existing server, its live managed process) + # instead of paying a full per-test session cold-start + dependency + # bootstrap. Isolation between a test and the live servers is provided + # at the *process* level: each server/test has its own process_id and a + # test only ever stops its own process_id (see cleanup_session), so it + # never disturbs another server's process or the shared session itself. return 'mcp-shared' def _rewrite_path(self, path: str, host_path: str | None) -> str: diff --git a/src/langbot/pkg/provider/tools/loaders/mcp_stdio.py b/src/langbot/pkg/provider/tools/loaders/mcp_stdio.py index 4bab3712f..120538cf6 100644 --- a/src/langbot/pkg/provider/tools/loaders/mcp_stdio.py +++ b/src/langbot/pkg/provider/tools/loaders/mcp_stdio.py @@ -370,16 +370,20 @@ class BoxStdioSessionRuntime: workspace = self._build_workspace(host_path=None) - # Transient test sessions own their isolated Box session, so tear the - # whole session down rather than leaking it. This cannot affect live - # servers because they live in the separate shared session. + # Transient config-page tests now share the same 'mcp-shared' Box + # session as live servers, so we must NOT tear the session down here — + # that would kill every other MCP server in the container. A test is + # isolated at the process level: it ran under its own process_id, so we + # stop only that process, exactly like a live server does below. The + # shared session and all other servers' live processes are untouched. + # (Staged per-test workspace files are still cleaned up.) if getattr(self.owner, 'is_transient', False): try: - await workspace.cleanup() + await workspace.stop_managed_process(self.process_id) except Exception as exc: self.ap.logger.warning( - f'MCP server {self.server_name}: failed to delete transient test session ' - f'{self.owner._build_box_session_id()}: {type(exc).__name__}: {exc}' + f'MCP server {self.server_name}: failed to stop transient test process ' + f'process_id={self.process_id}: {type(exc).__name__}: {exc}' ) await self._cleanup_staged_workspace() return diff --git a/tests/unit_tests/provider/test_mcp_box_integration.py b/tests/unit_tests/provider/test_mcp_box_integration.py index 74cd2487c..f0e06b2e8 100644 --- a/tests/unit_tests/provider/test_mcp_box_integration.py +++ b/tests/unit_tests/provider/test_mcp_box_integration.py @@ -639,10 +639,13 @@ class TestGetRuntimeInfoDict: assert info['box_session_id'] == 'mcp-shared' assert info['box_enabled'] is True - def test_transient_test_session_is_isolated_from_shared(self, mcp_module): - """A transient test session (config-page "test", no persisted UUID) - must NOT share the live "mcp-shared" Box session. Regression: a failing - test churned the shared session and tore down healthy live servers.""" + def test_transient_test_shares_session_but_isolated_by_process(self, mcp_module): + """A transient config-page "test" now shares the same 'mcp-shared' Box + session as live servers (so a test reuses the running container / live + process instead of a cold per-test session bootstrap). Isolation is at + the PROCESS level: the test runs under its own process_id and only ever + stops that process_id, so it cannot disturb another server's live + process or the shared session itself.""" ap = _make_ap() ap.box_service.available = True transient = _make_session( @@ -670,10 +673,12 @@ class TestGetRuntimeInfoDict: ) assert transient.is_transient is True assert live.is_transient is False - # Isolated session id for the test, shared for the live server. - assert transient._build_box_session_id() == 'mcp-test-gen-uuid-123' + # Both share ONE Box session ... + assert transient._build_box_session_id() == 'mcp-shared' assert live._build_box_session_id() == 'mcp-shared' - assert transient._build_box_session_id() != live._build_box_session_id() + assert transient._build_box_session_id() == live._build_box_session_id() + # ... but are isolated by distinct process_ids within that session. + assert transient._box_stdio_runtime.process_id != live._box_stdio_runtime.process_id def test_stdio_session_refuses_when_box_unavailable(self, mcp_module): """Policy: when Box is configured but unavailable (disabled in config diff --git a/uv.lock b/uv.lock index 7ecac56bd..4add38bc0 100644 --- a/uv.lock +++ b/uv.lock @@ -2123,7 +2123,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.8" }, + { name = "langbot-plugin", specifier = "==0.4.9" }, { name = "langchain", specifier = ">=1.3.9" }, { name = "langchain-core", specifier = ">=1.3.3" }, { name = "langchain-text-splitters", specifier = ">=1.1.2" }, @@ -2187,7 +2187,7 @@ dev = [ [[package]] name = "langbot-plugin" -version = "0.4.8" +version = "0.4.9" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "aiofiles" }, @@ -2208,9 +2208,9 @@ dependencies = [ { name = "watchdog" }, { name = "websockets" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/7f/f5/1b417a48d329f959a261eaf9a37b26197f18081d2b6f0ef3e4bbab687b81/langbot_plugin-0.4.8.tar.gz", hash = "sha256:e9614be248acf352acb9430d9e154e0abae981575de16812f68282c92a4e42d8", size = 334076, upload-time = "2026-07-03T04:15:28.233Z" } +sdist = { url = "https://files.pythonhosted.org/packages/54/52/f85939d3be929696dc5f9f3369f18bd6652981a8985c6e6f0dfe48e04f75/langbot_plugin-0.4.9.tar.gz", hash = "sha256:1cc2882f1be96cd52be6c392c3925e44cf785c97c10f82a321339636da318c6e", size = 334682, upload-time = "2026-07-03T04:57:42.813Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/40/21/46bfa911ab27d6856a3b4363805c4bb8a657500cbc12e1e719d0b3b349dc/langbot_plugin-0.4.8-py3-none-any.whl", hash = "sha256:87c12d9d5579b0ab4dc43c9e559b4d15d7387db8c8b52d476cf27a3de74aedf5", size = 221244, upload-time = "2026-07-03T04:15:27.087Z" }, + { url = "https://files.pythonhosted.org/packages/47/12/1cf535377a81cf01607bbee5e3f4745a8c9f6dbb561887467001a182a533/langbot_plugin-0.4.9-py3-none-any.whl", hash = "sha256:693225a089ba38bc8d0b58d9e0fdf07b027b5be6b5e858cf7b4e56d095b68ba5", size = 221722, upload-time = "2026-07-03T04:57:41.578Z" }, ] [[package]]