diff --git a/pkg/agentcompose/app/app.go b/pkg/agentcompose/app/app.go index 2e9264bf..7c13c930 100644 --- a/pkg/agentcompose/app/app.go +++ b/pkg/agentcompose/app/app.go @@ -53,8 +53,9 @@ func Register(di do.Injector) { func RegisterDependencies(di do.Injector) { do.Provide(di, func(do.Injector) (*sessions.LifecycleLocks, error) { return sessions.NewLifecycleLocks(), nil }) - do.Provide(di, sessionstore.NewStore) do.Provide(di, NewConfigStore) + do.MustAs[*configstore.ConfigStore, sessionstore.SandboxRecorder](di) + do.Provide(di, sessionstore.NewStore) do.Provide(di, NewWorkspaceProvisioner) do.MustAs[*workspaces.Provisioner, workspaces.WorkspaceEnsurer](di) do.Provide(di, NewRuntimeProvider) diff --git a/pkg/model/model.go b/pkg/model/model.go index e70de3e5..d0c81e18 100644 --- a/pkg/model/model.go +++ b/pkg/model/model.go @@ -130,6 +130,22 @@ type SandboxSummary struct { Tags []SandboxTag `json:"tags,omitempty"` } +// SandboxSummaryListOptions defines the indexed filters supported by the +// persisted Sandbox summary records. +type SandboxSummaryListOptions struct { + Driver string + VMStatus string + BeforeUpdatedAt time.Time + BeforeID string + Limit int +} + +// SandboxSummaryListResult contains one page of persisted Sandbox summaries. +type SandboxSummaryListResult struct { + Sandboxes []SandboxSummary + HasMore bool +} + type SandboxListOptions struct { SandboxType string TriggerSourceQuery string diff --git a/pkg/storage/configstore/core_store.go b/pkg/storage/configstore/core_store.go index 69ea09c8..90db40bc 100644 --- a/pkg/storage/configstore/core_store.go +++ b/pkg/storage/configstore/core_store.go @@ -55,6 +55,7 @@ type ConfigStore struct { *loaderStore *eventStore *projectStore + *sandboxStore *llmStore *capabilityGatewayStore *volumeStore @@ -87,6 +88,7 @@ func FromDB(db *sql.DB) *ConfigStore { loaderStore: &loaderStore{db: db}, eventStore: &eventStore{db: db}, projectStore: &projectStore{db: db}, + sandboxStore: &sandboxStore{db: db}, llmStore: &llmStore{db: db}, capabilityGatewayStore: &capabilityGatewayStore{db: db}, volumeStore: &volumeStore{db: db}, @@ -147,6 +149,9 @@ func (s *ConfigStore) initSchema(ctx context.Context) error { if err := s.ensureProjectSchema(ctx); err != nil { return err } + if err := s.ensureSandboxSchema(ctx); err != nil { + return err + } if err := s.ensureEventSchema(ctx); err != nil { return err } diff --git a/pkg/storage/configstore/sandbox_store.go b/pkg/storage/configstore/sandbox_store.go new file mode 100644 index 00000000..c2c1b7aa --- /dev/null +++ b/pkg/storage/configstore/sandbox_store.go @@ -0,0 +1,241 @@ +package configstore + +import ( + "context" + "database/sql" + "encoding/json" + "fmt" + "strings" + "time" + + domain "agent-compose/pkg/model" +) + +type sandboxStore struct { + db *sql.DB +} + +const sandboxSummaryColumns = `id, short_id, title, trigger_source, driver, + vm_status, guest_image, pull_policy, runtime_ref, workspace_path, proxy_path, + cell_count, event_count, tags_json, created_at, updated_at` + +func (s *sandboxStore) ensureSandboxSchema(ctx context.Context) error { + // created_at and updated_at are Unix nanoseconds. Nanosecond precision keeps + // database ordering identical to the time.Time values in metadata.json. + statements := []string{ + `CREATE TABLE IF NOT EXISTS sandboxes ( + id TEXT PRIMARY KEY, + short_id TEXT NOT NULL DEFAULT '', + title TEXT NOT NULL DEFAULT '', + trigger_source TEXT NOT NULL DEFAULT '', + driver TEXT NOT NULL DEFAULT '', + vm_status TEXT NOT NULL DEFAULT '', + guest_image TEXT NOT NULL DEFAULT '', + pull_policy TEXT NOT NULL DEFAULT '', + runtime_ref TEXT NOT NULL DEFAULT '', + workspace_path TEXT NOT NULL DEFAULT '', + proxy_path TEXT NOT NULL DEFAULT '', + cell_count INTEGER NOT NULL DEFAULT 0, + event_count INTEGER NOT NULL DEFAULT 0, + tags_json TEXT NOT NULL DEFAULT '[]', + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL + );`, + `CREATE INDEX IF NOT EXISTS idx_sandboxes_updated ON sandboxes(updated_at DESC, id DESC);`, + `CREATE INDEX IF NOT EXISTS idx_sandboxes_vm_status_updated ON sandboxes(vm_status, updated_at DESC, id DESC);`, + } + for _, statement := range statements { + if _, err := s.db.ExecContext(ctx, statement); err != nil { + return fmt.Errorf("create sandboxes schema: %w", err) + } + } + return nil +} + +// UpsertSandbox records the queryable Sandbox summary in SQLite. Sandbox +// directories remain the source for workspace and runtime files. +func (s *sandboxStore) UpsertSandbox(ctx context.Context, sandbox *domain.Sandbox) error { + if sandbox == nil { + return fmt.Errorf("sandbox is required") + } + summary := sandbox.Summary + if strings.TrimSpace(summary.ID) == "" { + return fmt.Errorf("sandbox id is required") + } + tags := summary.Tags + if tags == nil { + tags = []domain.SandboxTag{} + } + tagsJSON, err := json.Marshal(tags) + if err != nil { + return fmt.Errorf("encode sandbox tags: %w", err) + } + _, err = s.db.ExecContext(ctx, `INSERT INTO sandboxes ( + id, short_id, title, trigger_source, driver, vm_status, guest_image, + pull_policy, runtime_ref, workspace_path, proxy_path, cell_count, + event_count, tags_json, created_at, updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(id) DO UPDATE SET + short_id = excluded.short_id, + title = excluded.title, + trigger_source = excluded.trigger_source, + driver = excluded.driver, + vm_status = excluded.vm_status, + guest_image = excluded.guest_image, + pull_policy = excluded.pull_policy, + runtime_ref = excluded.runtime_ref, + workspace_path = excluded.workspace_path, + proxy_path = excluded.proxy_path, + cell_count = excluded.cell_count, + event_count = excluded.event_count, + tags_json = excluded.tags_json, + updated_at = excluded.updated_at + WHERE excluded.updated_at >= sandboxes.updated_at`, + summary.ID, + summary.ShortID, + summary.Title, + summary.TriggerSource, + summary.Driver, + summary.VMStatus, + summary.GuestImage, + summary.PullPolicy, + summary.RuntimeRef, + summary.WorkspacePath, + summary.ProxyPath, + summary.CellCount, + summary.EventCount, + string(tagsJSON), + sandboxTimestampValue(summary.CreatedAt), + sandboxTimestampValue(summary.UpdatedAt), + ) + if err != nil { + return fmt.Errorf("upsert sandbox %s: %w", summary.ID, err) + } + return nil +} + +func (s *sandboxStore) GetSandboxSummary(ctx context.Context, id string) (domain.SandboxSummary, error) { + id = strings.TrimSpace(id) + if id == "" { + return domain.SandboxSummary{}, fmt.Errorf("sandbox id is required") + } + row := s.db.QueryRowContext(ctx, `SELECT `+sandboxSummaryColumns+` FROM sandboxes WHERE id = ?`, id) + summary, err := scanSandboxSummary(row.Scan) + if err != nil { + return domain.SandboxSummary{}, fmt.Errorf("get sandbox %s: %w", id, err) + } + return summary, nil +} + +func (s *sandboxStore) ListSandboxSummaries(ctx context.Context, options domain.SandboxSummaryListOptions) (domain.SandboxSummaryListResult, error) { + beforeID := strings.TrimSpace(options.BeforeID) + if !options.BeforeUpdatedAt.IsZero() && beforeID == "" { + return domain.SandboxSummaryListResult{}, fmt.Errorf("sandbox before id is required when before updated at is set") + } + if options.BeforeUpdatedAt.IsZero() && beforeID != "" { + return domain.SandboxSummaryListResult{}, fmt.Errorf("sandbox before updated at is required when before id is set") + } + where := make([]string, 0, 3) + args := make([]any, 0, 6) + if driver := strings.ToLower(strings.TrimSpace(options.Driver)); driver != "" { + where = append(where, "driver = ?") + args = append(args, driver) + } + if status := strings.ToUpper(strings.TrimSpace(options.VMStatus)); status != "" { + where = append(where, "vm_status = ?") + args = append(args, status) + } + if !options.BeforeUpdatedAt.IsZero() { + where = append(where, "(updated_at < ? OR (updated_at = ? AND id < ?))") + updatedAt := sandboxTimestampValue(options.BeforeUpdatedAt) + args = append(args, updatedAt, updatedAt, beforeID) + } + query := `SELECT ` + sandboxSummaryColumns + ` FROM sandboxes` + if len(where) > 0 { + query += ` WHERE ` + strings.Join(where, ` AND `) + } + query += ` ORDER BY updated_at DESC, id DESC LIMIT ?` + limit := options.Limit + if limit <= 0 { + limit = domain.DefaultSandboxListLimit + } + args = append(args, limit+1) + rows, err := s.db.QueryContext(ctx, query, args...) + if err != nil { + return domain.SandboxSummaryListResult{}, fmt.Errorf("list sandboxes: %w", err) + } + defer func() { _ = rows.Close() }() + + summaries := make([]domain.SandboxSummary, 0, limit+1) + for rows.Next() { + summary, err := scanSandboxSummary(rows.Scan) + if err != nil { + return domain.SandboxSummaryListResult{}, err + } + summaries = append(summaries, summary) + } + if err := rows.Err(); err != nil { + return domain.SandboxSummaryListResult{}, fmt.Errorf("iterate sandboxes: %w", err) + } + hasMore := len(summaries) > limit + if hasMore { + summaries = summaries[:limit] + } + return domain.SandboxSummaryListResult{Sandboxes: summaries, HasMore: hasMore}, nil +} + +func (s *sandboxStore) DeleteSandbox(ctx context.Context, id string) error { + id = strings.TrimSpace(id) + if id == "" { + return fmt.Errorf("sandbox id is required") + } + if _, err := s.db.ExecContext(ctx, `DELETE FROM sandboxes WHERE id = ?`, id); err != nil { + return fmt.Errorf("delete sandbox %s: %w", id, err) + } + return nil +} + +type sandboxSummaryScanner func(dest ...any) error + +func scanSandboxSummary(scan sandboxSummaryScanner) (domain.SandboxSummary, error) { + var summary domain.SandboxSummary + var tagsJSON string + var createdAt, updatedAt int64 + if err := scan( + &summary.ID, + &summary.ShortID, + &summary.Title, + &summary.TriggerSource, + &summary.Driver, + &summary.VMStatus, + &summary.GuestImage, + &summary.PullPolicy, + &summary.RuntimeRef, + &summary.WorkspacePath, + &summary.ProxyPath, + &summary.CellCount, + &summary.EventCount, + &tagsJSON, + &createdAt, + &updatedAt, + ); err != nil { + return domain.SandboxSummary{}, fmt.Errorf("scan sandbox summary: %w", err) + } + if err := json.Unmarshal([]byte(tagsJSON), &summary.Tags); err != nil { + return domain.SandboxSummary{}, fmt.Errorf("decode sandbox %s tags: %w", summary.ID, err) + } + if summary.Tags == nil { + summary.Tags = []domain.SandboxTag{} + } + summary.CreatedAt = sandboxTimestampTime(createdAt) + summary.UpdatedAt = sandboxTimestampTime(updatedAt) + return summary, nil +} + +func sandboxTimestampValue(value time.Time) int64 { + return value.UnixNano() +} + +func sandboxTimestampTime(value int64) time.Time { + return time.Unix(0, value).UTC() +} diff --git a/pkg/storage/configstore/sandbox_store_test.go b/pkg/storage/configstore/sandbox_store_test.go new file mode 100644 index 00000000..e1b303fd --- /dev/null +++ b/pkg/storage/configstore/sandbox_store_test.go @@ -0,0 +1,178 @@ +package configstore + +import ( + "context" + "database/sql" + "errors" + "testing" + "time" + + domain "agent-compose/pkg/model" +) + +func TestSandboxStoreCreatesSchemaAndUpsertsSummary(t *testing.T) { + ctx := context.Background() + store := FromDB(newMemoryDB(t)) + if err := store.initSchema(ctx); err != nil { + t.Fatalf("init schema: %v", err) + } + + createdAt := time.Date(2026, time.July, 17, 8, 0, 0, 123456789, time.UTC) + updatedAt := createdAt.Add(time.Minute + 123*time.Nanosecond) + sandbox := &domain.Sandbox{Summary: domain.SandboxSummary{ + ID: "sandbox-record-id", + ShortID: "sandbox-reco", + Title: "recorded sandbox", + TriggerSource: "api", + Driver: "docker", + VMStatus: domain.VMStatusPending, + GuestImage: "guest:latest", + PullPolicy: "if-not-present", + RuntimeRef: "agent-compose-sandbox-reco", + WorkspacePath: "/data/sandboxes/sandbox-record-id/workspace", + ProxyPath: "/jupyter/sandbox-record-id/lab", + CreatedAt: createdAt, + UpdatedAt: updatedAt, + CellCount: 2, + EventCount: 3, + }} + if err := store.UpsertSandbox(ctx, sandbox); err != nil { + t.Fatalf("upsert sandbox: %v", err) + } + var initialTagsJSON string + if err := store.db.QueryRowContext(ctx, `SELECT tags_json FROM sandboxes WHERE id = ?`, sandbox.Summary.ID).Scan(&initialTagsJSON); err != nil { + t.Fatalf("query initial tags: %v", err) + } + if initialTagsJSON != `[]` { + t.Fatalf("initial tags_json = %q, want []", initialTagsJSON) + } + + originalCreatedAt := sandbox.Summary.CreatedAt + sandbox.Summary.VMStatus = domain.VMStatusRunning + sandbox.Summary.CreatedAt = createdAt.Add(time.Hour) + sandbox.Summary.UpdatedAt = updatedAt.Add(time.Minute + 321*time.Nanosecond) + sandbox.Summary.Tags = []domain.SandboxTag{{Name: "origin", Value: "test"}} + if err := store.UpsertSandbox(ctx, sandbox); err != nil { + t.Fatalf("update sandbox: %v", err) + } + newestUpdatedAt := sandbox.Summary.UpdatedAt + sandbox.Summary.VMStatus = domain.VMStatusStopped + sandbox.Summary.UpdatedAt = updatedAt.Add(30 * time.Second) + if err := store.UpsertSandbox(ctx, sandbox); err != nil { + t.Fatalf("apply stale sandbox update: %v", err) + } + + var status, tagsJSON string + var createdNanos, updatedNanos int64 + err := store.db.QueryRowContext(ctx, `SELECT vm_status, tags_json, created_at, updated_at FROM sandboxes WHERE id = ?`, sandbox.Summary.ID). + Scan(&status, &tagsJSON, &createdNanos, &updatedNanos) + if err != nil { + t.Fatalf("query sandbox: %v", err) + } + if status != domain.VMStatusRunning { + t.Fatalf("status = %q, want %q", status, domain.VMStatusRunning) + } + if tagsJSON != `[{"name":"origin","value":"test"}]` { + t.Fatalf("tags_json = %q", tagsJSON) + } + if createdNanos != originalCreatedAt.UnixNano() || updatedNanos != newestUpdatedAt.UnixNano() { + t.Fatalf("timestamps = (%d, %d), want (%d, %d)", createdNanos, updatedNanos, originalCreatedAt.UnixNano(), newestUpdatedAt.UnixNano()) + } + + loaded, err := store.GetSandboxSummary(ctx, sandbox.Summary.ID) + if err != nil { + t.Fatalf("get sandbox summary: %v", err) + } + if loaded.VMStatus != domain.VMStatusRunning || !loaded.CreatedAt.Equal(originalCreatedAt) || !loaded.UpdatedAt.Equal(newestUpdatedAt) { + t.Fatalf("loaded summary = %#v", loaded) + } +} + +func TestSandboxStoreRejectsInvalidSandbox(t *testing.T) { + store := FromDB(newMemoryDB(t)) + if err := store.initSchema(context.Background()); err != nil { + t.Fatalf("init schema: %v", err) + } + if err := store.UpsertSandbox(context.Background(), nil); err == nil { + t.Fatal("upsert nil sandbox succeeded") + } + if err := store.UpsertSandbox(context.Background(), &domain.Sandbox{}); err == nil { + t.Fatal("upsert sandbox without id succeeded") + } +} + +func TestSandboxStoreListsAndDeletesSummaries(t *testing.T) { + ctx := context.Background() + store := FromDB(newMemoryDB(t)) + if err := store.initSchema(ctx); err != nil { + t.Fatalf("init schema: %v", err) + } + baseTime := time.Date(2026, time.July, 17, 9, 0, 0, 0, time.UTC) + for index, item := range []struct { + id string + driver string + status string + }{ + {id: "sandbox-a", driver: "docker", status: domain.VMStatusRunning}, + {id: "sandbox-b", driver: "docker", status: domain.VMStatusRunning}, + {id: "sandbox-c", driver: "boxlite", status: domain.VMStatusStopped}, + } { + sandbox := &domain.Sandbox{Summary: domain.SandboxSummary{ + ID: item.id, + ShortID: item.id, + Driver: item.driver, + VMStatus: item.status, + CreatedAt: baseTime.Add(time.Duration(index) * time.Nanosecond), + UpdatedAt: baseTime.Add(time.Duration(index) * time.Nanosecond), + }} + if err := store.UpsertSandbox(ctx, sandbox); err != nil { + t.Fatalf("upsert %s: %v", item.id, err) + } + } + + first, err := store.ListSandboxSummaries(ctx, domain.SandboxSummaryListOptions{Driver: "DOCKER", VMStatus: "running", Limit: 1}) + if err != nil { + t.Fatalf("list first page: %v", err) + } + if len(first.Sandboxes) != 1 || first.Sandboxes[0].ID != "sandbox-b" || !first.HasMore { + t.Fatalf("first page = %#v", first) + } + second, err := store.ListSandboxSummaries(ctx, domain.SandboxSummaryListOptions{ + Driver: "docker", + VMStatus: domain.VMStatusRunning, + BeforeUpdatedAt: first.Sandboxes[0].UpdatedAt, + BeforeID: first.Sandboxes[0].ID, + Limit: 1, + }) + if err != nil { + t.Fatalf("list second page: %v", err) + } + if len(second.Sandboxes) != 1 || second.Sandboxes[0].ID != "sandbox-a" || second.HasMore { + t.Fatalf("second page = %#v", second) + } + + if err := store.DeleteSandbox(ctx, "sandbox-b"); err != nil { + t.Fatalf("delete sandbox: %v", err) + } + if _, err := store.GetSandboxSummary(ctx, "sandbox-b"); !errors.Is(err, sql.ErrNoRows) { + t.Fatalf("get deleted sandbox error = %v, want sql.ErrNoRows", err) + } +} + +func TestSandboxStoreListRejectsIncompleteCursor(t *testing.T) { + store := FromDB(newMemoryDB(t)) + if err := store.initSchema(context.Background()); err != nil { + t.Fatalf("init schema: %v", err) + } + + for name, options := range map[string]domain.SandboxSummaryListOptions{ + "missing id": {BeforeUpdatedAt: time.Now().UTC()}, + "missing timestamp": {BeforeID: "sandbox-cursor"}, + } { + t.Run(name, func(t *testing.T) { + if _, err := store.ListSandboxSummaries(context.Background(), options); err == nil { + t.Fatal("list sandboxes accepted an incomplete cursor") + } + }) + } +} diff --git a/pkg/storage/sessionstore/sandbox_record_test.go b/pkg/storage/sessionstore/sandbox_record_test.go new file mode 100644 index 00000000..600d36eb --- /dev/null +++ b/pkg/storage/sessionstore/sandbox_record_test.go @@ -0,0 +1,199 @@ +package sessionstore + +import ( + "context" + "fmt" + "os" + "path/filepath" + "testing" + + appconfig "agent-compose/pkg/config" + driverpkg "agent-compose/pkg/driver" + domain "agent-compose/pkg/model" +) + +type sandboxRecorderStub struct { + recorded map[string]*domain.Sandbox + deleted []string + err error + upsertErrors map[string]error +} + +func (s *sandboxRecorderStub) UpsertSandbox(_ context.Context, sandbox *domain.Sandbox) error { + if s.err != nil { + return s.err + } + if err := s.upsertErrors[sandbox.Summary.ID]; err != nil { + return err + } + copy := *sandbox + copy.Summary = sandbox.Summary + copy.Summary.Tags = append([]domain.SandboxTag(nil), sandbox.Summary.Tags...) + s.recorded[sandbox.Summary.ID] = © + return nil +} + +func (s *sandboxRecorderStub) DeleteSandbox(_ context.Context, id string) error { + if s.err != nil { + return s.err + } + delete(s.recorded, id) + s.deleted = append(s.deleted, id) + return nil +} + +func TestCreateSandboxRecordsNewSandbox(t *testing.T) { + recorder := &sandboxRecorderStub{recorded: make(map[string]*domain.Sandbox)} + store, err := NewWithConfigAndRecorder(sandboxRecordTestConfig(t), recorder) + if err != nil { + t.Fatalf("create store: %v", err) + } + + created, err := store.CreateSandbox(context.Background(), "record me", "", driverpkg.RuntimeDriverBoxlite, "guest:latest", "", "api", nil, nil, []domain.SandboxTag{{Name: "origin", Value: "test"}}) + if err != nil { + t.Fatalf("create sandbox: %v", err) + } + recorded := recorder.recorded[created.Summary.ID] + if recorded == nil { + t.Fatalf("sandbox %q was not recorded", created.Summary.ID) + } + if recorded.Summary.Title != created.Summary.Title || recorded.Summary.VMStatus != domain.VMStatusPending { + t.Fatalf("recorded summary = %#v, want %#v", recorded.Summary, created.Summary) + } + + created.Summary.VMStatus = domain.VMStatusRunning + created.Summary.CellCount = 1 + if err := store.UpdateSandbox(context.Background(), created); err != nil { + t.Fatalf("update sandbox: %v", err) + } + recorded = recorder.recorded[created.Summary.ID] + if recorded.Summary.VMStatus != domain.VMStatusRunning || recorded.Summary.CellCount != 1 { + t.Fatalf("recorded updated summary = %#v", recorded.Summary) + } + if err := store.AddEvent(context.Background(), created.Summary.ID, domain.SandboxEvent{ID: "event-1", Type: "test"}); err != nil { + t.Fatalf("add event: %v", err) + } + recorded = recorder.recorded[created.Summary.ID] + if recorded.Summary.EventCount != 1 { + t.Fatalf("recorded event count = %d, want 1", recorded.Summary.EventCount) + } + if err := store.RemoveSandbox(context.Background(), created.Summary.ID); err != nil { + t.Fatalf("remove sandbox: %v", err) + } + if recorder.recorded[created.Summary.ID] != nil || len(recorder.deleted) != 1 || recorder.deleted[0] != created.Summary.ID { + t.Fatalf("deleted records = %#v, remaining = %#v", recorder.deleted, recorder.recorded) + } +} + +func TestRemoveSandboxRetriesRecordDeletionAfterDirectoryRemoval(t *testing.T) { + recorder := &sandboxRecorderStub{recorded: make(map[string]*domain.Sandbox)} + store, err := NewWithConfigAndRecorder(sandboxRecordTestConfig(t), recorder) + if err != nil { + t.Fatalf("create store: %v", err) + } + created, err := store.CreateSandbox(context.Background(), "retry record deletion", "", driverpkg.RuntimeDriverBoxlite, "", "", "api", nil, nil, nil) + if err != nil { + t.Fatalf("create sandbox: %v", err) + } + + recorder.err = fmt.Errorf("database unavailable") + err = store.RemoveSandbox(context.Background(), created.Summary.ID) + if err == nil { + t.Fatal("remove sandbox returned nil error when record deletion failed") + } + if _, statErr := os.Stat(store.SandboxDir(created.Summary.ID)); !os.IsNotExist(statErr) { + t.Fatalf("sandbox directory stat error = %v, want not exist", statErr) + } + if recorder.recorded[created.Summary.ID] == nil { + t.Fatal("sandbox record was deleted despite recorder failure") + } + + recorder.err = nil + if err := store.RemoveSandbox(context.Background(), created.Summary.ID); err != nil { + t.Fatalf("retry remove sandbox: %v", err) + } + if recorder.recorded[created.Summary.ID] != nil { + t.Fatal("sandbox record remains after retry") + } +} + +func TestMigrateSandboxRecordsScansExistingMetadata(t *testing.T) { + config := sandboxRecordTestConfig(t) + filesystemStore, err := NewWithConfig(config) + if err != nil { + t.Fatalf("create filesystem store: %v", err) + } + existing, err := filesystemStore.CreateSandbox(context.Background(), "existing", "", driverpkg.RuntimeDriverBoxlite, "existing:latest", "", "api", nil, nil, nil) + if err != nil { + t.Fatalf("create existing sandbox: %v", err) + } + + recorder := &sandboxRecorderStub{recorded: make(map[string]*domain.Sandbox)} + if _, err := NewWithConfigAndRecorder(config, recorder); err != nil { + t.Fatalf("create store with migration recorder: %v", err) + } + recorded := recorder.recorded[existing.Summary.ID] + if recorded == nil { + t.Fatalf("existing sandbox %q was not migrated", existing.Summary.ID) + } + if recorded.Summary.GuestImage != "existing:latest" { + t.Fatalf("migrated image = %q, want existing:latest", recorded.Summary.GuestImage) + } +} + +func TestMigrateSandboxRecordsContinuesAfterUpsertFailure(t *testing.T) { + config := sandboxRecordTestConfig(t) + filesystemStore, err := NewWithConfig(config) + if err != nil { + t.Fatalf("create filesystem store: %v", err) + } + first, err := filesystemStore.CreateSandbox(context.Background(), "first", "", driverpkg.RuntimeDriverBoxlite, "", "", "api", nil, nil, nil) + if err != nil { + t.Fatalf("create first sandbox: %v", err) + } + second, err := filesystemStore.CreateSandbox(context.Background(), "second", "", driverpkg.RuntimeDriverBoxlite, "", "", "api", nil, nil, nil) + if err != nil { + t.Fatalf("create second sandbox: %v", err) + } + recorder := &sandboxRecorderStub{ + recorded: make(map[string]*domain.Sandbox), + upsertErrors: map[string]error{first.Summary.ID: fmt.Errorf("invalid record")}, + } + if _, err := NewWithConfigAndRecorder(config, recorder); err != nil { + t.Fatalf("migration blocked store initialization: %v", err) + } + if recorder.recorded[second.Summary.ID] == nil { + t.Fatalf("second sandbox %q was not migrated after first failure", second.Summary.ID) + } +} + +func TestCreateSandboxSucceedsWhenRecorderFails(t *testing.T) { + recorder := &sandboxRecorderStub{recorded: make(map[string]*domain.Sandbox), err: fmt.Errorf("database unavailable")} + store, err := NewWithConfigAndRecorder(sandboxRecordTestConfig(t), recorder) + if err != nil { + t.Fatalf("create store: %v", err) + } + created, err := store.CreateSandbox(context.Background(), "record failure", "", driverpkg.RuntimeDriverBoxlite, "", "", "api", nil, nil, nil) + if err != nil { + t.Fatalf("create sandbox: %v", err) + } + if _, err := store.GetSandbox(context.Background(), created.Summary.ID); err != nil { + t.Fatalf("load filesystem sandbox after recorder failure: %v", err) + } + if recorder.recorded[created.Summary.ID] != nil { + t.Fatal("failed recorder unexpectedly stored sandbox") + } +} + +func sandboxRecordTestConfig(t *testing.T) *appconfig.Config { + t.Helper() + root := t.TempDir() + return &appconfig.Config{ + SandboxRoot: filepath.Join(root, "sandboxes"), + RuntimeDriver: driverpkg.RuntimeDriverBoxlite, + DefaultImage: "default:latest", + MicrosandboxHome: filepath.Join(root, "microsandbox"), + JupyterGuestPort: 8888, + JupyterProxyBasePath: "/jupyter", + } +} diff --git a/pkg/storage/sessionstore/store.go b/pkg/storage/sessionstore/store.go index 4fb464c7..6dfc1207 100644 --- a/pkg/storage/sessionstore/store.go +++ b/pkg/storage/sessionstore/store.go @@ -9,6 +9,7 @@ import ( "fmt" "io" "io/fs" + "log/slog" "net" "os" "path/filepath" @@ -54,6 +55,7 @@ type ( type Store struct { config *appconfig.Config + recorder SandboxRecorder sandboxLocks sync.Map cacheDependencyMu sync.RWMutex cacheDependencyLocker CacheDependencyLocker @@ -63,19 +65,35 @@ type CacheDependencyLocker interface { WithLockContext(context.Context, func() error) error } +// SandboxRecorder stores the queryable record for a persisted Sandbox. +type SandboxRecorder interface { + UpsertSandbox(context.Context, *domain.Sandbox) error + DeleteSandbox(context.Context, string) error +} + func NewStore(di do.Injector) (*Store, error) { config := do.MustInvoke[*appconfig.Config](di) - return NewWithConfig(config) + recorder := do.MustInvoke[SandboxRecorder](di) + return NewWithConfigAndRecorder(config, recorder) } func NewWithConfig(config *appconfig.Config) (*Store, error) { + return NewWithConfigAndRecorder(config, nil) +} + +// NewWithConfigAndRecorder constructs a filesystem store that also records +// persisted Sandboxes for database queries. +func NewWithConfigAndRecorder(config *appconfig.Config, recorder SandboxRecorder) (*Store, error) { if err := os.MkdirAll(config.SandboxRoot, 0o755); err != nil { return nil, fmt.Errorf("create sandbox root: %w", err) } - store := &Store{config: config} + store := &Store{config: config, recorder: recorder} if err := store.backfillOwnershipRecords(); err != nil { return nil, err } + if err := store.MigrateSandboxRecords(context.Background()); err != nil { + return nil, err + } return store, nil } @@ -111,18 +129,18 @@ func (s *Store) createSandboxWithCacheDependencyLock(ctx context.Context, title, locker := s.cacheDependencyLocker s.cacheDependencyMu.RUnlock() if locker == nil { - return s.createSandboxWithOptions(title, baseWorkspace, driver, guestImage, workspaceID, triggerSource, workspace, envItems, tags, options) + return s.createSandboxWithOptions(ctx, title, baseWorkspace, driver, guestImage, workspaceID, triggerSource, workspace, envItems, tags, options) } var sandbox *Sandbox err := locker.WithLockContext(ctx, func() error { var err error - sandbox, err = s.createSandboxWithOptions(title, baseWorkspace, driver, guestImage, workspaceID, triggerSource, workspace, envItems, tags, options) + sandbox, err = s.createSandboxWithOptions(ctx, title, baseWorkspace, driver, guestImage, workspaceID, triggerSource, workspace, envItems, tags, options) return err }) return sandbox, err } -func (s *Store) createSandboxWithOptions(title, baseWorkspace, driver, guestImage, workspaceID, triggerSource string, workspace *SandboxWorkspace, envItems []SandboxEnvVar, tags []SandboxTag, options CreateSandboxOptions) (*Sandbox, error) { +func (s *Store) createSandboxWithOptions(ctx context.Context, title, baseWorkspace, driver, guestImage, workspaceID, triggerSource string, workspace *SandboxWorkspace, envItems []SandboxEnvVar, tags []SandboxTag, options CreateSandboxOptions) (*Sandbox, error) { now := time.Now().UTC() workspaceID = strings.TrimSpace(workspaceID) id := identity.NewRandomID(identity.ResourceSandbox) @@ -250,6 +268,11 @@ func (s *Store) createSandboxWithOptions(title, baseWorkspace, driver, guestImag }); err != nil { return nil, fmt.Errorf("write sandbox ownership record: %w", err) } + if s.recorder != nil { + if err := s.recorder.UpsertSandbox(ctx, session); err != nil { + slog.Warn("sandbox created but database record could not be written", "sandbox_id", session.Summary.ID, "error", err) + } + } return session, nil } @@ -260,6 +283,33 @@ func (s *Store) SetCacheDependencyLocker(locker CacheDependencyLocker) { s.cacheDependencyLocker = locker } +// MigrateSandboxRecords records existing filesystem Sandboxes in the database. +// It is safe to run repeatedly because the recorder performs an upsert. +func (s *Store) MigrateSandboxRecords(ctx context.Context) error { + if s.recorder == nil { + return nil + } + entries, err := os.ReadDir(s.config.SandboxRoot) + if err != nil { + return fmt.Errorf("read sandbox root for database migration: %w", err) + } + for _, entry := range entries { + if !entry.IsDir() || entry.Name() == ".lifecycle" { + continue + } + sandbox, err := s.loadSandbox(entry.Name()) + if err != nil { + slog.Warn("skip unrecognized sandbox during database migration", "directory", entry.Name(), "error", err) + continue + } + s.hydrateSandboxGuestImage(sandbox) + if err := s.recorder.UpsertSandbox(ctx, sandbox); err != nil { + slog.Warn("failed to record sandbox during database migration", "sandbox_id", sandbox.Summary.ID, "error", err) + } + } + return nil +} + func (s *Store) GetSandbox(_ context.Context, id string) (*Sandbox, error) { session, err := s.loadSandbox(strings.TrimSpace(id)) if err != nil { @@ -313,34 +363,61 @@ func (s *Store) ListSandboxes(_ context.Context, options SandboxListOptions) (Sa return result, nil } -func (s *Store) UpdateSandbox(_ context.Context, session *Sandbox) error { +func (s *Store) UpdateSandbox(ctx context.Context, session *Sandbox) error { s.hydrateSandboxGuestImage(session) - session.Summary.UpdatedAt = time.Now().UTC() unlock := s.lockSandbox(session.Summary.ID) - defer unlock() - return s.saveSandboxPreservingCounts(session) + // Assign the version timestamp while holding the filesystem lock so it + // reflects the serialization order used to persist metadata.json. + session.Summary.UpdatedAt = time.Now().UTC() + err := s.saveSandboxPreservingCounts(session) + unlock() + if err != nil { + return err + } + if s.recorder != nil { + if err := s.recorder.UpsertSandbox(ctx, session); err != nil { + return fmt.Errorf("record sandbox update: %w", err) + } + } + return nil } -func (s *Store) RemoveSandbox(_ context.Context, id string) error { +func (s *Store) RemoveSandbox(ctx context.Context, id string) error { id = strings.TrimSpace(id) if err := validateSandboxIDForRemove(id); err != nil { return err } path := s.sandboxDir(id) if _, err := os.Stat(path); err != nil { + if errors.Is(err, fs.ErrNotExist) && s.recorder != nil { + if err := s.recorder.DeleteSandbox(ctx, id); err != nil { + return fmt.Errorf("delete sandbox record after directory removal: %w", err) + } + return nil + } return fmt.Errorf("stat sandbox dir %s: %w", id, err) } - unlock := s.lockSandbox(id) - defer unlock() + if err := func() error { + unlock := s.lockSandbox(id) + defer unlock() - if err := driverpkg.CleanupBoxliteVolumeBridgeMounts(path); err != nil { - return fmt.Errorf("cleanup session mounts %s: %w", id, err) - } - if err := os.RemoveAll(path); err != nil { - return fmt.Errorf("remove sandbox dir %s: %w", id, err) + if err := driverpkg.CleanupBoxliteVolumeBridgeMounts(path); err != nil { + return fmt.Errorf("cleanup session mounts %s: %w", id, err) + } + if err := os.RemoveAll(path); err != nil { + return fmt.Errorf("remove sandbox dir %s: %w", id, err) + } + return nil + }(); err != nil { + return err } s.sandboxLocks.Delete(sandboxLockKey(id)) + if s.recorder != nil { + if err := s.recorder.DeleteSandbox(ctx, id); err != nil { + return fmt.Errorf("delete sandbox record: %w", err) + } + } return nil } @@ -443,29 +520,33 @@ func (s *Store) AddAgentRun(_ context.Context, sessionID string, run AgentRun) e return nil } -func (s *Store) AddEvent(_ context.Context, sessionID string, event SandboxEvent) error { - unlock := s.lockSandbox(sessionID) - defer unlock() +func (s *Store) AddEvent(ctx context.Context, sessionID string, event SandboxEvent) error { + var sandboxToRecord *Sandbox + if err := func() error { + unlock := s.lockSandbox(sessionID) + defer unlock() - jsonlExisted, err := s.eventsJSONLExists(sessionID) - if err != nil { - return err - } - legacyCount := 0 - if !jsonlExisted { - events, err := s.loadLegacyEvents(sessionID) + jsonlExisted, err := s.eventsJSONLExists(sessionID) if err != nil { return err } - legacyCount = len(events) - } + legacyCount := 0 + if !jsonlExisted { + events, err := s.loadLegacyEvents(sessionID) + if err != nil { + return err + } + legacyCount = len(events) + } - if err := s.appendEvent(sessionID, event); err != nil { - return err - } + if err := s.appendEvent(sessionID, event); err != nil { + return err + } - session, err := s.loadSandbox(sessionID) - if err == nil { + session, err := s.loadSandbox(sessionID) + if err != nil { + return err + } nextCount := session.Summary.EventCount + 1 if !jsonlExisted && legacyCount >= session.Summary.EventCount { nextCount = legacyCount + 1 @@ -473,7 +554,18 @@ func (s *Store) AddEvent(_ context.Context, sessionID string, event SandboxEvent session.Summary.EventCount = nextCount s.hydrateSandboxGuestImage(session) session.Summary.UpdatedAt = time.Now().UTC() - _ = s.saveSandboxPreservingCounts(session) + if err := s.saveSandboxPreservingCounts(session); err != nil { + return err + } + sandboxToRecord = session + return nil + }(); err != nil { + return err + } + if s.recorder != nil { + if err := s.recorder.UpsertSandbox(ctx, sandboxToRecord); err != nil { + return fmt.Errorf("record sandbox event count: %w", err) + } } return nil }