From 9ae17be43cdce13fa4c878ca8a4ea11d99e78c3f Mon Sep 17 00:00:00 2001 From: wizardchen Date: Thu, 16 Jul 2026 16:38:22 +0800 Subject: [PATCH] fix(chat): persist attachment content and fix fallback multi-turn context Persist resolved temporary-attachment content onto the stored user message so later turns can replay it in history. Close the retrieval progress spinner on early pipeline exits, and reshape model-fallback messages with a proper system role plus trailing user question. --- .../service/chat_pipeline/progress.go | 11 ++ .../service/chat_pipeline/progress_test.go | 15 +++ .../service/fallback_messages_test.go | 81 +++++++++++++ .../service/session_knowledge_qa.go | 32 +++++- .../service/session_knowledge_qa_test.go | 19 ++-- internal/handler/session/qa.go | 53 +++++++++ .../temporary_attachment_persist_test.go | 107 ++++++++++++++++++ 7 files changed, 308 insertions(+), 10 deletions(-) create mode 100644 internal/application/service/fallback_messages_test.go create mode 100644 internal/handler/session/temporary_attachment_persist_test.go diff --git a/internal/application/service/chat_pipeline/progress.go b/internal/application/service/chat_pipeline/progress.go index 4b46120e..4e5f8df4 100644 --- a/internal/application/service/chat_pipeline/progress.go +++ b/internal/application/service/chat_pipeline/progress.go @@ -65,6 +65,17 @@ func LastConsolidatedRetrievalStage(eventList []types.EventType, chatManage *typ return last } +// ShouldCloseRetrievalProgress reports whether the consolidated retrieval +// progress window must be closed after a pipeline stage. It returns true when +// either the planned last retrieval stage completed, or a retrieval stage +// short-circuited the pipeline with an error — including ErrSearchNothing, +// which routes into the fallback response. Closing on the error paths prevents +// the frontend "knowledge_search" spinner from hanging forever when the +// pipeline early-returns before reaching the last retrieval stage. +func ShouldCloseRetrievalProgress(stage, lastRetrievalStage types.EventType, stageErr *PluginError) bool { + return stage == lastRetrievalStage || stageErr != nil +} + // BeginRetrievalProgress emits a single pending knowledge_search tool_call. func BeginRetrievalProgress(ctx context.Context, chatManage *types.ChatManage) *StageProgress { if chatManage == nil || chatManage.EventBus == nil { diff --git a/internal/application/service/chat_pipeline/progress_test.go b/internal/application/service/chat_pipeline/progress_test.go index f28eeda7..d9770360 100644 --- a/internal/application/service/chat_pipeline/progress_test.go +++ b/internal/application/service/chat_pipeline/progress_test.go @@ -44,6 +44,21 @@ func TestLastConsolidatedRetrievalStage(t *testing.T) { assert.Equal(t, types.FILTER_TOP_K, LastConsolidatedRetrievalStage(pipeline, cm)) } +func TestShouldCloseRetrievalProgress(t *testing.T) { + last := types.FILTER_TOP_K + + // Normal completion: only the last retrieval stage closes the window. + assert.True(t, ShouldCloseRetrievalProgress(types.FILTER_TOP_K, last, nil)) + assert.False(t, ShouldCloseRetrievalProgress(types.CHUNK_SEARCH_PARALLEL, last, nil)) + + // ErrSearchNothing at an earlier retrieval stage must still close the + // window so the frontend stops spinning before the fallback answer streams. + assert.True(t, ShouldCloseRetrievalProgress(types.CHUNK_SEARCH_PARALLEL, last, ErrSearchNothing)) + + // A hard error at any retrieval stage must also close the window. + assert.True(t, ShouldCloseRetrievalProgress(types.CHUNK_RERANK, last, &PluginError{})) +} + func TestShouldEmitQueryUnderstandProgress(t *testing.T) { cm := &types.ChatManage{PipelineRequest: types.PipelineRequest{EnableRewrite: true}} assert.True(t, ShouldEmitQueryUnderstandProgress(cm)) diff --git a/internal/application/service/fallback_messages_test.go b/internal/application/service/fallback_messages_test.go new file mode 100644 index 00000000..2e226d7b --- /dev/null +++ b/internal/application/service/fallback_messages_test.go @@ -0,0 +1,81 @@ +package service + +import ( + "testing" + + "github.com/Tencent/WeKnora/internal/types" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestBuildFallbackMessages_PrependsSystemAndEndsWithUser guards the model +// fallback path: the LLM input must start with a system message (the fallback +// instruction) and end with the user's question, with history in between. +// Previously the fallback dropped the system role entirely, producing a +// [user, assistant, user] input with no system message. +func TestBuildFallbackMessages_PrependsSystemAndEndsWithUser(t *testing.T) { + cm := &types.ChatManage{} + cm.Query = "这个文件内容" + cm.History = []*types.History{ + {Query: "上一个问题", Answer: "上一个回答"}, + } + + msgs := buildFallbackMessages(cm, "No content directly matched...\n\nUser question: 这个文件内容") + + require.GreaterOrEqual(t, len(msgs), 2) + assert.Equal(t, "system", msgs[0].Role, "fallback input must start with a system message") + assert.Contains(t, msgs[0].Content, "No content directly matched") + + // History is replayed between system and the trailing user turn. + assert.Equal(t, "user", msgs[1].Role) + assert.Equal(t, "上一个问题", msgs[1].Content) + assert.Equal(t, "assistant", msgs[2].Role) + assert.Equal(t, "上一个回答", msgs[2].Content) + + last := msgs[len(msgs)-1] + assert.Equal(t, "user", last.Role, "generation must be prompted by a trailing user turn") + assert.Equal(t, "这个文件内容", last.Content) +} + +// TestBuildFallbackMessages_PrefersRewriteQuery verifies the trailing user turn +// uses the rewritten query when available. +func TestBuildFallbackMessages_PrefersRewriteQuery(t *testing.T) { + cm := &types.ChatManage{} + cm.Query = "它怎么样" + cm.RewriteQuery = "混元大模型性能怎么样" + + msgs := buildFallbackMessages(cm, "fallback instruction") + + last := msgs[len(msgs)-1] + assert.Equal(t, "user", last.Role) + assert.Equal(t, "混元大模型性能怎么样", last.Content) +} + +// TestBuildFallbackMessages_EmptyPromptSkipsSystem ensures we don't inject an +// empty system message when there is no fallback instruction to carry. +func TestBuildFallbackMessages_EmptyPromptSkipsSystem(t *testing.T) { + cm := &types.ChatManage{} + cm.Query = "hello" + + msgs := buildFallbackMessages(cm, " ") + + require.Len(t, msgs, 1) + assert.Equal(t, "user", msgs[0].Role) + assert.Equal(t, "hello", msgs[0].Content) +} + +// TestBuildFallbackMessages_AttachesImagesToUserTurn confirms images ride on the +// trailing user message only when the chat model supports vision. +func TestBuildFallbackMessages_AttachesImagesToUserTurn(t *testing.T) { + cm := &types.ChatManage{} + cm.Query = "看图" + cm.Images = []string{"https://example.com/a.png"} + + cm.ChatModelSupportsVision = false + noVision := buildFallbackMessages(cm, "fallback") + assert.Empty(t, noVision[len(noVision)-1].Images) + + cm.ChatModelSupportsVision = true + withVision := buildFallbackMessages(cm, "fallback") + assert.Equal(t, cm.Images, withVision[len(withVision)-1].Images) +} diff --git a/internal/application/service/session_knowledge_qa.go b/internal/application/service/session_knowledge_qa.go index 608e61b5..810be0ec 100644 --- a/internal/application/service/session_knowledge_qa.go +++ b/internal/application/service/session_knowledge_qa.go @@ -718,7 +718,14 @@ func (s *sessionService) KnowledgeQAByEvent(ctx context.Context, chatpipeline.EndQueryUnderstandProgress(stageCtx, chatManage, understandProgress, understandStart, err) understandProgress = nil } - if retrievalProgress != nil && eventType == lastRetrievalStage { + // Close the consolidated retrieval progress window as soon as retrieval + // is done: either the planned last retrieval stage completed, or a + // retrieval stage short-circuited the pipeline (ErrSearchNothing or a + // hard error). The early returns below (fallback / stage_failed) would + // otherwise skip EndRetrievalProgress, leaving the "knowledge_search" + // tool_call pending — so the frontend keeps spinning on "正在检索知识库" + // forever even though the fallback answer has already streamed. + if retrievalProgress != nil && chatpipeline.ShouldCloseRetrievalProgress(eventType, lastRetrievalStage, err) { chatpipeline.EndRetrievalProgress(stageCtx, chatManage, retrievalProgress, retrievalStart, err) retrievalProgress = nil } @@ -982,10 +989,29 @@ func (s *sessionService) handleModelFallback(ctx context.Context, chatManage *ty } func buildFallbackMessages(chatManage *types.ChatManage, promptContent string) []chat.Message { - messages := make([]chat.Message, 0, len(chatManage.History)*2+1) + messages := make([]chat.Message, 0, len(chatManage.History)*2+2) + + // The model-fallback prompt is a system-style instruction (KB document + // listing + "use general knowledge when nothing matched" guidance). Carry + // it in the system role so the LLM input keeps a proper system message + // instead of starting with a bare user turn. We deliberately do NOT reuse + // the RAG summary system prompt (SummaryConfig.Prompt) here: that template + // forbids prior knowledge ("reply ONLY based on retrieved information"), + // which directly contradicts the fallback's purpose. + if strings.TrimSpace(promptContent) != "" { + messages = append(messages, chat.Message{Role: "system", Content: promptContent}) + } + messages = chatpipeline.AppendHistoryMessages(messages, chatManage.History) - userMsg := chat.Message{Role: "user", Content: promptContent} + // End on the user's actual question so generation is prompted by a user + // turn (the query is also embedded in the system instruction above, but a + // trailing user message keeps the chat shape valid for all providers). + query := chatManage.Query + if rq := strings.TrimSpace(chatManage.RewriteQuery); rq != "" { + query = rq + } + userMsg := chat.Message{Role: "user", Content: query} if chatManage.ChatModelSupportsVision && len(chatManage.Images) > 0 { userMsg.Images = chatManage.Images } diff --git a/internal/application/service/session_knowledge_qa_test.go b/internal/application/service/session_knowledge_qa_test.go index 6800b1ca..b1479f40 100644 --- a/internal/application/service/session_knowledge_qa_test.go +++ b/internal/application/service/session_knowledge_qa_test.go @@ -139,11 +139,16 @@ func TestHandleModelFallback_IncludesHistoryMessages(t *testing.T) { svc.handleModelFallback(context.Background(), cm) - require.Len(t, chatModel.lastMessages, 3) - assert.Equal(t, "user", chatModel.lastMessages[0].Role) - assert.Equal(t, "先介绍一下 WeKnora", chatModel.lastMessages[0].Content) - assert.Equal(t, "assistant", chatModel.lastMessages[1].Role) - assert.Equal(t, "WeKnora 是一个知识库问答系统。", chatModel.lastMessages[1].Content) - assert.Equal(t, "user", chatModel.lastMessages[2].Role) - assert.Contains(t, chatModel.lastMessages[2].Content, "现在还能继续讲吗？") + // Corrected fallback shape: a system message carries the fallback + // instruction, history is replayed in the middle, and the turn ends on the + // user's question. Previously the system message was dropped entirely. + require.Len(t, chatModel.lastMessages, 4) + assert.Equal(t, "system", chatModel.lastMessages[0].Role) + assert.Contains(t, chatModel.lastMessages[0].Content, "Answer the latest user question") + assert.Equal(t, "user", chatModel.lastMessages[1].Role) + assert.Equal(t, "先介绍一下 WeKnora", chatModel.lastMessages[1].Content) + assert.Equal(t, "assistant", chatModel.lastMessages[2].Role) + assert.Equal(t, "WeKnora 是一个知识库问答系统。", chatModel.lastMessages[2].Content) + assert.Equal(t, "user", chatModel.lastMessages[3].Role) + assert.Contains(t, chatModel.lastMessages[3].Content, "现在还能继续讲吗？") } diff --git a/internal/handler/session/qa.go b/internal/handler/session/qa.go index fdf5664b..1c87a9c5 100644 --- a/internal/handler/session/qa.go +++ b/internal/handler/session/qa.go @@ -1146,6 +1146,12 @@ func (h *Handler) resolveTemporaryAttachments(streamCtx *sseStreamContext, reqCt attachments = filtered } reqCtx.attachments = append(reqCtx.attachments, attachments...) + // Persist the freshly selected content back onto the stored user message. + // The message was created with metadata-only attachment entries (content is + // selected here, after the SSE stream is live), so without this write a + // later Agent-mode turn rebuilds history from the Attachments column and + // sees empty attachments (see buildUserHistoryMessage in agent_history.go). + h.persistResolvedAttachmentContent(ctx, reqCtx, attachments) if reqCtx.customAgent != nil && reqCtx.customAgent.Config.ImageUploadEnabled { for _, imageURL := range temporaryResult.ImageURLs { reqCtx.images = append(reqCtx.images, ImageAttachment{URL: imageURL}) @@ -1153,6 +1159,53 @@ func (h *Handler) resolveTemporaryAttachments(streamCtx *sseStreamContext, reqCt } } +// persistResolvedAttachmentContent writes the parsed content of pre-uploaded +// attachments back onto the stored user message so multi-turn history can +// replay it. Entries are matched by their temporary-document ID; only those +// present on the message are enriched. Failures are logged but never bubble up +// — losing the write only degrades follow-up context, it must not fail the turn. +func (h *Handler) persistResolvedAttachmentContent( + ctx context.Context, reqCtx *qaRequestContext, resolved types.MessageAttachments, +) { + if reqCtx.userMessageID == "" || len(resolved) == 0 { + return + } + // Detach from the request/stream lifetime and pin the session tenant so a + // user-triggered stop (which cancels asyncCtx) does not drop the write. + updateCtx := context.WithValue( + context.WithoutCancel(ctx), types.TenantIDContextKey, reqCtx.session.TenantID, + ) + msg, err := h.messageService.GetMessage(updateCtx, reqCtx.sessionID, reqCtx.userMessageID) + if err != nil || msg == nil { + logger.Warnf(updateCtx, "persist attachment content: load user message %s failed: %v", + reqCtx.userMessageID, err) + return + } + byID := make(map[string]types.MessageAttachment, len(resolved)) + for _, att := range resolved { + if att.ID != "" { + byID[att.ID] = att + } + } + changed := false + for i := range msg.Attachments { + if msg.Attachments[i].ID == "" { + continue + } + if enriched, ok := byID[msg.Attachments[i].ID]; ok { + msg.Attachments[i] = enriched + changed = true + } + } + if !changed { + return + } + if err := h.messageService.UpdateMessage(updateCtx, msg); err != nil { + logger.Warnf(updateCtx, "persist attachment content: update user message %s failed: %v", + reqCtx.userMessageID, err) + } +} + // hasPendingAttachments reports whether any of the given documents is still // uploading or being parsed. func (h *Handler) hasPendingAttachments(ctx context.Context, tenantID uint64, sessionID string, ids []string) bool { diff --git a/internal/handler/session/temporary_attachment_persist_test.go b/internal/handler/session/temporary_attachment_persist_test.go new file mode 100644 index 00000000..79d00744 --- /dev/null +++ b/internal/handler/session/temporary_attachment_persist_test.go @@ -0,0 +1,107 @@ +package session + +import ( + "context" + "testing" + + "github.com/Tencent/WeKnora/internal/types" + "github.com/Tencent/WeKnora/internal/types/interfaces" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// attachmentMsgStub implements just the two MessageService methods used by +// persistResolvedAttachmentContent; the rest are inherited from the embedded +// interface and panic if unexpectedly called. +type attachmentMsgStub struct { + interfaces.MessageService + stored *types.Message + updated *types.Message + updateCalls int +} + +func (s *attachmentMsgStub) GetMessage(_ context.Context, _ string, _ string) (*types.Message, error) { + return s.stored, nil +} + +func (s *attachmentMsgStub) UpdateMessage(_ context.Context, message *types.Message) error { + s.updateCalls++ + s.updated = message + return nil +} + +func newAttachmentReqCtx() *qaRequestContext { + return &qaRequestContext{ + sessionID: "sess-1", + userMessageID: "user-msg-1", + session: &types.Session{TenantID: 42}, + } +} + +// TestPersistResolvedAttachmentContent_EnrichesMatchingAttachments is the core +// regression guard for PR #2086: the user message is created with +// metadata-only attachment entries, and the parsed content selected after the +// SSE stream starts must be written back so multi-turn (Agent-mode) history can +// replay it via the Attachments column. +func TestPersistResolvedAttachmentContent_EnrichesMatchingAttachments(t *testing.T) { + stub := &attachmentMsgStub{ + stored: &types.Message{ + ID: "user-msg-1", + SessionID: "sess-1", + Role: "user", + Attachments: types.MessageAttachments{ + {ID: "doc-1", FileName: "report.pdf", FileType: ".pdf"}, + // Legacy inline upload without an ID must stay untouched. + {FileName: "note.txt", FileType: ".txt", Content: "legacy inline"}, + }, + }, + } + h := &Handler{messageService: stub} + + resolved := types.MessageAttachments{ + { + ID: "doc-1", FileName: "report.pdf", FileType: ".pdf", + Content: "parsed body text", ContentMode: "selected_chunks", + SelectedChunks: 1, TotalChunks: 3, TokenCount: 12, + }, + } + h.persistResolvedAttachmentContent(context.Background(), newAttachmentReqCtx(), resolved) + + require.Equal(t, 1, stub.updateCalls, "matching attachment must trigger one persist") + require.NotNil(t, stub.updated) + require.Len(t, stub.updated.Attachments, 2) + assert.Equal(t, "parsed body text", stub.updated.Attachments[0].Content) + assert.Equal(t, "selected_chunks", stub.updated.Attachments[0].ContentMode) + assert.Equal(t, 3, stub.updated.Attachments[0].TotalChunks) + // Legacy inline attachment (no ID) is preserved verbatim. + assert.Equal(t, "legacy inline", stub.updated.Attachments[1].Content) +} + +func TestPersistResolvedAttachmentContent_NoMatchSkipsUpdate(t *testing.T) { + stub := &attachmentMsgStub{ + stored: &types.Message{ + ID: "user-msg-1", SessionID: "sess-1", Role: "user", + Attachments: types.MessageAttachments{{ID: "doc-1", FileName: "a.pdf"}}, + }, + } + h := &Handler{messageService: stub} + + resolved := types.MessageAttachments{{ID: "doc-other", Content: "x"}} + h.persistResolvedAttachmentContent(context.Background(), newAttachmentReqCtx(), resolved) + + assert.Equal(t, 0, stub.updateCalls, "no matching ID must not persist") +} + +func TestPersistResolvedAttachmentContent_GuardsEmptyInput(t *testing.T) { + stub := &attachmentMsgStub{} + h := &Handler{messageService: stub} + + // No resolved attachments. + h.persistResolvedAttachmentContent(context.Background(), newAttachmentReqCtx(), nil) + // No user message ID. + rc := newAttachmentReqCtx() + rc.userMessageID = "" + h.persistResolvedAttachmentContent(context.Background(), rc, types.MessageAttachments{{ID: "doc-1", Content: "x"}}) + + assert.Equal(t, 0, stub.updateCalls) +}