llmgw/internal/cli/ctl_mcp.go

package cli

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

	"github.com/spf13/cobra"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
)

func newMCPCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
	c := &cobra.Command{
		Use:   "mcp",
		Short: "The company's own MCP servers, served through llmgw's /mcp",
		Long: "An MCP server registered here — Jira, a file server, a database — is served\n" +
			"through llmgw's /mcp beside llmgw's own tools, its tools named SERVER.tool.\n" +
			"Its credentials stay on the control planes: a client gives only its llmgw\n" +
			"key. A key sees and calls only the tools granted it, as patterns, with\n" +
			"`llmgw keys limits --tools 'jira.*'` at any level of the tenancy tree;\n" +
			"none granted anywhere is none. Every call is in `llmgw ctl audit`.",
	}
	c.AddCommand(newMCPAddCmd(client, asJSON), newMCPListCmd(client, asJSON), newMCPRemoveCmd(client),
		newMCPTestCmd(client, asJSON))
	return c
}

func newMCPAddCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
	var req adminapi.MCPServerRequest
	var headers, refs, secretFiles []string
	var bearerFile string
	var disabled bool
	c := &cobra.Command{
		Use:   "add NAME --url URL",
		Short: "Register an MCP server, or change one",
		Example: "  llmgw ctl mcp add jira --url https://jira.internal/mcp --bearer-file /root/jira-token\n" +
			"  llmgw ctl mcp add llmgw --url llmgw:admin --description \"this cluster, read only\"\n" +
			"  llmgw ctl mcp add wiki --url https://wiki.internal/mcp --header-ref Authorization=vault:secret/mcp/wiki#token\n" +
			"  llmgw ctl mcp add jira --disabled      # off, kept",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			var err error
			if req.Headers, err = pairs(headers, "--header"); err != nil {
				return err
			}
			if req.HeaderRefs, err = pairs(refs, "--header-ref"); err != nil {
				return err
			}
			files, err := pairs(secretFiles, "--secret-file")
			if err != nil {
				return err
			}
			// A secret read here and kept in the store, sealed: the file
			// need not be on every control plane, and is not read again.
			for header, file := range files {
				b, err := os.ReadFile(file) //nolint:gosec // a file the operator names
				if err != nil {
					return err
				}
				if req.Secrets == nil {
					req.Secrets = map[string]string{}
				}
				req.Secrets[header] = strings.TrimSpace(string(b))
			}
			if bearerFile != "" {
				b, err := os.ReadFile(bearerFile) //nolint:gosec // a file the operator names
				if err != nil {
					return err
				}
				if req.Secrets == nil {
					req.Secrets = map[string]string{}
				}
				req.Secrets["Authorization"] = "Bearer " + strings.TrimSpace(string(b))
			}
			if cmd.Flags().Changed("disabled") {
				on := !disabled
				req.Enabled = &on
			}
			body, _ := json.Marshal(req)
			var v adminapi.MCPServerView
			if err := client().put(cmd.Context(), "/api/v1/mcp/servers/"+args[0], body, &v); err != nil {
				return err
			}
			if *asJSON {
				return json.NewEncoder(cmd.OutOrStdout()).Encode(v)
			}
			state := "on"
			if !v.Enabled {
				state = "off"
			}
			fmt.Fprintf(cmd.OutOrStdout(), "%s: %s, %s. Its tools are %s.TOOL in /mcp, to keys granted them "+
				"(llmgw keys limits --tools '%s.*'). Check it: llmgw ctl mcp test %s\n", v.Name, v.URL, state, v.Name, v.Name, v.Name)
			return nil
		},
	}
	f := c.Flags()
	f.StringVar(&req.URL, "url", "", "the server's streamable HTTP address, https://…/mcp; llmgw:admin for this cluster, read only")
	f.StringVar(&req.Description, "description", "", "what it is, shown beside its tools")
	f.StringArrayVar(&headers, "header", nil, "NAME=VALUE sent as written, for what is not secret; repeatable")
	f.StringArrayVar(&refs, "header-ref", nil, "NAME=REF read at each call: a file on every control plane, env:, vault: or store:; repeatable")
	f.StringArrayVar(&secretFiles, "secret-file", nil, "NAME=FILE read now and kept in the store, sealed: a token; repeatable")
	f.StringVar(&bearerFile, "bearer-file", "", "a token FILE, read now and kept sealed, sent as Authorization: Bearer")
	f.BoolVar(&disabled, "disabled", false, "keep it but serve none of its tools")
	f.IntVar(&req.TimeoutSeconds, "timeout", 0, "seconds one call may take (default 60)")
	f.IntVar(&req.RatePerMinute, "rate", 0, "calls a minute each caller may make to its tools (default: no bound; -1 takes one away)")
	return c
}

func newMCPListCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
	return &cobra.Command{
		Use:   "list",
		Short: "The MCP servers /mcp serves",
		Args:  cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			var list []adminapi.MCPServerView
			if err := client().get(cmd.Context(), "/api/v1/mcp/servers", &list); err != nil {
				return err
			}
			if *asJSON {
				return json.NewEncoder(cmd.OutOrStdout()).Encode(list)
			}
			if len(list) == 0 {
				fmt.Fprintln(cmd.OutOrStdout(), "no MCP servers; add one: llmgw ctl mcp add NAME --url URL")
				return nil
			}
			w := tabTo(cmd.OutOrStdout())
			fmt.Fprintln(w, "NAME\tSTATE\tURL\tCREDENTIALS\tDESCRIPTION")
			for _, s := range list {
				state := "on"
				if !s.Enabled {
					state = "off"
				}
				var creds []string
				for h := range s.HeaderRefs {
					creds = append(creds, h)
				}
				fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\n", s.Name, state, s.URL, dash(strings.Join(creds, ",")), s.Description)
			}
			return w.Flush()
		},
	}
}

func newMCPRemoveCmd(client func() *ctlClient) *cobra.Command {
	return &cobra.Command{
		Use:   "remove NAME",
		Short: "Remove an MCP server: its tools leave /mcp",
		Args:  cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			if err := client().delete(cmd.Context(), "/api/v1/mcp/servers/"+args[0], nil); err != nil {
				return err
			}
			fmt.Fprintf(cmd.OutOrStdout(), "%s removed; its tools leave /mcp within half a minute\n", args[0])
			return nil
		},
	}
}

func newMCPTestCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
	return &cobra.Command{
		Use:   "test NAME",
		Short: "Ask an MCP server for its tools, with its credentials, as /mcp would",
		Args:  cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			var v adminapi.MCPTestView
			if err := client().post(cmd.Context(), "/api/v1/mcp/servers/"+args[0]+"/test", nil, &v); err != nil {
				return err
			}
			if *asJSON {
				return json.NewEncoder(cmd.OutOrStdout()).Encode(v)
			}
			if !v.OK {
				return errors.New(args[0] + " did not answer: " + v.Error)
			}
			fmt.Fprintf(cmd.OutOrStdout(), "%s answers with %d tool(s):\n", args[0], len(v.Tools))
			for _, t := range v.Tools {
				fmt.Fprintln(cmd.OutOrStdout(), "  "+t)
			}
			return nil
		},
	}
}

// pairs reads NAME=VALUE flags.
func pairs(list []string, flag string) (map[string]string, error) {
	if len(list) == 0 {
		return nil, nil
	}
	out := map[string]string{}
	for _, p := range list {
		k, v, ok := strings.Cut(p, "=")
		if !ok || strings.TrimSpace(k) == "" {
			return nil, fmt.Errorf("%s %q is not NAME=VALUE", flag, p)
		}
		out[strings.TrimSpace(k)] = v
	}
	return out, nil
}