Repository: GoogleCloudPlatform/kubectl-ai
File: main.go (Go)

Code before:
func runCmd(cmd *cobra.Command, args []string) error {
	// Retrieving the context from cmd which passed in init()
	ctx := cmd.Context()

	// pull in flags/env
	maxIterations = viper.GetInt("max-iterations")
	kubeconfigPath = viper.GetString("kubeconfig-path")
	promptTemplateFilePath = viper.GetString("prompt-template-file-path")
	tracePath = viper.GetString("trace-path")
	removeWorkDir = viper.GetBool("remove-workdir")

	// Load YAML config values
	if err := opt.LoadConfigurationFile(); err != nil {
		return fmt.Errorf("loading config file: %w", err)
	}

	// override Options from viper (flags or env)
	opt.ProviderID = viper.GetString("llm-provider")
	opt.ModelID = viper.GetString("model")
	opt.SkipPermissions = viper.GetBool("skip-permissions")
	opt.MCPServer = viper.GetBool("mcp-server")
	opt.EnableToolUseShim = viper.GetBool("enable-tool-use-shim")
	opt.Quiet = viper.GetBool("quiet")

	// Validation of flag values passed by user, use this helper method
	_, err := isFlagValueCorrect(opt.ModelID, geminiModels)
	if err != nil {
		return fmt.Errorf("'%s' is not a valid value", opt.ModelID)

	}

	// Do this early, before the third-party code logs anything.
	redirectStdLogToKlog()

	// Handle kubeconfig with priority: flag/env > KUBECONFIG > default path
	if kubeconfigPath == "" {
		p, err := getDefaultKubeConfigPath()
		if err != nil {
			return fmt.Errorf("failed to get default kubeconfig path: %w", err)
		}
		kubeconfigPath = p
	}

	if opt.MCPServer {
		workDir := "/tmp/kubectl-ai-mcp"
		if err := os.MkdirAll(workDir, 0755); err != nil {
			return fmt.Errorf("error creating work directory: %w", err)
		}
		mcpServer, err := newKubectlMCPServer(ctx, kubeconfigPath, tools.Default(), workDir)
		if err != nil {
			return fmt.Errorf("creating mcp server: %w", err)

Review comment (suggestion):
Let's avoid validating models because it is not possible to do it correctly without fetching the supported models for a given llm-provider and some llm-providers don't even allow listing models.

Note that we support gemini, openai, vertexai, az-openai, ollama, llaama.cpp as llm-provider today and each support different models. 

Code after:
	f.BoolVar(&opt.RemoveWorkDir, "remove-workdir", opt.RemoveWorkDir, "remove the temporary working directory after execution")

	f.StringVar(&opt.ProviderID, "llm-provider", opt.ProviderID, "language model provider")
	f.StringVar(&opt.ModelID, "model", opt.ModelID, "language model e.g. gemini-2.0-flash-thinking-exp-01-21, gemini-2.0-flash")
	f.BoolVar(&opt.SkipPermissions, "skip-permissions", opt.SkipPermissions, "(dangerous) skip asking for confirmation before executing kubectl commands that modify resources")
	f.BoolVar(&opt.MCPServer, "mcp-server", opt.MCPServer, "run in MCP server mode")
	f.BoolVar(&opt.EnableToolUseShim, "enable-tool-use-shim", opt.EnableToolUseShim, "enable tool use shim")
	f.BoolVar(&opt.Quiet, "quiet", opt.Quiet, "run in non-interactive mode, requires a query to be provided as a positional argument")

	// viper binds and env var prefixes
	err := viper.BindPFlags(f)
	if err != nil {
		return fmt.Errorf("failed to bind viper flags")
	}

	return nil
}

func setupViperEnv() {
	viper.SetEnvPrefix(viperEnvPrefix)
	viper.SetEnvKeyReplacer(strings.NewReplacer("-", "_"))
	viper.AutomaticEnv()
}

func runCmd(cmd *cobra.Command, args []string) error {
	// Retrieving the context from cobra
	ctx := cmd.Context()

	var opt Options

	// load YAML config values
	if err := opt.LoadConfigurationFile(); err != nil {
		return fmt.Errorf("failed to load config file: %w", err)
	}

	// populate Options either from CLI flags or ENV vars
	populateOptionsFromViper(&opt)

	// do this early, before the third-party code logs anything.
	redirectStdLogToKlog()

	// resolve kubeconfig path with priority: flag/env > KUBECONFIG > default path
	if err := resolveKubeConfigPath(&opt); err != nil {
		return fmt.Errorf("failed to resolve kubeconfig path: %w", err)
	}

	if opt.MCPServer {
		if err := startMCPServer(ctx, opt); err != nil {
			return fmt.Errorf("failed to start MCP server: %w", err)
		}
	}