summaryrefslogtreecommitdiff
path: root/tests/test_tools.py
diff options
context:
space:
mode:
authortanaka takashi <isazaka@gmail.com>2026-09-13 13:35:03 +0900
committertanaka takashi <isazaka@gmail.com>2026-09-13 13:35:03 +0900
commitf2d24cb7d437ad099291262d7c1b021809cd66ae (patch)
treed9a11a3795e0296358cea7a497edffbe898f4583 /tests/test_tools.py
parentf36cc852cc1792bd624ff701fb2476d0b43fa1af (diff)
feat(multimedia): MCP画像・音声artifactの安全な表示経路を追加
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.
Diffstat (limited to 'tests/test_tools.py')
-rw-r--r--tests/test_tools.py74
1 files changed, 53 insertions, 21 deletions
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: