llmgw/internal/cli/tui_mcp_test.go

package cli

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"os"
	"slices"
	"strings"
	"testing"

	tea "github.com/charmbracelet/bubbletea"

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

func TestAnMCPServerIsAddedAndTestedFromTheConsole(t *testing.T) {
	m, s := consoleWith(t, cMCP)
	m.mcp = mcpState{loaded: true, grants: []mcpGrant{}, servers: []adminapi.MCPServerView{{
		Name: "jira", URL: "https://jira-mcp.internal/mcp", Enabled: true, RatePerMinute: 30,
		HeaderRefs: map[string]string{"Authorization": "store:mcp/jira/Authorization"},
	}}}
	var tm tea.Model = m
	view := tm.View()
	for _, want := range []string{"MCP servers", "jira", "Authorization (secret)",
		"30/min", "nobody yet", "claude mcp add"} {
		if !strings.Contains(view, want) {
			t.Errorf("missing %q in:\n%s", want, view)
		}
	}
	if strings.Contains(view, "store:mcp/jira") {
		t.Error("where the secret is kept is shown")
	}

	// A: a server with a bearer token, which is sent once and kept sealed.
	tm = press(tm, "a")
	f := tm.(consoleModel).formD
	if f == nil {
		t.Fatal("no form")
	}
	f.fields[0].text.SetValue("hr")
	f.fields[1].text.SetValue("https://hr-mcp.internal/mcp")
	f.fields[5].text.SetValue("t0ken")
	f.focus = len(f.fields)
	_ = pressRun(tm, "enter")
	got := s.last()
	var req adminapi.MCPServerRequest
	if !strings.HasPrefix(got, "PUT /api/v1/mcp/servers/hr ") ||
		json.Unmarshal([]byte(strings.SplitN(got, " ", 3)[2]), &req) != nil ||
		req.URL != "https://hr-mcp.internal/mcp" || req.Secrets["Authorization"] != "Bearer t0ken" ||
		req.Enabled == nil || !*req.Enabled {
		t.Fatalf("sent %q", got)
	}

	// Enter on jira: test it, then turn it off.
	tm = press(m, "enter")
	md := tm.(consoleModel).menuD
	if md == nil || md.items[0].label != "Test, and list its tools" || !strings.Contains(md.title, "https://jira-mcp.internal/mcp") {
		t.Fatalf("menu %+v", md)
	}
	_ = pressRun(tm, "enter")
	if got := s.last(); got != "POST /api/v1/mcp/servers/jira/test " {
		t.Fatalf("sent %q", got)
	}
}

func TestAnMCPServersToolsAreGrantedBesideOtherGrants(t *testing.T) {
	// The servers, from a control plane; the grants, from the store.
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_ = json.NewEncoder(w).Encode([]adminapi.MCPServerView{
			{Name: "hr", URL: "https://hr/mcp", Enabled: true}, {Name: "jira", URL: "https://jira/mcp", Enabled: true}})
	}))
	defer srv.Close()
	dir := t.TempDir()
	cfgPath := dir + "/server.yaml"
	if err := os.WriteFile(cfgPath, []byte("auth:\n  store: "+dir+"/llmgw.db\n"), 0o600); err != nil {
		t.Fatal(err)
	}
	if _, err := config.Load(cfgPath); err != nil {
		t.Fatal(err)
	}
	m := newConsoleModel(srv.URL, cfgPath)
	defer m.st.close()
	m.loaded, m.inPanel, m.menu = true, true, cMCP
	ctx := context.Background()
	st, err := m.st.get(ctx)
	if err != nil {
		t.Fatal(err)
	}
	org, _ := st.CreateOrg(ctx, "Acme")
	team, _ := st.CreateTeam(ctx, org, "ml")
	if err := st.PutLimits(ctx, store.Limits{Scope: store.ScopeTeam, ScopeID: team, AllowedTools: []string{"jira.*"}}); err != nil {
		t.Fatal(err)
	}
	toolsOf := func() []string {
		all, _ := st.AllLimits(ctx)
		return all[store.LimitOwner{Scope: store.ScopeTeam, ID: team}.CounterKey()].AllowedTools
	}
	refresh := func(tm tea.Model) tea.Model {
		tm, _ = tm.Update(tm.(consoleModel).fetchMCP()())
		return tm
	}

	var tm tea.Model = refresh(m)
	if v := tm.View(); !strings.Contains(v, "team Acme/ml") || !strings.Contains(v, "nobody yet") {
		t.Fatalf("jira's grant, and hr's none, should show:\n%s", v)
	}
	// hr is the first row: grant its tools to the team.
	tm = press(tm, "enter", "down", "enter")
	f := tm.(consoleModel).formD
	if f == nil || f.fields[2].value() != "hr.*" {
		t.Fatalf("no grant form, or not for hr: %+v", f)
	}
	f.fields[1].text.SetValue("Acme/ml")
	f.focus = len(f.fields)
	tm = pressRun(tm, "enter")
	if got := toolsOf(); !slices.Equal(got, []string{"jira.*", "hr.*"}) {
		t.Fatalf("the team's grants are %v; jira's should have stayed beside hr's", got)
	}

	// Revoke hr from the team: jira's stays.
	tm = refresh(press(tm, "x"))
	tm = press(tm, "enter", "down", "down", "enter")
	if c := tm.(consoleModel).confirmD; c == nil || !strings.Contains(c.title, "Revoke hr from team Acme/ml") {
		t.Fatalf("no revoke: %+v", c)
	}
	_ = pressRun(tm, "y")
	if got := toolsOf(); !slices.Equal(got, []string{"jira.*"}) {
		t.Fatalf("after the revoke the team's grants are %v", got)
	}
}