diff --git a/openviking/parse/accessors/base.py b/openviking/parse/accessors/base.py index 22feee6ec6..e9a4245588 100644 --- a/openviking/parse/accessors/base.py +++ b/openviking/parse/accessors/base.py @@ -68,19 +68,22 @@ def cleanup(self) -> None: if not self.is_temporary: return - if not self.path.exists(): + cleanup_path_value = self.meta.get("_cleanup_path") + cleanup_path = Path(cleanup_path_value) if cleanup_path_value else self.path + + if not cleanup_path.exists(): return try: - if self.path.is_dir(): - shutil.rmtree(self.path, ignore_errors=True) + if cleanup_path.is_dir(): + shutil.rmtree(cleanup_path, ignore_errors=True) else: - self.path.unlink(missing_ok=True) + cleanup_path.unlink(missing_ok=True) except Exception as e: from openviking_cli.utils.logger import get_logger logger = get_logger(__name__) - logger.warning(f"[LocalResource] Failed to cleanup resource {self.path}: {e}") + logger.warning(f"[LocalResource] Failed to cleanup resource {cleanup_path}: {e}") def __enter__(self) -> "LocalResource": """Support context manager protocol.""" diff --git a/openviking/parse/accessors/feishu_accessor.py b/openviking/parse/accessors/feishu_accessor.py index 16f561c024..e4a499dda8 100644 --- a/openviking/parse/accessors/feishu_accessor.py +++ b/openviking/parse/accessors/feishu_accessor.py @@ -10,7 +10,9 @@ Install with: pip install 'openviking[bot-feishu]' """ +import asyncio import os +import re import tempfile from dataclasses import dataclass from pathlib import Path @@ -20,9 +22,12 @@ from openviking_cli.utils.logger import get_logger from .base import DataAccessor, LocalResource, SourceType +from .mime_types import get_preferred_extension logger = get_logger(__name__) +_FEISHU_IMAGE_RE = re.compile(r"!\[([^\]]*)\]\(feishu://image/([^)]+)\)") + def _getattr_safe(obj, key: str, default=None): """Get attribute from SDK object or dict, with safe fallback.""" @@ -185,26 +190,48 @@ async def access(self, source: Union[str, Path], **kwargs) -> LocalResource: feishu_access_token=feishu_access_token, ) - # Create temporary file - temp_file = tempfile.NamedTemporaryFile( - mode="w", - suffix=".md", - prefix="ov_feishu_", - delete=False, + # lark-oapi media downloads are synchronous; run them off the event + # loop so a slow Feishu request cannot block unrelated async work. + markdown_content, downloaded_images = await asyncio.to_thread( + self._resolve_image_refs, + doc.markdown_content, + feishu_access_token=feishu_access_token, ) - temp_file.write(doc.markdown_content) - temp_file.close() # Build metadata meta = { "feishu_doc_type": doc.doc_type, "feishu_token": doc.token, "feishu_title": doc.title, + "original_filename": doc.title, **doc.meta, } + if downloaded_images: + temp_dir = Path(tempfile.mkdtemp(prefix="ov_feishu_")) + markdown_path = temp_dir / "document.md" + markdown_path.write_text(markdown_content, encoding="utf-8") + for rel_path, image_bytes in downloaded_images.items(): + image_path = temp_dir / rel_path + image_path.parent.mkdir(parents=True, exist_ok=True) + image_path.write_bytes(image_bytes) + meta["_cleanup_path"] = str(temp_dir) + local_path = markdown_path + else: + # Create temporary file + temp_file = tempfile.NamedTemporaryFile( + mode="w", + suffix=".md", + prefix="ov_feishu_", + delete=False, + encoding="utf-8", + ) + temp_file.write(markdown_content) + temp_file.close() + local_path = Path(temp_file.name) + return LocalResource( - path=Path(temp_file.name), + path=local_path, source_type=SourceType.FEISHU, original_source=source_str, meta=meta, @@ -226,8 +253,6 @@ async def _fetch_document( This method extracts and adapts the logic from FeishuParser.parse(). """ - import asyncio - doc_type, token = self._parse_feishu_url(url) title = None meta = {} @@ -590,6 +615,169 @@ def _handle_image(block, block_map: Dict = None, **_) -> Optional[str]: alt_text = getattr(image, "alt", "") or "image" return f"![{alt_text}](feishu://image/{file_token})" + # Image byte-magic signatures → file extension. Sniffed from the raw bytes + # first, since the actual content is authoritative over a (possibly generic + # or wrong) Content-Type header. + _IMAGE_MAGIC = ( + (b"\x89PNG\r\n\x1a\n", ".png"), + (b"\xff\xd8\xff", ".jpg"), + (b"GIF87a", ".gif"), + (b"GIF89a", ".gif"), + (b"BM", ".bmp"), + ) + + @classmethod + def _guess_image_ext(cls, content: bytes, content_type: Optional[str]) -> str: + """Infer an image file extension from the bytes, then Content-Type. + + Feishu media are not guaranteed to be PNG, so we avoid a hardcoded + extension that would misrepresent JPEG/WebP/GIF bytes to downstream + consumers (e.g. emitting JPEG bytes as ``data:image/png``). Byte magic + is checked first because the payload is authoritative; the response + Content-Type is only a fallback for formats we do not sniff here. + """ + # WebP: "RIFF....WEBP" + if len(content) >= 12 and content[:4] == b"RIFF" and content[8:12] == b"WEBP": + return ".webp" + for magic, ext in cls._IMAGE_MAGIC: + if content.startswith(magic): + return ext + if content_type: + ext = get_preferred_extension(content_type) + if ext: + return ext + return ".png" + + @staticmethod + def _image_filename(file_token: str, ext: str = ".png") -> str: + """Return a conservative local filename for a Feishu media token.""" + safe_token = re.sub(r"[^A-Za-z0-9_.-]+", "_", file_token).strip("._") + if not ext.startswith("."): + ext = f".{ext}" + return f"{safe_token or 'image'}{ext}" + + def _download_image( + self, + file_token: str, + *, + feishu_access_token: Optional[str] = None, + ) -> Optional[Tuple[bytes, Optional[str]]]: + """Download an image from Feishu Drive API by file token. + + Returns a ``(content, content_type)`` tuple, or ``None`` on failure. + """ + import lark_oapi as lark + + client = self._get_client(use_user_token=bool(feishu_access_token)) + # Match the auth mode used to fetch the document: with a user access + # token the request must advertise USER, otherwise lark-oapi never + # injects it (see lark_oapi.core.token.auth.verify) and the download + # silently fails — dropping images from user-token imports. + token_type = ( + lark.AccessTokenType.USER + if feishu_access_token + else lark.AccessTokenType.TENANT + ) + raw_req = ( + lark.BaseRequest.builder() + .http_method(lark.HttpMethod.GET) + .uri(f"/open-apis/drive/v1/medias/{file_token}/download") + .token_types({token_type}) + .build() + ) + option = self._user_request_option(feishu_access_token) + + try: + raw_resp = client.request(raw_req) if option is None else client.request(raw_req, option) + except Exception as exc: + logger.warning("[FeishuAccessor] Error downloading image %s: %s", file_token, exc) + return None + + if not raw_resp.success(): + raw = getattr(raw_resp, "raw", None) + http_status = getattr(raw, "status_code", None) + detail = getattr(raw_resp, "msg", "") or f"HTTP {http_status}" + if http_status == 403: + detail = ( + f"{detail} (missing Feishu permission docs:document.media:download)" + ) + logger.warning( + "[FeishuAccessor] Failed to download image %s: code=%s, http=%s, msg=%s", + file_token, + getattr(raw_resp, "code", None), + http_status, + detail, + ) + return None + + raw = getattr(raw_resp, "raw", None) + content = getattr(raw, "content", None) + if not content: + logger.warning("[FeishuAccessor] Empty image response for %s", file_token) + return None + return content, self._response_content_type(raw) + + @staticmethod + def _response_content_type(raw) -> Optional[str]: + """Best-effort extraction of the Content-Type header from a lark raw response.""" + headers = getattr(raw, "headers", None) + if not headers: + return None + # lark's raw.headers may be a plain dict or a case-insensitive mapping. + try: + get = headers.get + except AttributeError: + return None + return get("Content-Type") or get("content-type") + + def _resolve_image_refs( + self, + markdown: str, + *, + feishu_access_token: Optional[str] = None, + ) -> Tuple[str, Dict[str, bytes]]: + """Download Feishu image refs and rewrite them to local relative paths.""" + config = self._get_config() + if not getattr(config, "download_images", True): + return markdown, {} + + matches = list(_FEISHU_IMAGE_RE.finditer(markdown)) + if not matches: + return markdown, {} + + token_to_rel_path: Dict[str, str] = {} + downloaded_images: Dict[str, bytes] = {} + for match in matches: + file_token = match.group(2) + if file_token in token_to_rel_path: + continue + + downloaded = self._download_image( + file_token, + feishu_access_token=feishu_access_token, + ) + if downloaded is None: + continue + image_bytes, content_type = downloaded + + ext = self._guess_image_ext(image_bytes, content_type) + rel_path = f"images/{self._image_filename(file_token, ext)}" + token_to_rel_path[file_token] = rel_path + downloaded_images[rel_path] = image_bytes + + if not downloaded_images: + return markdown, {} + + def _replace(match: re.Match[str]) -> str: + alt_text = match.group(1) + file_token = match.group(2) + rel_path = token_to_rel_path.get(file_token) + if not rel_path: + return match.group(0) + return f"![{alt_text}]({rel_path})" + + return _FEISHU_IMAGE_RE.sub(_replace, markdown), downloaded_images + def _extract_block_text(self, block, attr_name: str) -> str: """Extract text from a block's named attribute (e.g. block.text, block.heading2).""" content_obj = getattr(block, attr_name, None) diff --git a/openviking_cli/utils/config/parser_config.py b/openviking_cli/utils/config/parser_config.py index 2aa84c4054..6d3d1a2c3e 100644 --- a/openviking_cli/utils/config/parser_config.py +++ b/openviking_cli/utils/config/parser_config.py @@ -507,9 +507,7 @@ class FeishuConfig(ParserConfig): domain: str = "https://open.feishu.cn" max_rows_per_sheet: int = 1000 max_records_per_table: int = 1000 - download_images: bool = ( - True # TODO: not yet implemented, reserved for future image download support - ) + download_images: bool = True request_timeout: float = ( 30.0 # TODO: not yet passed to lark-oapi client, reserved for future use ) diff --git a/tests/parse/test_feishu_accessor.py b/tests/parse/test_feishu_accessor.py index 77184abb70..bcb3fb0c08 100644 --- a/tests/parse/test_feishu_accessor.py +++ b/tests/parse/test_feishu_accessor.py @@ -1,7 +1,8 @@ # Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. # SPDX-License-Identifier: AGPL-3.0 -"""Tests for FeishuAccessor user token handling.""" +"""Tests for FeishuAccessor user token and image handling.""" +import asyncio import sys from types import ModuleType, SimpleNamespace from unittest.mock import MagicMock @@ -41,6 +42,50 @@ def build(self): return self._option +class _FakeBaseRequest: + @staticmethod + def builder(): + return _FakeBaseRequestBuilder() + + +class _FakeBaseRequestBuilder: + def __init__(self): + self._request = SimpleNamespace(http_method=None, uri=None, token_types=None) + + def http_method(self, method): + self._request.http_method = method + return self + + def uri(self, uri): + self._request.uri = uri + return self + + def token_types(self, token_types): + self._request.token_types = token_types + return self + + def build(self): + return self._request + + +class _FakeRawResponse: + def __init__(self, content=b"image-bytes", status_code=200, headers=None): + self.content = content + self.status_code = status_code + self.headers = headers or {} + + +class _FakeMediaResponse: + def __init__(self, content=b"image-bytes", success=True, code=0, msg="", headers=None): + self.raw = _FakeRawResponse(content, headers=headers) + self.code = code + self.msg = msg + self._success = success + + def success(self): + return self._success + + class _FakeListDocumentBlockRequest: @staticmethod def builder(): @@ -66,10 +111,15 @@ def build(self): def _install_fake_lark_modules(monkeypatch): + lark = ModuleType("lark_oapi") + lark.BaseRequest = _FakeBaseRequest + lark.HttpMethod = SimpleNamespace(GET="GET") + lark.AccessTokenType = SimpleNamespace(TENANT="tenant", USER="user") docx_v1 = ModuleType("lark_oapi.api.docx.v1") docx_v1.ListDocumentBlockRequest = _FakeListDocumentBlockRequest core_model = ModuleType("lark_oapi.core.model") core_model.RequestOption = _FakeRequestOption + monkeypatch.setitem(sys.modules, "lark_oapi", lark) monkeypatch.setitem(sys.modules, "lark_oapi.api.docx.v1", docx_v1) monkeypatch.setitem(sys.modules, "lark_oapi.core.model", core_model) @@ -92,3 +142,193 @@ def test_fetch_all_blocks_uses_user_access_token_option(monkeypatch): request, option = list_blocks.call_args.args assert request.document_id == "doc_token" assert option.user_access_token == "u-test" + + +def test_resolve_image_refs_respects_download_images_disabled(): + accessor = FeishuAccessor() + accessor._config = SimpleNamespace(download_images=False) + markdown = "![screenshot](feishu://image/img_token_123)" + + updated, images = accessor._resolve_image_refs(markdown) + + assert updated == markdown + assert images == {} + + +def test_resolve_image_refs_downloads_media_and_rewrites_markdown(monkeypatch): + _install_fake_lark_modules(monkeypatch) + request_media = MagicMock(return_value=_FakeMediaResponse(b"\x89PNG\r\n")) + accessor = FeishuAccessor() + accessor._config = SimpleNamespace(download_images=True) + accessor._client = SimpleNamespace(request=request_media) + + updated, images = accessor._resolve_image_refs( + "before ![screenshot](feishu://image/img_token_123) after" + ) + + assert updated == "before ![screenshot](images/img_token_123.png) after" + assert images == {"images/img_token_123.png": b"\x89PNG\r\n"} + request = request_media.call_args.args[0] + assert request.http_method == "GET" + assert request.uri == "/open-apis/drive/v1/medias/img_token_123/download" + + +def test_resolve_image_refs_uses_content_type_extension(monkeypatch): + _install_fake_lark_modules(monkeypatch) + request_media = MagicMock( + return_value=_FakeMediaResponse( + b"\xff\xd8\xff\xe0jpeg-bytes", + headers={"Content-Type": "image/jpeg"}, + ) + ) + accessor = FeishuAccessor() + accessor._config = SimpleNamespace(download_images=True) + accessor._client = SimpleNamespace(request=request_media) + + updated, images = accessor._resolve_image_refs( + "![j](feishu://image/img_token_jpeg)" + ) + + assert updated == "![j](images/img_token_jpeg.jpg)" + assert images == {"images/img_token_jpeg.jpg": b"\xff\xd8\xff\xe0jpeg-bytes"} + + +def test_resolve_image_refs_falls_back_to_byte_magic_extension(monkeypatch): + _install_fake_lark_modules(monkeypatch) + # No usable Content-Type header; extension must come from WebP byte magic. + webp_bytes = b"RIFF\x00\x00\x00\x00WEBPfake" + request_media = MagicMock(return_value=_FakeMediaResponse(webp_bytes, headers={})) + accessor = FeishuAccessor() + accessor._config = SimpleNamespace(download_images=True) + accessor._client = SimpleNamespace(request=request_media) + + updated, images = accessor._resolve_image_refs( + "![w](feishu://image/img_token_webp)" + ) + + assert updated == "![w](images/img_token_webp.webp)" + assert images == {"images/img_token_webp.webp": webp_bytes} + + +def test_download_image_uses_tenant_token_without_user_token(monkeypatch): + _install_fake_lark_modules(monkeypatch) + request_media = MagicMock(return_value=_FakeMediaResponse(b"\x89PNG\r\n")) + accessor = FeishuAccessor() + accessor._config = SimpleNamespace(download_images=True) + accessor._client = SimpleNamespace(request=request_media) + + accessor._download_image("img_token_123") + + request = request_media.call_args.args[0] + assert request.token_types == {"tenant"} + + +def test_download_image_advertises_user_token_when_provided(monkeypatch): + """With a user access token the media request must advertise USER, or + lark-oapi never injects it and the download silently fails.""" + _install_fake_lark_modules(monkeypatch) + request_media = MagicMock(return_value=_FakeMediaResponse(b"\x89PNG\r\n")) + accessor = FeishuAccessor() + accessor._config = SimpleNamespace(download_images=True) + accessor._user_token_client = SimpleNamespace(request=request_media) + + accessor._download_image("img_token_123", feishu_access_token="u-test") + + args = request_media.call_args.args + request = args[0] + assert request.token_types == {"user"} + # The user access token option must also be forwarded on the call. + assert len(args) == 2 + assert args[1].user_access_token == "u-test" + + +def test_guess_image_ext_defaults_to_png_when_unknown(): + accessor = FeishuAccessor() + assert accessor._guess_image_ext(b"not-an-image", None) == ".png" + assert accessor._guess_image_ext(b"\xff\xd8\xff", None) == ".jpg" + assert accessor._guess_image_ext(b"anything", "image/gif") == ".gif" + + +def test_access_offloads_synchronous_download_to_thread(monkeypatch): + """access() must not run the synchronous _resolve_image_refs on the event loop.""" + import threading + + _install_fake_lark_modules(monkeypatch) + accessor = FeishuAccessor() + accessor._config = SimpleNamespace(download_images=True) + + async def fake_fetch_document(*_args, **_kwargs): + from openviking.parse.accessors.feishu_accessor import FeishuDocument + + return FeishuDocument( + doc_type="docx", + token="doc_token", + markdown_content="![s](feishu://image/img_token_123)", + title="Test Doc", + meta={}, + ) + + monkeypatch.setattr(accessor, "_fetch_document", fake_fetch_document) + + main_thread = threading.get_ident() + ran_on = {} + + def fake_resolve(markdown, **_): + ran_on["thread"] = threading.get_ident() + return ( + "![s](images/img_token_123.png)", + {"images/img_token_123.png": b"\x89PNG\r\n"}, + ) + + monkeypatch.setattr(accessor, "_resolve_image_refs", fake_resolve) + + resource = asyncio.run(accessor.access("https://example.feishu.cn/docx/doc_token")) + try: + assert "thread" in ran_on, "_resolve_image_refs was never called" + assert ran_on["thread"] != main_thread, ( + "_resolve_image_refs ran on the event-loop thread; " + "it must be offloaded via asyncio.to_thread" + ) + finally: + resource.cleanup() + + +def test_access_writes_downloaded_images_next_to_markdown(monkeypatch): + accessor = FeishuAccessor() + accessor._config = SimpleNamespace(download_images=True) + + async def fake_fetch_document(*_args, **_kwargs): + from openviking.parse.accessors.feishu_accessor import FeishuDocument + + return FeishuDocument( + doc_type="docx", + token="doc_token", + markdown_content="![screenshot](feishu://image/img_token_123)", + title="Test Doc", + meta={}, + ) + + monkeypatch.setattr(accessor, "_fetch_document", fake_fetch_document) + monkeypatch.setattr( + accessor, + "_resolve_image_refs", + lambda markdown, **_: ( + "![screenshot](images/img_token_123.png)", + {"images/img_token_123.png": b"\x89PNG\r\n"}, + ), + ) + + resource = asyncio.run(accessor.access("https://example.feishu.cn/docx/doc_token")) + + try: + assert resource.path.name == "document.md" + assert resource.path.read_text(encoding="utf-8") == ( + "![screenshot](images/img_token_123.png)" + ) + image_path = resource.path.parent / "images" / "img_token_123.png" + assert image_path.read_bytes() == b"\x89PNG\r\n" + assert resource.meta["original_filename"] == "Test Doc" + finally: + resource.cleanup() + + assert not resource.path.parent.exists()