xinmi/internal/application/service/knowledge_version_cleanup.go

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
}