ragflow/internal/ingestion/knowledge_compile/consumer.go

//
//  Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
//
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.

package knowledge_compile

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"sort"
	"strings"
	"sync"
	"time"

	"ragflow/internal/common"
	"ragflow/internal/dao"
	"ragflow/internal/engine"
	kccommon "ragflow/internal/ingestion/component/knowledge_compiler/common"
	"ragflow/internal/service/nav"

	"go.uber.org/zap"
	"gorm.io/gorm"
)

// kcDB is the package-level MySQL handle installed by Provision. It backs the
// doc enumeration used by RebuildDataset; it is nil (and docLister returns no
// docs) when the scheduler was provisioned without a DB.
var kcDB *gorm.DB

// Consumer is the dataset-level post-processing worker (§11.5). Multiple
// instances compete on the MySQL scheduling rows; each KB is processed by at
// most one instance at a time via the per-KB claim (so the same KB is handled
// by a single worker that serializes the batch). The MySQL row — not the
// broker — is the scheduling system of record and the source of same-KB
// serialization.
type Consumer struct {
	scheduler Claimer
	reader    Reader
	writer    Writer
	factory   DeduperFactory

	ttl            time.Duration
	heartbeat      time.Duration
	pollInterval   time.Duration
	sweepInterval  time.Duration
	mergeThreshold float64 // KNN similarity threshold for "existing merged row is a duplicate"

	mu    sync.Mutex
	tombs map[string]map[string]uint64 // dataset -> docID -> delete marker (tombstone)

	// rebuildPause suppresses local claims while this consumer rebuilds a dataset.
	rebuildPause map[string]bool

	// docLister enumerates every doc id in a KB so a rewrite can republish them.
	docLister func(ctx context.Context, tenant, kb string) ([]string, error)
}

// rewriteScheduler is the subset of the scheduler API a dataset rewrite needs,
// satisfied by *mysqlScheduler (and FakeScheduler in tests). Using a narrow
// interface keeps the Claimer surface unchanged.
type rewriteScheduler interface {
	Publish(ctx context.Context, tenantID, datasetID, docID, eventType string, variants []string) error
	CancelInflight(ctx context.Context, datasetID, token string) error
}

// NewConsumer constructs a Consumer driven by the given Claimer. Tests pass a
// FakeScheduler and override the Reader/Writer/Deduper via options.
func NewConsumer(scheduler Claimer, opts ...Option) *Consumer {
	c := &Consumer{
		scheduler:      scheduler,
		reader:         engineReader{eng: engine.Get()},
		writer:         engineWriter{eng: engine.Get()},
		factory:        defaultDeduperFactory,
		ttl:            2 * time.Minute,
		heartbeat:      20 * time.Second,
		pollInterval:   2 * time.Second,
		sweepInterval:  30 * time.Second,
		mergeThreshold: 0.85,
		tombs:          map[string]map[string]uint64{},
		rebuildPause:   map[string]bool{},
		docLister:      defaultDocLister,
	}
	for _, o := range opts {
		o(c)
	}
	return c
}

// Run is one owned worker loop (Option E §11.5/§11.7): it wakes on NATS notify
// and otherwise polls the scheduling table for a claimable KB, then claims the
// closed batch, processes it, and acks. It returns when ctx is cancelled.
func (c *Consumer) Run(ctx context.Context) {
	if c.scheduler == nil {
		return
	}
	notifyCh, _ := c.scheduler.SubscribeNotify(ctx)
	poll := time.NewTicker(c.pollInterval)
	sweep := time.NewTicker(c.sweepInterval)
	defer poll.Stop()
	defer sweep.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case datasetID, ok := <-notifyCh:
			if !ok {
				notifyCh = nil
				continue
			}
			// The notify carries the dataset that just received backlog, so
			// claim that specific dataset directly instead of probing for an
			// arbitrary claimable one.
			c.claimAndProcess(ctx, datasetID)
		case <-poll.C:
			c.tryClaimAndProcess(ctx)
		case <-sweep.C:
			// Crash recovery: TryClaim reclaims expired inflight leases
			// (§11.5) before claiming any ready batch.
			c.tryClaimAndProcess(ctx)
		}
	}
}

// claimAndProcess claims the given dataset (a notify pointed at it) and processes
// the closed batch. ok=false means the dataset has no ready batch or a live lease
// already holds it (the race was lost).
func (c *Consumer) claimAndProcess(ctx context.Context, datasetID string) {
	cr, ok, err := c.scheduler.Claim(ctx, datasetID)
	if err != nil || !ok || len(cr.Entries) == 0 {
		return
	}
	c.processClaim(ctx, cr)
}

// tryClaimAndProcess claims one closed batch (ready or reclaimed) and processes
// it. ok=false means there was nothing to do this tick.
func (c *Consumer) tryClaimAndProcess(ctx context.Context) {
	cr, ok, err := c.scheduler.TryClaim(ctx)
	if err != nil || !ok || len(cr.Entries) == 0 {
		return // nothing to claim, or the race was lost
	}
	c.processClaim(ctx, cr)
}

