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

"""Exact artifact, real b5 discovery/controller/proxies and WebSocket RPC.

Two native InstallationBinding connections; fake Host storage/vector/model and
fake local HTTP provider. No Cloud, real provider, nsjail or certificate claims.
"""
import asyncio
import base64
import hashlib
import importlib.metadata
import json
import logging
import os
from pathlib import Path
import sys
from contextlib import asynccontextmanager

import websockets
from aiohttp import web
from langbot_plugin.utils.discover.engine import ComponentDiscoveryEngine
from langbot_plugin.cli.utils.page_components import discover_plugin_components
from langbot_plugin.cli.run.controller import PluginRuntimeController
from langbot_plugin.cli.run.handler import PluginRuntimeHandler
from langbot_plugin.runtime.io.handler import Handler, ActionResponse
from langbot_plugin.runtime.io.connections.ws import WebSocketConnection
from langbot_plugin.entities.io.actions.enums import RuntimeToPluginAction as R, PluginToRuntimeAction as P
from langbot_plugin.entities.io.context import InstallationBinding
from langbot_plugin.entities.io.errors import ActionCallError

CODE, OUT, NAME, DIGEST = Path(sys.argv[1]), Path(sys.argv[2]), sys.argv[3], sys.argv[4]
os.chdir(CODE)
sys.path.insert(0, str(CODE))


def _metadata_matches(metadata, filters):
    """Evaluate the Host's metadata filters (equality plus $and/$or terms)."""
    if not filters:
        return True
    if not isinstance(filters, dict):
        raise AssertionError(f'Unsupported vector filter: {filters!r}')
    for key, expected in filters.items():
        if key == '$and' and not all(_metadata_matches(metadata, term) for term in expected):
            return False
        if key == '$or' and not any(_metadata_matches(metadata, term) for term in expected):
            return False
        if key in ('$and', '$or'):
            continue
        if str(metadata.get(key, '')) != str(expected):
            return False
    return True


