LangBot/plugins/KnowledgeEngine/certification-tests/test_shared_state.py

"""Real SDK objects; storage/HTTP are explicit local fixtures, not live services."""
import asyncio
import importlib.util
import sys
from pathlib import Path

import httpx
import pytest

from langbot_plugin.api.entities.builtin.rag import (
    DocumentStatus,
    FileMetadata,
    FileObject,
    IngestionContext,
)

ROOT = Path(__file__).resolve().parents[1]
CONNECTORS = ['DifyDatasetsConnector', 'FastGPTConnector', 'RAGFlowConnector']


def load_plugin(name):
    # Runtime uses one process per installation. Clear package names for this test loader.
    for key in list(sys.modules):
        if key == 'components' or key.startswith('components.'):
            del sys.modules[key]
    sys.path.insert(0, str(ROOT / name))
    try:
        spec = importlib.util.spec_from_file_location('candidate_main', ROOT / name / 'main.py')
        mod = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(mod)
        from importlib import import_module
        engine = import_module('components.knowledge_engine.' + ('langrag' if name == 'LangRAG' else 'engine'))
        return getattr(mod, name), getattr(engine, name)
    finally:
        sys.path.pop(0)


class StorageFixture:
    """Explicit fixture for Host's installation-bound storage, not production storage."""
    def __init__(self):
        self.data = {}
        self.fail = False
        # Optional ``key -> exception`` fault injection for one durable write.
        self.write_fault = None
        self.started = asyncio.Event()
        self.release = asyncio.Event()
        self.release.set()
        # Host-read instrumentation: how many key listings were served, and
        # one-shot gates that suspend the next read of one key so a test can land
        # a concurrent write inside that reader's window.
        self.keys_reads = 0
        self.gates = {}
        self.gate_entered = asyncio.Event()
        self.gate_release = asyncio.Event()

    def gate_read(self, key):
        """Suspend the next ``get_plugin_storage(key)`` until ``gate_release``.

        The gated reader captures the value it was answered with before it
        blocks, so the caller can change that record while the read is in flight
        exactly like a response computed before a concurrent write. Only the
        first read of ``key`` is gated; later readers pass through.
        """
        self.gate_entered = asyncio.Event()
        self.gate_release = asyncio.Event()
        self.gates[key] = self.gate_entered

    async def get_plugin_storage_keys(self):
        if self.fail:
            raise RuntimeError('fixture storage offline')
        self.keys_reads += 1
        return list(self.data)

    async def get_plugin_storage(self, key):
        if self.fail:
            raise RuntimeError('fixture storage offline')
        value = self.data[key]
        entered = self.gates.pop(key, None)
        if entered is not None:
            entered.set()
            await self.gate_release.wait()
        return value

    async def set_plugin_storage(self, key, value):
        # A detached Host commit survives cancellation of the caller, like real RPC.
        async def commit():
            self.started.set()
            await self.release.wait()
            if self.write_fault is not None:
                fault = self.write_fault(key)
                if fault is not None:
                    raise fault
            if self.fail:
                raise RuntimeError('fixture ambiguous commit')
            self.data[key] = value
        await asyncio.shield(asyncio.create_task(commit()))


def bind(plugin_cls, storage):
    plugin = plugin_cls()
    for name in ['get_plugin_storage_keys', 'get_plugin_storage', 'set_plugin_storage']:
        setattr(plugin, name, getattr(storage, name))
    return plugin


def configuration(marker):
    return {'api_base_url': 'https://fixture.invalid', 'api_key': marker,
            'dify_apikey': marker, 'dataset_id': marker, 'dataset_ids': marker}


