xinmi/scripts/native_graph_backfill.go

// Run explicitly by filename, like scripts/wiki_backfill_trigger.go.
package main

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"net"
	"net/url"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/Tencent/WeKnora/internal/common"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/google/uuid"
	"github.com/hibiken/asynq"
	"gorm.io/driver/postgres"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
	"gorm.io/gorm/logger"
)

type graphRecoveryCandidate struct {
	ID          string `json:"chunk_id"`
	KnowledgeID string `json:"knowledge_id"`
	SeqID       int64  `json:"seq_id"`
	ChunkIndex  int    `json:"chunk_index"`
	ContentHash string `json:"content_sha256"`
}

// Failure records survive Redis retention and watcher restarts. Raw provider
// errors are deliberately excluded: state and stable identities are sufficient
// to find the task's private diagnostics and authorize a bounded manual retry.
type graphRecoveryFailure struct {
	TaskID          string    `gorm:"primaryKey" json:"task_id"`
	TenantID        uint64    `gorm:"index:idx_graph_backfill_scope;index:idx_graph_backfill_candidate,priority:1" json:"tenant_id"`
	KnowledgeBaseID string    `gorm:"index:idx_graph_backfill_scope;index:idx_graph_backfill_candidate,priority:2" json:"knowledge_base_id"`
	ChunkID         string    `gorm:"index:idx_graph_backfill_candidate,priority:3" json:"chunk_id"`
	KnowledgeID     string    `json:"knowledge_id"`
	ContentHash     string    `gorm:"index:idx_graph_backfill_candidate,priority:4" json:"content_sha256"`
	Revision        string    `gorm:"index:idx_graph_backfill_candidate,priority:5" json:"revision"`
	State           string    `gorm:"index:idx_graph_backfill_candidate,priority:6" json:"state"`
	Reason          string    `json:"reason"`
	ManualRetries   int       `json:"manual_retries"`
	RetryReason     string    `json:"retry_reason,omitempty"`
	FirstFailedAt   time.Time `json:"first_failed_at"`
	UpdatedAt       time.Time `json:"updated_at"`
}

func (graphRecoveryFailure) TableName() string { return "native_graph_backfill_failures" }

func graphFailureFilter(includeFailed bool) string {
	if includeFailed {
		return ""
	}
	return ` AND NOT EXISTS(SELECT 1 FROM native_graph_backfill_failures f WHERE f.tenant_id=c.tenant_id AND f.knowledge_base_id=c.knowledge_base_id AND f.chunk_id=c.id AND f.content_hash=encode(sha256(convert_to(c.content,'UTF8')),'hex') AND f.revision=? AND f.state='failed')`
}

const graphCurrentFailureSource = `EXISTS(SELECT 1 FROM chunks c JOIN knowledges k ON k.id=c.knowledge_id AND k.tenant_id=c.tenant_id AND k.knowledge_base_id=c.knowledge_base_id AND k.deleted_at IS NULL WHERE c.tenant_id=f.tenant_id AND c.knowledge_base_id=f.knowledge_base_id AND c.id=f.chunk_id AND c.knowledge_id=f.knowledge_id AND c.deleted_at IS NULL AND c.is_enabled AND c.chunk_type='text' AND encode(sha256(convert_to(c.content,'UTF8')),'hex')=f.content_hash)`
const graphFailureCheckpoint = `EXISTS(SELECT 1 FROM native_graph_records r WHERE r.tenant_id=f.tenant_id AND r.knowledge_base_id=f.knowledge_base_id AND r.source_table='knowledgegraphchunkstate' AND r.payload->>'origin'='native_extraction_v1' AND r.payload->>'extraction_status'='complete' AND CAST(COALESCE(r.payload->>'relation_count','0') AS INTEGER)>0 AND r.payload->>'native_chunk_id'=f.chunk_id AND r.payload->>'native_knowledge_id'=f.knowledge_id AND r.payload->>'source_content_hash'=f.content_hash)`

func graphUnresolvedFailures(db *gorm.DB, tenant uint64, kbID, revision string, reconcile bool) ([]graphRecoveryFailure, error) {
	if reconcile {
		// Keep historical rows, but a deleted/changed source or a newly verified
		// checkpoint must not leave a completed knowledge base blocked forever.
		for _, resolution := range []struct{ state, condition string }{
			{"superseded", "NOT " + graphCurrentFailureSource},
			{"resolved", graphCurrentFailureSource + " AND " + graphFailureCheckpoint},
		} {
			if err := db.Exec("UPDATE native_graph_backfill_failures AS f SET state=?,updated_at=? WHERE tenant_id=? AND knowledge_base_id=? AND revision=? AND state NOT IN ('resolved','superseded') AND "+resolution.condition, resolution.state, time.Now().UTC(), tenant, kbID, revision).Error; err != nil {
				return nil, err
			}
		}
	}
	var failures []graphRecoveryFailure
	err := db.Table("native_graph_backfill_failures AS f").Where("f.tenant_id=? AND f.knowledge_base_id=? AND f.revision=? AND f.state='failed'", tenant, kbID, revision).Where(graphCurrentFailureSource).Where("NOT " + graphFailureCheckpoint).Order("f.first_failed_at,f.task_id").Find(&failures).Error
	return failures, err
}

