molecule-cli/internal/cmd/workspace_mgmt.go

// Package cmd implements the CLI command tree.
package cmd

import (
	"fmt"

	"github.com/spf13/cobra"
	"go.moleculesai.app/cli/internal/client"
)

// ---------------------------------------------------------------------------
// Workspace management verbs (PLATFORM-MANAGEMENT-API.md §5(b)):
//   get, pause, resume, budget, billing-mode, token mint.
// list / create / delete / restart / inspect live in workspace.go.
// ---------------------------------------------------------------------------

// ===========================================================================
// molecule workspace get <id>  (alias of inspect)
// ===========================================================================
var workspaceGetCmd = &cobra.Command{
	Use:   "get <workspace-id>",
	Short: "Show full details for a workspace (alias of inspect)",
	Args:  cobra.ExactArgs(1),
	RunE:  runWorkspaceInspect,
}

// ===========================================================================
// molecule workspace pause <id>
// ===========================================================================
var workspacePauseCmd = &cobra.Command{
	Use:   "pause <workspace-id>",
	Short: "Pause a workspace",
	Args:  cobra.ExactArgs(1),
	RunE:  runWorkspacePause,
}

func runWorkspacePause(_ *cobra.Command, args []string) error {
	if err := newClient().PauseWorkspace(args[0]); err != nil {
		return fmt.Errorf("workspace pause: %w", err)
	}
	fmt.Printf("Pause triggered for workspace %q.\n", args[0])
	return nil
}

// ===========================================================================
// molecule workspace resume <id>
// ===========================================================================
var workspaceResumeCmd = &cobra.Command{
	Use:   "resume <workspace-id>",
	Short: "Resume a paused workspace",
	Args:  cobra.ExactArgs(1),
	RunE:  runWorkspaceResume,
}

func runWorkspaceResume(_ *cobra.Command, args []string) error {
	if err := newClient().ResumeWorkspace(args[0]); err != nil {
		return fmt.Errorf("workspace resume: %w", err)
	}
	fmt.Printf("Resume triggered for workspace %q.\n", args[0])
	return nil
}

// ===========================================================================
// molecule workspace budget <id> [--hourly|--daily|--weekly|--monthly cents]
// ===========================================================================
var budgetFlags struct {
	hourly  int64
	daily   int64
	weekly  int64
	monthly int64
}

var workspaceBudgetCmd = &cobra.Command{
	Use:   "budget <workspace-id> [--hourly|--daily|--weekly|--monthly <cents>]",
	Short: "Show or set a workspace's LLM budget (USD cents per period)",
	Long: `With no period flags, prints the current budget.
With one or more period flags, sets those period limits (USD cents).
Pass a period flag with value -1 to CLEAR that period's limit.

  molecule workspace budget ws_123                 # show
  molecule workspace budget ws_123 --monthly 50000 # $500/mo
  molecule workspace budget ws_123 --daily -1      # clear daily limit`,
	Args: cobra.ExactArgs(1),
	RunE: runWorkspaceBudget,
}

func init() {
	f := workspaceBudgetCmd.Flags()
	// Sentinel: math.MinInt64 means "flag not provided". -1 means "clear".
	f.Int64Var(&budgetFlags.hourly, "hourly", budgetUnset, "Hourly limit (USD cents); -1 to clear")
	f.Int64Var(&budgetFlags.daily, "daily", budgetUnset, "Daily limit (USD cents); -1 to clear")
	f.Int64Var(&budgetFlags.weekly, "weekly", budgetUnset, "Weekly limit (USD cents); -1 to clear")
	f.Int64Var(&budgetFlags.monthly, "monthly", budgetUnset, "Monthly limit (USD cents); -1 to clear")
}

const budgetUnset = -1 << 62

// budgetLimitsFromFlags translates the four --hourly/--daily/--weekly/--monthly
// flags into the budget_limits map sent to PATCH /workspaces/:id/budget:
// an unset flag (budgetUnset) is omitted; a negative value clears that period
// (nil); a non-negative value sets the limit. Shared by prod and tests so the
// translation logic has exactly one definition.
func budgetLimitsFromFlags(hourly, daily, weekly, monthly int64) map[string]*int64 {
	limits := map[string]*int64{}
	add := func(period string, v int64) {
		if v == budgetUnset {
			return
		}
		if v < 0 {
			limits[period] = nil // clear
			return
		}
		vv := v
		limits[period] = &vv
	}
	add("hourly", hourly)
	add("daily", daily)
	add("weekly", weekly)
	add("monthly", monthly)
	return limits
}

