From ca77637a71a6ed74bee37c28c16dd6c6aa76fd14 Mon Sep 17 00:00:00 2001 From: Ilker Manap Date: Tue, 20 Oct 2026 10:53:24 +0200 Subject: [PATCH] /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) --- internal/router/mcp_gateway.go | 305 ++++++++++++++++++++++++++++ internal/router/mcp_gateway_test.go | 112 ++++++++++ internal/router/memory.go | 31 ++- internal/router/router.go | 5 + 4 files changed, 452 insertions(+), 1 deletion(-) create mode 100644 internal/router/mcp_gateway.go create mode 100644 internal/router/mcp_gateway_test.go diff --git a/internal/router/mcp_gateway.go b/internal/router/mcp_gateway.go new file mode 100644 index 0000000..cfb6d3f --- /dev/null +++ b/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) + } +} diff --git a/internal/router/mcp_gateway_test.go b/internal/router/mcp_gateway_test.go new file mode 100644 index 0000000..b3d9614 --- /dev/null +++ b/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 +} diff --git a/internal/router/memory.go b/internal/router/memory.go index 7d9c4c5..739c16a 100644 --- a/internal/router/memory.go +++ b/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 diff --git a/internal/router/router.go b/internal/router/router.go index 03380ad..f23b788 100644 --- a/internal/router/router.go +++ b/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