WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes

internal/middleware/asynqdl/asynqdl.go

@@ -0,0 +1,255 @@
+// Package asynqdl provides an asynq middleware that records every task
+// whose retry budget is exhausted into the generic task_dead_letters table.
+//
+// The middleware is the catch-all observability path for asynq failures:
+// without it, archived tasks live only inside Redis with the asynq
+// inspector CLI as the sole readable surface. With it, operators can
+// SQL-query failures by task type, scope, or tenant.
+//
+// The package is intentionally separate from /internal/middleware (which
+// is gin-only) and from /internal/tracing/langfuse (which has its own
+// mux.Use). Putting it in its own package keeps the import graph tidy:
+// it depends on types/interfaces and asynq, nothing else inside the
+// project.
+package asynqdl
+
+import (
+	"context"
+	"encoding/json"
+	"time"
+
+	"github.com/Tencent/WeKnora/internal/logger"
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/Tencent/WeKnora/internal/types/interfaces"
+	"github.com/hibiken/asynq"
+)
+
+// Middleware returns an asynq.MiddlewareFunc that records dead-lettered
+// tasks. It is safe to install before any other asynq middleware: the
+// inner handler's return value is forwarded verbatim, so every layer
+// (langfuse tracing, business handlers) still sees the original
+// error/result.
+//
+// Behaviour:
+//
+//   - Tasks that succeed (handler returns nil) are passed through with
+//     zero overhead besides the function call.
+//
+//   - Tasks that fail (handler returns non-nil) trigger an inspection of
+//     asynq's retry counters. We only record on the FINAL retry —
+//     i.e. when GetRetryCount(ctx) >= GetMaxRetry(ctx). Earlier failures
+//     are still reported up to asynq for backoff/retry but do not
+//     create dead-letter rows; otherwise we'd write one row per
+//     transient hiccup.
+//
+//   - Insert is best-effort: a DB error is logged and swallowed. The
+//     middleware never alters the task error returned to the caller.
+//
+// This middleware should be installed once on the asynq mux, BEFORE
+// other middlewares that may transform errors. See router/task.go.
+// OnDeadLetter is an optional callback invoked AFTER the dead-letter row has
+// been recorded. It exists so callers can flip user-visible state alongside
+// the bookkeeping row — e.g. mark a Knowledge as failed when its document
+// processing task is permanently archived. Without this hook the only signal
+// to the UI is the timestamp on the knowledge row going stale, which is what
+// produced "stuck in processing forever" reports.
+//
+// The callback receives the raw asynq Task and the final error. It runs with
+// context.Background() (NOT the task ctx) so a cancelled task ctx during
+// shutdown doesn't also drop the state update. Errors are logged and
+// swallowed; the callback must NEVER alter the original task error.
+type OnDeadLetter func(ctx context.Context, t *asynq.Task, taskErr error)
+
+// MiddlewareWithCallback is the extended form of Middleware. The callback may
+// be nil, in which case behaviour matches Middleware exactly.
+func MiddlewareWithCallback(repo interfaces.TaskDeadLetterRepository, cb OnDeadLetter) asynq.MiddlewareFunc {
+	return func(next asynq.Handler) asynq.Handler {
+		return asynq.HandlerFunc(func(ctx context.Context, t *asynq.Task) error {
+			err := next.ProcessTask(ctx, t)
+			if err == nil {
+				return nil
+			}
+			if !isFinalAttempt(ctx) {
+				return err
+			}
+			attempts := 0
+			if retried, ok := asynq.GetRetryCount(ctx); ok {
+				attempts = retried + 1
+			}
+			if repo != nil {
+				dl := buildDeadLetter(t, err, attempts)
+				if dl != nil {
+					if insertErr := repo.Insert(context.Background(), dl); insertErr != nil {
+						logger.Warnf(ctx,
+							"asynq dead-letter: failed to record %s task: %v (original task error: %v)",
+							t.Type(), insertErr, err,
+						)
+					}
+				}
+			}
+			if cb != nil {
+				// Wrap callback so a panic in user code doesn't escape
+				// into asynq's worker goroutine.
+				safeInvokeCallback(ctx, t, err, cb)
+			}
+			return err
+		})
+	}
+}
+
+func safeInvokeCallback(ctx context.Context, t *asynq.Task, taskErr error, cb OnDeadLetter) {
+	defer func() {
+		if r := recover(); r != nil {
+			logger.Warnf(ctx, "asynq dead-letter callback panicked for %s: %v", t.Type(), r)
+		}
+	}()
+	cb(context.Background(), t, taskErr)
+}
+
+// Middleware preserves the original signature for callers that don't need
+// the dead-letter state-linkback hook. Equivalent to MiddlewareWithCallback
+// with a nil callback.
+func Middleware(repo interfaces.TaskDeadLetterRepository) asynq.MiddlewareFunc {
+	return MiddlewareWithCallback(repo, nil)
+}
+
+// isFinalAttempt reports whether the just-finished handler invocation
+// was the LAST one asynq will run before archiving the task. Asynq
+// increments retry_count *after* a successful retry, so on the final
+// attempt retry_count == max_retry. Both helpers can fail (returning
+// 0, false) when the middleware runs outside an asynq worker context;
+// in that case we conservatively report "final" so test setups that
+// invoke the middleware directly still exercise the insert path.
+func isFinalAttempt(ctx context.Context) bool {
+	retried, retriedOK := asynq.GetRetryCount(ctx)
+	maxRetry, maxOK := asynq.GetMaxRetry(ctx)
+	if !retriedOK || !maxOK {
+		return true
+	}
+	return retried >= maxRetry
+}
+
+// buildDeadLetter constructs a TaskDeadLetter from the task and the
+// final error. The payload extraction is deliberately schema-agnostic:
+// every project payload is a JSON object containing some subset of the
+// well-known fields {tenant_id, knowledge_base_id, kb_id, knowledge_id,
+// task_id}. We probe for those fields with a tolerant struct so a new
+// payload shape doesn't need to know about this middleware.
+//
+// Returns nil only if the task itself is nil — every other failure mode
+// (unparseable payload, unknown task type) still produces a row, with
+// scope="unknown" so operators can find it.
+func buildDeadLetter(t *asynq.Task, taskErr error, attempts int) *types.TaskDeadLetter {
+	if t == nil {
+		return nil
+	}
+	probe := payloadProbe{}
+	_ = json.Unmarshal(t.Payload(), &probe) // best-effort
+
+	scope, scopeID := inferScope(probe)
+	relatedID := probe.KnowledgeID
+
+	// asynq stores the payload as raw bytes; preserve it verbatim so a
+	// future requeue can reuse it directly.
+	rawPayload := t.Payload()
+	if len(rawPayload) == 0 {
+		rawPayload = []byte("{}")
+	}
+
+	return &types.TaskDeadLetter{
+		TenantID:  probe.TenantID,
+		TaskType:  t.Type(),
+		Scope:     scope,
+		ScopeID:   scopeID,
+		RelatedID: relatedID,
+		Payload:   json.RawMessage(rawPayload),
+		LastError: truncateError(taskErr.Error(), 8192),
+		FailCount: attempts,
+		FailedAt:  time.Now(),
+	}
+}
+
+// payloadProbe is a best-effort decoder that extracts the common
+// identifier fields from any project payload. We deliberately limit
+// the probe to field names that have a consistent meaning across every
+// existing payload type:
+//
+//   - tenant_id          — uniformly the tenant uint64
+//   - knowledge_base_id  — wiki / chunk / image / post-process / ...
+//   - kb_id              — FAQImportPayload's alias for the same thing
+//   - knowledge_id       — per-document scope
+//   - source_kb_id       — KnowledgeMovePayload (semantic: source KB)
+//
+// Other plausible-looking keys (`source_id`, `target_id`) are NOT
+// probed: different payloads use them with different meanings (KB
+// clone vs data-source sync), and a wrong scope inference would route
+// dead letters under a misleading bucket. New payloads that need a
+// specific scope should reuse one of the canonical keys above.
+//
+// Missing keys decode to zero values, which we then fall back through
+// in inferScope.
+type payloadProbe struct {
+	TenantID        uint64 `json:"tenant_id,omitempty"`
+	KnowledgeBaseID string `json:"knowledge_base_id,omitempty"`
+	KBID            string `json:"kb_id,omitempty"` // FAQImportPayload uses this name
+	KnowledgeID     string `json:"knowledge_id,omitempty"`
+	SourceKBID      string `json:"source_kb_id,omitempty"` // KnowledgeMovePayload
+}
+
+// inferScope picks the most-specific scope tuple available in the
+// probe. The order reflects the conventional "blast radius" of a
+// task: knowledge_base is broader than knowledge, etc. Wiki ingest is
+// scope="knowledge_base"; chunk extract is scope="knowledge"; and so on.
+//
+// Falls back to scope="unknown" with empty scope_id if the payload
+// carries no recognizable identifier — better than silently dropping
+// the row.
+func inferScope(p payloadProbe) (string, string) {
+	switch {
+	case p.KnowledgeBaseID != "":
+		return types.TaskScopeKnowledgeBase, p.KnowledgeBaseID
+	case p.KBID != "":
+		return types.TaskScopeKnowledgeBase, p.KBID
+	case p.SourceKBID != "":
+		// Knowledge move spans two KBs; the source is the better
+		// blast-radius indicator (it's where the work was being read
+		// from when the failure occurred).
+		return types.TaskScopeKnowledgeBase, p.SourceKBID
+	case p.KnowledgeID != "":
+		return types.TaskScopeKnowledge, p.KnowledgeID
+	case p.TenantID != 0:
+		return types.TaskScopeTenant, formatUint(p.TenantID)
+	default:
+		return types.TaskScopeUnknown, ""
+	}
+}
+
+// truncateError prevents a single runaway stack trace from inflating
+// the dead-letter table beyond the disk capacity reasonably allotted
+// to it.
+func truncateError(s string, max int) string {
+	if max <= 0 || len(s) <= max {
+		return s
+	}
+	const suffix = "...(truncated)"
+	if max <= len(suffix) {
+		return s[:max]
+	}
+	return s[:max-len(suffix)] + suffix
+}
+
+// formatUint inlines strconv.FormatUint without importing strconv just
+// for one call (keeps the package import set minimal).
+func formatUint(v uint64) string {
+	if v == 0 {
+		return "0"
+	}
+	var buf [20]byte
+	i := len(buf)
+	for v > 0 {
+		i--
+		buf[i] = byte('0' + v%10)
+		v /= 10
+	}
+	return string(buf[i:])
+}

