mcp/cmd/mcp/node.go

package main

import (
	"context"
	"fmt"
	"os"
	"text/tabwriter"
	"time"

	toml "github.com/pelletier/go-toml/v2"

	mcpv1 "git.wntrmute.dev/mc/mcp/gen/mcp/v1"
	"git.wntrmute.dev/mc/mcp/internal/config"
	"github.com/spf13/cobra"
)

func nodeCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "node",
		Short: "Node management",
	}

	list := &cobra.Command{
		Use:   "list",
		Short: "List registered nodes",
		RunE:  runNodeList,
	}

	add := &cobra.Command{
		Use:   "add <name> <address>",
		Short: "Register a node",
		Args:  cobra.ExactArgs(2),
		RunE:  runNodeAdd,
	}

	remove := &cobra.Command{
		Use:   "remove <name>",
		Short: "Deregister a node",
		Args:  cobra.ExactArgs(1),
		RunE:  runNodeRemove,
	}

	cmd.AddCommand(list, add, remove)
	return cmd
}

func runNodeList(_ *cobra.Command, _ []string) error {
	cfg, err := config.LoadCLIConfig(cfgPath)
	if err != nil {
		return fmt.Errorf("load config: %w", err)
	}

	w := tabwriter.NewWriter(os.Stdout, 0, 4, 2, ' ', 0)
	_, _ = fmt.Fprintln(w, "NAME\tADDRESS\tVERSION")
	for _, n := range cfg.Nodes {
		ver := queryAgentVersion(cfg, n.Address)
		_, _ = fmt.Fprintf(w, "%s\t%s\t%s\n", n.Name, n.Address, ver)
	}
	return w.Flush()
}

// queryAgentVersion dials the agent and returns its version, or an error indicator.
func queryAgentVersion(cfg *config.CLIConfig, address string) string {
	client, conn, err := dialAgent(address, cfg)
	if err != nil {
		return "error"
	}
	defer func() { _ = conn.Close() }()

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	resp, err := client.NodeStatus(ctx, &mcpv1.NodeStatusRequest{})
	if err != nil {
		return "error"
	}
	if resp.AgentVersion == "" {
		return "unknown"
	}
	return resp.AgentVersion
}

func runNodeAdd(_ *cobra.Command, args []string) error {
	cfg, err := config.LoadCLIConfig(cfgPath)
	if err != nil {
		return fmt.Errorf("load config: %w", err)
	}

	name := args[0]
	address := args[1]

	for _, n := range cfg.Nodes {
		if n.Name == name {
			return fmt.Errorf("node %q already exists", name)
		}
	}

	cfg.Nodes = append(cfg.Nodes, config.NodeConfig{
		Name:    name,
		Address: address,
	})

	if err := writeConfig(cfgPath, cfg); err != nil {
		return fmt.Errorf("write config: %w", err)
	}

	fmt.Printf("Added node %s (%s)\n", name, address)
	return nil
}

func runNodeRemove(_ *cobra.Command, args []string) error {
	cfg, err := config.LoadCLIConfig(cfgPath)
	if err != nil {
		return fmt.Errorf("load config: %w", err)
	}

	name := args[0]

	var found bool
	nodes := make([]config.NodeConfig, 0, len(cfg.Nodes))
	for _, n := range cfg.Nodes {
		if n.Name == name {
			found = true
			continue
		}
		nodes = append(nodes, n)
	}

	if !found {
		return fmt.Errorf("node %q not found", name)
	}

	cfg.Nodes = nodes

	if err := writeConfig(cfgPath, cfg); err != nil {
		return fmt.Errorf("write config: %w", err)
	}

	fmt.Printf("Removed node %s\n", name)
	return nil
}

// writeConfig serializes the CLIConfig to TOML and writes it back to the
// config file.
func writeConfig(path string, cfg *config.CLIConfig) error {
	data, err := toml.Marshal(cfg)
	if err != nil {
		return fmt.Errorf("marshal config: %w", err)
	}

	if err := os.WriteFile(path, data, 0o600); err != nil {
		return fmt.Errorf("write config %q: %w", path, err)
	}

	return nil
}