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

Code before:
	viper.SetEnvKeyReplacer(strings.NewReplacer("-", "_"))
	viper.AutomaticEnv()
}

func main() {
	ctx := context.Background()

	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)

	go func() {
		sig := <-sigCh
		fmt.Fprintf(os.Stderr, "Received signal, shutting down... %s\n", sig)
		klog.Flush()
		os.Exit(0)
	}()

	if err := rootCmd.ExecuteContext(ctx); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

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)

Review comment (refactor):
IMO, these vipers instructions should be under a func `configureViper` which is called by `runCmd` to abstract the CLI setting to the CLI logic.

From line 198 til 238

Code after:

	var opt Options
	opt.InitDefaults()

	if err := bindCLIFlagsToViper(opt); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	setupViperEnv()

	if err := rootCmd.ExecuteContext(ctx); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

func bindCLIFlagsToViper(opt Options) error {
	f := rootCmd.Flags()
	f.IntVar(&opt.MaxIterations, "max-iterations", opt.MaxIterations, "maximum number of iterations agent will try before giving up")
	f.StringVar(&opt.KubeConfigPath, "kubeconfig", opt.KubeConfigPath, "path to kubeconfig file")
	f.StringVar(&opt.PromptTemplateFilePath, "prompt-template-file-path", opt.PromptTemplateFilePath, "path to custom prompt template file")
	f.StringVar(&opt.TracePath, "trace-path", opt.TracePath, "path to the trace file")
	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()