internal/mcpgateway/service.go
Rama main · réplica de lectura del repositorio original.
package mcpgateway
import (
"context"
"crypto/rand"
"errors"
"fmt"
"log/slog"
"maps"
"net/http"
"slices"
"sort"
"strings"
"sync"
"time"
"github.com/modelcontextprotocol/go-sdk/mcp"
"golang.org/x/sync/singleflight"
"github.com/enterpilot/gomodel/config"
"github.com/enterpilot/gomodel/ext"
"github.com/enterpilot/gomodel/internal/core"
"github.com/enterpilot/gomodel/internal/encryption"
"github.com/enterpilot/gomodel/internal/usage"
"github.com/enterpilot/gomodel/internal/version"
)
// sessionIdleTimeout closes downstream sessions that stop sending requests,
// and bounds the session-binding registry.
const sessionIdleTimeout = 30 * time.Minute
// ScopeHeader restricts a request's visible servers to a comma-separated
// subset, so one gateway key can serve differently-scoped clients without
// extra endpoints. Unknown names are ignored; an empty header means all.
const ScopeHeader = "X-MCP-Servers"
// labelsHeader carries the request labels resolved by the gateway (tagging
// headers plus auth-key labels) from the HTTP layer into tool handlers, which
// only see per-message HTTP headers. It is internal: the HTTP entrypoint
// strips any client-sent value before stamping its own.
const labelsHeader = "X-Gomodel-Mcp-Labels"
// authKeyIDHeader carries the managed auth key id into tool handlers so MCP
// usage rows can be attributed per key. Internal like labelsHeader.
const authKeyIDHeader = "X-Gomodel-Mcp-Auth-Key-Id"
// pinnedServerKey carries the /mcp/{server} path pin through the request
// context into per-session server construction.
type pinnedServerKey struct{}
// virtualServerKey carries the /mcp/{name} virtual server the same way.
type virtualServerKey struct{}
// virtualBindingPrefix marks a session bound to a virtual server, so a session
// opened on /mcp/{name} stops matching if a real server takes that path. "/"
// never occurs in a slug, so the two kinds cannot collide.
const virtualBindingPrefix = "virtual/"
// Service is the MCP gateway: it merges declarative and admin-store server
// specs into the upstream manager and serves the downstream MCP endpoints.
type Service struct {
manager *Manager
store Store
usageLogger usage.LoggerInterface
userPathHeader string
configSpecs map[string]ServerSpec
// virtualConfig holds the declarative virtual servers; virtualStore the
// admin-managed ones. virtualSpecs is their merge, rebuilt on reload.
virtualConfig map[string]VirtualServerSpec
virtualStore VirtualStore
virtualMu sync.RWMutex
virtualSpecs map[string]VirtualServerSpec
// secrets resolves the secret references admin-managed headers hold.
secrets *config.Secrets
// oauth holds operator OAuth sign-ins to upstream servers; nil when the
// store cannot persist them.
oauth *oauthManager
// reloadMu serializes reloads, so a slower one never applies an older
// server set over a newer one. Close takes it too: an in-flight reload
// finishes before the upstreams are closed, and reloadStopped turns later
// ones into no-ops, so no upstream is dialed after shutdown.
reloadMu sync.Mutex
reloadStopped bool
// stopCtx is cancelled by Close, so a reload blocked on a store read,
// whether background, admin-triggered, or a secret rotation, does not
// hold up shutdown.
stopCtx context.Context
stopCancel context.CancelFunc
// mutateMu serializes admin saves and deletes from reading the stored
// row through releasing the secrets it held, so none releases a secret
// another one has just stored, or stores one another has just released.
mutateMu sync.Mutex
// recorded names the admin-managed servers whose references rotation
// watches.
recorded map[string]struct{}
// skipped holds, by name, why the last full reload left out stored
// servers whose references did not resolve. Guarded by reloadMu.
skipped map[string]error
// unresolved holds, by name, the last logged reason a stored server's
// references did not resolve, so a reload retried every few seconds logs
// a lasting failure once. Guarded by reloadMu.
unresolved map[string]string
// searchDiscovery is the default for sessions that do not send
// ToolDiscoveryHeader: serve search_tools/call_tool instead of the catalog.
searchDiscovery bool
handler http.Handler
origins *originGuard
bindMu sync.Mutex
bindings map[string]sessionBinding
// shared holds session bindings for the other replicas, so a session
// opened on one continues on any; nil keeps sessions on this instance.
// See sharedsessions.go.
shared ext.SharedStore
// resumes collapses concurrent resumes of one session ID.
resumes singleflight.Group
// now stamps session bindings and their shared records.
now func() time.Time
// postTouchInterval is how often a POST in flight extends its shared
// session's expiry; see trackPOST.
postTouchInterval time.Duration
requestMu sync.Mutex
requestCancels map[uint64]context.CancelFunc
nextRequestID uint64
closing bool
stopOnce sync.Once
stop chan struct{}
// changed is called after every successful store write; see
// SetChangeNotifier.
changed func(context.Context)
}
// SetChangeNotifier installs fn, called after every successful write of a
// server or virtual server so other gateway instances reload them. Call it
// before serving.
func (s *Service) SetChangeNotifier(fn func(context.Context)) {
if s != nil {
s.changed = fn
}
}
func (s *Service) notifyChanged(ctx context.Context) {
if s.changed != nil {
s.changed(ctx)
}
}
// errServiceClosed rejects a reload after Close.
var errServiceClosed = errors.New("mcp gateway is shut down")
// ErrServerNotVisible rejects a call from a session whose user path may no
// longer use the server.
var ErrServerNotVisible = errors.New("server is not available for this user path")
// sessionBinding pins a downstream MCP session to the user path it was
// initialized under. Bearer auth still runs on every request; the binding
// additionally stops a *different* principal from riding a leaked session ID,
// per the MCP session-hijacking guidance.
type sessionBinding struct {
authKeyID string
userPath string
// endpoint identifies the endpoint the session was opened on: a pinned
// server slug, virtualBindingPrefix plus a virtual server name, or "" for
// the aggregated endpoint.
endpoint string
lastSeen time.Time
// virtualIncarnation is the Incarnation of the virtual server the
// session opened on (see endpointReplaced); empty for any other endpoint.
virtualIncarnation string
// discovery and toolAliases let request logs name the tool a call runs:
// a discovery session's call_tool target, and the namespaced name an
// unambiguous bare name resolves to.
discovery bool
toolAliases map[string]string
// server hosts the session; closeEndpointSessions closes it through it.
server *mcp.Server
// resumed is set when this replica rebuilt the session from the shared
// store and serves it itself; nil when the SDK handler created it.
resumed *resumedSession
// shared reports that record, the session's encoded shared-store entry,
// is stored. checkedAt and touchedAt are when this replica last read the
// session's close marks and last extended the entry's expiry.
shared bool
record []byte
checkedAt time.Time
touchedAt time.Time
// verifiedAt is when this replica last read the close marks
// successfully, or opened or resumed the session; see
// sessionUnverifiedLimit.
verifiedAt time.Time
// posts counts this replica's POSTs in flight for the session, which
// pause its idle timeout, so sweepBindings keeps the binding.
posts int
}
// boundTo reports whether a request's identity and endpoint are the ones the
// session was opened under.
func (b sessionBinding) boundTo(authKeyID, userPath, endpoint string) bool {
return b.authKeyID == authKeyID && b.userPath == userPath && b.endpoint == endpoint
}
// Options configures NewService.
type Options struct {
// ConfigServers are the declarative servers from config.yaml / MCP_SERVERS.
ConfigServers map[string]ServerSpec
// VirtualServers are the declarative virtual servers, keyed by name.
VirtualServers map[string]VirtualServerSpec
// VirtualStore persists admin-managed virtual servers. Optional.
VirtualStore VirtualStore
// RefreshInterval re-reads the stores in the background, so admin edits
// made on another gateway instance sharing the database reach this one.
// Zero disables it.
RefreshInterval time.Duration
// Store persists admin-managed servers. Optional.
Store Store
// HTTPClient is the shared outbound HTTP client for http/sse upstreams.
HTTPClient *http.Client
// UsageLogger records one usage entry per tool call. Optional.
UsageLogger usage.LoggerInterface
// UserPathHeader is the configured user-path header name.
UserPathHeader string
// AllowedOrigins are the browser origins permitted to reach the endpoint.
// Empty trusts none, which is the default; see originGuard.
AllowedOrigins []string
// ToolDiscovery is the default discovery mode, config.MCPToolDiscoveryOff
// or config.MCPToolDiscoverySearch.
ToolDiscovery string
// Secrets resolves secret references in admin-managed header values.
// Nil resolves the built-in env and file schemes only.
Secrets *config.Secrets
// oauthStore persists operator OAuth sign-ins; nil disables them.
oauthStore *oauthStore
// SharedStore shares downstream session bindings between the replicas on
// one database, so a session continues on whichever replica a request
// reaches. Nil keeps every session on the instance that opened it.
SharedStore ext.SharedStore
}
// NewService builds the gateway service and starts connecting to the merged
// server set. Upstream connects are asynchronous; construction never blocks.
func NewService(ctx context.Context, opts Options) (*Service, error) {
s := &Service{
manager: NewManager(opts.HTTPClient),
store: opts.Store,
usageLogger: opts.UsageLogger,
userPathHeader: core.UserPathHeaderName(opts.UserPathHeader),
configSpecs: opts.ConfigServers,
virtualConfig: opts.VirtualServers,
virtualStore: opts.VirtualStore,
secrets: opts.Secrets,
recorded: make(map[string]struct{}),
unresolved: make(map[string]string),
searchDiscovery: opts.ToolDiscovery == config.MCPToolDiscoverySearch,
bindings: make(map[string]sessionBinding),
shared: opts.SharedStore,
now: time.Now,
// Half the idle timeout: the record never lapses mid-call.
postTouchInterval: sessionIdleTimeout / 2,
requestCancels: make(map[uint64]context.CancelFunc),
stop: make(chan struct{}),
}
if opts.oauthStore != nil {
s.oauth = newOAuthManager(opts.oauthStore, opts.HTTPClient)
}
s.stopCtx, s.stopCancel = context.WithCancel(context.Background())
guard, err := newOriginGuard(opts.AllowedOrigins)
if err != nil {
return nil, fmt.Errorf("mcp allowed origins: %w", err)
}
s.origins = guard
// The SDK's own DNS-rebinding guard stays enabled, but it only fires on
// loopback connections by design, so originGuard — applied in ServeHTTP
// ahead of this handler — is what defends every other bind address.
s.handler = mcp.NewStreamableHTTPHandler(s.getServer, &mcp.StreamableHTTPOptions{
SessionTimeout: sessionIdleTimeout,
Logger: slog.Default(),
})
if err := s.Reload(ctx); err != nil {
s.Close()
return nil, err
}
go s.sweepBindings()
if opts.RefreshInterval > 0 && (s.store != nil || s.virtualStore != nil) {
go s.refreshLoop(opts.RefreshInterval) //nolint:gosec // Runs for the service lifetime, not the constructor call.
}
return s, nil
}
// Refresh re-reads the stores and reconciles, like Reload, but also reports
// the stored servers whose references did not resolve, so the admin runtime
// refresh shows them and the cluster sync poller retries them.
func (s *Service) Refresh(ctx context.Context) error {
unresolved, err := s.reload(ctx, nil)
if err != nil {
return err
}
return unresolved
}
func (s *Service) refreshLoop(interval time.Duration) {
ticker := time.NewTicker(interval)
defer ticker.Stop()
for {
select {
case <-s.stop:
return
case <-ticker.C:
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
if err := s.Reload(ctx); err != nil && !errors.Is(err, errServiceClosed) && s.stopCtx.Err() == nil {
slog.Error("failed to refresh mcp servers", "error", err)
}
cancel()
}
}
}
// Reload re-merges declarative and store specs and reconciles the upstream
// set. Declarative entries shadow store rows with the same name, mirroring
// the tagging/virtual-models source precedence.
//
// Secret references in stored headers are resolved here. A row whose
// references do not resolve is logged and left out at startup; once its
// server runs, it keeps its current headers instead, so a secret backend
// outage never disconnects a working upstream.
func (s *Service) Reload(ctx context.Context) error {
_, err := s.reload(ctx, nil)
return err
}
// reload is Reload resolving only the rows rotated reports true for; every
// other running server keeps its current spec. A nil rotated resolves all.
// It reports the rows whose references did not resolve apart from err, which
// means nothing was applied.
func (s *Service) reload(ctx context.Context, rotated func(name string) bool) (unresolved, err error) {
// Close cancels every reload, not only background ones: an admin save
// or a rotation waiting on a slow store read must not hold shutdown open.
ctx, cancel := context.WithCancel(ctx)
defer cancel()
defer context.AfterFunc(s.stopCtx, cancel)()
s.reloadMu.Lock()
defer s.reloadMu.Unlock()
if s.reloadStopped || s.stopCtx.Err() != nil {
return nil, errServiceClosed
}
specs := make([]ServerSpec, 0, len(s.configSpecs))
seen := make(map[string]struct{}, len(s.configSpecs))
for name, spec := range s.configSpecs {
specs = append(specs, spec)
seen[name] = struct{}{}
}
var resolved []*config.ResolvedEntity
stored := make(map[string]struct{})
skipped := make(map[string]error)
var errs []error
if s.store != nil {
rows, err := s.store.List(ctx)
if err != nil {
return nil, fmt.Errorf("list managed mcp servers: %w", err)
}
for _, row := range rows {
if _, shadowed := seen[row.Name]; shadowed {
slog.Warn("mcp server from admin store is shadowed by config", "server", row.Name)
continue
}
stored[row.Name] = struct{}{}
current, running := s.manager.spec(row.Name)
if running && rotated != nil && !rotated(row.Name) {
specs = append(specs, current)
continue
}
spec, secrets, err := s.resolveServer(ctx, row)
if err != nil {
errs = append(errs, err)
repeated := s.unresolved[row.Name] == err.Error()
s.unresolved[row.Name] = err.Error()
if running {
if !repeated {
slog.Warn("mcp server secret references could not be resolved; keeping its current headers", "server", row.Name, "error", err)
}
specs = append(specs, current)
} else {
if !repeated {
slog.Error("mcp server secret references could not be resolved; server skipped", "server", row.Name, "error", err)
}
skipped[row.Name] = err
}
continue
}
delete(s.unresolved, row.Name)
specs = append(specs, spec)
resolved = append(resolved, secrets)
}
}
virtuals, err := s.mergedVirtualSpecs(ctx)
if err != nil {
return nil, err
}
s.virtualMu.Lock()
s.virtualSpecs = virtuals
s.virtualMu.Unlock()
s.attachOAuth(ctx, specs)
s.manager.Apply(specs)
s.logVirtualServerIssues()
for _, secrets := range resolved {
secrets.Record()
}
for name := range s.recorded {
if _, ok := stored[name]; !ok {
s.secrets.ForgetEntity(serverSecretEntity(name))
}
}
s.recorded = stored
for name := range s.unresolved {
if _, ok := stored[name]; !ok {
delete(s.unresolved, name)
}
}
if rotated == nil {
s.skipped = skipped
}
return errors.Join(errs...), nil
}
// mergedVirtualSpecs merges declarative and store virtual servers. Config
// entries shadow store rows with the same name, like servers.
func (s *Service) mergedVirtualSpecs(ctx context.Context) (map[string]VirtualServerSpec, error) {
merged := make(map[string]VirtualServerSpec, len(s.virtualConfig))
maps.Copy(merged, s.virtualConfig)
if s.virtualStore == nil {
return merged, nil
}
rows, err := s.virtualStore.ListVirtual(ctx)
if err != nil {
return nil, fmt.Errorf("list managed mcp virtual servers: %w", err)
}
for _, row := range rows {
if _, shadowed := merged[row.Name]; shadowed {
slog.Warn("mcp virtual server from admin store is shadowed by config", "virtual_server", row.Name)
continue
}
merged[row.Name] = row.Spec()
}
return merged, nil
}
// Skipped returns, by name, why the last full reload left out stored
// servers whose secret references did not resolve.
func (s *Service) Skipped() map[string]error {
s.reloadMu.Lock()
defer s.reloadMu.Unlock()
return maps.Clone(s.skipped)
}
// Running reports whether the server name is enabled in the running upstream
// set. A disabled server is kept there for the admin view but serves nothing.
func (s *Service) Running(name string) bool {
spec, ok := s.manager.spec(name)
return ok && spec.Enabled
}
// Views returns the current admin snapshot of all servers, with their OAuth
// sign-in state.
func (s *Service) Views() []ServerView {
views := s.manager.Views()
for i := range views {
views[i].OAuth = s.oauth.view(views[i].Spec)
}
return views
}
// IsManaged reports whether name is declared in config/env (read-only).
func (s *Service) IsManaged(name string) bool {
_, ok := s.configSpecs[name]
return ok
}
// Upsert validates and persists one admin-managed server, then reconciles.
// Config-declared names are read-only; stdio definitions are rejected at the
// store boundary (declarative-only by design — see the gateway spec).
func (s *Service) Upsert(ctx context.Context, server ManagedServer) error {
if s.store == nil {
return fmt.Errorf("mcp server persistence is unavailable")
}
if s.IsManaged(server.Name) {
return fmt.Errorf("mcp server %q is managed by config/env and is read-only", server.Name)
}
if err := server.Validate(); err != nil {
return err
}
// ErrSealUnconfirmed means the row was written: apply it like any other
// save, then report the error.
saveErr := s.persist(ctx, &server)
if saveErr != nil && !errors.Is(saveErr, encryption.ErrSealUnconfirmed) {
return saveErr
}
defer s.notifyChanged(ctx)
if err := s.Reload(ctx); err != nil {
// The row is persisted; only applying it to the running manager
// failed. Say so — a retry or restart picks the row up.
return errors.Join(saveErr, fmt.Errorf("mcp server %q was saved but not applied: %w", server.Name, err))
}
return saveErr
}
// persist stores server in place of the row it replaces and releases the
// secrets that row held and server no longer does. A name a virtual server
// takes is refused with VirtualNameTakenError unless the row already exists.
// An error wrapping encryption.ErrSealUnconfirmed means the row was written.
func (s *Service) persist(ctx context.Context, server *ManagedServer) error {
s.mutateMu.Lock()
defer s.mutateMu.Unlock()
previous, err := s.store.Get(ctx, server.Name)
if err != nil && !errors.Is(err, ErrNotFound) {
return err
}
write := s.store.Upsert
// Stored rows, not the loaded list: a virtual server another instance
// saved counts before this one refreshes. mutateMu, which virtual saves
// hold too, keeps an overlapping save on this instance from slipping in.
virtualTaken, err := s.virtualNameTaken(ctx, server.Name)
if err != nil {
return err
}
if virtualTaken {
// Only a new server is refused: one stored before the virtual server
// was declared keeps its endpoint and stays editable. Update fails
// atomically if the row is gone, so an edit racing a delete cannot
// recreate it under the virtual server's name.
if previous == nil {
return VirtualNameTakenError(server.Name)
}
write = s.store.Update
}
if err := config.CheckSecretDestinations(ctx, ServerSecretEntity, server.Name, serverDestinations(previous, server), headerValues(server)); err != nil {
return err
}
// Literal headers go to the secret writer first, when one is registered,
// so the row is stored with references only; then every reference must
// resolve before anything is persisted.
kept := headerValues(previous)
written, err := s.storeServerSecrets(ctx, server, kept)
if err != nil {
return err
}
if _, _, err := s.resolveServer(ctx, *server); err != nil {
// Nothing was stored, so whatever the stored row holds stays.
s.releaseSecrets(ctx, server.Name, written, kept)
return err
}
// ErrSealUnconfirmed means the row was written: treat it like any other
// save, then report the error.
saveErr := write(ctx, *server)
if saveErr != nil && !errors.Is(saveErr, encryption.ErrSealUnconfirmed) {
// The stored row is unchanged, so whatever it holds stays.
s.releaseSecrets(ctx, server.Name, written, kept)
if errors.Is(saveErr, ErrNotFound) {
return VirtualNameTakenError(server.Name)
}
return saveErr
}
// The new row is committed, so what it replaced is released even if
// applying it fails: a later edit only sees the new row.
s.releaseSecrets(ctx, server.Name, kept, headerValues(server))
return saveErr
}
// GetManaged returns one admin-managed server row from the store. Config-
// declared servers are not store rows, so they (and a missing store) report
// ErrNotFound.
func (s *Service) GetManaged(ctx context.Context, name string) (*ManagedServer, error) {
if s.store == nil {
return nil, ErrNotFound
}
return s.store.Get(ctx, name)
}
// Delete removes one admin-managed server, then reconciles.
func (s *Service) Delete(ctx context.Context, name string) error {
if s.store == nil {
return fmt.Errorf("mcp server persistence is unavailable")
}
if s.IsManaged(name) {
return fmt.Errorf("mcp server %q is managed by config/env and is read-only", name)
}
if err := s.remove(ctx, name); err != nil {
return err
}
defer s.notifyChanged(ctx)
if err := s.Reload(ctx); err != nil {
return fmt.Errorf("mcp server %q was deleted but the running set was not updated: %w", name, err)
}
return nil
}
// remove deletes the stored row name and releases the secrets it held.
func (s *Service) remove(ctx context.Context, name string) error {
s.mutateMu.Lock()
defer s.mutateMu.Unlock()
previous, err := s.store.Get(ctx, name)
if err != nil && !errors.Is(err, ErrNotFound) {
return err
}
if err := s.store.Delete(ctx, name); err != nil {
return err
}
s.releaseSecrets(ctx, name, headerValues(previous), nil)
s.forgetOAuth(ctx, name)
return nil
}
// Reconnect force-redials one server and returns its fresh state.
func (s *Service) Reconnect(ctx context.Context, name string) (ServerView, error) {
view, err := s.manager.Reconnect(ctx, name)
view.OAuth = s.oauth.view(view.Spec)
return view, err
}
// Close stops background work, terminates upstream sessions, and cancels
// downstream HTTP exchanges. Streamable HTTP clients keep a GET request open
// for server events, so those request contexts must be ended before the HTTP
// server can complete its graceful drain.
func (s *Service) Close() {
s.stopOnce.Do(func() {
close(s.stop)
s.requestMu.Lock()
s.closing = true
cancels := make([]context.CancelFunc, 0, len(s.requestCancels))
for _, cancel := range s.requestCancels {
cancels = append(cancels, cancel)
}
s.requestMu.Unlock()
for _, cancel := range cancels {
cancel()
}
// Wait out an in-flight Reload, then stop later ones, before closing
// the upstreams, so none is added after this point.
s.stopCancel()
s.reloadMu.Lock()
s.reloadStopped = true
s.reloadMu.Unlock()
s.closeResumedSessions()
s.manager.Close()
})
}
// ServeHTTP handles one downstream MCP HTTP exchange. pinnedServer is the
// /mcp/{server} path segment ("" for the aggregated endpoint): a server slug
// or a virtual server name; a server wins over a virtual server. Gateway
// authentication has already run; this layer enforces session-to-principal
// binding and stamps the internal identity headers tool handlers read.
func (s *Service) ServeHTTP(w http.ResponseWriter, r *http.Request, pinnedServer string) error {
// Rebinding check first: a rejected request must not touch a session, an
// upstream, or the request registry.
if err := s.origins.check(r); err != nil {
return err
}
if rejectUnsupportedProtocolVersion(w, r) {
return nil
}
requestCtx, cancel := context.WithCancel(r.Context())
s.requestMu.Lock()
if s.closing {
s.requestMu.Unlock()
cancel()
return fmt.Errorf("MCP gateway is shutting down")
}
requestID := s.nextRequestID
s.nextRequestID++
s.requestCancels[requestID] = cancel
s.requestMu.Unlock()
defer func() {
cancel()
s.requestMu.Lock()
delete(s.requestCancels, requestID)
s.requestMu.Unlock()
}()
r = r.WithContext(requestCtx)
userPath := core.UserPathFromContext(r.Context())
authKeyID := core.GetAuthKeyID(r.Context())
virtual := ""
if s.servedVirtual(pinnedServer) {
virtual = pinnedServer
} else if pinnedServer != "" {
s.closeStaleVirtualSessions(r, pinnedServer)
view, ok := s.findVisibleServer(pinnedServer, userPath)
if !ok {
return core.NewNotFoundError("unknown MCP server or virtual server: " + pinnedServer)
}
if !view.Spec.Enabled {
return core.NewNotFoundError("mcp server is disabled: " + pinnedServer)
}
}
bindingEndpoint := pinnedServer
if virtual != "" {
bindingEndpoint = virtualBindingPrefix + virtual
}
var resumed *resumedSession
sessionID := strings.TrimSpace(r.Header.Get("Mcp-Session-Id"))
if sessionID != "" {
admission, host := s.admitSession(r.Context(), sessionID, authKeyID, userPath, bindingEndpoint, r.Method == http.MethodDelete)
switch admission {
case sessionRejected:
// A different principal presented this session ID, or the session
// was closed on another replica. Report the session as gone (404
// per the transport spec) so the legitimate client's session stays
// unaffected and this caller re-initializes.
return core.NewNotFoundError("unknown MCP session")
case sessionEnded:
w.WriteHeader(http.StatusNoContent)
return nil
}
resumed = host
}
// The labels header is internal: never trust an inbound value.
r.Header.Del(labelsHeader)
if labels := core.RequestLabelsFromContext(r.Context()); len(labels) > 0 {
r.Header.Set(labelsHeader, strings.Join(labels, ","))
}
r.Header.Del(authKeyIDHeader)
if authKeyID != "" {
r.Header.Set(authKeyIDHeader, authKeyID)
}
ctx := r.Context()
if virtual != "" {
ctx = context.WithValue(ctx, virtualServerKey{}, virtual)
} else {
ctx = context.WithValue(ctx, pinnedServerKey{}, pinnedServer)
}
r = r.WithContext(ctx)
if sessionID != "" && r.Method == http.MethodPost {
defer s.trackPOST(sessionID)()
}
if resumed != nil {
resumed.ServeHTTP(w, r)
return nil
}
s.handler.ServeHTTP(w, r)
return nil
}
// requestScope captures the visibility inputs of one downstream session.
type requestScope struct {
authKeyID string
userPath string
pinned string
// virtual names the virtual server whose members bound the view.
virtual string
include map[string]struct{}
discovery bool
}
// endpoint is the /mcp/{name} segment the scope was opened on.
func (scope requestScope) endpoint() string {
if scope.virtual != "" {
return scope.virtual
}
return scope.pinned
}
// bindingEndpoint is the endpoint identity a session is bound to; see
// sessionBinding.endpoint.
func (scope requestScope) bindingEndpoint() string {
if scope.virtual != "" {
return virtualBindingPrefix + scope.virtual
}
return scope.pinned
}
func (s *Service) scopeFromRequest(r *http.Request) requestScope {
scope := requestScope{
authKeyID: core.GetAuthKeyID(r.Context()),
userPath: core.UserPathFromContext(r.Context()),
}
if pinned, ok := r.Context().Value(pinnedServerKey{}).(string); ok {
scope.pinned = pinned
}
defaultDiscovery := s.searchDiscovery
if virtual, ok := r.Context().Value(virtualServerKey{}).(string); ok {
scope.virtual = virtual
defaultDiscovery = s.virtualDiscovery(virtual)
}
scope.discovery = discoveryMode(r, defaultDiscovery)
if raw := r.Header.Get(ScopeHeader); raw != "" {
scope.include = make(map[string]struct{})
for name := range strings.SplitSeq(raw, ",") {
if trimmed := strings.TrimSpace(name); trimmed != "" {
scope.include[trimmed] = struct{}{}
}
}
}
return scope
}
// getServer builds the per-session MCP server for one downstream client: the
// SDK calls it when a new session initializes. The tool/prompt/resource view
// is filtered to what this principal may see — clients never discover tools
// they cannot call — and stays a snapshot while this replica hosts the
// session. A replica that resumes the session builds it from the catalog as
// it is then.
func (s *Service) getServer(r *http.Request) *mcp.Server {
server, _ := s.newSessionServer(s.scopeFromRequest(r))
return server
}
// newSessionServer builds the server for one session scope and returns it
// with the session's bare tool aliases. A new session binds itself when the
// SDK asks the server for its ID; the caller binds a resumed one.
func (s *Service) newSessionServer(scope requestScope) (*mcp.Server, map[string]string) {
views := s.visibleServers(scope)
// Assigned below; the SDK asks for the session ID only after getServer
// returns, so the binding sees the final map.
var aliases map[string]string
// Declared first so GetSessionID can bind the session to the server that
// hosts it; deleting a virtual server closes its sessions through it.
var server *mcp.Server
server = mcp.NewServer(&mcp.Implementation{
Name: "gomodel",
Title: "GoModel MCP Gateway",
Version: version.Version,
}, &mcp.ServerOptions{
Instructions: s.composeInstructions(scope, views),
SupportedProtocolVersions: sessionProtocolVersions,
// Advertise the full feature set even when the visible catalog is
// currently empty, so list calls return empty lists instead of
// method-not-supported errors.
Capabilities: &mcp.ServerCapabilities{
Tools: &mcp.ToolCapabilities{},
Prompts: &mcp.PromptCapabilities{},
Resources: &mcp.ResourceCapabilities{},
},
GetSessionID: func() string {
id := rand.Text()
s.bindSession(id, scope.authKeyID, scope.userPath, scope.bindingEndpoint(), scope.discovery, aliases, server)
return id
},
})
endpoint := "/mcp"
if name := scope.endpoint(); name != "" {
endpoint = "/mcp/" + name
}
prefixNames := scope.pinned == ""
var index *toolIndex
if scope.discovery {
index = &toolIndex{}
}
toolOwners := make(map[string]string)
promptOwners := make(map[string]string)
resourceOwners := make(map[string]string)
for _, view := range views {
snapshot, _ := s.upstreamCatalog(view.Spec.Name)
if snapshot == nil {
continue
}
s.registerTools(server, view.Spec.Name, snapshot, prefixNames, endpoint, toolOwners, index)
s.registerPrompts(server, view.Spec.Name, snapshot, prefixNames, promptOwners)
s.registerResources(server, view.Spec.Name, snapshot, resourceOwners)
}
if prefixNames {
aliases = bareToolAliases(toolOwners)
}
if index != nil {
index.aliases = aliases
s.registerDiscoveryTools(server, index, endpoint)
} else if len(aliases) > 0 {
server.AddReceivingMiddleware(bareToolCallMiddleware(aliases))
}
if s.shared != nil {
server.AddReceivingMiddleware(s.shareOnInitialize(scope))
}
return server, aliases
}
// bareToolAliases maps each unambiguous bare tool name to its namespaced
// name, so tools/call on the aggregated endpoint accepts a unique bare name
// (Postel, per the gateway spec). A bare name claimed by several servers, or
// one that collides with a registered namespaced name, gets no alias — the
// caller must use the namespaced form.
func bareToolAliases(owners map[string]string) map[string]string {
counts := make(map[string]int, len(owners))
aliases := make(map[string]string, len(owners))
for exposed, owner := range owners {
bare := strings.TrimPrefix(exposed, owner+namespaceSeparator)
counts[bare]++
aliases[bare] = exposed
}
for bare := range aliases {
if counts[bare] > 1 {
delete(aliases, bare)
continue
}
if _, taken := owners[bare]; taken {
delete(aliases, bare)
}
}
return aliases
}
// bareToolCallMiddleware rewrites a bare tools/call name to its namespaced
// registration before dispatch. Exact namespaced names never reach the alias
// map (they are excluded at build time), so registered names always win.
func bareToolCallMiddleware(aliases map[string]string) mcp.Middleware {
return func(next mcp.MethodHandler) mcp.MethodHandler {
return func(ctx context.Context, method string, req mcp.Request) (mcp.Result, error) {
if method == "tools/call" {
if call, ok := req.(*mcp.CallToolRequest); ok && call.Params != nil {
if exposed, ok := aliases[call.Params.Name]; ok {
call.Params.Name = exposed
}
}
}
return next(ctx, method, req)
}
}
}
func (s *Service) upstreamCatalog(name string) (*catalog, ServerStatus) {
u, ok := s.manager.get(name)
if !ok {
return nil, StatusDisabled
}
return u.snapshot()
}
// visibleServers filters the upstream set down to one request's view:
// enabled, allowed for the user path, and inside the pin, virtual server,
// and header scopes. A virtual server only narrows the view.
func (s *Service) visibleServers(scope requestScope) []ServerView {
var virtualMembers []string
if scope.virtual != "" {
spec, _ := s.virtualSpec(scope.virtual)
virtualMembers = spec.Servers
}
views := s.manager.Views()
visible := make([]ServerView, 0, len(views))
for _, view := range views {
if !view.Spec.Enabled {
continue
}
if scope.pinned != "" && view.Spec.Name != scope.pinned {
continue
}
if scope.virtual != "" && !slices.Contains(virtualMembers, view.Spec.Name) {
continue
}
if scope.include != nil {
if _, ok := scope.include[view.Spec.Name]; !ok {
continue
}
}
if !view.Spec.visibleTo(scope.userPath) {
continue
}
visible = append(visible, view)
}
return visible
}
func (s *Service) findVisibleServer(name, userPath string) (ServerView, bool) {
for _, view := range s.manager.Views() {
if view.Spec.Name != name {
continue
}
if !view.Spec.visibleTo(userPath) {
return ServerView{}, false
}
return view, true
}
return ServerView{}, false
}
// composeInstructions merges upstream instructions into the gateway's own,
// so guidance written for an upstream server survives aggregation.
func (s *Service) composeInstructions(scope requestScope, views []ServerView) string {
var b strings.Builder
if scope.pinned == "" {
names := make([]string, 0, len(views))
for _, view := range views {
names = append(names, view.Spec.Name)
}
sort.Strings(names)
if scope.virtual != "" {
fmt.Fprintf(&b, "GoModel MCP virtual server %q", scope.virtual)
} else {
b.WriteString("GoModel MCP gateway")
}
fmt.Fprintf(&b, " aggregating %d server(s): %s. Tools and prompts are namespaced as {server}%s{name}.",
len(names), strings.Join(names, ", "), namespaceSeparator)
if scope.virtual != "" && len(names) == 0 {
b.WriteString(" None of its member servers is available for this API key.")
}
}
if scope.discovery {
if b.Len() > 0 {
b.WriteString(" ")
}
fmt.Fprintf(&b, "Tools are not listed directly: find them with %s and run them with %s.", searchToolsName, callToolName)
}
for _, view := range views {
snapshot, _ := s.upstreamCatalog(view.Spec.Name)
if snapshot == nil || strings.TrimSpace(snapshot.instructions) == "" {
continue
}
if b.Len() > 0 {
displayName := strings.TrimSpace(view.Spec.DisplayName)
if displayName == "" || displayName == view.Spec.Name {
fmt.Fprintf(&b, "\n\n## %s\n", view.Spec.Name)
} else {
fmt.Fprintf(&b, "\n\n## %s (`%s`)\n", displayName, view.Spec.Name)
}
}
b.WriteString(strings.TrimSpace(snapshot.instructions))
}
return b.String()
}
// registerTools adds one upstream's tools to a session server. Tool metadata
// and valid schemas relay verbatim; only the name is prefixed on the
// aggregated endpoint. Arguments relay as raw JSON — validation belongs to
// the upstream. A non-nil index collects the tools for search discovery
// instead of listing them.
func (s *Service) registerTools(server *mcp.Server, upstreamName string, snapshot *catalog, prefix bool, endpoint string, owners map[string]string, index *toolIndex) {
for _, tool := range snapshot.tools {
exposed := tool.Name
if prefix {
exposed = NamespacedName(upstreamName, tool.Name)
}
if owner, taken := owners[exposed]; taken {
slog.Warn("mcp tool name collision; keeping first server",
"tool", exposed, "kept", owner, "skipped", upstreamName)
continue
}
owners[exposed] = upstreamName
if index != nil {
index.add(exposed, upstreamName, tool)
continue
}
clone := *tool
clone.Name = exposed
server.AddTool(&clone, s.toolHandler(upstreamName, tool.Name, exposed, endpoint))
}
}
func (s *Service) toolHandler(upstreamName, toolName, exposedName, endpoint string) mcp.ToolHandler {
return func(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
if err := s.authorizeSession(req.Session, upstreamName); err != nil {
return nil, err
}
started := time.Now()
result, err := s.manager.CallTool(ctx, upstreamName, toolName, req.Params.Arguments)
s.recordToolCall(req, upstreamName, exposedName, endpoint, started, result, err)
if err != nil {
return nil, fmt.Errorf("mcp server %q: %w", upstreamName, err)
}
return result, nil
}
}
// registerPrompts mirrors registerTools for prompts.
func (s *Service) registerPrompts(server *mcp.Server, upstreamName string, snapshot *catalog, prefix bool, owners map[string]string) {
for _, prompt := range snapshot.prompts {
exposed := prompt.Name
if prefix {
exposed = NamespacedName(upstreamName, prompt.Name)
}
if owner, taken := owners[exposed]; taken {
slog.Warn("mcp prompt name collision; keeping first server",
"prompt", exposed, "kept", owner, "skipped", upstreamName)
continue
}
owners[exposed] = upstreamName
clone := *prompt
clone.Name = exposed
originalName := prompt.Name
server.AddPrompt(&clone, func(ctx context.Context, req *mcp.GetPromptRequest) (*mcp.GetPromptResult, error) {
if err := s.authorizeSession(req.Session, upstreamName); err != nil {
return nil, err
}
params := *req.Params
params.Name = originalName
result, err := s.manager.GetPrompt(ctx, upstreamName, ¶ms)
if err != nil {
return nil, fmt.Errorf("mcp server %q: %w", upstreamName, err)
}
return result, nil
})
}
}
// registerResources adds resources and resource templates. Resource URIs are
// globally meaningful, so a URI already claimed by an earlier server is
// skipped with a warning rather than silently re-routed.
func (s *Service) registerResources(server *mcp.Server, upstreamName string, snapshot *catalog, owners map[string]string) {
read := func(ctx context.Context, req *mcp.ReadResourceRequest) (*mcp.ReadResourceResult, error) {
if err := s.authorizeSession(req.Session, upstreamName); err != nil {
return nil, err
}
result, err := s.manager.ReadResource(ctx, upstreamName, req.Params)
if err != nil {
return nil, fmt.Errorf("mcp server %q: %w", upstreamName, err)
}
return result, nil
}
for _, resource := range snapshot.resources {
if owner, taken := owners[resource.URI]; taken {
slog.Warn("mcp resource URI collision; keeping first server",
"uri", resource.URI, "kept", owner, "skipped", upstreamName)
continue
}
owners[resource.URI] = upstreamName
server.AddResource(resource, read)
}
for _, template := range snapshot.templates {
key := "template:" + template.URITemplate
if owner, taken := owners[key]; taken {
slog.Warn("mcp resource template collision; keeping first server",
"template", template.URITemplate, "kept", owner, "skipped", upstreamName)
continue
}
owners[key] = upstreamName
server.AddResourceTemplate(template, read)
}
}
// authorizeSession re-checks, on every call, that the user path a session
// was bound to may still use upstreamName. Sessions snapshot their catalog at
// initialize, so without this a narrowed user_paths or a new
// disallowed_user_paths entry would not reach sessions that are already open.
// A session without a binding is rejected: bindings live as long as the
// session, so a missing one means it was deleted or expired mid-call, and
// guessing a user path could slip past disallowed_user_paths. A session on a
// virtual server is also held to the virtual server's current members, so
// removing a member or the virtual server reaches open sessions too.
func (s *Service) authorizeSession(session *mcp.ServerSession, upstreamName string) error {
sessionID := ""
if session != nil {
sessionID = session.ID()
}
return s.authorizeSessionID(sessionID, upstreamName)
}
func (s *Service) authorizeSessionID(sessionID, upstreamName string) error {
s.bindMu.Lock()
binding, ok := s.bindings[sessionID]
s.bindMu.Unlock()
if !ok {
return fmt.Errorf("mcp server %q: %w", upstreamName, ErrServerNotVisible)
}
if _, visible := s.findVisibleServer(upstreamName, binding.userPath); !visible {
return fmt.Errorf("mcp server %q: %w", upstreamName, ErrServerNotVisible)
}
if virtual, ok := strings.CutPrefix(binding.endpoint, virtualBindingPrefix); ok && !s.virtualServes(virtual, upstreamName) {
return fmt.Errorf("mcp server %q: %w", upstreamName, ErrServerNotVisible)
}
return nil
}
// bindSession records the principal a new session was initialized under.
func (s *Service) bindSession(sessionID, authKeyID, userPath, endpoint string, discovery bool, toolAliases map[string]string, server *mcp.Server) {
now := s.now()
s.bindMu.Lock()
s.bindings[sessionID] = sessionBinding{
authKeyID: authKeyID,
userPath: userPath,
endpoint: endpoint,
lastSeen: now,
virtualIncarnation: s.endpointIncarnation(endpoint),
discovery: discovery,
toolAliases: toolAliases,
server: server,
}
s.bindMu.Unlock()
}
// closeStaleVirtualSessions closes the sessions of a virtual server that is no
// longer served when one of them sends a request. DeleteVirtual closes the
// sessions it finds, but one still initializing during the delete registers
// afterwards; this catches it on its next request instead of at the idle
// timeout.
func (s *Service) closeStaleVirtualSessions(r *http.Request, name string) {
sessionID := strings.TrimSpace(r.Header.Get("Mcp-Session-Id"))
if sessionID == "" {
return
}
endpoint := virtualBindingPrefix + name
s.bindMu.Lock()
binding, ok := s.bindings[sessionID]
s.bindMu.Unlock()
if ok && binding.endpoint == endpoint {
s.closeEndpointSessions(endpoint)
}
}
// closeEndpointSessions ends every session bound to endpoint (see
// sessionBinding.endpoint) and drops their bindings, so an endpoint that no
// longer exists does not keep idle sessions until the idle timeout. Closing
// waits for in-flight requests to drain, so it runs in the background.
func (s *Service) closeEndpointSessions(endpoint string) {
s.closeMatchingSessions(func(binding sessionBinding) bool { return binding.endpoint == endpoint })
}
// closeIncarnationSessions ends the sessions bound to one incarnation of a
// virtual server endpoint, and not those of a server recreated under its name.
func (s *Service) closeIncarnationSessions(endpoint, incarnation string) {
s.closeMatchingSessions(func(binding sessionBinding) bool {
return binding.endpoint == endpoint && binding.virtualIncarnation == incarnation
})
}
func (s *Service) closeMatchingSessions(match func(sessionBinding) bool) {
var servers []*mcp.Server
s.bindMu.Lock()
for id, binding := range s.bindings {
if !match(binding) {
continue
}
delete(s.bindings, id)
servers = append(servers, binding.server)
}
s.bindMu.Unlock()
for _, server := range servers {
closeSessions(server)
}
}
// ToolCallLabel names the tool a session's tools/call runs, matching its
// usage entry. name is the called tool and target the call's
// arguments.name. In a discovery session call_tool resolves to its target,
// and a bare name the session accepts resolves to its namespaced name. Only
// the session knows whether call_tool is the meta-tool: on a pinned endpoint
// without discovery it can be an upstream tool of that name.
func (s *Service) ToolCallLabel(sessionID, name, target string) string {
s.bindMu.Lock()
defer s.bindMu.Unlock()
binding := s.bindings[sessionID]
if binding.discovery && name == callToolName && target != "" {
name = target
}
if exposed, ok := binding.toolAliases[name]; ok {
return exposed
}
return name
}
// liveBinding returns the binding of a session this instance still hosts and
// refreshes its last-seen time. The binding of an SDK session that already
// ended (the idle timeout) is dropped, so the shared store can resume it.
func (s *Service) liveBinding(sessionID string) (sessionBinding, bool) {
s.bindMu.Lock()
defer s.bindMu.Unlock()
binding, ok := s.bindings[sessionID]
if !ok {
return sessionBinding{}, false
}
if binding.resumed == nil && binding.server != nil && !hasSessions(binding.server) {
delete(s.bindings, sessionID)
return sessionBinding{}, false
}
binding.lastSeen = s.now()
s.bindings[sessionID] = binding
return binding, true
}
func hasSessions(server *mcp.Server) bool {
for range server.Sessions() {
return true
}
return false
}
// sweepBindings evicts bindings whose sessions idled out, tracking the SDK's
// own session timeout.
func (s *Service) sweepBindings() {
ticker := time.NewTicker(sessionIdleTimeout / 2)
defer ticker.Stop()
for {
select {
case <-s.stop:
return
case <-ticker.C:
s.evictIdleBindings()
}
}
}
// evictIdleBindings drops the bindings of sessions idle past the idle
// timeout. A session with a POST in flight is not idle: its timeout is
// paused until the POST ends.
func (s *Service) evictIdleBindings() {
cutoff := s.now().Add(-sessionIdleTimeout - time.Minute)
s.bindMu.Lock()
defer s.bindMu.Unlock()
for id, binding := range s.bindings {
if binding.posts == 0 && binding.lastSeen.Before(cutoff) {
delete(s.bindings, id)
}
}
}