feat(tenancy): implement workspace isolation

src/langbot/pkg/rag/knowledge/base.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 import abc
 
 from langbot.pkg.core import app
+from langbot.pkg.api.http.context import ExecutionContext
 from langbot_plugin.api.entities.builtin.rag import context as rag_context
 
 
@@ -22,10 +23,16 @@ class KnowledgeBaseInterface(metaclass=abc.ABCMeta):
         pass
 
     @abc.abstractmethod
-    async def retrieve(self, query: str, settings: dict | None = None) -> list[rag_context.RetrievalResultEntry]:
+    async def retrieve(
+        self,
+        execution_context: ExecutionContext,
+        query: str,
+        settings: dict | None = None,
+    ) -> list[rag_context.RetrievalResultEntry]:
         """Retrieve relevant documents from the knowledge base
 
         Args:
+            execution_context: Trusted active Workspace placement.
             query: The query string
             settings: Optional per-request retrieval settings overrides
 
@@ -50,6 +57,6 @@ class KnowledgeBaseInterface(metaclass=abc.ABCMeta):
         pass
 
     @abc.abstractmethod
-    async def dispose(self):
+    async def dispose(self, execution_context: ExecutionContext):
         """Clean up resources"""
         pass

src/langbot/pkg/rag/knowledge/kbmgr.py

@@ -1,18 +1,22 @@
 from __future__ import annotations
+import io
 import mimetypes
 import os.path
 import traceback
 import uuid
 import zipfile
-import io
 from typing import Any
-from langbot.pkg.core import app
+
 import sqlalchemy
-
-
-from langbot.pkg.entity.persistence import rag as persistence_rag
-from langbot.pkg.core import taskmgr
 from langbot_plugin.api.entities.builtin.rag import context as rag_context
+
+from langbot.pkg.api.http.authz import WorkspaceRequiredError
+from langbot.pkg.api.http.context import ExecutionContext, RequestContext
+from langbot.pkg.api.http.service.tenant import TenantContext, require_workspace_uuid
+from langbot.pkg.core import app, taskmgr
+from langbot.pkg.entity.persistence import rag as persistence_rag
+from langbot.pkg.workspace.errors import WorkspaceNotFoundError
+
 from .base import KnowledgeBaseInterface
 
 
@@ -21,20 +25,78 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
 
     knowledge_base_entity: persistence_rag.KnowledgeBase
 
-    def __init__(self, ap: app.Application, knowledge_base_entity: persistence_rag.KnowledgeBase):
+    def __init__(
+        self,
+        ap: app.Application,
+        knowledge_base_entity: persistence_rag.KnowledgeBase,
+        execution_context: ExecutionContext,
+    ):
         super().__init__(ap)
         self.knowledge_base_entity = knowledge_base_entity
+        self.execution_context = execution_context
 
     async def initialize(self):
         pass
 
+    async def _assert_execution_context(self, execution_context: ExecutionContext) -> None:
+        """Reject stale or cross-Workspace runtime access."""
+
+        if not isinstance(execution_context, ExecutionContext):
+            raise WorkspaceRequiredError('ExecutionContext is required for knowledge runtime access')
+        if (
+            execution_context.instance_uuid != self.execution_context.instance_uuid
+            or execution_context.workspace_uuid != self.execution_context.workspace_uuid
+            or execution_context.placement_generation != self.execution_context.placement_generation
+        ):
+            raise WorkspaceNotFoundError('Knowledge base not found')
+        if self.knowledge_base_entity.workspace_uuid != execution_context.workspace_uuid:
+            raise WorkspaceNotFoundError('Knowledge base not found')
+        binding = await self.ap.workspace_service.get_execution_binding(
+            execution_context.workspace_uuid,
+            expected_generation=execution_context.placement_generation,
+        )
+        if binding.instance_uuid != execution_context.instance_uuid:
+            raise WorkspaceNotFoundError('Knowledge base not found')
+
+    async def _require_plugin_runtime_context(
+        self,
+        execution_context: ExecutionContext,
+    ) -> ExecutionContext:
+        """Fence every singleton Plugin Runtime call to this runtime KB."""
+
+        await self._assert_execution_context(execution_context)
+        return await self.ap.plugin_connector.require_workspace_context(execution_context)
+
+    def _require_upload_object_key(
+        self,
+        execution_context: ExecutionContext,
+        object_key: str,
+    ) -> None:
+        """Reject raw, cross-Workspace, stale, or non-upload object keys."""
+
+        try:
+            self.ap.storage_mgr.require_scoped_object_key(
+                execution_context,
+                object_key,
+                expected_owner_type='upload_document',
+            )
+        except (WorkspaceRequiredError, ValueError) as exc:
+            raise WorkspaceNotFoundError('Upload not found') from exc
+
     async def _store_file_task(
-        self, file: persistence_rag.File, task_context: taskmgr.TaskContext, parser_plugin_id: str | None = None
+        self,
+        execution_context: ExecutionContext,
+        file: persistence_rag.File,
+        task_context: taskmgr.TaskContext,
+        parser_plugin_id: str | None = None,
     ):
+        await self._assert_execution_context(execution_context)
+        self._require_upload_object_key(execution_context, file.file_name)
         try:
             # set file status to processing
             await self.ap.persistence_mgr.execute_async(
                 sqlalchemy.update(persistence_rag.File)
+                .where(persistence_rag.File.workspace_uuid == execution_context.workspace_uuid)
                 .where(persistence_rag.File.uuid == file.uuid)
                 .values(status='processing')
             )
@@ -42,7 +104,11 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
             task_context.set_current_action('Processing file')
 
             # Get file size from storage
-            file_size = await self.ap.storage_mgr.storage_provider.size(file.file_name)
+            file_size = await self.ap.storage_mgr.size_scoped_object_key(
+                execution_context,
+                file.file_name,
+                expected_owner_type='upload_document',
+            )
 
             # Detect MIME type from extension
             mime_type, _ = mimetypes.guess_type(file.file_name)
@@ -53,16 +119,22 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
             parsed_content = None
             if parser_plugin_id:
                 task_context.set_current_action('Parsing file')
-                file_bytes = await self.ap.storage_mgr.storage_provider.load(file.file_name)
+                file_bytes = await self.ap.storage_mgr.load_scoped_object_key(
+                    execution_context,
+                    file.file_name,
+                    expected_owner_type='upload_document',
+                )
                 parse_context = {
                     'mime_type': mime_type,
                     'filename': file.file_name,
                     'metadata': {},
                 }
+                await self._require_plugin_runtime_context(execution_context)
                 parsed_content = await self.ap.plugin_connector.call_parser(parser_plugin_id, parse_context, file_bytes)
 
             # Call plugin to ingest document
             result = await self._ingest_document(
+                execution_context,
                 {
                     'document_id': file.uuid,
                     'filename': file.file_name,
@@ -80,8 +152,10 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
                 raise Exception(error_msg)
 
             # set file status to completed
+            await self._assert_execution_context(execution_context)
             await self.ap.persistence_mgr.execute_async(
                 sqlalchemy.update(persistence_rag.File)
+                .where(persistence_rag.File.workspace_uuid == execution_context.workspace_uuid)
                 .where(persistence_rag.File.uuid == file.uuid)
                 .values(status='completed')
             )
@@ -89,35 +163,63 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
         except Exception as e:
             self.ap.logger.error(f'Error storing file {file.uuid}: {e}')
             traceback.print_exc()
-            # set file status to failed
-            await self.ap.persistence_mgr.execute_async(
-                sqlalchemy.update(persistence_rag.File)
-                .where(persistence_rag.File.uuid == file.uuid)
-                .values(status='failed')
-            )
+            # A stale placement is fenced from all writes, including failure
+            # status updates from an old background task.
+            try:
+                await self._assert_execution_context(execution_context)
+            except Exception:
+                self.ap.logger.warning(f'Skipping stale RAG task status update for file {file.uuid}')
+            else:
+                await self.ap.persistence_mgr.execute_async(
+                    sqlalchemy.update(persistence_rag.File)
+                    .where(persistence_rag.File.workspace_uuid == execution_context.workspace_uuid)
+                    .where(persistence_rag.File.uuid == file.uuid)
+                    .values(status='failed')
+                )
 
             raise
         finally:
-            # delete file from storage
-            await self.ap.storage_mgr.storage_provider.delete(file.file_name)
+            # An old background task must not touch an upload after its
+            # placement generation has been fenced off.
+            try:
+                await self._assert_execution_context(execution_context)
+                await self.ap.storage_mgr.delete_scoped_object_key(
+                    execution_context,
+                    file.file_name,
+                    expected_owner_type='upload_document',
+                )
+            except (WorkspaceRequiredError, WorkspaceNotFoundError):
+                self.ap.logger.warning(f'Skipping stale RAG upload cleanup for file {file.uuid}')
 
-    async def store_file(self, file_id: str, parser_plugin_id: str | None = None) -> str:
+    async def store_file(
+        self,
+        execution_context: ExecutionContext,
+        file_id: str,
+        parser_plugin_id: str | None = None,
+    ) -> str:
+        await self._assert_execution_context(execution_context)
+        self._require_upload_object_key(execution_context, file_id)
         # pre checking
-        if not await self.ap.storage_mgr.storage_provider.exists(file_id):
-            raise Exception(f'File {file_id} not found')
+        if not await self.ap.storage_mgr.exists_scoped_object_key(
+            execution_context,
+            file_id,
+            expected_owner_type='upload_document',
+        ):
+            raise WorkspaceNotFoundError('Upload not found')
 
         file_name = file_id
         _, ext = os.path.splitext(file_name)
         extension = ext.lstrip('.').lower() if ext else ''
 
         if extension == 'zip':
-            return await self._store_zip_file(file_id, parser_plugin_id=parser_plugin_id)
+            return await self._store_zip_file(execution_context, file_id, parser_plugin_id=parser_plugin_id)
 
         file_uuid = str(uuid.uuid4())
         kb_id = self.knowledge_base_entity.uuid
 
         file_obj_data = {
             'uuid': file_uuid,
+            'workspace_uuid': execution_context.workspace_uuid,
             'kb_id': kb_id,
             'file_name': file_name,
             'extension': extension,
@@ -131,19 +233,38 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
         # run background task asynchronously
         ctx = taskmgr.TaskContext.new()
         wrapper = self.ap.task_mgr.create_user_task(
-            self._store_file_task(file_obj, task_context=ctx, parser_plugin_id=parser_plugin_id),
+            self._store_file_task(
+                execution_context,
+                file_obj,
+                task_context=ctx,
+                parser_plugin_id=parser_plugin_id,
+            ),
             kind='knowledge-operation',
             name=f'knowledge-store-file-{file_id}',
             label=f'Store file {file_id}',
             context=ctx,
+            instance_uuid=execution_context.instance_uuid,
+            workspace_uuid=execution_context.workspace_uuid,
+            placement_generation=execution_context.placement_generation,
         )
         return wrapper.id
 
-    async def _store_zip_file(self, zip_file_id: str, parser_plugin_id: str | None = None) -> str:
+    async def _store_zip_file(
+        self,
+        execution_context: ExecutionContext,
+        zip_file_id: str,
+        parser_plugin_id: str | None = None,
+    ) -> str:
         """Handle ZIP file by extracting each document and storing them separately."""
+        await self._assert_execution_context(execution_context)
+        self._require_upload_object_key(execution_context, zip_file_id)
         self.ap.logger.info(f'Processing ZIP file: {zip_file_id}')
 
-        zip_bytes = await self.ap.storage_mgr.storage_provider.load(zip_file_id)
+        zip_bytes = await self.ap.storage_mgr.load_scoped_object_key(
+            execution_context,
+            zip_file_id,
+            expected_owner_type='upload_document',
+        )
 
         supported_extensions = {'txt', 'pdf', 'docx', 'md', 'html'}
         stored_file_tasks = []
@@ -173,15 +294,31 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
                             continue
 
                         extracted_file_id = file_stem + '_' + str(uuid.uuid4())[:8] + '.' + extension
-                        # save file to storage
+                        extracted_object_key = await self.ap.storage_mgr.save_scoped(
+                            execution_context,
+                            owner_type='upload_document',
+                            owner=f'knowledge-base:{self.knowledge_base_entity.uuid}',
+                            key=extracted_file_id,
+                            value=file_content,
+                        )
 
-                        await self.ap.storage_mgr.storage_provider.save(extracted_file_id, file_content)
-
-                        task_id = await self.store_file(extracted_file_id, parser_plugin_id=parser_plugin_id)
+                        try:
+                            task_id = await self.store_file(
+                                execution_context,
+                                extracted_object_key,
+                                parser_plugin_id=parser_plugin_id,
+                            )
+                        except Exception:
+                            await self.ap.storage_mgr.delete_scoped_object_key(
+                                execution_context,
+                                extracted_object_key,
+                                expected_owner_type='upload_document',
+                            )
+                            raise
                         stored_file_tasks.append(task_id)
 
                         self.ap.logger.info(
-                            f'Extracted and stored file from ZIP: {file_info.filename} -> {extracted_file_id}'
+                            f'Extracted and stored file from ZIP: {file_info.filename} -> {extracted_object_key}'
                         )
 
                     except Exception as e:
@@ -197,20 +334,33 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
             return stored_file_tasks[0] if stored_file_tasks else ''
         finally:
             try:
-                await self.ap.storage_mgr.storage_provider.delete(zip_file_id)
+                await self._assert_execution_context(execution_context)
+                await self.ap.storage_mgr.delete_scoped_object_key(
+                    execution_context,
+                    zip_file_id,
+                    expected_owner_type='upload_document',
+                )
             except FileNotFoundError:
                 pass
+            except (WorkspaceRequiredError, WorkspaceNotFoundError):
+                self.ap.logger.warning(f'Skipping stale RAG ZIP cleanup for upload {zip_file_id}')
             except Exception as e:
                 self.ap.logger.warning(f'Failed to cleanup ZIP file {zip_file_id}: {e}')
 
-    async def retrieve(self, query: str, settings: dict | None = None) -> list[rag_context.RetrievalResultEntry]:
+    async def retrieve(
+        self,
+        execution_context: ExecutionContext,
+        query: str,
+        settings: dict | None = None,
+    ) -> list[rag_context.RetrievalResultEntry]:
+        await self._assert_execution_context(execution_context)
         # Merge stored retrieval_settings with per-request overrides
         stored = self.knowledge_base_entity.retrieval_settings or {}
         merged = {**stored, **(settings or {})}
         if 'top_k' not in merged:
             merged['top_k'] = 5  # fallback default
 
-        response = await self._retrieve(query, merged)
+        response = await self._retrieve(execution_context, query, merged)
 
         results_data = response.get('results', [])
         entries = []
@@ -221,12 +371,25 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
                 entries.append(r)
         return entries
 
-    async def delete_file(self, file_id: str):
-        await self._delete_document(file_id)
+    async def delete_file(self, execution_context: ExecutionContext, file_id: str):
+        await self._assert_execution_context(execution_context)
+        result = await self.ap.persistence_mgr.execute_async(
+            sqlalchemy.select(persistence_rag.File.uuid)
+            .where(persistence_rag.File.workspace_uuid == execution_context.workspace_uuid)
+            .where(persistence_rag.File.kb_id == self.knowledge_base_entity.uuid)
+            .where(persistence_rag.File.uuid == file_id)
+            .limit(1)
+        )
+        if result.first() is None:
+            raise WorkspaceNotFoundError('Knowledge file not found')
+        await self._delete_document(execution_context, file_id)
 
         # Also cleanup DB record
         await self.ap.persistence_mgr.execute_async(
-            sqlalchemy.delete(persistence_rag.File).where(persistence_rag.File.uuid == file_id)
+            sqlalchemy.delete(persistence_rag.File)
+            .where(persistence_rag.File.workspace_uuid == execution_context.workspace_uuid)
+            .where(persistence_rag.File.kb_id == self.knowledge_base_entity.uuid)
+            .where(persistence_rag.File.uuid == file_id)
         )
 
     def get_uuid(self) -> str:
@@ -241,14 +404,16 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
         """Get the Knowledge Engine plugin ID"""
         return self.knowledge_base_entity.knowledge_engine_plugin_id or ''
 
-    async def dispose(self):
+    async def dispose(self, execution_context: ExecutionContext):
         """Dispose the knowledge base, notifying the plugin to cleanup."""
-        await self._on_kb_delete()
+        await self._assert_execution_context(execution_context)
+        await self._on_kb_delete(execution_context)
 
     # ========== Plugin Communication Methods ==========
 
-    async def _on_kb_create(self) -> None:
+    async def _on_kb_create(self, execution_context: ExecutionContext) -> None:
         """Notify plugin about KB creation."""
+        await self._assert_execution_context(execution_context)
         plugin_id = self.knowledge_base_entity.knowledge_engine_plugin_id
         if not plugin_id:
             return
@@ -258,17 +423,20 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
             self.ap.logger.info(
                 f'Calling RAG plugin {plugin_id}: on_knowledge_base_create(kb_id={self.knowledge_base_entity.uuid})'
             )
+            await self._require_plugin_runtime_context(execution_context)
             await self.ap.plugin_connector.rag_on_kb_create(plugin_id, self.knowledge_base_entity.uuid, config)
         except Exception as e:
             self.ap.logger.error(f'Failed to notify plugin {plugin_id} on KB create: {e}')
             raise
 
-    async def _on_kb_delete(self) -> None:
+    async def _on_kb_delete(self, execution_context: ExecutionContext) -> None:
         """Notify plugin about KB deletion."""
+        await self._assert_execution_context(execution_context)
         plugin_id = self.knowledge_base_entity.knowledge_engine_plugin_id
         if not plugin_id:
             return
 
+        await self._require_plugin_runtime_context(execution_context)
         try:
             self.ap.logger.info(
                 f'Calling RAG plugin {plugin_id}: on_knowledge_base_delete(kb_id={self.knowledge_base_entity.uuid})'
@@ -279,11 +447,13 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
 
     async def _ingest_document(
         self,
+        execution_context: ExecutionContext,
         file_metadata: dict[str, Any],
         storage_path: str,
         parsed_content: dict[str, Any] | None = None,
     ) -> dict[str, Any]:
         """Call plugin to ingest document."""
+        await self._assert_execution_context(execution_context)
         kb = self.knowledge_base_entity
         plugin_id = kb.knowledge_engine_plugin_id
         if not plugin_id:
@@ -306,6 +476,7 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
             'parsed_content': parsed_content,
         }
 
+        await self._require_plugin_runtime_context(execution_context)
         try:
             result = await self.ap.plugin_connector.call_rag_ingest(plugin_id, context_data)
             return result
@@ -315,6 +486,7 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
 
     async def _retrieve(
         self,
+        execution_context: ExecutionContext,
         query: str,
         settings: dict[str, Any],
     ) -> dict[str, Any]:
@@ -324,6 +496,7 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
             ValueError: If no RAG plugin is configured for this KB.
             Exception: If the plugin retrieval call fails.
         """