// processClaim processes an already-claimed batch (cr) and acks on success. It
// is the Option E replacement for the old processOnce (lease + drain + merge):
// the claim returned a frozen batch boundary, so there is no moving target and
// no Nak-churn routing.
func (c *Consumer) processClaim(ctx context.Context, cr ClaimResult) {
	datasetID := cr.DatasetID

	c.mu.Lock()
	paused := c.rebuildPause[datasetID]
	c.mu.Unlock()
	if paused {
		return
	}

	// Heartbeat refreshes the claim TTL while we process; a failed touch means
	// the lease was taken over (or reclaimed) and we must abort without acking.
	stopHb := make(chan struct{})
	hbFailed := make(chan struct{}, 1)
	go func() {
		t := time.NewTicker(c.heartbeat)
		defer t.Stop()
		for {
			select {
			case <-stopHb:
				return
			case <-ctx.Done():
				return
			case <-t.C:
				alive, e := c.scheduler.TouchClaim(ctx, datasetID, cr.Token, c.ttl)
				if e != nil || !alive {
					select {
					case hbFailed <- struct{}{}:
					default:
					}
					return
				}
			}
		}
	}()

	done := make(chan struct{})
	var batchErr error
	go func() {
		defer close(done)
		batchErr = c.processBatch(ctx, cr.TenantID, datasetID, cr.Token, cr.Entries)
	}()

	select {
	case <-ctx.Done():
		close(stopHb)
		<-done
		return // graceful shutdown: leave inflight for reclamation, do not ack
	case <-hbFailed:
		close(stopHb)
		<-done
		return // lease lost: do not ack; sweeper/redelivery reprocesses (idempotent)
	case <-done:
		close(stopHb)
	}

	// Ack only on success. A batch error (reader/dedup/writer failure) means we
	// must leave the claimed batch in the backlog for reclamation/retry rather
	// than silently dropping it (C5: never ack what we failed to merge).
	if batchErr != nil {
		if errors.Is(batchErr, errClaimSuperseded) {
			return
		}
		common.Error("knowledge_compile: batch processing failed, leaving batch for retry",
			batchErr,
			zap.String("dataset_id", datasetID),
			zap.Int("entries", len(cr.Entries)))
		if err := c.scheduler.SetError(ctx, datasetID, cr.Token, batchErr.Error()); err != nil {
			common.Warn("knowledge_compile: failed to record error_msg",
				zap.String("dataset_id", datasetID), zap.Error(err))
		}
		return
	}
	if _, err := c.scheduler.Ack(ctx, datasetID, cr.Token, cr.Entries); err != nil {
		common.Warn("knowledge_compile: ack failed",
			zap.String("dataset_id", datasetID), zap.Error(err))
	}
}

var errClaimSuperseded = errors.New("knowledge_compile: claim superseded by rewrite")

// withWriteLock runs a destructive side effect fn under the scheduler's
// per-dataset write/rebuild lock, verifying the claim token inside the lock
// immediately before fn. This closes the TOCTOU gap between cancellation and a
// writer side effect: a cancelled worker cannot write after the rebuild clears
// storage, and a rebuild cannot interleave while fn runs.
func (c *Consumer) withWriteLock(ctx context.Context, kb, token string, fn func() error) error {
	return c.scheduler.WithDatasetLock(ctx, kb, func(currentToken string) error {
		if currentToken != token {
			return errClaimSuperseded
		}
		return fn()
	})
}

// RebuildDataset performs a full incremental rewrite of a KB's dataset-level
// merged products (W1/W2/W5/W6). The order is fixed to close both the
// enumerate-then-clear race (M19/C-race) AND the cross-process stale-write
// window:
//  1. set the rewrite pause (in-process only; clears via defer on every path);
//  2. CancelInflight invalidates any in-flight claim before clearing storage.
//     Each worker compares its claim token under the write lock and therefore
//     cannot write after cancellation;
//  3. DeleteMerged + DropWikiGraph clear the old merged + graph state INSIDE the
//     same per-dataset row lock (WithDatasetLock). This is what actually closes
//     the TOCTOU gap: the generation check + write of every worker run under this
//     lock, so the rebuild's clear either runs after any in-flight worker write
//     has finished (it is later removed by the clear) or after a stale worker has
//     self-dropped — a worker can never repopulate cleared state with old results
//     because it cannot hold the lock concurrently with the clear;
//  4. enumerate every doc and republish;
//  5. clear the local pause.
//
// mode is "incremental" or "rewrite"; today both take the same clean-and-rebuild
// path, with mode retained for future differential strategies.
func (c *Consumer) RebuildDataset(ctx context.Context, tenant, kb, mode string) error {
	rs, ok := c.scheduler.(rewriteScheduler)
	if !ok {
		return fmt.Errorf("knowledge_compile: scheduler %T does not support rewrite", c.scheduler)
	}

	// 1. pause the rewrite window. The pause is cleared by a deferred cleanup so a
	// failure mid-rebuild (DeleteMerged/DropWikiGraph/CancelInflight/Bump/Publish)
	// can never leave the dataset permanently paused (which would drop every
	// future claim). Partially published docs are retried on the next claim.
	c.mu.Lock()
	c.rebuildPause[kb] = true
	c.mu.Unlock()
	defer func() {
		c.mu.Lock()
		c.rebuildPause[kb] = false
		c.mu.Unlock()
	}()

	// 2. Cancel any in-flight claim before clearing storage. A worker that was
	// already running observes its revoked claim token inside withWriteLock and
	// returns without writing.
	if err := rs.CancelInflight(ctx, kb, ""); err != nil {
		return fmt.Errorf("knowledge_compile: rebuild cancel inflight: %w", err)
	}

	// 3. clear old merged + graph state (structural filter is the source of
	// truth; no per-doc rows are touched) under the per-dataset row lock. This
	// guarantees mutual exclusion with every worker's destructive write (which
	// runs under the same lock via withWriteLock): the clear either waits for an
	// in-flight worker write to finish (then removes its old-generation result)
	// or runs after stale workers have self-dropped. A worker cannot hold the
	// lock while the clear runs, so it can never repopulate cleared state with
	// old results.
	if err := c.scheduler.WithDatasetLock(ctx, kb, func(_ string) error {
		// B4/B1b: a full rebuild clears EVERY variant the consumer manages
		// (wiki merged rows, structure scope_kwd="dataset" rows, and the nav
		// tree rows) via the fixed full-set clean (empty variants), not just the
		// wiki rows that DeleteMerged removes — otherwise removed-template ghost
		// rows survive. DropWikiGraph additionally sweeps the legacy wiki graph.
		if derr := c.writer.DeleteMergedForVariant(ctx, tenant, kb, nil); derr != nil {
			return derr
		}
		return c.writer.DropWikiGraph(ctx, tenant, kb)
	}); err != nil {
		return fmt.Errorf("knowledge_compile: rebuild clear all variants + graph: %w", err)
	}

	// 4. enumerate every doc and republish. Per B1, the republished events must
	// carry the doc's compile variants so the consumer can route to the correct
	// dataset-level path (a nil/empty Variants would fall back to the legacy
	// unified path and tree nav could not be rebuilt). The variants are recovered
	// from the doc-level compiled products' authoritative `compilation_template_kind_kwd`
	// (via KindToVariant, O2a whitelist hard-fail), not re-derived from `compile_kwd`.
	docs, err := c.docLister(ctx, tenant, kb)
	if err != nil {
		return fmt.Errorf("knowledge_compile: rebuild list docs: %w", err)
	}
	for _, docID := range docs {
		variants, rerr := c.recoverDocVariants(ctx, tenant, kb, docID)
		if rerr != nil {
			// O2a hard failure: abort the rebuild so an unknown template kind
			// cannot leave the doc's dataset-level products unrebuilt.
			return fmt.Errorf("knowledge_compile: rebuild recover variants %s: %w", docID, rerr)
		}
		if perr := rs.Publish(ctx, tenant, kb, docID, string(EventTypeCompleted), variants); perr != nil {
			return fmt.Errorf("knowledge_compile: rebuild republish %s: %w", docID, perr)
		}
	}

	// 6. resweep: the deferred cleanup clears the pause so any claim arriving
	// post-enumeration is processed (and stale ones self-drop in processClaim).
	common.Info("knowledge_compile: dataset rebuild complete",
		zap.String("dataset_id", kb),
		zap.String("mode", mode),
		zap.Int("docs", len(docs)))
	return nil
}

