obsidian-wiki/tests/test_mcp_server_remote.py

import asyncio
import hashlib
import json
import os
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch

from obsidian_wiki import mcp_server
from obsidian_wiki.mcp_ingest import IngestManager


class FakeApp:
    def __init__(self):
        self.called = False

    async def __call__(self, scope, receive, send):
        self.called = True
        await send({"type": "http.response.start", "status": 200, "headers": []})
        await send({"type": "http.response.body", "body": b"ok"})


def run_asgi_app(app, *, headers=None, scope_type="http"):
    messages = []
    scope = {"type": scope_type, "headers": headers or []}

    async def receive():
        return {"type": "http.request", "body": b"", "more_body": False}

    async def send(message):
        messages.append(message)

    asyncio.run(app(scope, receive, send))
    return messages


def run_asgi_http_request(app, *, method="GET", path="/", headers=None, body=b""):
    messages = []
    sent = False
    scope = {
        "type": "http",
        "method": method,
        "path": path,
        "raw_path": path.encode("ascii"),
        "query_string": b"",
        "headers": headers or [],
    }

    async def receive():
        nonlocal sent
        if sent:
            return {"type": "http.disconnect"}
        sent = True
        return {"type": "http.request", "body": body, "more_body": False}

    async def send(message):
        messages.append(message)

    asyncio.run(app(scope, receive, send))
    return messages


def multipart_body(fields, files, boundary="test-boundary"):
    chunks = []
    for name, value in fields.items():
        chunks.append(f"--{boundary}\r\n".encode("ascii"))
        chunks.append(f'Content-Disposition: form-data; name="{name}"\r\n\r\n'.encode("ascii"))
        chunks.append(str(value).encode("utf-8"))
        chunks.append(b"\r\n")
    for field_name, file_name, content_type, payload in files:
        chunks.append(f"--{boundary}\r\n".encode("ascii"))
        chunks.append(
            (
                f'Content-Disposition: form-data; name="{field_name}"; filename="{file_name}"\r\n'
                f"Content-Type: {content_type}\r\n\r\n"
            ).encode("ascii")
        )
        chunks.append(payload)
        chunks.append(b"\r\n")
    chunks.append(f"--{boundary}--\r\n".encode("ascii"))
    return b"".join(chunks), boundary