+        await self._assert_execution_context(execution_context)
         kb = self.knowledge_base_entity
         plugin_id = kb.knowledge_engine_plugin_id
         if not plugin_id:
@@ -333,25 +506,28 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
         # for plugins that need it. Do NOT move them into filters, as filters
         # are passed directly to vector_search by some plugins (e.g. LangRAG)
         # and would cause empty results when the metadata field doesn't exist.
-        filters = settings.pop('filters', {})
+        plugin_settings = dict(settings)
+        filters = plugin_settings.pop('filters', {})
 
         retrieval_context = {
             'query': query,
             'knowledge_base_id': kb.uuid,
             'collection_id': kb.collection_id or kb.uuid,
-            'retrieval_settings': settings,
+            'retrieval_settings': plugin_settings,
             'creation_settings': kb.creation_settings or {},
             'filters': filters,
         }
 
+        await self._require_plugin_runtime_context(execution_context)
         result = await self.ap.plugin_connector.call_rag_retrieve(
             plugin_id,
             retrieval_context,
         )
         return result
 
-    async def _delete_document(self, document_id: str) -> bool:
+    async def _delete_document(self, execution_context: ExecutionContext, document_id: str) -> bool:
         """Call plugin to delete document."""
+        await self._assert_execution_context(execution_context)
         kb = self.knowledge_base_entity
         plugin_id = kb.knowledge_engine_plugin_id
         if not plugin_id:
@@ -359,6 +535,7 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
 
         self.ap.logger.info(f'Calling RAG plugin {plugin_id}: delete_document(doc_id={document_id})')
 
+        await self._require_plugin_runtime_context(execution_context)
         try:
             return await self.ap.plugin_connector.call_rag_delete_document(plugin_id, document_id, kb.uuid)
         except Exception as e:
@@ -369,7 +546,7 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
 class RAGManager:
     ap: app.Application
 
-    knowledge_bases: dict[str, KnowledgeBaseInterface]
+    knowledge_bases: dict[tuple[str, str], RuntimeKnowledgeBase]
 
     def __init__(self, ap: app.Application):
         self.ap = ap
@@ -378,20 +555,50 @@ class RAGManager:
     async def initialize(self):
         await self.load_knowledge_bases_from_db()
 
-    async def get_all_knowledge_base_details(self) -> list[dict]:
+    async def _to_execution_context(
+        self,
+        context: RequestContext | ExecutionContext,
+    ) -> ExecutionContext:
+        if isinstance(context, RequestContext):
+            execution_context = ExecutionContext.from_request(context)
+        elif isinstance(context, ExecutionContext):
+            execution_context = context
+        else:
+            raise WorkspaceRequiredError('RequestContext or ExecutionContext is required')
+
+        binding = await self.ap.workspace_service.get_execution_binding(
+            execution_context.workspace_uuid,
+            expected_generation=execution_context.placement_generation,
+        )
+        if binding.instance_uuid != execution_context.instance_uuid:
+            raise WorkspaceNotFoundError('Workspace not found')
+        return execution_context
+
+    async def _get_engine_map(self, context: TenantContext) -> dict[str, dict]:
+        engine_map: dict[str, dict] = {}
+        connector = getattr(self.ap, 'plugin_connector', None)
+        if connector is not None and connector.is_enable_plugin:
+            await connector.require_workspace_context(context)
+            try:
+                engines = await connector.list_knowledge_engines()
+                engine_map = {engine['plugin_id']: engine for engine in engines}
+            except Exception as e:
+                self.ap.logger.warning(f'Failed to list Knowledge Engines: {e}')
+        return engine_map
+
+    async def get_all_knowledge_base_details(self, context: TenantContext) -> list[dict]:
         """Get all knowledge bases with enriched Knowledge Engine details."""