// defaultDocLister enumerates every doc id in a KB via the Document DAO. When
// no DB was provisioned it returns no docs (a rewrite becomes a no-op clean).
func defaultDocLister(ctx context.Context, tenant, kb string) ([]string, error) {
	if kcDB == nil {
		return nil, nil
	}
	return dao.NewDocumentDAO().ListIDsByKBIDWithOptions(ctx, kcDB, dao.DocumentListOptions{KbID: kb})
}

// recoverDocVariants reconstructs the compile-type set a document produced, by
// reading its persisted doc-level products and mapping each product's
// authoritative `compilation_template_kind_kwd` through common.KindToVariant.
// It is used by RebuildDataset (B1) so the republished completed events carry
// the variants needed to route the dataset-level re-compile (tree nav, structure
// merge, wiki merge) — a nil/empty Variants would fall back to the legacy
// unified path and tree nav could not be rebuilt.
//
// KindToVariant (O2a) is a fixed whitelist that HARD-FAILS on unknown kinds: a
// product whose kind is not in the whitelist is a corrupt/unsupported state and
// the rebuild must abort (returning an error) rather than silently drop the
// variant and leave the doc's dataset-level products unrebuilt. A product
// without an authoritative kind falls back to its reverse-mapped variant (which
// itself must be whitelist-valid, enforced by KwdToVariant in the Reader). The
// returned slice is sorted and de-duplicated.
func (c *Consumer) recoverDocVariants(ctx context.Context, tenant, kb, docID string) ([]string, error) {
	products, err := c.reader.LoadDocProducts(ctx, tenant, kb, docID)
	if err != nil {
		return nil, fmt.Errorf("knowledge_compile: recover variants load doc %s: %w", docID, err)
	}
	seen := map[string]struct{}{}
	var out []string
	for _, p := range products {
		// Authoritative: product.Kind (compilation_template_kind_kwd) when present.
		if p.Kind != "" {
			v, err := kccommon.KindToVariant(p.Kind)
			if err != nil {
				// O2a hard failure: do not continue with an incomplete variant set.
				return nil, fmt.Errorf("knowledge_compile: recover variants doc %s kind %q: %w", docID, p.Kind, err)
			}
			if _, dup := seen[string(v)]; dup {
				continue
			}
			seen[string(v)] = struct{}{}
			out = append(out, string(v))
			continue
		}
		// Fallback: the product's reverse-mapped variant (already whitelist-valid).
		if p.Variant == "" {
			// No authoritative kind and no reverse-mapped variant: this is a
			// legacy/unknown row. Skip it so we never republish an empty variant
			// (which the consumer cannot route and defeats "nil = legacy").
			continue
		}
		if _, dup := seen[string(p.Variant)]; dup {
			continue
		}
		seen[string(p.Variant)] = struct{}{}
		out = append(out, string(p.Variant))
	}
	sort.Strings(out)
	return out, nil
}

// filterWikiPageCandidates returns ONLY the wiki page products the dataset-level
// wiki merge may fold. It drops every non-wiki product (tree/structure/nav) so
// those never enter the wiki Dedup/KNN/WriteMerged path, and within the wiki
// variant keeps strictly page-kind products (Meta.kind == "page"). Sections are
// a doc-level concern and must never enter the dataset-level page bucket, and
// the deduper must not carry an implicit section-filter contract. This is the
// per-variant product selection (review issue 2): a completed wiki-only doc whose
// event carried tree/structure products on a prior run must not feed stale
// tree/structure rows into the wiki merge.
func filterWikiPageCandidates(candidates []kccommon.Product) []kccommon.Product {
	filtered := make([]kccommon.Product, 0, len(candidates))
	for _, cand := range candidates {
		if cand.Variant != kccommon.VariantWiki {
			continue // tree/structure products never enter the wiki merge
		}
		if metaString(cand.Meta, "kind") != "page" {
			continue
		}
		filtered = append(filtered, cand)
	}
	return filtered
}

