diff --git a/api/apps/restful_apis/agent_api.py b/api/apps/restful_apis/agent_api.py
index 33b6788436..8dd4a27a2d 100644
--- a/api/apps/restful_apis/agent_api.py
+++ b/api/apps/restful_apis/agent_api.py
@@ -173,7 +173,17 @@ def _normalize_agent_session(conv):
         messages = [message for i, message in enumerate(conv["message"]) if i != 0 and message["role"] != "user"]
         for message, reference in zip(messages, conv["reference"]):
             chunks = reference.get("chunks", [])
-            message["reference"] = [_normalize_agent_reference_chunk(chunk) for chunk in chunks]
+            if isinstance(chunks, dict):
+                refs = []
+                for citation_id, chunk in chunks.items():
+                    ref = _normalize_agent_reference_chunk(chunk)
+                    ref["citation_id"] = str(citation_id)
+                    refs.append(ref)
+                message["reference"] = refs
+            elif isinstance(chunks, list):
+                message["reference"] = [_normalize_agent_reference_chunk(chunk) for chunk in chunks]
+            else:
+                message["reference"] = []
     del conv["reference"]
     return conv
 
diff --git a/api/db/services/canvas_service.py b/api/db/services/canvas_service.py
index 1777a21b4d..c560b0bac8 100644
--- a/api/db/services/canvas_service.py
+++ b/api/db/services/canvas_service.py
@@ -368,7 +368,14 @@ async def completion(tenant_id, agent_id, session_id=None, **kwargs):
         yield "data:" + json.dumps(ans, ensure_ascii=False) + "\n\n"
 
     conv.message.append({"role": "assistant", "content": txt, "created_at": time.time(), "id": message_id})
-    conv.reference = canvas.get_reference()
+    current_reference = canvas.get_reference()
+    if not isinstance(current_reference, dict):
+        current_reference = {}
+    if not conv.reference:
+        conv.reference = []
+    if isinstance(conv.reference, dict):
+        conv.reference = [conv.reference]
+    conv.reference.append(current_reference)
     conv.errors = canvas.error
     conv.dsl = str(canvas)
     conv = conv.to_dict()