DocForge2/tests/test_mcp_server.py

from __future__ import annotations

import os
import shutil
import sys
import tempfile
import threading
import time
import unittest
from contextlib import contextmanager
from pathlib import Path

from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from mcp.shared.memory import create_connected_server_and_client_session

from docforge.index import ProjectIndex
from docforge.mcp_server import (
    ALL_TOOLS,
    CONTENT_WARNING,
    PROPOSAL_TOOLS,
    DocForgeService,
    _create_bound_server,
    create_server,
)
from docforge.project import Project
from docforge.viewer_manager import ViewerManager

ROOT = Path(__file__).resolve().parents[1]
FIXTURES = ROOT / "tests" / "fixtures"


class DocForgeMcpTests(unittest.IsolatedAsyncioTestCase):
    def copy_fixture(self, name: str, destination: Path) -> Path:
        root = destination / name
        shutil.copytree(FIXTURES / name, root)
        return root

    @contextmanager
    def running_manager(self, state_path: Path):
        manager = ViewerManager(state_path, check_interval_seconds=0.02)
        thread = threading.Thread(target=manager.serve_forever, daemon=True)
        previous = os.environ.get("DOCFORGE_VIEWER_MANAGER_STATE")
        os.environ["DOCFORGE_VIEWER_MANAGER_STATE"] = str(state_path)
        thread.start()
        deadline = time.monotonic() + 2
        while not state_path.exists() and time.monotonic() < deadline:
            time.sleep(0.01)
        self.assertTrue(state_path.exists())
        try:
            yield
        finally:
            manager.shutdown()
            thread.join(timeout=2)
            if previous is None:
                os.environ.pop("DOCFORGE_VIEWER_MANAGER_STATE", None)
            else:
                os.environ["DOCFORGE_VIEWER_MANAGER_STATE"] = previous

    async def test_protocol_lists_only_the_fixed_safe_surface(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = self.copy_fixture("alpha", Path(directory))
            ProjectIndex(Project.open(root)).build()
            async with create_connected_server_and_client_session(
                create_server(root), raise_exceptions=True
            ) as session:
                response = await session.list_tools()

            names = tuple(tool.name for tool in response.tools)
            self.assertEqual(ALL_TOOLS, names)
            self.assertEqual(10, len(PROPOSAL_TOOLS))
            self.assertFalse(
                any(
                    token in name
                    for name in names
                    for token in ("apply", "commit", "push", "deploy", "publish", "shell")
                )
            )

    async def test_every_read_tool_returns_scoped_structured_results(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = self.copy_fixture("alpha", Path(directory))
            ProjectIndex(Project.open(root)).build()
            calls = (
                ("docforge_project_info", {}),
                ("docforge_get_contract", {}),
                ("docforge_get_node", {"node_id": "guide.workflow"}),
                ("docforge_search", {"query": "canonical nodes", "limit": 5}),
                ("docforge_filter_nodes", {"family": "proof", "tag": "validation"}),
                ("docforge_backlinks", {"node_id": "guide.workflow"}),
                ("docforge_dependencies", {"node_id": "guide.workflow", "depth": 2}),
                ("docforge_impact", {"node_id": "guide.foundation", "depth": 2}),
                ("docforge_get_context", {"profile": "active", "budget": 180}),
                ("docforge_validate_project", {}),
                ("docforge_render_status", {}),
                ("docforge_visualize", {"node_id": "guide.workflow", "depth": 1}),
                ("docforge_stop_visualization", {}),
                ("docforge_visualization_status", {}),
            )
            with self.running_manager(Path(directory) / "viewer-manager.json"):
                service = DocForgeService(Project.open(root))
                try:
                    async with create_connected_server_and_client_session(
                        _create_bound_server(service, read_only=True), raise_exceptions=True
                    ) as session:
                        results = [
                            await session.call_tool(name, arguments) for name, arguments in calls
                        ]
                finally:
                    service.visualization.stop()

            for result in results:
                self.assertFalse(result.isError)
                self.assertIsNotNone(result.structuredContent)
                payload = result.structuredContent
                self.assertEqual("ok", payload["status"])
                self.assertEqual("alpha-docs", payload["project_id"])
                self.assertEqual(CONTENT_WARNING, payload["content_warning"])
                self.assertTrue(payload["project_root_fingerprint"])
                self.assertEqual("current", payload["staleness"])
            contract = results[1].structuredContent
            self.assertFalse(contract["canonical_writes_allowed"])
            self.assertFalse(contract["project_switching_allowed"])
            self.assertIn("canonical_writes", contract["excluded_operations"])
            self.assertIn("arbitrary_renderer_execution", contract["excluded_operations"])
            self.assertFalse(contract["isolated_changeset_writes_allowed"])
            self.assertFalse(contract["proposal_access"]["enabled"])
            self.assertTrue(results[10].structuredContent["configured"])
            self.assertEqual("stale", results[10].structuredContent["state"])
            visualization = results[11].structuredContent["visualization"]
            self.assertTrue(visualization["read_only"])
            self.assertTrue(visualization["project_bound"])
            self.assertEqual("graph-browser@8", visualization["template"])
            self.assertEqual("managed_idle", visualization["lifetime"]["policy"])
            self.assertEqual("docforge_stop_visualization", visualization["lifetime"]["stop_tool"])
            self.assertTrue(visualization["url"].startswith("http://127.0.0.1:"))
            self.assertEqual("stopped", results[12].structuredContent["state"])
            self.assertEqual("not_running", results[13].structuredContent["state"])
            context = results[8].structuredContent
            self.assertLessEqual(context["estimated_tokens"], 180)
            self.assertTrue(context["omissions"])

    async def test_missing_node_and_stale_index_are_structured_failures(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = self.copy_fixture("alpha", Path(directory))
            ProjectIndex(Project.open(root)).build()
            server = create_server(root)
            async with create_connected_server_and_client_session(
                server, raise_exceptions=True
            ) as session:
                missing = await session.call_tool(
                    "docforge_get_node", {"node_id": "research.question"}
                )
                workflow = root / "docs" / "content" / "workflow.md"
                workflow.write_text(
                    workflow.read_text(encoding="utf-8") + "\nChanged after startup.\n",
                    encoding="utf-8",
                )
                stale = await session.call_tool("docforge_get_node", {"node_id": "guide.workflow"})

            self.assertEqual("missing_node", missing.structuredContent["error"]["code"])
            self.assertEqual("stale_index", stale.structuredContent["error"]["code"])
            self.assertEqual("current", missing.structuredContent["staleness"])
            self.assertEqual("stale", stale.structuredContent["staleness"])
            self.assertTrue(missing.structuredContent["source_hash"])
            self.assertTrue(stale.structuredContent["source_hash"])
            self.assertFalse(missing.isError)
            self.assertFalse(stale.isError)

    async def test_output_limit_fails_without_returning_partial_content(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = self.copy_fixture("alpha", Path(directory))
            descriptor = root / ".docforge" / "project.toml"
            descriptor.write_text(
                descriptor.read_text(encoding="utf-8").replace(
                    "max_context_tokens = 2000",
                    "max_context_tokens = 2000\nmax_tool_output_chars = 700",
                ),
                encoding="utf-8",
            )
            ProjectIndex(Project.open(root)).build()
            async with create_connected_server_and_client_session(
                create_server(root), raise_exceptions=True
            ) as session:
                result = await session.call_tool("docforge_get_contract", {})

            payload = result.structuredContent
            self.assertEqual("error", payload["status"])
            self.assertEqual("result_too_large", payload["error"]["code"])
            self.assertNotIn("canonical_paths", payload)

    async def test_proposal_tools_use_fixed_writer_and_never_change_canonical_content(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = self.copy_fixture("alpha", Path(directory))
            project = Project.open(root)
            ProjectIndex(project).build()
            canonical_before = {
                path.relative_to(root).as_posix(): path.read_bytes()
                for path in (root / "docs/content").glob("*")
                if path.is_file()
            }
            node_hashes = {node.node_id: node.content_hash for node in project.load().nodes}
            async with create_connected_server_and_client_session(
                create_server(root, "alpha-editor"), raise_exceptions=True
            ) as session:
                contract = await session.call_tool("docforge_get_contract", {})
                created = await session.call_tool(
                    "docforge_create_changeset", {"changeset_id": "mcp-update"}
                )
                updated = await session.call_tool(
                    "docforge_propose_node_update",
                    {
                        "changeset_id": "mcp-update",
                        "expected_changeset_hash": created.structuredContent["changeset_hash"],
                        "node_id": "guide.workflow",
                        "expected_content_hash": node_hashes["guide.workflow"],
                        "metadata": {"summary": "A proposal written through MCP."},
                        "content": None,
                        "relationship_changes": [],
                        "rationale": "Prove fixed-writer isolated proposal access.",
                    },
                )
                created_node = await session.call_tool(
                    "docforge_propose_node_create",
                    {
                        "changeset_id": "mcp-update",
                        "expected_changeset_hash": updated.structuredContent["changeset_hash"],
                        "node_id": "guide.mcp-node",
                        "target_source": "docs/content/mcp-node.md",
                        "metadata": {
                            "title": "MCP proposal node",
                            "family": "guide",
                            "authority": "proposal",
                            "status": "active",
                            "tags": ["mcp", "proposal"],
                            "summary": "A node creation proposed through MCP.",
                        },
                        "content": "This node is not canonical until external integration.",
                        "relationship_changes": [],
                        "rationale": "Prove isolated MCP creation.",
                    },
                )
                moved = await session.call_tool(
                    "docforge_propose_node_move",
                    {
                        "changeset_id": "mcp-update",
                        "expected_changeset_hash": created_node.structuredContent["changeset_hash"],
                        "node_id": "guide.foundation",
                        "expected_content_hash": node_hashes["guide.foundation"],
                        "target_source": "docs/content/foundation-moved.md",
                        "rationale": "Prove isolated MCP movement.",
                    },
                )
                deleted = await session.call_tool(
                    "docforge_propose_node_delete",
                    {
                        "changeset_id": "mcp-update",
                        "expected_changeset_hash": moved.structuredContent["changeset_hash"],
                        "node_id": "proof.validation",
                        "expected_content_hash": node_hashes["proof.validation"],
                        "relationship_changes": [
                            {
                                "action": "remove",
                                "source_id": "proof.validation",
                                "relation": "proves",
                                "target_id": "guide.workflow",
                            }
                        ],
                        "rationale": "Prove isolated MCP deletion.",
                    },
                )
                validated = await session.call_tool(
                    "docforge_validate_changeset", {"changeset_id": "mcp-update"}
                )
                listed = await session.call_tool("docforge_list_changesets", {})
                inspected = await session.call_tool(
                    "docforge_get_changeset", {"changeset_id": "mcp-update"}
                )
                diff = await session.call_tool(
                    "docforge_get_changeset_diff", {"changeset_id": "mcp-update"}
                )
                preview = await session.call_tool(
                    "docforge_preview_changeset",
                    {"changeset_id": "mcp-update", "view_id": "manual"},
                )
                undeclared = await session.call_tool(
                    "docforge_preview_changeset",
                    {"changeset_id": "mcp-update", "view_id": "not-declared"},
                )

            self.assertTrue(contract.structuredContent["proposal_access"]["enabled"])
            self.assertEqual(
                "alpha-editor", contract.structuredContent["proposal_access"]["writer"]
            )
            self.assertTrue(contract.structuredContent["isolated_changeset_writes_allowed"])
            self.assertFalse(contract.structuredContent["canonical_writes_allowed"])
            self.assertEqual("alpha-editor", created.structuredContent["creator"])
            self.assertEqual(1, updated.structuredContent["operation_count"])
            self.assertEqual(4, deleted.structuredContent["operation_count"])
            self.assertTrue(validated.structuredContent["valid"])
            self.assertEqual(1, listed.structuredContent["count"])
            self.assertEqual(
                "mcp-update", listed.structuredContent["changesets"][0]["changeset_id"]
            )
            self.assertEqual("current", inspected.structuredContent["base_state"])
            self.assertEqual(4, inspected.structuredContent["operation_count"])
            self.assertEqual(
                ["update", "create", "move", "delete"],
                [change["operation"] for change in diff.structuredContent["changes"]],
            )
            self.assertEqual("current", preview.structuredContent["state"])
            self.assertEqual(
                ".docforge/previews/mcp-update/manual.html",
                preview.structuredContent["preview"]["path"],
            )
            self.assertTrue(preview.structuredContent["preview_identity"])
            self.assertEqual("error", undeclared.structuredContent["status"])
            self.assertEqual("unknown_render_view", undeclared.structuredContent["error"]["code"])
            canonical_after = {
                path.relative_to(root).as_posix(): path.read_bytes()
                for path in (root / "docs/content").glob("*")
                if path.is_file()
            }
            self.assertEqual(canonical_before, canonical_after)
            self.assertFalse((root / "docs/content/mcp-node.md").exists())
            self.assertTrue((root / ".docforge/changesets/mcp-update.json").is_file())
            self.assertTrue((root / ".docforge/previews/mcp-update/manual.html").is_file())
            self.assertFalse((root / ".docforge/rendered/manual.html").exists())

    async def test_server_without_writer_rejects_proposal_mutation_structurally(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = self.copy_fixture("alpha", Path(directory))
            ProjectIndex(Project.open(root)).build()
            async with create_connected_server_and_client_session(
                create_server(root), raise_exceptions=True
            ) as session:
                result = await session.call_tool(
                    "docforge_create_changeset", {"changeset_id": "disabled"}
                )

            self.assertFalse(result.isError)
            self.assertEqual("error", result.structuredContent["status"])
            self.assertEqual("proposal_access_disabled", result.structuredContent["error"]["code"])
            self.assertFalse((root / ".docforge/changesets/disabled.json").exists())

    async def test_stdio_transport_serves_the_same_project_bound_contract(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = self.copy_fixture("beta", Path(directory))
            ProjectIndex(Project.open(root)).build()
            parameters = StdioServerParameters(
                command=sys.executable,
                args=["-m", "docforge.mcp_server", "--project-root", str(root)],
            )
            async with (
                stdio_client(parameters) as (read, write),
                ClientSession(read, write) as session,
            ):
                await session.initialize()
                result = await session.call_tool("docforge_project_info", {})

            self.assertFalse(result.isError)
            self.assertEqual("beta-notes", result.structuredContent["project_id"])
            self.assertEqual(1, result.structuredContent["node_count"])


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