LangBot/tests/integration_tests/box/test_cloud_box_admission_integration.py

from __future__ import annotations

import asyncio
import datetime as dt
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock

import pytest

import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
from langbot_plugin.box.backend import BaseSandboxBackend
from langbot_plugin.box.client import ActionRPCBoxClient
from langbot_plugin.box.errors import BoxAdmissionError
from langbot_plugin.box.models import (
    BoxExecutionResult,
    BoxExecutionStatus,
    BoxNetworkMode,
    BoxSessionInfo,
    BoxSpec,
)
from langbot_plugin.box.runtime import BoxRuntime
from langbot_plugin.box.server import BoxServerHandler
from langbot_plugin.runtime.io.handler import Handler

from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.box.service import BoxService
from langbot.pkg.cloud.entitlements import (
    EntitlementResolver,
    EntitlementSnapshot,
    EntitlementUnavailableError,
)


pytestmark = pytest.mark.integration
_UTC = dt.timezone.utc


class _AdmissionBackend(BaseSandboxBackend):
    name = 'nsjail'

    def __init__(self, logger):
        super().__init__(logger)
        self.started_specs: list[BoxSpec] = []
        self.stopped_sessions: list[str] = []

    async def is_available(self) -> bool:
        return True

    async def get_readiness(self, *, workspace_path=None, strict=False) -> dict:
        return {
            'available': True,
            'cgroup_v2': True,
            'namespace_isolation': True,
            'mount_isolation': True,
            'network_isolation': True,
            'hard_workspace_quota': True,
            'hard_skill_storage_quota': True,
            'bounded_ephemeral_storage': True,
            'inode_quota': True,
        }

    async def start_session(self, spec: BoxSpec) -> BoxSessionInfo:
        self.started_specs.append(spec)
        now = dt.datetime.now(_UTC)
        return BoxSessionInfo(
            session_id=spec.session_id,
            backend_name=self.name,
            backend_session_id=f'jail-{len(self.started_specs)}',
            image=spec.image,
            network=spec.network,
            host_path=spec.host_path,
            host_path_mode=spec.host_path_mode,
            mount_path=spec.mount_path,
            persistent=spec.persistent,
            cpus=spec.cpus,
            memory_mb=spec.memory_mb,
            pids_limit=spec.pids_limit,
            read_only_rootfs=spec.read_only_rootfs,
            workspace_quota_mb=spec.workspace_quota_mb,
            created_at=now,
            last_used_at=now,
        )

    async def exec(self, session: BoxSessionInfo, spec: BoxSpec) -> BoxExecutionResult:
        await asyncio.sleep(0)
        return BoxExecutionResult(
            session_id=session.session_id,
            backend_name=self.name,
            status=BoxExecutionStatus.COMPLETED,
            exit_code=0,
            stdout=spec.cmd,
            stderr='',
            duration_ms=1,
        )

    async def stop_session(self, session: BoxSessionInfo):
        self.stopped_sessions.append(session.session_id)


class _QueueConnection:
    def __init__(self, rx: asyncio.Queue[str], tx: asyncio.Queue[str]):
        self._rx = rx
        self._tx = tx

    async def send(self, message: str) -> None:
        await self._tx.put(message)

    async def receive(self) -> str:
        return await self._rx.get()

    async def close(self) -> None:
        return None


async def _rpc_client(runtime: BoxRuntime):
    client_to_server: asyncio.Queue[str] = asyncio.Queue()
    server_to_client: asyncio.Queue[str] = asyncio.Queue()
    client_connection = _QueueConnection(server_to_client, client_to_server)
    server_connection = _QueueConnection(client_to_server, server_to_client)
    server_handler = BoxServerHandler(
        server_connection,
        runtime,
        host_control_authenticated=True,
        trusted_instance_uuid='instance-a',
    )
    server_task = asyncio.create_task(server_handler.run())
    client_handler = Handler(client_connection)
    client_task = asyncio.create_task(client_handler.run())
    client = ActionRPCBoxClient(logger=Mock())
    client.set_handler(client_handler)
    return client, server_task, client_task


class _Entitlements:
    def __init__(self):
        self.snapshots: dict[str, EntitlementSnapshot] = {}

    async def get_workspace_entitlement(self, workspace_uuid: str) -> EntitlementSnapshot:
        return self.snapshots[workspace_uuid]


def _snapshot(workspace_uuid: str, *, revision: int = 1, managed: bool = True) -> EntitlementSnapshot:
    return EntitlementSnapshot(
        instance_uuid='instance-a',
        workspace_uuid=workspace_uuid,
        entitlement_revision=revision,
        status='active',
        not_before=1,
        expires_at=4_000_000_000,
        features={'managed_sandbox': managed},
        limits={'managed_sandbox_sessions': 1 if managed else 0},
    )