// processBatch applies out-of-order / tombstone handling, then recomputes and
// writes the dataset-level merged products for the claimed closed batch. It
// returns an error if any reader/dedup/writer step fails so the caller can
// leave the batch for reclamation instead of acking dropped work.
func (c *Consumer) processBatch(ctx context.Context, tenant, kb, token string, entries []BacklogEntry) error {
	common.Info("knowledge_compile: processing claimed batch",
		zap.String("dataset_id", kb),
		zap.String("tenant_id", tenant),
		zap.Int("entries", len(entries)))
	c.mu.Lock()
	if c.tombs == nil {
		c.tombs = map[string]map[string]uint64{}
	}
	tomb := c.tombs[kb]
	var completed []BacklogEntry
	var deleted []string
	// Per-document tombstone clears are staged locally and committed only after
	// the write/delete paths below return without error, so a failed batch
	// leaves c.tombs untouched and can be retried on reclaim.
	var pendingTombClear []string
	// Reconcile each document to its last event before classifying it as
	// deleted or completed. A completion and a deletion for the same doc can
	// land in the same batch, and completions/deletions can also arrive out of
	// order across batches. The LAST event for a doc (by backlog append order)
	// decides its fate, so we must not delete a doc that was re-ingested after
	// its deletion, nor merge a completion shadowed by a later deletion. No
	// sequence number is needed: the backlog is an ordered log and the last
	// write wins, which is exactly the re-parse (completed last) vs delete
	// (deleted last) semantics we want.
	type docState struct {
		winEvent EventType
		variants []string
	}
	byDoc := make(map[string]*docState, len(entries))
	for _, e := range entries {
		et := EventType(e.EventType)
		st := byDoc[e.DocID]
		if st == nil {
			byDoc[e.DocID] = &docState{winEvent: et, variants: e.Variants}
			continue
		}
		// Later entries overwrite earlier ones: last event wins.
		st.winEvent = et
		st.variants = e.Variants
	}
	for docID, st := range byDoc {
		switch st.winEvent {
		case EventTypeDeleted:
			// Record the tombstone for this batch's deletion so a stale
			// completion is skipped. The tombstone is committed immediately
			// (deletions are not rolled back on a failed batch because
			// re-running the delete is idempotent).
			if tomb == nil {
				tomb = map[string]uint64{}
				c.tombs[kb] = tomb
			}
			tomb[docID] = 1
			deleted = append(deleted, docID)
		case EventTypeCompleted:
			// A re-ingest after an earlier deletion: clear the prior
			// tombstone (deferred until the batch succeeds).
			if _, hadTomb := tomb[docID]; hadTomb {
				pendingTombClear = append(pendingTombClear, docID)
			}
			// Preserve the winning event's variants so the dispatch below can
			// route this doc to the matching dataset-level compile path (A0-4).
			completed = append(completed, BacklogEntry{DocID: docID, EventType: string(EventTypeCompleted), Variants: st.variants})
		}
	}
	// Sort for deterministic iteration in the delete/load passes below.
	sort.Strings(deleted)
	c.mu.Unlock()

	// Per-entry detail (each claim carries a frozen, closed batch of doc events;
	// log them on a single line so a batch can be reconstructed from the ingestor
	// log alone).
	entryDetails := make([]string, 0, len(entries))
	for _, e := range entries {
		entryDetails = append(entryDetails, fmt.Sprintf("doc=%s event=%s", e.DocID, e.EventType))
	}

	common.Info("knowledge_compile: batch merge start",
		zap.String("dataset_id", kb),
		zap.Int("completed_docs", len(completed)),
		zap.Int("deleted_docs", len(deleted)),
		zap.Strings("entries", entryDetails))

	if len(deleted) == 0 && len(completed) == 0 {
		return nil
	}

	deduper, err := c.factory(tenant)
	if err != nil || deduper == nil {
		// A no-op deduper would silently disable dataset-level LLM merging: every
		// candidate (including e.g. a "吕布" wiki_page) would be written as its own
		// merged row and duplicates would accumulate across runs. That degradation
		// is never acceptable here, so fail loudly instead of papering over it.
		common.Fatal("knowledge_compile: dataset-level LLM deduper unavailable, refusing to continue with a no-op merge",
			zap.String("kb_id", kb),
			zap.String("tenant_id", tenant),
			zap.String("factory_err", func() string {
				if err != nil {
					return err.Error()
				}
				return "deduper factory returned nil"
			}()))
	}

	deletedSet := make(map[string]bool, len(deleted))
	for _, d := range deleted {
		deletedSet[d] = true
	}

	// --- Deletion (two sequential DocEngine calls, no in-memory load) ---
	// Deleted wins regardless of batch order, so we process deletions first.
	// The deleted docs' products are never loaded into memory: the DocEngine
	// does all the work in two calls:
	//   1. DeleteDocLevelForDocs drops every per-document (doc-level) product of
	//      the deleted docs in a single engine call.
	//   2. StripMergedSources removes the deleted doc ids from the source_doc_ids
	//      array of every dataset-level merged product and deletes any product
	//      whose array became empty.
	// A merged product referencing several deleted docs is pruned in one pass,
	// and an emptied product is removed exactly once.
	if len(deleted) > 0 {
		delIDs := make([]string, 0, len(deletedSet))
		for d := range deletedSet {
			delIDs = append(delIDs, d)
		}
		// Each destructive write runs under the scheduler's per-dataset
		// write/rebuild lock with the generation check performed INSIDE the lock,
		// so a rewrite can neither land between the check and the write nor
		// interleave with the write itself (it must wait for this lock). This
		// closes the TOCTOU window where a worker past its fence could repopulate
		// storage the rebuild just cleared.
		if err := c.withWriteLock(ctx, kb, token, func() error {
			return c.writer.DeleteDocLevelForDocs(ctx, tenant, kb, delIDs)
		}); err != nil {
			if errors.Is(err, errClaimSuperseded) {
				common.Info("knowledge_compile: batch stale before delete, aborting (rewrite barrier)",
					zap.String("dataset_id", kb))
			}
			return err
		}
		// The second destructive call is its own locked section: a rewrite can
		// land between the two, in which case the stale batch must not apply its
		// second side effect under the new generation.
		if err := c.withWriteLock(ctx, kb, token, func() error {
			return c.writer.StripMergedSources(ctx, tenant, kb, delIDs)
		}); err != nil {
			if errors.Is(err, errClaimSuperseded) {
				common.Info("knowledge_compile: batch stale before strip, aborting (rewrite barrier)",
					zap.String("dataset_id", kb))
			}
			return err
		}
		// G3: structure ghost cleanup — remove structure dataset rows whose only
		// source docs were just deleted. Also a locked destructive write.
		if err := c.withWriteLock(ctx, kb, token, func() error {
			return c.writer.DeleteStructureForDocs(ctx, tenant, kb, delIDs)
		}); err != nil {
			if errors.Is(err, errClaimSuperseded) {
				common.Info("knowledge_compile: batch stale before structure ghost cleanup, aborting (rewrite barrier)",
					zap.String("dataset_id", kb))
			}
			return err
		}
	}

	// --- Completion merge ---
	// Load only the per-document products of the completed (and not deleted)
	// docs — bounded by this batch, never the whole KB. A doc that is both
	// completed and deleted is a stale tombstone: the deletion wins, so we skip
	// its completion.
	var incoming []kccommon.Product
	for _, e := range completed {
		if deletedSet[e.DocID] {
			continue
		}
		docProducts, err := c.reader.LoadDocProducts(ctx, tenant, kb, e.DocID)
		if err != nil {
			return err
		}
		// A0-4: the event's Variants are the authoritative dispatch signal — route
		// ONLY the products whose variant the compiler actually produced for this
		// doc. A doc re-compiled under a new template may still hold stale
		// doc-level rows of an earlier variant in storage; without this gate those
		// would leak into the nav/structure paths despite the event asking only
		// for, say, wiki. Legacy events with empty Variants keep all products.
		incoming = append(incoming, productsForVariants(docProducts, e.Variants)...)
	}

	// A0-4 dispatch: route tree/structure products to the dataset-navigation tree
	// (NavService.UpsertDoc), and keep wiki (and only wiki) products on the
	// existing dataset-level merge path below. This is the per-variant dispatch
	// the plan requires: tree and structure both produce a nav by-product (B2),
	// so both upsert their summary into the cross-document nav tree; wiki products
	// continue to the replace-only merge.
	navIn := navInputFromProducts(kb, incoming)
	if len(navIn) > 0 {
		ns := nav.GetNavService()
		if ns == nil {
			common.Warn("knowledge_compile: nav service unavailable, skipping dataset-nav upsert for batch",
				zap.String("dataset_id", kb))
		} else if err := c.upsertNavLocked(ctx, tenant, kb, token, ns, navIn); err != nil {
			return err
		}
	}

	// G1+G2: dataset-level structure merge and build-time stamp, both under the
	// claim-fenced write lock (B7 applies to structure writes too): a worker
	// invalidated by CancelInflight must not write structure rows after a rebuild
	// has cleared them.
	if err := c.withWriteLock(ctx, kb, token, func() error {
		return c.mergeStructureDataset(ctx, tenant, kb, incoming)
	}); err != nil {
		if errors.Is(err, errClaimSuperseded) {
			common.Info("knowledge_compile: batch stale before structure merge, aborting (rewrite barrier)",
				zap.String("dataset_id", kb))
		}
		return err
	}

	// wiki_incremental port (M1): the dataset-level merge only processes wiki
	// PAGES. A wiki doc yields both page and section products (Meta.kind
	// "page"/"section"); sections are a doc-level concern and must never be folded
	// into the dataset-level page bucket. Filter BEFORE the in-memory dedup so
	// sections never enter the deduper (and never trigger LLM/embedding/alias
	// processing) — the decider must not carry an implicit section-filter
	// contract. Legacy rows whose kind is empty are derived in the Reader; only
	// truly page-kind wiki products proceed.
	candidates := filterWikiPageCandidates(incoming)
	// In-memory dedup among the completed batch first.
	candidates, dedupErr := deduper.Dedup(ctx, candidates)
	if dedupErr != nil {
		return dedupErr
	}
	// Diagnostics: after batch dedup, verify the per-doc products still carry
	// their embedding before the KNN/merge path consumes cand.Vector. If
	// candidates show 0 vectors while incoming had vectors, Dedup is dropping
	// them and both the KNN lookups and the merged rows lose embeddings.
	{
		incomingVec, candVec := 0, 0
		for _, p := range incoming {
			if len(p.Vector) > 0 {
				incomingVec++
			}
		}
		for _, p := range candidates {
			if len(p.Vector) > 0 {
				candVec++
			}
		}
		common.Info("knowledge_compile: dedup vector audit",
			zap.String("kb_id", kb),
			zap.Int("incoming", len(incoming)),
			zap.Int("incoming_with_vector", incomingVec),
			zap.Int("candidates", len(candidates)),
			zap.Int("candidates_with_vector", candVec))
	}

	// Diagnostics: break the KNN-input set down by variant and list the wiki_page
	// slugs, so the reader can reconcile the number of wiki_page candidates here
	// against the wiki_page rows WriteMerged emits (they differ when KNN merges a
	// candidate into an existing merged row, when DecideBatch returns distinct new
	// rows, or when candidates of the same variant collapse in-memory). This is
	// what explains "39 wiki_page written vs 21 KNN requests".
	{
		byVariant := map[string]int{}
		var wikiPageSlugs []string
		for _, cand := range candidates {
			v := string(cand.Variant)
			byVariant[v]++
			if v == "wiki_page" {
				wikiPageSlugs = append(wikiPageSlugs, candidateIdentity(cand))
			}
		}
		common.Info("knowledge_compile: candidates breakdown",
			zap.String("kb_id", kb),
			zap.Int("candidates", len(candidates)),
			zap.Any("by_variant", byVariant),
			zap.Strings("wiki_page_candidate_slugs", wikiPageSlugs))
	}

	// Then dedup each candidate against the DocEngine via KNN top1 + LLM judge,
	// mirroring Python _struct_doc_storage_dedup_batch. Candidates that KNN-hit
	// the same existing merged row are grouped so the row is merged/updated once
	// instead of being rewritten per candidate.
	type mergeGroup struct {
		existing   kccommon.Product
		candidates []kccommon.Product
		score      float64
	}
	groupsByID := make(map[string]*mergeGroup, len(candidates))
	var (
		unmatchedMu sync.Mutex
		unmatched   []kccommon.Product
		groupsMu    sync.Mutex
	)
	// The KNN pass is docengine-bounded (vector search), not CPU-bounded, so we
	// fan it out across the shared global compilerPool (vCPU-sized). Output order
	// is irrelevant: merged rows are upserted by their idempotent dataset-level
	// id, and each candidate lands in exactly one group / the unmatched set.
	jobs := make([]CompilerJob, 0, len(candidates))
	for _, cand := range candidates {
		cand := cand
		jobs = append(jobs, func() error {
			var vec64 []float64
			if len(cand.Vector) > 0 {
				vec64 = make([]float64, len(cand.Vector))
				for i, v := range cand.Vector {
					vec64[i] = float64(v)
				}
			}
			// wiki_incremental port (B2): request topN >= 2 so SearchSimilar can
			// apply its score-descending skip rule on a dirty top-1 (the reader
			// drops rows whose compile_kwd does not map to the searched variant;
			// a re-query with topN=2 lets the next clean candidate surface instead
			// of falling through to "no hit").
			hit, score, err := c.reader.SearchSimilar(ctx, tenant, kb, cand.Variant, vec64, 2, c.mergeThreshold)
			if err != nil {
				return err
			}
			if hit.ID == "" {
				// No sufficiently-similar merged row: insert the candidate as a new
				// merged row.
				common.Debug("knowledge_compile: KNN no hit, candidate becomes new merged row",
					zap.String("kb_id", kb),
					zap.String("candidate", candidateIdentity(cand)),
					zap.String("variant", string(cand.Variant)))
				cand.Merged = true
				cand.DocID = kb
				unmatchedMu.Lock()
				unmatched = append(unmatched, cand)
				unmatchedMu.Unlock()
				return nil
			}
			common.Debug("knowledge_compile: KNN hit",
				zap.String("kb_id", kb),
				zap.String("candidate", candidateIdentity(cand)),
				zap.String("variant", string(cand.Variant)),
				zap.String("hit", candidateIdentity(hit)),
				zap.Float64("score", score))
			groupsMu.Lock()
			g := groupsByID[hit.ID]
			if g == nil {
				g = &mergeGroup{existing: hit, score: score}
				groupsByID[hit.ID] = g
			}
			g.candidates = append(g.candidates, cand)
			groupsMu.Unlock()
			return nil
		})
	}
	if err := runCompilerJobs(ctx, jobs); err != nil {
		return err
	}

	// Fold every KNN group into the LLM in a single batch round-trip (one
	// DecideBatch call instead of one Decide per pair), then collect the
	// updated existing rows and the candidates judged distinct (new rows).
	var newMerged []kccommon.Product
	if len(groupsByID) > 0 {
		// Diagnostics: summarize the KNN groups before the LLM merge decision so
		// the reader can see which existing merged rows were hit and by how many
		// candidates (e.g. whether a "吕布" candidate hit an existing 吕布 row and
		// was then judged by the LLM).
		{
			type groupDump struct {
				Existing   string   `json:"existing"`
				Candidates []string `json:"candidates"`
				Score      float64  `json:"score"`
			}
			dumps := make([]groupDump, 0, len(groupsByID))
			for _, g := range groupsByID {
				cands := make([]string, 0, len(g.candidates))
				for _, cand := range g.candidates {
					cands = append(cands, candidateIdentity(cand))
				}
				dumps = append(dumps, groupDump{Existing: candidateIdentity(g.existing), Candidates: cands, Score: g.score})
			}
			common.Info("knowledge_compile: KNN groups for DecideBatch",
				zap.String("kb_id", kb),
				zap.Int("groups", len(groupsByID)),
				zap.Any("groups_detail", dumps))
		}
		batched := make([]MergeGroup, 0, len(groupsByID))
		for _, g := range groupsByID {
			batched = append(batched, MergeGroup{
				Existing:   g.existing,
				Candidates: g.candidates,
				Score:      g.score,
			})
		}
		batched, err = deduper.DecideBatch(ctx, batched)
		if err != nil {
			return err
		}
		for _, g := range batched {
			newMerged = append(newMerged, g.Merged)
			unmatched = append(unmatched, g.Distinct...)
		}
	}

	// Write the surviving merged set (updated existing + new distinct rows). The
	// generation check + WriteMerged run atomically under the per-dataset
	// write/rebuild lock, so a rewrite that lands during the (now long) KNN + LLM
	// merge cannot leak into the freshly rewritten index, and a rewrite that is
	// clearing storage cannot interleave with the write.
	mergedFinal := make([]kccommon.Product, 0, len(newMerged)+len(unmatched))
	mergedFinal = append(mergedFinal, newMerged...)
	mergedFinal = append(mergedFinal, unmatched...)
	if err := c.withWriteLock(ctx, kb, token, func() error {
		return c.writer.WriteMerged(ctx, tenant, kb, mergedFinal)
	}); err != nil {
		if errors.Is(err, errClaimSuperseded) {
			common.Info("knowledge_compile: batch stale before write, aborting (rewrite barrier)",
				zap.String("dataset_id", kb))
		}
		return err
	}

	// Re-materialize the wiki page graph from the now-consistent merged set. This
	// runs unconditionally after a merge: the graph is a global per-dataset
	// projection rebuilt from a consistent read each time, so it must run even
	// when the current batch only prunes (deletes) wiki pages, and even for a
	// non-wiki batch where the dataset already has no wiki pages (ProjectWikiGraph
	// then just drops any stale graph rows). The graph is reconstructible, so the
	// cost of an occasional no-op reprojection is accepted. It is also a locked
	// destructive side effect for the same TOCTOU reasons as WriteMerged.
	if err := c.withWriteLock(ctx, kb, token, func() error {
		return c.writer.ProjectWikiGraph(ctx, tenant, kb)
	}); err != nil {
		if errors.Is(err, errClaimSuperseded) {
			common.Info("knowledge_compile: batch stale before graph projection, aborting (rewrite barrier)",
				zap.String("dataset_id", kb))
		}
		return err
	}

	// All merge and delete paths succeeded: commit the staged tombstone
	// clears. These are only now persisted so a failed batch leaves c.tombs
	// untouched and a later reclaim can retry.
	c.mu.Lock()
	for _, docID := range pendingTombClear {
		delete(c.tombs[kb], docID)
	}
	c.mu.Unlock()
	common.Info("knowledge_compile: batch merge complete",
		zap.String("dataset_id", kb),
		zap.Int("completed_docs", len(completed)),
		zap.Int("deleted_docs", len(deleted)),
		zap.Int("merged_rows_written", len(mergedFinal)))
	return nil
}

