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

Code before:
	opt.InitDefaults()

	f := rootCmd.Flags()
	f.IntVar(&maxIterations, "max-iterations", 20, "maximum number of iterations agent will try before giving up")
	f.StringVar(&kubeconfigPath, "kubeconfig-path", "", "path to kubeconfig file")
	f.StringVar(&promptTemplateFilePath, "prompt-template-file-path", "", "path to custom prompt template file")
	f.StringVar(&tracePath, "trace-path", "trace.log", "path to the trace file")
	f.BoolVar(&removeWorkDir, "remove-workdir", false, "remove the temporary working directory after execution")

	// This part is gotten from Options struct? May be put all of them there?
	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 {
		print(fmt.Errorf("failed to bind viper flags"))
	}
	viper.SetEnvPrefix("KUBECTL_AI")
	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

Review comment (suggestion):
I see that `klog.InitFlags(nil)` is no longer here ? (that is useful for controlling the logging configuration).

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

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

	// klog setup must happen before Cobra parses any flags
	// add commandline flags for logging
	klog.InitFlags(nil)

	flag.Set("logtostderr", "false")
	flag.Set("log_file", filepath.Join(os.TempDir(), "kubectl-ai.log"))

	// cobra has to know that we pass pass flags with flag lib, otherwise it creates conflict with flags.parse() method
	rootCmd.PersistentFlags().AddGoFlagSet(flag.CommandLine)
	defer klog.Flush()

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

	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")