def _context(workspace_uuid: str, *, revision: int = 1) -> ExecutionContext:
    return ExecutionContext(
        instance_uuid='instance-a',
        workspace_uuid=workspace_uuid,
        placement_generation=1,
        entitlement_revision=revision,
    )


def _query(context: ExecutionContext, query_id: int):
    query = pipeline_query.Query.model_construct(
        query_id=query_id,
        bot_uuid='bot-a',
        pipeline_uuid='pipeline-a',
        launcher_type='person',
        launcher_id=f'user-{query_id}',
        variables={},
    )
    object.__setattr__(query, 'instance_uuid', context.instance_uuid)
    object.__setattr__(query, 'workspace_uuid', context.workspace_uuid)
    object.__setattr__(query, 'placement_generation', context.placement_generation)
    object.__setattr__(query, '_execution_context', context)
    return query


async def _stack(tmp_path):
    shared_root = tmp_path / 'shared-box'
    workspace_root = shared_root / 'workspaces'
    workspace_root.mkdir(parents=True)
    box_config = {
        'enabled': True,
        'backend': 'nsjail',
        'runtime': {'endpoint': 'ws://langbot-box:5410'},
        'local': {
            'profile': 'default',
            'host_root': str(shared_root),
            'default_workspace': str(workspace_root),
            'allowed_mount_roots': [str(shared_root)],
        },
        'admission': {
            'required': True,
            'logical_session_id': 'global',
            'required_backend': 'nsjail',
            'max_sessions': 1,
            'max_managed_processes': 0,
            'max_grant_ttl_sec': 300,
            'max_timeout_sec': 60,
            'cpus': 0.5,
            'memory_mb': 256,
            'pids_limit': 64,
            'read_only_rootfs': True,
            'workspace_quota_mb': 32,
            'readiness_cache_sec': 0,
        },
    }
    logger = Mock()
    backend = _AdmissionBackend(logger)
    runtime = BoxRuntime(logger, backends=[backend])
    runtime.init(box_config)
    await runtime.initialize()
    client, server_task, client_task = await _rpc_client(runtime)

    entitlements = _Entitlements()
    workspace_service = SimpleNamespace(
        instance_uuid='instance-a',
        get_execution_binding=AsyncMock(
            side_effect=lambda workspace_uuid, expected_generation: SimpleNamespace(
                instance_uuid='instance-a',
                workspace_uuid=workspace_uuid,
                placement_generation=expected_generation,
            )
        ),
    )
    app = SimpleNamespace(
        logger=logger,
        deployment=SimpleNamespace(multi_workspace_enabled=True),
        entitlement_resolver=EntitlementResolver('instance-a', entitlements),
        workspace_service=workspace_service,
        instance_config=SimpleNamespace(data={'box': box_config, 'system': {'limitation': {}}}),
    )
    service = BoxService(app, client=client)
    await service.initialize()
    return service, runtime, backend, entitlements, server_task, client_task


@pytest.mark.asyncio
async def test_concurrent_first_use_creates_one_persistent_global_session(tmp_path):
    service, runtime, backend, entitlements, server_task, client_task = await _stack(tmp_path)
    context = _context('workspace-a')
    entitlements.snapshots[context.workspace_uuid] = _snapshot(context.workspace_uuid)
    try:
        first, second = await asyncio.gather(
            service.execute_tool({'command': 'echo first'}, _query(context, 1)),
            service.execute_tool({'command': 'echo second'}, _query(context, 2)),
        )

        assert first['session_id'] == 'global'
        assert second['session_id'] == 'global'
        assert len(backend.started_specs) == 1
        spec = backend.started_specs[0]
        assert spec.persistent is True
        assert spec.network == BoxNetworkMode.OFF
        assert spec.cpus == 0.5
        assert spec.memory_mb == 256
        assert spec.pids_limit == 64
        assert spec.workspace_quota_mb == 32
    finally:
        server_task.cancel()
        client_task.cancel()
        await runtime.shutdown()


@pytest.mark.asyncio
async def test_entitlement_loss_revokes_and_closes_existing_global_session(tmp_path):
    service, runtime, backend, entitlements, server_task, client_task = await _stack(tmp_path)
    context = _context('workspace-a')
    entitlements.snapshots[context.workspace_uuid] = _snapshot(context.workspace_uuid, revision=1)
    try:
        await service.execute_tool({'command': 'true'}, _query(context, 1))
        assert len(runtime.get_sessions()) == 1

        entitlements.snapshots[context.workspace_uuid] = _snapshot(
            context.workspace_uuid,
            revision=2,
            managed=False,
        )
        with pytest.raises(EntitlementUnavailableError):
            await service.execute_tool({'command': 'true'}, _query(context, 2))

        assert runtime.get_sessions() == []
        assert len(backend.stopped_sessions) == 1
    finally:
        server_task.cancel()
        client_task.cancel()
        await runtime.shutdown()


