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 }