Agent-System/internal/workspace/registry.go

package workspace

import (
	"encoding/json"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"time"
)

// Record is one registered workspace directory.
type Record struct {
	ID        string             `json:"id"`
	RunID     string             `json:"run_id"`
	Bytes     int64              `json:"bytes"`
	CreatedAt time.Time          `json:"created_at"`
	LastSeen  time.Time          `json:"last_seen"`
	Frozen    *FrozenManifest    `json:"frozen,omitempty"`
	Receipts  []RetentionReceipt `json:"receipts,omitempty"`
}

// FileRegistry stores workspace records next to the environment data directory.
type FileRegistry struct {
	dataRoot string
	path     string
	mu       sync.Mutex
}

func NewFileRegistry(dataRoot string) *FileRegistry {
	return &FileRegistry{
		dataRoot: dataRoot,
		path:     filepath.Join(dataRoot, "workspace-registry.json"),
	}
}

func (r *FileRegistry) Upsert(rec Record) error {
	if err := validateID(rec.ID); err != nil {
		return err
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	all, err := r.load()
	if err != nil {
		return err
	}
	now := time.Now().UTC()
	if rec.CreatedAt.IsZero() {
		if prev, ok := all[rec.ID]; ok && !prev.CreatedAt.IsZero() {
			rec.CreatedAt = prev.CreatedAt
		} else {
			rec.CreatedAt = now
		}
	}
	if rec.LastSeen.IsZero() {
		rec.LastSeen = now
	}
	if rec.Bytes == 0 {
		rec.Bytes = dirBytes(filepath.Join(r.dataRoot, "workspaces", rec.ID))
	}
	all[rec.ID] = rec
	return r.save(all)
}

func (r *FileRegistry) Get(id string) (Record, bool, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	all, err := r.load()
	if err != nil {
		return Record{}, false, err
	}
	rec, ok := all[id]
	return rec, ok, nil
}

func (r *FileRegistry) List() ([]Record, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	all, err := r.load()
	if err != nil {
		return nil, err
	}
	out := make([]Record, 0, len(all))
	for _, rec := range all {
		out = append(out, rec)
	}
	sort.Slice(out, func(i, j int) bool { return out[i].ID < out[j].ID })
	return out, nil
}

func (r *FileRegistry) Delete(id string) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	all, err := r.load()
	if err != nil {
		return err
	}
	delete(all, id)
	return r.save(all)
}

// Discover registers workspace directories that are not yet in the registry.
func (r *FileRegistry) Discover() error {
	root := filepath.Join(r.dataRoot, "workspaces")
	entries, err := os.ReadDir(root)
	if os.IsNotExist(err) {
		return nil
	}
	if err != nil {
		return err
	}
	for _, e := range entries {
		if !e.IsDir() {
			continue
		}
		if err := validateID(e.Name()); err != nil {
			continue
		}
		info, err := e.Info()
		if err != nil {
			continue
		}
		_, ok, err := r.Get(e.Name())
		if err != nil {
			return err
		}
		if ok {
			continue
		}
		mod := info.ModTime().UTC()
		if err := r.Upsert(Record{ID: e.Name(), CreatedAt: mod, LastSeen: mod, Bytes: dirBytes(filepath.Join(root, e.Name()))}); err != nil {
			return err
		}
	}
	return nil
}

// NoteWorkspace registers a directory when it lives under dataRoot/workspaces/<id>.
func NoteWorkspace(dataRoot, workDir, runID string) error {
	if dataRoot == "" || workDir == "" {
		return nil
	}
	root, err := filepath.Abs(filepath.Join(dataRoot, "workspaces"))
	if err != nil {
		return err
	}
	abs, err := filepath.Abs(workDir)
	if err != nil {
		return err
	}
	rel, err := filepath.Rel(root, abs)
	if err != nil || rel == "." || rel == ".." || strings.HasPrefix(rel, ".."+string(os.PathSeparator)) {
		return nil
	}
	id := strings.Split(rel, string(os.PathSeparator))[0]
	return NewFileRegistry(dataRoot).Upsert(Record{ID: id, RunID: runID})
}

func (r *FileRegistry) load() (map[string]Record, error) {
	b, err := os.ReadFile(r.path)
	if os.IsNotExist(err) {
		return map[string]Record{}, nil
	}
	if err != nil {
		return nil, err
	}
	if len(strings.TrimSpace(string(b))) == 0 {
		return map[string]Record{}, nil
	}
	var list []Record
	if err := json.Unmarshal(b, &list); err != nil {
		return nil, err
	}
	out := make(map[string]Record, len(list))
	for _, rec := range list {
		out[rec.ID] = rec
	}
	return out, nil
}

func (r *FileRegistry) save(all map[string]Record) error {
	list := make([]Record, 0, len(all))
	for _, rec := range all {
		list = append(list, rec)
	}
	sort.Slice(list, func(i, j int) bool { return list[i].ID < list[j].ID })
	raw, err := json.MarshalIndent(list, "", "  ")
	if err != nil {
		return err
	}
	if err := os.MkdirAll(filepath.Dir(r.path), 0o755); err != nil {
		return err
	}
	tmp := r.path + ".tmp"
	if err := os.WriteFile(tmp, append(raw, '\n'), 0o644); err != nil {
		return err
	}
	return os.Rename(tmp, r.path)
}

func validateID(id string) error {
	if id == "" || strings.Contains(id, "/") || strings.Contains(id, "..") || id == "." {
		return ErrInvalidKey
	}
	return nil
}

func dirBytes(root string) int64 {
	var n int64
	_ = filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
		if err != nil || info == nil || info.IsDir() {
			return nil
		}
		n += info.Size()
		return nil
	})
	return n
}