llmgw/internal/cli/tui_mcp.go

package cli

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"slices"
	"strconv"
	"strings"
	"time"

	tea "github.com/charmbracelet/bubbletea"

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

// The company's own MCP servers, as the console shows them: each with its
// address, the names of the credentials it is sent (never their values),
// its rate, and who has been granted its tools; Enter for what can be done
// to one, A to add one. Clients reach them all through the gateway's /mcp
// with an llmgw key, as server.tool.

type mcpState struct {
	servers []adminapi.MCPServerView
	err     error
	loaded  bool
	// grants are the tool grants in the tenancy tree, nil when the store
	// could not be read from here; where says what each scope's ID is
	// called, and ids the other way, for the grant form.
	grants []mcpGrant
	where  map[string]string
	ids    map[string]store.LimitOwner
}

// mcpGrant is one level of the tree that grants tools, and which.
type mcpGrant struct {
	owner store.LimitOwner
	label string
	tools []string
}

type consoleMCPMsg mcpState

func (m consoleModel) fetchMCP() tea.Cmd {
	c, h := m.client, m.st
	return func() tea.Msg {
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		var msg consoleMCPMsg
		msg.err = c.get(ctx, "/api/v1/mcp/servers", &msg.servers)
		msg.loaded = true
		if h == nil {
			return msg
		}
		s, err := h.get(ctx)
		if err != nil {
			return msg
		}
		limits, err := s.AllLimits(ctx)
		if err != nil {
			return msg
		}
		msg.where, msg.ids = map[string]string{}, map[string]store.LimitOwner{}
		if tree, err := s.Tree(ctx); err == nil {
			for _, o := range tree {
				msg.name(store.ScopeOrg, o.ID, o.Name)
				for _, t := range o.Teams {
					msg.name(store.ScopeTeam, t.ID, o.Name+"/"+t.Name)
					for _, p := range t.Projects {
						msg.name(store.ScopeProject, p.ID, o.Name+"/"+t.Name+"/"+p.Name)
					}
				}
			}
		}
		if users, err := s.Users(ctx, ""); err == nil {
			for _, u := range users {
				msg.name(store.ScopeUser, u.ID, u.Email)
			}
		}
		msg.grants = []mcpGrant{}
		for _, l := range limits {
			if len(l.AllowedTools) == 0 {
				continue
			}
			o := store.LimitOwner{Scope: l.Scope, ID: l.ScopeID}
			label := msg.where[o.CounterKey()]
			if label == "" {
				label = string(l.Scope) + " " + l.ScopeID
			}
			msg.grants = append(msg.grants, mcpGrant{owner: o, label: string(l.Scope) + " " + label, tools: l.AllowedTools})
		}
		slices.SortFunc(msg.grants, func(a, b mcpGrant) int { return strings.Compare(a.label, b.label) })
		return msg
	}
}

func (s *consoleMCPMsg) name(scope store.Scope, id, label string) {
	o := store.LimitOwner{Scope: scope, ID: id}
	s.where[o.CounterKey()] = label
	s.ids[string(scope)+":"+label] = o
}

// mcpGrantsOf is who may use a server's tools, and which of them.
func (k mcpState) grantsOf(server string) []mcpGrant {
	var out []mcpGrant
	for _, g := range k.grants {
		var mine []string
		for _, t := range g.tools {
			if strings.HasPrefix(t, server+".") || t == "*" {
				mine = append(mine, t)
			}
		}
		if len(mine) > 0 {
			out = append(out, mcpGrant{owner: g.owner, label: g.label, tools: mine})
		}
	}
	return out
}

func (m consoleModel) mcpPanel(pw int) (string, string, string) {
	keys := "→ or Enter open · ↑↓ choose · Q leave"
	if m.inPanel {
		keys = "↑↓ server · Enter actions · A add a server · ← back · Q leave"
	}
	k := m.mcp
	var b strings.Builder
	b.WriteString("\n  " + sDim.Render("The company's MCP servers, offered to agents through the gateway's /mcp as server.tool.") + "\n\n")
	switch {
	case k.err != nil:
		b.WriteString("  " + sBad.Render(k.err.Error()))
		return "MCP servers", b.String(), keys
	case !k.loaded:
		b.WriteString("  " + sDim.Render("asking the control plane…"))
		return "MCP servers", b.String(), keys
	case len(k.servers) == 0:
		b.WriteString("  " + sDim.Render("None yet. A to add one: its address, and the credential it wants."))
		b.WriteString("\n\n" + mcpClientHint())
		return "MCP servers", b.String(), keys
	}
	var rows [][]string
	for _, s := range k.servers {
		state := sOK.Render("on")
		if !s.Enabled {
			state = sDim.Render("off")
		}
		granted := "—"
		if k.grants != nil {
			var who []string
			for _, g := range k.grantsOf(s.Name) {
				who = append(who, g.label)
			}
			granted = sBad.Render("nobody yet")
			if len(who) > 0 {
				granted = strings.Join(who, ", ")
			}
		}
		rate := "—"
		if s.RatePerMinute > 0 {
			rate = fmt.Sprintf("%d/min", s.RatePerMinute)
		}
		rows = append(rows, []string{s.Name, state, granted, rate, dash(strings.Join(mcpCredentials(s), ", ")), s.URL})
	}
	b.WriteString(table([]string{"SERVER", "STATE", "GRANTED TO", "RATE", "SENT", "ADDRESS"}, rows, m.row, pw, m.inPanel))
	b.WriteString("\n\n  " + sDim.Render("A server's tools are seen only where they are granted; nothing is granted by default.") + "\n\n")
	b.WriteString(mcpClientHint())
	return "MCP servers", b.String(), keys
}