func graphRecoveryCandidates(db *gorm.DB, tenant uint64, kbID, revision string, after int64, limit int, hasFailures, includeFailed bool, exact ...graphRecoveryCandidate) ([]graphRecoveryCandidate, error) {
	var candidates []graphRecoveryCandidate
	filter := ""
	queryArgs := []any{tenant, kbID, tenant, kbID, after}
	if hasFailures && !includeFailed {
		filter = graphFailureFilter(false)
		queryArgs = append(queryArgs, revision)
	}
	if len(exact) > 0 {
		filter += " AND c.id=? AND c.knowledge_id=? AND encode(sha256(convert_to(c.content,'UTF8')),'hex')=?"
		queryArgs = append(queryArgs, exact[0].ID, exact[0].KnowledgeID, exact[0].ContentHash)
	}
	queryArgs = append(queryArgs, limit)
	err := db.Raw(`WITH completed AS NOT MATERIALIZED(SELECT payload FROM native_graph_records WHERE tenant_id=? AND knowledge_base_id=? AND source_table='knowledgegraphchunkstate' AND payload->>'origin'='native_extraction_v1' AND payload->>'extraction_status'='complete' AND CAST(COALESCE(payload->>'relation_count','0') AS INTEGER)>0)
		SELECT c.id,c.knowledge_id,c.seq_id,c.chunk_index,encode(sha256(convert_to(c.content,'UTF8')),'hex') AS content_hash
		FROM chunks c JOIN knowledges k ON k.id=c.knowledge_id AND k.tenant_id=c.tenant_id AND k.knowledge_base_id=c.knowledge_base_id AND k.deleted_at IS NULL
		WHERE c.tenant_id=? AND c.knowledge_base_id=? AND c.deleted_at IS NULL AND c.is_enabled AND c.chunk_type='text' AND c.seq_id>? AND k.parse_status='completed'
		AND NOT EXISTS(SELECT 1 FROM completed r WHERE r.payload->>'native_chunk_id'=c.id AND r.payload->>'native_knowledge_id'=c.knowledge_id AND r.payload->>'source_content_hash'=encode(sha256(convert_to(c.content,'UTF8')),'hex'))
		`+filter+` ORDER BY c.seq_id LIMIT ?`, queryArgs...).Scan(&candidates).Error
	if err == nil && len(exact) > 0 && len(candidates) != 1 {
		return nil, errors.New("exact graph source unavailable, changed, completed or excluded by its failure ledger; no substitute was selected")
	}
	return candidates, err
}

func graphValidateExactScope(chunkID, knowledgeID, hash string, after int64, limit, maxOutstanding int) error {
	if chunkID == "" && knowledgeID == "" && hash == "" {
		return nil
	}
	chunk, chunkErr := uuid.Parse(chunkID)
	knowledge, knowledgeErr := uuid.Parse(knowledgeID)
	decoded, hashErr := hex.DecodeString(hash)
	if chunkErr != nil || knowledgeErr != nil || chunk == uuid.Nil || knowledge == uuid.Nil ||
		chunk.String() != chunkID || knowledge.String() != knowledgeID ||
		hashErr != nil || len(decoded) != sha256.Size || hex.EncodeToString(decoded) != hash ||
		after != 0 || limit != 1 || maxOutstanding != 1 {
		return errors.New("exact graph scope requires canonical chunk and knowledge UUIDs, lowercase source SHA-256, limit=1, max-outstanding=1 and no cursor")
	}
	return nil
}

func graphTaskMatchesReviewedSource(payload []byte, tenant uint64, kbID string, c graphRecoveryCandidate) bool {
	var task types.ExtractChunkPayload
	return json.Unmarshal(payload, &task) == nil && task.TenantID == tenant && task.ChunkID == c.ID &&
		(task.KnowledgeID == "" || task.KnowledgeID == c.KnowledgeID) &&
		task.ExpectedSource != nil && task.ExpectedSource.KnowledgeBaseID == kbID &&
		task.ExpectedSource.KnowledgeID == c.KnowledgeID && task.ExpectedSource.ContentSHA256 == c.ContentHash
}

type graphRecoveryQueue interface {
	RunTask(string, string) error
}

func graphQueueSlots(info *asynq.QueueInfo, limit int) int {
	if info == nil || info.Paused {
		return 0
	}
	remaining := limit - info.Active - info.Pending - info.Scheduled - info.Retry - info.Aggregating
	if remaining < 0 {
		return 0
	}
	return remaining
}

func graphResumeCursor(after int64, candidates []graphRecoveryCandidate, settled map[string]bool) int64 {
	for _, c := range candidates {
		if !settled[c.ID] {
			break
		}
		after = c.SeqID
	}
	return after
}

func graphRecoveryTaskID(tenant uint64, kbID string, candidate graphRecoveryCandidate, revision string) string {
	identity := sha256.Sum256([]byte(fmt.Sprintf("%d:%s:%s:%s:%s", tenant, kbID, candidate.ID, candidate.ContentHash, revision)))
	return "native-graph-" + hex.EncodeToString(identity[:])
}

func recordGraphFailure(db *gorm.DB, tenant uint64, kbID, revision, taskID string, c graphRecoveryCandidate, reason string) error {
	var failure graphRecoveryFailure
	err := db.First(&failure, "task_id=?", taskID).Error
	if errors.Is(err, gorm.ErrRecordNotFound) {
		failure = graphRecoveryFailure{TaskID: taskID, TenantID: tenant, KnowledgeBaseID: kbID, ChunkID: c.ID, KnowledgeID: c.KnowledgeID, ContentHash: c.ContentHash, Revision: revision, FirstFailedAt: time.Now().UTC()}
	} else if err != nil {
		return err
	}
	if failure.FirstFailedAt.IsZero() {
		failure.FirstFailedAt = time.Now().UTC()
	}
	failure.State, failure.Reason, failure.UpdatedAt = "failed", reason, time.Now().UTC()
	return db.Save(&failure).Error
}

func reserveGraphDispatch(db *gorm.DB, tenant uint64, kbID, revision, taskID string, c graphRecoveryCandidate) (bool, error) {
	var prior graphRecoveryFailure
	err := db.First(&prior, "task_id=?", taskID).Error
	if err == nil {
		// An absent Redis task is not permission to repeat an earlier attempt.
		return false, recordGraphFailure(db, tenant, kbID, revision, taskID, c, "task_missing_after_dispatch")
	}
	if !errors.Is(err, gorm.ErrRecordNotFound) {
		return false, err
	}
	now := time.Now().UTC()
	entry := graphRecoveryFailure{TaskID: taskID, TenantID: tenant, KnowledgeBaseID: kbID, ChunkID: c.ID, KnowledgeID: c.KnowledgeID, ContentHash: c.ContentHash, Revision: revision, State: "dispatch_requested", UpdatedAt: now}
	if err := db.Create(&entry).Error; err != nil {
		return false, err
	}
	return true, nil
}

func trackGraphTask(db *gorm.DB, tenant uint64, kbID, revision, taskID string, c graphRecoveryCandidate) error {
	entry := graphRecoveryFailure{TaskID: taskID, TenantID: tenant, KnowledgeBaseID: kbID, ChunkID: c.ID, KnowledgeID: c.KnowledgeID, ContentHash: c.ContentHash, Revision: revision, State: "dispatch_requested", UpdatedAt: time.Now().UTC()}
	return db.Clauses(clause.OnConflict{DoNothing: true}).Create(&entry).Error
}

func graphArchivedReason(lastError string) string {
	message := strings.ToLower(lastError)
	if strings.Contains(message, "graph_needs_review") {
		return "relation_coverage_review_required"
	}
	if strings.Contains(message, "read response") && (strings.Contains(message, "context deadline exceeded") || strings.Contains(message, "timeout")) {
		return "provider_response_timeout"
	}
	if strings.Contains(message, "context deadline exceeded") || strings.Contains(message, "deadline exceeded") || strings.Contains(message, "timeout") {
		return "task_timeout"
	}
	if strings.Contains(message, "context canceled") || strings.Contains(message, "context cancelled") {
		return "task_cancelled"
	}
	return "archived"
}

// Reconciliation receives only the Inspector read API, never a queue client.
type graphTaskReader interface {
	GetTaskInfo(string, string) (*asynq.TaskInfo, error)
}

type graphReconcileReport struct {
	Mode               string                 `json:"mode"`
	Apply              bool                   `json:"apply"`
	TenantID           uint64                 `json:"tenant_id"`
	KnowledgeBaseID    string                 `json:"knowledge_base_id"`
	Revision           string                 `json:"revision"`
	Source             graphRecoveryCandidate `json:"source"`
	TaskID             string                 `json:"task_id"`
	QueueState         string                 `json:"queue_state"`
	LedgerState        string                 `json:"ledger_state"`
	ProposedState      string                 `json:"proposed_state"`
	SourceCurrent      bool                   `json:"source_current"`
	CheckpointComplete bool                   `json:"checkpoint_complete"`
	Reason             string                 `json:"reason"`
	Blocked            bool                   `json:"blocked"`
	WouldUpdate        bool                   `json:"would_update"`
	LedgerUpdated      bool                   `json:"ledger_updated"`
	Enqueued           int                    `json:"enqueued"`
	Retried            int                    `json:"retried"`
	ModelCalls         int                    `json:"model_calls"`
}

const graphReconcileCurrentSource = graphCurrentFailureSource + ` AND EXISTS(SELECT 1 FROM knowledges k JOIN knowledge_bases b ON b.id=k.knowledge_base_id AND b.tenant_id=k.tenant_id WHERE k.id=f.knowledge_id AND k.tenant_id=f.tenant_id AND k.knowledge_base_id=f.knowledge_base_id AND k.deleted_at IS NULL AND k.parse_status='completed' AND b.deleted_at IS NULL)`

func graphValidateReconciliation(reconcile bool, chunkID, revision string, retry bool, retryReason string) error {
	if reconcile && (chunkID == "" || strings.TrimSpace(revision) == "" || retry || retryReason != "") {
		return errors.New("reconcile-only requires one exact source, GRAPH_BACKFILL_REVISION and no retry flags")
	}
	return nil
}

func graphReconcileSourceStatus(db *gorm.DB, taskID string) (current, complete bool, err error) {
	var status struct {
		Current, Complete bool
	}
	result := db.Raw("SELECT ("+graphReconcileCurrentSource+") AS current,("+graphFailureCheckpoint+") AS complete FROM native_graph_backfill_failures f WHERE f.task_id=?", taskID).Scan(&status)
	if result.Error != nil {
		return false, false, errors.New("cannot validate graph source and checkpoint")
	}
	if result.RowsAffected != 1 {
		return false, false, errors.New("graph reconciliation ledger disappeared")
	}
	return status.Current, status.Complete, nil
}

