LangBot/tests/unit_tests/rag/test_tenant_isolation.py

from __future__ import annotations

import datetime
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock

import pytest
import sqlalchemy
from sqlalchemy.ext.asyncio import create_async_engine

from langbot.pkg.api.http.authz import WorkspaceRequiredError
from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.api.http.service.knowledge import KnowledgeService
from langbot.pkg.entity.persistence.base import Base
from langbot.pkg.entity.persistence.rag import File, KnowledgeBase
from langbot.pkg.entity.persistence.workspace import Workspace, WorkspaceExecutionState
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.workspace.errors import WorkspaceNotFoundError
from langbot.pkg.workspace.policy import CloudWorkspacePolicy, SingleWorkspacePolicy
from langbot.pkg.workspace.service import WorkspaceService


pytestmark = pytest.mark.asyncio

INSTANCE_UUID = 'instance-rag-isolation'
WORKSPACE_A = '00000000-0000-0000-0000-00000000000a'
WORKSPACE_B = '00000000-0000-0000-0000-00000000000b'


class _PersistenceManager:
    def __init__(self, engine):
        self.engine = engine

    def get_db_engine(self):
        return self.engine

    async def execute_async(self, *args, **kwargs):
        async with self.engine.connect() as connection:
            result = await connection.execute(*args, **kwargs)
            await connection.commit()
            return result

    @staticmethod
    def serialize_model(model, row, masked_columns=()):
        return {
            column.name: (
                getattr(row, column.name).isoformat()
                if isinstance(getattr(row, column.name), datetime.datetime)
                else getattr(row, column.name)
            )
            for column in model.__table__.columns
            if column.name not in masked_columns
        }


class _RecordingVectorDatabase:
    def __init__(self):
        self.collections: list[str] = []
        self.metadatas: list[list[dict]] = []
        self.calls: list[tuple[str, str]] = []

    async def add_embeddings(self, *, collection, ids, embeddings_list, metadatas, documents):
        self.collections.append(collection)
        self.metadatas.append(metadatas)
        self.calls.append(('upsert', collection))

    async def search(self, *, collection, **_kwargs):
        self.calls.append(('search', collection))
        return {'ids': [[]], 'distances': [[]], 'metadatas': [[]]}

    async def delete_by_file_id(self, collection, _file_id):
        self.calls.append(('delete_by_file_id', collection))

    async def delete_collection(self, collection):
        self.calls.append(('delete_collection', collection))

    async def delete_by_filter(self, collection, _filter):
        self.calls.append(('delete_by_filter', collection))
        return 1

    async def list_by_filter(self, collection, _filter, _limit, _offset):
        self.calls.append(('list_by_filter', collection))
        return [], 0