@pytest.mark.asyncio
async def test_two_workspaces_get_isolated_physical_sessions_and_paths(tmp_path):
    service, runtime, backend, entitlements, server_task, client_task = await _stack(tmp_path)
    first = _context('workspace-a')
    second = _context('workspace-b')
    entitlements.snapshots[first.workspace_uuid] = _snapshot(first.workspace_uuid)
    entitlements.snapshots[second.workspace_uuid] = _snapshot(second.workspace_uuid)
    try:
        result_a = await service.execute_tool({'command': 'tenant-a'}, _query(first, 1))
        result_b = await service.execute_tool({'command': 'tenant-b'}, _query(second, 2))

        assert result_a['session_id'] == result_b['session_id'] == 'global'
        assert len(backend.started_specs) == 2
        assert backend.started_specs[0].session_id != backend.started_specs[1].session_id
        assert backend.started_specs[0].host_path != backend.started_specs[1].host_path
        assert len(runtime.get_sessions()) == 2
    finally:
        server_task.cancel()
        client_task.cancel()
        await runtime.shutdown()


@pytest.mark.asyncio
async def test_cloud_skills_reject_host_paths_and_require_managed_entitlement(tmp_path):
    service, runtime, backend, entitlements, server_task, client_task = await _stack(tmp_path)
    first = _context('workspace-a')
    second = _context('workspace-b')
    ineligible = _context('workspace-free')
    entitlements.snapshots[first.workspace_uuid] = _snapshot(first.workspace_uuid)
    entitlements.snapshots[second.workspace_uuid] = _snapshot(second.workspace_uuid)
    entitlements.snapshots[ineligible.workspace_uuid] = _snapshot(
        ineligible.workspace_uuid,
        managed=False,
    )
    try:
        private = await service.create_skill(
            second,
            {
                'name': 'private',
                'instructions': 'workspace-b secret',
            },
        )
        own_skill = await service.create_skill(
            first,
            {
                'name': 'runner',
                'instructions': 'Run scripts/main.py',
            },
        )
        await service.write_skill_file(first, 'runner', 'scripts/main.py', "print('ok')")
        await service.write_skill_file(first, 'runner', 'requirements.txt', 'requests==2.32.0\n')
        refreshed_skill = await service.get_skill(first, 'runner')
        assert refreshed_skill is not None
        assert refreshed_skill['python_project'] is True
        await service.execute_tool(
            {
                'command': 'python /workspace/.skills/runner/scripts/main.py',
                'workdir': '/workspace/.skills/runner',
            },
            _query(first, 91),
            skill_name='runner',
        )

        mounted_spec = backend.started_specs[-1]
        assert len(mounted_spec.extra_mounts) == 1
        assert mounted_spec.extra_mounts[0].host_path == own_skill['package_root']
        assert mounted_spec.extra_mounts[0].mount_path == '/workspace/.skills/runner'
        assert mounted_spec.extra_mounts[0].mode.value == 'ro'

        with pytest.raises(BoxAdmissionError, match='Scanning arbitrary host'):
            await service.scan_skill_directory(first, private['package_root'])
        with pytest.raises(BoxAdmissionError, match='package_root is runtime-owned'):
            await service.create_skill(
                first,
                {
                    'name': 'stolen',
                    'package_root': private['package_root'],
                },
            )

        assert await service.get_skill(first, 'private') is None
        with pytest.raises(EntitlementUnavailableError):
            await service.list_skills(ineligible)
    finally:
        server_task.cancel()
        client_task.cancel()
        await runtime.shutdown()


@pytest.mark.asyncio
async def test_forged_plan_network_session_and_managed_process_never_reach_runtime(tmp_path):
    service, runtime, backend, entitlements, server_task, client_task = await _stack(tmp_path)
    context = _context('workspace-a')
    entitlements.snapshots[context.workspace_uuid] = _snapshot(context.workspace_uuid)
    query = _query(context, 1)
    try:
        with pytest.raises(BoxAdmissionError, match='host-controlled'):
            await service.execute_spec_payload(
                {'cmd': 'true', 'session_id': 'global', 'plan': 'pro'},
                query,
            )
        with pytest.raises(BoxAdmissionError, match='network access is disabled'):
            await service.execute_spec_payload(
                {'cmd': 'true', 'session_id': 'global', 'network': 'on'},
                query,
            )
        with pytest.raises(BoxAdmissionError, match='session_id is runtime-owned'):
            await service.execute_spec_payload(
                {'cmd': 'true', 'session_id': 'attacker'},
                query,
            )
        with pytest.raises(BoxAdmissionError, match='Managed processes are disabled'):
            await service.start_managed_process(
                context,
                'global',
                {'command': 'sleep', 'args': ['60']},
            )

        assert backend.started_specs == []
        assert runtime.get_sessions() == []
    finally:
        server_task.cancel()
        client_task.cancel()
        await runtime.shutdown()