LangBot/tests/integration/plugin/test_rag_file_transfer_protocol.py

"""Real Core/SDK protocol regression tests; no subprocesses or external services.

Run against the intended local SDK (``uv run --no-sync`` after local install).
The in-memory transport carries JSON strings through Handler.run on both sides;
send_file, envelope validation, base64 decoding and transfer storage are real.
Only Core's database/object-storage services, parser dispatch/provider and host
sandbox prerequisite probing are doubles. Worker launch/registration is
represented by its already-registered state.
"""

from __future__ import annotations

import asyncio
import json
import logging
from contextlib import asynccontextmanager
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import AsyncMock

import pytest

from langbot.pkg.plugin.handler import RuntimeConnectionHandler
from langbot_plugin.entities.io.actions.enums import CommonAction, LangBotToRuntimeAction, PluginToRuntimeAction
from langbot_plugin.entities.io.context import ActionContext, InstallationBinding, PluginWorkerPolicy, RuntimeIdentity
from langbot_plugin.runtime.context import RuntimeContext
from langbot_plugin.runtime.io.connection import Connection
from langbot_plugin.entities.io.errors import ActionCallError, ConnectionClosedError
from langbot_plugin.runtime.io.handler import FILE_CHUNK_LENGTH, Handler
from langbot_plugin.runtime.io.handlers.control import ControlConnectionHandler
from langbot_plugin.runtime.io.handlers.plugin import PluginConnectionHandler
from langbot_plugin.runtime.plugin.mgr import PluginManager
from langbot_plugin.runtime.security import PLUGIN_FILE_STORAGE_DIR_ENV


pytestmark = pytest.mark.asyncio
PAYLOAD = bytes(range(256)) * 161 + b'\x00original RAG file\xff'
BINDING = InstallationBinding(
    instance_uuid='instance-a',
    workspace_uuid='workspace-a',
    placement_generation=7,
    installation_uuid='00000000-0000-4000-8000-000000000001',
    runtime_revision=3,
    artifact_digest='a' * 64,
)
LEGACY = ActionContext(**BINDING.model_dump(exclude={'runtime_revision', 'artifact_digest'}))


class QueueConnection(Connection):
    """Only the byte transport is replaced, not the request/response machinery."""

    def __init__(self):
        self.incoming = asyncio.Queue()
        self.sent = []
        self.peer = None

    async def send(self, message: str) -> None:
        assert isinstance(message, str)
        self.sent.append(json.loads(message))
        await self.peer.incoming.put(message)

    async def receive(self) -> str:
        message = await self.incoming.get()
        if message is None:
            raise ConnectionClosedError('test transport closed')
        return message

    async def close(self) -> None:
        await self.incoming.put(None)
        await self.peer.incoming.put(None)


def connection_pair():
    left, right = QueueConnection(), QueueConnection()
    left.peer, right.peer = right, left
    return left, right