// resolveBillingMode validates a billing-mode argument and returns the value
// to send to the tenant: the three real modes pass through; clear/null/"" map
// to "" (clear the override); anything else is an error. Shared by prod and
// tests.
func resolveBillingMode(in string) (string, error) {
	switch in {
	case "platform_managed", "byok", "disabled":
		return in, nil
	case "clear", "null", "":
		return "", nil
	default:
		return "", &exitError{code: 2, msg: "workspace billing-mode: mode must be platform_managed, byok, disabled, or clear"}
	}
}

func runWorkspaceBudget(cmd *cobra.Command, args []string) error {
	cl := newClient()
	id := args[0]

	limits := budgetLimitsFromFlags(
		budgetFlags.hourly, budgetFlags.daily, budgetFlags.weekly, budgetFlags.monthly,
	)

	if len(limits) == 0 {
		// Show current budget.
		raw, err := cl.GetBudget(id)
		if err != nil {
			return fmt.Errorf("workspace budget: %w", err)
		}
		return printRaw(raw)
	}
	raw, err := cl.SetBudget(id, limits)
	if err != nil {
		return fmt.Errorf("workspace budget: %w", err)
	}
	return printRaw(raw)
}

// ===========================================================================
// molecule workspace billing-mode <id> <mode|clear>
// ===========================================================================
var workspaceBillingModeCmd = &cobra.Command{
	Use:   "billing-mode <workspace-id> <platform_managed|byok|disabled|clear>",
	Short: "Set a workspace's LLM billing-mode override",
	Args:  cobra.ExactArgs(2),
	RunE:  runWorkspaceBillingMode,
}

func runWorkspaceBillingMode(_ *cobra.Command, args []string) error {
	id := args[0]
	mode, err := resolveBillingMode(args[1])
	if err != nil {
		return err
	}
	raw, err := newClient().SetBillingMode(id, mode)
	if err != nil {
		return fmt.Errorf("workspace billing-mode: %w", err)
	}
	return printRaw(raw)
}

// ===========================================================================
// molecule workspace migrate-provider <id> --to <p> [--from <p>]
// molecule workspace migration-status <id>
//
// Cross-cloud compute-provider migration (AWS ↔ Hetzner ↔ GCP). These hit the
// CONTROL-PLANE admin surface, so they use cpAdminClient() (CP-admin bearer +
// MOLECULE_CP_URL), NOT the tenant newClient() — the tenant Org API Key has no
// standing on the CP.
// ===========================================================================

var migrateFlags struct {
	to             string
	from           string
	fromInstanceID string
	orgID          string
	runtime        string
	confirm        bool
}

var workspaceMigrateProviderCmd = &cobra.Command{
	Use:   "migrate-provider <workspace-id> --to <aws|hetzner|gcp> [--from <provider>]",
	Short: "Migrate a workspace's compute box across clouds (AWS/Hetzner/GCP)",
	Long: `Start a data-safe cross-cloud provider migration for a workspace.

The control plane snapshots the source's /workspace to R2, provisions the
target (which restores on boot), verifies it's healthy, then retires the
source (verify-before-destroy + rollback live in the CP). It runs
asynchronously (~15-20 min) — poll 'molecule workspace migration-status'.

This hits the control-plane admin surface and requires MOLECULE_CP_ADMIN_TOKEN
and MOLECULE_CP_URL (distinct from the tenant MOLECULE_API_KEY).

  molecule workspace migrate-provider ws_123 --to hetzner --from aws --confirm
  # non-AWS sources need the current box id:
  molecule workspace migrate-provider ws_123 --to aws --from gcp \
    --from-instance-id gcp-box-9 --confirm

--confirm is REQUIRED: a real migration mutates two clouds.`,
	Args: cobra.ExactArgs(1),
	RunE: runWorkspaceMigrateProvider,
}

var workspaceMigrationStatusCmd = &cobra.Command{
	Use:   "migration-status <workspace-id>",
	Short: "Show the latest cross-cloud provider-migration status for a workspace",
	Long: `Read the most recent provider-migration record for a workspace
(state, from/to provider, detail, and whether it has reached a terminal state).

States: snapshotting → provisioning_target → target_healthy → retiring_source
→ completed (terminal also: failed, rolled_back).

Hits the control-plane admin surface — requires MOLECULE_CP_ADMIN_TOKEN and
MOLECULE_CP_URL.`,
	Args: cobra.ExactArgs(1),
	RunE: runWorkspaceMigrationStatus,
}