// candidateIdentity returns a compact identity string for a product, preferring
// the wiki slug (full "<page_type>/<slug>") when present, else the id/doc_id,
// so KNN/dedup diagnostics can tie a candidate to the page it belongs to.
func candidateIdentity(p kccommon.Product) string {
	if slug, _ := p.Meta["slug"].(string); slug != "" {
		return slug
	}
	if id := p.ID; id != "" {
		return id
	}
	if docID := p.DocID; docID != "" {
		return docID
	}
	return "<unknown>"
}

// productsForVariants filters a doc's loaded products down to those whose
// variant is in the requested set. An empty set (legacy backlog entries without
// variants, or a rebuild event that carries none) is treated as "all variants"
// to preserve existing behavior; otherwise a product is kept only when its
// variant appears in variants. This is what makes BacklogEntry.Variants actually
// drive per-variant dispatch instead of handing every doc's full product set to
// every dataset-level path.
func productsForVariants(products []kccommon.Product, variants []string) []kccommon.Product {
	if len(variants) == 0 {
		return products
	}
	want := make(map[string]struct{}, len(variants))
	for _, v := range variants {
		want[v] = struct{}{}
	}
	out := products[:0:0]
	for i := range products {
		if _, ok := want[string(products[i].Variant)]; ok {
			out = append(out, products[i])
		}
	}
	return out
}