async def main():
    assert importlib.metadata.version('langbot-plugin') == '0.7.4'
    before = {str(p.relative_to(CODE)): hashlib.sha256(p.read_bytes()).hexdigest() for p in CODE.rglob('*') if p.is_file()}
    discovery = ComponentDiscoveryEngine()
    manifest = discovery.load_component_manifest('manifest.yaml', no_save=True)
    components = discover_plugin_components(manifest, discovery)
    kinds = sorted(c.kind for c in components)
    if NAME == 'LongTermMemory':
        # LongTermMemory ships a KnowledgeEngine, an EventListener, four Tools,
        # a Command and a Page, so the exact kind count is not a fixed list. The
        # kinds are read from the manifest inside the exact archive instead.
        declared_kinds = sorted(manifest.manifest['spec']['components'])
        assert sorted(set(kinds)) == declared_kinds, (kinds, declared_kinds)
    else:
        assert kinds == (['KnowledgeEngine', 'Page'] if NAME == 'LangRAG' else ['KnowledgeEngine'])
    # Every verified candidate is a stateless shared-placement archive, so the
    # manifest inside the built artifact has to declare exactly that. The
    # declaration is not a certificate: this only checks it survived packaging.
    assert manifest.execution.shared_runtime == 'shared-runtime-v1'
    assert manifest.execution.component_model == 'stateless-v1'
    stores, vectors, calls = {}, {}, []
    searches = []
    # Filled only by the LongTermMemory branch, so the archived report records
    # the memory markers and retrieval scopes its assertions actually checked.
    ltm_evidence = {}
    http_calls = []
    async def provider(request):
        body = await request.read()
        marker = request.headers['Authorization'].split()[-1]
        http_calls.append((request.method, request.path, marker))
        await asyncio.sleep(.002)
        if request.path.endswith('/retrieve'):
            return web.json_response({'records': [{'segment': {'id': 'seg', 'content': marker}, 'score': .9}]})
        if request.path.endswith('/searchTest'):
            return web.json_response({'data': [{'id': 'seg', 'q': marker, 'score': .9}]})
        if request.path.endswith('/retrieval'):
            return web.json_response({'code': 0, 'data': {'chunks': [{'id': 'seg', 'content': marker, 'similarity': .9}]}})
        if request.path.endswith('/create-by-file'):
            assert b'fixture file' in body
            return web.json_response({'document': {'id': 'remote-doc'}})
        if request.path.endswith('/localFile'):
            assert b'fixture file' in body
            return web.json_response({'code': 200, 'data': {'collectionId': 'remote-doc'}})
        if request.method == 'POST' and request.path.endswith('/documents'):
            assert b'fixture file' in body
            return web.json_response({'code': 0, 'data': [{'id': 'remote-doc'}]})
        if request.method == 'GET':
            return web.json_response({'code': 0, 'data': []})
        return web.json_response({'code': 200 if NAME == 'FastGPTConnector' else 0})
    app = web.Application()
    app.router.add_route('*', '/{path:.*}', provider)
    runner = web.AppRunner(app)
    await runner.setup()
    site = web.TCPSite(runner, '127.0.0.1', 0)
    await site.start()
    url = f'http://127.0.0.1:{site._server.sockets[0].getsockname()[1]}'
    def binding(label):
        return InstallationBinding(instance_uuid='fixture-instance', workspace_uuid='workspace-'+label,
            placement_generation=1, installation_uuid='installation-'+label, runtime_revision=1, artifact_digest=DIGEST)
    a, b = binding('A'), binding('B')

    @asynccontextmanager
    async def installation(scope):
        label = scope.workspace_uuid
        store = stores.setdefault(label, {})
        rows = vectors.setdefault(label, {})
        controller = PluginRuntimeController(manifest, components, False, '')
        ready = asyncio.get_running_loop().create_future()
        async def connected(ws):
            os.environ['LANGBOT_PLUGIN_FILE_STORAGE_DIR'] = str(OUT / 'rpc-tmp' / NAME / label / 'worker')
            handler = PluginRuntimeHandler(WebSocketConnection(ws), controller.initialize)
            controller.handler = handler
            handler.plugin_container = controller.plugin_container
            # Stateless shared placement: the runtime drives the plugin through
            # per-slot lifecycle callbacks, not one bound plugin instance.
            handler._slot_initialize_callback = controller.initialize_slot
            handler._slot_detach_callback = controller.detach_slot
            handler._slot_cancel_callback = controller.invalidate_slot
            ready.set_result(handler)
            await handler.run()
        async with websockets.serve(connected, '127.0.0.1', 0) as server:
            async with websockets.connect(f'ws://127.0.0.1:{server.sockets[0].getsockname()[1]}') as ws:
                host = Handler(WebSocketConnection(ws), file_storage_dir=OUT / 'rpc-tmp' / NAME / label / 'host')
                def check(action, data):
                    assert host.current_action_context == scope, 'SDK proxy lost its immutable binding'
                    calls.append((label, action))
                @host.action(P.GET_PLUGIN_STORAGE_KEYS)
                async def keys(data):
                    check('keys', data)
                    return ActionResponse.success({'keys': list(store)})
                @host.action(P.GET_PLUGIN_STORAGE)
                async def get(data):
                    check('get', data)
                    # A key that was never written reads back as an empty object,
                    # matching the Host storage fixture the shared-placement probe
                    # already drives LongTermMemory with.
                    if data['key'] not in store:
                        store[data['key']] = base64.b64encode(b'{}').decode()
                    return ActionResponse.success({'value_base64': store[data['key']]})
                @host.action(P.SET_PLUGIN_STORAGE)
                async def set_value(data):
                    check('set', data)
                    await asyncio.sleep(.002)
                    store[data['key']] = data['value_base64']
                    return ActionResponse.success({})
                @host.action(P.INVOKE_EMBEDDING)
                async def embed(data):
                    check('embedding', data)
                    return ActionResponse.success({'vectors': [[1.0] for _ in data['texts']]})
                @host.action(P.VECTOR_UPSERT)
                async def upsert(data):
                    check('upsert', data)
                    for id_, meta in zip(data['ids'], data['metadata']):
                        rows[id_] = {'id': id_, 'distance': .1, 'metadata': meta}
                    return ActionResponse.success({})
                @host.action(P.VECTOR_SEARCH)
                async def search(data):
                    check('search', data)
                    searches.append((label, data.get('filters')))
                    return ActionResponse.success({'results': list(rows.values())[:data['top_k']]})
                @host.action(P.VECTOR_LIST)
                async def list_vectors(data):
                    check('vector_list', data)
                    hits = [row for row in rows.values() if _metadata_matches(row['metadata'], data.get('filters'))]
                    offset, limit = data.get('offset', 0), data.get('limit', 20)
                    return ActionResponse.success({'items': hits[offset:offset + limit], 'total': len(hits)})
                @host.action(P.VECTOR_DELETE)
                async def delete(data):
                    check('delete', data)
                    n = len(rows)
                    rows.clear()
                    return ActionResponse.success({'count': n})
                @host.action(P.GET_KNOWLEDEGE_FILE_STREAM)
                async def file_stream(data):
                    check('file', data)
                    key = await host.send_file(b'fixture file', 'txt', action_context=scope)
                    return ActionResponse.success({'file_key': key})
                task = asyncio.create_task(host.run())
                handler = await ready
                async def call(action, data, context=scope):
                    return await host.call_action(action, data, action_context=context, timeout=8)
                try:
                    settings = {'plugin_settings': {'enabled': True, 'priority': 0, 'plugin_config': {'marker': label}}}
                    await call(R.ATTACH_PLUGIN_SLOT, settings)
                    slot = await call(R.GET_PLUGIN_SLOT_CONTAINER, {})
                    assert slot['plugin_config'] == {'marker': label}
                    assert controller.plugin_container_for_slot(scope.installation_uuid).binding == scope
                    yield call, controller
                finally:
                    await host.close()
                    await handler.close()
                    task.cancel()
                    await asyncio.gather(task, return_exceptions=True)
                    await controller.cleanup_instances()

    def config(marker):
        return {'api_base_url': url, 'api_key': marker, 'dify_apikey': marker,
                'dataset_id': marker, 'dataset_ids': marker, 'embedding_model_uuid': 'fixture-embedding'}
    async def ingest(call, marker):
        ctx = {'knowledge_base_id': 'same-kb', 'creation_settings': config(marker),
               'file_object': {'metadata': {'filename': 'doc.txt', 'document_id': 'same-doc',
                 'knowledge_base_id': 'same-kb', 'file_size': 12, 'mime_type': 'text/plain'}, 'storage_path': 'fixture-file'}}
        if NAME == 'LangRAG':
            ctx['parsed_content'] = {'text': 'private-' + marker}
        result = await call(R.INGEST_DOCUMENT, {'context': ctx})
        assert result['status'] == ('completed' if NAME == 'LangRAG' else 'processing'), result
    async def retrieve(call, marker):
        result = await call(R.RETRIEVE_KNOWLEDGE, {'retriever_name': '', 'retrieval_context': {
            'query': 'fixture query', 'knowledge_base_id': 'same-kb', 'creation_settings': config(marker), 'retrieval_settings': {'top_k': 2}}})
        assert result['total_found'] == 1, result
        assert marker in result['results'][0]['content'][0]['text'], result
        return result
    ltm_kb, ltm_doc = 'ltm-kb', 'ltm-doc'
    ltm_creation = {'embedding_model_uuid': 'fixture-embedding', 'isolation': 'session'}
    page_id = next((c.metadata.name for c in components if c.kind == 'Page'), None)
    async def page(call, endpoint, method='GET', body=None):
        return await call(R.PAGE_API, {'page_id': page_id, 'endpoint': endpoint, 'method': method,
                                       'body': {} if body is None else body})
    def ltm_session(marker):
        return 'person_fixture-' + marker
    async def ltm_user_key(call, marker):
        resp = await page(call, '/derive-scope', 'POST', {'bot_uuid': 'fixture-bot',
            'session_name': ltm_session(marker), 'isolation': 'session'})
        assert not resp.get('error'), resp
        return resp['data']['user_key']
    async def ltm_ingest(call, marker, user_key):
        payload = json.dumps([{'content': 'ltm-' + marker + ' remembered fact', 'tags': ['fixture'],
                               'importance': 3, 'user_key': user_key}])
        ctx = {'knowledge_base_id': ltm_kb, 'creation_settings': ltm_creation,
               'parsed_content': {'text': payload},
               'file_object': {'metadata': {'filename': 'memories.json', 'document_id': ltm_doc,
                 'knowledge_base_id': ltm_kb, 'file_size': len(payload), 'mime_type': 'application/json'},
                 'storage_path': 'fixture-file'}}
        result = await call(R.INGEST_DOCUMENT, {'context': ctx})
        assert result['status'] == 'completed', result
        assert result['chunks_created'] == 1, result
    async def ltm_episodes(call, user_key):
        resp = await page(call, '/episodes/list', 'POST', {'collection_id': ltm_kb, 'user_key': user_key})
        assert not resp.get('error'), resp
        return resp['data']
    async def ltm_summary(call):
        resp = await page(call, '/summary')
        assert not resp.get('error'), resp
        return resp['data']
    async def ltm_retrieve(call, marker):
        result = await call(R.RETRIEVE_KNOWLEDGE, {'retriever_name': '', 'retrieval_context': {
            'query': 'remembered fact', 'knowledge_base_id': ltm_kb, 'creation_settings': ltm_creation,
            'retrieval_settings': {'top_k': 2, 'session_name': ltm_session(marker), 'bot_uuid': 'fixture-bot'}}})
        assert result['total_found'] == 1, result
        text = result['results'][0]['content'][0]['text']
        assert '<mem>ltm-' + marker + ' remembered fact</mem>' in text, text
    try:
        async with installation(a) as (ca, cta), installation(b) as (cb, ctb):
            if NAME == 'LongTermMemory':
                capabilities = await ca(R.GET_RAG_CAPABILITIES, {})
                assert capabilities['capabilities'] == ['doc_ingestion'], capabilities
                for call in (ca, cb):
                    created = await call(R.ON_KB_CREATE, {'kb_id': ltm_kb, 'config': dict(ltm_creation)})
                    assert created['success'], created
                user_keys = {}
                for call, marker in ((ca, 'A'), (cb, 'B')):
                    user_keys[marker] = await ltm_user_key(call, marker)
                    await ltm_ingest(call, marker, user_keys[marker])
                for call, marker, other in ((ca, 'A', 'B'), (cb, 'B', 'A')):
                    mine = await ltm_episodes(call, user_keys[marker])
                    assert mine['total'] == 1 and 'ltm-' + marker in mine['episodes'][0]['content'], mine
                    assert (await ltm_episodes(call, user_keys[other]))['total'] == 0, other
                    summary = await ltm_summary(call)
                    assert summary['kb_count'] == 1, summary
                    assert summary['kb_entries'][0]['id'] == ltm_kb, summary
                    assert summary['profile_key_count'] == 0, summary
                await asyncio.gather(*(ltm_retrieve(ca if i % 2 == 0 else cb, 'A' if i % 2 == 0 else 'B') for i in range(12)))
                # Writes and reads went to this installation's own Host storage:
                # each binding holds exactly its own episode, and the retrieval
                # scope filter was derived from that binding's session.
                assert [row['metadata']['user_key'] for row in vectors['workspace-A'].values()] == [user_keys['A']]
                assert [row['metadata']['user_key'] for row in vectors['workspace-B'].values()] == [user_keys['B']]
                for label, marker in (('workspace-A', 'A'), ('workspace-B', 'B')):
                    scopes = [f for l, f in searches if l == label]
                    assert scopes and all(f == {'user_key': user_keys[marker]} for f in scopes), scopes
                ltm_evidence['memory_markers'] = sorted('ltm-' + marker for marker in user_keys)
                ltm_evidence['retrieval_scope_filters'] = {
                    label: sorted({json.dumps(f) for l, f in searches if l == label})
                    for label in ('workspace-A', 'workspace-B')}
            else:
                await asyncio.gather(ingest(ca, 'A'), ingest(cb, 'B'))
                await asyncio.gather(*(retrieve(ca if i % 2 == 0 else cb, 'A' if i % 2 == 0 else 'B') for i in range(12)))
            try:
                await ca(R.ON_KB_DELETE, {'kb_id': 'same-kb'}, context=b)
            except ActionCallError:
                pass
            else:
                raise AssertionError('Cross-Workspace action was accepted')
            assert cta.plugin_container.plugin_instance is not ctb.plugin_container.plugin_instance
            if NAME == 'LangRAG':
                sa, sb = await page(ca, '/snapshot'), await page(cb, '/snapshot')
                assert sa['data']['counters']['ingest.total'] == sb['data']['counters']['ingest.total'] == 1
                await page(ca, '/clear', 'POST')
                assert (await page(ca, '/snapshot'))['data']['counters'] == {}
                assert (await page(cb, '/snapshot'))['data']['counters']['ingest.total'] == 1
        # Fresh real controller objects recover each installation's durable state.
        async with installation(a) as (ca, _), installation(b) as (cb, _):
            if NAME == 'LongTermMemory':
                for call, marker in ((ca, 'A'), (cb, 'B')):
                    assert (await ltm_summary(call))['kb_count'] == 1, marker
                    assert (await ltm_episodes(call, user_keys[marker]))['total'] == 1, marker
                deleted = await ca(R.DELETE_DOCUMENT, {'kb_id': ltm_kb, 'document_id': ltm_doc})
                assert deleted['success'], deleted
                assert (await ltm_episodes(ca, user_keys['A']))['total'] == 0
                assert (await ltm_episodes(cb, user_keys['B']))['total'] == 1
                assert (await ca(R.ON_KB_DELETE, {'kb_id': ltm_kb}))['success']
                assert (await ltm_summary(ca))['kb_count'] == 0
                assert (await ltm_summary(cb))['kb_count'] == 1
                await ltm_retrieve(cb, 'B')
            else:
                if NAME == 'LangRAG':
                    assert (await page(ca, '/snapshot'))['data']['counters'] == {}
                    assert (await page(cb, '/snapshot'))['data']['counters']['ingest.total'] == 1
                result = await ca(R.DELETE_DOCUMENT, {'kb_id': 'same-kb', 'document_id': 'same-doc'})
                assert result['success']
                if NAME != 'LangRAG':
                    assert http_calls[-1][2] == 'A'
                    await ca(R.ON_KB_DELETE, {'kb_id': 'same-kb'})
                    assert not (await ca(R.DELETE_DOCUMENT, {'kb_id': 'same-kb', 'document_id': 'same-doc'}))['success']
                await retrieve(cb, 'B')
        after = {str(p.relative_to(CODE)): hashlib.sha256(p.read_bytes()).hexdigest() for p in CODE.rglob('*') if p.is_file()}
        assert before == after, 'Runtime wrote into shared artifact tree'
        report = {'name': NAME, 'sdk': '0.7.4', 'artifact_sha256': DIGEST, 'components': kinds,
                  'native_installation_bindings': 2, 'concurrent_retrievals': 12,
                  'cross_workspace_denied': True, 'restart_state_recovered': True,
                  'artifact_unchanged': True, 'host_calls': len(calls), 'http_calls': len(http_calls),
                  'scope': 'Real SDK loopback RPC; fake Host/provider; no production or nsjail execution',
                  **ltm_evidence}
        (OUT / (NAME + '-rpc-results.json')).write_text(json.dumps(report, indent=2))
        print(json.dumps(report, indent=2))
    finally:
        await runner.cleanup()


if __name__ == '__main__':
    logging.basicConfig(level=logging.CRITICAL)
    asyncio.run(asyncio.wait_for(main(), 80))