Agent-System/internal/dataset/types.go

package dataset

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"sort"
	"strings"
	"time"

	"chat-agent-system/internal/artifactvault"
)

const (
	ClassPinned             = "pinned"
	ClassDemo               = "demo"
	ClassRegression         = "regression"
	ClassMigrationBaseline  = "migration_baseline"
	ClassEphemeral          = "ephemeral"
	ClassQuarantine         = "quarantine"

	JobRunning   = "running"
	JobPaused    = "paused"
	JobCompleted = "completed"
	JobFailed    = "failed"

	OldestSupportedSchema = artifactvault.OldestSupportedSchema
	MinQuarantineWindow   = 7 * 24 * time.Hour
)

type Dataset struct {
	ID               string `json:"dataset_id"`
	Class            string `json:"class"`
	DisplayName      string `json:"display_name"`
	StorageIdentity  string `json:"storage_identity"`
	SourceSchema     int    `json:"source_schema"`
	SourceCommit     string `json:"source_commit"`
	Status           string `json:"status"`
	ProjectCount     int    `json:"project_count"`
}

func ValidateClass(class string) error {
	switch class {
	case ClassPinned, ClassDemo, ClassRegression, ClassMigrationBaseline, ClassEphemeral, ClassQuarantine:
		return nil
	default:
		return fmt.Errorf("dataset: invalid class %q", class)
	}
}

type Checkpoint struct {
	JobID          string
	Shard          int
	SourceKind     string
	LastPK         int64
	Scanned        int
	Succeeded      int
	Quarantined    int
	RetryWatermark int
}

type JobStats struct {
	Scanned     int
	Succeeded   int
	Quarantined int
}

func Recalc(cps []Checkpoint) JobStats {
	var s JobStats
	for _, c := range cps {
		s.Scanned += c.Scanned
		s.Succeeded += c.Succeeded
		s.Quarantined += c.Quarantined
	}
	return s
}

type DeleteManifest struct {
	DatasetIDs       []string `json:"dataset_ids"`
	StorageIdentity  string   `json:"storage_identity"`
	Reason           string   `json:"reason"`
}

func CanonicalDestroyBytes(m DeleteManifest) []byte {
	ids := append([]string(nil), m.DatasetIDs...)
	sort.Strings(ids)
	raw, _ := json.Marshal(struct {
		DatasetIDs      []string `json:"dataset_ids"`
		StorageIdentity string   `json:"storage_identity"`
		Reason          string   `json:"reason"`
	}{ids, m.StorageIdentity, m.Reason})
	return raw
}

func DestroyManifestSHA(m DeleteManifest) string {
	sum := sha256.Sum256(CanonicalDestroyBytes(m))
	return hex.EncodeToString(sum[:])
}

type DestroyRequest struct {
	TargetIdentity        string
	TargetDSN             string
	Manifest              DeleteManifest
	ExpectedManifestSHA   string
	RecoverableLocation   string
	CreatedAt             time.Time
	QuarantineUntil       time.Time
	Approval1             string
	Approval2             string
	Coverage              artifactvault.Coverage
	BackupOK              bool
	CorpusOK              bool
	VaultReady            bool
	Now                   time.Time
}

func RefuseOperationalDSN(dsn, identity string) error {
	low := strings.ToLower(dsn + " " + identity)
	for _, ban := range []string{"agent_system_dev", "agent_system_test", "agent_system_main"} {
		if strings.Contains(low, ban) && !strings.Contains(low, "scratch") {
			return fmt.Errorf("dataset: refuse operational database %s", ban)
		}
	}
	if strings.Contains(low, ":5433/") || strings.Contains(low, ":5435/") {
		if !strings.Contains(low, "scratch") {
			return fmt.Errorf("dataset: refuse shared env port")
		}
	}
	if strings.Contains(low, ":5434/") && strings.Contains(low, "agent_system_test") {
		return fmt.Errorf("dataset: refuse shared test database")
	}
	return nil
}

func EvaluateDestroy(req DestroyRequest) error {
	if err := RefuseOperationalDSN(req.TargetDSN, req.TargetIdentity); err != nil {
		return err
	}
	if !strings.Contains(strings.ToLower(req.RecoverableLocation), "scratch") {
		return fmt.Errorf("dataset: recoverable location must be scratch")
	}
	if req.QuarantineUntil.Sub(req.CreatedAt) < MinQuarantineWindow {
		return fmt.Errorf("dataset: quarantine observation window must be at least 7 days")
	}
	now := req.Now
	if now.IsZero() {
		now = time.Now().UTC()
	}
	if now.Before(req.QuarantineUntil) {
		return fmt.Errorf("dataset: quarantine window still open")
	}
	if strings.TrimSpace(req.Approval1) == "" || strings.TrimSpace(req.Approval2) == "" || req.Approval1 == req.Approval2 {
		return fmt.Errorf("dataset: two distinct user approvals required")
	}
	got := DestroyManifestSHA(req.Manifest)
	if req.ExpectedManifestSHA != "" && req.ExpectedManifestSHA != got {
		return fmt.Errorf("dataset: destroy manifest sha mismatch")
	}
	if len(req.Manifest.DatasetIDs) == 0 || strings.TrimSpace(req.Manifest.StorageIdentity) == "" {
		return fmt.Errorf("dataset: destroy manifest incomplete")
	}
	return artifactvault.PreflightDestroy(req.Coverage, req.BackupOK, req.CorpusOK, req.VaultReady)
}

func RequireScratchRestore(dsn, storageIdentity, manifestSHA, expectedSHA string, hasAdmin bool) error {
	if err := RefuseOperationalDSN(dsn, storageIdentity); err != nil {
		return err
	}
	if !strings.Contains(strings.ToLower(dsn+" "+storageIdentity), "scratch") {
		return fmt.Errorf("dataset: restore target must be scratch")
	}
	if !hasAdmin {
		return fmt.Errorf("dataset: restore requires role-admin")
	}
	if manifestSHA == "" || expectedSHA == "" || manifestSHA != expectedSHA {
		return fmt.Errorf("dataset: restore manifest mismatch or missing")
	}
	return nil
}