path: root/tests/test_mcp.py

"""Tests for MCP client integration and settings API."""

from __future__ import annotations

import asyncio
import base64
import sys
import tempfile
from pathlib import Path

import pytest
from mcp.types import (
    AudioContent,
    BlobResourceContents,
    CallToolResult,
    EmbeddedResource,
    ImageContent,
)
from starlette.testclient import TestClient

from ollama_web.app import create_app
from ollama_web.config import settings
from ollama_web.mcp import (
    _clean_schema,
    _normalized_tool_name,
    _parse_tool_name,
    _parse_tool_result,
    load_mcp_config,
    make_mcp_executor,
    save_mcp_config,
)
from ollama_web.sessions import SessionStore
from ollama_web.tools.registry import ToolArtifact, ToolExecutionResult, ToolRegistry


@pytest.fixture
def client(monkeypatch):
    monkeypatch.setattr(settings, "pin", "123456")
    monkeypatch.setattr(settings, "secret_key", "test-secret")
    monkeypatch.setattr(settings, "allowed_origins", [])
    monkeypatch.setattr(settings, "mcp_stdio_allowlist", [sys.executable])
    monkeypatch.setattr(settings, "mcp_https_allowlist", [])
    with tempfile.TemporaryDirectory() as tmp:
        monkeypatch.setattr(settings, "data_dir", tmp)
        client = TestClient(create_app())
        login = client.post("/api/auth/login", json={"pin": "123456"})
        assert login.status_code == 200
        client.csrf = login.headers["x-csrf-token"]
        yield client


def test_load_save_mcp_config():
    with tempfile.TemporaryDirectory() as tmp:
        config = {
            "mcpServers": {
                "fs": {
                    "command": "npx",
                    "args": ["-y", "@modelcontextprotocol/server-filesystem", "/tmp"],
                },
                "remote": {
                    "url": "https://example.com/mcp",
                    "headers": {"Authorization": "Bearer x"},
                },
            }
        }
        save_mcp_config(tmp, config)
        loaded = load_mcp_config(tmp)
        assert loaded == config
        assert (Path(tmp) / "mcpServers.json").exists()


def test_load_mcp_config_missing_file():
    with tempfile.TemporaryDirectory() as tmp:
        loaded = load_mcp_config(tmp)
        assert loaded == {"mcpServers": {}}


def test_mcp_config_api(client):
    res = client.get("/api/mcp/servers")
    assert res.status_code == 200
    assert res.json() == {"mcpServers": {}}

    payload = {
        "mcpServers": {
            "fs": {
                "command": sys.executable,
                "args": ["mcp_servers/calc_server.py"],
            },
            "remote": {"url": "http://127.0.0.1:9000/mcp", "timeout": 60},
        }
    }
    res = client.put(
        "/api/mcp/servers",
        json=payload,
        headers={"X-CSRF-Token": client.csrf},
    )
    assert res.status_code == 200
    data = res.json()
    assert data["mcpServers"]["fs"]["command"] == sys.executable
    assert data["mcpServers"]["remote"]["timeout"] == 60.0

    res = client.get("/api/mcp/servers")
    assert res.json() == data


def test_mcp_config_api_rejects_unknown_fields(client):
    payload = {"mcpServers": {"fs": {"command": "npx", "invalid": 1}}}
    res = client.put(
        "/api/mcp/servers",
        json=payload,
        headers={"X-CSRF-Token": client.csrf},
    )
    assert res.status_code == 400


def test_mcp_config_api_requires_transport(client):
    payload = {"mcpServers": {"fs": {}}}
    res = client.put(
        "/api/mcp/servers",
        json=payload,
        headers={"X-CSRF-Token": client.csrf},
    )
    assert res.status_code == 400


def test_normalized_tool_name():
    assert _normalized_tool_name("server1", "tool_a") == "mcp__server1__tool_a"


def test_parse_tool_name():
    assert _parse_tool_name("mcp__server1__tool_a") == ("server1", "tool_a")
    assert _parse_tool_name("web_search") is None


def test_clean_schema():
    schema = {
        "type": "object",
        "properties": {"path": {"type": "string"}},
        "required": ["path"],
    }
    cleaned = _clean_schema(schema)
    assert cleaned["type"] == "object"
    assert cleaned["required"] == ["path"]


def test_tool_registry_mcp_tool():
    reg = ToolRegistry()

    def executor(full_name, arguments=""):
        return "ok"

    reg.register_mcp_tool(
        "mcp__s__t", {"type": "function", "function": {"name": "mcp__s__t"}}, executor
    )
    assert "mcp__s__t" in reg.names
    # Executor is registered internally for execution but must not be exposed
    # to ollama as a callable (ollama needs the dict definition only).
    assert len(reg.ollama_tools) == 1
    result = reg.execute("mcp__s__t", {})
    assert result == "ok"


