DocForge2/tests/test_adapter_launcher.py

from __future__ import annotations

import asyncio
import json
import os
import shutil
import subprocess
import sys
import tempfile
import tomllib
import unittest
import venv
from dataclasses import replace
from pathlib import Path
from unittest import mock

from jsonschema import Draft202012Validator
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from referencing import Registry, Resource

from docforge.adapter_contract import AdapterNode, AdapterProject, AdapterProjection
from docforge.adapter_launcher import AdapterLauncherV1
from docforge.client_config import (
    _validate_adapter_configuration_result,
    generate_adapter_client_configuration,
    generate_client_configuration,
)
from docforge.errors import DocForgeError
from docforge.models import Node
from docforge.project import Project
from docforge.reference_mcp import (
    REFERENCE_MCP_MODULE,
    create_reference_project,
)

ROOT = Path(__file__).resolve().parents[1]
FIXTURES = ROOT / "tests" / "fixtures"
SCHEMAS = ROOT / "schemas"
LAUNCHER_SCHEMA = json.loads((SCHEMAS / "adapter-launcher.schema.json").read_text(encoding="utf-8"))
ADAPTER_CONFIGURATION_SCHEMA = json.loads(
    (SCHEMAS / "adapter-client-configuration.schema.json").read_text(encoding="utf-8")
)
GENERIC_CONFIGURATION_SCHEMA = json.loads(
    (SCHEMAS / "client-configuration.schema.json").read_text(encoding="utf-8")
)
SCHEMA_REGISTRY = Registry().with_resources(
    (
        (LAUNCHER_SCHEMA["$id"], Resource.from_contents(LAUNCHER_SCHEMA)),
        (
            GENERIC_CONFIGURATION_SCHEMA["$id"],
            Resource.from_contents(GENERIC_CONFIGURATION_SCHEMA),
        ),
        (
            ADAPTER_CONFIGURATION_SCHEMA["$id"],
            Resource.from_contents(ADAPTER_CONFIGURATION_SCHEMA),
        ),
    )
)
ADAPTER_CONFIGURATION_VALIDATOR = Draft202012Validator(
    ADAPTER_CONFIGURATION_SCHEMA,
    registry=SCHEMA_REGISTRY,
)


class Loader:
    def __init__(self, projection: AdapterProjection) -> None:
        self.projection = projection
        self.load_calls = 0

    def load_projection(self) -> AdapterProjection:
        self.load_calls += 1
        return self.projection


class DriftingLoader(Loader):
    def __init__(
        self,
        projection: AdapterProjection,
        changed: AdapterProjection,
    ) -> None:
        super().__init__(projection)
        self.changed = changed

    def load_projection(self) -> AdapterProjection:
        self.load_calls += 1
        return self.projection if self.load_calls < 3 else self.changed


