From e77acfa3abd05afe6faee0b779b4a01f1ef1b62a Mon Sep 17 00:00:00 2001 From: RockChinQ Date: Fri, 18 Sep 2026 01:07:43 +0800 Subject: [PATCH] fix(box): align integration tests with runner-owned bindings --- src/langbot/pkg/box/service.py | 2 + .../box/test_box_integration.py | 15 +++++-- .../test_cloud_box_admission_integration.py | 42 +++++++++++++------ tests/unit_tests/box/test_box_service.py | 22 ++++++++-- tests/unit_tests/provider/test_skill_tools.py | 34 +++++++++++---- 5 files changed, 89 insertions(+), 26 deletions(-) diff --git a/src/langbot/pkg/box/service.py b/src/langbot/pkg/box/service.py index 6208f316d..25d0614c3 100644 --- a/src/langbot/pkg/box/service.py +++ b/src/langbot/pkg/box/service.py @@ -585,6 +585,8 @@ class BoxService: from .runner import binding_for binding = binding_for(query) + if spec_payload.get('session_id') not in (None, '', binding.session_id): + raise BoxValidationError('session_id must match the bound Box') spec_payload = {**binding.spec, **spec_payload, 'session_id': binding.session_id} execution_context = await self._validated_execution_context(self._query_execution_context(query)) spec_payload = self._managed_policy_payload(execution_context, spec_payload) diff --git a/tests/integration_tests/box/test_box_integration.py b/tests/integration_tests/box/test_box_integration.py index cda85815a..9358515cd 100644 --- a/tests/integration_tests/box/test_box_integration.py +++ b/tests/integration_tests/box/test_box_integration.py @@ -4,8 +4,7 @@ These tests verify the end-to-end behavior of the Box sandbox execution system. Tests decorated with ``requires_container`` need a real container runtime (Podman or Docker) and are skipped otherwise. -CI only runs ``tests/unit_tests/``, so these tests never execute in the -CI pipeline. Run them locally with:: +CI runs these tests in the Box Integration Tests job. Run them locally with:: pytest tests/integration_tests/ -v """ @@ -23,6 +22,8 @@ from unittest.mock import AsyncMock import pytest from langbot.pkg.box.service import BoxService +from langbot.pkg.box.runner import RunnerBoxService +from langbot.pkg.api.http.context import ExecutionContext from langbot_plugin.box.backend import BaseSandboxBackend from langbot_plugin.box.client import ActionRPCBoxClient from langbot_plugin.box.errors import BoxBackendUnavailableError @@ -358,6 +359,14 @@ async def test_full_service_to_remote_runtime(tmp_path): workspace_uuid=_ACTION_CONTEXT.workspace_uuid, placement_generation=_ACTION_CONTEXT.placement_generation, ) + context = ExecutionContext( + instance_uuid=_ACTION_CONTEXT.instance_uuid, + workspace_uuid=_ACTION_CONTEXT.workspace_uuid, + placement_generation=_ACTION_CONTEXT.placement_generation, + ) + runner_box = RunnerBoxService(service) + box = await runner_box.acquire(context, {'reuse_key': 'integration-test'}, query) + await runner_box.bind(context, query, 'run-42', box['id']) result = await service.execute_tool( {'command': 'echo service-path'}, query, @@ -366,7 +375,7 @@ async def test_full_service_to_remote_runtime(tmp_path): assert result['ok'] is True assert result['status'] == 'completed' assert 'service-path' in result['stdout'] - assert result['session_id'] == 'query_42' + assert result['session_id'] == box['id'] finally: server_task.cancel() client_task.cancel() diff --git a/tests/integration_tests/box/test_cloud_box_admission_integration.py b/tests/integration_tests/box/test_cloud_box_admission_integration.py index 96610d511..1080dc68d 100644 --- a/tests/integration_tests/box/test_cloud_box_admission_integration.py +++ b/tests/integration_tests/box/test_cloud_box_admission_integration.py @@ -10,7 +10,7 @@ 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.errors import BoxAdmissionError, BoxValidationError from langbot_plugin.box.models import ( BoxExecutionResult, BoxExecutionStatus, @@ -24,6 +24,7 @@ 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.box.runner import RunnerBoxService from langbot.pkg.cloud.entitlements import ( EntitlementResolver, EntitlementSnapshot, @@ -41,6 +42,7 @@ class _AdmissionBackend(BaseSandboxBackend): def __init__(self, logger): super().__init__(logger) self.started_specs: list[BoxSpec] = [] + self.executed_specs: list[BoxSpec] = [] self.stopped_sessions: list[str] = [] async def is_available(self) -> bool: @@ -82,6 +84,7 @@ class _AdmissionBackend(BaseSandboxBackend): ) async def exec(self, session: BoxSessionInfo, spec: BoxSpec) -> BoxExecutionResult: + self.executed_specs.append(spec) await asyncio.sleep(0) return BoxExecutionResult( session_id=session.session_id, @@ -177,6 +180,19 @@ def _query(context: ExecutionContext, query_id: int): return query +async def _bound_query(service, context, query_id): + query = _query(context, query_id) + runner_box = RunnerBoxService(service) + box = await runner_box.acquire(context, {'reuse_key': 'global'}, query) + await runner_box.bind(context, query, f'run-{query_id}', box['id']) + return query + + +async def _execute(service, context, query_id, command): + query = await _bound_query(service, context, query_id) + return await service.execute_tool({'command': command}, query) + + async def _stack(tmp_path): shared_root = tmp_path / 'shared-box' workspace_root = shared_root / 'workspaces' @@ -244,8 +260,8 @@ async def test_concurrent_first_use_creates_one_persistent_global_session(tmp_pa 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)), + _execute(service, context, 1, 'echo first'), + _execute(service, context, 2, 'echo second'), ) assert first['session_id'] == 'global' @@ -270,7 +286,8 @@ async def test_entitlement_loss_revokes_and_closes_existing_global_session(tmp_p 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)) + query = await _bound_query(service, context, 1) + await service.execute_tool({'command': 'true'}, query) assert len(runtime.get_sessions()) == 1 entitlements.snapshots[context.workspace_uuid] = _snapshot( @@ -279,7 +296,7 @@ async def test_entitlement_loss_revokes_and_closes_existing_global_session(tmp_p managed=False, ) with pytest.raises(EntitlementUnavailableError): - await service.execute_tool({'command': 'true'}, _query(context, 2)) + await service.execute_tool({'command': 'true'}, query) assert runtime.get_sessions() == [] assert len(backend.stopped_sessions) == 1 @@ -297,8 +314,8 @@ async def test_two_workspaces_get_isolated_physical_sessions_and_paths(tmp_path) 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)) + result_a = await _execute(service, first, 1, 'tenant-a') + result_b = await _execute(service, second, 2, 'tenant-b') assert result_a['session_id'] == result_b['session_id'] == 'global' assert len(backend.started_specs) == 2 @@ -348,7 +365,7 @@ async def test_cloud_skills_reject_host_paths_and_require_managed_entitlement(tm 'command': 'python /workspace/.skills/runner/scripts/main.py', 'workdir': '/workspace/.skills/runner', }, - _query(first, 91), + await _bound_query(service, first, 91), skill_name='runner', ) @@ -383,8 +400,8 @@ async def test_forged_plan_network_session_and_managed_process_never_reach_runti 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: + query = await _bound_query(service, context, 1) with pytest.raises(BoxAdmissionError, match='host-controlled'): await service.execute_spec_payload( {'cmd': 'true', 'session_id': 'global', 'plan': 'pro'}, @@ -395,7 +412,7 @@ async def test_forged_plan_network_session_and_managed_process_never_reach_runti {'cmd': 'true', 'session_id': 'global', 'network': 'on'}, query, ) - with pytest.raises(BoxAdmissionError, match='session_id is runtime-owned'): + with pytest.raises(BoxValidationError, match='session_id must match the bound Box'): await service.execute_spec_payload( {'cmd': 'true', 'session_id': 'attacker'}, query, @@ -407,8 +424,9 @@ async def test_forged_plan_network_session_and_managed_process_never_reach_runti {'command': 'sleep', 'args': ['60']}, ) - assert backend.started_specs == [] - assert runtime.get_sessions() == [] + assert len(backend.started_specs) == 1 + assert len(runtime.get_sessions()) == 1 + assert backend.executed_specs == [] finally: server_task.cancel() client_task.cancel() diff --git a/tests/unit_tests/box/test_box_service.py b/tests/unit_tests/box/test_box_service.py index 0762d2bd0..655eba7e1 100644 --- a/tests/unit_tests/box/test_box_service.py +++ b/tests/unit_tests/box/test_box_service.py @@ -670,7 +670,7 @@ async def test_box_service_allows_host_mount_under_configured_root(tmp_path): 'cmd': 'pwd', 'host_path': str(host_dir), 'host_path_mode': BoxHostMountMode.READ_WRITE.value, - 'session_id': '11', + 'session_id': 'person_test_user', }, make_query(11), ) @@ -962,7 +962,7 @@ async def test_profile_unlocked_field_can_be_overridden(): await service.initialize() result = await service.execute_spec_payload( - {'cmd': 'echo hi', 'timeout_sec': 60, 'network': 'on', 'session_id': '31'}, + {'cmd': 'echo hi', 'timeout_sec': 60, 'network': 'on', 'session_id': 'person_test_user'}, make_query(31), ) @@ -984,7 +984,7 @@ async def test_profile_locked_field_cannot_be_overridden(): await service.initialize() result = await service.execute_spec_payload( - {'cmd': 'echo hi', 'network': 'on', 'host_path_mode': 'rw', 'session_id': '32'}, + {'cmd': 'echo hi', 'network': 'on', 'host_path_mode': 'rw', 'session_id': 'person_test_user'}, make_query(32), ) @@ -1193,7 +1193,7 @@ async def test_profile_offline_readonly_locks_read_only_rootfs(): await service.initialize() await service.execute_spec_payload( - {'cmd': 'echo hi', 'read_only_rootfs': False, 'session_id': '41'}, make_query(41) + {'cmd': 'echo hi', 'read_only_rootfs': False, 'session_id': 'person_test_user'}, make_query(41) ) spec = backend.start_specs[0] @@ -2372,3 +2372,17 @@ class TestAttachmentHostPath: service.default_workspace = None # Must not raise. await service._purge_attachment_dirs() + + +@pytest.mark.asyncio +async def test_execution_cannot_override_runner_bound_session(): + logger = Mock() + backend = FakeBackend(logger) + runtime = BoxRuntime(logger=logger, backends=[backend], session_ttl_sec=300) + service = BoxService(make_app(logger), client=_InProcessBoxRuntimeClient(logger, runtime)) + await service.initialize() + + with pytest.raises(BoxValidationError, match='session_id must match the bound Box'): + await service.execute_spec_payload({'cmd': 'true', 'session_id': 'other-box'}, make_query()) + + assert backend.start_calls == [] diff --git a/tests/unit_tests/provider/test_skill_tools.py b/tests/unit_tests/provider/test_skill_tools.py index 133d926b0..b0e7288ca 100644 --- a/tests/unit_tests/provider/test_skill_tools.py +++ b/tests/unit_tests/provider/test_skill_tools.py @@ -8,6 +8,7 @@ from unittest.mock import AsyncMock, Mock import pytest from langbot.pkg.api.http.context import ExecutionContext +from langbot.pkg.box.runner import RunBoxBinding _CONTEXT = ExecutionContext( @@ -30,6 +31,12 @@ def _make_query(*, variables=None, **kwargs): ) +def _make_bound_query(**kwargs): + query = _make_query(**kwargs) + query._box_binding = RunBoxBinding('run-a', 'skill-box', {}, 'run-a') + return query + + def _make_skill_manager(skills: dict[str, dict], **kwargs): return SimpleNamespace( skills=skills, @@ -546,14 +553,14 @@ class TestNativeToolLoaderSkillPaths: await loader.invoke_tool( 'read', {'path': '/workspace/.skills/demo/SKILL.md'}, - _make_query(query_id='q1', variables={PIPELINE_BOUND_SKILLS_KEY: ['demo']}), + _make_bound_query(query_id='q1', variables={PIPELINE_BOUND_SKILLS_KEY: ['demo']}), ) return result = await loader.invoke_tool( 'read', {'path': '/workspace/.skills/demo/SKILL.md'}, - _make_query(query_id='q1', variables={PIPELINE_BOUND_SKILLS_KEY: ['demo']}), + _make_bound_query(query_id='q1', variables={PIPELINE_BOUND_SKILLS_KEY: ['demo']}), ) assert result['ok'] is True @@ -577,7 +584,7 @@ class TestNativeToolLoaderSkillPaths: ) ap.skill_mgr = _make_skill_manager({'demo': _make_skill_data(name='demo', package_root=tmpdir)}) loader = NativeToolLoader(ap) - query = _make_query( + query = _make_bound_query( query_id='q-external-read', variables={PIPELINE_BOUND_SKILLS_KEY: ['demo']}, ) @@ -609,7 +616,7 @@ class TestNativeToolLoaderSkillPaths: ) ap.skill_mgr = _make_skill_manager({'demo': _make_skill_data(name='demo', package_root=tmpdir)}) loader = NativeToolLoader(ap) - query = _make_query( + query = _make_bound_query( query_id='q-external-no-protocol', variables={PIPELINE_BOUND_SKILLS_KEY: ['demo']}, ) @@ -639,7 +646,7 @@ class TestNativeToolLoaderSkillPaths: ap.skill_mgr = SimpleNamespace(refresh_skill_from_disk=Mock()) loader = NativeToolLoader(ap) - query = _make_query(query_id='q1', launcher_type='person', launcher_id='123') + query = _make_bound_query(query_id='q1', launcher_type='person', launcher_id='123') register_activated_skill(query, _make_skill_data(name='demo', package_root=tmpdir)) result = await loader.invoke_tool( @@ -671,7 +678,7 @@ class TestNativeToolLoaderSkillPaths: ) ap.skill_mgr = SimpleNamespace(refresh_skill_from_disk=Mock()) loader = NativeToolLoader(ap) - query = _make_query(query_id='q-external', launcher_type='person', launcher_id='123') + query = _make_bound_query(query_id='q-external', launcher_type='person', launcher_id='123') register_activated_skill( query, _make_skill_data( @@ -709,7 +716,7 @@ class TestNativeToolLoaderSkillPaths: ap.skill_mgr = _make_skill_manager({'demo': _make_skill_data(name='demo', package_root=tmpdir)}) loader = NativeToolLoader(ap) - query = _make_query(query_id='q1', variables={PIPELINE_BOUND_SKILLS_KEY: ['demo']}) + query = _make_bound_query(query_id='q1', variables={PIPELINE_BOUND_SKILLS_KEY: ['demo']}) with pytest.raises(ValueError, match='Skill "demo" is not available at this path'): await loader.invoke_tool( @@ -717,3 +724,16 @@ class TestNativeToolLoaderSkillPaths: {'path': '/workspace/.skills/demo/notes.txt', 'content': 'hi'}, query, ) + + +@pytest.mark.asyncio +async def test_native_skill_tools_require_runner_box_binding(): + from langbot.pkg.provider.tools.loaders.native import NativeToolLoader + from langbot_plugin.box.errors import BoxValidationError + + ap = _make_ap() + ap.box_service = SimpleNamespace(available=True, execute_tool=AsyncMock()) + loader = NativeToolLoader(ap) + with pytest.raises(BoxValidationError, match='Runner must bind a Box'): + await loader.invoke_tool('exec', {'command': 'true'}, _make_query()) + ap.box_service.execute_tool.assert_not_awaited()