func reconcileGraphTask(db *gorm.DB, reader graphTaskReader, tenant uint64, kbID, revision string, c graphRecoveryCandidate, apply bool) (graphReconcileReport, error) {
	report := graphReconcileReport{Mode: "reconcile_only", Apply: apply, TenantID: tenant, KnowledgeBaseID: kbID, Revision: revision, Source: c,
		TaskID: graphRecoveryTaskID(tenant, kbID, c, revision), QueueState: "unknown", Blocked: true}
	var ledger graphRecoveryFailure
	err := db.First(&ledger, "task_id=? AND tenant_id=? AND knowledge_base_id=? AND chunk_id=? AND knowledge_id=? AND content_hash=? AND revision=?",
		report.TaskID, tenant, kbID, c.ID, c.KnowledgeID, c.ContentHash, revision).Error
	if errors.Is(err, gorm.ErrRecordNotFound) {
		report.Reason = "ledger_missing_or_identity_mismatch"
		return report, nil
	}
	if err != nil {
		return report, errors.New("cannot read exact graph ledger")
	}
	report.LedgerState, report.ProposedState = ledger.State, ledger.State
	report.SourceCurrent, report.CheckpointComplete, err = graphReconcileSourceStatus(db, report.TaskID)
	if err != nil {
		return report, err
	}
	if !report.SourceCurrent {
		report.Reason = "source_changed_or_unavailable"
		return report, nil
	}
	if ledger.State == "failed" && strings.HasPrefix(ledger.Reason, "semantic_review_rejected") {
		report.Reason = "manual_semantic_review_required"
		return report, nil
	}
	info, err := reader.GetTaskInfo(types.QueueGraph, report.TaskID)
	if errors.Is(err, asynq.ErrTaskNotFound) || errors.Is(err, asynq.ErrQueueNotFound) || (err == nil && info == nil) {
		report.QueueState, report.Reason = "missing", "task_missing_after_dispatch"
		return report, nil
	}
	if err != nil {
		report.Reason = "queue_inspection_failed"
		return report, errors.New("cannot inspect graph task; no reconciliation applied")
	}
	report.QueueState = info.State.String()
	if info.ID != report.TaskID || info.Queue != types.QueueGraph || info.Type != types.TypeChunkExtract ||
		!graphTaskMatchesReviewedSource(info.Payload, tenant, kbID, c) || info.MaxRetry != 0 {
		report.Reason = "queue_task_identity_mismatch"
		return report, nil
	}
	if info.State != asynq.TaskStateArchived && info.State != asynq.TaskStateCompleted {
		report.Blocked, report.Reason = false, "task_not_terminal"
		return report, nil
	}
	// A worker can commit between the first source read and terminal queue read.
	report.SourceCurrent, report.CheckpointComplete, err = graphReconcileSourceStatus(db, report.TaskID)
	if err != nil {
		return report, err
	}
	if !report.SourceCurrent {
		report.Reason = "source_changed_or_unavailable"
		return report, nil
	}
	if report.CheckpointComplete {
		report.ProposedState, report.Reason, report.Blocked = "resolved", "native_checkpoint_verified", false
	} else {
		report.ProposedState, report.Reason = "failed", "completed_without_native_checkpoint"
		if info.State == asynq.TaskStateArchived {
			report.Reason = graphArchivedReason(info.LastErr)
		}
	}
	if ledger.State == "superseded" || (ledger.State == "resolved" && report.ProposedState != "resolved") {
		report.ProposedState, report.Reason, report.Blocked = ledger.State, "ledger_terminal_conflict", true
		return report, nil
	}
	// Preserve earlier failure and retry history; refine only an empty/generic
	// reason. Resolution never erases the original failure reason or timestamp.
	reason := ledger.Reason
	if report.ProposedState == "failed" && (reason == "" || reason == "archived") {
		reason = report.Reason
	}
	report.WouldUpdate = ledger.State != report.ProposedState || ledger.Reason != reason ||
		(report.ProposedState == "failed" && ledger.FirstFailedAt.IsZero())
	if !apply || !report.WouldUpdate {
		return report, nil
	}
	updates := map[string]any{"state": report.ProposedState, "reason": reason, "updated_at": time.Now().UTC()}
	if report.ProposedState == "failed" && ledger.FirstFailedAt.IsZero() {
		failedAt := info.LastFailedAt
		if failedAt.IsZero() {
			failedAt = time.Now().UTC()
		}
		updates["first_failed_at"] = failedAt
	}
	checkpoint := graphFailureCheckpoint
	if !report.CheckpointComplete {
		checkpoint = "NOT " + checkpoint
	}
	// Revalidate source, checkpoint, and ledger state in the update itself so a
	// concurrent source edit or recovery cannot be overwritten by this snapshot.
	result := db.Table("native_graph_backfill_failures AS f").Where("f.task_id=? AND f.state=? AND f.reason=? AND f.manual_retries=? AND f.updated_at=?",
		report.TaskID, ledger.State, ledger.Reason, ledger.ManualRetries, ledger.UpdatedAt).
		Where("f.tenant_id=? AND f.knowledge_base_id=? AND f.chunk_id=? AND f.knowledge_id=? AND f.content_hash=? AND f.revision=?", tenant, kbID, c.ID, c.KnowledgeID, c.ContentHash, revision).
		Where(graphReconcileCurrentSource).Where(checkpoint).Updates(updates)
	if result.Error != nil {
		return report, errors.New("cannot update exact graph ledger")
	}
	if result.RowsAffected != 1 {
		report.Reason, report.Blocked = "reconciliation_state_changed", true
		return report, nil
	}
	report.LedgerUpdated = true
	return report, nil
}

