Merge commit 'f8fff54682' into feat/agent-v2

2026-07-15 17:37:38 +08:00 · 2026-07-09 16:32:22 +08:00
parent 819ecad3f7 f8fff54682
commit 3ee9d66aaa
10 changed files with 166 additions and 22 deletions
--- 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
--- 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)
--- 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
--- 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,
--- 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,
--- 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(
--- 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()

--- 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,
--- 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
--- 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"),
    [