diff --git a/deploy/postgres/goose/00062_2026-09-21-agentx-workspace-registry.sql b/deploy/postgres/goose/00062_2026-09-21-agentx-workspace-registry.sql new file mode 100644 index 00000000..7fc21318 --- /dev/null +++ b/deploy/postgres/goose/00062_2026-09-21-agentx-workspace-registry.sql @@ -0,0 +1,13 @@ +-- +goose Up +CREATE TABLE IF NOT EXISTS agentx_workspace_registry ( + id TEXT PRIMARY KEY, + run_id TEXT NOT NULL DEFAULT '', + bytes BIGINT NOT NULL DEFAULT 0, + created_at TIMESTAMPTZ NOT NULL, + last_seen TIMESTAMPTZ NOT NULL, + frozen JSONB, + receipts JSONB NOT NULL DEFAULT '[]'::jsonb +); + +-- +goose Down +DROP TABLE IF EXISTS agentx_workspace_registry; diff --git a/internal/agent/pinative.go b/internal/agent/pinative.go index 1ec2c1e7..c7610801 100644 --- a/internal/agent/pinative.go +++ b/internal/agent/pinative.go @@ -11,6 +11,7 @@ import ( "chat-agent-system/internal/agentx/session" "chat-agent-system/internal/agentx/tools" "chat-agent-system/internal/config" + "chat-agent-system/internal/workspace" ) // PiNativeBackend runs the in-process Go agent loop. The container pi-agent backend stays available. @@ -44,6 +45,9 @@ func (p *PiNativeBackend) Execute(ctx context.Context, req Request) (*Result, er if err := os.MkdirAll(workDir, 0o755); err != nil { return &Result{Success: false, ErrorReason: err.Error()}, nil } + if p.cfg != nil { + _ = workspace.NoteWorkspace(p.cfg.DataDir, workDir, filepath.Base(workDir)) + } if strings.TrimSpace(req.Instructions) != "" { _ = os.WriteFile(filepath.Join(workDir, "INSTRUCTIONS.md"), []byte(req.Instructions), 0o644) _ = os.WriteFile(filepath.Join(workDir, "AGENTS.md"), []byte(req.Instructions), 0o644) diff --git a/internal/server/server.go b/internal/server/server.go index e49b3584..30d8e420 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -67,6 +67,8 @@ type Server struct { freeDoc *freedoc.Service kb *kb.Client sysadmin *sysadmin.Service + wsReg *workspace.FileRegistry + wsGCStop chan struct{} httpServer *http.Server intentPlans intentflow.Service planAccessFn func(ctx context.Context, projectID, userID int64) error @@ -116,6 +118,7 @@ func NewFromRuntime(rt *app.Runtime) *Server { freeDoc: rt.FreeDoc, kb: rt.KB, sysadmin: rt.SysAdmin, + wsReg: workspaceRegistryFor(rt.Cfg), intentPlans: rt.IntentPlans, // 登录/注册限流：10 次 / 5 分钟 / (IP+username) 滑动窗口 loginLimiter: newLoginRateLimiter(10, 5*time.Minute), @@ -243,6 +246,8 @@ func (s *Server) Routes() http.Handler { mux.HandleFunc("POST /api/v1/tasks/runs/{run_id}/cancel", s.handleCancelTaskRun) mux.HandleFunc("POST /api/v1/tasks/runs/{run_id}/resume", s.handleResumeTaskRun) mux.HandleFunc("POST /api/v1/admin/workspaces/reconcile", s.handleReconcileWorkspace) + mux.HandleFunc("GET /api/v1/admin/workspaces", s.handleAdminWorkspaces) + mux.HandleFunc("POST /api/v1/admin/workspaces/{id}/cleanup", s.handleAdminWorkspaceCleanup) mux.HandleFunc("POST /api/v1/projects/{id}/assets", s.handleUploadProjectAsset) mux.HandleFunc("GET /api/v1/projects/{id}/assets", s.handleListProjectAssets) mux.HandleFunc("GET /api/v1/projects/{id}/assets/{asset_id}/raw", s.handleGetProjectAssetRaw) @@ -359,6 +364,7 @@ func (s *Server) Start() error { // Workers are started by cmd/server before Listen. Runtime.Start is // idempotent, but calling it here swallowed the first error and raced // with main. Do not start again. + s.startWorkspaceGC() addr := fmt.Sprintf("%s:%d", s.cfg.HTTPHost, s.cfg.HTTPPort) s.httpServer = &http.Server{ Addr: addr, @@ -374,6 +380,10 @@ func (s *Server) Start() error { // Shutdown 优雅停机 func (s *Server) Shutdown(ctx context.Context) error { + if s.wsGCStop != nil { + close(s.wsGCStop) + s.wsGCStop = nil + } if s.reminderStop != nil { close(s.reminderStop) s.reminderStop = nil diff --git a/internal/server/workspace_admin_handlers.go b/internal/server/workspace_admin_handlers.go new file mode 100644 index 00000000..844951ef --- /dev/null +++ b/internal/server/workspace_admin_handlers.go @@ -0,0 +1,133 @@ +package server + +import ( + "context" + "encoding/json" + "net/http" + "time" + + "chat-agent-system/internal/config" + "chat-agent-system/internal/workspace" +) + +func (s *Server) handleAdminWorkspaces(w http.ResponseWriter, r *http.Request) { + if _, ok := s.requireAdmin(w, r); !ok { + return + } + s.serveAdminWorkspaces(w, r) +} + +func (s *Server) handleAdminWorkspaceCleanup(w http.ResponseWriter, r *http.Request) { + if _, ok := s.requireAdmin(w, r); !ok { + return + } + s.serveAdminWorkspaceCleanup(w, r) +} + +func (s *Server) workspaceRegistry() *workspace.FileRegistry { + if s.wsReg != nil { + return s.wsReg + } + if s.cfg == nil || s.cfg.DataDir == "" { + return nil + } + s.wsReg = workspace.NewFileRegistry(s.cfg.DataDir) + return s.wsReg +} + +func (s *Server) serveAdminWorkspaces(w http.ResponseWriter, r *http.Request) { + reg := s.workspaceRegistry() + if reg == nil { + writeSafeError(w, r, http.StatusServiceUnavailable, "workspace registry unavailable") + return + } + if err := reg.Discover(); err != nil { + writeSafeError(w, r, http.StatusInternalServerError, err.Error()) + return + } + list, err := reg.List() + if err != nil { + writeSafeError(w, r, http.StatusInternalServerError, err.Error()) + return + } + var bytes int64 + for _, rec := range list { + bytes += rec.Bytes + } + writeJSONResponse(w, http.StatusOK, map[string]any{ + "workspaces": list, + "stats": map[string]any{ + "count": len(list), + "bytes": bytes, + }, + }) +} + +func (s *Server) serveAdminWorkspaceCleanup(w http.ResponseWriter, r *http.Request) { + reg := s.workspaceRegistry() + if reg == nil { + writeSafeError(w, r, http.StatusServiceUnavailable, "workspace registry unavailable") + return + } + id := r.PathValue("id") + var body struct { + Frozen *workspace.FrozenManifest `json:"frozen"` + Receipts []workspace.RetentionReceipt `json:"receipts"` + } + if r.Body != nil { + _ = json.NewDecoder(r.Body).Decode(&body) + } + if body.Frozen != nil { + if err := reg.SetRetention(id, *body.Frozen, body.Receipts); err != nil { + writeSafeError(w, r, http.StatusBadRequest, err.Error()) + return + } + } + cap := s.wsCap + if cap == nil && s.cfg != nil { + cap = workspace.NewCapability(s.cfg.DataDir, nil) + } + if err := reg.CleanupOne(r.Context(), id, cap); err != nil { + status := http.StatusConflict + if err == workspace.ErrWorkDirNotFound { + status = http.StatusNotFound + } + writeSafeError(w, r, status, err.Error()) + return + } + writeJSONResponse(w, http.StatusOK, map[string]any{"deleted": id}) +} + +func workspaceRegistryFor(cfg *config.Config) *workspace.FileRegistry { + if cfg == nil || cfg.DataDir == "" { + return nil + } + return workspace.NewFileRegistry(cfg.DataDir) +} + +func (s *Server) startWorkspaceGC() { + reg := s.workspaceRegistry() + if reg == nil || s.cfg == nil || s.wsGCStop != nil { + return + } + s.wsGCStop = make(chan struct{}) + ttl := s.cfg.WorkspaceTTL + quota := s.cfg.WorkspaceQuotaBytes + stop := s.wsGCStop + go func() { + ticker := time.NewTicker(time.Hour) + defer ticker.Stop() + for { + select { + case <-stop: + return + case <-ticker.C: + cap := s.wsCap + if cap == nil { + cap = workspace.NewCapability(s.cfg.DataDir, nil) + } + _, _ = reg.CollectGarbage(context.Background(), time.Now(), ttl, quota, cap) + } + } + }() +} diff --git a/internal/server/workspace_admin_handlers_test.go b/internal/server/workspace_admin_handlers_test.go new file mode 100644 index 00000000..e6c0fe71 --- /dev/null +++ b/internal/server/workspace_admin_handlers_test.go @@ -0,0 +1,54 @@ +package server + +import ( + "context" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "chat-agent-system/internal/config" + "chat-agent-system/internal/workspace" +) + +func TestServeAdminWorkspacesListsRegistry(t *testing.T) { + dir := t.TempDir() + id := "ws-admin" + if err := os.MkdirAll(filepath.Join(dir, "workspaces", id), 0o755); err != nil { + t.Fatal(err) + } + s := &Server{cfg: &config.Config{DataDir: dir}} + req := httptest.NewRequest(http.MethodGet, "/api/v1/admin/workspaces", nil) + rec := httptest.NewRecorder() + s.serveAdminWorkspaces(rec, req) + if rec.Code != http.StatusOK || !strings.Contains(rec.Body.String(), id) { + t.Fatalf("status %d body %s", rec.Code, rec.Body.String()) + } +} + +func TestServeAdminWorkspaceCleanupHonorsGate(t *testing.T) { + dir := t.TempDir() + id := "ws-clean" + ws := filepath.Join(dir, "workspaces", id) + if err := os.MkdirAll(ws, 0o755); err != nil { + t.Fatal(err) + } + reg := workspace.NewFileRegistry(dir) + if err := reg.Upsert(workspace.Record{ID: id, CreatedAt: time.Now()}); err != nil { + t.Fatal(err) + } + s := &Server{cfg: &config.Config{DataDir: dir}, wsReg: reg} + req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/workspaces/"+id+"/cleanup", nil) + req.SetPathValue("id", id) + rec := httptest.NewRecorder() + s.serveAdminWorkspaceCleanup(rec, req.WithContext(context.Background())) + if rec.Code != http.StatusConflict { + t.Fatalf("status %d body %s", rec.Code, rec.Body.String()) + } + if _, err := os.Stat(ws); err != nil { + t.Fatal(err) + } +} diff --git a/internal/workspace/gc.go b/internal/workspace/gc.go new file mode 100644 index 00000000..3e1369ff --- /dev/null +++ b/internal/workspace/gc.go @@ -0,0 +1,107 @@ +package workspace + +import ( + "context" + "sort" + "time" +) + +// GCReport is the outcome of one retention pass. Directories without a matching +// frozen manifest and verified receipts are left in place. +type GCReport struct { + Deleted []string `json:"deleted"` + Blocked []string `json:"blocked"` + Bytes int64 `json:"bytes"` +} + +// CollectGarbage removes expired or over-quota workspaces only when +// AssertCleanupAllowed accepts the stored manifest and receipts. +func (r *FileRegistry) CollectGarbage(ctx context.Context, now time.Time, ttl time.Duration, quota int64, cap *Capability) (GCReport, error) { + if err := r.Discover(); err != nil { + return GCReport{}, err + } + list, err := r.List() + if err != nil { + return GCReport{}, err + } + if cap == nil { + cap = NewCapability(r.dataRoot, nil) + } + var report GCReport + for _, rec := range list { + report.Bytes += rec.Bytes + } + expired := map[string]bool{} + if ttl > 0 { + for _, rec := range list { + if now.Sub(rec.CreatedAt) > ttl { + expired[rec.ID] = true + } + } + } + over := quota > 0 && report.Bytes > quota + if !over && len(expired) == 0 { + return report, nil + } + candidates := append([]Record(nil), list...) + sort.Slice(candidates, func(i, j int) bool { return candidates[i].LastSeen.Before(candidates[j].LastSeen) }) + for _, rec := range candidates { + if ctx.Err() != nil { + return report, ctx.Err() + } + need := expired[rec.ID] || (quota > 0 && report.Bytes > quota) + if !need { + continue + } + if err := r.deleteIfAllowed(ctx, cap, rec); err != nil { + report.Blocked = append(report.Blocked, rec.ID) + continue + } + report.Deleted = append(report.Deleted, rec.ID) + report.Bytes -= rec.Bytes + if report.Bytes < 0 { + report.Bytes = 0 + } + } + return report, nil +} + +// CleanupOne deletes one workspace when its manifest and receipts match. +func (r *FileRegistry) CleanupOne(ctx context.Context, id string, cap *Capability) error { + rec, ok, err := r.Get(id) + if err != nil { + return err + } + if !ok { + return ErrWorkDirNotFound + } + if cap == nil { + cap = NewCapability(r.dataRoot, nil) + } + return r.deleteIfAllowed(ctx, cap, rec) +} + +func (r *FileRegistry) deleteIfAllowed(ctx context.Context, cap *Capability, rec Record) error { + if rec.Frozen == nil { + return ErrCleanupBlocked + } + ref := WorkspaceRef{ID: rec.ID, SchemaVersion: SchemaVersion, RunID: rec.RunID} + if err := cap.Cleanup(ctx, ref, *rec.Frozen, rec.Receipts); err != nil { + return err + } + return r.Delete(rec.ID) +} + +// SetRetention stores the frozen manifest and receipts required before deletion. +func (r *FileRegistry) SetRetention(id string, frozen FrozenManifest, receipts []RetentionReceipt) error { + rec, ok, err := r.Get(id) + if err != nil { + return err + } + if !ok { + rec.ID = id + } + rec.Frozen = &frozen + rec.Receipts = receipts + return r.Upsert(rec) +} diff --git a/internal/workspace/registry.go b/internal/workspace/registry.go new file mode 100644 index 00000000..d1d80b42 --- /dev/null +++ b/internal/workspace/registry.go @@ -0,0 +1,219 @@ +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/. +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 +} diff --git a/internal/workspace/registry_pg.go b/internal/workspace/registry_pg.go new file mode 100644 index 00000000..a64bdaa1 --- /dev/null +++ b/internal/workspace/registry_pg.go @@ -0,0 +1,112 @@ +package workspace + +import ( + "context" + "encoding/json" + "fmt" + "time" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" +) + +// RegistryDDL is the table used by PGRegistry. Goose migration 00062 is the deploy path. +const RegistryDDL = ` +CREATE TABLE IF NOT EXISTS agentx_workspace_registry ( + id TEXT PRIMARY KEY, + run_id TEXT NOT NULL DEFAULT '', + bytes BIGINT NOT NULL DEFAULT 0, + created_at TIMESTAMPTZ NOT NULL, + last_seen TIMESTAMPTZ NOT NULL, + frozen JSONB, + receipts JSONB NOT NULL DEFAULT '[]'::jsonb +)` + +// PGRegistry persists workspace records in Postgres. +type PGRegistry struct { + pool *pgxpool.Pool +} + +func NewPGRegistry(pool *pgxpool.Pool) *PGRegistry { + return &PGRegistry{pool: pool} +} + +func (p *PGRegistry) Ensure(ctx context.Context) error { + _, err := p.pool.Exec(ctx, RegistryDDL) + return err +} + +func (p *PGRegistry) Upsert(ctx context.Context, rec Record) error { + if err := validateID(rec.ID); err != nil { + return err + } + if rec.CreatedAt.IsZero() { + rec.CreatedAt = time.Now().UTC() + } + if rec.LastSeen.IsZero() { + rec.LastSeen = rec.CreatedAt + } + frozen, err := json.Marshal(rec.Frozen) + if err != nil { + return err + } + if rec.Frozen == nil { + frozen = []byte("null") + } + receipts, err := json.Marshal(rec.Receipts) + if err != nil { + return err + } + if rec.Receipts == nil { + receipts = []byte("[]") + } + _, err = p.pool.Exec(ctx, ` + INSERT INTO agentx_workspace_registry (id, run_id, bytes, created_at, last_seen, frozen, receipts) + VALUES ($1,$2,$3,$4,$5,$6::jsonb,$7::jsonb) + ON CONFLICT (id) DO UPDATE SET + run_id=EXCLUDED.run_id, bytes=EXCLUDED.bytes, last_seen=EXCLUDED.last_seen, + frozen=EXCLUDED.frozen, receipts=EXCLUDED.receipts`, + rec.ID, rec.RunID, rec.Bytes, rec.CreatedAt, rec.LastSeen, frozen, receipts) + return err +} + +func (p *PGRegistry) Get(ctx context.Context, id string) (Record, bool, error) { + row := p.pool.QueryRow(ctx, `SELECT id, run_id, bytes, created_at, last_seen, frozen, receipts FROM agentx_workspace_registry WHERE id=$1`, id) + rec, err := scanRecord(row) + if err == pgx.ErrNoRows { + return Record{}, false, nil + } + if err != nil { + return Record{}, false, err + } + return rec, true, nil +} + +func scanRecord(row pgx.Row) (Record, error) { + var rec Record + var frozen, receipts []byte + if err := row.Scan(&rec.ID, &rec.RunID, &rec.Bytes, &rec.CreatedAt, &rec.LastSeen, &frozen, &receipts); err != nil { + return Record{}, err + } + if len(frozen) > 0 && string(frozen) != "null" { + rec.Frozen = &FrozenManifest{} + if err := json.Unmarshal(frozen, rec.Frozen); err != nil { + return Record{}, err + } + } + if len(receipts) > 0 { + if err := json.Unmarshal(receipts, &rec.Receipts); err != nil { + return Record{}, err + } + } + return rec, nil +} + +// TouchSchema creates an isolated schema for tests and points the pool at it. +func TouchSchema(ctx context.Context, pool *pgxpool.Pool, schema string) error { + if _, err := pool.Exec(ctx, fmt.Sprintf("CREATE SCHEMA IF NOT EXISTS %s", schema)); err != nil { + return err + } + _, err := pool.Exec(ctx, fmt.Sprintf("SET search_path TO %s", schema)) + return err +} diff --git a/internal/workspace/registry_pg_test.go b/internal/workspace/registry_pg_test.go new file mode 100644 index 00000000..7ce94fe5 --- /dev/null +++ b/internal/workspace/registry_pg_test.go @@ -0,0 +1,68 @@ +package workspace + +import ( + "context" + "fmt" + "os" + "testing" + "time" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" +) + +func TestPGRegistryRoundTrip(t *testing.T) { + url := os.Getenv("TEST_DATABASE_URL") + if url == "" { + t.Skip("set TEST_DATABASE_URL to run the workspace registry postgres test") + } + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + schema := fmt.Sprintf("agentx_reg_%d", time.Now().UnixNano()) + boot, err := pgxpool.New(ctx, url) + if err != nil { + t.Fatal(err) + } + defer boot.Close() + if _, err := boot.Exec(ctx, "CREATE SCHEMA "+schema); err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + cctx, ccancel := context.WithTimeout(context.Background(), 10*time.Second) + defer ccancel() + _, _ = boot.Exec(cctx, "DROP SCHEMA "+schema+" CASCADE") + }) + cfg, err := pgxpool.ParseConfig(url) + if err != nil { + t.Fatal(err) + } + cfg.AfterConnect = func(ctx context.Context, conn *pgx.Conn) error { + _, err := conn.Exec(ctx, "SET search_path TO "+schema) + return err + } + pool, err := pgxpool.NewWithConfig(ctx, cfg) + if err != nil { + t.Fatal(err) + } + defer pool.Close() + reg := NewPGRegistry(pool) + if err := reg.Ensure(ctx); err != nil { + t.Fatal(err) + } + frozen, receipts := sampleFrozen("ws-pg") + rec := Record{ + ID: "ws-pg", RunID: "run-pg", Bytes: 12, + CreatedAt: time.Now().UTC().Add(-time.Hour), + Frozen: &frozen, Receipts: receipts, + } + if err := reg.Upsert(ctx, rec); err != nil { + t.Fatal(err) + } + got, ok, err := reg.Get(ctx, "ws-pg") + if err != nil || !ok { + t.Fatalf("get %#v %v %v", got, ok, err) + } + if got.RunID != "run-pg" || got.Frozen == nil || got.Frozen.SHA256 != "abc" || len(got.Receipts) != 1 || !got.Receipts[0].Verified { + t.Fatalf("round trip %#v", got) + } +} diff --git a/internal/workspace/registry_test.go b/internal/workspace/registry_test.go new file mode 100644 index 00000000..c88d8e5a --- /dev/null +++ b/internal/workspace/registry_test.go @@ -0,0 +1,100 @@ +package workspace + +import ( + "context" + "os" + "path/filepath" + "testing" + "time" +) + +func sampleFrozen(id string) (FrozenManifest, []RetentionReceipt) { + frozen := FrozenManifest{ + SHA256: "abc", + StorageIdentity: "disk", + WorkspaceID: id, + Members: []FrozenMember{{LogicalID: "doc", Digest: "d1", Kind: "document"}}, + } + receipts := []RetentionReceipt{{LogicalID: "doc", Digest: "d1", StorageIdentity: "disk", Verified: true}} + return frozen, receipts +} + +func TestGCDeletesExpiredWhenReceiptsMatch(t *testing.T) { + root := t.TempDir() + id := "ws-old" + dir := filepath.Join(root, "workspaces", id) + if err := os.MkdirAll(dir, 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(dir, "output.md"), []byte("kept-until-gc"), 0o644); err != nil { + t.Fatal(err) + } + reg := NewFileRegistry(root) + frozen, receipts := sampleFrozen(id) + if err := reg.Upsert(Record{ID: id, CreatedAt: time.Now().Add(-8 * 24 * time.Hour), LastSeen: time.Now().Add(-8 * 24 * time.Hour)}); err != nil { + t.Fatal(err) + } + if err := reg.SetRetention(id, frozen, receipts); err != nil { + t.Fatal(err) + } + report, err := reg.CollectGarbage(context.Background(), time.Now(), 7*24*time.Hour, 0, nil) + if err != nil { + t.Fatal(err) + } + if len(report.Deleted) != 1 || report.Deleted[0] != id { + t.Fatalf("report %#v", report) + } + if _, err := os.Stat(dir); !os.IsNotExist(err) { + t.Fatal("workspace directory still exists") + } +} + +func TestGCBlocksWithoutReceipts(t *testing.T) { + root := t.TempDir() + id := "ws-blocked" + dir := filepath.Join(root, "workspaces", id) + if err := os.MkdirAll(dir, 0o755); err != nil { + t.Fatal(err) + } + reg := NewFileRegistry(root) + if err := reg.Upsert(Record{ID: id, CreatedAt: time.Now().Add(-30 * 24 * time.Hour)}); err != nil { + t.Fatal(err) + } + report, err := reg.CollectGarbage(context.Background(), time.Now(), time.Hour, 1, nil) + if err != nil { + t.Fatal(err) + } + if len(report.Deleted) != 0 || len(report.Blocked) != 1 { + t.Fatalf("report %#v", report) + } + if _, err := os.Stat(dir); err != nil { + t.Fatal(err) + } +} + +func TestDiscoverAndNote(t *testing.T) { + root := t.TempDir() + id := "ws-scan" + if err := os.MkdirAll(filepath.Join(root, "workspaces", id), 0o755); err != nil { + t.Fatal(err) + } + reg := NewFileRegistry(root) + if err := reg.Discover(); err != nil { + t.Fatal(err) + } + rec, ok, err := reg.Get(id) + if err != nil || !ok || rec.ID != id { + t.Fatalf("discovered %#v %v %v", rec, ok, err) + } + nested := filepath.Join(root, "workspaces", "ws-note", "tasks") + if err := os.MkdirAll(nested, 0o755); err != nil { + t.Fatal(err) + } + if err := NoteWorkspace(root, nested, "run-1"); err != nil { + t.Fatal(err) + } + noted, ok, err := reg.Get("ws-note") + if err != nil || !ok || noted.RunID != "run-1" { + t.Fatalf("noted %#v %v", noted, err) + } +}