+        workspace_uuid = require_workspace_uuid(context)
         # 1. Get raw KBs from DB
-        result = await self.ap.persistence_mgr.execute_async(sqlalchemy.select(persistence_rag.KnowledgeBase))
+        result = await self.ap.persistence_mgr.execute_async(
+            sqlalchemy.select(persistence_rag.KnowledgeBase).where(
+                persistence_rag.KnowledgeBase.workspace_uuid == workspace_uuid
+            )
+        )
         knowledge_bases = result.all()
 
         # 2. Get all available Knowledge Engines for enrichment
-        engine_map = {}
-        if self.ap.plugin_connector.is_enable_plugin:
-            try:
-                engines = await self.ap.plugin_connector.list_knowledge_engines()
-                engine_map = {e['plugin_id']: e for e in engines}
-            except Exception as e:
-                self.ap.logger.warning(f'Failed to list Knowledge Engines: {e}')
+        engine_map = await self._get_engine_map(context)
 
         # 3. Serialize and enrich
         kb_list = []
@@ -402,10 +609,13 @@ class RAGManager:
 
         return kb_list
 
-    async def get_knowledge_base_details(self, kb_uuid: str) -> dict | None:
+    async def get_knowledge_base_details(self, context: TenantContext, kb_uuid: str) -> dict | None:
         """Get specific knowledge base with enriched Knowledge Engine details."""