@pytest.mark.asyncio
@pytest.mark.parametrize('name', CONNECTORS)
async def test_connector_restart_isolation_and_delete(name, monkeypatch):
    pc, ec = load_plugin(name)
    calls = []
    async def service(request):
        calls.append(request)
        if request.method == 'GET':
            return httpx.Response(200, json={'code': 0, 'data': []})
        return httpx.Response(200, json={'code': 200 if name == 'FastGPTConnector' else 0})
    client = httpx.AsyncClient
    monkeypatch.setattr(httpx, 'AsyncClient', lambda **kw: client(transport=httpx.MockTransport(service), **kw))
    a, b = StorageFixture(), StorageFixture()
    ea, eb = ec(), ec()
    ea.plugin, eb.plugin = bind(pc, a), bind(pc, b)
    config = configuration('A')
    await ea.on_knowledge_base_create('same-kb', config)
    config['api_key'] = 'MUTATED'
    await eb.on_knowledge_base_create('same-kb', configuration('B'))
    restarted = ec()
    restarted.plugin = bind(pc, a)
    # No acknowledged ingestion mapping exists; guessing a provider ID is unsafe.
    assert await restarted.delete_document('same-kb', 'doc') is False
    assert await eb.delete_document('same-kb', 'doc') is False
    assert not any(request.method == 'DELETE' for request in calls)
    # Mappings record the upstream target they were created on; a delete only
    # replays against that same target.
    await restarted._save_document('same-kb', 'doc', {'upstream_id': 'upstream-A', 'dataset_id': 'A', 'status': 'created', 'api_base_url': 'https://fixture.invalid'})
    await eb._save_document('same-kb', 'doc', {'upstream_id': 'upstream-B', 'dataset_id': 'B', 'status': 'created', 'api_base_url': 'https://fixture.invalid'})
    assert await restarted.delete_document('same-kb', 'doc') is True
    assert calls[-1].headers['authorization'] == 'Bearer A'
    assert await eb.delete_document('same-kb', 'doc') is True
    assert calls[-1].headers['authorization'] == 'Bearer B'
    await restarted.on_knowledge_base_delete('same-kb')
    assert await restarted.delete_document('same-kb', 'doc') is False
    assert await eb.delete_document('same-kb', 'doc') is False


@pytest.mark.asyncio
@pytest.mark.parametrize('name', CONNECTORS)
async def test_cancelled_config_commit_cannot_resurrect_deleted_kb(name):
    pc, ec = load_plugin(name)
    store = StorageFixture()
    engine = ec()
    engine.plugin = bind(pc, store)
    store.release.clear()
    # Empty credentials avoid the optional RAGFlow validation request.
    write = asyncio.create_task(engine.on_knowledge_base_create('kb', {'marker': 'A'}))
    try:
        await asyncio.wait_for(store.started.wait(), .3)
        write.cancel()
        await asyncio.sleep(0)
        write.cancel()
        delete = asyncio.create_task(engine.on_knowledge_base_delete('kb'))
        await asyncio.sleep(.01)
        # The dispatched Host commit already holds the lock and is settled before
        # release, so the delete is ordered behind it rather than fenced: a
        # settled commit is not an ambiguous mutation. Ordinary cancellation of
        # the caller therefore does not fence the knowledge base.
        assert not delete.done()
        store.release.set()
        with pytest.raises(asyncio.CancelledError):
            await write
        await delete
    finally:
        store.release.set()
        await asyncio.gather(write, return_exceptions=True)
    # The tombstone is the last acknowledged commit: the cancelled create cannot
    # resurrect the deleted knowledge base, and the survivor is not fenced.
    restarted = ec()
    restarted.plugin = bind(pc, store)
    assert await restarted._load_config('kb') is None
    assert await restarted.delete_document('kb', 'doc') is False
    await restarted.on_knowledge_base_create('kb', {'marker': 'B'})
    assert (await restarted._load_config('kb'))['marker'] == 'B'


@pytest.mark.asyncio
async def test_langrag_telemetry_is_per_plugin_and_survives_restart():
    pc, _ = load_plugin('LangRAG')
    sa, sb = StorageFixture(), StorageFixture()
    a, b = bind(pc, sa), bind(pc, sb)
    await a.initialize()
    await b.initialize()
    assert hasattr(a, 'telemetry'), 'Telemetry must belong to the installation plugin'
    await a.telemetry.record_delete(collection_id='same', document_id='private-A', status='completed', duration_ms=1)
    assert (await b.telemetry.snapshot())['recent']['delete'] == []
    metrics = await b.telemetry.prometheus()
    warning_line = next(line for line in metrics.splitlines() if line.startswith('langrag_alerts_active{') and 'severity="warning"' in line)
    assert warning_line.endswith(' 0'), 'SDK persistence must not report JSONL-disabled warning'
    a2 = bind(pc, sa)
    await a2.initialize()
    assert (await a2.telemetry.snapshot())['recent']['delete'][0]['document_id'] == 'private-A'
    await a2.telemetry.clear()
    a3 = bind(pc, sa)
    await a3.initialize()
    assert (await a3.telemetry.snapshot())['recent']['delete'] == []