@pytest.fixture
async def tenant_rag(tmp_path):
    engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "rag-tenant.db"}')
    async with engine.begin() as connection:
        await connection.run_sync(Base.metadata.create_all)
        await connection.execute(
            sqlalchemy.insert(Workspace),
            [
                {
                    'uuid': WORKSPACE_A,
                    'instance_uuid': INSTANCE_UUID,
                    'name': 'Workspace A',
                    'slug': 'workspace-a',
                    'source': 'local',
                },
                {
                    'uuid': WORKSPACE_B,
                    'instance_uuid': INSTANCE_UUID,
                    'name': 'Workspace B',
                    'slug': 'workspace-b',
                    'source': 'cloud_projection',
                },
            ],
        )
        await connection.execute(
            sqlalchemy.insert(WorkspaceExecutionState),
            [
                {
                    'workspace_uuid': WORKSPACE_A,
                    'instance_uuid': INSTANCE_UUID,
                    'active_generation': 3,
                    'state': 'active',
                    'source': 'local',
                    'write_fenced': False,
                },
                {
                    'workspace_uuid': WORKSPACE_B,
                    'instance_uuid': INSTANCE_UUID,
                    'active_generation': 3,
                    'state': 'active',
                    'source': 'cloud',
                    'write_fenced': False,
                },
            ],
        )
        await connection.execute(
            sqlalchemy.insert(KnowledgeBase),
            [
                {
                    'uuid': 'kb-a',
                    'workspace_uuid': WORKSPACE_A,
                    'name': 'Same Knowledge Base',
                    'description': 'A',
                    'knowledge_engine_plugin_id': 'author/engine',
                    'collection_id': 'kb-a',
                    'creation_settings': {},
                    'retrieval_settings': {},
                },
                {
                    'uuid': 'kb-b',
                    'workspace_uuid': WORKSPACE_B,
                    'name': 'Same Knowledge Base',
                    'description': 'B',
                    'knowledge_engine_plugin_id': 'author/engine',
                    'collection_id': 'kb-b',
                    'creation_settings': {},
                    'retrieval_settings': {},
                },
            ],
        )
        await connection.execute(
            sqlalchemy.insert(File),
            [
                {
                    'uuid': 'file-a',
                    'workspace_uuid': WORKSPACE_A,
                    'kb_id': 'kb-a',
                    'file_name': 'a.pdf',
                    'extension': 'pdf',
                },
                {
                    'uuid': 'file-b',
                    'workspace_uuid': WORKSPACE_B,
                    'kb_id': 'kb-b',
                    'file_name': 'b.pdf',
                    'extension': 'pdf',
                },
            ],
        )

    app = SimpleNamespace()
    app.persistence_mgr = _PersistenceManager(engine)
    app.logger = Mock()
    app.workspace_policy = SingleWorkspacePolicy()
    app.plugin_connector = SimpleNamespace(
        is_enable_plugin=False,
        rag_on_kb_create=AsyncMock(),
        rag_on_kb_delete=AsyncMock(),
    )
    app.workspace_service = WorkspaceService(app, instance_uuid=INSTANCE_UUID)
    app.rag_mgr = RAGManager(app)
    await app.rag_mgr.initialize()
    app.knowledge_service = KnowledgeService(app)

    yield app, engine
    await engine.dispose()


def _context(workspace_uuid: str) -> ExecutionContext:
    return ExecutionContext(
        instance_uuid=INSTANCE_UUID,
        workspace_uuid=workspace_uuid,
        placement_generation=3,
    )


async def test_context_is_mandatory_and_same_names_are_isolated(tenant_rag):
    app, _engine = tenant_rag

    with pytest.raises(WorkspaceRequiredError):
        await app.knowledge_service.get_knowledge_bases(None)

    bases_a = await app.knowledge_service.get_knowledge_bases(_context(WORKSPACE_A))
    bases_b = await app.knowledge_service.get_knowledge_bases(_context(WORKSPACE_B))
    assert [(item['uuid'], item['name']) for item in bases_a] == [('kb-a', 'Same Knowledge Base')]
    assert [(item['uuid'], item['name']) for item in bases_b] == [('kb-b', 'Same Knowledge Base')]


async def test_cross_workspace_uuid_and_file_guessing_return_not_found(tenant_rag):
    app, engine = tenant_rag
    context_a = _context(WORKSPACE_A)

    assert await app.knowledge_service.get_knowledge_base(context_a, 'kb-b') is None
    with pytest.raises(WorkspaceNotFoundError):
        await app.knowledge_service.update_knowledge_base(context_a, 'kb-b', {'name': 'stolen'})
    with pytest.raises(WorkspaceNotFoundError):
        await app.knowledge_service.delete_knowledge_base(context_a, 'kb-b')
    with pytest.raises(WorkspaceNotFoundError):
        await app.knowledge_service.get_files_by_knowledge_base(context_a, 'kb-b')
    with pytest.raises(WorkspaceNotFoundError):
        await app.knowledge_service.delete_file(context_a, 'kb-b', 'file-b')

    async with engine.connect() as connection:
        assert (
            await connection.scalar(sqlalchemy.select(KnowledgeBase.name).where(KnowledgeBase.uuid == 'kb-b'))
            == 'Same Knowledge Base'
        )
        assert await connection.scalar(sqlalchemy.select(File.uuid).where(File.uuid == 'file-b')) == 'file-b'


async def test_runtime_rejects_cross_workspace_collection_reference(tenant_rag):
    app, _engine = tenant_rag
    app.vector_db_mgr = SimpleNamespace(upsert=AsyncMock())
    service = RAGRuntimeService(app)

    with pytest.raises(WorkspaceNotFoundError):
        await service.vector_upsert(
            _context(WORKSPACE_A),
            'kb-b',
            vectors=[[0.1, 0.2]],
            ids=['chunk-1'],
        )
    app.vector_db_mgr.upsert.assert_not_awaited()