@asynccontextmanager
async def protocol_stack(tmp_path, monkeypatch, profile='oss_dev', binding=LEGACY):
    monkeypatch.chdir(tmp_path)
    stored = tmp_path / 'original.bin'
    stored.write_bytes(PAYLOAD)
    storage_calls = []

    async def get_file_stream(execution_context, storage_path):
        storage_calls.append((execution_context, storage_path))
        assert execution_context.workspace_uuid == BINDING.workspace_uuid
        assert storage_path == 'knowledge/original.bin'
        return stored.read_bytes()

    async def get_execution_binding(workspace_uuid, expected_generation):
        assert workspace_uuid == BINDING.workspace_uuid
        assert expected_generation == BINDING.placement_generation
        return BINDING

    setting = SimpleNamespace(
        plugin_author='tester',
        plugin_name='engine',
        installation_uuid=BINDING.installation_uuid,
        runtime_revision=BINDING.runtime_revision,
        artifact_digest=BINDING.artifact_digest,
    )
    app = SimpleNamespace(
        deployment=SimpleNamespace(mode='oss' if profile == 'oss_dev' else 'cloud'),
        logger=logging.getLogger(__name__),
        persistence_mgr=SimpleNamespace(execute_async=AsyncMock(return_value=SimpleNamespace(first=lambda: setting))),
        workspace_service=SimpleNamespace(get_execution_binding=get_execution_binding),
        rag_runtime_service=SimpleNamespace(get_file_stream=get_file_stream),
    )
    core_conn, control_conn = connection_pair()
    monkeypatch.setenv(PLUGIN_FILE_STORAGE_DIR_ENV, str(tmp_path / 'core-transfer'))
    core = RuntimeConnectionHandler(core_conn, AsyncMock(return_value=False), app)
    core.register_installation_binding(BINDING, plugin_author='tester', plugin_name='engine')
    runtime = RuntimeContext()
    runtime.plugin_mgr = PluginManager(runtime)
    # No worker is launched: omit only host nsjail/cgroup prerequisite probing.
    monkeypatch.setattr(runtime.plugin_mgr.worker_launcher, 'configure', lambda policy, profile: None)
    monkeypatch.setenv(PLUGIN_FILE_STORAGE_DIR_ENV, str(tmp_path / 'runtime-transfer'))
    control = ControlConnectionHandler(control_conn, runtime)
    runtime.activate_control_handler(control)
    bridge_conn, plugin_conn = connection_pair()
    bridge = PluginConnectionHandler(bridge_conn, runtime, file_storage_dir=str(tmp_path / 'bridge-transfer'))
    plugin = Handler(plugin_conn, file_storage_dir=str(tmp_path / 'plugin-transfer'))
    # Trusted state left by registration, not plugin-supplied action data.
    bridge.bind_action_context(binding)
    runtime.plugin_mgr.plugin_handlers.append(bridge)
    runtime.plugin_mgr.plugins.append(SimpleNamespace(_runtime_plugin_handler=bridge))
    handlers = [core, control, bridge, plugin]
    tasks = [asyncio.create_task(handler.run()) for handler in handlers]
    try:
        await asyncio.wait_for(
            core.set_runtime_config(
                runtime_identity=RuntimeIdentity(instance_uuid='instance-a', runtime_id='test-runtime'),
                worker_policy=PluginWorkerPolicy(
                    max_cpus=1,
                    max_memory_mb=128,
                    max_pids=32,
                    max_open_files=64,
                    max_file_size_mb=8,
                    require_hard_limits=False,
                ),
                runtime_profile=profile,
                cloud_service_url=None,
            ),
            5,
        )
        if isinstance(binding, InstallationBinding):
            runtime.activate_installation_binding(binding)
        else:
            runtime.bind_workspace(binding)
        yield SimpleNamespace(
            core=core,
            control=control,
            runtime=runtime,
            bridge=bridge,
            plugin=plugin,
            core_conn=core_conn,
            control_conn=control_conn,
            bridge_conn=bridge_conn,
            plugin_conn=plugin_conn,
            app=app,
            storage_calls=storage_calls,
        )
    finally:
        for handler in handlers:
            await handler.close()
        await asyncio.wait_for(asyncio.gather(*tasks, return_exceptions=True), 5)


def assert_chunks(connection, binding, payload=PAYLOAD):
    chunks = [message for message in connection.sent if message.get('action') == CommonAction.FILE_CHUNK.value]
    expected = (len(payload) + FILE_CHUNK_LENGTH - 1) // FILE_CHUNK_LENGTH
    assert expected > 1
    assert len(chunks) == expected
    assert [chunk['data']['chunk_index'] for chunk in chunks] == list(range(expected))
    assert {chunk['data']['chunk_amount'] for chunk in chunks} == {expected}
    assert all(chunk['context'] == binding.model_dump() for chunk in chunks)
    assert len({chunk['data']['file_key'] for chunk in chunks}) == 1
    return chunks[0]['data']['file_key']


@pytest.mark.parametrize(
    'profile,binding',
    [('oss_dev', LEGACY), ('oss_dev', BINDING), ('shared', BINDING)],
    ids=['legacy-oss', 'managed-oss', 'managed-shared'],
)
async def test_knowledge_file_roundtrip_reaches_plugin_original_bytes(tmp_path, monkeypatch, profile, binding):
    async with protocol_stack(tmp_path, monkeypatch, profile, binding) as stack:
        # Legacy plugin API sends no authority; Runtime supplies its trusted binding.
        result = await asyncio.wait_for(
            stack.plugin.call_action(
                PluginToRuntimeAction.GET_KNOWLEDEGE_FILE_STREAM,
                {'storage_path': 'knowledge/original.bin'},
            ),
            5,
        )
        assert await stack.plugin.read_local_file(result['file_key']) == PAYLOAD
        assert len(stack.storage_calls) == 1
        core_key = assert_chunks(stack.core_conn, binding)
        plugin_key = assert_chunks(stack.bridge_conn, binding)
        assert result['file_key'] == plugin_key != core_key
        assert not (Path(stack.control.file_storage_dir) / core_key).exists()
        assert not stack.control._owned_transfer_files
        callbacks = [
            message
            for message in stack.control_conn.sent
            if message.get('action') == PluginToRuntimeAction.GET_KNOWLEDEGE_FILE_STREAM.value
        ]
        assert len(callbacks) == 1
        assert callbacks[0]['context'] == binding.model_dump()
        assert callbacks[0]['data'] == {'storage_path': 'knowledge/original.bin'}


