orchestrator

type AgentSpec struct {
	Files        map[string][]byte // arbitrary files to upload into the guest (path -> bytes)
	DownloadURL  string            // fetch the agent binary (e.g. GitHub releases)
	Command      string            // how to launch the daemon
	AuthToken    string            // pass-through bearer token
	Gateway      string            // pass-through gateway websocket URL
	GatewayToken string            // pass-through gateway token
	DrainCommand string            // command run inside the guest for graceful shutdown (” => skip)
	Expose       []ExposeSpec      // guest ports to publish as reachable endpoints, if any
}

type AttachSpec struct {
	// Cmd is the argv to run. Empty means the guest's login shell.
	Cmd []string

	// TTY requests a pty for the process. Rows and Cols seed its initial
	// window size; further resizes arrive as frames on the stream itself.
	TTY  bool
	Rows uint16
	Cols uint16
}

type Attacher interface {
	Attach(ctx context.Context, spec AttachSpec) (io.ReadWriteCloser, error)
}

type BootOptions struct {
	StartupScript string
	GatewayURL    string
	GatewayToken  string
	Expose        []ExposeSpec
}

type BootResult struct {
	Env                Environment
	BootTime           time.Duration
	FromCache          bool       // true if restored from checkpoint
	AuthTokenEncrypted []byte     // AES-GCM encrypted per-VM auth token for persistence
	DrainCommand       string     // graceful-shutdown command for the configured agent (” => skip)
	Endpoints          []Endpoint // published endpoints, if the provider reported any
}

type CapacityProber interface {
	Capacity(ctx context.Context) (HostCapacity, error)
}

type Checkpoint struct {
	ID        string
	Comment   string
	SizeBytes int64
	CreatedAt time.Time
}

type DeadLetterKind string

type DeadLetterRecord struct {
	ID          int64
	Kind        DeadLetterKind
	EntityID    string
	TaskID      string
	Reason      string
	RetryCount  int
	Payload     json.RawMessage
	FirstSeenAt time.Time
	LastSeenAt  time.Time
}

type Endpoint struct {
	As   string // caller-chosen label, e.g. "http"
	URL  string // reachable address, e.g. "http://203.0.113.5:41231"
	Port int    // the guest-side port this endpoint publishes
}

type EndpointReporter interface {
	Endpoints() []Endpoint
}

type Environment interface {
	// Name returns the sandbox identifier.
	Name() string

	// URL returns the sandbox's reachable address for the guest agent.
	URL() string

	// Exec runs argv inside the sandbox and reports the result. A non-zero
	// guest exit code is carried in ExecResult.ExitCode, not returned as an
	// error; err is reserved for transport and provider failures, so callers
	// must check both.
	Exec(ctx context.Context, cmd []string, opts ExecOptions) (ExecResult, error)

	// ExecStream runs a command with stdout/stderr wired to writers.
	ExecStream(ctx context.Context, stdout, stderr io.Writer, name string, args ...string) error

	// Upload writes data to a file path inside the sandbox.
	Upload(ctx context.Context, data []byte, path string) error

	// StartAgent launches the configured guest agent inside the sandbox.
	StartAgent(ctx context.Context, spec AgentSpec) error

	// Token returns the per-sandbox auth token that callers must include
	// when reaching URL(). Empty string for providers whose URL
	// self-authenticates.
	Token() string
}

type EnvironmentEvent struct {
	ID        string    `json:"id"`
	Kind      string    `json:"event"`
	VMID      string    `json:"vm_id"`
	State     VMState   `json:"state"`
	URL       string    `json:"url,omitempty"`
	Error     string    `json:"error,omitempty"`
	UpdatedAt time.Time `json:"updated_at"`
}

type EventRecord struct {
	ID         int64
	EntityType string
	EntityID   string
	EventType  string
	Payload    json.RawMessage
	CreatedAt  time.Time
}

