From 38118aff078992d9da7477bfb5d7c3daea971a4b Mon Sep 17 00:00:00 2001
From: minzifang <927410375@qq.com>
Date: Wed, 3 Jun 2026 16:28:11 +0800
Subject: [PATCH] fix: preserve annotation reply metadata

---
 api/core/app/apps/advanced_chat/generate_task_pipeline.py  | 3 ---
 api/core/app/apps/base_app_generate_response_converter.py  | 4 ----
 .../apps/advanced_chat/test_generate_task_pipeline_core.py | 7 +++++--
 .../test_agent_chat_generate_response_converter.py         | 4 ++--
 .../test_completion_generate_response_converter.py         | 4 ++--
 5 files changed, 9 insertions(+), 13 deletions(-)

diff --git a/api/core/app/apps/advanced_chat/generate_task_pipeline.py b/api/core/app/apps/advanced_chat/generate_task_pipeline.py
index 6f8117ee590..eb1f610e6c8 100644
--- a/api/core/app/apps/advanced_chat/generate_task_pipeline.py
+++ b/api/core/app/apps/advanced_chat/generate_task_pipeline.py
@@ -1078,9 +1078,6 @@ class AdvancedChatAppGenerateTaskPipeline(GraphRuntimeStateSupport):
         """
         extras = self._task_state.metadata.model_dump()
 
-        if self._task_state.metadata.annotation_reply:
-            del extras["annotation_reply"]
-
         return MessageEndStreamResponse(
             task_id=self._application_generate_entity.task_id,
             id=self._message_id,
diff --git a/api/core/app/apps/base_app_generate_response_converter.py b/api/core/app/apps/base_app_generate_response_converter.py
index abcbb2f9436..6c4556d4a0c 100644
--- a/api/core/app/apps/base_app_generate_response_converter.py
+++ b/api/core/app/apps/base_app_generate_response_converter.py
@@ -99,10 +99,6 @@ class AppGenerateResponseConverter[TBlockingResponse: AppBlockingResponse](ABC):
                 )
             metadata["retriever_resources"] = updated_resources
 
-        # show annotation reply
-        if "annotation_reply" in metadata:
-            del metadata["annotation_reply"]
-
         # show usage
         if "usage" in metadata:
             del metadata["usage"]
diff --git a/api/tests/unit_tests/core/app/apps/advanced_chat/test_generate_task_pipeline_core.py b/api/tests/unit_tests/core/app/apps/advanced_chat/test_generate_task_pipeline_core.py
index 9726c939e94..96c7a1d36c6 100644
--- a/api/tests/unit_tests/core/app/apps/advanced_chat/test_generate_task_pipeline_core.py
+++ b/api/tests/unit_tests/core/app/apps/advanced_chat/test_generate_task_pipeline_core.py
@@ -264,7 +264,7 @@ class TestAdvancedChatGenerateTaskPipeline:
         assert "UPDATE messages" in stmt_str
         assert "WHERE messages.id" in stmt_str
 
-    def test_message_end_to_stream_response_strips_annotation_reply(self):
+    def test_message_end_to_stream_response_preserves_annotation_reply(self):
         pipeline = _make_pipeline()
         pipeline._task_state.metadata.annotation_reply = AnnotationReply(
             id="ann",
@@ -273,7 +273,10 @@ class TestAdvancedChatGenerateTaskPipeline:
 
         response = pipeline._message_end_to_stream_response()
 
-        assert "annotation_reply" not in response.metadata
+        assert response.metadata["annotation_reply"] == {
+            "id": "ann",
+            "account": {"id": "acc", "name": "acc"},
+        }
 
     def test_handle_output_moderation_chunk_publishes_stop(self):
         pipeline = _make_pipeline()
diff --git a/api/tests/unit_tests/core/app/apps/agent_chat/test_agent_chat_generate_response_converter.py b/api/tests/unit_tests/core/app/apps/agent_chat/test_agent_chat_generate_response_converter.py
index e861a0c684b..de167bd7d00 100644
--- a/api/tests/unit_tests/core/app/apps/agent_chat/test_agent_chat_generate_response_converter.py
+++ b/api/tests/unit_tests/core/app/apps/agent_chat/test_agent_chat_generate_response_converter.py
@@ -71,7 +71,7 @@ class TestAgentChatAppGenerateResponseConverterBlocking:
 
         result = AgentChatAppGenerateResponseConverter.convert_blocking_simple_response(blocking)
 
-        assert "annotation_reply" not in result["metadata"]
+        assert result["metadata"]["annotation_reply"] == {"id": "a"}
         assert "usage" not in result["metadata"]
 
     def test_convert_blocking_simple_response_with_non_dict_metadata(self):
@@ -169,7 +169,7 @@ class TestAgentChatAppGenerateResponseConverterStream:
         assert items[2]["event"] == "message_end"
         assert "metadata" in items[2]
         metadata = items[2]["metadata"]
-        assert "annotation_reply" not in metadata
+        assert metadata["annotation_reply"] == {"id": "a"}
         assert "usage" not in metadata
         assert metadata["retriever_resources"] == [
             {
diff --git a/api/tests/unit_tests/core/app/apps/completion/test_completion_generate_response_converter.py b/api/tests/unit_tests/core/app/apps/completion/test_completion_generate_response_converter.py
index 0136dbf5ad6..f1ea01b27f3 100644
--- a/api/tests/unit_tests/core/app/apps/completion/test_completion_generate_response_converter.py
+++ b/api/tests/unit_tests/core/app/apps/completion/test_completion_generate_response_converter.py
@@ -75,7 +75,7 @@ class TestCompletionAppGenerateResponseConverter:
 
         result = CompletionAppGenerateResponseConverter.convert_blocking_simple_response(blocking)
 
-        assert "annotation_reply" not in result["metadata"]
+        assert result["metadata"]["annotation_reply"] == {"a": 1}
         assert "usage" not in result["metadata"]
         assert result["metadata"]["retriever_resources"][0]["dataset_id"] == "dataset-1"
         assert result["metadata"]["retriever_resources"][0]["document_id"] == "document-1"
@@ -164,6 +164,6 @@ class TestCompletionAppGenerateResponseConverter:
 
         assert result[0] == "ping"
         assert result[1]["event"] == "message_end"
-        assert "annotation_reply" not in result[1]["metadata"]
+        assert result[1]["metadata"]["annotation_reply"] == {"a": 1}
         assert "usage" not in result[1]["metadata"]
         assert result[2]["event"] == "error"