func init() {
	f := workspaceMigrateProviderCmd.Flags()
	f.StringVar(&migrateFlags.to, "to", "", "Target provider: aws|hetzner|gcp (required)")
	f.StringVar(&migrateFlags.from, "from", "", "Current provider: aws|hetzner|gcp (required)")
	f.StringVar(&migrateFlags.fromInstanceID, "from-instance-id", "", "Current box id to snapshot + retire (required for non-AWS sources)")
	f.StringVar(&migrateFlags.orgID, "org-id", "", "Org hint for non-AWS sources (usually unnecessary)")
	f.StringVar(&migrateFlags.runtime, "runtime", "", "Runtime hint for non-AWS sources (usually unnecessary)")
	f.BoolVar(&migrateFlags.confirm, "confirm", false, "Confirm the migration (REQUIRED — mutates two clouds)")
	workspaceMigrateProviderCmd.MarkFlagRequired("to")
}

// validMigrationProvider reports whether p is one of the three supported
// compute providers. Shared by the command and its tests.
func validMigrationProvider(p string) bool {
	switch p {
	case "aws", "hetzner", "gcp":
		return true
	default:
		return false
	}
}

func runWorkspaceMigrateProvider(_ *cobra.Command, args []string) error {
	id := args[0]

	// Validate providers client-side so a typo fails fast with a clear message
	// instead of a CP 400. --to is cobra-required; --from is required by the CP.
	if !validMigrationProvider(migrateFlags.to) {
		return &exitError{code: 2, msg: "workspace migrate-provider: --to must be one of aws|hetzner|gcp"}
	}
	if migrateFlags.from == "" {
		return &exitError{code: 2, msg: "workspace migrate-provider: --from is required (the current provider: aws|hetzner|gcp)"}
	}
	if !validMigrationProvider(migrateFlags.from) {
		return &exitError{code: 2, msg: "workspace migrate-provider: --from must be one of aws|hetzner|gcp"}
	}
	if migrateFlags.from == migrateFlags.to {
		return &exitError{code: 2, msg: "workspace migrate-provider: --from and --to are the same provider — nothing to migrate"}
	}
	// from-instance-id is required for non-AWS sources (no workspace→instance
	// resolver). Enforce it here rather than round-trip a CP 400.
	if migrateFlags.from != "aws" && migrateFlags.fromInstanceID == "" {
		return &exitError{code: 2, msg: fmt.Sprintf("workspace migrate-provider: --from-instance-id is required for a non-AWS (%s) source (the current box to snapshot + retire)", migrateFlags.from)}
	}
	// Never auto-confirm a destructive two-cloud op.
	if !migrateFlags.confirm {
		return &exitError{code: 2, msg: "workspace migrate-provider: refusing to migrate without --confirm (a real migration mutates two clouds: snapshot source → provision target → retire source)"}
	}

	cp, err := cpAdminClient()
	if err != nil {
		return err
	}
	raw, err := cp.MigrateProvider(id, client.MigrateProviderRequest{
		From:           migrateFlags.from,
		To:             migrateFlags.to,
		Confirm:        true,
		FromInstanceID: migrateFlags.fromInstanceID,
		OrgID:          migrateFlags.orgID,
		Runtime:        migrateFlags.runtime,
	})
	if err != nil {
		return fmt.Errorf("workspace migrate-provider: %w", err)
	}
	if outputFormat == "json" || outputFormat == "yaml" {
		return printRaw(raw)
	}
	fmt.Printf("Migration started for workspace %q: %s → %s (runs asynchronously, ~15-20 min).\n", id, migrateFlags.from, migrateFlags.to)
	fmt.Printf("Poll: molecule workspace migration-status %s\n", id)
	return nil
}

func runWorkspaceMigrationStatus(_ *cobra.Command, args []string) error {
	cp, err := cpAdminClient()
	if err != nil {
		return err
	}
	raw, err := cp.GetMigrationStatus(args[0])
	if err != nil {
		return fmt.Errorf("workspace migration-status: %w", err)
	}
	return printRaw(raw)
}

// ===========================================================================
// molecule workspace token mint <id>
// ===========================================================================
var workspaceTokenCmd = &cobra.Command{
	Use:   "token",
	Short: "Manage per-workspace auth tokens",
}

var workspaceTokenMintCmd = &cobra.Command{
	Use:   "mint <workspace-id>",
	Short: "Mint a new per-workspace auth token (plaintext shown once)",
	Args:  cobra.ExactArgs(1),
	RunE:  runWorkspaceTokenMint,
}

func runWorkspaceTokenMint(_ *cobra.Command, args []string) error {
	resp, err := newClient().MintWorkspaceToken(args[0])
	if err != nil {
		return fmt.Errorf("workspace token mint: %w", err)
	}
	if outputFormat == "json" {
		return printJSON(resp)
	}
	if outputFormat == "yaml" {
		return printYAML(resp)
	}
	fmt.Printf("Token minted for workspace %s\n", resp.WorkspaceID)
	fmt.Printf("auth_token: %s\n", resp.Token)
	if resp.Message != "" {
		fmt.Printf("%s\n", resp.Message)
	}
	return nil
}