From ed522558686b6f82c353ac50524056ffa3c2867f Mon Sep 17 00:00:00 2001
From: Wang Qi <wangq8@outlook.com>
Date: Mon, 6 Jul 2026 17:47:38 +0800
Subject: [PATCH] Fix Agent session lost <think></think> tag (#16670)

---
 api/apps/restful_apis/agent_api.py |  9 ++++++
 internal/service/agent.go          | 48 ++++++++++++++++++++++++++++++
 2 files changed, 57 insertions(+)

diff --git a/api/apps/restful_apis/agent_api.py b/api/apps/restful_apis/agent_api.py
index 06e8041a4..9a99fbabe 100644
--- a/api/apps/restful_apis/agent_api.py
+++ b/api/apps/restful_apis/agent_api.py
@@ -315,6 +315,10 @@ async def _run_workflow_session(
                     ans["session_id"] = session_id
                     if ans.get("event") == "message":
                         full_content += ans.get("data", {}).get("content", "")
+                        if ans.get("data", {}).get("start_to_think", False):
+                            full_content += "<think>"
+                        elif ans.get("data", {}).get("end_to_think", False):
+                            full_content += "</think>"
                     if ans.get("data", {}).get("reference", None):
                         reference.update(ans["data"]["reference"])
                     if ans.get("event") == "node_finished":
@@ -371,6 +375,10 @@ async def _run_workflow_session(
             ans["session_id"] = session_id
             if ans.get("event") == "message":
                 full_content += ans.get("data", {}).get("content", "")
+                if ans.get("data", {}).get("start_to_think", False):
+                    full_content += "<think>"
+                elif ans.get("data", {}).get("end_to_think", False):
+                    full_content += "</think>"
             if ans.get("data", {}).get("reference", None):
                 reference.update(ans["data"]["reference"])
             if ans.get("event") == "node_finished":
@@ -2467,6 +2475,7 @@ async def webhook_trace(agent_id: str):
         }
     )
 
+
 def _attachment_request_metadata():
     ext = request.args.get("ext")
     mime_type = request.args.get("mime_type")
diff --git a/internal/service/agent.go b/internal/service/agent.go
index c93928972..146708a94 100644
--- a/internal/service/agent.go
+++ b/internal/service/agent.go
@@ -1100,6 +1100,7 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv
 			if canvas.IsInterruptError(err) {
 				s.markRunFailed(ctx2, runID, "interrupt: "+err.Error())
 				if answer != "" {
+					s.persistAgentRunSession(canvasID, sessionID, messageID, userInput, answer, reference)
 					msgData, _ := json.Marshal(canvas.MessageEvent{
 						Content:   answer,
 						Reference: reference,
@@ -1114,6 +1115,7 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv
 				return state, err
 			}
 			if shouldTreatAsCompletedLoopRun(err, answer) {
+				s.persistAgentRunSession(canvasID, sessionID, messageID, userInput, answer, reference)
 				msgData, _ := json.Marshal(canvas.MessageEvent{
 					Content:   answer,
 					Reference: reference,
@@ -1145,6 +1147,7 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv
 		}
 
 		// Emit message + message_end (mirrors Python's ans dict).
+		s.persistAgentRunSession(canvasID, sessionID, messageID, userInput, answer, reference)
 		msgData, _ := json.Marshal(canvas.MessageEvent{
 			Content:   answer,
 			Reference: reference,
@@ -1186,6 +1189,51 @@ func runIDFor(canvasID string, root map[string]any) string {
 	return canvasID
 }
 
+func (s *AgentService) persistAgentRunSession(agentID, sessionID, messageID string, userInput any, answer string, reference []interface{}) {
+	if sessionID == "" || s == nil || s.api4ConversationDAO == nil || dao.DB == nil {
+		return
+	}
+	session, err := s.api4ConversationDAO.GetBySessionID(sessionID, agentID)
+	if err != nil {
+		common.Warn("agent run: load session for update failed", zap.String("agent_id", agentID), zap.String("session_id", sessionID), zap.Error(err))
+		return
+	}
+	if session == nil {
+		return
+	}
+	messages := parseAgentSessionMessages(session.Message)
+	now := time.Now().Unix()
+	if text := stringifyAgentUserInput(userInput); text != "" {
+		messages = append(messages, map[string]interface{}{"role": "user", "content": text, "id": strings.ReplaceAll(uuid.New().String(), "-", ""), "created_at": now})
+	}
+	messages = append(messages, map[string]interface{}{"role": "assistant", "content": answer, "id": messageID, "created_at": now})
+	if raw, err := json.Marshal(messages); err == nil {
+		session.Message = raw
+	}
+	references := parseAgentSessionReferences(session.Reference)
+	references = append(references, normalizeAgentReferenceEntry(map[string]interface{}{"chunks": reference}))
+	if raw, err := json.Marshal(references); err == nil {
+		session.Reference = raw
+	}
+	if err := s.api4ConversationDAO.Update(session); err != nil {
+		common.Warn("agent run: update session failed", zap.String("agent_id", agentID), zap.String("session_id", sessionID), zap.Error(err))
+	}
+}
+
+func stringifyAgentUserInput(userInput any) string {
+	switch v := userInput.(type) {
+	case nil:
+		return ""
+	case string:
+		return v
+	default:
+		if b, err := json.Marshal(v); err == nil {
+			return string(b)
+		}
+		return fmt.Sprint(v)
+	}
+}
+
 // tenantIDFromRoot returns the optional run-tracker tenant id that
 // RunAgent populated on the root map. Runtime components use
 // root["tenant_id"] / state.Sys["tenant_id"] for the caller tenant;