// navInputFromProducts extracts the dataset-navigation upsert inputs for the
// tree and structure products in a batch (B2: nav trigger is tree || structure).
// Summary extraction matches the component's by-product hooks:
//   - tree: the root product (Meta.kind=="root") carries the document summary in
//     Content and its vector in Vector.
//   - structure: the graph product (Meta.kind=="graph") carries the graph JSON in
//     Content; its entity descriptions are folded into a document-level summary
//     via pageIndexSummary.
//
// Each document contributes at most one nav input, keyed by DocID, so a doc that
// yields several tree/structure products upserts once. NavService embeds the
// summary itself when Embedd is empty, so the consumer needs no embedder.
func navInputFromProducts(kb string, products []kccommon.Product) []nav.UpsertDocInput {
	type acc struct {
		in nav.UpsertDocInput
	}
	byDoc := make(map[string]*acc, len(products))
	for i := range products {
		p := &products[i]
		switch p.Variant {
		case kccommon.VariantTree:
			if kind, _ := p.Meta["kind"].(string); kind != "root" {
				continue
			}
			a := byDoc[p.DocID]
			if a == nil {
				byDoc[p.DocID] = &acc{in: nav.UpsertDocInput{TenantID: p.TenantID, KbID: kb, DocID: p.DocID}}
				a = byDoc[p.DocID]
			}
			a.in.Summary = p.Content
			a.in.Embedd = p.Vector
		case kccommon.VariantStructure:
			if kind, _ := p.Meta["kind"].(string); kind != "graph" {
				continue
			}
			a := byDoc[p.DocID]
			if a == nil {
				byDoc[p.DocID] = &acc{in: nav.UpsertDocInput{TenantID: p.TenantID, KbID: kb, DocID: p.DocID}}
				a = byDoc[p.DocID]
			}
			// Graph vector is NOT the summary vector; leave Embedd empty so
			// NavService embeds the folded summary text.
			a.in.Summary = pageIndexSummary(p.Content)
		}
	}
	out := make([]nav.UpsertDocInput, 0, len(byDoc))
	for _, a := range byDoc {
		if strings.TrimSpace(a.in.Summary) == "" {
			continue
		}
		out = append(out, a.in)
	}
	return out
}