@pytest.mark.asyncio
async def test_text_parser_is_offloop(monkeypatch):
    load_plugin('LangRAG')
    from components.knowledge_engine.parser import FileParser
    import threading
    parser = FileParser()
    main_thread = threading.get_ident()
    monkeypatch.setattr(parser, '_decode_text', lambda data: str(threading.get_ident()))
    for filename in ['doc.txt', 'doc.unknown']:
        assert await parser.parse(b'x', filename) != str(main_thread)


def ingest_context(config):
    return IngestionContext(
        file_object=FileObject(
            metadata=FileMetadata(filename='doc.txt', file_size=12, mime_type='text/plain',
                                  document_id='local-doc', knowledge_base_id='kb'),
            storage_path='fixture-file'),
        knowledge_base_id='kb', creation_settings=config,
    )


def bind_mock_http(monkeypatch, handler):
    client = httpx.AsyncClient
    monkeypatch.setattr(
        httpx, 'AsyncClient',
        lambda **kw: client(transport=httpx.MockTransport(handler), **kw),
    )


@pytest.mark.asyncio
@pytest.mark.parametrize('name', CONNECTORS)
async def test_deterministic_validation_error_does_not_fence(name):
    pc, ec = load_plugin(name)
    store = StorageFixture()
    engine = ec()
    engine.plugin = bind(pc, store)
    # Raised before any Host or provider call: deterministic, must not fence.
    with pytest.raises(ValueError, match='64 KiB'):
        await engine.on_knowledge_base_create('kb', {'blob': 'x' * 70000})
    await engine.on_knowledge_base_create('kb', {'marker': 'ok'})
    restarted = ec()
    restarted.plugin = bind(pc, store)
    assert (await restarted._load_config('kb'))['marker'] == 'ok'


class GatedStorageFixture(StorageFixture):
    """Suspends the first key listing so cancellation lands before any dispatch."""
    def __init__(self):
        super().__init__()
        self.keys_entered = asyncio.Event()
        self.keys_gate = asyncio.Event()

    async def get_plugin_storage_keys(self):
        self.keys_entered.set()
        await self.keys_gate.wait()
        return await super().get_plugin_storage_keys()


@pytest.mark.asyncio
@pytest.mark.parametrize('name', CONNECTORS)
async def test_cancellation_before_dispatch_does_not_fence(name):
    pc, ec = load_plugin(name)
    store = GatedStorageFixture()
    engine = ec()
    engine.plugin = bind(pc, store)
    task = asyncio.create_task(engine.on_knowledge_base_create('kb', {'marker': 'A'}))
    await store.keys_entered.wait()
    task.cancel()
    with pytest.raises(asyncio.CancelledError):
        await task
    store.keys_gate.set()
    # No mutation was dispatched, so the knowledge base stays usable.
    await engine.on_knowledge_base_create('kb', {'marker': 'B'})
    assert (await engine._load_config('kb'))['marker'] == 'B'


@pytest.mark.asyncio
@pytest.mark.parametrize('name', CONNECTORS)
async def test_ambiguous_provider_response_fences_only_that_knowledge_base(name, monkeypatch):
    pc, ec = load_plugin(name)

    async def service(request):
        # Provider reports success but omits the upstream id: dispatched, unknown.
        if request.url.path.endswith('/create-by-file'):
            return httpx.Response(200, json={'document': {}})
        if request.url.path.endswith('/localFile'):
            return httpx.Response(200, json={'code': 200, 'data': {'collectionId': ''}})
        return httpx.Response(200, json={'code': 0, 'data': []})

    bind_mock_http(monkeypatch, service)
    store = StorageFixture()
    engine = ec()
    plugin = bind(pc, store)

    async def read_file(path):
        return b'fixture file'

    plugin.get_knowledge_file_stream = read_file
    engine.plugin = plugin
    cfg = configuration('A')
    with pytest.raises(RuntimeError, match='reconciliation'):
        await engine.ingest(ingest_context(cfg))
    # Only the failing knowledge base is fenced; its sibling still works.
    with pytest.raises(RuntimeError, match='fenced'):
        await engine.on_knowledge_base_create('kb', cfg)
    await engine.on_knowledge_base_create('other-kb', cfg)
    # The fence is persisted, so a restarted worker still refuses that KB only.
    restarted = ec()
    restarted.plugin = bind(pc, store)
    with pytest.raises(RuntimeError, match='fenced'):
        await restarted.on_knowledge_base_create('kb', cfg)
    assert (await restarted._load_config('other-kb')) is not None
    # Deleting the knowledge base is the operator path out of the fence.
    await restarted.on_knowledge_base_delete('kb')
    await restarted.on_knowledge_base_create('kb', cfg)