async def test_shared_control_rejects_legacy_chunks_before_storage(tmp_path, monkeypatch):
    async with protocol_stack(tmp_path, monkeypatch, 'shared', BINDING) as stack:
        with stack.core.installation_scope(LEGACY):
            with pytest.raises(ActionCallError, match='InstallationBinding|Legacy FILE_CHUNK'):
                await asyncio.wait_for(stack.core.send_file(PAYLOAD, ''), 5)
        assert not list(Path(stack.control.file_storage_dir).iterdir())
        assert not stack.control._owned_transfer_files


async def test_candidate_artifact_pretransfer_does_not_require_active_installation(tmp_path, monkeypatch):
    async with protocol_stack(tmp_path, monkeypatch, 'shared', BINDING) as stack:
        candidate = BINDING.model_copy(
            update={'installation_uuid': 'candidate-installation', 'runtime_revision': 1, 'artifact_digest': 'c' * 64}
        )
        assert not stack.runtime.is_current_installation_binding(candidate)
        with stack.core.installation_scope(candidate):
            key = await asyncio.wait_for(stack.core.send_file(PAYLOAD, 'lbp'), 5)
        assert_chunks(stack.core_conn, candidate)
        assert await stack.control.read_local_file(key) == PAYLOAD
        assert not stack.runtime.is_current_installation_binding(candidate)


async def test_nested_parser_target_owns_file_and_action_envelopes(tmp_path, monkeypatch):
    async with protocol_stack(tmp_path, monkeypatch, 'shared', BINDING) as stack:
        target = BINDING.model_copy(
            update={
                'installation_uuid': 'parser-installation',
                'runtime_revision': 2,
                'artifact_digest': 'b' * 64,
            }
        )
        stack.runtime.activate_installation_binding(target)
        parser_calls = []
        restored = []

        async def parse_document(author, name, context_data, file_bytes):
            parser_calls.append((stack.control.current_action_context, author, name, context_data, file_bytes))
            return {'documents': [{'text': 'parsed'}]}

        stack.runtime.plugin_mgr.parse_document = parse_document

        class ParserConnector:
            async def require_workspace_context(self, context):
                assert context.workspace_uuid == BINDING.workspace_uuid

            async def call_parser(self, plugin_name, context_data, file_bytes):
                assert plugin_name == 'tester/parser'
                assert stack.core.current_action_context == BINDING
                with stack.core.installation_scope(target):
                    result = await stack.core.parse_document('tester', 'parser', context_data, file_bytes)
                restored.append(stack.core.resolve_outbound_action_context(None))
                return result

        stack.app.plugin_connector = ParserConnector()
        result = await asyncio.wait_for(
            stack.plugin.call_action(
                PluginToRuntimeAction.INVOKE_PARSER,
                {
                    'plugin_author': 'tester',
                    'plugin_name': 'parser',
                    'storage_path': 'knowledge/original.bin',
                    'filename': 'original.bin',
                },
            ),
            5,
        )
        assert result == {'documents': [{'text': 'parsed'}]}
        key = assert_chunks(stack.core_conn, target)
        parse_requests = [
            message
            for message in stack.core_conn.sent
            if message.get('action') == LangBotToRuntimeAction.PARSE_DOCUMENT.value
        ]
        assert len(parse_requests) == 1
        assert parse_requests[0]['context'] == target.model_dump()
        assert parse_requests[0]['data']['context']['file_key'] == key
        assert parser_calls == [
            (
                target,
                'tester',
                'parser',
                {
                    'mime_type': 'application/octet-stream',
                    'filename': 'original.bin',
                    'metadata': {},
                },
                PAYLOAD,
            )
        ]
        assert restored == [BINDING]
        assert stack.core.current_action_context is None
        assert stack.core.resolve_outbound_action_context(None) is None
        assert not (Path(stack.control.file_storage_dir) / key).exists()