internal/middleware/asynqdl/asynqdl_test.go

@@ -0,0 +1,236 @@
+package asynqdl
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"sync"
+	"testing"
+
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/hibiken/asynq"
+)
+
+// fakeRepo records calls to Insert in memory so tests can assert how
+// many dead-letter rows would have been written. It implements only
+// the subset of TaskDeadLetterRepository the middleware uses; the
+// Lister methods are unused in the middleware path.
+type fakeRepo struct {
+	mu   sync.Mutex
+	rows []*types.TaskDeadLetter
+	fail error
+}
+
+func (f *fakeRepo) Insert(ctx context.Context, dl *types.TaskDeadLetter) error {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	if f.fail != nil {
+		return f.fail
+	}
+	cp := *dl // defensive copy: callers may reuse the struct
+	f.rows = append(f.rows, &cp)
+	return nil
+}
+
+func (f *fakeRepo) ListByScope(context.Context, string, string, string, int) ([]*types.TaskDeadLetter, string, error) {
+	return nil, "", nil
+}
+
+func (f *fakeRepo) ListByTaskType(context.Context, string, string, int) ([]*types.TaskDeadLetter, string, error) {
+	return nil, "", nil
+}
+
+func (f *fakeRepo) DeleteByID(context.Context, int64) error { return nil }
+
+// captureRow returns the i'th row safely under the lock.
+func (f *fakeRepo) captureRow(i int) *types.TaskDeadLetter {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	if i >= len(f.rows) {
+		return nil
+	}
+	return f.rows[i]
+}
+
+func (f *fakeRepo) rowCount() int {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	return len(f.rows)
+}
+
+// runMiddleware wires up the middleware around a handler that always
+// returns handlerErr. We invoke it directly (no asynq server) so retry
+// counters come from the helper-injected ctx values; isFinalAttempt
+// treats a missing pair as "final" for testability.
+func runMiddleware(repo *fakeRepo, taskType string, payload []byte, handlerErr error) error {
+	mw := Middleware(repo)
+	wrapped := mw(asynq.HandlerFunc(func(_ context.Context, _ *asynq.Task) error {
+		return handlerErr
+	}))
+	t := asynq.NewTask(taskType, payload)
+	return wrapped.ProcessTask(context.Background(), t)
+}
+
+func TestMiddleware_NilHandlerErr_NoInsert(t *testing.T) {
+	repo := &fakeRepo{}
+	if err := runMiddleware(repo, "any:task", []byte(`{"tenant_id":1}`), nil); err != nil {
+		t.Fatalf("expected nil error, got %v", err)
+	}
+	if got := repo.rowCount(); got != 0 {
+		t.Fatalf("expected zero inserts on success, got %d", got)
+	}
+}
+
+func TestMiddleware_FailureWithoutAsynqCtx_RecordsRow(t *testing.T) {
+	repo := &fakeRepo{}
+	payload, _ := json.Marshal(map[string]any{
+		"tenant_id":         42,
+		"knowledge_base_id": "kb-abc",
+		"knowledge_id":      "k-1",
+	})
+	wantErr := errors.New("boom")
+	gotErr := runMiddleware(repo, "wiki:ingest", payload, wantErr)
+	if !errors.Is(gotErr, wantErr) {
+		t.Fatalf("expected the original task error to propagate, got %v", gotErr)
+	}
+	if got := repo.rowCount(); got != 1 {
+		t.Fatalf("expected exactly one insert, got %d", got)
+	}
+	row := repo.captureRow(0)
+	if row.TaskType != "wiki:ingest" {
+		t.Errorf("task_type: got %q", row.TaskType)
+	}
+	if row.TenantID != 42 {
+		t.Errorf("tenant_id: got %d", row.TenantID)
+	}
+	if row.Scope != types.TaskScopeKnowledgeBase || row.ScopeID != "kb-abc" {
+		t.Errorf("scope tuple: got (%q, %q)", row.Scope, row.ScopeID)
+	}
+	if row.RelatedID != "k-1" {
+		t.Errorf("related_id: got %q", row.RelatedID)
+	}
+	if row.LastError != "boom" {
+		t.Errorf("last_error: got %q", row.LastError)
+	}
+	// Outside an asynq worker ctx GetRetryCount returns !ok, so the
+	// middleware records 0 attempts (clearer than a misleading 1).
+	if row.FailCount != 0 {
+		t.Errorf("fail_count: expected 0 outside worker ctx, got %d", row.FailCount)
+	}
+	if row.FailedAt.IsZero() {
+		t.Error("failed_at: expected non-zero timestamp")
+	}
+}
+
+func TestMiddleware_UnknownPayload_StillRecords(t *testing.T) {
+	repo := &fakeRepo{}
+	// Payload that doesn't carry any of the well-known identifier keys.
+	payload := []byte(`{"some_unrelated":"value"}`)
+	_ = runMiddleware(repo, "future:weird", payload, errors.New("nope"))
+	if got := repo.rowCount(); got != 1 {
+		t.Fatalf("expected one insert for unknown payload, got %d", got)
+	}
+	row := repo.captureRow(0)
+	if row.Scope != types.TaskScopeUnknown {
+		t.Errorf("expected scope=unknown for unparseable payload, got %q", row.Scope)
+	}
+	if row.ScopeID != "" {
+		t.Errorf("expected empty scope_id for unknown scope, got %q", row.ScopeID)
+	}
+	if string(row.Payload) != string(payload) {
+		t.Errorf("expected payload preserved verbatim, got %s", string(row.Payload))
+	}
+}
+
+func TestMiddleware_NilPayload_DefaultsToEmptyObject(t *testing.T) {
+	repo := &fakeRepo{}
+	_ = runMiddleware(repo, "any:task", nil, errors.New("err"))
+	row := repo.captureRow(0)
+	if row == nil {
+		t.Fatal("expected a row")
+	}
+	if string(row.Payload) != "{}" {
+		t.Errorf("expected payload normalized to {}, got %s", string(row.Payload))
+	}
+}
+
+func TestMiddleware_LongErrorTruncated(t *testing.T) {
+	repo := &fakeRepo{}
+	long := make([]byte, 16384)
+	for i := range long {
+		long[i] = 'x'
+	}
+	_ = runMiddleware(repo, "any", []byte("{}"), errors.New(string(long)))
+	row := repo.captureRow(0)
+	if got := len(row.LastError); got > 8192 {
+		t.Errorf("expected last_error truncated to <=8192 bytes, got %d", got)
+	}
+	const suffix = "...(truncated)"
+	if got := row.LastError[len(row.LastError)-len(suffix):]; got != suffix {
+		t.Errorf("expected truncation suffix %q, got %q", suffix, got)
+	}
+}
+
+func TestMiddleware_RepoFailure_DoesNotMaskTaskError(t *testing.T) {
+	repo := &fakeRepo{fail: errors.New("db down")}
+	wantErr := errors.New("task boom")
+	gotErr := runMiddleware(repo, "any", []byte(`{"tenant_id":1}`), wantErr)
+	if !errors.Is(gotErr, wantErr) {
+		t.Fatalf("repo failure must not mask task error; got %v", gotErr)
+	}
+}
+
+func TestInferScope_Priority(t *testing.T) {
+	tests := []struct {
+		name      string
+		probe     payloadProbe
+		wantScope string
+		wantID    string
+	}{
+		{
+			name:      "knowledge_base_id wins over knowledge_id",
+			probe:     payloadProbe{KnowledgeBaseID: "kb1", KnowledgeID: "k1"},
+			wantScope: types.TaskScopeKnowledgeBase,
+			wantID:    "kb1",
+		},
+		{
+			name:      "kb_id alias resolves to knowledge_base scope",
+			probe:     payloadProbe{KBID: "kb-faq"},
+			wantScope: types.TaskScopeKnowledgeBase,
+			wantID:    "kb-faq",
+		},
+		{
+			name:      "source_kb_id from KnowledgeMovePayload",
+			probe:     payloadProbe{SourceKBID: "src-kb"},
+			wantScope: types.TaskScopeKnowledgeBase,
+			wantID:    "src-kb",
+		},
+		{
+			name:      "knowledge_id only",
+			probe:     payloadProbe{KnowledgeID: "k-only"},
+			wantScope: types.TaskScopeKnowledge,
+			wantID:    "k-only",
+		},
+		{
+			name:      "tenant_id fallback",
+			probe:     payloadProbe{TenantID: 7},
+			wantScope: types.TaskScopeTenant,
+			wantID:    "7",
+		},
+		{
+			name:      "no identifiers — unknown scope",
+			probe:     payloadProbe{},
+			wantScope: types.TaskScopeUnknown,
+			wantID:    "",
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			gotScope, gotID := inferScope(tc.probe)
+			if gotScope != tc.wantScope || gotID != tc.wantID {
+				t.Errorf("inferScope: got (%q, %q), want (%q, %q)",
+					gotScope, gotID, tc.wantScope, tc.wantID)
+			}
+		})
+	}
+}