class AdapterLauncherTests(unittest.TestCase):
    def copy_fixture(self, destination: Path) -> Path:
        root = destination / "alpha"
        shutil.copytree(FIXTURES / "alpha", root)
        return root

    def copy_reference_fixture(self, destination: Path) -> Path:
        root = destination / "reference-python"
        shutil.copytree(FIXTURES / "reference-python", root)
        config = root / ".docforge" / "reference-adapter.toml"
        config.parent.mkdir(parents=True)
        config.write_text(
            "\n".join(
                (
                    "schema_version = 1",
                    'project_id = "reference-python"',
                    'title = "Runnable Python reference"',
                    'language = "python"',
                    'source_roots = ["src"]',
                    "",
                )
            ),
            encoding="utf-8",
        )
        return root.resolve()

    def projection(
        self,
        root: Path,
        *,
        source_hash: str = "3" * 64,
    ) -> AdapterProjection:
        node = Node(
            node_id="source.entry",
            title="Entry",
            family="source",
            authority="derived",
            status="active",
            tags=("python",),
            summary="Reference adapter entry.",
            content="Reference adapter content.",
            source_path="src/entry.py",
            source_anchor="L1",
            content_hash="2" * 64,
        )
        return AdapterProjection(
            project_id="adapter-client-fixture",
            title="Adapter client fixture",
            adapter_id="fixture-client",
            adapter_version="1",
            root=root,
            revision="fixture-revision",
            source_hash=source_hash,
            nodes=(AdapterNode(node),),
            edges=(),
        )

    def project(
        self,
        root: Path,
        *,
        loader: Loader | None = None,
    ) -> tuple[AdapterProject, Loader]:
        effective_loader = loader or Loader(self.projection(root))
        return (
            AdapterProject(
                effective_loader,
                cache_root=root / ".docforge-cache" / "adapter-client",
            ),
            effective_loader,
        )

    def test_launcher_and_all_client_fragments_are_deterministic_and_schema_valid(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory).resolve()
            project, _ = self.project(root)
            launcher = AdapterLauncherV1.for_project(
                project,
                module=REFERENCE_MCP_MODULE,
            )
            Draft202012Validator(LAUNCHER_SCHEMA).validate(launcher.as_dict())
            previous = os.environ.get("DOCFORGE_ADAPTER_LAUNCHER_SECRET")
            os.environ["DOCFORGE_ADAPTER_LAUNCHER_SECRET"] = "must-not-appear"
            try:
                with mock.patch("docforge.client_config.subprocess.run") as executed:
                    results: dict[str, dict[str, object]] = {}
                    for client in ("codex", "claude", "openclaw"):
                        first = generate_adapter_client_configuration(
                            project,
                            launcher,
                            client,
                            no_ast=True,
                        )
                        second = generate_adapter_client_configuration(
                            project,
                            launcher,
                            client,
                            no_ast=True,
                        )
                        self.assertEqual(first, second)
                        ADAPTER_CONFIGURATION_VALIDATOR.validate(first)
                        self.assertEqual(launcher.launcher_hash, first["launcher_hash"])
                        self.assertEqual(
                            launcher.launcher_hash,
                            first["binding"]["launcher_hash"],
                        )
                        self.assertEqual(
                            [
                                "-I",
                                "-m",
                                REFERENCE_MCP_MODULE,
                                "--project-root",
                                str(root),
                                "--capability-mode",
                                "read",
                                "--no-ast",
                            ],
                            first["binding"]["args"],
                        )
                        self.assertEqual({}, first["binding"]["environment"])
                        self.assertNotIn("cwd", first["binding"])
                        self.assertNotIn(
                            "must-not-appear",
                            json.dumps(first, sort_keys=True),
                        )
                        results[client] = first
                    executed.assert_not_called()
            finally:
                if previous is None:
                    os.environ.pop("DOCFORGE_ADAPTER_LAUNCHER_SECRET", None)
                else:
                    os.environ["DOCFORGE_ADAPTER_LAUNCHER_SECRET"] = previous

            self.assertIn(
                "mcp_servers",
                tomllib.loads(results["codex"]["artifact"]["content"]),
            )
            self.assertIn(
                "mcpServers",
                json.loads(results["claude"]["artifact"]["content"]),
            )
            self.assertIn(
                "mcp",
                json.loads(results["openclaw"]["artifact"]["content"]),
            )

    def test_launcher_rejects_module_path_entry_point_and_argument_injection(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory).resolve()
            project, _ = self.project(root)
            launcher = AdapterLauncherV1.for_project(
                project,
                module=REFERENCE_MCP_MODULE,
            )
            for module in (
                "-m",
                "fixture_adapter.mcp_server",
                "fixture_adapter.mcp_server --debug",
                "fixture_adapter:mcp_server",
                "fixture_adapter/mcp_server",
                ".fixture_adapter",
                "docforge.mcp_server",
                "fixture_adapter.$server",
            ):
                with self.subTest(module=module), self.assertRaises(DocForgeError) as captured:
                    replace(launcher, module=module)
                self.assertEqual("invalid_adapter_launcher", captured.exception.code)

            with self.assertRaises(DocForgeError):
                replace(launcher, project_root=Path("relative/project"))
            with self.assertRaises(DocForgeError):
                replace(launcher, entry_point="console-script")  # type: ignore[arg-type]
            payload = launcher.as_dict()
            payload["arguments"] = ["--shell", "command"]
            with self.assertRaises(TypeError):
                AdapterLauncherV1(**payload)  # type: ignore[arg-type]

            result = generate_adapter_client_configuration(project, launcher, "codex")
            result["binding"]["args"].append("--arbitrary")
            with self.assertRaisesRegex(AssertionError, "Generated adapter client"):
                _validate_adapter_configuration_result(
                    result,
                    project=project,
                    launcher=launcher,
                    source_availability=project_source_availability(result),
                )

    def test_wrong_project_descriptor_and_source_drift_fail_closed(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            parent = Path(directory)
            first_root = (parent / "first").resolve()
            second_root = (parent / "second").resolve()
            first_root.mkdir()
            second_root.mkdir()
            first, _ = self.project(first_root)
            second, _ = self.project(second_root)
            launcher = AdapterLauncherV1.for_project(
                first,
                module=REFERENCE_MCP_MODULE,
            )

            with self.assertRaises(DocForgeError) as wrong_project:
                generate_adapter_client_configuration(second, launcher, "codex")
            self.assertEqual("adapter_launcher_mismatch", wrong_project.exception.code)
            with self.assertRaises(DocForgeError) as descriptor_drift:
                generate_adapter_client_configuration(
                    first,
                    replace(launcher, descriptor_hash="0" * 64),
                    "codex",
                )
            self.assertEqual("adapter_launcher_mismatch", descriptor_drift.exception.code)

            base = self.projection(first_root)
            changed = replace(base, source_hash="4" * 64)
            drifting, _ = self.project(
                first_root,
                loader=DriftingLoader(base, changed),
            )
            drifting_launcher = AdapterLauncherV1.for_project(
                drifting,
                module=REFERENCE_MCP_MODULE,
            )
            with self.assertRaises(DocForgeError) as source_drift:
                generate_adapter_client_configuration(
                    drifting,
                    drifting_launcher,
                    "openclaw",
                )
            self.assertEqual("source_changed", source_drift.exception.code)

    def test_generic_configuration_remains_separate_and_compatible(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = self.copy_fixture(Path(directory))
            generic = Project.open(root)
            result = generate_client_configuration(generic, "codex", no_ast=True)
            Draft202012Validator(GENERIC_CONFIGURATION_SCHEMA).validate(result)
            self.assertEqual("client.configure", result["operation"])
            self.assertNotIn("launcher", result)
            with self.assertRaises(DocForgeError) as generic_launcher:
                AdapterLauncherV1.for_project(
                    generic,
                    module=REFERENCE_MCP_MODULE,
                )
            self.assertEqual("adapter_launcher_unavailable", generic_launcher.exception.code)

            custom_root = (Path(directory) / "custom").resolve()
            custom_root.mkdir()
            custom, _ = self.project(custom_root)
            with self.assertRaises(DocForgeError) as generic_api:
                generate_client_configuration(custom, "codex")
            self.assertEqual("missing_config", generic_api.exception.code)

    def test_installed_project_owned_top_level_module_launches_without_probe_execution(
        self,
    ) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory).resolve()
            environment = root / ".launcher-venv"
            venv.EnvBuilder(with_pip=False).create(environment)
            executable = environment / "bin" / "python"
            site_packages = Path(
                subprocess.run(
                    [
                        str(executable),
                        "-I",
                        "-c",
                        "import site; print(site.getsitepackages()[0])",
                    ],
                    check=True,
                    capture_output=True,
                    text=True,
                ).stdout.strip()
            )
            (site_packages / "fixture-adapter.pth").write_text(
                f"{root}\n",
                encoding="utf-8",
            )
            sentinel = root / "fixture-adapter-executed"
            (root / "fixture_adapter.py").write_text(
                "\n".join(
                    (
                        "import argparse",
                        "import json",
                        "from pathlib import Path",
                        "parser = argparse.ArgumentParser()",
                        'parser.add_argument("--project-root", required=True)',
                        (
                            'parser.add_argument("--capability-mode", '
                            'choices=("read",), required=True)'
                        ),
                        "arguments = parser.parse_args()",
                        f"Path({str(sentinel)!r}).write_text('executed', encoding='utf-8')",
                        (
                            "print(json.dumps({'project_root': arguments.project_root}, "
                            "sort_keys=True))"
                        ),
                        "",
                    )
                ),
                encoding="utf-8",
            )

            project, _ = self.project(root)
            with mock.patch.object(sys, "executable", str(executable)):
                launcher = AdapterLauncherV1.for_project(
                    project,
                    module="fixture_adapter",
                )
                result = generate_adapter_client_configuration(
                    project,
                    launcher,
                    "codex",
                )
            self.assertFalse(sentinel.exists())
            Draft202012Validator(LAUNCHER_SCHEMA).validate(launcher.as_dict())
            binding = result["binding"]
            completed = subprocess.run(
                [binding["command"], *binding["args"]],
                check=True,
                capture_output=True,
                text=True,
                timeout=10,
            )
            self.assertEqual({"project_root": str(root)}, json.loads(completed.stdout))
            self.assertEqual("executed", sentinel.read_text(encoding="utf-8"))

    def test_uninstalled_packages_and_non_project_modules_fail_closed(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory).resolve()
            project, _ = self.project(root)
            for module, code in (
                ("fixture_adapter", "adapter_launcher_unavailable"),
                ("os", "adapter_launcher_unavailable"),
                ("json", "invalid_adapter_launcher"),
            ):
                with self.subTest(module=module), self.assertRaises(DocForgeError) as captured:
                    AdapterLauncherV1.for_project(project, module=module)
                self.assertEqual(code, captured.exception.code)

    def test_generated_reference_binding_starts_a_real_isolated_mcp_process(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = self.copy_reference_fixture(Path(directory))
            project = create_reference_project(root)
            launcher = AdapterLauncherV1.for_project(
                project,
                module=REFERENCE_MCP_MODULE,
            )
            result = generate_adapter_client_configuration(
                project,
                launcher,
                "codex",
            )
            binding = result["binding"]

            async def inspect() -> tuple[str, ...]:
                parameters = StdioServerParameters(
                    command=binding["command"],
                    args=binding["args"],
                    env=binding["environment"],
                )
                async with (
                    stdio_client(parameters) as streams,
                    ClientSession(*streams) as session,
                ):
                    await session.initialize()
                    return tuple(tool.name for tool in (await session.list_tools()).tools)

            tools = asyncio.run(inspect())
            self.assertIn("docforge_bootstrap", tools)


def project_source_availability(result: dict[str, object]):
    from docforge.adapter_launcher import AdapterSourceAvailabilityV1

    payload = result["source_availability"]
    return AdapterSourceAvailabilityV1(
        schema_version=payload["schema_version"],
        status=payload["status"],
        method=payload["method"],
        revision=payload["revision"],
        source_hash=payload["source_hash"],
    )


if __name__ == "__main__":
    unittest.main()