// pageIndexSummary folds the entity descriptions of a structure graph JSON
// ({"entities":[{"name","description"},...]}) into a document-level summary for
// dataset navigation, matching the component's by-product logic.
func pageIndexSummary(graphJSON string) string {
	var graph struct {
		Entities []struct {
			Name        string `json:"name"`
			Description string `json:"description"`
		} `json:"entities"`
	}
	if err := json.Unmarshal([]byte(graphJSON), &graph); err != nil {
		return ""
	}
	var b strings.Builder
	for _, e := range graph.Entities {
		desc := strings.Join(strings.Fields(e.Description), " ")
		if desc == "" {
			continue
		}
		if e.Name != "" {
			b.WriteString(e.Name)
			b.WriteString(": ")
		}
		b.WriteString(desc)
		b.WriteString("\n")
	}
	return b.String()
}

// mergeStructureDataset performs the dataset-level structure merge (G1/G4): it
// groups the structure products of a batch by (name, type), folds their
// descriptions and source doc/chunk sets into one StructureBucket per group, and
// writes scope_kwd="dataset" rows via WriteMergedStructure. This mirrors Python
// dataset_structure_merger._merge_bucket but is a self-contained from-scratch
// path (B3: the legacy unified merge never aggregated structure by name/type).
func (c *Consumer) mergeStructureDataset(ctx context.Context, tenant, kb string, products []kccommon.Product) error {
	// Bucket by (lower(name), type) for entities and (lower(from), lower(to))
	// for relations — relations have no "name" and must not be silently dropped
	// (review issue 3).
	type ekey struct{ name, typ string }
	type rkey struct{ from, to string }
	entByKey := make(map[ekey]*StructureBucket, 16)
	relByKey := make(map[rkey]*StructureBucket, 16)
	for _, p := range products {
		if p.Variant != kccommon.VariantStructure {
			continue
		}
		kind := metaString(p.Meta, "kind")
		from := metaString(p.Meta, "from")
		to := metaString(p.Meta, "to")
		if kind == "relation" || (from != "" && to != "") {
			// Skip a relation with an incomplete endpoint pair: writing a bucket
			// like "A -> " would let multiple malformed relations collapse into
			// the same dataset row.
			if from == "" || to == "" {
				continue
			}
			k := rkey{from: strings.ToLower(from), to: strings.ToLower(to)}
			b := relByKey[k]
			if b == nil {
				b = &StructureBucket{Name: from + " -> " + to, Type: "relation", FromEntity: from, ToEntity: to}
				relByKey[k] = b
			}
			appendBucket(b, p)
			continue
		}
		// entity bucket
		name := metaString(p.Meta, "name")
		typ := metaString(p.Meta, "entity_type")
		if typ == "" {
			typ = metaString(p.Meta, "type")
		}
		if name == "" {
			continue
		}
		k := ekey{name: strings.ToLower(name), typ: typ}
		b := entByKey[k]
		if b == nil {
			b = &StructureBucket{Name: name, Type: typ}
			entByKey[k] = b
		}
		appendBucket(b, p)
	}
	buckets := make([]StructureBucket, 0, len(entByKey)+len(relByKey))
	for _, b := range entByKey {
		buckets = append(buckets, *b)
	}
	for _, b := range relByKey {
		buckets = append(buckets, *b)
	}
	if len(buckets) == 0 {
		return nil
	}
	sort.Slice(buckets, func(i, j int) bool {
		if buckets[i].Name != buckets[j].Name {
			return buckets[i].Name < buckets[j].Name
		}
		return buckets[i].Type < buckets[j].Type
	})
	return c.writer.WriteMergedStructure(ctx, tenant, kb, buckets)
}

