From f2d24cb7d437ad099291262d7c1b021809cd66ae Mon Sep 17 00:00:00 2001 From: tanaka takashi Date: Sun, 13 Sep 2026 13:35:03 +0900 Subject: feat(multimedia): MCP画像・音声artifactの安全な表示経路を追加 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit MCP Tool結果をテキストと型付きartifactへ分離し、画像・音声のbase64をMIME、サイズ、件数のallowlistで検証する。annotations.audienceを反映し、ユーザー向けartifactだけをセッションへ保存してSSEにはメタデータのみを公開する。 画像はvision capabilityが明示されたモデルだけへ渡し、非対応時は既存の多言語プロンプト機構による固定通知へ置き換える。MCP生成音声とユーザー添付音声はOllamaへ送らず、誤デコードしたテキストやbase64がモデル入力へ混入しないポリシーに統一する。 保存メディア用の認証済み同一オリジン配信ルートを追加し、セッションスコープ、MIME allowlist、nosniff、private/no-store、Rangeリクエストを適用する。UIはMCP由来HTMLを使わずDOM APIでimgとaudio controlsを生成し、CSPへmedia-src selfを追加する。 web_searchは文字列メタデータだけを返す既存仕様を維持する。scrape_urlとsearch_and_fetchは画像・音声・動画レスポンスを拒否し、search_and_fetchの対象拡張子を明示的allowlistへ変更して標準Webツールをテキスト専用の契約に揃える。 責務境界と導入しない過剰対策をdocs/multimedia.mdへ記録し、README、ツール仕様、CHANGELOGを更新する。MCP artifact変換・永続化、vision分岐、音声非送信、認証、Range配信、MIME偽装拒否をテストで追加確認した。 Validation: 148 tests passed; ruff check passed; mypy src passed; node --check src/ollama_web/static/app.js passed. --- tests/test_llm.py | 177 ++++++++++++++++++++++++++++++++++++++++++++++--- tests/test_mcp.py | 156 ++++++++++++++++++++++++++++++++++++++++++- tests/test_security.py | 92 +++++++++++++++++++++++++ tests/test_tools.py | 74 +++++++++++++++------ 4 files changed, 469 insertions(+), 30 deletions(-) (limited to 'tests') diff --git a/tests/test_llm.py b/tests/test_llm.py index edaaef2..0e4d0a2 100644 --- a/tests/test_llm.py +++ b/tests/test_llm.py @@ -2,7 +2,6 @@ from __future__ import annotations -import base64 from types import SimpleNamespace from typing import Any @@ -11,12 +10,12 @@ from ollama import Message from ollama._types import Image as OllamaImage from ollama_web.llm import ( - _embed_images_in_content, _is_image_message, + _remove_images_with_notice, _token_counts, stream_chat_with_tools, ) -from ollama_web.tools.registry import ToolRegistry +from ollama_web.tools.registry import ToolArtifact, ToolExecutionResult, ToolRegistry def _image_message() -> Message: @@ -35,11 +34,11 @@ def test_is_image_message_without_images() -> None: assert _is_image_message(Message(role="user", content="hello")) is False -def test_embed_images_in_content_adds_base64_markdown() -> None: - embedded = _embed_images_in_content(_image_message()) - encoded = base64.b64encode(b"fake-image-bytes").decode("ascii") - assert f"![image](data:image/png;base64,{encoded})" in embedded.content - assert not getattr(embedded, "images", None) +def test_remove_images_adds_localized_notice_without_base64() -> None: + converted = _remove_images_with_notice(_image_message(), "ja") + assert "画像入力に対応していません" in str(converted.content) + assert "ZmFrZS1pbWFnZS1ieXRlcw" not in str(converted.content) + assert not getattr(converted, "images", None) def test_stream_chat_with_tools_disables_unsupported_capabilities( @@ -82,6 +81,168 @@ def test_stream_chat_with_tools_disables_unsupported_capabilities( assert events[-1] == {"type": "done"} +@pytest.mark.parametrize( + ("capabilities", "expect_images", "notice"), + [ + ({"tools", "vision"}, True, None), + ({"tools"}, False, "画像の内容はモデルへ提供されていません"), + ], +) +def test_mcp_image_is_only_sent_to_vision_models( + monkeypatch: pytest.MonkeyPatch, + capabilities: set[str], + expect_images: bool, + notice: str | None, +) -> None: + from ollama_web import llm + + registry = ToolRegistry() + + def executor(full_name: str, arguments: dict[str, object]) -> ToolExecutionResult: + return ToolExecutionResult( + text="tool text", + artifacts=[ + ToolArtifact( + kind="image", + mime="image/png", + name="generated.png", + data=b"fake-image-bytes", + file_id="a" * 32, + ) + ], + ) + + registry.register_mcp_tool( + "mcp__test__image", + {"type": "function", "function": {"name": "mcp__test__image"}}, + executor, + ) + + class FakeClient: + def __init__(self) -> None: + self.calls: list[dict[str, Any]] = [] + + def chat(self, **kwargs: Any) -> list[SimpleNamespace]: + self.calls.append(kwargs) + if len(self.calls) == 1: + return [ + SimpleNamespace( + message={ + "content": "", + "tool_calls": [ + { + "function": { + "name": "mcp__test__image", + "arguments": {}, + } + } + ], + } + ) + ] + return [SimpleNamespace(message={"content": "done"})] + + client = FakeClient() + monkeypatch.setattr(llm, "get_client", lambda _host=None: client) + monkeypatch.setattr(llm, "resize_image", lambda data, name="": data) + + events = list( + llm.stream_chat_with_tools( + [Message(role="user", content="generate")], + model="test-model", + registry=registry, + capabilities=capabilities, + language="ja", + ) + ) + + tool_message = next( + message + for message in reversed(client.calls[1]["messages"]) + if getattr(message, "role", None) == "tool" + ) + assert bool(getattr(tool_message, "images", None)) is expect_images + if notice is not None: + assert notice in str(tool_message.content) + assert any( + event.get("media", [{}])[0].get("id") == "a" * 32 for event in events if event.get("media") + ) + + +def test_mcp_audio_is_never_sent_to_ollama( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from ollama_web import llm + + registry = ToolRegistry() + + def executor(full_name: str, arguments: dict[str, object]) -> ToolExecutionResult: + return ToolExecutionResult( + text="music generated", + artifacts=[ + ToolArtifact( + kind="audio", + mime="audio/mpeg", + name="generated.mp3", + data=b"audio-bytes", + file_id="b" * 32, + ) + ], + ) + + registry.register_mcp_tool( + "mcp__test__audio", + {"type": "function", "function": {"name": "mcp__test__audio"}}, + executor, + ) + + class FakeClient: + def __init__(self) -> None: + self.calls: list[dict[str, Any]] = [] + + def chat(self, **kwargs: Any) -> list[SimpleNamespace]: + self.calls.append(kwargs) + if len(self.calls) == 1: + return [ + SimpleNamespace( + message={ + "content": "", + "tool_calls": [ + { + "function": { + "name": "mcp__test__audio", + "arguments": {}, + } + } + ], + } + ) + ] + return [SimpleNamespace(message={"content": "done"})] + + client = FakeClient() + monkeypatch.setattr(llm, "get_client", lambda _host=None: client) + + list( + llm.stream_chat_with_tools( + [Message(role="user", content="generate")], + model="test-model", + registry=registry, + capabilities={"tools", "vision"}, + language="ja", + ) + ) + + tool_message = next( + message + for message in reversed(client.calls[1]["messages"]) + if getattr(message, "role", None) == "tool" + ) + assert not getattr(tool_message, "images", None) + assert "音声データそのものはモデルへ提供されていません" in str(tool_message.content) + assert "YXVkaW8tYnl0ZXM" not in str(tool_message.content) + + def test_token_counts_extracts_from_dict() -> None: assert _token_counts({"prompt_eval_count": 100, "eval_count": 50}) == { "prompt_eval_count": 100, diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 4f649d9..e928277 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -2,11 +2,20 @@ 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 @@ -15,10 +24,13 @@ 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.tools.registry import ToolRegistry +from ollama_web.sessions import SessionStore +from ollama_web.tools.registry import ToolArtifact, ToolExecutionResult, ToolRegistry @pytest.fixture @@ -146,3 +158,145 @@ def test_tool_registry_mcp_tool(): 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 diff --git a/tests/test_security.py b/tests/test_security.py index f929bb8..dea165b 100644 --- a/tests/test_security.py +++ b/tests/test_security.py @@ -111,3 +111,95 @@ def test_katex_local_font_assets_are_present() -> None: assert font_paths assert not [path for path in font_paths if not (vendor_dir / path).exists()] + + +def test_authenticated_audio_content_supports_range_requests( + workspace_tmp: Path, monkeypatch +) -> None: + client = _client(workspace_tmp, monkeypatch) + session = client.app.state.session_store.create() + audio = client.app.state.session_store.add_file( + session_id=session["id"], + name="generated.mp3", + data=b"0123456789", + mime="audio/mpeg", + text="", + source="mcp__music__generate", + ) + url = f"/api/sessions/{session['id']}/files/{audio.id}/content" + + assert client.get(url).status_code == 401 + assert client.post("/api/auth/login", json={"pin": "123456"}).status_code == 200 + response = client.get(url, headers={"Range": "bytes=2-5"}) + + assert response.status_code == 206 + assert response.content == b"2345" + assert response.headers["content-type"] == "audio/mpeg" + assert response.headers["content-range"] == "bytes 2-5/10" + assert response.headers["cache-control"] == "private, no-store" + assert response.headers["x-content-type-options"] == "nosniff" + assert "media-src 'self'" in response.headers["content-security-policy"] + + +def test_file_content_does_not_inline_non_media(workspace_tmp: Path, monkeypatch) -> None: + client = _client(workspace_tmp, monkeypatch) + session = client.app.state.session_store.create() + text_file = client.app.state.session_store.add_file( + session_id=session["id"], + name="notes.html", + data=b"", + mime="text/html", + text="", + ) + client.post("/api/auth/login", json={"pin": "123456"}) + + response = client.get(f"/api/sessions/{session['id']}/files/{text_file.id}/content") + assert response.status_code == 415 + + disguised_html = client.app.state.session_store.add_file( + session_id=session["id"], + name="disguised.png", + data=b"", + mime="text/html", + text="", + ) + response = client.get( + f"/api/sessions/{session['id']}/files/{disguised_html.id}/content" + ) + assert response.status_code == 415 + + +def test_uploaded_audio_is_playable_but_excluded_from_attachment_text( + workspace_tmp: Path, monkeypatch +) -> None: + from ollama_web.routes.chat import _attachment_context, _has_audio_attachments + + client = _client(workspace_tmp, monkeypatch) + login = client.post("/api/auth/login", json={"pin": "123456"}) + csrf = login.headers["x-csrf-token"] + session = client.app.state.session_store.create() + response = client.post( + f"/api/sessions/{session['id']}/files", + headers={"X-CSRF-Token": csrf}, + files={"files": ("recording.mp3", b"audio-bytes", "audio/mpeg")}, + ) + + assert response.status_code == 200 + audio = response.json()["files"][0] + assert audio["text"] == "" + assert _has_audio_attachments( + client.app.state.session_store, session["id"], [audio["id"]] + ) + assert ( + _attachment_context( + client.app.state.session_store, + session["id"], + [audio["id"]], + total_limit=4000, + language="ja", + ) + == "" + ) + playable = client.get(f"/api/sessions/{session['id']}/files/{audio['id']}/content") + assert playable.status_code == 200 + assert playable.content == b"audio-bytes" diff --git a/tests/test_tools.py b/tests/test_tools.py index 4960d1d..6ba537b 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -265,6 +265,19 @@ def test_scrape_url_short_text_no_next_token(): assert next_token is None +@pytest.mark.parametrize("content_type", ["image/png", "audio/mpeg", "video/mp4"]) +def test_scrape_url_rejects_media_content(content_type): + from ollama_web.tools import scrape as scrape_mod + + with patch.object( + scrape_mod, + "fetch_public_bytes", + _fake_fetch_public_bytes(b"binary", content_type), + ): + result = scrape_mod.scrape_url("https://example.com/media") + assert "Unsupported content type" in result + + # --------------------------------------------------------------------------- # pdf_to_text tests (pagination) # --------------------------------------------------------------------------- @@ -338,9 +351,7 @@ def test_pdf_to_text_pagination_roundtrip(monkeypatch): pages = [f"PAGE{i:02d}-" + ("X" * (7000 if i == 0 else 80)) for i in range(30)] data = _make_multipage_pdf(pages) - expected = "\n\n".join( - (page.extract_text() or "") for page in PdfReader(BytesIO(data)).pages - ) + expected = "\n\n".join((page.extract_text() or "") for page in PdfReader(BytesIO(data)).pages) # Patch fetch_public_bytes to return dummy bytes. monkeypatch.setattr( @@ -421,9 +432,7 @@ def test_session_pdf_tool_reads_original_bytes_and_is_session_scoped(tmp_path): from ollama_web.tools.pdf import make_session_pdf_tool data = _make_multipage_pdf(["FIRST-" + "A" * 7000, "SECOND-" + "B" * 200]) - expected = "\n\n".join( - (page.extract_text() or "") for page in PdfReader(BytesIO(data)).pages - ) + expected = "\n\n".join((page.extract_text() or "") for page in PdfReader(BytesIO(data)).pages) store = SessionStore(tmp_path) session = store.create() other_session = store.create() @@ -461,9 +470,7 @@ def test_search_and_fetch_pdf_pagination_is_lossless_and_persists_once(monkeypat from ollama_web.tools import fetch as fetch_mod data = _make_multipage_pdf([f"P{i:02d}-" + "Z" * 300 for i in range(30)]) - expected = "\n\n".join( - (page.extract_text() or "") for page in PdfReader(BytesIO(data)).pages - ) + expected = "\n\n".join((page.extract_text() or "") for page in PdfReader(BytesIO(data)).pages) monkeypatch.setattr( fetch_mod, "web_search", @@ -513,9 +520,11 @@ def test_search_and_fetch_accepts_legacy_v1_pdf_offset(monkeypatch): "file_idx": 0, "matched": matched, } - token = base64.urlsafe_b64encode( - json.dumps(payload, separators=(",", ":")).encode("ascii") - ).decode("ascii").rstrip("=") + token = ( + base64.urlsafe_b64encode(json.dumps(payload, separators=(",", ":")).encode("ascii")) + .decode("ascii") + .rstrip("=") + ) monkeypatch.setattr( fetch_mod, "fetch_public_bytes", lambda fetched_url, timeout=None: (data, url, "") ) @@ -535,9 +544,7 @@ def test_attachment_context_keeps_pdf_continuation_and_registry_can_resume(tmp_p chat_file = store.add_file( session["id"], "attached.pdf", data, "application/pdf", "old preview" ) - context = _attachment_context( - store, session["id"], [chat_file.id], total_limit=5000 - ) + context = _attachment_context(store, session["id"], [chat_file.id], total_limit=5000) assert len(context) <= 5000 assert f'uri="attachment:{chat_file.id}"' in context assert "next_token:" in context @@ -566,9 +573,7 @@ def test_uploaded_pdf_can_be_resumed_from_original_file(tmp_path, monkeypatch): client = TestClient(app) login = client.post("/api/auth/login", json={"pin": "123456"}) csrf = login.headers["x-csrf-token"] - session_id = client.post( - "/api/sessions", headers={"X-CSRF-Token": csrf} - ).json()["id"] + session_id = client.post("/api/sessions", headers={"X-CSRF-Token": csrf}).json()["id"] uploaded = client.post( f"/api/sessions/{session_id}/files", headers={"X-CSRF-Token": csrf}, @@ -601,6 +606,35 @@ def test_search_and_fetch_no_results(monkeypatch): assert "No pdf files found" in result +def test_search_and_fetch_rejects_unsupported_file_type(monkeypatch): + from ollama_web.tools import fetch as fetch_mod + + search = MagicMock(return_value=[]) + monkeypatch.setattr(fetch_mod, "web_search", search) + result = fetch_mod.search_and_fetch("query", "png") + assert result == "Invalid file type." + search.assert_not_called() + + +def test_search_and_fetch_rejects_media_response(monkeypatch): + from ollama_web.tools import fetch as fetch_mod + + monkeypatch.setattr( + fetch_mod, + "web_search", + lambda query, max_results=None: [ + {"title": "Fake text", "href": "https://example.com/file.txt", "body": ""} + ], + ) + monkeypatch.setattr( + fetch_mod, + "fetch_public_bytes", + lambda url, timeout=None: (b"binary", url, "image/png"), + ) + result = fetch_mod.search_and_fetch("query", "txt") + assert "unsupported media content type" in result + + def test_search_and_fetch_continue_from_invalid_token(): """Invalid continuation token returns an error.""" from ollama_web.tools.fetch import search_and_fetch @@ -687,9 +721,7 @@ def test_search_and_fetch_long_html_uses_full_cleaned_source(monkeypatch): token = None collected = "" for _ in range(20): - result = fetch_mod.search_and_fetch( - "query", "html", max_results=1, continue_from=token - ) + result = fetch_mod.search_and_fetch("query", "html", max_results=1, continue_from=token) body, has_more, token = parse_paged_result(result) collected += body.split("---\n", 1)[1].removesuffix("\n") if not has_more: -- cgit v1.2.3