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

internal/application/service/knowledge_version_cleanup.go

@@ -0,0 +1,365 @@
+package service
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"os"
+	"strconv"
+	"strings"
+	"time"
+
+	"github.com/Tencent/WeKnora/internal/application/service/retriever"
+	"github.com/Tencent/WeKnora/internal/knowledgeversion"
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/Tencent/WeKnora/internal/types/interfaces"
+	"gorm.io/gorm"
+)
+
+const knowledgeVersionCleanupKey = "_version_copy_cleanup"
+
+type knowledgeVersionObject struct {
+	Reference string `json:"reference"`
+	Physical  string `json:"physical"`
+	Done      bool   `json:"done,omitempty"`
+}
+
+type knowledgeVersionInventory struct {
+	OwnerID string                   `json:"owner_id"`
+	Objects []knowledgeVersionObject `json:"objects"`
+}
+
+func readKnowledgeVersionInventory(metadata types.JSON, owner string) (map[string]json.RawMessage, knowledgeVersionInventory, error) {
+	fields := map[string]json.RawMessage{}
+	if len(metadata) != 0 {
+		if err := json.Unmarshal(metadata, &fields); err != nil || fields == nil {
+			return nil, knowledgeVersionInventory{}, fmt.Errorf("invalid version cleanup metadata")
+		}
+	}
+	inventory := knowledgeVersionInventory{OwnerID: owner}
+	if data, ok := fields[knowledgeVersionCleanupKey]; ok {
+		if err := json.Unmarshal(data, &inventory); err != nil || inventory.OwnerID != owner {
+			return nil, inventory, fmt.Errorf("version cleanup inventory owner mismatch")
+		}
+	}
+	return fields, inventory, nil
+}
+
+func (s *knowledgeService) writeKnowledgeVersionInventory(ctx context.Context, doc *types.Knowledge, inventory knowledgeVersionInventory) error {
+	fields, _, err := readKnowledgeVersionInventory(doc.Metadata, doc.ID)
+	if err != nil {
+		return err
+	}
+	fields[knowledgeVersionCleanupKey], err = json.Marshal(inventory)
+	if err != nil {
+		return err
+	}
+	metadata, err := json.Marshal(fields)
+	if err != nil {
+		return err
+	}
+	if err := s.repo.UpdateKnowledgeColumns(ctx, doc.ID, map[string]interface{}{"metadata": types.JSON(metadata)}); err != nil {
+		return err
+	}
+	doc.Metadata = metadata
+	return nil
+}
+
+// versionTrackingFileService journals images immediately, before chunks are
+// created. A failed partial document therefore retains its object inventory.
+type versionTrackingFileService struct {
+	interfaces.FileService
+	service   *knowledgeService
+	knowledge *types.Knowledge
+}
+
+func (f *versionTrackingFileService) SaveBytes(ctx context.Context, data []byte, tenantID uint64, name string, temporary bool) (string, error) {
+	if !knowledgeversion.IsWorker(ctx) || f.knowledge == nil || f.knowledge.TenantID != tenantID {
+		return "", fmt.Errorf("version image tracking requires the owning worker workspace")
+	}
+	ref, err := f.FileService.SaveBytes(ctx, data, tenantID, name, temporary)
+	if err != nil {
+		return "", err
+	}
+	physical := ref
+	if _, catalogRef := types.ParseResourcePath(ref); catalogRef {
+		if f.service.resourceCatalog == nil {
+			err = fmt.Errorf("version object has no resource catalog")
+		} else {
+			var resource *types.StoredResource
+			physical, resource, err = f.service.resourceCatalog.ResolvePath(ctx, ref)
+			if err == nil && (resource == nil || resource.TenantID != tenantID) {
+				err = fmt.Errorf("version object belongs to another workspace")
+			}
+		}
+	}
+	if err == nil {
+		var latest *types.Knowledge
+		latest, err = f.service.repo.GetKnowledgeByID(ctx, tenantID, f.knowledge.ID)
+		if err == nil {
+			var inventory knowledgeVersionInventory
+			_, inventory, err = readKnowledgeVersionInventory(latest.Metadata, latest.ID)
+			if err == nil {
+				inventory.Objects = append(inventory.Objects, knowledgeVersionObject{Reference: ref, Physical: physical})
+				err = f.service.writeKnowledgeVersionInventory(ctx, latest, inventory)
+			}
+		}
+	}
+	if err != nil {
+		cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 20*time.Second)
+		defer cancel()
+		cleanupErr := f.FileService.DeleteFile(cleanupCtx, ref)
+		return "", errors.Join(fmt.Errorf("record copied object: %w", err), cleanupErr)
+	}
+	return ref, nil
+}
+
+// preserveKnowledgeVersionInventory must precede the final document update:
+// source metadata must never replace the worker's ownership journal.
+func (s *knowledgeService) preserveKnowledgeVersionInventory(ctx context.Context, doc *types.Knowledge) error {
+	latest, err := s.repo.GetKnowledgeByID(ctx, doc.TenantID, doc.ID)
+	if err != nil {
+		return err
+	}
+	_, inventory, err := readKnowledgeVersionInventory(latest.Metadata, doc.ID)
+	if err != nil {
+		return err
+	}
+	var fields map[string]json.RawMessage
+	if len(doc.Metadata) > 0 {
+		if err := json.Unmarshal(doc.Metadata, &fields); err != nil {
+			return err
+		}
+	}
+	if fields == nil {
+		fields = map[string]json.RawMessage{}
+	}
+	delete(fields, knowledgeVersionCleanupKey)
+	if len(inventory.Objects) > 0 {
+		fields[knowledgeVersionCleanupKey], err = json.Marshal(inventory)
+		if err != nil {
+			return err
+		}
+	}
+	doc.Metadata, err = json.Marshal(fields)
+	return err
+}
+
+// CleanupKnowledgeVersion removes only abandoned worker staging content. The
+// durable caller owns the lease, KB-shell deletion, and storage reconciliation
+// in its final transaction; this method never adjusts tenant storage counters.
+func (s *knowledgeService) CleanupKnowledgeVersion(ctx context.Context, target string) error {
+	tenantID, ok := types.TenantIDFromContext(ctx)
+	if !knowledgeversion.IsWorker(ctx) || !ok || tenantID == 0 || strings.TrimSpace(target) == "" {
+		return fmt.Errorf("knowledge version cleanup requires an isolated worker target")
+	}
+	kb, err := s.kbService.GetKnowledgeBaseByID(ctx, target)
+	if err != nil {
+		return err
+	}
+	if kb == nil || kb.TenantID != tenantID || kb.ID != target || kb.VersionState != "copying" || !kb.IsTemporary {
+		return fmt.Errorf("knowledge version cleanup requires a temporary copying base in the same workspace")
+	}
+	if kb.Type == types.KnowledgeBaseTypeWiki || kb.ExtractConfig != nil && kb.ExtractConfig.Enabled {
+		return fmt.Errorf("graph and Wiki version cleanup is not supported")
+	}
+	tenant, err := s.tenantRepo.GetTenantByID(ctx, tenantID)
+	if err != nil {
+		return err
+	}
+	if tenant == nil || tenant.ID != tenantID {
+		return fmt.Errorf("version cleanup workspace not found")
+	}
+	ctx = context.WithValue(ctx, types.TenantInfoContextKey, tenant)
+	docs, err := s.repo.ListKnowledgeByKnowledgeBaseID(ctx, tenantID, target)
+	if err != nil {
+		return err
+	}
+	for _, doc := range docs {
+		if doc.TenantID != tenantID || doc.KnowledgeBaseID != target {
+			return fmt.Errorf("version cleanup document scope mismatch")
+		}
+		if err := s.cleanupKnowledgeVersionDocument(ctx, kb, doc); err != nil {
+			return fmt.Errorf("cleanup knowledge version %s document %s: %w", target, doc.ID, err)
+		}
+	}
+	if _, err := s.tagRepo.DeleteUnusedTags(ctx, tenantID, target); err != nil {
+		return err
+	}
+	count, err := s.chunkRepo.CountChunksByKnowledgeBaseID(ctx, tenantID, target)
+	if err != nil {
+		return err
+	}
+	if count != 0 {
+		return fmt.Errorf("version cleanup left %d unowned chunks", count)
+	}
+	_, tagCount, err := s.tagRepo.ListByKB(ctx, tenantID, target, &types.Pagination{Page: 1, PageSize: 1}, "")
+	if err != nil {
+		return err
+	}
+	if tagCount != 0 {
+		return fmt.Errorf("version cleanup left %d referenced tags", tagCount)
+	}
+	return nil
+}
+
+func versionOwnsRawPath(value string, tenantID uint64, knowledgeID string) bool {
+	if _, inner, ok := types.ParseStorageBackendPath(value); ok {
+		value = inner
+	}
+	provider, rest, ok := strings.Cut(value, "://")
+	if !ok || provider == "" || strings.Contains(rest, "..") || strings.ContainsAny(rest, "\\?#") {
+		return false
+	}
+	parts := strings.Split(rest, "/")
+	if provider != "local" {
+		if len(parts) < 4 {
+			return false
+		}
+		parts = parts[1:]
+	}
+	return len(parts) == 3 && parts[0] == strconv.FormatUint(tenantID, 10) && parts[1] == knowledgeID && parts[2] != ""
+}
+
+func (s *knowledgeService) versionCleanupResource(ctx context.Context, reference string) (*types.StoredResource, error) {
+	if s.resourceCatalog == nil {
+		return nil, fmt.Errorf("version cleanup resource catalog is unavailable")
+	}
+	// Resolve cannot distinguish a removed handle from a database failure.
+	// The native repository provides that distinction without parsing errors.
+	if catalog, ok := s.resourceCatalog.(*resourceCatalog); ok {
+		handle, valid := types.ParseResourcePath(reference)
+		if !valid {
+			return nil, fmt.Errorf("invalid version cleanup resource")
+		}
+		return catalog.repo.GetByHandle(ctx, handle)
+	}
+	resource, err := s.resourceCatalog.Resolve(ctx, reference)
+	if errors.Is(err, os.ErrNotExist) || errors.Is(err, gorm.ErrRecordNotFound) {
+		return nil, nil
+	}
+	return resource, err
+}
+
+func (s *knowledgeService) cleanupKnowledgeVersionDocument(ctx context.Context, kb *types.KnowledgeBase, doc *types.Knowledge) error {
+	_, inventory, err := readKnowledgeVersionInventory(doc.Metadata, doc.ID)
+	if err != nil {
+		return err
+	}
+	owned := map[string]bool{}
+	for _, object := range inventory.Objects {
+		if object.Reference == "" || object.Physical == "" && !object.Done {
+			return fmt.Errorf("incomplete version object inventory")
+		}
+		owned[object.Reference] = true
+	}
+	if doc.FilePath != "" && !owned[doc.FilePath] {
+		physical := doc.FilePath
+		if _, ok := types.ParseResourcePath(doc.FilePath); ok {
+			resource, err := s.versionCleanupResource(ctx, doc.FilePath)
+			if err != nil {
+				return err
+			}
+			if resource == nil {
+				physical = ""
+			} else if resource.TenantID != doc.TenantID {
+				return fmt.Errorf("raw resource belongs to another workspace")
+			} else {
+				physical = resource.PhysicalPath
+			}
+		}
+		if physical != "" && !versionOwnsRawPath(physical, doc.TenantID, doc.ID) {
+			return fmt.Errorf("refusing to delete a raw file outside the staging document namespace")
+		}
+		inventory.Objects = append(inventory.Objects, knowledgeVersionObject{Reference: doc.FilePath, Physical: physical, Done: physical == ""})
+		owned[doc.FilePath] = true
+	}
+	images, err := s.chunkRepo.ListImageInfoByKnowledgeIDs(ctx, doc.TenantID, []string{doc.ID})
+	if err != nil {
+		return err
+	}
+	for _, chunk := range images {
+		var refs []*types.ImageInfo
+		if err := json.Unmarshal([]byte(chunk.ImageInfo), &refs); err != nil {
+			return fmt.Errorf("invalid staged image inventory: %w", err)
+		}
+		for _, image := range refs {
+			if image != nil && managedVersionReference(image.URL) && !owned[image.URL] {
+				return fmt.Errorf("staged image has no durable ownership record")
+			}
+		}
+	}
+	if err := s.writeKnowledgeVersionInventory(ctx, doc, inventory); err != nil {
+		return err
+	}
+	if strings.TrimSpace(doc.EmbeddingModelID) != "" {
+		engine, err := retriever.CreateRetrieveEngineForKB(ctx, s.retrieveEngine, s.ownership, doc.TenantID, kb.VectorStoreID)
+		if err != nil {
+			return err
+		}
+		model, err := s.modelService.GetEmbeddingModel(ctx, doc.EmbeddingModelID)
+		if err != nil {
+			return err
+		}
+		if err := engine.DeleteByKnowledgeIDList(ctx, []string{doc.ID}, model.GetDimensions(), doc.Type); err != nil {
+			return err
+		}
+	}
+	for i := range inventory.Objects {
+		object := &inventory.Objects[i]
+		if object.Done {
+			continue
+		}
+		if err := s.cleanupKnowledgeVersionObject(ctx, kb, doc, *object); err != nil {
+			return err
+		}
+		object.Done = true
+		if err := s.writeKnowledgeVersionInventory(ctx, doc, inventory); err != nil {
+			return err
+		}
+	}
+	if err := s.chunkRepo.DeleteChunksByKnowledgeID(ctx, doc.TenantID, doc.ID); err != nil {
+		return err
+	}
+	if err := s.repo.DeleteKnowledgeTagRelations(ctx, doc.ID); err != nil {
+		return err
+	}
+	// Keep sizes durable until the caller deletes rows and adjusts storage in
+	// one transaction. Retrying external cleanup is deliberately idempotent.
+	return nil
+}
+
+func (s *knowledgeService) cleanupKnowledgeVersionObject(ctx context.Context, kb *types.KnowledgeBase, doc *types.Knowledge, object knowledgeVersionObject) error {
+	_, catalogRef := types.ParseResourcePath(object.Reference)
+	if catalogRef {
+		resource, err := s.versionCleanupResource(ctx, object.Reference)
+		if err != nil {
+			return err
+		}
+		if resource == nil {
+			return nil
+		}
+		if resource.TenantID != doc.TenantID || resource.PhysicalPath != object.Physical {
+			return fmt.Errorf("version object ownership changed")
+		}
+		remaining, err := s.resourceCatalog.Release(ctx, object.Reference, types.ResourceOwnerKnowledge, doc.ID)
+		if err != nil {
+			return err
+		}
+		if remaining < 0 {
+			return fmt.Errorf("version object ownership could not be checked")
+		}
+		if remaining > 0 {
+			return nil
+		}
+	}
+	files := s.resolveFileServiceForPath(ctx, kb, object.Physical)
+	if err := files.DeleteFile(ctx, object.Physical); err != nil && !errors.Is(err, os.ErrNotExist) {
+		return fmt.Errorf("delete copied file: %w", err)
+	}
+	if catalogRef {
+		return s.resourceCatalog.MarkDeleted(ctx, object.Reference)
+	}
+	return nil
+}