async def test_physical_vector_handles_do_not_collide_across_workspaces(tenant_rag):
    app, _engine = tenant_rag
    database = _RecordingVectorDatabase()
    manager = VectorDBManager(app)
    manager.vector_db = database

    context_a = _context(WORKSPACE_A)
    context_b = _context(WORKSPACE_B)
    assert manager.physical_collection_name(context_a, 'same-kb-id') != manager.physical_collection_name(
        context_b,
        'same-kb-id',
    )

    await manager.upsert(context_a, 'kb-a', [[0.1]], ['a'], metadata=[{'source': 'client'}])
    await manager.upsert(context_b, 'kb-b', [[0.2]], ['b'], metadata=[{'source': 'client'}])
    assert len(set(database.collections)) == 2
    assert database.metadatas[0][0]['_langbot_workspace_uuid'] == WORKSPACE_A
    assert database.metadatas[1][0]['_langbot_workspace_uuid'] == WORKSPACE_B


async def test_migrated_local_kb_keeps_legacy_collection_for_every_vector_operation(tenant_rag):
    app, engine = tenant_rag
    legacy_collection = 'legacy-collection-kb-a'
    async with engine.begin() as connection:
        await connection.execute(
            sqlalchemy.update(KnowledgeBase)
            .where(KnowledgeBase.uuid == 'kb-a')
            .values(
                collection_id=legacy_collection,
                legacy_vector_collection=True,
            )
        )

    database = _RecordingVectorDatabase()
    manager = VectorDBManager(app)
    manager.vector_db = database
    context = _context(WORKSPACE_A)

    await manager.upsert(context, 'kb-a', [[0.1]], ['chunk-a'])
    await manager.search(context, 'kb-a', [0.1], 3)
    await manager.delete_by_file_id(context, 'kb-a', ['file-a'])
    assert await manager.delete_by_filter(context, 'kb-a', {'file_id': 'file-a'}) == 1
    assert await manager.list_by_filter(context, 'kb-a', {'file_id': 'file-a'}) == ([], 0)
    await manager.delete_collection(context, 'kb-a')

    assert database.calls == [
        ('upsert', legacy_collection),
        ('search', legacy_collection),
        ('delete_by_file_id', legacy_collection),
        ('delete_by_filter', legacy_collection),
        ('list_by_filter', legacy_collection),
        ('delete_collection', legacy_collection),
    ]


@pytest.mark.parametrize('deny_by', ['projected_workspace', 'multi_workspace_policy'])
async def test_legacy_marker_is_ignored_outside_single_local_workspace(tenant_rag, deny_by):
    app, engine = tenant_rag
    workspace_uuid = WORKSPACE_B if deny_by == 'projected_workspace' else WORKSPACE_A
    kb_uuid = 'kb-b' if deny_by == 'projected_workspace' else 'kb-a'
    legacy_collection = f'legacy-{deny_by}'
    async with engine.begin() as connection:
        await connection.execute(
            sqlalchemy.update(KnowledgeBase)
            .where(KnowledgeBase.uuid == kb_uuid)
            .values(
                collection_id=legacy_collection,
                legacy_vector_collection=True,
            )
        )
    if deny_by == 'multi_workspace_policy':
        app.workspace_policy = CloudWorkspacePolicy()

    database = _RecordingVectorDatabase()
    manager = VectorDBManager(app)
    manager.vector_db = database
    context = _context(workspace_uuid)
    await manager.search(context, kb_uuid, [0.1], 3)

    assert database.calls == [('search', manager.physical_collection_name(context, kb_uuid))]
    assert database.calls[0][1] != legacy_collection
    app.logger.warning.assert_called_once()


async def test_stale_generation_is_rejected_before_vector_access(tenant_rag):
    app, _engine = tenant_rag
    database = _RecordingVectorDatabase()
    manager = VectorDBManager(app)
    manager.vector_db = database
    stale = ExecutionContext(
        instance_uuid=INSTANCE_UUID,
        workspace_uuid=WORKSPACE_A,
        placement_generation=2,
    )

    with pytest.raises(Exception, match='generation'):
        await manager.upsert(stale, 'kb-a', [[0.1]], ['a'])
    assert database.collections == []