ai-gateway-go/internal/workbench/mcp.go

package workbench

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"

	"aigateway.local/core/internal/platform/cryptox"
	"aigateway.local/core/internal/provider"
	"github.com/jackc/pgx/v5"
)

type MCPServer struct {
	ID                string    `json:"id"`
	Code              string    `json:"code"`
	Name              string    `json:"name"`
	Description       string    `json:"description"`
	Transport         string    `json:"transport"`
	EndpointURL       string    `json:"endpoint_url"`
	Status            string    `json:"status"`
	CategoryID        *string   `json:"category_id,omitempty"`
	CategoryName      string    `json:"category_name"`
	Tags              []string  `json:"tags"`
	DepartmentIDs     []string  `json:"department_ids"`
	Enabled           bool      `json:"enabled"`
	HasSecretHeaders  bool      `json:"has_secret_headers"`
	Revision          int64     `json:"revision"`
	CreatedAt         time.Time `json:"created_at"`
	UpdatedAt         time.Time `json:"updated_at"`
	EncryptedHeaders  []byte    `json:"-"`
	HeadersKEKVersion int       `json:"-"`
}

type MCPServerInput struct {
	Code          string            `json:"code"`
	Name          string            `json:"name"`
	Description   string            `json:"description"`
	Transport     string            `json:"transport"`
	EndpointURL   string            `json:"endpoint_url"`
	Headers       map[string]string `json:"headers,omitempty"`
	Status        string            `json:"status"`
	CategoryID    *string           `json:"category_id,omitempty"`
	Tags          []string          `json:"tags"`
	DepartmentIDs []string          `json:"department_ids"`
	Enabled       bool              `json:"enabled"`
}

// MCPServerService manages registered MCP servers: lifecycle CRUD plus header
// decryption for the MCP client.
type MCPServerService struct {
	assets       *Service
	cipher       cryptox.Cipher
	allowPrivate bool
}

func NewMCPServerService(assets *Service, cipher cryptox.Cipher, allowPrivate bool) *MCPServerService {
	return &MCPServerService{assets: assets, cipher: cipher, allowPrivate: allowPrivate}
}

func (s *MCPServerService) validate(ctx context.Context, input *MCPServerInput, create bool) error {
	input.Code = strings.ToLower(strings.TrimSpace(input.Code))
	input.Name = strings.TrimSpace(input.Name)
	input.Description = strings.TrimSpace(input.Description)
	input.Transport = strings.ToLower(strings.TrimSpace(input.Transport))
	if !codeRE.MatchString(input.Code) || input.Name == "" || len(input.Name) > 128 || len(input.Description) > 4000 {
		return errors.New("MCP 服务器编码、名称或描述格式无效")
	}
	switch input.Transport {
	case "", "streamable-http":
		input.Transport = "streamable-http"
	case "sse":
	default:
		return errors.New("不支持的 MCP 传输方式")
	}
	validated, err := provider.ValidateBaseURL(ctx, input.EndpointURL, s.allowPrivate)
	if err != nil {
		return fmt.Errorf("MCP 端点校验失败: %w", err)
	}
	input.EndpointURL = validated
	switch input.Status {
	case "", "draft":
		input.Status = "draft"
	case "published", "archived":
	default:
		return errors.New("无效的资源状态")
	}
	if input.CategoryID != nil && strings.TrimSpace(*input.CategoryID) == "" {
		input.CategoryID = nil
	}
	var categoryErr error
	input.CategoryID, categoryErr = validCategoryID(ctx, s.assets.pool, input.CategoryID, "")
	if categoryErr != nil {
		return categoryErr
	}
	input.Tags, err = normalizeStrings(input.Tags, 30)
	if err != nil {
		return err
	}
	input.DepartmentIDs, err = normalizeStrings(input.DepartmentIDs, 100)
	if err != nil {
		return err
	}
	if create && input.Headers == nil {
		input.Headers = map[string]string{}
	}
	for key, value := range input.Headers {
		if strings.TrimSpace(key) == "" || len(key) > 128 || strings.ContainsAny(key, "\r\n") || len(value) > 8192 || strings.ContainsAny(value, "\r\n") {
			return errors.New("MCP 请求头格式无效")
		}
	}
	return nil
}

const mcpServerSelect = `SELECT m.id::text,m.code,m.name,m.description,m.transport,m.endpoint_url,m.status,m.category_id::text,coalesce(c.name,''),m.tags,m.department_ids::text[],m.enabled,octet_length(m.encrypted_headers)>0,m.revision,m.created_at,m.updated_at,m.encrypted_headers,m.headers_kek_version FROM gateway.mcp_servers m LEFT JOIN gateway.marketplace_categories c ON c.id=m.category_id`

func scanMCPServer(row pgx.Row) (MCPServer, error) {
	var s MCPServer
	err := row.Scan(&s.ID, &s.Code, &s.Name, &s.Description, &s.Transport, &s.EndpointURL, &s.Status, &s.CategoryID, &s.CategoryName, &s.Tags, &s.DepartmentIDs, &s.Enabled, &s.HasSecretHeaders, &s.Revision, &s.CreatedAt, &s.UpdatedAt, &s.EncryptedHeaders, &s.HeadersKEKVersion)
	return s, mapNotFound(err)
}

