feat(tenancy): implement workspace isolation

tests/unit_tests/plugin/test_handler_actions.py

@@ -8,13 +8,75 @@ from unittest.mock import AsyncMock, Mock
 
 import pytest
 from langbot_plugin.entities.io.actions.enums import PluginToRuntimeAction, RuntimeToLangBotAction
+from langbot_plugin.entities.io.context import ActionContext
+
+from langbot.pkg.api.http.context import ExecutionContext
+from langbot.pkg.storage.mgr import StorageMgr
 
 
-def make_handler(app):
+TEST_EXECUTION_CONTEXT = ExecutionContext(
+    instance_uuid='instance-a',
+    workspace_uuid='workspace-a',
+    placement_generation=1,
+)
+
+
+def canonical_binary_key(owner_type: str, owner: str, key: str) -> str:
+    return StorageMgr.canonical_binary_storage_key(
+        TEST_EXECUTION_CONTEXT,
+        owner_type=owner_type,
+        owner=owner,
+        key=key,
+    )
+
+
+def make_handler(app, workspace_context: ActionContext | None = None):
     """Create a RuntimeConnectionHandler with mocked external connection."""
     from langbot.pkg.plugin.handler import RuntimeConnectionHandler
 
-    return RuntimeConnectionHandler(Mock(), AsyncMock(return_value=True), app)
+    workspace_context = workspace_context or ActionContext(
+        instance_uuid='instance-a',
+        workspace_uuid='workspace-a',
+        placement_generation=1,
+    )
+    app.workspace_service = SimpleNamespace(
+        get_execution_binding=AsyncMock(
+            return_value=SimpleNamespace(
+                instance_uuid=workspace_context.instance_uuid,
+                workspace_uuid=workspace_context.workspace_uuid,
+                placement_generation=workspace_context.placement_generation,
+            )
+        )
+    )
+    runtime_handler = RuntimeConnectionHandler(
+        Mock(),
+        AsyncMock(return_value=True),
+        app,
+        workspace_context,
+    )
+    installation_uuid = runtime_handler._remember_installation(
+        workspace_context,
+        'test-author',
+        'test-plugin',
+    )
+    runtime_handler.bind_action_context(workspace_context.for_installation(installation_uuid))
+    query_pool = getattr(app, 'query_pool', None)
+    if query_pool is not None and hasattr(query_pool, 'cached_queries'):
+
+        def scoped_query(query):
+            if query is not None:
+                query.instance_uuid = workspace_context.instance_uuid
+                query.workspace_uuid = workspace_context.workspace_uuid
+                query.placement_generation = workspace_context.placement_generation
+            return query
+
+        query_pool.get_query = AsyncMock(
+            side_effect=lambda workspace_uuid, query_uuid: scoped_query(query_pool.cached_queries.get(query_uuid))
+        )
+        query_pool.get_query_by_legacy_id = AsyncMock(
+            side_effect=lambda workspace_uuid, query_id: scoped_query(query_pool.cached_queries.get(query_id))
+        )
+    return runtime_handler
 
 
 def make_result(first_item=None):
@@ -34,6 +96,8 @@ class TestRagRerankAction:
     def app(self):
         mock_app = Mock()
         mock_app.model_mgr = Mock()
+        mock_app.persistence_mgr = Mock()
+        mock_app.persistence_mgr.execute_async = AsyncMock(return_value=make_result(SimpleNamespace(uuid='rerank-1')))
         mock_app.logger = Mock()
         return mock_app
 
@@ -63,12 +127,16 @@ class TestRagRerankAction:
 
         assert response.code == 0
         assert response.data['results'] == [{'index': 1, 'relevance_score': 0.9}]
-        app.model_mgr.get_rerank_model_by_uuid.assert_awaited_once_with('rerank-1')
+        app.model_mgr.get_rerank_model_by_uuid.assert_awaited_once_with(
+            TEST_EXECUTION_CONTEXT,
+            'rerank-1',
+        )
         provider.invoke_rerank.assert_awaited_once_with(
             model=rerank_model,
             query='hello',
             documents=['a', 'b'],
             extra_args={'return_documents': False},
+            execution_context=TEST_EXECUTION_CONTEXT,
         )
 
     @pytest.mark.asyncio
