mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-28 14:10:37 +00:00
03afd05917
* fix(artifacts): authorize shared image reads and offload MIME lookup * fix(artifacts): skip disabled blob lookups and warn on failures
197 lines
8.7 KiB
Python
197 lines
8.7 KiB
Python
"""Regression anchor: serving artifacts must not block the event loop.
|
|
|
|
``get_artifact`` probes the artifact path (``exists`` / ``is_file``), sniffs
|
|
text-ness (``is_text_file_by_content``),
|
|
and extracts ``.skill`` archive members — all blocking filesystem IO. The
|
|
handler offloads each branch's IO via ``asyncio.to_thread``; if any regresses
|
|
back onto the event loop, the strict Blockbuster gate raises ``BlockingError``
|
|
and these tests fail.
|
|
|
|
Binary (non-text, non-active-content) artifacts still run the same
|
|
``exists`` / ``is_file`` / ``mimetypes.guess_type`` / ``is_text_file_by_content``
|
|
probes as the text branch, offloaded the same way via ``asyncio.to_thread``.
|
|
What differs is the payload: instead of a read, the handler returns a
|
|
``FileResponse`` that defers its own file IO to ASGI response time (streamed
|
|
via ``anyio.open_file`` when the response is actually sent, which this test
|
|
never triggers) — so awaiting ``get_artifact`` itself for a binary artifact
|
|
does no full-file read; ``FileResponse`` streams the bytes at ASGI send
|
|
time, off the loop, which is why the gate has nothing to catch there
|
|
either way.
|
|
|
|
The ``@require_permission`` decorator is bypassed via ``__wrapped__`` so the
|
|
anchor exercises the handler's own filesystem IO, not the authz layer. Imports
|
|
sit at module top so any import-time IO runs at collection, outside the gate.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import hashlib
|
|
import mimetypes
|
|
import zipfile
|
|
from contextlib import asynccontextmanager
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
from starlette.requests import Request
|
|
from starlette.responses import FileResponse
|
|
|
|
import app.gateway.routers.artifacts as artifacts_router
|
|
from app.gateway import services
|
|
from app.gateway.authz import AuthContext
|
|
from app.gateway.path_utils import resolve_thread_virtual_path
|
|
from app.gateway.routers.artifacts import ArtifactUpdateRequest, get_artifact, update_artifact
|
|
from deerflow.storage import BlobRef
|
|
|
|
pytestmark = pytest.mark.asyncio
|
|
|
|
# The undecorated coroutine (``require_permission`` uses ``functools.wraps``).
|
|
_get_artifact = get_artifact.__wrapped__
|
|
_update_artifact = update_artifact.__wrapped__
|
|
|
|
|
|
async def _seed(tmp_path: Path, monkeypatch, thread_id: str, virtual_path: str) -> Path:
|
|
monkeypatch.setenv("DEER_FLOW_HOME", str(tmp_path))
|
|
# Rebuild cached Paths against the tmp home so the artifact resolves under it.
|
|
import deerflow.config.paths as paths_mod
|
|
|
|
monkeypatch.setattr(paths_mod, "_paths", None)
|
|
# Test-side path resolution also touches the filesystem (`.resolve()`); offload
|
|
# it so this seeding helper doesn't itself trip the gate.
|
|
target = await asyncio.to_thread(resolve_thread_virtual_path, thread_id, virtual_path)
|
|
await asyncio.to_thread(target.parent.mkdir, parents=True, exist_ok=True)
|
|
return target
|
|
|
|
|
|
async def test_get_artifact_text_does_not_block_event_loop(tmp_path: Path, monkeypatch) -> None:
|
|
vpath = "mnt/user-data/outputs/notes.txt"
|
|
target = await _seed(tmp_path, monkeypatch, "t1", vpath)
|
|
await asyncio.to_thread(target.write_text, "hello world", encoding="utf-8")
|
|
|
|
resp = await _get_artifact("t1", vpath, request=None, download=False)
|
|
|
|
assert isinstance(resp, FileResponse)
|
|
assert resp.status_code == 200
|
|
assert Path(resp.path) == target
|
|
assert resp.headers.get("content-disposition", "").startswith("inline;")
|
|
|
|
|
|
async def test_get_artifact_binary_does_not_block_event_loop(tmp_path: Path, monkeypatch) -> None:
|
|
vpath = "mnt/user-data/outputs/blob.bin"
|
|
target = await _seed(tmp_path, monkeypatch, "t1", vpath)
|
|
payload = b"\x00\x01\x02PNGDATA" # null byte -> binary branch (inline FileResponse)
|
|
await asyncio.to_thread(target.write_bytes, payload)
|
|
|
|
resp = await _get_artifact("t1", vpath, request=None, download=False)
|
|
|
|
# Binary artifacts are streamed via FileResponse (so browsers can issue
|
|
# byte-Range requests) instead of being read into memory up front, so the
|
|
# file bytes are only touched later during ASGI send, never here.
|
|
assert isinstance(resp, FileResponse)
|
|
assert resp.status_code == 200
|
|
assert Path(resp.path) == target
|
|
assert resp.headers.get("content-disposition", "").startswith("inline;")
|
|
|
|
|
|
async def test_get_artifact_skill_archive_member_does_not_block_event_loop(tmp_path: Path, monkeypatch) -> None:
|
|
skill_vpath = "mnt/user-data/outputs/demo.skill"
|
|
target = await _seed(tmp_path, monkeypatch, "t1", skill_vpath)
|
|
|
|
def _build_skill_zip() -> None:
|
|
with zipfile.ZipFile(target, "w") as zf:
|
|
zf.writestr("SKILL.md", "# demo skill\n")
|
|
|
|
await asyncio.to_thread(_build_skill_zip)
|
|
|
|
resp = await _get_artifact("t1", f"{skill_vpath}/SKILL.md", request=None, download=False)
|
|
|
|
assert resp.status_code == 200
|
|
assert b"# demo skill" in resp.body
|
|
|
|
|
|
async def test_update_artifact_does_not_block_event_loop(tmp_path: Path, monkeypatch) -> None:
|
|
vpath = "/mnt/user-data/outputs/notes.txt"
|
|
target = await _seed(tmp_path, monkeypatch, "t1", vpath)
|
|
original = b"hello world"
|
|
await asyncio.to_thread(target.write_bytes, original)
|
|
|
|
@asynccontextmanager
|
|
async def allow_write(*_args, **_kwargs):
|
|
yield
|
|
|
|
class MountedProvider:
|
|
uses_thread_data_mounts = True
|
|
|
|
monkeypatch.setattr(artifacts_router, "reserve_artifact_write", allow_write)
|
|
monkeypatch.setattr(artifacts_router, "get_sandbox_provider", lambda: MountedProvider())
|
|
|
|
result = await _update_artifact(
|
|
"t1",
|
|
vpath,
|
|
ArtifactUpdateRequest(
|
|
content="updated",
|
|
expected_sha256=hashlib.sha256(original).hexdigest(),
|
|
),
|
|
request=None,
|
|
)
|
|
|
|
assert result.sha256 == hashlib.sha256(b"updated").hexdigest()
|
|
assert await asyncio.to_thread(target.read_bytes) == b"updated"
|
|
|
|
|
|
@pytest.mark.parametrize("filename", ["large.txt", "large.html", "large.bin"])
|
|
@pytest.mark.parametrize("download", [False, True])
|
|
async def test_large_artifact_validator_does_not_block_event_loop(tmp_path: Path, monkeypatch, filename: str, download: bool) -> None:
|
|
vpath = f"mnt/user-data/outputs/{filename}"
|
|
target = await _seed(tmp_path, monkeypatch, "t1", vpath)
|
|
await asyncio.to_thread(target.write_bytes, b"a" * (artifacts_router.MAX_EDITABLE_ARTIFACT_BYTES + 1))
|
|
|
|
def unexpected_hash(*_args, **_kwargs):
|
|
pytest.fail("Large artifact response planning must not hash the file")
|
|
|
|
monkeypatch.setattr(artifacts_router, "_sha256_of_file", unexpected_hash)
|
|
response = await _get_artifact("t1", vpath, request=None, download=download)
|
|
|
|
assert isinstance(response, FileResponse)
|
|
assert response.headers["etag"].startswith('"stat-')
|
|
|
|
|
|
async def test_shared_image_blob_cold_mime_lookup_does_not_block(tmp_path, monkeypatch):
|
|
payload = b"\x89PNG\r\n\x1a\nimage"
|
|
virtual_path = "/mnt/user-data/outputs/chart.png"
|
|
ref = BlobRef(sha256=hashlib.sha256(payload).hexdigest(), size=len(payload), kind="viewed-image", content_type="image/png")
|
|
mime_file = tmp_path / "mime.types"
|
|
await asyncio.to_thread(mime_file.write_text, "image/png png\n", encoding="utf-8")
|
|
# Restore all globals replaced by init(), not just its lazy-init flag.
|
|
for name in ("suffix_map", "encodings_map", "types_map", "common_types"):
|
|
monkeypatch.setattr(mimetypes, name, getattr(mimetypes, name))
|
|
monkeypatch.setattr(mimetypes, "inited", False)
|
|
monkeypatch.setattr(mimetypes, "_db", None)
|
|
monkeypatch.setattr(mimetypes, "knownfiles", [str(mime_file)])
|
|
|
|
class Accessor:
|
|
async def aget(self, _config):
|
|
return SimpleNamespace(values={"viewed_images": {virtual_path: {"mime_type": "image/png", "size": len(payload), "sha256": ref.sha256, "blob_ref": ref.model_dump(mode="json")}}})
|
|
|
|
async def build_accessor(_request, *, thread_id):
|
|
return Accessor(), {"configurable": {"thread_id": thread_id}}
|
|
|
|
class Store:
|
|
def get_bytes(self, requested_ref):
|
|
assert requested_ref == ref
|
|
return payload
|
|
|
|
monkeypatch.setattr(services, "build_thread_checkpoint_state_accessor", build_accessor)
|
|
monkeypatch.setattr("deerflow.storage.get_blob_store_if_enabled", lambda: Store())
|
|
monkeypatch.setattr(artifacts_router, "resolve_thread_virtual_path", lambda *_args, **_kwargs: tmp_path / "missing.png")
|
|
request = Request({"type": "http", "method": "GET", "path": "/", "headers": [], "query_string": b"", "app": SimpleNamespace(state=SimpleNamespace(thread_store=SimpleNamespace(check_access=AsyncMock(return_value=True))))})
|
|
user = SimpleNamespace(id="owner", system_role="user")
|
|
request.state.user = user
|
|
request.state.auth = AuthContext(user=user, permissions=["threads:read"])
|
|
|
|
response = await _get_artifact("thread-1", virtual_path.lstrip("/"), request=request)
|
|
assert response.status_code == 200
|
|
assert response.body == payload
|