Agent-System/internal/dataset/runner.go

package dataset

import (
	"context"
	"fmt"
	"sync"

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

type Enumerator interface {
	Next(ctx context.Context, kind string, lastPK int64, limit int) ([]artifactvault.Candidate, error)
}

type MemoryEnumerator struct {
	mu    sync.Mutex
	items map[string][]artifactvault.Candidate
}

func NewMemoryEnumerator() *MemoryEnumerator {
	return &MemoryEnumerator{items: map[string][]artifactvault.Candidate{}}
}

func (m *MemoryEnumerator) Seed(kind string, items []artifactvault.Candidate) {
	m.mu.Lock()
	defer m.mu.Unlock()
	cp := append([]artifactvault.Candidate(nil), items...)
	m.items[kind] = cp
}

func (m *MemoryEnumerator) Next(_ context.Context, kind string, lastPK int64, limit int) ([]artifactvault.Candidate, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	all := m.items[kind]
	out := make([]artifactvault.Candidate, 0, limit)
	for _, c := range all {
		if c.SourcePK <= lastPK {
			continue
		}
		out = append(out, c)
		if limit > 0 && len(out) >= limit {
			break
		}
	}
	return out, nil
}

type Job struct {
	ID               string
	Status           string
	SourceSchema     int
	SourceCommit     string
	StorageIdentity  string
	Checkpoints      map[string]Checkpoint
}

func NewJob(id, commit, identity string, schema int) *Job {
	return &Job{
		ID: id, Status: JobRunning, SourceSchema: schema, SourceCommit: commit,
		StorageIdentity: identity, Checkpoints: map[string]Checkpoint{},
	}
}

func (j *Job) Pause() {
	if j.Status == JobRunning {
		j.Status = JobPaused
	}
}

func (j *Job) Resume() error {
	if j.Status != JobPaused && j.Status != JobRunning {
		return fmt.Errorf("dataset: job %s cannot resume from %s", j.ID, j.Status)
	}
	j.Status = JobRunning
	return nil
}

type Runner struct {
	Enum Enumerator
	Env  string
	Tx   artifactvaultTx
}

type artifactvaultTx interface {
	Process(ctx context.Context, env string, c artifactvault.Candidate) (artifactvault.ProcessResult, error)
}

type memoryProcessor struct {
	reg *artifactvault.MemoryRegistry
}

func NewMemoryRunner(enum Enumerator, reg *artifactvault.MemoryRegistry) *Runner {
	if reg == nil {
		reg = artifactvault.NewMemoryRegistry()
	}
	return &Runner{Enum: enum, Env: "dev", Tx: &memoryProcessor{reg: reg}}
}

func (p *memoryProcessor) Process(ctx context.Context, env string, c artifactvault.Candidate) (artifactvault.ProcessResult, error) {
	if hz := artifactvault.ClassifyHazard(c); hz != artifactvault.HazardNone {
		_ = p.reg.Quarantine(ctx, hz, c.LogicalID)
		return artifactvault.ProcessResult{Status: "quarantine", Hazard: hz}, nil
	}
	if c.Kind == artifactvault.KindSnapshot {
		return artifactvault.ProcessResult{Status: "done"}, nil
	}
	ing := &artifactvault.Ingest{Env: env, Reg: p.reg}
	var err error
	switch c.Kind {
	case artifactvault.KindAgentCommit:
		err = ing.AgentCommit(ctx, c.ProjectID, c.LogicalID, nz(c.EventID, fmt.Sprintf("commit:%d", c.SourcePK)), c.SnapshotID, c.Payload)
	case artifactvault.KindApproved:
		err = ing.Approved(ctx, c.ProjectID, c.LogicalID, c.SnapshotID, c.ApprovedAt, c.Payload)
	case artifactvault.KindChatTurn:
		policy := c.ExportPolicy
		if policy == "" {
			policy = "full"
		}
		err = ing.AssistantReply(ctx, c.ProjectID, c.LogicalID, nz(c.EventID, fmt.Sprintf("turn:%d", c.SourcePK)), c.MessageID, c.RunID, policy, c.CompletedAt, c.Payload)
	default:
		return artifactvault.ProcessResult{}, fmt.Errorf("dataset: kind %s", c.Kind)
	}
	if err != nil {
		return artifactvault.ProcessResult{}, err
	}
	return artifactvault.ProcessResult{Status: "done"}, nil
}

func nz(v, def string) string {
	if v == "" {
		return def
	}
	return v
}

func (r *Runner) RunBatch(ctx context.Context, job *Job, kind string, shard, limit int) (JobStats, error) {
	if job.Status == JobPaused {
		return Recalc(checkpointsOf(job)), fmt.Errorf("dataset: job paused")
	}
	if job.Status != JobRunning {
		return JobStats{}, fmt.Errorf("dataset: job not running")
	}
	key := fmt.Sprintf("%s:%d", kind, shard)
	cp := job.Checkpoints[key]
	cp.JobID, cp.Shard, cp.SourceKind = job.ID, shard, kind
	items, err := r.Enum.Next(ctx, kind, cp.LastPK, limit)
	if err != nil {
		job.Status = JobFailed
		return Recalc(checkpointsOf(job)), err
	}
	env := r.Env
	if env == "" {
		env = "dev"
	}
	for _, c := range items {
		c.SourceSchema = job.SourceSchema
		res, err := r.Tx.Process(ctx, env, c)
		cp.Scanned++
		cp.LastPK = c.SourcePK
		if err != nil {
			job.Checkpoints[key] = cp
			job.Status = JobFailed
			return Recalc(checkpointsOf(job)), err
		}
		switch res.Status {
		case "quarantine":
			cp.Quarantined++
		default:
			cp.Succeeded++
		}
	}
	job.Checkpoints[key] = cp
	return Recalc(checkpointsOf(job)), nil
}

func checkpointsOf(job *Job) []Checkpoint {
	out := make([]Checkpoint, 0, len(job.Checkpoints))
	for _, c := range job.Checkpoints {
		out = append(out, c)
	}
	return out
}