diff --git a/CHANGELOG.md b/CHANGELOG.md index 5501bf6..c4a44b9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,18 @@ All notable changes to KTMS are listed here. The format follows ## [Unreleased] +### Added +- MCP gateway (experimental): MCP clients such as Claude Code reach registered MCP servers through + Keystone at `/mcp` (every allowed server, tools named `server__tool`) and `/mcp/{server}`, with + a KTMS key of the new `mcp` scope. Servers (Streamable HTTP or SSE) are managed with the + `UpsertMCPServer`, `DeleteMCPServer`, `SetMCPCredential`, `DeleteMCPCredential` and + `ListMCPTools` RPCs, `ktms mcp-servers`, and the `mcp_servers` and `mcp_credentials` file + sections. Each server's credential is encrypted with the vendor keys' keyring and never handed + to clients. `allowed_mcp` on a key's policy and on a team limits servers and tools, `tool_rpm` + limits calls per key and tool, private addresses are refused unless `allow_private` is set, and + every call is a usage event with `tool_call` and counted in `ktms_mcp_*` metrics. + `MCP_GATEWAY_ENABLED=false` turns it off. + ## [1.3.6] - 2026-10-27 ### Added diff --git a/README.md b/README.md index bcfc832..c3efd5e 100644 --- a/README.md +++ b/README.md @@ -46,6 +46,7 @@ vendor is down, enforces rate and spend limits, and records usage. - [Command-line tool](#command-line-tool) - [Calling Keystone](#calling-keystone) - [OpenAI-compatible HTTP API](#openai-compatible-http-api) +- [MCP gateway](#mcp-gateway) - [Concepts](#concepts) - [Error contract](#error-contract) - [API stability](#api-stability) @@ -453,6 +454,7 @@ Everything is an environment variable; see [`env.example`](env.example). | `GRPC_MAX_RECV_MSG_SIZE` | `16777216` | Maximum gRPC message receive size in bytes (16MB). | | `HTTP_ADDR` | `:8080` | HTTP listen address (health probes and the OpenAI-compatible API). | | `HTTP_API_ENABLED` | `true` | Serve the `/v1` API on `HTTP_ADDR`; `false` leaves only the health probes. | +| `MCP_GATEWAY_ENABLED` | `true` | Serve the [MCP gateway](#mcp-gateway) (`/mcp`, `/mcp/{server}`) on `HTTP_ADDR`. | | `HTTP_API_TIMEOUT` | `5m` | Read/write timeout for one HTTP API call (it waits on an upstream model). | | `METRICS_ADDR` | `:9091` | Internal HTTP listen address for Prometheus `/metrics` and `/healthz`. Empty disables the metrics server. | | `LOG_LEVEL` | `info` | zerolog level. | @@ -573,7 +575,8 @@ user_rate_limits: - Secrets never go in the file. A vendor key is read from the variable named in `plain_key_env` and compared with the stored one by fingerprint (`VendorKeyInfo.fingerprint`); a new client key is printed once, after apply. Exported vendor keys point at - `KTMS_VENDOR_KEY__`. + `KTMS_VENDOR_KEY__`. MCP credentials (`mcp_credentials`) work the same way with + `value_env`, exported as `KTMS_MCP_CREDENTIAL_`. - What the Admin API cannot change is refused before anything is applied: a key's scopes or team, or a revoked key named again. Mint a key under another name instead. - A team's cap is its `monthly_budget_usd`; `spend_limits` holds key and user caps. @@ -589,9 +592,9 @@ grpcurl -plaintext -H "authorization: Bearer $ADMIN_KEY" \ ``` Keys look like `ktms_` (keys minted before the rename start with `llm_` and keep working). Only a hash is stored, and the plaintext is returned once. -Scopes: `generate` (data plane), `admin` (implies everything) and `admin:read` (the Admin API -reads: `GetConfig`, `GetSpend`, `GetTeam`, `GetTeamSpend`, `ExplainRoute` and every `List*`; nothing that changes -state). A read key sees every tenant, so it cannot be bound to an organization or team. `GetConfig` +Scopes: `generate` (data plane), `mcp` (the [MCP gateway](#mcp-gateway)), `admin` (implies +everything) and `admin:read` (the Admin API reads: `GetConfig`, `GetSpend`, `GetTeam`, +`GetTeamSpend`, `ExplainRoute` and every `List*`; nothing that changes state). A read key sees every tenant, so it cannot be bound to an organization or team. `GetConfig` never returns a vendor key itself, only its fingerprint. `ListAPIKeys` and `RevokeAPIKey` do what they say. @@ -605,6 +608,8 @@ grpcurl -plaintext -H "authorization: Bearer $ADMIN_KEY" \ localhost:9090 ktms.v1.AdminService/SetAPIKeyPolicy ``` +- `allowed_mcp`: the MCP servers and tools a key with the `mcp` scope may use; see + [MCP gateway](#mcp-gateway). - `allowed_models`: the models the key may call; empty allows any. A chain keeps only the models both the key's and its team's lists allow; an override outside them, or a chain with none left, is `PERMISSION_DENIED` (the message names `key` or `team`). @@ -986,6 +991,7 @@ ktms nodes list; ktms nodes drain gpu-1 ktms node-keys mint gpu-1 --node-id gpu-1 --allowed-cidrs 10.0.0.0/8 ktms keys revoke billing # asks first; --yes for scripts ktms rotate-encryption +ktms mcp-servers list # the MCP gateway's servers, see MCP gateway ktms completion zsh > "${fpath[1]}/_ktms" # also bash and fish ktms mcp # tools for AI assistants, see below ``` @@ -1393,6 +1399,62 @@ Errors use OpenAI's shape, `{"error": {"message", "type", "param", "code"}}`, an OpenAI SDKs retry 429 and 5xx automatically; for `throttled:*` and `vendor_error` you may want to disable that (`max_retries=0`) and handle them yourself. +## MCP gateway + +Keystone is also a gateway for MCP servers. Claude Code, Cursor and other MCP clients reach the +servers you register through Keystone, with a KTMS key: the gateway holds each server's +credential, shows a key only the servers and tools it may use, limits calls per tool and +accounts for every call in light-tokenmeter and Prometheus. The MCP gateway is +[experimental](#experimental) in 1.x. + +- `/mcp` aggregates every server the key may use; tools are named `server__tool`. +- `/mcp/{server}` is one server, with its tools named as the server names them. + +Both are Streamable HTTP on `HTTP_ADDR`, stateless, so any replica answers any request. Set +`MCP_GATEWAY_ENABLED=false` to turn them off. Upstream servers may speak Streamable HTTP or SSE. +Only tools are proxied; resources and prompts are not. + +```bash +# register a server and give it the header the gateway sends it +ktms mcp-servers set github --url https://api.githubcopilot.com/mcp/ --tool-rpm 60 +echo "Bearer $GITHUB_TOKEN" | ktms mcp-servers credential set github +ktms mcp-servers tools github # what it offers, as allowed_mcp entries + +# a key for the agent: the mcp scope, two tools of github and all of docs +ktms keys mint coding-agent --scopes mcp --allowed-mcp github/create_issue,github/search_code,docs + +# Claude Code +claude mcp add --transport http ktms https://ktms.example.com/mcp --header "Authorization: Bearer ktms_" +``` + +- **Servers** are `MCPServerSpec`: `name` (1-32 of `a-z`, `0-9`, `-`), `url`, `transport` + (`streamable_http` or `sse`), `enabled`, `allow_private`, `timeout_ms` (one call, default + 60 s), `max_response_bytes` (one result, default 4 MiB) and `tool_rpm`. Manage them with + `UpsertMCPServer` and `DeleteMCPServer`, `ktms mcp-servers`, the `mcp_servers` section of a + [configuration file](#configuration-as-a-file) or the console. `GetConfig` lists them. +- **Credentials**: `SetMCPCredential` stores a header and its value (`Authorization: Bearer ...` + by default), encrypted with the vendor keys' keyring and re-encrypted by `RotateEncryption`. + The gateway adds it to every request it sends the server. It is never returned, logged, put in + usage events or passed to clients; `GetConfig` shows its header and fingerprint. In a file it is + `mcp_credentials` with `value_env`. +- **Access**: a key needs the `mcp` scope. `allowed_mcp` on the key's policy and on its team lists + servers (`github`) and single tools (`github/create_issue`); empty allows any, and a tool must + pass both lists. Others are not listed, and calling one is the JSON-RPC error `-32003` + `permission_denied` (`/mcp/{server}` for a server outside the lists is HTTP 403); the upstream + is not called. +- **Addresses**: a server whose host is or resolves to a loopback, private, link-local or + metadata address is refused at registration and at every connection, unless `allow_private` is + set for a server on your own network. `http://` is allowed only then. The gateway ignores + `HTTPS_PROXY` for these connections. +- **Limits and failures**: above `tool_rpm` (per client key and tool, needs Redis), past the + timeout, with a result over `max_response_bytes` or with the server down, the call returns a + tool error (`isError`) that says so, for the model to act on. A server that cannot list its + tools is left out of `/mcp`, not the others. +- **Accounting**: each forwarded call is a usage event with vendor `mcp`, model `server__tool`, + no tokens, its duration, status and `tool_call` (server, tool, result size); `user_id` comes + from the `X-KTMS-User` header, or is `svc:`. Metrics: `ktms_mcp_tool_calls_total`, + `ktms_mcp_tool_call_duration_seconds`, `ktms_mcp_response_bytes`. + ## Concepts **Effort** is the only routing field the proxy understands; everything else about a request's @@ -1537,6 +1599,8 @@ These work and are supported, but may change in a minor release. The change is t (`GUARDRAILS_LATENT_*`), the encoder settings (`LATENT_*`), and embeddings in judge state (`EmbeddingPart`, `input_embedding` over HTTP). - **`JudgeStream` and `EvaluateRollouts`**, and `POST /v1/decisions/rollouts` over HTTP. +- **The MCP gateway:** `/mcp`, the `MCPServerSpec` RPCs, `allowed_mcp`, the `mcp_servers` and + `mcp_credentials` file sections, the `ktms_mcp_*` metrics and `ToolCall` in usage events. ### Compatibility with ktms-admin and light-tokenmeter @@ -1642,17 +1706,20 @@ metrics are served next to them. | `ktms_tokens_total` | `vendor`, `model`, `type` (`input`, `output`) | Tokens consumed. | | `ktms_cache_tokens_total` | `vendor`, `model`, `type` (`read`, `write`) | The part of input read from or written to the prompt cache. | | `ktms_active_requests` | `vendor`, `model` | Upstream calls in flight. | -| `ktms_throttled_total` | `reason` (`rpm`, `key_rpm`, `key_tpm`, `budget_key`, `budget_user`, `spend_key_daily`, `spend_key_monthly`, `spend_user_daily`, `spend_user_monthly`, `spend_team_monthly`, `spend_unavailable`) | Calls refused by a limit; `spend_unavailable` is Redis down with `SPEND_LIMITS_FAIL_MODE=closed`. | +| `ktms_throttled_total` | `reason` (`rpm`, `key_rpm`, `key_tpm`, `mcp_tool_rpm`, `budget_key`, `budget_user`, `spend_key_daily`, `spend_key_monthly`, `spend_user_daily`, `spend_user_monthly`, `spend_team_monthly`, `spend_unavailable`) | Calls refused by a limit; `spend_unavailable` is Redis down with `SPEND_LIMITS_FAIL_MODE=closed`. | | `ktms_fallbacks_total` | `from_model`, `to_model`, `reason` | Moves down a fallback chain. | | `ktms_circuit_breaker_state` | `vendor`, `model` | `0` closed, `1` open, `2` half-open. | | `ktms_circuit_breaker_transitions_total`, `ktms_circuit_breaker_rejections_total` | `vendor`, `model` (+ `from`, `to`) | Breaker flaps and the calls it refused. | | `ktms_vendor_cooldowns_total` | `vendor`, `model` | Pauses started from a vendor `Retry-After`. | -| `ktms_dependency_errors_total` | `dependency` (`redis`, `postgres`, `usage_collector`, `keyring`), `operation` | Failed calls to a dependency; `keyring` is a stored vendor key that no keyring version decrypts. | +| `ktms_dependency_errors_total` | `dependency` (`redis`, `postgres`, `usage_collector`, `keyring`, `mcp`), `operation` | Failed calls to a dependency; `keyring` is a stored vendor key that no keyring version decrypts. | | `ktms_usage_events_dropped_total` | `reason` (`queue_full`, `rejected`, `shutdown`, `closed`) | Usage events that never reached light-tokenmeter. | | `ktms_batch_jobs_total`, `ktms_batch_items_total` | `vendor`, `state` | Batch progress. | | `ktms_batch_poll_duration_seconds` | `vendor` | Vendor batch polling, histogram. | | `ktms_shadow_requests_total`, `ktms_shadow_judge_evaluations_total`, `ktms_shadow_judge_scores` | | Shadow traffic and its judged quality. | | `ktms_legacy_headers_requests_total` | `header` | Clients still sending `X-LLM-Proxy-*` headers. | +| `ktms_mcp_tool_calls_total` | `server`, `tool`, `status` (`success`, `tool_error`, `denied`, `throttled`, `timeout`, `too_large`, `upstream_error`, `cancelled`) | Tool calls through the [MCP gateway](#mcp-gateway). | +| `ktms_mcp_tool_call_duration_seconds` | `server` | Tool call duration, histogram. | +| `ktms_mcp_response_bytes` | `server` | Tool result size, histogram. | The `throttled` and `usage_events_dropped` reasons are exported at 0 from startup, so `increase()` catches the first event after a restart. Suggested alerts (tune thresholds to your traffic; `job` diff --git a/cmd/ktms/admin_cmd.go b/cmd/ktms/admin_cmd.go index da51477..cc91533 100644 --- a/cmd/ktms/admin_cmd.go +++ b/cmd/ktms/admin_cmd.go @@ -228,16 +228,16 @@ func (a *app) keysCmd() *cobra.Command { } pol := k.GetPolicy() rows = append(rows, []string{k.GetName(), list(k.GetScopes()), tenant, list(pol.GetAllowedModels()), - limit(int64(pol.GetRpmLimit())), limit(pol.GetTpmLimit()), when(k.GetCreatedAt()), yesNo(k.GetDisabled())}) + limit(int64(pol.GetRpmLimit())), limit(pol.GetTpmLimit()), list(pol.GetAllowedMcp()), when(k.GetCreatedAt()), yesNo(k.GetDisabled())}) } - return p.rows([]string{"NAME", "SCOPES", "TENANT", "MODELS", "RPM", "TPM", "CREATED", "REVOKED"}, rows) + return p.rows([]string{"NAME", "SCOPES", "TENANT", "MODELS", "RPM", "TPM", "MCP", "CREATED", "REVOKED"}, rows) }) }, } outputFlag(ls, &format) ls.Flags().BoolVar(&all, "all", false, "include revoked keys") - var scopes, models []string + var scopes, models, mcpAllow []string var org, team string var rpm int32 var tpm int64 @@ -251,7 +251,7 @@ func (a *app) keysCmd() *cobra.Command { return err } return a.call(func(ctx context.Context, c pb.AdminServiceClient) error { - req, err := mintRequest(ctx, c, args[0], scopes, org, team, policy(models, rpm, tpm)) + req, err := mintRequest(ctx, c, args[0], scopes, org, team, policy(models, rpm, tpm, mcpAllow)) if err != nil { return err } @@ -269,10 +269,10 @@ func (a *app) keysCmd() *cobra.Command { }, } outputFlag(mint, &format) - mint.Flags().StringSliceVar(&scopes, "scopes", []string{"generate"}, "generate, admin or admin:read") + mint.Flags().StringSliceVar(&scopes, "scopes", []string{"generate"}, "generate, mcp, admin or admin:read") mint.Flags().StringVar(&org, "org", "", "organization name") mint.Flags().StringVar(&team, "team", "", "team as org/team") - policyFlags(mint, &models, &rpm, &tpm) + policyFlags(mint, &models, &rpm, &tpm, &mcpAllow) var yes bool revoke := &cobra.Command{ @@ -294,7 +294,7 @@ func (a *app) keysCmd() *cobra.Command { } revoke.Flags().BoolVar(&yes, "yes", false, "revoke without asking") - var pModels []string + var pModels, pMCP []string var pRPM int32 var pTPM int64 set := &cobra.Command{ @@ -303,7 +303,7 @@ func (a *app) keysCmd() *cobra.Command { Args: cobra.ExactArgs(1), RunE: func(_ *cobra.Command, args []string) error { return a.call(func(ctx context.Context, c pb.AdminServiceClient) error { - if _, err := c.SetAPIKeyPolicy(ctx, &pb.SetAPIKeyPolicyRequest{Name: args[0], Policy: policy(pModels, pRPM, pTPM)}); err != nil { + if _, err := c.SetAPIKeyPolicy(ctx, &pb.SetAPIKeyPolicyRequest{Name: args[0], Policy: policy(pModels, pRPM, pTPM, pMCP)}); err != nil { return err } fmt.Fprintf(a.stdout, "Policy of %s replaced.\n", args[0]) @@ -311,7 +311,7 @@ func (a *app) keysCmd() *cobra.Command { }) }, } - policyFlags(set, &pModels, &pRPM, &pTPM) + policyFlags(set, &pModels, &pRPM, &pTPM, &pMCP) pol := &cobra.Command{Use: "policy", Short: "A key's own limits"} pol.AddCommand(set) @@ -319,17 +319,18 @@ func (a *app) keysCmd() *cobra.Command { return cmd } -func policyFlags(cmd *cobra.Command, models *[]string, rpm *int32, tpm *int64) { +func policyFlags(cmd *cobra.Command, models *[]string, rpm *int32, tpm *int64, mcpAllow *[]string) { cmd.Flags().StringSliceVar(models, "allowed-models", nil, "model ids the key may call (default any)") + cmd.Flags().StringSliceVar(mcpAllow, "allowed-mcp", nil, "MCP servers (server) and tools (server/tool) the key may use (default any)") cmd.Flags().Int32Var(rpm, "rpm", 0, "calls per minute (default unlimited)") cmd.Flags().Int64Var(tpm, "tpm", 0, "tokens per minute (default unlimited)") } -func policy(models []string, rpm int32, tpm int64) *pb.APIKeyPolicy { - if len(models) == 0 && rpm == 0 && tpm == 0 { +func policy(models []string, rpm int32, tpm int64, mcpAllow []string) *pb.APIKeyPolicy { + if len(models) == 0 && rpm == 0 && tpm == 0 && len(mcpAllow) == 0 { return nil } - return &pb.APIKeyPolicy{AllowedModels: models, RpmLimit: rpm, TpmLimit: tpm} + return &pb.APIKeyPolicy{AllowedModels: models, RpmLimit: rpm, TpmLimit: tpm, AllowedMcp: mcpAllow} } func limit(n int64) string { diff --git a/cmd/ktms/integration_test.go b/cmd/ktms/integration_test.go index a7a6a17..bdee75e 100644 --- a/cmd/ktms/integration_test.go +++ b/cmd/ktms/integration_test.go @@ -112,6 +112,7 @@ func TestAgainstKeystone(t *testing.T) { t.Setenv("KTMS_ADDR", addr) t.Setenv("KTMS_ADMIN_KEY", bootstrapKey) t.Setenv("TEST_OPENAI_KEY", "sk-test") + t.Setenv("TEST_MCP_TOKEN", "Bearer ghp-test") must := func(want int, stdin string, args ...string) string { t.Helper() code, out, errs := callIn(stdin, args...) diff --git a/cmd/ktms/main.go b/cmd/ktms/main.go index 10f4dbd..fa16eaa 100644 --- a/cmd/ktms/main.go +++ b/cmd/ktms/main.go @@ -126,7 +126,7 @@ context; admin:read is enough for every command that only reads. a.initCmd(), a.configCmd(), a.contextCmd(), a.modelsCmd(), a.rulesCmd(), a.keysCmd(), a.vendorKeysCmd(), a.orgsCmd(), a.teamsCmd(), a.spendCmd(), a.limitsCmd(), a.nodesCmd(), a.nodeKeysCmd(), a.routeCmd(), a.rotateCmd(), - a.chatCmd(), a.healthCmd(), a.mcpCmd(), &cobra.Command{ + a.mcpServersCmd(), a.chatCmd(), a.healthCmd(), a.mcpCmd(), &cobra.Command{ Use: "version", Short: "Print the version", Args: cobra.NoArgs, RunE: func(*cobra.Command, []string) error { fmt.Fprintln(a.stdout, version) diff --git a/cmd/ktms/mcp.go b/cmd/ktms/mcp.go index 055c28d..dc29f0d 100644 --- a/cmd/ktms/mcp.go +++ b/cmd/ktms/mcp.go @@ -551,12 +551,13 @@ func (a *app) addWriteTools(srv *mcp.Server, tool func(name, title, desc string, type mintIn struct { Name string `json:"name" jsonschema:"key name, unique"` - Scopes []string `json:"scopes,omitempty" jsonschema:"generate (default), admin or admin:read"` + Scopes []string `json:"scopes,omitempty" jsonschema:"generate (default), mcp, admin or admin:read"` Organization string `json:"organization,omitempty" jsonschema:"organization name"` Team string `json:"team,omitempty" jsonschema:"team as org/team"` AllowedModels []string `json:"allowed_models,omitempty" jsonschema:"model ids the key may call; empty means any"` RPMLimit int32 `json:"rpm_limit,omitempty" jsonschema:"calls per minute, 0 unlimited"` TPMLimit int64 `json:"tpm_limit,omitempty" jsonschema:"tokens per minute, 0 unlimited"` + AllowedMCP []string `json:"allowed_mcp,omitempty" jsonschema:"MCP servers (server) and tools (server/tool) a key with the mcp scope may use; empty means any"` } mcp.AddTool(srv, tool("mint_key", "Mint a client key", "Mint a client API key. The key is returned once, for the user: do not repeat it later.", change), func(ctx context.Context, _ *mcp.CallToolRequest, in mintIn) (*mcp.CallToolResult, any, error) { @@ -566,7 +567,7 @@ func (a *app) addWriteTools(srv *mcp.Server, tool func(name, title, desc string, } var plain string err := a.callCtx(ctx, func(ctx context.Context, c pb.AdminServiceClient) error { - req, err := mintRequest(ctx, c, in.Name, scopes, in.Organization, in.Team, policy(in.AllowedModels, in.RPMLimit, in.TPMLimit)) + req, err := mintRequest(ctx, c, in.Name, scopes, in.Organization, in.Team, policy(in.AllowedModels, in.RPMLimit, in.TPMLimit, in.AllowedMCP)) if err != nil { return err } @@ -594,8 +595,9 @@ func (a *app) addWriteTools(srv *mcp.Server, tool func(name, title, desc string, AllowedModels []string `json:"allowed_models,omitempty" jsonschema:"model ids the key may call; empty means any"` RPMLimit int32 `json:"rpm_limit,omitempty" jsonschema:"calls per minute, 0 unlimited"` TPMLimit int64 `json:"tpm_limit,omitempty" jsonschema:"tokens per minute, 0 unlimited"` + AllowedMCP []string `json:"allowed_mcp,omitempty" jsonschema:"MCP servers (server) and tools (server/tool) a key with the mcp scope may use; empty means any"` }) (string, error) { - _, err := c.SetAPIKeyPolicy(ctx, &pb.SetAPIKeyPolicyRequest{Name: in.Name, Policy: policy(in.AllowedModels, in.RPMLimit, in.TPMLimit)}) + _, err := c.SetAPIKeyPolicy(ctx, &pb.SetAPIKeyPolicyRequest{Name: in.Name, Policy: policy(in.AllowedModels, in.RPMLimit, in.TPMLimit, in.AllowedMCP)}) return "Policy of " + in.Name + " replaced.", err })) diff --git a/cmd/ktms/mcp_servers_cmd.go b/cmd/ktms/mcp_servers_cmd.go new file mode 100644 index 0000000..0f983ec --- /dev/null +++ b/cmd/ktms/mcp_servers_cmd.go @@ -0,0 +1,172 @@ +package main + +import ( + "context" + "fmt" + "sort" + "strconv" + + "github.com/spf13/cobra" + + pb "github.com/ai-process/ktms/gen/ktms/v1" +) + +func (a *app) mcpServersCmd() *cobra.Command { + cmd := &cobra.Command{Use: "mcp-servers", Short: "MCP servers proxied at /mcp, and their credentials"} + + var format string + ls := &cobra.Command{ + Use: "list", Short: "List MCP servers (never their credentials)", Args: cobra.NoArgs, + RunE: func(*cobra.Command, []string) error { + p, err := newPrinter(a.stdout, format) + if err != nil { + return err + } + return a.call(func(ctx context.Context, c pb.AdminServiceClient) error { + cfg, err := c.GetConfig(ctx, &pb.GetConfigRequest{}) + if err != nil { + return err + } + ss := cfg.GetMcpServers() + sort.Slice(ss, func(i, j int) bool { return ss[i].GetName() < ss[j].GetName() }) + if !p.table() { + return p.data(map[string]any{"servers": messages(ss), "credentials": messages(cfg.GetMcpCredentials())}) + } + creds := map[string]*pb.MCPCredentialInfo{} + for _, mc := range cfg.GetMcpCredentials() { + creds[mc.GetServer()] = mc + } + var rows [][]string + for _, s := range ss { + cred := "-" + if mc := creds[s.GetName()]; mc != nil { + cred = mc.GetHeader() + " " + orDash(mc.GetFingerprint()) + } + rows = append(rows, []string{s.GetName(), s.GetUrl(), s.GetTransport(), yesNo(s.GetAllowPrivate()), + limit(int64(s.GetToolRpm())), cred, yesNo(s.GetEnabled())}) + } + return p.rows([]string{"NAME", "URL", "TRANSPORT", "PRIVATE", "TOOL RPM", "CREDENTIAL", "ENABLED"}, rows) + }) + }, + } + outputFlag(ls, &format) + + spec := &pb.MCPServerSpec{} + var disabled bool + set := &cobra.Command{ + Use: "set NAME --url URL [--transport sse] [--allow-private] [--timeout-ms N] [--max-response-bytes N] [--tool-rpm N]", + Short: "Register an MCP server or replace its settings", + Long: `Register an MCP server or replace its settings; flags left out take their defaults. +Its tools appear on /mcp as NAME__tool and on /mcp/NAME as they are. A host that is +or resolves to a private address is refused unless --allow-private is given.`, + Args: cobra.ExactArgs(1), + RunE: func(_ *cobra.Command, args []string) error { + if spec.GetUrl() == "" { + return usageError{fmt.Errorf("--url is required")} + } + spec.Name, spec.Enabled = args[0], !disabled + return a.call(func(ctx context.Context, c pb.AdminServiceClient) error { + if _, err := c.UpsertMCPServer(ctx, &pb.UpsertMCPServerRequest{Server: spec}); err != nil { + return err + } + fmt.Fprintf(a.stdout, "MCP server %s registered.\n", args[0]) + return nil + }) + }, + } + set.Flags().StringVar(&spec.Url, "url", "", "the server's MCP endpoint") + set.Flags().StringVar(&spec.Transport, "transport", "streamable_http", "streamable_http or sse") + set.Flags().BoolVar(&spec.AllowPrivate, "allow-private", false, "permit a private, loopback or link-local address") + set.Flags().Int32Var(&spec.TimeoutMs, "timeout-ms", 0, "one tool call (default 60000)") + set.Flags().Int64Var(&spec.MaxResponseBytes, "max-response-bytes", 0, "one tool result (default 4 MiB)") + set.Flags().Int32Var(&spec.ToolRpm, "tool-rpm", 0, "calls per minute per client key and tool (default unlimited)") + set.Flags().BoolVar(&disabled, "disabled", false, "register it switched off") + + var yes bool + del := &cobra.Command{ + Use: "delete NAME", Short: "Remove an MCP server and its credential", Args: cobra.ExactArgs(1), + RunE: func(_ *cobra.Command, args []string) error { + if !yes { + if err := a.confirm("Remove MCP server " + args[0] + "?"); err != nil { + return err + } + } + return a.call(func(ctx context.Context, c pb.AdminServiceClient) error { + if _, err := c.DeleteMCPServer(ctx, &pb.DeleteMCPServerRequest{Name: args[0]}); err != nil { + return err + } + fmt.Fprintf(a.stdout, "MCP server %s removed.\n", args[0]) + return nil + }) + }, + } + del.Flags().BoolVar(&yes, "yes", false, "remove without asking") + + var header string + credSet := &cobra.Command{ + Use: "set NAME [--header H]", + Short: "Give an MCP server the header value the gateway sends it, read from stdin", + Long: `Read the value from stdin (without echo on a terminal), such as "Bearer ", +and store it encrypted. The gateway adds it to every request to the server; clients +never see it.`, + Args: cobra.ExactArgs(1), + RunE: func(_ *cobra.Command, args []string) error { + plain, err := a.readSecret(header + ": ") + if err != nil { + return err + } + return a.call(func(ctx context.Context, c pb.AdminServiceClient) error { + if _, err := c.SetMCPCredential(ctx, &pb.SetMCPCredentialRequest{Server: args[0], Header: header, PlainValue: plain}); err != nil { + return err + } + fmt.Fprintf(a.stdout, "Credential of MCP server %s stored.\n", args[0]) + return nil + }) + }, + } + credSet.Flags().StringVar(&header, "header", "Authorization", "HTTP header to send") + credDel := &cobra.Command{ + Use: "delete NAME", Short: "Remove an MCP server's credential", Args: cobra.ExactArgs(1), + RunE: func(_ *cobra.Command, args []string) error { + return a.call(func(ctx context.Context, c pb.AdminServiceClient) error { + if _, err := c.DeleteMCPCredential(ctx, &pb.DeleteMCPCredentialRequest{Server: args[0]}); err != nil { + return err + } + fmt.Fprintf(a.stdout, "Credential of MCP server %s removed.\n", args[0]) + return nil + }) + }, + } + cred := &cobra.Command{Use: "credential", Short: "The header value the gateway sends a server"} + cred.AddCommand(credSet, credDel) + + var toolsFormat string + tools := &cobra.Command{ + Use: "tools NAME", Short: "Ask a server for its tools through the gateway", Args: cobra.ExactArgs(1), + RunE: func(_ *cobra.Command, args []string) error { + p, err := newPrinter(a.stdout, toolsFormat) + if err != nil { + return err + } + return a.call(func(ctx context.Context, c pb.AdminServiceClient) error { + res, err := c.ListMCPTools(ctx, &pb.ListMCPToolsRequest{Server: args[0]}) + if err != nil { + return err + } + if !p.table() { + return p.data(messages(res.GetTools())) + } + var rows [][]string + for _, t := range res.GetTools() { + rows = append(rows, []string{args[0] + "/" + t.GetName(), args[0] + "__" + t.GetName(), t.GetDescription()}) + } + fmt.Fprintln(a.stderr, strconv.Itoa(len(rows))+" tools; allowed_mcp takes the first column, /mcp shows the second.") + return p.rows([]string{"ALLOWED_MCP", "ON /mcp", "DESCRIPTION"}, rows) + }) + }, + } + outputFlag(tools, &toolsFormat) + + cmd.AddCommand(ls, set, del, cred, tools) + return cmd +} diff --git a/cmd/server/main.go b/cmd/server/main.go index a839d97..9756a27 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -26,6 +26,7 @@ import ( "github.com/ai-process/ktms/internal/latent" "github.com/ai-process/ktms/internal/llm" "github.com/ai-process/ktms/internal/logger" + "github.com/ai-process/ktms/internal/mcpgw" "github.com/ai-process/ktms/internal/nodes" "github.com/ai-process/ktms/internal/observability" "github.com/ai-process/ktms/internal/proxydb" @@ -295,7 +296,22 @@ func main() { if cfg.HTTPAPIEnabled { openaiAPI = server.NewOpenAIAPI(proxySrv, verifier) } - httpSrv := server.NewHTTPServer(pool, cfg.HTTPAddr, openaiAPI, cfg.HTTPAPITimeout). + var extra []server.Registrar + if cfg.MCPGatewayEnabled { + var toolLimiter mcpgw.ToolLimiter + if rl, ok := limiter.(*throttle.Redis); ok { + toolLimiter = rl + } + var recorder mcpgw.Recorder + if usageClient != nil { + recorder = usageClient + } + gw := mcpgw.New(loader, verifier, toolLimiter, recorder, server.Version) + defer gw.Close() + adminSrv.SetMCPToolLister(gw) + extra = append(extra, gw) + } + httpSrv := server.NewHTTPServer(pool, cfg.HTTPAddr, openaiAPI, cfg.HTTPAPITimeout, extra...). WithDegradedReadiness(cfg.AuthStaleIfError) grpcErrCh := make(chan error, 1) diff --git a/deploy/helm/ktms/templates/deployment.yaml b/deploy/helm/ktms/templates/deployment.yaml index b0f22a0..7226335 100644 --- a/deploy/helm/ktms/templates/deployment.yaml +++ b/deploy/helm/ktms/templates/deployment.yaml @@ -103,6 +103,8 @@ spec: {{- end }} - name: HTTP_API_ENABLED value: {{ .Values.config.httpApiEnabled | quote }} + - name: MCP_GATEWAY_ENABLED + value: {{ .Values.config.mcpGatewayEnabled | quote }} - name: LOG_LEVEL value: {{ .Values.config.logLevel | quote }} - name: AUTO_MIGRATE diff --git a/deploy/helm/ktms/values.schema.json b/deploy/helm/ktms/values.schema.json index b5292a4..d51daaa 100644 --- a/deploy/helm/ktms/values.schema.json +++ b/deploy/helm/ktms/values.schema.json @@ -56,6 +56,7 @@ "type": "object", "properties": { "httpApiEnabled": { "type": "boolean" }, + "mcpGatewayEnabled": { "type": "boolean" }, "logLevel": { "type": "string" }, "autoMigrate": { "type": "boolean" }, "autoSeed": { "type": "boolean" }, diff --git a/deploy/helm/ktms/values.yaml b/deploy/helm/ktms/values.yaml index 266b244..1373fe2 100644 --- a/deploy/helm/ktms/values.yaml +++ b/deploy/helm/ktms/values.yaml @@ -154,6 +154,8 @@ secrets: config: httpApiEnabled: true + # MCP gateway (/mcp) on the HTTP port for keys with the mcp scope. + mcpGatewayEnabled: true logLevel: "info" autoMigrate: true autoSeed: false diff --git a/env.example b/env.example index c56cd46..d065795 100644 --- a/env.example +++ b/env.example @@ -10,6 +10,8 @@ GRPC_MAX_RECV_MSG_SIZE=16777216 HTTP_ADDR=:8080 # OpenAI-compatible HTTP API (/v1/...) on HTTP_ADDR; false leaves only /healthz and /readyz. HTTP_API_ENABLED=true +# MCP gateway (/mcp, /mcp/{server}) on HTTP_ADDR for keys with the mcp scope; false turns it off. +MCP_GATEWAY_ENABLED=true HTTP_API_TIMEOUT=5m METRICS_ADDR=:9091 SHUTDOWN_TIMEOUT=30s diff --git a/gen/ktms/v1/admin.pb.go b/gen/ktms/v1/admin.pb.go index 5262907..fe6bbb5 100644 --- a/gen/ktms/v1/admin.pb.go +++ b/gen/ktms/v1/admin.pb.go @@ -891,6 +891,8 @@ type GetConfigResponse struct { Version int64 `protobuf:"varint,4,opt,name=version,proto3" json:"version,omitempty"` UserRateLimits []*UserRateLimitSpec `protobuf:"bytes,5,rep,name=user_rate_limits,json=userRateLimits,proto3" json:"user_rate_limits,omitempty"` SpendLimits []*SpendLimitSpec `protobuf:"bytes,6,rep,name=spend_limits,json=spendLimits,proto3" json:"spend_limits,omitempty"` + McpServers []*MCPServerSpec `protobuf:"bytes,7,rep,name=mcp_servers,json=mcpServers,proto3" json:"mcp_servers,omitempty"` + McpCredentials []*MCPCredentialInfo `protobuf:"bytes,8,rep,name=mcp_credentials,json=mcpCredentials,proto3" json:"mcp_credentials,omitempty"` // metadata only, never the value unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -967,10 +969,24 @@ func (x *GetConfigResponse) GetSpendLimits() []*SpendLimitSpec { return nil } +func (x *GetConfigResponse) GetMcpServers() []*MCPServerSpec { + if x != nil { + return x.McpServers + } + return nil +} + +func (x *GetConfigResponse) GetMcpCredentials() []*MCPCredentialInfo { + if x != nil { + return x.McpCredentials + } + return nil +} + type MintAPIKeyRequest struct { state protoimpl.MessageState `protogen:"open.v1"` Name string `protobuf:"bytes,1,opt,name=name,proto3" json:"name,omitempty"` - Scopes []string `protobuf:"bytes,2,rep,name=scopes,proto3" json:"scopes,omitempty"` // "generate" | "admin" | "admin:read" + Scopes []string `protobuf:"bytes,2,rep,name=scopes,proto3" json:"scopes,omitempty"` // "generate" | "mcp" | "admin" | "admin:read" OrgId string `protobuf:"bytes,3,opt,name=org_id,json=orgId,proto3" json:"org_id,omitempty"` TeamId string `protobuf:"bytes,4,opt,name=team_id,json=teamId,proto3" json:"team_id,omitempty"` Policy *APIKeyPolicy `protobuf:"bytes,5,opt,name=policy,proto3" json:"policy,omitempty"` @@ -1052,7 +1068,10 @@ type APIKeyPolicy struct { RpmLimit int32 `protobuf:"varint,2,opt,name=rpm_limit,json=rpmLimit,proto3" json:"rpm_limit,omitempty"` // calls per minute over all models, 0 = unlimited // Tokens per minute over all models, 0 = unlimited. A call books its input // estimate when admitted and its real input + output when it ends. - TpmLimit int64 `protobuf:"varint,3,opt,name=tpm_limit,json=tpmLimit,proto3" json:"tpm_limit,omitempty"` + TpmLimit int64 `protobuf:"varint,3,opt,name=tpm_limit,json=tpmLimit,proto3" json:"tpm_limit,omitempty"` + // MCP servers ("github") and single tools ("github/create_issue") a key with + // the mcp scope may use; empty = any. Others are hidden and refused. + AllowedMcp []string `protobuf:"bytes,4,rep,name=allowed_mcp,json=allowedMcp,proto3" json:"allowed_mcp,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -1108,6 +1127,13 @@ func (x *APIKeyPolicy) GetTpmLimit() int64 { return 0 } +func (x *APIKeyPolicy) GetAllowedMcp() []string { + if x != nil { + return x.AllowedMcp + } + return nil +} + // SetAPIKeyPolicy replaces a key's policy; it applies from the key's next call. type SetAPIKeyPolicyRequest struct { state protoimpl.MessageState `protogen:"open.v1"` @@ -3095,6 +3121,7 @@ type Team struct { MonthlyBudgetUsd float64 `protobuf:"fixed64,4,opt,name=monthly_budget_usd,json=monthlyBudgetUsd,proto3" json:"monthly_budget_usd,omitempty"` AllowedModels []string `protobuf:"bytes,5,rep,name=allowed_models,json=allowedModels,proto3" json:"allowed_models,omitempty"` CreatedAt *timestamppb.Timestamp `protobuf:"bytes,6,opt,name=created_at,json=createdAt,proto3" json:"created_at,omitempty"` + AllowedMcp []string `protobuf:"bytes,7,rep,name=allowed_mcp,json=allowedMcp,proto3" json:"allowed_mcp,omitempty"` // as APIKeyPolicy.allowed_mcp, for every key of the team unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -3171,6 +3198,13 @@ func (x *Team) GetCreatedAt() *timestamppb.Timestamp { return nil } +func (x *Team) GetAllowedMcp() []string { + if x != nil { + return x.AllowedMcp + } + return nil +} + type CreateOrganizationRequest struct { state protoimpl.MessageState `protogen:"open.v1"` Name string `protobuf:"bytes,1,opt,name=name,proto3" json:"name,omitempty"` @@ -3425,6 +3459,7 @@ type CreateTeamRequest struct { Name string `protobuf:"bytes,2,opt,name=name,proto3" json:"name,omitempty"` MonthlyBudgetUsd float64 `protobuf:"fixed64,3,opt,name=monthly_budget_usd,json=monthlyBudgetUsd,proto3" json:"monthly_budget_usd,omitempty"` AllowedModels []string `protobuf:"bytes,4,rep,name=allowed_models,json=allowedModels,proto3" json:"allowed_models,omitempty"` + AllowedMcp []string `protobuf:"bytes,5,rep,name=allowed_mcp,json=allowedMcp,proto3" json:"allowed_mcp,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -3487,6 +3522,13 @@ func (x *CreateTeamRequest) GetAllowedModels() []string { return nil } +func (x *CreateTeamRequest) GetAllowedMcp() []string { + if x != nil { + return x.AllowedMcp + } + return nil +} + type CreateTeamResponse struct { state protoimpl.MessageState `protogen:"open.v1"` Team *Team `protobuf:"bytes,1,opt,name=team,proto3" json:"team,omitempty"` @@ -3708,13 +3750,14 @@ func (x *ListTeamsResponse) GetTeams() []*Team { } // Replaces the team: an omitted budget removes the cap and omitted -// allowed_models lifts the restriction, so send every field. +// allowed_models or allowed_mcp lifts the restriction, so send every field. type UpdateTeamRequest struct { state protoimpl.MessageState `protogen:"open.v1"` Id string `protobuf:"bytes,1,opt,name=id,proto3" json:"id,omitempty"` Name string `protobuf:"bytes,2,opt,name=name,proto3" json:"name,omitempty"` MonthlyBudgetUsd float64 `protobuf:"fixed64,3,opt,name=monthly_budget_usd,json=monthlyBudgetUsd,proto3" json:"monthly_budget_usd,omitempty"` AllowedModels []string `protobuf:"bytes,4,rep,name=allowed_models,json=allowedModels,proto3" json:"allowed_models,omitempty"` + AllowedMcp []string `protobuf:"bytes,5,rep,name=allowed_mcp,json=allowedMcp,proto3" json:"allowed_mcp,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -3777,6 +3820,13 @@ func (x *UpdateTeamRequest) GetAllowedModels() []string { return nil } +func (x *UpdateTeamRequest) GetAllowedMcp() []string { + if x != nil { + return x.AllowedMcp + } + return nil +} + type UpdateTeamResponse struct { state protoimpl.MessageState `protogen:"open.v1"` Team *Team `protobuf:"bytes,1,opt,name=team,proto3" json:"team,omitempty"` @@ -4198,6 +4248,670 @@ func (x *GetSpendResponse) GetWindows() []*SpendWindow { return nil } +// MCPServerSpec is an upstream MCP server, proxied at /mcp/{name} and, with +// its tools named name__tool, at /mcp. Its credential is not part of the spec: +// SetMCPCredential stores it. Experimental: the MCP gateway may change in a 1.x minor release. +type MCPServerSpec struct { + state protoimpl.MessageState `protogen:"open.v1"` + Name string `protobuf:"bytes,1,opt,name=name,proto3" json:"name,omitempty"` // 1-32 of a-z, 0-9 and "-", starting with a letter + Url string `protobuf:"bytes,2,opt,name=url,proto3" json:"url,omitempty"` // https; http only with allow_private + Transport string `protobuf:"bytes,3,opt,name=transport,proto3" json:"transport,omitempty"` // "streamable_http" (default) | "sse" + Enabled bool `protobuf:"varint,4,opt,name=enabled,proto3" json:"enabled,omitempty"` + // Permit a private, loopback or link-local address, such as a server in the + // same cluster. Without it such an address is refused at registration and + // at every connection. + AllowPrivate bool `protobuf:"varint,5,opt,name=allow_private,json=allowPrivate,proto3" json:"allow_private,omitempty"` + TimeoutMs int32 `protobuf:"varint,6,opt,name=timeout_ms,json=timeoutMs,proto3" json:"timeout_ms,omitempty"` // one tool call, 0 = 60000 + MaxResponseBytes int64 `protobuf:"varint,7,opt,name=max_response_bytes,json=maxResponseBytes,proto3" json:"max_response_bytes,omitempty"` // one tool result, 0 = 4 MiB + ToolRpm int32 `protobuf:"varint,8,opt,name=tool_rpm,json=toolRpm,proto3" json:"tool_rpm,omitempty"` // calls per minute per client key and tool, 0 = unlimited; needs Redis + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *MCPServerSpec) Reset() { + *x = MCPServerSpec{} + mi := &file_ktms_v1_admin_proto_msgTypes[76] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *MCPServerSpec) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*MCPServerSpec) ProtoMessage() {} + +func (x *MCPServerSpec) ProtoReflect() protoreflect.Message { + mi := &file_ktms_v1_admin_proto_msgTypes[76] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use MCPServerSpec.ProtoReflect.Descriptor instead. +func (*MCPServerSpec) Descriptor() ([]byte, []int) { + return file_ktms_v1_admin_proto_rawDescGZIP(), []int{76} +} + +func (x *MCPServerSpec) GetName() string { + if x != nil { + return x.Name + } + return "" +} + +func (x *MCPServerSpec) GetUrl() string { + if x != nil { + return x.Url + } + return "" +} + +func (x *MCPServerSpec) GetTransport() string { + if x != nil { + return x.Transport + } + return "" +} + +func (x *MCPServerSpec) GetEnabled() bool { + if x != nil { + return x.Enabled + } + return false +} + +func (x *MCPServerSpec) GetAllowPrivate() bool { + if x != nil { + return x.AllowPrivate + } + return false +} + +func (x *MCPServerSpec) GetTimeoutMs() int32 { + if x != nil { + return x.TimeoutMs + } + return 0 +} + +func (x *MCPServerSpec) GetMaxResponseBytes() int64 { + if x != nil { + return x.MaxResponseBytes + } + return 0 +} + +func (x *MCPServerSpec) GetToolRpm() int32 { + if x != nil { + return x.ToolRpm + } + return 0 +} + +type UpsertMCPServerRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + Server *MCPServerSpec `protobuf:"bytes,1,opt,name=server,proto3" json:"server,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *UpsertMCPServerRequest) Reset() { + *x = UpsertMCPServerRequest{} + mi := &file_ktms_v1_admin_proto_msgTypes[77] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *UpsertMCPServerRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*UpsertMCPServerRequest) ProtoMessage() {} + +func (x *UpsertMCPServerRequest) ProtoReflect() protoreflect.Message { + mi := &file_ktms_v1_admin_proto_msgTypes[77] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use UpsertMCPServerRequest.ProtoReflect.Descriptor instead. +func (*UpsertMCPServerRequest) Descriptor() ([]byte, []int) { + return file_ktms_v1_admin_proto_rawDescGZIP(), []int{77} +} + +func (x *UpsertMCPServerRequest) GetServer() *MCPServerSpec { + if x != nil { + return x.Server + } + return nil +} + +type UpsertMCPServerResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *UpsertMCPServerResponse) Reset() { + *x = UpsertMCPServerResponse{} + mi := &file_ktms_v1_admin_proto_msgTypes[78] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *UpsertMCPServerResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*UpsertMCPServerResponse) ProtoMessage() {} + +func (x *UpsertMCPServerResponse) ProtoReflect() protoreflect.Message { + mi := &file_ktms_v1_admin_proto_msgTypes[78] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use UpsertMCPServerResponse.ProtoReflect.Descriptor instead. +func (*UpsertMCPServerResponse) Descriptor() ([]byte, []int) { + return file_ktms_v1_admin_proto_rawDescGZIP(), []int{78} +} + +// DeleteMCPServer removes the server and its credential. +type DeleteMCPServerRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + Name string `protobuf:"bytes,1,opt,name=name,proto3" json:"name,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *DeleteMCPServerRequest) Reset() { + *x = DeleteMCPServerRequest{} + mi := &file_ktms_v1_admin_proto_msgTypes[79] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *DeleteMCPServerRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*DeleteMCPServerRequest) ProtoMessage() {} + +func (x *DeleteMCPServerRequest) ProtoReflect() protoreflect.Message { + mi := &file_ktms_v1_admin_proto_msgTypes[79] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use DeleteMCPServerRequest.ProtoReflect.Descriptor instead. +func (*DeleteMCPServerRequest) Descriptor() ([]byte, []int) { + return file_ktms_v1_admin_proto_rawDescGZIP(), []int{79} +} + +func (x *DeleteMCPServerRequest) GetName() string { + if x != nil { + return x.Name + } + return "" +} + +type DeleteMCPServerResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *DeleteMCPServerResponse) Reset() { + *x = DeleteMCPServerResponse{} + mi := &file_ktms_v1_admin_proto_msgTypes[80] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *DeleteMCPServerResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*DeleteMCPServerResponse) ProtoMessage() {} + +func (x *DeleteMCPServerResponse) ProtoReflect() protoreflect.Message { + mi := &file_ktms_v1_admin_proto_msgTypes[80] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use DeleteMCPServerResponse.ProtoReflect.Descriptor instead. +func (*DeleteMCPServerResponse) Descriptor() ([]byte, []int) { + return file_ktms_v1_admin_proto_rawDescGZIP(), []int{80} +} + +// SetMCPCredential stores the header the gateway adds to every request it +// sends the server, such as Authorization: Bearer . Encrypted at rest +// with the vendor keys' keyring; never returned, logged or passed to clients. +type SetMCPCredentialRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + Server string `protobuf:"bytes,1,opt,name=server,proto3" json:"server,omitempty"` + Header string `protobuf:"bytes,2,opt,name=header,proto3" json:"header,omitempty"` // "" = Authorization + PlainValue string `protobuf:"bytes,3,opt,name=plain_value,json=plainValue,proto3" json:"plain_value,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *SetMCPCredentialRequest) Reset() { + *x = SetMCPCredentialRequest{} + mi := &file_ktms_v1_admin_proto_msgTypes[81] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *SetMCPCredentialRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*SetMCPCredentialRequest) ProtoMessage() {} + +func (x *SetMCPCredentialRequest) ProtoReflect() protoreflect.Message { + mi := &file_ktms_v1_admin_proto_msgTypes[81] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use SetMCPCredentialRequest.ProtoReflect.Descriptor instead. +func (*SetMCPCredentialRequest) Descriptor() ([]byte, []int) { + return file_ktms_v1_admin_proto_rawDescGZIP(), []int{81} +} + +func (x *SetMCPCredentialRequest) GetServer() string { + if x != nil { + return x.Server + } + return "" +} + +func (x *SetMCPCredentialRequest) GetHeader() string { + if x != nil { + return x.Header + } + return "" +} + +func (x *SetMCPCredentialRequest) GetPlainValue() string { + if x != nil { + return x.PlainValue + } + return "" +} + +type SetMCPCredentialResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *SetMCPCredentialResponse) Reset() { + *x = SetMCPCredentialResponse{} + mi := &file_ktms_v1_admin_proto_msgTypes[82] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *SetMCPCredentialResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*SetMCPCredentialResponse) ProtoMessage() {} + +func (x *SetMCPCredentialResponse) ProtoReflect() protoreflect.Message { + mi := &file_ktms_v1_admin_proto_msgTypes[82] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use SetMCPCredentialResponse.ProtoReflect.Descriptor instead. +func (*SetMCPCredentialResponse) Descriptor() ([]byte, []int) { + return file_ktms_v1_admin_proto_rawDescGZIP(), []int{82} +} + +type DeleteMCPCredentialRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + Server string `protobuf:"bytes,1,opt,name=server,proto3" json:"server,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *DeleteMCPCredentialRequest) Reset() { + *x = DeleteMCPCredentialRequest{} + mi := &file_ktms_v1_admin_proto_msgTypes[83] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *DeleteMCPCredentialRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*DeleteMCPCredentialRequest) ProtoMessage() {} + +func (x *DeleteMCPCredentialRequest) ProtoReflect() protoreflect.Message { + mi := &file_ktms_v1_admin_proto_msgTypes[83] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use DeleteMCPCredentialRequest.ProtoReflect.Descriptor instead. +func (*DeleteMCPCredentialRequest) Descriptor() ([]byte, []int) { + return file_ktms_v1_admin_proto_rawDescGZIP(), []int{83} +} + +func (x *DeleteMCPCredentialRequest) GetServer() string { + if x != nil { + return x.Server + } + return "" +} + +type DeleteMCPCredentialResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *DeleteMCPCredentialResponse) Reset() { + *x = DeleteMCPCredentialResponse{} + mi := &file_ktms_v1_admin_proto_msgTypes[84] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *DeleteMCPCredentialResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*DeleteMCPCredentialResponse) ProtoMessage() {} + +func (x *DeleteMCPCredentialResponse) ProtoReflect() protoreflect.Message { + mi := &file_ktms_v1_admin_proto_msgTypes[84] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use DeleteMCPCredentialResponse.ProtoReflect.Descriptor instead. +func (*DeleteMCPCredentialResponse) Descriptor() ([]byte, []int) { + return file_ktms_v1_admin_proto_rawDescGZIP(), []int{84} +} + +type MCPCredentialInfo struct { + state protoimpl.MessageState `protogen:"open.v1"` + Server string `protobuf:"bytes,1,opt,name=server,proto3" json:"server,omitempty"` + Header string `protobuf:"bytes,2,opt,name=header,proto3" json:"header,omitempty"` + KeyVersion int32 `protobuf:"varint,3,opt,name=key_version,json=keyVersion,proto3" json:"key_version,omitempty"` + CreatedAt *timestamppb.Timestamp `protobuf:"bytes,4,opt,name=created_at,json=createdAt,proto3" json:"created_at,omitempty"` + Fingerprint string `protobuf:"bytes,5,opt,name=fingerprint,proto3" json:"fingerprint,omitempty"` // as VendorKeyInfo.fingerprint + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *MCPCredentialInfo) Reset() { + *x = MCPCredentialInfo{} + mi := &file_ktms_v1_admin_proto_msgTypes[85] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *MCPCredentialInfo) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*MCPCredentialInfo) ProtoMessage() {} + +func (x *MCPCredentialInfo) ProtoReflect() protoreflect.Message { + mi := &file_ktms_v1_admin_proto_msgTypes[85] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use MCPCredentialInfo.ProtoReflect.Descriptor instead. +func (*MCPCredentialInfo) Descriptor() ([]byte, []int) { + return file_ktms_v1_admin_proto_rawDescGZIP(), []int{85} +} + +func (x *MCPCredentialInfo) GetServer() string { + if x != nil { + return x.Server + } + return "" +} + +func (x *MCPCredentialInfo) GetHeader() string { + if x != nil { + return x.Header + } + return "" +} + +func (x *MCPCredentialInfo) GetKeyVersion() int32 { + if x != nil { + return x.KeyVersion + } + return 0 +} + +func (x *MCPCredentialInfo) GetCreatedAt() *timestamppb.Timestamp { + if x != nil { + return x.CreatedAt + } + return nil +} + +func (x *MCPCredentialInfo) GetFingerprint() string { + if x != nil { + return x.Fingerprint + } + return "" +} + +// ListMCPTools asks the server for its tools through the gateway, with its +// credential, so an operator can write allowed_mcp. +type ListMCPToolsRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + Server string `protobuf:"bytes,1,opt,name=server,proto3" json:"server,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ListMCPToolsRequest) Reset() { + *x = ListMCPToolsRequest{} + mi := &file_ktms_v1_admin_proto_msgTypes[86] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ListMCPToolsRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ListMCPToolsRequest) ProtoMessage() {} + +func (x *ListMCPToolsRequest) ProtoReflect() protoreflect.Message { + mi := &file_ktms_v1_admin_proto_msgTypes[86] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ListMCPToolsRequest.ProtoReflect.Descriptor instead. +func (*ListMCPToolsRequest) Descriptor() ([]byte, []int) { + return file_ktms_v1_admin_proto_rawDescGZIP(), []int{86} +} + +func (x *ListMCPToolsRequest) GetServer() string { + if x != nil { + return x.Server + } + return "" +} + +type MCPToolInfo struct { + state protoimpl.MessageState `protogen:"open.v1"` + Name string `protobuf:"bytes,1,opt,name=name,proto3" json:"name,omitempty"` // as the server names it; /mcp shows server__name + Description string `protobuf:"bytes,2,opt,name=description,proto3" json:"description,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *MCPToolInfo) Reset() { + *x = MCPToolInfo{} + mi := &file_ktms_v1_admin_proto_msgTypes[87] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *MCPToolInfo) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*MCPToolInfo) ProtoMessage() {} + +func (x *MCPToolInfo) ProtoReflect() protoreflect.Message { + mi := &file_ktms_v1_admin_proto_msgTypes[87] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use MCPToolInfo.ProtoReflect.Descriptor instead. +func (*MCPToolInfo) Descriptor() ([]byte, []int) { + return file_ktms_v1_admin_proto_rawDescGZIP(), []int{87} +} + +func (x *MCPToolInfo) GetName() string { + if x != nil { + return x.Name + } + return "" +} + +func (x *MCPToolInfo) GetDescription() string { + if x != nil { + return x.Description + } + return "" +} + +type ListMCPToolsResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + Tools []*MCPToolInfo `protobuf:"bytes,1,rep,name=tools,proto3" json:"tools,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ListMCPToolsResponse) Reset() { + *x = ListMCPToolsResponse{} + mi := &file_ktms_v1_admin_proto_msgTypes[88] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ListMCPToolsResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ListMCPToolsResponse) ProtoMessage() {} + +func (x *ListMCPToolsResponse) ProtoReflect() protoreflect.Message { + mi := &file_ktms_v1_admin_proto_msgTypes[88] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ListMCPToolsResponse.ProtoReflect.Descriptor instead. +func (*ListMCPToolsResponse) Descriptor() ([]byte, []int) { + return file_ktms_v1_admin_proto_rawDescGZIP(), []int{88} +} + +func (x *ListMCPToolsResponse) GetTools() []*MCPToolInfo { + if x != nil { + return x.Tools + } + return nil +} + type ExplainRouteRequest struct { state protoimpl.MessageState `protogen:"open.v1"` ApiKeyName string `protobuf:"bytes,1,opt,name=api_key_name,json=apiKeyName,proto3" json:"api_key_name,omitempty"` // the client key the request would carry @@ -4213,7 +4927,7 @@ type ExplainRouteRequest struct { func (x *ExplainRouteRequest) Reset() { *x = ExplainRouteRequest{} - mi := &file_ktms_v1_admin_proto_msgTypes[76] + mi := &file_ktms_v1_admin_proto_msgTypes[89] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -4225,7 +4939,7 @@ func (x *ExplainRouteRequest) String() string { func (*ExplainRouteRequest) ProtoMessage() {} func (x *ExplainRouteRequest) ProtoReflect() protoreflect.Message { - mi := &file_ktms_v1_admin_proto_msgTypes[76] + mi := &file_ktms_v1_admin_proto_msgTypes[89] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -4238,7 +4952,7 @@ func (x *ExplainRouteRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use ExplainRouteRequest.ProtoReflect.Descriptor instead. func (*ExplainRouteRequest) Descriptor() ([]byte, []int) { - return file_ktms_v1_admin_proto_rawDescGZIP(), []int{76} + return file_ktms_v1_admin_proto_rawDescGZIP(), []int{89} } func (x *ExplainRouteRequest) GetApiKeyName() string { @@ -4305,7 +5019,7 @@ type RouteCandidate struct { func (x *RouteCandidate) Reset() { *x = RouteCandidate{} - mi := &file_ktms_v1_admin_proto_msgTypes[77] + mi := &file_ktms_v1_admin_proto_msgTypes[90] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -4317,7 +5031,7 @@ func (x *RouteCandidate) String() string { func (*RouteCandidate) ProtoMessage() {} func (x *RouteCandidate) ProtoReflect() protoreflect.Message { - mi := &file_ktms_v1_admin_proto_msgTypes[77] + mi := &file_ktms_v1_admin_proto_msgTypes[90] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -4330,7 +5044,7 @@ func (x *RouteCandidate) ProtoReflect() protoreflect.Message { // Deprecated: Use RouteCandidate.ProtoReflect.Descriptor instead. func (*RouteCandidate) Descriptor() ([]byte, []int) { - return file_ktms_v1_admin_proto_rawDescGZIP(), []int{77} + return file_ktms_v1_admin_proto_rawDescGZIP(), []int{90} } func (x *RouteCandidate) GetModelId() string { @@ -4382,7 +5096,7 @@ type ExplainRouteResponse struct { func (x *ExplainRouteResponse) Reset() { *x = ExplainRouteResponse{} - mi := &file_ktms_v1_admin_proto_msgTypes[78] + mi := &file_ktms_v1_admin_proto_msgTypes[91] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -4394,7 +5108,7 @@ func (x *ExplainRouteResponse) String() string { func (*ExplainRouteResponse) ProtoMessage() {} func (x *ExplainRouteResponse) ProtoReflect() protoreflect.Message { - mi := &file_ktms_v1_admin_proto_msgTypes[78] + mi := &file_ktms_v1_admin_proto_msgTypes[91] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -4407,7 +5121,7 @@ func (x *ExplainRouteResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use ExplainRouteResponse.ProtoReflect.Descriptor instead. func (*ExplainRouteResponse) Descriptor() ([]byte, []int) { - return file_ktms_v1_admin_proto_rawDescGZIP(), []int{78} + return file_ktms_v1_admin_proto_rawDescGZIP(), []int{91} } func (x *ExplainRouteResponse) GetRule() string { @@ -4561,7 +5275,7 @@ const file_ktms_v1_admin_proto_rawDesc = "" + "\n" + "created_at\x18\x04 \x01(\v2\x1a.google.protobuf.TimestampR\tcreatedAt\x12 \n" + "\vfingerprint\x18\x05 \x01(\tR\vfingerprint\"\x12\n" + - "\x10GetConfigRequest\"\xbd\x02\n" + + "\x10GetConfigRequest\"\xbb\x03\n" + "\x11GetConfigResponse\x12*\n" + "\x06models\x18\x01 \x03(\v2\x12.ktms.v1.ModelSpecR\x06models\x12'\n" + "\x05rules\x18\x02 \x03(\v2\x11.ktms.v1.RuleSpecR\x05rules\x127\n" + @@ -4569,17 +5283,22 @@ const file_ktms_v1_admin_proto_rawDesc = "" + "vendorKeys\x12\x18\n" + "\aversion\x18\x04 \x01(\x03R\aversion\x12D\n" + "\x10user_rate_limits\x18\x05 \x03(\v2\x1a.ktms.v1.UserRateLimitSpecR\x0euserRateLimits\x12:\n" + - "\fspend_limits\x18\x06 \x03(\v2\x17.ktms.v1.SpendLimitSpecR\vspendLimits\"\x9e\x01\n" + + "\fspend_limits\x18\x06 \x03(\v2\x17.ktms.v1.SpendLimitSpecR\vspendLimits\x127\n" + + "\vmcp_servers\x18\a \x03(\v2\x16.ktms.v1.MCPServerSpecR\n" + + "mcpServers\x12C\n" + + "\x0fmcp_credentials\x18\b \x03(\v2\x1a.ktms.v1.MCPCredentialInfoR\x0emcpCredentials\"\x9e\x01\n" + "\x11MintAPIKeyRequest\x12\x12\n" + "\x04name\x18\x01 \x01(\tR\x04name\x12\x16\n" + "\x06scopes\x18\x02 \x03(\tR\x06scopes\x12\x15\n" + "\x06org_id\x18\x03 \x01(\tR\x05orgId\x12\x17\n" + "\ateam_id\x18\x04 \x01(\tR\x06teamId\x12-\n" + - "\x06policy\x18\x05 \x01(\v2\x15.ktms.v1.APIKeyPolicyR\x06policy\"o\n" + + "\x06policy\x18\x05 \x01(\v2\x15.ktms.v1.APIKeyPolicyR\x06policy\"\x90\x01\n" + "\fAPIKeyPolicy\x12%\n" + "\x0eallowed_models\x18\x01 \x03(\tR\rallowedModels\x12\x1b\n" + "\trpm_limit\x18\x02 \x01(\x05R\brpmLimit\x12\x1b\n" + - "\ttpm_limit\x18\x03 \x01(\x03R\btpmLimit\"[\n" + + "\ttpm_limit\x18\x03 \x01(\x03R\btpmLimit\x12\x1f\n" + + "\vallowed_mcp\x18\x04 \x03(\tR\n" + + "allowedMcp\"[\n" + "\x16SetAPIKeyPolicyRequest\x12\x12\n" + "\x04name\x18\x01 \x01(\tR\x04name\x12-\n" + "\x06policy\x18\x02 \x01(\v2\x15.ktms.v1.APIKeyPolicyR\x06policy\"\x19\n" + @@ -4710,7 +5429,7 @@ const file_ktms_v1_admin_proto_rawDesc = "" + "\x02id\x18\x01 \x01(\tR\x02id\x12\x12\n" + "\x04name\x18\x02 \x01(\tR\x04name\x129\n" + "\n" + - "created_at\x18\x03 \x01(\v2\x1a.google.protobuf.TimestampR\tcreatedAt\"\xd1\x01\n" + + "created_at\x18\x03 \x01(\v2\x1a.google.protobuf.TimestampR\tcreatedAt\"\xf2\x01\n" + "\x04Team\x12\x0e\n" + "\x02id\x18\x01 \x01(\tR\x02id\x12\x15\n" + "\x06org_id\x18\x02 \x01(\tR\x05orgId\x12\x12\n" + @@ -4718,7 +5437,9 @@ const file_ktms_v1_admin_proto_rawDesc = "" + "\x12monthly_budget_usd\x18\x04 \x01(\x01R\x10monthlyBudgetUsd\x12%\n" + "\x0eallowed_models\x18\x05 \x03(\tR\rallowedModels\x129\n" + "\n" + - "created_at\x18\x06 \x01(\v2\x1a.google.protobuf.TimestampR\tcreatedAt\"/\n" + + "created_at\x18\x06 \x01(\v2\x1a.google.protobuf.TimestampR\tcreatedAt\x12\x1f\n" + + "\vallowed_mcp\x18\a \x03(\tR\n" + + "allowedMcp\"/\n" + "\x19CreateOrganizationRequest\x12\x12\n" + "\x04name\x18\x01 \x01(\tR\x04name\"W\n" + "\x1aCreateOrganizationResponse\x129\n" + @@ -4728,12 +5449,14 @@ const file_ktms_v1_admin_proto_rawDesc = "" + "\rorganizations\x18\x01 \x03(\v2\x15.ktms.v1.OrganizationR\rorganizations\"+\n" + "\x19DeleteOrganizationRequest\x12\x0e\n" + "\x02id\x18\x01 \x01(\tR\x02id\"\x1c\n" + - "\x1aDeleteOrganizationResponse\"\x93\x01\n" + + "\x1aDeleteOrganizationResponse\"\xb4\x01\n" + "\x11CreateTeamRequest\x12\x15\n" + "\x06org_id\x18\x01 \x01(\tR\x05orgId\x12\x12\n" + "\x04name\x18\x02 \x01(\tR\x04name\x12,\n" + "\x12monthly_budget_usd\x18\x03 \x01(\x01R\x10monthlyBudgetUsd\x12%\n" + - "\x0eallowed_models\x18\x04 \x03(\tR\rallowedModels\"7\n" + + "\x0eallowed_models\x18\x04 \x03(\tR\rallowedModels\x12\x1f\n" + + "\vallowed_mcp\x18\x05 \x03(\tR\n" + + "allowedMcp\"7\n" + "\x12CreateTeamResponse\x12!\n" + "\x04team\x18\x01 \x01(\v2\r.ktms.v1.TeamR\x04team\" \n" + "\x0eGetTeamRequest\x12\x0e\n" + @@ -4743,12 +5466,14 @@ const file_ktms_v1_admin_proto_rawDesc = "" + "\x10ListTeamsRequest\x12\x15\n" + "\x06org_id\x18\x01 \x01(\tR\x05orgId\"8\n" + "\x11ListTeamsResponse\x12#\n" + - "\x05teams\x18\x01 \x03(\v2\r.ktms.v1.TeamR\x05teams\"\x8c\x01\n" + + "\x05teams\x18\x01 \x03(\v2\r.ktms.v1.TeamR\x05teams\"\xad\x01\n" + "\x11UpdateTeamRequest\x12\x0e\n" + "\x02id\x18\x01 \x01(\tR\x02id\x12\x12\n" + "\x04name\x18\x02 \x01(\tR\x04name\x12,\n" + "\x12monthly_budget_usd\x18\x03 \x01(\x01R\x10monthlyBudgetUsd\x12%\n" + - "\x0eallowed_models\x18\x04 \x03(\tR\rallowedModels\"7\n" + + "\x0eallowed_models\x18\x04 \x03(\tR\rallowedModels\x12\x1f\n" + + "\vallowed_mcp\x18\x05 \x03(\tR\n" + + "allowedMcp\"7\n" + "\x12UpdateTeamResponse\x12!\n" + "\x04team\x18\x01 \x01(\v2\r.ktms.v1.TeamR\x04team\"#\n" + "\x11DeleteTeamRequest\x12\x0e\n" + @@ -4773,7 +5498,47 @@ const file_ktms_v1_admin_proto_rawDesc = "" + "\tresets_at\x18\x05 \x01(\v2\x1a.google.protobuf.TimestampR\bresetsAt\x12\x18\n" + "\ametered\x18\x06 \x01(\bR\ametered\"B\n" + "\x10GetSpendResponse\x12.\n" + - "\awindows\x18\x01 \x03(\v2\x14.ktms.v1.SpendWindowR\awindows\"\xf5\x02\n" + + "\awindows\x18\x01 \x03(\v2\x14.ktms.v1.SpendWindowR\awindows\"\xfa\x01\n" + + "\rMCPServerSpec\x12\x12\n" + + "\x04name\x18\x01 \x01(\tR\x04name\x12\x10\n" + + "\x03url\x18\x02 \x01(\tR\x03url\x12\x1c\n" + + "\ttransport\x18\x03 \x01(\tR\ttransport\x12\x18\n" + + "\aenabled\x18\x04 \x01(\bR\aenabled\x12#\n" + + "\rallow_private\x18\x05 \x01(\bR\fallowPrivate\x12\x1d\n" + + "\n" + + "timeout_ms\x18\x06 \x01(\x05R\ttimeoutMs\x12,\n" + + "\x12max_response_bytes\x18\a \x01(\x03R\x10maxResponseBytes\x12\x19\n" + + "\btool_rpm\x18\b \x01(\x05R\atoolRpm\"H\n" + + "\x16UpsertMCPServerRequest\x12.\n" + + "\x06server\x18\x01 \x01(\v2\x16.ktms.v1.MCPServerSpecR\x06server\"\x19\n" + + "\x17UpsertMCPServerResponse\",\n" + + "\x16DeleteMCPServerRequest\x12\x12\n" + + "\x04name\x18\x01 \x01(\tR\x04name\"\x19\n" + + "\x17DeleteMCPServerResponse\"j\n" + + "\x17SetMCPCredentialRequest\x12\x16\n" + + "\x06server\x18\x01 \x01(\tR\x06server\x12\x16\n" + + "\x06header\x18\x02 \x01(\tR\x06header\x12\x1f\n" + + "\vplain_value\x18\x03 \x01(\tR\n" + + "plainValue\"\x1a\n" + + "\x18SetMCPCredentialResponse\"4\n" + + "\x1aDeleteMCPCredentialRequest\x12\x16\n" + + "\x06server\x18\x01 \x01(\tR\x06server\"\x1d\n" + + "\x1bDeleteMCPCredentialResponse\"\xc1\x01\n" + + "\x11MCPCredentialInfo\x12\x16\n" + + "\x06server\x18\x01 \x01(\tR\x06server\x12\x16\n" + + "\x06header\x18\x02 \x01(\tR\x06header\x12\x1f\n" + + "\vkey_version\x18\x03 \x01(\x05R\n" + + "keyVersion\x129\n" + + "\n" + + "created_at\x18\x04 \x01(\v2\x1a.google.protobuf.TimestampR\tcreatedAt\x12 \n" + + "\vfingerprint\x18\x05 \x01(\tR\vfingerprint\"-\n" + + "\x13ListMCPToolsRequest\x12\x16\n" + + "\x06server\x18\x01 \x01(\tR\x06server\"C\n" + + "\vMCPToolInfo\x12\x12\n" + + "\x04name\x18\x01 \x01(\tR\x04name\x12 \n" + + "\vdescription\x18\x02 \x01(\tR\vdescription\"B\n" + + "\x14ListMCPToolsResponse\x12*\n" + + "\x05tools\x18\x01 \x03(\v2\x14.ktms.v1.MCPToolInfoR\x05tools\"\xf5\x02\n" + "\x13ExplainRouteRequest\x12 \n" + "\fapi_key_name\x18\x01 \x01(\tR\n" + "apiKeyName\x12'\n" + @@ -4810,7 +5575,7 @@ const file_ktms_v1_admin_proto_rawDesc = "" + "\x11shadow_percentage\x18\t \x01(\x01R\x10shadowPercentage\x12#\n" + "\rjudge_skipped\x18\n" + " \x01(\bR\fjudgeSkipped\x12\x14\n" + - "\x05error\x18\v \x01(\tR\x05error2\xb7\x14\n" + + "\x05error\x18\v \x01(\tR\x05error2\xeb\x17\n" + "\fAdminService\x12H\n" + "\vUpsertModel\x12\x1b.ktms.v1.UpsertModelRequest\x1a\x1c.ktms.v1.UpsertModelResponse\x12H\n" + "\vDeleteModel\x12\x1b.ktms.v1.DeleteModelRequest\x1a\x1c.ktms.v1.DeleteModelResponse\x12K\n" + @@ -4848,7 +5613,12 @@ const file_ktms_v1_admin_proto_rawDesc = "" + "DeleteTeam\x12\x1a.ktms.v1.DeleteTeamRequest\x1a\x1b.ktms.v1.DeleteTeamResponse\x12K\n" + "\fGetTeamSpend\x12\x1c.ktms.v1.GetTeamSpendRequest\x1a\x1d.ktms.v1.GetTeamSpendResponse\x12?\n" + "\bGetSpend\x12\x18.ktms.v1.GetSpendRequest\x1a\x19.ktms.v1.GetSpendResponse\x12K\n" + - "\fExplainRoute\x12\x1c.ktms.v1.ExplainRouteRequest\x1a\x1d.ktms.v1.ExplainRouteResponseB\x85\x01\n" + + "\fExplainRoute\x12\x1c.ktms.v1.ExplainRouteRequest\x1a\x1d.ktms.v1.ExplainRouteResponse\x12T\n" + + "\x0fUpsertMCPServer\x12\x1f.ktms.v1.UpsertMCPServerRequest\x1a .ktms.v1.UpsertMCPServerResponse\x12T\n" + + "\x0fDeleteMCPServer\x12\x1f.ktms.v1.DeleteMCPServerRequest\x1a .ktms.v1.DeleteMCPServerResponse\x12W\n" + + "\x10SetMCPCredential\x12 .ktms.v1.SetMCPCredentialRequest\x1a!.ktms.v1.SetMCPCredentialResponse\x12`\n" + + "\x13DeleteMCPCredential\x12#.ktms.v1.DeleteMCPCredentialRequest\x1a$.ktms.v1.DeleteMCPCredentialResponse\x12K\n" + + "\fListMCPTools\x12\x1c.ktms.v1.ListMCPToolsRequest\x1a\x1d.ktms.v1.ListMCPToolsResponseB\x85\x01\n" + "\vcom.ktms.v1B\n" + "AdminProtoP\x01Z-github.com/ai-process/ktms/gen/ktms/v1;ktmsv1\xa2\x02\x03KXX\xaa\x02\aKtms.V1\xca\x02\aKtms\\V1\xe2\x02\x13Ktms\\V1\\GPBMetadata\xea\x02\bKtms::V1b\x06proto3" @@ -4864,7 +5634,7 @@ func file_ktms_v1_admin_proto_rawDescGZIP() []byte { return file_ktms_v1_admin_proto_rawDescData } -var file_ktms_v1_admin_proto_msgTypes = make([]protoimpl.MessageInfo, 81) +var file_ktms_v1_admin_proto_msgTypes = make([]protoimpl.MessageInfo, 94) var file_ktms_v1_admin_proto_goTypes = []any{ (*ModelSpec)(nil), // 0: ktms.v1.ModelSpec (*RuleSpec)(nil), // 1: ktms.v1.RuleSpec @@ -4942,126 +5712,154 @@ var file_ktms_v1_admin_proto_goTypes = []any{ (*GetSpendRequest)(nil), // 73: ktms.v1.GetSpendRequest (*SpendWindow)(nil), // 74: ktms.v1.SpendWindow (*GetSpendResponse)(nil), // 75: ktms.v1.GetSpendResponse - (*ExplainRouteRequest)(nil), // 76: ktms.v1.ExplainRouteRequest - (*RouteCandidate)(nil), // 77: ktms.v1.RouteCandidate - (*ExplainRouteResponse)(nil), // 78: ktms.v1.ExplainRouteResponse - nil, // 79: ktms.v1.RuleSpec.AttributesEntry - nil, // 80: ktms.v1.ExplainRouteRequest.AttributesEntry - (Effort)(0), // 81: ktms.v1.Effort - (*timestamppb.Timestamp)(nil), // 82: google.protobuf.Timestamp - (*durationpb.Duration)(nil), // 83: google.protobuf.Duration + (*MCPServerSpec)(nil), // 76: ktms.v1.MCPServerSpec + (*UpsertMCPServerRequest)(nil), // 77: ktms.v1.UpsertMCPServerRequest + (*UpsertMCPServerResponse)(nil), // 78: ktms.v1.UpsertMCPServerResponse + (*DeleteMCPServerRequest)(nil), // 79: ktms.v1.DeleteMCPServerRequest + (*DeleteMCPServerResponse)(nil), // 80: ktms.v1.DeleteMCPServerResponse + (*SetMCPCredentialRequest)(nil), // 81: ktms.v1.SetMCPCredentialRequest + (*SetMCPCredentialResponse)(nil), // 82: ktms.v1.SetMCPCredentialResponse + (*DeleteMCPCredentialRequest)(nil), // 83: ktms.v1.DeleteMCPCredentialRequest + (*DeleteMCPCredentialResponse)(nil), // 84: ktms.v1.DeleteMCPCredentialResponse + (*MCPCredentialInfo)(nil), // 85: ktms.v1.MCPCredentialInfo + (*ListMCPToolsRequest)(nil), // 86: ktms.v1.ListMCPToolsRequest + (*MCPToolInfo)(nil), // 87: ktms.v1.MCPToolInfo + (*ListMCPToolsResponse)(nil), // 88: ktms.v1.ListMCPToolsResponse + (*ExplainRouteRequest)(nil), // 89: ktms.v1.ExplainRouteRequest + (*RouteCandidate)(nil), // 90: ktms.v1.RouteCandidate + (*ExplainRouteResponse)(nil), // 91: ktms.v1.ExplainRouteResponse + nil, // 92: ktms.v1.RuleSpec.AttributesEntry + nil, // 93: ktms.v1.ExplainRouteRequest.AttributesEntry + (Effort)(0), // 94: ktms.v1.Effort + (*timestamppb.Timestamp)(nil), // 95: google.protobuf.Timestamp + (*durationpb.Duration)(nil), // 96: google.protobuf.Duration } var file_ktms_v1_admin_proto_depIdxs = []int32{ - 81, // 0: ktms.v1.ModelSpec.efforts:type_name -> ktms.v1.Effort - 81, // 1: ktms.v1.RuleSpec.effort:type_name -> ktms.v1.Effort - 79, // 2: ktms.v1.RuleSpec.attributes:type_name -> ktms.v1.RuleSpec.AttributesEntry + 94, // 0: ktms.v1.ModelSpec.efforts:type_name -> ktms.v1.Effort + 94, // 1: ktms.v1.RuleSpec.effort:type_name -> ktms.v1.Effort + 92, // 2: ktms.v1.RuleSpec.attributes:type_name -> ktms.v1.RuleSpec.AttributesEntry 0, // 3: ktms.v1.UpsertModelRequest.model:type_name -> ktms.v1.ModelSpec 1, // 4: ktms.v1.ReplaceRulesRequest.rules:type_name -> ktms.v1.RuleSpec - 82, // 5: ktms.v1.VendorKeyInfo.created_at:type_name -> google.protobuf.Timestamp + 95, // 5: ktms.v1.VendorKeyInfo.created_at:type_name -> google.protobuf.Timestamp 0, // 6: ktms.v1.GetConfigResponse.models:type_name -> ktms.v1.ModelSpec 1, // 7: ktms.v1.GetConfigResponse.rules:type_name -> ktms.v1.RuleSpec 12, // 8: ktms.v1.GetConfigResponse.vendor_keys:type_name -> ktms.v1.VendorKeyInfo 39, // 9: ktms.v1.GetConfigResponse.user_rate_limits:type_name -> ktms.v1.UserRateLimitSpec 46, // 10: ktms.v1.GetConfigResponse.spend_limits:type_name -> ktms.v1.SpendLimitSpec - 16, // 11: ktms.v1.MintAPIKeyRequest.policy:type_name -> ktms.v1.APIKeyPolicy - 16, // 12: ktms.v1.SetAPIKeyPolicyRequest.policy:type_name -> ktms.v1.APIKeyPolicy - 82, // 13: ktms.v1.APIKeyInfo.created_at:type_name -> google.protobuf.Timestamp - 16, // 14: ktms.v1.APIKeyInfo.policy:type_name -> ktms.v1.APIKeyPolicy - 23, // 15: ktms.v1.ListAPIKeysResponse.keys:type_name -> ktms.v1.APIKeyInfo - 82, // 16: ktms.v1.NodeKeyInfo.created_at:type_name -> google.protobuf.Timestamp - 32, // 17: ktms.v1.ListNodeKeysResponse.keys:type_name -> ktms.v1.NodeKeyInfo - 82, // 18: ktms.v1.NodeInfo.last_heartbeat:type_name -> google.protobuf.Timestamp - 35, // 19: ktms.v1.ListNodesResponse.nodes:type_name -> ktms.v1.NodeInfo - 39, // 20: ktms.v1.SetUserRateLimitRequest.rate_limit:type_name -> ktms.v1.UserRateLimitSpec - 39, // 21: ktms.v1.ListUserRateLimitsResponse.rate_limits:type_name -> ktms.v1.UserRateLimitSpec - 46, // 22: ktms.v1.SetSpendLimitRequest.spend_limit:type_name -> ktms.v1.SpendLimitSpec - 46, // 23: ktms.v1.ListSpendLimitsResponse.spend_limits:type_name -> ktms.v1.SpendLimitSpec - 82, // 24: ktms.v1.Organization.created_at:type_name -> google.protobuf.Timestamp - 82, // 25: ktms.v1.Team.created_at:type_name -> google.protobuf.Timestamp - 53, // 26: ktms.v1.CreateOrganizationResponse.organization:type_name -> ktms.v1.Organization - 53, // 27: ktms.v1.ListOrganizationsResponse.organizations:type_name -> ktms.v1.Organization - 54, // 28: ktms.v1.CreateTeamResponse.team:type_name -> ktms.v1.Team - 54, // 29: ktms.v1.GetTeamResponse.team:type_name -> ktms.v1.Team - 54, // 30: ktms.v1.ListTeamsResponse.teams:type_name -> ktms.v1.Team - 54, // 31: ktms.v1.UpdateTeamResponse.team:type_name -> ktms.v1.Team - 82, // 32: ktms.v1.GetTeamSpendResponse.resets_at:type_name -> google.protobuf.Timestamp - 82, // 33: ktms.v1.SpendWindow.resets_at:type_name -> google.protobuf.Timestamp - 74, // 34: ktms.v1.GetSpendResponse.windows:type_name -> ktms.v1.SpendWindow - 81, // 35: ktms.v1.ExplainRouteRequest.effort:type_name -> ktms.v1.Effort - 80, // 36: ktms.v1.ExplainRouteRequest.attributes:type_name -> ktms.v1.ExplainRouteRequest.AttributesEntry - 83, // 37: ktms.v1.RouteCandidate.paused_for:type_name -> google.protobuf.Duration - 77, // 38: ktms.v1.ExplainRouteResponse.candidates:type_name -> ktms.v1.RouteCandidate - 2, // 39: ktms.v1.AdminService.UpsertModel:input_type -> ktms.v1.UpsertModelRequest - 4, // 40: ktms.v1.AdminService.DeleteModel:input_type -> ktms.v1.DeleteModelRequest - 6, // 41: ktms.v1.AdminService.ReplaceRules:input_type -> ktms.v1.ReplaceRulesRequest - 8, // 42: ktms.v1.AdminService.SetVendorKey:input_type -> ktms.v1.SetVendorKeyRequest - 10, // 43: ktms.v1.AdminService.DeleteVendorKey:input_type -> ktms.v1.DeleteVendorKeyRequest - 13, // 44: ktms.v1.AdminService.GetConfig:input_type -> ktms.v1.GetConfigRequest - 15, // 45: ktms.v1.AdminService.MintAPIKey:input_type -> ktms.v1.MintAPIKeyRequest - 20, // 46: ktms.v1.AdminService.RevokeAPIKey:input_type -> ktms.v1.RevokeAPIKeyRequest - 22, // 47: ktms.v1.AdminService.ListAPIKeys:input_type -> ktms.v1.ListAPIKeysRequest - 17, // 48: ktms.v1.AdminService.SetAPIKeyPolicy:input_type -> ktms.v1.SetAPIKeyPolicyRequest - 25, // 49: ktms.v1.AdminService.RotateEncryption:input_type -> ktms.v1.RotateEncryptionRequest - 27, // 50: ktms.v1.AdminService.MintNodeKey:input_type -> ktms.v1.MintNodeKeyRequest - 29, // 51: ktms.v1.AdminService.RevokeNodeKey:input_type -> ktms.v1.RevokeNodeKeyRequest - 31, // 52: ktms.v1.AdminService.ListNodeKeys:input_type -> ktms.v1.ListNodeKeysRequest - 34, // 53: ktms.v1.AdminService.ListNodes:input_type -> ktms.v1.ListNodesRequest - 37, // 54: ktms.v1.AdminService.DrainNode:input_type -> ktms.v1.DrainNodeRequest - 40, // 55: ktms.v1.AdminService.SetUserRateLimit:input_type -> ktms.v1.SetUserRateLimitRequest - 42, // 56: ktms.v1.AdminService.DeleteUserRateLimit:input_type -> ktms.v1.DeleteUserRateLimitRequest - 44, // 57: ktms.v1.AdminService.ListUserRateLimits:input_type -> ktms.v1.ListUserRateLimitsRequest - 47, // 58: ktms.v1.AdminService.SetSpendLimit:input_type -> ktms.v1.SetSpendLimitRequest - 49, // 59: ktms.v1.AdminService.DeleteSpendLimit:input_type -> ktms.v1.DeleteSpendLimitRequest - 51, // 60: ktms.v1.AdminService.ListSpendLimits:input_type -> ktms.v1.ListSpendLimitsRequest - 55, // 61: ktms.v1.AdminService.CreateOrganization:input_type -> ktms.v1.CreateOrganizationRequest - 57, // 62: ktms.v1.AdminService.ListOrganizations:input_type -> ktms.v1.ListOrganizationsRequest - 59, // 63: ktms.v1.AdminService.DeleteOrganization:input_type -> ktms.v1.DeleteOrganizationRequest - 61, // 64: ktms.v1.AdminService.CreateTeam:input_type -> ktms.v1.CreateTeamRequest - 63, // 65: ktms.v1.AdminService.GetTeam:input_type -> ktms.v1.GetTeamRequest - 65, // 66: ktms.v1.AdminService.ListTeams:input_type -> ktms.v1.ListTeamsRequest - 67, // 67: ktms.v1.AdminService.UpdateTeam:input_type -> ktms.v1.UpdateTeamRequest - 69, // 68: ktms.v1.AdminService.DeleteTeam:input_type -> ktms.v1.DeleteTeamRequest - 71, // 69: ktms.v1.AdminService.GetTeamSpend:input_type -> ktms.v1.GetTeamSpendRequest - 73, // 70: ktms.v1.AdminService.GetSpend:input_type -> ktms.v1.GetSpendRequest - 76, // 71: ktms.v1.AdminService.ExplainRoute:input_type -> ktms.v1.ExplainRouteRequest - 3, // 72: ktms.v1.AdminService.UpsertModel:output_type -> ktms.v1.UpsertModelResponse - 5, // 73: ktms.v1.AdminService.DeleteModel:output_type -> ktms.v1.DeleteModelResponse - 7, // 74: ktms.v1.AdminService.ReplaceRules:output_type -> ktms.v1.ReplaceRulesResponse - 9, // 75: ktms.v1.AdminService.SetVendorKey:output_type -> ktms.v1.SetVendorKeyResponse - 11, // 76: ktms.v1.AdminService.DeleteVendorKey:output_type -> ktms.v1.DeleteVendorKeyResponse - 14, // 77: ktms.v1.AdminService.GetConfig:output_type -> ktms.v1.GetConfigResponse - 19, // 78: ktms.v1.AdminService.MintAPIKey:output_type -> ktms.v1.MintAPIKeyResponse - 21, // 79: ktms.v1.AdminService.RevokeAPIKey:output_type -> ktms.v1.RevokeAPIKeyResponse - 24, // 80: ktms.v1.AdminService.ListAPIKeys:output_type -> ktms.v1.ListAPIKeysResponse - 18, // 81: ktms.v1.AdminService.SetAPIKeyPolicy:output_type -> ktms.v1.SetAPIKeyPolicyResponse - 26, // 82: ktms.v1.AdminService.RotateEncryption:output_type -> ktms.v1.RotateEncryptionResponse - 28, // 83: ktms.v1.AdminService.MintNodeKey:output_type -> ktms.v1.MintNodeKeyResponse - 30, // 84: ktms.v1.AdminService.RevokeNodeKey:output_type -> ktms.v1.RevokeNodeKeyResponse - 33, // 85: ktms.v1.AdminService.ListNodeKeys:output_type -> ktms.v1.ListNodeKeysResponse - 36, // 86: ktms.v1.AdminService.ListNodes:output_type -> ktms.v1.ListNodesResponse - 38, // 87: ktms.v1.AdminService.DrainNode:output_type -> ktms.v1.DrainNodeResponse - 41, // 88: ktms.v1.AdminService.SetUserRateLimit:output_type -> ktms.v1.SetUserRateLimitResponse - 43, // 89: ktms.v1.AdminService.DeleteUserRateLimit:output_type -> ktms.v1.DeleteUserRateLimitResponse - 45, // 90: ktms.v1.AdminService.ListUserRateLimits:output_type -> ktms.v1.ListUserRateLimitsResponse - 48, // 91: ktms.v1.AdminService.SetSpendLimit:output_type -> ktms.v1.SetSpendLimitResponse - 50, // 92: ktms.v1.AdminService.DeleteSpendLimit:output_type -> ktms.v1.DeleteSpendLimitResponse - 52, // 93: ktms.v1.AdminService.ListSpendLimits:output_type -> ktms.v1.ListSpendLimitsResponse - 56, // 94: ktms.v1.AdminService.CreateOrganization:output_type -> ktms.v1.CreateOrganizationResponse - 58, // 95: ktms.v1.AdminService.ListOrganizations:output_type -> ktms.v1.ListOrganizationsResponse - 60, // 96: ktms.v1.AdminService.DeleteOrganization:output_type -> ktms.v1.DeleteOrganizationResponse - 62, // 97: ktms.v1.AdminService.CreateTeam:output_type -> ktms.v1.CreateTeamResponse - 64, // 98: ktms.v1.AdminService.GetTeam:output_type -> ktms.v1.GetTeamResponse - 66, // 99: ktms.v1.AdminService.ListTeams:output_type -> ktms.v1.ListTeamsResponse - 68, // 100: ktms.v1.AdminService.UpdateTeam:output_type -> ktms.v1.UpdateTeamResponse - 70, // 101: ktms.v1.AdminService.DeleteTeam:output_type -> ktms.v1.DeleteTeamResponse - 72, // 102: ktms.v1.AdminService.GetTeamSpend:output_type -> ktms.v1.GetTeamSpendResponse - 75, // 103: ktms.v1.AdminService.GetSpend:output_type -> ktms.v1.GetSpendResponse - 78, // 104: ktms.v1.AdminService.ExplainRoute:output_type -> ktms.v1.ExplainRouteResponse - 72, // [72:105] is the sub-list for method output_type - 39, // [39:72] is the sub-list for method input_type - 39, // [39:39] is the sub-list for extension type_name - 39, // [39:39] is the sub-list for extension extendee - 0, // [0:39] is the sub-list for field type_name + 76, // 11: ktms.v1.GetConfigResponse.mcp_servers:type_name -> ktms.v1.MCPServerSpec + 85, // 12: ktms.v1.GetConfigResponse.mcp_credentials:type_name -> ktms.v1.MCPCredentialInfo + 16, // 13: ktms.v1.MintAPIKeyRequest.policy:type_name -> ktms.v1.APIKeyPolicy + 16, // 14: ktms.v1.SetAPIKeyPolicyRequest.policy:type_name -> ktms.v1.APIKeyPolicy + 95, // 15: ktms.v1.APIKeyInfo.created_at:type_name -> google.protobuf.Timestamp + 16, // 16: ktms.v1.APIKeyInfo.policy:type_name -> ktms.v1.APIKeyPolicy + 23, // 17: ktms.v1.ListAPIKeysResponse.keys:type_name -> ktms.v1.APIKeyInfo + 95, // 18: ktms.v1.NodeKeyInfo.created_at:type_name -> google.protobuf.Timestamp + 32, // 19: ktms.v1.ListNodeKeysResponse.keys:type_name -> ktms.v1.NodeKeyInfo + 95, // 20: ktms.v1.NodeInfo.last_heartbeat:type_name -> google.protobuf.Timestamp + 35, // 21: ktms.v1.ListNodesResponse.nodes:type_name -> ktms.v1.NodeInfo + 39, // 22: ktms.v1.SetUserRateLimitRequest.rate_limit:type_name -> ktms.v1.UserRateLimitSpec + 39, // 23: ktms.v1.ListUserRateLimitsResponse.rate_limits:type_name -> ktms.v1.UserRateLimitSpec + 46, // 24: ktms.v1.SetSpendLimitRequest.spend_limit:type_name -> ktms.v1.SpendLimitSpec + 46, // 25: ktms.v1.ListSpendLimitsResponse.spend_limits:type_name -> ktms.v1.SpendLimitSpec + 95, // 26: ktms.v1.Organization.created_at:type_name -> google.protobuf.Timestamp + 95, // 27: ktms.v1.Team.created_at:type_name -> google.protobuf.Timestamp + 53, // 28: ktms.v1.CreateOrganizationResponse.organization:type_name -> ktms.v1.Organization + 53, // 29: ktms.v1.ListOrganizationsResponse.organizations:type_name -> ktms.v1.Organization + 54, // 30: ktms.v1.CreateTeamResponse.team:type_name -> ktms.v1.Team + 54, // 31: ktms.v1.GetTeamResponse.team:type_name -> ktms.v1.Team + 54, // 32: ktms.v1.ListTeamsResponse.teams:type_name -> ktms.v1.Team + 54, // 33: ktms.v1.UpdateTeamResponse.team:type_name -> ktms.v1.Team + 95, // 34: ktms.v1.GetTeamSpendResponse.resets_at:type_name -> google.protobuf.Timestamp + 95, // 35: ktms.v1.SpendWindow.resets_at:type_name -> google.protobuf.Timestamp + 74, // 36: ktms.v1.GetSpendResponse.windows:type_name -> ktms.v1.SpendWindow + 76, // 37: ktms.v1.UpsertMCPServerRequest.server:type_name -> ktms.v1.MCPServerSpec + 95, // 38: ktms.v1.MCPCredentialInfo.created_at:type_name -> google.protobuf.Timestamp + 87, // 39: ktms.v1.ListMCPToolsResponse.tools:type_name -> ktms.v1.MCPToolInfo + 94, // 40: ktms.v1.ExplainRouteRequest.effort:type_name -> ktms.v1.Effort + 93, // 41: ktms.v1.ExplainRouteRequest.attributes:type_name -> ktms.v1.ExplainRouteRequest.AttributesEntry + 96, // 42: ktms.v1.RouteCandidate.paused_for:type_name -> google.protobuf.Duration + 90, // 43: ktms.v1.ExplainRouteResponse.candidates:type_name -> ktms.v1.RouteCandidate + 2, // 44: ktms.v1.AdminService.UpsertModel:input_type -> ktms.v1.UpsertModelRequest + 4, // 45: ktms.v1.AdminService.DeleteModel:input_type -> ktms.v1.DeleteModelRequest + 6, // 46: ktms.v1.AdminService.ReplaceRules:input_type -> ktms.v1.ReplaceRulesRequest + 8, // 47: ktms.v1.AdminService.SetVendorKey:input_type -> ktms.v1.SetVendorKeyRequest + 10, // 48: ktms.v1.AdminService.DeleteVendorKey:input_type -> ktms.v1.DeleteVendorKeyRequest + 13, // 49: ktms.v1.AdminService.GetConfig:input_type -> ktms.v1.GetConfigRequest + 15, // 50: ktms.v1.AdminService.MintAPIKey:input_type -> ktms.v1.MintAPIKeyRequest + 20, // 51: ktms.v1.AdminService.RevokeAPIKey:input_type -> ktms.v1.RevokeAPIKeyRequest + 22, // 52: ktms.v1.AdminService.ListAPIKeys:input_type -> ktms.v1.ListAPIKeysRequest + 17, // 53: ktms.v1.AdminService.SetAPIKeyPolicy:input_type -> ktms.v1.SetAPIKeyPolicyRequest + 25, // 54: ktms.v1.AdminService.RotateEncryption:input_type -> ktms.v1.RotateEncryptionRequest + 27, // 55: ktms.v1.AdminService.MintNodeKey:input_type -> ktms.v1.MintNodeKeyRequest + 29, // 56: ktms.v1.AdminService.RevokeNodeKey:input_type -> ktms.v1.RevokeNodeKeyRequest + 31, // 57: ktms.v1.AdminService.ListNodeKeys:input_type -> ktms.v1.ListNodeKeysRequest + 34, // 58: ktms.v1.AdminService.ListNodes:input_type -> ktms.v1.ListNodesRequest + 37, // 59: ktms.v1.AdminService.DrainNode:input_type -> ktms.v1.DrainNodeRequest + 40, // 60: ktms.v1.AdminService.SetUserRateLimit:input_type -> ktms.v1.SetUserRateLimitRequest + 42, // 61: ktms.v1.AdminService.DeleteUserRateLimit:input_type -> ktms.v1.DeleteUserRateLimitRequest + 44, // 62: ktms.v1.AdminService.ListUserRateLimits:input_type -> ktms.v1.ListUserRateLimitsRequest + 47, // 63: ktms.v1.AdminService.SetSpendLimit:input_type -> ktms.v1.SetSpendLimitRequest + 49, // 64: ktms.v1.AdminService.DeleteSpendLimit:input_type -> ktms.v1.DeleteSpendLimitRequest + 51, // 65: ktms.v1.AdminService.ListSpendLimits:input_type -> ktms.v1.ListSpendLimitsRequest + 55, // 66: ktms.v1.AdminService.CreateOrganization:input_type -> ktms.v1.CreateOrganizationRequest + 57, // 67: ktms.v1.AdminService.ListOrganizations:input_type -> ktms.v1.ListOrganizationsRequest + 59, // 68: ktms.v1.AdminService.DeleteOrganization:input_type -> ktms.v1.DeleteOrganizationRequest + 61, // 69: ktms.v1.AdminService.CreateTeam:input_type -> ktms.v1.CreateTeamRequest + 63, // 70: ktms.v1.AdminService.GetTeam:input_type -> ktms.v1.GetTeamRequest + 65, // 71: ktms.v1.AdminService.ListTeams:input_type -> ktms.v1.ListTeamsRequest + 67, // 72: ktms.v1.AdminService.UpdateTeam:input_type -> ktms.v1.UpdateTeamRequest + 69, // 73: ktms.v1.AdminService.DeleteTeam:input_type -> ktms.v1.DeleteTeamRequest + 71, // 74: ktms.v1.AdminService.GetTeamSpend:input_type -> ktms.v1.GetTeamSpendRequest + 73, // 75: ktms.v1.AdminService.GetSpend:input_type -> ktms.v1.GetSpendRequest + 89, // 76: ktms.v1.AdminService.ExplainRoute:input_type -> ktms.v1.ExplainRouteRequest + 77, // 77: ktms.v1.AdminService.UpsertMCPServer:input_type -> ktms.v1.UpsertMCPServerRequest + 79, // 78: ktms.v1.AdminService.DeleteMCPServer:input_type -> ktms.v1.DeleteMCPServerRequest + 81, // 79: ktms.v1.AdminService.SetMCPCredential:input_type -> ktms.v1.SetMCPCredentialRequest + 83, // 80: ktms.v1.AdminService.DeleteMCPCredential:input_type -> ktms.v1.DeleteMCPCredentialRequest + 86, // 81: ktms.v1.AdminService.ListMCPTools:input_type -> ktms.v1.ListMCPToolsRequest + 3, // 82: ktms.v1.AdminService.UpsertModel:output_type -> ktms.v1.UpsertModelResponse + 5, // 83: ktms.v1.AdminService.DeleteModel:output_type -> ktms.v1.DeleteModelResponse + 7, // 84: ktms.v1.AdminService.ReplaceRules:output_type -> ktms.v1.ReplaceRulesResponse + 9, // 85: ktms.v1.AdminService.SetVendorKey:output_type -> ktms.v1.SetVendorKeyResponse + 11, // 86: ktms.v1.AdminService.DeleteVendorKey:output_type -> ktms.v1.DeleteVendorKeyResponse + 14, // 87: ktms.v1.AdminService.GetConfig:output_type -> ktms.v1.GetConfigResponse + 19, // 88: ktms.v1.AdminService.MintAPIKey:output_type -> ktms.v1.MintAPIKeyResponse + 21, // 89: ktms.v1.AdminService.RevokeAPIKey:output_type -> ktms.v1.RevokeAPIKeyResponse + 24, // 90: ktms.v1.AdminService.ListAPIKeys:output_type -> ktms.v1.ListAPIKeysResponse + 18, // 91: ktms.v1.AdminService.SetAPIKeyPolicy:output_type -> ktms.v1.SetAPIKeyPolicyResponse + 26, // 92: ktms.v1.AdminService.RotateEncryption:output_type -> ktms.v1.RotateEncryptionResponse + 28, // 93: ktms.v1.AdminService.MintNodeKey:output_type -> ktms.v1.MintNodeKeyResponse + 30, // 94: ktms.v1.AdminService.RevokeNodeKey:output_type -> ktms.v1.RevokeNodeKeyResponse + 33, // 95: ktms.v1.AdminService.ListNodeKeys:output_type -> ktms.v1.ListNodeKeysResponse + 36, // 96: ktms.v1.AdminService.ListNodes:output_type -> ktms.v1.ListNodesResponse + 38, // 97: ktms.v1.AdminService.DrainNode:output_type -> ktms.v1.DrainNodeResponse + 41, // 98: ktms.v1.AdminService.SetUserRateLimit:output_type -> ktms.v1.SetUserRateLimitResponse + 43, // 99: ktms.v1.AdminService.DeleteUserRateLimit:output_type -> ktms.v1.DeleteUserRateLimitResponse + 45, // 100: ktms.v1.AdminService.ListUserRateLimits:output_type -> ktms.v1.ListUserRateLimitsResponse + 48, // 101: ktms.v1.AdminService.SetSpendLimit:output_type -> ktms.v1.SetSpendLimitResponse + 50, // 102: ktms.v1.AdminService.DeleteSpendLimit:output_type -> ktms.v1.DeleteSpendLimitResponse + 52, // 103: ktms.v1.AdminService.ListSpendLimits:output_type -> ktms.v1.ListSpendLimitsResponse + 56, // 104: ktms.v1.AdminService.CreateOrganization:output_type -> ktms.v1.CreateOrganizationResponse + 58, // 105: ktms.v1.AdminService.ListOrganizations:output_type -> ktms.v1.ListOrganizationsResponse + 60, // 106: ktms.v1.AdminService.DeleteOrganization:output_type -> ktms.v1.DeleteOrganizationResponse + 62, // 107: ktms.v1.AdminService.CreateTeam:output_type -> ktms.v1.CreateTeamResponse + 64, // 108: ktms.v1.AdminService.GetTeam:output_type -> ktms.v1.GetTeamResponse + 66, // 109: ktms.v1.AdminService.ListTeams:output_type -> ktms.v1.ListTeamsResponse + 68, // 110: ktms.v1.AdminService.UpdateTeam:output_type -> ktms.v1.UpdateTeamResponse + 70, // 111: ktms.v1.AdminService.DeleteTeam:output_type -> ktms.v1.DeleteTeamResponse + 72, // 112: ktms.v1.AdminService.GetTeamSpend:output_type -> ktms.v1.GetTeamSpendResponse + 75, // 113: ktms.v1.AdminService.GetSpend:output_type -> ktms.v1.GetSpendResponse + 91, // 114: ktms.v1.AdminService.ExplainRoute:output_type -> ktms.v1.ExplainRouteResponse + 78, // 115: ktms.v1.AdminService.UpsertMCPServer:output_type -> ktms.v1.UpsertMCPServerResponse + 80, // 116: ktms.v1.AdminService.DeleteMCPServer:output_type -> ktms.v1.DeleteMCPServerResponse + 82, // 117: ktms.v1.AdminService.SetMCPCredential:output_type -> ktms.v1.SetMCPCredentialResponse + 84, // 118: ktms.v1.AdminService.DeleteMCPCredential:output_type -> ktms.v1.DeleteMCPCredentialResponse + 88, // 119: ktms.v1.AdminService.ListMCPTools:output_type -> ktms.v1.ListMCPToolsResponse + 82, // [82:120] is the sub-list for method output_type + 44, // [44:82] is the sub-list for method input_type + 44, // [44:44] is the sub-list for extension type_name + 44, // [44:44] is the sub-list for extension extendee + 0, // [0:44] is the sub-list for field type_name } func init() { file_ktms_v1_admin_proto_init() } @@ -5076,7 +5874,7 @@ func file_ktms_v1_admin_proto_init() { GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: unsafe.Slice(unsafe.StringData(file_ktms_v1_admin_proto_rawDesc), len(file_ktms_v1_admin_proto_rawDesc)), NumEnums: 0, - NumMessages: 81, + NumMessages: 94, NumExtensions: 0, NumServices: 1, }, diff --git a/gen/ktms/v1/admin_grpc.pb.go b/gen/ktms/v1/admin_grpc.pb.go index f035053..818dd35 100644 --- a/gen/ktms/v1/admin_grpc.pb.go +++ b/gen/ktms/v1/admin_grpc.pb.go @@ -52,6 +52,11 @@ const ( AdminService_GetTeamSpend_FullMethodName = "/ktms.v1.AdminService/GetTeamSpend" AdminService_GetSpend_FullMethodName = "/ktms.v1.AdminService/GetSpend" AdminService_ExplainRoute_FullMethodName = "/ktms.v1.AdminService/ExplainRoute" + AdminService_UpsertMCPServer_FullMethodName = "/ktms.v1.AdminService/UpsertMCPServer" + AdminService_DeleteMCPServer_FullMethodName = "/ktms.v1.AdminService/DeleteMCPServer" + AdminService_SetMCPCredential_FullMethodName = "/ktms.v1.AdminService/SetMCPCredential" + AdminService_DeleteMCPCredential_FullMethodName = "/ktms.v1.AdminService/DeleteMCPCredential" + AdminService_ListMCPTools_FullMethodName = "/ktms.v1.AdminService/ListMCPTools" ) // AdminServiceClient is the client API for AdminService service. @@ -98,6 +103,11 @@ type AdminServiceClient interface { // and every model of its chain, kept or dropped with the reason. No vendor // or judge is called and nothing is spent; rate limits are not checked. ExplainRoute(ctx context.Context, in *ExplainRouteRequest, opts ...grpc.CallOption) (*ExplainRouteResponse, error) + UpsertMCPServer(ctx context.Context, in *UpsertMCPServerRequest, opts ...grpc.CallOption) (*UpsertMCPServerResponse, error) + DeleteMCPServer(ctx context.Context, in *DeleteMCPServerRequest, opts ...grpc.CallOption) (*DeleteMCPServerResponse, error) + SetMCPCredential(ctx context.Context, in *SetMCPCredentialRequest, opts ...grpc.CallOption) (*SetMCPCredentialResponse, error) + DeleteMCPCredential(ctx context.Context, in *DeleteMCPCredentialRequest, opts ...grpc.CallOption) (*DeleteMCPCredentialResponse, error) + ListMCPTools(ctx context.Context, in *ListMCPToolsRequest, opts ...grpc.CallOption) (*ListMCPToolsResponse, error) } type adminServiceClient struct { @@ -438,6 +448,56 @@ func (c *adminServiceClient) ExplainRoute(ctx context.Context, in *ExplainRouteR return out, nil } +func (c *adminServiceClient) UpsertMCPServer(ctx context.Context, in *UpsertMCPServerRequest, opts ...grpc.CallOption) (*UpsertMCPServerResponse, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(UpsertMCPServerResponse) + err := c.cc.Invoke(ctx, AdminService_UpsertMCPServer_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + +func (c *adminServiceClient) DeleteMCPServer(ctx context.Context, in *DeleteMCPServerRequest, opts ...grpc.CallOption) (*DeleteMCPServerResponse, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(DeleteMCPServerResponse) + err := c.cc.Invoke(ctx, AdminService_DeleteMCPServer_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + +func (c *adminServiceClient) SetMCPCredential(ctx context.Context, in *SetMCPCredentialRequest, opts ...grpc.CallOption) (*SetMCPCredentialResponse, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(SetMCPCredentialResponse) + err := c.cc.Invoke(ctx, AdminService_SetMCPCredential_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + +func (c *adminServiceClient) DeleteMCPCredential(ctx context.Context, in *DeleteMCPCredentialRequest, opts ...grpc.CallOption) (*DeleteMCPCredentialResponse, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(DeleteMCPCredentialResponse) + err := c.cc.Invoke(ctx, AdminService_DeleteMCPCredential_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + +func (c *adminServiceClient) ListMCPTools(ctx context.Context, in *ListMCPToolsRequest, opts ...grpc.CallOption) (*ListMCPToolsResponse, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(ListMCPToolsResponse) + err := c.cc.Invoke(ctx, AdminService_ListMCPTools_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + // AdminServiceServer is the server API for AdminService service. // All implementations must embed UnimplementedAdminServiceServer // for forward compatibility. @@ -482,6 +542,11 @@ type AdminServiceServer interface { // and every model of its chain, kept or dropped with the reason. No vendor // or judge is called and nothing is spent; rate limits are not checked. ExplainRoute(context.Context, *ExplainRouteRequest) (*ExplainRouteResponse, error) + UpsertMCPServer(context.Context, *UpsertMCPServerRequest) (*UpsertMCPServerResponse, error) + DeleteMCPServer(context.Context, *DeleteMCPServerRequest) (*DeleteMCPServerResponse, error) + SetMCPCredential(context.Context, *SetMCPCredentialRequest) (*SetMCPCredentialResponse, error) + DeleteMCPCredential(context.Context, *DeleteMCPCredentialRequest) (*DeleteMCPCredentialResponse, error) + ListMCPTools(context.Context, *ListMCPToolsRequest) (*ListMCPToolsResponse, error) mustEmbedUnimplementedAdminServiceServer() } @@ -591,6 +656,21 @@ func (UnimplementedAdminServiceServer) GetSpend(context.Context, *GetSpendReques func (UnimplementedAdminServiceServer) ExplainRoute(context.Context, *ExplainRouteRequest) (*ExplainRouteResponse, error) { return nil, status.Error(codes.Unimplemented, "method ExplainRoute not implemented") } +func (UnimplementedAdminServiceServer) UpsertMCPServer(context.Context, *UpsertMCPServerRequest) (*UpsertMCPServerResponse, error) { + return nil, status.Error(codes.Unimplemented, "method UpsertMCPServer not implemented") +} +func (UnimplementedAdminServiceServer) DeleteMCPServer(context.Context, *DeleteMCPServerRequest) (*DeleteMCPServerResponse, error) { + return nil, status.Error(codes.Unimplemented, "method DeleteMCPServer not implemented") +} +func (UnimplementedAdminServiceServer) SetMCPCredential(context.Context, *SetMCPCredentialRequest) (*SetMCPCredentialResponse, error) { + return nil, status.Error(codes.Unimplemented, "method SetMCPCredential not implemented") +} +func (UnimplementedAdminServiceServer) DeleteMCPCredential(context.Context, *DeleteMCPCredentialRequest) (*DeleteMCPCredentialResponse, error) { + return nil, status.Error(codes.Unimplemented, "method DeleteMCPCredential not implemented") +} +func (UnimplementedAdminServiceServer) ListMCPTools(context.Context, *ListMCPToolsRequest) (*ListMCPToolsResponse, error) { + return nil, status.Error(codes.Unimplemented, "method ListMCPTools not implemented") +} func (UnimplementedAdminServiceServer) mustEmbedUnimplementedAdminServiceServer() {} func (UnimplementedAdminServiceServer) testEmbeddedByValue() {} @@ -1206,6 +1286,96 @@ func _AdminService_ExplainRoute_Handler(srv interface{}, ctx context.Context, de return interceptor(ctx, in, info, handler) } +func _AdminService_UpsertMCPServer_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(UpsertMCPServerRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(AdminServiceServer).UpsertMCPServer(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: AdminService_UpsertMCPServer_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(AdminServiceServer).UpsertMCPServer(ctx, req.(*UpsertMCPServerRequest)) + } + return interceptor(ctx, in, info, handler) +} + +func _AdminService_DeleteMCPServer_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(DeleteMCPServerRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(AdminServiceServer).DeleteMCPServer(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: AdminService_DeleteMCPServer_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(AdminServiceServer).DeleteMCPServer(ctx, req.(*DeleteMCPServerRequest)) + } + return interceptor(ctx, in, info, handler) +} + +func _AdminService_SetMCPCredential_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(SetMCPCredentialRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(AdminServiceServer).SetMCPCredential(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: AdminService_SetMCPCredential_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(AdminServiceServer).SetMCPCredential(ctx, req.(*SetMCPCredentialRequest)) + } + return interceptor(ctx, in, info, handler) +} + +func _AdminService_DeleteMCPCredential_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(DeleteMCPCredentialRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(AdminServiceServer).DeleteMCPCredential(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: AdminService_DeleteMCPCredential_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(AdminServiceServer).DeleteMCPCredential(ctx, req.(*DeleteMCPCredentialRequest)) + } + return interceptor(ctx, in, info, handler) +} + +func _AdminService_ListMCPTools_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(ListMCPToolsRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(AdminServiceServer).ListMCPTools(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: AdminService_ListMCPTools_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(AdminServiceServer).ListMCPTools(ctx, req.(*ListMCPToolsRequest)) + } + return interceptor(ctx, in, info, handler) +} + // AdminService_ServiceDesc is the grpc.ServiceDesc for AdminService service. // It's only intended for direct use with grpc.RegisterService, // and not to be introspected or modified (even as a copy) @@ -1345,6 +1515,26 @@ var AdminService_ServiceDesc = grpc.ServiceDesc{ MethodName: "ExplainRoute", Handler: _AdminService_ExplainRoute_Handler, }, + { + MethodName: "UpsertMCPServer", + Handler: _AdminService_UpsertMCPServer_Handler, + }, + { + MethodName: "DeleteMCPServer", + Handler: _AdminService_DeleteMCPServer_Handler, + }, + { + MethodName: "SetMCPCredential", + Handler: _AdminService_SetMCPCredential_Handler, + }, + { + MethodName: "DeleteMCPCredential", + Handler: _AdminService_DeleteMCPCredential_Handler, + }, + { + MethodName: "ListMCPTools", + Handler: _AdminService_ListMCPTools_Handler, + }, }, Streams: []grpc.StreamDesc{}, Metadata: "ktms/v1/admin.proto", diff --git a/internal/apikeys/apikeys.go b/internal/apikeys/apikeys.go index 5dcb8cb..118c5ba 100644 --- a/internal/apikeys/apikeys.go +++ b/internal/apikeys/apikeys.go @@ -29,7 +29,9 @@ import ( const ( // ScopeGenerate allows the data-plane RPCs (GenerateText, ListModels). ScopeGenerate = "generate" - ScopeAdmin = "admin" + // ScopeMCP allows the MCP gateway at /mcp. + ScopeMCP = "mcp" + ScopeAdmin = "admin" // ScopeAdminRead allows the Admin API reads and nothing that changes state. ScopeAdminRead = "admin:read" ScopeNode = "node" @@ -49,7 +51,7 @@ var ( ) // ValidScopes is the closed set the admin endpoint accepts. -var ValidScopes = []string{ScopeGenerate, ScopeAdmin, ScopeAdminRead, ScopeNode} +var ValidScopes = []string{ScopeGenerate, ScopeMCP, ScopeAdmin, ScopeAdminRead, ScopeNode} // Store is the slice of proxydb the verifier needs, so tests can fake it. type Store interface { @@ -74,6 +76,9 @@ type Identity struct { KeyAllowedModels []string RPMLimit int32 TPMLimit int64 + // MCP servers ("server") and tools ("server/tool") allowed; empty = any. + KeyAllowedMCP []string + TeamAllowedMCP []string } // HasScope reports whether the caller may perform an action needing want. @@ -110,6 +115,33 @@ func (i *Identity) RestrictsModels() bool { return i != nil && (len(i.TeamAllowedModels) > 0 || len(i.KeyAllowedModels) > 0) } +// CanUseTool reports whether both the key's and its team's allowed_mcp permit +// the server's tool; tool "" asks whether any tool of the server may be. +func (i *Identity) CanUseTool(server, tool string) bool { + return i != nil && mcpAllows(i.TeamAllowedMCP, server, tool) && mcpAllows(i.KeyAllowedMCP, server, tool) +} + +// MCPDeniedBy names the allowlist that refuses the tool, as DeniedBy does. +func (i *Identity) MCPDeniedBy(server, tool string) string { + if i != nil && !mcpAllows(i.TeamAllowedMCP, server, tool) { + return "team" + } + return "key" +} + +func mcpAllows(list []string, server, tool string) bool { + if len(list) == 0 { + return true + } + for _, e := range list { + s, t, one := strings.Cut(e, "/") + if s == server && (!one || tool == "" || t == tool) { + return true + } + } + return false +} + // LookupHash is the indexed sha256 hex of a plain key. func LookupHash(plain string) string { sum := sha256.Sum256([]byte(plain)) @@ -343,6 +375,8 @@ func IdentityFor(row *proxydb.APIKey) *Identity { KeyAllowedModels: row.AllowedModels, RPMLimit: row.RPMLimit, TPMLimit: row.TPMLimit, + KeyAllowedMCP: row.AllowedMCP, + TeamAllowedMCP: row.TeamAllowedMCP, } } diff --git a/internal/config/config.go b/internal/config/config.go index cf095c1..8877922 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -25,6 +25,8 @@ type Config struct { HTTPAddr string // HTTPAPIEnabled mounts the OpenAI-compatible API on HTTPAddr; off leaves health only. HTTPAPIEnabled bool + // MCPGatewayEnabled mounts the MCP gateway (/mcp) on HTTPAddr. + MCPGatewayEnabled bool // HTTPAPITimeout bounds one HTTP API call, which waits on an upstream model. HTTPAPITimeout time.Duration LogLevel string @@ -155,6 +157,7 @@ func Load() *Config { GRPCMaxRecvMsgSize: getIntEnv("GRPC_MAX_RECV_MSG_SIZE", 16*1024*1024), HTTPAddr: getEnv("HTTP_ADDR", ":8080"), HTTPAPIEnabled: getBoolEnv("HTTP_API_ENABLED", true), + MCPGatewayEnabled: getBoolEnv("MCP_GATEWAY_ENABLED", true), HTTPAPITimeout: getDurationEnv("HTTP_API_TIMEOUT", 5*time.Minute), LogLevel: getEnv("LOG_LEVEL", "info"), ShutdownTimeout: getDurationEnv("SHUTDOWN_TIMEOUT", 30*time.Second), diff --git a/internal/declarative/config.go b/internal/declarative/config.go index fc7c040..6901c31 100644 --- a/internal/declarative/config.go +++ b/internal/declarative/config.go @@ -35,6 +35,8 @@ type Config struct { VendorKeys []VendorKey SpendLimits []SpendLimit UserRateLimits []*pb.UserRateLimitSpec + MCPServers []*pb.MCPServerSpec + MCPCredentials []MCPCredential // Which sections the file has, since an empty list differs from a missing one. managed map[string]bool @@ -51,6 +53,7 @@ type Team struct { Name string `json:"name"` MonthlyBudgetUSD float64 `json:"monthly_budget_usd,omitempty"` AllowedModels []string `json:"allowed_models,omitempty"` + AllowedMCP []string `json:"allowed_mcp,omitempty"` } // APIKey is a client key, matched by name. Its secret exists only in the @@ -71,6 +74,14 @@ type VendorKey struct { PlainKeyEnv string `json:"plain_key_env"` } +// MCPCredential is the header the MCP gateway adds to a server's requests. +// The value is read from an environment variable, never from the file. +type MCPCredential struct { + Server string `json:"server"` + Header string `json:"header,omitempty"` + ValueEnv string `json:"value_env"` +} + // SpendLimit caps a key's or a user's spend per UTC day or month. Team caps // belong to Team.MonthlyBudgetUSD. type SpendLimit struct { @@ -90,10 +101,11 @@ func (l SpendLimit) Nano() int64 { } // Manages reports whether the file has the section (models, rules, -// organizations, api_keys, vendor_keys, spend_limits, user_rate_limits). +// organizations, api_keys, vendor_keys, spend_limits, user_rate_limits, +// mcp_servers, mcp_credentials). func (c *Config) Manages(section string) bool { return c.managed[section] } -var sections = []string{"models", "rules", "organizations", "api_keys", "vendor_keys", "spend_limits", "user_rate_limits"} +var sections = []string{"models", "rules", "organizations", "api_keys", "vendor_keys", "spend_limits", "user_rate_limits", "mcp_servers", "mcp_credentials"} // Load reads and parses a configuration file. func Load(path string) (*Config, error) { @@ -145,6 +157,12 @@ func Parse(b []byte) (*Config, error) { if cfg.UserRateLimits, err = protoList(doc, "user_rate_limits", func() *pb.UserRateLimitSpec { return &pb.UserRateLimitSpec{} }); err != nil { return nil, err } + if cfg.MCPServers, err = protoList(doc, "mcp_servers", func() *pb.MCPServerSpec { return &pb.MCPServerSpec{} }); err != nil { + return nil, err + } + if err := jsonList(doc, "mcp_credentials", &cfg.MCPCredentials); err != nil { + return nil, err + } if err := jsonList(doc, "organizations", &cfg.Organizations); err != nil { return nil, err } diff --git a/internal/declarative/declarative_test.go b/internal/declarative/declarative_test.go index 00d7c10..6f6b0e8 100644 --- a/internal/declarative/declarative_test.go +++ b/internal/declarative/declarative_test.go @@ -75,7 +75,7 @@ func env(m map[string]string) func(string) (string, bool) { } func TestParseExample(t *testing.T) { - c := mustParse(t, example) + c := mustParse(t, example+mcpExample) if len(c.Models) != 2 || !c.Models[0].GetEnabled() { t.Fatalf("models = %v; a model without enabled must be enabled", c.Models) } diff --git a/internal/declarative/export.go b/internal/declarative/export.go index a6c6bf8..1e3e952 100644 --- a/internal/declarative/export.go +++ b/internal/declarative/export.go @@ -15,7 +15,8 @@ import ( ) // Export writes the live configuration in the file format. Secrets are not -// exported: each vendor key becomes a plain_key_env reference to fill in, and +// exported: each vendor key and MCP credential becomes a reference to an +// environment variable to fill in, and // client keys are listed without theirs. Revoked keys are left out. func Export(live *State, w io.Writer) error { doc, err := exportDoc(live) @@ -88,6 +89,9 @@ func exportDoc(live *State) (*yaml.Node, error) { if len(t.GetAllowedModels()) > 0 { add(tn, "allowed_models", stringSeq(t.GetAllowedModels())) } + if len(t.GetAllowedMcp()) > 0 { + add(tn, "allowed_mcp", stringSeq(t.GetAllowedMcp())) + } teams.Content = append(teams.Content, tn) } if len(teams.Content) > 0 { @@ -160,6 +164,26 @@ func exportDoc(live *State) (*yaml.Node, error) { } add(doc, "user_rate_limits", urls) + servers := seq() + for _, m := range live.Config.GetMcpServers() { + n, err := messageNode(m, true) + if err != nil { + return nil, err + } + servers.Content = append(servers.Content, n) + } + add(doc, "mcp_servers", servers) + + creds := seq() + for _, mc := range live.Config.GetMcpCredentials() { + cn := mapping() + add(cn, "server", scalarNode(mc.GetServer())) + add(cn, "header", scalarNode(mc.GetHeader())) + add(cn, "value_env", scalarNode(mcpEnvName(mc.GetServer()))) + creds.Content = append(creds.Content, cn) + } + add(doc, "mcp_credentials", creds) + return doc, nil } @@ -172,6 +196,12 @@ func encode(doc *yaml.Node, w io.Writer) error { return enc.Close() } +// mcpEnvName is the variable an exported MCP credential points at, e.g. +// KTMS_MCP_CREDENTIAL_GITHUB. +func mcpEnvName(server string) string { + return "KTMS_MCP_CREDENTIAL_" + strings.ToUpper(strings.ReplaceAll(server, "-", "_")) +} + // envName is the variable an exported vendor key points at, e.g. // KTMS_VENDOR_KEY_BILLING_SERVICE_OPENAI. func envName(key, vendor string) string { diff --git a/internal/declarative/mcp_test.go b/internal/declarative/mcp_test.go new file mode 100644 index 0000000..464f9a9 --- /dev/null +++ b/internal/declarative/mcp_test.go @@ -0,0 +1,91 @@ +package declarative + +import ( + "bytes" + "strings" + "testing" + + pb "github.com/ai-process/ktms/gen/ktms/v1" + "github.com/ai-process/ktms/internal/crypto" +) + +const mcpExample = `mcp_servers: + - name: github + url: https://api.githubcopilot.com/mcp/ + tool_rpm: 60 + - name: docs + url: http://docs-mcp:8000/sse + transport: sse + allow_private: true +mcp_credentials: + - server: github + value_env: GITHUB_MCP_TOKEN +` + +func TestPlanMCP(t *testing.T) { + c := mustParse(t, example+mcpExample) + envs := env(map[string]string{"TEST_OPENAI_KEY": "sk-1", "GITHUB_MCP_TOKEN": "Bearer ghp-1"}) + p, err := Build(c, emptyState(), Options{Env: envs}) + if err != nil { + t.Fatal(err) + } + var buf bytes.Buffer + p.Write(&buf) + out := buf.String() + if i, j := strings.Index(out, "+ mcp server github"), strings.Index(out, "+ mcp credential github"); i < 0 || j < i { + t.Fatalf("server before its credential:\n%s", out) + } + if strings.Contains(out, "ghp-1") { + t.Fatal("the plan shows the credential") + } + + live := liveOf(t, c, "sk-1") + live.Config.McpServers = []*pb.MCPServerSpec{ + {Name: "github", Url: "https://api.githubcopilot.com/mcp/", Transport: "streamable_http", Enabled: true, ToolRpm: 60}, + {Name: "docs", Url: "http://docs-mcp:8000/sse", Transport: "sse", Enabled: true, AllowPrivate: true}, + {Name: "old", Url: "https://old.example/mcp", Transport: "streamable_http", Enabled: true}, + } + live.Config.McpCredentials = []*pb.MCPCredentialInfo{{Server: "github", Header: "Authorization", Fingerprint: crypto.Fingerprint("Bearer ghp-0")}} + p, err = Build(c, live, Options{Env: envs, Prune: true}) + if err != nil { + t.Fatal(err) + } + buf.Reset() + p.Write(&buf) + out = buf.String() + for _, want := range []string{"~ mcp credential github", "changed in $GITHUB_MCP_TOKEN", "- mcp server old"} { + if !strings.Contains(out, want) { + t.Errorf("plan lacks %q:\n%s", want, out) + } + } + if strings.Contains(out, "mcp server github") || strings.Contains(out, "mcp server docs") { + t.Errorf("a default transport is not a change:\n%s", out) + } + + if _, err := Build(c, emptyState(), Options{Env: env(map[string]string{"TEST_OPENAI_KEY": "sk-1"})}); err == nil || !strings.Contains(err.Error(), "GITHUB_MCP_TOKEN is not set") { + t.Fatalf("missing variable: %v", err) + } + if _, err := Parse([]byte("api_version: " + APIVersion + "\nmcp_servers: []\nmcp_credentials:\n - server: nope\n value_env: X\n")); err == nil || !strings.Contains(err.Error(), "not in mcp_servers") { + t.Fatalf("credential of an unknown server: %v", err) + } +} + +func TestExportMCPParsesBack(t *testing.T) { + live := emptyState() + live.Config.McpServers = []*pb.MCPServerSpec{{Name: "github", Url: "https://x.example/mcp", Transport: "streamable_http", Enabled: false}} + live.Config.McpCredentials = []*pb.MCPCredentialInfo{{Server: "github", Header: "Authorization"}} + live.Teams = []*pb.Team{{Id: "t1", OrgId: "o1", Name: "t", AllowedMcp: []string{"github"}}} + live.Orgs = []*pb.Organization{{Id: "o1", Name: "o"}} + var buf bytes.Buffer + if err := Export(live, &buf); err != nil { + t.Fatal(err) + } + c, err := Parse(buf.Bytes()) + if err != nil { + t.Fatalf("%v\n%s", err, buf.String()) + } + if len(c.MCPServers) != 1 || c.MCPServers[0].GetEnabled() || c.MCPCredentials[0].ValueEnv != "KTMS_MCP_CREDENTIAL_GITHUB" || + c.Organizations[0].Teams[0].AllowedMCP[0] != "github" { + t.Fatalf("parsed back %+v\n%s", c, buf.String()) + } +} diff --git a/internal/declarative/plan.go b/internal/declarative/plan.go index bc5a80e..a2798be 100644 --- a/internal/declarative/plan.go +++ b/internal/declarative/plan.go @@ -144,7 +144,7 @@ func Build(cfg *Config, live *State, opt Options) (*Plan, error) { if lt == nil { p.add(Change{Op: Create, Kind: "team", Name: name, Details: teamDetails(nil, t)}, func(ctx context.Context, c pb.AdminServiceClient, a *applied) error { r, err := c.CreateTeam(ctx, &pb.CreateTeamRequest{OrgId: a.orgIDs[o.Name], Name: t.Name, - MonthlyBudgetUsd: t.MonthlyBudgetUSD, AllowedModels: t.AllowedModels}) + MonthlyBudgetUsd: t.MonthlyBudgetUSD, AllowedModels: t.AllowedModels, AllowedMcp: t.AllowedMCP}) if err != nil { return err } @@ -156,7 +156,7 @@ func Build(cfg *Config, live *State, opt Options) (*Plan, error) { if d := teamDetails(lt, t); len(d) > 0 { p.add(Change{Op: Update, Kind: "team", Name: name, Details: d}, func(ctx context.Context, c pb.AdminServiceClient, _ *applied) error { _, err := c.UpdateTeam(ctx, &pb.UpdateTeamRequest{Id: lt.GetId(), Name: t.Name, - MonthlyBudgetUsd: t.MonthlyBudgetUSD, AllowedModels: t.AllowedModels}) + MonthlyBudgetUsd: t.MonthlyBudgetUSD, AllowedModels: t.AllowedModels, AllowedMcp: t.AllowedMCP}) return err }) } @@ -211,6 +211,92 @@ func Build(cfg *Config, live *State, opt Options) (*Plan, error) { } } + // MCP servers before their credentials. + liveMCP := map[string]*pb.MCPServerSpec{} + for _, m := range live.Config.GetMcpServers() { + liveMCP[m.GetName()] = m + } + if cfg.Manages("mcp_servers") { + for _, m := range cfg.MCPServers { + m := m + upsert := func(ctx context.Context, c pb.AdminServiceClient, _ *applied) error { + _, err := c.UpsertMCPServer(ctx, &pb.UpsertMCPServerRequest{Server: m}) + return err + } + lm := liveMCP[m.GetName()] + if lm == nil { + p.add(Change{Op: Create, Kind: "mcp server", Name: m.GetName(), Details: []string{m.GetUrl()}}, upsert) + } else if d := fieldDiff(normMCP(lm), normMCP(m)); len(d) > 0 { + p.add(Change{Op: Update, Kind: "mcp server", Name: m.GetName(), Details: d}, upsert) + } + } + if opt.Prune { + for name := range liveMCP { + name := name + if slices.ContainsFunc(cfg.MCPServers, func(m *pb.MCPServerSpec) bool { return m.GetName() == name }) { + continue + } + deletes = append(deletes, func() { + p.add(Change{Op: Delete, Kind: "mcp server", Name: name, Details: []string{"its credential goes with it"}}, func(ctx context.Context, c pb.AdminServiceClient, _ *applied) error { + _, err := c.DeleteMCPServer(ctx, &pb.DeleteMCPServerRequest{Name: name}) + return err + }) + }) + } + } + } + if cfg.Manages("mcp_credentials") { + liveCred := map[string]*pb.MCPCredentialInfo{} + for _, mc := range live.Config.GetMcpCredentials() { + liveCred[mc.GetServer()] = mc + } + for _, mc := range cfg.MCPCredentials { + mc := mc + plain, ok := opt.Env(mc.ValueEnv) + if !ok || plain == "" { + conflict("mcp credential %s: environment variable %s is not set", mc.Server, mc.ValueEnv) + continue + } + if liveMCP[mc.Server] == nil && !slices.ContainsFunc(cfg.MCPServers, func(m *pb.MCPServerSpec) bool { return m.GetName() == mc.Server }) { + conflict("mcp credential %s: the server exists neither in Keystone nor in mcp_servers", mc.Server) + continue + } + header := mc.Header + if header == "" { + header = "Authorization" + } + set := func(ctx context.Context, c pb.AdminServiceClient, _ *applied) error { + _, err := c.SetMCPCredential(ctx, &pb.SetMCPCredentialRequest{Server: mc.Server, Header: header, PlainValue: plain}) + return err + } + lc := liveCred[mc.Server] + switch { + case lc == nil: + p.add(Change{Op: Create, Kind: "mcp credential", Name: mc.Server, Details: []string{header + " from $" + mc.ValueEnv}}, set) + case !strings.EqualFold(lc.GetHeader(), header): + p.add(Change{Op: Update, Kind: "mcp credential", Name: mc.Server, Details: []string{"header " + lc.GetHeader() + " → " + header}}, set) + case lc.GetFingerprint() == "": + p.add(Change{Op: Update, Kind: "mcp credential", Name: mc.Server, Details: []string{"set again from $" + mc.ValueEnv + ": Keystone cannot tell whether it changed"}}, set) + case lc.GetFingerprint() != crypto.Fingerprint(plain): + p.add(Change{Op: Update, Kind: "mcp credential", Name: mc.Server, Details: []string{"changed in $" + mc.ValueEnv}}, set) + } + } + if opt.Prune { + for _, lc := range live.Config.GetMcpCredentials() { + lc := lc + if slices.ContainsFunc(cfg.MCPCredentials, func(mc MCPCredential) bool { return mc.Server == lc.GetServer() }) { + continue + } + deletes = append(deletes, func() { + p.add(Change{Op: Delete, Kind: "mcp credential", Name: lc.GetServer()}, func(ctx context.Context, c pb.AdminServiceClient, _ *applied) error { + _, err := c.DeleteMCPCredential(ctx, &pb.DeleteMCPCredentialRequest{Server: lc.GetServer()}) + return err + }) + }) + } + } + } + // Keys: mint the missing ones, update policies, refuse what the API cannot change. if cfg.Manages("api_keys") { for _, k := range cfg.APIKeys { @@ -458,7 +544,7 @@ func Build(cfg *Config, live *State, opt Options) (*Plan, error) { // refers to it: limits and keys, then models (the new rules no longer name // them), then teams, then organizations. func sortDeletes(p *Plan) { - order := map[string]int{"user rate limit": 1, "spend limit": 1, "vendor key": 2, "api key": 3, "model": 4, "team": 5, "organization": 6} + order := map[string]int{"user rate limit": 1, "spend limit": 1, "vendor key": 2, "mcp credential": 2, "api key": 3, "model": 4, "mcp server": 4, "team": 5, "organization": 6} rank := func(s step) int { if s.change.Op != Delete { return 0 @@ -576,6 +662,9 @@ func teamDetails(live *pb.Team, t Team) []string { if len(t.AllowedModels) > 0 { d = append(d, "allowed models "+strings.Join(t.AllowedModels, ", ")) } + if len(t.AllowedMCP) > 0 { + d = append(d, "allowed mcp "+strings.Join(t.AllowedMCP, ", ")) + } return d } if live.GetMonthlyBudgetUsd() != t.MonthlyBudgetUSD { @@ -584,6 +673,9 @@ func teamDetails(live *pb.Team, t Team) []string { if !sameSet(live.GetAllowedModels(), t.AllowedModels) { d = append(d, fmt.Sprintf("allowed models %v → %v", live.GetAllowedModels(), t.AllowedModels)) } + if !sameSet(live.GetAllowedMcp(), t.AllowedMCP) { + d = append(d, fmt.Sprintf("allowed mcp %v → %v", live.GetAllowedMcp(), t.AllowedMCP)) + } return d } @@ -598,6 +690,9 @@ func policyDiff(live, want *pb.APIKeyPolicy) []string { if live.GetTpmLimit() != want.GetTpmLimit() { d = append(d, fmt.Sprintf("tpm_limit %d → %d", live.GetTpmLimit(), want.GetTpmLimit())) } + if !sameSet(live.GetAllowedMcp(), want.GetAllowedMcp()) { + d = append(d, fmt.Sprintf("allowed mcp %v → %v", live.GetAllowedMcp(), want.GetAllowedMcp())) + } return d } @@ -638,6 +733,15 @@ func normModel(m *pb.ModelSpec) *pb.ModelSpec { return n } +// normMCP spells out the default transport, so leaving it out is no difference. +func normMCP(m *pb.MCPServerSpec) *pb.MCPServerSpec { + n := proto.Clone(m).(*pb.MCPServerSpec) + if n.Transport == "" { + n.Transport = "streamable_http" + } + return n +} + // fieldDiff lists the fields of two messages of one type that differ, as // "field old → new". func fieldDiff(a, b proto.Message) []string { diff --git a/internal/declarative/schema.go b/internal/declarative/schema.go index bc74932..cf2224b 100644 --- a/internal/declarative/schema.go +++ b/internal/declarative/schema.go @@ -43,6 +43,7 @@ func Schema() []byte { "name": str, "monthly_budget_usd": map[string]any{"type": "number", "minimum": 0}, "allowed_models": strs, + "allowed_mcp": map[string]any{"type": "array", "items": str, "description": "MCP servers (server) and tools (server/tool) the team's keys may use; empty means any."}, }, "name")), }, "name")), "api_keys": list(obj(map[string]any{ @@ -57,6 +58,12 @@ func Schema() []byte { "vendor": str, "plain_key_env": map[string]any{"type": "string", "description": "Environment variable holding the vendor key; the key itself never goes in the file."}, }, "api_key_name", "vendor", "plain_key_env")), + "mcp_servers": list(messageSchema((&pb.MCPServerSpec{}).ProtoReflect().Descriptor(), "name", "url")), + "mcp_credentials": list(obj(map[string]any{ + "server": str, + "header": map[string]any{"type": "string", "description": "HTTP header the gateway sends the server; Authorization when left out."}, + "value_env": map[string]any{"type": "string", "description": "Environment variable holding the header's value, such as Bearer ; the value itself never goes in the file."}, + }, "server", "value_env")), "spend_limits": list(obj(map[string]any{ "subject_type": map[string]any{"enum": []string{"key", "user"}}, "subject_id": map[string]any{"type": "string", "description": "Key name or user id, * for every subject of that type."}, diff --git a/internal/declarative/validate.go b/internal/declarative/validate.go index 280889f..9634736 100644 --- a/internal/declarative/validate.go +++ b/internal/declarative/validate.go @@ -11,7 +11,7 @@ import ( ) // KeyScopes are the scopes a client key in the file may carry. -var KeyScopes = []string{apikeys.ScopeGenerate, apikeys.ScopeAdmin, apikeys.ScopeAdminRead} +var KeyScopes = []string{apikeys.ScopeGenerate, apikeys.ScopeMCP, apikeys.ScopeAdmin, apikeys.ScopeAdminRead} var efforts = []pb.Effort{pb.Effort_EFFORT_LOW, pb.Effort_EFFORT_MEDIUM, pb.Effort_EFFORT_HIGH} @@ -182,6 +182,29 @@ func (c *Config) Validate() error { sl[id] = true } + servers := map[string]bool{} + for i, m := range c.MCPServers { + switch { + case m.GetName() == "" || m.GetUrl() == "": + bad("mcp_servers[%d]: name and url are required", i) + case servers[m.GetName()]: + bad("mcp_servers[%d]: duplicate name %q", i, m.GetName()) + } + servers[m.GetName()] = true + } + creds := map[string]bool{} + for i, mc := range c.MCPCredentials { + switch { + case mc.Server == "" || mc.ValueEnv == "": + bad("mcp_credentials[%d]: server and value_env are required", i) + case creds[mc.Server]: + bad("mcp_credentials[%d]: duplicate server %q", i, mc.Server) + case c.Manages("mcp_servers") && !servers[mc.Server]: + bad("mcp_credentials[%d]: server %q is not in mcp_servers", i, mc.Server) + } + creds[mc.Server] = true + } + ur := map[string]bool{} for i, u := range c.UserRateLimits { id := u.GetUserId() + "/" + u.GetModelId() diff --git a/internal/grpcapi/admin.go b/internal/grpcapi/admin.go index 3987e31..dd5f83b 100644 --- a/internal/grpcapi/admin.go +++ b/internal/grpcapi/admin.go @@ -35,10 +35,11 @@ type AdminServer struct { spendReader SpendReader spendLimits func(subjectType, subjectID, period string) int64 pb.UnimplementedAdminServiceServer - db *proxydb.DB - keyring *crypto.Keyring - nodeMgr *nodes.Manager - explain RouteExplainer + db *proxydb.DB + keyring *crypto.Keyring + nodeMgr *nodes.Manager + explain RouteExplainer + mcpTools MCPToolLister // onChange is called after a successful config mutation so this replica // picks the change up immediately instead of waiting for the poll. onChange func() @@ -153,12 +154,26 @@ func (s *AdminServer) GetConfig(ctx context.Context, _ *pb.GetConfigRequest) (*p if err != nil { return nil, dbStatus(err) } + mcpServers, err := s.db.LoadMCPServers(ctx) + if err != nil { + return nil, dbStatus(err) + } + mcpCreds, err := s.db.LoadMCPCredentials(ctx) + if err != nil { + return nil, dbStatus(err) + } version, err := s.db.ConfigVersion(ctx) if err != nil { return nil, dbStatus(err) } resp := &pb.GetConfigResponse{Version: version} + for _, m := range mcpServers { + resp.McpServers = append(resp.McpServers, mcpServerToProto(m)) + } + for _, c := range mcpCreds { + resp.McpCredentials = append(resp.McpCredentials, s.mcpCredentialInfo(c)) + } for _, m := range models { resp.Models = append(resp.Models, modelToProto(m)) } @@ -368,11 +383,15 @@ func (s *AdminServer) MintAPIKey(ctx context.Context, req *pb.MintAPIKeyRequest) if err != nil { return nil, err } + allowedMCP, err := normalizeAllowedMCP(req.GetPolicy().GetAllowedMcp()) + if err != nil { + return nil, err + } plain, row, err := apikeys.Generate(req.GetName(), scopes) if err != nil { return nil, status.Error(codes.Internal, "key generation failed") } - row.AllowedModels, row.RPMLimit, row.TPMLimit = allowed, rpm, tpm + row.AllowedModels, row.RPMLimit, row.TPMLimit, row.AllowedMCP = allowed, rpm, tpm, allowedMCP var orgID, teamID *string if req.GetOrgId() != "" { @@ -452,8 +471,8 @@ func (s *AdminServer) ListAPIKeys(ctx context.Context, _ *pb.ListAPIKeysRequest) if k.TeamID != nil { info.TeamId = *k.TeamID } - if len(k.AllowedModels) > 0 || k.RPMLimit > 0 || k.TPMLimit > 0 { - info.Policy = &pb.APIKeyPolicy{AllowedModels: k.AllowedModels, RpmLimit: k.RPMLimit, TpmLimit: k.TPMLimit} + if len(k.AllowedModels) > 0 || k.RPMLimit > 0 || k.TPMLimit > 0 || len(k.AllowedMCP) > 0 { + info.Policy = &pb.APIKeyPolicy{AllowedModels: k.AllowedModels, RpmLimit: k.RPMLimit, TpmLimit: k.TPMLimit, AllowedMcp: k.AllowedMCP} } out = append(out, info) } @@ -468,7 +487,11 @@ func (s *AdminServer) SetAPIKeyPolicy(ctx context.Context, req *pb.SetAPIKeyPoli if err != nil { return nil, err } - if err := s.db.SetAPIKeyPolicy(ctx, req.GetName(), allowed, rpm, tpm); err != nil { + allowedMCP, err := normalizeAllowedMCP(req.GetPolicy().GetAllowedMcp()) + if err != nil { + return nil, err + } + if err := s.db.SetAPIKeyPolicy(ctx, req.GetName(), allowed, rpm, tpm, allowedMCP); err != nil { return nil, dbStatus(err) } return &pb.SetAPIKeyPolicyResponse{}, nil @@ -552,7 +575,7 @@ func (s *AdminServer) RotateEncryption(ctx context.Context, _ *pb.RotateEncrypti for _, vk := range keys { plain, err := s.keyring.Decrypt(vk.KeyCiphertext, int(vk.KeyVersion)) if err != nil { - return &unreadableKeyError{id: vk.ID, version: vk.KeyVersion} + return &unreadableKeyError{what: "vendor key " + vk.ID, version: vk.KeyVersion, fix: "SetVendorKey"} } ciphertext, version, err := s.keyring.Encrypt(plain) if err != nil { @@ -563,7 +586,9 @@ func (s *AdminServer) RotateEncryption(ctx context.Context, _ *pb.RotateEncrypti } count++ } - return nil + n, err := s.rotateMCPCredentials(ctx, tx, active) + count += n + return err }, nil) var unreadable *unreadableKeyError if errors.As(err, &unreadable) { @@ -584,12 +609,13 @@ func (s *AdminServer) RotateEncryption(ctx context.Context, _ *pb.RotateEncrypti // unreadableKeyError stops a rotation that would otherwise report success // while a key still needs a version the operator is about to drop. type unreadableKeyError struct { - id string + what string // "vendor key " or "mcp credential of " version int32 + fix string // the RPC that sets it again } func (e *unreadableKeyError) Error() string { - return fmt.Sprintf("vendor key %s is encrypted with version %d, which KTMS_ENCRYPTION_KEYS lacks: add that version back, or set the key again with SetVendorKey", e.id, e.version) + return fmt.Sprintf("%s is encrypted with version %d, which KTMS_ENCRYPTION_KEYS lacks: add that version back, or set it again with %s", e.what, e.version, e.fix) } // rotateNodeTokens re-encrypts the node auth tokens kept in Redis with the @@ -726,12 +752,17 @@ func (s *AdminServer) CreateTeam(ctx context.Context, req *pb.CreateTeamRequest) if allowed == nil { allowed = []string{} } + allowedMCP, err := normalizeAllowedMCP(req.GetAllowedMcp()) + if err != nil { + return nil, err + } team := &proxydb.Team{ ID: uuid.NewString(), OrgID: req.GetOrgId(), Name: req.GetName(), MonthlyBudgetUSD: req.GetMonthlyBudgetUsd(), AllowedModels: allowed, + AllowedMCP: allowedMCP, CreatedAt: time.Now().UTC(), } _, err = s.db.Mutate(ctx, func(tx pgx.Tx) error { @@ -796,11 +827,16 @@ func (s *AdminServer) UpdateTeam(ctx context.Context, req *pb.UpdateTeamRequest) if allowed == nil { allowed = []string{} } + allowedMCP, err := normalizeAllowedMCP(req.GetAllowedMcp()) + if err != nil { + return nil, err + } t := &proxydb.Team{ ID: req.GetId(), Name: req.GetName(), MonthlyBudgetUSD: req.GetMonthlyBudgetUsd(), AllowedModels: allowed, + AllowedMCP: allowedMCP, } _, err = s.db.Mutate(ctx, func(tx pgx.Tx) error { if err := proxydb.UpdateTeamTx(ctx, tx, t); err != nil { @@ -928,5 +964,6 @@ func teamToProto(t *proxydb.Team) *pb.Team { MonthlyBudgetUsd: t.MonthlyBudgetUSD, AllowedModels: t.AllowedModels, CreatedAt: timestamppb.New(t.CreatedAt), + AllowedMcp: t.AllowedMCP, } } diff --git a/internal/grpcapi/mcp.go b/internal/grpcapi/mcp.go new file mode 100644 index 0000000..855021a --- /dev/null +++ b/internal/grpcapi/mcp.go @@ -0,0 +1,224 @@ +package grpcapi + +import ( + "context" + "errors" + "net/url" + "slices" + "strings" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" + "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/rs/zerolog/log" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + "google.golang.org/protobuf/types/known/timestamppb" + + pb "github.com/ai-process/ktms/gen/ktms/v1" + "github.com/ai-process/ktms/internal/crypto" + "github.com/ai-process/ktms/internal/proxydb" + "github.com/ai-process/ktms/internal/registry" + "github.com/ai-process/ktms/internal/safefetch" +) + +// MCPToolLister asks a registered MCP server for its tools (mcpgw.Gateway). +type MCPToolLister interface { + ListTools(ctx context.Context, server string) ([]*mcp.Tool, error) +} + +// SetMCPToolLister lets ListMCPTools reach the servers; without it the RPC is +// UNAVAILABLE (the gateway is off). +func (s *AdminServer) SetMCPToolLister(l MCPToolLister) { s.mcpTools = l } + +// checkHost resolves an MCP server's host; tests replace it. +var checkHost = safefetch.CheckHost + +func (s *AdminServer) UpsertMCPServer(ctx context.Context, req *pb.UpsertMCPServerRequest) (*pb.UpsertMCPServerResponse, error) { + m := mcpServerFromProto(req.GetServer()) + if err := registry.ValidateMCPServer(m); err != nil { + return nil, status.Error(codes.InvalidArgument, err.Error()) + } + if !m.AllowPrivate { + u, _ := url.Parse(m.URL) // ValidateMCPServer parsed it + if err := checkHost(ctx, u.Hostname()); err != nil { + if errors.Is(err, safefetch.ErrProhibitedIP) { + return nil, status.Errorf(codes.InvalidArgument, "mcp server %s: %v; set allow_private for a server on a private network", m.Name, err) + } + return nil, status.Errorf(codes.InvalidArgument, "mcp server %s: %v", m.Name, err) + } + } + _, err := s.db.Mutate(ctx, func(tx pgx.Tx) error { + return proxydb.UpsertMCPServer(ctx, tx, m) + }, nil) + if err != nil { + return nil, mutationStatus(err) + } + s.onChange() + return &pb.UpsertMCPServerResponse{}, nil +} + +func (s *AdminServer) DeleteMCPServer(ctx context.Context, req *pb.DeleteMCPServerRequest) (*pb.DeleteMCPServerResponse, error) { + if req.GetName() == "" { + return nil, status.Error(codes.InvalidArgument, "name is required") + } + _, err := s.db.Mutate(ctx, func(tx pgx.Tx) error { + return proxydb.DeleteMCPServer(ctx, tx, req.GetName()) + }, nil) + if err != nil { + return nil, mutationStatus(err) + } + s.onChange() + return &pb.DeleteMCPServerResponse{}, nil +} + +func (s *AdminServer) SetMCPCredential(ctx context.Context, req *pb.SetMCPCredentialRequest) (*pb.SetMCPCredentialResponse, error) { + header := strings.TrimSpace(req.GetHeader()) + if header == "" { + header = "Authorization" + } + if !validHeaderName(header) { + return nil, status.Errorf(codes.InvalidArgument, "header %q is not a valid HTTP header name", header) + } + if req.GetPlainValue() == "" || strings.ContainsAny(req.GetPlainValue(), "\r\n") { + return nil, status.Error(codes.InvalidArgument, "plain_value is required and must be one line") + } + ciphertext, version, err := s.keyring.Encrypt(req.GetPlainValue()) + if err != nil { + return nil, status.Error(codes.Internal, "encryption failed") + } + _, err = s.db.Mutate(ctx, func(tx pgx.Tx) error { + return proxydb.SetMCPCredential(ctx, tx, &proxydb.MCPCredential{ + Server: req.GetServer(), Header: header, ValueCiphertext: ciphertext, ValueVersion: int32(version), + }) + }, nil) + var pgErr *pgconn.PgError + if errors.As(err, &pgErr) && pgErr.Code == "23503" { + return nil, status.Errorf(codes.NotFound, "no mcp server %q", req.GetServer()) + } + if err != nil { + return nil, mutationStatus(err) + } + s.onChange() + return &pb.SetMCPCredentialResponse{}, nil +} + +func (s *AdminServer) DeleteMCPCredential(ctx context.Context, req *pb.DeleteMCPCredentialRequest) (*pb.DeleteMCPCredentialResponse, error) { + _, err := s.db.Mutate(ctx, func(tx pgx.Tx) error { + return proxydb.DeleteMCPCredential(ctx, tx, req.GetServer()) + }, nil) + if err != nil { + return nil, mutationStatus(err) + } + s.onChange() + return &pb.DeleteMCPCredentialResponse{}, nil +} + +func (s *AdminServer) ListMCPTools(ctx context.Context, req *pb.ListMCPToolsRequest) (*pb.ListMCPToolsResponse, error) { + if s.mcpTools == nil { + return nil, status.Error(codes.Unavailable, "the MCP gateway is off (MCP_GATEWAY_ENABLED=false)") + } + tools, err := s.mcpTools.ListTools(ctx, req.GetServer()) + if err != nil { + // The error names the server and its trouble, never its credential. + return nil, status.Error(codes.FailedPrecondition, err.Error()) + } + resp := &pb.ListMCPToolsResponse{} + for _, t := range tools { + resp.Tools = append(resp.Tools, &pb.MCPToolInfo{Name: t.Name, Description: t.Description}) + } + return resp, nil +} + +// rotateMCPCredentials re-encrypts the MCP credentials not at the active +// keyring version, inside the vendor keys' rotation transaction. +func (s *AdminServer) rotateMCPCredentials(ctx context.Context, tx pgx.Tx, active int32) (int32, error) { + creds, err := proxydb.LockStaleMCPCredentialsTx(ctx, tx, active) + if err != nil { + return 0, err + } + var n int32 + for _, c := range creds { + plain, err := s.keyring.Decrypt(c.ValueCiphertext, int(c.ValueVersion)) + if err != nil { + return n, &unreadableKeyError{what: "mcp credential of " + c.Server, version: c.ValueVersion, fix: "SetMCPCredential"} + } + ciphertext, version, err := s.keyring.Encrypt(plain) + if err != nil { + return n, err + } + if err := proxydb.UpdateMCPCredentialCiphertext(ctx, tx, c.Server, ciphertext, int32(version)); err != nil { + return n, err + } + n++ + } + return n, nil +} + +func (s *AdminServer) mcpCredentialInfo(c *proxydb.MCPCredential) *pb.MCPCredentialInfo { + info := &pb.MCPCredentialInfo{Server: c.Server, Header: c.Header, KeyVersion: c.ValueVersion, CreatedAt: timestamppb.New(c.CreatedAt)} + if s.keyring != nil { + if plain, err := s.keyring.Decrypt(c.ValueCiphertext, int(c.ValueVersion)); err == nil { + info.Fingerprint = crypto.Fingerprint(plain) + } else { + log.Warn().Err(err).Str("mcp_server", c.Server).Msg("mcp credential unreadable for its fingerprint") + } + } + return info +} + +func mcpServerFromProto(p *pb.MCPServerSpec) *proxydb.MCPServer { + t := p.GetTransport() + if t == "" { + t = registry.MCPStreamableHTTP + } + return &proxydb.MCPServer{ + Name: p.GetName(), + URL: p.GetUrl(), + Transport: t, + Enabled: p.GetEnabled(), + AllowPrivate: p.GetAllowPrivate(), + TimeoutMs: p.GetTimeoutMs(), + MaxResponseBytes: p.GetMaxResponseBytes(), + ToolRPM: p.GetToolRpm(), + } +} + +func mcpServerToProto(m *proxydb.MCPServer) *pb.MCPServerSpec { + return &pb.MCPServerSpec{ + Name: m.Name, + Url: m.URL, + Transport: m.Transport, + Enabled: m.Enabled, + AllowPrivate: m.AllowPrivate, + TimeoutMs: m.TimeoutMs, + MaxResponseBytes: m.MaxResponseBytes, + ToolRpm: m.ToolRPM, + } +} + +// normalizeAllowedMCP checks entries "server" and "server/tool", dropping +// duplicates. Servers are not checked against the registry, like models. +func normalizeAllowedMCP(in []string) ([]string, error) { + out := []string{} + for _, e := range in { + e = strings.TrimSpace(e) + server, tool, hasTool := strings.Cut(e, "/") + if !registry.ValidMCPName(server) || (hasTool && tool == "") { + return nil, status.Errorf(codes.InvalidArgument, "allowed_mcp entry %q: want server or server/tool", e) + } + if !slices.Contains(out, e) { + out = append(out, e) + } + } + return out, nil +} + +func validHeaderName(h string) bool { + for _, r := range h { + if r > 127 || !(r == '-' || r == '_' || r >= '0' && r <= '9' || r >= 'a' && r <= 'z' || r >= 'A' && r <= 'Z') { + return false + } + } + return h != "" && !strings.EqualFold(h, "host") && !strings.EqualFold(h, "content-length") +} diff --git a/internal/grpcapi/mcp_test.go b/internal/grpcapi/mcp_test.go new file mode 100644 index 0000000..1152a20 --- /dev/null +++ b/internal/grpcapi/mcp_test.go @@ -0,0 +1,148 @@ +package grpcapi + +import ( + "context" + "slices" + "strings" + "testing" + + "github.com/google/uuid" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + "google.golang.org/protobuf/encoding/protojson" + + pb "github.com/ai-process/ktms/gen/ktms/v1" + "github.com/ai-process/ktms/internal/apikeys" + "github.com/ai-process/ktms/internal/registry" +) + +func TestMCPServerRegistryAndCredential(t *testing.T) { + db := rotateTestDB(t) + ctx := context.Background() + const secret = "Bearer ghp_upstream_secret" + s := NewAdminServer(db, keyring(t, 1), nil) + + for _, c := range []struct { + spec *pb.MCPServerSpec + code codes.Code + }{ + {&pb.MCPServerSpec{Name: "local", Url: "https://127.0.0.1/mcp", Enabled: true}, codes.InvalidArgument}, + {&pb.MCPServerSpec{Name: "meta", Url: "https://169.254.169.254/", Enabled: true}, codes.InvalidArgument}, + {&pb.MCPServerSpec{Name: "plain", Url: "http://93.184.215.14/mcp", Enabled: true}, codes.InvalidArgument}, + {&pb.MCPServerSpec{Name: "Bad_Name", Url: "https://93.184.215.14/mcp"}, codes.InvalidArgument}, + {&pb.MCPServerSpec{Name: "sse", Url: "https://93.184.215.14/sse", Transport: "ws"}, codes.InvalidArgument}, + {&pb.MCPServerSpec{Name: "creds", Url: "https://u:p@93.184.215.14/mcp"}, codes.InvalidArgument}, + {&pb.MCPServerSpec{Name: "inside", Url: "http://10.0.0.7:8000/mcp", Enabled: true, AllowPrivate: true}, codes.OK}, + {&pb.MCPServerSpec{Name: "github", Url: "https://93.184.215.14/mcp", Enabled: true, ToolRpm: 30}, codes.OK}, + } { + _, err := s.UpsertMCPServer(ctx, &pb.UpsertMCPServerRequest{Server: c.spec}) + if status.Code(err) != c.code { + t.Fatalf("%s: %v, want %s", c.spec.GetName(), err, c.code) + } + } + if _, err := s.SetMCPCredential(ctx, &pb.SetMCPCredentialRequest{Server: "nope", PlainValue: "x"}); status.Code(err) != codes.NotFound { + t.Fatalf("credential of an unknown server: %v", err) + } + if _, err := s.SetMCPCredential(ctx, &pb.SetMCPCredentialRequest{Server: "github", Header: "Bad Header", PlainValue: "x"}); status.Code(err) != codes.InvalidArgument { + t.Fatalf("bad header: %v", err) + } + if _, err := s.SetMCPCredential(ctx, &pb.SetMCPCredentialRequest{Server: "github", PlainValue: secret}); err != nil { + t.Fatal(err) + } + + cfg, err := s.GetConfig(ctx, &pb.GetConfigRequest{}) + if err != nil { + t.Fatal(err) + } + b, _ := protojson.Marshal(cfg) + if strings.Contains(string(b), "ghp_upstream_secret") { + t.Fatalf("GetConfig returned the credential: %s", b) + } + if len(cfg.GetMcpServers()) != 2 || cfg.GetMcpServers()[0].GetTransport() != registry.MCPStreamableHTTP { + t.Fatalf("servers = %v", cfg.GetMcpServers()) + } + creds := cfg.GetMcpCredentials() + if len(creds) != 1 || creds[0].GetHeader() != "Authorization" || creds[0].GetFingerprint() == "" { + t.Fatalf("credentials = %v", creds) + } + + // The loader hands the gateway the decrypted credential. + l := registry.NewLoader(db, keyring(t, 1), nil) + if err := l.Load(ctx); err != nil { + t.Fatal(err) + } + if gh := l.Snapshot().MCPServers["github"]; gh == nil || gh.Value != secret || gh.ToolRPM != 30 { + t.Fatalf("snapshot server = %+v", gh) + } + + // Rotation re-encrypts MCP credentials with the vendor keys. + both := NewAdminServer(db, keyring(t, 1, 2), nil) + if res, err := both.RotateEncryption(ctx, &pb.RotateEncryptionRequest{}); err != nil || res.GetReencrypted() != 1 { + t.Fatalf("rotation = %v, %v", res, err) + } + l = registry.NewLoader(db, keyring(t, 2), nil) + if err := l.Load(ctx); err != nil { + t.Fatal(err) + } + if gh := l.Snapshot().MCPServers["github"]; gh == nil || gh.Value != secret { + t.Fatalf("after rotation, version 2 alone reads %+v", gh) + } + + if _, err := s.DeleteMCPServer(ctx, &pb.DeleteMCPServerRequest{Name: "github"}); err != nil { + t.Fatal(err) + } + if creds, _ := db.LoadMCPCredentials(ctx); len(creds) != 0 { + t.Fatalf("credential outlived its server: %v", creds) + } + if _, err := s.ListMCPTools(ctx, &pb.ListMCPToolsRequest{Server: "inside"}); status.Code(err) != codes.Unavailable { + t.Fatalf("ListMCPTools without a gateway: %v", err) + } +} + +func TestKeyAndTeamAllowedMCP(t *testing.T) { + db := rotateTestDB(t) + ctx := context.Background() + s := NewAdminServer(db, keyring(t, 1), nil) + org, err := s.CreateOrganization(ctx, &pb.CreateOrganizationRequest{Name: "acme"}) + if err != nil { + t.Fatal(err) + } + if _, err := s.CreateTeam(ctx, &pb.CreateTeamRequest{OrgId: org.GetOrganization().GetId(), Name: "t", AllowedMcp: []string{"github/"}}); status.Code(err) != codes.InvalidArgument { + t.Fatalf("bad entry: %v", err) + } + team, err := s.CreateTeam(ctx, &pb.CreateTeamRequest{OrgId: org.GetOrganization().GetId(), Name: "t", AllowedMcp: []string{"github", "docs"}}) + if err != nil { + t.Fatal(err) + } + name := "agent-" + uuid.NewString()[:8] + if _, err := s.MintAPIKey(ctx, &pb.MintAPIKeyRequest{Name: name, Scopes: []string{"mcp"}, TeamId: team.GetTeam().GetId(), + Policy: &pb.APIKeyPolicy{AllowedMcp: []string{"github/create_issue", "github/create_issue"}}}); err != nil { + t.Fatal(err) + } + row, err := db.GetAPIKeyByName(ctx, name) + if err != nil { + t.Fatal(err) + } + id := apikeys.IdentityFor(row) + if !slices.Equal(row.AllowedMCP, []string{"github/create_issue"}) || !id.CanUseTool("github", "create_issue") || + id.CanUseTool("github", "delete_repo") || id.CanUseTool("docs", "search") || !id.HasScope(apikeys.ScopeMCP) { + t.Fatalf("identity = %+v", id) + } + + if _, err := s.SetAPIKeyPolicy(ctx, &pb.SetAPIKeyPolicyRequest{Name: name}); err != nil { + t.Fatal(err) + } + row, _ = db.GetAPIKeyByName(ctx, name) + id = apikeys.IdentityFor(row) + if !id.CanUseTool("docs", "search") || id.CanUseTool("slack", "post") || id.MCPDeniedBy("slack", "post") != "team" { + t.Fatalf("without a key list the team's applies: %+v", id) + } + keys, err := s.ListAPIKeys(ctx, &pb.ListAPIKeysRequest{}) + if err != nil || len(keys.GetKeys()) != 1 || keys.GetKeys()[0].GetPolicy() != nil { + t.Fatalf("keys = %v, %v", keys, err) + } + got, err := s.GetTeam(ctx, &pb.GetTeamRequest{Id: team.GetTeam().GetId()}) + if err != nil || !slices.Equal(got.GetTeam().GetAllowedMcp(), []string{"github", "docs"}) { + t.Fatalf("team = %v, %v", got, err) + } +} diff --git a/internal/mcpgw/gateway.go b/internal/mcpgw/gateway.go new file mode 100644 index 0000000..01a512c --- /dev/null +++ b/internal/mcpgw/gateway.go @@ -0,0 +1,339 @@ +package mcpgw + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "slices" + "sort" + "strings" + "sync" + "time" + + "github.com/google/uuid" + "github.com/modelcontextprotocol/go-sdk/jsonrpc" + "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/rs/zerolog/log" + + "github.com/ai-process/ktms/internal/apikeys" + "github.com/ai-process/ktms/internal/observability" + "github.com/ai-process/ktms/internal/registry" + "github.com/ai-process/ktms/internal/router" + "github.com/ai-process/ktms/internal/usagereport" + usagev1 "github.com/ai-process/ktms/internal/usagereport/gen/usage/v1" +) + +// Separator joins a server and a tool name in the /mcp aggregate. +const Separator = "__" + +// CodeForbidden is the JSON-RPC error of a tool or server the key may not use. +const CodeForbidden = -32003 + +// UserHeader names the end user a call is made for, as user_id does on /v1. +const UserHeader = "X-KTMS-User" + +// Snapshots is the live configuration (registry.Loader). +type Snapshots interface { + Snapshot() *registry.Snapshot +} + +// Verifier authenticates a client key for a scope (apikeys.Verifier). +type Verifier interface { + VerifyScope(ctx context.Context, plain, scope string) (*apikeys.Identity, error) +} + +// ToolLimiter enforces tool_rpm (throttle.Redis). +type ToolLimiter interface { + AdmitTool(ctx context.Context, keyName, server, tool string, rpm int32) (bool, time.Duration) +} + +// Recorder accounts for tool calls (usagereport.UsageCollectorClient). +type Recorder interface { + RecordToolCall(usagereport.ToolCall) error +} + +// Gateway serves /mcp and /mcp/{server}. +type Gateway struct { + snaps Snapshots + verifier Verifier + limiter ToolLimiter + recorder Recorder + version string + pool pool + sdk *mcp.StreamableHTTPHandler +} + +// New builds a gateway; limiter and recorder may be nil. +func New(snaps Snapshots, verifier Verifier, limiter ToolLimiter, recorder Recorder, version string) *Gateway { + g := &Gateway{snaps: snaps, verifier: verifier, limiter: limiter, recorder: recorder, version: version} + // Stateless: any replica answers any request, and every request carries + // the key, so a session cannot outlive the key's revocation. + g.sdk = mcp.NewStreamableHTTPHandler(g.server, &mcp.StreamableHTTPOptions{ + Stateless: true, + JSONResponse: true, + DisableLocalhostProtection: true, // a key is required; KTMS often sits behind a local proxy + }) + return g +} + +// Register mounts the gateway on mux. +func (g *Gateway) Register(mux *http.ServeMux) { + mux.HandleFunc("/mcp", g.serveHTTP) + mux.HandleFunc("/mcp/{server}", g.serveHTTP) +} + +// Close ends every upstream session. +func (g *Gateway) Close() { g.pool.closeAll() } + +type ctxKey struct{} + +type caller struct { + id *apikeys.Identity + server string // "" on the aggregate + user string +} + +func (g *Gateway) serveHTTP(w http.ResponseWriter, r *http.Request) { + plain, err := apikeys.ParseBearer(r.Header.Get("Authorization")) + if err != nil { + w.Header().Set("WWW-Authenticate", `Bearer realm="ktms"`) + http.Error(w, "unauthenticated: send Authorization: Bearer ", http.StatusUnauthorized) + return + } + id, err := g.verifier.VerifyScope(r.Context(), plain, apikeys.ScopeMCP) + switch { + case errors.Is(err, apikeys.ErrUnauthenticated): + w.Header().Set("WWW-Authenticate", `Bearer realm="ktms"`) + http.Error(w, "unauthenticated", http.StatusUnauthorized) + return + case errors.Is(err, apikeys.ErrForbidden): + http.Error(w, "permission_denied: the key needs the mcp scope", http.StatusForbidden) + return + case err != nil: + http.Error(w, "unavailable", http.StatusServiceUnavailable) + return + } + c := &caller{id: id, server: r.PathValue("server"), user: r.Header.Get(UserHeader)} + if c.server != "" { + if g.snaps.Snapshot().MCPServers[c.server] == nil { + http.Error(w, "not_found: no mcp server "+c.server, http.StatusNotFound) + return + } + if !id.CanUseTool(c.server, "") { + observability.MCPToolCallsTotal.WithLabelValues(c.server, "", "denied").Inc() + http.Error(w, fmt.Sprintf("permission_denied: mcp server %s is not in the %s's allowed_mcp", c.server, id.MCPDeniedBy(c.server, "")), http.StatusForbidden) + return + } + } + g.sdk.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), ctxKey{}, c))) +} + +// server is the MCP server one request talks to: tools only, the ones its key +// may use, answered from the upstreams. +func (g *Gateway) server(r *http.Request) *mcp.Server { + c, _ := r.Context().Value(ctxKey{}).(*caller) + name, instructions := "ktms", "Tools of the MCP servers registered in Keystone (KTMS), named server__tool." + if c != nil && c.server != "" { + name, instructions = "ktms-"+c.server, "" + } + s := mcp.NewServer(&mcp.Implementation{Name: name, Version: g.version}, &mcp.ServerOptions{ + Instructions: instructions, + Capabilities: &mcp.ServerCapabilities{Tools: &mcp.ToolCapabilities{}}, + }) + s.AddReceivingMiddleware(func(next mcp.MethodHandler) mcp.MethodHandler { + return func(ctx context.Context, method string, req mcp.Request) (mcp.Result, error) { + if c == nil { + return next(ctx, method, req) + } + // No typed nil may reach the SDK as a result alongside an error. + switch method { + case "tools/list": + if res, err := g.listTools(ctx, c); err != nil { + return nil, err + } else { + return res, nil + } + case "tools/call": + p, _ := req.GetParams().(*mcp.CallToolParamsRaw) + if p == nil { + return nil, &jsonrpc.Error{Code: jsonrpc.CodeInvalidParams, Message: "missing params"} + } + if res, err := g.callTool(ctx, c, p); err != nil { + return nil, err + } else { + return res, nil + } + } + return next(ctx, method, req) + } + }) + return s +} + +// servers lists the enabled servers the caller may see, in name order. +func (g *Gateway) servers(c *caller) []*registry.MCPServer { + all := g.snaps.Snapshot().MCPServers + g.pool.prune(all) + var out []*registry.MCPServer + for name, srv := range all { + if (c.server == "" || c.server == name) && c.id.CanUseTool(name, "") { + out = append(out, srv) + } + } + sort.Slice(out, func(i, j int) bool { return out[i].Name < out[j].Name }) + return out +} + +func (g *Gateway) listTools(ctx context.Context, c *caller) (*mcp.ListToolsResult, error) { + servers := g.servers(c) + lists := make([][]*mcp.Tool, len(servers)) + var wg sync.WaitGroup + for i, srv := range servers { + wg.Add(1) + go func() { + defer wg.Done() + tools, err := g.pool.get(srv).listTools(ctx, g.version) + if err != nil { + // One server down must not hide the others' tools. + observability.DependencyErrorsTotal.WithLabelValues("mcp", "tools_list").Inc() + log.Warn().Err(err).Str("mcp_server", srv.Name).Msg("mcp server tools/list failed; its tools are left out") + return + } + lists[i] = tools + }() + } + wg.Wait() + if c.server != "" && len(servers) == 1 && lists[0] == nil { + return nil, &jsonrpc.Error{Code: jsonrpc.CodeInternalError, Message: "mcp server " + c.server + " is unavailable"} + } + // The list depends on the key: no intermediary may serve it to another. + res := &mcp.ListToolsResult{Tools: []*mcp.Tool{}, Cacheable: mcp.Cacheable{CacheScope: "private"}} + for i, srv := range servers { + for _, t := range lists[i] { + if !c.id.CanUseTool(srv.Name, t.Name) { + continue + } + out := *t + if c.server == "" { + out.Name = srv.Name + Separator + t.Name + } + res.Tools = append(res.Tools, &out) + } + } + return res, nil +} + +func (g *Gateway) callTool(ctx context.Context, c *caller, p *mcp.CallToolParamsRaw) (*mcp.CallToolResult, error) { + server, tool := c.server, p.Name + if server == "" { + var ok bool + if server, tool, ok = strings.Cut(p.Name, Separator); !ok { + return nil, &jsonrpc.Error{Code: jsonrpc.CodeInvalidParams, Message: fmt.Sprintf("unknown tool %q: tools here are named server%stool", p.Name, Separator)} + } + } + srv := g.snaps.Snapshot().MCPServers[server] + if srv == nil { + return nil, &jsonrpc.Error{Code: jsonrpc.CodeInvalidParams, Message: fmt.Sprintf("unknown tool %q: no mcp server %s", p.Name, server)} + } + // Only the server's own tool names become metric labels and rate keys. + if known, err := g.pool.get(srv).listTools(ctx, g.version); err == nil && !slices.ContainsFunc(known, func(t *mcp.Tool) bool { return t.Name == tool }) { + return nil, &jsonrpc.Error{Code: jsonrpc.CodeInvalidParams, Message: fmt.Sprintf("unknown tool %q", p.Name)} + } + if !c.id.CanUseTool(server, tool) { + observability.MCPToolCallsTotal.WithLabelValues(server, tool, "denied").Inc() + return nil, &jsonrpc.Error{Code: CodeForbidden, Message: fmt.Sprintf("permission_denied: tool %s of mcp server %s is not in the %s's allowed_mcp", tool, server, c.id.MCPDeniedBy(server, tool))} + } + if g.limiter != nil && srv.ToolRPM > 0 { + if ok, retry := g.limiter.AdmitTool(ctx, c.id.Name, server, tool, srv.ToolRPM); !ok { + observability.MCPToolCallsTotal.WithLabelValues(server, tool, "throttled").Inc() + observability.ThrottledTotal.WithLabelValues("mcp_tool_rpm").Inc() + return toolError(fmt.Sprintf("throttled: %s allows %d calls a minute per key; retry in %ds", p.Name, srv.ToolRPM, int(retry.Seconds()+0.999))), nil + } + } + + start := time.Now() + cctx, cancel := context.WithTimeout(ctx, srv.Timeout()) + defer cancel() + cctx, cut := withCutMark(cctx) + res, err := g.pool.get(srv).call(cctx, g.version, &mcp.CallToolParams{Name: tool, Arguments: p.Arguments}) + elapsed := time.Since(start) + + var size int64 + st, class := usagev1.Status_STATUS_SUCCESS, "" + var rpcErr *jsonrpc.Error + switch { + case err == nil: + b, merr := json.Marshal(res) + size = int64(len(b)) + switch { + case merr != nil || size > srv.MaxResponse(): + st, class = usagev1.Status_STATUS_ERROR, "too_large" + res = toolError(fmt.Sprintf("the result of %s exceeds max_response_bytes (%d)", p.Name, srv.MaxResponse())) + case res.IsError: + st, class = usagev1.Status_STATUS_ERROR, "tool_error" + } + case errors.As(err, &rpcErr): + st, class = usagev1.Status_STATUS_ERROR, "tool_error" + case cut.Load() || errors.Is(err, errTooLarge): + st, class = usagev1.Status_STATUS_ERROR, "too_large" + res, err = toolError(fmt.Sprintf("the result of %s exceeds max_response_bytes (%d)", p.Name, srv.MaxResponse())), nil + case ctx.Err() != nil: + st, class = usagev1.Status_STATUS_CANCELLED, "cancelled" + case errors.Is(err, context.DeadlineExceeded) || cctx.Err() != nil: + st, class = usagev1.Status_STATUS_TIMEOUT, "timeout" + res, err = toolError(fmt.Sprintf("mcp server %s did not answer within %s", server, srv.Timeout())), nil + default: + st, class = usagev1.Status_STATUS_ERROR, "upstream_error" + log.Warn().Err(err).Str("mcp_server", server).Str("tool", tool).Msg("mcp tool call failed") + res, err = toolError(fmt.Sprintf("mcp server %s is unavailable", server)), nil + } + g.account(c, server, tool, st, class, elapsed, size) + return res, err +} + +func (g *Gateway) account(c *caller, server, tool string, st usagev1.Status, class string, elapsed time.Duration, size int64) { + status := class + if status == "" { + status = "success" + } + observability.MCPToolCallsTotal.WithLabelValues(server, tool, status).Inc() + observability.MCPToolCallDuration.WithLabelValues(server).Observe(elapsed.Seconds()) + if size > 0 { + observability.MCPResponseBytes.WithLabelValues(server).Observe(float64(size)) + } + if g.recorder == nil { + return + } + actionID := uuid.Must(uuid.NewV7()).String() + err := g.recorder.RecordToolCall(usagereport.ToolCall{ + UserID: router.UserID(map[string]string{"user_id": c.user}, c.id.Name), + ActionID: actionID, + Server: server, + Tool: tool, + ErrorClass: class, + DurationMs: elapsed.Milliseconds(), + ResponseBytes: size, + Status: st, + Meta: map[string]string{"api_key": c.id.Name, "mcp_server": server, "mcp_tool": tool}, + }) + if err != nil { + log.Warn().Err(err).Str("action_id", actionID).Str("mcp_server", server).Str("tool", tool).Msg("tool call not accounted") + } +} + +// toolError is a failure the model sees and may act on, as MCP wants for +// errors of the tool rather than of the protocol. +func toolError(msg string) *mcp.CallToolResult { + return &mcp.CallToolResult{IsError: true, Content: []mcp.Content{&mcp.TextContent{Text: msg}}} +} + +// ListTools is the server's tools for an operator (Admin API ListMCPTools). +func (g *Gateway) ListTools(ctx context.Context, server string) ([]*mcp.Tool, error) { + srv := g.snaps.Snapshot().MCPServers[server] + if srv == nil { + return nil, fmt.Errorf("mcp server %s is not registered or not enabled", server) + } + return g.pool.get(srv).listTools(ctx, g.version) +} diff --git a/internal/mcpgw/gateway_test.go b/internal/mcpgw/gateway_test.go new file mode 100644 index 0000000..b354c32 --- /dev/null +++ b/internal/mcpgw/gateway_test.go @@ -0,0 +1,422 @@ +package mcpgw + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "slices" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/modelcontextprotocol/go-sdk/mcp" + + "github.com/ai-process/ktms/internal/apikeys" + "github.com/ai-process/ktms/internal/proxydb" + "github.com/ai-process/ktms/internal/registry" + "github.com/ai-process/ktms/internal/usagereport" + usagev1 "github.com/ai-process/ktms/internal/usagereport/gen/usage/v1" +) + +type snaps struct { + s atomic.Pointer[registry.Snapshot] +} + +func (f *snaps) Snapshot() *registry.Snapshot { return f.s.Load() } + +type verifier map[string]*apikeys.Identity + +func (v verifier) VerifyScope(_ context.Context, plain, scope string) (*apikeys.Identity, error) { + id := v[plain] + if id == nil { + return nil, apikeys.ErrUnauthenticated + } + if !id.HasScope(scope) { + return nil, apikeys.ErrForbidden + } + return id, nil +} + +type recorder struct { + mu sync.Mutex + calls []usagereport.ToolCall +} + +func (r *recorder) RecordToolCall(t usagereport.ToolCall) error { + r.mu.Lock() + defer r.mu.Unlock() + r.calls = append(r.calls, t) + return nil +} + +func (r *recorder) all() []usagereport.ToolCall { + r.mu.Lock() + defer r.mu.Unlock() + return slices.Clone(r.calls) +} + +type limiter struct{ allow bool } + +func (l limiter) AdmitTool(context.Context, string, string, string, int32) (bool, time.Duration) { + return l.allow, 2 * time.Second +} + +type echoArgs struct { + Text string `json:"text"` +} + +// upstreamServer is an MCP server with the given tools; each echoes its +// text prefixed by the tool name and reports the credential header it saw. +func upstreamServer(t *testing.T, sse bool, tools ...string) (url string, seen *atomic.Value, calls *atomic.Int32) { + t.Helper() + seen, calls = &atomic.Value{}, &atomic.Int32{} + s := mcp.NewServer(&mcp.Implementation{Name: "up", Version: "1"}, nil) + for _, name := range tools { + mcp.AddTool(s, &mcp.Tool{Name: name, Description: "the " + name + " tool"}, + func(_ context.Context, _ *mcp.CallToolRequest, in echoArgs) (*mcp.CallToolResult, any, error) { + calls.Add(1) + if in.Text == "huge" { + return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: strings.Repeat("x", envelope+8192)}}}, nil, nil + } + if in.Text == "big" { + return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: strings.Repeat("x", 4096)}}}, nil, nil + } + if in.Text == "slow" { + time.Sleep(300 * time.Millisecond) + } + return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: name + ":" + in.Text}}}, nil, nil + }) + } + var h http.Handler + if sse { + h = mcp.NewSSEHandler(func(*http.Request) *mcp.Server { return s }, nil) + } else { + h = mcp.NewStreamableHTTPHandler(func(*http.Request) *mcp.Server { return s }, nil) + } + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if v := r.Header.Get("X-Upstream-Token"); v != "" { + seen.Store(v) + } + h.ServeHTTP(w, r) + })) + t.Cleanup(ts.Close) + return ts.URL, seen, calls +} + +type fixture struct { + gw *Gateway + ts *httptest.Server + rec *recorder + snaps *snaps + githubN *atomic.Int32 + docsN *atomic.Int32 + seen *atomic.Value +} + +func newFixture(t *testing.T, lim ToolLimiter) *fixture { + t.Helper() + ghURL, seen, ghN := upstreamServer(t, false, "create_issue", "delete_repo") + docsURL, _, docsN := upstreamServer(t, true, "search") + f := &fixture{rec: &recorder{}, snaps: &snaps{}, githubN: ghN, docsN: docsN, seen: seen} + f.snaps.s.Store(®istry.Snapshot{MCPServers: map[string]*registry.MCPServer{ + "github": {MCPServer: proxydb.MCPServer{Name: "github", URL: ghURL, Transport: registry.MCPStreamableHTTP, Enabled: true, AllowPrivate: true, MaxResponseBytes: 2048, TimeoutMs: 200, ToolRPM: 5}, + Header: "X-Upstream-Token", Value: "upstream-secret"}, + "docs": {MCPServer: proxydb.MCPServer{Name: "docs", URL: docsURL + "/", Transport: registry.MCPSSE, Enabled: true, AllowPrivate: true}}, + }}) + v := verifier{ + "all": {Name: "all", Scopes: []string{apikeys.ScopeMCP}}, + "narrow": {Name: "narrow", Scopes: []string{apikeys.ScopeMCP}, KeyAllowedMCP: []string{"github/create_issue", "docs"}}, + "team": {Name: "team", Scopes: []string{apikeys.ScopeMCP}, TeamAllowedMCP: []string{"docs"}}, + "gen": {Name: "gen", Scopes: []string{apikeys.ScopeGenerate}}, + } + f.gw = New(f.snaps, v, lim, f.rec, "test") + t.Cleanup(f.gw.Close) + mux := http.NewServeMux() + f.gw.Register(mux) + f.ts = httptest.NewServer(mux) + t.Cleanup(f.ts.Close) + return f +} + +type bearer struct{ key string } + +func (b bearer) RoundTrip(r *http.Request) (*http.Response, error) { + r = r.Clone(r.Context()) + r.Header.Set("Authorization", "Bearer "+b.key) + r.Header.Set(UserHeader, "alice") + return http.DefaultTransport.RoundTrip(r) +} + +func (f *fixture) connect(t *testing.T, key, path string) *mcp.ClientSession { + t.Helper() + c := mcp.NewClient(&mcp.Implementation{Name: "test"}, nil) + s, err := c.Connect(context.Background(), &mcp.StreamableClientTransport{Endpoint: f.ts.URL + path, HTTPClient: &http.Client{Transport: bearer{key}}}, nil) + if err != nil { + t.Fatalf("connect %s as %s: %v", path, key, err) + } + t.Cleanup(func() { _ = s.Close() }) + return s +} + +func toolNames(t *testing.T, s *mcp.ClientSession) []string { + t.Helper() + res, err := s.ListTools(context.Background(), nil) + if err != nil { + t.Fatalf("tools/list: %v", err) + } + var names []string + for _, tool := range res.Tools { + names = append(names, tool.Name) + } + slices.Sort(names) + return names +} + +func text(t *testing.T, r *mcp.CallToolResult) string { + t.Helper() + if len(r.Content) != 1 { + t.Fatalf("content = %v", r.Content) + } + return r.Content[0].(*mcp.TextContent).Text +} + +// Claude Code refuses a tools/list result whose cacheScope is not public or private. +func TestToolsListCacheScope(t *testing.T) { + f := newFixture(t, nil) + req, _ := http.NewRequest(http.MethodPost, f.ts.URL+"/mcp", strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"tools/list"}`)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + req.Header.Set("Authorization", "Bearer all") + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + var body struct { + Result struct { + CacheScope string `json:"cacheScope"` + } `json:"result"` + } + if err := json.NewDecoder(resp.Body).Decode(&body); err != nil || body.Result.CacheScope != "private" { + t.Fatalf("cacheScope = %q, %v", body.Result.CacheScope, err) + } +} + +func TestAggregateListsAllowedToolsOfTwoServers(t *testing.T) { + f := newFixture(t, nil) + if got := toolNames(t, f.connect(t, "all", "/mcp")); !slices.Equal(got, []string{"docs__search", "github__create_issue", "github__delete_repo"}) { + t.Fatalf("all: %v", got) + } + if got := toolNames(t, f.connect(t, "narrow", "/mcp")); !slices.Equal(got, []string{"docs__search", "github__create_issue"}) { + t.Fatalf("narrow: %v", got) + } + if got := toolNames(t, f.connect(t, "team", "/mcp")); !slices.Equal(got, []string{"docs__search"}) { + t.Fatalf("team: %v", got) + } +} + +func TestCallsReachUpstreamWithCredentialAndAreAccounted(t *testing.T) { + f := newFixture(t, nil) + s := f.connect(t, "narrow", "/mcp") + for _, c := range []struct{ name, want string }{ + {"github__create_issue", "create_issue:hi"}, + {"docs__search", "search:hi"}, + } { + res, err := s.CallTool(context.Background(), &mcp.CallToolParams{Name: c.name, Arguments: map[string]any{"text": "hi"}}) + if err != nil || res.IsError { + t.Fatalf("%s: %v %v", c.name, err, res) + } + if got := text(t, res); got != c.want { + t.Fatalf("%s = %q", c.name, got) + } + } + if v, _ := f.seen.Load().(string); v != "upstream-secret" { + t.Fatalf("upstream saw credential %q", v) + } + calls := f.rec.all() + if len(calls) != 2 { + t.Fatalf("accounted %d calls", len(calls)) + } + c := calls[0] + if c.Server != "github" || c.Tool != "create_issue" || c.Status != usagev1.Status_STATUS_SUCCESS || c.UserID != "alice" || c.ResponseBytes == 0 || c.Meta["api_key"] != "narrow" { + t.Fatalf("event = %+v", c) + } + for _, c := range calls { + b, _ := json.Marshal(c) + if strings.Contains(string(b), "upstream-secret") { + t.Fatalf("credential in usage event: %s", b) + } + } +} + +func TestForbiddenToolAndServerAreRefusedWithoutCallingUpstream(t *testing.T) { + f := newFixture(t, nil) + s := f.connect(t, "narrow", "/mcp") + _, err := s.CallTool(context.Background(), &mcp.CallToolParams{Name: "github__delete_repo", Arguments: map[string]any{"text": "x"}}) + if err == nil || !strings.Contains(err.Error(), "permission_denied") { + t.Fatalf("delete_repo: %v", err) + } + if _, err := s.CallTool(context.Background(), &mcp.CallToolParams{Name: "docs__search", Arguments: map[string]any{"text": "x"}}); err != nil { + t.Fatalf("a refusal must not end the session: %v", err) + } + if n := f.githubN.Load(); n != 0 { + t.Fatalf("upstream called %d times", n) + } + if calls := f.rec.all(); len(calls) != 1 || calls[0].Server != "docs" { + t.Fatalf("only the docs call is accounted, got %+v", calls) + } + + // The team's list refuses github as a whole, on /mcp/github too. + c := mcp.NewClient(&mcp.Implementation{Name: "test"}, nil) + _, err = c.Connect(context.Background(), &mcp.StreamableClientTransport{Endpoint: f.ts.URL + "/mcp/github", HTTPClient: &http.Client{Transport: bearer{"team"}}}, nil) + if err == nil || !strings.Contains(err.Error(), "Forbidden") { + t.Fatalf("team on /mcp/github: %v", err) + } +} + +func TestSingleServerEndpointKeepsToolNames(t *testing.T) { + f := newFixture(t, nil) + s := f.connect(t, "narrow", "/mcp/github") + if got := toolNames(t, s); !slices.Equal(got, []string{"create_issue"}) { + t.Fatalf("tools: %v", got) + } + res, err := s.CallTool(context.Background(), &mcp.CallToolParams{Name: "create_issue", Arguments: map[string]any{"text": "a"}}) + if err != nil || text(t, res) != "create_issue:a" { + t.Fatalf("call: %v %v", err, res) + } +} + +func TestAuthentication(t *testing.T) { + f := newFixture(t, nil) + for _, c := range []struct { + key string + want int + }{{"", http.StatusUnauthorized}, {"wrong", http.StatusUnauthorized}, {"gen", http.StatusForbidden}} { + req, _ := http.NewRequest(http.MethodPost, f.ts.URL+"/mcp", strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"tools/list"}`)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + if c.key != "" { + req.Header.Set("Authorization", "Bearer "+c.key) + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + resp.Body.Close() + if resp.StatusCode != c.want { + t.Fatalf("key %q: status %d, want %d", c.key, resp.StatusCode, c.want) + } + } + req, _ := http.NewRequest(http.MethodPost, f.ts.URL+"/mcp/nope", strings.NewReader(`{}`)) + req.Header.Set("Authorization", "Bearer all") + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + resp.Body.Close() + if resp.StatusCode != http.StatusNotFound { + t.Fatalf("unknown server: %d", resp.StatusCode) + } +} + +func TestLimits(t *testing.T) { + f := newFixture(t, nil) + s := f.connect(t, "all", "/mcp") + res, err := s.CallTool(context.Background(), &mcp.CallToolParams{Name: "github__create_issue", Arguments: map[string]any{"text": "big"}}) + if err != nil || !res.IsError || !strings.Contains(text(t, res), "max_response_bytes") { + t.Fatalf("big: %v %v", err, res) + } + res, err = s.CallTool(context.Background(), &mcp.CallToolParams{Name: "github__create_issue", Arguments: map[string]any{"text": "slow"}}) + if err != nil || !res.IsError || !strings.Contains(text(t, res), "did not answer") { + t.Fatalf("slow: %v %v", err, res) + } + calls := f.rec.all() + if len(calls) != 2 || calls[0].ErrorClass != "too_large" || calls[1].Status != usagev1.Status_STATUS_TIMEOUT { + t.Fatalf("accounted %+v", calls) + } + + f = newFixture(t, limiter{allow: false}) + s = f.connect(t, "all", "/mcp") + res, err = s.CallTool(context.Background(), &mcp.CallToolParams{Name: "github__create_issue", Arguments: map[string]any{"text": "a"}}) + if err != nil || !res.IsError || !strings.Contains(text(t, res), "throttled") || f.githubN.Load() != 0 { + t.Fatalf("throttled: %v %v", err, res) + } +} + +func TestReconnectsWhenTheCredentialChanges(t *testing.T) { + f := newFixture(t, nil) + s := f.connect(t, "all", "/mcp") + call := func() { + if _, err := s.CallTool(context.Background(), &mcp.CallToolParams{Name: "github__create_issue", Arguments: map[string]any{"text": "a"}}); err != nil { + t.Fatal(err) + } + } + call() + old := f.snaps.Snapshot() + gh := *old.MCPServers["github"] + gh.Value = "rotated" + f.snaps.s.Store(®istry.Snapshot{MCPServers: map[string]*registry.MCPServer{"github": &gh, "docs": old.MCPServers["docs"]}}) + call() + if v, _ := f.seen.Load().(string); v != "rotated" { + t.Fatalf("upstream saw %q after rotation", v) + } +} + +// A redirect must not carry the upstream credential to another host. +func TestRedirectsAreNotFollowed(t *testing.T) { + var leaked atomic.Value + evil := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + leaked.Store(r.Header.Get("X-Upstream-Token")) + })) + defer evil.Close() + redirect := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Redirect(w, r, evil.URL+"/steal", http.StatusTemporaryRedirect) + })) + defer redirect.Close() + g := New(&snaps{}, verifier{}, nil, nil, "test") + defer g.Close() + srv := ®istry.MCPServer{MCPServer: proxydb.MCPServer{Name: "r", URL: redirect.URL, Transport: registry.MCPStreamableHTTP, Enabled: true, AllowPrivate: true, TimeoutMs: 2000}, + Header: "X-Upstream-Token", Value: "upstream-secret"} + if _, err := g.pool.get(srv).listTools(context.Background(), "test"); err == nil { + t.Fatal("a redirecting server listed tools") + } + if v, _ := leaked.Load().(string); v != "" { + t.Fatalf("credential sent to the redirect target: %q", v) + } +} + +func TestUnknownToolsAndEmptyServers(t *testing.T) { + f := newFixture(t, nil) + s := f.connect(t, "all", "/mcp") + _, err := s.CallTool(context.Background(), &mcp.CallToolParams{Name: "github__no_such_tool", Arguments: map[string]any{}}) + if err == nil || !strings.Contains(err.Error(), "unknown tool") { + t.Fatalf("unknown tool: %v", err) + } + if f.githubN.Load() != 0 || len(f.rec.all()) != 0 { + t.Fatal("an unknown tool reached the upstream or the accounting") + } + + emptyURL, _, _ := upstreamServer(t, false) + snap := f.snaps.Snapshot() + snap.MCPServers["empty"] = ®istry.MCPServer{MCPServer: proxydb.MCPServer{Name: "empty", URL: emptyURL, Transport: registry.MCPStreamableHTTP, Enabled: true, AllowPrivate: true}} + if got := toolNames(t, f.connect(t, "all", "/mcp/empty")); len(got) != 0 { + t.Fatalf("empty server lists %v", got) + } +} + +// A body past max_response_bytes and the envelope is cut off while read; the +// call is still a too_large tool error, not an unavailable server. +func TestOversizedBodyIsTooLarge(t *testing.T) { + f := newFixture(t, nil) + s := f.connect(t, "all", "/mcp") + res, err := s.CallTool(context.Background(), &mcp.CallToolParams{Name: "github__create_issue", Arguments: map[string]any{"text": "huge"}}) + if err != nil || !res.IsError || !strings.Contains(text(t, res), "max_response_bytes") { + t.Fatalf("huge: %v %v", err, res) + } + if calls := f.rec.all(); len(calls) != 1 || calls[0].ErrorClass != "too_large" { + t.Fatalf("accounted %+v", calls) + } +} diff --git a/internal/mcpgw/upstream.go b/internal/mcpgw/upstream.go new file mode 100644 index 0000000..20e21ef --- /dev/null +++ b/internal/mcpgw/upstream.go @@ -0,0 +1,320 @@ +// Package mcpgw is Keystone's MCP gateway: clients reach registered MCP +// servers through /mcp with a KTMS key, and the gateway holds the servers' +// credentials, enforces each key's allowed_mcp and tool_rpm, and accounts for +// every tool call. +package mcpgw + +import ( + "context" + "errors" + "io" + "net" + "net/http" + "sync" + "sync/atomic" + "time" + + "github.com/modelcontextprotocol/go-sdk/jsonrpc" + "github.com/modelcontextprotocol/go-sdk/mcp" + + "github.com/ai-process/ktms/internal/registry" + "github.com/ai-process/ktms/internal/safefetch" +) + +// toolsTTL is how long a server's tool list is trusted; the gateway asks for +// no change notifications, so a new tool shows up within this. +const toolsTTL = time.Minute + +// envelope is room for the JSON-RPC wrapping around a result of max size. +const envelope = 64 << 10 + +var ( + errTooLarge = errors.New("mcp server response exceeds max_response_bytes") + errRedirect = errors.New("mcp server redirects are not followed") +) + +// cutKey marks a call's context: a reply body cut at the size cap sets it. The +// SDK reports a body read error without wrapping it, so only the mark tells. +type cutKey struct{} + +func withCutMark(ctx context.Context) (context.Context, *atomic.Bool) { + cut := &atomic.Bool{} + return context.WithValue(ctx, cutKey{}, cut), cut +} + +// upstream is one shared client session to a server, made with the spec and +// credential its fingerprint names. Sessions are reconnected on demand. +type upstream struct { + srv *registry.MCPServer + fp string + + mu sync.Mutex + session *mcp.ClientSession + cancel context.CancelFunc // ends the session's streams + tools []*mcp.Tool + toolsAt time.Time // zero: no list cached + retired bool // replaced or removed: never connects again +} + +var errRetired = errors.New("mcp server was reconfigured or removed during the call") + +// connect returns the live session, making one if there is none. The session +// outlives the request that made it: an SSE stream is bound to the context +// it was opened with, so that one ends only when the session is dropped. +func (u *upstream) connect(version string) (*mcp.ClientSession, error) { + u.mu.Lock() + defer u.mu.Unlock() + if u.retired { + return nil, errRetired + } + if u.session != nil { + return u.session, nil + } + client := mcp.NewClient(&mcp.Implementation{Name: "ktms-mcp-gateway", Version: version}, nil) + // A redirect would carry the credential to whatever host it names. + hc := &http.Client{Transport: &credentialTransport{srv: u.srv, base: transport(u.srv)}, + CheckRedirect: func(*http.Request, []*http.Request) error { return errRedirect }} + var t mcp.Transport + if u.srv.Transport == registry.MCPSSE { + t = &mcp.SSEClientTransport{Endpoint: u.srv.URL, HTTPClient: hc, MaxEventSize: int(u.srv.MaxResponse() + envelope)} + } else { + t = &mcp.StreamableClientTransport{Endpoint: u.srv.URL, HTTPClient: hc, MaxRetries: -1, + DisableStandaloneSSE: true, MaxEventSize: int(u.srv.MaxResponse() + envelope)} + } + ctx, cancel := context.WithCancel(context.Background()) + timer := time.AfterFunc(u.srv.Timeout(), cancel) + s, err := client.Connect(ctx, t, nil) + if !timer.Stop() && err == nil { + err = context.DeadlineExceeded + } + if err != nil { + cancel() + if s != nil { + _ = s.Close() + } + return nil, err + } + u.session, u.cancel = s, cancel + return s, nil +} + +// drop forgets the session s, if it is still the current one. +func (u *upstream) drop(s *mcp.ClientSession) { + u.mu.Lock() + if u.session != s { + u.mu.Unlock() + return + } + cancel := u.cancel + u.session, u.cancel, u.tools, u.toolsAt = nil, nil, nil, time.Time{} + u.mu.Unlock() + _ = s.Close() // the session is unusable either way + if cancel != nil { + cancel() + } +} + +// retire closes the upstream for good, when the pool replaces or drops it. +func (u *upstream) retire() { + u.mu.Lock() + u.retired = true + u.mu.Unlock() + u.close() +} + +func (u *upstream) close() { + u.mu.Lock() + s, cancel := u.session, u.cancel + u.session, u.cancel, u.tools, u.toolsAt = nil, nil, nil, time.Time{} + u.mu.Unlock() + if s != nil { + _ = s.Close() // the session is unusable either way + } + if cancel != nil { + cancel() + } +} + +// listTools returns the server's tools, from the cache while it is fresh. +func (u *upstream) listTools(ctx context.Context, version string) ([]*mcp.Tool, error) { + u.mu.Lock() + if !u.toolsAt.IsZero() && time.Since(u.toolsAt) < toolsTTL { + t := u.tools + u.mu.Unlock() + return t, nil + } + u.mu.Unlock() + + // A list is safe to ask twice: one retry covers a session the server dropped. + var err error + for attempt := 0; attempt < 2; attempt++ { + var s *mcp.ClientSession + if s, err = u.connect(version); err != nil { + return nil, err + } + tools := []*mcp.Tool{} + lctx, cancel := context.WithTimeout(ctx, u.srv.Timeout()) + for t, lerr := range s.Tools(lctx, nil) { + if lerr != nil { + err = lerr + break + } + tools = append(tools, t) + } + cancel() + if err == nil { + u.mu.Lock() + if u.session == s { + u.tools, u.toolsAt = tools, time.Now() + } + u.mu.Unlock() + return tools, nil + } + if !transportError(err) { + return nil, err + } + u.drop(s) + } + return nil, err +} + +// call forwards one tools/call. Only a call the server refused for an unknown +// session is retried: a tool may not be safe to run twice, and a closed +// connection does not say whether the call ran. +func (u *upstream) call(ctx context.Context, version string, p *mcp.CallToolParams) (*mcp.CallToolResult, error) { + for attempt := 0; ; attempt++ { + s, err := u.connect(version) + if err != nil { + return nil, err + } + res, err := s.CallTool(ctx, p) + if err == nil { + return res, nil + } + if transportError(err) { + u.drop(s) + if attempt == 0 && errors.Is(err, mcp.ErrSessionMissing) { + continue + } + } + return nil, err + } +} + +// transportError reports an error that is not the server's JSON-RPC answer: +// the session is then suspect and is made again. +func transportError(err error) bool { + var rpc *jsonrpc.Error + return !errors.As(err, &rpc) && !errors.Is(err, context.Canceled) && !errors.Is(err, context.DeadlineExceeded) +} + +// transport dials only public addresses unless the server allows private ones, +// and ignores proxy variables, which would dial past that check. +func transport(srv *registry.MCPServer) http.RoundTripper { + d := &net.Dialer{Timeout: 10 * time.Second, KeepAlive: 30 * time.Second} + dial := d.DialContext + if !srv.AllowPrivate { + dial = safefetch.GuardedDial(d) + } + return &http.Transport{ + DialContext: dial, + TLSHandshakeTimeout: 10 * time.Second, + MaxIdleConnsPerHost: 16, + IdleConnTimeout: 90 * time.Second, + ForceAttemptHTTP2: true, + } +} + +// credentialTransport adds the server's credential to every request and caps +// what a reply body may hold. The standalone stream (GET) is never opened. +type credentialTransport struct { + srv *registry.MCPServer + base http.RoundTripper +} + +func (t *credentialTransport) RoundTrip(r *http.Request) (*http.Response, error) { + if t.srv.Header != "" { + r = r.Clone(r.Context()) + r.Header.Set(t.srv.Header, t.srv.Value) + } + resp, err := t.base.RoundTrip(r) + if err != nil { + return nil, err + } + // An SSE transport's event stream is long-lived; its events are capped by MaxEventSize. + if r.Method == http.MethodPost { + cut, _ := r.Context().Value(cutKey{}).(*atomic.Bool) + resp.Body = &limitedBody{rc: resp.Body, left: t.srv.MaxResponse() + envelope, cut: cut} + } + return resp, nil +} + +type limitedBody struct { + rc io.ReadCloser + left int64 + cut *atomic.Bool +} + +func (b *limitedBody) Read(p []byte) (int, error) { + if b.left <= 0 { + if b.cut != nil { + b.cut.Store(true) + } + return 0, errTooLarge + } + if int64(len(p)) > b.left { + p = p[:b.left] + } + n, err := b.rc.Read(p) + b.left -= int64(n) + return n, err +} + +func (b *limitedBody) Close() error { return b.rc.Close() } + +// pool keeps one upstream per server and replaces it when the server's spec +// or credential changes. +type pool struct { + mu sync.Mutex + m map[string]*upstream +} + +func (p *pool) get(srv *registry.MCPServer) *upstream { + fp := srv.Fingerprint() + p.mu.Lock() + defer p.mu.Unlock() + if p.m == nil { + p.m = map[string]*upstream{} + } + u := p.m[srv.Name] + if u != nil && u.fp == fp { + return u + } + if u != nil { + go u.retire() + } + u = &upstream{srv: srv, fp: fp} + p.m[srv.Name] = u + return u +} + +// prune closes the sessions of servers no longer registered or enabled. +func (p *pool) prune(servers map[string]*registry.MCPServer) { + p.mu.Lock() + defer p.mu.Unlock() + for name, u := range p.m { + if servers[name] == nil { + delete(p.m, name) + go u.retire() + } + } +} + +func (p *pool) closeAll() { + p.mu.Lock() + defer p.mu.Unlock() + for name, u := range p.m { + delete(p.m, name) + u.retire() + } +} diff --git a/internal/observability/metrics.go b/internal/observability/metrics.go index 1c5fb57..e68753d 100644 --- a/internal/observability/metrics.go +++ b/internal/observability/metrics.go @@ -38,6 +38,39 @@ var ( []string{"vendor", "model"}, ) + // MCPToolCallsTotal counts tool calls through the MCP gateway by outcome: + // success | tool_error | denied | throttled | timeout | too_large | upstream_error. + MCPToolCallsTotal = promauto.NewCounterVec( + prometheus.CounterOpts{ + Namespace: "ktms", + Name: "mcp_tool_calls_total", + Help: "Tool calls through the MCP gateway by server, tool and status.", + }, + []string{"server", "tool", "status"}, + ) + + // MCPToolCallDuration measures a tool call from the gateway to the server and back. + MCPToolCallDuration = promauto.NewHistogramVec( + prometheus.HistogramOpts{ + Namespace: "ktms", + Name: "mcp_tool_call_duration_seconds", + Help: "Latency of tool calls through the MCP gateway in seconds.", + Buckets: []float64{0.05, 0.1, 0.25, 0.5, 1, 2.5, 5, 10, 30, 60, 120}, + }, + []string{"server"}, + ) + + // MCPResponseBytes is the size of tool results, as JSON. + MCPResponseBytes = promauto.NewHistogramVec( + prometheus.HistogramOpts{ + Namespace: "ktms", + Name: "mcp_response_bytes", + Help: "Size of tool results through the MCP gateway in bytes.", + Buckets: prometheus.ExponentialBuckets(256, 4, 9), + }, + []string{"server"}, + ) + // TokensTotal tracks billed input and output tokens. TokensTotal = promauto.NewCounterVec( prometheus.CounterOpts{ diff --git a/internal/proxydb/apikey_policy_test.go b/internal/proxydb/apikey_policy_test.go index 7ab7378..f34579a 100644 --- a/internal/proxydb/apikey_policy_test.go +++ b/internal/proxydb/apikey_policy_test.go @@ -25,7 +25,7 @@ func TestAPIKeyPolicyRoundTrip(t *testing.T) { if !slices.Equal(k.AllowedModels, []string{"m1"}) || k.RPMLimit != 60 || k.TPMLimit != 100000 { t.Fatalf("minted policy = %v %d %d", k.AllowedModels, k.RPMLimit, k.TPMLimit) } - if err := d.SetAPIKeyPolicy(ctx, name, nil, 0, 5); err != nil { + if err := d.SetAPIKeyPolicy(ctx, name, nil, 0, 5, nil); err != nil { t.Fatal(err) } k, err = d.GetAPIKeyByName(ctx, name) @@ -35,7 +35,7 @@ func TestAPIKeyPolicyRoundTrip(t *testing.T) { if len(k.AllowedModels) != 0 || k.RPMLimit != 0 || k.TPMLimit != 5 { t.Fatalf("replaced policy = %v %d %d", k.AllowedModels, k.RPMLimit, k.TPMLimit) } - if err := d.SetAPIKeyPolicy(ctx, "missing-"+name, nil, 0, 0); !errors.Is(err, ErrNotFound) { + if err := d.SetAPIKeyPolicy(ctx, "missing-"+name, nil, 0, 0, nil); !errors.Is(err, ErrNotFound) { t.Fatalf("unknown key: %v, want ErrNotFound", err) } } diff --git a/internal/proxydb/db.go b/internal/proxydb/db.go index 1d5fb14..64eb08b 100644 --- a/internal/proxydb/db.go +++ b/internal/proxydb/db.go @@ -30,7 +30,8 @@ func (d *DB) GetAPIKeyByLookup(ctx context.Context, lookup string) (*APIKey, err return scanAPIKey(d.pool.QueryRow(ctx, `SELECT k.id, k.name, k.key_lookup, k.key_hash, k.scopes, k.disabled_at, k.created_at, k.org_id, k.team_id, COALESCE(t.allowed_models, '{}'), - k.allowed_models, k.rpm_limit, k.tpm_limit + k.allowed_models, k.rpm_limit, k.tpm_limit, + COALESCE(t.allowed_mcp, '{}'), k.allowed_mcp FROM api_key k LEFT JOIN team t ON t.id = k.team_id WHERE k.key_lookup = $1`, lookup)) @@ -40,7 +41,8 @@ func (d *DB) GetAPIKeyByName(ctx context.Context, name string) (*APIKey, error) return scanAPIKey(d.pool.QueryRow(ctx, `SELECT k.id, k.name, k.key_lookup, k.key_hash, k.scopes, k.disabled_at, k.created_at, k.org_id, k.team_id, COALESCE(t.allowed_models, '{}'), - k.allowed_models, k.rpm_limit, k.tpm_limit + k.allowed_models, k.rpm_limit, k.tpm_limit, + COALESCE(t.allowed_mcp, '{}'), k.allowed_mcp FROM api_key k LEFT JOIN team t ON t.id = k.team_id WHERE k.name = $1`, name)) @@ -49,7 +51,8 @@ func (d *DB) GetAPIKeyByName(ctx context.Context, name string) (*APIKey, error) func scanAPIKey(row pgx.Row) (*APIKey, error) { var k APIKey err := row.Scan(&k.ID, &k.Name, &k.KeyLookup, &k.KeyHash, &k.Scopes, &k.DisabledAt, &k.CreatedAt, - &k.OrgID, &k.TeamID, &k.TeamAllowedModels, &k.AllowedModels, &k.RPMLimit, &k.TPMLimit) + &k.OrgID, &k.TeamID, &k.TeamAllowedModels, &k.AllowedModels, &k.RPMLimit, &k.TPMLimit, + &k.TeamAllowedMCP, &k.AllowedMCP) if errors.Is(err, pgx.ErrNoRows) { return nil, ErrNotFound } @@ -61,17 +64,17 @@ func scanAPIKey(row pgx.Row) (*APIKey, error) { func (d *DB) InsertAPIKey(ctx context.Context, k *APIKey) error { _, err := d.pool.Exec(ctx, - `INSERT INTO api_key (id, name, key_lookup, key_hash, scopes, org_id, team_id, allowed_models, rpm_limit, tpm_limit) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)`, - k.ID, k.Name, k.KeyLookup, k.KeyHash, k.Scopes, k.OrgID, k.TeamID, nonNil(k.AllowedModels), k.RPMLimit, k.TPMLimit) + `INSERT INTO api_key (id, name, key_lookup, key_hash, scopes, org_id, team_id, allowed_models, rpm_limit, tpm_limit, allowed_mcp) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)`, + k.ID, k.Name, k.KeyLookup, k.KeyHash, k.Scopes, k.OrgID, k.TeamID, nonNil(k.AllowedModels), k.RPMLimit, k.TPMLimit, nonNil(k.AllowedMCP)) return err } // SetAPIKeyPolicy replaces the policy of the key with this name. -func (d *DB) SetAPIKeyPolicy(ctx context.Context, name string, allowedModels []string, rpm int32, tpm int64) error { +func (d *DB) SetAPIKeyPolicy(ctx context.Context, name string, allowedModels []string, rpm int32, tpm int64, allowedMCP []string) error { tag, err := d.pool.Exec(ctx, - `UPDATE api_key SET allowed_models = $2, rpm_limit = $3, tpm_limit = $4 WHERE name = $1`, - name, nonNil(allowedModels), rpm, tpm) + `UPDATE api_key SET allowed_models = $2, rpm_limit = $3, tpm_limit = $4, allowed_mcp = $5 WHERE name = $1`, + name, nonNil(allowedModels), rpm, tpm, nonNil(allowedMCP)) if err != nil { return err } @@ -105,7 +108,8 @@ func (d *DB) ListAPIKeys(ctx context.Context) ([]*APIKey, error) { rows, err := d.pool.Query(ctx, `SELECT k.id, k.name, k.key_lookup, k.key_hash, k.scopes, k.disabled_at, k.created_at, k.org_id, k.team_id, COALESCE(t.allowed_models, '{}'), - k.allowed_models, k.rpm_limit, k.tpm_limit + k.allowed_models, k.rpm_limit, k.tpm_limit, + COALESCE(t.allowed_mcp, '{}'), k.allowed_mcp FROM api_key k LEFT JOIN team t ON t.id = k.team_id ORDER BY k.created_at`) @@ -559,18 +563,18 @@ func SetTeamBudgetTx(ctx context.Context, tx pgx.Tx, teamID string, limitNanoUSD func CreateTeamTx(ctx context.Context, tx pgx.Tx, team *Team) error { _, err := tx.Exec(ctx, - `INSERT INTO team (id, org_id, name, monthly_budget_usd, allowed_models, created_at) - VALUES ($1, $2, $3, $4, $5, $6)`, - team.ID, team.OrgID, team.Name, team.MonthlyBudgetUSD, team.AllowedModels, team.CreatedAt) + `INSERT INTO team (id, org_id, name, monthly_budget_usd, allowed_models, created_at, allowed_mcp) + VALUES ($1, $2, $3, $4, $5, $6, $7)`, + team.ID, team.OrgID, team.Name, team.MonthlyBudgetUSD, team.AllowedModels, team.CreatedAt, nonNil(team.AllowedMCP)) return err } func (d *DB) GetTeam(ctx context.Context, id string) (*Team, error) { var t Team err := d.pool.QueryRow(ctx, - `SELECT id, org_id, name, monthly_budget_usd, allowed_models, created_at + `SELECT id, org_id, name, monthly_budget_usd, allowed_models, created_at, allowed_mcp FROM team WHERE id = $1`, id). - Scan(&t.ID, &t.OrgID, &t.Name, &t.MonthlyBudgetUSD, &t.AllowedModels, &t.CreatedAt) + Scan(&t.ID, &t.OrgID, &t.Name, &t.MonthlyBudgetUSD, &t.AllowedModels, &t.CreatedAt, &t.AllowedMCP) if errors.Is(err, pgx.ErrNoRows) { return nil, ErrNotFound } @@ -587,11 +591,11 @@ func (d *DB) ListTeams(ctx context.Context, orgID string) ([]*Team, error) { ) if orgID != "" { rows, err = d.pool.Query(ctx, - `SELECT id, org_id, name, monthly_budget_usd, allowed_models, created_at + `SELECT id, org_id, name, monthly_budget_usd, allowed_models, created_at, allowed_mcp FROM team WHERE org_id = $1 ORDER BY created_at DESC`, orgID) } else { rows, err = d.pool.Query(ctx, - `SELECT id, org_id, name, monthly_budget_usd, allowed_models, created_at + `SELECT id, org_id, name, monthly_budget_usd, allowed_models, created_at, allowed_mcp FROM team ORDER BY created_at DESC`) } if err != nil { @@ -601,7 +605,7 @@ func (d *DB) ListTeams(ctx context.Context, orgID string) ([]*Team, error) { var out []*Team for rows.Next() { var t Team - if err := rows.Scan(&t.ID, &t.OrgID, &t.Name, &t.MonthlyBudgetUSD, &t.AllowedModels, &t.CreatedAt); err != nil { + if err := rows.Scan(&t.ID, &t.OrgID, &t.Name, &t.MonthlyBudgetUSD, &t.AllowedModels, &t.CreatedAt, &t.AllowedMCP); err != nil { return nil, err } out = append(out, &t) @@ -611,9 +615,9 @@ func (d *DB) ListTeams(ctx context.Context, orgID string) ([]*Team, error) { func UpdateTeamTx(ctx context.Context, tx pgx.Tx, team *Team) error { tag, err := tx.Exec(ctx, - `UPDATE team SET name = $2, monthly_budget_usd = $3, allowed_models = $4 + `UPDATE team SET name = $2, monthly_budget_usd = $3, allowed_models = $4, allowed_mcp = $5 WHERE id = $1`, - team.ID, team.Name, team.MonthlyBudgetUSD, team.AllowedModels) + team.ID, team.Name, team.MonthlyBudgetUSD, team.AllowedModels, nonNil(team.AllowedMCP)) if err != nil { return err } diff --git a/internal/proxydb/mcp.go b/internal/proxydb/mcp.go new file mode 100644 index 0000000..32c8c4e --- /dev/null +++ b/internal/proxydb/mcp.go @@ -0,0 +1,139 @@ +package proxydb + +import ( + "context" + "time" + + "github.com/jackc/pgx/v5" +) + +// MCPServer is an upstream MCP server the gateway proxies. +type MCPServer struct { + Name string + URL string + Transport string // "streamable_http" | "sse" + Enabled bool + AllowPrivate bool + TimeoutMs int32 + MaxResponseBytes int64 + ToolRPM int32 +} + +// MCPCredential is the header the gateway adds to a server's requests. +type MCPCredential struct { + Server string + Header string + ValueCiphertext []byte + ValueVersion int32 + CreatedAt time.Time +} + +func (d *DB) LoadMCPServers(ctx context.Context) ([]*MCPServer, error) { + rows, err := d.pool.Query(ctx, + `SELECT name, url, transport, enabled, allow_private, timeout_ms, max_response_bytes, tool_rpm + FROM mcp_server ORDER BY name`) + if err != nil { + return nil, err + } + defer rows.Close() + var out []*MCPServer + for rows.Next() { + var m MCPServer + if err := rows.Scan(&m.Name, &m.URL, &m.Transport, &m.Enabled, &m.AllowPrivate, + &m.TimeoutMs, &m.MaxResponseBytes, &m.ToolRPM); err != nil { + return nil, err + } + out = append(out, &m) + } + return out, rows.Err() +} + +func (d *DB) LoadMCPCredentials(ctx context.Context) ([]*MCPCredential, error) { + rows, err := d.pool.Query(ctx, + `SELECT server, header, value_ciphertext, value_version, created_at FROM mcp_credential ORDER BY server`) + if err != nil { + return nil, err + } + defer rows.Close() + var out []*MCPCredential + for rows.Next() { + var c MCPCredential + if err := rows.Scan(&c.Server, &c.Header, &c.ValueCiphertext, &c.ValueVersion, &c.CreatedAt); err != nil { + return nil, err + } + out = append(out, &c) + } + return out, rows.Err() +} + +func UpsertMCPServer(ctx context.Context, tx pgx.Tx, m *MCPServer) error { + _, err := tx.Exec(ctx, + `INSERT INTO mcp_server (name, url, transport, enabled, allow_private, timeout_ms, max_response_bytes, tool_rpm) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8) + ON CONFLICT (name) DO UPDATE SET url = $2, transport = $3, enabled = $4, allow_private = $5, + timeout_ms = $6, max_response_bytes = $7, tool_rpm = $8, updated_at = NOW()`, + m.Name, m.URL, m.Transport, m.Enabled, m.AllowPrivate, m.TimeoutMs, m.MaxResponseBytes, m.ToolRPM) + return err +} + +// DeleteMCPServer removes the server; its credential goes with it (cascade). +func DeleteMCPServer(ctx context.Context, tx pgx.Tx, name string) error { + tag, err := tx.Exec(ctx, `DELETE FROM mcp_server WHERE name = $1`, name) + if err != nil { + return err + } + if tag.RowsAffected() == 0 { + return ErrNotFound + } + return nil +} + +// SetMCPCredential stores a server's credential; an unknown server is a +// foreign key violation (23503). +func SetMCPCredential(ctx context.Context, tx pgx.Tx, c *MCPCredential) error { + _, err := tx.Exec(ctx, + `INSERT INTO mcp_credential (server, header, value_ciphertext, value_version) + VALUES ($1, $2, $3, $4) + ON CONFLICT (server) DO UPDATE SET header = $2, value_ciphertext = $3, value_version = $4, created_at = NOW()`, + c.Server, c.Header, c.ValueCiphertext, c.ValueVersion) + return err +} + +func DeleteMCPCredential(ctx context.Context, tx pgx.Tx, server string) error { + tag, err := tx.Exec(ctx, `DELETE FROM mcp_credential WHERE server = $1`, server) + if err != nil { + return err + } + if tag.RowsAffected() == 0 { + return ErrNotFound + } + return nil +} + +// LockStaleMCPCredentialsTx is LockStaleVendorKeysTx for MCP credentials. +func LockStaleMCPCredentialsTx(ctx context.Context, tx pgx.Tx, version int32) ([]*MCPCredential, error) { + rows, err := tx.Query(ctx, + `SELECT server, header, value_ciphertext, value_version FROM mcp_credential + WHERE value_version <> $1 ORDER BY server FOR UPDATE`, version) + if err != nil { + return nil, err + } + defer rows.Close() + var out []*MCPCredential + for rows.Next() { + var c MCPCredential + if err := rows.Scan(&c.Server, &c.Header, &c.ValueCiphertext, &c.ValueVersion); err != nil { + return nil, err + } + out = append(out, &c) + } + return out, rows.Err() +} + +// UpdateMCPCredentialCiphertext is the rotation path: same value, newer keyring version. +func UpdateMCPCredentialCiphertext(ctx context.Context, tx pgx.Tx, server string, ciphertext []byte, version int32) error { + _, err := tx.Exec(ctx, + `UPDATE mcp_credential SET value_ciphertext = $2, value_version = $3 WHERE server = $1`, + server, ciphertext, version) + return err +} diff --git a/internal/proxydb/models.go b/internal/proxydb/models.go index f95442b..011d390 100644 --- a/internal/proxydb/models.go +++ b/internal/proxydb/models.go @@ -25,6 +25,7 @@ type Team struct { Name string MonthlyBudgetUSD float64 AllowedModels []string + AllowedMCP []string CreatedAt time.Time } @@ -43,6 +44,9 @@ type APIKey struct { AllowedModels []string RPMLimit int32 TPMLimit int64 + // MCP servers and tools the key, and its team, may use; empty = any. + AllowedMCP []string + TeamAllowedMCP []string } func (k *APIKey) Active() bool { return k.DisabledAt == nil } diff --git a/internal/rediskeys/rediskeys.go b/internal/rediskeys/rediskeys.go index b790f70..8d06d2f 100644 --- a/internal/rediskeys/rediskeys.go +++ b/internal/rediskeys/rediskeys.go @@ -40,6 +40,11 @@ func (k Keys) KeyTPM(keyName string) string { return fmt.Sprintf("%sllmproxy:gcra:ktpm:%s", k.prefix, keyName) } +// ToolRPM is the GCRA arrival time of a client key's calls to one MCP tool. +func (k Keys) ToolRPM(keyName, server, tool string) string { + return fmt.Sprintf("%sllmproxy:gcra:mcp:%s:%s:%s", k.prefix, keyName, server, tool) +} + // BudgetKey is the daily token budget per client key and model (UTC day). func (k Keys) BudgetKey(keyName, modelID, day string) string { return fmt.Sprintf("%sllmproxy:bud:k:%s:%s:%s", k.prefix, keyName, modelID, day) diff --git a/internal/registry/loader.go b/internal/registry/loader.go index 7efe7da..feed2c1 100644 --- a/internal/registry/loader.go +++ b/internal/registry/loader.go @@ -145,6 +145,14 @@ func (l *Loader) build(ctx context.Context, version int64) (*Snapshot, error) { if err != nil { return nil, fmt.Errorf("load spend limits: %w", err) } + mcpServers, err := l.db.LoadMCPServers(ctx) + if err != nil { + return nil, fmt.Errorf("load mcp servers: %w", err) + } + mcpCreds, err := l.db.LoadMCPCredentials(ctx) + if err != nil { + return nil, fmt.Errorf("load mcp credentials: %w", err) + } snap := &Snapshot{ Version: version, @@ -205,9 +213,44 @@ func (l *Loader) build(ctx context.Context, version int64) (*Snapshot, error) { } snap.vendorKeys[vk.APIKeyID][vk.Vendor] = vc } + snap.MCPServers = l.mcpServers(mcpServers, mcpCreds, prev) // Not one readable key means a wrong keyring, not a rotation in progress. if len(vendorKeys) > 0 && unreadable == len(vendorKeys) { return nil, fmt.Errorf("no vendor key can be decrypted with KTMS_ENCRYPTION_KEYS (%d keys): wrong keyring?", unreadable) } return snap, nil } + +// mcpServers joins the enabled servers with their credentials. One that cannot +// be decrypted keeps the previous snapshot's, as a vendor key does; without +// one the server is left out rather than called unauthenticated. +func (l *Loader) mcpServers(servers []*proxydb.MCPServer, creds []*proxydb.MCPCredential, prev *Snapshot) map[string]*MCPServer { + byServer := make(map[string]*proxydb.MCPCredential, len(creds)) + for _, c := range creds { + byServer[c.Server] = c + } + out := make(map[string]*MCPServer, len(servers)) + for _, s := range servers { + if !s.Enabled { + continue + } + m := &MCPServer{MCPServer: *s} + if c := byServer[s.Name]; c != nil { + plain, err := l.keyring.Decrypt(c.ValueCiphertext, int(c.ValueVersion)) + if err != nil { + observability.DependencyErrorsTotal.WithLabelValues("keyring", "decrypt").Inc() + old := prev.MCPServers[s.Name] + log.Error().Err(err).Str("mcp_server", s.Name).Int32("key_version", c.ValueVersion).Bool("kept_previous", old != nil && old.Header != ""). + Msg("mcp credential cannot be decrypted with KTMS_ENCRYPTION_KEYS") + if old == nil || old.Header == "" { + continue + } + m.Header, m.Value = old.Header, old.Value + } else { + m.Header, m.Value = c.Header, plain + } + } + out[s.Name] = m + } + return out +} diff --git a/internal/registry/mcp.go b/internal/registry/mcp.go new file mode 100644 index 0000000..d4dd625 --- /dev/null +++ b/internal/registry/mcp.go @@ -0,0 +1,93 @@ +package registry + +import ( + "errors" + "fmt" + "net/url" + "regexp" + "time" + + "github.com/ai-process/ktms/internal/crypto" + "github.com/ai-process/ktms/internal/proxydb" +) + +// MCP transports. +const ( + MCPStreamableHTTP = "streamable_http" + MCPSSE = "sse" +) + +// MCP server defaults for the 0 values of the spec. +const ( + DefaultMCPTimeout = 60 * time.Second + DefaultMCPMaxResponseBytes = 4 << 20 +) + +// mcpName keeps server names out of the "__" and "/" that separate them from +// tool names, and short enough for usage events and client tool name limits. +var mcpName = regexp.MustCompile(`^[a-z][a-z0-9-]{0,31}$`) + +// MCPServer is an enabled MCP server with its decrypted credential. +type MCPServer struct { + proxydb.MCPServer + Header string // "" when the server has no credential + Value string +} + +// Timeout bounds one tool call. +func (m *MCPServer) Timeout() time.Duration { + if m.TimeoutMs > 0 { + return time.Duration(m.TimeoutMs) * time.Millisecond + } + return DefaultMCPTimeout +} + +// MaxResponse bounds one tool result. +func (m *MCPServer) MaxResponse() int64 { + if m.MaxResponseBytes > 0 { + return m.MaxResponseBytes + } + return DefaultMCPMaxResponseBytes +} + +// Fingerprint changes whenever anything a connection depends on does, so the +// gateway knows to drop a connection made with an older spec or credential. +func (m *MCPServer) Fingerprint() string { + return crypto.Fingerprint(fmt.Sprintf("%s\x00%s\x00%t\x00%d\x00%d\x00%s\x00%s", + m.URL, m.Transport, m.AllowPrivate, m.TimeoutMs, m.MaxResponseBytes, m.Header, m.Value)) +} + +// ValidMCPName reports whether name may name an MCP server. +func ValidMCPName(name string) bool { return mcpName.MatchString(name) } + +// ValidateMCPServer checks a spec on its own; whether its host resolves to a +// private address is checked where the network is (the Admin API). +func ValidateMCPServer(m *proxydb.MCPServer) error { + if !mcpName.MatchString(m.Name) { + return fmt.Errorf("mcp server name %q: 1-32 of a-z, 0-9 and -, starting with a letter", m.Name) + } + switch m.Transport { + case MCPStreamableHTTP, MCPSSE: + default: + return fmt.Errorf("mcp server %s: transport must be %s or %s, got %q", m.Name, MCPStreamableHTTP, MCPSSE, m.Transport) + } + u, err := url.Parse(m.URL) + if err != nil || u.Host == "" { + return fmt.Errorf("mcp server %s: url %q is not an absolute URL", m.Name, m.URL) + } + switch { + case u.Scheme == "https": + case u.Scheme == "http" && m.AllowPrivate: + case u.Scheme == "http": + return fmt.Errorf("mcp server %s: http is allowed only with allow_private (a server on a private network); use https", m.Name) + default: + return fmt.Errorf("mcp server %s: url scheme must be https", m.Name) + } + if u.User != nil { + return errors.New("mcp server " + m.Name + ": keep credentials out of the url, set them with SetMCPCredential") + } + if m.TimeoutMs < 0 || m.MaxResponseBytes < 0 || m.ToolRPM < 0 { + return fmt.Errorf("mcp server %s: timeout_ms, max_response_bytes and tool_rpm must be >= 0", m.Name) + } + return nil +} diff --git a/internal/registry/registry.go b/internal/registry/registry.go index bf699b1..758fbdd 100644 --- a/internal/registry/registry.go +++ b/internal/registry/registry.go @@ -24,6 +24,8 @@ type Snapshot struct { UserRateLimits map[string]map[string]*proxydb.UserRateLimit // SpendLimits[subjectType+"\x00"+subjectID][period] = cap in nano-USD. SpendLimits map[string]map[string]int64 + // MCPServers[name] = enabled MCP server with its credential. + MCPServers map[string]*MCPServer // adapters caches built vendor adapters per (apiKeyID, modelID) for this // snapshot's lifetime; dropped wholesale on config change. adapters sync.Map diff --git a/internal/safefetch/fetch.go b/internal/safefetch/fetch.go index 2b20c6c..885e003 100644 --- a/internal/safefetch/fetch.go +++ b/internal/safefetch/fetch.go @@ -198,3 +198,61 @@ func FetchURL(ctx context.Context, rawURL string, maxBytes int64) ([]byte, strin return data, mimeType, nil } + +// CheckHost fails with ErrProhibitedIP when host is, or resolves to, an +// address IsProhibitedIP refuses. +func CheckHost(ctx context.Context, host string) error { + ips, err := lookup(ctx, host) + if err != nil { + return err + } + for _, ip := range ips { + if IsProhibitedIP(ip) { + return fmt.Errorf("%w: %s resolves to %s", ErrProhibitedIP, host, ip) + } + } + return nil +} + +// GuardedDial dials like d but refuses a prohibited address. It connects to +// the address it checked, so a DNS answer cannot change in between. +func GuardedDial(d *net.Dialer) func(ctx context.Context, network, addr string) (net.Conn, error) { + return func(ctx context.Context, network, addr string) (net.Conn, error) { + host, port, err := net.SplitHostPort(addr) + if err != nil { + return nil, err + } + ips, err := lookup(ctx, host) + if err != nil { + return nil, err + } + for _, ip := range ips { + if IsProhibitedIP(ip) { + return nil, fmt.Errorf("%w: %s resolves to %s", ErrProhibitedIP, host, ip) + } + } + var lastErr error + for _, ip := range ips { + conn, err := d.DialContext(ctx, network, net.JoinHostPort(ip.String(), port)) + if err == nil { + return conn, nil + } + lastErr = err + } + return nil, lastErr + } +} + +func lookup(ctx context.Context, host string) ([]net.IP, error) { + if ip := net.ParseIP(host); ip != nil { + return []net.IP{ip}, nil + } + ips, err := net.DefaultResolver.LookupIP(ctx, "ip", host) + if err != nil { + return nil, fmt.Errorf("DNS resolution failed for %s: %w", host, err) + } + if len(ips) == 0 { + return nil, fmt.Errorf("DNS resolution failed for %s: no addresses", host) + } + return ips, nil +} diff --git a/internal/safefetch/fetch_test.go b/internal/safefetch/fetch_test.go index 5ce8798..1364fb4 100644 --- a/internal/safefetch/fetch_test.go +++ b/internal/safefetch/fetch_test.go @@ -90,3 +90,18 @@ func TestValidateMediaMimeType(t *testing.T) { } } } + +func TestCheckHostAndGuardedDial(t *testing.T) { + for _, host := range []string{"127.0.0.1", "10.1.2.3", "169.254.169.254", "::1"} { + if err := CheckHost(context.Background(), host); !errors.Is(err, ErrProhibitedIP) { + t.Errorf("CheckHost(%s) = %v, want ErrProhibitedIP", host, err) + } + } + if err := CheckHost(context.Background(), "93.184.215.14"); err != nil { + t.Errorf("public address refused: %v", err) + } + dial := GuardedDial(&net.Dialer{}) + if _, err := dial(context.Background(), "tcp", "127.0.0.1:80"); !errors.Is(err, ErrProhibitedIP) { + t.Errorf("GuardedDial to loopback = %v, want ErrProhibitedIP", err) + } +} diff --git a/internal/server/auth_interceptor.go b/internal/server/auth_interceptor.go index dfee362..fdd82a4 100644 --- a/internal/server/auth_interceptor.go +++ b/internal/server/auth_interceptor.go @@ -57,17 +57,22 @@ var methodScopes = map[string]string{ method(pb.AdminService_DeleteSpendLimit_FullMethodName): apikeys.ScopeAdmin, method(pb.AdminService_ListSpendLimits_FullMethodName): apikeys.ScopeAdminRead, - method(pb.AdminService_CreateOrganization_FullMethodName): apikeys.ScopeAdmin, - method(pb.AdminService_ListOrganizations_FullMethodName): apikeys.ScopeAdminRead, - method(pb.AdminService_DeleteOrganization_FullMethodName): apikeys.ScopeAdmin, - method(pb.AdminService_CreateTeam_FullMethodName): apikeys.ScopeAdmin, - method(pb.AdminService_GetTeam_FullMethodName): apikeys.ScopeAdminRead, - method(pb.AdminService_ListTeams_FullMethodName): apikeys.ScopeAdminRead, - method(pb.AdminService_UpdateTeam_FullMethodName): apikeys.ScopeAdmin, - method(pb.AdminService_DeleteTeam_FullMethodName): apikeys.ScopeAdmin, - method(pb.AdminService_GetTeamSpend_FullMethodName): apikeys.ScopeAdminRead, - method(pb.AdminService_GetSpend_FullMethodName): apikeys.ScopeAdminRead, - method(pb.AdminService_ExplainRoute_FullMethodName): apikeys.ScopeAdminRead, + method(pb.AdminService_CreateOrganization_FullMethodName): apikeys.ScopeAdmin, + method(pb.AdminService_ListOrganizations_FullMethodName): apikeys.ScopeAdminRead, + method(pb.AdminService_DeleteOrganization_FullMethodName): apikeys.ScopeAdmin, + method(pb.AdminService_CreateTeam_FullMethodName): apikeys.ScopeAdmin, + method(pb.AdminService_GetTeam_FullMethodName): apikeys.ScopeAdminRead, + method(pb.AdminService_ListTeams_FullMethodName): apikeys.ScopeAdminRead, + method(pb.AdminService_UpdateTeam_FullMethodName): apikeys.ScopeAdmin, + method(pb.AdminService_DeleteTeam_FullMethodName): apikeys.ScopeAdmin, + method(pb.AdminService_GetTeamSpend_FullMethodName): apikeys.ScopeAdminRead, + method(pb.AdminService_GetSpend_FullMethodName): apikeys.ScopeAdminRead, + method(pb.AdminService_ExplainRoute_FullMethodName): apikeys.ScopeAdminRead, + method(pb.AdminService_UpsertMCPServer_FullMethodName): apikeys.ScopeAdmin, + method(pb.AdminService_DeleteMCPServer_FullMethodName): apikeys.ScopeAdmin, + method(pb.AdminService_SetMCPCredential_FullMethodName): apikeys.ScopeAdmin, + method(pb.AdminService_DeleteMCPCredential_FullMethodName): apikeys.ScopeAdmin, + method(pb.AdminService_ListMCPTools_FullMethodName): apikeys.ScopeAdminRead, method(pb.NodeRegistryService_RegisterNode_FullMethodName): apikeys.ScopeNode, method(pb.NodeRegistryService_ValidateNode_FullMethodName): apikeys.ScopeNode, diff --git a/internal/server/auth_interceptor_test.go b/internal/server/auth_interceptor_test.go index 6af7867..1019059 100644 --- a/internal/server/auth_interceptor_test.go +++ b/internal/server/auth_interceptor_test.go @@ -188,7 +188,7 @@ func TestStreamAuthInterceptor(t *testing.T) { // adminReads are the AdminService methods a read key may call. Listed by hand // so a new RPC is an admin-only mutation until someone decides otherwise. var adminReads = []string{"GetConfig", "ListAPIKeys", "ListNodeKeys", "ListNodes", "ListUserRateLimits", - "ListSpendLimits", "ListOrganizations", "GetTeam", "ListTeams", "GetTeamSpend", "GetSpend", "ExplainRoute"} + "ListSpendLimits", "ListOrganizations", "GetTeam", "ListTeams", "GetTeamSpend", "GetSpend", "ExplainRoute", "ListMCPTools"} // A read key reaches every Admin API read and none of the methods that change state. func TestAdminReadKeyOnlyReads(t *testing.T) { diff --git a/internal/server/http.go b/internal/server/http.go index 72b8ec2..10f55d3 100644 --- a/internal/server/http.go +++ b/internal/server/http.go @@ -18,6 +18,11 @@ type Pinger interface { Ping(ctx context.Context) error } +// Registrar mounts more handlers on the HTTP server, such as the MCP gateway. +type Registrar interface { + Register(mux *http.ServeMux) +} + // HTTPServer serves the health endpoints and, when configured, the OpenAI-compatible API. type HTTPServer struct { degradedFor time.Duration // how long after the last good ping a down database still reads as degraded @@ -29,7 +34,7 @@ type HTTPServer struct { // NewHTTPServer creates the HTTP server. api may be nil (health only); // requestTimeout bounds one API call, which waits on an upstream model. -func NewHTTPServer(pinger Pinger, addr string, api *OpenAIAPI, requestTimeout time.Duration) *HTTPServer { +func NewHTTPServer(pinger Pinger, addr string, api *OpenAIAPI, requestTimeout time.Duration, more ...Registrar) *HTTPServer { s := &HTTPServer{pinger: pinger, addr: addr} mux := http.NewServeMux() @@ -38,6 +43,9 @@ func NewHTTPServer(pinger Pinger, addr string, api *OpenAIAPI, requestTimeout ti if api != nil { api.Register(mux) } + for _, r := range more { + r.Register(mux) + } s.server = &http.Server{ Addr: addr, diff --git a/internal/throttle/keylimit.go b/internal/throttle/keylimit.go index 42d2e32..c25e05a 100644 --- a/internal/throttle/keylimit.go +++ b/internal/throttle/keylimit.go @@ -135,3 +135,28 @@ func (t *Redis) AdmitKey(ctx context.Context, keyName string, rpm int32, tpm, es } }, "", 0 } + +// KindToolRPM refuses a call above an MCP server's tool_rpm. +const KindToolRPM = "mcp_tool_rpm" + +// AdmitTool books one call of a key to an MCP tool under rpm calls a minute, +// or says how long until one fits. Fails open, like AdmitKey. +func (t *Redis) AdmitTool(ctx context.Context, keyName, server, tool string, rpm int32) (bool, time.Duration) { + if t.client == nil || rpm <= 0 { + return true, 0 + } + var nowMs int64 + if t.now != nil { + nowMs = t.now().UTC().UnixMilli() + } + out, err := keyAdmitScript.Run(ctx, t.client, []string{t.keys.ToolRPM(keyName, server, tool), ""}, nowMs, rpm, 0, 0).Int64Slice() + if err != nil || len(out) != 2 { + observability.DependencyErrorsTotal.WithLabelValues("redis", "mcp_tool_rate").Inc() + redisguard.Log(err).Err(err).Str("key", keyName).Str("mcp_server", server).Msg("throttle: tool rate check failed, allowing (fail-open)") + return true, 0 + } + if out[0] != 0 { + return false, time.Duration(out[1]) * time.Millisecond + } + return true, 0 +} diff --git a/internal/throttle/throttle.go b/internal/throttle/throttle.go index 3b2cb04..def31bd 100644 --- a/internal/throttle/throttle.go +++ b/internal/throttle/throttle.go @@ -586,7 +586,7 @@ func (t *Redis) ChargeDailyBudget(ctx context.Context, m *proxydb.Model, keyName // Every refusal reason is exported at 0 from the start, so an alert on // increase() also catches the first refusal after a restart. func init() { - reasons := []string{"rpm", "budget_key", "budget_user", KindKeyRPM, KindKeyTPM, spendlimit.ReasonUnavailable} + reasons := []string{"rpm", "budget_key", "budget_user", KindKeyRPM, KindKeyTPM, KindToolRPM, spendlimit.ReasonUnavailable} for _, e := range spendlimit.Enforced { reasons = append(reasons, spendlimit.Reason(e.Subject, e.Period)) } diff --git a/internal/usagereport/client.go b/internal/usagereport/client.go index dbc5d92..0b98179 100644 --- a/internal/usagereport/client.go +++ b/internal/usagereport/client.go @@ -387,21 +387,80 @@ func (c *UsageCollectorClient) RecordUsage( // light-tokenmeter refuses cache counts beyond tokens_input, and the batch with them. clamped := usage.Clamped() event.CacheReadTokens, event.CacheWriteTokens = int64(clamped.CacheRead), int64(clamped.CacheWrite) - // Never block a request on accounting: a full queue means the collector has - // been down long enough to fill it, and the loss is counted, not hidden. + c.enqueue(event) + return nil +} + +// enqueue never blocks a request on accounting: a full queue means the +// collector has been down long enough to fill it, and the loss is counted. +func (c *UsageCollectorClient) enqueue(event *usagev1.UsageEvent) { c.closeMu.RLock() defer c.closeMu.RUnlock() if c.closed { observability.UsageEventsDroppedTotal.WithLabelValues("closed").Inc() - return nil + return } select { case c.queue <- event: default: observability.UsageEventsDroppedTotal.WithLabelValues("queue_full").Inc() - log.Warn().Str("actionID", actionID).Str("model", model).Msg("usage event queue full, event dropped") + log.Warn().Str("actionID", event.GetActionId()).Str("model", event.GetModel()).Msg("usage event queue full, event dropped") } +} +// ToolCall is one MCP tool call to account for. +type ToolCall struct { + UserID, ActionID, Server, Tool, ErrorClass string + DurationMs, ResponseBytes int64 + Status usagev1.Status + Meta map[string]string +} + +// RecordToolCall queues a tool call through the MCP gateway: vendor "mcp", +// model server__tool and no tokens, so summaries by model count the calls. +func (c *UsageCollectorClient) RecordToolCall(t ToolCall) error { + if c == nil { + return errors.New("usage collector client is nil") + } + projectID := c.projectID + if c.projectFor != nil { + if p := c.projectFor(t.Meta); p != "" { + projectID = p + } + } + model := t.Server + "__" + t.Tool + if len(model) > maxModelLength { + model = model[:maxModelLength] + } + // light-tokenmeter refuses longer names, and the event with them. + server, tool := t.Server, t.Tool + if len(server) > maxVendorLength { + server = server[:maxVendorLength] + } + if len(tool) > maxModelLength { + tool = tool[:maxModelLength] + } + if t.UserID == "" || t.ActionID == "" || projectID == "" { + return fmt.Errorf("tool call %s needs user, action and project ids", model) + } + if err := validateLimits(t.UserID, projectID, t.ActionID, "mcp", model, 0, t.Status); err != nil { + return err + } + event := &usagev1.UsageEvent{ + EventId: uuid.Must(uuid.NewV7()).String(), + OccurredAt: timestamppb.Now(), + UserId: t.UserID, + ProjectId: projectID, + ActionId: t.ActionID, + Vendor: "mcp", + Model: model, + Status: t.Status, + DurationMs: t.DurationMs, + ErrorClass: t.ErrorClass, + Meta: fitMeta(t.Meta, t.ActionID), + ToolCall: &usagev1.ToolCall{Server: server, Tool: tool, ResponseBytes: t.ResponseBytes}, + } + c.enqueue(event) return nil } diff --git a/internal/usagereport/gen/usage/v1/usage.pb.go b/internal/usagereport/gen/usage/v1/usage.pb.go index 4878e88..e5cf8d6 100644 --- a/internal/usagereport/gen/usage/v1/usage.pb.go +++ b/internal/usagereport/gen/usage/v1/usage.pb.go @@ -120,8 +120,11 @@ type UsageEvent struct { // Input tokens written to the vendor's prompt cache (optional, >= 0). Part of // tokens_input; cache_read_tokens + cache_write_tokens must not exceed it. CacheWriteTokens int64 `protobuf:"varint,18,opt,name=cache_write_tokens,json=cacheWriteTokens,proto3" json:"cache_write_tokens,omitempty"` - unknownFields protoimpl.UnknownFields - sizeCache protoimpl.SizeCache + // Set on a tool call through an MCP gateway. vendor is then "mcp", model is + // server__tool (cut to fit) and the token counts are 0. + ToolCall *ToolCall `protobuf:"bytes,19,opt,name=tool_call,json=toolCall,proto3" json:"tool_call,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache } func (x *UsageEvent) Reset() { @@ -280,6 +283,77 @@ func (x *UsageEvent) GetCacheWriteTokens() int64 { return 0 } +func (x *UsageEvent) GetToolCall() *ToolCall { + if x != nil { + return x.ToolCall + } + return nil +} + +// ToolCall is one MCP tools/call an MCP gateway forwarded. +type ToolCall struct { + state protoimpl.MessageState `protogen:"open.v1"` + // MCP server name, as the gateway registered it (max 64 chars) + Server string `protobuf:"bytes,1,opt,name=server,proto3" json:"server,omitempty"` + // Tool name, as the server names it (max 128 chars) + Tool string `protobuf:"bytes,2,opt,name=tool,proto3" json:"tool,omitempty"` + // Size of the tool result in bytes, as JSON + ResponseBytes int64 `protobuf:"varint,3,opt,name=response_bytes,json=responseBytes,proto3" json:"response_bytes,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ToolCall) Reset() { + *x = ToolCall{} + mi := &file_usage_v1_usage_proto_msgTypes[1] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ToolCall) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ToolCall) ProtoMessage() {} + +func (x *ToolCall) ProtoReflect() protoreflect.Message { + mi := &file_usage_v1_usage_proto_msgTypes[1] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ToolCall.ProtoReflect.Descriptor instead. +func (*ToolCall) Descriptor() ([]byte, []int) { + return file_usage_v1_usage_proto_rawDescGZIP(), []int{1} +} + +func (x *ToolCall) GetServer() string { + if x != nil { + return x.Server + } + return "" +} + +func (x *ToolCall) GetTool() string { + if x != nil { + return x.Tool + } + return "" +} + +func (x *ToolCall) GetResponseBytes() int64 { + if x != nil { + return x.ResponseBytes + } + return 0 +} + // RecordUsageEventRequest is the request for recording a single usage event type RecordUsageEventRequest struct { state protoimpl.MessageState `protogen:"open.v1"` @@ -290,7 +364,7 @@ type RecordUsageEventRequest struct { func (x *RecordUsageEventRequest) Reset() { *x = RecordUsageEventRequest{} - mi := &file_usage_v1_usage_proto_msgTypes[1] + mi := &file_usage_v1_usage_proto_msgTypes[2] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -302,7 +376,7 @@ func (x *RecordUsageEventRequest) String() string { func (*RecordUsageEventRequest) ProtoMessage() {} func (x *RecordUsageEventRequest) ProtoReflect() protoreflect.Message { - mi := &file_usage_v1_usage_proto_msgTypes[1] + mi := &file_usage_v1_usage_proto_msgTypes[2] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -315,7 +389,7 @@ func (x *RecordUsageEventRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use RecordUsageEventRequest.ProtoReflect.Descriptor instead. func (*RecordUsageEventRequest) Descriptor() ([]byte, []int) { - return file_usage_v1_usage_proto_rawDescGZIP(), []int{1} + return file_usage_v1_usage_proto_rawDescGZIP(), []int{2} } func (x *RecordUsageEventRequest) GetEvent() *UsageEvent { @@ -336,7 +410,7 @@ type RecordUsageEventResponse struct { func (x *RecordUsageEventResponse) Reset() { *x = RecordUsageEventResponse{} - mi := &file_usage_v1_usage_proto_msgTypes[2] + mi := &file_usage_v1_usage_proto_msgTypes[3] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -348,7 +422,7 @@ func (x *RecordUsageEventResponse) String() string { func (*RecordUsageEventResponse) ProtoMessage() {} func (x *RecordUsageEventResponse) ProtoReflect() protoreflect.Message { - mi := &file_usage_v1_usage_proto_msgTypes[2] + mi := &file_usage_v1_usage_proto_msgTypes[3] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -361,7 +435,7 @@ func (x *RecordUsageEventResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use RecordUsageEventResponse.ProtoReflect.Descriptor instead. func (*RecordUsageEventResponse) Descriptor() ([]byte, []int) { - return file_usage_v1_usage_proto_rawDescGZIP(), []int{2} + return file_usage_v1_usage_proto_rawDescGZIP(), []int{3} } func (x *RecordUsageEventResponse) GetCreated() bool { @@ -382,7 +456,7 @@ type RecordUsageEventsRequest struct { func (x *RecordUsageEventsRequest) Reset() { *x = RecordUsageEventsRequest{} - mi := &file_usage_v1_usage_proto_msgTypes[3] + mi := &file_usage_v1_usage_proto_msgTypes[4] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -394,7 +468,7 @@ func (x *RecordUsageEventsRequest) String() string { func (*RecordUsageEventsRequest) ProtoMessage() {} func (x *RecordUsageEventsRequest) ProtoReflect() protoreflect.Message { - mi := &file_usage_v1_usage_proto_msgTypes[3] + mi := &file_usage_v1_usage_proto_msgTypes[4] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -407,7 +481,7 @@ func (x *RecordUsageEventsRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use RecordUsageEventsRequest.ProtoReflect.Descriptor instead. func (*RecordUsageEventsRequest) Descriptor() ([]byte, []int) { - return file_usage_v1_usage_proto_rawDescGZIP(), []int{3} + return file_usage_v1_usage_proto_rawDescGZIP(), []int{4} } func (x *RecordUsageEventsRequest) GetEvents() []*UsageEvent { @@ -430,7 +504,7 @@ type RecordUsageEventsResponse struct { func (x *RecordUsageEventsResponse) Reset() { *x = RecordUsageEventsResponse{} - mi := &file_usage_v1_usage_proto_msgTypes[4] + mi := &file_usage_v1_usage_proto_msgTypes[5] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -442,7 +516,7 @@ func (x *RecordUsageEventsResponse) String() string { func (*RecordUsageEventsResponse) ProtoMessage() {} func (x *RecordUsageEventsResponse) ProtoReflect() protoreflect.Message { - mi := &file_usage_v1_usage_proto_msgTypes[4] + mi := &file_usage_v1_usage_proto_msgTypes[5] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -455,7 +529,7 @@ func (x *RecordUsageEventsResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use RecordUsageEventsResponse.ProtoReflect.Descriptor instead. func (*RecordUsageEventsResponse) Descriptor() ([]byte, []int) { - return file_usage_v1_usage_proto_rawDescGZIP(), []int{4} + return file_usage_v1_usage_proto_rawDescGZIP(), []int{5} } func (x *RecordUsageEventsResponse) GetCreatedCount() int32 { @@ -499,7 +573,7 @@ type GetUsageSummaryRequest struct { func (x *GetUsageSummaryRequest) Reset() { *x = GetUsageSummaryRequest{} - mi := &file_usage_v1_usage_proto_msgTypes[5] + mi := &file_usage_v1_usage_proto_msgTypes[6] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -511,7 +585,7 @@ func (x *GetUsageSummaryRequest) String() string { func (*GetUsageSummaryRequest) ProtoMessage() {} func (x *GetUsageSummaryRequest) ProtoReflect() protoreflect.Message { - mi := &file_usage_v1_usage_proto_msgTypes[5] + mi := &file_usage_v1_usage_proto_msgTypes[6] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -524,7 +598,7 @@ func (x *GetUsageSummaryRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use GetUsageSummaryRequest.ProtoReflect.Descriptor instead. func (*GetUsageSummaryRequest) Descriptor() ([]byte, []int) { - return file_usage_v1_usage_proto_rawDescGZIP(), []int{5} + return file_usage_v1_usage_proto_rawDescGZIP(), []int{6} } func (x *GetUsageSummaryRequest) GetFrom() *timestamppb.Timestamp { @@ -654,7 +728,7 @@ type UsageSummaryRow struct { func (x *UsageSummaryRow) Reset() { *x = UsageSummaryRow{} - mi := &file_usage_v1_usage_proto_msgTypes[6] + mi := &file_usage_v1_usage_proto_msgTypes[7] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -666,7 +740,7 @@ func (x *UsageSummaryRow) String() string { func (*UsageSummaryRow) ProtoMessage() {} func (x *UsageSummaryRow) ProtoReflect() protoreflect.Message { - mi := &file_usage_v1_usage_proto_msgTypes[6] + mi := &file_usage_v1_usage_proto_msgTypes[7] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -679,7 +753,7 @@ func (x *UsageSummaryRow) ProtoReflect() protoreflect.Message { // Deprecated: Use UsageSummaryRow.ProtoReflect.Descriptor instead. func (*UsageSummaryRow) Descriptor() ([]byte, []int) { - return file_usage_v1_usage_proto_rawDescGZIP(), []int{6} + return file_usage_v1_usage_proto_rawDescGZIP(), []int{7} } func (x *UsageSummaryRow) GetBucket() *timestamppb.Timestamp { @@ -798,7 +872,7 @@ type GetUsageSummaryResponse struct { func (x *GetUsageSummaryResponse) Reset() { *x = GetUsageSummaryResponse{} - mi := &file_usage_v1_usage_proto_msgTypes[7] + mi := &file_usage_v1_usage_proto_msgTypes[8] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -810,7 +884,7 @@ func (x *GetUsageSummaryResponse) String() string { func (*GetUsageSummaryResponse) ProtoMessage() {} func (x *GetUsageSummaryResponse) ProtoReflect() protoreflect.Message { - mi := &file_usage_v1_usage_proto_msgTypes[7] + mi := &file_usage_v1_usage_proto_msgTypes[8] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -823,7 +897,7 @@ func (x *GetUsageSummaryResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use GetUsageSummaryResponse.ProtoReflect.Descriptor instead. func (*GetUsageSummaryResponse) Descriptor() ([]byte, []int) { - return file_usage_v1_usage_proto_rawDescGZIP(), []int{7} + return file_usage_v1_usage_proto_rawDescGZIP(), []int{8} } func (x *GetUsageSummaryResponse) GetRows() []*UsageSummaryRow { @@ -859,7 +933,7 @@ type GetUsageByMetaRequest struct { func (x *GetUsageByMetaRequest) Reset() { *x = GetUsageByMetaRequest{} - mi := &file_usage_v1_usage_proto_msgTypes[8] + mi := &file_usage_v1_usage_proto_msgTypes[9] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -871,7 +945,7 @@ func (x *GetUsageByMetaRequest) String() string { func (*GetUsageByMetaRequest) ProtoMessage() {} func (x *GetUsageByMetaRequest) ProtoReflect() protoreflect.Message { - mi := &file_usage_v1_usage_proto_msgTypes[8] + mi := &file_usage_v1_usage_proto_msgTypes[9] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -884,7 +958,7 @@ func (x *GetUsageByMetaRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use GetUsageByMetaRequest.ProtoReflect.Descriptor instead. func (*GetUsageByMetaRequest) Descriptor() ([]byte, []int) { - return file_usage_v1_usage_proto_rawDescGZIP(), []int{8} + return file_usage_v1_usage_proto_rawDescGZIP(), []int{9} } func (x *GetUsageByMetaRequest) GetMetaKey() string { @@ -937,7 +1011,7 @@ type UsageByMetaRow struct { func (x *UsageByMetaRow) Reset() { *x = UsageByMetaRow{} - mi := &file_usage_v1_usage_proto_msgTypes[9] + mi := &file_usage_v1_usage_proto_msgTypes[10] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -949,7 +1023,7 @@ func (x *UsageByMetaRow) String() string { func (*UsageByMetaRow) ProtoMessage() {} func (x *UsageByMetaRow) ProtoReflect() protoreflect.Message { - mi := &file_usage_v1_usage_proto_msgTypes[9] + mi := &file_usage_v1_usage_proto_msgTypes[10] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -962,7 +1036,7 @@ func (x *UsageByMetaRow) ProtoReflect() protoreflect.Message { // Deprecated: Use UsageByMetaRow.ProtoReflect.Descriptor instead. func (*UsageByMetaRow) Descriptor() ([]byte, []int) { - return file_usage_v1_usage_proto_rawDescGZIP(), []int{9} + return file_usage_v1_usage_proto_rawDescGZIP(), []int{10} } func (x *UsageByMetaRow) GetGroupValue() string { @@ -1046,7 +1120,7 @@ type GetUsageByMetaResponse struct { func (x *GetUsageByMetaResponse) Reset() { *x = GetUsageByMetaResponse{} - mi := &file_usage_v1_usage_proto_msgTypes[10] + mi := &file_usage_v1_usage_proto_msgTypes[11] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1058,7 +1132,7 @@ func (x *GetUsageByMetaResponse) String() string { func (*GetUsageByMetaResponse) ProtoMessage() {} func (x *GetUsageByMetaResponse) ProtoReflect() protoreflect.Message { - mi := &file_usage_v1_usage_proto_msgTypes[10] + mi := &file_usage_v1_usage_proto_msgTypes[11] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1071,7 +1145,7 @@ func (x *GetUsageByMetaResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use GetUsageByMetaResponse.ProtoReflect.Descriptor instead. func (*GetUsageByMetaResponse) Descriptor() ([]byte, []int) { - return file_usage_v1_usage_proto_rawDescGZIP(), []int{10} + return file_usage_v1_usage_proto_rawDescGZIP(), []int{11} } func (x *GetUsageByMetaResponse) GetRows() []*UsageByMetaRow { @@ -1092,7 +1166,7 @@ var File_usage_v1_usage_proto protoreflect.FileDescriptor const file_usage_v1_usage_proto_rawDesc = "" + "\n" + - "\x14usage/v1/usage.proto\x12\busage.v1\x1a\x1fgoogle/protobuf/timestamp.proto\"\xd6\x05\n" + + "\x14usage/v1/usage.proto\x12\busage.v1\x1a\x1fgoogle/protobuf/timestamp.proto\"\x87\x06\n" + "\n" + "UsageEvent\x12\x19\n" + "\bevent_id\x18\x01 \x01(\tR\aeventId\x12;\n" + @@ -1117,10 +1191,15 @@ const file_usage_v1_usage_proto_rawDesc = "" + "\x04meta\x18\x0f \x03(\v2\x1e.usage.v1.UsageEvent.MetaEntryR\x04meta\x12!\n" + "\fcost_nanousd\x18\x10 \x01(\x03R\vcostNanousd\x12*\n" + "\x11cache_read_tokens\x18\x11 \x01(\x03R\x0fcacheReadTokens\x12,\n" + - "\x12cache_write_tokens\x18\x12 \x01(\x03R\x10cacheWriteTokens\x1a7\n" + + "\x12cache_write_tokens\x18\x12 \x01(\x03R\x10cacheWriteTokens\x12/\n" + + "\ttool_call\x18\x13 \x01(\v2\x12.usage.v1.ToolCallR\btoolCall\x1a7\n" + "\tMetaEntry\x12\x10\n" + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + - "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"E\n" + + "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"]\n" + + "\bToolCall\x12\x16\n" + + "\x06server\x18\x01 \x01(\tR\x06server\x12\x12\n" + + "\x04tool\x18\x02 \x01(\tR\x04tool\x12%\n" + + "\x0eresponse_bytes\x18\x03 \x01(\x03R\rresponseBytes\"E\n" + "\x17RecordUsageEventRequest\x12*\n" + "\x05event\x18\x01 \x01(\v2\x14.usage.v1.UsageEventR\x05event\"4\n" + "\x18RecordUsageEventResponse\x12\x18\n" + @@ -1218,49 +1297,51 @@ func file_usage_v1_usage_proto_rawDescGZIP() []byte { } var file_usage_v1_usage_proto_enumTypes = make([]protoimpl.EnumInfo, 1) -var file_usage_v1_usage_proto_msgTypes = make([]protoimpl.MessageInfo, 12) +var file_usage_v1_usage_proto_msgTypes = make([]protoimpl.MessageInfo, 13) var file_usage_v1_usage_proto_goTypes = []any{ (Status)(0), // 0: usage.v1.Status (*UsageEvent)(nil), // 1: usage.v1.UsageEvent - (*RecordUsageEventRequest)(nil), // 2: usage.v1.RecordUsageEventRequest - (*RecordUsageEventResponse)(nil), // 3: usage.v1.RecordUsageEventResponse - (*RecordUsageEventsRequest)(nil), // 4: usage.v1.RecordUsageEventsRequest - (*RecordUsageEventsResponse)(nil), // 5: usage.v1.RecordUsageEventsResponse - (*GetUsageSummaryRequest)(nil), // 6: usage.v1.GetUsageSummaryRequest - (*UsageSummaryRow)(nil), // 7: usage.v1.UsageSummaryRow - (*GetUsageSummaryResponse)(nil), // 8: usage.v1.GetUsageSummaryResponse - (*GetUsageByMetaRequest)(nil), // 9: usage.v1.GetUsageByMetaRequest - (*UsageByMetaRow)(nil), // 10: usage.v1.UsageByMetaRow - (*GetUsageByMetaResponse)(nil), // 11: usage.v1.GetUsageByMetaResponse - nil, // 12: usage.v1.UsageEvent.MetaEntry - (*timestamppb.Timestamp)(nil), // 13: google.protobuf.Timestamp + (*ToolCall)(nil), // 2: usage.v1.ToolCall + (*RecordUsageEventRequest)(nil), // 3: usage.v1.RecordUsageEventRequest + (*RecordUsageEventResponse)(nil), // 4: usage.v1.RecordUsageEventResponse + (*RecordUsageEventsRequest)(nil), // 5: usage.v1.RecordUsageEventsRequest + (*RecordUsageEventsResponse)(nil), // 6: usage.v1.RecordUsageEventsResponse + (*GetUsageSummaryRequest)(nil), // 7: usage.v1.GetUsageSummaryRequest + (*UsageSummaryRow)(nil), // 8: usage.v1.UsageSummaryRow + (*GetUsageSummaryResponse)(nil), // 9: usage.v1.GetUsageSummaryResponse + (*GetUsageByMetaRequest)(nil), // 10: usage.v1.GetUsageByMetaRequest + (*UsageByMetaRow)(nil), // 11: usage.v1.UsageByMetaRow + (*GetUsageByMetaResponse)(nil), // 12: usage.v1.GetUsageByMetaResponse + nil, // 13: usage.v1.UsageEvent.MetaEntry + (*timestamppb.Timestamp)(nil), // 14: google.protobuf.Timestamp } var file_usage_v1_usage_proto_depIdxs = []int32{ - 13, // 0: usage.v1.UsageEvent.occurred_at:type_name -> google.protobuf.Timestamp + 14, // 0: usage.v1.UsageEvent.occurred_at:type_name -> google.protobuf.Timestamp 0, // 1: usage.v1.UsageEvent.status:type_name -> usage.v1.Status - 12, // 2: usage.v1.UsageEvent.meta:type_name -> usage.v1.UsageEvent.MetaEntry - 1, // 3: usage.v1.RecordUsageEventRequest.event:type_name -> usage.v1.UsageEvent - 1, // 4: usage.v1.RecordUsageEventsRequest.events:type_name -> usage.v1.UsageEvent - 13, // 5: usage.v1.GetUsageSummaryRequest.from:type_name -> google.protobuf.Timestamp - 13, // 6: usage.v1.GetUsageSummaryRequest.to:type_name -> google.protobuf.Timestamp - 0, // 7: usage.v1.GetUsageSummaryRequest.status:type_name -> usage.v1.Status - 13, // 8: usage.v1.UsageSummaryRow.bucket:type_name -> google.protobuf.Timestamp - 0, // 9: usage.v1.UsageSummaryRow.status:type_name -> usage.v1.Status - 7, // 10: usage.v1.GetUsageSummaryResponse.rows:type_name -> usage.v1.UsageSummaryRow - 10, // 11: usage.v1.GetUsageByMetaResponse.rows:type_name -> usage.v1.UsageByMetaRow - 2, // 12: usage.v1.UsageService.RecordUsageEvent:input_type -> usage.v1.RecordUsageEventRequest - 4, // 13: usage.v1.UsageService.RecordUsageEvents:input_type -> usage.v1.RecordUsageEventsRequest - 6, // 14: usage.v1.UsageService.GetUsageSummary:input_type -> usage.v1.GetUsageSummaryRequest - 9, // 15: usage.v1.UsageService.GetUsageByMeta:input_type -> usage.v1.GetUsageByMetaRequest - 3, // 16: usage.v1.UsageService.RecordUsageEvent:output_type -> usage.v1.RecordUsageEventResponse - 5, // 17: usage.v1.UsageService.RecordUsageEvents:output_type -> usage.v1.RecordUsageEventsResponse - 8, // 18: usage.v1.UsageService.GetUsageSummary:output_type -> usage.v1.GetUsageSummaryResponse - 11, // 19: usage.v1.UsageService.GetUsageByMeta:output_type -> usage.v1.GetUsageByMetaResponse - 16, // [16:20] is the sub-list for method output_type - 12, // [12:16] is the sub-list for method input_type - 12, // [12:12] is the sub-list for extension type_name - 12, // [12:12] is the sub-list for extension extendee - 0, // [0:12] is the sub-list for field type_name + 13, // 2: usage.v1.UsageEvent.meta:type_name -> usage.v1.UsageEvent.MetaEntry + 2, // 3: usage.v1.UsageEvent.tool_call:type_name -> usage.v1.ToolCall + 1, // 4: usage.v1.RecordUsageEventRequest.event:type_name -> usage.v1.UsageEvent + 1, // 5: usage.v1.RecordUsageEventsRequest.events:type_name -> usage.v1.UsageEvent + 14, // 6: usage.v1.GetUsageSummaryRequest.from:type_name -> google.protobuf.Timestamp + 14, // 7: usage.v1.GetUsageSummaryRequest.to:type_name -> google.protobuf.Timestamp + 0, // 8: usage.v1.GetUsageSummaryRequest.status:type_name -> usage.v1.Status + 14, // 9: usage.v1.UsageSummaryRow.bucket:type_name -> google.protobuf.Timestamp + 0, // 10: usage.v1.UsageSummaryRow.status:type_name -> usage.v1.Status + 8, // 11: usage.v1.GetUsageSummaryResponse.rows:type_name -> usage.v1.UsageSummaryRow + 11, // 12: usage.v1.GetUsageByMetaResponse.rows:type_name -> usage.v1.UsageByMetaRow + 3, // 13: usage.v1.UsageService.RecordUsageEvent:input_type -> usage.v1.RecordUsageEventRequest + 5, // 14: usage.v1.UsageService.RecordUsageEvents:input_type -> usage.v1.RecordUsageEventsRequest + 7, // 15: usage.v1.UsageService.GetUsageSummary:input_type -> usage.v1.GetUsageSummaryRequest + 10, // 16: usage.v1.UsageService.GetUsageByMeta:input_type -> usage.v1.GetUsageByMetaRequest + 4, // 17: usage.v1.UsageService.RecordUsageEvent:output_type -> usage.v1.RecordUsageEventResponse + 6, // 18: usage.v1.UsageService.RecordUsageEvents:output_type -> usage.v1.RecordUsageEventsResponse + 9, // 19: usage.v1.UsageService.GetUsageSummary:output_type -> usage.v1.GetUsageSummaryResponse + 12, // 20: usage.v1.UsageService.GetUsageByMeta:output_type -> usage.v1.GetUsageByMetaResponse + 17, // [17:21] is the sub-list for method output_type + 13, // [13:17] is the sub-list for method input_type + 13, // [13:13] is the sub-list for extension type_name + 13, // [13:13] is the sub-list for extension extendee + 0, // [0:13] is the sub-list for field type_name } func init() { file_usage_v1_usage_proto_init() } @@ -1274,7 +1355,7 @@ func file_usage_v1_usage_proto_init() { GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: unsafe.Slice(unsafe.StringData(file_usage_v1_usage_proto_rawDesc), len(file_usage_v1_usage_proto_rawDesc)), NumEnums: 1, - NumMessages: 12, + NumMessages: 13, NumExtensions: 0, NumServices: 1, }, diff --git a/migrations/20261028120000_mcp_gateway.sql b/migrations/20261028120000_mcp_gateway.sql new file mode 100644 index 0000000..656cecb --- /dev/null +++ b/migrations/20261028120000_mcp_gateway.sql @@ -0,0 +1,26 @@ +-- Upstream MCP servers proxied at /mcp, and the header the gateway adds to +-- every request it sends one. Keys and teams name the servers and tools they +-- may use in allowed_mcp ("server" or "server/tool"); empty means any. +CREATE TABLE IF NOT EXISTS mcp_server ( + name TEXT PRIMARY KEY, + url TEXT NOT NULL, + transport TEXT NOT NULL DEFAULT 'streamable_http', + enabled BOOLEAN NOT NULL DEFAULT TRUE, + allow_private BOOLEAN NOT NULL DEFAULT FALSE, + timeout_ms INT NOT NULL DEFAULT 0, + max_response_bytes BIGINT NOT NULL DEFAULT 0, + tool_rpm INT NOT NULL DEFAULT 0, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +CREATE TABLE IF NOT EXISTS mcp_credential ( + server TEXT PRIMARY KEY REFERENCES mcp_server(name) ON DELETE CASCADE, + header TEXT NOT NULL, + value_ciphertext BYTEA NOT NULL, -- AES-256-GCM, as vendor_key + value_version INT NOT NULL, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +ALTER TABLE api_key ADD COLUMN IF NOT EXISTS allowed_mcp TEXT[] NOT NULL DEFAULT '{}'; +ALTER TABLE team ADD COLUMN IF NOT EXISTS allowed_mcp TEXT[] NOT NULL DEFAULT '{}'; diff --git a/proto/ktms/v1/admin.proto b/proto/ktms/v1/admin.proto index 9c59582..721abab 100644 --- a/proto/ktms/v1/admin.proto +++ b/proto/ktms/v1/admin.proto @@ -106,11 +106,13 @@ message GetConfigResponse { int64 version = 4; repeated UserRateLimitSpec user_rate_limits = 5; repeated SpendLimitSpec spend_limits = 6; + repeated MCPServerSpec mcp_servers = 7; + repeated MCPCredentialInfo mcp_credentials = 8; // metadata only, never the value } message MintAPIKeyRequest { string name = 1; - repeated string scopes = 2; // "generate" | "admin" | "admin:read" + repeated string scopes = 2; // "generate" | "mcp" | "admin" | "admin:read" string org_id = 3; string team_id = 4; APIKeyPolicy policy = 5; @@ -125,6 +127,9 @@ message APIKeyPolicy { // Tokens per minute over all models, 0 = unlimited. A call books its input // estimate when admitted and its real input + output when it ends. int64 tpm_limit = 3; + // MCP servers ("github") and single tools ("github/create_issue") a key with + // the mcp scope may use; empty = any. Others are hidden and refused. + repeated string allowed_mcp = 4; } // SetAPIKeyPolicy replaces a key's policy; it applies from the key's next call. @@ -327,6 +332,7 @@ message Team { double monthly_budget_usd = 4; repeated string allowed_models = 5; google.protobuf.Timestamp created_at = 6; + repeated string allowed_mcp = 7; // as APIKeyPolicy.allowed_mcp, for every key of the team } message CreateOrganizationRequest { @@ -354,6 +360,7 @@ message CreateTeamRequest { string name = 2; double monthly_budget_usd = 3; repeated string allowed_models = 4; + repeated string allowed_mcp = 5; } message CreateTeamResponse { @@ -377,12 +384,13 @@ message ListTeamsResponse { } // Replaces the team: an omitted budget removes the cap and omitted -// allowed_models lifts the restriction, so send every field. +// allowed_models or allowed_mcp lifts the restriction, so send every field. message UpdateTeamRequest { string id = 1; string name = 2; double monthly_budget_usd = 3; repeated string allowed_models = 4; + repeated string allowed_mcp = 5; } message UpdateTeamResponse { @@ -429,6 +437,76 @@ message GetSpendResponse { repeated SpendWindow windows = 1; } +// MCPServerSpec is an upstream MCP server, proxied at /mcp/{name} and, with +// its tools named name__tool, at /mcp. Its credential is not part of the spec: +// SetMCPCredential stores it. Experimental: the MCP gateway may change in a 1.x minor release. +message MCPServerSpec { + string name = 1; // 1-32 of a-z, 0-9 and "-", starting with a letter + string url = 2; // https; http only with allow_private + string transport = 3; // "streamable_http" (default) | "sse" + bool enabled = 4; + // Permit a private, loopback or link-local address, such as a server in the + // same cluster. Without it such an address is refused at registration and + // at every connection. + bool allow_private = 5; + int32 timeout_ms = 6; // one tool call, 0 = 60000 + int64 max_response_bytes = 7; // one tool result, 0 = 4 MiB + int32 tool_rpm = 8; // calls per minute per client key and tool, 0 = unlimited; needs Redis +} + +message UpsertMCPServerRequest { + MCPServerSpec server = 1; +} + +message UpsertMCPServerResponse {} + +// DeleteMCPServer removes the server and its credential. +message DeleteMCPServerRequest { + string name = 1; +} + +message DeleteMCPServerResponse {} + +// SetMCPCredential stores the header the gateway adds to every request it +// sends the server, such as Authorization: Bearer . Encrypted at rest +// with the vendor keys' keyring; never returned, logged or passed to clients. +message SetMCPCredentialRequest { + string server = 1; + string header = 2; // "" = Authorization + string plain_value = 3; +} + +message SetMCPCredentialResponse {} + +message DeleteMCPCredentialRequest { + string server = 1; +} + +message DeleteMCPCredentialResponse {} + +message MCPCredentialInfo { + string server = 1; + string header = 2; + int32 key_version = 3; + google.protobuf.Timestamp created_at = 4; + string fingerprint = 5; // as VendorKeyInfo.fingerprint +} + +// ListMCPTools asks the server for its tools through the gateway, with its +// credential, so an operator can write allowed_mcp. +message ListMCPToolsRequest { + string server = 1; +} + +message MCPToolInfo { + string name = 1; // as the server names it; /mcp shows server__name + string description = 2; +} + +message ListMCPToolsResponse { + repeated MCPToolInfo tools = 1; +} + // Admin RPCs require the "admin" scope; the Get* and List* reads and // ExplainRoute also accept "admin:read". Every config mutation validates the // resulting full config in one transaction and bumps the config version. @@ -471,6 +549,12 @@ service AdminService { // and every model of its chain, kept or dropped with the reason. No vendor // or judge is called and nothing is spent; rate limits are not checked. rpc ExplainRoute(ExplainRouteRequest) returns (ExplainRouteResponse); + + rpc UpsertMCPServer(UpsertMCPServerRequest) returns (UpsertMCPServerResponse); + rpc DeleteMCPServer(DeleteMCPServerRequest) returns (DeleteMCPServerResponse); + rpc SetMCPCredential(SetMCPCredentialRequest) returns (SetMCPCredentialResponse); + rpc DeleteMCPCredential(DeleteMCPCredentialRequest) returns (DeleteMCPCredentialResponse); + rpc ListMCPTools(ListMCPToolsRequest) returns (ListMCPToolsResponse); } message ExplainRouteRequest { diff --git a/schema/example.yaml b/schema/example.yaml index 122ec36..3a3e70e 100644 --- a/schema/example.yaml +++ b/schema/example.yaml @@ -28,6 +28,7 @@ organizations: - name: search monthly_budget_usd: 100 allowed_models: [gpt-4o-mini] + allowed_mcp: [github] api_keys: - name: billing scopes: [generate] @@ -37,6 +38,10 @@ api_keys: - name: search-bot organization: acme team: search + - name: coding-agent + scopes: [mcp] + policy: + allowed_mcp: [github/create_issue, github/search_code] vendor_keys: - api_key_name: billing vendor: openai @@ -50,3 +55,11 @@ user_rate_limits: - user_id: "*" model_id: gpt-4o rpm: 30 +mcp_servers: + - name: github + url: http://github-mcp.tools.svc:8080/mcp + allow_private: true + tool_rpm: 60 +mcp_credentials: + - server: github + value_env: TEST_MCP_TOKEN diff --git a/schema/ktms-config.schema.json b/schema/ktms-config.schema.json index 1159cb3..10a26c7 100644 --- a/schema/ktms-config.schema.json +++ b/schema/ktms-config.schema.json @@ -16,6 +16,12 @@ "policy": { "additionalProperties": false, "properties": { + "allowed_mcp": { + "items": { + "type": "string" + }, + "type": "array" + }, "allowed_models": { "items": { "type": "string" @@ -38,6 +44,7 @@ "items": { "enum": [ "generate", + "mcp", "admin", "admin:read" ] @@ -58,6 +65,70 @@ "api_version": { "const": "ktms.ai-process.dev/v1alpha1" }, + "mcp_credentials": { + "items": { + "additionalProperties": false, + "properties": { + "header": { + "description": "HTTP header the gateway sends the server; Authorization when left out.", + "type": "string" + }, + "server": { + "type": "string" + }, + "value_env": { + "description": "Environment variable holding the header's value, such as Bearer \u003ctoken\u003e; the value itself never goes in the file.", + "type": "string" + } + }, + "required": [ + "server", + "value_env" + ], + "type": "object" + }, + "type": "array" + }, + "mcp_servers": { + "items": { + "additionalProperties": false, + "properties": { + "allow_private": { + "type": "boolean" + }, + "enabled": { + "type": "boolean" + }, + "max_response_bytes": { + "type": [ + "integer", + "string" + ] + }, + "name": { + "type": "string" + }, + "timeout_ms": { + "type": "integer" + }, + "tool_rpm": { + "type": "integer" + }, + "transport": { + "type": "string" + }, + "url": { + "type": "string" + } + }, + "required": [ + "name", + "url" + ], + "type": "object" + }, + "type": "array" + }, "models": { "items": { "additionalProperties": false, @@ -159,6 +230,13 @@ "items": { "additionalProperties": false, "properties": { + "allowed_mcp": { + "description": "MCP servers (server) and tools (server/tool) the team's keys may use; empty means any.", + "items": { + "type": "string" + }, + "type": "array" + }, "allowed_models": { "items": { "type": "string" diff --git a/usage-proto/usage/v1/usage.proto b/usage-proto/usage/v1/usage.proto index 02b2888..6fd77ca 100644 --- a/usage-proto/usage/v1/usage.proto +++ b/usage-proto/usage/v1/usage.proto @@ -73,6 +73,22 @@ message UsageEvent { // Input tokens written to the vendor's prompt cache (optional, >= 0). Part of // tokens_input; cache_read_tokens + cache_write_tokens must not exceed it. int64 cache_write_tokens = 18; + + // Set on a tool call through an MCP gateway. vendor is then "mcp", model is + // server__tool (cut to fit) and the token counts are 0. + ToolCall tool_call = 19; +} + +// ToolCall is one MCP tools/call an MCP gateway forwarded. +message ToolCall { + // MCP server name, as the gateway registered it (max 64 chars) + string server = 1; + + // Tool name, as the server names it (max 128 chars) + string tool = 2; + + // Size of the tool result in bytes, as JSON + int64 response_bytes = 3; } // RecordUsageEventRequest is the request for recording a single usage event