LangBot/tests/unit_tests/rag/test_tenant_isolation.py

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 AsyncSession, 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.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


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

    @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 {
            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 == []


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)