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

"""SDK 0.7.4 real child, WebSocket RPC, two slots and one component graph.

LangRAG/LongTermMemory: local Host vector/storage fixture only; no Cloud/sandbox
claim. Connectors: local Host storage/file fixture plus a loopback HTTP provider
fixture; no upstream/Cloud/sandbox claim.
"""
import asyncio
import base64
import hashlib
import json
import os
from pathlib import Path
import sys

from aiohttp import web
import websockets
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

CODE = Path(sys.argv[1]).resolve()
os.chdir(CODE)
sys.path.insert(0, str(CODE))

CONNECTORS = ('DifyDatasetsConnector', 'FastGPTConnector', 'RAGFlowConnector')

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='a'*64)

async def child(url):
    discovery = ComponentDiscoveryEngine()
    manifest = discovery.load_component_manifest('manifest.yaml', no_save=True)
    # Probe shared eligibility without changing the dedicated release manifest.
    if CODE.name == 'LongTermMemory':
        manifest.execution.shared_runtime = 'shared-runtime-v1'
        manifest.execution.component_model = 'stateless-v1'
    components = discover_plugin_components(manifest, discovery)
    controller = PluginRuntimeController(manifest, components, False, '')
    original_attach = controller.initialize_slot
    identities = []
    async def attach(binding, settings):
        slot = await original_attach(binding, settings)
        identity = (id(slot.plugin_container.plugin_instance),
                    tuple(id(c.component_instance) for c in slot.plugin_container.components))
        identities.append(identity)
        if len(identities) == 2:
            assert identities[0] == identities[1], 'Slots have different object graphs'
            print('SAME_OBJECT_GRAPH=true', file=sys.stderr, flush=True)
        return slot
    async with websockets.connect(url) as ws:
        handler = PluginRuntimeHandler(WebSocketConnection(ws), controller.initialize)
        handler.plugin_container = controller.plugin_container
        handler._slot_initialize_callback = attach
        handler._slot_detach_callback = controller.detach_slot
        handler._slot_cancel_callback = controller.invalidate_slot
        controller.handler = handler
        await handler.run()

