From e1330423acea5fcc72beb5e01c45143afd110bc3 Mon Sep 17 00:00:00 2001 From: euvre <93761161+euvre@users.noreply.github.com> Date: Fri, 14 Aug 2026 01:00:08 -0700 Subject: [PATCH] fix(agent): keep think tags in persisted agent session messages (#18222) --- internal/handler/agent.go | 8 +++ internal/handler/agent_test.go | 40 +++++++++++++ internal/service/agent.go | 22 +++++-- internal/service/agent_persistence_test.go | 69 ++++++++++++++++++++++ 4 files changed, 135 insertions(+), 4 deletions(-) diff --git a/internal/handler/agent.go b/internal/handler/agent.go index b3dac9f881..761f25dd02 100644 --- a/internal/handler/agent.go +++ b/internal/handler/agent.go @@ -1395,6 +1395,14 @@ func (h *AgentHandler) AgentChatCompletions(c *gin.Context) { if c, ok := evData["content"].(string); ok { fullContent += c } + // Mirror Python agent_api.py: the reasoning segment stays + // wrapped in tags in the aggregated answer so the + // chat UI can render the "thought" section. + if st, _ := evData["start_to_think"].(bool); st { + fullContent += "" + } else if et, _ := evData["end_to_think"].(bool); et { + fullContent += "" + } } if ref, _ := evData["reference"].(map[string]any); ref != nil { for k, v := range ref { diff --git a/internal/handler/agent_test.go b/internal/handler/agent_test.go index 77ce4e9cb5..5baad8e1d8 100644 --- a/internal/handler/agent_test.go +++ b/internal/handler/agent_test.go @@ -866,6 +866,46 @@ func TestAgentChatCompletions_DefaultBranchNonStreaming(t *testing.T) { } } +// TestAgentChatCompletions_NonStreamingPreservesThinkMarkers covers the +// non-streaming aggregation: start_to_think/end_to_think message events must +// survive as tags in the final content, mirroring Python +// agent_api.py, so clients can render the "thought" section. +func TestAgentChatCompletions_NonStreamingPreservesThinkMarkers(t *testing.T) { + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("POST", "/api/v1/agents/chat/completions", + strings.NewReader(`{"agent_id":"a1","query":"hello","stream":false}`)) + c.Request.Header.Set("Content-Type", "application/json") + c.Set("user", &entity.User{ID: "u1"}) + c.Set("user_id", "u1") + + runner := &stubChatRunner{events: []canvas.RunEvent{ + {Type: "message", MessageID: "msg-1", SessionID: "sess-1", Data: `{"content":"","start_to_think":true}`}, + {Type: "message", MessageID: "msg-1", SessionID: "sess-1", Data: `{"content":"reasoning trace"}`}, + {Type: "message", MessageID: "msg-1", SessionID: "sess-1", Data: `{"content":"","end_to_think":true}`}, + {Type: "message", MessageID: "msg-1", SessionID: "sess-1", Data: `{"content":"final answer"}`}, + {Type: "message_end", MessageID: "msg-1", SessionID: "sess-1", Data: `{}`}, + {Type: "done", Data: ""}, + }} + h := &AgentHandler{chatRunner: runner} + h.AgentChatCompletions(c) + + body := w.Body.String() + if !strings.Contains(body, `"code":0`) { + t.Fatalf("body should contain success code, got %q", body) + } + var response map[string]any + if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil { + t.Fatalf("decode response: %v", err) + } + data, _ := response["data"].(map[string]any) + ansData, _ := data["data"].(map[string]any) + if got, _ := ansData["content"].(string); got != "reasoning tracefinal answer" { + t.Errorf("aggregated content = %q, want the think section preserved as %q", got, "reasoning tracefinal answer") + } +} + type emptySessionCaptureRunner struct { sessionID string } diff --git a/internal/service/agent.go b/internal/service/agent.go index 7e4dabe513..638ed602d5 100644 --- a/internal/service/agent.go +++ b/internal/service/agent.go @@ -253,6 +253,19 @@ func splitInlineThink(answer, thinking string) (string, string) { return answer, thinking } +// agentSessionMessageContent mirrors how Python's canvas_service.completion +// accumulates the persisted assistant message: reasoning streams first, +// bracketed by start_to_think/end_to_think markers, so the stored content +// keeps the thinking segment wrapped in tags. The chat UI refetches +// the session message list after streaming and needs those tags for +// MarkdownContent to render the "thought" section (chat.thought). +func agentSessionMessageContent(answer, thinking string) string { + if thinking == "" { + return answer + } + return "" + thinking + "" + answer +} + func splitMessageContent(content string) []string { if content == "" { return nil @@ -2110,7 +2123,7 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv if answer != "" { appendAssistantHistory(state, partialAssistantOutput(answer, downloads)) } - if persistErr := s.persistAgentRunSession(ctx, canvasID, userID, sessionID, messageID, userInput, answer, referencePayload, dsl, state, answer != ""); persistErr != nil { + if persistErr := s.persistAgentRunSession(ctx, canvasID, userID, sessionID, messageID, userInput, answer, thinking, referencePayload, dsl, state, answer != ""); persistErr != nil { return nil, fmt.Errorf("persist interrupted agent session: %w: %w", persistErr, ErrAgentStorageError) } if answer != "" && shouldEmitMessage { @@ -2127,7 +2140,7 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv } if shouldTreatAsCompletedLoopRun(err, answer) { appendAssistantHistory(state, assistantOutput) - if persistErr := s.persistAgentRunSession(ctx, canvasID, userID, sessionID, messageID, userInput, answer, referencePayload, dsl, state, true); persistErr != nil { + if persistErr := s.persistAgentRunSession(ctx, canvasID, userID, sessionID, messageID, userInput, answer, thinking, referencePayload, dsl, state, true); persistErr != nil { s.markRunFailed(ctx2, runID, "persist session: "+persistErr.Error()) return nil, fmt.Errorf("persist agent session: %w: %w", persistErr, ErrAgentStorageError) } @@ -2163,7 +2176,7 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv // Emit message + message_end (mirrors Python's ans dict). appendAssistantHistory(state, assistantOutput) - if persistErr := s.persistAgentRunSession(ctx, canvasID, userID, sessionID, messageID, userInput, answer, referencePayload, dsl, state, true); persistErr != nil { + if persistErr := s.persistAgentRunSession(ctx, canvasID, userID, sessionID, messageID, userInput, answer, thinking, referencePayload, dsl, state, true); persistErr != nil { s.markRunFailed(ctx2, runID, "persist session: "+persistErr.Error()) return nil, fmt.Errorf("persist agent session: %w: %w", persistErr, ErrAgentStorageError) } @@ -2291,6 +2304,7 @@ func (s *AgentService) persistAgentRunSession( agentID, userID, sessionID, messageID string, userInput any, answer string, + thinking string, reference map[string]interface{}, runDSL map[string]any, state *canvas.CanvasState, @@ -2313,7 +2327,7 @@ func (s *AgentService) persistAgentRunSession( messages = append(messages, map[string]interface{}{"role": "user", "content": text, "id": utility.GenerateToken(), "created_at": now}) } if appendAssistantMessage { - messages = append(messages, map[string]interface{}{"role": "assistant", "content": answer, "id": messageID, "created_at": now}) + messages = append(messages, map[string]interface{}{"role": "assistant", "content": agentSessionMessageContent(answer, thinking), "id": messageID, "created_at": now}) } if raw, err := json.Marshal(messages); err == nil { session.Message = raw diff --git a/internal/service/agent_persistence_test.go b/internal/service/agent_persistence_test.go index 9c9910bd09..7a094bc213 100644 --- a/internal/service/agent_persistence_test.go +++ b/internal/service/agent_persistence_test.go @@ -90,3 +90,72 @@ func TestAgentRunSessionUpdateFailurePreventsSuccessEvents(t *testing.T) { } } } + +func TestAgentSessionMessageContent(t *testing.T) { + tests := []struct { + name string + answer string + thinking string + want string + }{ + {"no thinking keeps answer as-is", "final answer", "", "final answer"}, + {"thinking wrapped before answer", "final answer", "reasoning trace", "reasoning tracefinal answer"}, + {"empty answer keeps thinking section", "", "reasoning trace", "reasoning trace"}, + {"inline think tags preserved when no separate thinking", "pre inline post", "", "pre inline post"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := agentSessionMessageContent(tt.answer, tt.thinking); got != tt.want { + t.Errorf("agentSessionMessageContent(%q, %q) = %q, want %q", tt.answer, tt.thinking, got, tt.want) + } + }) + } +} + +// TestPersistAgentRunSessionPreservesThinking guards the chat.thought +// ("思考完成") indicator: the assistant message stored on the agent session +// must keep the reasoning segment wrapped in tags, mirroring Python +// canvas_service.completion. The chat UI refetches the session message list +// after streaming and renders the thought section from those tags. +func TestPersistAgentRunSessionPreservesThinking(t *testing.T) { + testDB := setupServiceTestDB(t) + if err := testDB.AutoMigrate(&entity.API4Conversation{}); err != nil { + t.Fatalf("migrate: %v", err) + } + originalDB := dao.DB + dao.DB = testDB + t.Cleanup(func() { dao.DB = originalDB }) + + if err := testDB.Create(&entity.API4Conversation{ + ID: "session-think", + DialogID: "canvas-think", + UserID: "user-1", + Message: json.RawMessage(`[]`), + Reference: json.RawMessage(`[]`), + }).Error; err != nil { + t.Fatalf("create session: %v", err) + } + + svc := NewAgentService() + if err := svc.persistAgentRunSession(context.Background(), "canvas-think", "user-1", "session-think", "msg-think-1", "question", "final answer", "reasoning trace", map[string]interface{}{}, nil, nil, true); err != nil { + t.Fatalf("persist: %v", err) + } + + var conv entity.API4Conversation + if err := testDB.First(&conv, "id = ?", "session-think").Error; err != nil { + t.Fatalf("reload session: %v", err) + } + var messages []map[string]any + if err := json.Unmarshal(conv.Message, &messages); err != nil { + t.Fatalf("decode messages: %v", err) + } + if len(messages) != 2 { + t.Fatalf("messages = %d, want 2 (user + assistant)", len(messages)) + } + if role, _ := messages[1]["role"].(string); role != "assistant" { + t.Fatalf("second message role = %q, want assistant", role) + } + if got, _ := messages[1]["content"].(string); got != "reasoning tracefinal answer" { + t.Fatalf("persisted assistant content = %q, want %q", got, "reasoning tracefinal answer") + } +}