diff options
Diffstat (limited to 'tests')
| -rw-r--r-- | tests/test_llm.py | 177 | ||||
| -rw-r--r-- | tests/test_mcp.py | 156 | ||||
| -rw-r--r-- | tests/test_security.py | 92 | ||||
| -rw-r--r-- | tests/test_tools.py | 74 |
4 files changed, 469 insertions, 30 deletions
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"" 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"<script>alert(1)</script>", + 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"<script>alert(1)</script>", + 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: |