// mcpClientHint is how a person's agent is pointed at the gateway.
func mcpClientHint() string {
	return "  " + sDim.Render("In a person's agent: the gateway's https://…/mcp, with their key from /profile/keys") + "\n" +
		"  " + sDim.Render("  claude mcp add --transport http company https://HOST/mcp --header \"Authorization: Bearer sk-llmgw-…\"")
}

// mcpCredentials are the names of the headers a server is sent: which are
// kept as secrets and which as written. Never their values.
func mcpCredentials(s adminapi.MCPServerView) []string {
	var out []string
	for h := range s.HeaderRefs {
		out = append(out, h+" (secret)")
	}
	for h := range s.Headers {
		out = append(out, h)
	}
	slices.Sort(out)
	return out
}

// mcpPath is a server's admin API address.
func mcpPath(name string) string { return "/api/v1/mcp/servers/" + url.PathEscape(name) }

// mcpAuth turns the form's credential fields into what the server is sent:
// a bearer token, another header, or nothing.
func mcpAuth(kind, header, secret string) (map[string]string, error) {
	switch kind {
	case "none":
		return nil, nil
	case "bearer":
		if secret == "" {
			return nil, nil
		}
		return map[string]string{"Authorization": "Bearer " + strings.TrimPrefix(secret, "Bearer ")}, nil
	default:
		if header == "" {
			return nil, fmt.Errorf("which header carries the credential?")
		}
		if secret == "" {
			return nil, nil
		}
		return map[string]string{header: secret}, nil
	}
}

func atoiOr(s string, def int) (int, error) {
	if strings.TrimSpace(s) == "" {
		return def, nil
	}
	return strconv.Atoi(strings.TrimSpace(s))
}

// mcpAddForm registers a server.
func (m *consoleModel) mcpAddForm() {
	m.openForm(&formDialog{title: "Add an MCP server",
		note: "Reached by the control plane over HTTP (streamable HTTP or SSE). Its credential is kept in\n" +
			"the store, sealed, and never reaches a client. Its tools are seen nowhere until granted.",
		fields: []formField{
			textF("name", "", "lowercase, the tools' prefix: jira → jira.search"),
			textF("address", "", "https://jira-mcp.internal/mcp"),
			textF("about", "", "what it is for, in a line"),
			choiceF("credential", []string{"bearer", "header", "none"}, "bearer"),
			textF("header", "", "for credential: header, its name (X-Api-Key)"),
			pwF("secret", "the token or key; kept sealed in the store"),
			textF("timeout", "", "seconds a call may take; empty is the default"),
			textF("rate", "", "calls a minute per caller; empty is no bound"),
		},
		submit: func(m *consoleModel, v []string) (tea.Cmd, error) {
			name, addr := strings.TrimSpace(v[0]), strings.TrimSpace(v[1])
			if name == "" || addr == "" {
				return nil, fmt.Errorf("a name and an address")
			}
			if u, err := url.Parse(addr); err != nil || (u.Scheme != "http" && u.Scheme != "https") || u.Host == "" {
				return nil, fmt.Errorf("the address is http(s)://host/path")
			}
			secrets, err := mcpAuth(v[3], strings.TrimSpace(v[4]), v[5])
			if err != nil {
				return nil, err
			}
			timeout, err := atoiOr(v[6], 0)
			if err != nil || timeout < 0 {
				return nil, fmt.Errorf("timeout is whole seconds")
			}
			rate, err := atoiOr(v[7], 0)
			if err != nil || rate < 0 {
				return nil, fmt.Errorf("rate is calls a minute")
			}
			on := true
			req := adminapi.MCPServerRequest{URL: addr, Description: v[2], Secrets: secrets, Enabled: &on,
				TimeoutSeconds: timeout, RatePerMinute: rate}
			return m.call("Add "+name, http.MethodPut, mcpPath(name), req, func(json.RawMessage) string {
				return "Added " + name + ". Test it from its menu, then grant its tools: nobody sees them yet."
			}), nil
		}})
}

