internal/mcpgateway/factory.go

Rama main · réplica de lectura del repositorio original.

package mcpgateway

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"net/http"
	"sync"
	"time"

	"go.mongodb.org/mongo-driver/v2/mongo"

	"github.com/enterpilot/gomodel/config"
	"github.com/enterpilot/gomodel/ext"
	"github.com/enterpilot/gomodel/internal/encryption"
	"github.com/enterpilot/gomodel/internal/httpclient"
	"github.com/enterpilot/gomodel/internal/storage"
	"github.com/enterpilot/gomodel/internal/storage/sqlx"
	"github.com/enterpilot/gomodel/internal/usage"
)

// Result holds the initialized MCP gateway and any owned resources.
type Result struct {
	Service *Service
	Store   Store

	closeOnce sync.Once
	closeErr  error
}

// Close releases resources held by the MCP gateway subsystem.
func (r *Result) Close() error {
	if r == nil {
		return nil
	}
	r.closeOnce.Do(func() {
		if r.Service != nil {
			r.Service.Close()
		}
		var errs []error
		if r.Store != nil {
			if err := r.Store.Close(); err != nil {
				errs = append(errs, fmt.Errorf("store close: %w", err))
			}
		}
		if len(errs) > 0 {
			r.closeErr = fmt.Errorf("close errors: %w", errors.Join(errs...))
		}
	})
	return r.closeErr
}

// New creates the MCP gateway subsystem using an existing
// storage connection. Header values of managed servers are sealed with box;
// a nil box stores them in plaintext. secrets resolves the secret references
// admin-managed servers hold. providerServers are the provider-hosted servers
// mcp.provider_servers enables; a server in cfg.MCP.Servers with the same slug
// replaces one. sessions shares downstream sessions between replicas; nil
// keeps them on this instance.
func New(ctx context.Context, cfg *config.Config, shared storage.Storage, box *encryption.Box, httpClient *http.Client, usageLogger usage.LoggerInterface, secrets *config.Secrets, providerServers map[string]config.MCPServerConfig, sessions ext.SharedStore) (*Result, error) {
	if shared == nil {
		return nil, fmt.Errorf("shared storage is required")
	}
	if cfg == nil {
		return nil, fmt.Errorf("config is required")
	}
	return newResult(ctx, cfg, shared, box, httpClient, usageLogger, secrets, providerServers, sessions)
}

// Reencrypt seals every managed server header that is in plaintext or sealed
// with a data key other than box's active one.
func Reencrypt(ctx context.Context, shared storage.Storage, box *encryption.Box) (encryption.Report, error) {
	store, err := createStore(ctx, shared)
	if err != nil {
		return encryption.Report{Entity: "mcp_servers"}, err
	}
	defer func() { _ = store.Close() }()
	swap, ok := store.(headerSwapper)
	if !ok {
		return encryption.Report{Entity: "mcp_servers"}, fmt.Errorf("mcp server store %T cannot re-encrypt", store)
	}
	return (&sealedStore{Store: store, box: box, swap: swap}).reencrypt(ctx, swap)
}

// ReencryptOAuth seals every stored MCP OAuth sign-in that is in plaintext
// or sealed with a data key other than box's active one.
func ReencryptOAuth(ctx context.Context, shared storage.Storage, box *encryption.Box) (encryption.Report, error) {
	store, err := createStore(ctx, shared)
	if err != nil {
		return encryption.Report{Entity: "mcp_server_oauth"}, err
	}
	defer func() { _ = store.Close() }()
	oauth := newOAuthStore(store, box)
	if oauth == nil {
		return encryption.Report{Entity: "mcp_server_oauth"}, fmt.Errorf("mcp server store %T cannot hold oauth records", store)
	}
	return oauth.reencrypt(ctx)
}

func newResult(ctx context.Context, cfg *config.Config, storeConn storage.Storage, box *encryption.Box, httpClient *http.Client, usageLogger usage.LoggerInterface, secrets *config.Secrets, providerServers map[string]config.MCPServerConfig, sessions ext.SharedStore) (*Result, error) {
	rawStore, err := createStore(ctx, storeConn)
	if err != nil {
		return nil, err
	}
	store := sealStore(rawStore, box)
	if httpClient == nil {
		httpClient = defaultUpstreamHTTPClient()
	}

	configSpecs := configServerSpecs(cfg.MCP.Servers, providerServers)

	virtualSpecs := make(map[string]VirtualServerSpec, len(cfg.MCP.VirtualServers))
	for name, virtualCfg := range cfg.MCP.VirtualServers {
		virtualSpecs[name] = VirtualFromConfig(name, virtualCfg)
	}

	service, err := NewService(ctx, Options{
		ConfigServers:  configSpecs,
		VirtualServers: virtualSpecs,
		// Virtual server rows hold no secrets, so they use the unsealed store;
		// the sealing wrapper only implements Store.
		VirtualStore: virtualStoreOf(rawStore),
		// Admin-managed servers and virtual servers share the model-config
		// refresh cadence, like virtual models; operators tune
		// CACHE_REFRESH_INTERVAL for faster cross-instance propagation.
		RefreshInterval: time.Duration(cfg.Cache.Model.RefreshInterval) * time.Second,
		Store:           store,
		HTTPClient:      httpClient,
		UsageLogger:     usageLogger,
		UserPathHeader:  cfg.Server.UserPathHeader,
		AllowedOrigins:  cfg.MCP.AllowedOrigins,
		ToolDiscovery:   cfg.MCP.ToolDiscovery,
		Secrets:         secrets,
		SharedStore:     sessions,
		// OAuth records are sealed as a whole: they hold tokens and client
		// secrets.
		oauthStore: newOAuthStore(rawStore, box),
	})
	if err != nil {
		return nil, err
	}

	return &Result{
		Service: service,
		Store:   store,
	}, nil
}

// configServerSpecs merges the declared servers with the provider-hosted ones.
// Both are read-only in the dashboard; a declared server wins a slug clash, so
// an operator can replace a provider-hosted server with their own definition.
func configServerSpecs(declared, providerServers map[string]config.MCPServerConfig) map[string]ServerSpec {
	specs := make(map[string]ServerSpec, len(declared)+len(providerServers))
	for name, serverCfg := range declared {
		specs[name] = SpecFromConfig(name, serverCfg)
	}
	for name, serverCfg := range providerServers {
		if _, clash := specs[name]; clash {
			slog.Info("mcp server in mcp.servers replaces the provider-hosted server with the same slug", "server", name)
			continue
		}
		specs[name] = SpecFromConfig(name, serverCfg)
	}
	return specs
}

// defaultUpstreamHTTPClient is the shared pooled client for http/sse
// upstreams. It carries no whole-request timeout: the standalone SSE
// notification stream is expected to outlive HTTP_TIMEOUT, and every
// individual operation is already bounded by a context deadline.
func defaultUpstreamHTTPClient() *http.Client {
	cfg := httpclient.DefaultConfig()
	cfg.Timeout = 0
	return httpclient.NewHTTPClient(&cfg)
}

func createStore(ctx context.Context, store storage.Storage) (Store, error) {
	return storage.ResolveSQLBackend[Store](
		ctx,
		store,
		func(db sqlx.DB) (Store, error) { return NewSQLStore(ctx, db) },
		func(db *mongo.Database) (Store, error) { return NewMongoDBStore(db) },
	)
}

// virtualStoreOf returns the store's virtual-server persistence; both
// backends provide it.
func virtualStoreOf(store Store) VirtualStore {
	virtualStore, _ := store.(VirtualStore)
	return virtualStore
}