+        workspace_uuid = require_workspace_uuid(context)
         result = await self.ap.persistence_mgr.execute_async(
-            sqlalchemy.select(persistence_rag.KnowledgeBase).where(persistence_rag.KnowledgeBase.uuid == kb_uuid)
+            sqlalchemy.select(persistence_rag.KnowledgeBase)
+            .where(persistence_rag.KnowledgeBase.workspace_uuid == workspace_uuid)
+            .where(persistence_rag.KnowledgeBase.uuid == kb_uuid)
         )
         kb = result.first()
         if not kb:
@@ -414,13 +624,7 @@ class RAGManager:
         kb_dict = self.ap.persistence_mgr.serialize_model(persistence_rag.KnowledgeBase, kb)
 
         # Fetch engines
-        engine_map = {}
-        if self.ap.plugin_connector.is_enable_plugin:
-            try:
-                engines = await self.ap.plugin_connector.list_knowledge_engines()
-                engine_map = {e['plugin_id']: e for e in engines}
-            except Exception as e:
-                self.ap.logger.warning(f'Failed to list Knowledge Engines: {e}')
+        engine_map = await self._get_engine_map(context)
 
         self._enrich_kb_dict(kb_dict, engine_map)
         return kb_dict
@@ -465,6 +669,7 @@ class RAGManager:
 
     async def create_knowledge_base(
         self,
+        context: RequestContext | ExecutionContext,
         name: str,
         knowledge_engine_plugin_id: str,
         creation_settings: dict,
@@ -472,8 +677,10 @@ class RAGManager:
         description: str = '',
     ) -> persistence_rag.KnowledgeBase:
         """Create a new knowledge base using a RAG plugin."""