func graphRecoveryRedisOptions() (asynq.RedisClientOpt, error) {
	redisDB, err := strconv.Atoi(os.Getenv("REDIS_DB"))
	if err != nil || os.Getenv("REDIS_ADDR") == "" {
		return asynq.RedisClientOpt{}, errors.New("runtime Redis configuration is missing")
	}
	return asynq.RedisClientOpt{Addr: os.Getenv("REDIS_ADDR"), Username: os.Getenv("REDIS_USERNAME"), Password: os.Getenv("REDIS_PASSWORD"), DB: redisDB, TLSConfig: common.RedisTLSConfig()}, nil
}

// The checkpoint is queried again after observing a completed task because a
// worker may have committed between candidate selection and Redis inspection.
func graphCandidateCompleted(db *gorm.DB, tenant uint64, kbID string, c graphRecoveryCandidate) (bool, error) {
	var count int64
	err := db.Table("native_graph_records").Where("tenant_id=? AND knowledge_base_id=? AND source_table='knowledgegraphchunkstate' AND payload->>'origin'='native_extraction_v1' AND payload->>'extraction_status'='complete' AND CAST(COALESCE(payload->>'relation_count','0') AS INTEGER)>0 AND payload->>'native_chunk_id'=? AND payload->>'native_knowledge_id'=? AND payload->>'source_content_hash'=?", tenant, kbID, c.ID, c.KnowledgeID, c.ContentHash).Count(&count).Error
	return count > 0, err
}

func graphFailureRetry(db *gorm.DB, queue graphRecoveryQueue, taskID, reason string, slots int) (bool, error) {
	if slots <= 0 {
		return false, nil
	}
	if len(strings.TrimSpace(reason)) < 12 {
		return false, errors.New("manual retry requires a reason describing the investigated failure and fix")
	}
	var failure graphRecoveryFailure
	if err := db.First(&failure, "task_id=?", taskID).Error; err != nil {
		return false, err
	}
	if failure.ManualRetries >= 1 {
		return false, nil
	}
	// Reserve the one allowed manual retry before sending to Redis. A crash
	// here fails closed and leaves a visible record instead of retrying forever.
	failure.ManualRetries++
	failure.State, failure.RetryReason, failure.UpdatedAt = "retry_requested", reason, time.Now().UTC()
	if err := db.Save(&failure).Error; err != nil {
		return false, err
	}
	if err := queue.RunTask(types.QueueGraph, taskID); err != nil {
		return false, err
	}
	return true, nil
}

func graphRecoveryDSN() (string, error) {
	if dsn := os.Getenv("GRAPH_BACKFILL_DSN"); dsn != "" {
		return dsn, nil
	}
	if os.Getenv("DB_HOST") == "" || os.Getenv("DB_NAME") == "" || os.Getenv("DB_USER") == "" {
		return "", errors.New("runtime PostgreSQL configuration is missing")
	}
	port := os.Getenv("DB_PORT")
	if port == "" {
		port = "5432"
	}
	u := url.URL{Scheme: "postgresql", User: url.UserPassword(os.Getenv("DB_USER"), os.Getenv("DB_PASSWORD")), Host: net.JoinHostPort(os.Getenv("DB_HOST"), port), Path: "/" + os.Getenv("DB_NAME")}
	q := u.Query()
	ssl := os.Getenv("DB_SSL_MODE")
	if ssl == "" {
		ssl = "disable"
	}
	q.Set("sslmode", ssl)
	u.RawQuery = q.Encode()
	return u.String(), nil
}

