feat(tenancy): connect cloud workspace control plane

tests/unit_tests/provider/test_localagent_no_duplicate.py

@@ -10,6 +10,7 @@ import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
 import langbot_plugin.api.entities.builtin.provider.message as provider_message
 import langbot_plugin.api.entities.builtin.provider.session as provider_session
 
+from langbot.pkg.api.http.context import ExecutionContext, PrincipalContext, PrincipalType
 from langbot.pkg.provider.runners.localagent import LocalAgentRunner
 
 
@@ -97,7 +98,7 @@ def make_query() -> pipeline_query.Query:
     adapter = AsyncMock()
     adapter.is_stream_output_supported = AsyncMock(return_value=False)
 
-    return pipeline_query.Query.model_construct(
+    query = pipeline_query.Query.model_construct(
         query_id='no-dup-query',
         launcher_type=provider_session.LauncherTypes.PERSON,
         launcher_id=12345,
@@ -124,6 +125,17 @@ def make_query() -> pipeline_query.Query:
         use_llm_model_uuid='test-model-uuid',
         variables={},
     )
+    object.__setattr__(
+        query,
+        '_execution_context',
+        ExecutionContext(
+            instance_uuid='instance-test',
+            workspace_uuid='workspace-test',
+            placement_generation=1,
+            trigger_principal=PrincipalContext(PrincipalType.SYSTEM),
+        ),
+    )
+    return query
 
 
 def _make_app(provider) -> SimpleNamespace:

tests/unit_tests/provider/test_mcp_box_integration.py

@@ -829,7 +829,7 @@ class TestGetRuntimeInfoDict:
         assert ap.box_service.available is True
 
     @pytest.mark.asyncio
-    async def test_enabled_box_timeout_does_not_exhaust_mcp_retry_budget(self, mcp_module, monkeypatch):
+    async def test_enabled_box_timeout_does_not_exhaust_mcp_retry_budget(self, mcp_module):
         ap = _make_ap()
         ap.box_service.available = False
         ap.box_service.enabled = True
@@ -854,14 +854,13 @@ class TestGetRuntimeInfoDict:
                 raise RuntimeError('Box runtime is not available after 1 seconds')
 
         session._lifecycle_loop = lifecycle
-        sleep = AsyncMock()
-        monkeypatch.setattr(mcp_module.asyncio, 'sleep', sleep)
+        session._sleep_with_execution_fence = AsyncMock()
 
         await session._lifecycle_loop_with_retry()
 
         assert attempts == 2
         assert session.retry_count == 0
-        sleep.assert_awaited_once_with(1)
+        session._sleep_with_execution_fence.assert_awaited_once_with(1)
 
     @pytest.mark.asyncio
     async def test_disabled_box_still_stops_mcp_retry_loop(self, mcp_module):

tests/unit_tests/provider/test_mcp_remote_transport.py

@@ -13,7 +13,7 @@ from aiohttp import web
 from mcp import types as mcp_types
 
 from langbot.pkg.api.http.context import ExecutionContext
-from langbot.pkg.provider.tools.loaders.mcp import RuntimeMCPSession
+from langbot.pkg.provider.tools.loaders.mcp import MCPToolCallTimeoutError, RuntimeMCPSession
 
 
 TEST_EXECUTION_CONTEXT = ExecutionContext(
@@ -159,7 +159,21 @@ def _session(
     timeout: float = 2,
     tool_call_timeout_sec: float = 300,
 ) -> RuntimeMCPSession:
-    app = cast(Any, SimpleNamespace(logger=Mock()))
+    app = cast(
+        Any,
+        SimpleNamespace(
+            logger=Mock(),
+            workspace_service=SimpleNamespace(
+                get_execution_binding=AsyncMock(
+                    return_value=SimpleNamespace(
+                        instance_uuid=TEST_EXECUTION_CONTEXT.instance_uuid,
+                        workspace_uuid=TEST_EXECUTION_CONTEXT.workspace_uuid,
+                        placement_generation=TEST_EXECUTION_CONTEXT.placement_generation,
+                    )
+                )
+            ),
+        ),
+    )
     return RuntimeMCPSession(
         'remote-transport-test',
         {

tests/unit_tests/provider/test_mcp_resources.py

@@ -124,6 +124,7 @@ async def test_invoke_mcp_tool_uses_configurable_request_timeout():
         },
         True,
         _app(),
+        TEST_EXECUTION_CONTEXT,
     )
     session.session = SimpleNamespace(call_tool=AsyncMock(return_value=_tool_result()))
 
@@ -148,6 +149,7 @@ async def test_invoke_mcp_tool_zero_timeout_disables_request_deadline():
         },
         True,
         _app(),
+        TEST_EXECUTION_CONTEXT,
     )
     session.session = SimpleNamespace(call_tool=AsyncMock(return_value=_tool_result()))
 
@@ -171,6 +173,7 @@ async def test_invoke_mcp_tool_timeout_is_not_retried_and_session_remains_usable
         },
         True,
         _app(),
+        TEST_EXECUTION_CONTEXT,
     )
     timeout = McpError(
         mcp_types.ErrorData(
@@ -205,6 +208,7 @@ def test_invalid_tool_call_timeout_falls_back_to_default(invalid_timeout):
         },
         True,
         ap,
+        TEST_EXECUTION_CONTEXT,
     )
 
     assert session.tool_call_timeout_sec == MCP_TOOL_CALL_TIMEOUT_DEFAULT_SECONDS
@@ -525,7 +529,11 @@ async def test_mcp_tool_result_is_discarded_when_generation_changes_during_call(
     with pytest.raises(WorkspaceGenerationMismatchError):
         await session.invoke_mcp_tool('side_effecting_tool', {})
 
-    session.session.call_tool.assert_awaited_once_with('side_effecting_tool', {})
+    session.session.call_tool.assert_awaited_once_with(
+        'side_effecting_tool',
+        {},
+        read_timeout_seconds=timedelta(seconds=MCP_TOOL_CALL_TIMEOUT_DEFAULT_SECONDS),
+    )
 
 
 @pytest.mark.asyncio
@@ -567,3 +575,43 @@ async def test_mcp_idle_lifecycle_stops_without_retry_after_generation_bump():
     assert session.status == MCPSessionStatus.ERROR
     assert session.error_message == 'Workspace execution binding is stale'
     assert session._shutdown_event.is_set()
+
+
+@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
+            self.server_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 == {}

tests/unit_tests/provider/test_skill_tools.py

@@ -508,7 +508,7 @@ class TestSkillToolLoader:
         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}}),
+            get_backend_status=AsyncMock(return_value={'backend': {'available': True}}),
         )
 
         loader = SkillToolLoader(ap)

tests/unit_tests/provider/test_tool_manager_native.py

@@ -206,6 +206,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_backend_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']
+
+
 @pytest.mark.asyncio
 async def test_native_tool_loader_rechecks_admission_at_the_final_invoke_boundary():
     box_service = SimpleNamespace(