/mcp serves the company's own MCP servers' tools to keys granted them

Each registered server's tools are listed as server.tool beside llmgw's
own, to a key whose tool grants match, and called with the server's
credentials, read from their secret references at each call; the
server is told the key, person and project calling. A call is bounded
in time and size, its progress passed on to a streaming client, and
recorded in the audit log with its arguments' size and digest, never
their text. A key granted nothing sees none of them.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

internal/router/mcp_gateway.go

@@ -0,0 +1,305 @@
+package router
+
+import (
+	"context"
+	"crypto/sha256"
+	"encoding/hex"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"log/slog"
+	"maps"
+	"net/http"
+	"regexp"
+	"slices"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/prometheus/client_golang/prometheus"
+
+	"gitea.mynodes.xyz/ilker/llmgw/internal/auth"
+	"gitea.mynodes.xyz/ilker/llmgw/internal/mcpclient"
+	"gitea.mynodes.xyz/ilker/llmgw/internal/secrets"
+	storepkg "gitea.mynodes.xyz/ilker/llmgw/internal/store"
+)
+
+// The company's own MCP servers behind llmgw's /mcp. An administrator
+// registers each once (llmgw ctl mcp add); /mcp then lists its tools as
+// server.tool beside llmgw's own, to a caller whose key is granted them, and
+// calls them with the server's credentials, which never leave the control
+// plane. Every call is in the audit log.
+
+// mcpServerName is what a server may be called: its tools are name.tool, so
+// no dot, and short enough to read in a tool list.
+var mcpServerName = regexp.MustCompile(`^[a-z][a-z0-9_-]{0,31}$`)
+
+// ValidMCPServerName reports whether a name can be a server's.
+func ValidMCPServerName(name string) bool { return mcpServerName.MatchString(name) }
+
+const (
+	// mcpServersFresh is how long the server list read from the store is
+	// used before it is read again: a server added on one control plane is
+	// served by every one within it.
+	mcpServersFresh = 30 * time.Second
+	// mcpToolsFresh is how long a server's tool list is used. A server that
+	// cannot be asked keeps the last list it gave.
+	mcpToolsFresh = 5 * time.Minute
+	// mcpCallTimeout is a call's bound when the server's own is not set.
+	mcpCallTimeout = 60 * time.Second
+	// mcpAnswerBytes bounds one tool's answer: it becomes a prompt.
+	mcpAnswerBytes = 4 << 20
+)
+
+type mcpStore interface {
+	MCPServers(ctx context.Context) ([]storepkg.MCPServer, error)
+	WriteAudit(ctx context.Context, e storepkg.AuditEvent) error
+}
+
+// mcpGateway is the company servers as /mcp serves them.
+type mcpGateway struct {
+	st   mcpStore
+	log  *slog.Logger
+	http *http.Client
+	now  func() time.Time
+
+	calls    *prometheus.CounterVec
+	duration *prometheus.HistogramVec
+
+	mu       sync.Mutex
+	readAt   time.Time
+	upstream map[string]*mcpUpstream
+}
+
+// mcpUpstream is one server: its settings, its client, its tools.
+type mcpUpstream struct {
+	cfg    storepkg.MCPServer
+	client *mcpclient.Client
+
+	mu      sync.Mutex
+	tools   []mcpclient.Tool
+	toolsAt time.Time
+	toolErr error
+}
+
+func newMCPGateway(st mcpStore, log *slog.Logger) *mcpGateway {
+	return &mcpGateway{
+		st: st, log: log, http: &http.Client{}, now: time.Now, upstream: map[string]*mcpUpstream{},
+		calls: prometheus.NewCounterVec(prometheus.CounterOpts{
+			Name: "llmgw_mcp_tool_calls_total",
+			Help: "Calls to the company MCP servers' tools through /mcp, by server, tool and outcome (ok, tool_error, error, denied).",
+		}, []string{"server", "tool", "outcome"}),
+		duration: prometheus.NewHistogramVec(prometheus.HistogramOpts{
+			Name:    "llmgw_mcp_tool_call_seconds",
+			Help:    "How long a company MCP server took to answer a tool call.",
+			Buckets: prometheus.ExponentialBuckets(0.05, 2, 12),
+		}, []string{"server", "tool"}),
+	}
+}
+
+func (g *mcpGateway) collectors() []prometheus.Collector {
+	return []prometheus.Collector{g.calls, g.duration}
+}
+
+// servers are the enabled servers, read from the store when the list is
+// stale; a store that does not answer leaves the last list.
+func (g *mcpGateway) servers(ctx context.Context) []*mcpUpstream {
+	g.mu.Lock()
+	defer g.mu.Unlock()
+	if g.now().Sub(g.readAt) >= mcpServersFresh {
+		list, err := g.st.MCPServers(ctx)
+		if err != nil {
+			g.log.Warn("could not read the MCP servers; serving the last list", "err", err)
+		} else {
+			g.readAt = g.now()
+			next := map[string]*mcpUpstream{}
+			for _, s := range list {
+				if old, ok := g.upstream[s.Name]; ok && sameMCPServer(old.cfg, s) {
+					old.cfg = s
+					next[s.Name] = old
+					continue
+				}
+				next[s.Name] = g.upstreamFor(s)
+			}
+			g.upstream = next
+		}
+	}
+	var out []*mcpUpstream
+	for _, name := range slices.Sorted(maps.Keys(g.upstream)) {
+		if u := g.upstream[name]; u.cfg.Enabled {
+			out = append(out, u)
+		}
+	}
+	return out
+}
+
+// sameMCPServer is whether a client made for a can serve b: the same address
+// and headers. A changed description or switch keeps the session.
+func sameMCPServer(a, b storepkg.MCPServer) bool {
+	return a.URL == b.URL && maps.Equal(a.Headers, b.Headers) && maps.Equal(a.HeaderRefs, b.HeaderRefs)
+}
+
+func (g *mcpGateway) upstreamFor(s storepkg.MCPServer) *mcpUpstream {
+	cfg := s
+	return &mcpUpstream{cfg: s, client: &mcpclient.Client{
+		URL: s.URL, HTTP: g.http, MaxBytes: mcpAnswerBytes,
+		// Read at each request, so a rotated credential is used at once.
+		Header: func(ctx context.Context) (http.Header, error) {
+			h := http.Header{}
+			for k, v := range cfg.Headers {
+				h.Set(k, v)
+			}
+			for k, ref := range cfg.HeaderRefs {
+				v, err := secrets.ReadString(ctx, ref)
+				if err != nil {
+					return nil, fmt.Errorf("the %s header of MCP server %s: %w", k, cfg.Name, err)
+				}
+				h.Set(k, v)
+			}
+			return h, nil
+		},
+	}}
+}
+
+// toolsOf is a server's tools, asked for again when stale.
+func (u *mcpUpstream) toolsOf(ctx context.Context, now time.Time, log *slog.Logger) []mcpclient.Tool {
+	u.mu.Lock()
+	defer u.mu.Unlock()
+	if now.Sub(u.toolsAt) < mcpToolsFresh && (u.tools != nil || u.toolErr != nil) {
+		return u.tools
+	}
+	ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
+	defer cancel()
+	tools, err := u.client.Tools(ctx)
+	u.toolsAt, u.toolErr = now, err
+	if err != nil {
+		log.Warn("could not list an MCP server's tools; serving the last list", "server", u.cfg.Name, "err", err)
+		return u.tools
+	}
+	u.tools = tools
+	return tools
+}
+
+// list is the company tools this caller may use, as /mcp lists them. No
+// caller (authentication off) is granted none: they are granted, not assumed.
+func (g *mcpGateway) list(req *http.Request) []map[string]any {
+	p, ok := auth.FromContext(req.Context())
+	if g == nil || !ok || len(p.Limits.ToolGrants) == 0 {
+		return nil
+	}
+	var out []map[string]any
+	for _, u := range g.servers(req.Context()) {
+		for _, t := range u.toolsOf(req.Context(), g.now(), g.log) {
+			name := u.cfg.Name + "." + t.Name
+			if !p.Limits.AllowsTool(name) {
+				continue
+			}
+			desc := t.Description
+			if u.cfg.Description != "" {
+				desc = "[" + u.cfg.Description + "] " + desc
+			}
+			tool := map[string]any{"name": name, "description": desc}
+			if len(t.InputSchema) > 0 {
+				tool["inputSchema"] = t.InputSchema
+			} else {
+				tool["inputSchema"] = map[string]any{"type": "object"}
+			}
+			if t.Title != "" {
+				tool["title"] = t.Title
+			}
+			if len(t.Annotations) > 0 {
+				tool["annotations"] = t.Annotations
+			}
+			out = append(out, tool)
+		}
+	}
+	return out
+}
+
+// isCompanyTool is a tool name with a server's prefix: llmgw's own have no dot.
+func isCompanyTool(name string) bool { return strings.Contains(name, ".") }
+
+// call calls a company tool for the caller and gives the tool's result as
+// /mcp answers it: the server's own, or a tool error saying what went wrong.
+func (g *mcpGateway) call(req *http.Request, name string, args json.RawMessage, progress func(json.RawMessage)) any {
+	server, tool, _ := strings.Cut(name, ".")
+	p, ok := auth.FromContext(req.Context())
+	if g == nil || !ok || !p.Limits.AllowsTool(name) {
+		if g != nil {
+			g.calls.WithLabelValues(server, tool, "denied").Inc()
+			g.audit(req.Context(), p, ok, name, server, args, 0, http.StatusForbidden, "not granted")
+		}
+		return toolText(name+" is not a tool this key may call", true)
+	}
+	var u *mcpUpstream
+	for _, s := range g.servers(req.Context()) {
+		if s.cfg.Name == server {
+			u = s
+		}
+	}
+	if u == nil {
+		return toolText("no MCP server "+server+" is registered and on", true)
+	}
+	timeout := u.cfg.Timeout
+	if timeout <= 0 {
+		timeout = mcpCallTimeout
+	}
+	ctx, cancel := context.WithTimeout(req.Context(), timeout)
+	defer cancel()
+	// Who is calling, for the server's own log; its credentials are llmgw's.
+	who := http.Header{"X-Llmgw-Key": {p.KeyID}}
+	if p.UserID != "" {
+		who.Set("X-Llmgw-User", p.UserID)
+	}
+	if p.ProjectID != "" {
+		who.Set("X-Llmgw-Project", p.ProjectID)
+	}
+	start := g.now()
+	raw, err := u.client.Call(ctx, tool, args, who, progress)
+	took := g.now().Sub(start)
+	g.duration.WithLabelValues(server, tool).Observe(took.Seconds())
+	switch {
+	case err != nil:
+		g.calls.WithLabelValues(server, tool, "error").Inc()
+		g.audit(req.Context(), p, true, name, server, args, took, http.StatusBadGateway, err.Error())
+		msg := err.Error()
+		if errors.Is(ctx.Err(), context.DeadlineExceeded) {
+			msg = fmt.Sprintf("the MCP server %s did not answer within %s", server, timeout)
+		}
+		return toolText(msg, true)
+	}
+	var result map[string]any
+	if json.Unmarshal(raw, &result) != nil {
+		g.calls.WithLabelValues(server, tool, "error").Inc()
+		g.audit(req.Context(), p, true, name, server, args, took, http.StatusBadGateway, "the answer was not a tool result")
+		return toolText("the MCP server "+server+" answered something that is not a tool result", true)
+	}
+	outcome, status := "ok", http.StatusOK
+	if isErr, _ := result["isError"].(bool); isErr {
+		outcome, status = "tool_error", http.StatusUnprocessableEntity
+	}
+	g.calls.WithLabelValues(server, tool, outcome).Inc()
+	g.audit(req.Context(), p, true, name, server, args, took, status, "")
+	return result
+}
+
+// audit records a call: who, which tool, the arguments' size and a digest of
+// them, how long, and how it ended. Never the arguments themselves: they are
+// the company's data.
+func (g *mcpGateway) audit(ctx context.Context, p auth.Principal, authed bool, tool, server string,
+	args json.RawMessage, took time.Duration, status int, why string,
+) {
+	actor := "anonymous"
+	if authed {
+		actor = kbActor(p)
+	}
+	sum := sha256.Sum256(args)
+	detail := fmt.Sprintf("arguments %d bytes, sha256 %s, %d ms", len(args), hex.EncodeToString(sum[:6]), took.Milliseconds())
+	if why != "" {
+		detail += ": " + why
+	}
+	if err := g.st.WriteAudit(context.WithoutCancel(ctx), storepkg.AuditEvent{Actor: actor, Source: "mcp",
+		Action: "call " + tool, Target: "mcp:" + server, Detail: detail, Status: status}); err != nil {
+		g.log.Warn("could not record an MCP tool call in the audit log", "tool", tool, "err", err)
+	}
+}