func runGraphRecovery() error {
	tenant := flag.Uint64("tenant", 0, "required tenant")
	kbID := flag.String("kb", "", "required knowledge base UUID")
	limit := flag.Int("limit", 20, "bounded chunks to enqueue (1-200)")
	after := flag.Int64("after", 0, "resume after this chunk sequence ID")
	apply := flag.Bool("apply", false, "apply the selected mode; otherwise report only")
	reconcileOnly := flag.Bool("reconcile-only", false, "reconcile one existing exact-source ledger only; never enqueue or retry")
	retry := flag.Bool("retry-failed", false, "one audited manual retry for tasks still archived in Redis")
	retryReason := flag.String("retry-reason", "", "audited reason for the single permitted manual retry")
	maxOutstanding := flag.Int("max-outstanding", 2, "maximum outstanding tasks in the shared graph queue (1-20)")
	chunkID := flag.String("chunk", "", "only this reviewed chunk; requires knowledge and source-sha256")
	knowledgeID := flag.String("knowledge", "", "reviewed source document UUID")
	sourceHash := flag.String("source-sha256", "", "reviewed exact source SHA-256")
	flag.Parse()
	revision := os.Getenv("GRAPH_BACKFILL_REVISION")
	if flag.NArg() != 0 {
		return errors.New("unexpected positional arguments")
	}
	if err := graphValidateReconciliation(*reconcileOnly, *chunkID, revision, *retry, *retryReason); err != nil {
		return err
	}
	if *tenant == 0 || *limit < 1 || *limit > 200 || *after < 0 || *maxOutstanding < 1 || *maxOutstanding > 20 {
		return errors.New("tenant, limit 1-200 and non-negative cursor are required")
	}
	if *retry && len(strings.TrimSpace(*retryReason)) < 12 {
		return errors.New("retry-failed requires an audited retry-reason describing the investigated failure and fix")
	}
	if _, err := uuid.Parse(*kbID); err != nil {
		return errors.New("invalid knowledge base UUID")
	}
	if err := graphValidateExactScope(*chunkID, *knowledgeID, *sourceHash, *after, *limit, *maxOutstanding); err != nil {
		return err
	}
	var exact []graphRecoveryCandidate
	if *chunkID != "" {
		exact = append(exact, graphRecoveryCandidate{ID: *chunkID, KnowledgeID: *knowledgeID, ContentHash: *sourceHash})
	}
	dsn, err := graphRecoveryDSN()
	if err != nil {
		return err
	}
	db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)})
	if err != nil {
		return errors.New("cannot open graph recovery database")
	}
	sqlDB, err := db.DB()
	if err != nil {
		return err
	}
	defer sqlDB.Close()
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
	defer cancel()
	db = db.WithContext(ctx)
	if *apply {
		if *reconcileOnly && os.Getenv("GRAPH_LEDGER_RECONCILE_VERIFIED") != "1" {
			return errors.New("review the reconciliation preview first; then set GRAPH_LEDGER_RECONCILE_VERIFIED=1")
		}
		if !*reconcileOnly && os.Getenv("GRAPH_NATIVE_WRITE_VERIFIED") != "1" {
			return errors.New("deploy and verify native graph persistence first; then set GRAPH_NATIVE_WRITE_VERIFIED=1")
		}
		// Hold a session lock on a dedicated connection while ledger writes use
		// autocommit. Redis calls must never precede a durable attempt record.
		lockConn, err := sqlDB.Conn(ctx)
		if err != nil {
			return err
		}
		defer lockConn.Close()
		var locked bool
		if err := lockConn.QueryRowContext(ctx, "SELECT pg_try_advisory_lock(20260919, 7351)").Scan(&locked); err != nil {
			return err
		}
		if !locked {
			return errors.New("another graph backfill dispatcher is active")
		}
		defer func() {
			unlockCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
			defer cancel()
			_, _ = lockConn.ExecContext(unlockCtx, "SELECT pg_advisory_unlock(20260919, 7351)")
		}()
	}
	if *reconcileOnly {
		if !db.Migrator().HasTable(&graphRecoveryFailure{}) {
			return errors.New("graph failure ledger is missing; reconciliation will not create it")
		}
		redisOpt, err := graphRecoveryRedisOptions()
		if err != nil {
			return err
		}
		inspector := asynq.NewInspector(redisOpt)
		defer inspector.Close()
		report, err := reconcileGraphTask(db, inspector, *tenant, *kbID, revision, exact[0], *apply)
		if err != nil {
			return err
		}
		return json.NewEncoder(os.Stdout).Encode(report)
	}
	var kb struct {
		Name, SummaryModelID string
		Enabled              bool
	}
	result := db.Raw(`SELECT name,summary_model_id,COALESCE((extract_config->>'enabled')::boolean,false) AS enabled FROM knowledge_bases WHERE tenant_id=? AND id=? AND deleted_at IS NULL`, *tenant, *kbID).Scan(&kb)
	if result.Error != nil {
		return result.Error
	}
	if result.RowsAffected != 1 {
		return errors.New("knowledge base not found in tenant")
	}
	if *apply && (!strings.EqualFold(os.Getenv("LEGACY_GRAPH_ENABLE"), "true") || !kb.Enabled || kb.SummaryModelID == "") {
		return errors.New("native graph backend, knowledge base extraction and summary model must be enabled first")
	}
	hasFailures := db.Migrator().HasTable(&graphRecoveryFailure{})
	if *apply && !hasFailures {
		if err := db.AutoMigrate(&graphRecoveryFailure{}); err != nil {
			return err
		}
		hasFailures = true
	}
	if *apply {
		if _, err := graphUnresolvedFailures(db, *tenant, *kbID, revision, true); err != nil {
			return err
		}
	}
	var coverage struct{ Documents, TextChunks, NativeCompleted, MappedChunks int64 }
	err = db.Raw(`WITH completed AS NOT MATERIALIZED(SELECT payload FROM native_graph_records WHERE tenant_id=? AND knowledge_base_id=? AND source_table='knowledgegraphchunkstate' AND payload->>'origin'='native_extraction_v1' AND payload->>'extraction_status'='complete' AND CAST(COALESCE(payload->>'relation_count','0') AS INTEGER)>0)
		SELECT count(DISTINCT c.knowledge_id) AS documents,count(*) AS text_chunks,
		count(*) FILTER(WHERE EXISTS(SELECT 1 FROM completed r WHERE r.payload->>'native_chunk_id'=c.id AND r.payload->>'native_knowledge_id'=c.knowledge_id AND r.payload->>'source_content_hash'=encode(sha256(convert_to(c.content,'UTF8')),'hex'))) AS native_completed,
		count(*) FILTER(WHERE EXISTS(SELECT 1 FROM native_graph_sources s WHERE s.tenant_id=c.tenant_id AND s.knowledge_base_id=c.knowledge_base_id AND s.native_chunk_id=c.id AND s.native_content_sha256=encode(sha256(convert_to(c.content,'UTF8')),'hex'))) AS mapped_chunks
		FROM chunks c JOIN knowledges k ON k.id=c.knowledge_id AND k.tenant_id=c.tenant_id AND k.knowledge_base_id=c.knowledge_base_id AND k.deleted_at IS NULL
		WHERE c.tenant_id=? AND c.knowledge_base_id=? AND c.deleted_at IS NULL AND c.is_enabled AND c.chunk_type='text'`, *tenant, *kbID, *tenant, *kbID).Scan(&coverage).Error
	if err != nil {
		return err
	}
	candidates, err := graphRecoveryCandidates(db, *tenant, *kbID, revision, *after, *limit, hasFailures, *retry, exact...)
	if err != nil {
		return err
	}
	report := map[string]any{"tenant_id": *tenant, "knowledge_base_id": *kbID, "knowledge_base_name": kb.Name, "graph_enabled": kb.Enabled, "coverage": coverage, "candidate_count": len(candidates), "after": *after, "next_cursor": *after, "apply": *apply, "enqueued": 0, "already_queued": 0, "retried": 0, "archived": 0, "candidates": candidates}
	if !*apply {
		if hasFailures {
			failures, err := graphUnresolvedFailures(db, *tenant, *kbID, revision, false)
			if err != nil {
				return err
			}
			report["failures"], report["failure_count"] = failures, len(failures)
		}
		return json.NewEncoder(os.Stdout).Encode(report)
	}
	redisOpt, err := graphRecoveryRedisOptions()
	if err != nil {
		return err
	}
	client := asynq.NewClient(redisOpt)
	defer client.Close()
	inspector := asynq.NewInspector(redisOpt)
	defer inspector.Close()
	enqueued, existing, retried, archived, deferred := 0, 0, 0, 0, 0
	queueInfo, err := inspector.GetQueueInfo(types.QueueGraph)
	if err != nil && !errors.Is(err, asynq.ErrQueueNotFound) {
		return err
	}
	if queueInfo == nil {
		queueInfo = &asynq.QueueInfo{}
	}
	slots := graphQueueSlots(queueInfo, *maxOutstanding)
	settled := make(map[string]bool)
	for _, c := range candidates {
		if len(exact) > 0 {
			// Recheck immediately before dispatch. The worker checks the same
			// immutable identity before calling the model if a source changes later.
			if _, err := graphRecoveryCandidates(db, *tenant, *kbID, revision, 0, 1, hasFailures, *retry, exact...); err != nil {
				return err
			}
		}
		taskID := graphRecoveryTaskID(*tenant, *kbID, c, revision)
		info, inspectErr := inspector.GetTaskInfo(types.QueueGraph, taskID)
		if inspectErr != nil && !errors.Is(inspectErr, asynq.ErrTaskNotFound) && !errors.Is(inspectErr, asynq.ErrQueueNotFound) {
			return inspectErr
		}
		if info != nil {
			if len(exact) > 0 && (!graphTaskMatchesReviewedSource(info.Payload, *tenant, *kbID, c) || info.MaxRetry != 0) {
				return errors.New("existing graph task does not have the reviewed source and single-execution guard; inspect it without resetting the failure ledger")
			}
			existing++
			if err := trackGraphTask(db, *tenant, *kbID, revision, taskID, c); err != nil {
				return err
			}
			if info.State == asynq.TaskStateArchived || info.State == asynq.TaskStateCompleted {
				verified, err := graphCandidateCompleted(db, *tenant, *kbID, c)
				if err != nil {
					return err
				}
				if verified {
					settled[c.ID] = true
					continue
				}
				reason := "completed_without_native_checkpoint"
				if info.State == asynq.TaskStateArchived {
					reason = graphArchivedReason(info.LastErr)
					archived++
				}
				if err := recordGraphFailure(db, *tenant, *kbID, revision, taskID, c, reason); err != nil {
					return err
				}
				settled[c.ID] = true
				if *retry && info.State == asynq.TaskStateArchived {
					didRetry, err := graphFailureRetry(db, inspector, taskID, *retryReason, slots)
					if err != nil {
						return err
					}
					if didRetry {
						retried++
						slots--
						settled[c.ID] = false
					}
				}
			}
			continue
		}
		if slots == 0 {
			deferred++
			continue
		}
		reserved, err := reserveGraphDispatch(db, *tenant, *kbID, revision, taskID, c)
		if err != nil {
			return err
		}
		if !reserved {
			settled[c.ID] = true
			continue
		}
		// KnowledgeID is intentionally empty: recovery has no ingestion slot
		// to decrement and must not finalize a concurrent document upload.
		payload, err := json.Marshal(types.ExtractChunkPayload{TenantID: *tenant, ChunkID: c.ID, ModelID: kb.SummaryModelID, ChunkIndex: c.ChunkIndex})
		if len(exact) > 0 {
			payload, err = json.Marshal(types.ExtractChunkPayload{TenantID: *tenant, ChunkID: c.ID, ModelID: kb.SummaryModelID, ChunkIndex: c.ChunkIndex,
				ExpectedSource: &types.GraphSourceIdentity{KnowledgeBaseID: *kbID, KnowledgeID: c.KnowledgeID, ContentSHA256: c.ContentHash}})
		}
		if err != nil {
			return err
		}
		maxRetry, timeout := 3, 10*time.Minute
		if len(exact) > 0 {
			maxRetry, timeout = 0, 3*time.Minute
		}
		_, err = client.EnqueueContext(ctx, asynq.NewTask(types.TypeChunkExtract, payload), asynq.Queue(types.QueueGraph), asynq.TaskID(taskID), asynq.MaxRetry(maxRetry), asynq.Timeout(timeout), asynq.Retention(24*time.Hour))
		if errors.Is(err, asynq.ErrTaskIDConflict) {
			existing++
			// Another producer raced with this dispatcher. Count the slot
			// conservatively; the next poll will inspect its actual state.
			slots--
			continue
		}
		if err != nil {
			return err
		}
		enqueued++
		slots--
	}
	report["enqueued"], report["already_queued"], report["retried"] = enqueued, existing, retried
	// Do not let a manual resume skip deferred or still-running work. Watch
	// always starts at zero and additionally uses persisted failure exclusions.
	report["next_cursor"] = graphResumeCursor(*after, candidates, settled)
	report["archived"] = archived
	report["deferred_for_queue_limit"], report["max_outstanding"] = deferred, *maxOutstanding
	report["queue_paused"] = queueInfo.Paused
	report["queue_limit_scope"] = "serialized_backfill_dispatchers; other graph producers may enqueue concurrently"
	failures, err := graphUnresolvedFailures(db, *tenant, *kbID, revision, true)
	if err != nil {
		return err
	}
	report["failures"], report["failure_count"] = failures, len(failures)
	return json.NewEncoder(os.Stdout).Encode(report)
}