async def parent():
    stores = {label: {} for label in 'AB'}
    vectors = {label: {} for label in 'AB'}
    actions = []
    provider_calls = []

    async def provider(request):
        """Loopback stand-in for Dify/FastGPT/RAGFlow; echoes the caller's own credentials."""
        body = await request.read()
        auth = request.headers.get('Authorization', '')
        marker = auth.removeprefix('Bearer ')
        provider_calls.append((request.method, request.path_qs, auth, body))
        if request.method == 'DELETE':
            # FastGPT acknowledges with code 200; Dify accepts any 2xx; RAGFlow reads code 0.
            return web.json_response({'code': 200 if request.path.endswith('/collection/delete') else 0})
        if request.path.endswith('/retrieve'):
            return web.json_response({'records': [{'segment': {'id': 'seg', 'content': auth}, 'score': .8}]})
        if request.path.endswith('/searchTest'):
            return web.json_response({'code': 200, 'data': [{'id': 'seg', 'q': auth, 'a': 'answer', 'score': .8}]})
        if request.path.endswith('/retrieval'):
            return web.json_response({'code': 0, 'data': {'chunks': [{'id': 'seg', 'content': auth, 'similarity': .8}]}})
        if request.path.endswith('/create-by-file'):
            return web.json_response({'document': {'id': 'remote-doc-' + marker}})
        if request.path.endswith('/localFile'):
            return web.json_response({'code': 200, 'data': {'collectionId': 'remote-doc-' + marker,
                'results': {'insertLen': 2}}})
        if request.path.endswith(('/chunks', '/run_graphrag', '/run_raptor')):
            return web.json_response({'code': 0, 'data': {}})
        if request.method == 'POST' and request.path.endswith('/documents'):
            return web.json_response({'code': 0, 'data': [{'id': 'remote-doc-' + marker}]})
        # Dataset validation/listing.
        return web.json_response({'code': 0, 'data': []})

    # The loopback provider fixture is only needed by the connector branch.
    provider_runner = None
    base_url = ''
    if CODE.name in CONNECTORS:
        app = web.Application()
        app.router.add_route('*', '/{path:.*}', provider)
        provider_runner = web.AppRunner(app)
        await provider_runner.setup()
        site = web.TCPSite(provider_runner, '127.0.0.1', 0)
        await site.start()
        base_url = f"http://127.0.0.1:{site._server.sockets[0].getsockname()[1]}"

    connected = asyncio.get_running_loop().create_future()
    async def accept(ws):
        handler = Handler(WebSocketConnection(ws), file_storage_dir=Path('/tmp/ke-shared-probe-host'))
        def label():
            ctx = handler.current_action_context
            assert ctx in (binding('A'), binding('B')), ctx
            return ctx.workspace_uuid[-1]
        @handler.action(P.GET_PLUGIN_STORAGE_KEYS)
        async def keys(data):
            l = label(); actions.append((l, 'keys'))
            return ActionResponse.success({'keys': list(stores[l])})
        @handler.action(P.GET_PLUGIN_STORAGE)
        async def get(data):
            l = label(); actions.append((l, 'get'))
            if data['key'] not in stores[l]:
                stores[l][data['key']] = base64.b64encode(b'{}').decode()
            return ActionResponse.success({'value_base64': stores[l][data['key']]})
        @handler.action(P.SET_PLUGIN_STORAGE)
        async def set_value(data):
            l = label(); actions.append((l, 'set'))
            stores[l][data['key']] = data['value_base64']
            return ActionResponse.success({})
        @handler.action(P.GET_KNOWLEDEGE_FILE_STREAM)
        async def file_stream(data):
            # Per-binding fixture file: a binding must never read its sibling's bytes.
            l = label(); actions.append((l, 'file'))
            assert data['storage_path'] == 'fixture-file', data
            key = await host.send_file(b'private-' + l.encode(), '')
            return ActionResponse.success({'file_key': key})
        @handler.action(P.INVOKE_EMBEDDING)
        async def embed(data):
            l = label(); actions.append((l, 'embed'))
            return ActionResponse.success({'vectors': [[1.0] for _ in data['texts']]})
        @handler.action(P.VECTOR_UPSERT)
        async def upsert(data):
            l = label(); actions.append((l, 'upsert'))
            for id_, meta in zip(data['ids'], data['metadata']):
                vectors[l][id_] = {'id': id_, 'distance': .1, 'metadata': meta}
            return ActionResponse.success({})
        @handler.action(P.VECTOR_SEARCH)
        async def search(data):
            l = label(); actions.append((l, 'search'))
            return ActionResponse.success({'results': list(vectors[l].values())[:data['top_k']]})
        @handler.action(P.VECTOR_DELETE)
        async def delete(data):
            l = label(); actions.append((l, 'delete'))
            n = len(vectors[l]); vectors[l].clear()
            return ActionResponse.success({'count': n})
        connected.set_result(handler)
        await handler.run()
    async with websockets.serve(accept, '127.0.0.1', 0) as server:
        url = f'ws://127.0.0.1:{server.sockets[0].getsockname()[1]}'
        proc = await asyncio.create_subprocess_exec(sys.executable, __file__, str(CODE), '--child', url,
            stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE,
            env={**os.environ, 'LANGBOT_PLUGIN_FILE_STORAGE_DIR': '/tmp/ke-shared-probe-child-files'})
        try:
            host = await asyncio.wait_for(connected, 15)
            async def call(label, action, data):
                return await host.call_action(action, data, action_context=binding(label), timeout=20)
            def settings(label):
                if CODE.name in CONNECTORS:
                    # Connectors declare no plugin config; the per-installation
                    # settings still differ between the two bindings.
                    return {'enabled': True, 'priority': 0 if label == 'A' else 1, 'plugin_config': {}}
                return {'enabled': True, 'priority': 0, 'plugin_config': {
                    'max_profile_traits': 1 if label == 'A' else 4,
                    'max_profile_preferences': 1 if label == 'A' else 4}}
            for label in 'AB':
                await call(label, R.ATTACH_PLUGIN_SLOT, {'plugin_settings': settings(label)})
            a = await call('A', R.GET_PLUGIN_SLOT_CONTAINER, {})
            b = await call('B', R.GET_PLUGIN_SLOT_CONTAINER, {})
            if CODE.name in CONNECTORS:
                assert a['priority'] != b['priority'], (a['priority'], b['priority'])
            else:
                assert a['plugin_config'] != b['plugin_config']
            if CODE.name == 'LangRAG':
                for label in 'AB':
                    result = await call(label, R.INGEST_DOCUMENT, {'context': {
                        'knowledge_base_id': 'same-kb', 'creation_settings': {'embedding_model_uuid': 'fixture-embedding'},
                        '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'},
                        'parsed_content': {'text': 'private-'+label}}})
                    assert result['status'] == 'completed', result
                for label in 'AB':
                    result = await call(label, R.RETRIEVE_KNOWLEDGE, {'retriever_name': '', 'retrieval_context': {
                        'query': 'fixture query', 'knowledge_base_id': 'same-kb',
                        'creation_settings': {'embedding_model_uuid': 'fixture-embedding'},
                        'retrieval_settings': {'top_k': 2}}})
                    assert result['total_found'] == 1, result
                    assert 'private-'+label in str(result), result
                for label in 'AB':
                    page = await call(label, R.PAGE_API, {'page_id': 'observability', 'endpoint': '/snapshot',
                        'method': 'GET', 'body': {}})
                    assert page['data']['counters']['ingest.total'] == 1, page
            elif CODE.name in CONNECTORS:
                def connector_config(label):
                    return {'api_base_url': base_url, 'api_key': label,
                        'dify_apikey': label, 'dataset_id': label, 'dataset_ids': label}
                for label in 'AB':
                    # Host notification for the same knowledge-base ID on each binding.
                    await call(label, R.ON_KB_CREATE, {'kb_id': 'same-kb', 'config': connector_config(label)})
                for label in 'AB':
                    result = await call(label, R.INGEST_DOCUMENT, {'context': {
                        'knowledge_base_id': 'same-kb', 'creation_settings': connector_config(label),
                        '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'}}})
                    # The provider acknowledged a per-binding upstream document.
                    assert result['status'] == 'processing', result
                    assert result['document_id'] == 'remote-doc-'+label, result
                for label in 'AB':
                    result = await call(label, R.RETRIEVE_KNOWLEDGE, {'retriever_name': '', 'retrieval_context': {
                        'query': 'fixture query', 'knowledge_base_id': 'same-kb',
                        'creation_settings': connector_config(label), 'retrieval_settings': {'top_k': 2}}})
                    assert result['total_found'] == 1, result
                    other = 'B' if label == 'A' else 'A'
                    assert 'Bearer '+label in str(result), result
                    assert 'Bearer '+other not in str(result), result
                # Same Host storage key per binding, distinct durable per-binding value.
                for label in 'AB':
                    other = 'B' if label == 'A' else 'A'
                    key = 'ke.config.v1.' + hashlib.sha256(b'same-kb').hexdigest()
                    assert json.loads(base64.b64decode(stores[label][key]))['api_key'] == label
                    assert all('"'+other+'"' not in base64.b64decode(value).decode()
                        for value in stores[label].values()), stores[label]
                for label in 'AB':
                    # Delete resolves the stored per-binding credentials, not the sibling's.
                    result = await call(label, R.DELETE_DOCUMENT, {'kb_id': 'same-kb', 'document_id': 'same-doc'})
                    assert result['success'] is True, result
                for label in 'AB':
                    other = 'B' if label == 'A' else 'A'
                    own = [c for c in provider_calls if c[2] == 'Bearer '+label]
                    assert own, 'no provider call carried binding '+label+' credentials'
                    assert any(c[0] == 'DELETE' for c in own), own
                    assert 'remote-doc-'+label in str(own) and 'remote-doc-'+other not in str(own), own
                    assert 'private-'+label in str([c[3] for c in own])
                    assert 'private-'+other not in str([c[3] for c in own]), own
                # Each binding's ingest read its own Host file stream.
                assert [l for l, action in actions if action == 'file'] == ['A', 'B'], actions
            else:
                for label in 'AB':
                    page = await call(label, R.PAGE_API, {'page_id': 'memory_console', 'endpoint': '/summary',
                        'method': 'GET', 'body': {}})
                    assert page.get('data') and page['data']['plugin_config']['max_profile_traits'] == (1 if label == 'A' else 4), page
            summary = {'plugin': CODE.name, 'sdk': '0.7.4', 'worker_pid': proc.pid,
                'same_worker_two_workspace_slots': True, 'slot_settings_distinct': True,
                'successful_representative_invocation': True,
                'host_actions': actions, 'scope': 'real SDK child RPC, local Host fixture; no production/upstream/sandbox proof'}
            if CODE.name not in CONNECTORS:
                # Connectors declare no plugin config; their slots differ in
                # per-installation settings only.
                summary['slot_configs_distinct'] = True
            print(json.dumps(summary))
        finally:
            proc.terminate()
            try:
                stdout, stderr = await asyncio.wait_for(proc.communicate(), 5)
            except asyncio.TimeoutError:
                proc.kill(); stdout, stderr = await proc.communicate()
            if provider_runner is not None:
                await provider_runner.cleanup()
            # A terminated worker is SIGTERM on POSIX and exit code 1 on Windows.
            terminated = (0, -15) if os.name != 'nt' else (0, 1)
            if proc.returncode not in terminated:
                print('child stderr:', stderr.decode()[-3500:], file=sys.stderr)
            assert b'SAME_OBJECT_GRAPH=true' in stderr, stderr.decode()[-3500:]
            assert proc.returncode in terminated, proc.returncode

if __name__ == '__main__':
    if len(sys.argv) > 2 and sys.argv[2] == '--child':
        asyncio.run(child(sys.argv[3]))
    else:
        asyncio.run(asyncio.wait_for(parent(), 60))