internal/router/mcp_gateway_test.go

@@ -0,0 +1,112 @@
+package router
+
+import (
+	"crypto/sha256"
+	"strings"
+	"testing"
+	"time"
+
+	"gitea.mynodes.xyz/ilker/llmgw/internal/mcpclient/mcptest"
+	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
+)
+
+// companyRouter is a gateway with one company MCP server, jira, and two keys:
+// one granted jira.echo, one granted nothing.
+func companyRouter(t *testing.T, srv *mcptest.Server) (*Router, string, string) {
+	t.Helper()
+	cfg := baseConfig()
+	cfg.Auth.Enabled = true
+	cfg.Auth.Store = t.TempDir() + "/store.db"
+	rt := newTestRouter(t, cfg)
+	ctx := t.Context()
+	org, _ := rt.store.CreateOrg(ctx, "acme")
+	team, _ := rt.store.CreateTeam(ctx, org, "eng")
+	project, _ := rt.store.CreateProject(ctx, team, "agents")
+	key := func(id, letter string) string {
+		secret := store.KeyPrefix + strings.Repeat(letter, 43)
+		sum := sha256.Sum256([]byte(secret))
+		if err := rt.store.CreateKey(ctx, store.Key{ID: id, Hash: sum[:], Prefix: letter, Name: id,
+			ProjectID: project, CreatedAt: time.Now()}); err != nil {
+			t.Fatal(err)
+		}
+		return secret
+	}
+	granted, plain := key("key-g", "g"), key("key-p", "p")
+	if err := rt.store.PutLimits(ctx, store.Limits{Scope: store.ScopeKey, ScopeID: "key-g", AllowedTools: []string{"jira.echo"}}); err != nil {
+		t.Fatal(err)
+	}
+	if err := rt.store.SetClusterSecret(ctx, "mcp/jira/Authorization", "Bearer service-token"); err != nil {
+		t.Fatal(err)
+	}
+	if err := rt.store.PutMCPServer(ctx, store.MCPServer{Name: "jira", URL: srv.URL, Description: "tickets", Enabled: true,
+		Headers: map[string]string{"X-Org": "acme"}, HeaderRefs: map[string]string{"Authorization": "store:mcp/jira/Authorization"}}); err != nil {
+		t.Fatal(err)
+	}
+	return rt, granted, plain
+}
+
+func TestACompanyToolIsServedThroughMCPToAKeyGrantedIt(t *testing.T) {
+	srv := mcptest.New()
+	defer srv.Close()
+	rt, granted, plain := companyRouter(t, srv)
+	h := rt.Handler()
+
+	list, _ := mcp(t, h, granted, "tools/list", map[string]any{}, false)
+	names := strings.Join(toolNames(list), ",")
+	if !strings.Contains(names, "jira.echo") || strings.Contains(names, "jira.fail") || !strings.Contains(names, "remember") {
+		t.Fatalf("the granted key sees %s", names)
+	}
+	res, _ := mcp(t, h, granted, "tools/call", map[string]any{"name": "jira.echo", "arguments": map[string]any{"text": "PROJ-1"}}, false)
+	if res["isError"] == true || !strings.Contains(toolResultText(res), "PROJ-1") {
+		t.Fatalf("the call: %v", res)
+	}
+	// With llmgw's credentials and the caller named, never the caller's key.
+	hdr := srv.Headers[0]
+	if hdr.Get("Authorization") != "Bearer service-token" || hdr.Get("X-Org") != "acme" ||
+		hdr.Get("X-Llmgw-Key") != "key-g" || strings.Contains(hdr.Get("Authorization"), store.KeyPrefix) {
+		t.Errorf("headers to the server: %v", hdr)
+	}
+
+	// A key granted nothing sees none of it, and cannot call it by name.
+	list, _ = mcp(t, h, plain, "tools/list", map[string]any{}, false)
+	if strings.Contains(strings.Join(toolNames(list), ","), "jira.") {
+		t.Errorf("an ungranted key sees %v", toolNames(list))
+	}
+	res, _ = mcp(t, h, plain, "tools/call", map[string]any{"name": "jira.echo", "arguments": map[string]any{"text": "x"}}, false)
+	if res["isError"] != true || len(srv.Calls) != 1 {
+		t.Errorf("an ungranted call: %v (server saw %d)", res, len(srv.Calls))
+	}
+
+	// Each is in the audit log, with the arguments' size and never their text.
+	events, err := rt.store.AuditEvents(t.Context(), store.AuditFilter{Contains: "jira.echo"})
+	if err != nil || len(events) != 2 {
+		t.Fatalf("audit %+v %v", events, err)
+	}
+	for _, e := range events {
+		if e.Source != "mcp" || strings.Contains(e.Detail, "PROJ-1") || !strings.Contains(e.Detail, "bytes") {
+			t.Errorf("audit event %+v", e)
+		}
+	}
+}
+
+func TestACompanyToolsProgressIsPassedOnToAStreamingClient(t *testing.T) {
+	srv := mcptest.New()
+	defer srv.Close()
+	srv.Stream = true
+	rt, granted, _ := companyRouter(t, srv)
+	res, events := mcp(t, rt.Handler(), granted, "tools/call", map[string]any{"name": "jira.echo",
+		"arguments": map[string]any{"text": "streamed"}, "_meta": map[string]any{"progressToken": "p1"}}, true)
+	if !strings.Contains(toolResultText(res), "streamed") || !strings.Contains(strings.Join(events, "\n"), "halfway") {
+		t.Errorf("result %v, events %v", res, events)
+	}
+}
+
+// toolResultText is a tool result's first text.
+func toolResultText(res map[string]any) string {
+	content, _ := res["content"].([]any)
+	if len(content) == 0 {
+		return ""
+	}
+	text, _ := content[0].(map[string]any)["text"].(string)
+	return text
+}

