diff --git a/internal/application/service/image_multimodal.go b/internal/application/service/image_multimodal.go index 64e471f6bc..46b4d892c0 100644 --- a/internal/application/service/image_multimodal.go +++ b/internal/application/service/image_multimodal.go @@ -401,7 +401,7 @@ func (s *ImageMultimodalService) shouldDropOrphanedMultimodal( } // isFinalAsynqAttempt reports whether the current task context belongs to the -// last retry attempt before asynq archives the task as a dead-letter. We use +// last retry attempt before Asynq (or the Lite executor) archives the task. We use // this to flip multimodal finalize semantics: during normal retries we skip // counter decrement (the retry might still succeed), but on the final attempt // we count the image regardless of outcome so a permanently-failing image @@ -412,14 +412,14 @@ func (s *ImageMultimodalService) shouldDropOrphanedMultimodal( // "not final" keeps test ergonomics — tests should drive finalize explicitly. func isFinalAsynqAttempt(ctx context.Context) bool { retried, ok := asynq.GetRetryCount(ctx) - if !ok { - return false - } - maxRetry, ok := asynq.GetMaxRetry(ctx) - if !ok { - return false + if ok { + maxRetry, maxRetryOK := asynq.GetMaxRetry(ctx) + if maxRetryOK { + return retried >= maxRetry + } } - return retried >= maxRetry + retried, maxRetry, ok := types.TaskRetryMetadataFromContext(ctx) + return ok && retried >= maxRetry } // indexChunks indexes the newly created multimodal chunks into the retrieval engine diff --git a/internal/application/service/knowledge_process.go b/internal/application/service/knowledge_process.go index 8fbdbb5dfd..5debd9cc92 100644 --- a/internal/application/service/knowledge_process.go +++ b/internal/application/service/knowledge_process.go @@ -712,7 +712,105 @@ const imageDominatedTextThreshold = 200 // (typical for scanned PDFs where VLM OCR yielded nothing). Callers should // mark the knowledge's summary as failed instead of falling back to the first // chunk's raw content (which would just be a bare image reference). -var errInsufficientSummaryContent = errors.New("insufficient text content for summary generation") +var ( + errInsufficientSummaryContent = errors.New("insufficient text content for summary generation") + errEmptySummaryOutput = errors.New("summary model returned empty output") +) + +const summaryFallbackMaxRunes = 500 + +// validateSummaryOutput rejects successful model responses that contain no +// user-visible text. Treating whitespace-only output as an error lets Asynq +// retry the summary task instead of persisting description="" as completed. +func validateSummaryOutput(response *types.ChatResponse) (string, error) { + if response == nil { + return "", errEmptySummaryOutput + } + content := strings.TrimSpace(response.Content) + if content == "" { + return "", errEmptySummaryOutput + } + return content, nil +} + +// firstTextChunkSummaryFallback preserves the existing deterministic fallback: +// use the first already-ordered text chunk and cap it by runes so Chinese and +// emoji are never cut in the middle of a UTF-8 sequence. +func firstTextChunkSummaryFallback(textChunks []*types.Chunk) string { + if len(textChunks) == 0 || textChunks[0] == nil { + return "" + } + fallback := strings.TrimSpace(textChunks[0].Content) + runes := []rune(fallback) + if len(runes) > summaryFallbackMaxRunes { + fallback = string(runes[:summaryFallbackMaxRunes]) + } + return fallback +} + +// applyRetryableSummaryFailureState keeps an existing description visible +// while another attempt is queued, then publishes the deterministic fallback +// and marks only the summary subtask failed after the retry budget is exhausted. +func applyRetryableSummaryFailureState( + knowledge *types.Knowledge, textChunks []*types.Chunk, willRetry bool, +) string { + knowledge.UpdatedAt = time.Now() + if willRetry { + knowledge.SummaryStatus = types.SummaryStatusPending + return "" + } + fallback := firstTextChunkSummaryFallback(textChunks) + knowledge.Description = fallback + knowledge.SummaryStatus = types.SummaryStatusFailed + return fallback +} + +// handleSummaryRefreshFailure maps a failed refresh onto the value returned to +// the task executor, and guarantees a terminal delivery never leaves the row in +// pending/processing — the frontend renders both as "generating summary", so a +// refresh that dies on a path which never writes a status (tenant lookup, KB +// lookup, unconfigured summary model, chunk listing, freshness verification) +// would otherwise spin the placeholder forever. +// +// Stale and insufficient-content outcomes are deliberately swallowed: a newer +// refresh owns the status in the first case, and the second already persisted +// failed before returning. +func (s *knowledgeService) handleSummaryRefreshFailure( + ctx context.Context, knowledgeID string, err error, +) error { + if errors.Is(err, ErrSummaryRefreshStale) { + logger.Infof(ctx, "Discarding stale summary refresh for knowledge %s", knowledgeID) + return nil + } + logger.Warnf(ctx, "Summary refresh failed for knowledge %s: %v", knowledgeID, err) + if errors.Is(err, errInsufficientSummaryContent) { + return nil + } + if !summaryTaskWillRetry(ctx) && s.repo != nil { + // Column update rather than a full row write: RegenerateKnowledgeSummary + // may already have published a first-chunk fallback description that + // must survive this status write. + if updateErr := s.repo.UpdateKnowledgeColumn( + ctx, knowledgeID, "summary_status", types.SummaryStatusFailed, + ); updateErr != nil { + logger.Warnf(ctx, "Failed to mark summary refresh failed for knowledge %s: %v", + knowledgeID, updateErr) + } + } + return err +} + +// summaryTaskWillRetry reports whether the current Asynq delivery has another +// configured attempt remaining. Calls outside an Asynq worker are terminal. +func summaryTaskWillRetry(ctx context.Context) bool { + retried, retryOK := asynq.GetRetryCount(ctx) + maxRetry, maxRetryOK := asynq.GetMaxRetry(ctx) + if retryOK && maxRetryOK { + return retried < maxRetry + } + retried, maxRetry, ok := types.TaskRetryMetadataFromContext(ctx) + return ok && retried < maxRetry +} // checkSufficientSummaryContent returns errInsufficientSummaryContent if the // given content does not carry enough real text (after stripping image markup) @@ -891,8 +989,13 @@ func (s *knowledgeService) getSummary(ctx context.Context, logger.GetLogger(ctx).WithField("error", err).Errorf("GetSummary failed") return "", err } - logger.GetLogger(ctx).WithField("summary", summary.Content).Infof("GetSummary success") - return summary.Content, nil + content, err := validateSummaryOutput(summary) + if err != nil { + logger.GetLogger(ctx).WithField("error", err).Warnf("GetSummary returned no usable content") + return "", err + } + logger.GetLogger(ctx).WithField("summary", content).Infof("GetSummary success") + return content, nil } // sampleLongContent returns content that fits within maxChars. @@ -970,12 +1073,7 @@ func (s *knowledgeService) ProcessSummaryGeneration(ctx context.Context, t *asyn _, err = s.RegenerateKnowledgeSummary(ctx, payload.KnowledgeID) } if err != nil { - if errors.Is(err, ErrSummaryRefreshStale) { - logger.Infof(ctx, "Discarding stale summary refresh for knowledge %s", payload.KnowledgeID) - return nil - } - logger.Warnf(ctx, "Summary refresh failed for knowledge %s: %v", payload.KnowledgeID, err) - _ = s.repo.UpdateKnowledgeColumn(ctx, payload.KnowledgeID, "summary_status", types.SummaryStatusFailed) + return s.handleSummaryRefreshFailure(ctx, payload.KnowledgeID, err) } return nil } @@ -1087,8 +1185,8 @@ func (s *knowledgeService) ProcessSummaryGeneration(ctx context.Context, t *asyn if len(textChunks) == 0 { logger.Infof(ctx, "No text chunks found for knowledge: %s", payload.KnowledgeID) - // Mark as completed since there's nothing to summarize - knowledge.SummaryStatus = types.SummaryStatusCompleted + knowledge.Description = "" + knowledge.SummaryStatus = types.SummaryStatusFailed knowledge.UpdatedAt = time.Now() s.repo.UpdateKnowledge(ctx, knowledge) summaryOut["skipped"] = "no_text_chunks" @@ -1100,17 +1198,64 @@ func (s *knowledgeService) ProcessSummaryGeneration(ctx context.Context, t *asyn return textChunks[i].ChunkIndex < textChunks[j].ChunkIndex }) - // Initialize chat model for summary + summaryMetadataVersion := string(knowledge.CustomMetadata) + handleRetryableSummaryFailure := func(generationErr error) error { + summaryErr = generationErr + summaryOut["error"] = previewText(generationErr.Error(), 500) + summaryOut["error_type"] = fmt.Sprintf("%T", generationErr) + + if summaryTaskWillRetry(ctx) { + applyRetryableSummaryFailureState(knowledge, textChunks, true) + if updateErr := s.repo.UpdateKnowledge(ctx, knowledge); updateErr != nil { + logger.Warnf(ctx, "Failed to mark summary pending for retry: %v", updateErr) + } + summaryOut["retrying"] = true + return fmt.Errorf("summary generation attempt failed: %w", generationErr) + } + + // Before publishing the terminal fallback, make sure its source still + // matches the chunks and metadata captured for this attempt. + stale, staleErr := summarySourceChanged( + ctx, s.repo, s.chunkRepo, payload.TenantID, payload.KnowledgeID, + summaryMetadataVersion, textChunks, + ) + if staleErr != nil { + logger.Errorf(ctx, "Failed to verify summary fallback freshness for knowledge %s: %v", + payload.KnowledgeID, staleErr) + markSummaryFailed() + summaryErr = staleErr + return fmt.Errorf("verify summary fallback freshness: %w", staleErr) + } + if stale { + logger.Infof(ctx, "Discarding stale summary fallback for knowledge %s", payload.KnowledgeID) + summaryOut["skipped"] = "content_revision_changed" + return nil + } + + fallback := applyRetryableSummaryFailureState(knowledge, textChunks, false) + if updateErr := s.repo.UpdateKnowledge(ctx, knowledge); updateErr != nil { + logger.Errorf(ctx, "Failed to save terminal summary fallback: %v", updateErr) + summaryErr = updateErr + return fmt.Errorf("save terminal summary fallback: %w", updateErr) + } + if fallback == "" { + summaryOut["fallback"] = "empty" + } else { + summaryOut["fallback"] = "first_chunk" + } + summaryOut["fallback_chars"] = len([]rune(fallback)) + return fmt.Errorf("summary generation exhausted retries: %w", generationErr) + } + + // Initialize chat model for summary. Model resolution failures use the same + // retry budget and terminal first-chunk fallback as LLM request failures. chatModel, err := s.modelService.GetChatModel(ctx, kb.SummaryModelID) if err != nil { logger.Errorf(ctx, "Failed to get chat model: %v", err) - markSummaryFailed() - summaryErr = err - return fmt.Errorf("failed to get chat model: %w", err) + return handleRetryableSummaryFailure(fmt.Errorf("get chat model: %w", err)) } // Generate summary - summaryMetadataVersion := string(knowledge.CustomMetadata) summary, err := s.getSummary(ctx, chatModel, knowledge, textChunks) if err != nil { logger.Errorf(ctx, "Failed to generate summary for knowledge %s: %v", payload.KnowledgeID, err) @@ -1138,17 +1283,7 @@ func (s *knowledgeService) ProcessSummaryGeneration(ctx context.Context, t *asyn summaryErr = err return nil } - // For other errors (LLM API issues etc.), fall back to the first chunk. - if len(textChunks) > 0 { - summary = textChunks[0].Content - if len(summary) > 500 { - runes := []rune(summary) - if len(runes) > 500 { - summary = string(runes[:500]) - } - } - summaryOut["fallback"] = "first_chunk" - } + return handleRetryableSummaryFailure(err) } // Do not publish an answer derived from a superseded chunk or metadata // version. A user can explicitly refresh again from the latest revision. @@ -2177,22 +2312,66 @@ func (s *knowledgeService) RegenerateKnowledgeSummary( } } if len(textChunks) == 0 { - return nil, fmt.Errorf("no enabled text chunks to summarize") - } - chatModel, err := s.modelService.GetChatModel(ctx, kb.SummaryModelID) - if err != nil { - return nil, err + knowledge.Description = "" + knowledge.SummaryStatus = types.SummaryStatusFailed + knowledge.UpdatedAt = time.Now() + if updateErr := s.repo.UpdateKnowledge(ctx, knowledge); updateErr != nil { + return knowledge, updateErr + } + return knowledge, errInsufficientSummaryContent } + sort.Slice(textChunks, func(i, j int) bool { + return textChunks[i].ChunkIndex < textChunks[j].ChunkIndex + }) metadataVersion := string(knowledge.CustomMetadata) knowledge.SummaryStatus = types.SummaryStatusProcessing if err := s.repo.UpdateKnowledge(ctx, knowledge); err != nil { return nil, err } + handleGenerationFailure := func(generationErr error) (*types.Knowledge, error) { + if errors.Is(generationErr, errInsufficientSummaryContent) { + knowledge.Description = "" + knowledge.SummaryStatus = types.SummaryStatusFailed + knowledge.UpdatedAt = time.Now() + if updateErr := s.repo.UpdateKnowledge(ctx, knowledge); updateErr != nil { + return knowledge, updateErr + } + return knowledge, generationErr + } + if summaryTaskWillRetry(ctx) { + applyRetryableSummaryFailureState(knowledge, textChunks, true) + if updateErr := s.repo.UpdateKnowledge(ctx, knowledge); updateErr != nil { + logger.Warnf(ctx, "Failed to mark summary refresh pending for retry: %v", updateErr) + } + return knowledge, generationErr + } + + stale, staleErr := summarySourceChanged( + ctx, s.repo, s.chunkRepo, tenantID, knowledgeID, metadataVersion, textChunks, + ) + if staleErr != nil { + knowledge.SummaryStatus = types.SummaryStatusFailed + _ = s.repo.UpdateKnowledge(ctx, knowledge) + return knowledge, fmt.Errorf("verify summary fallback freshness: %w", staleErr) + } + if stale { + return knowledge, ErrSummaryRefreshStale + } + + applyRetryableSummaryFailureState(knowledge, textChunks, false) + if updateErr := s.repo.UpdateKnowledge(ctx, knowledge); updateErr != nil { + return knowledge, updateErr + } + return knowledge, generationErr + } + + chatModel, err := s.modelService.GetChatModel(ctx, kb.SummaryModelID) + if err != nil { + return handleGenerationFailure(fmt.Errorf("get chat model: %w", err)) + } summary, err := s.getSummary(ctx, chatModel, knowledge, textChunks) if err != nil { - knowledge.SummaryStatus = types.SummaryStatusFailed - _ = s.repo.UpdateKnowledge(ctx, knowledge) - return nil, err + return handleGenerationFailure(err) } stale, err := summarySourceChanged( ctx, s.repo, s.chunkRepo, tenantID, knowledgeID, metadataVersion, textChunks, diff --git a/internal/application/service/knowledge_summary_refresh_failure_test.go b/internal/application/service/knowledge_summary_refresh_failure_test.go new file mode 100644 index 0000000000..fee936fdd0 --- /dev/null +++ b/internal/application/service/knowledge_summary_refresh_failure_test.go @@ -0,0 +1,111 @@ +package service + +import ( + "context" + "errors" + "fmt" + "testing" + + "github.com/Tencent/WeKnora/internal/types" + "github.com/Tencent/WeKnora/internal/types/interfaces" +) + +type summaryRefreshStatusRepo struct { + interfaces.KnowledgeRepository + columnWrites map[string]interface{} + updateErr error +} + +func (r *summaryRefreshStatusRepo) UpdateKnowledgeColumn( + _ context.Context, _ string, column string, value interface{}, +) error { + if r.columnWrites == nil { + r.columnWrites = map[string]interface{}{} + } + r.columnWrites[column] = value + return r.updateErr +} + +func TestHandleSummaryRefreshFailure(t *testing.T) { + llmErr := errors.New("upstream 502") + + t.Run("retryable delivery leaves status untouched", func(t *testing.T) { + repo := &summaryRefreshStatusRepo{} + svc := &knowledgeService{repo: repo} + ctx := types.WithTaskRetryMetadata(context.Background(), 1, 3) + + err := svc.handleSummaryRefreshFailure(ctx, "knowledge-1", llmErr) + if !errors.Is(err, llmErr) { + t.Fatalf("error = %v, want the generation error so the executor retries", err) + } + if len(repo.columnWrites) != 0 { + t.Fatalf("retryable delivery wrote %v, want no status write", repo.columnWrites) + } + }) + + t.Run("terminal delivery marks summary failed", func(t *testing.T) { + repo := &summaryRefreshStatusRepo{} + svc := &knowledgeService{repo: repo} + ctx := types.WithTaskRetryMetadata(context.Background(), 3, 3) + + err := svc.handleSummaryRefreshFailure(ctx, "knowledge-1", llmErr) + if !errors.Is(err, llmErr) { + t.Fatalf("error = %v, want the generation error", err) + } + if got := repo.columnWrites["summary_status"]; got != types.SummaryStatusFailed { + t.Fatalf("summary_status = %v, want %q", got, types.SummaryStatusFailed) + } + }) + + t.Run("wrapped terminal failure still marks summary failed", func(t *testing.T) { + repo := &summaryRefreshStatusRepo{} + svc := &knowledgeService{repo: repo} + ctx := types.WithTaskRetryMetadata(context.Background(), 3, 3) + + wrapped := fmt.Errorf("get chat model: %w", llmErr) + if err := svc.handleSummaryRefreshFailure(ctx, "knowledge-1", wrapped); err == nil { + t.Fatal("expected the failure to propagate") + } + if got := repo.columnWrites["summary_status"]; got != types.SummaryStatusFailed { + t.Fatalf("summary_status = %v, want %q", got, types.SummaryStatusFailed) + } + }) + + t.Run("stale refresh is swallowed without a status write", func(t *testing.T) { + repo := &summaryRefreshStatusRepo{} + svc := &knowledgeService{repo: repo} + ctx := types.WithTaskRetryMetadata(context.Background(), 3, 3) + + if err := svc.handleSummaryRefreshFailure(ctx, "knowledge-1", ErrSummaryRefreshStale); err != nil { + t.Fatalf("stale refresh should not be retried or reported, got %v", err) + } + if len(repo.columnWrites) != 0 { + t.Fatalf("stale refresh wrote %v, want the newer refresh to own the status", repo.columnWrites) + } + }) + + t.Run("insufficient content is terminal without an extra write", func(t *testing.T) { + repo := &summaryRefreshStatusRepo{} + svc := &knowledgeService{repo: repo} + ctx := types.WithTaskRetryMetadata(context.Background(), 0, 3) + + if err := svc.handleSummaryRefreshFailure( + ctx, "knowledge-1", errInsufficientSummaryContent, + ); err != nil { + t.Fatalf("insufficient content should not be retried, got %v", err) + } + if len(repo.columnWrites) != 0 { + t.Fatalf("insufficient content wrote %v, want the caller's write to stand", repo.columnWrites) + } + }) + + t.Run("status write failure does not mask the generation error", func(t *testing.T) { + repo := &summaryRefreshStatusRepo{updateErr: errors.New("database unavailable")} + svc := &knowledgeService{repo: repo} + ctx := types.WithTaskRetryMetadata(context.Background(), 3, 3) + + if err := svc.handleSummaryRefreshFailure(ctx, "knowledge-1", llmErr); !errors.Is(err, llmErr) { + t.Fatalf("error = %v, want the generation error", err) + } + }) +} diff --git a/internal/application/service/knowledge_summary_test.go b/internal/application/service/knowledge_summary_test.go index f75f477e66..69fb5927f1 100644 --- a/internal/application/service/knowledge_summary_test.go +++ b/internal/application/service/knowledge_summary_test.go @@ -3,7 +3,10 @@ package service import ( "context" "errors" + "strings" "testing" + + "github.com/Tencent/WeKnora/internal/types" ) // TestCheckSufficientSummaryContent verifies the gate that prevents getSummary @@ -41,8 +44,8 @@ func TestCheckSufficientSummaryContent(t *testing.T) { wantError: true, }, { - name: "scanned PDF with empty wrapper rejected", - content: `![a](x)`, + name: "scanned PDF with empty wrapper rejected", + content: `![a](x)`, wantError: true, }, { @@ -103,3 +106,136 @@ func TestCheckSufficientSummaryContent_ThresholdOverride(t *testing.T) { t.Fatalf("tightened threshold: expected errInsufficientSummaryContent, got %v", err) } } + +func TestValidateSummaryOutput(t *testing.T) { + tests := []struct { + name string + response *types.ChatResponse + want string + wantError bool + }{ + {name: "nil response rejected", response: nil, wantError: true}, + {name: "empty response rejected", response: &types.ChatResponse{}, wantError: true}, + { + name: "whitespace response rejected", + response: &types.ChatResponse{Content: " \n\t "}, + wantError: true, + }, + { + name: "valid response is trimmed", + response: &types.ChatResponse{Content: " useful summary \n"}, + want: "useful summary", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := validateSummaryOutput(tt.response) + if tt.wantError { + if !errors.Is(err, errEmptySummaryOutput) { + t.Fatalf("expected errEmptySummaryOutput, got %v", err) + } + return + } + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if got != tt.want { + t.Fatalf("summary = %q, want %q", got, tt.want) + } + }) + } +} + +func TestFirstTextChunkSummaryFallback(t *testing.T) { + t.Run("uses only the first chunk", func(t *testing.T) { + got := firstTextChunkSummaryFallback([]*types.Chunk{ + {Content: " first chunk "}, + {Content: "second chunk"}, + }) + if got != "first chunk" { + t.Fatalf("fallback = %q, want first chunk", got) + } + }) + + t.Run("does not skip an empty first chunk", func(t *testing.T) { + got := firstTextChunkSummaryFallback([]*types.Chunk{ + {Content: " \n\t "}, + {Content: "second chunk"}, + }) + if got != "" { + t.Fatalf("fallback = %q, want empty", got) + } + }) + + t.Run("caps unicode content by runes", func(t *testing.T) { + got := firstTextChunkSummaryFallback([]*types.Chunk{{ + Content: strings.Repeat("摘", summaryFallbackMaxRunes+25), + }}) + if len([]rune(got)) != summaryFallbackMaxRunes { + t.Fatalf("fallback rune count = %d, want %d", len([]rune(got)), summaryFallbackMaxRunes) + } + }) + + t.Run("empty input stays empty", func(t *testing.T) { + if got := firstTextChunkSummaryFallback(nil); got != "" { + t.Fatalf("fallback = %q, want empty", got) + } + }) +} + +func TestApplyRetryableSummaryFailureState(t *testing.T) { + chunks := []*types.Chunk{{Content: "first body chunk"}} + + t.Run("retry keeps existing description", func(t *testing.T) { + knowledge := &types.Knowledge{ + Description: "previous summary", + SummaryStatus: types.SummaryStatusProcessing, + } + fallback := applyRetryableSummaryFailureState(knowledge, chunks, true) + if fallback != "" { + t.Fatalf("retry fallback = %q, want empty", fallback) + } + if knowledge.Description != "previous summary" { + t.Fatalf("retry changed description to %q", knowledge.Description) + } + if knowledge.SummaryStatus != types.SummaryStatusPending { + t.Fatalf("retry status = %q, want pending", knowledge.SummaryStatus) + } + }) + + t.Run("terminal failure publishes fallback and fails summary", func(t *testing.T) { + knowledge := &types.Knowledge{ + Description: "previous summary", + SummaryStatus: types.SummaryStatusProcessing, + } + fallback := applyRetryableSummaryFailureState(knowledge, chunks, false) + if fallback != "first body chunk" { + t.Fatalf("terminal fallback = %q", fallback) + } + if knowledge.Description != fallback { + t.Fatalf("description = %q, want %q", knowledge.Description, fallback) + } + if knowledge.SummaryStatus != types.SummaryStatusFailed { + t.Fatalf("terminal status = %q, want failed", knowledge.SummaryStatus) + } + }) +} + +func TestSummaryRetryStateSupportsLiteExecutorContext(t *testing.T) { + retryCtx := types.WithTaskRetryMetadata(context.Background(), 1, 3) + if !summaryTaskWillRetry(retryCtx) { + t.Fatal("attempt 1 of maxRetry 3 should have another retry") + } + if isFinalAsynqAttempt(retryCtx) { + t.Fatal("attempt 1 of maxRetry 3 should not be final") + } + + finalCtx := types.WithTaskRetryMetadata(context.Background(), 3, 3) + if summaryTaskWillRetry(finalCtx) { + t.Fatal("attempt 3 of maxRetry 3 should not retry") + } + if !isFinalAsynqAttempt(finalCtx) { + t.Fatal("attempt 3 of maxRetry 3 should be final") + } +} diff --git a/internal/router/sync_task.go b/internal/router/sync_task.go index 8b901ea3bb..c316500326 100644 --- a/internal/router/sync_task.go +++ b/internal/router/sync_task.go @@ -99,7 +99,8 @@ func (e *SyncTaskExecutor) Enqueue(task *asynq.Task, opts ...asynq.Option) (*asy time.Sleep(backoff) } - lastErr = handler(ctx, task) + attemptCtx := types.WithTaskRetryMetadata(ctx, attempt, maxRetry) + lastErr = handler(attemptCtx, task) if lastErr == nil { logger.Infof(ctx, "[SyncTask] Task completed type=%s id=%s elapsed=%v", task.Type(), taskID, time.Since(start)) diff --git a/internal/router/sync_task_retry_test.go b/internal/router/sync_task_retry_test.go new file mode 100644 index 0000000000..3f4113a7bb --- /dev/null +++ b/internal/router/sync_task_retry_test.go @@ -0,0 +1,38 @@ +package router + +import ( + "context" + "testing" + "time" + + "github.com/Tencent/WeKnora/internal/types" + "github.com/hibiken/asynq" +) + +func TestSyncTaskExecutorInjectsRetryMetadata(t *testing.T) { + executor := NewSyncTaskExecutor() + observed := make(chan [2]int, 1) + executor.RegisterHandler("test:retry-metadata", func(ctx context.Context, _ *asynq.Task) error { + retried, maxRetry, ok := types.TaskRetryMetadataFromContext(ctx) + if !ok { + observed <- [2]int{-1, -1} + return nil + } + observed <- [2]int{retried, maxRetry} + return nil + }) + + task := asynq.NewTask("test:retry-metadata", nil) + if _, err := executor.Enqueue(task, asynq.MaxRetry(3)); err != nil { + t.Fatalf("enqueue: %v", err) + } + + select { + case got := <-observed: + if got != [2]int{0, 3} { + t.Fatalf("retry metadata = %v, want [0 3]", got) + } + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for sync task") + } +} diff --git a/internal/types/context_helpers.go b/internal/types/context_helpers.go index 9ee804960f..921c5679c5 100644 --- a/internal/types/context_helpers.go +++ b/internal/types/context_helpers.go @@ -197,6 +197,34 @@ func IsBackgroundTask(ctx context.Context) bool { return v } +type taskRetryMetadata struct { + retried int + maxRetry int +} + +// WithTaskRetryMetadata records retry counters for task executors that do not +// provide Asynq's native worker context, notably the Lite synchronous executor. +func WithTaskRetryMetadata(ctx context.Context, retried, maxRetry int) context.Context { + return context.WithValue(ctx, taskRetryMetadataContextKey{}, taskRetryMetadata{ + retried: retried, maxRetry: maxRetry, + }) +} + +// TaskRetryMetadataFromContext returns retry counters supplied by a non-Asynq +// task executor. The boolean is false for ordinary request contexts. +func TaskRetryMetadataFromContext(ctx context.Context) (retried, maxRetry int, ok bool) { + if ctx == nil { + return 0, 0, false + } + metadata, ok := ctx.Value(taskRetryMetadataContextKey{}).(taskRetryMetadata) + if !ok { + return 0, 0, false + } + return metadata.retried, metadata.maxRetry, true +} + +type taskRetryMetadataContextKey struct{} + // WithLLMCallMetadata annotates a provider call for cache observability. The // fingerprint must be a hash, never raw prompt content. func WithLLMCallMetadata(ctx context.Context, purpose, prefixFingerprint string) context.Context { diff --git a/internal/types/context_helpers_test.go b/internal/types/context_helpers_test.go index 1d037d98a4..844450962f 100644 --- a/internal/types/context_helpers_test.go +++ b/internal/types/context_helpers_test.go @@ -153,6 +153,21 @@ func TestLLMCallMetadataContext(t *testing.T) { } } +func TestTaskRetryMetadataContext(t *testing.T) { + if _, _, ok := TaskRetryMetadataFromContext(nil); ok { + t.Fatal("nil context should not contain task retry metadata") + } + if _, _, ok := TaskRetryMetadataFromContext(context.Background()); ok { + t.Fatal("background context should not contain task retry metadata") + } + + ctx := WithTaskRetryMetadata(context.Background(), 2, 3) + retried, maxRetry, ok := TaskRetryMetadataFromContext(ctx) + if !ok || retried != 2 || maxRetry != 3 { + t.Fatalf("retry metadata = (%d, %d, %v), want (2, 3, true)", retried, maxRetry, ok) + } +} + // BenchmarkLanguageLocaleName benchmarks the language name lookup func BenchmarkLanguageLocaleName(b *testing.B) { testCases := []string{"zh", "en", "zh-CN", "ko", "unknown"}