// mcpMenu is what can be done to one server.
func (m consoleModel) mcpMenu(s adminapi.MCPServerView) *menuDialog {
	items := []menuItem{
		{"Test, and list its tools", func(m *consoleModel) tea.Cmd { return m.mcpTest(s) }},
		{"Grant its tools to…", func(m *consoleModel) tea.Cmd { m.mcpGrantForm(s); return nil }},
	}
	for _, g := range m.mcp.grantsOf(s.Name) {
		if !slices.ContainsFunc(g.tools, func(t string) bool { return t != "*" }) {
			continue // every tool, granted with the Limits form; revoked there too
		}
		items = append(items, menuItem{"Revoke from " + g.label + " (" + strings.Join(g.tools, ", ") + ")…",
			func(m *consoleModel) tea.Cmd {
				m.confirmD = &confirmDialog{title: "Revoke " + s.Name + " from " + g.label,
					text: "It and everything below it then no longer see " + strings.Join(g.tools, ", ") +
						",\nunless another level grants them.",
					yes: func(m *consoleModel) tea.Cmd { return m.mcpSetTools(s.Name, g.owner, nil) }}
				return nil
			}})
	}
	items = append(items,
		menuItem{"Change…", func(m *consoleModel) tea.Cmd { m.mcpEditForm(s); return nil }},
		menuItem{"Replace its credential…", func(m *consoleModel) tea.Cmd { m.mcpCredentialForm(s); return nil }},
	)
	if s.Enabled {
		items = append(items, menuItem{"Turn off", func(m *consoleModel) tea.Cmd {
			off := false
			return m.call("Turn off "+s.Name, http.MethodPut, mcpPath(s.Name), adminapi.MCPServerRequest{URL: s.URL, Enabled: &off},
				func(json.RawMessage) string {
					return s.Name + " is off: its tools are listed nowhere and every call is refused."
				})
		}})
	} else {
		items = append(items, menuItem{"Turn on", func(m *consoleModel) tea.Cmd {
			on := true
			return m.call("Turn on "+s.Name, http.MethodPut, mcpPath(s.Name), adminapi.MCPServerRequest{URL: s.URL, Enabled: &on}, nil)
		}})
	}
	items = append(items, menuItem{"Remove…", func(m *consoleModel) tea.Cmd {
		m.confirmD = &confirmDialog{title: "Remove " + s.Name,
			text: "Its tools go from every agent at once, and its sealed credential with it.\nThe grants stay in the tree until revoked.",
			yes: func(m *consoleModel) tea.Cmd {
				return m.call("Remove "+s.Name, http.MethodDelete, mcpPath(s.Name), nil, nil)
			}}
		return nil
	}})
	return &menuDialog{title: s.Name + " · " + s.URL, items: items}
}

// mcpTest asks the server for its tools, with its credentials, as a call
// through the gateway would.
func (m *consoleModel) mcpTest(s adminapi.MCPServerView) tea.Cmd {
	return m.call("Test "+s.Name, http.MethodPost, mcpPath(s.Name)+"/test", nil, func(raw json.RawMessage) string {
		var t adminapi.MCPTestView
		_ = json.Unmarshal(raw, &t)
		if !t.OK {
			return "It did not answer: " + t.Error
		}
		if len(t.Tools) == 0 {
			return "It answers, and offers no tools."
		}
		return fmt.Sprintf("It answers, with %d tools:\n\n  %s", len(t.Tools), strings.Join(t.Tools, "\n  "))
	})
}

func (m *consoleModel) mcpEditForm(s adminapi.MCPServerView) {
	rate := ""
	if s.RatePerMinute > 0 {
		rate = strconv.Itoa(s.RatePerMinute)
	}
	timeout := ""
	if s.TimeoutSeconds > 0 {
		timeout = strconv.Itoa(s.TimeoutSeconds)
	}
	m.openForm(&formDialog{title: "Change " + s.Name,
		note: "Its credential stays as it is; Replace its credential for that.",
		fields: []formField{
			textF("address", s.URL, ""),
			textF("about", s.Description, ""),
			textF("timeout", timeout, "seconds a call may take"),
			textF("rate", rate, "calls a minute per caller; 0 is no bound"),
		},
		submit: func(m *consoleModel, v []string) (tea.Cmd, error) {
			if u, err := url.Parse(v[0]); err != nil || (u.Scheme != "http" && u.Scheme != "https") || u.Host == "" {
				return nil, fmt.Errorf("the address is http(s)://host/path")
			}
			timeout, err := atoiOr(v[2], 0)
			if err != nil || timeout < 0 {
				return nil, fmt.Errorf("timeout is whole seconds")
			}
			rate, err := atoiOr(v[3], 0)
			if err != nil || rate < 0 {
				return nil, fmt.Errorf("rate is calls a minute")
			}
			if rate == 0 && s.RatePerMinute > 0 {
				rate = -1 // the API's way of taking the bound away
			}
			req := adminapi.MCPServerRequest{URL: v[0], Description: v[1], TimeoutSeconds: timeout, RatePerMinute: rate}
			return m.call("Change "+s.Name, http.MethodPut, mcpPath(s.Name), req, nil), nil
		}})
}

