// 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) } }