@@ -122,6 +190,7 @@ class TestInitializePluginSettings:
         assert app.persistence_mgr.execute_async.await_count == 2
         insert_params = compiled_params(app.persistence_mgr.execute_async.await_args_list[1].args[0])
         assert insert_params == {
+            'workspace_uuid': 'workspace-a',
             'plugin_author': 'test-author',
             'plugin_name': 'test-plugin',
             'install_source': 'local',
@@ -218,7 +287,12 @@ class TestSetBinaryStorage:
         assert response.code == 0
         assert app.persistence_mgr.execute_async.await_count == 2
         insert_params = compiled_params(app.persistence_mgr.execute_async.await_args_list[1].args[0])
-        assert insert_params['unique_key'] == 'plugin:test-owner:test-key'
+        assert insert_params['workspace_uuid'] == 'workspace-a'
+        assert insert_params['unique_key'] == canonical_binary_key(
+            'plugin',
+            'test-author/test-plugin',
+            'test-key',
+        )
         assert insert_params['value'] == b'x' * 512
 
     @pytest.mark.asyncio
@@ -231,7 +305,15 @@ class TestSetBinaryStorage:
 
         assert response.code == 0
         assert app.persistence_mgr.execute_async.await_count == 2
+        select_params = compiled_params(app.persistence_mgr.execute_async.await_args_list[0].args[0])
         update_params = compiled_params(app.persistence_mgr.execute_async.await_args_list[1].args[0])
+        expected_key = canonical_binary_key(
+            'plugin',
+            'test-author/test-plugin',
+            'test-key',
+        )
+        assert expected_key in select_params.values()
+        assert expected_key in update_params.values()
         assert update_params['value'] == b'new'
 
     @pytest.mark.asyncio
@@ -304,6 +386,11 @@ class TestGetPluginSettings:
             'plugin_config': {},
             'install_source': 'local',
             'install_info': {},
+            'installation_uuid': runtime_handler.derive_installation_uuid(
+                runtime_handler.require_bound_action_context(),
+                'test-author',
+                'test-plugin',
+            ),
         }
 
     @pytest.mark.asyncio
@@ -333,9 +420,90 @@ class TestGetPluginSettings:
             'plugin_config': {'custom': 'config'},
             'install_source': 'github',
             'install_info': {'repo': 'test/repo'},
+            'installation_uuid': runtime_handler.derive_installation_uuid(
+                runtime_handler.require_bound_action_context(),
+                'test-author',
+                'test-plugin',
+            ),
         }
 
 
+class TestGetConfigFile:
+    """Plugin config files remain bound to the trusted runtime placement."""
+
+    WORKSPACE_A = '11111111-1111-4111-8111-111111111111'
+    WORKSPACE_B = '22222222-2222-4222-8222-222222222222'
+
+    @pytest.fixture
+    def app(self):
+        mock_app = Mock()
+        mock_app.persistence_mgr = Mock()
+        mock_app.persistence_mgr.execute_async = AsyncMock()
+        mock_app.storage_mgr = StorageMgr(mock_app)
+        mock_app.storage_mgr.storage_provider = SimpleNamespace(load=AsyncMock(return_value=b'plugin config bytes'))
+        mock_app.logger = Mock()
+        return mock_app
+
+    @staticmethod
+    def object_key(
+        *,
+        workspace_uuid: str = WORKSPACE_A,
+        placement_generation: int = 1,
+        owner_type: str = 'plugin_config',
+    ) -> str:
+        return StorageMgr.scoped_object_key(
+            ExecutionContext(
+                instance_uuid='instance-a',
+                workspace_uuid=workspace_uuid,
+                placement_generation=placement_generation,
+            ),
+            owner_type=owner_type,
+            owner=workspace_uuid,
+            key='config.json',
+        )
+
+    async def invoke(self, app, file_key: str):
+        runtime_handler = make_handler(
+            app,
+            ActionContext(
+                instance_uuid='instance-a',
+                workspace_uuid=self.WORKSPACE_A,
+                placement_generation=1,
+            ),
+        )
+        app.persistence_mgr.execute_async.return_value = make_result(
+            SimpleNamespace(config={'uploaded_file': file_key})
+        )
+        return await runtime_handler.actions[PluginToRuntimeAction.GET_CONFIG_FILE.value]({'file_key': file_key})
+
+    @pytest.mark.asyncio
+    async def test_loads_config_file_from_same_workspace_generation(self, app):
+        file_key = self.object_key()
+
+        response = await self.invoke(app, file_key)
+
+        assert response.code == 0
+        assert base64.b64decode(response.data['file_base64']) == b'plugin config bytes'
+        app.storage_mgr.storage_provider.load.assert_awaited_once_with(file_key)
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        'file_key',
+        [
+            object_key.__func__(workspace_uuid=WORKSPACE_B),
+            object_key.__func__(placement_generation=2),
+            object_key.__func__(owner_type='upload_document'),
+        ],
+        ids=['other-workspace', 'stale-generation', 'wrong-owner-type'],
+    )
+    async def test_rejects_config_file_outside_trusted_scope(self, app, file_key):
+        response = await self.invoke(app, file_key)
+
+        assert response.code != 0
+        assert 'Failed to load config file' in response.message
+        app.storage_mgr.storage_provider.load.assert_not_awaited()
+
+
 class TestGetBinaryStorage:
     """Tests for get_binary_storage action handler."""
 
