"""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