From f8fff54682ccdf29e650164f966de3db45e8c8ce Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=9B=90=E7=B2=92=20Yanli?= Date: Thu, 9 Jul 2026 16:30:25 +0800 Subject: [PATCH] fix(api): prefer filename extension for tool files --- api/controllers/files/upload.py | 6 +- api/core/tools/tool_file_manager.py | 15 ++++- api/core/tools/utils/message_transformer.py | 9 +-- .../nodes/agent_v2/output_file_rebacker.py | 4 +- api/factories/file_factory/builders.py | 10 +++- .../controllers/files/test_upload.py | 15 +++-- .../core/tools/test_tool_file_manager.py | 57 +++++++++++++++++++ .../tools/utils/test_message_transformer.py | 28 ++++++++- .../agent_v2/test_output_file_rebacker.py | 11 ++++ .../factories/test_build_from_mapping.py | 33 +++++++++++ 10 files changed, 166 insertions(+), 22 deletions(-) diff --git a/api/controllers/files/upload.py b/api/controllers/files/upload.py index 249784dee92..3de92589f6a 100644 --- a/api/controllers/files/upload.py +++ b/api/controllers/files/upload.py @@ -1,5 +1,3 @@ -from mimetypes import guess_extension - from flask import request from flask_restx import Resource from flask_restx.api import HTTPStatus @@ -8,7 +6,7 @@ from werkzeug.exceptions import Forbidden import services from core.tools.signature import verify_plugin_file_signature -from core.tools.tool_file_manager import ToolFileManager +from core.tools.tool_file_manager import ToolFileManager, resolve_extension from core.workflow.file_reference import build_file_reference from fields.file_fields import FileResponse @@ -110,7 +108,7 @@ class PluginUploadFileApi(Resource): conversation_id=args.conversation_id, ) - extension = guess_extension(tool_file.mimetype) or ".bin" + extension = resolve_extension(filename=tool_file.name, mimetype=tool_file.mimetype) preview_url = ToolFileManager.sign_file(tool_file_id=tool_file.id, extension=extension) # Create a dictionary with all the necessary attributes diff --git a/api/core/tools/tool_file_manager.py b/api/core/tools/tool_file_manager.py index 25313dda49b..b6c1f080e19 100644 --- a/api/core/tools/tool_file_manager.py +++ b/api/core/tools/tool_file_manager.py @@ -4,6 +4,7 @@ import hmac import logging import os import time +import urllib.parse from collections.abc import Generator from mimetypes import guess_extension, guess_type from uuid import uuid4 @@ -26,7 +27,7 @@ logger = logging.getLogger(__name__) class ToolFileManager: @staticmethod def _build_graph_file_reference(tool_file: ToolFile) -> File: - extension = guess_extension(tool_file.mimetype) or ".bin" + extension = resolve_extension(filename=tool_file.name, mimetype=tool_file.mimetype) return File( file_type=get_file_type_by_mime_type(tool_file.mimetype), transfer_method=FileTransferMethod.TOOL_FILE, @@ -70,7 +71,7 @@ class ToolFileManager: mimetype: str, filename: str | None = None, ) -> ToolFile: - extension = guess_extension(mimetype) or ".bin" + extension = resolve_extension(filename=filename, mimetype=mimetype) unique_name = uuid4().hex unique_filename = f"{unique_name}{extension}" # default just as before @@ -120,7 +121,8 @@ class ToolFileManager: or response.headers.get("Content-Type", "").split(";")[0].strip() or "application/octet-stream" ) - extension = guess_extension(mimetype) or ".bin" + url_filename = os.path.basename(urllib.parse.urlparse(file_url).path) + extension = resolve_extension(filename=url_filename, mimetype=mimetype) unique_name = uuid4().hex filename = f"{unique_name}{extension}" filepath = f"tools/{tenant_id}/{filename}" @@ -220,4 +222,11 @@ def _factory() -> ToolFileManager: return ToolFileManager() +def resolve_extension(*, filename: str | None, mimetype: str) -> str: + filename_extension = os.path.splitext(filename or "")[1].lower() + if filename_extension: + return filename_extension + return guess_extension(mimetype) or ".bin" + + set_tool_file_manager_factory(_factory) diff --git a/api/core/tools/utils/message_transformer.py b/api/core/tools/utils/message_transformer.py index f2e4f0ea5fd..7511071ae31 100644 --- a/api/core/tools/utils/message_transformer.py +++ b/api/core/tools/utils/message_transformer.py @@ -3,7 +3,6 @@ import re from collections.abc import Generator from datetime import date, datetime from decimal import Decimal -from mimetypes import guess_extension from typing import Any from uuid import UUID @@ -11,7 +10,7 @@ import numpy as np import pytz from core.tools.entities.tool_entities import ToolInvokeMessage -from core.tools.tool_file_manager import ToolFileManager +from core.tools.tool_file_manager import ToolFileManager, resolve_extension from core.workflow.file_reference import parse_file_reference from graphon.file import File, FileTransferMethod, FileType from libs.login import current_user @@ -91,7 +90,8 @@ class ToolFileMessageTransformer: conversation_id=conversation_id, ) - url = f"/files/tools/{tool_file.id}{guess_extension(tool_file.mimetype) or '.png'}" + extension = resolve_extension(filename=tool_file.name, mimetype=tool_file.mimetype) + url = cls.get_tool_file_url(tool_file_id=tool_file.id, extension=extension) meta = cls._with_tool_file_meta( message.meta, tool_file_id=str(tool_file.id), @@ -136,7 +136,8 @@ class ToolFileMessageTransformer: filename=filename, ) - url = cls.get_tool_file_url(tool_file_id=tool_file.id, extension=guess_extension(tool_file.mimetype)) + extension = resolve_extension(filename=tool_file.name, mimetype=tool_file.mimetype) + url = cls.get_tool_file_url(tool_file_id=tool_file.id, extension=extension) meta = cls._with_tool_file_meta(meta, tool_file_id=str(tool_file.id)) # check if file is image diff --git a/api/core/workflow/nodes/agent_v2/output_file_rebacker.py b/api/core/workflow/nodes/agent_v2/output_file_rebacker.py index ac39583626a..90a73bc6246 100644 --- a/api/core/workflow/nodes/agent_v2/output_file_rebacker.py +++ b/api/core/workflow/nodes/agent_v2/output_file_rebacker.py @@ -10,13 +10,13 @@ trustworthy metadata. from __future__ import annotations -from mimetypes import guess_extension from uuid import UUID from sqlalchemy import select from sqlalchemy.exc import DataError, SQLAlchemyError from core.db.session_factory import session_factory +from core.tools.tool_file_manager import resolve_extension from core.workflow.file_reference import build_file_reference from graphon.file import File, FileTransferMethod, get_file_type_by_mime_type from models.tools import ToolFile @@ -46,7 +46,7 @@ def reback_tool_file_output(*, tenant_id: str, tool_file_id: str) -> File | None return None mime_type = tool_file.mimetype or "" - extension = guess_extension(mime_type) or ".bin" + extension = resolve_extension(filename=tool_file.name, mimetype=mime_type) return File( type=get_file_type_by_mime_type(mime_type), transfer_method=FileTransferMethod.TOOL_FILE, diff --git a/api/factories/file_factory/builders.py b/api/factories/file_factory/builders.py index 67b1c646db3..91ae554cffe 100644 --- a/api/factories/file_factory/builders.py +++ b/api/factories/file_factory/builders.py @@ -3,6 +3,7 @@ from __future__ import annotations import mimetypes +import os import uuid from collections.abc import Mapping, Sequence from typing import Any, Literal, NotRequired, TypedDict, assert_never, cast @@ -285,7 +286,7 @@ def _build_from_remote_url( raise ValueError("Invalid file url") mime_type, filename, file_size = get_remote_file_info(url) - extension = mimetypes.guess_extension(mime_type) or ("." + filename.split(".")[-1] if "." in filename else ".bin") + extension = os.path.splitext(filename)[1].lower() or mimetypes.guess_extension(mime_type) or ".bin" detected_file_type = standardize_file_type(extension=extension, mime_type=mime_type) file_type = _resolve_file_type( detected_file_type=detected_file_type, @@ -326,7 +327,12 @@ def _build_from_tool_file( if tool_file is None: raise ValueError(f"ToolFile {tool_file_id} not found") - extension = "." + tool_file.file_key.split(".")[-1] if "." in tool_file.file_key else ".bin" + extension = ( + os.path.splitext(tool_file.name)[1].lower() + or mimetypes.guess_extension(tool_file.mimetype) + or os.path.splitext(tool_file.file_key)[1].lower() + or ".bin" + ) detected_file_type = standardize_file_type(extension=extension, mime_type=tool_file.mimetype) file_type = _resolve_file_type( detected_file_type=detected_file_type, diff --git a/api/tests/unit_tests/controllers/files/test_upload.py b/api/tests/unit_tests/controllers/files/test_upload.py index 4d5522c0c6c..4f6fcd6f3e0 100644 --- a/api/tests/unit_tests/controllers/files/test_upload.py +++ b/api/tests/unit_tests/controllers/files/test_upload.py @@ -34,11 +34,11 @@ class DummyFile: class DummyToolFile: - def __init__(self): + def __init__(self, name="test.txt", mimetype="text/plain"): self.id = "file-id" - self.name = "test.txt" + self.name = name self.size = 10 - self.mimetype = "text/plain" + self.mimetype = mimetype self.original_url = "http://original" self.user_id = "user-1" self.tenant_id = "tenant-1" @@ -56,7 +56,7 @@ class TestPluginUploadFileApi: mock_get_user, mock_verify_signature, ): - dummy_file = DummyFile() + dummy_file = DummyFile(filename="report.docx", mimetype="application/octet-stream") module.request = fake_request( { @@ -71,7 +71,10 @@ class TestPluginUploadFileApi: ) tool_file_manager_instance = mock_tool_file_manager.return_value - tool_file_manager_instance.create_file_by_raw.return_value = DummyToolFile() + tool_file_manager_instance.create_file_by_raw.return_value = DummyToolFile( + name="report.docx", + mimetype="application/octet-stream", + ) mock_tool_file_manager.sign_file.return_value = "signed-url" @@ -84,10 +87,12 @@ class TestPluginUploadFileApi: assert result["id"] == "file-id" assert result["reference"] == build_file_reference(record_id="file-id") assert result["preview_url"] == "signed-url" + assert result["extension"] == ".docx" mock_verify_signature.assert_called_once() assert mock_verify_signature.call_args.kwargs["conversation_id"] == "conversation-1" tool_file_manager_instance.create_file_by_raw.assert_called_once() assert tool_file_manager_instance.create_file_by_raw.call_args.kwargs["conversation_id"] == "conversation-1" + mock_tool_file_manager.sign_file.assert_called_once_with(tool_file_id="file-id", extension=".docx") def test_missing_file(self): module.request = fake_request( diff --git a/api/tests/unit_tests/core/tools/test_tool_file_manager.py b/api/tests/unit_tests/core/tools/test_tool_file_manager.py index c44ac0d13ef..7f7997b8884 100644 --- a/api/tests/unit_tests/core/tools/test_tool_file_manager.py +++ b/api/tests/unit_tests/core/tools/test_tool_file_manager.py @@ -60,6 +60,37 @@ def test_create_file_by_raw_stores_file_and_persists_record() -> None: session.refresh.assert_called_once_with(file_model) +def test_create_file_by_raw_prefers_filename_extension_over_mimetype() -> None: + manager = ToolFileManager() + session = Mock() + session.refresh.side_effect = lambda model: setattr(model, "id", "tf-docx") + + def tool_file_factory(**kwargs): + return SimpleNamespace(**kwargs) + + with ( + patch("core.tools.tool_file_manager.storage") as storage, + patch("core.tools.tool_file_manager.ToolFile", side_effect=tool_file_factory), + patch("core.tools.tool_file_manager.uuid4", return_value=SimpleNamespace(hex="abc")), + _patch_session_factory(session), + ): + file_model = manager.create_file_by_raw( + user_id="u1", + tenant_id="t1", + conversation_id="c1", + file_binary=b"docx", + mimetype="application/octet-stream", + filename="report.docx", + ) + + assert file_model.name == "report.docx" + assert file_model.file_key == "tools/t1/abc.docx" + storage.save.assert_called_once_with("tools/t1/abc.docx", b"docx") + session.add.assert_called_once_with(file_model) + session.commit.assert_called_once() + session.refresh.assert_called_once_with(file_model) + + def test_create_file_by_url_downloads_and_persists_record() -> None: manager = ToolFileManager() response = Mock() @@ -88,6 +119,32 @@ def test_create_file_by_url_downloads_and_persists_record() -> None: session.refresh.assert_called_once_with(file_model) +def test_create_file_by_url_prefers_url_extension_over_mimetype() -> None: + manager = ToolFileManager() + response = Mock() + response.content = b"docx" + response.headers = {"Content-Type": "application/octet-stream"} + response.raise_for_status.return_value = None + session = Mock() + + def tool_file_factory(**kwargs): + return SimpleNamespace(**kwargs) + + session.refresh.side_effect = lambda model: setattr(model, "id", "tf-docx") + with ( + patch("core.tools.tool_file_manager.storage") as storage, + patch("core.tools.tool_file_manager.ToolFile", side_effect=tool_file_factory), + patch("core.tools.tool_file_manager.uuid4", return_value=SimpleNamespace(hex="urlabc")), + _patch_session_factory(session), + patch("core.tools.tool_file_manager.remote_fetcher.make_request", return_value=response), + ): + file_model = manager.create_file_by_url("u1", "t1", "https://example.com/report.docx?download=1", "c1") + + assert file_model.file_key == "tools/t1/urlabc.docx" + assert file_model.name == "urlabc.docx" + storage.save.assert_called_once_with("tools/t1/urlabc.docx", b"docx") + + def test_create_file_by_url_raises_on_timeout() -> None: manager = ToolFileManager() diff --git a/api/tests/unit_tests/core/tools/utils/test_message_transformer.py b/api/tests/unit_tests/core/tools/utils/test_message_transformer.py index 354b3955042..c5e3f32cca6 100644 --- a/api/tests/unit_tests/core/tools/utils/test_message_transformer.py +++ b/api/tests/unit_tests/core/tools/utils/test_message_transformer.py @@ -7,9 +7,10 @@ from core.tools.entities.tool_entities import ToolInvokeMessage class _FakeToolFile: - def __init__(self, mimetype: str): + def __init__(self, mimetype: str, name: str | None): self.id = "fake-tool-file-id" self.mimetype = mimetype + self.name = name or "fake-tool-file.bin" class _FakeToolFileManager: @@ -38,7 +39,7 @@ class _FakeToolFileManager: "mimetype": mimetype, "filename": filename, } - return _FakeToolFile(mimetype) + return _FakeToolFile(mimetype, filename) @pytest.fixture(autouse=True) @@ -89,6 +90,29 @@ def test_transform_tool_invoke_messages_mimetype_key_present_but_none(): assert o.meta["tool_file_id"] == "fake-tool-file-id" +def test_transform_tool_invoke_messages_prefers_filename_extension_over_mimetype(): + msg = ToolInvokeMessage( + type=ToolInvokeMessage.MessageType.BLOB, + message=ToolInvokeMessage.BlobMessage(blob=b"docx"), + meta={"mime_type": "application/octet-stream", "filename": "report.docx"}, + ) + + out = list( + mt.ToolFileMessageTransformer.transform_tool_invoke_messages( + messages=_gen([msg]), + user_id="u1", + tenant_id="t1", + conversation_id="c1", + ) + ) + + assert _FakeToolFileManager.last_call is not None + assert _FakeToolFileManager.last_call["filename"] == "report.docx" + assert len(out) == 1 + assert isinstance(out[0].message, ToolInvokeMessage.TextMessage) + assert out[0].message.text.endswith(".docx") + + def test_transform_tool_invoke_messages_parses_existing_tool_file_link_meta(): msg = ToolInvokeMessage( type=ToolInvokeMessage.MessageType.IMAGE_LINK, diff --git a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_output_file_rebacker.py b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_output_file_rebacker.py index f72b7f806aa..5299fe462c8 100644 --- a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_output_file_rebacker.py +++ b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_output_file_rebacker.py @@ -57,6 +57,17 @@ def test_reback_resolves_tenant_tool_file_to_file(): assert file.extension == ".png" +def test_reback_prefers_filename_extension_over_mimetype(): + tf = _seed(mimetype="application/octet-stream", name="report.docx", size=99) + file = reback_tool_file_output(tenant_id=TENANT, tool_file_id=tf) + + assert file is not None + assert file.filename == "report.docx" + assert file.mime_type == "application/octet-stream" + assert file.extension == ".docx" + assert file.type == FileType.CUSTOM + + def test_reback_other_tenant_returns_none(): tf = _seed() assert reback_tool_file_output(tenant_id="33333333-3333-3333-3333-333333333333", tool_file_id=tf) is None diff --git a/api/tests/unit_tests/factories/test_build_from_mapping.py b/api/tests/unit_tests/factories/test_build_from_mapping.py index e2dd0efee32..90d92989c61 100644 --- a/api/tests/unit_tests/factories/test_build_from_mapping.py +++ b/api/tests/unit_tests/factories/test_build_from_mapping.py @@ -165,6 +165,20 @@ def test_build_from_mapping_accepts_opaque_related_id_for_tool_file(mock_tool_fi assert file.storage_key == "tool_file.pdf" +def test_build_from_mapping_prefers_tool_filename_extension_over_mimetype(mock_tool_file): + mock_tool_file.name = "report.docx" + mock_tool_file.file_key = "tools/test_tenant_id/file.bin" + mock_tool_file.mimetype = "application/octet-stream" + mapping = tool_file_mapping(file_type="document") + + file = build_from_mapping(mapping=mapping, tenant_id=TEST_TENANT_ID) + + assert file.extension == ".docx" + assert file.filename == "report.docx" + assert file.mime_type == "application/octet-stream" + assert file.storage_key == "tools/test_tenant_id/file.bin" + + @pytest.mark.parametrize( ("file_type", "should_pass", "expected_error"), [ @@ -213,6 +227,25 @@ def test_build_from_remote_url(mock_http_head): assert file.size == 2048 +def test_build_from_remote_url_prefers_filename_extension_over_mimetype(): + mapping = { + "transfer_method": "remote_url", + "url": TEST_REMOTE_URL, + "type": "document", + } + + with patch( + "factories.file_factory.builders.get_remote_file_info", + return_value=("application/octet-stream", "report.docx", 99), + ): + file = build_from_mapping(mapping=mapping, tenant_id=TEST_TENANT_ID) + + assert file.filename == "report.docx" + assert file.extension == ".docx" + assert file.mime_type == "application/octet-stream" + assert file.size == 99 + + @pytest.mark.parametrize( ("file_type", "should_pass", "expected_error"), [