func (m *consoleModel) mcpCredentialForm(s adminapi.MCPServerView) {
	m.openForm(&formDialog{title: "Replace the credential " + s.Name + " is sent",
		note: "Kept in the store, sealed. The old value is overwritten, not shown.",
		fields: []formField{
			choiceF("credential", []string{"bearer", "header"}, "bearer"),
			textF("header", "", "for credential: header, its name (X-Api-Key)"),
			pwF("secret", "the token or key"),
		},
		submit: func(m *consoleModel, v []string) (tea.Cmd, error) {
			if v[2] == "" {
				return nil, fmt.Errorf("the new secret?")
			}
			secrets, err := mcpAuth(v[0], strings.TrimSpace(v[1]), v[2])
			if err != nil {
				return nil, err
			}
			return m.call("Credential of "+s.Name, http.MethodPut, mcpPath(s.Name),
				adminapi.MCPServerRequest{URL: s.URL, Secrets: secrets}, nil), nil
		}})
}

// mcpGrantForm grants a server's tools to part of the tree, a person or a
// key, beside whatever else that level is granted.
func (m *consoleModel) mcpGrantForm(s adminapi.MCPServerView) {
	if m.mcp.ids == nil {
		m.dialog, m.dialogTitle = "error", "Grant "+s.Name
		m.dialogText = "The store cannot be read from this machine, so its grants cannot be changed here.\n" +
			"On the control plane: llmgw keys limits --scope team --id ORG/TEAM --tools '" + s.Name + ".*'"
		return
	}
	m.openForm(&formDialog{title: "Grant " + s.Name + "'s tools",
		note: "It and everything below it then see them, unless a level below grants\n" +
			"tools and not these. Other grants at that level stay.",
		fields: []formField{
			choiceF("to", []string{"team", "project", "org", "user", "key"}, "team"),
			textF("which", "", "ORG/TEAM/PROJECT, ORG/TEAM, ORG, an e-mail, or a key's ID"),
			textF("tools", s.Name+".*", "patterns: "+s.Name+".* for all, or "+s.Name+".search,"+s.Name+".get"),
		},
		submit: func(m *consoleModel, v []string) (tea.Cmd, error) {
			which := strings.TrimSpace(v[1])
			if which == "" {
				return nil, fmt.Errorf("which %s?", v[0])
			}
			owner, ok := m.mcp.ids[v[0]+":"+which]
			if v[0] == "key" {
				owner, ok = store.LimitOwner{Scope: store.ScopeKey, ID: which}, true
			}
			if !ok {
				return nil, fmt.Errorf("no %s %q", v[0], which)
			}
			var tools []string
			for _, t := range strings.Split(v[2], ",") {
				if t = strings.TrimSpace(t); t != "" {
					if !strings.HasPrefix(t, s.Name+".") {
						return nil, fmt.Errorf("%q is not one of %s's tools: they start %s.", t, s.Name, s.Name)
					}
					tools = append(tools, t)
				}
			}
			if len(tools) == 0 {
				return nil, fmt.Errorf("which tools? %s.* for all", s.Name)
			}
			return m.mcpSetTools(s.Name, owner, tools), nil
		}})
}

// mcpSetTools makes a level's grants of a server's tools these, leaving its
// grants of other servers' as they are. None revokes them.
func (m *consoleModel) mcpSetTools(server string, owner store.LimitOwner, tools []string) tea.Cmd {
	var keep []string
	for _, g := range m.mcp.grants {
		if g.owner != owner {
			continue
		}
		for _, t := range g.tools {
			if !strings.HasPrefix(t, server+".") {
				keep = append(keep, t)
			}
		}
	}
	all := append(keep, tools...)
	title := "Grant " + server
	if len(tools) == 0 {
		title = "Revoke " + server
	}
	return m.runCLI(title, "keys", "limits", "--scope="+string(owner.Scope), "--id="+owner.ID,
		"--tools="+strings.Join(all, ","))
}