feat(tenancy): harden shared cloud runtime boundaries

tests/unit_tests/rag/test_file_storage.py

@@ -2,8 +2,10 @@
 
 from __future__ import annotations
 
+import contextvars
 import io
 import zipfile
+from contextlib import asynccontextmanager
 from types import SimpleNamespace
 from unittest.mock import AsyncMock, Mock
 
@@ -54,7 +56,7 @@ def _make_app() -> Mock:
     storage_mgr.storage_provider.delete = AsyncMock()
     app.storage_mgr = storage_mgr
     app.persistence_mgr = Mock()
-    app.persistence_mgr.execute_async = AsyncMock()
+    app.persistence_mgr.execute_async = AsyncMock(return_value=SimpleNamespace(rowcount=1))
     app.plugin_connector = Mock()
     app.plugin_connector.require_workspace_context = AsyncMock(side_effect=lambda context: context)
     app.workspace_service = SimpleNamespace(
@@ -202,6 +204,36 @@ class TestStoreZipFile:
 
 
 class TestStoreFileTask:
+    @pytest.mark.asyncio
+    async def test_store_file_task_opens_uow_before_first_database_helper(self):
+        kb = _make_kb()
+        active_workspace = contextvars.ContextVar('rag_task_workspace', default=None)
+        observed = []
+
+        @asynccontextmanager
+        async def tenant_uow(workspace_uuid):
+            token = active_workspace.set(workspace_uuid)
+            try:
+                yield
+            finally:
+                active_workspace.reset(token)
+
+        kb.ap.persistence_mgr.mode = SimpleNamespace(value='cloud_runtime')
+        kb.ap.persistence_mgr.tenant_uow = tenant_uow
+
+        async def assert_execution_context(_context):
+            observed.append(active_workspace.get())
+
+        kb._assert_execution_context = AsyncMock(side_effect=assert_execution_context)
+        kb._set_file_status = AsyncMock(side_effect=[True, True])
+        kb._ingest_document = AsyncMock(return_value={'status': 'completed'})
+        object_key = _upload_key('scoped.pdf')
+        file_obj = SimpleNamespace(uuid='file-uuid', file_name=object_key, extension='pdf')
+
+        await kb._store_file_task(CONTEXT, file_obj, Mock())
+
+        assert observed[0] == WORKSPACE_A
+
     @pytest.mark.asyncio
     async def test_store_file_task_marks_completed_and_cleans_storage(self):
         kb = _make_kb()

tests/unit_tests/rag/test_tenant_isolation.py

@@ -1,12 +1,13 @@
 from __future__ import annotations
 
 import datetime
+from contextlib import asynccontextmanager
 from types import SimpleNamespace
 from unittest.mock import AsyncMock, Mock
 
 import pytest
 import sqlalchemy
-from sqlalchemy.ext.asyncio import create_async_engine
+from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
 
 from langbot.pkg.api.http.authz import WorkspaceRequiredError
 from langbot.pkg.api.http.context import ExecutionContext
@@ -17,6 +18,7 @@ from langbot.pkg.entity.persistence.workspace import Workspace, WorkspaceExecuti
 from langbot.pkg.rag.knowledge.kbmgr import RAGManager
 from langbot.pkg.rag.service.runtime import RAGRuntimeService
 from langbot.pkg.vector.mgr import VectorDBManager
+from langbot.pkg.vector.vdbs.pgvector_db import PgVectorDatabase
 from langbot.pkg.workspace.errors import WorkspaceNotFoundError
 from langbot.pkg.workspace.policy import CloudWorkspacePolicy, SingleWorkspacePolicy
 from langbot.pkg.workspace.service import WorkspaceService
@@ -42,6 +44,13 @@ class _PersistenceManager:
             await connection.commit()
             return result
 
+    @asynccontextmanager
+    async def tenant_uow(self, _workspace_uuid):
+        # This lightweight fixture does not emulate PostgreSQL RLS; production
+        # persistence tests cover the transaction-bound unit of work itself.
+        async with AsyncSession(self.engine, expire_on_commit=False) as session, session.begin():
+            yield SimpleNamespace(session=session)
+
     @staticmethod
     def serialize_model(model, row, masked_columns=()):
         return {
@@ -348,3 +357,34 @@ async def test_stale_generation_is_rejected_before_vector_access(tenant_rag):
     with pytest.raises(Exception, match='generation'):
         await manager.upsert(stale, 'kb-a', [[0.1]], ['a'])
     assert database.collections == []
+
+
+async def test_pgvector_first_write_binds_dimension_and_later_mismatch_fails(tenant_rag):
+    app, engine = tenant_rag
+    manager = VectorDBManager(app)
+    pgvector = object.__new__(PgVectorDatabase)
+    pgvector.allowed_dimensions = frozenset({1, 2})
+    pgvector.add_embeddings = AsyncMock()
+    pgvector.search = AsyncMock(return_value={'ids': [[]], 'distances': [[]], 'metadatas': [[]]})
+    manager.vector_db = pgvector
+    context = _context(WORKSPACE_A)
+
+    await manager.upsert(context, 'kb-a', [[0.1]], ['chunk-a'])
+    scope = pgvector.add_embeddings.await_args.kwargs['scope']
+    assert scope.workspace_uuid == WORKSPACE_A
+    assert scope.knowledge_base_uuid == 'kb-a'
+    assert scope.embedding_dimension == 1
+
+    async with engine.connect() as connection:
+        selected_dimension = await connection.scalar(
+            sqlalchemy.select(KnowledgeBase.embedding_dimension).where(
+                KnowledgeBase.workspace_uuid == WORKSPACE_A,
+                KnowledgeBase.uuid == 'kb-a',
+            )
+        )
+    assert selected_dimension == 1
+
+    with pytest.raises(ValueError, match='dimension is 1, not 2'):
+        await manager.upsert(context, 'kb-a', [[0.1, 0.2]], ['chunk-b'])
+    with pytest.raises(ValueError, match='not enabled'):
+        await manager.search(context, 'kb-a', [0.1, 0.2, 0.3], 3)