def test_mcp_image_and_audio_are_preserved_as_artifacts_not_text():
    image_data = b"png-bytes"
    audio_data = b"mp3-bytes"
    result = CallToolResult(
        content=[
            ImageContent(
                type="image",
                data=base64.b64encode(image_data).decode("ascii"),
                mimeType="image/png",
            ),
            AudioContent(
                type="audio",
                data=base64.b64encode(audio_data).decode("ascii"),
                mimeType="audio/mpeg",
            ),
        ]
    )

    parsed = _parse_tool_result(result, tool_name="generate")

    assert [artifact.kind for artifact in parsed.artifacts] == ["image", "audio"]
    assert parsed.artifacts[0].data == image_data
    assert parsed.artifacts[1].data == audio_data
    assert base64.b64encode(image_data).decode("ascii") not in parsed.text
    assert base64.b64encode(audio_data).decode("ascii") not in parsed.text


def test_mcp_blob_resource_does_not_leak_base64_into_text():
    encoded = base64.b64encode(b"audio-bytes").decode("ascii")
    result = CallToolResult(
        content=[
            EmbeddedResource(
                type="resource",
                resource=BlobResourceContents(
                    uri="file:///generated.mp3",
                    mimeType="audio/mpeg",
                    blob=encoded,
                ),
            )
        ]
    )

    parsed = _parse_tool_result(result, tool_name="music")

    assert len(parsed.artifacts) == 1
    assert parsed.artifacts[0].kind == "audio"
    assert encoded not in parsed.text


def test_mcp_artifact_audience_is_preserved():
    encoded = base64.b64encode(b"image-bytes").decode("ascii")
    parsed = _parse_tool_result(
        {
            "content": [
                {
                    "type": "image",
                    "data": encoded,
                    "mimeType": "image/png",
                    "annotations": {"audience": ["assistant"]},
                }
            ]
        }
    )

    assert len(parsed.artifacts) == 1
    assert parsed.artifacts[0].for_user is False
    assert parsed.artifacts[0].for_assistant is True


def test_mcp_executor_runs_artifact_handler_without_flattening(monkeypatch):
    from ollama_web import mcp as mcp_mod

    artifact = ToolArtifact(
        kind="audio",
        mime="audio/mpeg",
        name="music.mp3",
        data=b"audio",
    )
    detailed = ToolExecutionResult("result text", [artifact])
    monkeypatch.setattr(mcp_mod, "call_mcp_tool_sync", lambda *_args, **_kwargs: detailed)
    seen = []
    executor = make_mcp_executor(
        {"mcpServers": {}},
        artifact_handler=lambda name, item: seen.append((name, item)),
    )

    returned = executor("mcp__music__generate", {})

    assert returned is detailed
    assert seen == [("mcp__music__generate", artifact)]


def test_chat_registry_persists_only_user_visible_mcp_artifacts(tmp_path, monkeypatch):
    from ollama_web.routes import chat as chat_route

    tool_name = "mcp__music__generate"
    tool_definition = {
        "type": "function",
        "function": {
            "name": tool_name,
            "parameters": {"type": "object", "properties": {}},
        },
    }

    async def fake_collect():
        return [tool_definition]

    artifacts = [
        ToolArtifact("audio", "audio/mpeg", "public.mp3", b"public"),
        ToolArtifact(
            "image",
            "image/png",
            "assistant-only.png",
            b"private",
            for_user=False,
        ),
    ]

    def fake_make_executor(*, artifact_handler=None):
        def executor(full_name, arguments=""):
            assert full_name == tool_name
            for artifact in artifacts:
                artifact_handler(tool_name, artifact)
            return ToolExecutionResult("done", artifacts)

        return executor

    monkeypatch.setattr(chat_route, "collect_mcp_tools", fake_collect)
    monkeypatch.setattr(chat_route, "make_mcp_executor", fake_make_executor)
    store = SessionStore(tmp_path)
    session = store.create()

    registry = asyncio.run(chat_route._build_registry(session["id"], store))
    result = registry.execute_detailed(tool_name, {})
    saved = store.get(session["id"])["files"]

    assert result.artifacts[0].file_id is not None
    assert result.artifacts[1].file_id is None
    assert [item["name"] for item in saved] == ["public.mp3"]
    assert saved[0]["source"] == tool_name