molecule-cli/internal/cmd/workspace_runtime_model.go

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

import (
	"encoding/json"
	"errors"
	"fmt"
	"strings"

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

// ---------------------------------------------------------------------------
// Workspace runtime / model management
//
// molecule workspace set-runtime <id> <runtime>
// molecule workspace set-model   <id> <model>
//
// Tied to molecule-core#2056 (OpenAPI management surface). The server-side
// SetModel handler validates (runtime, model) against the provider registry and
// rejects invalid combos (fail-closed). set-runtime guards against orphaning
// the workspace's current model by consulting the offered-models menu for the
// target runtime before applying the switch.
// ---------------------------------------------------------------------------

func init() {
	workspaceCmd.AddCommand(workspaceSetRuntimeCmd, workspaceSetModelCmd)
}

// ===========================================================================
// molecule workspace set-runtime <id> <runtime>
// ===========================================================================
var workspaceSetRuntimeCmd = &cobra.Command{
	Use:   "set-runtime <workspace-id> <runtime>",
	Short: "Change a workspace's runtime",
	Long: `Changes the runtime image a workspace uses (PATCH /workspaces/:id).

Before applying the switch, the CLI checks that the workspace's CURRENT model is
compatible with the target runtime. If it is not, the command fails and tells
you to run "molecule workspace set-model" first. This prevents the runtime
switch from orphaning a model the target runtime cannot route.

After a successful switch, run "molecule workspace restart <id>" to boot into
the new runtime.`,
	Args: cobra.ExactArgs(2),
	RunE: runWorkspaceSetRuntime,
}

func runWorkspaceSetRuntime(_ *cobra.Command, args []string) error {
	id, runtime := args[0], args[1]
	cl := newClient()

	ws, err := cl.GetWorkspace(id)
	if err != nil {
		return fmt.Errorf("workspace set-runtime: %w", err)
	}

	if ws.Model != "" {
		if err := requireModelCompatibleWithRuntime(cl, ws.Model, runtime); err != nil {
			return fmt.Errorf("workspace set-runtime: %w", err)
		}
	}

	raw, err := cl.SetRuntime(id, runtime)
	if err != nil {
		return fmt.Errorf("workspace set-runtime: %w", err)
	}

	var resp struct {
		Status       string `json:"status"`
		NeedsRestart bool   `json:"needs_restart"`
	}
	_ = json.Unmarshal(raw, &resp) // best-effort pretty print

	if outputFormat == "json" || outputFormat == "yaml" {
		return printRaw(raw)
	}

	fmt.Printf("Runtime for workspace %q updated to %q.\n", id, runtime)
	if resp.NeedsRestart {
		fmt.Printf("Restart required: run `molecule workspace restart %s`.\n", id)
	}
	return nil
}

// ===========================================================================
// molecule workspace set-model <id> <model>
// ===========================================================================
var workspaceSetModelCmd = &cobra.Command{
	Use:   "set-model <workspace-id> <model>",
	Short: "Change a workspace's LLM model override",
	Long: `Sets the model override for a workspace (PUT /workspaces/:id/model).

The workspace-server validates the (runtime, model) pair against the provider
registry and rejects incompatible combinations fail-closed (e.g. claude-code +
gpt-5.5). The server auto-restarts the workspace so the new model takes effect.`,
	Args: cobra.ExactArgs(2),
	RunE: runWorkspaceSetModel,
}

func runWorkspaceSetModel(_ *cobra.Command, args []string) error {
	id, model := args[0], args[1]
	cl := newClient()

	raw, err := cl.SetModel(id, model)
	if err != nil {
		return fmt.Errorf("workspace set-model: %w", err)
	}

	if outputFormat == "json" || outputFormat == "yaml" {
		return printRaw(raw)
	}

	var resp struct {
		Status string `json:"status"`
		Model  string `json:"model"`
	}
	if err := json.Unmarshal(raw, &resp); err == nil {
		switch resp.Status {
		case "saved":
			fmt.Printf("Model for workspace %q updated to %q.\n", id, resp.Model)
		case "cleared":
			fmt.Printf("Model override for workspace %q cleared.\n", id)
		default:
			return printRaw(raw)
		}
		return nil
	}
	return printRaw(raw)
}

// ===========================================================================
// Compatibility guard
// ===========================================================================

// requireModelCompatibleWithRuntime queries the target runtime's offered-model
// menu and rejects the current model if it is not on that menu. Unknown
// runtimes (federated / non-first-party) are allowed through: the registry
// cannot speak for them, so we mirror the server's fail-open federation
// contract while still failing closed on known-invalid combos.
func requireModelCompatibleWithRuntime(cl *client.Platform, currentModel, targetRuntime string) error {
	currentModel = strings.TrimSpace(currentModel)
	if currentModel == "" {
		return nil
	}

	offered, err := cl.ListOfferedModels(targetRuntime)
	if err != nil {
		// Only federated / non-first-party runtimes (unknown to the registry)
		// are allowed through. Every other fetch failure is treated as
		// ambiguous and therefore fail-closed — we cannot prove the switch
		// is safe, so we reject it.
		if errors.Is(err, client.ErrRuntimeNotInRegistry) {
			return nil
		}
		return &exitError{
			code: 1,
			msg: fmt.Sprintf(
				"could not verify model compatibility for runtime %q: %v; "+
					"set a compatible model first with `molecule workspace set-model <id> <model>` and retry",
				targetRuntime, err,
			),
		}
	}

	for _, m := range offered.Models {
		if m.Model == currentModel {
			return nil
		}
	}

	valid := make([]string, 0, len(offered.Models))
	for _, m := range offered.Models {
		valid = append(valid, m.Model)
	}

	return &exitError{
		code: 1,
		msg: fmt.Sprintf(
			"model %q is not compatible with runtime %q; set a compatible model first with `molecule workspace set-model <id> <%s-model>` (valid examples: %s)",
			currentModel, targetRuntime, targetRuntime, joinFirstN(valid, 5),
		),
	}
}

func joinFirstN(items []string, n int) string {
	if len(items) == 0 {
		return "(none)"
	}
	if len(items) > n {
		return strings.Join(items[:n], ", ") + ", ..."
	}
	return strings.Join(items, ", ")
}