@pytest.mark.asyncio
@pytest.mark.parametrize('name', CONNECTORS)
async def test_ambiguous_transport_failure_is_reported_and_fenced(name, monkeypatch):
    pc, ec = load_plugin(name)

    async def service(request):
        raise httpx.ReadTimeout('fixture timeout')

    bind_mock_http(monkeypatch, service)
    store = StorageFixture()
    engine = ec()
    plugin = bind(pc, store)

    async def read_file(path):
        return b'fixture file'

    plugin.get_knowledge_file_stream = read_file
    engine.plugin = plugin
    cfg = configuration('A')
    result = await engine.ingest(ingest_context(cfg))
    assert result.status == DocumentStatus.FAILED
    # The caller still gets a FAILED result, but the dispatched upload is unknown.
    with pytest.raises(RuntimeError, match='fenced'):
        await engine.on_knowledge_base_create('kb', cfg)


@pytest.mark.asyncio
@pytest.mark.parametrize('name', CONNECTORS + ['LangRAG'])
async def test_installation_revocation_releases_process_local_state(name):
    pc, ec = load_plugin(name)
    plugin = pc()
    engine = ec()
    engine.plugin = plugin
    await engine.initialize()
    binding = ('workspace', 'installation')
    lock = asyncio.Lock()
    engine._state._locks[(binding, 'kb')] = lock
    engine._state._fenced_keys(binding).add('ke.fence.v1.deadbeef')
    engine._state._fence_loaded.add(binding)
    if hasattr(engine._state, '_fence_load_locks'):
        # The single-flight load lock and the fence generation belong to the
        # binding too, so a revoked installation must not keep a stale pair.
        engine._state._fence_load_locks[binding] = asyncio.Lock()
        engine._state._fence_generation[binding] = 3
    await plugin.on_installation_revoked(binding)
    assert (binding, 'kb') not in engine._state._locks
    assert binding not in engine._state._fenced
    assert binding not in engine._state._fence_loaded
    if hasattr(engine._state, '_fence_load_locks'):
        assert binding not in engine._state._fence_load_locks
        assert binding not in engine._state._fence_generation


@pytest.mark.asyncio
async def test_langrag_fence_is_scoped_to_one_collection_and_survives_restart():
    """One knowledge base's ambiguous mutation must not block its siblings."""

    _, ec = load_plugin('LangRAG')
    engine = ec()
    plugin = StorageFixture()
    engine.plugin = plugin
    await engine.initialize()
    binding = engine._state.binding(plugin)

    async def boom():
        raise RuntimeError('host offline after dispatch')

    async def ok():
        return 7

    with pytest.raises(RuntimeError, match='host offline'):
        await engine._state.mutate(plugin, 'kb-a', boom, 'Host vector_upsert outcome unknown')

    # Only the failing knowledge base is fenced.
    assert engine._state.is_fenced(binding, 'kb-a')
    assert not engine._state.is_fenced(binding, 'kb-b')
    with pytest.raises(RuntimeError, match='fenced'):
        await engine._state.run(plugin, 'kb-a', ok)
    assert await engine._state.mutate(plugin, 'kb-b', ok, 'unrelated') == 7

    # A restarted worker still refuses the fenced knowledge base.
    restarted = ec()
    restarted.plugin = plugin
    with pytest.raises(RuntimeError, match='fenced'):
        await restarted._state.run(plugin, 'kb-a', ok)
    assert not restarted._state.is_fenced(restarted._state.binding(plugin), 'kb-b')

    # Deleting the knowledge base clears the fence.
    await restarted._state.clear_fence(plugin, 'kb-a')
    assert await restarted._state.run(plugin, 'kb-a', ok) == 7