molecule-cli/internal/cmd/workspace_mgmt.go

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

import (
	"fmt"

	"github.com/spf13/cobra"
)

// ---------------------------------------------------------------------------
// 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 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
}