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, &params)
			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)
		}
	}
}