internal/router/memory.go

@@ -297,7 +297,11 @@ func (r *Router) handleMCP(w http.ResponseWriter, req *http.Request) {
 	case "ping":
 		resp.Result = map[string]any{}
 	case "tools/list":
-		resp.Result = map[string]any{"tools": append(slices.Clone(mcpTools), r.mediaTools()...)}
+		tools := append(slices.Clone(mcpTools), r.mediaTools()...)
+		for _, t := range r.mcpgw.list(req) {
+			tools = append(tools, t)
+		}
+		resp.Result = map[string]any{"tools": tools}
 	case "tools/call":
 		var call struct {
 			Name      string          `json:"name"`
@@ -307,6 +311,31 @@ func (r *Router) handleMCP(w http.ResponseWriter, req *http.Request) {
 			} `json:"_meta"`
 		}
 		_ = json.Unmarshal(rpc.Params, &call)
+		if isCompanyTool(call.Name) {
+			if acceptsStream(req) {
+				// The server's progress passed on while it works.
+				w.Header().Set("Content-Type", "text/event-stream")
+				w.Header().Set("Cache-Control", "no-cache")
+				w.WriteHeader(http.StatusOK)
+				sw := &streamWriter{w: w, token: call.Meta.ProgressToken}
+				stop := make(chan struct{})
+				go sw.keepAlive(stop, 15*time.Second)
+				resp.Result = r.mcpgw.call(req, call.Name, call.Arguments, func(p json.RawMessage) {
+					var pr struct {
+						Progress, Total float64
+						Message         string
+					}
+					if json.Unmarshal(p, &pr) == nil && pr.Total > 0 {
+						sw.progress(100*pr.Progress/pr.Total, pr.Message)
+					}
+				})
+				close(stop)
+				sw.event(resp)
+				return
+			}
+			resp.Result = r.mcpgw.call(req, call.Name, call.Arguments, nil)
+			break
+		}
 		if !isMediaTool(call.Name) {
 			resp.Result = r.callMemoryTool(req, rpc.Params)
 			break

internal/router/router.go

@@ -94,6 +94,9 @@ type Router struct {
 
 	// extraMetrics are collectors other components publish on /metrics.
 	extraMetrics []prometheus.Collector
+	// mcpgw serves the company's own MCP servers through /mcp; nil without
+	// a store to register them in.
+	mcpgw *mcpGateway
 
 	// activator loads a model that is configured but not currently running.
 	activator Activator
@@ -366,6 +369,8 @@ func New(cfg config.Server, log *slog.Logger) (*Router, error) {
 		// Secrets the settings name as store:… are the same on every
 		// control plane: read from here.
 		secrets.SetStore(st.StoredSecret)
+		r.mcpgw = newMCPGateway(st, log)
+		r.extraMetrics = append(r.extraMetrics, r.mcpgw.collectors()...)
 		r.usageRetention = cfg.Auth.UsageRetention
 		// Applied where no scope set one, so a project created without limits
 		// is still metered. Zero leaves it unlimited, which was the old