func graphBatchStatus(report map[string]any) (complete, blocked bool) {
	coverage, _ := report["coverage"].(map[string]any)
	textChunks, totalOK := coverage["TextChunks"].(float64)
	completedChunks, completedOK := coverage["NativeCompleted"].(float64)
	candidates, candidatesOK := report["candidate_count"].(float64)
	archived, _ := report["archived"].(float64)
	failures, _ := report["failure_count"].(float64)
	paused, _ := report["queue_paused"].(bool)
	valid := totalOK && completedOK && candidatesOK && textChunks >= 0 && completedChunks >= 0 && completedChunks <= textChunks && candidates >= 0 && archived >= 0 && failures >= 0
	complete = valid && completedChunks == textChunks && candidates == 0 && failures == 0 && archived == 0
	blocked = archived > 0 || failures > 0 || !valid || (!complete && (candidates == 0 || paused))
	return
}

// A restart resumes from verified database checkpoints. Each knowledge base has
// only one bounded window in flight; failed windows require explicit recovery.
func watchGraphRecovery() error {
	flags := flag.NewFlagSet("watch", flag.ContinueOnError)
	tenant := flags.Uint64("tenant", 0, "required tenant")
	kbs := flags.String("kbs", "", "comma-separated knowledge base UUIDs")
	limit := flags.Int("limit", 2, "maximum candidate window per knowledge base (1-20)")
	maxOutstanding := flags.Int("max-outstanding", 2, "maximum outstanding tasks across the shared graph queue (1-20)")
	interval := flags.Duration("interval", time.Minute, "poll interval (at least 30s)")
	stateFile := flags.String("state", "/data/progress.json", "durable progress report")
	if err := flags.Parse(os.Args[2:]); err != nil {
		return err
	}
	if *tenant == 0 || *limit < 1 || *limit > 20 || *interval < 30*time.Second || *kbs == "" || *maxOutstanding < 1 || *maxOutstanding > 20 {
		return errors.New("watch requires tenant, knowledge bases, limit 1-20 and interval >=30s")
	}
	ids := strings.Split(*kbs, ",")
	if len(ids) > 10 {
		return errors.New("watch supports at most 10 knowledge bases")
	}
	for _, id := range ids {
		if _, err := uuid.Parse(id); err != nil {
			return errors.New("invalid knowledge base in watch scope")
		}
	}
	executable, err := os.Executable()
	if err != nil {
		return err
	}
	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
	defer stop()
	for {
		reports := []map[string]any{}
		complete, blocked := true, false
		for _, id := range ids {
			commandCtx, cancel := context.WithTimeout(ctx, 6*time.Minute)
			command := exec.CommandContext(commandCtx, executable, "-tenant", strconv.FormatUint(*tenant, 10), "-kb", id, "-limit", strconv.Itoa(*limit), "-max-outstanding", strconv.Itoa(*maxOutstanding), "-apply")
			output, commandErr := command.Output()
			cancel()
			var report map[string]any
			if commandErr != nil || json.Unmarshal(output, &report) != nil {
				report = map[string]any{"knowledge_base_id": id, "error": "batch command failed; check runtime and model configuration"}
				blocked, complete = true, false
			} else {
				delete(report, "candidates")
				batchComplete, batchBlocked := graphBatchStatus(report)
				complete = complete && batchComplete
				blocked = blocked || batchBlocked
				if batchBlocked {
					report["attention"] = "Native coverage is incomplete or extraction failures need recovery"
				}
			}
			reports = append(reports, report)
			if ctx.Err() != nil {
				return nil
			}
		}
		status := "running"
		if complete {
			status = "complete"
		} else if blocked {
			status = "needs_attention"
		}
		state := map[string]any{"status": status, "checked_at": time.Now().UTC(), "tenant_id": *tenant, "batch_limit_per_kb": *limit, "max_outstanding": *maxOutstanding, "knowledge_bases": reports}
		data, err := json.MarshalIndent(state, "", "  ")
		if err != nil {
			return err
		}
		if err = os.MkdirAll(filepath.Dir(*stateFile), 0700); err != nil {
			return err
		}
		if err = os.WriteFile(*stateFile+".tmp", data, 0600); err != nil {
			return err
		}
		if err = os.Rename(*stateFile+".tmp", *stateFile); err != nil {
			return err
		}
		fmt.Printf("{\"status\":%q,\"knowledge_bases\":%d}\n", status, len(reports))
		if complete {
			// Keep the successful checkpoint visible without restart-policy loops.
			<-ctx.Done()
			return nil
		}
		select {
		case <-ctx.Done():
			return nil
		case <-time.After(*interval):
		}
	}
}

func main() {
	var err error
	if len(os.Args) > 1 && os.Args[1] == "--watch" {
		err = watchGraphRecovery()
	} else {
		err = runGraphRecovery()
	}
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}