func (s *MCPServerService) List(ctx context.Context) ([]MCPServer, error) {
	rows, err := s.assets.pool.Query(ctx, mcpServerSelect+` ORDER BY m.updated_at DESC`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	items := []MCPServer{}
	for rows.Next() {
		server, err := scanMCPServer(rows)
		if err != nil {
			return nil, err
		}
		items = append(items, server)
	}
	return items, rows.Err()
}

func (s *MCPServerService) Get(ctx context.Context, id string) (MCPServer, error) {
	return scanMCPServer(s.assets.pool.QueryRow(ctx, mcpServerSelect+` WHERE m.id=$1`, id))
}

func (s *MCPServerService) GetByCode(ctx context.Context, code string) (MCPServer, error) {
	return scanMCPServer(s.assets.pool.QueryRow(ctx, mcpServerSelect+` WHERE m.code=$1`, code))
}

// GetPublishedByCode returns a published, enabled server by code.
func (s *MCPServerService) GetPublishedByCode(ctx context.Context, code string) (MCPServer, error) {
	return scanMCPServer(s.assets.pool.QueryRow(ctx, mcpServerSelect+` WHERE m.code=$1 AND m.status='published' AND m.enabled`, code))
}

func (s *MCPServerService) Save(ctx context.Context, id string, input MCPServerInput, actorID string, create bool) (MCPServer, error) {
	if err := s.validate(ctx, &input, create); err != nil {
		return MCPServer{}, err
	}
	tx, err := s.assets.pool.Begin(ctx)
	if err != nil {
		return MCPServer{}, err
	}
	defer rollback(ctx, tx)
	var encrypted []byte
	var version int
	if input.Headers != nil {
		raw, _ := json.Marshal(input.Headers)
		encrypted, version, err = s.cipher.Encrypt(raw)
		if err != nil {
			return MCPServer{}, err
		}
	}
	if create {
		id, err = newUUID()
		if err != nil {
			return MCPServer{}, err
		}
		_, err = tx.Exec(ctx, `INSERT INTO gateway.mcp_servers(id,code,name,description,transport,endpoint_url,encrypted_headers,headers_kek_version,status,category_id,tags,department_ids,enabled,created_by) VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14)`, id, input.Code, input.Name, input.Description, input.Transport, input.EndpointURL, encrypted, version, input.Status, input.CategoryID, input.Tags, input.DepartmentIDs, input.Enabled, actorID)
	} else if input.Headers == nil {
		tag, updateErr := tx.Exec(ctx, `UPDATE gateway.mcp_servers SET code=$2,name=$3,description=$4,transport=$5,endpoint_url=$6,status=$7,category_id=$8,tags=$9,department_ids=$10,enabled=$11,revision=revision+1,updated_at=clock_timestamp() WHERE id=$1`, id, input.Code, input.Name, input.Description, input.Transport, input.EndpointURL, input.Status, input.CategoryID, input.Tags, input.DepartmentIDs, input.Enabled)
		err = updateErr
		if err == nil && tag.RowsAffected() == 0 {
			return MCPServer{}, ErrNotFound
		}
	} else {
		tag, updateErr := tx.Exec(ctx, `UPDATE gateway.mcp_servers SET code=$2,name=$3,description=$4,transport=$5,endpoint_url=$6,encrypted_headers=$7,headers_kek_version=$8,status=$9,category_id=$10,tags=$11,department_ids=$12,enabled=$13,revision=revision+1,updated_at=clock_timestamp() WHERE id=$1`, id, input.Code, input.Name, input.Description, input.Transport, input.EndpointURL, encrypted, version, input.Status, input.CategoryID, input.Tags, input.DepartmentIDs, input.Enabled)
		err = updateErr
		if err == nil && tag.RowsAffected() == 0 {
			return MCPServer{}, ErrNotFound
		}
	}
	if err != nil {
		return MCPServer{}, err
	}
	event := "mcp_server.updated"
	if create {
		event = "mcp_server.created"
	}
	if err = emit(ctx, tx, event, "mcp_server", id, actorID, nil); err != nil {
		return MCPServer{}, err
	}
	if err = tx.Commit(ctx); err != nil {
		return MCPServer{}, err
	}
	return s.Get(ctx, id)
}

func (s *MCPServerService) Delete(ctx context.Context, id, actorID string) error {
	tx, err := s.assets.pool.Begin(ctx)
	if err != nil {
		return err
	}
	defer rollback(ctx, tx)
	var used bool
	if err = tx.QueryRow(ctx, `SELECT EXISTS(
		SELECT 1 FROM gateway.skills WHERE $1 = ANY(mcp_server_ids)
		UNION ALL SELECT 1 FROM gateway.digital_employees WHERE $1 = ANY(mcp_server_ids)
		UNION ALL SELECT 1 FROM gateway.marketplace_installations WHERE resource_type='mcp_server' AND resource_id=$1
	)`, id).Scan(&used); err != nil {
		return err
	}
	if used {
		return ErrConflict
	}
	tag, err := tx.Exec(ctx, `DELETE FROM gateway.mcp_servers WHERE id=$1`, id)
	if err != nil {
		return err
	}
	if tag.RowsAffected() == 0 {
		return ErrNotFound
	}
	if err = emit(ctx, tx, "mcp_server.deleted", "mcp_server", id, actorID, nil); err != nil {
		return err
	}
	return tx.Commit(ctx)
}

// Headers decrypts the stored request headers for an MCP server so the client
// can authenticate to it.
func (s *MCPServerService) Headers(server MCPServer) (map[string]string, error) {
	if len(server.EncryptedHeaders) == 0 {
		return map[string]string{}, nil
	}
	plain, err := s.cipher.Decrypt(server.EncryptedHeaders, server.HeadersKEKVersion)
	if err != nil {
		return nil, err
	}
	headers := map[string]string{}
	if err = json.Unmarshal(plain, &headers); err != nil {
		return nil, errors.New("MCP 请求头密文内容无效")
	}
	return headers, nil
}