// appendBucket folds a structure product into a bucket: concatenates its
// description, unions its source doc/chunk ids, and folds its vector into a true
// element-wise running mean (review issue 13). VecCount tracks how many vectors
// have been folded so the mean is order-independent.
func appendBucket(b *StructureBucket, p kccommon.Product) {
	if strings.TrimSpace(p.Content) != "" {
		if b.Description != "" {
			b.Description += "\n"
		}
		b.Description += strings.TrimSpace(p.Content)
	}
	b.SourceDocIDs = appendUnique(b.SourceDocIDs, []string{p.DocID})
	b.SourceChunkIDs = appendUnique(b.SourceChunkIDs, metaStringSlice(p.Meta, "source_chunk_ids"))
	if len(p.Vector) > 0 {
		b.VecCount++
		if len(b.Vector) == 0 {
			b.Vector = append([]float32(nil), p.Vector...)
		} else {
			n := len(p.Vector)
			if len(b.Vector) < n {
				b.Vector = append(b.Vector, make([]float32, n-len(b.Vector))...)
			}
			for i := 0; i < n; i++ {
				// running mean: acc + (next - acc)/count
				b.Vector[i] += (p.Vector[i] - b.Vector[i]) / float32(b.VecCount)
			}
		}
	}
}

// appendUnique appends only values not already present, preserving order.
func appendUnique(dst, values []string) []string {
	seen := make(map[string]bool, len(dst)+len(values))
	for _, v := range dst {
		seen[v] = true
	}
	for _, v := range values {
		if v == "" || seen[v] {
			continue
		}
		seen[v] = true
		dst = append(dst, v)
	}
	return dst
}

// upsertNavLocked places the batch's tree/structure summaries into the dataset
// navigation tree under the per-dataset write/rebuild lock, with the claim-token
// check inside the lock (B7: nav writes are claim-fenced like any other
// destructive write so a stale worker cannot repopulate a rebuilt tree).
func (c *Consumer) upsertNavLocked(ctx context.Context, tenant, kb, token string, ns nav.NavService, inputs []nav.UpsertDocInput) error {
	return c.withWriteLock(ctx, kb, token, func() error {
		for i := range inputs {
			// The minimal-loop NavService.UpsertDoc is idempotent by summary; a
			// failure here aborts the batch so the claim is not acked and the
			// batch is retried.
			if err := ns.UpsertDoc(ctx, inputs[i]); err != nil {
				return fmt.Errorf("knowledge_compile: nav upsert %s: %w", inputs[i].DocID, err)
			}
		}
		return nil
	})
}