type ExecOptions struct {
	// Timeout bounds how long the guest command may run. Zero means
	// DefaultExecTimeout. Values above MaxExecTimeout are clamped.
	Timeout time.Duration
}

type ExecResult struct {
	ExitCode int
	Stdout   []byte
	Stderr   []byte
}

type ExposeSpec struct {
	Port int
	As   string
}

type FleetConfig struct {
	Provider          Provider
	StateStore        StateStore
	Prefix            string        // VM name prefix, e.g. "fuse-"
	ReconcileInterval time.Duration // default 30s

	// TaskStuckTimeout is the maximum age of a Running VM with no state
	// transitions before it is considered stuck. This is a leak-detection
	// ceiling, NOT a heartbeat — healthy long-running tasks must set
	// Spec.MaxRuntime to override this. Default 2h.
	TaskStuckTimeout time.Duration

	// OrphanDestroyMaxRetries is the number of consecutive reconcile cycles
	// an orphan VM may fail to destroy before being dead-lettered and
	// skipped on subsequent cycles. Default 5.
	OrphanDestroyMaxRetries int

	// DefaultSnapshotRetention applies to snapshots created without an
	// explicit retention window. Zero leaves snapshots unbounded until
	// explicitly deleted.
	DefaultSnapshotRetention time.Duration

	// SnapshotQuotaMaxCount limits ready/in-flight snapshots per tenant.
	// Zero disables count-based enforcement.
	SnapshotQuotaMaxCount int

	// SnapshotQuotaMaxBytes limits the aggregate size of ready/in-flight
	// snapshots per tenant. Zero disables byte-based enforcement.
	SnapshotQuotaMaxBytes int64

	// PlacementPolicy controls how the scheduler picks among hosts
	// when multiple are registered. Default is spread.
	PlacementPolicy PlacementPolicy

	// TokenEncryptionKey is the 32-byte AES-256 key used to encrypt
	// per-VM auth tokens before storing them in the state store.
	// When nil, VM credential generation is skipped (insecure mode).
	TokenEncryptionKey []byte

	// HostProviderFactory builds a Provider for a given (url, token,
	// backend) triple. The fleet uses it during recoverState to
	// rehydrate the per-host providers that RegisterHost would have
	// created during normal operation. Without it, an orchestrator
	// restart loses the scheduler's host registry and falls back to
	// single-provider mode for all subsequent placements.
	HostProviderFactory func(url, token string, backend HostBackend) Provider

	Metrics ReconcileMetrics
	Logger  *slog.Logger
}

type FleetManager struct {
	// contains filtered or unexported fields
}

type ForkOptions struct {
	// Comment is attached to the seed snapshot when ForkEnvironment
	// creates one (empty ReuseSnapshotID). ignored when reusing an
	// existing snapshot.
	Comment string

	// ReuseSnapshotID selects an existing ready snapshot of the source
	// vm to seed the fork from. empty means snapshot the source first.
	ReuseSnapshotID string
}

type Host struct {
	ID        string
	URL       string // base URL of the host agent (e.g. https://agent-1.local)
	Token     string // bearer token for this host's agent
	Region    string
	Backend   HostBackend // "firecracker" or "qemu"; empty means firecracker (default)
	Capacity  HostCapacity
	Allocated HostCapacity
	State     HostState
	LastSeen  time.Time
	CreatedAt time.Time
	UpdatedAt time.Time
}

type HostBackend string

type HostCapacity struct {
	CPUs      int `json:"cpus"`
	RamMB     int `json:"ram_mb"`
	StorageGB int `json:"storage_gb"`
	VMCount   int `json:"vm_count"` // max concurrent VMs

	// GPUs is the count of whole GPU devices available on the host.
	// Zero means no GPUs. Only qemu-backed hosts may report GPUs > 0
	// (enforced at registration, see internal/api registerHost).
	GPUs int `json:"gpus,omitempty"`

	// GPUKind identifies the GPU model (e.g. "a100"). Empty when GPUs is 0.
	GPUKind string `json:"gpu_kind,omitempty"`
}

