feat: add artifact

2026-03-27 15:10:33 +08:00
parent ec585ae730
commit 37070d7896
10 changed files with 768 additions and 31 deletions
@@ -0,0 +1,131 @@
+from pathlib import Path
+
+import pytest
+from fastapi.testclient import TestClient
+
+from main import app
+
+
+def _backend_data_root() -> Path:
+    return Path(__file__).resolve().parents[1] / "data"
+
+
+def test_download_artifact_within_whitelist() -> None:
+    uploads_dir = _backend_data_root() / "uploads"
+    uploads_dir.mkdir(parents=True, exist_ok=True)
+    sample = uploads_dir / "task2-download.csv"
+    sample.write_text("id,name\n1,a\n", encoding="utf-8")
+
+    client = TestClient(app)
+    response = client.get("/nanobot/artifacts/download", params={"target": "local://task2-download.csv"})
+
+    assert response.status_code == 200
+    assert response.headers["content-type"].startswith("application/octet-stream")
+    assert response.headers["content-disposition"].startswith("attachment;")
+    assert response.content == sample.read_bytes()
+
+
+def test_download_artifact_rejects_outside_paths() -> None:
+    client = TestClient(app)
+    response = client.get("/nanobot/artifacts/download", params={"target": "/etc/hosts"})
+
+    assert response.status_code == 403
+    assert response.json()["detail"] == "非法路径访问"
+
+
+def test_preview_artifact_returns_unsupported_for_binary() -> None:
+    uploads_dir = _backend_data_root() / "uploads"
+    uploads_dir.mkdir(parents=True, exist_ok=True)
+    sample = uploads_dir / "task2-unsupported.bin"
+    sample.write_bytes(b"\x00\x01\x02")
+
+    client = TestClient(app)
+    response = client.get("/nanobot/artifacts/preview", params={"target": f"local://{sample.name}"})
+
+    assert response.status_code == 415
+    assert response.json()["detail"] == "当前文件类型不支持预览，请使用下载"
+    download = client.get("/nanobot/artifacts/download", params={"target": f"local://{sample.name}"})
+    assert download.status_code == 200
+    assert download.content == sample.read_bytes()
+
+
+def test_preview_html_supports_directory_resources() -> None:
+    web_dir = _backend_data_root() / "workspace" / "task2-web"
+    web_dir.mkdir(parents=True, exist_ok=True)
+    html_file = web_dir / "index.html"
+    css_file = web_dir / "styles.css"
+    html_file.write_text("<html><head><link rel='stylesheet' href='styles.css'></head><body>ok</body></html>", encoding="utf-8")
+    css_file.write_text("body{color:#333;}", encoding="utf-8")
+
+    client = TestClient(app)
+    preview = client.get(
+        "/nanobot/artifacts/preview",
+        params={"target": str(html_file)},
+        follow_redirects=False,
+    )
+
+    assert preview.status_code == 307
+    location = preview.headers["location"]
+    assert location.startswith("/nanobot/artifacts/web/")
+
+    html_response = client.get(location)
+    assert html_response.status_code == 200
+    assert "text/html" in html_response.headers["content-type"]
+    assert "styles.css" in html_response.text
+
+    css_response = client.get(location.replace("index.html", "styles.css"))
+    assert css_response.status_code == 200
+    assert "text/css" in css_response.headers["content-type"]
+    assert "color:#333" in css_response.text
+
+
+@pytest.mark.parametrize(
+    ("filename", "payload", "expected_mime"),
+    [
+        ("task4-image.png", b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR", "image/png"),
+        ("task4-preview.pdf", b"%PDF-1.4\n1 0 obj\n<<>>\nendobj\n", "application/pdf"),
+        (
+            "task4-preview.pptx",
+            b"PK\x03\x04\x14\x00\x00\x00\x08\x00\x00\x00!\x00",
+            "application/vnd.openxmlformats-officedocument.presentationml.presentation",
+        ),
+    ],
+)
+def test_preview_and_download_supported_types(filename: str, payload: bytes, expected_mime: str) -> None:
+    uploads_dir = _backend_data_root() / "uploads"
+    uploads_dir.mkdir(parents=True, exist_ok=True)
+    sample = uploads_dir / filename
+    sample.write_bytes(payload)
+
+    client = TestClient(app)
+    preview = client.get("/nanobot/artifacts/preview", params={"target": f"local://{filename}"})
+    assert preview.status_code == 200
+    assert preview.headers["content-type"].startswith(expected_mime)
+
+    download = client.get("/nanobot/artifacts/download", params={"target": f"local://{filename}"})
+    assert download.status_code == 200
+    assert download.content == sample.read_bytes()
+
+
+def test_web_preview_missing_resource_returns_error_and_download_still_works() -> None:
+    web_dir = _backend_data_root() / "workspace" / "task4-web-missing"
+    web_dir.mkdir(parents=True, exist_ok=True)
+    html_file = web_dir / "index.html"
+    html_file.write_text("<html><head><script src='missing.js'></script></head><body>ok</body></html>", encoding="utf-8")
+
+    client = TestClient(app)
+    preview = client.get(
+        "/nanobot/artifacts/preview",
+        params={"target": str(html_file)},
+        follow_redirects=False,
+    )
+    assert preview.status_code == 307
+    location = preview.headers["location"]
+
+    missing = client.get(location.replace("index.html", "missing.js"))
+    assert missing.status_code == 404
+    assert missing.json()["detail"] == "Web 资源不存在"
+
+    download = client.get("/nanobot/artifacts/download", params={"target": str(html_file)})
+    assert download.status_code == 200
+    assert download.content == html_file.read_bytes()
@@ -0,0 +1,55 @@
+from pathlib import Path
+
+from app.core.artifacts import extract_artifacts
+
+
+def _backend_data_root() -> Path:
+    return Path(__file__).resolve().parents[1] / "data"
+
+
+def test_extract_artifacts_from_local_and_tool_paths() -> None:
+    data_root = _backend_data_root()
+    uploads_dir = data_root / "uploads"
+    workspace_dir = data_root / "workspace" / "reports"
+    uploads_dir.mkdir(parents=True, exist_ok=True)
+    workspace_dir.mkdir(parents=True, exist_ok=True)
+
+    upload_file = uploads_dir / "task1-sample.csv"
+    upload_file.write_text("a,b\n1,2\n", encoding="utf-8")
+    report_file = workspace_dir / "task1-report.html"
+    report_file.write_text("<html><body>ok</body></html>", encoding="utf-8")
+
+    content = "请下载 local://task1-sample.csv"
+    session_messages = [
+        {"role": "user", "content": "生成报告"},
+        {"role": "tool", "content": f"输出文件：{report_file}"},
+    ]
+
+    artifacts = extract_artifacts(content, session_messages)
+
+    by_name = {item["name"]: item for item in artifacts}
+    assert "task1-sample.csv" in by_name
+    assert "task1-report.html" in by_name
+    assert by_name["task1-sample.csv"]["download_url"].startswith("/nanobot/artifacts/download?target=")
+    assert by_name["task1-sample.csv"]["previewable"] is True
+    assert by_name["task1-report.html"]["previewable"] is True
+    assert by_name["task1-report.html"]["preview_url"].startswith("/nanobot/artifacts/preview?target=")
+
+
+def test_extract_artifacts_deduplicate_and_skip_missing() -> None:
+    data_root = _backend_data_root()
+    workspace_dir = data_root / "workspace"
+    workspace_dir.mkdir(parents=True, exist_ok=True)
+
+    pdf_file = workspace_dir / "task1-dedup.pdf"
+    pdf_file.write_bytes(b"%PDF-1.4 test")
+    missing_file = workspace_dir / "task1-missing.pdf"
+
+    content = f"{pdf_file} and {pdf_file} and {missing_file}"
+    artifacts = extract_artifacts(content, [])
+
+    assert len(artifacts) == 1
+    item = artifacts[0]
+    assert item["name"] == "task1-dedup.pdf"
+    assert item["mime_type"] == "application/pdf"
+    assert item["previewable"] is True