internal/mcpgateway/upstream.go
Rama main · réplica de lectura del repositorio original.
package mcpgateway
import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"log/slog"
"net/http"
"net/url"
"os"
"os/exec"
"regexp"
"slices"
"strings"
"sync"
"sync/atomic"
"time"
"github.com/modelcontextprotocol/go-sdk/mcp"
"github.com/enterpilot/gomodel/internal/version"
)
// connectTimeout bounds one upstream dial + initialize handshake.
const connectTimeout = 15 * time.Second
// listTimeout bounds one full catalog listing pass.
const listTimeout = 30 * time.Second
// ErrToolExcluded rejects a call to a tool the operator filters hide.
var ErrToolExcluded = errors.New("tool is excluded by the gateway tool filters")
// upstream owns the client session and catalog snapshot for one server. One
// shared session serves all downstream sessions: v1 forwards no per-user
// upstream credentials and bridges no server-to-client requests, so
// multiplexing is safe and avoids a session per client (or worse, per call).
type upstream struct {
spec ServerSpec
httpClient *http.Client
// connectMu serializes dial/refresh so concurrent callers cannot race a
// reconnect. stateMu guards the fields below and is never held across IO.
connectMu sync.Mutex
stateMu sync.Mutex
session *mcp.ClientSession
catalog *catalog
status ServerStatus
lastErr string
connectedAt time.Time
authRequired bool
// closed marks a permanently disposed upstream (removed or replaced by a
// reconcile). A stale background refresh must not redial it and leak an
// untracked session.
closed bool
// life bounds every HTTP request the upstream makes; close cancels it, so
// a dial hung on an unresponsive server ends with the upstream.
life context.Context
lifeCancel context.CancelFunc
}
func newUpstream(spec ServerSpec, httpClient *http.Client) *upstream {
u := &upstream{spec: spec, httpClient: httpClient, status: StatusConnecting}
u.life, u.lifeCancel = context.WithCancel(context.Background())
if !spec.Enabled {
u.status = StatusDisabled
}
return u
}
// view snapshots the upstream state for admin/dashboard consumption.
func (u *upstream) view() ServerView {
u.stateMu.Lock()
defer u.stateMu.Unlock()
return ServerView{
Spec: u.spec,
Status: u.status,
LastError: u.lastErr,
ToolCount: u.catalog.toolCount(),
ExcludedToolCount: u.catalog.excludedToolCount(),
PromptCount: u.catalog.promptCount(),
ResourceCount: u.catalog.resourceCount(),
ConnectedAt: u.connectedAt,
AuthRequired: u.authRequired,
}
}
// currentSpec returns the spec under the state lock; the access policy can
// change in place (setAccessPolicy), so readers of it must not race it.
func (u *upstream) currentSpec() ServerSpec {
u.stateMu.Lock()
defer u.stateMu.Unlock()
return u.spec
}
// setAccessPolicy swaps the tool filters and user-path scopes from spec and
// re-filters the cached catalog without touching the upstream session. They
// are gateway policy, so changing them never needs a redial.
func (u *upstream) setAccessPolicy(spec ServerSpec) {
u.stateMu.Lock()
defer u.stateMu.Unlock()
u.spec.AllowedTools = spec.AllowedTools
u.spec.DisallowedTools = spec.DisallowedTools
u.spec.UserPaths = spec.UserPaths
u.spec.DisallowedUserPaths = spec.DisallowedUserPaths
u.catalog = u.catalog.withToolFilters(spec.AllowedTools, spec.DisallowedTools)
}
// toolExposed reports whether the current filters expose name. Downstream
// sessions snapshot their tool list at initialize, so calls re-check here to
// make an exclusion effective for sessions that are already open.
func (u *upstream) toolExposed(name string) bool {
u.stateMu.Lock()
defer u.stateMu.Unlock()
return toolAllowed(name, u.spec.AllowedTools, u.spec.DisallowedTools)
}
// snapshot returns the current catalog (nil when never listed).
func (u *upstream) snapshot() (*catalog, ServerStatus) {
u.stateMu.Lock()
defer u.stateMu.Unlock()
return u.catalog, u.status
}
// refresh (re)establishes the session if needed and rebuilds the catalog. On
// failure the server is marked degraded and any previous catalog is kept, so
// a transient upstream hiccup does not blank out tools that were working — a
// failed listing must never register an "empty but connected" server.
func (u *upstream) refresh(ctx context.Context) error {
u.connectMu.Lock()
defer u.connectMu.Unlock()
if !u.spec.Enabled || u.isClosed() {
return nil
}
session, err := u.ensureSessionLocked(ctx)
if err != nil {
u.markDegraded(err)
return err
}
listCtx, cancel := context.WithTimeout(ctx, listTimeout)
defer cancel()
fresh, err := u.list(listCtx, session)
if err != nil {
// A session whose listing failed is suspect (dead transport, wedged
// stream). Drop it so the next re-probe starts from a fresh dial
// instead of reusing it forever; refreshes are infrequent, so the
// redial cost is negligible next to guaranteed recovery.
u.dropSession(session)
u.markDegraded(err)
return err
}
u.stateMu.Lock()
u.catalog = fresh.withToolFilters(u.spec.AllowedTools, u.spec.DisallowedTools)
u.status = StatusConnected
u.lastErr = ""
u.stateMu.Unlock()
return nil
}
// ensureSessionLocked returns the live session, dialing when necessary.
// connectMu must be held.
func (u *upstream) ensureSessionLocked(ctx context.Context) (*mcp.ClientSession, error) {
u.stateMu.Lock()
session := u.session
closed := u.closed
u.stateMu.Unlock()
if closed {
return nil, fmt.Errorf("mcp server %q was removed", u.spec.Name)
}
if session != nil {
return session, nil
}
dialCtx, cancel := context.WithTimeout(ctx, connectTimeout)
defer cancel()
// The SDK's Connect closes the session when the handshake fails and waits
// for in-flight requests, so cancelling dialCtx alone cannot end a dial
// whose server never answers. The dial's HTTP requests are therefore bound
// to requests, which the timeout, the caller, and close all cancel. A
// connected session keeps it for its own requests until close.
requests, abort := context.WithCancel(u.life)
stopAbort := context.AfterFunc(dialCtx, abort)
transport, probe, err := u.transport(requests)
if err != nil {
abort()
return nil, err
}
client := mcp.NewClient(&mcp.Implementation{
Name: "gomodel",
Title: "GoModel MCP Gateway",
Version: version.Version,
}, u.clientOptions())
session, err = client.Connect(dialCtx, transport, nil)
if !stopAbort() && err == nil {
// The dial ended as Connect returned, so its requests may already be
// aborted; the session is not trustworthy.
_ = session.Close()
err = dialCtx.Err()
}
u.stateMu.Lock()
u.authRequired = err != nil && probe.unauthorized()
u.stateMu.Unlock()
if err != nil {
abort()
return nil, u.connectError(err, probe)
}
u.stateMu.Lock()
// close() may have disposed the upstream while the dial was in flight
// (close deliberately does not wait on connectMu). Storing the session
// then would leak it forever, so discard it instead.
if u.closed {
u.stateMu.Unlock()
_ = session.Close()
return nil, fmt.Errorf("mcp server %q was removed", u.spec.Name)
}
u.session = session
u.connectedAt = time.Now().UTC()
u.stateMu.Unlock()
return session, nil
}
// clientOptions wires upstream change notifications into catalog refreshes.
// No sampling/elicitation/roots handlers are set, so those capabilities are
// not advertised upstream and servers legally never send such requests.
func (u *upstream) clientOptions() *mcp.ClientOptions {
relist := func() {
go func() {
if err := u.refresh(context.Background()); err != nil {
slog.Warn("mcp catalog refresh after list_changed failed",
"server", u.spec.Name, "error", err)
}
}()
}
return &mcp.ClientOptions{
ToolListChangedHandler: func(context.Context, *mcp.ToolListChangedRequest) { relist() },
PromptListChangedHandler: func(context.Context, *mcp.PromptListChangedRequest) { relist() },
ResourceListChangedHandler: func(context.Context, *mcp.ResourceListChangedRequest) { relist() },
}
}
// transport builds a fresh transport for one dial attempt. HTTP transports
// also return the probe watching that dial (nil for stdio); see connectError.
// transport builds the dial's transport; requests bounds its HTTP requests.
func (u *upstream) transport(requests context.Context) (mcp.Transport, *connectProbe, error) {
switch u.spec.Transport {
case "http", "":
probe := &connectProbe{method: http.MethodPost}
return &mcp.StreamableClientTransport{
Endpoint: u.spec.URL,
HTTPClient: u.dialClient(probe, requests),
}, probe, nil
case "sse":
probe := &connectProbe{method: http.MethodGet}
return &mcp.SSEClientTransport{
Endpoint: u.spec.URL,
HTTPClient: u.dialClient(probe, requests),
}, probe, nil
case "stdio":
cmd := exec.Command(u.spec.Command, u.spec.Args...) //nolint:gosec // stdio servers come only from startup config.
// Start from a minimal environment, not os.Environ(): the gateway
// process holds every provider API key and the master key, and a
// compromised MCP server binary must not inherit them. Operators pass
// anything else explicitly via the server's env map (${VAR} expands
// in config), on top of the basics process launchers need. The
// non-nil initialization matters: a nil cmd.Env would fall back to
// inheriting the full parent environment.
cmd.Env = []string{}
for _, key := range []string{"PATH", "HOME", "TMPDIR", "USER", "LANG"} {
if value := os.Getenv(key); value != "" {
cmd.Env = append(cmd.Env, key+"="+value)
}
}
for key, value := range u.spec.Env {
cmd.Env = append(cmd.Env, key+"="+value)
}
cmd.Stderr = os.Stderr
return &mcp.CommandTransport{Command: cmd}, nil, nil
default:
return nil, nil, fmt.Errorf("mcp server %q: unsupported transport %q", u.spec.Name, u.spec.Transport)
}
}
// dialClient overlays the connect probe and the configured static headers on
// the shared HTTP client. The headers carry the upstream credential; the
// client's own bearer token was terminated at the gateway and is never
// forwarded.
func (u *upstream) dialClient(probe *connectProbe, requests context.Context) *http.Client {
base := u.httpClient
if base == nil {
base = http.DefaultClient
}
transport := base.Transport
if transport == nil {
transport = http.DefaultTransport
}
// The OAuth transport sits inside the static headers so it runs last: a
// signed-in token replaces a configured Authorization header.
if u.spec.oauth != nil {
transport = &bearerRoundTripper{
base: transport,
source: u.spec.oauth,
origin: requestOrigin(u.spec.URL),
}
}
if len(u.spec.Headers) > 0 {
transport = &headerRoundTripper{
base: transport,
headers: u.spec.Headers,
origin: requestOrigin(u.spec.URL),
}
}
probe.base = transport
clone := *base
clone.Transport = &boundRoundTripper{base: probe, bound: requests}
return &clone
}
// boundRoundTripper cancels each request when bound ends, on top of the
// request's own context. The binding lasts until the response body closes,
// so a streamed response is cut off too.
type boundRoundTripper struct {
base http.RoundTripper
bound context.Context
}
func (t *boundRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
ctx, cancel := context.WithCancel(req.Context())
stop := context.AfterFunc(t.bound, cancel)
release := func() {
stop()
cancel()
}
resp, err := t.base.RoundTrip(req.WithContext(ctx))
if err != nil {
release()
return resp, err
}
resp.Body = &releasingBody{ReadCloser: resp.Body, release: sync.OnceFunc(release)}
return resp, nil
}
type releasingBody struct {
io.ReadCloser
release func()
}
func (b *releasingBody) Close() error {
err := b.ReadCloser.Close()
b.release()
return err
}
// connectProbe watches one dial and remembers how the URL answered the first
// request of its handshake: the POST of a streamable transport, the GET of an
// SSE one. A 404 or 405 there means something listens at the URL but no MCP
// endpoint does — almost always a wrong path, since servers mount their
// endpoint under different paths. The SDK reports that as a bare "Not Found",
// which reads like a session problem; connectError turns it into a hint about
// the URL. Only that first request counts: a later 404 on a request carrying
// a session ID is the server dropping that session, not a missing endpoint,
// and stateless servers legitimately answer the standalone SSE GET with 405.
// A 401 there means the server wants a credential, which the admin view
// reports as AuthRequired so the dashboard can offer an OAuth sign-in. A
// redirect is not an answer: the request it leads to is.
type connectProbe struct {
base http.RoundTripper
method string
decided atomic.Bool
status atomic.Int32
}
func (p *connectProbe) RoundTrip(req *http.Request) (*http.Response, error) {
resp, err := p.base.RoundTrip(req)
if err != nil || req.Method != p.method || isRedirect(resp.StatusCode) || !p.decided.CompareAndSwap(false, true) {
return resp, err
}
switch resp.StatusCode {
case http.StatusNotFound, http.StatusMethodNotAllowed, http.StatusUnauthorized:
p.status.Store(int32(resp.StatusCode))
}
return resp, err
}
func isRedirect(status int) bool {
return status >= 300 && status < 400
}
// unauthorized reports whether the dial's first handshake request was
// answered 401 Unauthorized.
func (p *connectProbe) unauthorized() bool {
return p != nil && p.status.Load() == http.StatusUnauthorized
}
// handshakeStatus returns the status of the dial's first handshake request
// when it showed no MCP endpoint at the URL or asked for credentials, and 0
// otherwise.
func (p *connectProbe) handshakeStatus() int {
if p == nil {
return 0
}
return int(p.status.Load())
}
// connectError wraps a failed dial, pointing at the URL path or transport
// when the probe saw that the handshake never reached an MCP endpoint. A 404
// means nothing is mounted at the path. A 405 means the path exists but
// rejects the handshake method, which is what a streamable HTTP endpoint
// answers to the SSE transport's GET (and an SSE endpoint to a POST), so the
// transport is the first thing to check. Only the origin is ever named: the
// path and query stay out of errors and logs because some servers carry
// credentials in them.
func (u *upstream) connectError(err error, probe *connectProbe) error {
err = redactDialError(err)
switch status := probe.handshakeStatus(); status {
case http.StatusNotFound:
return fmt.Errorf("connect to mcp server %q: %s answered HTTP 404 Not Found; no MCP endpoint at that path, check the url: %w",
u.spec.Name, requestOrigin(u.spec.URL), err)
case http.StatusMethodNotAllowed:
return fmt.Errorf("connect to mcp server %q: %s answered HTTP 405 Method Not Allowed to the %s handshake; the endpoint likely speaks the other transport, check the transport and the url: %w",
u.spec.Name, requestOrigin(u.spec.URL), u.transportName(), err)
case http.StatusUnauthorized:
return fmt.Errorf("connect to mcp server %q: %s answered HTTP 401 Unauthorized; sign in from the dashboard or configure an Authorization header: %w",
u.spec.Name, requestOrigin(u.spec.URL), err)
}
return fmt.Errorf("connect to mcp server %q: %w", u.spec.Name, err)
}
// transportName is the configured transport with the default made explicit.
func (u *upstream) transportName() string {
if u.spec.Transport == "" {
return "http"
}
return u.spec.Transport
}
// redactDialError trims every URL quoted in a transport failure to its
// origin. The HTTP client's *url.Error carries the request URL, and Go masks
// only its password, leaving userinfo, path, and query — all places a
// credential can sit. The SDK formats that error into a string several
// layers up, so the text is rewritten rather than the chain unwrapped; dial
// errors are only ever logged and displayed.
func redactDialError(err error) error {
text := err.Error()
redacted := quotedURL.ReplaceAllStringFunc(text, func(match string) string {
if origin := requestOrigin(match[1 : len(match)-1]); origin != "" {
return `"` + origin + `"`
}
return `"redacted"`
})
if redacted == text {
return err
}
return errors.New(redacted)
}
// quotedURL matches the %q-formatted URL inside a *url.Error message.
var quotedURL = regexp.MustCompile(`"https?://[^"\s]+"`)
type headerRoundTripper struct {
base http.RoundTripper
headers map[string]string
origin string
}
func (t *headerRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
clone := req.Clone(req.Context())
if t.origin != "" && requestOrigin(req.URL.String()) == t.origin {
for key, value := range t.headers {
clone.Header.Set(key, value)
}
}
return t.base.RoundTrip(clone)
}
func requestOrigin(raw string) string {
parsed, err := url.Parse(raw)
if err != nil || parsed.Scheme == "" || parsed.Host == "" {
return ""
}
return strings.ToLower(parsed.Scheme) + "://" + strings.ToLower(parsed.Host)
}
// list rebuilds the catalog from the upstream's declared capabilities. Valid
// tool schemas and metadata pass through untouched; malformed schemas are
// made safe for the stricter downstream SDK. Tool filters are applied by the
// caller, under the state lock.
func (u *upstream) list(ctx context.Context, session *mcp.ClientSession) (*catalog, error) {
fresh := &catalog{}
init := session.InitializeResult()
var caps *mcp.ServerCapabilities
if init != nil {
fresh.instructions = init.Instructions
caps = init.Capabilities
}
if caps == nil || caps.Tools != nil {
var tools []*mcp.Tool
for tool, err := range session.Tools(ctx, nil) {
if err != nil {
return nil, fmt.Errorf("list tools from %q: %w", u.spec.Name, err)
}
tools = append(tools, tool)
}
fresh.discovered = discoverTools(tools)
}
if caps != nil && caps.Prompts != nil {
for prompt, err := range session.Prompts(ctx, nil) {
if err != nil {
return nil, fmt.Errorf("list prompts from %q: %w", u.spec.Name, err)
}
if prompt == nil || prompt.Name == "" {
continue
}
fresh.prompts = append(fresh.prompts, prompt)
}
}
if caps != nil && caps.Resources != nil {
for resource, err := range session.Resources(ctx, nil) {
if err != nil {
return nil, fmt.Errorf("list resources from %q: %w", u.spec.Name, err)
}
if resource == nil || resource.URI == "" {
continue
}
fresh.resources = append(fresh.resources, resource)
}
for template, err := range session.ResourceTemplates(ctx, nil) {
if err != nil {
return nil, fmt.Errorf("list resource templates from %q: %w", u.spec.Name, err)
}
if template == nil || template.URITemplate == "" {
continue
}
fresh.templates = append(fresh.templates, template)
}
}
// Deterministic ordering across the whole catalog, not just tools: clients
// that embed capability lists in prompts get stable prompt-cache keys.
slices.SortFunc(fresh.prompts, func(a, b *mcp.Prompt) int {
return strings.Compare(a.Name, b.Name)
})
slices.SortFunc(fresh.resources, func(a, b *mcp.Resource) int {
return strings.Compare(a.URI, b.URI)
})
slices.SortFunc(fresh.templates, func(a, b *mcp.ResourceTemplate) int {
return strings.Compare(a.URITemplate, b.URITemplate)
})
return fresh, nil
}
// callTool forwards one tools/call with the original tool name. A session
// that died since the last call is redialed once, transparently.
func (u *upstream) callTool(ctx context.Context, name string, args json.RawMessage) (*mcp.CallToolResult, error) {
if !u.toolExposed(name) {
return nil, fmt.Errorf("%w: %q", ErrToolExcluded, name)
}
params := &mcp.CallToolParams{Name: name}
if len(args) > 0 {
params.Arguments = args
}
var result *mcp.CallToolResult
err := u.forward(ctx, func(ctx context.Context, session *mcp.ClientSession) error {
var callErr error
result, callErr = session.CallTool(ctx, params)
return callErr
})
return result, err
}
// getPrompt forwards one prompts/get with the original prompt name.
func (u *upstream) getPrompt(ctx context.Context, params *mcp.GetPromptParams) (*mcp.GetPromptResult, error) {
var result *mcp.GetPromptResult
err := u.forward(ctx, func(ctx context.Context, session *mcp.ClientSession) error {
var callErr error
result, callErr = session.GetPrompt(ctx, params)
return callErr
})
return result, err
}
// readResource forwards one resources/read by URI.
func (u *upstream) readResource(ctx context.Context, params *mcp.ReadResourceParams) (*mcp.ReadResourceResult, error) {
var result *mcp.ReadResourceResult
err := u.forward(ctx, func(ctx context.Context, session *mcp.ClientSession) error {
var callErr error
result, callErr = session.ReadResource(ctx, params)
return callErr
})
return result, err
}
// forward runs one upstream operation under the per-server timeout, redialing
// once when the shared session turned out to be dead.
func (u *upstream) forward(ctx context.Context, op func(context.Context, *mcp.ClientSession) error) error {
if !u.spec.Enabled {
return fmt.Errorf("mcp server %q is disabled", u.spec.Name)
}
timeout := u.spec.ToolTimeout
if timeout <= 0 {
timeout = 30 * time.Second
}
opCtx, cancel := context.WithTimeout(ctx, timeout)
defer cancel()
u.connectMu.Lock()
session, err := u.ensureSessionLocked(opCtx)
u.connectMu.Unlock()
if err != nil {
u.markDegraded(err)
return err
}
err = op(opCtx, session)
if err == nil || !isConnectionClosed(err) {
return err
}
// The shared session died underneath us; drop it and retry once on a
// fresh dial so one upstream restart does not fail a client call.
u.dropSession(session)
u.connectMu.Lock()
session, dialErr := u.ensureSessionLocked(opCtx)
u.connectMu.Unlock()
if dialErr != nil {
u.markDegraded(dialErr)
return dialErr
}
return op(opCtx, session)
}
func isConnectionClosed(err error) bool {
return errors.Is(err, mcp.ErrConnectionClosed)
}
// dropSession forgets the shared session if it is still the given one.
func (u *upstream) dropSession(session *mcp.ClientSession) {
u.stateMu.Lock()
if u.session == session {
u.session = nil
}
u.stateMu.Unlock()
_ = session.Close()
}
func (u *upstream) markDegraded(err error) {
u.stateMu.Lock()
u.status = StatusDegraded
u.lastErr = err.Error()
u.stateMu.Unlock()
}
func (u *upstream) isClosed() bool {
u.stateMu.Lock()
defer u.stateMu.Unlock()
return u.closed
}
// reset drops the shared session so the next use redials. The upstream stays
// alive — this is the "force reconnect" path, not disposal.
func (u *upstream) reset() {
u.stateMu.Lock()
session := u.session
u.session = nil
u.stateMu.Unlock()
if session != nil {
_ = session.Close()
}
}
// close permanently disposes the upstream: the session is terminated and any
// in-flight background refresh is refused a redial.
func (u *upstream) close() {
u.stateMu.Lock()
u.closed = true
session := u.session
u.session = nil
u.stateMu.Unlock()
// Cancel first: it ends requests a dial is stuck on, which the session
// close below would otherwise wait for.
u.lifeCancel()
if session != nil {
_ = session.Close()
}
}