type HostRecord struct {
	ID             string
	URL            string
	TokenEncrypted []byte
	Region         string
	State          HostState
	TenantID       string
	Backend        HostBackend
	Capacity       HostCapacity
	Allocated      HostCapacity
	LastSeen       time.Time
	CreatedAt      time.Time
	UpdatedAt      time.Time
}

type HostState string

type MemoryStateStore struct {
	// contains filtered or unexported fields
}

type PlacementDecision struct {
	HostID       string          `json:"host_id"`
	Policy       PlacementPolicy `json:"policy"`
	Candidates   int             `json:"candidates"`    // eligible hosts considered
	HeadroomCPUs int             `json:"headroom_cpus"` // CPUs remaining after placement
	HeadroomRam  int             `json:"headroom_ram"`  // RAM remaining after placement
}

type PlacementPolicy string

type PostgresStateStore struct {
	// contains filtered or unexported fields
}

type Provider interface {
	// Create provisions a new sandbox.
	Create(ctx context.Context, spec Spec) (Environment, error)

	// Get returns a handle to an existing sandbox by name.
	Get(ctx context.Context, name string) (Environment, error)

	// Destroy tears down a sandbox.
	Destroy(ctx context.Context, name string) error

	// List returns all sandboxes matching the given prefix.
	List(ctx context.Context, prefix string) ([]Environment, error)

	// Close releases provider resources.
	Close() error
}

type ReconcileMetrics interface {
	ReconcileCompleted(summary ReconcileSummary)
}

type ReconcileSummary struct {
	TrackedVMs          int
	ProviderVMs         int
	OrphansDestroyed    int
	OrphansFailed       int
	OrphansDeadLettered int
	StuckTasksSuspected int
	StuckTasksFailed    int
	VMsMissingProvider  int
	Duration            time.Duration
}

type SnapshotCapable interface {
	Checkpoint(ctx context.Context, comment string) (string, error)
	Restore(ctx context.Context, checkpointID string) error
	ListCheckpoints(ctx context.Context) ([]Checkpoint, error)
}

type SnapshotDeleter interface {
	DeleteCheckpoint(ctx context.Context, checkpointID string) error
}

type SnapshotExportRecord struct {
	Destination string
	Status      SnapshotExportStatus
	RequestedAt time.Time
	UpdatedAt   time.Time
	LastError   string
}

type SnapshotExportStatus string

type SnapshotFilter struct {
	VMID     string
	TaskID   string
	TenantID string
	State    SnapshotState
}

type SnapshotForkable interface {
	CreateFromCheckpoint(ctx context.Context, spec Spec, srcVMID, checkpointID string) (Environment, error)
}

type SnapshotMode string

type SnapshotOptions struct {
	// Comment is a free-form note attached to the snapshot. Passed
	// through to the underlying Environment.Checkpoint implementation
	// and persisted in the snapshot record.
	Comment string

	// Mode records whether this snapshot was taken manually (via the
	// REST API) or by an automated process. Defaults to
	// SnapshotModeManual when empty.
	Mode SnapshotMode

	// RetentionUntil, if non-nil, records when the snapshot is
	// eligible for garbage collection.
	RetentionUntil *time.Time

	// Metadata augments the persisted metadata blob. Reserved fields
	// such as "comment" are set by the orchestrator.
	Metadata map[string]string

	// Exports records optional object-storage export intents/status.
	Exports []SnapshotExportRecord
}

type SnapshotRecord struct {
	SnapshotID       string
	VMID             string
	TaskID           string
	HostID           string
	TenantID         string
	ParentSnapshotID string
	Mode             SnapshotMode
	State            SnapshotState
	SizeBytes        int64
	RetentionUntil   *time.Time
	Metadata         json.RawMessage
	Exports          []SnapshotExportRecord
	LastError        string
	CreatedAt        time.Time
	UpdatedAt        time.Time
}

