feat(cloud): harden multi-tenant runtime resources

tests/unit_tests/rag/test_file_storage.py

@@ -12,6 +12,7 @@ from unittest.mock import AsyncMock, Mock
 import pytest
 
 from langbot.pkg.api.http.context import ExecutionContext
+from langbot.pkg.core.taskmgr import TaskCapacityError
 from langbot.pkg.rag.knowledge.kbmgr import RuntimeKnowledgeBase
 from langbot.pkg.storage.mgr import StorageMgr
 from langbot.pkg.workspace.errors import WorkspaceNotFoundError
@@ -34,9 +35,9 @@ def _upload_key(logical_key: str, *, context: ExecutionContext = CONTEXT) -> str
     )
 
 
-def _make_zip_bytes(entries: dict[str, bytes]) -> bytes:
+def _make_zip_bytes(entries: dict[str, bytes], *, compression: int = zipfile.ZIP_STORED) -> bytes:
     buffer = io.BytesIO()
-    with zipfile.ZipFile(buffer, 'w') as zf:
+    with zipfile.ZipFile(buffer, 'w', compression=compression) as zf:
         for name, content in entries.items():
             zf.writestr(name, content)
         zf.mkdir('emptydir')
@@ -151,6 +152,24 @@ class TestStoreFile:
 
         kb.ap.storage_mgr.storage_provider.exists.assert_not_awaited()
 
+    @pytest.mark.asyncio
+    async def test_store_file_rolls_back_pending_record_when_task_capacity_is_exhausted(self):
+        kb = _make_kb()
+        object_key = _upload_key('queued.pdf')
+
+        def reject(coro, **_kwargs):
+            coro.close()
+            raise TaskCapacityError('capacity')
+
+        kb.ap.task_mgr.create_user_task.side_effect = reject
+
+        with pytest.raises(TaskCapacityError, match='capacity'):
+            await kb.store_file(CONTEXT, object_key)
+
+        statements = [str(call.args[0]) for call in kb.ap.persistence_mgr.execute_async.await_args_list]
+        assert any(statement.startswith('INSERT') for statement in statements)
+        assert any(statement.startswith('DELETE') for statement in statements)
+
 
 class TestStoreZipFile:
     @pytest.mark.asyncio
@@ -202,6 +221,39 @@ class TestStoreZipFile:
         kb.store_file.assert_not_awaited()
         kb.ap.storage_mgr.storage_provider.delete.assert_awaited_once_with(_upload_key('archive.zip'))
 
+    @pytest.mark.asyncio
+    async def test_store_zip_file_rejects_too_many_documents_before_extracting(self):
+        kb = _make_kb()
+        kb.ap.storage_mgr.storage_provider.load = AsyncMock(
+            return_value=_make_zip_bytes({f'doc-{index}.txt': b'text' for index in range(9)})
+        )
+        kb.store_file = AsyncMock()
+
+        with pytest.raises(ValueError, match='too many supported documents'):
+            await kb._store_zip_file(CONTEXT, _upload_key('archive.zip'))
+
+        kb.store_file.assert_not_awaited()
+        kb.ap.storage_mgr.storage_provider.save.assert_not_awaited()
+        kb.ap.storage_mgr.storage_provider.delete.assert_awaited_once_with(_upload_key('archive.zip'))
+
+    @pytest.mark.asyncio
+    async def test_store_zip_file_rejects_extreme_compression_ratio_before_extracting(self):
+        kb = _make_kb()
+        kb.ap.storage_mgr.storage_provider.load = AsyncMock(
+            return_value=_make_zip_bytes(
+                {'bomb.txt': b'A' * (1024 * 1024)},
+                compression=zipfile.ZIP_DEFLATED,
+            )
+        )
+        kb.store_file = AsyncMock()
+
+        with pytest.raises(ValueError, match='compression-ratio limit'):
+            await kb._store_zip_file(CONTEXT, _upload_key('archive.zip'))
+
+        kb.store_file.assert_not_awaited()
+        kb.ap.storage_mgr.storage_provider.save.assert_not_awaited()
+        kb.ap.storage_mgr.storage_provider.delete.assert_awaited_once_with(_upload_key('archive.zip'))
+
 
 class TestStoreFileTask:
     @pytest.mark.asyncio

tests/unit_tests/rag/test_kbmgr.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import dataclasses
 from types import SimpleNamespace
 from unittest.mock import AsyncMock, Mock
 
@@ -10,7 +11,8 @@ import pytest
 from langbot.pkg.api.http.context import ExecutionContext
 from langbot.pkg.entity.persistence.rag import KnowledgeBase
 from langbot.pkg.rag.knowledge.kbmgr import RAGManager, RuntimeKnowledgeBase
-from langbot.pkg.workspace.errors import WorkspaceNotFoundError
+from langbot.pkg.workspace.entities import WorkspaceExecutionBinding
+from langbot.pkg.workspace.errors import WorkspaceInvariantError, WorkspaceNotFoundError
 
 
 CONTEXT_A = ExecutionContext(
@@ -511,6 +513,37 @@ class TestRAGManagerLoadKnowledgeBasesFromDB:
             ('workspace-a', 'kb-2'),
         }
 
+    @pytest.mark.asyncio
+    async def test_cloud_startup_reuses_validated_binding(self):
+        class TenantUow:
+            async def __aenter__(self):
+                return self
+
+            async def __aexit__(self, *_args):
+                return False
+
+        app = _app()
+        binding = WorkspaceExecutionBinding(
+            instance_uuid='instance-a',
+            workspace_uuid='workspace-a',
+            placement_generation=5,
+            write_fenced=False,
+            state='active',
+        )
+        app.persistence_mgr.mode = SimpleNamespace(value='cloud_runtime')
+        app.persistence_mgr.tenant_uow = lambda _workspace_uuid: TenantUow()
+        app.persistence_mgr.execute_async.return_value = _Result([_entity()])
+        app.workspace_service.list_active_execution_bindings = AsyncMock(return_value=[binding])
+        app.workspace_service.get_execution_binding = AsyncMock(
+            side_effect=AssertionError('startup RAG loader repeated a validated binding lookup')
+        )
+        manager = RAGManager(app)
+
+        await manager.load_knowledge_bases_from_db()
+
+        assert set(manager.knowledge_bases) == {('workspace-a', 'kb-a')}
+        app.workspace_service.get_execution_binding.assert_not_awaited()
+
     @pytest.mark.asyncio
     async def test_handles_load_error_gracefully(self):
         app = _app()
@@ -643,6 +676,29 @@ class TestRAGManagerInit:
     def test_init_creates_empty_knowledge_bases_dict(self):
         assert RAGManager(_app()).knowledge_bases == {}
 
+    def test_generation_advance_prunes_superseded_runtime_knowledge_bases(self):
+        class NoGlobalItemsScan(dict):
+            def items(self):
+                raise AssertionError('generation advance scanned every knowledge runtime')
+
+        app = _app()
+        manager = RAGManager(app)
+        manager._cache_runtime(
+            RuntimeKnowledgeBase(
+                app,
+                _entity(),
+                CONTEXT_A,
+            )
+        )
+        manager.knowledge_bases = NoGlobalItemsScan(manager.knowledge_bases)
+
+        next_context = dataclasses.replace(CONTEXT_A, placement_generation=6)
+        manager._observe_execution_context(next_context)
+
+        assert manager.knowledge_bases == {}
+        with pytest.raises(WorkspaceInvariantError, match='rolled back'):
+            manager._observe_execution_context(CONTEXT_A)
+
 
 class TestRAGManagerGetKnowledgeBase:
     @pytest.mark.asyncio