class RemoteMcpServerTests(unittest.TestCase):
    def test_bearer_token_middleware_accepts_matching_token(self):
        inner = FakeApp()
        middleware = mcp_server.BearerTokenAuthMiddleware(inner, "secret-token")

        messages = run_asgi_app(
            middleware,
            headers=[(b"authorization", b"Bearer secret-token")],
        )

        self.assertTrue(inner.called)
        self.assertEqual(messages[0]["status"], 200)

    def test_bearer_token_middleware_rejects_missing_or_wrong_token(self):
        inner = FakeApp()
        middleware = mcp_server.BearerTokenAuthMiddleware(inner, "secret-token")

        messages = run_asgi_app(
            middleware,
            headers=[(b"authorization", b"Bearer wrong-token")],
        )

        self.assertFalse(inner.called)
        self.assertEqual(messages[0]["status"], 401)
        self.assertIn((b"www-authenticate", b"Bearer"), messages[0]["headers"])

    def test_bearer_token_middleware_requires_non_empty_token(self):
        with self.assertRaises(ValueError):
            mcp_server.BearerTokenAuthMiddleware(FakeApp(), "")

    def test_read_token_from_env_trims_and_requires_value(self):
        with patch.dict(os.environ, {"TEST_MCP_TOKEN": "  secret-token  "}, clear=False):
            self.assertEqual(mcp_server.read_token_from_env("TEST_MCP_TOKEN"), "secret-token")

        with patch.dict(os.environ, {}, clear=True):
            with self.assertRaises(RuntimeError):
                mcp_server.read_token_from_env("TEST_MCP_TOKEN")

    def test_runtime_config_reads_environment_with_defaults(self):
        with patch.dict(os.environ, {"OBSIDIAN_MCP_HOST": " 0.0.0.0 ", "OBSIDIAN_MCP_PORT": "9876"}, clear=True):
            self.assertEqual(mcp_server.runtime_config_value("OBSIDIAN_MCP_HOST", "127.0.0.1"), "0.0.0.0")
            self.assertEqual(mcp_server.runtime_config_int("OBSIDIAN_MCP_PORT", 8765), 9876)

        with patch.dict(os.environ, {"OBSIDIAN_MCP_PORT": "not-an-int"}, clear=True):
            self.assertEqual(mcp_server.runtime_config_int("OBSIDIAN_MCP_PORT", 8765), 8765)

    def test_enabled_tools_config_prefers_tool_list_over_legacy_ingest_flag(self):
        with patch.dict(
            os.environ,
            {
                "OBSIDIAN_MCP_ENABLED_TOOLS": "file-assets",
                "OBSIDIAN_MCP_ENABLE_INGEST": "1",
            },
            clear=True,
        ):
            self.assertEqual(mcp_server.enabled_tool_names(), {"file-assets"})

        with patch.dict(os.environ, {"OBSIDIAN_MCP_ENABLE_INGEST": "1"}, clear=True):
            self.assertEqual(mcp_server.enabled_tool_names(), {"file-assets", "ingest"})

        with patch.dict(os.environ, {"OBSIDIAN_MCP_ENABLED_TOOLS": "file-assets,query"}, clear=True):
            self.assertEqual(mcp_server.enabled_tool_names(), {"file-assets", "query"})

    def test_create_http_app_wraps_streamable_http_app_with_auth(self):
        class FakeServer:
            def __init__(self):
                self.inner = FakeApp()

            def streamable_http_app(self):
                return self.inner

        fake_server = FakeServer()
        app = mcp_server.create_http_app(fake_server, "secret-token")

        messages = run_asgi_app(
            app,
            headers=[(b"authorization", b"Bearer secret-token")],
        )

        self.assertTrue(fake_server.inner.called)
        self.assertEqual(messages[0]["status"], 200)

    def test_http_upload_endpoint_accepts_multipart_when_ingest_enabled(self):
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            manager = IngestManager(
                vault_path=root / "vault",
                project_root=root / "project",
                runtime_root=root / "runtime",
            )
            app = mcp_server.create_http_app(FakeApp(), "secret-token", ingest_manager=manager)
            payload = b"uploaded drawing"
            body, boundary = multipart_body(
                {"project": "bjd", "doc_type_hint": "drawing", "start_job": "false"},
                [("files", "drawing.pdf", "application/pdf", payload)],
            )

            messages = run_asgi_http_request(
                app,
                method="POST",
                path="/mcp-ingest/uploads",
                headers=[
                    (b"authorization", b"Bearer secret-token"),
                    (b"content-type", f"multipart/form-data; boundary={boundary}".encode("ascii")),
                ],
                body=body,
            )

        self.assertEqual(messages[0]["status"], 200)
        response = b"".join(message.get("body", b"") for message in messages if message["type"] == "http.response.body")
        data = json.loads(response.decode("utf-8"))
        self.assertEqual(data["job"], None)
        self.assertEqual(data["uploads"][0]["state"], "complete")
        self.assertEqual(data["uploads"][0]["file_name"], "drawing.pdf")
        self.assertEqual(data["uploads"][0]["sha256"], "sha256:" + hashlib.sha256(payload).hexdigest())

    def test_http_upload_endpoint_accepts_upload_ticket_without_main_bearer_token(self):
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            manager = IngestManager(
                vault_path=root / "vault",
                project_root=root / "project",
                runtime_root=root / "runtime",
            )
            ticket = manager.create_upload_ticket(
                expires_seconds=60,
                max_files=1,
                project="bjd",
                doc_type_hint="drawing",
                start_job=False,
            )
            app = mcp_server.create_http_app(FakeApp(), "secret-token", ingest_manager=manager)
            payload = b"ticket uploaded drawing"
            body, boundary = multipart_body(
                {},
                [("files", "ticket-drawing.pdf", "application/pdf", payload)],
            )

            messages = run_asgi_http_request(
                app,
                method="POST",
                path="/mcp-ingest/uploads",
                headers=[
                    (b"x-obsidian-wiki-upload-token", ticket["upload_token"].encode("ascii")),
                    (b"content-type", f"multipart/form-data; boundary={boundary}".encode("ascii")),
                ],
                body=body,
            )

            reused = run_asgi_http_request(
                app,
                method="POST",
                path="/mcp-ingest/uploads",
                headers=[
                    (b"x-obsidian-wiki-upload-token", ticket["upload_token"].encode("ascii")),
                    (b"content-type", f"multipart/form-data; boundary={boundary}".encode("ascii")),
                ],
                body=body,
            )

        self.assertEqual(messages[0]["status"], 200)
        response = b"".join(message.get("body", b"") for message in messages if message["type"] == "http.response.body")
        data = json.loads(response.decode("utf-8"))
        self.assertEqual(data["job"], None)
        self.assertEqual(data["uploads"][0]["project"], "bjd")
        self.assertEqual(data["uploads"][0]["doc_type_hint"], "drawing")
        self.assertEqual(data["uploads"][0]["sha256"], "sha256:" + hashlib.sha256(payload).hexdigest())
        self.assertEqual(reused[0]["status"], 401)

    def test_http_upload_endpoint_requires_ingest_and_bearer_token(self):
        no_ingest_app = mcp_server.create_http_app(FakeApp(), "secret-token")
        missing_token = run_asgi_http_request(no_ingest_app, method="POST", path="/mcp-ingest/uploads")
        self.assertEqual(missing_token[0]["status"], 401)

        disabled = run_asgi_http_request(
            no_ingest_app,
            method="POST",
            path="/mcp-ingest/uploads",
            headers=[(b"authorization", b"Bearer secret-token")],
        )
        self.assertEqual(disabled[0]["status"], 404)

    def test_create_server_uses_file_asset_service_name(self):
        class FakeMCP:
            names = []
            tool_names = []
            prompt_names = []

            def __init__(self, name, **kwargs):
                self.names.append(name)

            def tool(self):
                def decorator(func):
                    self.tool_names.append(func.__name__)
                    return func

                return decorator

            def resource(self, *_args, **_kwargs):
                def decorator(func):
                    return func

                return decorator

            def prompt(self, name=None, **_kwargs):
                def decorator(func):
                    self.prompt_names.append(name or func.__name__)
                    return func

                return decorator

        with (
            tempfile.TemporaryDirectory() as tmp,
            patch.dict("sys.modules", {"mcp.server.fastmcp": type("Mod", (), {"FastMCP": FakeMCP})()}),
        ):
            server = mcp_server.create_server(vault_path=Path(tmp))

        self.assertIsInstance(server, FakeMCP)
        self.assertEqual(FakeMCP.names[-1], "wiki-file-assets")
        self.assertIn("search_file_assets", FakeMCP.tool_names)
        self.assertIn("list_file_assets", FakeMCP.tool_names)
        self.assertIn("get_file_asset_metadata", FakeMCP.tool_names)
        self.assertIn("how_to_use_wiki_file_assets", FakeMCP.prompt_names)
        self.assertNotIn("create_ingest_upload", FakeMCP.tool_names)

    def test_create_server_registers_ingest_tools_when_enabled(self):
        class FakeMCP:
            tool_names = []

            def __init__(self, _name, **_kwargs):
                pass

            def tool(self):
                def decorator(func):
                    self.tool_names.append(func.__name__)
                    return func

                return decorator

            def resource(self, *_args, **_kwargs):
                def decorator(func):
                    return func

                return decorator

            def prompt(self, *_args, **_kwargs):
                def decorator(func):
                    return func

                return decorator

        with (
            tempfile.TemporaryDirectory() as tmp,
            patch.dict("sys.modules", {"mcp.server.fastmcp": type("Mod", (), {"FastMCP": FakeMCP})()}),
        ):
            root = Path(tmp)
            server = mcp_server.create_server(
                vault_path=root / "vault",
                enable_ingest=True,
                project_root=root / "project",
                ingest_runtime_root=root / "runtime",
            )

        self.assertIsInstance(server, FakeMCP)
        self.assertNotIn("create_ingest_upload", FakeMCP.tool_names)
        self.assertNotIn("append_ingest_upload_chunk", FakeMCP.tool_names)
        self.assertNotIn("complete_ingest_upload", FakeMCP.tool_names)
        self.assertIn("start_ingest_job", FakeMCP.tool_names)
        self.assertIn("get_ingest_job_status", FakeMCP.tool_names)
        self.assertIn("get_ingest_job_result", FakeMCP.tool_names)
        self.assertIn("cancel_ingest_job", FakeMCP.tool_names)
        self.assertIn("check_ingest_opencode_readiness", FakeMCP.tool_names)
        self.assertIn("get_ingest_upload_instructions", FakeMCP.tool_names)
        self.assertIn("create_ingest_upload_ticket", FakeMCP.tool_names)
        self.assertNotIn("ingest_file_base64", FakeMCP.tool_names)

    def test_create_server_registers_query_tools_only_when_enabled(self):
        class FakeMCP:
            tool_names = []

            def __init__(self, _name, **_kwargs):
                pass

            def tool(self):
                def decorator(func):
                    self.tool_names.append(func.__name__)
                    return func

                return decorator

            def resource(self, *_args, **_kwargs):
                def decorator(func):
                    return func

                return decorator

            def prompt(self, *_args, **_kwargs):
                def decorator(func):
                    return func

                return decorator

        class FakeQueryManager:
            def start_query_job(self, query, mode="normal", public_only=False, note=None):
                return {"job_id": "query-job", "state": "queued"}

            def get_query_job_status(self, job_id):
                return {"job_id": job_id, "state": "running"}

            def get_query_job_result(self, job_id):
                return {"job_id": job_id, "state": "succeeded", "answer": "answer"}

            def cancel_query_job(self, job_id):
                return {"job_id": job_id, "state": "cancelled"}

            def check_opencode_readiness(self, *, run_smoke_test=False):
                return {"ok": True, "runtime_smoke_test": "skipped"}

        with (
            tempfile.TemporaryDirectory() as tmp,
            patch.dict("sys.modules", {"mcp.server.fastmcp": type("Mod", (), {"FastMCP": FakeMCP})()}),
        ):
            root = Path(tmp)
            disabled_server = mcp_server.create_server(vault_path=root / "vault")
            disabled_tool_names = list(FakeMCP.tool_names)
            FakeMCP.tool_names.clear()
            enabled_server = mcp_server.create_server(
                vault_path=root / "vault",
                enable_query=True,
                project_root=root / "project",
                query_runtime_root=root / "runtime",
                query_manager=FakeQueryManager(),
            )

        self.assertIsInstance(disabled_server, FakeMCP)
        self.assertIsInstance(enabled_server, FakeMCP)
        self.assertNotIn("start_wiki_query", disabled_tool_names)
        self.assertIn("start_wiki_query", FakeMCP.tool_names)
        self.assertIn("get_wiki_query_status", FakeMCP.tool_names)
        self.assertIn("get_wiki_query_result", FakeMCP.tool_names)
        self.assertIn("cancel_wiki_query", FakeMCP.tool_names)
        self.assertIn("check_query_opencode_readiness", FakeMCP.tool_names)

    def test_file_index_status_reports_opencode_model_settings(self):
        class FakeMCP:
            tools = {}

            def __init__(self, _name, **_kwargs):
                pass

            def tool(self):
                def decorator(func):
                    self.tools[func.__name__] = func
                    return func

                return decorator

            def resource(self, *_args, **_kwargs):
                def decorator(func):
                    return func

                return decorator

            def prompt(self, *_args, **_kwargs):
                def decorator(func):
                    return func

                return decorator

        with (
            tempfile.TemporaryDirectory() as tmp,
            patch.dict("sys.modules", {"mcp.server.fastmcp": type("Mod", (), {"FastMCP": FakeMCP})()}),
        ):
            root = Path(tmp)
            mcp_server.create_server(
                vault_path=root / "vault",
                enable_ingest=True,
                enable_query=True,
                project_root=root / "project",
                ingest_runtime_root=root / "ingest-runtime",
                query_runtime_root=root / "query-runtime",
                ingest_opencode_model="openai/gpt-5.5",
                ingest_opencode_variant="high",
                query_opencode_model="openai/gpt-5.5",
                query_opencode_variant="minimal",
                ingest_manager=IngestManager(
                    vault_path=root / "vault",
                    project_root=root / "project",
                    runtime_root=root / "fake-ingest",
                ),
                query_manager=type(
                    "FakeQueryManager",
                    (),
                    {
                        "start_query_job": lambda self, query, mode="normal", public_only=False, note=None: {},
                        "get_query_job_status": lambda self, job_id: {},
                        "get_query_job_result": lambda self, job_id: {},
                        "cancel_query_job": lambda self, job_id: {},
                        "check_opencode_readiness": lambda self, run_smoke_test=False: {},
                    },
                )(),
            )
            status = FakeMCP.tools["file_index_status"]()

        self.assertEqual(status["enabled_tools"], ["file-assets", "ingest", "query"])
        self.assertEqual(status["ingest_opencode_model"], "openai/gpt-5.5")
        self.assertEqual(status["ingest_opencode_variant"], "high")
        self.assertEqual(status["query_opencode_model"], "openai/gpt-5.5")
        self.assertEqual(status["query_opencode_variant"], "minimal")

    def test_create_ingest_upload_ticket_tool_returns_http_only_token(self):
        class FakeMCP:
            tools = {}

            def __init__(self, _name, **_kwargs):
                pass

            def tool(self):
                def decorator(func):
                    self.tools[func.__name__] = func
                    return func

                return decorator

            def resource(self, *_args, **_kwargs):
                def decorator(func):
                    return func

                return decorator

            def prompt(self, *_args, **_kwargs):
                def decorator(func):
                    return func

                return decorator

        with (
            tempfile.TemporaryDirectory() as tmp,
            patch.dict("sys.modules", {"mcp.server.fastmcp": type("Mod", (), {"FastMCP": FakeMCP})()}),
        ):
            root = Path(tmp)
            mcp_server.create_server(
                vault_path=root / "vault",
                enable_ingest=True,
                project_root=root / "project",
                ingest_runtime_root=root / "runtime",
            )
            ticket = FakeMCP.tools["create_ingest_upload_ticket"](
                expires_seconds=60,
                max_files=1,
                project="bjd",
                doc_type_hint="drawing",
                start_job=False,
            )

        self.assertIn("upload_token", ticket)
        self.assertEqual(ticket["token_header"], "X-Obsidian-Wiki-Upload-Token")
        self.assertEqual(ticket["upload_path"], "/mcp-ingest/uploads")
        self.assertEqual(ticket["max_files"], 1)
        self.assertFalse(ticket["start_job"])

    def test_create_server_can_disable_file_asset_tools(self):
        class FakeMCP:
            tool_names = []
            prompt_names = []

            def __init__(self, _name, **_kwargs):
                pass

            def tool(self):
                def decorator(func):
                    self.tool_names.append(func.__name__)
                    return func

                return decorator

            def resource(self, *_args, **_kwargs):
                def decorator(func):
                    return func

                return decorator

            def prompt(self, name=None, **_kwargs):
                def decorator(func):
                    self.prompt_names.append(name or func.__name__)
                    return func

                return decorator

        with (
            tempfile.TemporaryDirectory() as tmp,
            patch.dict("sys.modules", {"mcp.server.fastmcp": type("Mod", (), {"FastMCP": FakeMCP})()}),
        ):
            server = mcp_server.create_server(vault_path=Path(tmp), enable_file_assets=False)

        self.assertIsInstance(server, FakeMCP)
        self.assertIn("file_index_status", FakeMCP.tool_names)
        self.assertNotIn("search_file_assets", FakeMCP.tool_names)
        self.assertNotIn("list_file_assets", FakeMCP.tool_names)
        self.assertNotIn("get_file_asset_metadata", FakeMCP.tool_names)
        self.assertNotIn("how_to_use_wiki_file_assets", FakeMCP.prompt_names)

    def test_run_http_server_fails_before_server_creation_without_token(self):
        with tempfile.TemporaryDirectory() as tmp:
            with patch.dict(os.environ, {}, clear=True):
                with patch.object(mcp_server, "create_server") as create_server:
                    with self.assertRaises(RuntimeError):
                        mcp_server.run_http_server(
                            vault_path=Path(tmp),
                            db_path=None,
                            profile_name="bjd-general-bot",
                            max_file_bytes=1024,
                            expose_paths=False,
                            host="127.0.0.1",
                            port=8765,
                            token_env="TEST_MCP_TOKEN",
                            streamable_http_path="/mcp",
                            allowed_hosts=None,
                            allowed_origins=None,
                        )
                    create_server.assert_not_called()


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