feat(tenancy): connect cloud workspace control plane

2026-07-25 22:06:06 +00:00 · 2026-07-24 19:11:33 +08:00
parent d7cdd206c2
commit 98f45aa88e
40 changed files with 4159 additions and 452 deletions
@@ -6,6 +6,7 @@ from unittest.mock import AsyncMock, Mock

 import pytest

+from langbot.pkg.box import connector as connector_module
 from langbot_plugin.box.client import ActionRPCBoxClient
 from langbot_plugin.box.errors import BoxRuntimeUnavailableError
 from langbot_plugin.box.security import (
@@ -121,6 +122,139 @@ def test_box_runtime_connector_dispose_terminates_subprocess(monkeypatch: pytest
    assert connector._ctrl_task is None


+@pytest.mark.asyncio
+async def test_box_runtime_connector_cleans_partial_transport_on_connect_failure(
+    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'))
+    connector._stop_transport = AsyncMock()
+    connector._close_managed_subprocess = AsyncMock()
+
+    with pytest.raises(RuntimeError, match='bind failed'):
+        await connector.initialize()
+
+    assert connector._stop_transport.await_count == 2
+    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,
+):
+    monkeypatch.setattr('langbot.pkg.utils.platform.get_platform', lambda: 'linux')
+    monkeypatch.setattr('langbot.pkg.utils.platform.standalone_box', False)
+    created = {}
+
+    class FakeHandler:
+        def __init__(self, connection):
+            self.release = asyncio.Event()
+
+        async def call_action(self, action, data):
+            return None
+
+        async def run(self):
+            await self.release.wait()
+
+        async def close(self):
+            self.release.set()
+
+    class FakeController:
+        def __init__(self, **kwargs):
+            created.update(kwargs)
+            self.process = SimpleNamespace(returncode=0)
+
+        async def run(self, callback):
+            await callback(object())
+
+        async def close(self):
+            return None
+
+    monkeypatch.setattr(connector_module, 'Handler', FakeHandler)
+    monkeypatch.setattr(
+        'langbot_plugin.runtime.io.controllers.stdio.client.StdioClientController',
+        FakeController,
+    )
+    connector = BoxRuntimeConnector(make_app(Mock()))
+
+    await connector.initialize()
+
+    assert created['capture_stderr'] is False
+    assert connector._handler is not None
+    await connector.aclose()
+
+
+@pytest.mark.asyncio
+async def test_box_disconnect_notifies_once_and_clears_handler(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    monkeypatch.setattr('langbot.pkg.utils.platform.get_platform', lambda: 'linux')
+    monkeypatch.setattr('langbot.pkg.utils.platform.standalone_box', False)
+    disconnect = AsyncMock()
+
+    class FakeHandler:
+        def __init__(self, connection):
+            pass
+
+        async def call_action(self, action, data):
+            return None
+
+        async def run(self):
+            return None
+
+        async def close(self):
+            return None
+
+    class FakeController:
+        def __init__(self, **kwargs):
+            self.process = SimpleNamespace(returncode=0)
+
+        async def run(self, callback):
+            await callback(object())
+
+        async def close(self):
+            return None
+
+    monkeypatch.setattr(connector_module, 'Handler', FakeHandler)
+    monkeypatch.setattr(
+        'langbot_plugin.runtime.io.controllers.stdio.client.StdioClientController',
+        FakeController,
+    )
+    connector = BoxRuntimeConnector(make_app(Mock()), runtime_disconnect_callback=disconnect)
+
+    await connector.initialize()
+    await asyncio.sleep(0)
+
+    disconnect.assert_awaited_once_with(connector)
+    assert connector._handler is None
+    await connector.aclose()
+
+
 def test_box_runtime_connector_builds_host_control_headers(monkeypatch: pytest.MonkeyPatch):
    monkeypatch.setenv(BOX_CONTROL_TOKEN_ENV, _CONTROL_TOKEN)
    connector = BoxRuntimeConnector(make_app(Mock(), runtime_endpoint='http://box-runtime:5410'))
@@ -200,7 +334,7 @@ async def test_local_stdio_injects_generated_token_and_trusted_instance(
    )
    connector = BoxRuntimeConnector(make_app(Mock()))

-    def fake_callback(_transport_name, connected, _connect_error):
+    def fake_callback(_transport_name, connected, _connect_error, _generation):
        async def callback(_connection):
            connected.set()

@@ -231,7 +365,7 @@ async def test_websocket_controller_receives_control_headers(monkeypatch: pytest
    )
    connector = BoxRuntimeConnector(make_app(Mock(), runtime_endpoint='http://box-runtime:5410'))

-    def fake_callback(_transport_name, connected, _connect_error):
+    def fake_callback(_transport_name, connected, _connect_error, _generation):
        async def callback(_connection):
            connected.set()

@@ -253,6 +253,46 @@ async def test_box_service_without_explicit_client_initializes_internal_connecto
    connector.initialize.assert_awaited_once()


+@pytest.mark.asyncio
+async def test_cloud_initialize_validation_failure_closes_connector_and_cancels_reconnect(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    logger = Mock()
+    app = make_app(logger)
+    app.deployment = SimpleNamespace(multi_workspace_enabled=True)
+    app.instance_config.data['box'].update(
+        {
+            'backend': 'nsjail',
+            'admission': {'required': True, 'workspace_quota_mb': 32},
+        }
+    )
+    connector = Mock()
+    connector.client = Mock(spec=BoxRuntimeClient)
+    connector.initialize = AsyncMock()
+    connector.aclose = AsyncMock()
+    connector.runtime_disconnect_callback = Mock()
+    monkeypatch.setattr('langbot.pkg.box.service.BoxRuntimeConnector', Mock(return_value=connector))
+    service = BoxService(app)
+    service._ensure_default_workspace = Mock()
+    readiness_error = BoxValidationError('Cloud Box nsjail isolation readiness failed')
+    service._verify_cloud_runtime = AsyncMock(side_effect=readiness_error)
+    reconnect_task = asyncio.create_task(asyncio.Event().wait())
+    service._reconnect_task = reconnect_task
+    service._reconnecting = True
+
+    with pytest.raises(BoxValidationError) as exc_info:
+        await service.initialize()
+
+    assert exc_info.value is readiness_error
+    assert service.available is False
+    assert service._closing is True
+    assert service._reconnecting is False
+    assert service._reconnect_task is None
+    assert reconnect_task.cancelled()
+    assert connector.runtime_disconnect_callback is None
+    connector.aclose.assert_awaited_once()
+
+
 class TestSharesFilesystemWithBox:
    """``shares_filesystem_with_box`` must reflect the real LangBot<->Box
    filesystem topology, which is derived from the connector transport:
@@ -1788,6 +1828,57 @@ class TestBoxDisabledByConfig:
        assert service._reconnecting is False


+@pytest.mark.asyncio
+async def test_disconnect_callback_does_not_schedule_on_closing_event_loop(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    service = BoxService(make_app(Mock()), client=Mock(spec=BoxRuntimeClient))
+    closed_loop = Mock()
+    closed_loop.is_closed.return_value = True
+    monkeypatch.setattr('langbot.pkg.box.service.asyncio.get_running_loop', Mock(return_value=closed_loop))
+
+    await service._on_runtime_disconnect(connector=Mock())
+
+    closed_loop.create_task.assert_not_called()
+    assert service._reconnect_task is None
+    assert service._reconnecting is False
+
+
+@pytest.mark.asyncio
+async def test_disconnect_callback_closes_reconnect_coroutine_when_task_creation_races_with_loop_close(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    service = BoxService(make_app(Mock()), client=Mock(spec=BoxRuntimeClient))
+    loop = Mock()
+    loop.is_closed.return_value = False
+    loop.create_task.side_effect = RuntimeError('event loop is closed')
+    monkeypatch.setattr('langbot.pkg.box.service.asyncio.get_running_loop', Mock(return_value=loop))
+
+    async def reconnect():
+        await asyncio.Event().wait()
+
+    reconnect_coroutine = reconnect()
+    service._reconnect_loop = Mock(return_value=reconnect_coroutine)
+
+    await service._on_runtime_disconnect(connector=Mock())
+
+    loop.create_task.assert_called_once_with(reconnect_coroutine)
+    assert reconnect_coroutine.cr_frame is None
+    assert service._reconnect_task is None
+    assert service._reconnecting is False
+
+
+def test_disconnect_callback_does_not_schedule_without_running_event_loop():
+    service = BoxService(make_app(Mock()), client=Mock(spec=BoxRuntimeClient))
+    callback = service._on_runtime_disconnect(connector=Mock())
+
+    with pytest.raises(StopIteration):
+        callback.send(None)
+
+    assert service._reconnect_task is None
+    assert service._reconnecting is False
+
+
 class TestBuildSkillExtraMounts:
    """Robustness of skill mount construction against a stale skill cache.