+        execution_context = await self._to_execution_context(context)
         # Validate that the Knowledge Engine plugin exists
         if self.ap.plugin_connector.is_enable_plugin:
+            await self.ap.plugin_connector.require_workspace_context(execution_context)
             try:
                 engines = await self.ap.plugin_connector.list_knowledge_engines()
                 engine_ids = [e.get('plugin_id') for e in engines]
@@ -490,6 +697,7 @@ class RAGManager:
 
         kb_data = {
             'uuid': kb_uuid,
+            'workspace_uuid': execution_context.workspace_uuid,
             'name': name,
             'description': description,
             'knowledge_engine_plugin_id': knowledge_engine_plugin_id,
@@ -505,15 +713,17 @@ class RAGManager:
         await self.ap.persistence_mgr.execute_async(sqlalchemy.insert(persistence_rag.KnowledgeBase).values(kb_data))
 
         # Load into Runtime
-        runtime_kb = await self.load_knowledge_base(kb)
+        runtime_kb = await self.load_knowledge_base(execution_context, kb)
 
         # Notify Plugin — rollback DB record and runtime entry on failure
         try:
-            await runtime_kb._on_kb_create()
+            await runtime_kb._on_kb_create(execution_context)
         except Exception:
-            self.knowledge_bases.pop(kb_uuid, None)
+            self.knowledge_bases.pop((execution_context.workspace_uuid, kb_uuid), None)
             await self.ap.persistence_mgr.execute_async(
-                sqlalchemy.delete(persistence_rag.KnowledgeBase).where(persistence_rag.KnowledgeBase.uuid == kb_uuid)
+                sqlalchemy.delete(persistence_rag.KnowledgeBase)
+                .where(persistence_rag.KnowledgeBase.workspace_uuid == execution_context.workspace_uuid)
+                .where(persistence_rag.KnowledgeBase.uuid == kb_uuid)
             )
             raise
 
@@ -531,7 +741,13 @@ class RAGManager:
 
         for knowledge_base in knowledge_bases:
             try:
-                await self.load_knowledge_base(knowledge_base)
+                binding = await self.ap.workspace_service.get_execution_binding(knowledge_base.workspace_uuid)
+                execution_context = ExecutionContext(
+                    instance_uuid=binding.instance_uuid,
+                    workspace_uuid=binding.workspace_uuid,
+                    placement_generation=binding.placement_generation,
+                )
+                await self.load_knowledge_base(execution_context, knowledge_base)
             except Exception as e:
                 self.ap.logger.error(
                     f'Error loading knowledge base {knowledge_base.uuid}: {e}\n{traceback.format_exc()}'
@@ -539,6 +755,7 @@ class RAGManager:
 
     async def load_knowledge_base(
         self,
+        context: RequestContext | ExecutionContext,
         knowledge_base_entity: persistence_rag.KnowledgeBase | sqlalchemy.Row | dict,
     ) -> RuntimeKnowledgeBase:
         if isinstance(knowledge_base_entity, sqlalchemy.Row):
@@ -551,23 +768,47 @@ class RAGManager:
             }
             knowledge_base_entity = persistence_rag.KnowledgeBase(**filtered_dict)
 
-        runtime_knowledge_base = RuntimeKnowledgeBase(ap=self.ap, knowledge_base_entity=knowledge_base_entity)
+        execution_context = await self._to_execution_context(context)
+        if knowledge_base_entity.workspace_uuid != execution_context.workspace_uuid:
+            raise WorkspaceNotFoundError('Knowledge base not found')
+        runtime_knowledge_base = RuntimeKnowledgeBase(
+            ap=self.ap,
+            knowledge_base_entity=knowledge_base_entity,
+            execution_context=execution_context,
+        )
 
         await runtime_knowledge_base.initialize()
 
-        self.knowledge_bases[runtime_knowledge_base.get_uuid()] = runtime_knowledge_base
+        self.knowledge_bases[(execution_context.workspace_uuid, runtime_knowledge_base.get_uuid())] = (
+            runtime_knowledge_base
+        )
 
         return runtime_knowledge_base
 
-    async def get_knowledge_base_by_uuid(self, kb_uuid: str) -> KnowledgeBaseInterface | None:
-        return self.knowledge_bases.get(kb_uuid)
+    async def get_knowledge_base_by_uuid(
+        self,
+        context: RequestContext | ExecutionContext,
+        kb_uuid: str,
+    ) -> RuntimeKnowledgeBase | None:
+        execution_context = await self._to_execution_context(context)
+        return self.knowledge_bases.get((execution_context.workspace_uuid, kb_uuid))
 
-    async def remove_knowledge_base_from_runtime(self, kb_uuid: str):
-        self.knowledge_bases.pop(kb_uuid, None)
+    async def remove_knowledge_base_from_runtime(
+        self,
+        context: RequestContext | ExecutionContext,
+        kb_uuid: str,
+    ) -> None:
+        execution_context = await self._to_execution_context(context)
+        self.knowledge_bases.pop((execution_context.workspace_uuid, kb_uuid), None)
 
-    async def delete_knowledge_base(self, kb_uuid: str):
-        kb = self.knowledge_bases.pop(kb_uuid, None)
+    async def delete_knowledge_base(
+        self,
+        context: RequestContext | ExecutionContext,
+        kb_uuid: str,
+    ) -> None:
+        execution_context = await self._to_execution_context(context)
+        kb = self.knowledge_bases.pop((execution_context.workspace_uuid, kb_uuid), None)
         if kb is not None:
-            await kb.dispose()
+            await kb.dispose(execution_context)
         else:
             self.ap.logger.warning(f'Knowledge base {kb_uuid} not found in runtime, skipping plugin notification')

src/langbot/pkg/rag/service/runtime.py

@@ -5,6 +5,13 @@ import re
 from typing import TYPE_CHECKING, Any
 from urllib.parse import unquote
 
+import sqlalchemy
+
+from langbot.pkg.api.http.authz import WorkspaceRequiredError
+from langbot.pkg.api.http.context import ExecutionContext
+from langbot.pkg.entity.persistence import rag as persistence_rag
+from langbot.pkg.workspace.errors import WorkspaceNotFoundError
+
 if TYPE_CHECKING:
     from langbot.pkg.core import app
 
@@ -19,8 +26,54 @@ class RAGRuntimeService:
     def __init__(self, ap: app.Application):
         self.ap = ap
 
+    async def _validate_execution_context(self, execution_context: ExecutionContext) -> None:
+        if not isinstance(execution_context, ExecutionContext):
+            raise WorkspaceRequiredError('ExecutionContext is required for RAG runtime access')
+        if (
+            not execution_context.instance_uuid.strip()
+            or not execution_context.workspace_uuid.strip()
+            or execution_context.placement_generation <= 0
+        ):
+            raise WorkspaceRequiredError('A complete active ExecutionContext is required')
+        workspace_service = getattr(self.ap, 'workspace_service', None)
+        if workspace_service is None:
+            raise WorkspaceRequiredError('Workspace execution service is unavailable')
+        binding = await workspace_service.get_execution_binding(
+            execution_context.workspace_uuid,
+            expected_generation=execution_context.placement_generation,
+        )
+        if binding.instance_uuid != execution_context.instance_uuid:
+            raise WorkspaceRequiredError('ExecutionContext belongs to another LangBot instance')
+
+    async def _resolve_knowledge_base_uuid(
+        self,
+        execution_context: ExecutionContext,
+        collection_id: str,
+    ) -> str:
+        """Resolve a plugin logical handle to a Workspace-owned KB UUID."""
+
+        await self._validate_execution_context(execution_context)
+        if not isinstance(collection_id, str) or not collection_id.strip():
+            raise WorkspaceNotFoundError('Knowledge base not found')
+        result = await self.ap.persistence_mgr.execute_async(
+            sqlalchemy.select(persistence_rag.KnowledgeBase.uuid)
+            .where(persistence_rag.KnowledgeBase.workspace_uuid == execution_context.workspace_uuid)
+            .where(
+                sqlalchemy.or_(
+                    persistence_rag.KnowledgeBase.uuid == collection_id,
+                    persistence_rag.KnowledgeBase.collection_id == collection_id,
+                )
+            )
+            .limit(1)
+        )
+        kb_uuid = result.scalar_one_or_none()
+        if kb_uuid is None:
+            raise WorkspaceNotFoundError('Knowledge base not found')
+        return kb_uuid
+
     async def vector_upsert(
         self,
+        execution_context: ExecutionContext,
         collection_id: str,
         vectors: list[list[float]],
         ids: list[str],
@@ -28,9 +81,17 @@ class RAGRuntimeService:
         documents: list[str] | None = None,
     ) -> None:
         """Handle VECTOR_UPSERT action."""
+        knowledge_base_uuid = await self._resolve_knowledge_base_uuid(execution_context, collection_id)
+        if len(vectors) != len(ids):
+            raise ValueError('vectors and ids must have the same length')
+        if metadata is not None and len(metadata) != len(vectors):
+            raise ValueError('metadata must have the same length as vectors')
+        if documents is not None and len(documents) != len(vectors):
+            raise ValueError('documents must have the same length as vectors')
         metadatas = metadata if metadata else [{} for _ in vectors]
         await self.ap.vector_db_mgr.upsert(
-            collection_name=collection_id,
+            execution_context=execution_context,
+            knowledge_base_uuid=knowledge_base_uuid,
             vectors=vectors,
             ids=ids,
             metadata=metadatas,
@@ -39,6 +100,7 @@ class RAGRuntimeService:
 
     async def vector_search(
         self,
+        execution_context: ExecutionContext,
         collection_id: str,
         query_vector: list[float],
         top_k: int,
@@ -48,8 +110,10 @@ class RAGRuntimeService:
         vector_weight: float | None = None,
     ) -> list[dict[str, Any]]:
         """Handle VECTOR_SEARCH action."""
+        knowledge_base_uuid = await self._resolve_knowledge_base_uuid(execution_context, collection_id)
         return await self.ap.vector_db_mgr.search(
-            collection_name=collection_id,
+            execution_context=execution_context,
+            knowledge_base_uuid=knowledge_base_uuid,
             query_vector=query_vector,
             limit=top_k,
             filter=filters,
@@ -59,7 +123,11 @@ class RAGRuntimeService:
         )
 
     async def vector_delete(
-        self, collection_id: str, file_ids: list[str] | None = None, filters: dict[str, Any] | None = None
+        self,
+        execution_context: ExecutionContext,
+        collection_id: str,
+        file_ids: list[str] | None = None,
+        filters: dict[str, Any] | None = None,
     ) -> int:
         """Handle VECTOR_DELETE action.
 
@@ -73,16 +141,26 @@ class RAGRuntimeService:
                 in their metadata.
             filters: Filter-based deletion (not yet supported, will raise).
         """
+        knowledge_base_uuid = await self._resolve_knowledge_base_uuid(execution_context, collection_id)
         count = 0
         if file_ids:
-            await self.ap.vector_db_mgr.delete_by_file_id(collection_name=collection_id, file_ids=file_ids)
+            await self.ap.vector_db_mgr.delete_by_file_id(
+                execution_context=execution_context,
+                knowledge_base_uuid=knowledge_base_uuid,
+                file_ids=file_ids,
+            )
             count = len(file_ids)
         elif filters:
-            count = await self.ap.vector_db_mgr.delete_by_filter(collection_name=collection_id, filter=filters)
+            count = await self.ap.vector_db_mgr.delete_by_filter(
+                execution_context=execution_context,
+                knowledge_base_uuid=knowledge_base_uuid,
+                filter=filters,
+            )
         return count
 
     async def vector_list(
         self,
+        execution_context: ExecutionContext,
         collection_id: str,
         filters: dict[str, Any] | None = None,
         limit: int = 20,
@@ -99,14 +177,20 @@ class RAGRuntimeService:
         Returns:
             Tuple of (items, total).
         """
+        knowledge_base_uuid = await self._resolve_knowledge_base_uuid(execution_context, collection_id)
         return await self.ap.vector_db_mgr.list_by_filter(
-            collection_name=collection_id,
+            execution_context=execution_context,
+            knowledge_base_uuid=knowledge_base_uuid,
             filter=filters,
             limit=limit,
             offset=offset,
         )
 
-    async def get_file_stream(self, storage_path: str) -> bytes:
+    async def get_file_stream(
+        self,
+        execution_context: ExecutionContext,
+        storage_path: str,
+    ) -> bytes:
         """Handle GET_KNOWLEDEGE_FILE_STREAM action.
 
         Uses the storage manager abstraction to load file content,
@@ -125,5 +209,18 @@ class RAGRuntimeService:
             or re.match(r'^[A-Za-z]:/', normalized)
         ):
             raise ValueError('Invalid storage path')
-        content_bytes = await self.ap.storage_mgr.storage_provider.load(normalized)
+        await self._validate_execution_context(execution_context)
+        result = await self.ap.persistence_mgr.execute_async(
+            sqlalchemy.select(persistence_rag.File.uuid)
+            .where(persistence_rag.File.workspace_uuid == execution_context.workspace_uuid)
+            .where(persistence_rag.File.file_name == normalized)
+            .limit(1)
+        )
+        if result.first() is None:
+            raise WorkspaceNotFoundError('Knowledge file not found')
+        content_bytes = await self.ap.storage_mgr.load_scoped_object_key(
+            execution_context,
+            normalized,
+            expected_owner_type='upload_document',
+        )
         return content_bytes if content_bytes else b''