@@ -364,6 +532,16 @@ class TestGetBinaryStorage:
         assert response.data == {
             'value_base64': base64.b64encode(b'test binary content').decode('utf-8'),
         }
+        statement_params = compiled_params(app.persistence_mgr.execute_async.await_args.args[0])
+        assert 'workspace-a' in statement_params.values()
+        assert (
+            canonical_binary_key(
+                'plugin',
+                'test-author/test-plugin',
+                'test-key',
+            )
+            in statement_params.values()
+        )
 
     @pytest.mark.asyncio
     async def test_returns_error_when_not_found(self, app):
@@ -383,6 +561,63 @@ class TestGetBinaryStorage:
         assert 'Storage with key test-key not found' in response.message
 
 
+class TestDeleteAndListBinaryStorage:
+    """Delete/list remain fenced to the trusted canonical owner scope."""
+
+    @pytest.fixture
+    def app(self):
+        mock_app = Mock()
+        mock_app.persistence_mgr = Mock()
+        mock_app.persistence_mgr.execute_async = AsyncMock()
+        return mock_app
+
+    @pytest.mark.asyncio
+    async def test_delete_uses_workspace_and_canonical_unique_key(self, app):
+        runtime_handler = make_handler(app)
+
+        response = await runtime_handler.actions[RuntimeToLangBotAction.DELETE_BINARY_STORAGE.value](
+            {
+                'key': 'test-key',
+                'owner_type': 'plugin',
+                'owner': 'forged-owner',
+            }
+        )
+
+        assert response.code == 0
+        statement_params = compiled_params(app.persistence_mgr.execute_async.await_args.args[0])
+        assert 'workspace-a' in statement_params.values()
+        assert (
+            canonical_binary_key(
+                'plugin',
+                'test-author/test-plugin',
+                'test-key',
+            )
+            in statement_params.values()
+        )
+        assert 'forged-owner' not in statement_params.values()
+
+    @pytest.mark.asyncio
+    async def test_list_keys_uses_trusted_plugin_owner(self, app):
+        result = Mock()
+        result.scalars.return_value.all.return_value = ['first', 'second']
+        app.persistence_mgr.execute_async.return_value = result
+        runtime_handler = make_handler(app)
+
+        response = await runtime_handler.actions[RuntimeToLangBotAction.GET_BINARY_STORAGE_KEYS.value](
+            {
+                'owner_type': 'plugin',
+                'owner': 'forged-owner',
+            }
+        )
+
+        assert response.code == 0
+        assert response.data == {'keys': ['first', 'second']}
+        statement_params = compiled_params(app.persistence_mgr.execute_async.await_args.args[0])
+        assert 'workspace-a' in statement_params.values()
+        assert 'test-author/test-plugin' in statement_params.values()
+        assert 'forged-owner' not in statement_params.values()
+
+
 class TestHandlerQueryLookup:
     """Tests for query lookup in cached_queries."""