type SnapshotState string

type Spec struct {
	Name      string // e.g. "fuse-{task-id}"
	CPUs      int
	RamMB     int
	StorageGB int
	Region    string

	// GPUs is the count of whole GPU devices requested. Zero means no GPU.
	GPUs int32

	// GPUKind identifies the requested GPU model (e.g. "a100"). Empty when
	// GPUs is 0.
	GPUKind string

	// MaxRuntime overrides FleetConfig.TaskStuckTimeout for this task.
	// Zero means "use the fleet default". This is a leak-detection ceiling,
	// not a liveness check — set it higher than any plausible healthy runtime.
	MaxRuntime time.Duration

	// Image names a base rootfs for the provider to boot from, resolved by
	// the provider (e.g. the firecracker host agent looks it up in its own
	// named-rootfs directory). Empty means the provider's default base.
	Image string
}

type StateStore interface {
	UpsertVM(ctx context.Context, vm VMRecord) error
	DeleteVM(ctx context.Context, vmID string) error
	ListVMs(ctx context.Context) ([]VMRecord, error)

	UpsertTask(ctx context.Context, task TaskRecord) error
	ListTasks(ctx context.Context) ([]TaskRecord, error)

	UpsertSnapshot(ctx context.Context, snapshot SnapshotRecord) error
	GetSnapshot(ctx context.Context, snapshotID string) (SnapshotRecord, error)
	ListSnapshots(ctx context.Context) ([]SnapshotRecord, error)
	DeleteSnapshot(ctx context.Context, snapshotID string) error

	AppendEvent(ctx context.Context, event EventRecord) error

	// UpsertDeadLetter inserts or updates a dead-letter entry keyed by
	// (Kind, EntityID). On update, RetryCount and LastSeenAt are advanced.
	UpsertDeadLetter(ctx context.Context, entry DeadLetterRecord) error

	// ListDeadLetters returns all dead-letter entries. Implementations
	// may order arbitrarily.
	ListDeadLetters(ctx context.Context) ([]DeadLetterRecord, error)

	// UpsertHost inserts or updates a host registration.
	UpsertHost(ctx context.Context, host HostRecord) error

	// DeleteHost removes a host from the registry. No-op if absent.
	DeleteHost(ctx context.Context, hostID string) error

	// ListHosts returns all registered hosts.
	ListHosts(ctx context.Context) ([]HostRecord, error)

	// GetHost returns a single host by ID, or an error if not found.
	GetHost(ctx context.Context, hostID string) (HostRecord, error)
}

type TaskRecord struct {
	TaskID     string
	VMID       string
	RunStatus  TaskRunStatus
	RetryCount int
	LastError  string
	AssignedAt time.Time
	UpdatedAt  time.Time
}

type TaskRunStatus string

type TokenSetter interface {
	SetToken(token string)
}

type VMFilter struct {
	TaskID string
	State  VMState
	HostID string
}

type VMInfo struct {
	ID        string
	State     VMState
	TaskID    string
	HostID    string
	URL       string
	Spec      Spec
	CreatedAt time.Time
	UpdatedAt time.Time
	Error     string
	Endpoints []Endpoint
}

type VMRecord struct {
	ID                 string
	HostID             string
	NetworkHost        string
	State              VMState
	URL                string
	TaskID             string
	TenantID           string
	Spec               Spec
	LastError          string
	AuthTokenEncrypted []byte     // AES-GCM encrypted per-VM auth token (nil for legacy VMs)
	SecretsEncrypted   []byte     // AES-GCM encrypted JSON of the secret map (nil when no secrets supplied)
	Endpoints          []Endpoint // published endpoints (e.g. ingress), if any
	CreatedAt          time.Time
	UpdatedAt          time.Time
}

type VMState string