feat(tenancy): implement workspace isolation

This commit is contained in:
Junyan Qin
2026-07-19 09:58:59 +08:00
parent 9eb292992d
commit c6f826fe2d
271 changed files with 31162 additions and 6106 deletions
+141
View File
@@ -6,12 +6,26 @@ from unittest.mock import Mock
import pytest
from langbot_plugin.box.client import ActionRPCBoxClient
from langbot_plugin.box.errors import BoxRuntimeUnavailableError
from langbot_plugin.box.security import (
BOX_CONTROL_TOKEN_ENV,
BOX_CONTROL_TOKEN_HEADER,
BOX_INSTANCE_HEADER,
BOX_PLACEMENT_GENERATION_HEADER,
BOX_TRUSTED_INSTANCE_ENV,
BOX_WORKSPACE_HEADER,
)
from langbot_plugin.entities.io.context import ActionContext
from langbot.pkg.box.connector import BoxRuntimeConnector
_CONTROL_TOKEN = 'box-control-token-that-is-longer-than-32-bytes'
def make_app(logger: Mock, runtime_endpoint: str = ''):
return SimpleNamespace(
logger=logger,
workspace_service=SimpleNamespace(instance_uuid='instance-a'),
instance_config=SimpleNamespace(
data={
'box': {
@@ -104,3 +118,130 @@ def test_box_runtime_connector_dispose_terminates_subprocess(monkeypatch: pytest
ctrl_task.cancel.assert_called_once()
assert connector._handler_task is None
assert connector._ctrl_task is None
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'))
headers = connector.get_control_headers()
assert headers == {
BOX_CONTROL_TOKEN_HEADER: _CONTROL_TOKEN,
BOX_INSTANCE_HEADER: 'instance-a',
}
assert _CONTROL_TOKEN not in connector._resolve_rpc_ws_url()
def test_box_runtime_connector_builds_placement_scoped_relay_headers(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setenv(BOX_CONTROL_TOKEN_ENV, _CONTROL_TOKEN)
connector = BoxRuntimeConnector(make_app(Mock(), runtime_endpoint='http://box-runtime:5410'))
headers = connector.get_relay_headers(
ActionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=7,
)
)
assert headers == {
BOX_CONTROL_TOKEN_HEADER: _CONTROL_TOKEN,
BOX_INSTANCE_HEADER: 'instance-a',
BOX_WORKSPACE_HEADER: 'workspace-a',
BOX_PLACEMENT_GENERATION_HEADER: '7',
}
def test_box_runtime_connector_rejects_relay_context_from_other_instance(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setenv(BOX_CONTROL_TOKEN_ENV, _CONTROL_TOKEN)
connector = BoxRuntimeConnector(make_app(Mock(), runtime_endpoint='http://box-runtime:5410'))
with pytest.raises(BoxRuntimeUnavailableError, match='another LangBot instance'):
connector.get_relay_headers(
ActionContext(
instance_uuid='instance-b',
workspace_uuid='workspace-a',
placement_generation=1,
)
)
def test_external_box_runtime_fails_closed_without_control_token(monkeypatch: pytest.MonkeyPatch):
monkeypatch.delenv(BOX_CONTROL_TOKEN_ENV, raising=False)
connector = BoxRuntimeConnector(make_app(Mock(), runtime_endpoint='http://box-runtime:5410'))
with pytest.raises(BoxRuntimeUnavailableError, match=BOX_CONTROL_TOKEN_ENV):
connector.get_control_headers()
async def test_local_stdio_injects_generated_token_and_trusted_instance(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.delenv(BOX_CONTROL_TOKEN_ENV, raising=False)
captured = {}
class FakeStdioClientController:
def __init__(self, **kwargs):
captured.update(kwargs)
self.process = Mock()
async def run(self, callback):
await callback(None)
monkeypatch.setattr(
'langbot_plugin.runtime.io.controllers.stdio.client.StdioClientController',
FakeStdioClientController,
)
connector = BoxRuntimeConnector(make_app(Mock()))
def fake_callback(_transport_name, connected, _connect_error):
async def callback(_connection):
connected.set()
return callback
monkeypatch.setattr(connector, '_make_connection_callback', fake_callback)
await connector._start_local_stdio()
assert len(captured['env'][BOX_CONTROL_TOKEN_ENV]) >= 32
assert captured['env'][BOX_TRUSTED_INSTANCE_ENV] == 'instance-a'
async def test_websocket_controller_receives_control_headers(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv(BOX_CONTROL_TOKEN_ENV, _CONTROL_TOKEN)
captured = {}
class FakeWebSocketClientController:
def __init__(self, **kwargs):
captured.update(kwargs)
async def run(self, callback):
await callback(None)
monkeypatch.setattr(
'langbot_plugin.runtime.io.controllers.ws.client.WebSocketClientController',
FakeWebSocketClientController,
)
connector = BoxRuntimeConnector(make_app(Mock(), runtime_endpoint='http://box-runtime:5410'))
def fake_callback(_transport_name, connected, _connect_error):
async def callback(_connection):
connected.set()
return callback
monkeypatch.setattr(connector, '_make_connection_callback', fake_callback)
await connector._connect_ws('ws://box-runtime:5410/rpc/ws', 'WebSocket')
assert captured['additional_headers'] == {
BOX_CONTROL_TOKEN_HEADER: _CONTROL_TOKEN,
BOX_INSTANCE_HEADER: 'instance-a',
}
assert _CONTROL_TOKEN not in captured['ws_url']
@@ -0,0 +1,65 @@
from pathlib import Path
_REPO_ROOT = Path(__file__).resolve().parents[3]
def test_compose_injects_the_same_box_control_token_into_host_and_runtime():
compose = (_REPO_ROOT / 'docker' / 'docker-compose.yaml').read_text(encoding='utf-8')
box_service = compose.split(' langbot_box:', 1)[1].split(' langbot:', 1)[0]
langbot_service = compose.split(' langbot:', 1)[1]
token_env = 'LANGBOT_BOX_CONTROL_TOKEN=${LANGBOT_BOX_CONTROL_TOKEN:-}'
assert token_env in box_service
assert token_env in langbot_service
def test_kubernetes_uses_one_secret_for_box_runtime_and_langbot():
manifest = (_REPO_ROOT / 'docker' / 'kubernetes.yaml').read_text(encoding='utf-8')
box_deployment = manifest.split('name: langbot-box', 1)[1].split('# Service for LangBot Box runtime', 1)[0]
langbot_deployment = manifest.split('# Deployment for LangBot\n', 1)[1]
secret_reference = '\n'.join(
[
'- name: LANGBOT_BOX_CONTROL_TOKEN',
' valueFrom:',
' secretKeyRef:',
' name: langbot-box-control',
' key: token',
]
)
assert secret_reference in box_deployment
assert secret_reference in langbot_deployment
assert '--from-literal=token="$(openssl rand -hex 32)"' in manifest
def test_compose_injects_same_plugin_runtime_control_token_into_both_services():
compose = (_REPO_ROOT / 'docker' / 'docker-compose.yaml').read_text(encoding='utf-8')
runtime_service = compose.split(' langbot_plugin_runtime:', 1)[1].split(' langbot_box:', 1)[0]
langbot_service = compose.split(' langbot:', 1)[1]
token_env = 'LANGBOT_PLUGIN_RUNTIME_CONTROL_TOKEN=${LANGBOT_PLUGIN_RUNTIME_CONTROL_TOKEN:-}'
assert token_env in runtime_service
assert token_env in langbot_service
def test_kubernetes_uses_one_secret_for_plugin_runtime_and_langbot():
manifest = (_REPO_ROOT / 'docker' / 'kubernetes.yaml').read_text(encoding='utf-8')
runtime_deployment = manifest.split('# Deployment for LangBot Plugin Runtime', 1)[1].split(
'# Service for LangBot Plugin Runtime',
1,
)[0]
langbot_deployment = manifest.split('# Deployment for LangBot\n', 1)[1]
secret_reference = '\n'.join(
[
'- name: LANGBOT_PLUGIN_RUNTIME_CONTROL_TOKEN',
' valueFrom:',
' secretKeyRef:',
' name: langbot-plugin-runtime-control',
' key: token',
]
)
assert secret_reference in runtime_deployment
assert secret_reference in langbot_deployment
assert 'create secret generic langbot-plugin-runtime-control' in manifest
+239 -53
View File
@@ -15,6 +15,7 @@ from langbot_plugin.box.backend import BaseSandboxBackend
from langbot_plugin.box.client import BoxRuntimeClient, ActionRPCBoxClient
from langbot_plugin.box.errors import (
BoxBackendUnavailableError,
BoxError,
BoxSessionConflictError,
BoxSessionNotFoundError,
BoxValidationError,
@@ -30,9 +31,27 @@ from langbot_plugin.box.models import (
BoxSpec,
)
from langbot_plugin.box.runtime import BoxRuntime
from langbot_plugin.box.security import (
BOX_CONTROL_TOKEN_HEADER,
BOX_INSTANCE_HEADER,
BOX_PLACEMENT_GENERATION_HEADER,
BOX_WORKSPACE_HEADER,
)
from langbot_plugin.entities.io.context import ActionContext
from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.box.service import BoxService
_UTC = dt.timezone.utc
_CONTEXT = ExecutionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=1,
)
_ACTION_CONTEXT = ActionContext(
instance_uuid=_CONTEXT.instance_uuid,
workspace_uuid=_CONTEXT.workspace_uuid,
placement_generation=_CONTEXT.placement_generation,
)
class _InProcessBoxRuntimeClient(BoxRuntimeClient):
@@ -44,37 +63,55 @@ class _InProcessBoxRuntimeClient(BoxRuntimeClient):
async def initialize(self):
await self._runtime.initialize()
async def execute(self, spec):
async def execute(self, spec, *, action_context=None):
return await self._runtime.execute(spec)
async def shutdown(self):
await self._runtime.shutdown()
async def get_status(self):
async def get_status(self, *, action_context=None):
return await self._runtime.get_status()
async def get_sessions(self):
async def get_sessions(self, *, action_context=None):
return self._runtime.get_sessions()
async def get_backend_info(self):
return await self._runtime.get_backend_info()
async def delete_session(self, session_id):
async def delete_session(self, session_id, *, action_context=None):
await self._runtime.delete_session(session_id)
async def create_session(self, spec):
async def create_session(self, spec, *, action_context=None):
return await self._runtime.create_session(spec)
async def start_managed_process(self, session_id: str, spec: BoxManagedProcessSpec):
async def start_managed_process(
self,
session_id: str,
spec: BoxManagedProcessSpec,
*,
action_context=None,
):
return await self._runtime.start_managed_process(session_id, spec)
async def get_managed_process(self, session_id: str, process_id: str = 'default'):
async def get_managed_process(
self,
session_id: str,
process_id: str = 'default',
*,
action_context=None,
):
return self._runtime.get_managed_process(session_id, process_id)
async def stop_managed_process(self, session_id: str, process_id: str = 'default'):
async def stop_managed_process(
self,
session_id: str,
process_id: str = 'default',
*,
action_context=None,
):
await self._runtime.stop_managed_process(session_id, process_id)
async def get_session(self, session_id: str):
async def get_session(self, session_id: str, *, action_context=None):
return self._runtime.get_session(session_id)
async def init(self, config: dict) -> None:
@@ -134,6 +171,12 @@ class FakeBackend(BaseSandboxBackend):
def make_query(query_id: int = 42) -> pipeline_query.Query:
return pipeline_query.Query.model_construct(
query_id=query_id,
query_uuid=f'query-{query_id}',
instance_uuid=_CONTEXT.instance_uuid,
workspace_uuid=_CONTEXT.workspace_uuid,
placement_generation=_CONTEXT.placement_generation,
bot_uuid='bot-a',
pipeline_uuid='pipeline-a',
launcher_type='person',
launcher_id='test_user',
sender_id='test_user',
@@ -170,8 +213,19 @@ def make_app(
if workspace_quota_mb is not None:
box_config['local']['workspace_quota_mb'] = workspace_quota_mb
workspace_service = SimpleNamespace(
instance_uuid=_CONTEXT.instance_uuid,
get_execution_binding=AsyncMock(
return_value=SimpleNamespace(
instance_uuid=_CONTEXT.instance_uuid,
workspace_uuid=_CONTEXT.workspace_uuid,
placement_generation=_CONTEXT.placement_generation,
)
),
)
return SimpleNamespace(
logger=logger,
workspace_service=workspace_service,
instance_config=SimpleNamespace(
data={
'box': box_config,
@@ -289,12 +343,77 @@ async def test_box_service_get_sessions_delegates_to_client():
service = BoxService(make_app(Mock()), client=client)
service._available = True
sessions = await service.get_sessions()
sessions = await service.get_sessions(_CONTEXT)
assert sessions == [{'session_id': 'test-session'}]
client.get_sessions.assert_awaited_once()
@pytest.mark.asyncio
async def test_box_service_relay_connection_is_binding_checked_and_scoped():
app = make_app(Mock())
client = Mock()
client.get_managed_process_websocket_url = Mock(
return_value='ws://box/v1/sessions/physical/managed-process/server-a/ws'
)
connector = Mock()
connector.ws_relay_base_url = 'http://box:5410'
connector.get_relay_headers = Mock(
return_value={
BOX_CONTROL_TOKEN_HEADER: 'secret',
BOX_INSTANCE_HEADER: _CONTEXT.instance_uuid,
BOX_WORKSPACE_HEADER: _CONTEXT.workspace_uuid,
BOX_PLACEMENT_GENERATION_HEADER: '1',
}
)
service = BoxService(app, client=client)
service._runtime_connector = connector
url, headers = await service.get_managed_process_websocket_connection(
_CONTEXT,
'mcp-shared',
'server-a',
)
assert url == 'ws://box/v1/sessions/physical/managed-process/server-a/ws'
assert headers[BOX_WORKSPACE_HEADER] == _CONTEXT.workspace_uuid
assert headers[BOX_PLACEMENT_GENERATION_HEADER] == '1'
app.workspace_service.get_execution_binding.assert_awaited_once_with(
_CONTEXT.workspace_uuid,
expected_generation=_CONTEXT.placement_generation,
)
action_context = connector.get_relay_headers.call_args.args[0]
assert action_context == _ACTION_CONTEXT
client.get_managed_process_websocket_url.assert_called_once_with(
'mcp-shared',
'http://box:5410',
'server-a',
action_context=_ACTION_CONTEXT,
)
@pytest.mark.asyncio
async def test_box_service_relay_connection_rejects_stale_binding():
app = make_app(Mock())
app.workspace_service.get_execution_binding.return_value = SimpleNamespace(
instance_uuid=_CONTEXT.instance_uuid,
workspace_uuid=_CONTEXT.workspace_uuid,
placement_generation=2,
)
client = Mock()
service = BoxService(app, client=client)
service._runtime_connector = Mock()
with pytest.raises(BoxValidationError, match='stale Workspace placement'):
await service.get_managed_process_websocket_connection(
_CONTEXT,
'mcp-shared',
'server-a',
)
service._runtime_connector.get_relay_headers.assert_not_called()
def test_box_service_dispose_delegates_to_internal_connector(monkeypatch: pytest.MonkeyPatch):
connector = Mock()
connector.client = Mock()
@@ -369,7 +488,14 @@ async def test_box_service_session_id_uses_query_attributes_without_variables():
service = BoxService(make_app(logger), client=_InProcessBoxRuntimeClient(logger, runtime))
await service.initialize()
query = pipeline_query.Query.model_construct(query_id=7, launcher_type='group', launcher_id='room-1')
query = pipeline_query.Query.model_construct(
query_id=7,
instance_uuid=_CONTEXT.instance_uuid,
workspace_uuid=_CONTEXT.workspace_uuid,
placement_generation=_CONTEXT.placement_generation,
launcher_type='group',
launcher_id='room-1',
)
result = await service.execute_tool({'command': 'pwd'}, query)
assert result['session_id'] == 'group_room-1'
@@ -385,7 +511,12 @@ async def test_box_service_session_id_falls_back_to_query_id_for_synthetic_queri
service = BoxService(make_app(logger), client=_InProcessBoxRuntimeClient(logger, runtime))
await service.initialize()
query = pipeline_query.Query.model_construct(query_id=7)
query = pipeline_query.Query.model_construct(
query_id=7,
instance_uuid=_CONTEXT.instance_uuid,
workspace_uuid=_CONTEXT.workspace_uuid,
placement_generation=_CONTEXT.placement_generation,
)
result = await service.execute_tool({'command': 'pwd'}, query)
assert result['session_id'] == 'query_7'
@@ -407,8 +538,22 @@ async def test_box_service_forced_global_scope_overrides_pipeline_template():
await service.initialize()
# Two distinct callers that would otherwise get separate sandboxes.
q1 = pipeline_query.Query.model_construct(query_id=1, launcher_type='group', launcher_id='room-1')
q2 = pipeline_query.Query.model_construct(query_id=2, launcher_type='person', launcher_id='alice')
q1 = pipeline_query.Query.model_construct(
query_id=1,
instance_uuid=_CONTEXT.instance_uuid,
workspace_uuid=_CONTEXT.workspace_uuid,
placement_generation=_CONTEXT.placement_generation,
launcher_type='group',
launcher_id='room-1',
)
q2 = pipeline_query.Query.model_construct(
query_id=2,
instance_uuid=_CONTEXT.instance_uuid,
workspace_uuid=_CONTEXT.workspace_uuid,
placement_generation=_CONTEXT.placement_generation,
launcher_type='person',
launcher_id='alice',
)
r1 = await service.execute_tool({'command': 'pwd'}, q1)
r2 = await service.execute_tool({'command': 'pwd'}, q2)
@@ -511,7 +656,7 @@ async def test_box_service_uses_default_workspace_when_host_path_omitted(tmp_pat
assert result['ok'] is True
assert backend.start_calls == ['person_test_user']
assert backend.exec_calls == [('person_test_user', 'pwd')]
assert backend.start_specs[0].host_path == os.path.realpath(host_dir)
assert backend.start_specs[0].host_path == service._tenant_workspace(_CONTEXT)
@pytest.mark.asyncio
@@ -954,11 +1099,16 @@ async def test_box_service_rejects_execution_when_workspace_already_exceeds_quot
runtime = BoxRuntime(logger=logger, backends=[backend], session_ttl_sec=300)
host_dir = tmp_path / 'quota-workspace'
host_dir.mkdir()
(host_dir / 'already-too-large.bin').write_bytes(b'x' * (2 * 1024 * 1024))
app = make_app(logger, [str(tmp_path)], workspace_quota_mb=1)
app.instance_config.data['box']['local']['default_workspace'] = str(host_dir)
service = BoxService(app, client=_InProcessBoxRuntimeClient(logger, runtime))
tenant_host_dir = service._tenant_workspace(_CONTEXT)
assert tenant_host_dir is not None
os.makedirs(tenant_host_dir, exist_ok=True)
with open(os.path.join(tenant_host_dir, 'already-too-large.bin'), 'wb') as handle:
handle.write(b'x' * (2 * 1024 * 1024))
await service.initialize()
with pytest.raises(BoxValidationError, match='workspace quota exceeded before execution'):
@@ -1097,7 +1247,7 @@ async def test_service_records_errors_on_failure():
with pytest.raises(Exception):
await service.execute_tool({'command': 'echo hello'}, make_query(50))
errors = service.get_recent_errors()
errors = service.get_recent_errors(_CONTEXT)
assert len(errors) == 1
assert errors[0]['type'] == 'BoxBackendUnavailableError'
assert errors[0]['query_id'] == '50'
@@ -1116,7 +1266,7 @@ async def test_service_error_ring_buffer_capped():
with pytest.raises(Exception):
await service.execute_tool({'command': 'fail'}, make_query(100 + i))
errors = service.get_recent_errors()
errors = service.get_recent_errors(_CONTEXT)
assert len(errors) == 50
# Oldest should have been evicted, newest kept
assert errors[0]['query_id'] == '110'
@@ -1131,7 +1281,7 @@ async def test_service_get_status_aggregates_runtime_and_profile():
service = BoxService(make_app(logger), client=_InProcessBoxRuntimeClient(logger, runtime))
await service.initialize()
status = await service.get_status()
status = await service.get_status(_CONTEXT)
assert status['profile'] == 'default'
assert status['backend']['name'] == 'fake'
assert status['backend']['available'] is True
@@ -1175,7 +1325,12 @@ async def _make_rpc_pair(runtime: BoxRuntime):
client_conn, server_conn = _make_queue_connection_pair()
server_handler = BoxServerHandler(server_conn, runtime)
server_handler = BoxServerHandler(
server_conn,
runtime,
host_control_authenticated=True,
trusted_instance_uuid=_CONTEXT.instance_uuid,
)
server_task = asyncio.create_task(server_handler.run())
client_handler = Handler.__new__(Handler)
@@ -1199,7 +1354,7 @@ async def test_rpc_client_execute():
client, server_task, client_task = await _make_rpc_pair(runtime)
try:
spec = BoxSpec.model_validate({'cmd': 'echo remote', 'session_id': 'r-1'})
result = await client.execute(spec)
result = await client.execute(spec, action_context=_ACTION_CONTEXT)
assert result.session_id == 'r-1'
assert result.status == BoxExecutionStatus.COMPLETED
@@ -1221,9 +1376,9 @@ async def test_rpc_client_get_sessions():
client, server_task, client_task = await _make_rpc_pair(runtime)
try:
spec = BoxSpec.model_validate({'cmd': 'echo hi', 'session_id': 'r-2'})
await client.execute(spec)
await client.execute(spec, action_context=_ACTION_CONTEXT)
sessions = await client.get_sessions()
sessions = await client.get_sessions(action_context=_ACTION_CONTEXT)
assert len(sessions) == 1
assert sessions[0]['session_id'] == 'r-2'
finally:
@@ -1232,6 +1387,34 @@ async def test_rpc_client_get_sessions():
await runtime.shutdown()
@pytest.mark.asyncio
async def test_rpc_generation_advance_retires_old_session_and_rejects_old_context():
logger = Mock()
backend = FakeBackend(logger)
runtime = BoxRuntime(logger=logger, backends=[backend], session_ttl_sec=300)
await runtime.initialize()
second_context = _ACTION_CONTEXT.model_copy(update={'placement_generation': 2})
client, server_task, client_task = await _make_rpc_pair(runtime)
try:
spec = BoxSpec.model_validate({'cmd': 'echo generation', 'session_id': 'shared'})
await client.execute(spec, action_context=_ACTION_CONTEXT)
await client.execute(spec, action_context=second_context)
assert len(backend.start_calls) == 2
assert backend.start_calls[0] != backend.start_calls[1]
assert backend.stop_calls == [backend.start_calls[0]]
assert [session['session_id'] for session in await client.get_sessions(action_context=second_context)] == [
'shared'
]
with pytest.raises(BoxError, match='Stale Box placement generation'):
await client.execute(spec, action_context=_ACTION_CONTEXT)
finally:
server_task.cancel()
client_task.cancel()
await runtime.shutdown()
@pytest.mark.asyncio
async def test_rpc_client_get_status():
logger = Mock()
@@ -1241,7 +1424,7 @@ async def test_rpc_client_get_status():
client, server_task, client_task = await _make_rpc_pair(runtime)
try:
status = await client.get_status()
status = await client.get_status(action_context=_ACTION_CONTEXT)
assert 'backend' in status
assert 'active_sessions' in status
@@ -1283,11 +1466,11 @@ async def test_rpc_client_delete_session():
client, server_task, client_task = await _make_rpc_pair(runtime)
try:
spec = BoxSpec.model_validate({'cmd': 'echo hi', 'session_id': 'r-del-1'})
await client.execute(spec)
await client.execute(spec, action_context=_ACTION_CONTEXT)
await client.delete_session('r-del-1')
await client.delete_session('r-del-1', action_context=_ACTION_CONTEXT)
sessions = await client.get_sessions()
sessions = await client.get_sessions(action_context=_ACTION_CONTEXT)
assert len(sessions) == 0
finally:
server_task.cancel()
@@ -1305,7 +1488,7 @@ async def test_rpc_client_delete_session_raises_not_found():
client, server_task, client_task = await _make_rpc_pair(runtime)
try:
with pytest.raises(BoxSessionNotFoundError):
await client.delete_session('nonexistent')
await client.delete_session('nonexistent', action_context=_ACTION_CONTEXT)
finally:
server_task.cancel()
client_task.cancel()
@@ -1322,11 +1505,11 @@ async def test_rpc_client_create_session():
client, server_task, client_task = await _make_rpc_pair(runtime)
try:
spec = BoxSpec.model_validate({'cmd': 'placeholder', 'session_id': 'r-create-1'})
info = await client.create_session(spec)
info = await client.create_session(spec, action_context=_ACTION_CONTEXT)
assert info['session_id'] == 'r-create-1'
assert info['backend_name'] == 'fake'
sessions = await client.get_sessions()
sessions = await client.get_sessions(action_context=_ACTION_CONTEXT)
assert len(sessions) == 1
finally:
server_task.cancel()
@@ -1344,11 +1527,11 @@ async def test_rpc_client_exec_raises_conflict_error():
client, server_task, client_task = await _make_rpc_pair(runtime)
try:
spec1 = BoxSpec.model_validate({'cmd': 'echo first', 'session_id': 'r-conflict-1', 'network': 'off'})
await client.execute(spec1)
await client.execute(spec1, action_context=_ACTION_CONTEXT)
spec2 = BoxSpec.model_validate({'cmd': 'echo second', 'session_id': 'r-conflict-1', 'network': 'on'})
with pytest.raises(BoxSessionConflictError):
await client.execute(spec2)
await client.execute(spec2, action_context=_ACTION_CONTEXT)
finally:
server_task.cancel()
client_task.cancel()
@@ -1447,7 +1630,7 @@ class TestBoxDisabledByConfig:
service = BoxService(make_app(logger, enabled=False), client=Mock(spec=BoxRuntimeClient))
await service.initialize()
status = await service.get_status()
status = await service.get_status(_CONTEXT)
assert status['available'] is False
assert status['enabled'] is False
@@ -1462,7 +1645,7 @@ class TestBoxDisabledByConfig:
await service.initialize()
status = await service.get_status()
status = await service.get_status(_CONTEXT)
assert status['available'] is False
assert status['enabled'] is True
assert 'docker daemon' in status['connector_error']
@@ -1486,7 +1669,7 @@ class TestBoxDisabledByConfig:
service = BoxService(make_app(logger, enabled=True), client=client)
await service.initialize()
status = await service.get_status()
status = await service.get_status(_CONTEXT)
assert status['available'] is False
assert status['enabled'] is True
# The detailed backend object is preserved for the dialog
@@ -1507,7 +1690,7 @@ class TestBoxDisabledByConfig:
service = BoxService(make_app(logger, enabled=True), client=client)
await service.initialize()
status = await service.get_status()
status = await service.get_status(_CONTEXT)
assert status['available'] is True
assert status['backend'] == {'name': 'docker', 'available': True}
# No spurious connector_error overlay when everything is healthy
@@ -1537,7 +1720,7 @@ class TestBuildSkillExtraMounts:
def _make_service(self, logger, skills, *, shares_filesystem=True):
app = make_app(logger)
app.skill_mgr = SimpleNamespace(skills=skills)
app.skill_mgr = SimpleNamespace(skills=skills, get_skills=Mock(return_value=skills))
client = Mock(spec=BoxRuntimeClient)
service = BoxService(app, client=client)
# Tests construct BoxService with an injected client (no connector), so
@@ -1795,14 +1978,17 @@ class TestAttachmentHostPath:
"""
def _service_with_workspace(self, tmp_path):
ws = str(tmp_path / 'box' / 'default')
os.makedirs(ws, exist_ok=True)
default_workspace = str(tmp_path / 'box' / 'default')
os.makedirs(default_workspace, exist_ok=True)
app = make_app(Mock(), allowed_mount_roots=[str(tmp_path)], host_root=str(tmp_path / 'box'))
service = BoxService(app, client=Mock(spec=BoxRuntimeClient))
service._available = True
# Force the default_workspace to our tmp dir so _host_query_dir resolves.
service.default_workspace = ws
return service, ws
service.default_workspace = default_workspace
tenant_workspace = service._tenant_workspace(_CONTEXT)
assert tenant_workspace is not None
os.makedirs(tenant_workspace, exist_ok=True)
return service, tenant_workspace
@pytest.mark.asyncio
async def test_inbound_writes_to_host_no_exec(self, tmp_path):
@@ -1923,9 +2109,10 @@ class TestAttachmentHostPath:
assert os.path.isdir(ws)
@pytest.mark.asyncio
async def test_purge_attachment_dirs_falls_back_to_exec_for_root_owned(self, tmp_path, monkeypatch):
# When the host delete cannot remove a dir (root-owned container output),
# purge must fall back to deleting from inside the sandbox via exec.
async def test_purge_attachment_dirs_never_uses_unscoped_exec_for_root_owned(self, tmp_path, monkeypatch):
# Startup has no trusted Workspace context. If host deletion cannot
# remove root-owned output, cleanup must fail closed instead of issuing
# an unscoped Box exec that could cross a tenant boundary.
service, ws = self._service_with_workspace(tmp_path)
outbox = os.path.join(ws, 'outbox')
os.makedirs(os.path.join(outbox, '0'), exist_ok=True)
@@ -1942,18 +2129,17 @@ class TestAttachmentHostPath:
monkeypatch.setattr(_shutil, 'rmtree', fake_rmtree)
executed = {}
spec_obj = object()
service.build_spec = Mock(return_value=spec_obj)
service.client.execute = AsyncMock(side_effect=lambda s: executed.setdefault('spec', s))
service.build_spec = Mock()
service.client.execute = AsyncMock()
await service._purge_attachment_dirs()
# build_spec was asked to rm the surviving outbox via exec.
cmd = service.build_spec.call_args.args[0]['cmd']
assert 'rm -rf' in cmd and '/workspace/outbox' in cmd
assert '/workspace/inbox' not in cmd # inbox was host-deletable
service.client.execute.assert_awaited_once_with(spec_obj)
assert os.path.isdir(outbox)
service.build_spec.assert_not_called()
service.client.execute.assert_not_awaited()
assert any(
'no trusted Workspace context' in str(call.args[0]) for call in service.ap.logger.warning.call_args_list
)
@pytest.mark.asyncio
async def test_purge_attachment_dirs_noop_without_workspace(self):
+44 -2
View File
@@ -7,6 +7,7 @@ from unittest.mock import AsyncMock, Mock
import pytest
from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.box.workspace import (
BoxWorkspaceSession,
classify_python_workspace,
@@ -16,6 +17,13 @@ from langbot.pkg.box.workspace import (
)
_CONTEXT = ExecutionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=1,
)
def test_rewrite_mounted_path_translates_host_prefix():
result = rewrite_mounted_path('/tmp/demo/project/app.py', '/tmp/demo/project')
assert result == '/workspace/app.py'
@@ -66,6 +74,7 @@ async def test_workspace_session_execute_for_query_uses_session_payload():
box_service = SimpleNamespace(execute_spec_payload=AsyncMock(return_value={'ok': True}))
workspace = BoxWorkspaceSession(
box_service,
_CONTEXT,
'skill-person_123-demo',
host_path='/tmp/project',
host_path_mode='rw',
@@ -94,6 +103,7 @@ async def test_workspace_session_start_managed_process_rewrites_command_and_args
box_service = SimpleNamespace(start_managed_process=AsyncMock(return_value={'status': 'running'}))
workspace = BoxWorkspaceSession(
box_service,
_CONTEXT,
'mcp-u1',
host_path='/tmp/project',
host_path_mode='ro',
@@ -106,8 +116,10 @@ async def test_workspace_session_start_managed_process_rewrites_command_and_args
)
assert result == {'status': 'running'}
session_id = box_service.start_managed_process.await_args.args[0]
payload = box_service.start_managed_process.await_args.args[1]
execution_context = box_service.start_managed_process.await_args.args[0]
session_id = box_service.start_managed_process.await_args.args[1]
payload = box_service.start_managed_process.await_args.args[2]
assert execution_context == _CONTEXT
assert session_id == 'mcp-u1'
assert payload == {
'command': 'python',
@@ -118,9 +130,39 @@ async def test_workspace_session_start_managed_process_rewrites_command_and_args
}
@pytest.mark.asyncio
async def test_workspace_session_relay_connection_keeps_execution_context():
box_service = SimpleNamespace(
get_managed_process_websocket_connection=AsyncMock(
return_value=(
'ws://box/relay',
{'X-LangBot-Placement-Generation': '1'},
)
)
)
workspace = BoxWorkspaceSession(
box_service,
_CONTEXT,
'mcp-shared',
)
connection = await workspace.get_managed_process_websocket_connection('server-a')
assert connection == (
'ws://box/relay',
{'X-LangBot-Placement-Generation': '1'},
)
box_service.get_managed_process_websocket_connection.assert_awaited_once_with(
_CONTEXT,
'mcp-shared',
'server-a',
)
def test_workspace_session_build_session_payload_keeps_generic_workspace_shape():
workspace = BoxWorkspaceSession(
Mock(),
_CONTEXT,
'workspace-1',
host_path='/tmp/project',
host_path_mode='rw',