From a49339843c41c5b9be0978f77c3f60cf5bd3a9a8 Mon Sep 17 00:00:00 2001 From: kfjie Date: Sat, 18 Jul 2026 12:15:33 +0000 Subject: [PATCH 1/3] fix: preserve context after successful assistant retry --- .../conversation/service_branch.go | 38 +++++++++-- .../conversation/service_branch_test.go | 63 +++++++++++++++++++ 2 files changed, 97 insertions(+), 4 deletions(-) diff --git a/backend/internal/application/conversation/service_branch.go b/backend/internal/application/conversation/service_branch.go index 69bc9941..d2bf00cf 100644 --- a/backend/internal/application/conversation/service_branch.go +++ b/backend/internal/application/conversation/service_branch.go @@ -147,8 +147,21 @@ func normalizeDefaultBranchContext( return nil, nil } - end := len(ancestors) - for end > 0 && !isContextMessage(&ancestors[end-1]) { + contextMessages := append([]model.Message(nil), ancestors...) + for index := range contextMessages { + if isRecoveredAssistantRetryUser(contextMessages, index) { + // A successful assistant retry makes the reused user message valid context. + // Promote only the in-memory copy so existing conversations recover without + // rewriting historical rows, while downstream attachment collection also + // sees the retried user message as usable context. + contextMessages[index].Status = "success" + contextMessages[index].ErrorCode = "" + contextMessages[index].ErrorMessage = "" + } + } + + end := len(contextMessages) + for end > 0 && !isContextMessage(&contextMessages[end-1]) { end-- } if end == 0 { @@ -157,13 +170,13 @@ func normalizeDefaultBranchContext( start := 0 for index := end - 1; index >= 0; index-- { - if !isContextMessage(&ancestors[index]) { + if !isContextMessage(&contextMessages[index]) { start = index + 1 break } } - normalized := append([]model.Message(nil), ancestors[start:end]...) + normalized := append([]model.Message(nil), contextMessages[start:end]...) if len(normalized) == 0 { return nil, nil } @@ -171,6 +184,23 @@ func normalizeDefaultBranchContext( return normalized, &nextParent } +func isRecoveredAssistantRetryUser(messages []model.Message, index int) bool { + if index < 0 || index+1 >= len(messages) { + return false + } + userMessage := &messages[index] + retryAssistant := &messages[index+1] + if userMessage.Role != "user" || !strings.EqualFold(strings.TrimSpace(userMessage.Status), "error") { + return false + } + if retryAssistant.Role != "assistant" || + !strings.EqualFold(strings.TrimSpace(retryAssistant.BranchReason), "retry") || + !isContextMessage(retryAssistant) { + return false + } + return retryAssistant.ParentMessageID != nil && *retryAssistant.ParentMessageID == userMessage.ID +} + func isContextMessage(item *model.Message) bool { if item == nil { return false diff --git a/backend/internal/application/conversation/service_branch_test.go b/backend/internal/application/conversation/service_branch_test.go index 03d11e8e..2627245f 100644 --- a/backend/internal/application/conversation/service_branch_test.go +++ b/backend/internal/application/conversation/service_branch_test.go @@ -1,11 +1,74 @@ package conversation import ( + "encoding/json" "testing" model "github.com/DEEIX-AI/DEEIX-Chat/backend/internal/domain/conversation" ) +func TestNormalizeDefaultBranchContextKeepsHistoryAfterSuccessfulAssistantRetry(t *testing.T) { + rootUserID := uint(1) + rootAssistantID := uint(2) + failedUserID := uint(3) + retryAssistantID := uint(4) + attachments, err := json.Marshal([]attachmentSnapshotRef{{FileID: "file_from_history"}}) + if err != nil { + t.Fatalf("marshal attachments: %v", err) + } + ancestors := []model.Message{ + {ID: rootUserID, PublicID: "msg_root_user", Role: "user", Status: "success", Attachments: string(attachments)}, + {ID: rootAssistantID, PublicID: "msg_root_assistant", ParentMessageID: &rootUserID, Role: "assistant", Status: "success"}, + {ID: failedUserID, PublicID: "msg_failed_user", ParentMessageID: &rootAssistantID, Role: "user", Status: "error"}, + {ID: retryAssistantID, PublicID: "msg_retry_assistant", ParentMessageID: &failedUserID, Role: "assistant", BranchReason: "retry", Status: "success"}, + } + + normalized, parent := normalizeDefaultBranchContext(ancestors, &ancestors[3]) + + if parent == nil || parent.ID != retryAssistantID { + t.Fatalf("expected successful retry assistant as parent, got %#v", parent) + } + if len(normalized) != len(ancestors) { + t.Fatalf("expected complete history after successful retry, got %#v", normalized) + } + for index, expected := range ancestors { + if normalized[index].ID != expected.ID { + t.Fatalf("expected message %d at index %d, got %#v", expected.ID, index, normalized) + } + } + if normalized[2].Status != "success" { + t.Fatalf("expected recovered user to be usable as context, got status %q", normalized[2].Status) + } + if ancestors[2].Status != "error" { + t.Fatalf("expected normalization to leave persisted message state unchanged, got status %q", ancestors[2].Status) + } + fileIDs := collectConversationFileIDs(normalized, nil) + if len(fileIDs) != 1 || fileIDs[0] != "file_from_history" { + t.Fatalf("expected historical attachment to remain in context, got %#v", fileIDs) + } +} + +func TestNormalizeDefaultBranchContextDoesNotRecoverFailedRetry(t *testing.T) { + rootUserID := uint(1) + rootAssistantID := uint(2) + failedUserID := uint(3) + ancestors := []model.Message{ + {ID: rootUserID, PublicID: "msg_root_user", Role: "user", Status: "success"}, + {ID: rootAssistantID, PublicID: "msg_root_assistant", ParentMessageID: &rootUserID, Role: "assistant", Status: "success"}, + {ID: failedUserID, PublicID: "msg_failed_user", ParentMessageID: &rootAssistantID, Role: "user", Status: "error"}, + {ID: 4, PublicID: "msg_failed_retry", ParentMessageID: &failedUserID, Role: "assistant", BranchReason: "retry", Status: "error"}, + } + + normalized, parent := normalizeDefaultBranchContext(ancestors, &ancestors[3]) + + if parent == nil || parent.ID != rootAssistantID { + t.Fatalf("expected latest successful ancestor as parent, got %#v", parent) + } + if len(normalized) != 2 || normalized[0].ID != rootUserID || normalized[1].ID != rootAssistantID { + t.Fatalf("expected failed retry tail removed from context, got %#v", normalized) + } +} + func TestNormalizeDefaultBranchContextSkipsFailedTail(t *testing.T) { rootID := uint(1) assistantID := uint(2) From 9209a42b6740b0f39980d31c99ea1a4cacd0e9df Mon Sep 17 00:00:00 2001 From: kfjie Date: Sat, 18 Jul 2026 12:40:17 +0000 Subject: [PATCH 2/3] fix: validate recovered assistant retry lineage --- .../conversation/service_branch.go | 1 + .../conversation/service_branch_test.go | 77 +++++++++++++++++-- 2 files changed, 72 insertions(+), 6 deletions(-) diff --git a/backend/internal/application/conversation/service_branch.go b/backend/internal/application/conversation/service_branch.go index d2bf00cf..f1d406c2 100644 --- a/backend/internal/application/conversation/service_branch.go +++ b/backend/internal/application/conversation/service_branch.go @@ -195,6 +195,7 @@ func isRecoveredAssistantRetryUser(messages []model.Message, index int) bool { } if retryAssistant.Role != "assistant" || !strings.EqualFold(strings.TrimSpace(retryAssistant.BranchReason), "retry") || + retryAssistant.SourceMessageID == nil || !isContextMessage(retryAssistant) { return false } diff --git a/backend/internal/application/conversation/service_branch_test.go b/backend/internal/application/conversation/service_branch_test.go index 2627245f..ff771634 100644 --- a/backend/internal/application/conversation/service_branch_test.go +++ b/backend/internal/application/conversation/service_branch_test.go @@ -12,15 +12,20 @@ func TestNormalizeDefaultBranchContextKeepsHistoryAfterSuccessfulAssistantRetry( rootAssistantID := uint(2) failedUserID := uint(3) retryAssistantID := uint(4) - attachments, err := json.Marshal([]attachmentSnapshotRef{{FileID: "file_from_history"}}) + failedAssistantID := uint(5) + historicalAttachments, err := json.Marshal([]attachmentSnapshotRef{{FileID: "file_from_history"}}) if err != nil { t.Fatalf("marshal attachments: %v", err) } + retriedAttachments, err := json.Marshal([]attachmentSnapshotRef{{FileID: "file_from_retried_user"}}) + if err != nil { + t.Fatalf("marshal retried user attachments: %v", err) + } ancestors := []model.Message{ - {ID: rootUserID, PublicID: "msg_root_user", Role: "user", Status: "success", Attachments: string(attachments)}, + {ID: rootUserID, PublicID: "msg_root_user", Role: "user", Status: "success", Attachments: string(historicalAttachments)}, {ID: rootAssistantID, PublicID: "msg_root_assistant", ParentMessageID: &rootUserID, Role: "assistant", Status: "success"}, - {ID: failedUserID, PublicID: "msg_failed_user", ParentMessageID: &rootAssistantID, Role: "user", Status: "error"}, - {ID: retryAssistantID, PublicID: "msg_retry_assistant", ParentMessageID: &failedUserID, Role: "assistant", BranchReason: "retry", Status: "success"}, + {ID: failedUserID, PublicID: "msg_failed_user", ParentMessageID: &rootAssistantID, Role: "user", Status: "error", Attachments: string(retriedAttachments)}, + {ID: retryAssistantID, PublicID: "msg_retry_assistant", ParentMessageID: &failedUserID, SourceMessageID: &failedAssistantID, Role: "assistant", BranchReason: "retry", Status: "success"}, } normalized, parent := normalizeDefaultBranchContext(ancestors, &ancestors[3]) @@ -43,7 +48,7 @@ func TestNormalizeDefaultBranchContextKeepsHistoryAfterSuccessfulAssistantRetry( t.Fatalf("expected normalization to leave persisted message state unchanged, got status %q", ancestors[2].Status) } fileIDs := collectConversationFileIDs(normalized, nil) - if len(fileIDs) != 1 || fileIDs[0] != "file_from_history" { + if len(fileIDs) != 2 || fileIDs[0] != "file_from_history" || fileIDs[1] != "file_from_retried_user" { t.Fatalf("expected historical attachment to remain in context, got %#v", fileIDs) } } @@ -52,11 +57,12 @@ func TestNormalizeDefaultBranchContextDoesNotRecoverFailedRetry(t *testing.T) { rootUserID := uint(1) rootAssistantID := uint(2) failedUserID := uint(3) + failedAssistantID := uint(5) ancestors := []model.Message{ {ID: rootUserID, PublicID: "msg_root_user", Role: "user", Status: "success"}, {ID: rootAssistantID, PublicID: "msg_root_assistant", ParentMessageID: &rootUserID, Role: "assistant", Status: "success"}, {ID: failedUserID, PublicID: "msg_failed_user", ParentMessageID: &rootAssistantID, Role: "user", Status: "error"}, - {ID: 4, PublicID: "msg_failed_retry", ParentMessageID: &failedUserID, Role: "assistant", BranchReason: "retry", Status: "error"}, + {ID: 4, PublicID: "msg_failed_retry", ParentMessageID: &failedUserID, SourceMessageID: &failedAssistantID, Role: "assistant", BranchReason: "retry", Status: "error"}, } normalized, parent := normalizeDefaultBranchContext(ancestors, &ancestors[3]) @@ -69,6 +75,65 @@ func TestNormalizeDefaultBranchContextDoesNotRecoverFailedRetry(t *testing.T) { } } +func TestIsRecoveredAssistantRetryUserRequiresGenuineUsableRetry(t *testing.T) { + failedUserID := uint(10) + sourceAssistantID := uint(11) + otherUserID := uint(12) + tests := []struct { + name string + user model.Message + assistant model.Message + expected bool + }{ + { + name: "successful retry", + user: model.Message{ID: failedUserID, Role: "user", Status: "error"}, + assistant: model.Message{Role: "assistant", ParentMessageID: &failedUserID, SourceMessageID: &sourceAssistantID, BranchReason: "retry", Status: "success"}, + expected: true, + }, + { + name: "interrupted retry", + user: model.Message{ID: failedUserID, Role: "user", Status: "error"}, + assistant: model.Message{Role: "assistant", ParentMessageID: &failedUserID, SourceMessageID: &sourceAssistantID, BranchReason: "retry", Status: "interrupted"}, + expected: true, + }, + { + name: "failed retry", + user: model.Message{ID: failedUserID, Role: "user", Status: "error"}, + assistant: model.Message{Role: "assistant", ParentMessageID: &failedUserID, SourceMessageID: &sourceAssistantID, BranchReason: "retry", Status: "error"}, + }, + { + name: "non-retry assistant", + user: model.Message{ID: failedUserID, Role: "user", Status: "error"}, + assistant: model.Message{Role: "assistant", ParentMessageID: &failedUserID, SourceMessageID: &sourceAssistantID, BranchReason: "default", Status: "success"}, + }, + { + name: "retry without source", + user: model.Message{ID: failedUserID, Role: "user", Status: "error"}, + assistant: model.Message{Role: "assistant", ParentMessageID: &failedUserID, BranchReason: "retry", Status: "success"}, + }, + { + name: "retry with different parent", + user: model.Message{ID: failedUserID, Role: "user", Status: "error"}, + assistant: model.Message{Role: "assistant", ParentMessageID: &otherUserID, SourceMessageID: &sourceAssistantID, BranchReason: "retry", Status: "success"}, + }, + { + name: "already successful user", + user: model.Message{ID: failedUserID, Role: "user", Status: "success"}, + assistant: model.Message{Role: "assistant", ParentMessageID: &failedUserID, SourceMessageID: &sourceAssistantID, BranchReason: "retry", Status: "success"}, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + messages := []model.Message{test.user, test.assistant} + if actual := isRecoveredAssistantRetryUser(messages, 0); actual != test.expected { + t.Fatalf("expected recovered=%v, got %v", test.expected, actual) + } + }) + } +} + func TestNormalizeDefaultBranchContextSkipsFailedTail(t *testing.T) { rootID := uint(1) assistantID := uint(2) From 896addcbd016cddcd7528ed0be527d80477d43f6 Mon Sep 17 00:00:00 2001 From: Chenyme <118253778+chenyme@users.noreply.github.com> Date: Sun, 19 Jul 2026 13:56:09 +0800 Subject: [PATCH 3/3] feat: implement recovery of assistant retry user states in conversation context --- .../conversation/service_branch.go | 35 +++++++++++------ .../conversation/service_branch_test.go | 38 +++++++++++++++++++ .../conversation/service_message_send.go | 2 + 3 files changed, 63 insertions(+), 12 deletions(-) diff --git a/backend/internal/application/conversation/service_branch.go b/backend/internal/application/conversation/service_branch.go index f1d406c2..9fe74ae3 100644 --- a/backend/internal/application/conversation/service_branch.go +++ b/backend/internal/application/conversation/service_branch.go @@ -147,18 +147,7 @@ func normalizeDefaultBranchContext( return nil, nil } - contextMessages := append([]model.Message(nil), ancestors...) - for index := range contextMessages { - if isRecoveredAssistantRetryUser(contextMessages, index) { - // A successful assistant retry makes the reused user message valid context. - // Promote only the in-memory copy so existing conversations recover without - // rewriting historical rows, while downstream attachment collection also - // sees the retried user message as usable context. - contextMessages[index].Status = "success" - contextMessages[index].ErrorCode = "" - contextMessages[index].ErrorMessage = "" - } - } + contextMessages := recoverAssistantRetryUserStates(ancestors) end := len(contextMessages) for end > 0 && !isContextMessage(&contextMessages[end-1]) { @@ -184,6 +173,28 @@ func normalizeDefaultBranchContext( return normalized, &nextParent } +// recoverAssistantRetryUserStates makes a reused user message valid context +// after its assistant retry produced usable output. Persisted history remains +// unchanged so the original failed run can still be diagnosed. +func recoverAssistantRetryUserStates(messages []model.Message) []model.Message { + var recovered []model.Message + for index := range messages { + if !isRecoveredAssistantRetryUser(messages, index) { + continue + } + if recovered == nil { + recovered = append([]model.Message(nil), messages...) + } + recovered[index].Status = "success" + recovered[index].ErrorCode = "" + recovered[index].ErrorMessage = "" + } + if recovered == nil { + return messages + } + return recovered +} + func isRecoveredAssistantRetryUser(messages []model.Message, index int) bool { if index < 0 || index+1 >= len(messages) { return false diff --git a/backend/internal/application/conversation/service_branch_test.go b/backend/internal/application/conversation/service_branch_test.go index ff771634..8ef4847c 100644 --- a/backend/internal/application/conversation/service_branch_test.go +++ b/backend/internal/application/conversation/service_branch_test.go @@ -53,6 +53,44 @@ func TestNormalizeDefaultBranchContextKeepsHistoryAfterSuccessfulAssistantRetry( } } +func TestRecoverAssistantRetryUserStatesRestoresExpandedContext(t *testing.T) { + rootUserID := uint(1) + rootAssistantID := uint(2) + failedUserID := uint(3) + sourceAssistantID := uint(4) + retryAssistantID := uint(5) + retriedAttachments, err := json.Marshal([]attachmentSnapshotRef{{FileID: "file_from_retried_user"}}) + if err != nil { + t.Fatalf("marshal retried user attachments: %v", err) + } + messages := []model.Message{ + {ID: rootUserID, Role: "user", Status: "success"}, + {ID: rootAssistantID, ParentMessageID: &rootUserID, Role: "assistant", Status: "success"}, + {ID: failedUserID, ParentMessageID: &rootAssistantID, Role: "user", Status: "error", Attachments: string(retriedAttachments)}, + {ID: retryAssistantID, ParentMessageID: &failedUserID, SourceMessageID: &sourceAssistantID, Role: "assistant", BranchReason: "retry", Status: "success"}, + {ID: 6, ParentMessageID: &retryAssistantID, Role: "user", Status: "pending"}, + } + + recovered := recoverAssistantRetryUserStates(messages) + + if len(recovered) != len(messages) { + t.Fatalf("expected expanded context length %d, got %#v", len(messages), recovered) + } + if recovered[2].Status != "success" { + t.Fatalf("expected retried user to be recovered, got status %q", recovered[2].Status) + } + if recovered[4].Status != "pending" { + t.Fatalf("expected unrelated message state to remain pending, got %q", recovered[4].Status) + } + if messages[2].Status != "error" { + t.Fatalf("expected persisted state copy to remain unchanged, got %q", messages[2].Status) + } + fileIDs := collectConversationFileIDs(recovered, nil) + if len(fileIDs) != 1 || fileIDs[0] != "file_from_retried_user" { + t.Fatalf("expected recovered attachment in expanded context, got %#v", fileIDs) + } +} + func TestNormalizeDefaultBranchContextDoesNotRecoverFailedRetry(t *testing.T) { rootUserID := uint(1) rootAssistantID := uint(2) diff --git a/backend/internal/application/conversation/service_message_send.go b/backend/internal/application/conversation/service_message_send.go index 0963246b..b5cf7bdb 100644 --- a/backend/internal/application/conversation/service_message_send.go +++ b/backend/internal/application/conversation/service_message_send.go @@ -472,6 +472,8 @@ func (s *Service) sendMessageInternal( // 收集并行预取结果，再规划本轮可发送的 PromptScope。 prefetch := <-prefetchCh contextMessages = s.expandContextMessagesToSnapshotBoundary(ctx, input.ConversationID, userMessage.ID, contextMessages, prefetch.snapshot, compactPolicy) + // 快照扩展可能重新加载数据库中的原始 error 状态；在最终分支路径上统一恢复可用的重试上下文。 + contextMessages = recoverAssistantRetryUserStates(contextMessages) promptScope := buildPromptScope(contextMessages, prefetch.snapshot, compactPolicy) promptMessages := s.applyContextTokenBudget(promptScope.activeMessages(), route.UpstreamModel, route.ModelCapabilitiesJSON, reasoningContentPassback) ragQuery := buildRAGQuery(promptMessages, input.Content, cfg.RAGQueryHistoryTurns)